From 3e6920627c3157e0967e29e41aa615fd793eb51c Mon Sep 17 00:00:00 2001 From: James Dawdy Date: Fri, 12 Jun 2026 02:11:05 -0500 Subject: [PATCH 01/81] fix(studio): load run.py by path for editable installs (#5909) * fix(studio): load run.py by path for editable installs `studio update` can leave a partial site-packages/studio/backend/ tree (plugin build artefacts only). That shadowed tree wins over an editable install and breaks `from studio.backend.run import ...`. Loading run.py by file path via importlib sidesteps the conflict. The module is cached in _RUN_MODULE so repeated calls are cheap. If exec_module fails, the module is removed from sys.modules before re-raising so a subsequent retry starts clean. Co-Authored-By: Claude Sonnet 4.6 * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * Handle None __file__ when checking cached run module for PR #5909 * Harden _load_backend_auth_storage against None __file__ and resolve cache-key path (PR #5909) * Adapt studio run/cloudflare in-venv tests to _load_run_module loader (PR #5909) --------- Co-authored-by: Jim Dawdy Co-authored-by: Claude Sonnet 4.6 Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com> Co-authored-by: Daniel Han --- unsloth_cli/commands/studio.py | 81 ++++++++++++++----- .../tests/test_studio_cloudflare_flag.py | 6 ++ .../tests/test_studio_run_parallel_flag.py | 4 + 3 files changed, 70 insertions(+), 21 deletions(-) diff --git a/unsloth_cli/commands/studio.py b/unsloth_cli/commands/studio.py index fb63fbf208..ffdf36bedf 100644 --- a/unsloth_cli/commands/studio.py +++ b/unsloth_cli/commands/studio.py @@ -184,6 +184,47 @@ def _find_run_py() -> Optional[Path]: return None +_RUN_MODULE = None + + +def _load_run_module(): + """Import studio.backend.run without relying on package resolution. + + `studio update` can leave a partial ``site-packages/studio/backend/`` + tree (plugin build artefacts only). That shadowed tree wins over an + editable install and breaks ``from studio.backend.run import ...``. + Loading by file path sidesteps the conflict. + """ + global _RUN_MODULE + if _RUN_MODULE is not None: + return _RUN_MODULE + + run_py = _find_run_py() + if run_py is None: + raise ImportError("Could not find studio/backend/run.py. Re-run: unsloth studio setup") + + loaded = sys.modules.get("studio.backend.run") + if loaded is not None: + # __file__ can be None for namespace packages from partial trees. + loaded_path = Path(getattr(loaded, "__file__", None) or "").resolve() + if loaded_path == run_py.resolve(): + _RUN_MODULE = loaded + return _RUN_MODULE + + spec = importlib.util.spec_from_file_location("studio.backend.run", run_py) + if spec is None or spec.loader is None: + raise ImportError(f"Could not load studio backend from {run_py}") + module = importlib.util.module_from_spec(spec) + sys.modules["studio.backend.run"] = module + try: + spec.loader.exec_module(module) + except Exception: + sys.modules.pop("studio.backend.run", None) + raise + _RUN_MODULE = module + return _RUN_MODULE + + def _find_setup_script() -> Optional[Path]: """Find studio/setup.sh or studio/setup.ps1. @@ -329,9 +370,11 @@ def _load_backend_auth_storage(): auth_dir = backend_dir / "auth" storage_py = auth_dir / "storage.py" loaded = sys.modules.get("auth.storage") - loaded_path = Path(getattr(loaded, "__file__", "")).resolve() - if loaded is not None and loaded_path == storage_py: - return loaded + if loaded is not None: + # __file__ can be None for namespace packages from partial trees. + loaded_path = Path(getattr(loaded, "__file__", None) or "").resolve() + if loaded_path == storage_py.resolve(): + return loaded package = sys.modules.get("auth") package_paths = [Path(path).resolve() for path in getattr(package, "__path__", [])] @@ -706,11 +749,11 @@ def studio_default( typer.echo("Studio not set up. Run install.sh first.") raise typer.Exit(1) - from studio.backend.run import run_server + run_mod = _load_run_module() + run_server = run_mod.run_server if not silent: - from studio.backend.run import _resolve_external_ip - display_host = _resolve_external_ip() if host == "0.0.0.0" else host + display_host = run_mod._resolve_external_ip() if host == "0.0.0.0" else host typer.echo(f"Starting Unsloth Studio on http://{display_host}:{port}") run_kwargs = dict( @@ -725,20 +768,17 @@ def studio_default( run_kwargs["frontend_path"] = frontend run_server(**run_kwargs) - from studio.backend.run import _shutdown_event - try: - if _shutdown_event is not None: + if run_mod._shutdown_event is not None: # Event.wait() with no timeout blocks at C-level on Linux # and swallows SIGINT; loop with a 1s timeout instead. - while not _shutdown_event.is_set(): - _shutdown_event.wait(timeout = 1) + while not run_mod._shutdown_event.is_set(): + run_mod._shutdown_event.wait(timeout = 1) else: while True: time.sleep(1) except KeyboardInterrupt: - from studio.backend.run import _graceful_shutdown, _server - _graceful_shutdown(_server) + run_mod._graceful_shutdown(run_mod._server) typer.echo("\nShutting down...") @@ -1036,7 +1076,8 @@ def run( os.execvp(str(studio_bin), args) # ── 2. Start server (always suppress built-in banner) ───────────── - from studio.backend.run import run_server, _resolve_external_ip + run_mod = _load_run_module() + run_server = run_mod.run_server run_kwargs = dict( host = host, @@ -1099,7 +1140,7 @@ def run( context_length_line = _format_context_length_line(result) # 6. Print banner. - display_host = _resolve_external_ip() if host == "0.0.0.0" else host + display_host = run_mod._resolve_external_ip() if host == "0.0.0.0" else host base_url = f"http://{display_host}:{actual_port}" sdk_base_url = f"{base_url}/v1" # run_server started the tunnel during the silent run above (0.0.0.0 only). @@ -1178,17 +1219,15 @@ def run( typer.secho(_tool_notice, fg = _tool_notice_fg, bold = True) # 7. Wait for Ctrl+C. - from studio.backend.run import _shutdown_event, _graceful_shutdown, _server - try: - if _shutdown_event is not None: - while not _shutdown_event.is_set(): - _shutdown_event.wait(timeout = 1) + if run_mod._shutdown_event is not None: + while not run_mod._shutdown_event.is_set(): + run_mod._shutdown_event.wait(timeout = 1) else: while True: time.sleep(1) except KeyboardInterrupt: - _graceful_shutdown(_server) + run_mod._graceful_shutdown(run_mod._server) typer.echo("\nShutting down...") diff --git a/unsloth_cli/tests/test_studio_cloudflare_flag.py b/unsloth_cli/tests/test_studio_cloudflare_flag.py index 5c2a039547..6ab1ce21ff 100644 --- a/unsloth_cli/tests/test_studio_cloudflare_flag.py +++ b/unsloth_cli/tests/test_studio_cloudflare_flag.py @@ -210,6 +210,9 @@ def test_run_in_venv_passes_cloudflare_to_run_server(monkeypatch, user_flag, exp ) fake_backend_run.run_server = fake_run_server fake_backend_run._resolve_external_ip = lambda: "127.0.0.1" + # run() loads the backend via _load_run_module() (by file path); inject the + # mock as the cached run module so the stubbed run_server is used. + monkeypatch.setattr(studio_mod, "_RUN_MODULE", fake_backend_run) import typer as _typer @@ -270,6 +273,9 @@ def test_run_in_venv_shuts_down_on_startup_abort(monkeypatch): backend._server = object() backend._shutdown_event = None backend._graceful_shutdown = lambda server: shutdown_calls.append(server) + # run() loads the backend via _load_run_module() (by file path); inject the + # mock as the cached run module so the stubbed symbols are used. + monkeypatch.setattr(studio_mod, "_RUN_MODULE", backend) # set_tool_policy is imported as `from state.tool_policy import set_tool_policy`. state_mod = sys.modules.setdefault("state", types.ModuleType("state")) diff --git a/unsloth_cli/tests/test_studio_run_parallel_flag.py b/unsloth_cli/tests/test_studio_run_parallel_flag.py index 8870fa348d..8caf31a432 100644 --- a/unsloth_cli/tests/test_studio_run_parallel_flag.py +++ b/unsloth_cli/tests/test_studio_run_parallel_flag.py @@ -426,6 +426,10 @@ def test_in_venv_path_passes_parallel_to_run_server(monkeypatch, value): ) fake_backend_run.run_server = fake_run_server fake_backend_run._resolve_external_ip = lambda: "127.0.0.1" + # run() loads the backend via _load_run_module() (by file path), which + # ignores a sys.modules mock with no matching __file__; inject it as the + # cached run module so the stubbed run_server is used. + monkeypatch.setattr(studio_mod, "_RUN_MODULE", fake_backend_run) import typer as _typer From f22e890ab81ebe3fd22505c954fe1df7c61e152d Mon Sep 17 00:00:00 2001 From: James Dawdy Date: Fri, 12 Jun 2026 02:27:04 -0500 Subject: [PATCH 02/81] fix(studio): inherit llama_extra_args and honor --no-mmproj (#5902) * fix(studio): inherit llama_extra_args and honor --no-mmproj Reloading the same GGUF from the UI without gguf_variant no longer drops CLI pass-through args like --no-mmproj. Skip mmproj download and launch when --no-mmproj is present in llama_extra_args. Co-authored-by: Cursor * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * fix(studio): tighten GGUF llama_extra_args variant inheritance guard Reject inherited CLI args when the request changes gguf_variant or when omitted variant resolves differently from the stored extra_args source. Co-authored-by: Cursor * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * Treat --no-mmproj-auto and --mmproj-auto with last-wins parsing for PR #5902 --------- Co-authored-by: Cursor Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com> Co-authored-by: Daniel Han --- studio/backend/core/inference/llama_cpp.py | 15 ++++++----- .../core/inference/llama_server_args.py | 22 +++++++++++++++ studio/backend/routes/inference.py | 27 +++++++++++-------- .../backend/tests/test_llama_server_args.py | 17 ++++++++++++ 4 files changed, 64 insertions(+), 17 deletions(-) diff --git a/studio/backend/core/inference/llama_cpp.py b/studio/backend/core/inference/llama_cpp.py index b83e4e6961..80b4862aa8 100644 --- a/studio/backend/core/inference/llama_cpp.py +++ b/studio/backend/core/inference/llama_cpp.py @@ -28,6 +28,7 @@ from typing import Callable, Generator, Iterable, List, Optional import httpx from core.inference.llama_server_args import ( + extra_args_disable_mmproj, parse_cache_override, parse_ctx_override, resolve_cache_type_kv, @@ -2929,8 +2930,8 @@ class LlamaCppBackend: hf_variant = hf_variant, hf_token = hf_token, ) - # Auto-download mmproj for vision models - if is_vision and not mmproj_path: + # Auto-download mmproj for vision models unless opted out. + if is_vision and not mmproj_path and not extra_args_disable_mmproj(extra_args): mmproj_path = self._download_mmproj( hf_repo = hf_repo, hf_token = hf_token, @@ -3191,10 +3192,12 @@ class LlamaCppBackend: gpu_indices, use_fit = None, True effective_ctx = requested_ctx # fall back to original - launch_mmproj_path = self._resolve_launch_mmproj_path( - model_path = model_path, - mmproj_path = mmproj_path, - ) + launch_mmproj_path = None + if not extra_args_disable_mmproj(extra_args): + launch_mmproj_path = self._resolve_launch_mmproj_path( + model_path = model_path, + mmproj_path = mmproj_path, + ) # Need both a resolved mmproj AND the config vision flag; a stray # mmproj passing the family-name heuristic must not flip a non-VLM # GGUF into vision mode. diff --git a/studio/backend/core/inference/llama_server_args.py b/studio/backend/core/inference/llama_server_args.py index 1862c9c5de..00f8c66d5c 100644 --- a/studio/backend/core/inference/llama_server_args.py +++ b/studio/backend/core/inference/llama_server_args.py @@ -260,6 +260,28 @@ def resolve_cache_type_kv( return override if override is not None else fallback_cache_type_kv +_MMPROJ_DISABLE_FLAGS: frozenset[str] = frozenset({"--no-mmproj", "--no-mmproj-auto"}) +_MMPROJ_ENABLE_FLAGS: frozenset[str] = frozenset({"--mmproj-auto"}) + + +def extra_args_disable_mmproj(args: Optional[Iterable[str]]) -> bool: + """True when pass-through args opt out of vision mmproj loading. + + llama-server parses --mmproj-auto / --no-mmproj / --no-mmproj-auto as one + boolean with last-wins semantics; mirror that here. + """ + if not args: + return False + disabled = False + for raw in args: + flag = _flag_name(str(raw)) + if flag in _MMPROJ_DISABLE_FLAGS: + disabled = True + elif flag in _MMPROJ_ENABLE_FLAGS: + disabled = False + return disabled + + def strip_shadowing_flags( args: Iterable[str], *, diff --git a/studio/backend/routes/inference.py b/studio/backend/routes/inference.py index 3850d0cbb2..856d9bacf8 100644 --- a/studio/backend/routes/inference.py +++ b/studio/backend/routes/inference.py @@ -1449,18 +1449,23 @@ async def load_model( # parse against a freshly-supplied first-class field. if request.llama_extra_args is None and llama_backend.extra_args: source = llama_backend.extra_args_source - # Compare against the resolved variant, not the request field: - # callers commonly omit gguf_variant for local ``.gguf`` paths - # and HF auto-pick flows. ``config.gguf_variant`` is the variant - # load_model was actually invoked with (see HF / local branches - # below), so both sides key off the same string. - resolved_variant = config.gguf_variant - same_source = bool( - source - and source[0] - and source[0].lower() == model_identifier.lower() - and (source[1] or "").lower() == (resolved_variant or "").lower() + # Compare against the resolved variant, not the request + # field: callers commonly omit gguf_variant for local + # ``.gguf`` paths and HF auto-pick flows. ``config.gguf_ + # variant`` is the variant load_model was actually + # invoked with (see the HF / local branches below), so + # both sides of the comparison key off the same string. + resolved_variant = (config.gguf_variant or "").lower() + request_variant = (request.gguf_variant or "").lower() + stored_variant = (source[1] or "").lower() if source else "" + same_model = bool( + source and source[0] and source[0].lower() == model_identifier.lower() ) + if request.gguf_variant: + variant_mismatch = request_variant != stored_variant + else: + variant_mismatch = bool(stored_variant and resolved_variant != stored_variant) + same_source = same_model and not variant_mismatch if not same_source: logger.info( "Not inheriting llama_extra_args: stored args came from %s, loading %s", diff --git a/studio/backend/tests/test_llama_server_args.py b/studio/backend/tests/test_llama_server_args.py index 09775707ee..6ae9d21e47 100644 --- a/studio/backend/tests/test_llama_server_args.py +++ b/studio/backend/tests/test_llama_server_args.py @@ -27,6 +27,7 @@ parse_cache_override = _lsa.parse_cache_override parse_ctx_override = _lsa.parse_ctx_override resolve_cache_type_kv = _lsa.resolve_cache_type_kv strip_shadowing_flags = _lsa.strip_shadowing_flags +extra_args_disable_mmproj = _lsa.extra_args_disable_mmproj validate_extra_args = _lsa.validate_extra_args @@ -509,6 +510,22 @@ def test_strip_shadowing_flags_defaults_strip_everything(): assert out == [] +def test_extra_args_disable_mmproj_detects_flag(): + assert extra_args_disable_mmproj(["--no-mmproj"]) is True + assert extra_args_disable_mmproj(["--threads", "12", "--no-mmproj"]) is True + assert extra_args_disable_mmproj(["--no-mmproj-auto"]) is True + + +def test_extra_args_disable_mmproj_false_when_absent(): + assert extra_args_disable_mmproj(None) is False + assert extra_args_disable_mmproj(["--threads", "12"]) is False + + +def test_extra_args_disable_mmproj_last_wins(): + assert extra_args_disable_mmproj(["--no-mmproj", "--mmproj-auto"]) is False + assert extra_args_disable_mmproj(["--mmproj-auto", "--no-mmproj-auto"]) is True + + def test_strip_shadowing_flags_drops_model_draft_with_spec(): # --model-draft (and aliases) are Studio-managed since the separate # MTP drafter support: an inherited copy must not last-wins-override From 515abca84efaf0e85723249f7a67fef6c634c2a1 Mon Sep 17 00:00:00 2001 From: James Dawdy Date: Fri, 12 Jun 2026 02:27:18 -0500 Subject: [PATCH 03/81] fix(studio): adopt server-loaded model before chat auto-load (#5900) * fix(studio): adopt server-loaded model before chat auto-load When the user starts Studio via `studio run -m`, the web UI could still auto-load a different cached GGUF on the first message because the chat checkpoint was empty. Sync from /api/inference/status before falling back to autoLoadSmallestModel so CLI-loaded models are not replaced. Co-authored-by: Cursor * fix(studio): hydrate adopted CLI model and harden auto-load errors Extract shared inference-status hydration for refresh() and CLI adopt paths so the first chat turn gets reasoning/tools flags. Wrap auto-load (including adopt) in try/catch for image-edit cleanup, and drop the redundant adopt call in run(). Co-authored-by: Cursor * Guard model adoption against status failures and mid-flight selection for PR #5900 * ci: trigger pre-commit.ci after main merge Co-authored-by: Cursor --------- Co-authored-by: Cursor Co-authored-by: Daniel Han --- .../src/features/chat/api/chat-adapter.ts | 6 + .../chat/hooks/use-chat-model-runtime.ts | 177 ++----------- .../lib/apply-inference-status-to-store.ts | 236 ++++++++++++++++++ 3 files changed, 257 insertions(+), 162 deletions(-) create mode 100644 studio/frontend/src/features/chat/lib/apply-inference-status-to-store.ts diff --git a/studio/frontend/src/features/chat/api/chat-adapter.ts b/studio/frontend/src/features/chat/api/chat-adapter.ts index d944c7ab74..983d4f5aa2 100644 --- a/studio/frontend/src/features/chat/api/chat-adapter.ts +++ b/studio/frontend/src/features/chat/api/chat-adapter.ts @@ -19,6 +19,7 @@ import { toExternalBackendProviderType, } from "../external-providers"; import { pickFriendlyContainerName } from "../lib/friendly-names"; +import { tryAdoptServerActiveModel } from "../lib/apply-inference-status-to-store"; import { clampReasoningEffortToLevels, getExternalMaxOutputTokens, @@ -1129,6 +1130,10 @@ async function autoLoadSmallestModel(): Promise<{ loaded: boolean; blockedByTrustRemoteCode: boolean; }> { + if (await tryAdoptServerActiveModel()) { + return { loaded: true, blockedByTrustRemoteCode: false }; + } + const store = useChatRuntimeStore.getState(); const hfToken = store.hfToken || null; const trustRemoteCode = store.params.trustRemoteCode ?? false; @@ -1472,6 +1477,7 @@ export function createOpenAIStreamAdapter(): ChatModelAdapter { } if (!useChatRuntimeStore.getState().params.checkpoint) { + // Prefer a model already loaded by the CLI/API before auto-loading. let loaded: boolean; let blockedByTrustRemoteCode: boolean; try { 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 4bd36e00d4..23ae1d80aa 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 @@ -24,12 +24,15 @@ import { } from "../api/chat-api"; import { formatEta, formatRate } from "../utils/format-transfer"; import { - CHAT_REASONING_ENABLED_KEY, - loadOptionalBool, - type ReasoningEffort, resolveToolsEnabledOnLoad, useChatRuntimeStore, } from "../stores/chat-runtime-store"; +import { + applyActiveModelStatusToStore, + clampLocalReasoningEffort, + normalizeSpeculativeType, + resolveInferenceCheckpointId, +} from "../lib/apply-inference-status-to-store"; import { mergeBackendRecommendedInference, resolveLoadMaxSeqLength, @@ -211,39 +214,6 @@ function getTrustRemoteCodeRequiredMessage(modelName: string): string { return `${modelName} needs custom code enabled to load. Turn on "Enable custom code" in Chat Settings, then try again.`; } -// Canonicalises any backend/persisted value onto the Speculative Decoding -// dropdown's modes ("auto"/"mtp"/"ngram"/"mtp+ngram"/"off"/null). Mirrors -// backend _canonicalize_spec_mode so legacy persisted values round-trip. -function normalizeSpeculativeType(v: string | null | undefined): string | null { - if (v == null) return null; - const s = String(v).trim().toLowerCase(); - if (!s) return null; - if (s === "auto" || s === "default") return "auto"; - if (s === "off") return "off"; - if (s === "ngram-simple") return "ngram-simple"; - if (s === "mtp" || s === "draft-mtp") return "mtp"; - if (s === "ngram" || s === "ngram-mod") return "ngram"; - if (s === "mtp+ngram") return "mtp+ngram"; - // Comma-chained legacy values (e.g. from older persisted state). - const parts = s.split(",").map((p) => p.trim()).filter(Boolean); - const hasMtp = parts.some((p) => p === "mtp" || p === "draft-mtp"); - const hasNgram = parts.some((p) => p === "ngram" || p === "ngram-mod"); - if (hasMtp && hasNgram) return "mtp+ngram"; - if (hasMtp) return "mtp"; - if (hasNgram) return "ngram"; - // Unknown -> safe fallback to Auto so the dropdown stays controlled. - return "auto"; -} - -type LocalReasoningEffort = Extract; - -function clampLocalReasoningEffort(value: ReasoningEffort): LocalReasoningEffort { - if (value === "low" || value === "medium" || value === "high") { - return value; - } - return "low"; -} - export function useChatModelRuntime() { const params = useChatRuntimeStore((state) => state.params); const models = useChatRuntimeStore((state) => state.models); @@ -328,132 +298,15 @@ export function useChatModelRuntime() { const selectedCheckpoint = useChatRuntimeStore.getState().params.checkpoint; const isExternalSelectionActive = isExternalModelId(selectedCheckpoint); if (statusRes.active_model && !isExternalSelectionActive) { - setCheckpoint(statusRes.active_model, statusRes.gguf_variant); - - // Apply inference defaults on reconnect (page refresh with model already loaded) - if (statusRes.inference) { - const currentParams = useChatRuntimeStore.getState().params; - setParams( - mergeBackendRecommendedInference({ - current: currentParams, - response: statusRes, - modelId: statusRes.active_model, - presetSource: useChatRuntimeStore.getState().activePresetSource, - }), - ); - } - - // Restore reasoning/tools support flags and context length - const hydratingExistingModel = - selectedCheckpoint !== statusRes.active_model || - useChatRuntimeStore.getState().activeGgufVariant !== - (statusRes.gguf_variant ?? null); - const supportsReasoning = statusRes.supports_reasoning ?? false; - const reasoningAlwaysOn = statusRes.reasoning_always_on ?? false; - const reasoningStyle = statusRes.reasoning_style ?? "enable_thinking"; - const reasoningEffortLevels = - reasoningStyle === "reasoning_effort" - ? (["low", "medium", "high"] as const) - : (["low", "medium", "high"] as const); - const supportsPreserveThinking = statusRes.supports_preserve_thinking ?? false; - const supportsTools = statusRes.supports_tools ?? false; - const storedReasoningEnabled = loadOptionalBool( - CHAT_REASONING_ENABLED_KEY, - ); - const currentGgufContextLength = statusRes.is_gguf - ? (statusRes.context_length ?? null) - : null; - const ggufMaxContextLength = statusRes.is_gguf - ? (statusRes.max_context_length ?? null) - : null; - const ggufNativeContextLength = statusRes.is_gguf - ? (statusRes.native_context_length ?? null) - : null; - const currentSpecType = normalizeSpeculativeType( - statusRes.speculative_type, - ); - // Refresh runs on F5 (needs hydration) and right after a load (store - // already set). For user-configurable params, only hydrate when the - // shadow `loaded*` field is null ("not yet hydrated"); otherwise we'd - // clobber what the load path just applied and revert the user. - const prevState = useChatRuntimeStore.getState(); - const clampedReasoningEffort = clampLocalReasoningEffort( - prevState.reasoningEffort, - ); - const nextDefaultChatTemplate = - statusRes.chat_template === undefined - ? prevState.defaultChatTemplate - : statusRes.chat_template; - useChatRuntimeStore.setState({ - supportsReasoning, - reasoningAlwaysOn, - reasoningStyle, - supportsReasoningOff: reasoningStyle !== "reasoning_effort", - reasoningEffortLevels, - reasoningEffort: clampedReasoningEffort, - supportsPreserveThinking, - supportsTools, - // Reset per-turn reasoning flag so: - // 1. non-reasoning models don't inherit a stale off state, and - // 2. local reasoning-effort models (Off hidden via - // supportsReasoningOff=false) don't carry reasoningEnabled=false - // from an external model where Off was selected -- the composer - // would still show "Think: " but the adapter would omit - // the kwarg, so Harmony falls back to its default effort. - reasoningEnabled: supportsReasoning - ? reasoningStyle === "reasoning_effort" - ? true - : useChatRuntimeStore.getState().reasoningEnabled - : true, - ggufContextLength: currentGgufContextLength, - ggufMaxContextLength, - ggufNativeContextLength, - modelRequiresTrustRemoteCode: - statusRes.requires_trust_remote_code ?? false, - defaultChatTemplate: nextDefaultChatTemplate, - loadedIsMultimodal: isMultimodalResponse(statusRes), - specFallbackReason: statusRes.spec_fallback_reason ?? null, - ...(prevState.loadedSpeculativeType === null && { - speculativeType: currentSpecType, - loadedSpeculativeType: currentSpecType, - }), - ...(statusRes.spec_draft_n_max !== undefined && - prevState.loadedSpecDraftNMax === null && - prevState.specDraftNMax === null && { - specDraftNMax: statusRes.spec_draft_n_max ?? null, - loadedSpecDraftNMax: statusRes.spec_draft_n_max ?? null, - }), - ...(statusRes.cache_type_kv !== undefined && - prevState.loadedKvCacheDtype === null && { - kvCacheDtype: statusRes.cache_type_kv, - loadedKvCacheDtype: statusRes.cache_type_kv, - }), - ...(statusRes.chat_template_override !== undefined && - prevState.loadedChatTemplateOverride === null && - prevState.chatTemplateOverride === null && { - chatTemplateOverride: statusRes.chat_template_override, - loadedChatTemplateOverride: statusRes.chat_template_override, - }), - }); - // setModels(listRes...) above used catalog data, which omits audio - // capability. Re-apply live status so attach gates survive a refresh. - syncModelCapabilities(statusRes.active_model, statusRes); - - // Set reasoning default for Qwen3.5/3.6 small models - if ( - supportsReasoning && - hydratingExistingModel && - storedReasoningEnabled === null - ) { - let reasoningDefault = true; - const mid = statusRes.active_model.toLowerCase(); - if (mid.includes("qwen3.5") || mid.includes("qwen3.6")) { - const sizeMatch = mid.match(/(\d+\.?\d*)\s*b/); - if (sizeMatch && parseFloat(sizeMatch[1]) < 9) { - reasoningDefault = false; - } - } - useChatRuntimeStore.setState({ reasoningEnabled: reasoningDefault }); + const checkpointId = resolveInferenceCheckpointId(statusRes); + if (checkpointId) { + setCheckpoint(checkpointId, statusRes.gguf_variant); + applyActiveModelStatusToStore(statusRes, { + previousCheckpoint: selectedCheckpoint, + }); + // setModels(listRes...) above used catalog data, which omits audio + // capability. Re-apply live status so attach gates survive a refresh. + syncModelCapabilities(checkpointId, statusRes); } } else if (!statusRes.active_model && !isExternalSelectionActive) { useChatRuntimeStore.setState({ diff --git a/studio/frontend/src/features/chat/lib/apply-inference-status-to-store.ts b/studio/frontend/src/features/chat/lib/apply-inference-status-to-store.ts new file mode 100644 index 0000000000..c02d304dcb --- /dev/null +++ b/studio/frontend/src/features/chat/lib/apply-inference-status-to-store.ts @@ -0,0 +1,236 @@ +// SPDX-License-Identifier: AGPL-3.0-only +// Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0 + +import { getInferenceStatus } from "../api/chat-api"; +import { mergeBackendRecommendedInference } from "../presets/preset-policy"; +import { + CHAT_REASONING_ENABLED_KEY, + loadOptionalBool, + type ReasoningEffort, + resolveToolsEnabledOnLoad, + useChatRuntimeStore, +} from "../stores/chat-runtime-store"; +import { isMultimodalResponse, type InferenceStatusResponse } from "../types/api"; +import type { ChatModelSummary } from "../types/runtime"; + +type LocalReasoningEffort = Extract; + +// Canonicalises backend / persisted speculative mode values onto the UI modes. +export function normalizeSpeculativeType( + v: string | null | undefined, +): string | null { + if (v == null) return null; + const s = String(v).trim().toLowerCase(); + if (!s) return null; + if (s === "auto" || s === "default") return "auto"; + if (s === "off") return "off"; + if (s === "ngram-simple") return "ngram-simple"; + if (s === "mtp" || s === "draft-mtp") return "mtp"; + if (s === "ngram" || s === "ngram-mod") return "ngram"; + if (s === "mtp+ngram") return "mtp+ngram"; + const parts = s.split(",").map((p) => p.trim()).filter(Boolean); + const hasMtp = parts.some((p) => p === "mtp" || p === "draft-mtp"); + const hasNgram = parts.some((p) => p === "ngram" || p === "ngram-mod"); + if (hasMtp && hasNgram) return "mtp+ngram"; + if (hasMtp) return "mtp"; + if (hasNgram) return "ngram"; + return "auto"; +} + +export function clampLocalReasoningEffort( + value: ReasoningEffort, +): LocalReasoningEffort { + if (value === "low" || value === "medium" || value === "high") { + return value; + } + return "low"; +} + +export function resolveInferenceCheckpointId( + status: InferenceStatusResponse, +): string | null { + if (!status.active_model) return null; + return status.model_identifier ?? status.active_model; +} + +function ensureActiveModelInStoreList( + status: InferenceStatusResponse, + checkpointId: string, +): void { + const store = useChatRuntimeStore.getState(); + if (store.models.some((model) => model.id === checkpointId)) { + return; + } + const summary: ChatModelSummary = { + id: checkpointId, + name: status.active_model ?? checkpointId, + isVision: status.is_vision ?? false, + isLora: false, + isGguf: status.is_gguf ?? false, + isAudio: status.is_audio ?? false, + audioType: status.audio_type ?? null, + hasAudioInput: status.has_audio_input ?? false, + }; + store.setModels([...store.models, summary]); +} + +export type ApplyInferenceStatusOptions = { + previousCheckpoint?: string; +}; + +/** Mirror refresh() hydration so adopted CLI models get reasoning/tools flags. */ +export function applyActiveModelStatusToStore( + status: InferenceStatusResponse, + options: ApplyInferenceStatusOptions = {}, +): void { + const checkpointId = resolveInferenceCheckpointId(status); + if (!checkpointId) return; + + const store = useChatRuntimeStore.getState(); + const previousCheckpoint = + options.previousCheckpoint ?? store.params.checkpoint; + + if (status.inference) { + store.setParams( + mergeBackendRecommendedInference({ + current: store.params, + response: status, + modelId: checkpointId, + presetSource: store.activePresetSource, + }), + ); + } + + const hydratingExistingModel = + previousCheckpoint !== checkpointId || + store.activeGgufVariant !== (status.gguf_variant ?? null); + const supportsReasoning = status.supports_reasoning ?? false; + const reasoningAlwaysOn = status.reasoning_always_on ?? false; + const reasoningStyle = status.reasoning_style ?? "enable_thinking"; + const reasoningEffortLevels = + reasoningStyle === "reasoning_effort" + ? (["low", "medium", "high"] as const) + : (["low", "medium", "high"] as const); + const supportsPreserveThinking = status.supports_preserve_thinking ?? false; + const supportsTools = status.supports_tools ?? false; + const storedReasoningEnabled = loadOptionalBool(CHAT_REASONING_ENABLED_KEY); + const currentGgufContextLength = status.is_gguf + ? (status.context_length ?? null) + : null; + const ggufMaxContextLength = status.is_gguf + ? (status.max_context_length ?? null) + : null; + const ggufNativeContextLength = status.is_gguf + ? (status.native_context_length ?? null) + : null; + const currentSpecType = normalizeSpeculativeType(status.speculative_type); + const prevState = useChatRuntimeStore.getState(); + const clampedReasoningEffort = clampLocalReasoningEffort( + prevState.reasoningEffort, + ); + const nextDefaultChatTemplate = + status.chat_template === undefined + ? prevState.defaultChatTemplate + : status.chat_template; + + useChatRuntimeStore.setState({ + supportsReasoning, + reasoningAlwaysOn, + reasoningStyle, + supportsReasoningOff: reasoningStyle !== "reasoning_effort", + reasoningEffortLevels, + reasoningEffort: clampedReasoningEffort, + supportsPreserveThinking, + supportsTools, + ...resolveToolsEnabledOnLoad(supportsTools), + reasoningEnabled: supportsReasoning + ? reasoningStyle === "reasoning_effort" + ? true + : useChatRuntimeStore.getState().reasoningEnabled + : true, + ggufContextLength: currentGgufContextLength, + ggufMaxContextLength, + ggufNativeContextLength, + modelRequiresTrustRemoteCode: status.requires_trust_remote_code ?? false, + defaultChatTemplate: nextDefaultChatTemplate, + loadedIsMultimodal: isMultimodalResponse(status), + specFallbackReason: status.spec_fallback_reason ?? null, + ...(prevState.loadedSpeculativeType === null && { + speculativeType: currentSpecType, + loadedSpeculativeType: currentSpecType, + }), + ...(status.spec_draft_n_max !== undefined && + prevState.loadedSpecDraftNMax === null && + prevState.specDraftNMax === null && { + specDraftNMax: status.spec_draft_n_max ?? null, + loadedSpecDraftNMax: status.spec_draft_n_max ?? null, + }), + ...(status.cache_type_kv !== undefined && + prevState.loadedKvCacheDtype === null && { + kvCacheDtype: status.cache_type_kv, + loadedKvCacheDtype: status.cache_type_kv, + }), + ...(status.chat_template_override !== undefined && + prevState.loadedChatTemplateOverride === null && + prevState.chatTemplateOverride === null && { + chatTemplateOverride: status.chat_template_override, + loadedChatTemplateOverride: status.chat_template_override, + }), + }); + + ensureActiveModelInStoreList(status, checkpointId); + + if ( + supportsReasoning && + hydratingExistingModel && + storedReasoningEnabled === null + ) { + let reasoningDefault = true; + const mid = checkpointId.toLowerCase(); + if (mid.includes("qwen3.5") || mid.includes("qwen3.6")) { + const sizeMatch = mid.match(/(\d+\.?\d*)\s*b/); + if (sizeMatch && parseFloat(sizeMatch[1]) < 9) { + reasoningDefault = false; + } + } + useChatRuntimeStore.setState({ reasoningEnabled: reasoningDefault }); + } +} + +/** + * Adopt the model already loaded on the inference server (e.g. via + * ``unsloth studio run -m``) into the chat UI checkpoint without + * triggering a new /api/inference/load. + */ +export async function tryAdoptServerActiveModel(): Promise { + const store = useChatRuntimeStore.getState(); + if (store.params.checkpoint) { + return true; + } + + let status: InferenceStatusResponse; + try { + status = await getInferenceStatus(); + } catch { + // Status endpoint unavailable: fall back to the normal auto-load path. + return false; + } + if (!status.active_model) { + return false; + } + + const checkpointId = resolveInferenceCheckpointId(status); + if (!checkpointId) { + return false; + } + + // Re-check after the await: keep a checkpoint the user picked meanwhile. + const previousCheckpoint = + useChatRuntimeStore.getState().params.checkpoint; + if (previousCheckpoint) { + return true; + } + store.setCheckpoint(checkpointId, status.gguf_variant); + applyActiveModelStatusToStore(status, { previousCheckpoint }); + return true; +} From 2fadc7b22c72863f91e9a79a91ecbdce0aae8959 Mon Sep 17 00:00:00 2001 From: Daniel Han Date: Fri, 12 Jun 2026 00:53:13 -0700 Subject: [PATCH 04/81] Fix stale sidebar regression test to match the gap-px markup (#6232) test_sidebar_account_block_uses_leading_tight hardcoded gap-0.5 in its selector, but the sidebar account-block div moved to gap-px during UI polish (#6196), so the regex stopped matching and the test failed across every studio PR's Repo tests (CPU). Match the gap utility loosely (gap-\S+) since this guard is about the leading-* class for descender clipping, not the spacing. --- tests/studio/test_studio_text_descender_clipping.py | 4 +++- 1 file changed, 3 insertions(+), 1 deletion(-) diff --git a/tests/studio/test_studio_text_descender_clipping.py b/tests/studio/test_studio_text_descender_clipping.py index 0b805cbd6f..ac26a10f06 100644 --- a/tests/studio/test_studio_text_descender_clipping.py +++ b/tests/studio/test_studio_text_descender_clipping.py @@ -36,8 +36,10 @@ def test_model_selector_trigger_label_uses_leading_tight(): def test_sidebar_account_block_uses_leading_tight(): src = _read(APP_SIDEBAR) + # Match the account-block parent div regardless of its gap utility (gap-0.5, + # gap-px, ...); this guard is about the leading-* class, not the spacing. pattern = re.compile( - r'', + r'', ) matches = pattern.findall(src) assert matches, "could not find sidebar account-block parent div" From a24c9987ca6a10537e59cadff8cc1bac9e38e19a Mon Sep 17 00:00:00 2001 From: Daniel Han Date: Fri, 12 Jun 2026 01:12:20 -0700 Subject: [PATCH 05/81] Studio: gate the staged prebuilt runtime validation behind a flag (off by default) (#6216) The post-download llama-quantize / llama-server smoke test JIT-compiles CUDA kernels on the first GPU forward pass and stalls every install and update by minutes on Blackwell (sm_100). Gate it behind _RUN_STAGED_PREBUILT_VALIDATION, disabled for now, keeping the smoke test and the source-build fallback it triggers fully intact so it can be restored by flipping the flag to True. Hashless external prebuilts (e.g. lemonade) are not in the approved-sha256 manifest and rely on the functional smoke test as their only integrity gate, so they are always validated regardless of the flag; only approved bundles, already proven by the sha256 manifest, skip it. The sha256 archive verification and the static Linux/macOS preflights are unchanged and still run for every install. --- studio/install_llama_prebuilt.py | 48 +++++--- .../test_install_llama_prebuilt_logic.py | 106 ++++++++++++++++++ 2 files changed, 137 insertions(+), 17 deletions(-) diff --git a/studio/install_llama_prebuilt.py b/studio/install_llama_prebuilt.py index bed7a27a63..e3b2d493fb 100644 --- a/studio/install_llama_prebuilt.py +++ b/studio/install_llama_prebuilt.py @@ -174,6 +174,12 @@ TEST_MODEL_URL = "https://huggingface.co/ggml-org/models/resolve/main/tinyllamas TEST_MODEL_SHA256 = "270cba1bd5109f42d03350f60406024560464db173c0e387d91f0426d3bd256d" VALIDATION_MODEL_CACHE_DIRNAME = ".cache" VALIDATION_MODEL_CACHE_FILENAME = "stories260K.gguf" +# Master switch for the staged runtime smoke test (llama-quantize + llama-server) +# in validate_prebuilt_choice. Disabled for now: the llama-server GPU forward pass +# JIT-compiles CUDA kernels on first load and stalls every install and update by +# minutes on Blackwell (sm_100). The check and the source-build fallback it triggers +# are kept intact -- set this to True to re-enable them. +_RUN_STAGED_PREBUILT_VALIDATION = False INSTALL_LOCK_TIMEOUT_SECONDS = 300 INSTALL_STAGING_ROOT_NAME = ".staging" GITHUB_AUTH_HOSTS = {"api.github.com", "github.com"} @@ -6551,23 +6557,31 @@ def validate_prebuilt_choice( approved_checksums = approved_checksums, prebuilt_fallback_used = prebuilt_fallback_used, ) - validate_quantize( - quantize_path, - probe_path, - quantized_path, - install_dir, - host, - runtime_line = choice.runtime_line, - ) - validate_server( - server_path, - probe_path, - host, - install_dir, - runtime_line = choice.runtime_line, - install_kind = choice.install_kind, - ) - log(f"staged prebuilt validation succeeded for {choice.name}") + # Hashless external prebuilts (e.g. lemonade) are not in the approved-sha256 + # manifest and rely on the functional smoke test as their only integrity gate, + # so they are always validated. For an approved bundle the sha256 manifest + # already proves integrity, so its runtime smoke test -- a cold CUDA-JIT pass + # costing minutes on Blackwell sm_100 -- is gated behind + # _RUN_STAGED_PREBUILT_VALIDATION, disabled for now. The check and the + # source-build fallback it triggers are kept intact; flip the flag to restore it. + if choice.expected_sha256 is None or _RUN_STAGED_PREBUILT_VALIDATION: + validate_quantize( + quantize_path, + probe_path, + quantized_path, + install_dir, + host, + runtime_line = choice.runtime_line, + ) + validate_server( + server_path, + probe_path, + host, + install_dir, + runtime_line = choice.runtime_line, + install_kind = choice.install_kind, + ) + log(f"staged prebuilt validation succeeded for {choice.name}") return server_path, quantize_path diff --git a/tests/studio/install/test_install_llama_prebuilt_logic.py b/tests/studio/install/test_install_llama_prebuilt_logic.py index 000c5a7ace..313cfa9225 100644 --- a/tests/studio/install/test_install_llama_prebuilt_logic.py +++ b/tests/studio/install/test_install_llama_prebuilt_logic.py @@ -2761,3 +2761,109 @@ def test_python_runtime_dirs_covers_cu13_and_library_bin(monkeypatch, tmp_path: assert str(cu13_arch) in dirs assert str(library_bin) in dirs assert str(torch_lib) in dirs + + +def _nvidia_linux_host(): + return HostInfo( + system = "Linux", + machine = "x86_64", + is_windows = False, + is_linux = True, + is_macos = False, + is_x86_64 = True, + is_arm64 = False, + nvidia_smi = None, + driver_cuda_version = None, + compute_caps = ["10.0"], + visible_cuda_devices = None, + has_physical_nvidia = True, + has_usable_nvidia = True, + ) + + +def _run_validate_prebuilt_choice(monkeypatch, tmp_path, *, expected_sha256): + """Drive validate_prebuilt_choice with every heavy install step stubbed and + return how many times the functional quantize/server smoke tests ran.""" + calls = {"quantize": 0, "server": 0} + server_path = tmp_path / "install" / "build" / "bin" / "llama-server" + quantize_path = tmp_path / "install" / "build" / "bin" / "llama-quantize" + + src = INSTALL_LLAMA_PREBUILT + monkeypatch.setattr( + src, "preferred_source_archive", lambda *a, **k: ("repo", "ref", None, False) + ) + monkeypatch.setattr(src, "hydrate_source_tree", lambda *a, **k: None) + monkeypatch.setattr(src, "install_from_archives", lambda *a, **k: (server_path, quantize_path)) + monkeypatch.setattr(src, "preflight_linux_installed_binaries", lambda *a, **k: None) + monkeypatch.setattr(src, "preflight_macos_installed_binaries", lambda *a, **k: None) + monkeypatch.setattr(src, "ensure_repo_shape", lambda *a, **k: None) + monkeypatch.setattr(src, "write_prebuilt_metadata", lambda *a, **k: None) + monkeypatch.setattr( + src, + "validate_quantize", + lambda *a, **k: calls.__setitem__("quantize", calls["quantize"] + 1), + ) + monkeypatch.setattr( + src, "validate_server", lambda *a, **k: calls.__setitem__("server", calls["server"] + 1) + ) + + bundle_name = "app-b9998-linux-x64-cuda13-newer.tar.gz" + source_archive = tmp_path / "source.tar.gz" + bundle_archive = tmp_path / "bundle.tar.gz" + source_archive.write_bytes(b"source") + bundle_archive.write_bytes(b"bundle") + + choice = AssetChoice( + repo = "local", + tag = "b9998", + name = bundle_name, + url = "file://bundle", + source_label = "local", + is_ready_bundle = True, + install_kind = "linux-cuda", + bundle_profile = "cuda13-newer", + runtime_line = "cuda13", + expected_sha256 = expected_sha256, + ) + src.validate_prebuilt_choice( + choice, + _nvidia_linux_host(), + tmp_path / "install", + tmp_path / "work", + tmp_path / "stories260K.gguf", + requested_tag = "b9998", + llama_tag = "b9998", + release_tag = "b9998", + approved_checksums = approved_checksums_for( + "b9998", + source_archive = source_archive, + bundle_archive = bundle_archive, + bundle_name = bundle_name, + ), + prebuilt_fallback_used = False, + quantized_path = tmp_path / "stories260K-q4.gguf", + ) + return calls + + +def test_validate_prebuilt_choice_approved_validation_skipped_when_flag_off(tmp_path, monkeypatch): + # An approved (sha256-verified) bundle skips the staged smoke test while the + # flag is off: the manifest hash is its integrity gate. + calls = _run_validate_prebuilt_choice(monkeypatch, tmp_path, expected_sha256 = "ab" * 32) + assert calls == {"quantize": 0, "server": 0} + + +def test_validate_prebuilt_choice_hashless_build_always_validated(tmp_path, monkeypatch): + # A hashless external build (e.g. lemonade) has no approved sha256, so the + # functional smoke test is its only integrity gate and must run even while the + # flag is off -- otherwise a corrupted/replaced archive could be activated. + calls = _run_validate_prebuilt_choice(monkeypatch, tmp_path, expected_sha256 = None) + assert calls == {"quantize": 1, "server": 1} + + +def test_validate_prebuilt_choice_approved_validation_runs_when_flag_enabled(tmp_path, monkeypatch): + # Flipping _RUN_STAGED_PREBUILT_VALIDATION back on restores the full smoke test + # for approved bundles too, proving the check is kept intact, only gated off. + monkeypatch.setattr(INSTALL_LLAMA_PREBUILT, "_RUN_STAGED_PREBUILT_VALIDATION", True) + calls = _run_validate_prebuilt_choice(monkeypatch, tmp_path, expected_sha256 = "ab" * 32) + assert calls == {"quantize": 1, "server": 1} From 14ed91e39a30e013928983bde3409ef1847ec36e Mon Sep 17 00:00:00 2001 From: alkinun Date: Fri, 12 Jun 2026 11:15:37 +0300 Subject: [PATCH 06/81] Fix FastModel config passthrough for sequence classification (#6203) * add FastModel config passthrough * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * fix fastmodel config passthrough for task configs * fix config-driven FastModel task model selection * fix text only fastmodel task config selection * fix fastmodel task config inference from user configs * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * fix fastmodel problem_type config passthrough * fix fastlanguagemodel config passthrough: FastLlamaModel owns user config * fix fastlanguagemodel config passthrough: forward user config to causal loads and keep checkpoint quantization_config --------- Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com> Co-authored-by: Etherll <61019402+Etherll@users.noreply.github.com> --- .../test_fast_model_config_passthrough.py | 215 ++++++++++++++++++ unsloth/models/_utils.py | 23 ++ unsloth/models/llama.py | 74 ++++-- unsloth/models/loader.py | 98 ++++++-- unsloth/models/vision.py | 14 +- 5 files changed, 389 insertions(+), 35 deletions(-) create mode 100644 tests/python/test_fast_model_config_passthrough.py diff --git a/tests/python/test_fast_model_config_passthrough.py b/tests/python/test_fast_model_config_passthrough.py new file mode 100644 index 0000000000..b2ba3d2eef --- /dev/null +++ b/tests/python/test_fast_model_config_passthrough.py @@ -0,0 +1,215 @@ +"""FastModel config passthrough and nested task config handling.""" + +import ast +from pathlib import Path + + +REPO_ROOT = Path(__file__).resolve().parents[2] +LOADER_PATH = REPO_ROOT / "unsloth" / "models" / "loader.py" +VISION_PATH = REPO_ROOT / "unsloth" / "models" / "vision.py" +UTILS_PATH = REPO_ROOT / "unsloth" / "models" / "_utils.py" +LLAMA_PATH = REPO_ROOT / "unsloth" / "models" / "llama.py" + + +def _source(path): + return path.read_text() + + +def _class_method(tree, class_name, method_name): + for node in tree.body: + if isinstance(node, ast.ClassDef) and node.name == class_name: + for item in node.body: + if isinstance(item, ast.FunctionDef) and item.name == method_name: + return item + raise AssertionError(f"{class_name}.{method_name} not found") + + +def _assigns_from_kwargs_pop(method, target_name, key_name): + for node in ast.walk(method): + if not isinstance(node, ast.Assign): + continue + if not any( + isinstance(target, ast.Name) and target.id == target_name for target in node.targets + ): + continue + value = node.value + if not ( + isinstance(value, ast.Call) + and isinstance(value.func, ast.Attribute) + and value.func.attr == "pop" + and isinstance(value.func.value, ast.Name) + and value.func.value.id == "kwargs" + and value.args + and isinstance(value.args[0], ast.Constant) + and value.args[0].value == key_name + ): + continue + return True + return False + + +def _calls_name(method, name): + return any( + isinstance(node, ast.Call) and isinstance(node.func, ast.Name) and node.func.id == name + for node in ast.walk(method) + ) + + +def _load_task_attr_helper(): + source = _source(UTILS_PATH) + funcs = { + node.name: ast.get_source_segment(source, node) + for node in ast.parse(source).body + if isinstance(node, ast.FunctionDef) + } + ns = {} + for name in ("_config_set", "set_task_config_attr"): + exec(funcs[name], ns) + return ns["set_task_config_attr"] + + +def _load_loader_task_helpers(): + source = _source(LOADER_PATH) + funcs = { + node.name: ast.get_source_segment(source, node) + for node in ast.parse(source).body + if isinstance(node, ast.FunctionDef) + } + ns = {} + for name in ( + "_config_get", + "_config_diff", + "_has_sequence_classification_architecture", + "_get_user_task_config_attrs", + ): + exec(funcs[name], ns) + return ns["_get_user_task_config_attrs"] + + +def test_fast_model_consumes_user_config_kwarg(): + tree = ast.parse(_source(LOADER_PATH)) + method = _class_method(tree, "FastModel", "from_pretrained") + + assert _assigns_from_kwargs_pop(method, "user_config", "config") + + +def test_fast_base_model_consumes_user_config_kwarg(): + tree = ast.parse(_source(VISION_PATH)) + method = _class_method(tree, "FastBaseModel", "from_pretrained") + + assert _assigns_from_kwargs_pop(method, "user_config", "config") + + +def test_fast_llama_model_consumes_user_config_kwarg(): + tree = ast.parse(_source(LLAMA_PATH)) + method = _class_method(tree, "FastLlamaModel", "from_pretrained") + + assert _assigns_from_kwargs_pop(method, "user_config", "config") + + +def test_fast_base_model_sets_task_attrs_on_nested_text_config(): + tree = ast.parse(_source(VISION_PATH)) + method = _class_method(tree, "FastBaseModel", "from_pretrained") + + assert _calls_name(method, "set_task_config_attr") + + +def test_fast_base_model_pops_problem_type_as_config_attr(): + source = _source(VISION_PATH) + + assert '("id2label", "label2id", "problem_type")' in source + + +def test_fast_model_uses_user_config_num_labels_for_task_model_selection(): + tree = ast.parse(_source(LOADER_PATH)) + method = _class_method(tree, "FastModel", "from_pretrained") + + assert _calls_name(method, "_get_user_task_config_attrs") + + +def test_fast_model_captures_user_config_num_labels_before_text_only_switch(): + source = _source(LOADER_PATH) + + fallback = source.index("task_config_attrs = _get_user_task_config_attrs(user_config)") + text_only_switch = source.index("model_config = text_config") + + assert fallback < text_only_switch + + +def test_user_task_config_attrs_ignore_default_num_labels(): + get_user_task_config_attrs = _load_loader_task_helpers() + + class Config: + num_labels = 2 + id2label = {0: "LABEL_0", 1: "LABEL_1"} + label2id = {"LABEL_0": 0, "LABEL_1": 1} + + def to_diff_dict(self): + return {} + + assert get_user_task_config_attrs(Config()) == {} + + +def test_user_task_config_attrs_preserve_custom_label_maps(): + get_user_task_config_attrs = _load_loader_task_helpers() + + class Config: + num_labels = 2 + id2label = {0: "negative", 1: "positive"} + label2id = {"negative": 0, "positive": 1} + + def to_diff_dict(self): + return {"id2label": self.id2label, "label2id": self.label2id} + + attrs = get_user_task_config_attrs(Config()) + + assert attrs["num_labels"] == 2 + assert attrs["id2label"] == {0: "negative", 1: "positive"} + assert attrs["label2id"] == {"negative": 0, "positive": 1} + + +def test_user_task_config_attrs_preserve_explicit_dict_num_labels(): + get_user_task_config_attrs = _load_loader_task_helpers() + + assert get_user_task_config_attrs({"num_labels": 2}) == {"num_labels": 2} + + +def test_task_config_attr_updates_parent_and_text_config_objects(): + set_task_config_attr = _load_task_attr_helper() + + class TextConfig: + pass + + class ParentConfig: + def __init__(self): + self.text_config = TextConfig() + + def get_text_config(self): + return self.text_config + + config = ParentConfig() + + set_task_config_attr(config, "num_labels", 3) + + assert config.num_labels == 3 + assert config.text_config.num_labels == 3 + + +def test_task_config_attr_updates_parent_and_text_config_dicts(): + set_task_config_attr = _load_task_attr_helper() + config = {"text_config": {}} + + set_task_config_attr(config, "label2id", {"negative": 0, "positive": 1}) + + assert config["label2id"] == {"negative": 0, "positive": 1} + assert config["text_config"]["label2id"] == {"negative": 0, "positive": 1} + + +def test_task_config_attr_ignores_primitive_text_config(): + set_task_config_attr = _load_task_attr_helper() + config = {"text_config": "not-a-config"} + + set_task_config_attr(config, "num_labels", 2) + + assert config["num_labels"] == 2 + assert config["text_config"] == "not-a-config" diff --git a/unsloth/models/_utils.py b/unsloth/models/_utils.py index 6baa9a1398..2f4e3a069e 100644 --- a/unsloth/models/_utils.py +++ b/unsloth/models/_utils.py @@ -69,6 +69,7 @@ __all__ = [ "resolve_attention_implementation", "resolve_encoder_attention_implementation", "_set_attn_impl", + "set_task_config_attr", "patch_fast_lora", "validate_loftq_config", "RaiseUninitialized", @@ -306,6 +307,28 @@ def _config_set(config, field_name, value): setattr(config, field_name, value) +def set_task_config_attr(config, field_name, value): + _config_set(config, field_name, value) + text_config = None + if isinstance(config, dict): + text_config = config.get("text_config", None) + elif config is not None: + get_text_config = getattr(config, "get_text_config", None) + if callable(get_text_config): + try: + text_config = get_text_config() + except Exception: + text_config = None + if text_config is None: + text_config = getattr(config, "text_config", None) + if ( + text_config is not None + and text_config is not config + and (isinstance(text_config, dict) or hasattr(text_config, "__dict__")) + ): + _config_set(text_config, field_name, value) + + def _iter_attention_configs(config, seen = None): if config is None or (not isinstance(config, dict) and not hasattr(config, "__dict__")): return diff --git a/unsloth/models/llama.py b/unsloth/models/llama.py index a60edf3fbe..08802f030e 100644 --- a/unsloth/models/llama.py +++ b/unsloth/models/llama.py @@ -2383,11 +2383,30 @@ class FastLlamaModel: assert dtype == torch.float16 or dtype == torch.bfloat16 or dtype == torch.float32 # RoPE Scaling - model_config = AutoConfig.from_pretrained( - model_name, - token = token, - attn_implementation = "sdpa", - ) + # Respect a user-provided config so it is the single config object used + # everywhere below; otherwise HF would receive it again through **kwargs + # alongside our own config= and fail with a duplicate-kwarg TypeError. + user_config = kwargs.pop("config", None) + if user_config is not None: + model_config = user_config + # model_name may have been remapped to a prequantized repo whose + # checkpoint needs its quantization_config; graft it onto the user + # config or the 4bit weights load without their quant state. + if getattr(model_config, "quantization_config", None) is None: + _checkpoint_config = AutoConfig.from_pretrained( + model_name, + token = token, + attn_implementation = "sdpa", + ) + _checkpoint_quant = getattr(_checkpoint_config, "quantization_config", None) + if _checkpoint_quant is not None: + model_config.quantization_config = _checkpoint_quant + else: + model_config = AutoConfig.from_pretrained( + model_name, + token = token, + attn_implementation = "sdpa", + ) model_config.model_name = model_name model_max_seq_length = model_config.max_position_embeddings @@ -2504,14 +2523,17 @@ class FastLlamaModel: # Transformers 5.x @strict config classes reject unexpected kwargs # like num_labels and max_position_embeddings. Set on the config # object directly and pass config= instead. - model_config.num_labels = num_labels + set_task_config_attr(model_config, "num_labels", num_labels) if max_position_embeddings is not None: model_config.max_position_embeddings = max_position_embeddings # Pop config-level attrs that would be rejected by @strict model init for _cfg_key in ("id2label", "label2id", "rope_scaling"): _cfg_val = kwargs.pop(_cfg_key, None) if _cfg_val is not None: - setattr(model_config, _cfg_key, _cfg_val) + if _cfg_key in ("id2label", "label2id"): + set_task_config_attr(model_config, _cfg_key, _cfg_val) + else: + setattr(model_config, _cfg_key, _cfg_val) model = AutoModelForSequenceClassification.from_pretrained( model_name, config = model_config, @@ -2544,17 +2566,33 @@ class FastLlamaModel: fast_inference = fast_inference, ) elif not fast_inference: - model = AutoModelForCausalLM.from_pretrained( - model_name, - device_map = device_map, - # torch_dtype = dtype, # transformers changed torch_dtype to dtype - # quantization_config = bnb_config, - token = token, - max_position_embeddings = max_position_embeddings, - trust_remote_code = trust_remote_code, - attn_implementation = preferred_attn_impl, - **kwargs, - ) + if user_config is not None: + # Transformers 5.x @strict model init rejects extra kwargs next + # to config=; set the override on the config and pass the single + # config object through so user overrides reach the actual load. + if max_position_embeddings is not None: + model_config.max_position_embeddings = max_position_embeddings + model = AutoModelForCausalLM.from_pretrained( + model_name, + config = model_config, + device_map = device_map, + token = token, + trust_remote_code = trust_remote_code, + attn_implementation = preferred_attn_impl, + **kwargs, + ) + else: + model = AutoModelForCausalLM.from_pretrained( + model_name, + device_map = device_map, + # torch_dtype = dtype, # transformers changed torch_dtype to dtype + # quantization_config = bnb_config, + token = token, + max_position_embeddings = max_position_embeddings, + trust_remote_code = trust_remote_code, + attn_implementation = preferred_attn_impl, + **kwargs, + ) # Attach dispatch hooks for bnb multi-device loads. from unsloth.models.vision import _attach_bnb_multidevice_hooks diff --git a/unsloth/models/loader.py b/unsloth/models/loader.py index 4dc3046928..cfcfcae505 100644 --- a/unsloth/models/loader.py +++ b/unsloth/models/loader.py @@ -101,6 +101,7 @@ from ._utils import ( resolve_model_class, _is_family_text_decoder, _apply_text_only_key_mapping, + set_task_config_attr, ) # Single source of truth is unsloth_zoo.model_lists. Re-exported so callers @@ -133,6 +134,57 @@ def _strip_unsloth_bnb_4bit_suffix(model_name: str) -> str: return s +def _config_get( + config, + field_name, + default = None, +): + if isinstance(config, dict): + return config.get(field_name, default) + return getattr(config, field_name, default) + + +def _config_diff(config): + if isinstance(config, dict): + return config + to_diff_dict = getattr(config, "to_diff_dict", None) + if callable(to_diff_dict): + try: + diff = to_diff_dict() + if isinstance(diff, dict): + return diff + except Exception: + pass + return {} + + +def _has_sequence_classification_architecture(config): + architectures = _config_get(config, "architectures", None) or [] + return any(str(arch).endswith("ForSequenceClassification") for arch in architectures) + + +def _get_user_task_config_attrs(user_config): + if user_config is None: + return {} + diff = _config_diff(user_config) + attrs = {} + for key in ("id2label", "label2id", "problem_type"): + if key in diff: + attrs[key] = _config_get(user_config, key, diff.get(key)) + if isinstance(user_config, dict) and "num_labels" in user_config: + attrs["num_labels"] = user_config["num_labels"] + elif _has_sequence_classification_architecture(user_config): + num_labels = _config_get(user_config, "num_labels", None) + if num_labels is not None: + attrs["num_labels"] = num_labels + elif "id2label" in attrs: + try: + attrs["num_labels"] = len(attrs["id2label"]) + except TypeError: + pass + return attrs + + DISABLE_COMPILE_MODEL_NAMES = [ "aya_vision", "modernbert", @@ -907,6 +959,7 @@ class FastModel(FastBaseModel): *args, **kwargs, ): + user_config = kwargs.pop("config", None) # Respect user-provided quantization_config (e.g. BitsAndBytesConfig) quantization_config = kwargs.get("quantization_config", None) if quantization_config is not None: @@ -1104,13 +1157,15 @@ class FastModel(FastBaseModel): ) try: - model_config = AutoConfig.from_pretrained( - model_name, - token = token, - revision = revision, - trust_remote_code = trust_remote_code, - local_files_only = local_files_only, - ) + model_config = user_config + if model_config is None: + model_config = AutoConfig.from_pretrained( + model_name, + token = token, + revision = revision, + trust_remote_code = trust_remote_code, + local_files_only = local_files_only, + ) is_model = True except ImportError: raise @@ -1384,12 +1439,15 @@ class FastModel(FastBaseModel): load_in_fp8 = False load_in_16bit = True - model_config = AutoConfig.from_pretrained( - model_name, - token = token, - trust_remote_code = trust_remote_code, - local_files_only = local_files_only, - ) + if user_config is not None: + model_config = user_config + else: + model_config = AutoConfig.from_pretrained( + model_name, + token = token, + trust_remote_code = trust_remote_code, + local_files_only = local_files_only, + ) if not was_disabled: enable_progress_bars() @@ -1469,6 +1527,17 @@ class FastModel(FastBaseModel): else: tokenizer_name = kwargs.pop("tokenizer_name", None) + # Capture task intent before text_only can replace a parent VLM config + # with its nested text config. + task_config_attrs = _get_user_task_config_attrs(user_config) + for _cfg_key in ("num_labels", "id2label", "label2id", "problem_type"): + _cfg_val = kwargs.get(_cfg_key, None) + if _cfg_val is not None: + task_config_attrs[_cfg_key] = _cfg_val + _num_labels = task_config_attrs.get("num_labels", None) + for _cfg_key, _cfg_val in task_config_attrs.items(): + set_task_config_attr(model_config, _cfg_key, _cfg_val) + # Check if VLM architectures = getattr(model_config, "architectures", None) if architectures is None: @@ -1499,7 +1568,8 @@ class FastModel(FastBaseModel): else: is_vlm = False # If num_labels is set, use AutoModelForSequenceClassification - _num_labels = kwargs.get("num_labels", None) + for _cfg_key, _cfg_val in task_config_attrs.items(): + set_task_config_attr(model_config, _cfg_key, _cfg_val) if auto_model is None: if _num_labels is not None: from transformers import AutoModelForSequenceClassification diff --git a/unsloth/models/vision.py b/unsloth/models/vision.py index a52ab359bd..e8161427d5 100644 --- a/unsloth/models/vision.py +++ b/unsloth/models/vision.py @@ -37,6 +37,7 @@ from ._utils import ( _get_text_only_config, _is_family_text_decoder, _apply_text_only_key_mapping, + set_task_config_attr, ) from ._utils import * from .loader_utils import _get_fp8_mode_and_check_settings @@ -592,6 +593,10 @@ class FastBaseModel: text_only = False, **kwargs, ): + user_config = kwargs.pop("config", None) + if auto_config is None and user_config is not None: + auto_config = user_config + if unsloth_vllm_standby and os.environ.get("UNSLOTH_VLLM_STANDBY", "0") != "1": raise RuntimeError( "Unsloth: UNSLOTH_VLLM_STANDBY is True, but UNSLOTH_VLLM_STANDBY is not set to 1!" @@ -950,11 +955,14 @@ class FastBaseModel: # Move config-level attributes onto the config object directly. _num_labels = kwargs.pop("num_labels", None) if _num_labels is not None: - model_config.num_labels = _num_labels - for _cfg_key in ("id2label", "label2id", "max_position_embeddings"): + set_task_config_attr(model_config, "num_labels", _num_labels) + for _cfg_key in ("id2label", "label2id", "problem_type"): _cfg_val = kwargs.pop(_cfg_key, None) if _cfg_val is not None: - setattr(model_config, _cfg_key, _cfg_val) + set_task_config_attr(model_config, _cfg_key, _cfg_val) + _cfg_val = kwargs.pop("max_position_embeddings", None) + if _cfg_val is not None: + setattr(model_config, "max_position_embeddings", _cfg_val) model = auto_model.from_pretrained( model_name, config = model_config, From 51f1c8732dba3cb6076a5929142fc2f933e6b353 Mon Sep 17 00:00:00 2001 From: dylanschroers <60888108+dylanschroers@users.noreply.github.com> Date: Fri, 12 Jun 2026 04:31:31 -0400 Subject: [PATCH 07/81] fix: decode subprocess output as UTF-8 in save.py on Windows (#6218) * Fix UnicodeDecodeError on Windows reading subprocess output in save path On Windows the default text encoding is the locale code page (cp1252), not UTF-8. The text-mode subprocess calls in save.py (text=True / universal_newlines=True) set no explicit encoding, so they decode llama.cpp / Ollama output with cp1252. When a child process emits a byte undefined in cp1252 -- e.g. 0x9d, which appears inside the UTF-8 encoding of common punctuation / box-drawing glyphs and in non-ASCII file paths -- the read raises UnicodeDecodeError and aborts GGUF export. Add encoding="utf-8", errors="replace" to all 8 text-mode subprocess calls. errors="replace" also avoids silent mojibake for inputs whose bytes happen to be valid-but-wrong in cp1252. Add tests/saving/test_save_subprocess_utf8_encoding.py: - an AST drift detector asserting every text-mode subprocess call in save.py pins encoding="utf-8" (runs without importing torch/unsloth_zoo) - a behavioural test reproducing the cp1252 failure and the utf-8 fix Relates-to: #2660 Co-Authored-By: Claude Opus 4.8 * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci --------- Co-authored-by: Claude Opus 4.8 Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com> --- .../test_save_subprocess_utf8_encoding.py | 148 ++++++++++++++++++ unsloth/save.py | 16 ++ 2 files changed, 164 insertions(+) create mode 100644 tests/saving/test_save_subprocess_utf8_encoding.py diff --git a/tests/saving/test_save_subprocess_utf8_encoding.py b/tests/saving/test_save_subprocess_utf8_encoding.py new file mode 100644 index 0000000000..4a609cd7b7 --- /dev/null +++ b/tests/saving/test_save_subprocess_utf8_encoding.py @@ -0,0 +1,148 @@ +# Unsloth - 2x faster, 60% less VRAM LLM training and finetuning +# Copyright 2023-present Daniel Han-Chen, Michael Han-Chen & the Unsloth team. All rights reserved. +# +# This program is free software: you can redistribute it and/or modify +# it under the terms of the GNU Lesser General Public License as published by +# the Free Software Foundation, either version 3 of the License, or +# (at your option) any later version. +# +# This program is distributed in the hope that it will be useful, +# but WITHOUT ANY WARRANTY; without even the implied warranty of +# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the +# GNU Lesser General Public License for more details. + +"""Regression tests for unslothai/unsloth#2660. + +On Windows the default text encoding is the locale code page (e.g. cp1252), +not UTF-8. ``subprocess.Popen`` / ``subprocess.run`` opened in text mode +(``text=True`` / ``universal_newlines=True``) without an explicit +``encoding`` therefore decode child-process output with cp1252. When +llama.cpp / Ollama emit a byte that is undefined in cp1252 (e.g. ``0x9d``, +which appears inside the UTF-8 encoding of common punctuation and box-drawing +glyphs), the read raises ``UnicodeDecodeError`` and aborts the GGUF export. + +Two checks: + +* ``test_save_subprocess_text_calls_declare_utf8_encoding`` -- a source-level + drift detector. It parses ``unsloth/save.py`` (no import, so it runs under + the GPU/torch-free harness) and fails if any text-mode subprocess call is + missing ``encoding="utf-8"``. This is the regression guard: it is red + before the fix and green after. +* ``test_utf8_replace_decodes_non_cp1252_subprocess_output`` -- a behavioural + check that documents the bug and the fix deterministically on any platform: + raw child output that is invalid under cp1252 raises, while the + ``encoding="utf-8", errors="replace"`` kwargs used by the fix read it + cleanly. +""" + +from __future__ import annotations + +import ast +import subprocess +import sys +from pathlib import Path + +import pytest + +SAVE_PY = Path(__file__).resolve().parents[2] / "unsloth" / "save.py" + + +def _is_subprocess_call(node: ast.Call) -> bool: + """True for ``subprocess.Popen(...)`` / ``subprocess.run(...)``.""" + func = node.func + return ( + isinstance(func, ast.Attribute) + and func.attr in {"Popen", "run"} + and isinstance(func.value, ast.Name) + and func.value.id == "subprocess" + ) + + +def _kw(node: ast.Call, name: str): + for kw in node.keywords: + if kw.arg == name: + return kw.value + return None + + +def _is_true(value) -> bool: + return isinstance(value, ast.Constant) and value.value is True + + +def _is_text_mode(node: ast.Call) -> bool: + """Text mode = ``text=True`` or ``universal_newlines=True``.""" + return _is_true(_kw(node, "text")) or _is_true(_kw(node, "universal_newlines")) + + +def _collect_text_mode_subprocess_calls() -> list[ast.Call]: + tree = ast.parse(SAVE_PY.read_text(encoding = "utf-8"), filename = str(SAVE_PY)) + return [ + node + for node in ast.walk(tree) + if isinstance(node, ast.Call) and _is_subprocess_call(node) and _is_text_mode(node) + ] + + +def test_text_mode_subprocess_calls_exist(): + """Guard the guard: if save.py stops using text-mode subprocess calls the + drift test below would vacuously pass, so make sure we are actually + inspecting something.""" + calls = _collect_text_mode_subprocess_calls() + assert len(calls) >= 6, ( + f"Expected several text-mode subprocess calls in {SAVE_PY.name}, " + f"found {len(calls)} -- has the file been restructured?" + ) + + +def test_save_subprocess_text_calls_declare_utf8_encoding(): + """Every text-mode subprocess call in save.py must pin encoding='utf-8'. + + Without it, reading llama.cpp/Ollama output crashes on Windows (cp1252). + Fails before the #2660 fix, passes after. + """ + offenders = [] + for node in _collect_text_mode_subprocess_calls(): + enc = _kw(node, "encoding") + ok = isinstance(enc, ast.Constant) and enc.value == "utf-8" + if not ok: + offenders.append(node.lineno) + + assert not offenders, ( + "Text-mode subprocess call(s) in unsloth/save.py missing " + 'encoding="utf-8" (UnicodeDecodeError on Windows, #2660) at line(s): ' + + ", ".join(map(str, sorted(offenders))) + ) + + +def test_utf8_replace_decodes_non_cp1252_subprocess_output(): + """Document the failure and the fix with a real subprocess. + + The child emits U+201D (right double quote), whose UTF-8 encoding + ``E2 80 9D`` contains byte 0x9D -- undefined in cp1252. Decoding the raw + bytes as cp1252 raises (the bug); the fix's kwargs read it cleanly. + """ + # All-ASCII argv; the child builds the non-ASCII char itself so this is + # deterministic regardless of the parent's locale. + child = ( + "import sys; " + "sys.stdout.buffer.write(('tensor ' + chr(0x201D) + ' x\\n').encode('utf-8'))" + ) + + raw = subprocess.run([sys.executable, "-c", child], capture_output = True).stdout + assert b"\x9d" in raw # precondition: output carries the cp1252-undefined byte + + # Failing behaviour before the fix: cp1252 (the Windows default) cannot + # decode this output. + with pytest.raises(UnicodeDecodeError): + raw.decode("cp1252") + + # Correct behaviour after the fix: the exact kwargs save.py now uses. + result = subprocess.run( + [sys.executable, "-c", child], + capture_output = True, + text = True, + encoding = "utf-8", + errors = "replace", + ) + assert result.stdout.startswith("tensor ") + assert "”" in result.stdout diff --git a/unsloth/save.py b/unsloth/save.py index 629cbb9548..a6cc665d3e 100644 --- a/unsloth/save.py +++ b/unsloth/save.py @@ -186,6 +186,8 @@ def _quantize_q2_k_l( command, shell = False, text = True, + encoding = "utf-8", + errors = "replace", stdout = subprocess.PIPE, stderr = subprocess.STDOUT, bufsize = 1, @@ -206,6 +208,8 @@ def _quantize_q2_k_l( check = True, capture_output = True, text = True, + encoding = "utf-8", + errors = "replace", ) except subprocess.CalledProcessError as e: if print_output and hasattr(e, "stdout") and e.stdout: @@ -1989,6 +1993,8 @@ def create_ollama_model(username: str, model_name: str, tag: str, modelfile_path ["curl", "http://localhost:11434"], capture_output = True, text = True, + encoding = "utf-8", + errors = "replace", timeout = 3, ) if init_check.returncode == 0: @@ -2011,6 +2017,8 @@ def create_ollama_model(username: str, model_name: str, tag: str, modelfile_path text = True, bufsize = 1, universal_newlines = True, + encoding = "utf-8", + errors = "replace", ) for line in iter(process.stdout.readline, ""): @@ -2031,6 +2039,8 @@ def push_to_ollama_hub(username: str, model_name: str, tag: str): ["curl", "http://localhost:11434"], capture_output = True, text = True, + encoding = "utf-8", + errors = "replace", timeout = 3, ) if init_check.returncode == 0: @@ -2047,6 +2057,8 @@ def push_to_ollama_hub(username: str, model_name: str, tag: str): text = True, bufsize = 1, universal_newlines = True, + encoding = "utf-8", + errors = "replace", ) for line in iter(process.stdout.readline, ""): @@ -2758,6 +2770,8 @@ def unsloth_convert_lora_to_ggml_and_push_to_hub( stderr = subprocess.PIPE, bufsize = 1, universal_newlines = True, + encoding = "utf-8", + errors = "replace", ) as sp: for line in sp.stdout: print(line, end = "", flush = True) @@ -2838,6 +2852,8 @@ def unsloth_convert_lora_to_ggml_and_save_locally( stderr = subprocess.PIPE, bufsize = 1, universal_newlines = True, + encoding = "utf-8", + errors = "replace", ) as sp: for line in sp.stdout: print(line, end = "", flush = True) From 514850fb327330448a285fa20e88834e1b9d035a Mon Sep 17 00:00:00 2001 From: Mohammad Hussian Date: Fri, 12 Jun 2026 14:15:19 +0530 Subject: [PATCH 08/81] patch: fix EmptyLogits gathering in nested payloads and Accelerate recursively_apply (#6092) * Fix EmptyLogits gathering in nested structure and patch recursively_apply on accelerator module * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * Wire EmptyLogits Accelerate patch into startup and fix find_device, pickling, tests for PR #6092 - Call patch_accelerate_recursively_apply() in _gpu_init.py so real imports install it; previously it was only invoked by the tests - Make both wrappers idempotent so repeated calls do not stack - Rework find_device: skip EmptyLogits while still finding real tensors in any order, keep returning None for tensor-free payloads (AlignDevicesHook relies on None), fall back to PartialState().device only for sentinel-only payloads - Give EmptyLogits stateless __reduce__ and drop the stomped pickle stubs on EMPTY_LOGITS so debug mode gather_object works in real distributed runs - Put test tensors on PartialState().device so the debug mode test also passes on GPU machines, and add drift tests for startup wiring, idempotency and find_device ordering Verified on 2x B200: ACCELERATE_DEBUG_MODE=1 torchrun gather/broadcast/pad of sentinel and mixed payloads all pass, training losses unchanged, full drift suite 25/25. * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * Define EmptyLogits equality on the class for PR #6092 Gathered sentinel copies must compare equal in accelerate debug mode regardless of whether the patched recursively_apply saw the sentinel first in that process. Class body __eq__ requires restoring __hash__ explicitly. Verified: 123 case simulation battery on accelerate 0.34.2 through latest, 2 process gloo CPU and NCCL GPU debug mode runs, drift suite 25/25. --------- Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com> Co-authored-by: Daniel Han --- tests/test_import_fixes_drift.py | 158 +++++++++++++++++++++++++++++++ unsloth/_gpu_init.py | 3 + unsloth/import_fixes.py | 90 ++++++++++++++++++ unsloth/models/_utils.py | 17 ++++ 4 files changed, 268 insertions(+) diff --git a/tests/test_import_fixes_drift.py b/tests/test_import_fixes_drift.py index 6c75a4b886..c9f971c891 100644 --- a/tests/test_import_fixes_drift.py +++ b/tests/test_import_fixes_drift.py @@ -586,6 +586,164 @@ def test_accelerate_utils_imports_module_present(): ) +def test_accelerate_recursively_apply_empty_logits_patch(): + """Verify patch_accelerate_recursively_apply overrides recursively_apply to bypass EmptyLogits.""" + pytest.importorskip("accelerate") + + import accelerate.utils.operations as acc_ops + from unsloth.import_fixes import patch_accelerate_recursively_apply + + class EmptyLogits: + pass + + e = EmptyLogits() + patch_accelerate_recursively_apply() + + res = acc_ops.recursively_apply(lambda x: x, e, error_on_other_type = True) + assert res is e + + +def test_accelerate_gather_empty_logits_debug_mode_patch(): + """Verify gather and broadcast bypass EmptyLogits when debug mode is enabled.""" + pytest.importorskip("accelerate") + from accelerate.state import PartialState, DistributedType + import accelerate.utils.operations as acc_ops + from unsloth.import_fixes import patch_accelerate_recursively_apply + import unittest.mock as mock + import torch + + class EmptyLogits: + pass + + e = EmptyLogits() + patch_accelerate_recursively_apply() + + # Enable debug mode and mock distributed state + state = PartialState() + orig_debug = state.debug + orig_dist_type = state.distributed_type + orig_num_processes = state.num_processes + + state.debug = True + state.distributed_type = DistributedType.MULTI_GPU + state.num_processes = 2 + + # Mock gather_object to return [obj] * num_processes + def mock_gather_object(obj, *args, **kwargs): + return [obj] * state.num_processes + + # Mock _gpu_gather to recursively apply replication of tensors + def mock_gpu_gather(tensor, *args, **kwargs): + def _gather_one(t): + if t.ndim == 0: + t = t.clone()[None] + return torch.cat([t] * state.num_processes, dim = 0) + + return acc_ops.recursively_apply(_gather_one, tensor, error_on_other_type = True) + + # Mock _gpu_broadcast to return data unchanged + def mock_gpu_broadcast(data, *args, **kwargs): + return data + + try: + with ( + mock.patch( + "accelerate.utils.operations.gather_object", + side_effect = mock_gather_object, + ), + mock.patch("accelerate.utils.operations._gpu_gather", side_effect = mock_gpu_gather), + mock.patch( + "accelerate.utils.operations._gpu_broadcast", + side_effect = mock_gpu_broadcast, + ), + ): + # 1. Top-level EmptyLogits should gather correctly (returns e) + res = acc_ops.gather(e) + assert res is e + + # 2. Nested EmptyLogits alone + res_nested = acc_ops.gather([e]) + assert isinstance(res_nested, list) and res_nested[0] is e + + # 3. Mixed payload with real tensor and EmptyLogits + # Real tensor should be gathered (concatenated across processes). + # Tensors must live on state.device or the debug-mode device + # check fails on GPU machines. + real_tensor = torch.tensor([42], device = state.device) + payload = {"labels": real_tensor, "logits": e} + res_mixed = acc_ops.gather(payload) + + assert isinstance(res_mixed, dict) + assert res_mixed["logits"] is e + # Since num_processes = 2, it should be gathered to [42, 42] + assert torch.equal(res_mixed["labels"], torch.tensor([42, 42], device = state.device)) + + # 4. Broadcast with EmptyLogits + res_broadcast = acc_ops.broadcast(e) + assert res_broadcast is e + + # 5. Mixed payload with broadcast + res_broadcast_mixed = acc_ops.broadcast(payload) + assert isinstance(res_broadcast_mixed, dict) + assert res_broadcast_mixed["logits"] is e + assert torch.equal(res_broadcast_mixed["labels"], real_tensor) + finally: + state.debug = orig_debug + state.distributed_type = orig_dist_type + state.num_processes = orig_num_processes + + +def test_accelerate_patch_is_idempotent(): + """Calling patch_accelerate_recursively_apply twice must not stack wrappers.""" + pytest.importorskip("accelerate") + import accelerate.utils.operations as acc_ops + from unsloth.import_fixes import patch_accelerate_recursively_apply + + patch_accelerate_recursively_apply() + recursively_apply = acc_ops.recursively_apply + find_device = acc_ops.find_device + patch_accelerate_recursively_apply() + assert ( + acc_ops.recursively_apply is recursively_apply + ), "DRIFT DETECTED: recursively_apply was wrapped twice." + assert acc_ops.find_device is find_device, "DRIFT DETECTED: find_device was wrapped twice." + + +def test_accelerate_find_device_skips_empty_logits(): + """find_device must search past EmptyLogits and keep None for tensor-free data.""" + pytest.importorskip("accelerate") + import torch + import accelerate.utils.operations as acc_ops + from accelerate.state import PartialState + from unsloth.import_fixes import patch_accelerate_recursively_apply + + class EmptyLogits: + pass + + patch_accelerate_recursively_apply() + tensor = torch.tensor([1.0]) + # Sentinel first must not stop the search before the real tensor + assert acc_ops.find_device({"logits": EmptyLogits(), "labels": tensor}) == tensor.device + # Tensor-free payloads without the sentinel keep returning None + # (AlignDevicesHook relies on None to skip output device moves) + assert acc_ops.find_device({"a": 1}) is None + # Sentinel-only payloads fall back to the current device so that + # debug mode find_device(...).type does not raise AttributeError + assert acc_ops.find_device(EmptyLogits()) == PartialState().device + + +def test_accelerate_patch_wired_into_gpu_init(): + """The patch must be installed at startup, not only importable.""" + import pathlib + import unsloth.import_fixes as import_fixes + + source = pathlib.Path(import_fixes.__file__).with_name("_gpu_init.py").read_text() + assert "patch_accelerate_recursively_apply()" in source, ( + "DRIFT DETECTED: patch_accelerate_recursively_apply is defined but " + "never called in _gpu_init.py, so real imports never install it." + ) + + # =========================================================================== # bitsandbytes -- ROCm arch / warp-size detection shape # =========================================================================== diff --git a/unsloth/_gpu_init.py b/unsloth/_gpu_init.py index 5da1e27a9d..86f675c281 100644 --- a/unsloth/_gpu_init.py +++ b/unsloth/_gpu_init.py @@ -189,6 +189,7 @@ from .import_fixes import ( fix_trl_vllm_ascend, fix_peft_transformers_weight_conversion_import, patch_peft_weight_converter_compatibility, + patch_accelerate_recursively_apply, ) fix_xformers_performance_issue() @@ -217,6 +218,7 @@ disable_broken_wandb() # build_peft_weight_mapping instead of being swallowed by its ImportError. fix_peft_transformers_weight_conversion_import() patch_peft_weight_converter_compatibility() +patch_accelerate_recursively_apply() del fix_xformers_performance_issue del fix_vllm_aimv2_issue @@ -240,6 +242,7 @@ del disable_torchcodec_if_broken del disable_broken_wandb del fix_peft_transformers_weight_conversion_import del patch_peft_weight_converter_compatibility +del patch_accelerate_recursively_apply # Torch 2.4 has including_emulation if DEVICE_TYPE == "cuda": diff --git a/unsloth/import_fixes.py b/unsloth/import_fixes.py index 695a6a577a..a100e11f0c 100644 --- a/unsloth/import_fixes.py +++ b/unsloth/import_fixes.py @@ -2584,3 +2584,93 @@ def maybe_set_windows_rocm_bnb_version(): "(detected from the installed bitsandbytes ROCm wheel on Windows)." ) return version + + +def patch_accelerate_recursively_apply(): + """ + Make Accelerate's recursive utilities tolerate Unsloth's EmptyLogits + sentinel. recursively_apply returns the sentinel unchanged instead of + raising TypeError, and find_device skips it while still finding real + tensors, falling back to PartialState().device only for sentinel-only + payloads. Both wrappers are idempotent and are propagated to every + already imported accelerate namespace. + """ + try: + import accelerate.utils.operations as acc_ops + except Exception: + return + + original_recursively_apply = getattr(acc_ops, "recursively_apply", None) + if original_recursively_apply is not None and not getattr( + original_recursively_apply, "__unsloth_patched__", False + ): + + @functools.wraps(original_recursively_apply) + def _patched_recursively_apply(func, data, *args, **kwargs): + if type(data).__name__ == "EmptyLogits": + cls = type(data) + if cls.__eq__ is object.__eq__: + # Debug mode compares gathered metadata across ranks with == + cls.__eq__ = lambda self, other: type(other).__name__ == "EmptyLogits" + return data + return original_recursively_apply(func, data, *args, **kwargs) + + _patched_recursively_apply.__unsloth_patched__ = True + + for mod_name, mod in tuple(sys.modules.items()): + if mod_name.startswith("accelerate") and mod is not None: + if getattr(mod, "recursively_apply", None) is original_recursively_apply: + try: + setattr(mod, "recursively_apply", _patched_recursively_apply) + except Exception: + pass + + original_find_device = getattr(acc_ops, "find_device", None) + if original_find_device is not None and not getattr( + original_find_device, "__unsloth_patched__", False + ): + from collections.abc import Mapping + + @functools.wraps(original_find_device) + def _patched_find_device(data): + import torch + + found_sentinel = False + + def _search(obj): + nonlocal found_sentinel + if type(obj).__name__ == "EmptyLogits": + found_sentinel = True + elif isinstance(obj, Mapping): + for value in obj.values(): + device = _search(value) + if device is not None: + return device + elif isinstance(obj, (tuple, list)): + for value in obj: + device = _search(value) + if device is not None: + return device + elif isinstance(obj, torch.Tensor): + return obj.device + return None + + device = _search(data) + if device is None and found_sentinel: + # Debug mode calls find_device(...).type on gather/broadcast inputs + try: + from accelerate.state import PartialState + return PartialState().device + except Exception: + pass + return device + + _patched_find_device.__unsloth_patched__ = True + + for mod_name, mod in tuple(sys.modules.items()): + if mod_name.startswith("accelerate") and mod is not None: + if getattr(mod, "find_device", None) is original_find_device: + try: + setattr(mod, "find_device", _patched_find_device) + except Exception: + pass diff --git a/unsloth/models/_utils.py b/unsloth/models/_utils.py index 2f4e3a069e..6024a02c2c 100644 --- a/unsloth/models/_utils.py +++ b/unsloth/models/_utils.py @@ -2703,6 +2703,16 @@ class EmptyLogits: def __str__(self): return LOGITS_ERROR_STRING + def __reduce__(self): + # Stateless pickling so gather_object works on the sentinel + return (type(self), ()) + + def __eq__(self, other): + # Gathered copies must compare equal in accelerate debug mode + return type(other).__name__ == "EmptyLogits" + + __hash__ = object.__hash__ + EMPTY_LOGITS = EmptyLogits() functions = dir(torch.Tensor) @@ -2713,6 +2723,13 @@ for j, function in enumerate(functions): exec(f"EMPTY_LOGITS.{function} = raise_{j}", globals(), locals()) except: continue +# The loop above stomps pickle hooks with stubs returning None, which breaks +# gather_object on EMPTY_LOGITS in distributed runs. Restore default pickling. +for function in ("__reduce__", "__reduce_ex__", "__getstate__", "__setstate__"): + try: + delattr(EMPTY_LOGITS, function) + except Exception: + pass def validate_loftq_config(loftq_config, lora_dropout, bias, init_lora_weights, model): From de0c5a2f09ba8230987dc6a580fe77c1a6da9bba Mon Sep 17 00:00:00 2001 From: Ban <3637117+Ban921@users.noreply.github.com> Date: Fri, 12 Jun 2026 16:50:45 +0800 Subject: [PATCH 09/81] Studio: show Apple GPU temperature and power in the GPU monitor (macOS) (#6187) * Studio: show Apple GPU temperature and power in the GPU monitor (macOS) The GPU monitor on Apple Silicon always showed -- for Temperature and Power: the MLX branch of get_gpu_utilization() hardcoded None because ioreg's AGXAccelerator PerformanceStatistics carries neither metric. Add utils/hardware/apple.py, mirroring macmon's no-sudo approach: - Temperature: average of the AppleSMC "Tg*" float keys via the AppleSMCKeysEndpoint user client (ctypes/IOKit, macOS 14+). - Power: IOReport "Energy Model" group, "GPU Energy" channels; each poll diffs the energy counter against the previous poll's sample, so the value is the average wattage over the polling window. The first poll only sets the baseline and returns None. Both readers latch to None on first failure and never raise, so non-Mac platforms and locked-down hosts keep the previous behavior. * Sample IOReport with the subscribed channels descriptor for PR #6187 IOReportCreateSubscription writes the channel descriptor that later samples must use; sampling with the original requested group can return no Energy Model entries on hosts that normalize the channel set, leaving power_draw_w null after the baseline. Use the subscribed descriptor (matching macmon) and fall back to the requested channels if the OS leaves it unset. --------- Co-authored-by: Lee Jackson <130007945+Imagineer99@users.noreply.github.com> Co-authored-by: danielhanchen <23090290+danielhanchen@users.noreply.github.com> --- .../backend/tests/test_apple_gpu_sensors.py | 74 +++ studio/backend/utils/hardware/apple.py | 430 ++++++++++++++++++ studio/backend/utils/hardware/hardware.py | 6 +- 3 files changed, 508 insertions(+), 2 deletions(-) create mode 100644 studio/backend/tests/test_apple_gpu_sensors.py create mode 100644 studio/backend/utils/hardware/apple.py diff --git a/studio/backend/tests/test_apple_gpu_sensors.py b/studio/backend/tests/test_apple_gpu_sensors.py new file mode 100644 index 0000000000..50947d1380 --- /dev/null +++ b/studio/backend/tests/test_apple_gpu_sensors.py @@ -0,0 +1,74 @@ +# SPDX-License-Identifier: AGPL-3.0-only +# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0 + +"""Tests for Apple Silicon GPU sensors (SMC temperature + IOReport power).""" + +import ctypes +import platform +import time + +import pytest + +from utils.hardware import apple + +_IS_APPLE_SILICON = platform.system() == "Darwin" and platform.machine() == "arm64" + + +class TestFourcc: + def test_roundtrip(self): + for key in ("#KEY", "Tg0D", "flt "): + assert apple._fourcc_str(apple._fourcc(key)) == key + + def test_known_value(self): + # "flt " FourCC, same constant macmon uses. + assert apple._fourcc("flt ") == 1718383648 + + +class TestWatts: + def test_millijoules(self): + assert apple._watts(2000, "mJ", 2.0) == pytest.approx(1.0) + + def test_microjoules(self): + assert apple._watts(5_000_000, "uJ", 1.0) == pytest.approx(5.0) + + def test_nanojoules(self): + assert apple._watts(1_500_000_000, "nJ", 1.0) == pytest.approx(1.5) + + def test_unknown_unit_returns_none(self): + assert apple._watts(1000, "J", 1.0) is None + + def test_zero_elapsed_returns_none(self): + assert apple._watts(1000, "mJ", 0.0) is None + + +class TestAverageValidTemps: + def test_averages_and_rounds(self): + assert apple._average_valid_temps([40.0, 50.0, 60.05]) == 50.0 + + def test_filters_invalid(self): + assert apple._average_valid_temps([-1.0, 0.0, 151.0, 42.0]) == 42.0 + + def test_empty_returns_none(self): + assert apple._average_valid_temps([]) is None + assert apple._average_valid_temps([0.0, 200.0]) is None + + +class TestSmcStructLayout: + def test_key_data_matches_smc_protocol_size(self): + # The AppleSMC user client rejects calls whose struct size differs. + assert ctypes.sizeof(apple._SMCKeyData) == 80 + + +@pytest.mark.skipif(not _IS_APPLE_SILICON, reason = "requires Apple Silicon") +class TestLiveSensors: + def test_gpu_temperature_in_plausible_range(self): + temp = apple.read_gpu_temperature_c() + assert temp is not None + assert 0.0 < temp <= 150.0 + + def test_gpu_power_after_baseline(self): + apple.read_gpu_power_w() # first call only sets the baseline + time.sleep(0.3) + power = apple.read_gpu_power_w() + assert power is not None + assert power >= 0.0 diff --git a/studio/backend/utils/hardware/apple.py b/studio/backend/utils/hardware/apple.py new file mode 100644 index 0000000000..3f14af60ad --- /dev/null +++ b/studio/backend/utils/hardware/apple.py @@ -0,0 +1,430 @@ +# SPDX-License-Identifier: AGPL-3.0-only +# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0 + +"""Apple Silicon GPU temperature and power -- no sudo required. + +Mirrors macmon's approach (https://github.com/vladkens/macmon): + * Temperature: average of the AppleSMC "Tg*" float keys (available since + macOS 14; on older systems the keys are absent and this returns None). + * Power: IOReport "Energy Model" group, "GPU Energy" channel. Each poll + diffs the energy counter against the previous poll's sample, so the + result is the average wattage over the polling window. The first poll + only sets the baseline and returns None. + +Public API (never raises; returns None when sensors are unavailable): + read_gpu_temperature_c() + read_gpu_power_w() +""" + +import ctypes +import struct +import time +from typing import Iterable, Optional + +from loggers import get_logger + +logger = get_logger(__name__) + +_IOKIT_PATH = "/System/Library/Frameworks/IOKit.framework/IOKit" +_CF_PATH = "/System/Library/Frameworks/CoreFoundation.framework/CoreFoundation" +_IOREPORT_PATH = "/usr/lib/libIOReport.dylib" + +# AppleSMC user-client protocol (same constants as macmon / SMCKit). +_SMC_SELECTOR_HANDLE_EVENT = 2 +_SMC_CMD_READ_BYTES = 5 +_SMC_CMD_KEY_AT_INDEX = 8 +_SMC_CMD_KEY_INFO = 9 + +_MAX_VALID_TEMP_C = 150.0 +_CF_STRING_ENCODING_UTF8 = 0x08000100 +_ENERGY_UNIT_DIVISORS = {"mJ": 1e3, "uJ": 1e6, "nJ": 1e9} + + +# ========== Pure helpers ========== + + +def _fourcc(key: str) -> int: + """Encode a 4-char SMC key/type name as a big-endian integer.""" + return int.from_bytes(key.encode("ascii"), "big") + + +def _fourcc_str(value: int) -> str: + return value.to_bytes(4, "big").decode("ascii", errors = "replace") + + +def _watts(energy: int, unit: str, elapsed_s: float) -> Optional[float]: + """Convert an IOReport energy counter delta into average watts.""" + divisor = _ENERGY_UNIT_DIVISORS.get(unit.strip()) + if divisor is None or elapsed_s <= 0: + return None + return energy / divisor / elapsed_s + + +def _average_valid_temps(values: Iterable[float]) -> Optional[float]: + valid = [v for v in values if 0.0 < v <= _MAX_VALID_TEMP_C] + if not valid: + return None + return round(sum(valid) / len(valid), 1) + + +def _is_gpu_energy_channel(name: str) -> bool: + # Exact "GPU Energy" plus "DIE_N_GPU Energy" on Ultra chips; the separate + # "GPU SRAM*" channels are not GPU core power. + return name.endswith("GPU Energy") and "SRAM" not in name + + +# ========== AppleSMC structs (layout must match the kernel exactly) ========== + + +class _SMCKeyDataVers(ctypes.Structure): + _fields_ = [ + ("major", ctypes.c_uint8), + ("minor", ctypes.c_uint8), + ("build", ctypes.c_uint8), + ("reserved", ctypes.c_uint8), + ("release", ctypes.c_uint16), + ] + + +class _SMCPLimitData(ctypes.Structure): + _fields_ = [ + ("version", ctypes.c_uint16), + ("length", ctypes.c_uint16), + ("cpu_p_limit", ctypes.c_uint32), + ("gpu_p_limit", ctypes.c_uint32), + ("mem_p_limit", ctypes.c_uint32), + ] + + +class _SMCKeyInfo(ctypes.Structure): + _fields_ = [ + ("data_size", ctypes.c_uint32), + ("data_type", ctypes.c_uint32), + ("data_attributes", ctypes.c_uint8), + ] + + +class _SMCKeyData(ctypes.Structure): + _fields_ = [ + ("key", ctypes.c_uint32), + ("vers", _SMCKeyDataVers), + ("p_limit_data", _SMCPLimitData), + ("key_info", _SMCKeyInfo), + ("result", ctypes.c_uint8), + ("status", ctypes.c_uint8), + ("data8", ctypes.c_uint8), + ("data32", ctypes.c_uint32), + ("bytes", ctypes.c_uint8 * 32), + ] + + +# ========== Library loaders ========== + + +def _load_iokit() -> ctypes.CDLL: + iokit = ctypes.CDLL(_IOKIT_PATH) + iokit.IOServiceMatching.restype = ctypes.c_void_p + iokit.IOServiceMatching.argtypes = [ctypes.c_char_p] + iokit.IOServiceGetMatchingServices.argtypes = [ + ctypes.c_uint32, + ctypes.c_void_p, + ctypes.POINTER(ctypes.c_uint32), + ] + iokit.IOIteratorNext.restype = ctypes.c_uint32 + iokit.IOIteratorNext.argtypes = [ctypes.c_uint32] + iokit.IORegistryEntryGetName.argtypes = [ctypes.c_uint32, ctypes.c_char_p] + iokit.IOServiceOpen.argtypes = [ + ctypes.c_uint32, + ctypes.c_uint32, + ctypes.c_uint32, + ctypes.POINTER(ctypes.c_uint32), + ] + iokit.IOObjectRelease.argtypes = [ctypes.c_uint32] + iokit.IOConnectCallStructMethod.argtypes = [ + ctypes.c_uint32, + ctypes.c_uint32, + ctypes.c_void_p, + ctypes.c_size_t, + ctypes.c_void_p, + ctypes.POINTER(ctypes.c_size_t), + ] + return iokit + + +def _load_cf() -> ctypes.CDLL: + cf = ctypes.CDLL(_CF_PATH) + cf.CFStringCreateWithCString.restype = ctypes.c_void_p + cf.CFStringCreateWithCString.argtypes = [ctypes.c_void_p, ctypes.c_char_p, ctypes.c_uint32] + cf.CFStringGetCString.restype = ctypes.c_bool + cf.CFStringGetCString.argtypes = [ + ctypes.c_void_p, + ctypes.c_char_p, + ctypes.c_long, + ctypes.c_uint32, + ] + cf.CFRelease.argtypes = [ctypes.c_void_p] + cf.CFDictionaryGetValue.restype = ctypes.c_void_p + cf.CFDictionaryGetValue.argtypes = [ctypes.c_void_p, ctypes.c_void_p] + cf.CFArrayGetCount.restype = ctypes.c_long + cf.CFArrayGetCount.argtypes = [ctypes.c_void_p] + cf.CFArrayGetValueAtIndex.restype = ctypes.c_void_p + cf.CFArrayGetValueAtIndex.argtypes = [ctypes.c_void_p, ctypes.c_long] + return cf + + +def _load_ioreport() -> ctypes.CDLL: + ior = ctypes.CDLL(_IOREPORT_PATH) + ior.IOReportCopyChannelsInGroup.restype = ctypes.c_void_p + ior.IOReportCopyChannelsInGroup.argtypes = [ + ctypes.c_void_p, + ctypes.c_void_p, + ctypes.c_uint64, + ctypes.c_uint64, + ctypes.c_uint64, + ] + ior.IOReportCreateSubscription.restype = ctypes.c_void_p + ior.IOReportCreateSubscription.argtypes = [ + ctypes.c_void_p, + ctypes.c_void_p, + ctypes.POINTER(ctypes.c_void_p), + ctypes.c_uint64, + ctypes.c_void_p, + ] + ior.IOReportCreateSamples.restype = ctypes.c_void_p + ior.IOReportCreateSamples.argtypes = [ctypes.c_void_p, ctypes.c_void_p, ctypes.c_void_p] + ior.IOReportCreateSamplesDelta.restype = ctypes.c_void_p + ior.IOReportCreateSamplesDelta.argtypes = [ctypes.c_void_p, ctypes.c_void_p, ctypes.c_void_p] + ior.IOReportChannelGetChannelName.restype = ctypes.c_void_p + ior.IOReportChannelGetChannelName.argtypes = [ctypes.c_void_p] + ior.IOReportChannelGetUnitLabel.restype = ctypes.c_void_p + ior.IOReportChannelGetUnitLabel.argtypes = [ctypes.c_void_p] + ior.IOReportSimpleGetIntegerValue.restype = ctypes.c_int64 + ior.IOReportSimpleGetIntegerValue.argtypes = [ctypes.c_void_p, ctypes.c_int32] + return ior + + +def _cfstr(cf: ctypes.CDLL, text: str) -> int: + return cf.CFStringCreateWithCString(None, text.encode("utf-8"), _CF_STRING_ENCODING_UTF8) + + +def _from_cfstr(cf: ctypes.CDLL, ref: Optional[int]) -> str: + if not ref: + return "" + buf = ctypes.create_string_buffer(128) + if not cf.CFStringGetCString(ref, buf, len(buf), _CF_STRING_ENCODING_UTF8): + return "" + return buf.value.decode("utf-8", errors = "replace").strip() + + +# ========== SMC connection (GPU temperature) ========== + + +class _SMCConnection: + """Connection to AppleSMCKeysEndpoint; discovers "Tg*" GPU temp keys once.""" + + def __init__(self): + self._iokit = _load_iokit() + self._conn = self._open() + self._key_info_cache: dict[int, _SMCKeyInfo] = {} + self.gpu_keys = [ + key + for key in self._all_keys() + if key.startswith("Tg") and self.read_float(key) is not None + ] + + def _open(self) -> int: + iterator = ctypes.c_uint32(0) + matching = self._iokit.IOServiceMatching(b"AppleSMC") + if self._iokit.IOServiceGetMatchingServices(0, matching, ctypes.byref(iterator)) != 0: + raise OSError("AppleSMC service not found") + try: + conn = self._open_keys_endpoint(iterator.value) + finally: + self._iokit.IOObjectRelease(iterator.value) + if conn is None: + raise OSError("AppleSMCKeysEndpoint not found") + return conn + + def _open_keys_endpoint(self, iterator: int) -> Optional[int]: + task = ctypes.CDLL(None).mach_task_self() + while device := self._iokit.IOIteratorNext(iterator): + name = ctypes.create_string_buffer(128) + self._iokit.IORegistryEntryGetName(device, name) + if name.value != b"AppleSMCKeysEndpoint": + self._iokit.IOObjectRelease(device) + continue + conn = ctypes.c_uint32(0) + status = self._iokit.IOServiceOpen(device, task, 0, ctypes.byref(conn)) + self._iokit.IOObjectRelease(device) + if status != 0: + raise OSError(f"IOServiceOpen(AppleSMCKeysEndpoint) failed: {status}") + return conn.value + return None + + def _call(self, ival: _SMCKeyData) -> _SMCKeyData: + oval = _SMCKeyData() + olen = ctypes.c_size_t(ctypes.sizeof(_SMCKeyData)) + status = self._iokit.IOConnectCallStructMethod( + self._conn, + _SMC_SELECTOR_HANDLE_EVENT, + ctypes.byref(ival), + ctypes.sizeof(_SMCKeyData), + ctypes.byref(oval), + ctypes.byref(olen), + ) + if status != 0: + raise OSError(f"IOConnectCallStructMethod failed: {status}") + if oval.result != 0: + raise OSError(f"SMC result code: {oval.result}") + return oval + + def _read_key_info(self, key_id: int) -> _SMCKeyInfo: + cached = self._key_info_cache.get(key_id) + if cached is not None: + return cached + oval = self._call(_SMCKeyData(key = key_id, data8 = _SMC_CMD_KEY_INFO)) + self._key_info_cache[key_id] = oval.key_info + return oval.key_info + + def _read_bytes(self, key: str) -> Optional[bytes]: + try: + key_id = _fourcc(key) + info = self._read_key_info(key_id) + oval = self._call(_SMCKeyData(key = key_id, data8 = _SMC_CMD_READ_BYTES, key_info = info)) + return bytes(oval.bytes[: info.data_size]) + except OSError: + return None + + def read_float(self, key: str) -> Optional[float]: + try: + info = self._read_key_info(_fourcc(key)) + except OSError: + return None + if info.data_size != 4 or info.data_type != _fourcc("flt "): + return None + data = self._read_bytes(key) + if data is None or len(data) != 4: + return None + return struct.unpack(" Optional[str]: + try: + oval = self._call(_SMCKeyData(data8 = _SMC_CMD_KEY_AT_INDEX, data32 = index)) + return oval.key.to_bytes(4, "big").decode("ascii") + except (OSError, UnicodeDecodeError): + return None + + def _all_keys(self) -> list[str]: + count_bytes = self._read_bytes("#KEY") + if count_bytes is None or len(count_bytes) != 4: + return [] + count = int.from_bytes(count_bytes, "big") + names = (self._key_name_at(i) for i in range(count)) + return [name for name in names if name is not None] + + def gpu_temperature_c(self) -> Optional[float]: + readings = (self.read_float(key) for key in self.gpu_keys) + return _average_valid_temps(value for value in readings if value is not None) + + +# ========== IOReport subscription (GPU power) ========== + + +class _IOReportEnergy: + """Persistent subscription to the "Energy Model" group for GPU wattage.""" + + def __init__(self): + self._cf = _load_cf() + self._ior = _load_ioreport() + self._channels = self._ior.IOReportCopyChannelsInGroup( + _cfstr(self._cf, "Energy Model"), None, 0, 0, 0 + ) + if not self._channels: + raise OSError("IOReport 'Energy Model' channel group unavailable") + subscribed = ctypes.c_void_p() + self._sub = self._ior.IOReportCreateSubscription( + None, self._channels, ctypes.byref(subscribed), 0, None + ) + if not self._sub: + raise OSError("IOReportCreateSubscription failed") + # Sample with the channels IOReport subscribes us to, not the requested + # group (matches macmon); fall back if the OS leaves it unset. + self._sample_channels = subscribed if subscribed else self._channels + self._channels_key = _cfstr(self._cf, "IOReportChannels") + self._prev: Optional[tuple[int, float]] = None # (sample ref, monotonic s) + + def gpu_power_w(self) -> Optional[float]: + sample = self._ior.IOReportCreateSamples(self._sub, self._sample_channels, None) + if not sample: + return None + now = time.monotonic() + prev, self._prev = self._prev, (sample, now) + if prev is None: + return None + prev_sample, prev_time = prev + delta = self._ior.IOReportCreateSamplesDelta(prev_sample, sample, None) + self._cf.CFRelease(prev_sample) + if not delta: + return None + try: + return self._gpu_watts_from_delta(delta, now - prev_time) + finally: + self._cf.CFRelease(delta) + + def _gpu_watts_from_delta(self, delta: int, elapsed_s: float) -> Optional[float]: + items = self._cf.CFDictionaryGetValue(delta, self._channels_key) + if not items: + return None + total: Optional[float] = None + for i in range(self._cf.CFArrayGetCount(items)): + item = self._cf.CFArrayGetValueAtIndex(items, i) + name = _from_cfstr(self._cf, self._ior.IOReportChannelGetChannelName(item)) + if not _is_gpu_energy_channel(name): + continue + unit = _from_cfstr(self._cf, self._ior.IOReportChannelGetUnitLabel(item)) + energy = self._ior.IOReportSimpleGetIntegerValue(item, 0) + watts = _watts(energy, unit, elapsed_s) + if watts is not None: + total = (total or 0.0) + watts + return round(total, 1) if total is not None else None + + +# ========== Public API (module singletons, failure-latched) ========== + +_smc: Optional[_SMCConnection] = None +_smc_failed = False +_energy: Optional[_IOReportEnergy] = None +_energy_failed = False + + +def read_gpu_temperature_c() -> Optional[float]: + """Average Apple GPU die temperature in degrees C, or None if unavailable.""" + global _smc, _smc_failed + if _smc_failed: + return None + try: + if _smc is None: + _smc = _SMCConnection() + return _smc.gpu_temperature_c() + except Exception as e: + _smc_failed = True + logger.warning("Apple SMC GPU temperature unavailable: %s", e) + return None + + +def read_gpu_power_w() -> Optional[float]: + """Average GPU power in watts since the previous call, or None. + + The first call establishes the baseline sample and returns None. + """ + global _energy, _energy_failed + if _energy_failed: + return None + try: + if _energy is None: + _energy = _IOReportEnergy() + return _energy.gpu_power_w() + except Exception as e: + _energy_failed = True + logger.warning("Apple IOReport GPU power unavailable: %s", e) + return None diff --git a/studio/backend/utils/hardware/hardware.py b/studio/backend/utils/hardware/hardware.py index 893d0364f7..86dfa8a93f 100644 --- a/studio/backend/utils/hardware/hardware.py +++ b/studio/backend/utils/hardware/hardware.py @@ -714,17 +714,19 @@ def get_gpu_utilization() -> Dict[str, Any]: except Exception: pass + from . import apple + return { "available": True, "backend": device.value, "gpu_utilization_pct": agx.get("utilization_pct") if agx else None, - "temperature_c": None, + "temperature_c": apple.read_gpu_temperature_c(), "vram_used_gb": round(vram_used_gb, 2), "vram_total_gb": round(total_gb, 2), "vram_utilization_pct": ( round((vram_used_gb / total_gb) * 100, 1) if total_gb > 0 else None ), - "power_draw_w": None, + "power_draw_w": apple.read_gpu_power_w(), "power_limit_w": None, "power_utilization_pct": None, } From 7f2986a413c0ddd2bbc077b58dbd82a750337e45 Mon Sep 17 00:00:00 2001 From: oobabooga Date: Fri, 12 Jun 2026 05:55:26 -0300 Subject: [PATCH 10/81] Studio: Add inline confirmation (Allow/Always allow/Deny) for tool calls (#5869) * Studio: Add inline confirmation (Allow/Always allow/Deny) for tool calls * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * Fix race in tool-call confirmation gate * Studio: gate built-in tool calls and harden the confirmation handshake The Allow / Always allow / Deny controls only lived in the fallback tool card, but the built-in tools (web search, python, terminal, code execution, image generation) render with their own components and so never showed the buttons. Those calls paused after tool_start with no way to approve them, hanging until the 1 hour timeout. Only MCP tools, which use the fallback renderer, actually worked. Render the controls for every tool card by wrapping each registered tool component (and the fallback) in thread.tsx with a shared ToolConfirmationControls, so the gate applies uniformly. Also make the handshake robust: - The gate keys on a per-call approval_id minted by the backend and echoed in tool_start, instead of session_id alone, so a stale or concurrent confirmation can no longer resolve the wrong call. - The approval slot is registered before tool_start is yielded, closing the race where a fast click or an auto "Always allow" could reach the backend before the waiter existed. - The frontend resolves with the same session id the request was sent with (plus the approval_id), fixing the new-thread mismatch where the confirmation targeted a different session than the blocked stream. - The confirm endpoint returns {resolved}; the UI keeps the buttons and shows a retry hint until the backend confirms a match, instead of hiding them on a failed or mistargeted post. - The gate runs after the disabled-tool and duplicate-call checks, so a call that will not execute is not put up for approval. A denied call is still excluded from duplicate detection, so re-issuing and approving it works. - "Always allow" is scoped per session to match the backend gate. Add backend tests for the approval registry, the SSE no-deadlock handshake, and the loop integration (allow, deny, disabled, duplicate, re-issue after deny). * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * Move "Confirm tool calls" to the Tools section * Studio: Keep tool group open while a tool call awaits confirmation * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * Fix tool confirmation session scope for PR #5869 * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * Fix confirmation follow-ups for PR #5869 * Apply pre-commit formatting for PR #5869 * Fix confirmation cleanup for PR #5869 * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * Harden confirmation lookups for PR #5869 * Studio: make the tool-call confirmation decision immutable resolve_tool_decision accepted a second confirmation for the same approval_id and overwrote slot["decision"] in the window before the waiter reads it and pops the slot, so a duplicate or out-of-order POST could flip an Allow to Deny (and returned a misleading resolved:true). Reject once the slot's event is already set so the first decision wins. Adds a regression test. * Fix/adjust tool confirmations for PR #5869 * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci --------- Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com> Co-authored-by: Daniel Han Co-authored-by: wasimysaid --- studio/backend/core/inference/llama_cpp.py | 71 ++++- studio/backend/core/inference/orchestrator.py | 2 + .../core/inference/safetensors_agentic.py | 53 +++- studio/backend/models/inference.py | 10 + studio/backend/routes/inference.py | 78 +++++- studio/backend/state/tool_approvals.py | 139 ++++++++++ .../backend/tests/test_anthropic_messages.py | 13 + .../backend/tests/test_llama_cpp_tool_loop.py | 171 ++++++++++++ .../tests/test_openai_tool_passthrough.py | 101 +++++++ .../tests/test_safetensors_tool_loop.py | 49 +++- studio/backend/tests/test_tool_approvals.py | 261 ++++++++++++++++++ .../backend/tests/test_tool_confirm_loop.py | 170 ++++++++++++ .../backend/tests/test_tool_confirm_stream.py | 219 +++++++++++++++ .../src/components/assistant-ui/thread.tsx | 30 +- .../tool-confirmation-controls.tsx | 157 +++++++++++ .../components/assistant-ui/tool-fallback.tsx | 3 + .../components/assistant-ui/tool-group.tsx | 28 +- .../src/features/chat/api/chat-adapter.ts | 46 ++- .../src/features/chat/api/chat-api.ts | 25 ++ .../src/features/chat/chat-settings-sheet.tsx | 25 ++ .../chat/stores/chat-runtime-store.ts | 70 +++++ .../frontend/src/features/chat/types/api.ts | 2 + 22 files changed, 1694 insertions(+), 29 deletions(-) create mode 100644 studio/backend/state/tool_approvals.py create mode 100644 studio/backend/tests/test_tool_approvals.py create mode 100644 studio/backend/tests/test_tool_confirm_loop.py create mode 100644 studio/backend/tests/test_tool_confirm_stream.py create mode 100644 studio/frontend/src/components/assistant-ui/tool-confirmation-controls.tsx diff --git a/studio/backend/core/inference/llama_cpp.py b/studio/backend/core/inference/llama_cpp.py index 80b4862aa8..8cf37ed9ec 100644 --- a/studio/backend/core/inference/llama_cpp.py +++ b/studio/backend/core/inference/llama_cpp.py @@ -60,6 +60,13 @@ from core.inference.tool_loop_controller import ( ToolLoopController, tool_event_provenance, ) +from state.tool_approvals import ( + TOOL_REJECTED_MESSAGE, + abort_tool_decision, + begin_tool_decision, + new_approval_id, + wait_tool_decision, +) logger = get_logger(__name__) @@ -2192,8 +2199,9 @@ class LlamaCppBackend: else None ), ( - f"{general['general.organization']}/" - f"{general['general.basename']}".replace(" ", "-") + f"{general['general.organization']}/{general['general.basename']}".replace( + " ", "-" + ) if general.get("general.organization") and general.get("general.basename") else None ), @@ -3748,7 +3756,7 @@ class LlamaCppBackend: ) logger.info( - f"llama-server ready on port {self._port} " f"for model '{model_identifier}'" + f"llama-server ready on port {self._port} for model '{model_identifier}'" ) # Probe outside _lock (interruptible by /unload); init inside. @@ -4373,7 +4381,7 @@ class LlamaCppBackend: proc.kill() logger.info( - f"Killed orphaned llama-server process " f"(pid={proc.info['pid']})" + f"Killed orphaned llama-server process (pid={proc.info['pid']})" ) except ( psutil.NoSuchProcess, @@ -4910,6 +4918,7 @@ class LlamaCppBackend: rag_scope: Optional[dict] = None, seed: Optional[int] = None, disable_parallel_tool_use: bool = False, + confirm_tool_calls: bool = False, ) -> Generator[dict, None, None]: """ Agentic loop: let the model call tools, execute them, and continue. @@ -4928,7 +4937,7 @@ class LlamaCppBackend: # Forced first-pass RAG so a doc question doesn't lose to web_search. Emits # the same tool card + citations a real call would. - _auto = build_rag_autoinject(conversation, rag_scope) + _auto = None if confirm_tool_calls else build_rag_autoinject(conversation, rag_scope) if _auto: for _ev in _auto["events"]: yield _ev @@ -5077,7 +5086,7 @@ class LlamaCppBackend: if response.status_code != 200: error_body = response.read().decode() raise RuntimeError( - f"llama-server returned {response.status_code}: " f"{error_body}" + f"llama-server returned {response.status_code}: {error_body}" ) raw_buf = "" @@ -5488,8 +5497,7 @@ class LlamaCppBackend: force = True, ) logger.info( - f"Safety net: parsed {len(tool_calls)} tool call(s) " - f"from streamed content" + f"Safety net: parsed {len(tool_calls)} tool call(s) from streamed content" ) else: # ── DRAINING path: assemble tool_calls ── @@ -5609,8 +5617,51 @@ class LlamaCppBackend: decision.as_assistant_tool_call() ) - yield {"type": "status", "text": decision.status_text} - yield decision.tool_start_event() + needs_confirm = bool(confirm_tool_calls) + approval_id = new_approval_id() if needs_confirm else "" + decision_slot = ( + begin_tool_decision(session_id, approval_id) if needs_confirm else None + ) + start_event = decision.tool_start_event() + start_event["approval_id"] = approval_id + start_event["awaiting_confirmation"] = needs_confirm + + try: + yield {"type": "status", "text": decision.status_text} + yield start_event + + if ( + decision_slot is not None + and wait_tool_decision( + decision_slot, + approval_id, + cancel_event = cancel_event, + ) + == "deny" + ): + decision_slot = None + yield { + "type": "tool_end", + "tool_name": decision.tool_name, + "tool_call_id": decision.tool_call_id, + "result": TOOL_REJECTED_MESSAGE, + "provenance": decision.provenance, + } + denied_message = { + "role": "tool", + "name": decision.tool_name, + "content": TOOL_REJECTED_MESSAGE, + } + if decision.tool_call_id: + denied_message["tool_call_id"] = decision.tool_call_id + conversation.append(denied_message) + if _forced_tool_call_pending: + _forced_tool_call_pending = False + continue + decision_slot = None + finally: + if decision_slot is not None: + abort_tool_decision(decision_slot, approval_id) _effective_timeout = None if tool_call_timeout >= 9999 else tool_call_timeout # RAG: cap paraphrased KB re-searches that slip past the dup guard. diff --git a/studio/backend/core/inference/orchestrator.py b/studio/backend/core/inference/orchestrator.py index 4e4788ace9..4ccac2912e 100644 --- a/studio/backend/core/inference/orchestrator.py +++ b/studio/backend/core/inference/orchestrator.py @@ -861,6 +861,7 @@ class InferenceOrchestrator: tool_call_timeout: int = 300, session_id: Optional[str] = None, rag_scope: Optional[dict] = None, + confirm_tool_calls: bool = False, use_adapter: Optional[Union[bool, str]] = None, stats_holder: Optional[dict] = None, **_unused, @@ -922,6 +923,7 @@ class InferenceOrchestrator: tool_call_timeout = tool_call_timeout, session_id = session_id, rag_scope = rag_scope, + confirm_tool_calls = confirm_tool_calls, ) def generate_with_adapter_control( diff --git a/studio/backend/core/inference/safetensors_agentic.py b/studio/backend/core/inference/safetensors_agentic.py index 7942edb6d7..3b6a393f3d 100644 --- a/studio/backend/core/inference/safetensors_agentic.py +++ b/studio/backend/core/inference/safetensors_agentic.py @@ -35,6 +35,13 @@ from core.inference.tool_loop_controller import ( status_for_tool, tool_event_provenance, ) +from state.tool_approvals import ( + TOOL_REJECTED_MESSAGE, + abort_tool_decision, + begin_tool_decision, + new_approval_id, + wait_tool_decision, +) logger = get_logger(__name__) @@ -146,6 +153,7 @@ def run_safetensors_tool_loop( tool_call_timeout: int = 300, session_id: Optional[str] = None, rag_scope: Optional[dict] = None, + confirm_tool_calls: bool = False, ) -> Generator[dict, None, None]: """Drive an agentic tool loop on top of a cumulative-text generator. @@ -174,7 +182,7 @@ def run_safetensors_tool_loop( # Forced first-pass RAG (mirrors the GGUF loop) so doc Qs don't lose to web_search. from core.inference.tools import build_rag_autoinject - _auto = build_rag_autoinject(conversation, rag_scope) + _auto = None if confirm_tool_calls else build_rag_autoinject(conversation, rag_scope) if _auto: for _ev in _auto["events"]: yield _ev @@ -509,8 +517,47 @@ def run_safetensors_tool_loop( else: assistant_msg.setdefault("tool_calls", []).append(decision.as_assistant_tool_call()) - yield {"type": "status", "text": decision.status_text} - yield decision.tool_start_event() + needs_confirm = bool(confirm_tool_calls) + approval_id = new_approval_id() if needs_confirm else "" + decision_slot = begin_tool_decision(session_id, approval_id) if needs_confirm else None + start_event = decision.tool_start_event() + start_event["approval_id"] = approval_id + start_event["awaiting_confirmation"] = needs_confirm + + try: + yield {"type": "status", "text": decision.status_text} + yield start_event + + if ( + decision_slot is not None + and wait_tool_decision( + decision_slot, + approval_id, + cancel_event = cancel_event, + ) + == "deny" + ): + decision_slot = None + yield { + "type": "tool_end", + "tool_name": decision.tool_name, + "tool_call_id": decision.tool_call_id, + "result": TOOL_REJECTED_MESSAGE, + "provenance": decision.provenance, + } + denied_message = { + "role": "tool", + "name": decision.tool_name, + "content": TOOL_REJECTED_MESSAGE, + } + if decision.tool_call_id: + denied_message["tool_call_id"] = decision.tool_call_id + conversation.append(denied_message) + continue + decision_slot = None + finally: + if decision_slot is not None: + abort_tool_decision(decision_slot, approval_id) eff_timeout = None if tool_call_timeout >= 9999 else tool_call_timeout # RAG: cap paraphrased KB re-searches that slip past the dup guard. diff --git a/studio/backend/models/inference.py b/studio/backend/models/inference.py index b70202d6ff..039a5d75dd 100644 --- a/studio/backend/models/inference.py +++ b/studio/backend/models/inference.py @@ -690,6 +690,10 @@ class ChatCompletionRequest(BaseModel): None, description = "[x-unsloth] When true, append tools from every enabled MCP server to this request's tool list.", ) + confirm_tool_calls: Optional[bool] = Field( + None, + description = "[x-unsloth] When true, pause before each tool call and wait for the user to allow/deny it via POST /api/inference/tool-confirm.", + ) auto_heal_tool_calls: Optional[bool] = Field( True, description = "[x-unsloth] Auto-detect and fix malformed tool calls from model output.", @@ -926,6 +930,12 @@ class ChatCompletionRequest(BaseModel): return self +class ToolConfirmRequest(BaseModel): + session_id: Optional[str] = None + approval_id: Optional[str] = None + decision: Literal["allow", "deny"] = "deny" + + # ── OpenAI shell-tool container management ───────────────────── diff --git a/studio/backend/routes/inference.py b/studio/backend/routes/inference.py index 856d9bacf8..9ff30cd6db 100644 --- a/studio/backend/routes/inference.py +++ b/studio/backend/routes/inference.py @@ -655,6 +655,7 @@ from models.inference import ( ChatCompletionRequest, ChatCompletionChunk, ChatCompletion, + ToolConfirmRequest, ChatMessage, ChunkChoice, ChoiceDelta, @@ -702,6 +703,7 @@ from core.inference.anthropic_compat import ( AnthropicPassthroughEmitter, ) from auth.authentication import get_current_subject +from state.tool_approvals import resolve_tool_decision from core.inference.key_exchange import decrypt_api_key from core.inference.providers import get_provider_info, get_base_url @@ -1495,8 +1497,7 @@ async def load_model( # Shouldn't happen on already-validated args; degrade to # no-extras rather than 400 if managed flags changed. logger.warning( - "Stored llama_extra_args failed revalidation; " - "loading without them: %s", + "Stored llama_extra_args failed revalidation; loading without them: %s", stripped, ) extra_llama_args = [] @@ -1938,6 +1939,20 @@ async def cancel_inference(request: Request, current_subject: str = Depends(get_ return {"cancelled": n} +@studio_router.post("/tool-confirm") +async def confirm_tool_call( + request: ToolConfirmRequest, current_subject: str = Depends(get_current_subject) +): + matched = resolve_tool_decision( + request.approval_id, + request.decision, + session_id = request.session_id, + ) + if not matched: + raise HTTPException(status_code = 404, detail = "No pending tool call confirmation") + return {"resolved": True} + + @router.post("/generate/stream") async def generate_stream( request: GenerateRequest, current_subject: str = Depends(get_current_subject) @@ -3196,6 +3211,22 @@ async def openai_chat_completions( # ── External provider routing ──────────────────────────────── # encrypted_api_key is optional -- local providers (llama.cpp / vLLM / Ollama) may run without auth. if payload.provider_id or payload.provider_type: + if payload.confirm_tool_calls and ( + payload.enable_tools is True + or bool(payload.enabled_tools) + or bool(payload.tools) + or bool(payload.openai_code_exec_container_id) + or bool(payload.anthropic_code_exec_container_id) + ): + raise HTTPException( + status_code = 400, + detail = openai_error_body( + "confirm_tool_calls is only supported for local streaming tools.", + status = 400, + code = "invalid_request_error", + param = "confirm_tool_calls", + ), + ) if _wants_multiple_choices(payload): _raise_unsupported_n("external provider chat completions") return await _proxy_to_external_provider(payload, request) @@ -3567,6 +3598,16 @@ async def openai_chat_completions( use_tools = False if use_tools: + if payload.confirm_tool_calls and not payload.stream: + raise HTTPException( + status_code = 400, + detail = openai_error_body( + "confirm_tool_calls requires stream=true for local tool execution.", + status = 400, + code = "invalid_request_error", + param = "confirm_tool_calls", + ), + ) if _wants_multiple_choices(payload): _raise_unsupported_n("GGUF tool chat completions") # ── Tool-use system prompt nudge ────────────────────── @@ -3637,6 +3678,7 @@ async def openai_chat_completions( session_id = payload.session_id, rag_scope = payload.rag_scope, disable_parallel_tool_use = payload.parallel_tool_calls is False, + confirm_tool_calls = bool(payload.confirm_tool_calls), ) _tool_sentinel = object() @@ -3646,6 +3688,7 @@ async def openai_chat_completions( _tracker.__enter__() async def gguf_tool_stream(): + gen = None try: first_chunk = ChatCompletionChunk( id = completion_id, @@ -3768,6 +3811,11 @@ async def openai_chat_completions( error_chunk = _openai_stream_error_chunk(e) yield f"data: {json.dumps(error_chunk)}\n\n" finally: + if gen is not None: + try: + gen.close() + except (RuntimeError, ValueError): + pass _tracker.__exit__(None, None, None) return StreamingResponse( @@ -4091,6 +4139,16 @@ async def openai_chat_completions( _sf_use_tools = False if _sf_use_tools: + if payload.confirm_tool_calls and not payload.stream: + raise HTTPException( + status_code = 400, + detail = openai_error_body( + "confirm_tool_calls requires stream=true for local tool execution.", + status = 400, + code = "invalid_request_error", + param = "confirm_tool_calls", + ), + ) _sf_nudge = _build_tool_action_nudge( tools = _sf_tools_to_use, model_name = model_name, @@ -4167,6 +4225,7 @@ async def openai_chat_completions( else 300, session_id = payload.session_id, rag_scope = payload.rag_scope, + confirm_tool_calls = bool(payload.confirm_tool_calls), use_adapter = payload.use_adapter, stats_holder = _sf_stats_holder, ) @@ -4177,6 +4236,7 @@ async def openai_chat_completions( _sf_tracker.__enter__() async def sf_tool_stream(): + gen = None try: first_chunk = ChatCompletionChunk( id = completion_id, @@ -4293,6 +4353,11 @@ async def openai_chat_completions( } yield f"data: {json.dumps(error_chunk)}\n\n" finally: + if gen is not None: + try: + gen.close() + except (RuntimeError, ValueError): + pass _sf_tracker.__exit__(None, None, None) if payload.stream: @@ -5934,6 +5999,15 @@ async def anthropic_messages( ) if server_tools: + if bool(getattr(payload, "confirm_tool_calls", False)): + raise HTTPException( + status_code = 400, + detail = anthropic_error_body( + "confirm_tool_calls is not supported for Anthropic Messages server tools.", + status = 400, + err_type = "invalid_request_error", + ), + ) from core.inference.tools import ALL_TOOLS openai_tools = _select_anthropic_server_tools( diff --git a/studio/backend/state/tool_approvals.py b/studio/backend/state/tool_approvals.py new file mode 100644 index 0000000000..f66226b61d --- /dev/null +++ b/studio/backend/state/tool_approvals.py @@ -0,0 +1,139 @@ +# SPDX-License-Identifier: AGPL-3.0-only +# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. + +"""Per-call tool-call confirmation gate. + +When a chat request sets ``confirm_tool_calls``, the agentic loop pauses +before executing each tool and waits here for the user's decision, which +arrives via ``POST /api/inference/tool-confirm`` on a separate connection. + +Each gated call is identified by a unique ``approval_id`` (minted with +``new_approval_id``) that the loop both registers here and echoes in the +``tool_start`` stream event. The frontend sends that exact id back, so a +stale or duplicate confirmation -- or a second tool awaiting a decision in +the same session -- can never resolve the wrong call. ``session_id`` is +kept alongside purely as a scope check. + +The slot is registered with ``begin_tool_decision`` *before* the loop +yields ``tool_start``, closing the race where a fast confirmation (or an +auto "Always allow") could otherwise arrive before the waiter exists. +``wait_tool_decision`` then blocks and cleans up its own slot. +""" + +import secrets +import threading +from typing import Optional + +# Generous ceiling so a user can deliberate; cancellation (stop button / +# disconnect) still breaks the wait early via ``cancel_event``. +_DECISION_TIMEOUT = 3600.0 + +# Fed to the model as the tool result when the user denies a call, so it +# can adapt and keep responding instead of the turn ending abruptly. +TOOL_REJECTED_MESSAGE = "The user declined to run this tool call." + +_lock = threading.Lock() +# approval_id -> {"event": threading.Event, "decision": str|None, "session": str} +_pending: dict[str, dict] = {} + + +def new_approval_id() -> str: + """Mint an unguessable id for one pending tool-call confirmation.""" + return secrets.token_urlsafe(16) + + +def begin_tool_decision(session_id, approval_id) -> dict: + """Register a pending decision slot and return it. + + Call this *before* yielding the ``tool_start`` event so the waiter + always exists by the time the user's confirmation can arrive. + """ + slot = { + "event": threading.Event(), + "decision": None, + "session": session_id or "", + } + with _lock: + _pending[approval_id] = slot + return slot + + +def wait_tool_decision( + slot, + approval_id, + cancel_event = None, + timeout = _DECISION_TIMEOUT, +): + """Block on a slot from ``begin_tool_decision`` until the user decides. + + Returns ``"allow"`` or ``"deny"``. Falls back to ``"deny"`` if the wait + times out or generation is cancelled before the user decides. Always + removes its own slot on exit. + """ + try: + waited = 0.0 + while not slot["event"].wait(timeout = 0.5): + if cancel_event is not None and cancel_event.is_set(): + return "deny" + waited += 0.5 + if waited >= timeout: + return "deny" + return slot["decision"] or "deny" + finally: + with _lock: + if _pending.get(approval_id) is slot: + _pending.pop(approval_id, None) + + +def abort_tool_decision(slot, approval_id) -> None: + """Remove a slot that was announced but never entered ``wait_tool_decision``. + + Streaming wrappers may stop after ``tool_start`` is yielded and before + the loop resumes into ``wait_tool_decision``. In that case there is no + waiter to run the normal cleanup path, so the generator close path calls + this explicitly. + """ + with _lock: + if _pending.get(approval_id) is slot: + _pending.pop(approval_id, None) + + +def request_tool_decision( + session_id, + approval_id, + cancel_event = None, + timeout = _DECISION_TIMEOUT, +): + """Register and wait in one call (when the slot is not needed early).""" + slot = begin_tool_decision(session_id, approval_id) + return wait_tool_decision(slot, approval_id, cancel_event = cancel_event, timeout = timeout) + + +def resolve_tool_decision( + approval_id, + decision, + session_id = None, +) -> bool: + """Record the user's "allow"/"deny" decision and unblock the loop. + + Returns ``True`` if a pending call matched, ``False`` otherwise (e.g. a + stale or duplicate confirmation, or a session-scope mismatch). + + The first decision wins: once a slot's event is set, a later (duplicate or + out-of-order) confirmation for the same id is rejected without mutating the + recorded decision, so an Allow can never be flipped to Deny in the window + before the waiter reads ``slot["decision"]`` and pops the slot. + """ + if not approval_id: + return False + with _lock: + slot = _pending.get(approval_id) + if not slot: + return False + if session_id is not None and slot["session"] != (session_id or ""): + return False + if slot["event"].is_set(): + return False + slot["decision"] = decision + slot["event"].set() + return True diff --git a/studio/backend/tests/test_anthropic_messages.py b/studio/backend/tests/test_anthropic_messages.py index e1230ae113..f8d3f44d4b 100644 --- a/studio/backend/tests/test_anthropic_messages.py +++ b/studio/backend/tests/test_anthropic_messages.py @@ -1543,6 +1543,19 @@ class TestAnthropicMessagesToolRouting: _drive(anthropic_messages(payload, request = None, current_subject = "t")) assert backend.calls[0][0] == "tools" + def test_confirm_tool_calls_rejected_for_server_tools(self, monkeypatch): + backend = _mock_backend(monkeypatch) + payload = _basic_payload( + confirm_tool_calls = True, + tools = [{"type": "web_search_20250305", "name": "web_search"}], + ) + + with pytest.raises(HTTPException) as exc: + _drive(anthropic_messages(payload, request = None, current_subject = "t")) + assert exc.value.status_code == 400 + assert "confirm_tool_calls is not supported" in exc.value.detail["error"]["message"] + assert backend.calls == [] + def test_per_request_enable_tools_false_blocks_server_tool_alias(self, monkeypatch): backend = _mock_backend(monkeypatch) payload = _basic_payload( diff --git a/studio/backend/tests/test_llama_cpp_tool_loop.py b/studio/backend/tests/test_llama_cpp_tool_loop.py index fa583ef53d..3c121c281d 100644 --- a/studio/backend/tests/test_llama_cpp_tool_loop.py +++ b/studio/backend/tests/test_llama_cpp_tool_loop.py @@ -21,6 +21,8 @@ if _BACKEND_DIR not in sys.path: sys.path.insert(0, _BACKEND_DIR) from core.inference.llama_cpp import LlamaCppBackend +from state import tool_approvals +from state.tool_approvals import TOOL_REJECTED_MESSAGE, resolve_tool_decision def _sse(delta: dict) -> str: @@ -70,6 +72,27 @@ def _tool_names(payload: dict) -> list[str]: ] +def _structured_tool_call(tool_name: str, arguments: dict, call_id: str) -> list[str]: + return [ + _sse( + { + "tool_calls": [ + { + "index": 0, + "id": call_id, + "type": "function", + "function": { + "name": tool_name, + "arguments": json.dumps(arguments), + }, + } + ] + } + ), + _done(), + ] + + def test_structured_tool_call_after_visible_preface_is_executed(monkeypatch): """llama-server may emit content first and then native delta.tool_calls. @@ -1149,3 +1172,151 @@ def test_reprompted_tool_call_still_streams_final_answer(monkeypatch): content_texts = [event.get("text", "") for event in events if event.get("type") == "content"] assert content_texts == ["I will use render_html now.", "Final note after tool."] assert len(payloads) == 3 + + +def test_confirm_tool_calls_allow_executes_gguf_tool(monkeypatch): + streams = [ + _structured_tool_call("python", {"code": "print(1)"}, "call_py"), + [_sse({"content": "Done."}), _done()], + ] + payloads: list[dict] = [] + backend = _make_backend(monkeypatch, streams, payloads) + + calls: list[tuple[str, dict]] = [] + + def fake_execute_tool(name, arguments, **_kwargs): + calls.append((name, arguments)) + return "OK" + + monkeypatch.setattr("core.inference.tools.execute_tool", fake_execute_tool) + monkeypatch.setattr("core.inference.llama_cpp.new_approval_id", lambda: "approval-1") + monkeypatch.setattr( + "core.inference.llama_cpp.begin_tool_decision", + lambda *_a, **_k: object(), + ) + monkeypatch.setattr("core.inference.llama_cpp.wait_tool_decision", lambda *_a, **_k: "allow") + + events = list( + backend.generate_chat_completion_with_tools( + messages = [{"role": "user", "content": "run python"}], + tools = [{"type": "function", "function": {"name": "python"}}], + max_tool_iterations = 1, + confirm_tool_calls = True, + session_id = "sess", + ) + ) + + starts = [event for event in events if event.get("type") == "tool_start"] + assert len(starts) == 1 + assert starts[0]["approval_id"] + assert starts[0]["awaiting_confirmation"] is True + assert calls == [("python", {"code": "print(1)"})] + assert any(event.get("type") == "tool_end" and event.get("result") == "OK" for event in events) + + +def test_confirm_tool_calls_close_after_prompt_cleans_gguf_slot(monkeypatch): + approval_id = "approval-close" + streams = [_structured_tool_call("python", {"code": "print(1)"}, "call_py")] + payloads: list[dict] = [] + backend = _make_backend(monkeypatch, streams, payloads) + + monkeypatch.setattr( + "core.inference.tools.execute_tool", + lambda *_a, **_k: (_ for _ in ()).throw(AssertionError("tool should not run")), + ) + monkeypatch.setattr("core.inference.llama_cpp.new_approval_id", lambda: approval_id) + + with tool_approvals._lock: + tool_approvals._pending.clear() + + gen = backend.generate_chat_completion_with_tools( + messages = [{"role": "user", "content": "run python"}], + tools = [{"type": "function", "function": {"name": "python"}}], + max_tool_iterations = 1, + confirm_tool_calls = True, + session_id = "sess", + ) + try: + assert next(gen)["type"] == "status" + start = next(gen) + assert start["type"] == "tool_start" + assert start["approval_id"] == approval_id + with tool_approvals._lock: + assert approval_id in tool_approvals._pending + finally: + gen.close() + + with tool_approvals._lock: + assert approval_id not in tool_approvals._pending + assert resolve_tool_decision(approval_id, "allow", session_id = "sess") is False + + +def test_confirm_tool_calls_skips_gguf_rag_autoinject(monkeypatch): + streams = [[_sse({"content": "Done."}), _done()]] + payloads: list[dict] = [] + backend = _make_backend(monkeypatch, streams, payloads) + + def fail_autoinject(*_args, **_kwargs): + raise AssertionError("RAG autoinject must not run before approval") + + monkeypatch.setattr("core.inference.tools.build_rag_autoinject", fail_autoinject) + + events = list( + backend.generate_chat_completion_with_tools( + messages = [{"role": "user", "content": "use docs"}], + tools = [{"type": "function", "function": {"name": "search_knowledge_base"}}], + max_tool_iterations = 1, + confirm_tool_calls = True, + session_id = "sess", + rag_scope = {"thread_id": "t1"}, + ) + ) + + assert any(event.get("type") == "content" and event.get("text") == "Done." for event in events) + + +def test_confirm_tool_calls_deny_skips_gguf_tool_and_retry_can_execute(monkeypatch): + same_call = _structured_tool_call("python", {"code": "print(1)"}, "call_py") + streams = [ + same_call, + _structured_tool_call("python", {"code": "print(1)"}, "call_py_retry"), + [_sse({"content": "Done."}), _done()], + ] + payloads: list[dict] = [] + backend = _make_backend(monkeypatch, streams, payloads) + + calls: list[tuple[str, dict]] = [] + + def fake_execute_tool(name, arguments, **_kwargs): + calls.append((name, arguments)) + return "OK" + + decisions = iter(["deny", "allow"]) + approvals = iter(["approval-1", "approval-2"]) + + monkeypatch.setattr("core.inference.tools.execute_tool", fake_execute_tool) + monkeypatch.setattr("core.inference.llama_cpp.new_approval_id", lambda: next(approvals)) + monkeypatch.setattr( + "core.inference.llama_cpp.begin_tool_decision", + lambda *_a, **_k: object(), + ) + monkeypatch.setattr( + "core.inference.llama_cpp.wait_tool_decision", + lambda *_a, **_k: next(decisions), + ) + + events = list( + backend.generate_chat_completion_with_tools( + messages = [{"role": "user", "content": "run python"}], + tools = [{"type": "function", "function": {"name": "python"}}], + max_tool_iterations = 2, + confirm_tool_calls = True, + session_id = "sess", + ) + ) + + starts = [event for event in events if event.get("type") == "tool_start"] + ends = [event for event in events if event.get("type") == "tool_end"] + assert len(starts) == 2 + assert [event["result"] for event in ends] == [TOOL_REJECTED_MESSAGE, "OK"] + assert calls == [("python", {"code": "print(1)"})] diff --git a/studio/backend/tests/test_openai_tool_passthrough.py b/studio/backend/tests/test_openai_tool_passthrough.py index 1d994b46c0..79e8c9ae49 100644 --- a/studio/backend/tests/test_openai_tool_passthrough.py +++ b/studio/backend/tests/test_openai_tool_passthrough.py @@ -401,6 +401,28 @@ class TestChatCompletionRequestToolFields: ) self._assert_unsupported_n(resp) + def test_confirm_tool_calls_rejected_for_provider_tools(self, monkeypatch): + class _UnusedBackend: + is_loaded = False + + client = self._v1_client(monkeypatch, _UnusedBackend()) + resp = client.post( + "/v1/chat/completions", + json = { + "messages": [{"role": "user", "content": "hi"}], + "provider_type": "openai", + "external_model": "gpt-4.1", + "enable_tools": True, + "enabled_tools": ["web_search"], + "confirm_tool_calls": True, + }, + ) + + assert resp.status_code == 400 + body = resp.json() + assert body["error"]["param"] == "confirm_tool_calls" + assert "only supported for local streaming tools" in body["error"]["message"] + def test_logprobs_rejected_until_supported(self, monkeypatch): class _UnusedBackend: is_loaded = False @@ -480,6 +502,7 @@ class TestChatCompletionRequestToolFields: def test_n_rejected_for_non_gguf_path(self, monkeypatch): class _NoGGUFBackend: is_loaded = False + supports_tools = False class _InferenceBackend: active_model_name = "test-model" @@ -495,6 +518,45 @@ class TestChatCompletionRequestToolFields: ) self._assert_unsupported_n(resp) + def test_confirm_tool_calls_requires_streaming_for_safetensors_tools(self, monkeypatch): + import routes.inference as inference_route + + class _NoGGUFBackend: + is_loaded = False + supports_tools = False + + class _InferenceBackend: + active_model_name = "test-model" + models = {"test-model": {"chat_template_info": {"template": "chatml"}}} + + def generate_chat_completion_with_tools(self, **kwargs): + raise AssertionError("tool loop should be rejected before starting") + + def generate_chat_completion(self, **kwargs): + raise AssertionError("plain path should not be used") + + monkeypatch.setattr( + inference_route, + "_detect_safetensors_features", + lambda backend, chat_template: {"supports_tools": True}, + ) + client = self._v1_client(monkeypatch, _NoGGUFBackend(), _InferenceBackend()) + resp = client.post( + "/v1/chat/completions", + json = { + "messages": [{"role": "user", "content": "hi"}], + "enable_tools": True, + "enabled_tools": ["web_search"], + "confirm_tool_calls": True, + "stream": False, + }, + ) + + assert resp.status_code == 400 + body = resp.json() + assert body["error"]["param"] == "confirm_tool_calls" + assert "requires stream=true" in body["error"]["message"] + def test_multiturn_tool_loop_messages(self): req = ChatCompletionRequest( messages = [ @@ -1206,6 +1268,45 @@ class TestGgufVisionToolRouting: assert captured["kwargs"]["disable_parallel_tool_use"] is True + def test_confirm_tool_calls_requires_streaming_for_gguf_tools(self, monkeypatch): + import routes.inference as inf_mod + + def _plain(**kwargs): + raise AssertionError("plain GGUF path should not be used") + + def _tools(**kwargs): + raise AssertionError("tool loop should be rejected before starting") + + backend = SimpleNamespace( + is_loaded = True, + is_vision = False, + supports_tools = True, + model_identifier = "test-gguf", + generate_chat_completion = _plain, + generate_chat_completion_with_tools = _tools, + ) + monkeypatch.setattr(inf_mod, "get_llama_cpp_backend", lambda: backend) + + payload = ChatCompletionRequest( + model = "default", + enable_tools = True, + enabled_tools = ["web_search"], + confirm_tool_calls = True, + stream = False, + messages = [{"role": "user", "content": "search once"}], + ) + + with pytest.raises(HTTPException) as exc: + self._drive( + openai_chat_completions( + payload, + request = self._Request(), + current_subject = "test", + ) + ) + assert exc.value.status_code == 400 + assert "requires stream=true" in exc.value.detail["error"]["message"] + def test_standard_gguf_merges_system_and_developer_messages(self, monkeypatch): import routes.inference as inf_mod diff --git a/studio/backend/tests/test_safetensors_tool_loop.py b/studio/backend/tests/test_safetensors_tool_loop.py index 8aa6e5df4e..12731783a0 100644 --- a/studio/backend/tests/test_safetensors_tool_loop.py +++ b/studio/backend/tests/test_safetensors_tool_loop.py @@ -28,6 +28,8 @@ from core.inference.tool_call_parser import ( parse_tool_calls_from_text, strip_tool_markup, ) +from state import tool_approvals +from state.tool_approvals import resolve_tool_decision from utils.datasets import is_gpt_oss_model_name @@ -84,8 +86,7 @@ class TestParser: # A code parameter with a literal must not truncate: the # parser uses end-of-body as the only boundary for single-param calls. text = ( - "html = ''\n" - "print('hi')" + "html = ''\nprint('hi')" ) result = parse_tool_calls_from_text(text) assert len(result) == 1 @@ -1033,6 +1034,50 @@ class TestGuardrails: _collect_events(loop) assert exec_fn.calls == [("web_search", {"query": "x"})] + def test_confirm_tool_calls_close_after_prompt_cleans_slot(self, monkeypatch): + approval_id = "approval-close-sf" + monkeypatch.setattr(safetensors_agentic, "new_approval_id", lambda: approval_id) + + loop, exec_fn = _make_loop( + turns = [['{"name":"python","arguments":{"code":"print(1)"}}']], + exec_results = ["OK"], + confirm_tool_calls = True, + session_id = "sess", + max_tool_iterations = 1, + ) + + with tool_approvals._lock: + tool_approvals._pending.clear() + + try: + assert next(loop)["type"] == "status" + start = next(loop) + assert start["type"] == "tool_start" + assert start["approval_id"] == approval_id + with tool_approvals._lock: + assert approval_id in tool_approvals._pending + finally: + loop.close() + + with tool_approvals._lock: + assert approval_id not in tool_approvals._pending + assert resolve_tool_decision(approval_id, "allow", session_id = "sess") is False + assert exec_fn.calls == [] + + def test_confirm_tool_calls_skips_rag_autoinject(self, monkeypatch): + def fail_autoinject(*_args, **_kwargs): + raise AssertionError("RAG autoinject must not run before approval") + + monkeypatch.setattr("core.inference.tools.build_rag_autoinject", fail_autoinject) + loop, exec_fn = _make_loop( + turns = [["plain answer"]], + confirm_tool_calls = True, + rag_scope = {"thread_id": "t1"}, + ) + events = _collect_events(loop) + assert any(e.get("type") == "content" and e.get("text") == "plain answer" for e in events) + assert exec_fn.calls == [] + def test_auto_heal_disabled_preserves_xml_on_final_no_tools_pass(self): turns = iter( [ diff --git a/studio/backend/tests/test_tool_approvals.py b/studio/backend/tests/test_tool_approvals.py new file mode 100644 index 0000000000..af792e652c --- /dev/null +++ b/studio/backend/tests/test_tool_approvals.py @@ -0,0 +1,261 @@ +# SPDX-License-Identifier: AGPL-3.0-only +# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0 + +"""Concurrency tests for the per-call tool-call confirmation gate. + +``state.tool_approvals`` coordinates two threads: the agentic loop thread +blocked in ``wait_tool_decision`` and the request thread that delivers the +user's choice through ``resolve_tool_decision``. Each gated call carries a +unique ``approval_id`` so a stale or concurrent confirmation can never +resolve the wrong call. These tests exercise that handshake directly -- +no model, no server -- so the race windows are fast and deterministic. +""" + +import threading +import time + +import pytest + +from state import tool_approvals +from state.tool_approvals import ( + TOOL_REJECTED_MESSAGE, + abort_tool_decision, + begin_tool_decision, + new_approval_id, + request_tool_decision, + resolve_tool_decision, + wait_tool_decision, +) + + +@pytest.fixture(autouse = True) +def _clear_pending(): + """Each test starts and ends with an empty ``_pending`` map.""" + with tool_approvals._lock: + tool_approvals._pending.clear() + yield + with tool_approvals._lock: + tool_approvals._pending.clear() + + +class _Waiter: + """Run ``request_tool_decision`` in a thread and capture its result.""" + + def __init__( + self, + session_id, + approval_id, + cancel_event = None, + timeout = None, + ): + self.session_id = session_id + self.approval_id = approval_id + self.cancel_event = cancel_event + self.timeout = timeout + self.result = None + self._thread = threading.Thread(target = self._run, daemon = True) + + def _run(self): + kwargs = {"cancel_event": self.cancel_event} + if self.timeout is not None: + kwargs["timeout"] = self.timeout + self.result = request_tool_decision(self.session_id, self.approval_id, **kwargs) + + def start(self): + self._thread.start() + _wait_until(lambda: _has_pending(self.approval_id)) + return self + + def join(self, timeout = 5.0): + self._thread.join(timeout = timeout) + assert not self._thread.is_alive(), "waiter thread did not finish" + return self.result + + +def _has_pending(approval_id) -> bool: + with tool_approvals._lock: + return approval_id in tool_approvals._pending + + +def _wait_until( + pred, + timeout = 2.0, + interval = 0.005, +) -> bool: + deadline = time.monotonic() + timeout + while time.monotonic() < deadline: + if pred(): + return True + time.sleep(interval) + return False + + +# ── Basic allow / deny ─────────────────────────────────────────────── + + +def test_allow_decision(): + aid = new_approval_id() + w = _Waiter("sess", aid).start() + assert resolve_tool_decision(aid, "allow", session_id = "sess") is True + assert w.join() == "allow" + + +def test_deny_decision(): + aid = new_approval_id() + w = _Waiter("sess", aid).start() + assert resolve_tool_decision(aid, "deny", session_id = "sess") is True + assert w.join() == "deny" + + +def test_slot_cleaned_up_after_decision(): + aid = new_approval_id() + w = _Waiter("sess", aid).start() + resolve_tool_decision(aid, "allow") + w.join() + assert _wait_until(lambda: not _has_pending(aid)) + + +def test_abort_tool_decision_removes_unwaited_slot(): + aid = new_approval_id() + slot = begin_tool_decision("sess", aid) + abort_tool_decision(slot, aid) + assert not _has_pending(aid) + assert resolve_tool_decision(aid, "allow", session_id = "sess") is False + + +def test_approval_ids_are_unique(): + ids = {new_approval_id() for _ in range(1000)} + assert len(ids) == 1000 + + +# ── Pre-registration race (begin before wait) ──────────────────────── + + +def test_resolve_before_wait_is_not_lost(): + """A decision delivered after ``begin`` but before ``wait`` survives. + + The loop registers the slot before it yields ``tool_start``, so even a + confirmation that races ahead of the blocking ``wait`` is recorded on + the slot and returned -- never dropped. + """ + aid = new_approval_id() + slot = begin_tool_decision("sess", aid) + assert resolve_tool_decision(aid, "allow", session_id = "sess") is True + # wait() is only entered now, after the decision already landed. + assert wait_tool_decision(slot, aid) == "allow" + assert not _has_pending(aid) + + +# ── Resolver edge cases ────────────────────────────────────────────── + + +def test_resolve_unknown_approval_returns_false(): + assert resolve_tool_decision(new_approval_id(), "allow") is False + + +def test_resolve_empty_approval_returns_false(): + assert resolve_tool_decision("", "allow") is False + assert resolve_tool_decision(None, "allow") is False + + +def test_resolve_wrong_session_scope_returns_false(): + aid = new_approval_id() + w = _Waiter("sess-a", aid).start() + # Correct approval_id but the wrong session must not resolve it. + assert resolve_tool_decision(aid, "allow", session_id = "sess-b") is False + assert _has_pending(aid) + # The right session still works. + assert resolve_tool_decision(aid, "allow", session_id = "sess-a") is True + assert w.join() == "allow" + + +def test_duplicate_resolve_after_completion_returns_false(): + aid = new_approval_id() + w = _Waiter("sess", aid).start() + assert resolve_tool_decision(aid, "allow") is True + w.join() + assert _wait_until(lambda: not _has_pending(aid)) + assert resolve_tool_decision(aid, "deny") is False + + +def test_first_decision_is_immutable(): + """A second confirmation cannot flip an already-recorded decision. + + The waiter reads ``slot["decision"]`` outside the lock and then cleans up, + so a duplicate or out-of-order POST that lands in that window must be + rejected and must not overwrite the first decision -- an Allow can never + become a Deny. Distinct from the after-completion case above: here the slot + is still pending (no waiter has consumed it yet). + """ + aid = new_approval_id() + slot = begin_tool_decision("sess", aid) + assert resolve_tool_decision(aid, "allow", session_id = "sess") is True + # Second decision, same id, before any waiter consumes/cleans the slot. + assert resolve_tool_decision(aid, "deny", session_id = "sess") is False + assert slot["decision"] == "allow" + # The waiter still observes the first (immutable) decision. + assert wait_tool_decision(slot, aid) == "allow" + assert not _has_pending(aid) + + +# ── Cancellation and timeout ───────────────────────────────────────── + + +def test_cancel_event_breaks_wait_as_deny(): + cancel = threading.Event() + aid = new_approval_id() + w = _Waiter("sess", aid, cancel_event = cancel).start() + cancel.set() + assert w.join(timeout = 3.0) == "deny" + assert _wait_until(lambda: not _has_pending(aid)) + + +def test_timeout_returns_deny(): + aid = new_approval_id() + start = time.monotonic() + result = request_tool_decision("sess", aid, timeout = 0.1) + assert result == "deny" + assert time.monotonic() - start < 2.0 + assert not _has_pending(aid) + + +# ── Independence across concurrent calls ───────────────────────────── + + +def test_two_pending_calls_same_session_are_independent(): + """Keying on approval_id, not session, keeps concurrent calls distinct. + + Resolving the first call's id must not unblock or alter the second + call pending in the same session. + """ + a1, a2 = new_approval_id(), new_approval_id() + w1 = _Waiter("sess", a1).start() + w2 = _Waiter("sess", a2).start() + + assert resolve_tool_decision(a1, "deny", session_id = "sess") is True + assert w1.join() == "deny" + # w2 is still waiting on its own id. + assert _has_pending(a2) + assert resolve_tool_decision(a2, "allow", session_id = "sess") is True + assert w2.join() == "allow" + + +def test_concurrent_distinct_calls_route_their_own_decisions(): + n = 25 + waiters = {} + for i in range(n): + aid = new_approval_id() + waiters[aid] = _Waiter(f"s{i}", aid).start() + expected = {aid: ("allow" if i % 2 == 0 else "deny") for i, aid in enumerate(waiters)} + for aid, decision in expected.items(): + assert resolve_tool_decision(aid, decision) is True + for aid, w in waiters.items(): + assert w.join() == expected[aid] + + +# ── Constants ──────────────────────────────────────────────────────── + + +def test_rejected_message_is_user_facing_text(): + assert isinstance(TOOL_REJECTED_MESSAGE, str) + assert TOOL_REJECTED_MESSAGE.strip() diff --git a/studio/backend/tests/test_tool_confirm_loop.py b/studio/backend/tests/test_tool_confirm_loop.py new file mode 100644 index 0000000000..ce7852c95f --- /dev/null +++ b/studio/backend/tests/test_tool_confirm_loop.py @@ -0,0 +1,170 @@ +# SPDX-License-Identifier: AGPL-3.0-only +# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0 + +"""Integration tests for the confirmation gate inside the real tool loop. + +These drive ``run_safetensors_tool_loop`` (no model -- hand-crafted fake +generators) with ``confirm_tool_calls=True`` and resolve each pending +decision inline. The slot is registered before ``tool_start`` is yielded, +so resolving right after receiving that event always lands before the +loop blocks. Covers: allow executes once, deny skips execution and feeds +back the rejection, disabled/duplicate calls are not prompted, and a +denied call does not pollute duplicate detection. +""" + +import pytest + +from core.inference.safetensors_agentic import run_safetensors_tool_loop +from state import tool_approvals +from state.tool_approvals import TOOL_REJECTED_MESSAGE, resolve_tool_decision + +_SESSION = "loop-session" + + +@pytest.fixture(autouse = True) +def _clear_pending(): + with tool_approvals._lock: + tool_approvals._pending.clear() + yield + with tool_approvals._lock: + tool_approvals._pending.clear() + + +class _FakeExecuteTool: + def __init__(self): + self.calls = [] + + def __call__( + self, + name, + arguments, + *, + cancel_event = None, + timeout = None, + session_id = None, + rag_scope = None, + ): + self.calls.append((name, arguments)) + return f"RESULT[{name}]" + + +def _tool_call(name, args_json): + return f'{{"name": "{name}", "arguments": {args_json}}}' + + +def _multi_turn(turns): + """A single_turn generator that yields one full snapshot per turn.""" + turn_iter = iter(turns) + + def _gen(_messages): + try: + yield next(turn_iter) + except StopIteration: + return + + return _gen + + +_DEFAULT_TOOLS = [ + {"type": "function", "function": {"name": "python"}}, + {"type": "function", "function": {"name": "web_search"}}, +] + + +def _drive( + turns, + decisions, + *, + tools = None, +): + """Run the loop, resolving each gated tool_start with the next decision. + + The advertised ``tools`` list drives the loop's enabled-tool filter + (pass a list omitting a tool to make a call to it "disabled"). + Returns (events, execute_calls). + """ + decision_iter = iter(decisions) + exec_fn = _FakeExecuteTool() + gen = run_safetensors_tool_loop( + single_turn = _multi_turn(turns), + messages = [{"role": "user", "content": "hi"}], + tools = _DEFAULT_TOOLS if tools is None else tools, + execute_tool = exec_fn, + session_id = _SESSION, + confirm_tool_calls = True, + ) + events = [] + for ev in gen: + events.append(ev) + if ev["type"] == "tool_start" and ev.get("awaiting_confirmation"): + # Slot is already registered (begin ran before this yield), so + # the decision lands before the loop enters its blocking wait. + resolve_tool_decision(ev["approval_id"], next(decision_iter), session_id = _SESSION) + return events, exec_fn.calls + + +def _tool_starts(events): + return [e for e in events if e["type"] == "tool_start"] + + +def _tool_ends(events): + return [e for e in events if e["type"] == "tool_end"] + + +def test_allow_executes_the_tool_once(): + events, calls = _drive( + [_tool_call("python", '{"code": "print(1)"}'), "final answer"], + ["allow"], + ) + starts = _tool_starts(events) + assert len(starts) == 1 + assert starts[0]["awaiting_confirmation"] is True + assert starts[0]["approval_id"] + assert calls == [("python", {"code": "print(1)"})] + assert _tool_ends(events)[0]["result"] == "RESULT[python]" + + +def test_deny_skips_execution_and_feeds_rejection(): + events, calls = _drive( + [_tool_call("python", '{"code": "print(1)"}'), "final answer"], + ["deny"], + ) + assert calls == [] # tool never ran + assert _tool_ends(events)[0]["result"] == TOOL_REJECTED_MESSAGE + + +def test_disabled_tool_is_not_prompted(): + events, calls = _drive( + [_tool_call("python", '{"code": "print(1)"}'), "final answer"], + [], + tools = [{"type": "function", "function": {"name": "web_search"}}], + ) + assert _tool_starts(events) == [] + assert _tool_ends(events) == [] + assert calls == [] + + +def test_duplicate_call_is_not_prompted(): + same = _tool_call("python", '{"code": "print(1)"}') + events, calls = _drive([same, same, "final answer"], ["allow"]) + starts = _tool_starts(events) + assert len(starts) == 1 + assert starts[0]["awaiting_confirmation"] is True + assert calls == [("python", {"code": "print(1)"})] + assert len(_tool_ends(events)) == 1 + + +def test_denied_call_can_be_reissued_and_approved(): + # Deny, then the model re-issues the identical call -> approving it must + # execute, not get suppressed as a duplicate (denied calls are not added + # to the duplicate-detection history). + same = _tool_call("python", '{"code": "print(1)"}') + events, calls = _drive([same, same, "final answer"], ["deny", "allow"]) + starts = _tool_starts(events) + assert len(starts) == 2 + assert starts[0]["awaiting_confirmation"] is True + assert starts[1]["awaiting_confirmation"] is True # not treated as dup + assert calls == [("python", {"code": "print(1)"})] # ran once, on approve + ends = _tool_ends(events) + assert ends[0]["result"] == TOOL_REJECTED_MESSAGE + assert ends[1]["result"] == "RESULT[python]" diff --git a/studio/backend/tests/test_tool_confirm_stream.py b/studio/backend/tests/test_tool_confirm_stream.py new file mode 100644 index 0000000000..b8e0472e12 --- /dev/null +++ b/studio/backend/tests/test_tool_confirm_stream.py @@ -0,0 +1,219 @@ +# SPDX-License-Identifier: AGPL-3.0-only +# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0 + +"""End-to-end handshake test for the tool-confirmation gate, no model. + +The real Studio stream wrappers in ``routes/inference.py`` drive the +synchronous agentic generator with ``await asyncio.to_thread(next, gen, +...)`` so the blocking ``threading.Event`` wait runs off the event loop. +This test rebuilds that exact pattern around the real +``state.tool_approvals`` functions, served by a real uvicorn process on +loopback (the same server Studio uses), and proves the load-bearing +property: + +* ``tool_start`` reaches the client before the gate blocks, and +* the separate ``/tool-confirm`` POST is served *while* the stream + connection is blocked, after which the stream resumes with the executed + (allow) or rejected (deny) result -- i.e. no deadlock. + +Each scenario runs under a socket-level timeout, so a regression that +reintroduces a deadlock fails fast instead of hanging the suite. +""" + +import asyncio +import json +import socket +import threading +import time + +import httpx +import pytest +import uvicorn +from fastapi import FastAPI, Request +from fastapi.responses import StreamingResponse + +from state import tool_approvals +from state.tool_approvals import ( + TOOL_REJECTED_MESSAGE, + begin_tool_decision, + new_approval_id, + resolve_tool_decision, + wait_tool_decision, +) + +_EXECUTED_RESULT = "tool executed: 2" + + +@pytest.fixture(autouse = True) +def _clear_pending(): + with tool_approvals._lock: + tool_approvals._pending.clear() + yield + with tool_approvals._lock: + tool_approvals._pending.clear() + + +def _build_app() -> FastAPI: + """Minimal app mirroring the real stream/confirm wiring.""" + app = FastAPI() + + def agentic_gen(session_id, cancel_event): + # Same shape as the real loops: register the approval slot, announce + # the call (echoing approval_id), gate on the decision, then either + # execute or feed back the rejection. + approval_id = new_approval_id() + slot = begin_tool_decision(session_id, approval_id) + yield { + "type": "tool_start", + "tool_name": "python", + "approval_id": approval_id, + "awaiting_confirmation": True, + } + denied = wait_tool_decision(slot, approval_id, cancel_event = cancel_event) == "deny" + result = TOOL_REJECTED_MESSAGE if denied else _EXECUTED_RESULT + yield {"type": "tool_end", "tool_name": "python", "result": result} + + @app.post("/stream") + async def stream(req: Request): + body = await req.json() + session_id = body.get("session_id") + cancel_event = threading.Event() + sentinel = object() + + async def wrapper(): + gen = agentic_gen(session_id, cancel_event) + while True: + event = await asyncio.to_thread(next, gen, sentinel) + if event is sentinel: + break + yield f"data: {json.dumps(event)}\n\n" + + return StreamingResponse(wrapper(), media_type = "text/event-stream") + + @app.post("/tool-confirm") + async def tool_confirm(req: Request): + body = await req.json() + resolved = resolve_tool_decision( + body.get("approval_id"), + body.get("decision"), + session_id = body.get("session_id"), + ) + return {"resolved": resolved} + + return app + + +def _free_port() -> int: + s = socket.socket(socket.AF_INET, socket.SOCK_STREAM) + s.bind(("127.0.0.1", 0)) + port = s.getsockname()[1] + s.close() + return port + + +class _Server: + """Run a uvicorn server in a background thread for the test's lifetime.""" + + def __init__(self, app): + self.port = _free_port() + config = uvicorn.Config(app, host = "127.0.0.1", port = self.port, log_level = "warning") + self.server = uvicorn.Server(config) + self._thread = threading.Thread(target = self.server.run, daemon = True) + + def __enter__(self): + self._thread.start() + deadline = time.monotonic() + 10.0 + while time.monotonic() < deadline: + if self.server.started: + return self + time.sleep(0.02) + raise AssertionError("uvicorn did not start in time") + + def __exit__(self, *exc): + self.server.should_exit = True + self._thread.join(timeout = 10.0) + + @property + def base_url(self) -> str: + return f"http://127.0.0.1:{self.port}" + + +async def _gate_is_blocking(approval_id) -> None: + """Wait until the stream thread is parked on this approval's slot. + + The slot is registered before ``tool_start`` is yielded, so it exists + by the time the client receives the event -- exactly as in reality, + where the confirm POST only arrives after the card renders. + """ + for _ in range(400): + with tool_approvals._lock: + slot = tool_approvals._pending.get(approval_id) + if slot is not None and not slot["event"].is_set(): + return + await asyncio.sleep(0.005) + raise AssertionError("gate never started waiting") + + +async def _drive(base_url, session_id, decision): + events = [] + resolved = None + timeout = httpx.Timeout(10.0) + async with httpx.AsyncClient(base_url = base_url, timeout = timeout) as client: + async with client.stream("POST", "/stream", json = {"session_id": session_id}) as resp: + assert resp.status_code == 200 + async for line in resp.aiter_lines(): + if not line.startswith("data: "): + continue + event = json.loads(line[len("data: ") :]) + events.append(event) + if event["type"] == "tool_start": + # The stream is now blocked on the gate; the confirm + # POST (echoing approval_id) must still be served over a + # second connection. + approval_id = event["approval_id"] + await _gate_is_blocking(approval_id) + r = await client.post( + "/tool-confirm", + json = { + "session_id": session_id, + "approval_id": approval_id, + "decision": decision, + }, + ) + resolved = r.json()["resolved"] + return events, resolved + + +def _run(session_id, decision): + with _Server(_build_app()) as srv: + return asyncio.run( + asyncio.wait_for(_drive(srv.base_url, session_id, decision), timeout = 15.0) + ) + + +def _types(events): + return [e["type"] for e in events] + + +def test_allow_resumes_stream_with_executed_result(): + events, resolved = _run("sess-allow", "allow") + assert resolved is True + assert _types(events) == ["tool_start", "tool_end"] + assert events[-1]["result"] == _EXECUTED_RESULT + + +def test_deny_resumes_stream_with_rejection_result(): + events, resolved = _run("sess-deny", "deny") + assert resolved is True + assert _types(events) == ["tool_start", "tool_end"] + assert events[-1]["result"] == TOOL_REJECTED_MESSAGE + + +def test_tool_start_precedes_the_block_and_carries_approval_id(): + # The first streamed event is always tool_start, proving the buttons + # can render before the backend pauses for the decision -- and it + # carries the approval_id / awaiting_confirmation the UI needs. + events, _ = _run("sess-order", "allow") + assert events[0]["type"] == "tool_start" + assert events[0]["awaiting_confirmation"] is True + assert events[0]["approval_id"] diff --git a/studio/frontend/src/components/assistant-ui/thread.tsx b/studio/frontend/src/components/assistant-ui/thread.tsx index 84e1654079..100143ed33 100644 --- a/studio/frontend/src/components/assistant-ui/thread.tsx +++ b/studio/frontend/src/components/assistant-ui/thread.tsx @@ -19,6 +19,7 @@ import { thinkEffortAriaLabel, thinkToggleAriaLabel, } from "@/components/assistant-ui/think-aria-label"; +import { withToolConfirmation } from "@/components/assistant-ui/tool-confirmation-controls"; import { ToolFallback } from "@/components/assistant-ui/tool-fallback"; import { ToolGroup } from "@/components/assistant-ui/tool-group"; import { CodeExecutionToolUI } from "@/components/assistant-ui/tool-ui-code-execution"; @@ -2520,6 +2521,19 @@ const CancelledIndicator: FC = () => { ); }; +const WebSearchToolUIConfirmable = withToolConfirmation(WebSearchToolUI); +const KnowledgeBaseToolUIConfirmable = + withToolConfirmation(KnowledgeBaseToolUI); +const PythonToolUIConfirmable = withToolConfirmation(PythonToolUI); +const TerminalToolUIConfirmable = withToolConfirmation(TerminalToolUI); +const CodeExecutionToolUIConfirmable = + withToolConfirmation(CodeExecutionToolUI); +const ImageGenerationToolUIConfirmable = withToolConfirmation( + ImageGenerationToolUI, +); +const RenderHtmlToolUIConfirmable = withToolConfirmation(RenderHtmlToolUI); +const ToolFallbackConfirmable = withToolConfirmation(ToolFallback); + const AssistantMessage: FC = () => { return ( { ToolGroup: ToolGroup, tools: { by_name: { - web_search: WebSearchToolUI, - search_knowledge_base: KnowledgeBaseToolUI, - python: PythonToolUI, - terminal: TerminalToolUI, - code_execution: CodeExecutionToolUI, - image_generation: ImageGenerationToolUI, - render_html: RenderHtmlToolUI, + web_search: WebSearchToolUIConfirmable, + search_knowledge_base: KnowledgeBaseToolUIConfirmable, + python: PythonToolUIConfirmable, + terminal: TerminalToolUIConfirmable, + code_execution: CodeExecutionToolUIConfirmable, + image_generation: ImageGenerationToolUIConfirmable, + render_html: RenderHtmlToolUIConfirmable, }, - Fallback: ToolFallback, + Fallback: ToolFallbackConfirmable, }, }} /> diff --git a/studio/frontend/src/components/assistant-ui/tool-confirmation-controls.tsx b/studio/frontend/src/components/assistant-ui/tool-confirmation-controls.tsx new file mode 100644 index 0000000000..89a15c61ab --- /dev/null +++ b/studio/frontend/src/components/assistant-ui/tool-confirmation-controls.tsx @@ -0,0 +1,157 @@ +// SPDX-License-Identifier: AGPL-3.0-only +// Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0 + +"use client"; + +import { Button } from "@/components/ui/button"; +import { resolveToolConfirmation } from "@/features/chat/api/chat-api"; +import { useChatRuntimeStore } from "@/features/chat/stores/chat-runtime-store"; +import type { + ToolCallMessagePartComponent, + ToolCallMessagePartStatus, +} from "@assistant-ui/react"; +import { useCallback, useEffect, useState } from "react"; + +/** + * Allow / Always allow / Deny controls for a tool call paused awaiting the + * user's confirmation. Rendered alongside every tool card (built-in and + * MCP) so the gate works for all tools, not just the ones using the + * fallback renderer. + * + * A card is "awaiting" only when the adapter registered a backend-gated + * pending call for it (see `toolConfirmations` in the runtime store), so + * non-gated cards -- toggle off, or external-provider tools that already + * ran -- never show controls. + */ +export function ToolConfirmationControls({ + toolCallId, + toolName, + result, + status, +}: { + toolCallId?: string; + toolName: string; + result: unknown; + status?: ToolCallMessagePartStatus; +}) { + const confirmation = useChatRuntimeStore((s) => + toolCallId && + Object.prototype.hasOwnProperty.call(s.toolConfirmations, toolCallId) + ? s.toolConfirmations[toolCallId] + : undefined, + ); + const allowToolAlways = useChatRuntimeStore((s) => s.allowToolAlways); + const clearToolConfirmation = useChatRuntimeStore( + (s) => s.clearToolConfirmation, + ); + const autoAllowKey = confirmation?.autoAllowKey ?? ""; + const autoAllowed = useChatRuntimeStore( + (s) => + s.alwaysAllowToolsBySession.get(autoAllowKey)?.has(toolName) ?? false, + ); + + const [decided, setDecided] = useState(false); + const [pending, setPending] = useState<"allow" | "deny" | null>(null); + const [failed, setFailed] = useState(false); + + // Still awaiting our decision: a gated pending entry exists, the tool has + // not produced a result, and the card is in its running state. + const awaiting = + confirmation !== undefined && + result === undefined && + status?.type === "running"; + const showControls = awaiting && !decided; + + const resolve = useCallback( + async (decision: "allow" | "deny") => { + if (!toolCallId || !confirmation) return; + setPending(decision); + setFailed(false); + try { + const ok = await resolveToolConfirmation( + confirmation.sessionId, + confirmation.approvalId, + decision, + ); + if (ok) { + // Only hide the controls once the backend confirms it matched the + // pending call -- otherwise the generation would stay blocked with + // no way to retry. + setDecided(true); + clearToolConfirmation(toolCallId); + } else { + setFailed(true); + } + } catch { + setFailed(true); + } finally { + setPending(null); + } + }, + [toolCallId, confirmation, clearToolConfirmation], + ); + + // Tools the user marked "Always allow" (this session) approve themselves. + useEffect(() => { + if (showControls && autoAllowed && pending === null && !failed) { + void resolve("allow"); + } + }, [showControls, autoAllowed, pending, failed, resolve]); + + if (!showControls) return null; + // Auto-approved tools resolve silently unless the post fails. + if (autoAllowed && !failed) return null; + + return ( +
+ + + + {failed ? ( + + Could not send your decision. Try again. + + ) : null} +
+ ); +} + +export function withToolConfirmation( + Component: ToolCallMessagePartComponent, +): ToolCallMessagePartComponent { + const WithToolConfirmation: ToolCallMessagePartComponent = (props) => ( + <> + + + + ); + return WithToolConfirmation; +} diff --git a/studio/frontend/src/components/assistant-ui/tool-fallback.tsx b/studio/frontend/src/components/assistant-ui/tool-fallback.tsx index eeb4c2059f..8930b6386f 100644 --- a/studio/frontend/src/components/assistant-ui/tool-fallback.tsx +++ b/studio/frontend/src/components/assistant-ui/tool-fallback.tsx @@ -325,6 +325,9 @@ const ToolFallbackImpl: ToolCallMessagePartComponent = ({ result, status, }) => { + // Allow/Deny confirmation controls are rendered uniformly for every tool + // card (built-in and fallback) by the `withToolConfirmation` wrapper in + // thread.tsx, so this renderer stays purely presentational. const isCancelled = status?.type === "incomplete" && status.reason === "cancelled"; diff --git a/studio/frontend/src/components/assistant-ui/tool-group.tsx b/studio/frontend/src/components/assistant-ui/tool-group.tsx index 569e673da3..3ce65627e2 100644 --- a/studio/frontend/src/components/assistant-ui/tool-group.tsx +++ b/studio/frontend/src/components/assistant-ui/tool-group.tsx @@ -9,6 +9,7 @@ import { type PropsWithChildren, } from "react"; import { useAuiState } from "@assistant-ui/react"; +import { useChatRuntimeStore } from "@/features/chat/stores/chat-runtime-store"; import { ChevronDownIcon } from "lucide-react"; import { Wrench01Icon } from "@hugeicons/core-free-icons"; import { HugeiconsIcon } from "@hugeicons/react"; @@ -216,6 +217,31 @@ const ToolGroupImpl: FC< (part) => part.type === "tool-call" && part.toolName === "render_html", ), ); + // A blocking allow/deny prompt must never be hidden inside a collapsed + // group, so force the group open while any of its calls awaits confirmation. + const toolConfirmations = useChatRuntimeStore((s) => s.toolConfirmations); + const hasPendingConfirmation = useAuiState(({ message }) => + message.parts + .slice(startIndex, endIndex + 1) + .some( + (part) => + part.type === "tool-call" && + Object.prototype.hasOwnProperty.call( + toolConfirmations, + part.toolCallId, + ), + ), + ); + const messageRunning = useAuiState( + ({ message }) => message.status?.type === "running", + ); + // Keep the group open once a confirmation forced it open, so answering an + // allow/deny doesn't snap it shut between sequential tool calls. It reverts + // to the default collapsed state once the turn finishes. + const forcedOpenRef = useRef(false); + if (hasPendingConfirmation) forcedOpenRef.current = true; + const forceOpen = + hasPendingConfirmation || (forcedOpenRef.current && messageRunning); // Render single tool calls and artifacts directly so cards never hide in a // collapsed group. @@ -224,7 +250,7 @@ const ToolGroupImpl: FC< } return ( - + {children} diff --git a/studio/frontend/src/features/chat/api/chat-adapter.ts b/studio/frontend/src/features/chat/api/chat-adapter.ts index 983d4f5aa2..ce9f1c6544 100644 --- a/studio/frontend/src/features/chat/api/chat-adapter.ts +++ b/studio/frontend/src/features/chat/api/chat-adapter.ts @@ -1440,6 +1440,10 @@ export function createOpenAIStreamAdapter(): ChatModelAdapter { const resolvedThreadId = (unstable_threadId ?? runtime.activeThreadId) || undefined; const sandboxSessionId = await resolveSandboxSessionId(resolvedThreadId); + const toolConfirmationScopeId = resolvedThreadId + ? `${sandboxSessionId || "_default"}:${resolvedThreadId}` + : sandboxSessionId || "_default"; + const toolConfirmationIdsByBackendId = new Map(); const resolvedThreadKey = resolvedThreadId ?? null; const pendingImageEditReferenceForRun = runtime.pendingImageEditReference; const selectedImageEditReference = @@ -1513,6 +1517,7 @@ export function createOpenAIStreamAdapter(): ChatModelAdapter { imageToolsEnabled, artifactsEnabled, mcpEnabledForChat, + confirmToolCalls, webFetchToolsEnabled, ragEnabled, ragSource, @@ -2435,6 +2440,7 @@ export function createOpenAIStreamAdapter(): ChatModelAdapter { : []), ], mcp_enabled: mcpEnabledForChat, + confirm_tool_calls: confirmToolCalls, // Scope: thread_id = this thread's docs, kb_id = a KB. ...(ragEnabled ? { @@ -2560,9 +2566,20 @@ export function createOpenAIStreamAdapter(): ChatModelAdapter { toolEvent.provenance, ); if (toolEvent.type === "tool_start") { + const backendToolCallId = + (toolEvent.tool_call_id as string) || ""; + const approvalId = (toolEvent.approval_id as string) || ""; + const awaitingConfirmation = + toolEvent.awaiting_confirmation === true; const id = - (toolEvent.tool_call_id as string) || - `${toolEvent.tool_name}_${Date.now()}`; + awaitingConfirmation && approvalId + ? `${toolConfirmationScopeId}:${approvalId}` + : backendToolCallId || + approvalId || + `${toolEvent.tool_name}_${Date.now()}`; + if (awaitingConfirmation && backendToolCallId) { + toolConfirmationIdsByBackendId.set(backendToolCallId, id); + } const toolArgs = (toolEvent.arguments ?? {}) as ToolCallMessagePart["args"]; const idx = toolCallParts.findIndex( @@ -2593,11 +2610,30 @@ export function createOpenAIStreamAdapter(): ChatModelAdapter { ...(toolProvenance ? { provenance: toolProvenance } : {}), } as PositionedToolCallPart); } + if (awaitingConfirmation) { + useChatRuntimeStore + .getState() + .setToolConfirmation( + id, + approvalId, + sandboxSessionId ?? "", + toolConfirmationScopeId, + ); + } } else if (toolEvent.type === "tool_end") { + const backendToolCallId = + (toolEvent.tool_call_id as string) || ""; const id = - (toolEvent.tool_call_id as string) || + (backendToolCallId + ? toolConfirmationIdsByBackendId.get(backendToolCallId) + : undefined) || + backendToolCallId || toolCallParts[toolCallParts.length - 1]?.toolCallId || ""; + if (backendToolCallId) { + toolConfirmationIdsByBackendId.delete(backendToolCallId); + } + useChatRuntimeStore.getState().clearToolConfirmation(id); const idx = toolCallParts.findIndex( (p) => p.toolCallId === id, ); @@ -3149,6 +3185,10 @@ export function createOpenAIStreamAdapter(): ChatModelAdapter { throw err; } finally { abortSignal.removeEventListener("abort", onAbortCancel); + const confirmStore = useChatRuntimeStore.getState(); + for (const part of toolCallParts) { + confirmStore.clearToolConfirmation(part.toolCallId); + } runtime.setGeneratingStatus(null); runtime.setToolStatus(null); clearTimeout(warmupTimer); diff --git a/studio/frontend/src/features/chat/api/chat-api.ts b/studio/frontend/src/features/chat/api/chat-api.ts index 9d0939c37d..110d573571 100644 --- a/studio/frontend/src/features/chat/api/chat-api.ts +++ b/studio/frontend/src/features/chat/api/chat-api.ts @@ -110,6 +110,31 @@ export async function unloadModel(payload: UnloadModelRequest): Promise { await parseJsonOrThrow(response); } +/** + * Allow or deny a tool call that is paused awaiting user confirmation + * (when the "Confirm tool calls" toggle is on). The call is identified by + * the backend ``approvalId`` echoed in the tool_start event; ``sessionId`` + * is a scope check. Resolves to ``true`` only when the backend matched a + * pending call, so the caller can surface a retry on a stale/failed post. + */ +export async function resolveToolConfirmation( + sessionId: string, + approvalId: string, + decision: "allow" | "deny", +): Promise { + const response = await authFetch("/api/inference/tool-confirm", { + method: "POST", + headers: { "Content-Type": "application/json" }, + body: JSON.stringify({ + session_id: sessionId, + approval_id: approvalId, + decision, + }), + }); + const parsed = await parseJsonOrThrow<{ resolved?: boolean }>(response); + return parsed.resolved === true; +} + export interface CachedGgufRepo { repo_id: string; size_bytes: number; diff --git a/studio/frontend/src/features/chat/chat-settings-sheet.tsx b/studio/frontend/src/features/chat/chat-settings-sheet.tsx index 5751a1a426..59ac2ae6b4 100644 --- a/studio/frontend/src/features/chat/chat-settings-sheet.tsx +++ b/studio/frontend/src/features/chat/chat-settings-sheet.tsx @@ -1496,6 +1496,7 @@ export function ChatSettingsPanel({
+
@@ -1682,6 +1683,30 @@ function AutoHealToolCallsToggle() { ); } +function ConfirmToolCallsToggle() { + const confirmToolCalls = useChatRuntimeStore((s) => s.confirmToolCalls); + const setConfirmToolCalls = useChatRuntimeStore((s) => s.setConfirmToolCalls); + + return ( +
+
+ + Confirm tool calls + + + When on, local Studio tool calls pause for your approval before they + run. Provider-hosted tools are not gated here. + +
+ +
+ ); +} + function ChatTemplateFields() { const defaultTemplate = useChatRuntimeStore((s) => s.defaultChatTemplate); const override = useChatRuntimeStore((s) => s.chatTemplateOverride); 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 3ef63d2945..2f0e3027a8 100644 --- a/studio/frontend/src/features/chat/stores/chat-runtime-store.ts +++ b/studio/frontend/src/features/chat/stores/chat-runtime-store.ts @@ -34,6 +34,7 @@ export const CHAT_COLLAPSE_HTML_ARTIFACTS_KEY = export const CHAT_ALLOW_ARTIFACT_NETWORK_ACCESS_KEY = "unsloth_chat_allow_artifact_network_access"; export const CHAT_MCP_ENABLED_KEY = "unsloth_chat_mcp_enabled"; +export const CHAT_CONFIRM_TOOL_CALLS_KEY = "unsloth_chat_confirm_tool_calls"; export const CHAT_WEB_FETCH_TOOLS_ENABLED_KEY = "unsloth_chat_web_fetch_tools_enabled"; export const CHAT_RAG_SOURCE_KEY = "unsloth_chat_rag_source"; @@ -402,6 +403,29 @@ type ChatRuntimeStore = { // autoInject = forced first-pass retrieval before answering. ragAutoInject: RagAutoInject; ragAutoInjectMinScore: number; + /** + * When on, local Studio tool calls pause for an explicit allow/deny in the + * chat before they run. + */ + confirmToolCalls: boolean; + /** + * Per-chat set of tool names the user chose to auto-approve via "Always + * allow". Keyed by UI confirmation scope, not necessarily the backend + * sandbox session id. Not persisted across reloads. + */ + alwaysAllowToolsBySession: Map>; + /** + * Tool calls currently paused awaiting the user's allow/deny decision, + * keyed by the scoped frontend tool-call id. Each entry carries the backend + * ``approvalId`` to echo back and the ``sessionId`` the generation runs + * under, so the confirmation always resolves the exact pending call. The + * ``autoAllowKey`` scopes the UI-only "Always allow" bucket per chat. + * Only backend-gated local tool calls are added here. + */ + toolConfirmations: Record< + string, + { approvalId: string; sessionId: string; autoAllowKey: string } + >; /** * Fetch pill state, independent of `toolsEnabled` (Search). Only * consulted when `providerSupportsBuiltinWebFetch` is true. @@ -483,6 +507,15 @@ type ChatRuntimeStore = { setCollapseHtmlArtifacts: (enabled: boolean) => void; setAllowArtifactNetworkAccess: (enabled: boolean) => void; setMcpEnabledForChat: (enabled: boolean) => void; + setConfirmToolCalls: (enabled: boolean) => void; + allowToolAlways: (sessionId: string, toolName: string) => void; + setToolConfirmation: ( + toolCallId: string, + approvalId: string, + sessionId: string, + autoAllowKey: string, + ) => void; + clearToolConfirmation: (toolCallId: string) => void; setWebFetchToolsEnabled: (enabled: boolean) => void; setRagEnabled: (enabled: boolean) => void; setRagSource: (source: RagSource) => void; @@ -749,6 +782,9 @@ export const useChatRuntimeStore = create((set, get) => ({ false, ), mcpEnabledForChat: loadBool(CHAT_MCP_ENABLED_KEY, false), + confirmToolCalls: loadBool(CHAT_CONFIRM_TOOL_CALLS_KEY, false), + alwaysAllowToolsBySession: new Map>(), + toolConfirmations: {}, webFetchToolsEnabled: loadBool(CHAT_WEB_FETCH_TOOLS_ENABLED_KEY, false), // RAG is opt-in per session: always starts off, never restored from storage. ragEnabled: false, @@ -1074,6 +1110,40 @@ export const useChatRuntimeStore = create((set, get) => ({ saveBool(CHAT_MCP_ENABLED_KEY, mcpEnabledForChat); return { mcpEnabledForChat }; }), + setConfirmToolCalls: (confirmToolCalls) => + set(() => { + saveBool(CHAT_CONFIRM_TOOL_CALLS_KEY, confirmToolCalls); + return { confirmToolCalls }; + }), + allowToolAlways: (sessionId, toolName) => + set((state) => { + const current = state.alwaysAllowToolsBySession.get(sessionId); + if (current?.has(toolName)) return state; + const next = new Map(state.alwaysAllowToolsBySession); + next.set(sessionId, new Set(current ?? []).add(toolName)); + return { alwaysAllowToolsBySession: next }; + }), + setToolConfirmation: (toolCallId, approvalId, sessionId, autoAllowKey) => + set((state) => ({ + toolConfirmations: { + ...state.toolConfirmations, + [toolCallId]: { approvalId, sessionId, autoAllowKey }, + }, + })), + clearToolConfirmation: (toolCallId) => + set((state) => { + if ( + !Object.prototype.hasOwnProperty.call( + state.toolConfirmations, + toolCallId, + ) + ) { + return state; + } + const next = { ...state.toolConfirmations }; + delete next[toolCallId]; + return { toolConfirmations: next }; + }), setWebFetchToolsEnabled: (webFetchToolsEnabled) => set(() => { saveBool(CHAT_WEB_FETCH_TOOLS_ENABLED_KEY, webFetchToolsEnabled); diff --git a/studio/frontend/src/features/chat/types/api.ts b/studio/frontend/src/features/chat/types/api.ts index a0e5d5355e..96caa38bd4 100644 --- a/studio/frontend/src/features/chat/types/api.ts +++ b/studio/frontend/src/features/chat/types/api.ts @@ -282,6 +282,8 @@ export interface OpenAIChatCompletionsRequest { enabled_tools?: string[]; /** Local models + enable_tools only. */ mcp_enabled?: boolean; + /** Local models + enable_tools only. */ + confirm_tool_calls?: boolean; /** Exactly one of `kb_id` (a KB) or `thread_id` (thread docs). */ rag_scope?: { kb_id?: string; From 6b62b2b5c04c640271b25746ac191cff37c66189 Mon Sep 17 00:00:00 2001 From: Daniel Han Date: Fri, 12 Jun 2026 01:56:05 -0700 Subject: [PATCH 11/81] Guard Apple GPU power against negative counter-reset readings (#6235) IOReport energy counters can reset (sleep/wake, power gating), making a poll delta negative. Return None for a negative total so the monitor shows -- for that poll instead of a bogus negative wattage; it self-corrects next poll. --- studio/backend/utils/hardware/apple.py | 4 +++- 1 file changed, 3 insertions(+), 1 deletion(-) diff --git a/studio/backend/utils/hardware/apple.py b/studio/backend/utils/hardware/apple.py index 3f14af60ad..62dbd10b8d 100644 --- a/studio/backend/utils/hardware/apple.py +++ b/studio/backend/utils/hardware/apple.py @@ -386,7 +386,9 @@ class _IOReportEnergy: watts = _watts(energy, unit, elapsed_s) if watts is not None: total = (total or 0.0) + watts - return round(total, 1) if total is not None else None + if total is None or total < 0: # negative = counter reset; show -- not a bogus draw + return None + return round(total, 1) # ========== Public API (module singletons, failure-latched) ========== From 95a2627bf6d96efebd64681a555bf8b9ab90e022 Mon Sep 17 00:00:00 2001 From: Irakli <39024518+IrakliXYZ@users.noreply.github.com> Date: Fri, 12 Jun 2026 13:30:53 +0400 Subject: [PATCH 12/81] Fix step count mismatch when sequence packing is enabled (#5967) * Fix step count mismatch when sequence packing is enabled * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * Emit a single step-0 progress event and guard applyStatus totalSteps Merge the two consecutive _update_progress calls before train() so the step-0 gate in _on_progress fires once instead of twice, avoiding a duplicate startup event and a null-metric step-0 row in training_metrics. Apply the same positive-number guard to applyStatus that applyProgress uses, so a stale or startup status poll can no longer overwrite the packed step count with 0 or replace it with a stale total. * Log debug message when train_dataset length is unavailable The TypeError fallback for length-less datasets (e.g. streaming IterableDataset) was silent, leaving no trace that the step estimate came from the raw dataset rather than the packed one. * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci --------- Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com> Co-authored-by: Etherll <61019402+Etherll@users.noreply.github.com> --- studio/backend/core/training/trainer.py | 18 ++++++++++++++---- studio/backend/core/training/worker.py | 2 +- .../training/stores/training-runtime-store.ts | 9 ++++++--- 3 files changed, 21 insertions(+), 8 deletions(-) diff --git a/studio/backend/core/training/trainer.py b/studio/backend/core/training/trainer.py index 57342b2453..085f999dd6 100644 --- a/studio/backend/core/training/trainer.py +++ b/studio/backend/core/training/trainer.py @@ -3367,7 +3367,19 @@ class UnslothTrainer: # ========== PROGRESS TRACKING ========== self.trainer.add_callback(self._create_progress_callback()) - num_samples = len(dataset["dataset"] if isinstance(dataset, dict) else dataset) + num_samples = None + if hasattr(self.trainer, "train_dataset") and self.trainer.train_dataset is not None: + try: + num_samples = len(self.trainer.train_dataset) + except TypeError: + logger.debug( + "train_dataset does not support len(); falling back to " + "raw dataset size for step estimation." + ) + + if num_samples is None: + num_samples = len(dataset["dataset"] if isinstance(dataset, dict) else dataset) + batch_size = training_args.get("batch_size", 2) total_steps = self._calculate_total_steps( num_samples, @@ -3376,10 +3388,8 @@ class UnslothTrainer: training_args.get("num_epochs", 3), training_args.get("max_steps", 0), ) - self._update_progress(total_steps = total_steps) - # ========== START TRAINING ========== - self._update_progress(status_message = "Starting training...") + self._update_progress(total_steps = total_steps, status_message = "Starting training...") logger.info("Starting training...\n") self.trainer.train(resume_from_checkpoint = training_args.get("resume_from_checkpoint")) diff --git a/studio/backend/core/training/worker.py b/studio/backend/core/training/worker.py index 18b25cb4fe..c7c9003a04 100644 --- a/studio/backend/core/training/worker.py +++ b/studio/backend/core/training/worker.py @@ -2462,7 +2462,7 @@ def run_training_process(*, event_queue: Any, stop_queue: Any, config: dict) -> def _on_progress(progress: TrainingProgress): has_train_loss = progress.step > 0 and progress.loss is not None has_eval_loss = progress.eval_loss is not None - if has_train_loss or has_eval_loss: + if (progress.step == 0 and progress.total_steps > 0) or has_train_loss or has_eval_loss: event_queue.put( { "type": "progress", diff --git a/studio/frontend/src/features/training/stores/training-runtime-store.ts b/studio/frontend/src/features/training/stores/training-runtime-store.ts index 97fbd32d57..9eaaa98c0e 100644 --- a/studio/frontend/src/features/training/stores/training-runtime-store.ts +++ b/studio/frontend/src/features/training/stores/training-runtime-store.ts @@ -209,8 +209,8 @@ export const useTrainingRuntimeStore = create()((set) => ( currentStep: typeof detailStep === "number" ? Math.max(detailStep, 0) : state.currentStep, totalSteps: - typeof detailTotal === "number" - ? Math.max(detailTotal, 0) + typeof detailTotal === "number" && detailTotal > 0 + ? detailTotal : state.totalSteps, currentLoss: typeof detailLoss === "number" ? detailLoss : state.currentLoss, @@ -273,7 +273,10 @@ export const useTrainingRuntimeStore = create()((set) => ( ...state, jobId: payload.job_id || state.jobId, currentStep: step, - totalSteps: Math.max(payload.total_steps, state.totalSteps), + totalSteps: + typeof payload.total_steps === "number" && payload.total_steps > 0 + ? payload.total_steps + : state.totalSteps, // A null loss at a new step means the backend reported a non-finite // loss; clear the display instead of keeping the stale value. currentLoss: From e59ce0db0477b4b3161be2ce63f199b5e270b726 Mon Sep 17 00:00:00 2001 From: alkinun Date: Fri, 12 Jun 2026 12:37:51 +0300 Subject: [PATCH 13/81] fix/uv-bytecode-timeout (#6166) * fix/uv-bytecode-timeout * make sure that win installer upgrades uv for bytecode timeout * Clarify uv bytecode timeout comment in install.sh and install.ps1 * Read installer scripts as UTF-8 in parity test so it runs on Windows * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * Prefer freshly installed uv when an older one shadows it on PATH --------- Co-authored-by: Daniel Han Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com> --- install.ps1 | 57 +++++++++++++++++++--- install.sh | 6 ++- tests/python/test_cross_platform_parity.py | 56 ++++++++++++++++++--- 3 files changed, 104 insertions(+), 15 deletions(-) diff --git a/install.ps1 b/install.ps1 index cf7bb63cdf..9abddc9ce2 100644 --- a/install.ps1 +++ b/install.ps1 @@ -1177,14 +1177,38 @@ shell.Run cmd, 0, False if ($SkipTorch) { $InitialGpuBranch = "no_torch" } Write-TauriDiag -GpuBranch $InitialGpuBranch -TorchIndexFamily "none" -PythonVersionForDiag $DiagPythonVersion - # ── Install uv if not present ── + # ── Install uv ── Write-TauriLog "STEP" "Installing uv package manager" - if (-not (Get-Command uv -ErrorAction SilentlyContinue)) { - substep "installing uv package manager..." + $UvMinVersion = "0.7.22" + function Test-UvVersionOk { + $cmd = Get-Command uv -ErrorAction SilentlyContinue + if (-not $cmd) { return $false } + try { + $raw = (& uv --version 2>$null | Select-Object -First 1) + } catch { + return $false + } + if ($raw -notmatch 'uv\s+([0-9]+(?:\.[0-9]+)+)') { return $false } + try { + return ([version]$Matches[1] -ge [version]$UvMinVersion) + } catch { + return $false + } + } + + if (-not (Test-UvVersionOk)) { + if (Get-Command uv -ErrorAction SilentlyContinue) { + substep "updating uv package manager..." + } else { + substep "installing uv package manager..." + } if ($script:WingetAvailable) { $prevEAP = $ErrorActionPreference $ErrorActionPreference = "Continue" - try { winget install --id=astral-sh.uv -e --source winget --accept-package-agreements --accept-source-agreements } catch {} + try { winget upgrade --id=astral-sh.uv -e --source winget --accept-package-agreements --accept-source-agreements } catch {} + if (-not (Test-UvVersionOk)) { + try { winget install --id=astral-sh.uv -e --source winget --accept-package-agreements --accept-source-agreements } catch {} + } $ErrorActionPreference = $prevEAP Refresh-SessionPath } @@ -1192,19 +1216,40 @@ shell.Run cmd, 0, False # use Astral's official PowerShell installer. This is the only # supported path on hosts without winget (Windows ARM64 runners, # corporate machines without the Store, etc.). - if (-not (Get-Command uv -ErrorAction SilentlyContinue)) { + if (-not (Test-UvVersionOk)) { substep "installing uv via https://astral.sh/uv/install.ps1..." "Yellow" Invoke-Expression (Invoke-RestMethod -Uri "https://astral.sh/uv/install.ps1") Refresh-SessionPath } } - if (-not (Get-Command uv -ErrorAction SilentlyContinue)) { + # A freshly installed uv can sit later on PATH than an older one (active + # venv, Scoop/pipx shim). Prefer a just-installed uv from a known location. + if (-not (Test-UvVersionOk)) { + $origPath = $env:PATH + foreach ($d in @($env:UV_INSTALL_DIR, $env:XDG_BIN_HOME, + (Join-Path $env:USERPROFILE ".local\bin"), + (Join-Path $env:LOCALAPPDATA "Microsoft\WinGet\Links"))) { + if ($d -and (Test-Path $d)) { + $env:PATH = "$d;$origPath" + if (Test-UvVersionOk) { break } + $env:PATH = $origPath + } + } + } + + if (-not (Test-UvVersionOk)) { step "uv" "could not be installed" "Red" substep "Install it from https://docs.astral.sh/uv/" "Yellow" return (Exit-InstallFailure "uv could not be installed") } + # When bytecode compilation is enabled, large installs can exceed uv's 60s + # default on slow machines. Default to 180s, preserving overrides ("0" disables). + if (-not $env:UV_COMPILE_BYTECODE_TIMEOUT) { + $env:UV_COMPILE_BYTECODE_TIMEOUT = "180" + } + # ── Create venv (migrate old layout if possible, otherwise fresh) ── # Pass the resolved executable path to uv so it does not re-resolve # a version string back to a conda interpreter. diff --git a/install.sh b/install.sh index 532ac61bc0..eba92d0746 100755 --- a/install.sh +++ b/install.sh @@ -1456,7 +1456,11 @@ fi # ── Install uv ── tauri_log "STEP" "Installing uv package manager" -UV_MIN_VERSION="0.7.14" +UV_MIN_VERSION="0.7.22" + +# When bytecode compilation is enabled, large installs can exceed uv's 60s default on slow machines. Default to 180s, preserving overrides ("0" disables). +: "${UV_COMPILE_BYTECODE_TIMEOUT:=180}" +export UV_COMPILE_BYTECODE_TIMEOUT version_ge() { # returns 0 if $1 >= $2 diff --git a/tests/python/test_cross_platform_parity.py b/tests/python/test_cross_platform_parity.py index 0f2e73257a..34f984714e 100644 --- a/tests/python/test_cross_platform_parity.py +++ b/tests/python/test_cross_platform_parity.py @@ -20,7 +20,7 @@ class TestNoTorchBackendAutoInInstallSh: """ def test_no_torch_backend_auto_outside_fallback(self): - lines = INSTALL_SH.read_text().splitlines() + lines = INSTALL_SH.read_text(encoding = "utf-8").splitlines() # Fallback block: from "GPU detection failed" to the next "fi". fallback_start = None fallback_end = None @@ -48,7 +48,7 @@ class TestNoTorchBackendAutoInInstallSh: def test_fallback_uses_torch_backend_auto(self): """The fallback branch should use --torch-backend=auto as recovery.""" - text = INSTALL_SH.read_text() + text = INSTALL_SH.read_text(encoding = "utf-8") assert ( "GPU detection failed" in text ), "install.sh should have a fallback branch for when GPU detection fails" @@ -58,13 +58,13 @@ class TestInstallShHasGpuDetection: """install.sh must contain the get_torch_index_url function.""" def test_function_exists(self): - text = INSTALL_SH.read_text() + text = INSTALL_SH.read_text(encoding = "utf-8") assert ( "get_torch_index_url()" in text ), "install.sh is missing the get_torch_index_url() function" def test_torch_index_url_assigned(self): - text = INSTALL_SH.read_text() + text = INSTALL_SH.read_text(encoding = "utf-8") assert ( "TORCH_INDEX_URL=$(get_torch_index_url)" in text ), "install.sh should assign TORCH_INDEX_URL from get_torch_index_url()" @@ -115,8 +115,8 @@ class TestCudaMappingParity: def test_same_cuda_suffixes(self): """Both scripts should produce the same ordered list of CUDA index suffixes.""" - sh_text = INSTALL_SH.read_text() - ps1_text = INSTALL_PS1.read_text() + sh_text = INSTALL_SH.read_text(encoding = "utf-8") + ps1_text = INSTALL_PS1.read_text(encoding = "utf-8") sh_thresholds = self._extract_cuda_thresholds_sh(sh_text) ps1_thresholds = self._extract_cuda_thresholds_ps1(ps1_text) @@ -134,13 +134,53 @@ class TestPyTorchMirrorEnvVar: """Both install scripts must support the UNSLOTH_PYTORCH_MIRROR env var.""" def test_install_sh_has_mirror_var(self): - text = INSTALL_SH.read_text() + text = INSTALL_SH.read_text(encoding = "utf-8") assert ( "UNSLOTH_PYTORCH_MIRROR" in text ), "install.sh should reference UNSLOTH_PYTORCH_MIRROR" def test_install_ps1_has_mirror_var(self): - text = INSTALL_PS1.read_text() + text = INSTALL_PS1.read_text(encoding = "utf-8") assert ( "UNSLOTH_PYTORCH_MIRROR" in text ), "install.ps1 should reference UNSLOTH_PYTORCH_MIRROR" + + +class TestUvBytecodeCompileTimeout: + """Installers should relax uv bytecode compilation timeout by default.""" + + @staticmethod + def _version_tuple(version: str) -> tuple[int, ...]: + return tuple(int(part) for part in version.split(".")) + + def test_install_sh_uses_uv_version_with_timeout_env(self): + text = INSTALL_SH.read_text(encoding = "utf-8") + match = re.search(r'^UV_MIN_VERSION="([^"]+)"$', text, re.MULTILINE) + assert match, "install.sh should declare UV_MIN_VERSION" + assert self._version_tuple(match.group(1)) >= self._version_tuple("0.7.22") + + def test_install_ps1_uses_uv_version_with_timeout_env(self): + text = INSTALL_PS1.read_text(encoding = "utf-8") + match = re.search(r'^\s*\$UvMinVersion = "([^"]+)"$', text, re.MULTILINE) + assert match, "install.ps1 should declare $UvMinVersion" + assert self._version_tuple(match.group(1)) >= self._version_tuple("0.7.22") + assert "function Test-UvVersionOk" in text + assert "if (-not (Test-UvVersionOk))" in text + + def test_install_sh_preserves_timeout_override(self): + text = INSTALL_SH.read_text(encoding = "utf-8") + assert ( + ': "${UV_COMPILE_BYTECODE_TIMEOUT:=180}"' in text + ), "install.sh should default UV_COMPILE_BYTECODE_TIMEOUT without overwriting callers" + assert ( + "export UV_COMPILE_BYTECODE_TIMEOUT" in text + ), "install.sh should export UV_COMPILE_BYTECODE_TIMEOUT for uv subprocesses" + + def test_install_ps1_preserves_timeout_override(self): + text = INSTALL_PS1.read_text(encoding = "utf-8") + assert ( + "if (-not $env:UV_COMPILE_BYTECODE_TIMEOUT)" in text + ), "install.ps1 should preserve caller UV_COMPILE_BYTECODE_TIMEOUT overrides" + assert ( + '$env:UV_COMPILE_BYTECODE_TIMEOUT = "180"' in text + ), "install.ps1 should default UV_COMPILE_BYTECODE_TIMEOUT" From 25ccfebc0b218fe8321c023d25f7363aad8b7016 Mon Sep 17 00:00:00 2001 From: Daniel Han Date: Fri, 12 Jun 2026 02:39:01 -0700 Subject: [PATCH 14/81] Studio: tune llama.cpp env for data-center GPUs (#6098) * Studio: tune llama.cpp env for data-center GPUs Detect datacenter/professional NVIDIA GPUs at llama-server launch and set the llama.cpp env flags that help them, gated so consumer GeForce, AMD/ROCm, CPU and macOS are never touched. - GGML_CUDA_FORCE_CUBLAS_COMPUTE_32F=1 for any DC GPU (FP32 cuBLAS accumulation). On a B200 this is ~0% throughput cost with identical perplexity (7.3230 wikitext-2-raw, baseline and on), where on GeForce the same flag costs real throughput, hence the gate. - GGML_CUDA_P2P=1 and CUDA_SCALE_LAUNCH_QUEUES=4x for multi-GPU DC boxes. Benchmarked on 6x B200: +33-51% prompt processing on tensor (row) split and +8-16% on the default pipeline (layer) split, with no regression on the other split or on token generation. Detection uses torch device names (A100/A30/H100/H200/H800/GH200/B200/GB200/ GB300/L40/L4/RTX PRO 6000/RTX 6000 Ada). A mixed box with one consumer GPU in the selection is treated as non-DC. All writes are setdefault so a user value always wins, and UNSLOTH_DISABLE_DC_TUNING=1 turns the whole thing off. 37 unit tests cover detection, multi-GPU gating, user-override precedence, the disable flag and fail-open on error. * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * Studio: fix data-center GPU detection false positives and physical-id mapping Two issues in the data-center llama.cpp env tuning gate: - _is_datacenter_gpu matched the marker allowlist as unbounded substrings, so workstation/laptop parts "NVIDIA RTX A1000" and "NVIDIA RTX A3000" matched "a100"/"a30" and were wrongly tuned as data-center GPUs (forcing FP32 cuBLAS accumulation and the multi-GPU env, which carry a real cost on those cards). Switch to a word-boundary regex. - gpu_indices carries physical GPU ids (translated from torch ordinals by _get_gpu_free_memory via CUDA_VISIBLE_DEVICES), but they were passed straight into torch.cuda.get_device_properties, which expects mask-relative ordinals. On a masked host (e.g. CUDA_VISIBLE_DEVICES=4,5,6,7) a selection like [4,5] fell out of range and silently dropped the tuning, and on a mixed mask it could probe the wrong GPU class. Build a physical-id to device-name map mirroring _get_gpu_free_memory, then look up the selection by physical id. Add regression tests for the A1000/A3000 false positives and for masked-host physical-id selection (reordered and mixed-class masks included). * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * Studio: tighten data-center GPU tuning comments Comment-only pass over the DC tuning block and its tests: shorten verbose docstrings/comments, drop ones that restate the code, collapse multi-line blocks. Keep the load-bearing rationale (physical-id vs ordinal mapping, the word-boundary reason, the B200 benchmark numbers). No code change: verified with comment_tools.py check --strip-docstrings (code unchanged, comments only). --------- Co-authored-by: danielhanchen Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com> --- studio/backend/core/inference/llama_cpp.py | 107 +++++++ .../tests/test_datacenter_gpu_tuning.py | 278 ++++++++++++++++++ 2 files changed, 385 insertions(+) create mode 100644 studio/backend/tests/test_datacenter_gpu_tuning.py diff --git a/studio/backend/core/inference/llama_cpp.py b/studio/backend/core/inference/llama_cpp.py index 8cf37ed9ec..3393f2b4be 100644 --- a/studio/backend/core/inference/llama_cpp.py +++ b/studio/backend/core/inference/llama_cpp.py @@ -1329,6 +1329,105 @@ class LlamaCppBackend: return False return False + # Datacenter / professional NVIDIA parts that benefit from the llama.cpp + # FP32-accum / P2P tunings. Whole-word (\b) so short markers don't match + # workstation parts as substrings: "a100" must not fire on "RTX A1000". + _DATACENTER_GPU_RE = re.compile( + r"\b(?:a100|a30|h100|h200|h800|gh200|b200|b100|b300|gb200|gb300|" + r"l40s?|l4|rtx pro 6000|rtx 6000 ada)\b" + ) + + @staticmethod + def _is_datacenter_gpu(gpu_indices = None) -> bool: + """True iff every selected NVIDIA GPU is a datacenter/professional part. + NVIDIA-only, fails open to False (consumer GeForce, ROCm, CPU and errors + are left untouched); a mixed DC+consumer selection counts as non-DC. + + gpu_indices are PHYSICAL ids (see _get_gpu_free_memory), but + get_device_properties wants mask-relative ordinals, so we rebuild the + ordinal->physical map from CUDA_VISIBLE_DEVICES and key names by physical + id. Otherwise a masked host (CUDA_VISIBLE_DEVICES=4,5,6,7, selection [4,5]) + would drop the tuning or probe the wrong GPU.""" + try: + import torch + + if getattr(torch.version, "hip", None) is not None: + return False # ROCm reuses torch.cuda.*; not a CUDA part + if not (hasattr(torch, "cuda") and torch.cuda.is_available()): + return False + count = torch.cuda.device_count() + + # Mirror _get_gpu_free_memory: map visible ordinal -> physical id via + # CUDA_VISIBLE_DEVICES; unset/unparsable leaves physical id == ordinal. + physical_ids: Optional[list[int]] = None + cvd = os.environ.get("CUDA_VISIBLE_DEVICES") + if cvd is not None: + try: + physical_ids = [int(x.strip()) for x in cvd.split(",") if x.strip()] + except ValueError: + physical_ids = None + + pattern = LlamaCppBackend._DATACENTER_GPU_RE + names_by_id: dict[int, str] = {} + for ordinal in range(count): + try: + name = (torch.cuda.get_device_properties(ordinal).name or "").lower() + except Exception: + continue + pid = ( + physical_ids[ordinal] + if physical_ids is not None and ordinal < len(physical_ids) + else ordinal + ) + names_by_id[pid] = name + + indices = list(gpu_indices) if gpu_indices else list(names_by_id) + saw = False + for _i in indices: + name = names_by_id.get(_i) + if name is None: + continue # not visible -> skip (fail conservative) + saw = True + if not pattern.search(name): + return False + return saw + except Exception: + return False + + @staticmethod + def _effective_gpu_count(gpu_indices = None) -> int: + """GPUs llama-server will use: len(selection), else the visible CUDA + device count (None = every visible GPU). 0 on error so multi-GPU tuning + stays off when the count is unknown.""" + if gpu_indices is not None: + return len(gpu_indices) + try: + import torch + if hasattr(torch, "cuda") and torch.cuda.is_available(): + return torch.cuda.device_count() + except Exception: + return 0 + return 0 + + @staticmethod + def _apply_datacenter_env(env: dict, gpu_indices = None) -> bool: + """Inject DC llama.cpp tuning into env in place via setdefault (user + values win); return whether the box qualified. Opt out with + UNSLOTH_DISABLE_DC_TUNING=1; only datacenter NVIDIA parts qualify + (consumer/ROCm/CPU/error are a no-op). Sets GGML_CUDA_FORCE_CUBLAS_COMPUTE_32F + for any qualifying GPU (FP32 accum: ~0% cost on B200, real cost on GeForce), + plus GGML_CUDA_P2P + CUDA_SCALE_LAUNCH_QUEUES=4x for multi-GPU (+33-51% pp + tensor-split, +8-16% pipeline split on B200).""" + if os.environ.get("UNSLOTH_DISABLE_DC_TUNING") == "1": + return False + if not LlamaCppBackend._is_datacenter_gpu(gpu_indices): + return False + env.setdefault("GGML_CUDA_FORCE_CUBLAS_COMPUTE_32F", "1") + if LlamaCppBackend._effective_gpu_count(gpu_indices) > 1: + env.setdefault("GGML_CUDA_P2P", "1") + env.setdefault("CUDA_SCALE_LAUNCH_QUEUES", "4x") + return True + @staticmethod def _get_gpu_free_memory() -> list[tuple[int, int]]: """Query free memory per GPU. @@ -3406,6 +3505,14 @@ class LlamaCppBackend: env.setdefault("GGML_CUDA_ENABLE_UNIFIED_MEMORY", "1") logger.info("AMD unified-memory APU: set GGML_CUDA_ENABLE_UNIFIED_MEMORY=1") + # DC NVIDIA GPUs: FP32 accum (+ P2P / launch queues for multi-GPU). + # See _apply_datacenter_env; opt out with UNSLOTH_DISABLE_DC_TUNING=1. + if self._apply_datacenter_env(env, gpu_indices): + multi_gpu = self._effective_gpu_count(gpu_indices) > 1 + logger.info( + f"Data-center GPU detected: applied DC llama.cpp env tuning (multi_gpu={multi_gpu})" + ) + if sys.platform == "win32": # Ordering: see _build_windows_path_dirs. #5106. path_dirs = self._build_windows_path_dirs( diff --git a/studio/backend/tests/test_datacenter_gpu_tuning.py b/studio/backend/tests/test_datacenter_gpu_tuning.py new file mode 100644 index 0000000000..fd9b291e8a --- /dev/null +++ b/studio/backend/tests/test_datacenter_gpu_tuning.py @@ -0,0 +1,278 @@ +# SPDX-License-Identifier: AGPL-3.0-only +# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0 + +"""Data-center llama.cpp env tuning: FP32 accum (+ P2P / launch queues for +multi-GPU) must apply only to datacenter NVIDIA parts, never consumer GeForce, +AMD/ROCm, CPU or macOS. User values win; UNSLOTH_DISABLE_DC_TUNING=1 disables. +""" + +from __future__ import annotations + +import sys +import types + +import pytest + +from core.inference.llama_cpp import LlamaCppBackend + + +def _fake_torch( + names, + *, + hip = None, + cuda_ok = True, +): + """torch stub: version.hip, cuda.is_available/device_count, get_device_properties(i).name.""" + t = types.ModuleType("torch") + t.version = types.SimpleNamespace(hip = hip) + t.cuda = types.SimpleNamespace( + is_available = lambda: cuda_ok, + device_count = lambda: len(names), + get_device_properties = lambda i: types.SimpleNamespace(name = names[i]), + ) + return t + + +@pytest.fixture(autouse = True) +def _clear_cuda_visible_devices(monkeypatch): + """Detection reads CUDA_VISIBLE_DEVICES, so clear it by default (run unmasked, + physical id == ordinal) regardless of host; masked tests set it explicitly.""" + monkeypatch.delenv("CUDA_VISIBLE_DEVICES", raising = False) + + +# --------------------------------------------------------------------------- +# _is_datacenter_gpu +# --------------------------------------------------------------------------- + + +@pytest.mark.parametrize( + "names,expected", + [ + # Datacenter / professional parts. + (["NVIDIA A100-SXM4-80GB"], True), + (["NVIDIA A30"], True), + (["NVIDIA H100 80GB HBM3"], True), + (["NVIDIA H200"], True), + (["NVIDIA H800"], True), + (["NVIDIA GH200 480GB"], True), + (["NVIDIA B200"], True), + (["NVIDIA GB200"], True), + (["NVIDIA L40S"], True), + (["NVIDIA L4"], True), + (["NVIDIA RTX PRO 6000 Blackwell Server Edition"], True), + (["NVIDIA RTX 6000 Ada Generation"], True), + # Consumer GeForce: never. + (["NVIDIA GeForce RTX 4090"], False), + (["NVIDIA GeForce RTX 5090"], False), + (["NVIDIA GeForce RTX 3090"], False), + (["NVIDIA GeForce RTX 2080 Ti"], False), + (["NVIDIA GeForce GTX 1080"], False), + # Workstation/laptop: short markers must not match as substrings + # ("a100" in "A1000", "a30" in "A3000"). + (["NVIDIA RTX A1000 Laptop GPU"], False), + (["NVIDIA RTX A1000 6GB Laptop GPU"], False), + (["NVIDIA RTX A3000 Laptop GPU"], False), + # Homogeneous multi-DC: all must match. + (["NVIDIA B200", "NVIDIA B200"], True), + (["NVIDIA H100 80GB HBM3", "NVIDIA H100 80GB HBM3"], True), + # Mixed DC + consumer: non-DC, so tuning never lands on the GeForce. + (["NVIDIA B200", "NVIDIA GeForce RTX 4090"], False), + (["NVIDIA GeForce RTX 4090", "NVIDIA B200"], False), + ], +) +def test_is_datacenter_gpu(monkeypatch, names, expected): + monkeypatch.setitem(sys.modules, "torch", _fake_torch(names)) + assert LlamaCppBackend._is_datacenter_gpu() is expected + + +def test_is_datacenter_gpu_respects_selection(monkeypatch): + # A mixed box where only the DC GPU is selected -> True; only consumer -> False. + monkeypatch.setitem( + sys.modules, + "torch", + _fake_torch(["NVIDIA B200", "NVIDIA GeForce RTX 4090"]), + ) + assert LlamaCppBackend._is_datacenter_gpu([0]) is True + assert LlamaCppBackend._is_datacenter_gpu([1]) is False + assert LlamaCppBackend._is_datacenter_gpu([0, 1]) is False + + +def test_is_datacenter_gpu_out_of_range_indices_skipped(monkeypatch): + monkeypatch.setitem(sys.modules, "torch", _fake_torch(["NVIDIA B200"])) + # Out-of-range / negative indices are skipped; the one valid DC GPU still wins. + assert LlamaCppBackend._is_datacenter_gpu([0, 5, -1]) is True + # Only invalid indices -> nothing seen -> False (fail closed for the flag). + assert LlamaCppBackend._is_datacenter_gpu([5, 9]) is False + + +def test_is_datacenter_gpu_masked_host_physical_ids(monkeypatch): + # Mask 4,5,6,7 -> ordinals 0..3 == physical 4..7. PHYSICAL selection [4,5] + # must resolve, not index out of range (the pre-fix bug: 4 >= device_count). + monkeypatch.setenv("CUDA_VISIBLE_DEVICES", "4,5,6,7") + monkeypatch.setitem(sys.modules, "torch", _fake_torch(["NVIDIA B200"] * 4)) + assert LlamaCppBackend._is_datacenter_gpu([4, 5]) is True + assert LlamaCppBackend._is_datacenter_gpu([4, 5, 6, 7]) is True + assert LlamaCppBackend._is_datacenter_gpu(None) is True + assert LlamaCppBackend._is_datacenter_gpu([0, 1]) is False # not visible -> skip + + +def test_is_datacenter_gpu_masked_host_reordered(monkeypatch): + # Reordered mask preserves order: ordinal 0 -> physical 7, 1 -> 4, ... + monkeypatch.setenv("CUDA_VISIBLE_DEVICES", "7,4,5,6") + monkeypatch.setitem(sys.modules, "torch", _fake_torch(["NVIDIA H100 80GB HBM3"] * 4)) + assert LlamaCppBackend._is_datacenter_gpu([7, 4]) is True + + +def test_is_datacenter_gpu_masked_host_mixed_class(monkeypatch): + # Mask 4,5: physical 4 = GeForce, physical 5 = B200. Detection must follow the + # selected physical GPU, not a same-numbered ordinal. + monkeypatch.setenv("CUDA_VISIBLE_DEVICES", "4,5") + monkeypatch.setitem( + sys.modules, + "torch", + _fake_torch(["NVIDIA GeForce RTX 4090", "NVIDIA B200"]), + ) + assert LlamaCppBackend._is_datacenter_gpu([4]) is False + assert LlamaCppBackend._is_datacenter_gpu([5]) is True + assert LlamaCppBackend._is_datacenter_gpu([4, 5]) is False + + +def test_is_datacenter_gpu_unparsable_mask_falls_back(monkeypatch): + # Unparsable (UUID) mask falls back to physical id == ordinal (mirrors + # _get_gpu_free_memory), so ordinal lookup still classifies the device. + monkeypatch.setenv("CUDA_VISIBLE_DEVICES", "GPU-abcdef12") + monkeypatch.setitem(sys.modules, "torch", _fake_torch(["NVIDIA B200"])) + assert LlamaCppBackend._is_datacenter_gpu([0]) is True + + +def test_is_datacenter_gpu_rocm_is_false(monkeypatch): + # ROCm reuses torch.cuda.*; an MI300X must not qualify. + monkeypatch.setitem( + sys.modules, + "torch", + _fake_torch(["AMD Instinct MI300X"], hip = "6.2.0"), + ) + assert LlamaCppBackend._is_datacenter_gpu() is False + + +def test_is_datacenter_gpu_no_cuda_is_false(monkeypatch): + monkeypatch.setitem(sys.modules, "torch", _fake_torch([], cuda_ok = False)) + assert LlamaCppBackend._is_datacenter_gpu() is False + + +def test_is_datacenter_gpu_missing_torch_is_false(monkeypatch): + monkeypatch.setitem(sys.modules, "torch", None) + assert LlamaCppBackend._is_datacenter_gpu() is False + + +# --------------------------------------------------------------------------- +# _effective_gpu_count +# --------------------------------------------------------------------------- + + +def test_effective_gpu_count_explicit_selection(monkeypatch): + monkeypatch.setitem(sys.modules, "torch", _fake_torch(["NVIDIA B200"] * 4)) + assert LlamaCppBackend._effective_gpu_count([0]) == 1 + assert LlamaCppBackend._effective_gpu_count([0, 1, 2]) == 3 + + +def test_effective_gpu_count_none_uses_visible(monkeypatch): + # None -> visible device count. + monkeypatch.setitem(sys.modules, "torch", _fake_torch(["NVIDIA B200"] * 4)) + assert LlamaCppBackend._effective_gpu_count(None) == 4 + + +def test_effective_gpu_count_no_cuda_is_zero(monkeypatch): + monkeypatch.setitem(sys.modules, "torch", _fake_torch([], cuda_ok = False)) + assert LlamaCppBackend._effective_gpu_count(None) == 0 + + +def test_effective_gpu_count_missing_torch_is_zero(monkeypatch): + monkeypatch.setitem(sys.modules, "torch", None) + assert LlamaCppBackend._effective_gpu_count(None) == 0 + + +# --------------------------------------------------------------------------- +# _apply_datacenter_env (the env-injection decision) +# --------------------------------------------------------------------------- + + +def test_apply_env_single_dc_gpu_sets_only_fp32(monkeypatch): + monkeypatch.delenv("UNSLOTH_DISABLE_DC_TUNING", raising = False) + monkeypatch.setitem(sys.modules, "torch", _fake_torch(["NVIDIA B200"])) + env: dict = {} + assert LlamaCppBackend._apply_datacenter_env(env, [0]) is True + assert env == {"GGML_CUDA_FORCE_CUBLAS_COMPUTE_32F": "1"} + assert "GGML_CUDA_P2P" not in env # no multi-GPU flags on one GPU + assert "CUDA_SCALE_LAUNCH_QUEUES" not in env + + +def test_apply_env_multi_dc_gpu_sets_all(monkeypatch): + monkeypatch.delenv("UNSLOTH_DISABLE_DC_TUNING", raising = False) + monkeypatch.setitem(sys.modules, "torch", _fake_torch(["NVIDIA B200"] * 4)) + env: dict = {} + assert LlamaCppBackend._apply_datacenter_env(env, [0, 1]) is True + assert env["GGML_CUDA_FORCE_CUBLAS_COMPUTE_32F"] == "1" + assert env["GGML_CUDA_P2P"] == "1" + assert env["CUDA_SCALE_LAUNCH_QUEUES"] == "4x" + + +def test_apply_env_none_indices_uses_visible_count(monkeypatch): + # None on a 2x DC box -> multi-GPU flags applied. + monkeypatch.delenv("UNSLOTH_DISABLE_DC_TUNING", raising = False) + monkeypatch.setitem(sys.modules, "torch", _fake_torch(["NVIDIA H100", "NVIDIA H100"])) + env: dict = {} + assert LlamaCppBackend._apply_datacenter_env(env, None) is True + assert env["GGML_CUDA_P2P"] == "1" + assert env["CUDA_SCALE_LAUNCH_QUEUES"] == "4x" + + +def test_apply_env_consumer_gpu_is_noop(monkeypatch): + monkeypatch.delenv("UNSLOTH_DISABLE_DC_TUNING", raising = False) + monkeypatch.setitem(sys.modules, "torch", _fake_torch(["NVIDIA GeForce RTX 4090"] * 2)) + env: dict = {} + assert LlamaCppBackend._apply_datacenter_env(env, [0, 1]) is False + assert env == {} + + +def test_apply_env_user_value_wins(monkeypatch): + monkeypatch.delenv("UNSLOTH_DISABLE_DC_TUNING", raising = False) + monkeypatch.setitem(sys.modules, "torch", _fake_torch(["NVIDIA B200"] * 2)) + env = { + "GGML_CUDA_FORCE_CUBLAS_COMPUTE_32F": "0", # user explicitly disabled + "CUDA_SCALE_LAUNCH_QUEUES": "8x", # user override + } + assert LlamaCppBackend._apply_datacenter_env(env, [0, 1]) is True + # setdefault must not clobber user values; the unset one still defaults. + assert env["GGML_CUDA_FORCE_CUBLAS_COMPUTE_32F"] == "0" + assert env["CUDA_SCALE_LAUNCH_QUEUES"] == "8x" + assert env["GGML_CUDA_P2P"] == "1" + + +def test_apply_env_disable_flag_respected(monkeypatch): + monkeypatch.setenv("UNSLOTH_DISABLE_DC_TUNING", "1") + monkeypatch.setitem(sys.modules, "torch", _fake_torch(["NVIDIA B200"] * 2)) + env: dict = {} + assert LlamaCppBackend._apply_datacenter_env(env, [0, 1]) is False + assert env == {} + + +def test_apply_env_fail_open_on_detection_error(monkeypatch): + monkeypatch.delenv("UNSLOTH_DISABLE_DC_TUNING", raising = False) + monkeypatch.setitem(sys.modules, "torch", None) # detection raises -> False + env: dict = {} + assert LlamaCppBackend._apply_datacenter_env(env, [0]) is False + assert env == {} + + +def test_apply_env_masked_host_multi_dc(monkeypatch): + # End-to-end masked host (mask 4,5,6,7, physical selection [4,5]): pre-fix + # applied no tuning; now all three multi-GPU flags must be set. + monkeypatch.delenv("UNSLOTH_DISABLE_DC_TUNING", raising = False) + monkeypatch.setenv("CUDA_VISIBLE_DEVICES", "4,5,6,7") + monkeypatch.setitem(sys.modules, "torch", _fake_torch(["NVIDIA B200"] * 4)) + env: dict = {} + assert LlamaCppBackend._apply_datacenter_env(env, [4, 5]) is True + assert env["GGML_CUDA_FORCE_CUBLAS_COMPUTE_32F"] == "1" + assert env["GGML_CUDA_P2P"] == "1" + assert env["CUDA_SCALE_LAUNCH_QUEUES"] == "4x" From 6a0a62ef65f56510d6475f35690ae2a937c3d4b5 Mon Sep 17 00:00:00 2001 From: Daniel Han Date: Fri, 12 Jun 2026 02:43:55 -0700 Subject: [PATCH 15/81] Studio: drop the on-disk freshness cache after a llama.cpp update (#6234) The post-install path cleared only the in-memory freshness caches and then re-primed the 24h disk cache with a forced GitHub refresh. When that refresh cannot reach GitHub, latest_published_release falls back to the last-good disk value, so a still-fresh same-base mix tag cached before the swap (b9596-mix-aaa vs the just-installed b9596-mix-bbb) is replayed and the prebuilt reads as behind, surfacing a false update banner that points back at the build that was just replaced. Give reset_caches a drop_disk option and use it on the update path: with the disk cache gone, an offline post-install refresh leaves latest as None and the banner fails open (off) instead of lingering on the stale same-base value. The no-arg form stays in-memory only. Adds regression coverage for the drop, the default no-op, and the fail-open vs stale-replay contrast. --- .../backend/tests/test_llama_cpp_freshness.py | 87 +++++++++++++++++++ studio/backend/utils/llama_cpp_freshness.py | 19 +++- studio/backend/utils/llama_cpp_update.py | 11 ++- 3 files changed, 111 insertions(+), 6 deletions(-) diff --git a/studio/backend/tests/test_llama_cpp_freshness.py b/studio/backend/tests/test_llama_cpp_freshness.py index f8e4619ded..f90c4ba0e7 100644 --- a/studio/backend/tests/test_llama_cpp_freshness.py +++ b/studio/backend/tests/test_llama_cpp_freshness.py @@ -433,3 +433,90 @@ def test_fetch_latest_release_tag_uses_publish_time(monkeypatch): ] monkeypatch.setattr(urllib.request, "urlopen", lambda req, timeout = 5.0: _Resp(payload)) assert fr._fetch_latest_release_tag("unslothai/llama.cpp") == "b9596-mix-e6f2453" + + +# reset_caches(drop_disk=...) -- post-update stale same-base mix disk cache. + + +def _seed_disk_cache(tmp_path: Path, latest_tag: str) -> Path: + # Matches _cache_path_for under the fixture's stubbed _cache_dir. + cache_dir = tmp_path / ".freshness" + cache_dir.mkdir(exist_ok = True) + cache_file = cache_dir / "unslothai__llama.cpp.json" + cache_file.write_text(json.dumps({"fetched_at": time.time(), "latest_tag": latest_tag})) + return cache_file + + +def test_reset_caches_drop_disk_removes_disk_cache(tmp_path): + cache_file = _seed_disk_cache(tmp_path, "b9596-mix-aaa") + assert cache_file.exists() + fr.reset_caches(drop_disk = True) + assert not cache_file.exists() + + +def test_reset_caches_default_keeps_disk_cache(tmp_path): + # The no-arg form is in-memory only (its existing test-only contract); it + # must not delete the on-disk cache. + cache_file = _seed_disk_cache(tmp_path, "b9596-mix-aaa") + fr.reset_caches() + assert cache_file.exists() + + +def test_reset_caches_drop_disk_on_missing_dir_is_noop(tmp_path): + # Fresh machine, no cache dir yet: drop_disk must be a quiet no-op. + assert not (tmp_path / ".freshness").exists() + fr.reset_caches(drop_disk = True) # must not raise + + +def test_drop_disk_lets_banner_fail_open_after_same_base_mix_swap(monkeypatch, tmp_path): + # P2 #2: the disk cache holds a still-fresh same-base mix (b9596-mix-aaa) + # from before an update to a *different* same-base mix (b9596-mix-bbb). + # The post-install path drops the disk cache; if the forced refresh is then + # offline, latest reads as None and the banner fails open -- instead of + # replaying the stale b9596-mix-aaa and falsely reading "behind". + _seed_disk_cache(tmp_path, "b9596-mix-aaa") + install_dir = tmp_path / "llama.cpp" + _write_marker( + install_dir, + tag = "b9596", + release_tag = "b9596-mix-bbb", + installed_at_utc = (datetime.now(tz = timezone.utc) - timedelta(days = 5)) + .isoformat() + .replace("+00:00", "Z"), + ) + bin_path = _fake_binary(install_dir, layout = "root") + # GitHub unreachable for the rest of the test (the offline post-install + # refresh, and the later status check). + monkeypatch.setattr(fr, "_fetch_latest_release_tag", lambda repo, timeout = 5.0: None) + + fr.reset_caches(drop_disk = True) # exactly what the apply path now does + info = fr.check_prebuilt_freshness(str(bin_path)) + assert info["latest_tag"] is None + assert info["behind"] is False + assert info["stale"] is False + + +def test_in_memory_only_reset_replays_stale_same_base_mix(monkeypatch, tmp_path): + # Contrast/guard for the case above: an in-memory-only reset leaves the + # stale same-base mix on disk, so an offline check replays it and falsely + # reads behind/stale. This is exactly the failure drop_disk removes; if a + # future change makes the no-arg reset also clear disk, the apply-path call + # and this guard should be revisited together. + _seed_disk_cache(tmp_path, "b9596-mix-aaa") + install_dir = tmp_path / "llama.cpp" + _write_marker( + install_dir, + tag = "b9596", + release_tag = "b9596-mix-bbb", + installed_at_utc = (datetime.now(tz = timezone.utc) - timedelta(days = 5)) + .isoformat() + .replace("+00:00", "Z"), + ) + bin_path = _fake_binary(install_dir, layout = "root") + monkeypatch.setattr(fr, "_fetch_latest_release_tag", lambda repo, timeout = 5.0: None) + + fr.reset_caches() # in-memory only -> stale disk value survives + info = fr.check_prebuilt_freshness(str(bin_path)) + assert info["latest_tag"] == "b9596-mix-aaa" + assert info["behind"] is True + assert info["stale"] is True diff --git a/studio/backend/utils/llama_cpp_freshness.py b/studio/backend/utils/llama_cpp_freshness.py index f5fd745334..87d0d2ec01 100644 --- a/studio/backend/utils/llama_cpp_freshness.py +++ b/studio/backend/utils/llama_cpp_freshness.py @@ -301,7 +301,22 @@ def format_stale_warning(info: dict) -> str: ) -def reset_caches() -> None: - """Test-only: drop all in-memory caches.""" +def reset_caches(*, drop_disk: bool = False) -> None: + """Drop the in-memory freshness caches. The no-arg form is test-only. + + With ``drop_disk = True`` also delete the on-disk 24h release cache. Used by + the post-install/update path: in-memory clearing alone leaves the stale + same-base value on disk, so if the post-install GitHub refresh can't reach + the network, ``latest_published_release`` would replay that stale disk value + (see its last-good fallback) and the banner could linger. Dropping the disk + cache makes latest read as None in that offline case, so the banner fails + open (off) instead of pointing at the just-replaced build.""" _marker_cache.clear() _release_memo.clear() + if drop_disk: + import shutil + + # _cache_dir() is a dedicated freshness-only subdir; it is re-created on + # the next _save_disk_cache. ignore_errors so a missing/locked dir is a + # no-op rather than breaking an otherwise successful install. + shutil.rmtree(_cache_dir(), ignore_errors = True) diff --git a/studio/backend/utils/llama_cpp_update.py b/studio/backend/utils/llama_cpp_update.py index 654ade6cd4..6eb34ffe34 100644 --- a/studio/backend/utils/llama_cpp_update.py +++ b/studio/backend/utils/llama_cpp_update.py @@ -406,10 +406,13 @@ def _run_update(install_dir: Path, repo: str, asset: Optional[str], script: Path tail = "".join(tail_lines).strip()[-1500:] raise RuntimeError(f"installer exited {returncode}: {tail or 'no output'}") - # New UNSLOTH_PREBUILT_INFO.json is on disk; drop in-memory caches and - # re-prime the 24h disk freshness cache with the true newest, so the - # banner can't linger on a stale same-base value after the swap. - reset_caches() + # New UNSLOTH_PREBUILT_INFO.json is on disk; drop the in-memory AND the + # on-disk freshness caches, then re-prime the 24h disk cache with the + # true newest, so the banner can't linger on a stale same-base value + # after the swap. drop_disk matters when the refresh below can't reach + # GitHub: without it, latest_published_release would replay the stale + # disk value; with it, latest reads as None and the banner fails open. + reset_caches(drop_disk = True) try: latest_published_release(repo, force_refresh = True) except Exception as exc: # pragma: no cover - network defensive From e0d6674ff6f19498f6825f2ff2c38e8091674fd2 Mon Sep 17 00:00:00 2001 From: Daniel Han Date: Fri, 12 Jun 2026 02:54:26 -0700 Subject: [PATCH 16/81] Add RAG runtime deps to no-torch-runtime.txt (#6236) The --local / GGUF-only install resolves its Python deps from no-torch-runtime.txt, installed with --no-deps. That file was missing the RAG group that studio.txt declares (sqlite-vec, pymupdf, python-docx), so a fresh `unsloth studio` came up with RAG disabled: rag_db.py cannot import sqlite_vec and logs "RAG unavailable: sqlite-vec extension could not be loaded", and the knowledge-base routes return 503. python-docx was also absent, so DOCX ingestion failed. Add the three RAG store and document-parsing deps with the same pins as studio.txt so knowledge bases work out of the box on the no-torch path. sentence-transformers (dense embeddings) was already present. --- studio/backend/requirements/no-torch-runtime.txt | 6 ++++++ 1 file changed, 6 insertions(+) diff --git a/studio/backend/requirements/no-torch-runtime.txt b/studio/backend/requirements/no-torch-runtime.txt index 85294114b1..6efe91d448 100644 --- a/studio/backend/requirements/no-torch-runtime.txt +++ b/studio/backend/requirements/no-torch-runtime.txt @@ -68,3 +68,9 @@ trl>=0.18.2,!=0.19.0,<=0.24.0 sentence-transformers cut_cross_entropy pillow + +# RAG store + document parsing, mirroring studio.txt. Pinned here because +# this file installs --no-deps; without them Studio runs with RAG disabled. +sqlite-vec==0.1.9 +pymupdf==1.27.2.3 +python-docx==1.2.0 From 068b2c120fa389f93e3027def6f44bd4cc98f39e Mon Sep 17 00:00:00 2001 From: Michael Han <107991372+shimmyshimmer@users.noreply.github.com> Date: Fri, 12 Jun 2026 03:03:24 -0700 Subject: [PATCH 17/81] Studio: rounded rectangle hover states for menu items instead of pills (#6210) * Studio: use rounded rectangles for menu item hover states instead of pills Dropdown, select, and model picker items previously used fully rounded pill highlights. Switch them to an 11px rounded rectangle so hover and selected states match across the plus menu, profile menu, run settings, selects, and the model picker. Also add a small side gutter to the plus menu so item highlights sit slightly inset from the menu edge. * Studio: concentric menu corners, wider gutters, single-item pill menus Container radius now equals the item hover radius plus the side gutter (12px + 10px = 22px) so the curves run parallel. Menus with a single item render as fully rounded pills. The profile menu gets the same gutter and hover radius. Model picker rows go back to their original fully rounded hover. --- .../frontend/src/components/app-sidebar.tsx | 2 +- .../src/components/ui/dropdown-menu.tsx | 8 ++-- studio/frontend/src/components/ui/select.tsx | 2 +- studio/frontend/src/index.css | 40 +++++++++++++------ 4 files changed, 34 insertions(+), 18 deletions(-) diff --git a/studio/frontend/src/components/app-sidebar.tsx b/studio/frontend/src/components/app-sidebar.tsx index f1a200b09b..faa98abeea 100644 --- a/studio/frontend/src/components/app-sidebar.tsx +++ b/studio/frontend/src/components/app-sidebar.tsx @@ -1099,7 +1099,7 @@ export function AppSidebar() { side="top" align="center" sideOffset={8} - className="app-user-menu menu-soft-surface-up ring-0 w-[16rem] px-1.5 py-2.5 font-heading rounded-[20px] border-0" + className="app-user-menu menu-soft-surface-up ring-0 w-[16rem] px-2.5 py-2.5 font-heading rounded-[20px] border-0" > :nth-child(2))) { + border-radius: 9999px !important; + } + .unsloth-plus-menu[data-slot]:not(:has(> :nth-child(2))) + :is([data-slot="dropdown-menu-item"], [data-slot="dropdown-menu-sub-trigger"]) { + border-radius: 9999px; } .dark .unsloth-plus-menu[data-slot] { @@ -1319,14 +1334,15 @@ [data-slot="dropdown-menu-item"], [data-slot="dropdown-menu-sub-trigger"] ) { - @apply gap-3 pl-4 pr-3 py-2 text-[14px]; + @apply gap-3 pl-3 pr-3 py-2 text-[14px]; cursor: pointer; - /* Pin hover-box radius so dark matches light (same as the container). */ - border-radius: 1.1rem; + /* Pin hover-box radius so dark matches light (container radius minus the + side gutter keeps the curves concentric). */ + border-radius: 12px; } .unsloth-plus-menu [data-slot="dropdown-menu-label"] { - @apply pl-4 pr-3 py-1.5 text-[12px]; + @apply pl-3 pr-3 py-1.5 text-[12px]; } /* Active (green) items keep their primary text and icon color on hover. */ @@ -1370,8 +1386,8 @@ [data-slot="dropdown-menu-sub-trigger"] ) svg { - width: 1.05rem; - height: 1.05rem; + width: 1.15rem; + height: 1.15rem; } /* Destructive items keep red text and a red-tinted hover, not the grey one. */ From c773d45a2ead76167c5354154fbe78f02462ad27 Mon Sep 17 00:00:00 2001 From: Agnibha Mukherjee Date: Fri, 12 Jun 2026 15:37:04 +0530 Subject: [PATCH 18/81] docs: repository cleanup (#5617) * docs: small repository cleanup * docs: improve contribution guidelines --------- Co-authored-by: Agnibha007 Co-authored-by: Roland Tannous <115670425+rolandtannous@users.noreply.github.com> Co-authored-by: Lee Jackson <130007945+Imagineer99@users.noreply.github.com> --- CONTRIBUTING.md | 6 ++++++ 1 file changed, 6 insertions(+) diff --git a/CONTRIBUTING.md b/CONTRIBUTING.md index eb60a5a201..6eb8d1bc6e 100644 --- a/CONTRIBUTING.md +++ b/CONTRIBUTING.md @@ -27,3 +27,9 @@ Your support extends beyond code: Finally, please be mindful of our [Code of Conduct](https://github.com/unslothai/unsloth/blob/main/CODE_OF_CONDUCT.md) to ensure a welcoming and inclusive environment for everyone. Thank you so much for reading and we hope you have lots of fun using Unsloth! 🦥 + + +## Pull Request Guidelines +- Keep PRs focused on a single change +- Include a concise description and motivation +- Link related issues when applicable From 36ea9a9196938fcc18a5690ab0687295a1715669 Mon Sep 17 00:00:00 2001 From: Daniel Han Date: Fri, 12 Jun 2026 03:40:50 -0700 Subject: [PATCH 19/81] Run cross-platform parity test on Windows and macOS in CI (#6241) --- .../workflows/cross-platform-parity-ci.yml | 61 +++++++++++++++++++ 1 file changed, 61 insertions(+) create mode 100644 .github/workflows/cross-platform-parity-ci.yml diff --git a/.github/workflows/cross-platform-parity-ci.yml b/.github/workflows/cross-platform-parity-ci.yml new file mode 100644 index 0000000000..4632794587 --- /dev/null +++ b/.github/workflows/cross-platform-parity-ci.yml @@ -0,0 +1,61 @@ +# SPDX-License-Identifier: AGPL-3.0-only +# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. + +# Runs tests/python/test_cross_platform_parity.py on Windows and macOS. +# +# Why: that test is the guard that install.sh and install.ps1 stay in +# sync, but today it only runs on ubuntu-latest (auto-discovered by +# studio-backend-ci.yml's "Repo tests (CPU)" job). The test reads both +# installer scripts, and on Windows Path.read_text() defaults to the +# cp1252 locale encoding, so a non-cp1252 byte in install.sh (it already +# contains a U+274C) raises UnicodeDecodeError there even though Linux and +# macOS default to UTF-8. The reads were pinned to encoding="utf-8" in +# #6166; this job keeps that from silently regressing by exercising the +# test on the platforms it claims parity for. Pure pytest, no GPU, +# sub-second, so the matrix is cheap. + +name: Cross-platform parity + +on: + pull_request: + paths: + - 'install.sh' + - 'install.ps1' + - 'tests/python/test_cross_platform_parity.py' + - '.github/workflows/cross-platform-parity-ci.yml' + push: + branches: [main] + paths: + - 'install.sh' + - 'install.ps1' + - 'tests/python/test_cross_platform_parity.py' + - '.github/workflows/cross-platform-parity-ci.yml' + workflow_dispatch: + +concurrency: + group: ${{ github.workflow }}-${{ github.ref }} + cancel-in-progress: true + +permissions: + contents: read + +jobs: + parity: + name: parity (${{ matrix.os }}) + strategy: + fail-fast: false + matrix: + os: [windows-latest, macos-latest] + runs-on: ${{ matrix.os }} + timeout-minutes: 10 + steps: + - uses: actions/checkout@de0fac2e4500dabe0009e67214ff5f5447ce83dd # v6.0.2 + with: + persist-credentials: false + - uses: actions/setup-python@a309ff8b426b58ec0e2a45f0f869d46889d02405 # v6.2.0 + with: + python-version: '3.12' + cache: 'pip' + - run: python -m pip install -U pip pytest + - name: Cross-platform parity test + run: python -m pytest tests/python/test_cross_platform_parity.py -q From 6d206b488c46d8407336c7f762cd72ed3b9b687b Mon Sep 17 00:00:00 2001 From: Daniel Han Date: Fri, 12 Jun 2026 03:51:59 -0700 Subject: [PATCH 20/81] chore(studio/frontend): normalize line endings to LF (#6012) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit * chore(studio/frontend): normalize line endings to LF 45 source files under studio/frontend/ were committed with CRLF or mixed line endings while the rest of the repo and the JS/TS tooling assume LF. Add a scoped `studio/frontend/** text=auto eol=lf` rule to .gitattributes and run `git add --renormalize studio/frontend` so these files are stored with LF in the index. The rule is scoped to the frontend tree (not a repo-wide *.ts/*.tsx/... policy) so it cannot force LF on files elsewhere; text=auto leaves binary assets (logos, fonts) untouched. This commit is whitespace-only (CRLF -> LF) — no source content changed (verified with `git diff --ignore-cr-at-eol`). It is intentionally isolated so it can be listed in .git-blame-ignore-revs and skipped by reviewers and `git blame`. Co-Authored-By: Claude Opus 4.8 * chore: ignore the frontend LF-normalization commit in git blame Add .git-blame-ignore-revs listing the whitespace-only line-ending normalization commit so it doesn't pollute `git blame` output. GitHub applies this file automatically; locally run `git config blame.ignoreRevsFile .git-blame-ignore-revs`. Co-Authored-By: Claude Opus 4.8 --------- Co-authored-by: Claude Opus 4.8 --- .git-blame-ignore-revs | 8 + .gitattributes | 6 + studio/frontend/index.html | 26 +- .../frontend/public/hub/profile/logo/meta.svg | 36 +- .../public/provider-logos/misc/meta.svg | 36 +- .../frontend/src/components/ui/accordion.tsx | 190 ++--- .../src/components/ui/alert-dialog.tsx | 368 ++++----- studio/frontend/src/components/ui/alert.tsx | 152 ++-- .../src/components/ui/animated-shiny-text.tsx | 76 +- .../src/components/ui/aspect-ratio.tsx | 18 +- studio/frontend/src/components/ui/avatar.tsx | 220 ++--- studio/frontend/src/components/ui/badge.tsx | 102 +-- .../frontend/src/components/ui/breadcrumb.tsx | 246 +++--- .../frontend/src/components/ui/calendar.tsx | 468 +++++------ studio/frontend/src/components/ui/card.tsx | 200 ++--- studio/frontend/src/components/ui/chart.tsx | 718 ++++++++-------- .../frontend/src/components/ui/checkbox.tsx | 62 +- .../frontend/src/components/ui/combobox.tsx | 764 +++++++++--------- studio/frontend/src/components/ui/command.tsx | 414 +++++----- .../src/components/ui/context-menu.tsx | 528 ++++++------ .../src/components/ui/dropdown-menu.tsx | 558 ++++++------- studio/frontend/src/components/ui/field.tsx | 472 +++++------ .../frontend/src/components/ui/hover-card.tsx | 90 +-- .../src/components/ui/input-group.tsx | 306 +++---- studio/frontend/src/components/ui/input.tsx | 38 +- studio/frontend/src/components/ui/label.tsx | 48 +- .../frontend/src/components/ui/light-rays.tsx | 286 +++---- studio/frontend/src/components/ui/menubar.tsx | 562 ++++++------- .../src/components/ui/navigation-menu.tsx | 348 ++++---- .../frontend/src/components/ui/pagination.tsx | 276 +++---- studio/frontend/src/components/ui/popover.tsx | 184 ++--- .../frontend/src/components/ui/progress.tsx | 74 +- .../src/components/ui/radio-group.tsx | 96 +-- .../src/components/ui/scroll-area.tsx | 110 +-- .../frontend/src/components/ui/separator.tsx | 52 +- .../frontend/src/components/ui/skeleton.tsx | 26 +- studio/frontend/src/components/ui/sonner.tsx | 168 ++-- .../src/components/ui/sparkles-text.tsx | 308 +++---- studio/frontend/src/components/ui/switch.tsx | 62 +- studio/frontend/src/components/ui/table.tsx | 228 +++--- studio/frontend/src/components/ui/tabs.tsx | 268 +++--- .../frontend/src/components/ui/textarea.tsx | 10 +- .../src/components/ui/toggle-group.tsx | 178 ++-- studio/frontend/src/components/ui/toggle.tsx | 92 +-- .../inline/inline-category-badges.tsx | 150 ++-- studio/frontend/tsconfig.app.json | 62 +- studio/frontend/tsconfig.json | 26 +- studio/frontend/tsconfig.node.json | 52 +- 48 files changed, 4891 insertions(+), 4877 deletions(-) create mode 100644 .git-blame-ignore-revs diff --git a/.git-blame-ignore-revs b/.git-blame-ignore-revs new file mode 100644 index 0000000000..17d96cd0f5 --- /dev/null +++ b/.git-blame-ignore-revs @@ -0,0 +1,8 @@ +# Commits listed here are skipped by `git blame` so that bulk, whitespace-only +# changes don't obscure the real authorship of a line. +# +# GitHub honors this file automatically. To use it locally, run once: +# git config blame.ignoreRevsFile .git-blame-ignore-revs + +# chore(studio/frontend): normalize line endings to LF +c50b8ab910f5aa56dd7ae0022d2c7b96bfe3384a diff --git a/.gitattributes b/.gitattributes index 75fba5d6ab..5f04b5e9d1 100644 --- a/.gitattributes +++ b/.gitattributes @@ -5,3 +5,9 @@ # clone (core.autocrlf=true) rewrites them to CRLF, and the trailing \r breaks # them when run in WSL/Linux (e.g. `set -e` -> "set: Illegal option -"). *.sh text eol=lf + +# Normalize Studio frontend sources to LF. Scoped to the frontend tree (rather +# than repo-wide *.ts/*.tsx/... rules) so the policy can't force LF on files +# elsewhere. text=auto lets Git detect and leave binary assets (logos, fonts) +# untouched while text files (.ts/.tsx/.json/.html/.svg/...) are stored as LF. +studio/frontend/** text=auto eol=lf diff --git a/studio/frontend/index.html b/studio/frontend/index.html index 4f81ffd4ff..0fbb4eaeeb 100644 --- a/studio/frontend/index.html +++ b/studio/frontend/index.html @@ -1,16 +1,16 @@ - + - - - - - - Unsloth Studio - - -
- - - + + + + + + Unsloth Studio + + +
+ + + diff --git a/studio/frontend/public/hub/profile/logo/meta.svg b/studio/frontend/public/hub/profile/logo/meta.svg index 9fa656bd6b..fe3709aeea 100644 --- a/studio/frontend/public/hub/profile/logo/meta.svg +++ b/studio/frontend/public/hub/profile/logo/meta.svg @@ -1,19 +1,19 @@ - - -Logo of Meta Platforms -- Graphic created by Detmar Owen - - - - - - - - - - - - - - - + + +Logo of Meta Platforms -- Graphic created by Detmar Owen + + + + + + + + + + + + + + + \ No newline at end of file diff --git a/studio/frontend/public/provider-logos/misc/meta.svg b/studio/frontend/public/provider-logos/misc/meta.svg index 9fa656bd6b..fe3709aeea 100644 --- a/studio/frontend/public/provider-logos/misc/meta.svg +++ b/studio/frontend/public/provider-logos/misc/meta.svg @@ -1,19 +1,19 @@ - - -Logo of Meta Platforms -- Graphic created by Detmar Owen - - - - - - - - - - - - - - - + + +Logo of Meta Platforms -- Graphic created by Detmar Owen + + + + + + + + + + + + + + + \ No newline at end of file diff --git a/studio/frontend/src/components/ui/accordion.tsx b/studio/frontend/src/components/ui/accordion.tsx index 7754c78a11..35de233858 100644 --- a/studio/frontend/src/components/ui/accordion.tsx +++ b/studio/frontend/src/components/ui/accordion.tsx @@ -1,98 +1,98 @@ // SPDX-License-Identifier: AGPL-3.0-only // Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0 -"use client"; - -import { Accordion as AccordionPrimitive } from "radix-ui"; -import type * as React from "react"; - -import { cn } from "@/lib/utils"; -import { ArrowDown01Icon, ArrowUp01Icon } from "@hugeicons/core-free-icons"; -import { HugeiconsIcon } from "@hugeicons/react"; - -function Accordion({ - className, - ...props -}: React.ComponentProps) { - return ( - - ); -} - -function AccordionItem({ - className, - ...props -}: React.ComponentProps) { - return ( - - ); -} - -function AccordionTrigger({ - className, - children, - ...props -}: React.ComponentProps) { - return ( - - - {children} - - - - - ); -} - -function AccordionContent({ - className, - children, - ...props -}: React.ComponentProps) { - return ( - -
- {children} -
- - ); -} - -export { Accordion, AccordionItem, AccordionTrigger, AccordionContent }; +"use client"; + +import { Accordion as AccordionPrimitive } from "radix-ui"; +import type * as React from "react"; + +import { cn } from "@/lib/utils"; +import { ArrowDown01Icon, ArrowUp01Icon } from "@hugeicons/core-free-icons"; +import { HugeiconsIcon } from "@hugeicons/react"; + +function Accordion({ + className, + ...props +}: React.ComponentProps) { + return ( + + ); +} + +function AccordionItem({ + className, + ...props +}: React.ComponentProps) { + return ( + + ); +} + +function AccordionTrigger({ + className, + children, + ...props +}: React.ComponentProps) { + return ( + + + {children} + + + + + ); +} + +function AccordionContent({ + className, + children, + ...props +}: React.ComponentProps) { + return ( + +
+ {children} +
+ + ); +} + +export { Accordion, AccordionItem, AccordionTrigger, AccordionContent }; diff --git a/studio/frontend/src/components/ui/alert-dialog.tsx b/studio/frontend/src/components/ui/alert-dialog.tsx index f5c1dbacca..97f4be7f44 100644 --- a/studio/frontend/src/components/ui/alert-dialog.tsx +++ b/studio/frontend/src/components/ui/alert-dialog.tsx @@ -1,50 +1,50 @@ // SPDX-License-Identifier: AGPL-3.0-only // Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0 -import { AlertDialog as AlertDialogPrimitive } from "radix-ui"; -import type * as React from "react"; - -import { Button } from "@/components/ui/button"; -import { cn } from "@/lib/utils"; - -function AlertDialog({ - ...props -}: React.ComponentProps) { - return ; -} - -function AlertDialogTrigger({ - ...props -}: React.ComponentProps) { - return ( - - ); -} - -function AlertDialogPortal({ - ...props -}: React.ComponentProps) { - return ( - - ); -} - -function AlertDialogOverlay({ - className, - ...props -}: React.ComponentProps) { - return ( - - ); -} - +import { AlertDialog as AlertDialogPrimitive } from "radix-ui"; +import type * as React from "react"; + +import { Button } from "@/components/ui/button"; +import { cn } from "@/lib/utils"; + +function AlertDialog({ + ...props +}: React.ComponentProps) { + return ; +} + +function AlertDialogTrigger({ + ...props +}: React.ComponentProps) { + return ( + + ); +} + +function AlertDialogPortal({ + ...props +}: React.ComponentProps) { + return ( + + ); +} + +function AlertDialogOverlay({ + className, + ...props +}: React.ComponentProps) { + return ( + + ); +} + function AlertDialogContent({ className, size = "default", @@ -60,143 +60,143 @@ function AlertDialogContent({ - - ); -} - -function AlertDialogHeader({ - className, - ...props -}: React.ComponentProps<"div">) { - return ( -
- ); -} - -function AlertDialogFooter({ - className, - ...props -}: React.ComponentProps<"div">) { - return ( -
- ); -} - -function AlertDialogMedia({ - className, - ...props -}: React.ComponentProps<"div">) { - return ( -
- ); -} - -function AlertDialogTitle({ - className, - ...props -}: React.ComponentProps) { - return ( - - ); -} - -function AlertDialogDescription({ - className, - ...props -}: React.ComponentProps) { - return ( - - ); -} - -function AlertDialogAction({ - className, - variant = "default", - size = "default", - ...props -}: React.ComponentProps & - Pick, "variant" | "size">) { - return ( - - ); -} - -function AlertDialogCancel({ - className, - variant = "outline", - size = "default", - ...props -}: React.ComponentProps & - Pick, "variant" | "size">) { - return ( - - ); -} - -export { - AlertDialog, - AlertDialogAction, - AlertDialogCancel, - AlertDialogContent, - AlertDialogDescription, - AlertDialogFooter, - AlertDialogHeader, - AlertDialogMedia, - AlertDialogOverlay, - AlertDialogPortal, - AlertDialogTitle, - AlertDialogTrigger, -}; + className={cn( + "data-open:animate-in data-closed:animate-out data-closed:fade-out-0 data-open:fade-in-0 data-closed:zoom-out-95 data-open:zoom-in-95 bg-background ring-foreground/5 gap-6 rounded-4xl p-6 ring-1 duration-100 data-[size=default]:max-w-xs data-[size=sm]:max-w-xs data-[size=default]:sm:max-w-md group/alert-dialog-content fixed top-1/2 left-1/2 z-50 grid w-full -translate-x-1/2 -translate-y-1/2 outline-none", + className, + )} + {...props} + /> + + ); +} + +function AlertDialogHeader({ + className, + ...props +}: React.ComponentProps<"div">) { + return ( +
+ ); +} + +function AlertDialogFooter({ + className, + ...props +}: React.ComponentProps<"div">) { + return ( +
+ ); +} + +function AlertDialogMedia({ + className, + ...props +}: React.ComponentProps<"div">) { + return ( +
+ ); +} + +function AlertDialogTitle({ + className, + ...props +}: React.ComponentProps) { + return ( + + ); +} + +function AlertDialogDescription({ + className, + ...props +}: React.ComponentProps) { + return ( + + ); +} + +function AlertDialogAction({ + className, + variant = "default", + size = "default", + ...props +}: React.ComponentProps & + Pick, "variant" | "size">) { + return ( + + ); +} + +function AlertDialogCancel({ + className, + variant = "outline", + size = "default", + ...props +}: React.ComponentProps & + Pick, "variant" | "size">) { + return ( + + ); +} + +export { + AlertDialog, + AlertDialogAction, + AlertDialogCancel, + AlertDialogContent, + AlertDialogDescription, + AlertDialogFooter, + AlertDialogHeader, + AlertDialogMedia, + AlertDialogOverlay, + AlertDialogPortal, + AlertDialogTitle, + AlertDialogTrigger, +}; diff --git a/studio/frontend/src/components/ui/alert.tsx b/studio/frontend/src/components/ui/alert.tsx index a4a5f4c4b7..094b607d9a 100644 --- a/studio/frontend/src/components/ui/alert.tsx +++ b/studio/frontend/src/components/ui/alert.tsx @@ -1,79 +1,79 @@ // SPDX-License-Identifier: AGPL-3.0-only // Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0 -import { type VariantProps, cva } from "class-variance-authority"; -import type * as React from "react"; - -import { cn } from "@/lib/utils"; - -const alertVariants = cva( - "grid gap-0.5 rounded-lg border px-4 py-3 text-left text-sm has-data-[slot=alert-action]:relative has-data-[slot=alert-action]:pr-18 has-[>svg]:grid-cols-[auto_1fr] has-[>svg]:gap-x-2.5 *:[svg]:row-span-2 *:[svg]:translate-y-0.5 *:[svg]:text-current *:[svg:not([class*='size-'])]:size-4 w-full relative group/alert", - { - variants: { - variant: { - default: "bg-card text-card-foreground", - destructive: - "text-destructive bg-card *:data-[slot=alert-description]:text-destructive/90 *:[svg]:text-current", - }, - }, - defaultVariants: { - variant: "default", - }, - }, -); - -function Alert({ - className, - variant, - ...props -}: React.ComponentProps<"div"> & VariantProps) { - return ( -
- ); -} - -function AlertTitle({ className, ...props }: React.ComponentProps<"div">) { - return ( -
svg]/alert:col-start-2 [&_a]:hover:text-foreground [&_a]:underline [&_a]:underline-offset-3", - className, - )} - {...props} - /> - ); -} - -function AlertDescription({ - className, - ...props -}: React.ComponentProps<"div">) { - return ( -
- ); -} - -function AlertAction({ className, ...props }: React.ComponentProps<"div">) { - return ( -
- ); -} - -export { Alert, AlertTitle, AlertDescription, AlertAction }; +import { type VariantProps, cva } from "class-variance-authority"; +import type * as React from "react"; + +import { cn } from "@/lib/utils"; + +const alertVariants = cva( + "grid gap-0.5 rounded-lg border px-4 py-3 text-left text-sm has-data-[slot=alert-action]:relative has-data-[slot=alert-action]:pr-18 has-[>svg]:grid-cols-[auto_1fr] has-[>svg]:gap-x-2.5 *:[svg]:row-span-2 *:[svg]:translate-y-0.5 *:[svg]:text-current *:[svg:not([class*='size-'])]:size-4 w-full relative group/alert", + { + variants: { + variant: { + default: "bg-card text-card-foreground", + destructive: + "text-destructive bg-card *:data-[slot=alert-description]:text-destructive/90 *:[svg]:text-current", + }, + }, + defaultVariants: { + variant: "default", + }, + }, +); + +function Alert({ + className, + variant, + ...props +}: React.ComponentProps<"div"> & VariantProps) { + return ( +
+ ); +} + +function AlertTitle({ className, ...props }: React.ComponentProps<"div">) { + return ( +
svg]/alert:col-start-2 [&_a]:hover:text-foreground [&_a]:underline [&_a]:underline-offset-3", + className, + )} + {...props} + /> + ); +} + +function AlertDescription({ + className, + ...props +}: React.ComponentProps<"div">) { + return ( +
+ ); +} + +function AlertAction({ className, ...props }: React.ComponentProps<"div">) { + return ( +
+ ); +} + +export { Alert, AlertTitle, AlertDescription, AlertAction }; diff --git a/studio/frontend/src/components/ui/animated-shiny-text.tsx b/studio/frontend/src/components/ui/animated-shiny-text.tsx index 4c650f1003..8d366ca3d6 100644 --- a/studio/frontend/src/components/ui/animated-shiny-text.tsx +++ b/studio/frontend/src/components/ui/animated-shiny-text.tsx @@ -1,41 +1,41 @@ // SPDX-License-Identifier: AGPL-3.0-only // Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0 -import type { ComponentPropsWithoutRef, CSSProperties, FC } from "react" - -import { cn } from "@/lib/utils" - -export interface AnimatedShinyTextProps extends ComponentPropsWithoutRef<"span"> { - shimmerWidth?: number -} - -export const AnimatedShinyText: FC = ({ - children, - className, - shimmerWidth = 100, - ...props -}) => { - return ( - - {children} - - ) -} +import type { ComponentPropsWithoutRef, CSSProperties, FC } from "react" + +import { cn } from "@/lib/utils" + +export interface AnimatedShinyTextProps extends ComponentPropsWithoutRef<"span"> { + shimmerWidth?: number +} + +export const AnimatedShinyText: FC = ({ + children, + className, + shimmerWidth = 100, + ...props +}) => { + return ( + + {children} + + ) +} diff --git a/studio/frontend/src/components/ui/aspect-ratio.tsx b/studio/frontend/src/components/ui/aspect-ratio.tsx index cb605f01eb..2471f4333d 100644 --- a/studio/frontend/src/components/ui/aspect-ratio.tsx +++ b/studio/frontend/src/components/ui/aspect-ratio.tsx @@ -1,12 +1,12 @@ // SPDX-License-Identifier: AGPL-3.0-only // Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0 -import { AspectRatio as AspectRatioPrimitive } from "radix-ui"; - -function AspectRatio({ - ...props -}: React.ComponentProps) { - return ; -} - -export { AspectRatio }; +import { AspectRatio as AspectRatioPrimitive } from "radix-ui"; + +function AspectRatio({ + ...props +}: React.ComponentProps) { + return ; +} + +export { AspectRatio }; diff --git a/studio/frontend/src/components/ui/avatar.tsx b/studio/frontend/src/components/ui/avatar.tsx index 2250bb849a..31262b32f7 100644 --- a/studio/frontend/src/components/ui/avatar.tsx +++ b/studio/frontend/src/components/ui/avatar.tsx @@ -1,113 +1,113 @@ // SPDX-License-Identifier: AGPL-3.0-only // Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0 -import { Avatar as AvatarPrimitive } from "radix-ui"; -import type * as React from "react"; - -import { cn } from "@/lib/utils"; - -function Avatar({ - className, - size = "default", - ...props -}: React.ComponentProps & { - size?: "default" | "sm" | "lg"; -}) { - return ( - - ); -} - -function AvatarImage({ - className, - ...props -}: React.ComponentProps) { - return ( - - ); -} - -function AvatarFallback({ - className, - ...props -}: React.ComponentProps) { - return ( - - ); -} - -function AvatarBadge({ className, ...props }: React.ComponentProps<"span">) { - return ( - svg]:hidden", - "group-data-[size=default]/avatar:size-2.5 group-data-[size=default]/avatar:[&>svg]:size-2", - "group-data-[size=lg]/avatar:size-3 group-data-[size=lg]/avatar:[&>svg]:size-2", - className, - )} - {...props} - /> - ); -} - -function AvatarGroup({ className, ...props }: React.ComponentProps<"div">) { - return ( -
- ); -} - -function AvatarGroupCount({ - className, - ...props -}: React.ComponentProps<"div">) { - return ( -
svg]:size-4 group-has-data-[size=lg]/avatar-group:[&>svg]:size-5 group-has-data-[size=sm]/avatar-group:[&>svg]:size-3 ring-background relative flex shrink-0 items-center justify-center ring-2", - className, - )} - {...props} - /> - ); -} - -export { - Avatar, - AvatarImage, - AvatarFallback, - AvatarGroup, - AvatarGroupCount, - AvatarBadge, -}; +import { Avatar as AvatarPrimitive } from "radix-ui"; +import type * as React from "react"; + +import { cn } from "@/lib/utils"; + +function Avatar({ + className, + size = "default", + ...props +}: React.ComponentProps & { + size?: "default" | "sm" | "lg"; +}) { + return ( + + ); +} + +function AvatarImage({ + className, + ...props +}: React.ComponentProps) { + return ( + + ); +} + +function AvatarFallback({ + className, + ...props +}: React.ComponentProps) { + return ( + + ); +} + +function AvatarBadge({ className, ...props }: React.ComponentProps<"span">) { + return ( + svg]:hidden", + "group-data-[size=default]/avatar:size-2.5 group-data-[size=default]/avatar:[&>svg]:size-2", + "group-data-[size=lg]/avatar:size-3 group-data-[size=lg]/avatar:[&>svg]:size-2", + className, + )} + {...props} + /> + ); +} + +function AvatarGroup({ className, ...props }: React.ComponentProps<"div">) { + return ( +
+ ); +} + +function AvatarGroupCount({ + className, + ...props +}: React.ComponentProps<"div">) { + return ( +
svg]:size-4 group-has-data-[size=lg]/avatar-group:[&>svg]:size-5 group-has-data-[size=sm]/avatar-group:[&>svg]:size-3 ring-background relative flex shrink-0 items-center justify-center ring-2", + className, + )} + {...props} + /> + ); +} + +export { + Avatar, + AvatarImage, + AvatarFallback, + AvatarGroup, + AvatarGroupCount, + AvatarBadge, +}; diff --git a/studio/frontend/src/components/ui/badge.tsx b/studio/frontend/src/components/ui/badge.tsx index 3951ae9de0..0f2f334986 100644 --- a/studio/frontend/src/components/ui/badge.tsx +++ b/studio/frontend/src/components/ui/badge.tsx @@ -1,54 +1,54 @@ // SPDX-License-Identifier: AGPL-3.0-only // Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0 -/* eslint-disable react-refresh/only-export-components */ - -import { type VariantProps, cva } from "class-variance-authority"; -import { Slot } from "radix-ui"; -import type * as React from "react"; - -import { cn } from "@/lib/utils"; - -export const badgeVariants = cva( - "h-5 gap-1 rounded-4xl border border-transparent px-2 py-0.5 text-xs font-medium transition-all has-data-[icon=inline-end]:pr-1.5 has-data-[icon=inline-start]:pl-1.5 [&>svg]:size-3! inline-flex items-center justify-center w-fit whitespace-nowrap shrink-0 [&>svg]:pointer-events-none focus-visible:border-ring focus-visible:ring-ring/50 focus-visible:ring-[3px] aria-invalid:ring-destructive/20 dark:aria-invalid:ring-destructive/40 aria-invalid:border-destructive overflow-hidden group/badge", - { - variants: { - variant: { - default: "bg-primary text-primary-foreground [a]:hover:bg-primary/80", - secondary: - "bg-secondary text-secondary-foreground [a]:hover:bg-secondary/80", - destructive: - "bg-destructive/10 [a]:hover:bg-destructive/20 focus-visible:ring-destructive/20 dark:focus-visible:ring-destructive/40 text-destructive dark:bg-destructive/20", - outline: - "border-border text-foreground [a]:hover:bg-muted [a]:hover:text-muted-foreground bg-input/30", - ghost: - "hover:bg-muted hover:text-muted-foreground dark:hover:bg-muted/50", - link: "text-primary underline-offset-4 hover:underline", - }, - }, - defaultVariants: { - variant: "default", - }, - }, -); - -export function Badge({ - className, - variant = "default", - asChild = false, - ...props -}: React.ComponentProps<"span"> & - VariantProps & { - asChild?: boolean; - }): React.ReactElement { - const Comp = asChild ? Slot.Root : "span"; - - return ( - - ); -} +/* eslint-disable react-refresh/only-export-components */ + +import { type VariantProps, cva } from "class-variance-authority"; +import { Slot } from "radix-ui"; +import type * as React from "react"; + +import { cn } from "@/lib/utils"; + +export const badgeVariants = cva( + "h-5 gap-1 rounded-4xl border border-transparent px-2 py-0.5 text-xs font-medium transition-all has-data-[icon=inline-end]:pr-1.5 has-data-[icon=inline-start]:pl-1.5 [&>svg]:size-3! inline-flex items-center justify-center w-fit whitespace-nowrap shrink-0 [&>svg]:pointer-events-none focus-visible:border-ring focus-visible:ring-ring/50 focus-visible:ring-[3px] aria-invalid:ring-destructive/20 dark:aria-invalid:ring-destructive/40 aria-invalid:border-destructive overflow-hidden group/badge", + { + variants: { + variant: { + default: "bg-primary text-primary-foreground [a]:hover:bg-primary/80", + secondary: + "bg-secondary text-secondary-foreground [a]:hover:bg-secondary/80", + destructive: + "bg-destructive/10 [a]:hover:bg-destructive/20 focus-visible:ring-destructive/20 dark:focus-visible:ring-destructive/40 text-destructive dark:bg-destructive/20", + outline: + "border-border text-foreground [a]:hover:bg-muted [a]:hover:text-muted-foreground bg-input/30", + ghost: + "hover:bg-muted hover:text-muted-foreground dark:hover:bg-muted/50", + link: "text-primary underline-offset-4 hover:underline", + }, + }, + defaultVariants: { + variant: "default", + }, + }, +); + +export function Badge({ + className, + variant = "default", + asChild = false, + ...props +}: React.ComponentProps<"span"> & + VariantProps & { + asChild?: boolean; + }): React.ReactElement { + const Comp = asChild ? Slot.Root : "span"; + + return ( + + ); +} diff --git a/studio/frontend/src/components/ui/breadcrumb.tsx b/studio/frontend/src/components/ui/breadcrumb.tsx index dc026994ce..a2dad8783f 100644 --- a/studio/frontend/src/components/ui/breadcrumb.tsx +++ b/studio/frontend/src/components/ui/breadcrumb.tsx @@ -1,126 +1,126 @@ // SPDX-License-Identifier: AGPL-3.0-only // Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0 -import { Slot } from "radix-ui"; -import type * as React from "react"; - -import { cn } from "@/lib/utils"; -import { - ArrowRight01Icon, - MoreHorizontalCircle01Icon, -} from "@hugeicons/core-free-icons"; -import { HugeiconsIcon } from "@hugeicons/react"; - -function Breadcrumb({ className, ...props }: React.ComponentProps<"nav">) { - return ( -