From 3927ec0a6875ec3b596a6828a556902f2e1f6352 Mon Sep 17 00:00:00 2001 From: Daniel Han Date: Tue, 5 May 2026 18:48:29 -0700 Subject: [PATCH 001/159] Bump installer floor to 2026.5.2 (#5297) Bumps the unsloth>= install floor in install.sh and install.ps1 from 2026.5.1 to 2026.5.2 so fresh curl/iwr installs pull the just-released PyPI version that ships PR #5296: Studio chat history and image attachments work again with newer @assistant-ui/react. --- install.ps1 | 10 +++++----- install.sh | 10 +++++----- 2 files changed, 10 insertions(+), 10 deletions(-) diff --git a/install.ps1 b/install.ps1 index a02cef0f6a..d5db1785a1 100644 --- a/install.ps1 +++ b/install.ps1 @@ -1068,7 +1068,7 @@ shell.Run cmd, 0, False if ($SkipTorch) { # No-torch: install unsloth + unsloth-zoo with --no-deps, then # runtime deps (typer, safetensors, transformers, etc.) with --no-deps. - $baseInstallExit = Invoke-InstallCommand { uv pip install --python $VenvPython --no-deps --reinstall-package unsloth --reinstall-package unsloth-zoo "unsloth>=2026.5.1" unsloth-zoo } + $baseInstallExit = Invoke-InstallCommand { uv pip install --python $VenvPython --no-deps --reinstall-package unsloth --reinstall-package unsloth-zoo "unsloth>=2026.5.2" unsloth-zoo } if ($baseInstallExit -eq 0) { $NoTorchReq = Find-NoTorchRuntimeFile if ($NoTorchReq) { @@ -1076,7 +1076,7 @@ shell.Run cmd, 0, False } } } else { - $baseInstallExit = Invoke-InstallCommand { uv pip install --python $VenvPython --reinstall-package unsloth --reinstall-package unsloth-zoo "unsloth>=2026.5.1" unsloth-zoo } + $baseInstallExit = Invoke-InstallCommand { uv pip install --python $VenvPython --reinstall-package unsloth --reinstall-package unsloth-zoo "unsloth>=2026.5.2" unsloth-zoo } } if ($baseInstallExit -ne 0) { Write-Host "[ERROR] Failed to install unsloth (exit code $baseInstallExit)" -ForegroundColor Red @@ -1114,7 +1114,7 @@ shell.Run cmd, 0, False if ($SkipTorch) { # No-torch: install unsloth + unsloth-zoo with --no-deps, then # runtime deps (typer, safetensors, transformers, etc.) with --no-deps. - $baseInstallExit = Invoke-InstallCommand { uv pip install --python $VenvPython --no-deps --upgrade-package unsloth --upgrade-package unsloth-zoo "unsloth>=2026.5.1" unsloth-zoo } + $baseInstallExit = Invoke-InstallCommand { uv pip install --python $VenvPython --no-deps --upgrade-package unsloth --upgrade-package unsloth-zoo "unsloth>=2026.5.2" unsloth-zoo } if ($baseInstallExit -eq 0) { $NoTorchReq = Find-NoTorchRuntimeFile if ($NoTorchReq) { @@ -1122,7 +1122,7 @@ shell.Run cmd, 0, False } } } elseif ($StudioLocalInstall) { - $baseInstallExit = Invoke-InstallCommand { uv pip install --python $VenvPython --upgrade-package unsloth "unsloth>=2026.5.1" unsloth-zoo } + $baseInstallExit = Invoke-InstallCommand { uv pip install --python $VenvPython --upgrade-package unsloth "unsloth>=2026.5.2" unsloth-zoo } } else { $baseInstallExit = Invoke-InstallCommand { uv pip install --python $VenvPython --upgrade-package unsloth -- "$PackageName" } } @@ -1150,7 +1150,7 @@ shell.Run cmd, 0, False Write-TauriLog "STEP" "Installing unsloth" substep "installing unsloth (this may take a few minutes)..." if ($StudioLocalInstall) { - $baseInstallExit = Invoke-InstallCommand { uv pip install --python $VenvPython unsloth-zoo "unsloth>=2026.5.1" --torch-backend=auto } + $baseInstallExit = Invoke-InstallCommand { uv pip install --python $VenvPython unsloth-zoo "unsloth>=2026.5.2" --torch-backend=auto } if ($baseInstallExit -ne 0) { Write-Host "[ERROR] Failed to install unsloth (exit code $baseInstallExit)" -ForegroundColor Red return (Exit-InstallFailure "Failed to install unsloth (exit code $baseInstallExit)" $baseInstallExit) diff --git a/install.sh b/install.sh index fc8bd27fb8..1d21117d16 100755 --- a/install.sh +++ b/install.sh @@ -1486,7 +1486,7 @@ if [ "$_MIGRATED" = true ]; then # to prevent transitive torch resolution. run_install_cmd "install unsloth (migrated no-torch)" uv pip install --python "$_VENV_PY" --no-deps \ --reinstall-package unsloth --reinstall-package unsloth-zoo \ - "unsloth>=2026.5.1" unsloth-zoo + "unsloth>=2026.5.2" unsloth-zoo _NO_TORCH_RT="$(_find_no_torch_runtime)" if [ -n "$_NO_TORCH_RT" ]; then run_install_cmd "install no-torch runtime deps" uv pip install --python "$_VENV_PY" --no-deps -r "$_NO_TORCH_RT" @@ -1494,7 +1494,7 @@ if [ "$_MIGRATED" = true ]; then else run_install_cmd "install unsloth (migrated)" uv pip install --python "$_VENV_PY" \ --reinstall-package unsloth --reinstall-package unsloth-zoo \ - "unsloth>=2026.5.1" unsloth-zoo + "unsloth>=2026.5.2" unsloth-zoo fi if [ "$STUDIO_LOCAL_INSTALL" = true ]; then substep "overlaying local repo (editable)..." @@ -1662,7 +1662,7 @@ elif [ -n "$TORCH_INDEX_URL" ]; then # runtime deps (typer, safetensors, transformers, etc.) with --no-deps. run_install_cmd "install unsloth (no-torch)" uv pip install --python "$_VENV_PY" --no-deps \ --upgrade-package unsloth --upgrade-package unsloth-zoo \ - "unsloth>=2026.5.1" unsloth-zoo + "unsloth>=2026.5.2" unsloth-zoo _NO_TORCH_RT="$(_find_no_torch_runtime)" if [ -n "$_NO_TORCH_RT" ]; then run_install_cmd "install no-torch runtime deps" uv pip install --python "$_VENV_PY" --no-deps -r "$_NO_TORCH_RT" @@ -1677,7 +1677,7 @@ elif [ -n "$TORCH_INDEX_URL" ]; then fi elif [ "$STUDIO_LOCAL_INSTALL" = true ]; then run_install_cmd "install unsloth (local)" uv pip install --python "$_VENV_PY" \ - --upgrade-package unsloth "unsloth>=2026.5.1" unsloth-zoo + --upgrade-package unsloth "unsloth>=2026.5.2" unsloth-zoo substep "overlaying local repo (editable)..." run_install_cmd "overlay local repo" uv pip install --python "$_VENV_PY" -e "$_REPO_ROOT" --no-deps substep "overlaying unsloth-zoo from git main..." @@ -1709,7 +1709,7 @@ else tauri_log "STEP" "Installing Unsloth" substep "installing unsloth (this may take a few minutes)..." if [ "$STUDIO_LOCAL_INSTALL" = true ]; then - run_install_cmd "install unsloth (auto torch backend)" uv pip install --python "$_VENV_PY" unsloth-zoo "unsloth>=2026.5.1" --torch-backend=auto + run_install_cmd "install unsloth (auto torch backend)" uv pip install --python "$_VENV_PY" unsloth-zoo "unsloth>=2026.5.2" --torch-backend=auto substep "overlaying local repo (editable)..." run_install_cmd "overlay local repo" uv pip install --python "$_VENV_PY" -e "$_REPO_ROOT" --no-deps substep "overlaying unsloth-zoo from git main..." From 7be10852cbd415f582b6ae79ccd85e2a01ae14a5 Mon Sep 17 00:00:00 2001 From: Daniel Han Date: Tue, 5 May 2026 23:17:40 -0700 Subject: [PATCH 002/159] install: support STUDIO_HOME / UNSLOTH_STUDIO_HOME for custom install paths (#5190) * install: support STUDIO_HOME / UNSLOTH_STUDIO_HOME for custom install paths Currently install.sh and install.ps1 hardcode all install paths off $HOME / $env:USERPROFILE with no env-var fallback. This blocks workspace-isolated installs (CI sandboxes, per-PR test environments, multi-tenant boxes) unless the entire HOME / USERPROFILE is faked, which also relocates ~/.gitconfig, ~/.ssh, and other unrelated state. Add an opt-in env-var override that does only what is needed. Resolution priority (highest first): 1. HOME / USERPROFILE explicitly redirected vs the password-database default. Detected via getent (Linux), dscl (macOS), or [Environment]::GetFolderPath (Windows). Best-effort: when the detection mechanism is unavailable the check is skipped and we fall through to step 2. 2. UNSLOTH_STUDIO_HOME, if set. 3. STUDIO_HOME, if set (alias for convenience; the variable name already matches the internal var install.sh sets). 4. Default: legacy $HOME/.unsloth/studio (or $USERPROFILE\.unsloth\studio on Windows). Identical to today's behavior when no env var is set. When an env var override fires: * DATA_DIR is nested inside ($STUDIO_HOME/share, or $StudioHome\share on Windows) so the runtime launcher and shortcuts find studio.conf in the same place install-time wrote it. * The unsloth CLI shim lands at $STUDIO_HOME/bin/unsloth (Unix) or $StudioHome\bin\unsloth.exe (Windows). On Windows the shim already lives under $StudioHome; the change only redirects DATA_DIR and skips the persistent registry PATH update. * Persistent shell PATH modifications are skipped (no .bashrc / .zshrc / .profile append on Unix; no Add-ToUserPath on Windows). Caller is expected to invoke via absolute path or add the bin dir to PATH explicitly. Avoids polluting the user's profile with a workspace-scoped path that may be deleted. The Unix launcher script is the only piece that must read DATA_DIR at runtime (it sources studio.conf from there). The hardcoded DATA_DIR inside the LAUNCHER_EOF heredoc is replaced with an @@DATA_DIR@@ placeholder substituted via sed at install time, using the same approach the script already uses for other install-time substitutions. Default path behavior is unchanged: when no env var is set and HOME is not redirected, install.sh / install.ps1 produce exactly the same file layout as today. Test scenarios verified locally on install.sh: * Default (no env vars) -> $HOME/.unsloth/studio (legacy) * HOME=/tmp/x -> /tmp/x/.unsloth/studio * UNSLOTH_STUDIO_HOME=/tmp/y -> /tmp/y as STUDIO_HOME root * STUDIO_HOME=/tmp/z (alias) -> /tmp/z as STUDIO_HOME root * HOME redirect + env var (HOME wins) -> install follows HOME * Unwritable override -> exits with clear ERROR message * install: priority change -- env vars now win over HOME redirect Flip the resolution order so explicit env vars take precedence over HOME / USERPROFILE redirection. New priority (highest first): 1. UNSLOTH_STUDIO_HOME, if set. 2. STUDIO_HOME, if set. 3. HOME / USERPROFILE explicitly redirected. 4. Default. Rationale: the env vars are explicit single-purpose signals (the user typed UNSLOTH_STUDIO_HOME=... specifically to redirect Studio). HOME redirection is broader and incidental -- the user may have redirected HOME for unrelated reasons (workspace tools, container builds) without wanting Studio to follow it. When both are set, the more specific signal should win. When only HOME is redirected (no env var), behavior is unchanged from the previous commit: install follows $HOME. * install: address review feedback (sed escape, downstream propagation, edge cases) Fixes from gemini-code-assist + chatgpt-codex-connector + reviewer.py 20-parallel run on the open PR. install.sh: * Escape sed replacement metacharacters before substituting @@DATA_DIR@@. Two-stage escape: ' -> '\'' for safe single-quote shell embedding, then \, &, | for sed replacement string + chosen delimiter. Heredoc switched to single-quoted DATA_DIR='@@DATA_DIR@@' so we only need single-quote escaping at runtime. Verified end-to-end with paths containing & and | (the sed delimiter). * Pass UNSLOTH_STUDIO_HOME into both setup.sh invocations (--local and PyPI paths) so the downstream install resolves the same Studio root install.sh picked. * macOS .app stub: replace hardcoded exec "$HOME/.local/share/unsloth/launch-studio.sh" with exec "$_css_data_dir/launch-studio.sh" so the .app launches the resolved launcher even in env-override mode. * Use mkdir -p -- and cd -- when validating the env override so paths starting with - cannot be misread as flags. install.ps1: * Drop .Guid from [guid]::NewGuid().Guid: the property does not exist; the probe filename was always identical and not unique. Default ToString() on System.Guid produces the canonical UUID string we want. * Guard LOCALAPPDATA before Join-Path to avoid aborting the installer in service / CI contexts where LOCALAPPDATA is unset (Join-Path under $ErrorActionPreference='Stop' would otherwise throw). Computed once into $defaultDataDir; both 'profile' and 'default' branches reuse it. * Set $env:UNSLOTH_STUDIO_HOME for the duration of the 'unsloth studio setup' subprocess so studio/setup.ps1 and unsloth_cli see the same install root install.ps1 picked. Restored in a finally block. studio/setup.sh: * Honor UNSLOTH_STUDIO_HOME / STUDIO_HOME (alias) when resolving STUDIO_HOME, VENV_DIR, VENV_T5_*_DIR. Falls back to the legacy $HOME/.unsloth/studio when no override is set. studio/setup.ps1: * Same change in PowerShell: honor $env:UNSLOTH_STUDIO_HOME / $env:STUDIO_HOME for $StudioHome / $VenvDir resolution. unsloth_cli/commands/studio.py: * Replace the module-level constant STUDIO_HOME = Path.home() / ".unsloth" / "studio" with a resolver that honors UNSLOTH_STUDIO_HOME / STUDIO_HOME before falling through to the legacy default. Same precedence the installers use. Verified locally: 6 install.sh scenarios still produce correct paths (default, HOME redirect, env var, alias, both, bad override). New sed-escape unit tests pass for paths containing & and |. Python resolver matches priority: UNSLOTH_STUDIO_HOME > STUDIO_HOME > default. * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * install.sh: portable sed (no -i.bak) per gemini review feedback GNU sed -i.bak vs BSD/macOS sed -i.bak vs BusyBox sed have subtly different semantics. Use the POSIX-portable redirect-then-mv pattern instead. Functionally identical, runs everywhere. * studio: persist UNSLOTH_STUDIO_HOME so fresh shells find custom installs Without this, a custom-root install (UNSLOTH_STUDIO_HOME=/work/studio bash install.sh --local) only worked in the same shell that ran the installer. Closing the terminal and reopening lost the env var, the PATH was deliberately not persisted, and the Python CLI fell back to ~/.unsloth/studio. Result: 'Studio not set up' or quietly operating on a stale legacy install. Three persistence layers, all backwards-compatible (default installs emit zero changes): 1. Unix studio.conf install.sh now writes 'export UNSLOTH_STUDIO_HOME=...' next to UNSLOTH_EXE in studio.conf when in env-override mode. The launcher sources studio.conf at startup so the exec'd binary gets the var. Default installs do not write this line; studio.conf stays byte-identical to before. 2. Windows launch-studio.ps1 install.ps1 prepends '$env:UNSLOTH_STUDIO_HOME = ...' to the generated launcher when in env-override mode. Default installs produce the same launcher content as before. 3. Python sys.prefix inference storage_roots.studio_root() and unsloth_cli/commands/studio.py now infer the install root from sys.prefix when no env var is set (Path(sys.prefix).parent for unsloth_studio venvs). Catches direct invocations of /bin/unsloth that bypass the launcher entirely. unsloth_cli/commands/studio.py also re-exports the resolved UNSLOTH_STUDIO_HOME via os.environ.setdefault so child processes (setup script, backend run.py) inherit it. Backend storage roots (storage_roots.studio_root, cache_root) now respect the env var via the shared resolver. run.py PID file, transformers_version.py T5 venvs, and model_config.py vision-check venv all switch to studio_root() so custom installs are self-contained. studio/setup.ps1: T5 sidecar venvs now resolve under $StudioHome (was $env:USERPROFILE\.unsloth\studio\.venv_t5_*). studio/setup.sh + studio/setup.ps1: llama.cpp build dir nests under $STUDIO_HOME / $StudioHome when env-override is active, otherwise keeps the legacy ~/.unsloth/llama.cpp. Verified locally: * studio.conf write block: env-override mode emits the export line; default mode does not (byte-identical to today). * PowerShell heredoc interpolation: correct output for both modes. * studio_root() resolver: default, UNSLOTH_STUDIO_HOME, STUDIO_HOME alias, and sys.prefix-based inference all return correct paths. * cache_root() now derives from studio_root(). * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * install: tilde expansion + macOS .app stub safe-quoting Two fixes from running a 25-scenario simulation sweep against install.sh across path edge cases (spaces, apostrophes, ampersands, pipes, backslashes, dollar signs, Unicode, trailing slash, relative paths). 1. UNSLOTH_STUDIO_HOME=~/foo was landing as literal '~/foo' (env vars are not subject to tilde expansion). Added a POSIX-portable case block in install.sh, install.ps1, studio/setup.sh, studio/setup.ps1 that expands a leading ~ or ~/ to $HOME / $env:USERPROFILE. The prefix-removal pattern is single-quoted ('${var#'~/'}') so the shell does not tilde-expand the pattern back to $HOME/ before matching -- a subtle dash/bash gotcha. 2. macOS .app stub used an unquoted heredoc ('<< STUB_EOF'), so any $VAR / backtick / etc in the path would expand at .app launch time. Switched to single-quoted heredoc ('<< 'STUB_EOF'') with a placeholder + sed substitution + single-quoted shell embedding, matching the @@DATA_DIR@@ pattern already used for launch-studio.sh. Verified: 25/25 simulation scenarios pass on Linux dash + bash, including paths with $VAR, &, |, \\, ', spaces, and Unicode. End-to-end install in env-mode + fresh-shell launcher invocation confirmed: studio binds to /api/health from a clean env, and sys.prefix-based inference correctly returns the workspace root. * install: stop accidentally treating default installs as env-override Reviewer.py 20-runs cycle 1 found a unanimous P1 regression: a default 'unsloth studio update' relocates llama.cpp from ~/.unsloth/llama.cpp to ~/.unsloth/studio/llama.cpp, because the CLI was re-exporting UNSLOTH_STUDIO_HOME unconditionally and install.sh / install.ps1 were passing it into setup.{sh,ps1} unconditionally. The setup scripts treated the var's mere presence as "env-override mode" and relocated the llama.cpp build dir away from the legacy path, breaking the runtime backend's _find_llama_server_binary lookup on default installs. Fixes: * unsloth_cli/commands/studio.py: _resolve_studio_home now returns (path, is_custom). Re-export only when is_custom -- a real env override or a sys.prefix inference that resolves to a non-legacy path. Default installs leave UNSLOTH_STUDIO_HOME unset. * install.sh: gate UNSLOTH_STUDIO_HOME on $_STUDIO_HOME_REDIRECT == env before calling setup.sh. Use 'env $VARS bash setup.sh' so the var is set only for the subprocess, never leaked. * install.ps1: gate $env:UNSLOTH_STUDIO_HOME on $StudioRedirectMode -eq 'env' before invoking 'unsloth studio setup'. Restore prior value in finally block (unset if it wasn't set). * studio/setup.sh + setup.ps1: decide llama.cpp install root from the resolved $STUDIO_HOME (not from env-var presence). If the resolved path equals the legacy default ($HOME/.unsloth/studio), fall back to ~/.unsloth/llama.cpp. This makes setup robust against a stale UNSLOTH_STUDIO_HOME inherited from a parent process that happens to point at the legacy default. * studio/backend/core/inference/llama_cpp.py: - _find_llama_server_binary() now searches studio_root() / llama.cpp AND the legacy ~/.unsloth/llama.cpp (de-duped). Custom-root installs become discoverable; default installs unaffected. - kill_orphaned_servers ownership allowlist also includes studio_root() / llama.cpp so custom-root processes are cleanable. Verified locally: * 25/25 sim scenarios still pass (path edge cases unchanged). * setup.sh unit test: default-mode lands UNSLOTH_HOME at $HOME/.unsloth; env-mode lands at $STUDIO_HOME. * Python CLI unit test: default-mode returns is_custom=False and does NOT setdefault UNSLOTH_STUDIO_HOME; env-mode sets is_custom=True. * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * install: || exit 1 on STUDIO_HOME subshell (dash set -e gap) Gemini review feedback: in dash, set -e does not trigger on subshell failures inside variable assignments. If 'cd -- "$_override" && pwd' fails, STUDIO_HOME stays empty and DATA_DIR collapses to /share. Add explicit '|| exit 1' on both install.sh:187 and setup.sh:413. * install.sh: argv-safe setup invocation for paths with spaces Cycle 2 reviewer.py 20-runs found a unanimous P1: passing the env-var through 'env $_STUDIO_ENV_FOR_SETUP' word-splits on whitespace, so a custom root like '/tmp/Unsloth Studio' becomes 'UNSLOTH_STUDIO_HOME= /tmp/Unsloth' followed by env trying to exec 'Studio'. Replaced with a tiny helper that prepends the env-var directly to the argv (no string-form intermediary), so spaces are preserved as a single argument. Default-mode invocation skips the env-var entirely. Verified: 'UNSLOTH_STUDIO_HOME=/tmp/test space/studio' now reaches setup.sh as a single value. * studio: tighten sys.prefix inference + Tauri env handling + llama.cpp env Cycle 3 reviewer.py findings (3 P1s converging): * sys.prefix inference too broad: a developer venv named 'unsloth_studio' was being treated as a custom Studio root. Narrow with an installer- sentinel check (presence of share/studio.conf or bin/unsloth shim inside the parent dir) in both unsloth_cli/commands/studio.py and studio/backend/utils/paths/storage_roots.py. * Tauri studio/src-tauri/src/process.rs::find_unsloth_binary() hardcoded ~/.unsloth/studio. Honor UNSLOTH_STUDIO_HOME / STUDIO_HOME (in that priority order) before falling back to legacy. * unsloth-zoo's GGUF export binds LLAMA_CPP_DEFAULT_DIR at import time from UNSLOTH_LLAMA_CPP_PATH. For env-override installs, persist UNSLOTH_LLAMA_CPP_PATH alongside UNSLOTH_STUDIO_HOME in studio.conf (Unix), in the generated PowerShell launcher (Windows), and via os.environ.setdefault in the Python CLI when running on a custom root, so GGUF export uses the custom-root llama.cpp build instead of the legacy ~/.unsloth/llama.cpp. Default behaviour unchanged: no env vars are written to studio.conf in default mode, no LLAMA_CPP_PATH is set, and the dev-venv inference falls through to legacy when no installer sentinels are present. * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * studio: desktop_auth env-aware + legacy-root llama.cpp consistency - desktop_auth.rs: honor UNSLOTH_STUDIO_HOME / STUDIO_HOME for the .desktop_secret path so Tauri desktop login works against custom-root installs instead of always reading ~/.unsloth/studio/auth/. - install.sh / install.ps1 / unsloth_cli/commands/studio.py: when an env override resolves to the legacy default ($HOME/.unsloth/studio), set UNSLOTH_LLAMA_CPP_PATH to ~/.unsloth/llama.cpp (matching setup.sh / setup.ps1's legacy-equality branch). Previously the persisted value pointed at $STUDIO_HOME/llama.cpp, which was a non-existent location and broke unsloth-zoo's import-time GGUF binding for that edge case. * studio: tauri studio_root helper + marker-file persistence + ~ expansion Address cycle-5 reviewer findings: - Add studio/src-tauri/src/studio_root.rs: shared resolver with UNSLOTH_STUDIO_HOME / STUDIO_HOME (priority order), tilde expansion (~, ~/..., ~\...), installer-written marker fallback, then ~/.unsloth/studio. 5 unit tests cover the expansion paths. - Tauri lookups now go through the shared resolver: - process.rs::find_unsloth_binary - desktop_auth.rs::desktop_secret_path - main.rs::setup_logging (tauri.log under custom root) - commands.rs::open_logs_dir (opens custom root dir) - install.rs work_dir uses parent of resolved root (avoids creating a stray ~/.unsloth on a custom-root install) - install.sh / install.ps1 (env-mode only): write ~/.unsloth/studio-home marker so the desktop app launched from Finder/Start Menu (no shell env inheritance) still resolves the custom root. - install.sh / install.ps1 non-interactive completion: when StudioRedirectMode=env, print the absolute custom-root shim path since the persistent rc/registry PATH update is intentionally skipped in env-override mode. - unsloth_cli/commands/studio.py: replace setdefault() with truthy-check so a blank UNSLOTH_STUDIO_HOME / UNSLOTH_LLAMA_CPP_PATH in the parent env doesn't suppress the inferred custom root. 40/40 cargo test --bins pass. * studio: validate marker file + write in --tauri mode + propagate to subprocess Cycle-6 reviewer follow-ups: - studio_root.rs marker resolver now validates the persisted path before using it. A stale ~/.unsloth/studio-home pointing at a deleted/moved workspace is ignored (resolution falls back to the legacy default rather than hijacking it). Validation accepts share/studio.conf sentinel or bin/unsloth shim. Trailing newline strip uses trim_end_matches(['\n','\r']) so paths whose content legitimately has leading/trailing spaces survive. - install.sh / install.ps1: marker write moved out of the launcher generation path so it runs before the Tauri-mode early exit. Both shell-launcher and Tauri-installed env-mode roots now persist the marker. Removed the duplicate marker write that was previously inside install.ps1's $studioHomeExport block. - studio/src-tauri/src/install.rs: pass UNSLOTH_STUDIO_HOME to the installer subprocess (when not already in scope) so app-initiated repair / update flows reach the same root the running app uses. cargo test --bins -- --test-threads=1: 44/44 pass (4 new tests for marker validation: sentinel accepted, bin shim accepted, empty dir rejected, missing path rejected). * studio: fix Tauri legacy-fallback regression + stale marker cleanup Cycle-7 reviewer follow-ups (regression I introduced in cycle 6): - studio_root.rs: add StudioRootSource enum + resolve_studio_root_with_source(). Lets callers distinguish a real custom override (Env / Marker) from the legacy fallback (Default). - studio/src-tauri/src/install.rs: only forward UNSLOTH_STUDIO_HOME to the installer subprocess when the resolution source is Env or Marker. The Default fallback must NOT be passed -- install.sh / install.ps1 treat any non-empty UNSLOTH_STUDIO_HOME as env-override mode and would relocate DATA_DIR to $STUDIO_HOME/share and _LOCAL_BIN to $STUDIO_HOME/bin (regressing default Tauri repair / update flows from the legacy ~/.local/share/unsloth and ~/.local/bin). - install.sh / install.ps1: clear stale marker on default / HOME-redirect installs. A user who first installed with UNSLOTH_STUDIO_HOME=/work/studio then later reinstalls without env vars no longer has the desktop app hijacked by ~/.unsloth/studio-home pointing at the old custom root. - install.sh / install.ps1: when env mode wins over a redirected HOME / USERPROFILE, write the marker into the OS-reported real profile home (getent / dscl on Unix; [Environment]::GetFolderPath on Windows) so a later desktop launch from the user's normal session still finds it. Falls back to the current HOME / USERPROFILE. cargo test --bins -- --test-threads=1: 45/45 pass (1 new for the source enum invariants). * install: scrub stale marker from real-home on HOME-redirect cleanup Cycle-8 reviewer follow-up: the previous cleanup branch only removed \$HOME/.unsloth/studio-home, leaving a stale marker in the real password-database home after a prior env-mode install. A later default install with redirected HOME / USERPROFILE would still see the desktop app resolving the old custom root. - install.sh: compute the real password-database home (via getent / dscl) unconditionally, and scrub markers from BOTH \$HOME and the real-home in the default / HOME-redirect cleanup branch. - install.ps1: build a profile-candidate list (current USERPROFILE + OS-reported real profile) and remove markers from EVERY candidate in the default / profile-redirect cleanup branch. bash -n + cleanup smoke verified. * revert: drop Tauri env-var support + marker file mechanism Keep this PR scoped to shell installer + Python backend env-var support. Tauri desktop integration with custom Studio roots is deferred to a separate, focused PR. Reverts to pre-PR state: - studio/src-tauri/src/process.rs (find_unsloth_binary) - studio/src-tauri/src/desktop_auth.rs (auth_secret_path) - studio/src-tauri/src/main.rs (setup_logging tauri.log path) - studio/src-tauri/src/commands.rs (open_logs_dir) - studio/src-tauri/src/install.rs (work_dir + subprocess env) - studio/src-tauri/src/studio_root.rs DELETED Removes from install.sh / install.ps1: - ~/.unsloth/studio-home marker write/read/cleanup - HOME-redirect-aware marker location logic What this PR keeps (the original scope): - install.sh / install.ps1: UNSLOTH_STUDIO_HOME / STUDIO_HOME env-var resolver with HOME-redirect detection, tilde expansion, legacy fallback. Default installs are byte-identical to pre-PR. - studio/setup.sh / studio/setup.ps1: legacy-equality llama.cpp path. - studio.conf / launcher persists UNSLOTH_STUDIO_HOME + UNSLOTH_LLAMA_CPP_PATH for fresh shells (env-mode only). - unsloth_cli/commands/studio.py: env > sys.prefix sentinel > legacy resolver, conditional re-export. - studio/backend/utils/paths/storage_roots.py: same resolver. - Backend modules use storage_roots (run.py, model_config.py, transformers_version.py, llama_cpp.py). cargo test --bins -- --test-threads=1: 34/34 pass (pre-PR baseline). bash -n install.sh: clean. * install: cycle-10 fixes (default launcher, --tauri guard, env-mode shortcuts, win PATH) - install.sh launcher: default and HOME-redirect installs keep the legacy DATA_DIR=\"\$HOME/.local/share/unsloth\" runtime form so a later shell with a different \$HOME still resolves DATA_DIR. Only env-mode bakes the resolved absolute path. Restores byte-identical default behavior. - install.sh / install.ps1: fail fast when --tauri is combined with UNSLOTH_STUDIO_HOME / STUDIO_HOME. The desktop app still resolves the legacy ~/.unsloth/studio root, so a custom-root --tauri install would yield a desktop app that cannot find its binary or auth secret. Print the right alternative. - install.sh / install.ps1: skip persistent desktop / Start-Menu shortcuts in env-override mode. Workspace-scoped installs would otherwise leave launchers pointing at a path the user may delete. Default and HOME/profile-redirect installs keep the shortcut. - install.ps1: re-prepend env-override \$ShimDir AFTER Refresh-SessionPath. Refresh rebuilds PATH as Machine > User > current \$env:Path, so a previously-installed legacy User PATH entry would otherwise win precedence over the current-session env-override shim. bash -n install.sh, pwsh parser install.ps1 + setup.ps1: clean. cargo test --bins -- --test-threads=1: 34/34 (Tauri unchanged). * install: cycle-11 fixes (env-mode launcher writes, --tauri legacy passthrough, run.py llama path) - install.sh / install.ps1: env-mode no longer skips the entire create_studio_shortcuts / New-StudioShortcuts function. Move the early-return INSIDE those functions, just before the persistent desktop / Start-Menu shortcut creation. The runtime launcher (launch-studio.sh / launch-studio.ps1), studio.conf with UNSLOTH_STUDIO_HOME / UNSLOTH_LLAMA_CPP_PATH exports, and the icon ARE always written so env-mode shims can resolve via fresh shells. - install.sh / install.ps1: --tauri guard passes through when the override resolves to the legacy default ($HOME/.unsloth/studio / %USERPROFILE%\.unsloth\studio). The desktop app already uses that path, so explicit-equality is a supported edge case (matches the llama.cpp legacy-equality branch). - studio/backend/run.py: when launched directly (bypassing the unsloth CLI), set UNSLOTH_STUDIO_HOME and UNSLOTH_LLAMA_CPP_PATH before the rest of import chain runs so unsloth-zoo's import-time LLAMA_CPP_DEFAULT_DIR binding picks up the custom-root build. Only set when STUDIO_ROOT is a real custom override; legacy default installs leave them unset. bash -n install.sh, pwsh parser install.ps1: clean. python ast parse studio/backend/run.py: clean. cargo test --bins -- --test-threads=1: 34/34 pass (Tauri unchanged). * install: cycle-12 fixes (--tauri trailing slash + main.py uvicorn env) - install.sh / install.ps1 --tauri legacy passthrough: strip trailing separators before comparing the override to the legacy default. Previously UNSLOTH_STUDIO_HOME=\"\$HOME/.unsloth/studio/\" (with trailing slash) was rejected even though it resolves to the supported legacy root. - studio/backend/main.py: when launched directly via \`uvicorn main:app\` from a custom-root venv (bypassing both unsloth_cli and run.py), export UNSLOTH_STUDIO_HOME and UNSLOTH_LLAMA_CPP_PATH before any unsloth-zoo import so its import-time LLAMA_CPP_DEFAULT_DIR binding picks up the custom-root build. Only sets when STUDIO_ROOT is a real custom override. bash -n install.sh, pwsh parser install.ps1, python ast main.py: clean. Smoke probe: UNSLOTH_STUDIO_HOME=\$HOME/.unsloth/studio/ install.sh --tauri no longer exits with the unsupported-custom-root error. * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * install.ps1: skip CWD-relative venv migration in env-override mode The legacy ~/unsloth_studio venv migration path on Windows reads %USERPROFILE%\unsloth_studio\Scripts\python.exe (a fixed home-relative path). Under env-override mode this would Move-Item the user's pre-existing default-install venv into $StudioHome\unsloth_studio, breaking the default install and contaminating the workspace root. Gate the migration on $StudioRedirectMode -ne 'env' so workspace-scoped installs leave the user's default-install venv untouched. No Linux equivalent: install.sh migrates from \$STUDIO_HOME/.venv which is already env-mode-aware (points at the workspace root, not \$HOME). * install: cycle-14 fixes (Tauri env scrub + setup.ps1 missing-root error) Tauri does not honor UNSLOTH_STUDIO_HOME / STUDIO_HOME / UNSLOTH_LLAMA_CPP_PATH yet -- the desktop app's Rust paths use the legacy ~/.unsloth/studio root. If the user's shell has these env vars set, spawned Python subprocesses would diverge from the Rust paths (custom-root Python <-> legacy-root Rust). Scrub the three env vars at all Tauri subprocess spawn sites: - process.rs: backend launch - desktop_auth.rs: provision-desktop-auth subprocess - install.rs: install.sh / install.ps1 invoked from the desktop app (also prevents the --tauri guard from rejecting an inherited override). setup.ps1: when UNSLOTH_STUDIO_HOME points at a non-existent directory, 'Resolve-Path -LiteralPath' threw a confusing PSObject error under $ErrorActionPreference = "Stop". Test-Path the override first and emit a friendly "run install.ps1 to create the install root" message instead. * install: cycle-15 fixes (preserve UNSLOTH_LLAMA_CPP_PATH + add update.rs scrub) UNSLOTH_LLAMA_CPP_PATH is a pre-existing custom-llama.cpp-directory override the Python backend (studio/backend/core/inference/llama_cpp.py) and unsloth-zoo intentionally support. It is unrelated to the Studio install root. Cycle 14 over-scrubbed it from the Tauri spawn sites, regressing desktop GGUF/llama.cpp workflows for users who set it in their shell. - process.rs / desktop_auth.rs / install.rs: stop scrubbing UNSLOTH_LLAMA_CPP_PATH; only scrub UNSLOTH_STUDIO_HOME and STUDIO_HOME. - update.rs: missed Tauri spawn site -- add the same UNSLOTH_STUDIO_HOME / STUDIO_HOME scrub so 'unsloth studio update' from the desktop app updates the legacy-root install Tauri actually manages. Verified: cargo test --bins -- --test-threads=1 -> 34/34 pass. * install.sh: document apostrophe-escape derivation inline The shell quoting at install.sh:642 / 659 / 679 / 680 / 823 has been flagged as broken across multiple review cycles, but every end-to-end verification (DATA_DIR=\"a b's&c|d\$e\" -> generated launcher -> source -> recovered exact input) passes. The proposed "8 backslash" fix would double the escape and actually break what currently works. Strengthen the inline comments to spell out the derivation: - shell pattern \"s/'/'\\\\''/g\" passes \"s/'/'\\''/g\" to sed (\\\\ -> \\) - sed replacement '\\'' yields close-quote / escaped-quote / open-quote - stage 2 (\\, &, |) only needed where the value is then sed-replaced into a launcher template via s|@@DATA_DIR@@|VALUE|g studio.conf is written via printf, not sed, so it only needs stage 1. No behavior change, only inline doc to head off future false positives. * install/setup .ps1: use -LiteralPath for $StudioHome-derived paths Pre-PR, $StudioHome was hardcoded to %USERPROFILE%\.unsloth\studio -- no wildcard characters possible. The PR introduces UNSLOTH_STUDIO_HOME / STUDIO_HOME, so $StudioHome (and every path derived from it: $VenvDir, $VenvPyExe, $UnslothExe, $UnslothHome, $LlamaCppDir, $VenvT5_*, etc.) can now contain bracket characters that PowerShell would interpret as wildcards. Reproducer (from cycle 17 review 20): pwsh> Test-Path 'studio[abc]/Scripts/python.exe' False pwsh> Test-Path -LiteralPath 'studio[abc]/Scripts/python.exe' True Switch the relevant Test-Path / Remove-Item / New-Item / Move-Item calls in install.ps1 and studio/setup.ps1 to -LiteralPath. Sites where the path is fixed (the shim under %LOCALAPPDATA%\Microsoft\WindowsApps, $RepoRoot from -PSCommandPath) keep the wildcard-aware form. * install/setup .ps1: fix New-Item -LiteralPath regression from cycle 17 Cycle 17 added -LiteralPath to all $StudioHome-derived path operations, but New-Item has no -LiteralPath parameter (verified pwsh 7.6 syntax: "New-Item [-Path] [-ItemType ] ..."). Every directory- creation site would throw "A parameter cannot be found that matches parameter name 'LiteralPath'" at runtime, blocking T5 sidecar setup, llama.cpp parent creation, and StudioHome creation. Likewise, "Split-Path -LiteralPath $X -Parent" cannot mix LiteralPath with -Parent (separate parameter sets). The default LiteralPath mode already returns the parent. Switch to [System.IO.Directory]::CreateDirectory($X), which natively takes a literal path, and drop the trailing -Parent on Split-Path. Verified end-to-end on a bracketed path "/tmp/...[abc]": - CreateDirectory: created - Test-Path -LiteralPath: detects - nested CreateDirectory(Split-Path -LiteralPath ...): works * install/setup .ps1: extend -LiteralPath sweep to remaining \$StudioHome paths Cycle 17/18 missed several wildcard-aware operations on user-controlled \$StudioHome-derived paths. Reviewers identified remaining sites: install.ps1: - \$UnslothExePath (Test-Path / Resolve-Path) at the shortcut creator - \$VenvDir (Get-ChildItem) at the no-torch-runtime resolver - \$ShimDir (New-Item Directory -- replaced with .NET CreateDirectory) - \$ShimExe (Test-Path / Remove-Item / re-prepend guards) -- the shim lives at \$StudioHome\\bin\\unsloth.exe in env-override mode, so it inherits bracket sensitivity from \$StudioHome. - \$UnslothExe (Copy-Item fallback) when HardLink fails. studio/setup.ps1: - \$LlamaServerBin (Test-Path) at the prebuilt-bundle / source-build validation gates (3 sites). \$LlamaServerBin lives under \$BuildDir under \$LlamaCppDir under \$UnslothHome under \$StudioHome. New-Item HardLink keeps -Path because creating a non-existent target with brackets succeeds (verified via direct pwsh smoke test). * install: cycle-20 fixes (more setup.ps1 -LiteralPath + shell-quote launch hints) setup.ps1: extend -LiteralPath sweep to remaining \$BuildDir-derived paths that the cycle-19 commit missed: - \$CmakeCacheFile (Test-Path + Select-String -Path) - \$buildTmp (10 Test-Path / Remove-Item sites in source-build cleanup) - \$QuantizeBin (Test-Path) - \$altBin (Test-Path) These all live under \$BuildDir -> \$LlamaCppDir -> \$UnslothHome -> \$StudioHome, which is now user-controlled via UNSLOTH_STUDIO_HOME. Bracket characters in the override would silently skip rebuild detection or leave stale build artifacts. install.sh: shell-quote the launch-instruction substep lines for env- override mode. UNSLOTH_STUDIO_HOME values containing spaces or apostrophes (e.g. "/tmp/O'Brien Studio") would print copy-paste- unsafe commands -- the install succeeded but the printed launch instructions split at the space. Now wraps with the canonical '\\''-style escape so the printed lines parse with bash -n. Verified end-to-end: - printed shim line: '/tmp/O'\''Brien Studio/bin/unsloth' studio ... - bash -n on the printed line passes. * install.ps1: -LiteralPath for macOS-stub-launcher \$appDir-derived paths The shortcut/launcher generator at install.ps1:418-693 writes the stub launcher, .vbs, and icon under \$appDir = \$StudioDataDir, which in env-override mode is \$StudioHome\share. Cycle 17/19/20 missed the following wildcard-aware ops on these paths: - Test-Path \$appDir (with New-Item Directory swap to .NET CreateDirectory) - Set-Content -Path \$launcherVbs (for the WSH .vbs stub) - Test-Path / Copy-Item \$bundledIcon (bundled icon copy) - Test-Path / Remove-Item \$iconPath (icon header validation) In env-override mode \$StudioHome can contain bracket characters; without -LiteralPath the .vbs write fails outright and the icon validation can either skip a present icon or fail to delete a malformed one. (The COM shortcut creation downstream returns early in env-override mode, so its path values don't need this treatment.) * install: don't override pre-existing UNSLOTH_LLAMA_CPP_PATH in launchers Cycle 14/15 established UNSLOTH_LLAMA_CPP_PATH as a pre-existing custom-llama.cpp-directory override the Python backend and unsloth-zoo intentionally support, independent of the Studio install root. The launchers (studio.conf sourced by Unix launch-studio.sh, and the PowerShell launch-studio.ps1) were unconditionally re-exporting it, which silently overrides a user's pre-existing value when they invoke the launcher from a shell where UNSLOTH_LLAMA_CPP_PATH is already set. Make the assignment conditional in both launchers: install.sh studio.conf: if [ -z "\${UNSLOTH_LLAMA_CPP_PATH:-}" ]; then export UNSLOTH_LLAMA_CPP_PATH='...' fi install.ps1 launch-studio.ps1: if (-not \$env:UNSLOTH_LLAMA_CPP_PATH) { \$env:UNSLOTH_LLAMA_CPP_PATH = '...' } UNSLOTH_STUDIO_HOME stays unconditional: the launcher is bound to a specific install, so its STUDIO_HOME must always match that install. * install.sh: harden --tauri legacy resolver against CDPATH and symlinks Reviewer cycle 23 (inst 19) noted that the bare \`cd -- ... && pwd\` form in the --tauri legacy comparison can echo a CDPATH-prefixed path when the user has CDPATH set in their environment, contaminating the resolved absolute path used in the legacy-equality check. Switch to \`CDPATH= cd -P -- ... && pwd -P\` so: - CDPATH= clears the cd-prefix-echo behavior - -P / pwd -P resolves any symlinks to a canonical path No behavior change for users without CDPATH set; correctness fix for users who have it set in their shell. * install + llama_cpp backend: cycle-24 hardening Three real findings from cycle 24 reviewers: 1. install.sh:231 + studio/setup.sh:413 -- main \$STUDIO_HOME resolvers used the same bare \`cd -- ... && pwd\` form that cycle 23 only fixed for the --tauri guard. Switch both to: \$(CDPATH= cd -P -- "\$override" && pwd -P) so relative custom-root values don't get CDPATH-prefixed or have the cd-on-CDPATH stdout newline contaminate the captured value. 2. install.sh --tauri legacy root used logical \$HOME/.unsloth/studio while the override side was canonicalized via pwd -P. A symlinked \$HOME (e.g. /home/alice -> /u/alice) made the comparison fail even when both sides pointed at the same directory. Canonicalize the legacy side too when the dir exists. 3. studio/backend/core/inference/llama_cpp.py:_find_llama_server_binary searched \$STUDIO_HOME/llama.cpp first then ~/.unsloth/llama.cpp in default-mode installs. setup.sh / setup.ps1 only install llama.cpp under \$STUDIO_HOME/llama.cpp in env-override mode; in default mode it always lives at ~/.unsloth/llama.cpp. The post-PR search would pick up a stale partial install at ~/.unsloth/studio/llama.cpp over the real legacy binary. Mirror setup's legacy-equality check: when studio_root() resolves equal to ~/.unsloth/studio, search ONLY the legacy ~/.unsloth/llama.cpp. Otherwise (env-override custom root), search custom first, legacy fallback. * install + setup: canonicalize legacy-equality comparison sites Cycle 24 made \$STUDIO_HOME canonical via 'CDPATH= cd -P -- ... && pwd -P', but the legacy-equality comparison sites still used the bare logical "\$HOME/.unsloth/studio" string. With a symlinked \$HOME (e.g. /home/alice -> /u/alice), the comparison fails even when both sides point at the same dir, and llama.cpp ends up under a custom-root path the Python backend's legacy comparison cannot find. Reviewer cycle 25 inst 2 reproduced this with HOME=/tmp/link -> /tmp/real and UNSLOTH_STUDIO_HOME=\$HOME/.unsloth/studio: setup.sh resolves UNSLOTH_HOME to /tmp/real/.unsloth/studio while the backend search resolves both physically equal and looks at /tmp/link/.unsloth/llama.cpp. Canonicalize the legacy side at all four sites: - install.sh:695 (create_studio_shortcuts llama.cpp path) - studio/setup.sh:577 (UNSLOTH_HOME selection) - install.ps1:462 (launcher UNSLOTH_LLAMA_CPP_PATH path) - studio/setup.ps1:1829 (UnslothHome selection) Apply CDPATH= cd -P -- ... && pwd -P (Unix) or Resolve-Path -LiteralPath (Windows) when the legacy dir exists. unsloth_cli/commands/studio.py already does this via Path.resolve(). * llama_cpp: gate _kill_orphaned_servers studio-root allowlist on env-override Cycle 24 fixed _find_llama_server_binary to only search \$STUDIO_HOME/llama.cpp when STUDIO_HOME is a real env override (not the legacy default), but the symmetric _kill_orphaned_servers allowlist still appended _sr() / "llama.cpp" unconditionally. In default mode _sr() resolves to ~/.unsloth/studio, so ~/.unsloth/studio/llama.cpp would be treated as a Studio-owned install root for the orphan-kill scan even though the default installer does not own that path. A llama-server process running there from a different tool or a stale partial install would be killed. Apply the same legacy-equality check used in _find_llama_server_binary and the install/setup scripts: only add _sr()/"llama.cpp" to the allowlist when STUDIO_HOME != legacy default. * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * setup.sh + setup.ps1: canonicalize both sides of legacy-equality check Proactive audit pass found one real asymmetry the cycle-by-cycle review process had not yet flagged: - install.sh:704 / install.ps1:469 are gated on env-mode and only run when STUDIO_HOME has already been canonicalized (cycle 24). Symmetric. - studio/setup.sh:577 / studio/setup.ps1:1829 run UNCONDITIONALLY, including in default mode. In default mode STUDIO_HOME is set to the bare logical \$HOME/.unsloth/studio (setup.sh:416) or Join-Path \$env:USERPROFILE ".unsloth\\studio" (setup.ps1:1480). Cycle 25 canonicalized only the legacy side, creating an asymmetry under symlinked \$HOME / junctioned %USERPROFILE%. Result of the asymmetry: a default-mode install on a host with \$HOME=/tmp/link -> /tmp/real treats the legacy default as a custom root, putting llama.cpp at \$STUDIO_HOME/llama.cpp instead of ~/.unsloth/llama.cpp -- and the Python backend's _find_llama_server_binary (which uses .resolve() on both sides) then can't find the install. Fix: canonicalize STUDIO_HOME on the fly at the comparison site, in both setup.sh and setup.ps1. Symmetric with the now-canonicalized legacy side from cycle 25, regardless of which mode set STUDIO_HOME. The other two comparison sites (install.sh:704, install.ps1:469) are already symmetric because they only run when STUDIO_HOME comes from the env-override resolution path that already does pwd -P / Resolve-Path. unsloth_cli/commands/studio.py + studio/backend/run.py + main.py + llama_cpp.py already use .resolve() on both sides -- symmetric. * install.ps1: env-override resolution uses .NET API for literal paths Gemini code-review (review 4177641398, commit 2ea2c91) caught two remaining New-Item -Path sites in the env-override resolution block that the cycle 18 sweep missed: - Line 123: New-Item -ItemType Directory -Path \$envOverride - Line 132: New-Item -ItemType File -Path \$probe (writability test) Both use -Path which interprets square brackets as wildcards. For a user with UNSLOTH_STUDIO_HOME=C:\\workspaces\\studio[abc], both calls would fail before the install starts. New-Item also has no -LiteralPath in PowerShell 5.1. Replace both with the .NET API: - [System.IO.Directory]::CreateDirectory(\$envOverride) - [System.IO.File]::WriteAllText(\$probe, "") -- closes the file handle before the Remove-Item below. End-to-end verified with /tmp/test-envoverride-[abc]-* path: CreateDirectory + WriteAllText + Test-Path -LiteralPath all work. * comments: condense multiline blocks added by this PR Across the 27-cycle review process, comments accumulated as multiline blocks explaining each fix's history (cycle numbers, prior bugs, reviewer rationale). Compress every block to 1-2 lines that capture just the WHY, dropping cycle references and history that belongs in the PR description / commit log instead. Net: 268 deletions / 124 insertions (-144 lines) of comments only. Behavior unchanged. Verified: bash -n, pwsh parser, python ast.parse, cargo check all pass. * install.ps1: use 'return' over 'exit 1' for Install-UnslothStudio bail-outs Per Gemini review #4177659001: when users run install.ps1 via 'irm ... | iex', 'exit 1' inside the function terminates the entire PowerShell process and closes the user's terminal. 'return' bails out of the function while keeping the shell open, matching existing error sites at lines 34, 50, 57. Three sites fixed: --tauri+env-override guard, env-override mkdir/access failure, and write-probe failure. The 'exit' calls at lines 591/611 are inside a generated launcher here-string (a separate top-level .ps1 that runs as its own process), so they correctly stay as 'exit'. * install.{sh,ps1}: address Gemini review #4177680451 Three medium fixes: 1. install.sh redirection detection: canonicalize both sides of the $HOME vs passwd-DB comparison via 'CDPATH= cd -P -- ... && pwd -P' so a trailing slash on $HOME (or symlink-vs-realpath mismatch with getent/dscl output) doesn't misfire the redirection branch. 2. install.sh shim symlink: 'ln -sf' into an existing directory creates the link INSIDE it ($_LOCAL_BIN/unsloth/unsloth instead of the intended file). Pre-strip a real (non-symlink) directory at $_LOCAL_BIN/unsloth before linking. 3. install.ps1 ShimExe: add -Recurse to Remove-Item so the launcher refresh recovers if $ShimExe somehow exists as a directory rather than a file (would otherwise drop into the catch and skip the shim update). * install.ps1: use 'throw' over 'return' for fatal validation failures Cycle 28 reviewer.py (12/8 RC/APPROVE) caught a regression introduced by the previous Gemini-review fix (#4177659001 -> commit 393e676b). 'return' inside Install-UnslothStudio kept iex'd terminals alive but made 'pwsh -File install.ps1' exit with code 0 on fatal validation failures (--tauri+custom-root rejected, STUDIO_HOME unwritable, etc.), so CI / wrapper scripts treated failed installs as successful. 'throw' satisfies both constraints: - pwsh -File install.ps1: exits with code 1 (CI sees failure) - irm | iex: shows error to user, does NOT close the host terminal Three sites: --tauri+env-override guard, mkdir/access failure, write-probe failure. Verified throw -> exit code 1 under pwsh -File. * install.ps1 launcher: single-quote child -Command path Cycle 28 P2 finding: the generated launch-studio.ps1 builds the child PowerShell -Command string with the executable path inside double quotes, so a custom Studio root containing PowerShell metacharacters (\$, backtick) re-expands in the child shell. Example: D:\work\\\$job\studio -> child reparses \$job and runs the wrong path. Fix: single-quote the path inside the child command and double any apostrophes (PowerShell's literal-quote-escape form) so paths like "O'Brien Studio & x|y" or "C:\work\\\$bad\studio" survive verbatim. * install: harden custom Studio root handling - install.sh shim refresh: refuse to recursively delete a real directory at $_LOCAL_BIN/unsloth before creating the symlink. The previous rm -rf could destroy unrelated user data living at that path. - install.ps1 shim refresh: drop -Recurse from Remove-Item on $ShimExe and refuse early when the shim path is a directory; mirrors the install.sh guard so a directory at $StudioHome\bin\unsloth.exe is not blown away. - install.ps1 PATH wiring: remove the redundant first $ShimDir prepend in env-override mode; the post-Refresh-SessionPath prepend is the one that takes effect, and the duplicate left $ShimDir in $env:Path twice. - install.ps1 manual launch instructions: single-quote the printed shim and Activate.ps1 paths so '$' / backtick metacharacters in custom roots do not reparse when the user copies and pastes the command. - studio/setup.sh: validate writability of UNSLOTH_STUDIO_HOME with the same [ -w ] check install.sh already has, so a read-only override fails with a clear message instead of an obscure uv pip permission error. - Drop the STUDIO_HOME alias everywhere (storage_roots.py, studio.py, install.sh, studio/setup.sh, install.ps1, studio/setup.ps1). The name is too generic and an ambient STUDIO_HOME from unrelated tooling could silently redirect the install. Only UNSLOTH_STUDIO_HOME is honored. - unsloth_cli/commands/studio.py: defer UNSLOTH_STUDIO_HOME / UNSLOTH_LLAMA_CPP_PATH re-export from import time into a helper invoked by the studio app callback. Importing the module no longer mutates os.environ as a side effect, so test runners and CLI introspection stop leaking those vars into unrelated subprocesses. - studio/backend/core/inference/llama_cpp.py: replace set-mutation inside list comprehension with an explicit dedup loop for readability. * install: harden custom Studio root edge cases - install.ps1 shim refresh: move the directory-collision preflight outside the lock-handling try/catch. The previous throw inside the try block was swallowed by the surrounding catch and downgraded to a "Continuing with the existing launcher" warning, leaving the install in a broken state with no usable shim on disk. - storage_roots.py / unsloth_cli/commands/studio.py: tighten the bin-shim sentinel from .exists() to .is_file(). A directory at the candidate bin/unsloth (or bin/unsloth.exe) path would otherwise false-positive the venv inference and pick the wrong Studio root. - storage_roots.py / unsloth_cli/commands/studio.py: wrap the env-var override Path(...).expanduser().resolve() in try/except (OSError, ValueError), matching the defensive pattern already used in studio/backend/main.py and studio/backend/run.py. An invalid override (unresolvable network drive, bad characters) now falls back to the un-resolved path instead of crashing at import time. * install: fail fast on missing custom root, allow brackets in shim path - install.ps1 shim hardlink: switch the New-Item -ItemType HardLink call from -Path to -LiteralPath so a custom Studio root containing bracket characters does not fail under PowerShell's wildcard-aware -Path parameter. Matches the -LiteralPath usage on every other Test-Path / Remove-Item / Copy-Item call against the same shim path. - studio/setup.sh override branch: replace the silent mkdir -p of the override directory with an existence check that exits 1 with a clear message. setup.sh runs against an existing install (via 'unsloth studio update'), so a typo in UNSLOTH_STUDIO_HOME must not materialize an empty workspace dir. Brings the Unix flow in line with setup.ps1, which already errors on a missing override root. * llama_cpp: scope orphan-server kill to the active install root _kill_orphaned_servers used to unconditionally include the legacy ~/.unsloth/llama.cpp tree in install_roots, even when the running Studio is in env-override mode and operates out of a custom root. On a single OS user running both a default-install Studio and a custom-root Studio concurrently, the custom Studio would kill the default Studio's llama-server during startup orphan cleanup. Hoist _is_custom_root out of the import try/catch so the legacy- append decision sees it (default to False on ImportError so default mode behaviour is unchanged), and gate the legacy ~/.unsloth/llama.cpp append on `not _is_custom_root`. * install: harden custom-root .venv migration and shim hardlink - install.sh / install.ps1 OLD-layout .venv migration: gate on default-mode only. Without the guard, pointing UNSLOTH_STUDIO_HOME at a workspace that already has .venv (e.g. an unrelated Python project) caused the torch validation to fail and the installer to recursively remove the user's project venv. Mirrors the existing env-mode skip on the CWD-relative venv migration immediately below. - install.ps1 shim hardlink: revert to New-Item -ItemType HardLink -Path. -LiteralPath is not accepted on the HardLink ItemType in any PowerShell version, so the previous form always threw and silently fell back to Copy-Item, breaking hardlink-update propagation. Bracket characters in $ShimExe are still defended by the directory-collision preflight added earlier. - storage_roots.py / unsloth_cli/commands/studio.py: strip whitespace from the UNSLOTH_STUDIO_HOME env var before the truthy check so a blank " " override does not become a real path with trailing spaces (which would silently break every downstream Studio path operation). * Studio paths: tolerate stat / resolve failures during root inference - storage_roots._infer_studio_home_from_venv: wrap the share/studio.conf and bin/shim is_file() sentinel checks in try/except OSError. A PermissionError on a restricted candidate dir would otherwise propagate out of studio_root() and crash module import in run.py / main.py / transformers_version.py / model_config.py at server startup. - llama_cpp._kill_orphaned_servers: broaden the studio_root() guard from ImportError-only to (ImportError, OSError, ValueError) so transient resolve / sentinel failures do not crash the orphan-killer at server startup. Matches _find_llama_server_binary's existing pattern. - llama_cpp._find_llama_server_binary: nest the inner resolve() in its own try/except and fall back to unresolved-path comparison instead of dropping the custom search root entirely. A transient resolve() error on the legacy path no longer loses the custom-root llama.cpp lookup. * Add Studio install-root resilience tests * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * Studio: isolate custom-root installs from default-install state - llama.cpp discovery in env-override mode no longer falls back to the legacy ~/.unsloth/llama.cpp tree. The orphan-cleanup path already excludes that root in custom mode; aligning discovery prevents a custom-root Studio from launching a sibling install's binary it then refuses to manage. Users who want a shared build set UNSLOTH_LLAMA_CPP_PATH explicitly. - Generated POSIX launcher (install.sh heredoc) namespaces LOCK_DIR with a hash of DATA_DIR and persists the launched port to $DATA_DIR/studio.port; in env-override mode the fast-path attaches only to a port we ourselves wrote, never to a sibling Studio that happens to be healthy on 8888..8908. - Generated Windows launcher (install.ps1 heredoc) bakes a per-install $portFile and SHA-256-suffixed mutex name, mirroring the POSIX side; Find-HealthyStudioPort uses the port file in env-override mode. - studio/setup.sh and studio/setup.ps1 require an .unsloth-studio-owned marker before deleting $STUDIO_HOME/.venv_t5*, $STUDIO_HOME/llama.cpp, and the sidecar T5 venvs in env-override mode. The marker is dropped after fresh creation so subsequent runs of 'unsloth studio update' proceed cleanly. Mirrors the existing .venv guard in install.sh. - Wrap bare Path.resolve() calls on the legacy STUDIO_HOME constant in studio/backend/main.py, studio/backend/run.py, and unsloth_cli/commands/studio.py in the same try/except (OSError, ValueError) used adjacently, so a restricted parent or recursive symlink on $HOME does not crash module import / CLI startup. * Studio: guard env-mode workspace against destructive cleanup - install.sh and install.ps1 unconditionally rm -rf / Remove-Item the new-layout $STUDIO_HOME/unsloth_studio when it has a python; in env-override mode that path is a user-chosen workspace, mirroring the .venv migration concern the .venv branch already guards. Refuse to remove an existing $STUDIO_HOME/unsloth_studio that lacks Studio sentinels (share/studio.conf or bin/unsloth). - studio/setup.ps1 only checked Test-Path -PathType Container on the custom root; setup.sh and install.ps1 both also write-probe via WriteAllText / Remove-Item. Add the matching probe so 'unsloth studio update' against an ACL-restricted root fails fast with a clear message instead of erroring later while creating sidecar venvs. * Add Studio install/setup workspace-isolation tests * Studio: tighten installer rationale comments - install.sh: collapse a 5-line restatement into 3 lines, naming env-mode behavior up front and the byte-identical pre-override fallback after. - install.ps1: correct misleading hardlink comment that claimed the directory-collision preflight guards against wildcard expansion; bracket characters in $ShimExe still glob-expand here, with the Copy-Item -LiteralPath fallback handling them. * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * Split: keep only 2 file(s) * Studio: harden env-mode workspace guards across installers and update path Tightens the UNSLOTH_STUDIO_HOME custom-root protections so destructive installer paths cannot displace unrelated user data when the override points at a workspace. install.sh / install.ps1: env-mode sentinel that gates rm -rf $VENV_DIR / Remove-Item $VenvDir now requires share/studio.conf or the bin/unsloth(.exe) shim to be a real file or symlink. Previously a directory at bin/unsloth or bin\unsloth.exe satisfied the check (-e and bare Test-Path accept any path type), so a workspace with unrelated content under unsloth_studio plus a sibling directory at bin/unsloth could be wiped. studio/setup.ps1: stale-venv rebuild branch now mirrors install.ps1's env-mode guard before Remove-Item -LiteralPath $VenvDir -Recurse -Force. Without this, "unsloth studio update" pointed at a custom workspace whose unsloth_studio venv fails torch validation deletes the venv even when the root carries no Studio sentinels. studio/setup.sh / studio/setup.ps1: prebuilt llama.cpp install path now calls _assert_studio_owned_or_absent / Assert-StudioOwnedOrAbsent before invoking install_llama_prebuilt.py, and writes the .unsloth-studio-owned marker on success. install_llama_prebuilt.py uses os.replace() to move any existing install_dir aside before staging, so an unrelated $STUDIO_HOME/llama.cpp could otherwise be displaced before the existing source-build ownership guard ever ran. * Studio: gate ownership guards on canonical custom-root and add venv marker Tightens UNSLOTH_STUDIO_HOME ownership semantics so they fire only for a genuinely custom root, never for an explicit override that resolves to the legacy default. Adds an in-VENV marker that lets a partial install be repaired and provides a strong primary sentinel for the deletion guard. studio/setup.sh + studio/setup.ps1: hoist the canonical $STUDIO_HOME vs legacy-default comparison so it sits next to the marker definition, derive _STUDIO_HOME_IS_CUSTOM / $StudioHomeIsCustom once, and gate the _assert_studio_owned_or_absent / Assert-StudioOwnedOrAbsent helpers and the prebuilt llama.cpp marker writes on that flag instead of raw env-var presence. UNSLOTH_STUDIO_HOME=$HOME/.unsloth/studio (legacy override) no longer trips the guard for pre-PR T5 sidecar venvs or llama.cpp dirs that predate the .unsloth-studio-owned marker. The duplicate canonical block inside the llama.cpp section is removed; the new flag is reused. studio/setup.ps1: Assert-StudioOwnedOrAbsent's marker check now requires -PathType Leaf so a directory at .unsloth-studio-owned cannot satisfy it. The in-place git-sync branch in the source-build path now calls Mark-StudioOwned after a successful sync so a later prebuilt-update path does not fail Assert-StudioOwnedOrAbsent on the same root. install.sh + install.ps1: write $VENV_DIR/.unsloth-studio-owned right after uv venv succeeds and accept it as the primary sentinel in the env-mode deletion guard. This recovers from a partial install that was previously unrepairable, and is a stronger sentinel than sibling shim files (the marker is inside the venv that is about to be wiped, so an unrelated workspace cannot accidentally satisfy it). install.sh: drop the standalone -L test on $STUDIO_HOME/bin/unsloth in the deletion guard. -L returns true for any symlink including symlinks to directories and broken symlinks; -f already accepts the legitimate file-targeted symlink shape created by ln -s at install.sh:1864. * Studio: close residual workspace-isolation gaps for custom roots Four follow-on hardenings that close the remaining cross-root leaks the custom-root install plumbing still left open. studio/setup.ps1 in-place git-sync: when the source-build path finds an existing $LlamaCppDir/.git, it ran git remote set-url, checkout -B, and clean -fdx in place before any ownership check. The previous fix marked the tree as Studio-owned AFTER the sync but did not guard the BEFORE case, so an unrelated workspace .git could be silently rewritten on the first source-build under a custom UNSLOTH_STUDIO_HOME. Add the same Assert-StudioOwnedOrAbsent guard already used by the prebuilt path and the temp-dir swap path (gated on $StudioHomeIsCustom for parity). Launcher port-file workspace isolation: the env-mode launchers' fast path attached to any backend listening on the cached port that returned a healthy /api/health, even when that backend belonged to a different install root. studio/backend/main.py /api/health now returns the resolved studio_root; install.sh _check_health and install.ps1 Test-StudioHealth verify it against UNSLOTH_STUDIO_HOME when set, so a stale studio.port pointing at a sibling Studio is rejected instead of opening the wrong UI. studio/src-tauri preflight + commands: the Tauri desktop app stays on the legacy root by design. process.rs / install.rs / desktop_auth.rs / update.rs already strip UNSLOTH_STUDIO_HOME and STUDIO_HOME from their CLI subprocesses, but preflight.rs run_cli_probe / probe_cli_capability and commands.rs check_install_status did not, so a desktop launch from a shell carrying those env vars produced status reflecting a different root than the desktop manages. Mirror the existing scrub. install.sh shim install: the previous `rm -f -- $_shim_path; ln -s ...` pair leaves a window with no shim if interrupted. Use ln -sfn for an atomic replace; the -n flag prevents descent into a symlink-to-directory target (the existing directory guard above already rejects a real dir). * Studio: replace launcher root verify with hex digest baked at install time The previous launcher identity check returned the absolute resolved Studio install root from /api/health and matched it against $UNSLOTH_STUDIO_HOME in the launcher. Three problems that this commit closes: - POSIX launcher used a raw bash `case` against the JSON-encoded value, so paths containing characters that JSON escapes (e.g. /tmp/back\slash, /tmp/O"Brien) caused the launcher to reject its own healthy backend. - /api/health is unauthenticated and Studio supports `-H 0.0.0.0`, so any reachable client could read the absolute install path (username, home dir, workspace name, CI checkout path). - The verification was gated on $UNSLOTH_STUDIO_HOME being set at runtime, so a default-mode launcher would attach to a sibling env-mode Studio listening on the same port instead of starting its own. The fix replaces the raw path with a SHA-256 hex digest computed at install time and baked into the generated launcher (mirroring how @@DATA_DIR@@ is substituted today): studio/backend/main.py: /api/health now returns `studio_root_id = sha256(str(_studio_root()))` instead of the raw `studio_root` path. install.sh: computes `_css_studio_root_id` once from $STUDIO_HOME using python3, bakes `_EXPECTED_STUDIO_ROOT_ID='@@STUDIO_ROOT_ID@@'` into the launcher heredoc, and adds `s|@@STUDIO_ROOT_ID@@|...|g` to the existing sed pipeline for ALL modes (env / home / default). _check_health verifies the baked id substring-matches the JSON response. Hex-only so no shell or sed escape corner cases. install.ps1: same shape on Windows. SHA256 the $StudioHome bytes, lower hex, bake `$_ExpectedStudioRootId = '...'` into the launcher heredoc. Test-StudioHealth now compares `$resp.studio_root_id -eq $_ExpectedStudioRootId` unconditionally (no special-case for env-mode). Default-mode launchers also bake their expected id, so two coexisting Studio installs on the same machine can no longer cross-attach. * Studio: harden launcher root-id and split install-time mode from runtime env - install.sh launcher: compute studio_root_id with the venv Python (uv-managed systems may not have system python3) and canonicalize STUDIO_HOME with cd -P/pwd -P so default and home-redirect modes match the backend's Path(sys.prefix).resolve() canonicalization. Fail fast instead of silently baking an empty discriminator. - install.sh launcher heredoc: gate PORT_FILE / namespaced LOCK_DIR on a baked install-time mode flag (@@INSTALLED_IS_ENV_MODE@@) instead of the runtime UNSLOTH_STUDIO_HOME variable so a sourced custom-root studio.conf cannot flip a default-mode launcher into env-mode behavior with stale state. - studio/backend/main.py: cache the studio_root_id digest at module load so /api/health does not recompute hashlib + filesystem probes on every poll. - studio/backend/core/inference/llama_cpp.py: widen the studio_root() probe except clause from ImportError to (ImportError, OSError, ValueError) so it matches the sibling _kill_orphaned_servers handler and tolerates Path.resolve failures from broken symlinks or odd codecs. * Studio: align launcher root-id digest with backend canonicalization - studio/backend/main.py: hash the already-resolved _STUDIO_ROOT_RESOLVED instead of recomputing str(_studio_root()); the default fallback in storage_roots returns Path.home()/.unsloth/studio without .resolve(), so on systems where $HOME is a symlink (NFS / AFS / Docker) the cached digest now matches install.sh's cd -P/pwd -P canonicalization and the launcher no longer rejects its own healthy backend. - install.ps1: canonicalize $StudioHome via Resolve-Path before the SHA256 compute (env-mode already resolves at line 121, only default and profile branches were raw); a junctioned USERPROFILE now produces the same digest the backend computes via Path.resolve() for the same install. - install.sh launcher template: substitute the non-user-controlled @@STUDIO_ROOT_ID@@ and @@INSTALLED_IS_ENV_MODE@@ placeholders before the user-controlled @@DATA_DIR@@ pass so a $DATA_DIR that contains the literal placeholder text cannot be mutated by the second sed. * Studio: tighten installer rationale comments * Studio install: extend workspace-guard test coverage Add behavioral coverage for env-mode workspace guards across install.sh, install.ps1, studio/setup.sh, studio/setup.ps1, the launcher root-id discriminator, and the backend's /api/health response. Also refresh the custom-mode llama.cpp resilience assertion so it matches the implementation that intentionally excludes the legacy tree from search_roots. * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * Honor STUDIO_HOME alias, fix workspace-guard test harness, harden rollback The PR title and description promise STUDIO_HOME as a priority-2 alias to UNSLOTH_STUDIO_HOME, but the implementation only read the longer name in all six resolution sites. Wire the alias through install.sh, install.ps1, studio/setup.sh, studio/setup.ps1, the Python storage_roots resolver, and the unsloth_cli studio resolver. UNSLOTH_STUDIO_HOME wins when both are set (more specific signal beats the generic alias). Whitespace-only values are now treated as unset to match the Python resolvers' .strip() semantics, preventing install/runtime layout drift where the installer would create a literal " " directory while the backend fell through to the legacy default. Error messages and the substep status line report the env-var name the user actually set ("UNSLOTH_STUDIO_HOME=..." vs "STUDIO_HOME=...") so diagnostics stay accurate under either spelling. Test harness fix: tests/test_studio_install_workspace_guard.py extracted the install.sh venv-replacement block, but after the merge that block delegates to _start_studio_venv_replacement (defined further up in install.sh, not in the extracted snippet). Five sentinel-positive tests echoed RESULT=ok but never moved $VENV_DIR. Add a single _INSTALL_GUARD_STUBS constant that stands in a minimal mv-based stub plus a no-op substep, and route every inline test script through a new _build_install_guard_script() helper. All 50 tests now pass (was 45/50). Rollback hardening: Start-StudioVenvRollback / Restore-StudioVenvRollback / Complete-StudioVenvRollback in install.ps1 used plain Test-Path, Move-Item, Remove-Item against paths derived from $StudioHome. With a custom UNSLOTH_STUDIO_HOME containing brackets (the very motivation for the broader -LiteralPath sweep this PR set out to do), rollback would silently misbehave under wildcard interpretation, turning a recoverable install error into a destroyed env. Same fix for the --local Tauri overlay block (Test-Path / Copy-Item / Get-FileHash on $VenvDir-derived paths). * Replace studio_root_id path-hash with per-install opaque id The previous design computed studio_root_id as sha256 of the resolved $STUDIO_HOME path, both at install time (baked into the launcher) and at backend startup (returned via /api/health). This worked but had three weaknesses: 1. Information disclosure on -H 0.0.0.0: anyone reaching /api/health could confirm a guessed install path (username, workspace name, etc.) by replaying the same hash. 2. Canonicalization brittleness: launcher (cd -P/pwd -P) and backend (Path.resolve()) had to produce identical strings, which required careful symlink/junction handling on every site (cycles 17-27 of the PR review history were entirely about closing this drift). 3. Stale-launcher attach: an uninstall + reinstall at the same path produced the same hash, so a launcher from the previous install would silently attach to the new (incompatible) backend. Replace the path-hash with a per-install opaque id: - install.sh and install.ps1 generate 32 bytes from the platform CSPRNG (/dev/urandom on POSIX with a python3 secrets fallback; RandomNumberGenerator.Create().GetBytes on Windows) and persist it to $STUDIO_HOME/share/studio_install_id with mode 0600. Atomic temp-file-rename so a crash mid-install can't leave a half-written id. The check 'if [ ! -s "$_css_id_file" ]' / Test-Path makes generation idempotent across re-runs (so re-running install.sh doesn't invalidate previously-baked launchers in the same install root). - studio/backend/main.py replaces hashlib.sha256 with _read_studio_install_id(), which reads $STUDIO_HOME/share/studio_install_id once at module load. Validates the content against ^[0-9a-f]{64}$ so malformed/truncated/uppercase/wrong-length content returns "" and triggers the launcher's existing "no baked id, accept any healthy Unsloth backend" fallback path. - /api/health field name (studio_root_id) and wire format (64 hex chars) preserved for compatibility with launchers already shipped via earlier PR iterations. Tests: - Drop test_install_sh_root_id_matches_backend_resolved_under_symlinked_home and test_install_ps1_canonicalizes_studio_home_before_root_id_hash -- the entire reason these existed (cd -P/Resolve-Path/Path.resolve() digest agreement under symlinks/junctions) is moot when the id comes from a file rather than from the path. - Drop test_main_py_studio_root_id_hashes_resolved_root_not_unresolved (no more hashing). - Rewrite test_main_py_studio_root_id_caches_at_module_load to assert the file-read pattern; add test_main_py_read_studio_install_id_validates_hex_and_handles_missing to pin the exact rejection rules (empty / non-hex / wrong case / wrong length all -> ""). - Rewrite test_install_sh_create_shortcuts_uses_venv_python_first as test_install_sh_create_shortcuts_seeds_id_from_csprng_with_python_fallback with a behavioral subprocess check that re-invocation is idempotent. - Rename test_check_health_handles_path_with_backslash_via_hash to test_check_health_handles_arbitrary_id_token (the JSON-escape concern it pinned is preserved -- ids are hex-only by construction -- but the test no longer derives the id from a path). - Add test_install_sh_install_id_survives_symlinked_studio_home as a regression test pinning that the new design has zero canonicalization drift across symlinked parents. - Update test_install_sh_bakes_studio_root_id_into_launcher and test_install_ps1_bakes_studio_root_id_into_launcher to assert the CSPRNG seed and the file location. 49/49 tests pass. Behavioral verification: install.sh-style generation is idempotent across runs, three parallel installs at different roots get distinct ids, reinstall at the same path produces a new id (so stale launchers correctly fail to attach to the new backend), and symlinked-\$HOME no longer causes launcher/backend disagreement. * [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 --- install.ps1 | 390 ++++++- install.sh | 415 ++++++- studio/backend/core/inference/llama_cpp.py | 77 +- studio/backend/main.py | 59 + studio/backend/run.py | 22 +- studio/backend/utils/models/model_config.py | 4 +- studio/backend/utils/paths/storage_roots.py | 44 +- studio/backend/utils/transformers_version.py | 8 +- studio/setup.ps1 | 211 +++- studio/setup.sh | 84 +- studio/src-tauri/src/commands.rs | 5 + studio/src-tauri/src/desktop_auth.rs | 5 + studio/src-tauri/src/install.rs | 5 + studio/src-tauri/src/preflight.rs | 10 + studio/src-tauri/src/process.rs | 6 + studio/src-tauri/src/update.rs | 5 + tests/test_studio_install_workspace_guard.py | 1021 ++++++++++++++++++ tests/test_studio_root_resilience.py | 154 +++ unsloth_cli/commands/studio.py | 70 +- 19 files changed, 2423 insertions(+), 172 deletions(-) create mode 100644 tests/test_studio_install_workspace_guard.py create mode 100644 tests/test_studio_root_resilience.py diff --git a/install.ps1 b/install.ps1 index d5db1785a1..ef87c5ed08 100644 --- a/install.ps1 +++ b/install.ps1 @@ -3,6 +3,11 @@ # Local: Set-ExecutionPolicy -Scope Process -ExecutionPolicy Bypass; .\install.ps1 --local # NoTorch: .\install.ps1 --no-torch (skip PyTorch, GGUF-only mode) # Test: .\install.ps1 --package roland-sloth +# +# Env vars (priority: UNSLOTH_STUDIO_HOME > STUDIO_HOME > USERPROFILE-redirect > default): +# UNSLOTH_STUDIO_HOME / STUDIO_HOME = path -> install under that path +# (DataDir nests inside; user PATH not modified persistently). +# Default ($USERPROFILE\.unsloth\studio) is preserved when no env var is set. function Install-UnslothStudio { $ErrorActionPreference = "Stop" @@ -126,7 +131,94 @@ function Install-UnslothStudio { } $PythonVersion = "3.13" - $StudioHome = Join-Path $env:USERPROFILE ".unsloth\studio" + + # Resolve install destinations. Priority: UNSLOTH_STUDIO_HOME, then + # STUDIO_HOME alias, then USERPROFILE-redirect, then default. + # Reject whitespace-only values so " " is treated as unset (matches the + # Python resolvers' .strip()), preventing install/runtime layout drift. + $envOverrideVar = $null + $envOverride = $null + if (-not [string]::IsNullOrWhiteSpace($env:UNSLOTH_STUDIO_HOME)) { + $envOverrideVar = "UNSLOTH_STUDIO_HOME" + $envOverride = $env:UNSLOTH_STUDIO_HOME.Trim() + } elseif (-not [string]::IsNullOrWhiteSpace($env:STUDIO_HOME)) { + $envOverrideVar = "STUDIO_HOME" + $envOverride = $env:STUDIO_HOME.Trim() + } + + # Custom Studio roots are not supported with --tauri (desktop app still + # resolves %USERPROFILE%\.unsloth\studio). Pass through if override == legacy. + if ($TauriMode -and $envOverride) { + $_tauriOverride = $envOverride + if ($_tauriOverride -eq "~" -or $_tauriOverride -like "~/*" -or $_tauriOverride -like "~\*") { + $_tauriOverride = (Join-Path $env:USERPROFILE $_tauriOverride.Substring(1).TrimStart('/','\')) + } + try { + $_tauriOverride = [System.IO.Path]::GetFullPath($_tauriOverride) + } catch {} + $_legacyTauriRoot = Join-Path $env:USERPROFILE ".unsloth\studio" + try { + $_legacyTauriRoot = [System.IO.Path]::GetFullPath($_legacyTauriRoot) + } catch {} + # Strip trailing separators so ".../studio\" matches ".../studio". + $_trimSeps = @( + [System.IO.Path]::DirectorySeparatorChar, + [System.IO.Path]::AltDirectorySeparatorChar + ) + $_tauriOverride = $_tauriOverride.TrimEnd($_trimSeps) + $_legacyTauriRoot = $_legacyTauriRoot.TrimEnd($_trimSeps) + if ($_tauriOverride -ne $_legacyTauriRoot) { + Write-Host "ERROR: $envOverrideVar is not supported with --tauri." -ForegroundColor Red + Write-Host " The desktop app still uses the legacy %USERPROFILE%\.unsloth\studio root." -ForegroundColor Red + Write-Host " Run install.ps1 without --tauri for custom-root shell installs," -ForegroundColor Yellow + Write-Host " or unset the env var for default desktop installs." -ForegroundColor Yellow + throw "$envOverrideVar is not supported with --tauri." + } + } + + $defaultProfile = $null + try { $defaultProfile = [Environment]::GetFolderPath("UserProfile") } catch {} + + # LOCALAPPDATA may be unset in service / CI contexts; Join-Path would abort + # under ErrorActionPreference=Stop without this guard. + $defaultDataDir = if ($env:LOCALAPPDATA -and -not [string]::IsNullOrWhiteSpace($env:LOCALAPPDATA)) { + Join-Path $env:LOCALAPPDATA "Unsloth Studio" + } else { $null } + + if ($envOverride) { + # Tilde expansion: env vars aren't subject to it when quoted on assignment. + if ($envOverride -eq "~" -or $envOverride -like "~/*" -or $envOverride -like "~\*") { + $envOverride = (Join-Path $env:USERPROFILE $envOverride.Substring(1).TrimStart('/','\')) + } + try { + # .NET API: New-Item -Path treats brackets as wildcards and has no + # -LiteralPath in PS 5.1, so a root like C:\studio[abc] would fail. + [System.IO.Directory]::CreateDirectory($envOverride) | Out-Null + $StudioHome = (Resolve-Path -LiteralPath $envOverride).Path + } catch { + Write-Host "ERROR: $envOverrideVar=$envOverride cannot be created or accessed." -ForegroundColor Red + throw "$envOverrideVar=$envOverride cannot be created or accessed." + } + $probe = Join-Path $StudioHome (".unsloth-write-probe-" + [guid]::NewGuid()) + try { + # WriteAllText: literal-path safe + closes handle so Remove-Item works. + [System.IO.File]::WriteAllText($probe, "") + Remove-Item -LiteralPath $probe -Force -ErrorAction SilentlyContinue + } catch { + Write-Host "ERROR: $envOverrideVar=$StudioHome is not writable." -ForegroundColor Red + throw "$envOverrideVar=$StudioHome is not writable." + } + $StudioDataDir = Join-Path $StudioHome "share" + $StudioRedirectMode = 'env' + } elseif ($defaultProfile -and $env:USERPROFILE -and ($env:USERPROFILE -ne $defaultProfile)) { + $StudioHome = Join-Path $env:USERPROFILE ".unsloth\studio" + $StudioDataDir = $defaultDataDir + $StudioRedirectMode = 'profile' + } else { + $StudioHome = Join-Path $env:USERPROFILE ".unsloth\studio" + $StudioDataDir = $defaultDataDir + $StudioRedirectMode = 'default' + } $VenvDir = Join-Path $StudioHome "unsloth_studio" $Rule = [string]::new([char]0x2500, 52) @@ -378,24 +470,24 @@ function Install-UnslothStudio { [Parameter(Mandatory = $true)][string]$UnslothExePath ) - if (-not (Test-Path $UnslothExePath)) { + if (-not (Test-Path -LiteralPath $UnslothExePath)) { substep "cannot create shortcuts, unsloth.exe not found at $UnslothExePath" "Yellow" return } try { # Persist an absolute path in launcher scripts so shortcut working # directory changes do not break process startup. - $UnslothExePath = (Resolve-Path $UnslothExePath).Path + $UnslothExePath = (Resolve-Path -LiteralPath $UnslothExePath).Path # Escape for single-quoted embedding in generated launcher script. # This prevents runtime variable expansion for paths containing '$'. $SingleQuotedExePath = $UnslothExePath -replace "'", "''" - $localAppDataDir = $env:LOCALAPPDATA - if (-not $localAppDataDir -or [string]::IsNullOrWhiteSpace($localAppDataDir)) { - substep "LOCALAPPDATA path unavailable; skipped shortcut creation" "Yellow" + # $StudioDataDir = LOCALAPPDATA\Unsloth Studio, or $StudioHome\share in env-mode. + if (-not $StudioDataDir -or [string]::IsNullOrWhiteSpace($StudioDataDir)) { + substep "DataDir path unavailable; skipped shortcut creation" "Yellow" return } - $appDir = Join-Path $localAppDataDir "Unsloth Studio" + $appDir = $StudioDataDir $launcherPs1 = Join-Path $appDir "launch-studio.ps1" $launcherVbs = Join-Path $appDir "launch-studio.vbs" $desktopDir = [Environment]::GetFolderPath("Desktop") @@ -427,23 +519,89 @@ function Install-UnslothStudio { } $iconUrl = "https://raw.githubusercontent.com/unslothai/unsloth/main/studio/frontend/public/unsloth.ico" - if (-not (Test-Path $appDir)) { - New-Item -ItemType Directory -Path $appDir -Force | Out-Null + if (-not (Test-Path -LiteralPath $appDir)) { + [System.IO.Directory]::CreateDirectory($appDir) | Out-Null + } + + # Same-install discriminator: per-install opaque id written once at + # install time and read by both this launcher and the backend + # (/api/health). Replaces the older sha256(resolved $StudioHome) + # scheme to (a) avoid leaking the install path on -H 0.0.0.0 + # deployments and (b) sidestep launcher/backend canonicalization + # drift (Resolve-Path vs Path.resolve() junction handling). Lives + # at $StudioHome\share\ (not $appDir) so the backend can find it + # via _STUDIO_ROOT_RESOLVED / "share" / "studio_install_id" + # regardless of mode. 32 bytes of crypto random -> 64 hex chars. + $_studioIdDir = Join-Path $StudioHome "share" + if (-not (Test-Path -LiteralPath $_studioIdDir)) { + [System.IO.Directory]::CreateDirectory($_studioIdDir) | Out-Null + } + $_studioIdFile = Join-Path $_studioIdDir "studio_install_id" + $_studioRootId = "" + if ((Test-Path -LiteralPath $_studioIdFile) -and ` + ((Get-Item -LiteralPath $_studioIdFile).Length -gt 0)) { + $_studioRootId = ([System.IO.File]::ReadAllText($_studioIdFile)).Trim() + } + if (-not $_studioRootId) { + $_idBytes = New-Object byte[] 32 + [Security.Cryptography.RandomNumberGenerator]::Create().GetBytes($_idBytes) + $_studioRootId = -join ($_idBytes | ForEach-Object { $_.ToString('x2') }) + # Atomic write: write to a temp sibling then rename, so a partial + # install cannot leave a half-written id. + $_idTmp = $_studioIdFile + ".$PID.tmp" + [System.IO.File]::WriteAllText($_idTmp, $_studioRootId) + Move-Item -LiteralPath $_idTmp -Destination $_studioIdFile -Force + } + + # Env-mode: persist UNSLOTH_STUDIO_HOME (and llama path) so fresh + # shells don't need to re-export, and bake per-install $portFile / + # $mutexName so concurrent custom-root launchers cannot serialize + # through one global mutex on 8888..8908. Default installs get an + # empty prefix to match pre-PR behavior. + $studioHomeExport = if ($StudioRedirectMode -eq 'env') { + # When override == legacy default, llama.cpp stays at + # ~/.unsloth/llama.cpp (one shared build). Canonicalize the + # legacy side so the comparison survives path normalization. + $_legacyStudio = Join-Path $env:USERPROFILE ".unsloth\studio" + if (Test-Path -LiteralPath $_legacyStudio -PathType Container) { + $_legacyStudio = (Resolve-Path -LiteralPath $_legacyStudio).Path + } + $_llamaPath = if ($StudioHome -eq $_legacyStudio) { + Join-Path $env:USERPROFILE ".unsloth\llama.cpp" + } else { + Join-Path $StudioHome "llama.cpp" + } + $_sq = $StudioHome -replace "'", "''" + $_llama = $_llamaPath -replace "'", "''" + $_appDirSq = $appDir -replace "'", "''" + $_appBytes = [Text.Encoding]::UTF8.GetBytes($appDir) + $_appHash = ([BitConverter]::ToString( + [Security.Cryptography.SHA256]::Create().ComputeHash($_appBytes) + ) -replace '-', '').Substring(0, 16) + # UNSLOTH_LLAMA_CPP_PATH is a pre-existing user override; only default if unset. + "`$env:UNSLOTH_STUDIO_HOME = '$_sq'`nif (-not `$env:UNSLOTH_LLAMA_CPP_PATH) {`n `$env:UNSLOTH_LLAMA_CPP_PATH = '$_llama'`n}`n`$portFile = '$_appDirSq\studio.port'`n`$mutexName = 'Local\UnslothStudioLauncher-$_appHash'`n" + } else { + "`$portFile = `$null`n`$mutexName = 'Local\UnslothStudioLauncher'`n" } $launcherContent = @" -`$ErrorActionPreference = 'Stop' +$studioHomeExport`$ErrorActionPreference = 'Stop' `$basePort = 8888 `$maxPortOffset = 20 `$timeoutSec = 60 `$pollIntervalMs = 1000 +`$_ExpectedStudioRootId = '$_studioRootId' function Test-StudioHealth { param([Parameter(Mandatory = `$true)][int]`$Port) try { `$url = "http://127.0.0.1:`$Port/api/health" `$resp = Invoke-RestMethod -Uri `$url -TimeoutSec 1 -Method Get - return (`$resp -and `$resp.status -eq 'healthy' -and `$resp.service -eq 'Unsloth UI Backend') + if (-not (`$resp -and `$resp.status -eq 'healthy' -and `$resp.service -eq 'Unsloth UI Backend')) { return `$false } + # why: verify the backend belongs to THIS install via the install-time + # hex digest; raw path is not leaked over /api/health. + if (`$_ExpectedStudioRootId -and `$resp.studio_root_id -ne `$_ExpectedStudioRootId) { return `$false } + return `$true } catch { return `$false } @@ -469,6 +627,17 @@ function Get-CandidatePorts { } function Find-HealthyStudioPort { + if (`$portFile) { + if (Test-Path -LiteralPath `$portFile) { + `$cached = Get-Content -LiteralPath `$portFile -ErrorAction SilentlyContinue | Select-Object -First 1 + if (`$cached -match '^\d+`$') { + `$cachedPort = [int]`$cached + if (Test-StudioHealth -Port `$cachedPort) { return `$cachedPort } + Remove-Item -LiteralPath `$portFile -Force -ErrorAction SilentlyContinue + } + } + return `$null + } foreach (`$candidate in (Get-CandidatePorts)) { if (Test-StudioHealth -Port `$candidate) { return `$candidate @@ -522,7 +691,7 @@ if (`$existingPort) { exit 0 } -`$launchMutex = [System.Threading.Mutex]::new(`$false, 'Local\UnslothStudioLauncher') +`$launchMutex = [System.Threading.Mutex]::new(`$false, `$mutexName) `$haveMutex = `$false try { try { @@ -552,7 +721,9 @@ try { } catch {} exit 1 } - `$studioCommand = '& "' + `$studioExe + '" studio -p ' + `$launchPort + # Single-quote the path in the child -Command so `$` / backtick in custom + # roots don't get reparsed; double any apostrophes so 'O''Brien' survives. + `$studioCommand = "& '" + (`$studioExe -replace "'", "''") + "' studio -p " + `$launchPort `$launchArgs = @( '-NoExit', '-NoProfile', @@ -576,9 +747,13 @@ try { `$browserOpened = `$false `$deadline = (Get-Date).AddSeconds(`$timeoutSec) while ((Get-Date) -lt `$deadline) { - `$healthyPort = Find-HealthyStudioPort - if (`$healthyPort) { - Start-Process "http://localhost:`$healthyPort" + if (Test-StudioHealth -Port `$launchPort) { + if (`$portFile) { + try { + [System.IO.File]::WriteAllText(`$portFile, "`$launchPort`n") + } catch {} + } + Start-Process "http://localhost:`$launchPort" `$browserOpened = `$true break } @@ -613,19 +788,19 @@ cmd = "powershell -NoProfile -ExecutionPolicy Bypass -WindowStyle Hidden -File " shell.Run cmd, 0, False "@ # WSH handles UTF-16LE reliably for .vbs files with non-ASCII paths. - Set-Content -Path $launcherVbs -Value $vbsContent -Encoding Unicode -Force + Set-Content -LiteralPath $launcherVbs -Value $vbsContent -Encoding Unicode -Force # Prefer bundled icon from local clone/dev installs. # If not available, best-effort download from raw GitHub. # We only attach the icon if the resulting file has a valid ICO header. $hasValidIcon = $false - if ($bundledIcon -and (Test-Path $bundledIcon)) { + if ($bundledIcon -and (Test-Path -LiteralPath $bundledIcon)) { try { - Copy-Item -Path $bundledIcon -Destination $iconPath -Force + Copy-Item -LiteralPath $bundledIcon -Destination $iconPath -Force } catch { Write-Host "[DEBUG] Error copying bundled icon: $($_.Exception.Message)" -ForegroundColor DarkGray } - } elseif (-not (Test-Path $iconPath)) { + } elseif (-not (Test-Path -LiteralPath $iconPath)) { try { Invoke-WebRequest -Uri $iconUrl -OutFile $iconPath -UseBasicParsing } catch { @@ -633,7 +808,7 @@ shell.Run cmd, 0, False } } - if (Test-Path $iconPath) { + if (Test-Path -LiteralPath $iconPath) { try { $bytes = [System.IO.File]::ReadAllBytes($iconPath) if ( @@ -645,14 +820,21 @@ shell.Run cmd, 0, False ) { $hasValidIcon = $true } else { - Remove-Item $iconPath -Force -ErrorAction SilentlyContinue + Remove-Item -LiteralPath $iconPath -Force -ErrorAction SilentlyContinue } } catch { Write-Host "[DEBUG] Error validating or removing icon: $($_.Exception.Message)" -ForegroundColor DarkGray - Remove-Item $iconPath -Force -ErrorAction SilentlyContinue + Remove-Item -LiteralPath $iconPath -Force -ErrorAction SilentlyContinue } } + # Env-mode: skip persistent Desktop / Start Menu .lnk shortcuts + # that may point at a deleted workspace; launcher + icon stay. + if ($StudioRedirectMode -eq 'env') { + substep "wrote launcher at $launcherPs1 (persistent shortcuts skipped in env-override mode)" + return + } + $wscriptExe = Join-Path $env:SystemRoot "System32\wscript.exe" $shortcutArgs = "//B //Nologo `"$launcherVbs`"" @@ -850,8 +1032,9 @@ shell.Run cmd, 0, False # Pass the resolved executable path to uv so it does not re-resolve # a version string back to a conda interpreter. Write-TauriLog "STEP" "Creating virtual environment" - if (-not (Test-Path $StudioHome)) { - New-Item -ItemType Directory -Path $StudioHome -Force | Out-Null + if (-not (Test-Path -LiteralPath $StudioHome)) { + # .NET API: New-Item -Path treats brackets as wildcards. + [System.IO.Directory]::CreateDirectory($StudioHome) | Out-Null } $VenvPython = Join-Path $VenvDir "Scripts\python.exe" @@ -865,11 +1048,13 @@ shell.Run cmd, 0, False $stamp = Get-Date -Format "yyyyMMddHHmmss" $candidate = Join-Path $StudioHome "unsloth_studio.rollback.$stamp.$PID" $suffix = 0 - while (Test-Path $candidate) { + # -LiteralPath: a custom $StudioHome may contain [ ] * ? which + # plain Test-Path / Move-Item would interpret as wildcards. + while (Test-Path -LiteralPath $candidate) { $suffix++ $candidate = Join-Path $StudioHome "unsloth_studio.rollback.$stamp.$PID.$suffix" } - Move-Item -Path $ExistingDir -Destination $candidate -ErrorAction Stop + Move-Item -LiteralPath $ExistingDir -Destination $candidate -ErrorAction Stop $script:StudioVenvRollbackDir = $candidate $script:StudioVenvRollbackTarget = $ExistingDir $script:StudioVenvRollbackActive = $true @@ -880,16 +1065,16 @@ shell.Run cmd, 0, False if (-not $script:StudioVenvRollbackActive) { return } $backup = $script:StudioVenvRollbackDir $target = $script:StudioVenvRollbackTarget - if (-not $backup -or -not (Test-Path $backup)) { + if (-not $backup -or -not (Test-Path -LiteralPath $backup)) { $script:StudioVenvRollbackActive = $false return } substep "restoring previous environment after failed install..." "Yellow" try { - if (Test-Path $target) { - Remove-Item -Recurse -Force $target -ErrorAction SilentlyContinue + if (Test-Path -LiteralPath $target) { + Remove-Item -LiteralPath $target -Recurse -Force -ErrorAction SilentlyContinue } - Move-Item -Path $backup -Destination $target -Force -ErrorAction Stop + Move-Item -LiteralPath $backup -Destination $target -Force -ErrorAction Stop substep "restored previous environment" $script:StudioVenvRollbackActive = $false $script:StudioVenvRollbackDir = $null @@ -902,14 +1087,29 @@ shell.Run cmd, 0, False function Complete-StudioVenvRollback { if (-not $script:StudioVenvRollbackActive) { return } $backup = $script:StudioVenvRollbackDir - if ($backup -and (Test-Path $backup)) { - Remove-Item -Recurse -Force $backup -ErrorAction SilentlyContinue + if ($backup -and (Test-Path -LiteralPath $backup)) { + Remove-Item -LiteralPath $backup -Recurse -Force -ErrorAction SilentlyContinue } $script:StudioVenvRollbackActive = $false $script:StudioVenvRollbackDir = $null } - if (Test-Path $VenvPython) { + if (Test-Path -LiteralPath $VenvPython) { + # why: matching guard to the .venv branch below -- in env-mode + # $StudioHome is a user-chosen workspace, so refuse to nuke an + # existing $StudioHome\unsloth_studio that lacks Studio sentinels. + # -PathType Leaf rejects a directory at the sentinel path. Accept the + # in-VENV ownership marker so partial-install retries are not blocked. + if ( + $StudioRedirectMode -eq 'env' -and + -not (Test-Path -LiteralPath (Join-Path $VenvDir ".unsloth-studio-owned") -PathType Leaf) -and + -not (Test-Path -LiteralPath (Join-Path $StudioHome "share\studio.conf") -PathType Leaf) -and + -not (Test-Path -LiteralPath (Join-Path $StudioHome "bin\unsloth.exe") -PathType Leaf) + ) { + Write-Host "[ERROR] $VenvDir already exists but does not look like an Unsloth Studio install." -ForegroundColor Red + Write-Host " Move it aside or choose an empty UNSLOTH_STUDIO_HOME." -ForegroundColor Yellow + throw "Refusing to delete non-Studio venv at $VenvDir" + } # New layout already exists -- replace only after preserving rollback copy. substep "preserving existing environment for rollback..." try { @@ -918,8 +1118,13 @@ shell.Run cmd, 0, False Write-Host "[ERROR] Could not prepare existing environment for reinstall: $($_.Exception.Message)" -ForegroundColor Red return (Exit-InstallFailure "Could not prepare existing environment for reinstall") } - } elseif (Test-Path (Join-Path $StudioHome ".venv\Scripts\python.exe")) { - # Old layout (~/.unsloth/studio/.venv) exists -- validate before migrating + } elseif ( + $StudioRedirectMode -ne 'env' ` + -and (Test-Path -LiteralPath (Join-Path $StudioHome ".venv\Scripts\python.exe")) + ) { + # Old layout (~/.unsloth/studio/.venv) exists -- validate before migrating. + # Skip in env-mode so we don't blow away an unrelated .venv at the + # workspace root (e.g. user's existing project Python venv). $OldVenv = Join-Path $StudioHome ".venv" $OldPy = Join-Path $OldVenv "Scripts\python.exe" substep "found legacy Studio environment, validating..." @@ -936,24 +1141,29 @@ shell.Run cmd, 0, False $ErrorActionPreference = $prevEAP2 if ($legacyOk) { substep "legacy environment is healthy -- migrating..." - Move-Item -Path $OldVenv -Destination $VenvDir -Force + Move-Item -LiteralPath $OldVenv -Destination $VenvDir -Force substep "moved .venv -> unsloth_studio" $_Migrated = $true } else { substep "legacy environment failed validation -- creating fresh environment" "Yellow" $invalidVenv = Join-Path $StudioHome (".venv.invalid.{0}.{1}" -f (Get-Date -Format "yyyyMMddHHmmss"), $PID) - Move-Item -Path $OldVenv -Destination $invalidVenv -Force -ErrorAction SilentlyContinue + Move-Item -LiteralPath $OldVenv -Destination $invalidVenv -Force -ErrorAction SilentlyContinue } - } elseif (Test-Path (Join-Path $env:USERPROFILE "unsloth_studio\Scripts\python.exe")) { - # CWD-relative venv from old install.ps1 -- migrate to absolute path + } elseif ( + $StudioRedirectMode -ne 'env' ` + -and (Test-Path -LiteralPath (Join-Path $env:USERPROFILE "unsloth_studio\Scripts\python.exe")) + ) { + # CWD-relative venv from old install.ps1 -> migrate to absolute path. + # Skip in env-mode so we don't relocate the default-install venv into + # the workspace root. $CwdVenv = Join-Path $env:USERPROFILE "unsloth_studio" substep "found CWD-relative Studio environment, migrating to $VenvDir..." - Move-Item -Path $CwdVenv -Destination $VenvDir -Force + Move-Item -LiteralPath $CwdVenv -Destination $VenvDir -Force substep "moved ~/unsloth_studio -> ~/.unsloth/studio/unsloth_studio" $_Migrated = $true } - if (-not (Test-Path $VenvPython)) { + if (-not (Test-Path -LiteralPath $VenvPython)) { step "venv" "creating Python $($DetectedPython.Version) virtual environment" substep "$VenvDir" $venvExit = Invoke-InstallCommand { uv venv $VenvDir --python "$($DetectedPython.Path)" } @@ -966,6 +1176,13 @@ shell.Run cmd, 0, False substep "$VenvDir" } + # Mark the freshly-created venv as Studio-owned so a partial install can be + # repaired by re-running install.ps1; the env-mode deletion guard above + # accepts this marker as the primary sentinel. + if (Test-Path -LiteralPath $VenvDir -PathType Container) { + try { [System.IO.File]::WriteAllText((Join-Path $VenvDir ".unsloth-studio-owned"), "") } catch {} + } + # ── Detect GPU (robust: PATH + hardcoded fallback paths, mirrors setup.ps1) ── $HasNvidiaSmi = $false $NvidiaSmiExe = $null @@ -1054,7 +1271,7 @@ shell.Run cmd, 0, False if ($StudioLocalInstall -and (Test-Path (Join-Path $RepoRoot "studio\backend\requirements\no-torch-runtime.txt"))) { return Join-Path $RepoRoot "studio\backend\requirements\no-torch-runtime.txt" } - $installed = Get-ChildItem -Path $VenvDir -Recurse -Filter "no-torch-runtime.txt" -ErrorAction SilentlyContinue | + $installed = Get-ChildItem -LiteralPath $VenvDir -Recurse -Filter "no-torch-runtime.txt" -ErrorAction SilentlyContinue | Where-Object { $_.FullName -like "*studio*backend*requirements*no-torch-runtime.txt" } | Select-Object -ExpandProperty FullName -First 1 return $installed @@ -1192,23 +1409,25 @@ shell.Run cmd, 0, False foreach ($rel in $overlayMap.Keys) { $src = Join-Path $scriptDir $rel $dst = Join-Path $VenvDir $overlayMap[$rel] - if (-not (Test-Path $src)) { continue } + # -LiteralPath: $VenvDir derives from $StudioHome which may + # contain [ ] * ? when the user overrode UNSLOTH_STUDIO_HOME. + if (-not (Test-Path -LiteralPath $src)) { continue } $dstParent = Split-Path -Parent $dst - if (-not (Test-Path $dstParent)) { + if (-not (Test-Path -LiteralPath $dstParent)) { Write-Host "[WARN] Overlay target dir missing: $dstParent; studio setup may use stale bundled file" -ForegroundColor Yellow continue } try { - if (-not (Test-Path $dst)) { + if (-not (Test-Path -LiteralPath $dst)) { # Backfill: target file missing but parent dir exists. - Copy-Item $src $dst -Force + Copy-Item -LiteralPath $src -Destination $dst -Force substep ("backfilled bundled " + (Split-Path -Leaf $rel)) } else { # Hash-compare so re-runs are no-ops when files already match. - $srcHash = (Get-FileHash $src -Algorithm SHA256).Hash - $dstHash = (Get-FileHash $dst -Algorithm SHA256).Hash + $srcHash = (Get-FileHash -LiteralPath $src -Algorithm SHA256).Hash + $dstHash = (Get-FileHash -LiteralPath $dst -Algorithm SHA256).Hash if ($srcHash -ne $dstHash) { - Copy-Item $src $dst -Force + Copy-Item -LiteralPath $src -Destination $dst -Force substep ("applied bundled " + (Split-Path -Leaf $rel)) } } @@ -1225,7 +1444,8 @@ shell.Run cmd, 0, False Write-TauriLog "STEP" "Running studio setup" step "setup" "running unsloth studio setup..." $UnslothExe = Join-Path $VenvDir "Scripts\unsloth.exe" - if (-not (Test-Path $UnslothExe)) { + if (-not (Test-Path -LiteralPath $UnslothExe)) { + Write-TauriLog "ERROR" "unsloth CLI was not installed correctly" Write-Host "[ERROR] unsloth CLI was not installed correctly." -ForegroundColor Red Write-Host " Expected: $UnslothExe" -ForegroundColor Yellow Write-Host " This usually means an older unsloth version was installed that does not include the Studio CLI." -ForegroundColor Yellow @@ -1250,6 +1470,15 @@ shell.Run cmd, 0, False # Use 'studio setup' (not 'studio update') because 'update' pops # SKIP_STUDIO_BASE, which would cause redundant package reinstallation # and bypass the fast-path version check from PR #4667. + # Propagate UNSLOTH_STUDIO_HOME only for env-override installs; otherwise + # an inherited value would put llama.cpp in the wrong place. + $previousUnslothStudioHome = $env:UNSLOTH_STUDIO_HOME + $hadPreviousUnslothStudioHome = ($null -ne $previousUnslothStudioHome) + if ($StudioRedirectMode -eq 'env') { + $env:UNSLOTH_STUDIO_HOME = $StudioHome + } else { + Remove-Item Env:UNSLOTH_STUDIO_HOME -ErrorAction SilentlyContinue + } $studioArgs = @('studio', 'setup') if ($script:UnslothVerbose) { $studioArgs += '--verbose' } $env:UNSLOTH_INSTALL_ROLLBACK_MANAGED = "1" @@ -1257,6 +1486,11 @@ shell.Run cmd, 0, False & $UnslothExe @studioArgs $setupExit = $LASTEXITCODE } finally { + if ($hadPreviousUnslothStudioHome) { + $env:UNSLOTH_STUDIO_HOME = $previousUnslothStudioHome + } else { + Remove-Item Env:UNSLOTH_STUDIO_HOME -ErrorAction SilentlyContinue + } Remove-Item Env:UNSLOTH_INSTALL_ROLLBACK_MANAGED -ErrorAction SilentlyContinue } if ($setupExit -ne 0) { @@ -1301,20 +1535,32 @@ shell.Run cmd, 0, False } } catch { } $ShimDir = Join-Path $StudioHome "bin" - New-Item -ItemType Directory -Force -Path $ShimDir | Out-Null + [System.IO.Directory]::CreateDirectory($ShimDir) | Out-Null $ShimExe = Join-Path $ShimDir "unsloth.exe" + # Fatal preflight outside the lock-handling try/catch -- a directory at + # the shim path must not be downgraded to "Continuing with the existing + # launcher", or the install finishes with no usable shim. + if (Test-Path -LiteralPath $ShimExe -PathType Container) { + Write-Host "[ERROR] Cannot create unsloth launcher: $ShimExe is a directory." -ForegroundColor Red + Write-Host " Move or remove it manually, then re-run the installer." -ForegroundColor Yellow + throw "Cannot create unsloth launcher: $ShimExe is a directory." + } # try/catch: if unsloth.exe is locked (Studio running), keep the old shim. $shimUpdated = $false try { - if (Test-Path $ShimExe) { Remove-Item $ShimExe -Force -ErrorAction Stop } + if (Test-Path -LiteralPath $ShimExe) { Remove-Item -LiteralPath $ShimExe -Force -ErrorAction Stop } try { + # New-Item -ItemType HardLink does NOT accept -LiteralPath in any + # PowerShell version, so use -Path. Wildcards in $ShimExe (e.g. + # brackets in custom roots) glob-expand here and fall through to + # the Copy-Item -LiteralPath fallback below. New-Item -ItemType HardLink -Path $ShimExe -Target $UnslothExe -ErrorAction Stop | Out-Null } catch { - Copy-Item -Path $UnslothExe -Destination $ShimExe -Force -ErrorAction Stop # fallback: copy + Copy-Item -LiteralPath $UnslothExe -Destination $ShimExe -Force -ErrorAction Stop # fallback: copy } $shimUpdated = $true } catch { - if (Test-Path $ShimExe) { + if (Test-Path -LiteralPath $ShimExe) { Write-Host "[WARN] Could not refresh unsloth launcher at $ShimExe." -ForegroundColor Yellow Write-Host " This usually means a running 'unsloth studio' process still holds the file open." -ForegroundColor Yellow Write-Host " Close Studio and re-run the installer to pick up the latest launcher." -ForegroundColor Yellow @@ -1325,10 +1571,13 @@ shell.Run cmd, 0, False Write-Host " Launch unsloth studio directly via '$UnslothExe' until the next successful install." -ForegroundColor Yellow } } - # Only add to PATH when the launcher actually exists on disk. + # Add to PATH only when launcher exists. Env-mode: session-only export, + # no registry change (workspace path may be deleted later). $pathAdded = $false - if (Test-Path $ShimExe) { - $pathAdded = Add-ToUserPath -Directory $ShimDir -Position 'Prepend' + if (Test-Path -LiteralPath $ShimExe) { + if ($StudioRedirectMode -ne 'env') { + $pathAdded = Add-ToUserPath -Directory $ShimDir -Position 'Prepend' + } } if ($shimUpdated -and $pathAdded) { step "path" "added unsloth launcher to PATH" @@ -1336,12 +1585,20 @@ shell.Run cmd, 0, False Refresh-SessionPath # sync current session with registry Complete-StudioVenvRollback + # Env-mode session export AFTER Refresh-SessionPath; otherwise a legacy + # User PATH entry (Machine > User > current $env:Path) would win. + if ($StudioRedirectMode -eq 'env' -and (Test-Path -LiteralPath $ShimExe)) { + $env:Path = "$ShimDir;$env:Path" + step "path" "exported $ShimDir for this session (no registry PATH change in env-override mode)" + } + # ── Tauri mode: done, skip shortcuts and auto-launch ── if ($TauriMode) { Write-TauriLog "DONE" "" return } + # New-StudioShortcuts gates the .lnk shortcuts on env-mode internally. New-StudioShortcuts -UnslothExePath $UnslothExe # In interactive terminals, ask the user before starting Studio. @@ -1360,8 +1617,21 @@ shell.Run cmd, 0, False } } else { step "launch" "manual commands:" - substep "& `"$VenvDir\Scripts\Activate.ps1`"" - substep "unsloth studio -p 8888" + # Single-quote the printed paths so $-vars / backticks in custom roots + # do not reparse when the user pastes the command. + $_actLiteral = "'" + ((Join-Path $VenvDir "Scripts\Activate.ps1") -replace "'", "''") + "'" + if ($StudioRedirectMode -eq 'env') { + # Env-mode skips registry PATH; print the absolute shim path. + $_shim = Join-Path $StudioHome "bin\unsloth.exe" + $_shimLiteral = "'" + ($_shim -replace "'", "''") + "'" + substep "& $_shimLiteral studio -p 8888" + substep "or activate env first:" + substep "& $_actLiteral" + substep "unsloth studio -p 8888" + } else { + substep "& $_actLiteral" + substep "unsloth studio -p 8888" + } substep "(add -H 0.0.0.0 to allow network / cloud access)" Write-Host "" } diff --git a/install.sh b/install.sh index 1d21117d16..ec47d016eb 100755 --- a/install.sh +++ b/install.sh @@ -6,6 +6,12 @@ # Usage (no-torch): ./install.sh --no-torch (skip PyTorch, GGUF-only mode) # Usage (test): ./install.sh --package roland-sloth (install a different package name) # Usage (py): ./install.sh --python 3.12 (override auto-detected Python version) +# +# Env vars (priority: UNSLOTH_STUDIO_HOME > STUDIO_HOME > HOME-redirect > default): +# UNSLOTH_STUDIO_HOME=/abs/path -> install under that path +# STUDIO_HOME=/abs/path -> alias, same effect (UNSLOTH_STUDIO_HOME wins) +# (DATA_DIR + unsloth CLI shim nest inside; no shell rc-file append.) +# Default ($HOME/.unsloth/studio) is preserved when no env var is set. set -e # ── Output style (aligned with studio/setup.sh) ── @@ -66,6 +72,56 @@ if [ "$_VERBOSE" = true ]; then export UNSLOTH_VERBOSE=1 fi +# Custom Studio roots are not supported with --tauri (desktop app still +# resolves ~/.unsloth/studio). Pass through if the override == legacy default. +if [ "$TAURI_MODE" = true ]; then + _tauri_override_var="" + _tauri_override="${UNSLOTH_STUDIO_HOME:-}" + if [ -n "$_tauri_override" ]; then + _tauri_override_var="UNSLOTH_STUDIO_HOME" + else + _tauri_override="${STUDIO_HOME:-}" + [ -n "$_tauri_override" ] && _tauri_override_var="STUDIO_HOME" + fi + # Strip whitespace so " " is treated as unset (matches Python .strip()). + _tauri_override=$(printf '%s' "$_tauri_override" | sed -e 's/^[[:space:]]*//' -e 's/[[:space:]]*$//') + if [ -n "$_tauri_override" ]; then + case "$_tauri_override" in + "~") _tauri_override="$HOME" ;; + "~/"*) _tauri_override="$HOME/${_tauri_override#'~/'}" ;; + esac + # Canonicalize both sides (CDPATH=, -P) so a CDPATH-set env or + # symlinked $HOME doesn't break the legacy-equality comparison. + if [ -d "$_tauri_override" ]; then + _tauri_override_abs=$(CDPATH= cd -P -- "$_tauri_override" 2>/dev/null && pwd -P) \ + || _tauri_override_abs="$_tauri_override" + else + _tauri_override_abs="$_tauri_override" + fi + # Strip trailing separators so ".../studio/" matches ".../studio". + while [ "$_tauri_override_abs" != "/" ] \ + && [ "${_tauri_override_abs%/}" != "$_tauri_override_abs" ]; do + _tauri_override_abs=${_tauri_override_abs%/} + done + _tauri_legacy_root="$HOME/.unsloth/studio" + if [ -d "$_tauri_legacy_root" ]; then + _tauri_legacy_root=$(CDPATH= cd -P -- "$_tauri_legacy_root" 2>/dev/null && pwd -P) \ + || _tauri_legacy_root="$HOME/.unsloth/studio" + fi + while [ "$_tauri_legacy_root" != "/" ] \ + && [ "${_tauri_legacy_root%/}" != "$_tauri_legacy_root" ]; do + _tauri_legacy_root=${_tauri_legacy_root%/} + done + if [ "$_tauri_override_abs" != "$_tauri_legacy_root" ]; then + echo "ERROR: $_tauri_override_var is not supported with --tauri." >&2 + echo " The desktop app still uses the legacy ~/.unsloth/studio root." >&2 + echo " Run install.sh without --tauri for custom-root shell installs," >&2 + echo " or unset the env var for default desktop installs." >&2 + exit 1 + fi + fi +fi + _is_verbose() { [ "${UNSLOTH_VERBOSE:-0}" = "1" ] } @@ -219,7 +275,67 @@ _tauri_gpu_branch() { } PYTHON_VERSION="" # resolved after platform detection -STUDIO_HOME="$HOME/.unsloth/studio" + +# Resolve install destinations: env override, HOME-redirect (best-effort +# via getent/dscl), or default. Env-var priority: UNSLOTH_STUDIO_HOME wins +# over STUDIO_HOME (the more specific signal beats the generic alias). +_resolve_studio_destinations() { + _override_var="" + _override="${UNSLOTH_STUDIO_HOME:-}" + if [ -n "$_override" ]; then + _override_var="UNSLOTH_STUDIO_HOME" + else + _override="${STUDIO_HOME:-}" + [ -n "$_override" ] && _override_var="STUDIO_HOME" + fi + # Strip surrounding whitespace so " " is treated as unset (matches the + # Python resolvers' .strip()), preventing install/runtime layout drift. + _override=$(printf '%s' "$_override" | sed -e 's/^[[:space:]]*//' -e 's/[[:space:]]*$//') + # Tilde expansion: env vars are not subject to it when quoted on assignment. + case "$_override" in + "~") _override="$HOME" ;; + "~/"*) _override="$HOME/${_override#'~/'}" ;; + esac + if [ -n "$_override" ]; then + mkdir -p -- "$_override" 2>/dev/null || { echo "ERROR: $_override_var=$_override cannot be created." >&2; exit 1; } + [ -w "$_override" ] || { echo "ERROR: $_override_var=$_override is not writable." >&2; exit 1; } + STUDIO_HOME="$(CDPATH= cd -P -- "$_override" && pwd -P)" || exit 1 + DATA_DIR="$STUDIO_HOME/share" + _LOCAL_BIN="$STUDIO_HOME/bin" + _STUDIO_HOME_REDIRECT=env + substep "custom $_override_var=$STUDIO_HOME" + return 0 + fi + _default_home="" + if command -v getent >/dev/null 2>&1; then + _default_home=$(getent passwd "${USER:-$(whoami)}" 2>/dev/null | cut -d: -f6) + elif [ "$(uname)" = "Darwin" ] && command -v dscl >/dev/null 2>&1; then + _default_home=$(dscl . -read "/Users/${USER:-$(whoami)}" NFSHomeDirectory 2>/dev/null | awk '{print $2}') + fi + # Canonicalize both sides so a trailing slash on $HOME (or symlink mismatch + # with passwd-DB output) doesn't misfire the redirection branch. + _home_canon="$HOME" + if [ -d "$_home_canon" ]; then + _home_canon=$(CDPATH= cd -P -- "$_home_canon" 2>/dev/null && pwd -P) || _home_canon="$HOME" + fi + _default_home_canon="$_default_home" + if [ -n "$_default_home_canon" ] && [ -d "$_default_home_canon" ]; then + _default_home_canon=$(CDPATH= cd -P -- "$_default_home_canon" 2>/dev/null && pwd -P) || _default_home_canon="$_default_home" + fi + if [ -n "$_default_home_canon" ] && [ "$_home_canon" != "$_default_home_canon" ]; then + STUDIO_HOME="$HOME/.unsloth/studio" + DATA_DIR="$HOME/.local/share/unsloth" + _LOCAL_BIN="$HOME/.local/bin" + _STUDIO_HOME_REDIRECT=home + substep "HOME redirected ($HOME); install follows \$HOME" + return 0 + fi + STUDIO_HOME="$HOME/.unsloth/studio" + DATA_DIR="$HOME/.local/share/unsloth" + _LOCAL_BIN="$HOME/.local/bin" + _STUDIO_HOME_REDIRECT=default +} +_resolve_studio_destinations VENV_DIR="$STUDIO_HOME/unsloth_studio" _VENV_ROLLBACK_DIR="" _VENV_ROLLBACK_TARGET="$VENV_DIR" @@ -383,23 +499,65 @@ create_studio_shortcuts() { _css_exe_dir=$(cd "$(dirname "$_css_exe")" && pwd) _css_exe="$_css_exe_dir/$(basename "$_css_exe")" - _css_data_dir="$HOME/.local/share/unsloth" + _css_data_dir="$DATA_DIR" _css_launcher="$_css_data_dir/launch-studio.sh" _css_icon_png="$_css_data_dir/unsloth-studio.png" _css_gem_png="$_css_data_dir/unsloth-gem.png" mkdir -p "$_css_data_dir" + # Same-install discriminator: per-install opaque id written once at install + # time and read by both this launcher and the backend (/api/health). Replaces + # the older sha256(canonical $STUDIO_HOME) scheme to (a) avoid leaking the + # install path on -H 0.0.0.0 deployments and (b) sidestep launcher/backend + # canonicalization drift (cd -P vs Path.resolve() symlink/junction handling). + # Lives at $STUDIO_HOME/share/ (not $DATA_DIR) so the backend can find it + # via _STUDIO_ROOT_RESOLVED / "share" / "studio_install_id" regardless of + # mode (in env-mode $STUDIO_HOME/share == $DATA_DIR; in default mode they + # diverge but the backend only knows the studio_root). 32 bytes of urandom + # -> 64 hex chars, byte-compatible with the prior digest so launcher + # placeholder, _check_health, and tests stay length-agnostic. + _css_id_dir="$STUDIO_HOME/share" + mkdir -p "$_css_id_dir" + _css_id_file="$_css_id_dir/studio_install_id" + if [ ! -s "$_css_id_file" ]; then + if [ -r /dev/urandom ]; then + _css_new_id=$(od -An -N32 -tx1 /dev/urandom 2>/dev/null | tr -d ' \n') + fi + if [ -z "${_css_new_id:-}" ] && command -v python3 >/dev/null 2>&1; then + _css_new_id=$(python3 -c 'import secrets; print(secrets.token_hex(32))' 2>/dev/null) + fi + if [ -z "${_css_new_id:-}" ]; then + echo "[WARN] Cannot create launcher: no entropy source for studio_install_id" >&2 + return 1 + fi + # Atomic write so a partial install can't leave a half-written id. + _css_id_tmp="$_css_id_file.$$.tmp" + printf '%s' "$_css_new_id" > "$_css_id_tmp" \ + && mv "$_css_id_tmp" "$_css_id_file" + chmod 600 "$_css_id_file" 2>/dev/null || true + unset _css_new_id _css_id_tmp + fi + _css_studio_root_id=$(cat "$_css_id_file" 2>/dev/null) + if [ -z "$_css_studio_root_id" ]; then + echo "[WARN] Cannot create launcher: failed to read $_css_id_file" >&2 + return 1 + fi + _css_is_env_mode=false + [ "$_STUDIO_HOME_REDIRECT" = "env" ] && _css_is_env_mode=true + # ── Write launcher script ── - # The launcher is Bash (not POSIX sh). - # We write it with a placeholder and substitute the exe path via sed. + # Single-quoted heredoc; @@DATA_DIR@@, @@STUDIO_ROOT_ID@@, and + # @@INSTALLED_IS_ENV_MODE@@ are substituted via sed below. cat > "$_css_launcher" << 'LAUNCHER_EOF' #!/usr/bin/env bash # Unsloth Studio Launcher # Auto-generated by install.sh -- do not edit manually. set -euo pipefail -DATA_DIR="$HOME/.local/share/unsloth" +DATA_DIR='@@DATA_DIR@@' +_EXPECTED_STUDIO_ROOT_ID='@@STUDIO_ROOT_ID@@' +_INSTALLED_IS_ENV_MODE='@@INSTALLED_IS_ENV_MODE@@' # Read exe path from config written at install time. # Sourcing is safe: the config file is written by install.sh, not user input. @@ -416,7 +574,23 @@ MAX_PORT_OFFSET=20 TIMEOUT_SEC=60 POLL_INTERVAL_SEC=1 LOG_FILE="$DATA_DIR/studio.log" +# why: in env-override mode multiple installs share an OS user; namespace the +# lock and remember our own healthy port so we never attach to an unrelated +# Studio listening on the global 8888..8908 range. LOCK_DIR="${XDG_RUNTIME_DIR:-/tmp}/unsloth-studio-launcher-$(id -u).lock" +PORT_FILE="" +# why: gate on the install-time mode (baked above) instead of the runtime env +# var; sourcing a custom-root studio.conf in shell must not flip a default-mode +# launcher into env-mode behavior with stale state. +if [ "$_INSTALLED_IS_ENV_MODE" = "true" ]; then + if command -v cksum >/dev/null 2>&1; then + _LOCK_KEY=$(printf '%s' "$DATA_DIR" | cksum | awk '{print $1}') + else + _LOCK_KEY="" + fi + [ -n "$_LOCK_KEY" ] && LOCK_DIR="${XDG_RUNTIME_DIR:-/tmp}/unsloth-studio-launcher-$(id -u)-${_LOCK_KEY}.lock" + PORT_FILE="$DATA_DIR/studio.port" +fi # ── HTTP GET helper (supports curl and wget) ── _http_get() { @@ -435,10 +609,20 @@ _check_health() { _port=$1 _resp=$(_http_get "http://127.0.0.1:$_port/api/health") || return 1 case "$_resp" in - *'"status"'*'"healthy"'*'"service"'*'"Unsloth UI Backend"'*) return 0 ;; - *'"service"'*'"Unsloth UI Backend"'*'"status"'*'"healthy"'*) return 0 ;; + *'"status"'*'"healthy"'*'"service"'*'"Unsloth UI Backend"'*) ;; + *'"service"'*'"Unsloth UI Backend"'*'"status"'*'"healthy"'*) ;; + *) return 1 ;; esac - return 1 + # why: verify the backend belongs to THIS install. Baked hex digest avoids + # JSON-escape mismatches on paths with `\`/`"` and avoids leaking the raw + # install path to unauthenticated callers. + if [ -n "$_EXPECTED_STUDIO_ROOT_ID" ]; then + case "$_resp" in + *"\"studio_root_id\":\"$_EXPECTED_STUDIO_ROOT_ID\""*|*"\"studio_root_id\": \"$_EXPECTED_STUDIO_ROOT_ID\""*) return 0 ;; + *) return 1 ;; + esac + fi + return 0 } # ── Port scanning ── @@ -461,6 +645,25 @@ _candidate_ports() { } _find_healthy_port() { + if [ -n "$PORT_FILE" ] && [ -f "$PORT_FILE" ]; then + # why: env-mode installs only attach to a port we previously launched + # ourselves; never to a sibling Studio that happens to be healthy. + _p=$(cat "$PORT_FILE" 2>/dev/null || true) + case "$_p" in + ''|*[!0-9]*) ;; + *) + if _check_health "$_p"; then + echo "$_p" + return 0 + fi + rm -f "$PORT_FILE" + ;; + esac + return 1 + fi + if [ -n "$PORT_FILE" ]; then + return 1 + fi for _p in $(_candidate_ports | sort -un); do if _check_health "$_p"; then echo "$_p" @@ -611,6 +814,7 @@ if [ -t 1 ]; then _obwr_deadline=$(($(date +%s) + TIMEOUT_SEC)) while [ "$(date +%s)" -lt "$_obwr_deadline" ]; do if _check_health "$_launch_port"; then + [ -n "$PORT_FILE" ] && printf '%s\n' "$_launch_port" > "$PORT_FILE" 2>/dev/null || true _release_lock _open_browser "http://localhost:$_launch_port" exit 0 @@ -634,6 +838,7 @@ else _deadline=$(($(date +%s) + TIMEOUT_SEC)) while [ "$(date +%s)" -lt "$_deadline" ]; do if _check_health "$_launch_port"; then + [ -n "$PORT_FILE" ] && printf '%s\n' "$_launch_port" > "$PORT_FILE" 2>/dev/null || true _open_browser "http://localhost:$_launch_port" exit 0 fi @@ -646,13 +851,62 @@ else fi LAUNCHER_EOF + # why: bake non-user-controlled placeholders FIRST so a literal + # `@@STUDIO_ROOT_ID@@` inside $DATA_DIR cannot be rewritten below. + sed -e "s|@@STUDIO_ROOT_ID@@|$_css_studio_root_id|g" \ + -e "s|@@INSTALLED_IS_ENV_MODE@@|$_css_is_env_mode|g" \ + "$_css_launcher" > "$_css_launcher.tmp" \ + && mv "$_css_launcher.tmp" "$_css_launcher" + + # Env-mode bakes an absolute DATA_DIR (root fixed at install time); + # default / HOME-redirect keeps the literal $HOME/.local/share/unsloth + # so behavior is byte-identical to pre-override. + if [ "$_STUDIO_HOME_REDIRECT" = "env" ]; then + # Two-stage escape: (1) `'` -> `'\''` for shell single-quote embedding, + # (2) backslash/&/| escape so the value survives the s|...|VALUE| sed + # below. Verified end-to-end with apostrophes, spaces, &, |, $. + _sq_escaped=$(printf '%s' "$DATA_DIR" | sed "s/'/'\\\\''/g") + _sed_safe=$(printf '%s' "$_sq_escaped" | sed 's/[\\&|]/\\&/g') + sed "s|@@DATA_DIR@@|$_sed_safe|g" "$_css_launcher" > "$_css_launcher.tmp" \ + && mv "$_css_launcher.tmp" "$_css_launcher" + else + sed "s|DATA_DIR='@@DATA_DIR@@'|DATA_DIR=\"\$HOME/.local/share/unsloth\"|" \ + "$_css_launcher" > "$_css_launcher.tmp" \ + && mv "$_css_launcher.tmp" "$_css_launcher" + fi + chmod +x "$_css_launcher" - # Write the exe path to a separate conf file sourced by the launcher. - # Using single-quote wrapping with the standard '\'' escape for any - # embedded apostrophes. This avoids all sed metacharacter issues. + # studio.conf: exe path + (env-mode only) persisted env vars so fresh + # shells launch the right install without re-exporting. _css_quoted_exe=$(printf '%s' "$_css_exe" | sed "s/'/'\\\\''/g") - printf '%s\n' "UNSLOTH_EXE='$_css_quoted_exe'" > "$_css_data_dir/studio.conf" + { + printf '%s\n' "UNSLOTH_EXE='$_css_quoted_exe'" + if [ "$_STUDIO_HOME_REDIRECT" = "env" ]; then + # When an override resolves to the legacy default, llama.cpp + # still lives at ~/.unsloth/llama.cpp (one shared build). + # Canonicalize the legacy side so a symlinked $HOME doesn't + # break the comparison. + _css_legacy_studio="$HOME/.unsloth/studio" + if [ -d "$_css_legacy_studio" ]; then + _css_legacy_studio=$(CDPATH= cd -P -- "$_css_legacy_studio" 2>/dev/null && pwd -P) \ + || _css_legacy_studio="$HOME/.unsloth/studio" + fi + if [ "$STUDIO_HOME" = "$_css_legacy_studio" ]; then + _css_llama_path="$HOME/.unsloth/llama.cpp" + else + _css_llama_path="$STUDIO_HOME/llama.cpp" + fi + _css_quoted_home=$(printf '%s' "$STUDIO_HOME" | sed "s/'/'\\\\''/g") + _css_quoted_llama=$(printf '%s' "$_css_llama_path" | sed "s/'/'\\\\''/g") + printf '%s\n' "export UNSLOTH_STUDIO_HOME='$_css_quoted_home'" + # UNSLOTH_LLAMA_CPP_PATH is a pre-existing user-controlled + # llama.cpp dir override; only default it if unset. + printf '%s\n' 'if [ -z "${UNSLOTH_LLAMA_CPP_PATH:-}" ]; then' + printf '%s\n' " export UNSLOTH_LLAMA_CPP_PATH='$_css_quoted_llama'" + printf '%s\n' 'fi' + fi + } > "$_css_data_dir/studio.conf" # ── Icon: try bundled, then download ── # rounded-512.png used for both Linux and macOS icons @@ -698,6 +952,14 @@ LAUNCHER_EOF fi # ── Platform-specific shortcuts ── + # Env-mode installs are workspace-scoped: skip persistent desktop / + # Start-Menu / dock launchers that may point at a deleted workspace. + # Runtime launcher + studio.conf + icon are still written above. + if [ "$_STUDIO_HOME_REDIRECT" = "env" ]; then + substep "wrote launcher at $_css_launcher (persistent shortcuts skipped in env-override mode)" + return 0 + fi + _css_created=0 if [ "$_css_os" = "linux" ]; then @@ -775,11 +1037,18 @@ DESKTOP_EOF PLIST_EOF - # Executable stub - cat > "$_css_macos_dir/launch-studio" << STUB_EOF + # Executable stub: same single-quoted-heredoc + sed-substitute + # pattern as launch-studio.sh so $-vars in $_css_data_dir don't + # expand at .app launch time. + _css_sq_dir=$(printf '%s' "$_css_data_dir" | sed "s/'/'\\\\''/g") + _css_sed_dir=$(printf '%s' "$_css_sq_dir" | sed 's/[\\&|]/\\&/g') + cat > "$_css_macos_dir/launch-studio" << 'STUB_EOF' #!/bin/sh -exec "$HOME/.local/share/unsloth/launch-studio.sh" "\$@" +exec '@@DATA_DIR@@/launch-studio.sh' "$@" STUB_EOF + sed "s|@@DATA_DIR@@|$_css_sed_dir|g" "$_css_macos_dir/launch-studio" \ + > "$_css_macos_dir/launch-studio.tmp" \ + && mv "$_css_macos_dir/launch-studio.tmp" "$_css_macos_dir/launch-studio" chmod +x "$_css_macos_dir/launch-studio" # Build AppIcon.icns from unsloth-gem.png (2240x2240) @@ -1079,11 +1348,28 @@ mkdir -p "$STUDIO_HOME" _MIGRATED=false if [ -x "$VENV_DIR/bin/python" ]; then + # why: matching guard to the .venv branch below -- in env-mode + # $STUDIO_HOME is a user-chosen workspace, so refuse to nuke an + # existing $STUDIO_HOME/unsloth_studio that lacks Studio sentinels. + # Accept the in-VENV ownership marker so partial-install retries are + # not blocked. Sentinels must be regular files: -f follows symlinks + # to files (the legitimate ln -s shim shape) but rejects directories + # and broken/dir-targeted symlinks. + if [ "$_STUDIO_HOME_REDIRECT" = "env" ] \ + && [ ! -f "$VENV_DIR/.unsloth-studio-owned" ] \ + && [ ! -f "$STUDIO_HOME/share/studio.conf" ] \ + && [ ! -f "$STUDIO_HOME/bin/unsloth" ]; then + echo "ERROR: $VENV_DIR already exists but does not look like an Unsloth Studio install." >&2 + echo " Move it aside or choose an empty UNSLOTH_STUDIO_HOME." >&2 + exit 1 + fi # New layout already exists — replace only after preserving rollback copy. substep "preserving existing environment for rollback..." _start_studio_venv_replacement "$VENV_DIR" -elif [ -x "$STUDIO_HOME/.venv/bin/python" ]; then +elif [ "$_STUDIO_HOME_REDIRECT" != "env" ] && [ -x "$STUDIO_HOME/.venv/bin/python" ]; then # Old layout exists — validate before migrating. + # Skip in env-mode so we don't rm -rf an unrelated .venv at the + # workspace root (e.g. user's existing project Python venv). # In no-torch mode, a missing torch package is expected; validate Python only. substep "found legacy Studio environment, validating..." _legacy_ok=false @@ -1132,6 +1418,13 @@ if [ ! -x "$VENV_DIR/bin/python" ]; then run_install_cmd "create venv" uv venv "$VENV_DIR" --python "$PYTHON_VERSION" fi +# Mark the freshly-created venv as Studio-owned so a partial install can be +# repaired by re-running install.sh; the env-mode deletion guard above accepts +# this marker as the primary sentinel. +if [ -x "$VENV_DIR/bin/python" ]; then + : > "$VENV_DIR/.unsloth-studio-owned" 2>/dev/null || true +fi + # Guard against Python 3.13.8 torch import bug on Apple Silicon # (skip when the user explicitly chose a version via --python) if [ -z "$_USER_PYTHON" ] && [ "$OS" = "macos" ] && [ "$_ARCH" = "arm64" ]; then @@ -1143,6 +1436,9 @@ if [ -z "$_USER_PYTHON" ] && [ "$OS" = "macos" ] && [ "$_ARCH" = "arm64" ]; then rm -rf "$VENV_DIR" PYTHON_VERSION="3.12" run_install_cmd "recreate venv" uv venv "$VENV_DIR" --python "$PYTHON_VERSION" + if [ -x "$VENV_DIR/bin/python" ]; then + : > "$VENV_DIR/.unsloth-studio-owned" 2>/dev/null || true + fi fi fi @@ -1768,7 +2064,17 @@ _SKIP_FRONTEND=0 if [ "$TAURI_MODE" = true ]; then _SKIP_FRONTEND=1 fi +# Prepend UNSLOTH_STUDIO_HOME=$STUDIO_HOME to "$@" for env-override installs +# without word-splitting on whitespace paths. +_run_setup_with_studio_home() { + if [ "$_STUDIO_HOME_REDIRECT" = "env" ]; then + UNSLOTH_STUDIO_HOME="$STUDIO_HOME" "$@" + else + "$@" + fi +} if [ "$STUDIO_LOCAL_INSTALL" = true ]; then + _run_setup_with_studio_home env \ SKIP_STUDIO_BASE="$_SKIP_BASE" \ SKIP_STUDIO_FRONTEND="$_SKIP_FRONTEND" \ STUDIO_PACKAGE_NAME="$PACKAGE_NAME" \ @@ -1782,6 +2088,7 @@ else # the same session) does not silently flip a normal install onto the # local-dev path in setup.sh and install_python_stack.py. Mirrors the # reset already done in install.ps1 for PowerShell. + _run_setup_with_studio_home env \ SKIP_STUDIO_BASE="$_SKIP_BASE" \ SKIP_STUDIO_FRONTEND="$_SKIP_FRONTEND" \ STUDIO_PACKAGE_NAME="$PACKAGE_NAME" \ @@ -1791,36 +2098,53 @@ else bash "$SETUP_SH" &2 + echo " Move or remove it manually, then re-run the installer." >&2 + exit 1 +fi +# why: -sfn is atomic and -n prevents descent into a symlink-to-directory at +# the shim path (the directory guard above already rejects a real directory). +ln -sfn "$VENV_DIR/bin/unsloth" "$_shim_path" -_LOCAL_BIN="$HOME/.local/bin" case ":$PATH:" in *":$_LOCAL_BIN:"*) ;; # already on PATH *) - _SHELL_PROFILE="" - if [ -n "${ZSH_VERSION:-}" ] || [ "$(basename "${SHELL:-}")" = "zsh" ]; then - _SHELL_PROFILE="$HOME/.zshrc" - elif [ -f "$HOME/.bashrc" ]; then - _SHELL_PROFILE="$HOME/.bashrc" - elif [ -f "$HOME/.profile" ]; then - _SHELL_PROFILE="$HOME/.profile" - fi - - if [ -n "$_SHELL_PROFILE" ]; then - if ! grep -q '\.local/bin' "$_SHELL_PROFILE" 2>/dev/null; then - echo '' >> "$_SHELL_PROFILE" - echo '# Added by Unsloth installer' >> "$_SHELL_PROFILE" - echo 'export PATH="$HOME/.local/bin:$PATH"' >> "$_SHELL_PROFILE" - step "path" "added ~/.local/bin to PATH in $_SHELL_PROFILE" + if [ "$_STUDIO_HOME_REDIRECT" = "env" ]; then + export PATH="$_LOCAL_BIN:$PATH" + step "path" "exported $_LOCAL_BIN for this session (no rc-file append in env-override mode)" + else + _SHELL_PROFILE="" + if [ -n "${ZSH_VERSION:-}" ] || [ "$(basename "${SHELL:-}")" = "zsh" ]; then + _SHELL_PROFILE="$HOME/.zshrc" + elif [ -f "$HOME/.bashrc" ]; then + _SHELL_PROFILE="$HOME/.bashrc" + elif [ -f "$HOME/.profile" ]; then + _SHELL_PROFILE="$HOME/.profile" fi + if [ -n "$_SHELL_PROFILE" ]; then + if ! grep -q '\.local/bin' "$_SHELL_PROFILE" 2>/dev/null; then + echo '' >> "$_SHELL_PROFILE" + echo '# Added by Unsloth installer' >> "$_SHELL_PROFILE" + echo 'export PATH="$HOME/.local/bin:$PATH"' >> "$_SHELL_PROFILE" + step "path" "added ~/.local/bin to PATH in $_SHELL_PROFILE" + fi + fi + export PATH="$_LOCAL_BIN:$PATH" fi - export PATH="$_LOCAL_BIN:$PATH" ;; esac # Non-Tauri installs keep shortcuts even if setup reports failure. +# create_studio_shortcuts gates persistent menu shortcuts on env-mode; +# launcher + studio.conf + icon are always written. if [ "$TAURI_MODE" != true ]; then create_studio_shortcuts "$VENV_ABS_BIN/unsloth" "$OS" fi @@ -1883,10 +2207,21 @@ if [ -t 1 ]; then esac else step "launch" "manual commands:" - substep "unsloth studio -p 8888" - substep "or activate env first:" - substep "source ${VENV_DIR}/bin/activate" - substep "unsloth studio -p 8888" + # Single-quote-escape so paths with spaces / apostrophes copy-paste cleanly. + _li_shim_q="'$(printf '%s' "${_LOCAL_BIN}/unsloth" | sed "s/'/'\\\\''/g")'" + _li_act_q="'$(printf '%s' "${VENV_DIR}/bin/activate" | sed "s/'/'\\\\''/g")'" + if [ "$_STUDIO_HOME_REDIRECT" = "env" ]; then + # Env-mode skips the rc PATH append, so print the absolute shim path. + substep "$_li_shim_q studio -p 8888" + substep "or activate env first:" + substep "source $_li_act_q" + substep "unsloth studio -p 8888" + else + substep "unsloth studio -p 8888" + substep "or activate env first:" + substep "source $_li_act_q" + substep "unsloth studio -p 8888" + fi substep "(add -H 0.0.0.0 to allow network / cloud access)" echo "" fi diff --git a/studio/backend/core/inference/llama_cpp.py b/studio/backend/core/inference/llama_cpp.py index f768764c22..8da836de38 100644 --- a/studio/backend/core/inference/llama_cpp.py +++ b/studio/backend/core/inference/llama_cpp.py @@ -732,22 +732,46 @@ class LlamaCppBackend: if win_bin.is_file(): return str(win_bin) - # 2–4. ~/.unsloth/llama.cpp (primary — setup.sh / setup.ps1 build here) - unsloth_home = Path.home() / ".unsloth" / "llama.cpp" - # Root dir (make builds copy binaries here) - home_root = unsloth_home / binary_name - if home_root.is_file(): - return str(home_root) - # build/bin/ (cmake builds on Linux) - home_linux = unsloth_home / "build" / "bin" / binary_name - if home_linux.is_file(): - return str(home_linux) + # 2-4. Match installer layout: env-mode -> $STUDIO_HOME/llama.cpp; + # default/HOME-redirect -> ~/.unsloth/llama.cpp (sibling of studio). + legacy_llama = Path.home() / ".unsloth" / "llama.cpp" + try: + from utils.paths.storage_roots import studio_root as _sr # noqa: WPS433 - # 3. Windows MSVC build has Release subdir - if sys.platform == "win32": - home_win = unsloth_home / "build" / "bin" / "Release" / binary_name - if home_win.is_file(): - return str(home_win) + _resolved_sr = _sr() + _legacy_studio = Path.home() / ".unsloth" / "studio" + try: + _is_legacy = _resolved_sr.resolve() == _legacy_studio.resolve() + except (OSError, ValueError): + _is_legacy = _resolved_sr == _legacy_studio + if _is_legacy: + search_roots = [legacy_llama] + else: + # why: _kill_orphaned_servers excludes the legacy root in custom + # mode; discovery must match so we never spawn a server we then + # refuse to clean up. UNSLOTH_LLAMA_CPP_PATH (handled earlier) + # is the explicit way to share a build across roots. + search_roots = [_resolved_sr / "llama.cpp"] + except (ImportError, OSError, ValueError): + search_roots = [legacy_llama] + _seen_roots: set[str] = set() + _unique_roots: list[Path] = [] + for r in search_roots: + k = str(r) + if k not in _seen_roots: + _seen_roots.add(k) + _unique_roots.append(r) + for unsloth_home in _unique_roots: + home_root = unsloth_home / binary_name + if home_root.is_file(): + return str(home_root) + home_linux = unsloth_home / "build" / "bin" / binary_name + if home_linux.is_file(): + return str(home_linux) + if sys.platform == "win32": + home_win = unsloth_home / "build" / "bin" / "Release" / binary_name + if home_win.is_file(): + return str(home_win) # 5–6. Legacy: in-tree build (older setup.sh / setup.ps1 versions) project_root = Path(__file__).resolve().parents[4] @@ -2592,8 +2616,27 @@ class LlamaCppBackend: # (binary must be *under* one of these) install_roots: list[Path] = [] - # Primary install dir (setup.sh / prebuilt installer) - install_roots.append(Path.home() / ".unsloth" / "llama.cpp") + # Env-mode custom root (mirrors _find_llama_server_binary). + _is_custom_root = False + try: + from utils.paths.storage_roots import studio_root as _sr # noqa: WPS433 + + _resolved_sr = _sr() + _legacy_studio = Path.home() / ".unsloth" / "studio" + try: + _is_custom_root = _resolved_sr.resolve() != _legacy_studio.resolve() + except (OSError, ValueError): + _is_custom_root = _resolved_sr != _legacy_studio + if _is_custom_root: + install_roots.append(_resolved_sr / "llama.cpp") + except (ImportError, OSError, ValueError): + pass + + # Primary install dir (default mode only). Env-mode skips this so + # a custom-root Studio cannot kill a concurrent default-install + # Studio's llama-server (same OS user, different install). + if not _is_custom_root: + install_roots.append(Path.home() / ".unsloth" / "llama.cpp") # Legacy in-tree build dirs (older setup.sh versions) project_root = Path(__file__).resolve().parents[4] diff --git a/studio/backend/main.py b/studio/backend/main.py index 0958094ff0..cd901327db 100644 --- a/studio/backend/main.py +++ b/studio/backend/main.py @@ -23,12 +23,67 @@ if _backend_dir not in sys.path: # See: https://github.com/python/cpython/issues/102396 import _platform_compat # noqa: F401 +# Direct `uvicorn main:app` launches bypass run.py, so re-export here too +# (mirrors run.py). Required BEFORE the unsloth-zoo import below, since +# its LLAMA_CPP_DEFAULT_DIR binding is import-time. +from utils.paths.storage_roots import studio_root as _studio_root + +try: + _LEGACY_STUDIO_ROOT = (_Path.home() / ".unsloth" / "studio").resolve() +except (OSError, ValueError): + _LEGACY_STUDIO_ROOT = _Path.home() / ".unsloth" / "studio" +try: + _STUDIO_ROOT_RESOLVED = _studio_root().resolve() +except (OSError, ValueError): + _STUDIO_ROOT_RESOLVED = _studio_root() +if _STUDIO_ROOT_RESOLVED != _LEGACY_STUDIO_ROOT: + if not os.environ.get("UNSLOTH_STUDIO_HOME"): + os.environ["UNSLOTH_STUDIO_HOME"] = str(_STUDIO_ROOT_RESOLVED) + if not os.environ.get("UNSLOTH_LLAMA_CPP_PATH"): + os.environ["UNSLOTH_LLAMA_CPP_PATH"] = str(_STUDIO_ROOT_RESOLVED / "llama.cpp") + import mimetypes +import re as _re import shutil import warnings from contextlib import asynccontextmanager from importlib.metadata import PackageNotFoundError, version as package_version + +_STUDIO_INSTALL_ID_RE = _re.compile(r"^[0-9a-f]{64}$") + + +def _read_studio_install_id() -> str: + """Per-install opaque id written by install.sh / install.ps1 at + $STUDIO_HOME/share/studio_install_id. Returns "" when the file is + absent (pre-PR install, fresh tree never run through the installer) + or contains anything other than a 64-char lowercase-hex token -- + in which case /api/health emits "" and the launcher's _check_health + falls back to the existing "no baked id, accept any healthy + Unsloth backend" path. This intentionally replaces a previous + sha256(resolved_install_path) so the field carries no install-path + information for callers reaching /api/health (relevant when Studio + is run with -H 0.0.0.0).""" + try: + token = ( + (_STUDIO_ROOT_RESOLVED / "share" / "studio_install_id").read_text().strip() + ) + except (OSError, ValueError): + return "" + return token if _STUDIO_INSTALL_ID_RE.fullmatch(token) else "" + + +_STUDIO_ROOT_ID_CACHE: str = _read_studio_install_id() + + +def _studio_root_id() -> str: + """Same-install discriminator for /api/health: a per-install opaque + token written once by the installer and read once at module import. + Empty when no installer-written token is present; the launcher + contract treats "" as "no baked id, accept any healthy backend".""" + return _STUDIO_ROOT_ID_CACHE + + # Fix broken Windows registry MIME types. Some Windows installs map .js to # "text/plain" in the registry (HKCR\.js\Content Type). Python's mimetypes # module reads from the registry, and FastAPI/Starlette's StaticFiles uses @@ -245,6 +300,10 @@ async def health_check(): "chat_only": _hw_module.CHAT_ONLY, "desktop_protocol_version": 1, "supports_desktop_auth": True, + # why: launchers compare against an install-time hash so a sibling + # Studio on the same port is rejected; hex digest avoids leaking the + # raw install path on -H 0.0.0.0. + "studio_root_id": _studio_root_id(), "native_path_leases_supported": native_path_leases_supported(), } diff --git a/studio/backend/run.py b/studio/backend/run.py index c5b103ff70..1dd1230a17 100644 --- a/studio/backend/run.py +++ b/studio/backend/run.py @@ -159,7 +159,27 @@ def _find_free_port(host: str, start: int, max_attempts: int = 20) -> int: ) -_PID_FILE = Path.home() / ".unsloth" / "studio" / "studio.pid" +from utils.paths.storage_roots import studio_root as _studio_root + +_PID_FILE = _studio_root() / "studio.pid" + +# Direct backend launches bypass the CLI's env re-export; do it here for +# real custom roots so unsloth-zoo's import-time LLAMA_CPP_DEFAULT_DIR +# picks up the custom build. Skip for legacy-default to avoid flipping +# default-mode installs into env-override. +try: + _LEGACY_STUDIO_ROOT = (Path.home() / ".unsloth" / "studio").resolve() +except (OSError, ValueError): + _LEGACY_STUDIO_ROOT = Path.home() / ".unsloth" / "studio" +try: + _STUDIO_ROOT_RESOLVED = _studio_root().resolve() +except (OSError, ValueError): + _STUDIO_ROOT_RESOLVED = _studio_root() +if _STUDIO_ROOT_RESOLVED != _LEGACY_STUDIO_ROOT: + if not os.environ.get("UNSLOTH_STUDIO_HOME"): + os.environ["UNSLOTH_STUDIO_HOME"] = str(_STUDIO_ROOT_RESOLVED) + if not os.environ.get("UNSLOTH_LLAMA_CPP_PATH"): + os.environ["UNSLOTH_LLAMA_CPP_PATH"] = str(_STUDIO_ROOT_RESOLVED / "llama.cpp") def _write_pid_file(): diff --git a/studio/backend/utils/models/model_config.py b/studio/backend/utils/models/model_config.py index 16f6d21edb..dc8dd08315 100644 --- a/studio/backend/utils/models/model_config.py +++ b/studio/backend/utils/models/model_config.py @@ -500,7 +500,9 @@ _VLM_MODEL_TYPES = { # Pre-computed .venv_t5 paths and backend dir for subprocess version switching. # Vision check uses 5.5.0 (newest, recognizes all architectures). -_VENV_T5_DIR = str(Path.home() / ".unsloth" / "studio" / ".venv_t5_550") +from utils.paths.storage_roots import studio_root as _studio_root # noqa: E402 + +_VENV_T5_DIR = str(_studio_root() / ".venv_t5_550") _BACKEND_DIR = str(Path(__file__).resolve().parent.parent.parent) # Inline script executed in a subprocess with transformers 5.x activated. diff --git a/studio/backend/utils/paths/storage_roots.py b/studio/backend/utils/paths/storage_roots.py index b52609b06b..58a4d7967c 100644 --- a/studio/backend/utils/paths/storage_roots.py +++ b/studio/backend/utils/paths/storage_roots.py @@ -5,17 +5,59 @@ from __future__ import annotations import json import os +import sys from pathlib import Path import tempfile +def _infer_studio_home_from_venv() -> Path | None: + """Return parent dir of sys.prefix as STUDIO_HOME if running from an + installer-managed unsloth_studio venv. Sentinel-gated (share/studio.conf + or bin shim) so a developer venv named unsloth_studio is not misidentified. + """ + try: + prefix = Path(sys.prefix).resolve() + except (OSError, ValueError): + return None + if prefix.name != "unsloth_studio": + return None + candidate = prefix.parent + shim_name = "unsloth.exe" if os.name == "nt" else "unsloth" + try: + has_sentinel = (candidate / "share" / "studio.conf").is_file() or ( + candidate / "bin" / shim_name + ).is_file() + except OSError: + return None + if has_sentinel: + return candidate + return None + + def studio_root() -> Path: + """Studio install root. + + Priority: UNSLOTH_STUDIO_HOME, then STUDIO_HOME alias, then sys.prefix + inference, then legacy ~/.unsloth/studio. UNSLOTH_STUDIO_HOME wins when + both are set (the more specific signal beats the generic alias). + """ + override = (os.environ.get("UNSLOTH_STUDIO_HOME") or "").strip() + if not override: + override = (os.environ.get("STUDIO_HOME") or "").strip() + if override: + try: + return Path(override).expanduser().resolve() + except (OSError, ValueError): + return Path(override).expanduser() + inferred = _infer_studio_home_from_venv() + if inferred is not None: + return inferred return Path.home() / ".unsloth" / "studio" def cache_root() -> Path: """Central cache directory for all studio downloads (models, datasets, etc.).""" - return Path.home() / ".unsloth" / "studio" / "cache" + return studio_root() / "cache" def assets_root() -> Path: diff --git a/studio/backend/utils/transformers_version.py b/studio/backend/utils/transformers_version.py index 17af40f663..9075c590ca 100644 --- a/studio/backend/utils/transformers_version.py +++ b/studio/backend/utils/transformers_version.py @@ -95,9 +95,11 @@ TRANSFORMERS_DEFAULT_VERSION = "4.57.6" # Consumers should prefer TRANSFORMERS_530_VERSION / TRANSFORMERS_550_VERSION. TRANSFORMERS_5_VERSION = TRANSFORMERS_550_VERSION -# Pre-installed directories — created by setup.sh / setup.ps1 -_VENV_T5_530_DIR = str(Path.home() / ".unsloth" / "studio" / ".venv_t5_530") -_VENV_T5_550_DIR = str(Path.home() / ".unsloth" / "studio" / ".venv_t5_550") +# Pre-installed directories — created by setup.sh / setup.ps1. +from utils.paths.storage_roots import studio_root as _studio_root # noqa: E402 + +_VENV_T5_530_DIR = str(_studio_root() / ".venv_t5_530") +_VENV_T5_550_DIR = str(_studio_root() / ".venv_t5_550") # Backwards-compat alias _VENV_T5_DIR = _VENV_T5_550_DIR diff --git a/studio/setup.ps1 b/studio/setup.ps1 index 3d082aa70d..f2753d5c88 100644 --- a/studio/setup.ps1 +++ b/studio/setup.ps1 @@ -1492,9 +1492,79 @@ if (-not $PythonCmd) { substep "Using $PythonCmd ($(& $PythonCmd --version 2>&1))" -# The venv must already exist (created by install.ps1). -# This script (setup.ps1 / "unsloth studio update") only updates packages. -$VenvDir = Join-Path $env:USERPROFILE ".unsloth\studio\unsloth_studio" +# The venv must already exist (created by install.ps1); this script only +# updates packages. UNSLOTH_STUDIO_HOME (or STUDIO_HOME alias) overrides the +# root. UNSLOTH_STUDIO_HOME wins when both are set. Whitespace-only values +# are treated as unset to match Python .strip() semantics. +$_studioOverrideVar = $null +$_studioOverride = $null +if (-not [string]::IsNullOrWhiteSpace($env:UNSLOTH_STUDIO_HOME)) { + $_studioOverrideVar = "UNSLOTH_STUDIO_HOME" + $_studioOverride = $env:UNSLOTH_STUDIO_HOME.Trim() +} elseif (-not [string]::IsNullOrWhiteSpace($env:STUDIO_HOME)) { + $_studioOverrideVar = "STUDIO_HOME" + $_studioOverride = $env:STUDIO_HOME.Trim() +} +if ($_studioOverride) { + if ($_studioOverride -eq "~" -or $_studioOverride -like "~/*" -or $_studioOverride -like "~\*") { + $_studioOverride = (Join-Path $env:USERPROFILE $_studioOverride.Substring(1).TrimStart('/','\')) + } + if (Test-Path -LiteralPath $_studioOverride -PathType Container) { + $StudioHome = (Resolve-Path -LiteralPath $_studioOverride).Path + # why: mirror setup.sh:417 and install.ps1:130 -- fail fast when the + # custom root is read-only instead of erroring later while creating + # sidecar venvs / installing packages. + $_setupWriteProbe = Join-Path $StudioHome (".unsloth-write-probe-" + [guid]::NewGuid()) + try { + [System.IO.File]::WriteAllText($_setupWriteProbe, "") + Remove-Item -LiteralPath $_setupWriteProbe -Force -ErrorAction SilentlyContinue + } catch { + Write-Host "ERROR: $_studioOverrideVar=$StudioHome is not writable." -ForegroundColor Red + exit 1 + } + } else { + Write-Host "ERROR: $_studioOverrideVar=$_studioOverride does not exist." -ForegroundColor Red + Write-Host " Run install.ps1 to create the install root before 'unsloth studio update'." -ForegroundColor Red + exit 1 + } +} else { + $StudioHome = Join-Path $env:USERPROFILE ".unsloth\studio" +} +$VenvDir = Join-Path $StudioHome "unsloth_studio" + +# why: in env-override mode $StudioHome is user-chosen; require the +# ownership marker before Remove-Item so unrelated dirs survive. Gated on +# the canonical comparison so an override pointing at the legacy default +# still behaves like a default install. +$StudioOwnedMarker = ".unsloth-studio-owned" +$LegacyStudioHome = Join-Path $env:USERPROFILE ".unsloth\studio" +$_studioHomeCanon = $StudioHome +if (Test-Path -LiteralPath $_studioHomeCanon -PathType Container) { + $_studioHomeCanon = (Resolve-Path -LiteralPath $_studioHomeCanon).Path +} +if (Test-Path -LiteralPath $LegacyStudioHome -PathType Container) { + $LegacyStudioHome = (Resolve-Path -LiteralPath $LegacyStudioHome).Path +} +$StudioHomeIsCustom = ($_studioHomeCanon -ne $LegacyStudioHome) +function Assert-StudioOwnedOrAbsent { + param( + [Parameter(Mandatory = $true)][string]$Path, + [Parameter(Mandatory = $true)][string]$Label + ) + if (-not (Test-Path -LiteralPath $Path -PathType Container)) { return } + if ($StudioHomeIsCustom -and -not (Test-Path -LiteralPath (Join-Path $Path $StudioOwnedMarker) -PathType Leaf)) { + Write-Host "[ERROR] $Path already exists and is not marked as a Studio-owned $Label." -ForegroundColor Red + Write-Host " Move it aside or choose an empty UNSLOTH_STUDIO_HOME before re-running." -ForegroundColor Yellow + exit 1 + } +} +function Mark-StudioOwned { + param([Parameter(Mandatory = $true)][string]$Path) + if (-not (Test-Path -LiteralPath $Path -PathType Container)) { return } + try { + [System.IO.File]::WriteAllText((Join-Path $Path $StudioOwnedMarker), "") + } catch {} +} # Stale-venv detection: if the venv exists but its torch flavor no longer # matches the current machine, repair according to invocation context. @@ -1504,12 +1574,12 @@ $VenvDir = Join-Path $env:USERPROFILE ".unsloth\studio\unsloth_studio" # In no-torch mode, a missing torch package is expected. $NoTorchMode = $env:UNSLOTH_NO_TORCH -match '^(?i:true|1|yes)$' $InstallerManagedSetup = $env:UNSLOTH_INSTALL_ROLLBACK_MANAGED -match '^(?i:true|1|yes)$' -if ((Test-Path $VenvDir -PathType Container) -and -not $NoTorchMode) { +if ((Test-Path -LiteralPath $VenvDir -PathType Container) -and -not $NoTorchMode) { $VenvPyExe = Join-Path $VenvDir "Scripts\python.exe" $installedTorchTag = $null $shouldRebuild = $false - if (Test-Path $VenvPyExe) { + if (Test-Path -LiteralPath $VenvPyExe) { try { $psi = New-Object System.Diagnostics.ProcessStartInfo $psi.FileName = $VenvPyExe @@ -1558,8 +1628,21 @@ if ((Test-Path $VenvDir -PathType Container) -and -not $NoTorchMode) { exit 1 } substep "Stale venv detected ($reason) -- rebuilding..." "Yellow" + # why: mirror install.ps1 env-mode guard so an update against a custom + # UNSLOTH_STUDIO_HOME never wipes an unrelated unsloth_studio venv; + # -PathType Leaf rejects a directory masquerading as the sentinel. + if ( + $StudioHomeIsCustom -and + -not (Test-Path -LiteralPath (Join-Path $VenvDir $StudioOwnedMarker) -PathType Leaf) -and + -not (Test-Path -LiteralPath (Join-Path $StudioHome "share\studio.conf") -PathType Leaf) -and + -not (Test-Path -LiteralPath (Join-Path $StudioHome "bin\unsloth.exe") -PathType Leaf) + ) { + Write-Host "[ERROR] $VenvDir already exists but does not look like an Unsloth Studio install." -ForegroundColor Red + Write-Host " Move it aside or choose an empty UNSLOTH_STUDIO_HOME before re-running." -ForegroundColor Yellow + exit 1 + } try { - Remove-Item $VenvDir -Recurse -Force -ErrorAction Stop + Remove-Item -LiteralPath $VenvDir -Recurse -Force -ErrorAction Stop } catch { Write-Host " [ERROR] Could not remove stale venv: $($_.Exception.Message)" -ForegroundColor Red Write-Host " Close any running Studio/Python processes and re-run setup." -ForegroundColor Red @@ -1568,7 +1651,7 @@ if ((Test-Path $VenvDir -PathType Container) -and -not $NoTorchMode) { } } -if (-not (Test-Path $VenvDir)) { +if (-not (Test-Path -LiteralPath $VenvDir)) { Write-Host "[ERROR] Virtual environment not found at $VenvDir" -ForegroundColor Red Write-Host " Run install.ps1 first to create the environment:" -ForegroundColor Yellow Write-Host " irm https://unsloth.ai/install.ps1 | iex" -ForegroundColor Yellow @@ -1759,17 +1842,19 @@ if ($stackExit -ne 0) { # ── Pre-install transformers 5.x into .venv_t5_530/ and .venv_t5_550/ ── # Runs outside the deps fast-path gate so that upgrades from the legacy # single .venv_t5 are always migrated to the tiered layout. -$VenvT5_530Dir = Join-Path $env:USERPROFILE ".unsloth\studio\.venv_t5_530" -$VenvT5_550Dir = Join-Path $env:USERPROFILE ".unsloth\studio\.venv_t5_550" -$VenvT5Legacy = Join-Path $env:USERPROFILE ".unsloth\studio\.venv_t5" +# T5 sidecar venvs live under the resolved $StudioHome so custom installs are self-contained. +$VenvT5_530Dir = Join-Path $StudioHome ".venv_t5_530" +$VenvT5_550Dir = Join-Path $StudioHome ".venv_t5_550" +$VenvT5Legacy = Join-Path $StudioHome ".venv_t5" $_NeedT5Install = $false -if (Test-Path $VenvT5Legacy) { - Remove-Item -Recurse -Force $VenvT5Legacy +if (Test-Path -LiteralPath $VenvT5Legacy) { + Assert-StudioOwnedOrAbsent -Path $VenvT5Legacy -Label "legacy transformers sidecar venv" + Remove-Item -LiteralPath $VenvT5Legacy -Recurse -Force $_NeedT5Install = $true } -if (-not (Test-Path $VenvT5_530Dir)) { $_NeedT5Install = $true } -if (-not (Test-Path $VenvT5_550Dir)) { $_NeedT5Install = $true } +if (-not (Test-Path -LiteralPath $VenvT5_530Dir)) { $_NeedT5Install = $true } +if (-not (Test-Path -LiteralPath $VenvT5_550Dir)) { $_NeedT5Install = $true } # Also reinstall when python deps were updated if (-not $SkipPythonDeps) { $_NeedT5Install = $true } @@ -1781,8 +1866,10 @@ $ErrorActionPreference = "Continue" # --- .venv_t5_530 (transformers 5.3.0) --- substep "pre-installing transformers 5.3.0 for newer model support..." -if (Test-Path $VenvT5_530Dir) { Remove-Item -Recurse -Force $VenvT5_530Dir } -New-Item -ItemType Directory -Path $VenvT5_530Dir -Force | Out-Null +Assert-StudioOwnedOrAbsent -Path $VenvT5_530Dir -Label "transformers 5.3 sidecar venv" +if (Test-Path -LiteralPath $VenvT5_530Dir) { Remove-Item -LiteralPath $VenvT5_530Dir -Recurse -Force } +[System.IO.Directory]::CreateDirectory($VenvT5_530Dir) | Out-Null +Mark-StudioOwned -Path $VenvT5_530Dir foreach ($pkg in @("transformers==5.3.0", "huggingface_hub==1.8.0", "hf_xet==1.4.2")) { if ($script:UnslothVerbose) { Fast-Install --target $VenvT5_530Dir --no-deps $pkg @@ -1814,8 +1901,10 @@ step "transformers" "5.3.0 pre-installed" # --- .venv_t5_550 (transformers 5.5.0) --- substep "pre-installing transformers 5.5.0 for Gemma 4 support..." -if (Test-Path $VenvT5_550Dir) { Remove-Item -Recurse -Force $VenvT5_550Dir } -New-Item -ItemType Directory -Path $VenvT5_550Dir -Force | Out-Null +Assert-StudioOwnedOrAbsent -Path $VenvT5_550Dir -Label "transformers 5.5 sidecar venv" +if (Test-Path -LiteralPath $VenvT5_550Dir) { Remove-Item -LiteralPath $VenvT5_550Dir -Recurse -Force } +[System.IO.Directory]::CreateDirectory($VenvT5_550Dir) | Out-Null +Mark-StudioOwned -Path $VenvT5_550Dir foreach ($pkg in @("transformers==5.5.0", "huggingface_hub==1.8.0", "hf_xet==1.4.2")) { if ($script:UnslothVerbose) { Fast-Install --target $VenvT5_550Dir --no-deps $pkg @@ -1851,8 +1940,15 @@ step "transformers" "5.5.0 pre-installed" # ========================================================================== # PHASE 3.4: Prefer prebuilt llama.cpp bundles before source build # ========================================================================== -$UnslothHome = Join-Path $env:USERPROFILE ".unsloth" -if (-not (Test-Path $UnslothHome)) { New-Item -ItemType Directory -Force $UnslothHome | Out-Null } +# Nest llama.cpp under $StudioHome only for real env-overrides, never the +# legacy default. Reuses $StudioHomeIsCustom from the canonical comparison +# computed above so the llama.cpp nest matches ownership-guard semantics. +if ($StudioHomeIsCustom) { + $UnslothHome = $StudioHome +} else { + $UnslothHome = Join-Path $env:USERPROFILE ".unsloth" +} +if (-not (Test-Path -LiteralPath $UnslothHome)) { [System.IO.Directory]::CreateDirectory($UnslothHome) | Out-Null } $LlamaCppDir = Join-Path $UnslothHome "llama.cpp" $NeedLlamaSourceBuild = $false $SkipPrebuiltInstall = $false @@ -1954,9 +2050,15 @@ if ($env:UNSLOTH_LLAMA_FORCE_COMPILE -eq "1") { } else { Write-Host "" substep "installing prebuilt llama.cpp bundle (preferred path)..." - if (Test-Path $LlamaCppDir) { + if (Test-Path -LiteralPath $LlamaCppDir) { substep "Existing llama.cpp install detected -- validating staged prebuilt update before replacement" } + # why: install_llama_prebuilt.py uses os.replace(), which would displace + # an unrelated $env:UNSLOTH_STUDIO_HOME\llama.cpp before the source-build + # ownership check below ever runs. + if ($StudioHomeIsCustom) { + Assert-StudioOwnedOrAbsent -Path $LlamaCppDir -Label "llama.cpp install" + } $prebuiltArgs = @( "$PSScriptRoot\install_llama_prebuilt.py", "--install-dir", $LlamaCppDir, @@ -2001,6 +2103,9 @@ if ($env:UNSLOTH_LLAMA_FORCE_COMPILE -eq "1") { } else { step "llama.cpp" "prebuilt installed and validated" } + if ($StudioHomeIsCustom -and (Test-Path -LiteralPath $LlamaCppDir -PathType Container)) { + Mark-StudioOwned -Path $LlamaCppDir + } $installedRelease = Get-InstalledLlamaPrebuiltRelease -InstallDir $LlamaCppDir if ($installedRelease) { substep $installedRelease @@ -2008,7 +2113,7 @@ if ($env:UNSLOTH_LLAMA_FORCE_COMPILE -eq "1") { } elseif ($prebuiltExit -eq 3) { step "llama.cpp" "install blocked by active llama.cpp process" "Yellow" Write-LlamaFailureLog -Output $prebuiltOutput - if (Test-Path $LlamaCppDir) { + if (Test-Path -LiteralPath $LlamaCppDir) { substep "Existing install was restored" "Yellow" } substep "Close Studio or other llama.cpp users and retry" "Yellow" @@ -2016,7 +2121,7 @@ if ($env:UNSLOTH_LLAMA_FORCE_COMPILE -eq "1") { } else { step "llama.cpp" "prebuilt install failed (continuing)" "Yellow" Write-LlamaFailureLog -Output $prebuiltOutput - if (Test-Path $LlamaCppDir) { + if (Test-Path -LiteralPath $LlamaCppDir) { substep "Prebuilt update failed; existing install was restored or cleaned before source build fallback" "Yellow" } substep "Prebuilt llama.cpp path unavailable or failed validation -- falling back to source build" "Yellow" @@ -2092,10 +2197,10 @@ $HasCmakeForBuild = $null -ne (Get-Command cmake -ErrorAction SilentlyContinue) # Check if existing llama-server matches current GPU mode. A CUDA-built binary # on a now-CPU-only machine (or vice versa) needs to be rebuilt. $NeedRebuild = $false -if (Test-Path $LlamaServerBin) { +if (Test-Path -LiteralPath $LlamaServerBin) { $CmakeCacheFile = Join-Path $BuildDir "CMakeCache.txt" - if (Test-Path $CmakeCacheFile) { - $cachedCuda = Select-String -Path $CmakeCacheFile -Pattern 'GGML_CUDA:BOOL=ON' -Quiet + if (Test-Path -LiteralPath $CmakeCacheFile) { + $cachedCuda = Select-String -LiteralPath $CmakeCacheFile -Pattern 'GGML_CUDA:BOOL=ON' -Quiet if ($HasNvidiaSmi -and -not $cachedCuda) { Write-Host " Existing llama-server is CPU-only but GPU is available -- rebuilding" -ForegroundColor Yellow $NeedRebuild = $true @@ -2109,7 +2214,7 @@ if (Test-Path $LlamaServerBin) { if (-not $NeedLlamaSourceBuild) { Write-Host "" step "llama.cpp" "prebuilt (validated)" -} elseif ((Test-Path $LlamaServerBin) -and -not $NeedRebuild -and $RequestedLlamaTag -ne "master") { +} elseif ((Test-Path -LiteralPath $LlamaServerBin) -and -not $NeedRebuild -and $RequestedLlamaTag -ne "master") { # Skip rebuild only for pinned tags (e.g. b8635). When the requested # tag is "master" (a moving target), always rebuild so the binary picks # up new model architecture support (e.g. Gemma 4). @@ -2211,7 +2316,13 @@ if (-not $NeedLlamaSourceBuild) { $UseConcreteRef = ($ResolvedSourceRef -ne "latest" -and -not [string]::IsNullOrWhiteSpace($ResolvedSourceRef)) - if (Test-Path (Join-Path $LlamaCppDir ".git")) { + if (Test-Path -LiteralPath (Join-Path $LlamaCppDir ".git")) { + # why: in-place git mutation (remote set-url, checkout -B, clean -fdx) + # rewrites $LlamaCppDir; mirror the prebuilt and temp-dir-swap guards + # so an unrelated workspace .git tree is never silently overwritten. + if ($StudioHomeIsCustom) { + Assert-StudioOwnedOrAbsent -Path $LlamaCppDir -Label "llama.cpp install" + } Write-Host " Syncing llama.cpp to $ResolvedSourceRef..." -ForegroundColor Gray # Always sync the remote URL so switching between default/fork sources works Invoke-SetupCommand -AlwaysQuiet { git -C $LlamaCppDir remote set-url origin "$ResolvedSourceUrl.git" } | Out-Null @@ -2282,24 +2393,30 @@ if (-not $NeedLlamaSourceBuild) { } } } + # why: in-place git-sync (the temp-dir clone path calls Mark-StudioOwned + # at swap-time) must mark the existing tree so a subsequent prebuilt + # update path's Assert-StudioOwnedOrAbsent does not exit on the same root. + if ($BuildOk -and $StudioHomeIsCustom) { + Mark-StudioOwned -Path $LlamaCppDir + } } else { Write-Host " Cloning llama.cpp @ $ResolvedSourceRef..." -ForegroundColor Gray $buildTmp = "$LlamaCppDir.build.$PID" - $null = New-Item -ItemType Directory -Force -Path (Split-Path $LlamaCppDir -Parent) - if (Test-Path $buildTmp) { Remove-Item -Recurse -Force $buildTmp } + $null = [System.IO.Directory]::CreateDirectory((Split-Path -LiteralPath $LlamaCppDir)) + if (Test-Path -LiteralPath $buildTmp) { Remove-Item -LiteralPath $buildTmp -Recurse -Force } if ($LlamaPr) { $cloneExit = Invoke-SetupCommand -AlwaysQuiet { git clone --depth 1 "$LlamaSource.git" $buildTmp } if ($cloneExit -ne 0) { $BuildOk = $false $FailedStep = "git clone" - if (Test-Path $buildTmp) { Remove-Item -Recurse -Force $buildTmp } + if (Test-Path -LiteralPath $buildTmp) { Remove-Item -LiteralPath $buildTmp -Recurse -Force } } if ($BuildOk) { $fetchExit = Invoke-SetupCommand -AlwaysQuiet { git -C $buildTmp fetch --depth 1 origin "pull/$LlamaPr/head:pr-$LlamaPr" } if ($fetchExit -ne 0) { $BuildOk = $false $FailedStep = "git fetch PR #$LlamaPr" - if (Test-Path $buildTmp) { Remove-Item -Recurse -Force $buildTmp } + if (Test-Path -LiteralPath $buildTmp) { Remove-Item -LiteralPath $buildTmp -Recurse -Force } } } if ($BuildOk) { @@ -2307,7 +2424,7 @@ if (-not $NeedLlamaSourceBuild) { if ($checkoutExit -ne 0) { $BuildOk = $false $FailedStep = "git checkout PR #$LlamaPr" - if (Test-Path $buildTmp) { Remove-Item -Recurse -Force $buildTmp } + if (Test-Path -LiteralPath $buildTmp) { Remove-Item -LiteralPath $buildTmp -Recurse -Force } } } } elseif ($ResolvedSourceRefKind -eq "pull") { @@ -2315,14 +2432,14 @@ if (-not $NeedLlamaSourceBuild) { if ($cloneExit -ne 0) { $BuildOk = $false $FailedStep = "git clone" - if (Test-Path $buildTmp) { Remove-Item -Recurse -Force $buildTmp } + if (Test-Path -LiteralPath $buildTmp) { Remove-Item -LiteralPath $buildTmp -Recurse -Force } } if ($BuildOk) { $fetchExit = Invoke-SetupCommand -AlwaysQuiet { git -C $buildTmp fetch --depth 1 origin $ResolvedSourceRef } if ($fetchExit -ne 0) { $BuildOk = $false $FailedStep = "git fetch source PR ref" - if (Test-Path $buildTmp) { Remove-Item -Recurse -Force $buildTmp } + if (Test-Path -LiteralPath $buildTmp) { Remove-Item -LiteralPath $buildTmp -Recurse -Force } } } if ($BuildOk) { @@ -2330,7 +2447,7 @@ if (-not $NeedLlamaSourceBuild) { if ($checkoutExit -ne 0) { $BuildOk = $false $FailedStep = "git checkout source PR ref" - if (Test-Path $buildTmp) { Remove-Item -Recurse -Force $buildTmp } + if (Test-Path -LiteralPath $buildTmp) { Remove-Item -LiteralPath $buildTmp -Recurse -Force } } } } elseif ($ResolvedSourceRefKind -eq "commit") { @@ -2338,14 +2455,14 @@ if (-not $NeedLlamaSourceBuild) { if ($cloneExit -ne 0) { $BuildOk = $false $FailedStep = "git clone" - if (Test-Path $buildTmp) { Remove-Item -Recurse -Force $buildTmp } + if (Test-Path -LiteralPath $buildTmp) { Remove-Item -LiteralPath $buildTmp -Recurse -Force } } if ($BuildOk) { $fetchExit = Invoke-SetupCommand -AlwaysQuiet { git -C $buildTmp fetch --depth 1 origin $ResolvedSourceRef } if ($fetchExit -ne 0) { $BuildOk = $false $FailedStep = "git fetch source commit" - if (Test-Path $buildTmp) { Remove-Item -Recurse -Force $buildTmp } + if (Test-Path -LiteralPath $buildTmp) { Remove-Item -LiteralPath $buildTmp -Recurse -Force } } } if ($BuildOk) { @@ -2353,7 +2470,7 @@ if (-not $NeedLlamaSourceBuild) { if ($checkoutExit -ne 0) { $BuildOk = $false $FailedStep = "git checkout source commit" - if (Test-Path $buildTmp) { Remove-Item -Recurse -Force $buildTmp } + if (Test-Path -LiteralPath $buildTmp) { Remove-Item -LiteralPath $buildTmp -Recurse -Force } } } } else { @@ -2366,7 +2483,7 @@ if (-not $NeedLlamaSourceBuild) { if ($cloneExit -ne 0) { $BuildOk = $false $FailedStep = "git clone" - if (Test-Path $buildTmp) { Remove-Item -Recurse -Force $buildTmp } + if (Test-Path -LiteralPath $buildTmp) { Remove-Item -LiteralPath $buildTmp -Recurse -Force } } } # Use temp dir for build; swap into $LlamaCppDir only after build succeeds @@ -2482,14 +2599,16 @@ if (-not $NeedLlamaSourceBuild) { # Swap temp build dir into final location (only if we built in a temp dir) if ($BuildOk -and $LlamaCppDir -ne $OriginalLlamaCppDir) { - if (Test-Path $OriginalLlamaCppDir) { Remove-Item -Recurse -Force $OriginalLlamaCppDir } - Move-Item $LlamaCppDir $OriginalLlamaCppDir + Assert-StudioOwnedOrAbsent -Path $OriginalLlamaCppDir -Label "llama.cpp install" + if (Test-Path -LiteralPath $OriginalLlamaCppDir) { Remove-Item -LiteralPath $OriginalLlamaCppDir -Recurse -Force } + Move-Item -LiteralPath $LlamaCppDir -Destination $OriginalLlamaCppDir $LlamaCppDir = $OriginalLlamaCppDir $BuildDir = Join-Path $LlamaCppDir "build" $LlamaServerBin = Join-Path $BuildDir "bin\Release\llama-server.exe" + Mark-StudioOwned -Path $LlamaCppDir } elseif (-not $BuildOk -and $LlamaCppDir -ne $OriginalLlamaCppDir) { # Build failed -- clean up temp dir, preserve existing install - if (Test-Path $LlamaCppDir) { Remove-Item -Recurse -Force $LlamaCppDir } + if (Test-Path -LiteralPath $LlamaCppDir) { Remove-Item -LiteralPath $LlamaCppDir -Recurse -Force } $LlamaCppDir = $OriginalLlamaCppDir $BuildDir = Join-Path $LlamaCppDir "build" $LlamaServerBin = Join-Path $BuildDir "bin\Release\llama-server.exe" @@ -2504,16 +2623,16 @@ if (-not $NeedLlamaSourceBuild) { $totalSec = [math]::Round($totalSw.Elapsed.TotalSeconds % 60, 1) # -- Summary -- - if ($BuildOk -and (Test-Path $LlamaServerBin)) { + if ($BuildOk -and (Test-Path -LiteralPath $LlamaServerBin)) { step "llama.cpp" "built" $QuantizeBin = Join-Path $BuildDir "bin\Release\llama-quantize.exe" - if (Test-Path $QuantizeBin) { + if (Test-Path -LiteralPath $QuantizeBin) { step "llama-quantize" "built" } step "build time" "${totalMin}m ${totalSec}s" "DarkGray" } else { $altBin = Join-Path $BuildDir "bin\llama-server.exe" - if ($BuildOk -and (Test-Path $altBin)) { + if ($BuildOk -and (Test-Path -LiteralPath $altBin)) { step "llama.cpp" "built" step "build time" "${totalMin}m ${totalSec}s" "DarkGray" } else { diff --git a/studio/setup.sh b/studio/setup.sh index 3e875eed30..ff93d2d41d 100755 --- a/studio/setup.sh +++ b/studio/setup.sh @@ -417,7 +417,36 @@ if [ -d "$SCRIPT_DIR/backend/core/data_recipe/oxc-validator" ] && command -v npm fi # ── Python venv + deps ── -STUDIO_HOME="$HOME/.unsloth/studio" +# UNSLOTH_STUDIO_HOME (or STUDIO_HOME alias) overrides the install root +# (mirrors install.sh). UNSLOTH_STUDIO_HOME wins when both are set. +_studio_override_var="" +_studio_override="${UNSLOTH_STUDIO_HOME:-}" +if [ -n "$_studio_override" ]; then + _studio_override_var="UNSLOTH_STUDIO_HOME" +else + _studio_override="${STUDIO_HOME:-}" + [ -n "$_studio_override" ] && _studio_override_var="STUDIO_HOME" +fi +# Strip whitespace so " " is treated as unset (matches Python .strip()). +_studio_override=$(printf '%s' "$_studio_override" | sed -e 's/^[[:space:]]*//' -e 's/[[:space:]]*$//') +case "$_studio_override" in + "~") _studio_override="$HOME" ;; + "~/"*) _studio_override="$HOME/${_studio_override#'~/'}" ;; +esac +if [ -n "$_studio_override" ]; then + # setup.sh runs against an existing install (via 'unsloth studio update'); + # a typo in the override must fail fast instead of materializing an + # empty workspace dir. Mirrors setup.ps1 behavior. + if [ ! -d "$_studio_override" ]; then + echo "ERROR: $_studio_override_var=$_studio_override does not exist." >&2 + echo " Run install.sh to create the install root before 'unsloth studio update'." >&2 + exit 1 + fi + [ -w "$_studio_override" ] || { echo "ERROR: $_studio_override_var=$_studio_override is not writable." >&2; exit 1; } + STUDIO_HOME="$(CDPATH= cd -P -- "$_studio_override" && pwd -P)" || exit 1 +else + STUDIO_HOME="$HOME/.unsloth/studio" +fi VENV_DIR="$STUDIO_HOME/unsloth_studio" VENV_T5_530_DIR="$STUDIO_HOME/.venv_t5_530" VENV_T5_550_DIR="$STUDIO_HOME/.venv_t5_550" @@ -542,9 +571,39 @@ fi # # Runs outside the _SKIP_PYTHON_DEPS gate so that upgrades from legacy # single .venv_t5 are always migrated to the tiered layout. +# why: in env-override mode $STUDIO_HOME is user-chosen; require the +# ownership marker before rm -rf so unrelated dirs survive. Gated on the +# canonical comparison so an override pointing at the legacy default still +# behaves like a default install. +_STUDIO_OWNED_MARKER=".unsloth-studio-owned" +_LEGACY_STUDIO_HOME="$HOME/.unsloth/studio" +_studio_home_canon="$STUDIO_HOME" +if [ -d "$_studio_home_canon" ]; then + _studio_home_canon=$(CDPATH= cd -P -- "$_studio_home_canon" 2>/dev/null && pwd -P) \ + || _studio_home_canon="$STUDIO_HOME" +fi +if [ -d "$_LEGACY_STUDIO_HOME" ]; then + _LEGACY_STUDIO_HOME=$(CDPATH= cd -P -- "$_LEGACY_STUDIO_HOME" 2>/dev/null && pwd -P) \ + || _LEGACY_STUDIO_HOME="$HOME/.unsloth/studio" +fi +_STUDIO_HOME_IS_CUSTOM=false +if [ "$_studio_home_canon" != "$_LEGACY_STUDIO_HOME" ]; then + _STUDIO_HOME_IS_CUSTOM=true +fi +_assert_studio_owned_or_absent() { + _aso_dir="$1" + _aso_label="$2" + [ -d "$_aso_dir" ] || return 0 + if [ "$_STUDIO_HOME_IS_CUSTOM" = true ] && [ ! -f "$_aso_dir/$_STUDIO_OWNED_MARKER" ]; then + echo "ERROR: $_aso_dir already exists and is not marked as a Studio-owned $_aso_label." >&2 + echo " Move it aside or choose an empty UNSLOTH_STUDIO_HOME before re-running." >&2 + exit 1 + fi +} _NEED_T5_INSTALL=false if [ -d "$STUDIO_HOME/.venv_t5" ]; then # Legacy layout — migrate + _assert_studio_owned_or_absent "$STUDIO_HOME/.venv_t5" "legacy transformers sidecar venv" rm -rf "$STUDIO_HOME/.venv_t5" _NEED_T5_INSTALL=true fi @@ -554,16 +613,20 @@ fi [ "$_SKIP_PYTHON_DEPS" = false ] && _NEED_T5_INSTALL=true if [ "$_NEED_T5_INSTALL" = true ]; then + _assert_studio_owned_or_absent "$VENV_T5_530_DIR" "transformers 5.3 sidecar venv" [ -d "$VENV_T5_530_DIR" ] && rm -rf "$VENV_T5_530_DIR" mkdir -p "$VENV_T5_530_DIR" + : > "$VENV_T5_530_DIR/$_STUDIO_OWNED_MARKER" 2>/dev/null || true run_quiet "install transformers 5.3.0" fast_install --target "$VENV_T5_530_DIR" --no-deps "transformers==5.3.0" run_quiet "install huggingface_hub for t5_530" fast_install --target "$VENV_T5_530_DIR" --no-deps "huggingface_hub==1.8.0" run_quiet "install hf_xet for t5_530" fast_install --target "$VENV_T5_530_DIR" --no-deps "hf_xet==1.4.2" run_quiet "install tiktoken for t5_530" fast_install --target "$VENV_T5_530_DIR" "tiktoken" step "transformers" "5.3.0 pre-installed" + _assert_studio_owned_or_absent "$VENV_T5_550_DIR" "transformers 5.5 sidecar venv" [ -d "$VENV_T5_550_DIR" ] && rm -rf "$VENV_T5_550_DIR" mkdir -p "$VENV_T5_550_DIR" + : > "$VENV_T5_550_DIR/$_STUDIO_OWNED_MARKER" 2>/dev/null || true run_quiet "install transformers 5.5.0" fast_install --target "$VENV_T5_550_DIR" --no-deps "transformers==5.5.0" run_quiet "install huggingface_hub for t5_550" fast_install --target "$VENV_T5_550_DIR" --no-deps "huggingface_hub==1.8.0" run_quiet "install hf_xet for t5_550" fast_install --target "$VENV_T5_550_DIR" --no-deps "hf_xet==1.4.2" @@ -573,7 +636,13 @@ fi fi # ── 7. Prefer prebuilt llama.cpp bundles before any source build path ── -UNSLOTH_HOME="$HOME/.unsloth" +# Nest llama.cpp under $STUDIO_HOME only for real env-overrides; legacy +# default keeps ~/.unsloth/llama.cpp so pre-PR builds are still discovered. +if [ "$_STUDIO_HOME_IS_CUSTOM" = true ]; then + UNSLOTH_HOME="$STUDIO_HOME" +else + UNSLOTH_HOME="$HOME/.unsloth" +fi mkdir -p "$UNSLOTH_HOME" LLAMA_CPP_DIR="$UNSLOTH_HOME/llama.cpp" LLAMA_SERVER_BIN="$LLAMA_CPP_DIR/build/bin/llama-server" @@ -635,6 +704,12 @@ else if [ -d "$LLAMA_CPP_DIR" ]; then substep "existing install detected -- validating update" fi + # why: install_llama_prebuilt.py uses os.replace(), which would displace + # an unrelated $UNSLOTH_STUDIO_HOME/llama.cpp before the source-build + # ownership check below ever runs. + if [ "$_STUDIO_HOME_IS_CUSTOM" = true ]; then + _assert_studio_owned_or_absent "$LLAMA_CPP_DIR" "llama.cpp install" + fi _PREBUILT_CMD=( python "$SCRIPT_DIR/install_llama_prebuilt.py" --install-dir "$LLAMA_CPP_DIR" @@ -662,6 +737,9 @@ else else step "llama.cpp" "prebuilt installed and validated" fi + if [ "$_STUDIO_HOME_IS_CUSTOM" = true ] && [ -d "$LLAMA_CPP_DIR" ]; then + : > "$LLAMA_CPP_DIR/$_STUDIO_OWNED_MARKER" 2>/dev/null || true + fi print_installed_llama_prebuilt_release "$LLAMA_CPP_DIR" verbose_substep "llama.cpp install dir: $LLAMA_CPP_DIR" rm -f "$_PREBUILT_LOG" @@ -1032,8 +1110,10 @@ else # Swap only after build succeeds -- preserves existing install on failure if [ "$BUILD_OK" = true ]; then + _assert_studio_owned_or_absent "$LLAMA_CPP_DIR" "llama.cpp install" rm -rf "$LLAMA_CPP_DIR" mv "$_BUILD_TMP" "$LLAMA_CPP_DIR" + : > "$LLAMA_CPP_DIR/$_STUDIO_OWNED_MARKER" 2>/dev/null || true # Symlink to llama.cpp root -- check_llama_cpp() looks for the binary there QUANTIZE_BIN="$LLAMA_CPP_DIR/build/bin/llama-quantize" if [ -f "$QUANTIZE_BIN" ]; then diff --git a/studio/src-tauri/src/commands.rs b/studio/src-tauri/src/commands.rs index 48a0af6e48..e9a27644df 100644 --- a/studio/src-tauri/src/commands.rs +++ b/studio/src-tauri/src/commands.rs @@ -60,6 +60,11 @@ pub async fn check_install_status() -> bool { cmd.env_remove("PYTHONPATH"); } + // Tauri uses the legacy root regardless of UNSLOTH_STUDIO_HOME / STUDIO_HOME; + // probe subprocesses must follow the same isolation as process.rs. + cmd.env_remove("UNSLOTH_STUDIO_HOME"); + cmd.env_remove("STUDIO_HOME"); + let mut child = match cmd.spawn() { Ok(c) => c, Err(e) => { diff --git a/studio/src-tauri/src/desktop_auth.rs b/studio/src-tauri/src/desktop_auth.rs index 483e7c0432..49b19008fb 100644 --- a/studio/src-tauri/src/desktop_auth.rs +++ b/studio/src-tauri/src/desktop_auth.rs @@ -203,6 +203,11 @@ async fn provision_desktop_auth() -> Result<(), String> { cmd.env_remove("PYTHONHOME"); cmd.env_remove("PYTHONPATH"); } + + // Tauri uses the legacy root regardless of UNSLOTH_STUDIO_HOME / STUDIO_HOME. + // Scrub so provisioning writes match what the Rust auth code reads. + cmd.env_remove("UNSLOTH_STUDIO_HOME"); + cmd.env_remove("STUDIO_HOME"); #[cfg(windows)] { use std::os::windows::process::CommandExt; diff --git a/studio/src-tauri/src/install.rs b/studio/src-tauri/src/install.rs index 9d672f5e73..024b730735 100644 --- a/studio/src-tauri/src/install.rs +++ b/studio/src-tauri/src/install.rs @@ -196,6 +196,11 @@ fn spawn_script( cmd.env_remove("PYTHONPATH"); } + // Tauri only does default-root installs; install.sh / install.ps1 reject + // these under --tauri. Scrub so an inherited value can't trip the guard. + cmd.env_remove("UNSLOTH_STUDIO_HOME"); + cmd.env_remove("STUDIO_HOME"); + // On Windows, launch the installer directly with CREATE_NO_WINDOW. // The app process is assigned to a KILL_ON_JOB_CLOSE job in main.rs, so // child cleanup on crash comes from inherited job membership instead. diff --git a/studio/src-tauri/src/preflight.rs b/studio/src-tauri/src/preflight.rs index d3df06d057..c0bbb07b36 100644 --- a/studio/src-tauri/src/preflight.rs +++ b/studio/src-tauri/src/preflight.rs @@ -102,6 +102,11 @@ async fn run_cli_probe(bin: &std::path::Path, args: &[&str]) -> bool { cmd.env_remove("PYTHONPATH"); } + // Tauri uses the legacy root regardless of UNSLOTH_STUDIO_HOME / STUDIO_HOME; + // probe subprocesses must follow the same isolation as process.rs. + cmd.env_remove("UNSLOTH_STUDIO_HOME"); + cmd.env_remove("STUDIO_HOME"); + #[cfg(windows)] { use std::os::windows::process::CommandExt; @@ -135,6 +140,11 @@ async fn probe_cli_capability(bin: &std::path::Path) -> Option = { use std::os::windows::process::CommandExt; diff --git a/tests/test_studio_install_workspace_guard.py b/tests/test_studio_install_workspace_guard.py new file mode 100644 index 0000000000..d077cdb824 --- /dev/null +++ b/tests/test_studio_install_workspace_guard.py @@ -0,0 +1,1021 @@ +"""install.sh / install.ps1 must refuse to rm -rf an existing +$STUDIO_HOME/unsloth_studio in env-override mode unless the directory +carries a Studio sentinel (share/studio.conf or bin/unsloth). Also +asserts studio/setup.ps1 has the matching writability probe that +setup.sh:417 already performs.""" + +from __future__ import annotations + +import re +import subprocess +from pathlib import Path + +REPO_ROOT = Path(__file__).resolve().parents[1] +INSTALL_SH = REPO_ROOT / "install.sh" +INSTALL_PS1 = REPO_ROOT / "install.ps1" +SETUP_PS1 = REPO_ROOT / "studio" / "setup.ps1" +SETUP_SH = REPO_ROOT / "studio" / "setup.sh" + +# Stubs for helpers that the extracted install.sh guard block calls in real +# installs (`substep` for status output, `_start_studio_venv_replacement` for +# the rollback-managed move). The tests run the block in isolation, so we +# stand in a minimal `mv`-based replacement that exercises the same observable +# effect (venv directory is no longer present at $VENV_DIR after a permitted +# cleanup) without dragging in install.sh's full rollback machinery. +_INSTALL_GUARD_STUBS = ( + "substep() { :; }\n" + "_start_studio_venv_replacement() {\n" + ' mv -- "$1" "$1.replaced"\n' + "}\n" +) + + +def _extract_install_sh_guard_block() -> str: + """Pull the `if [ -x "$VENV_DIR/bin/python" ]; then ... fi` block out + of install.sh as a self-contained snippet. Stops at the first elif so + the block can be paired with a synthetic else and run in isolation.""" + src = INSTALL_SH.read_text() + m = re.search( + r'(if \[ -x "\$VENV_DIR/bin/python" \]; then\n.*?)elif \[ "\$_STUDIO_HOME_REDIRECT" != "env"', + src, + re.DOTALL, + ) + assert m, "install.sh venv guard block not found" + return m.group(1) + "fi\n" + + +def _build_install_guard_script( + studio_home: Path, redirect: str, block: str | None = None +) -> str: + """Build a self-contained bash script that exercises the extracted + guard block. Includes stubs for substep / _start_studio_venv_replacement + so the snippet runs without install.sh's full rollback machinery.""" + if block is None: + block = _extract_install_sh_guard_block() + return ( + _INSTALL_GUARD_STUBS + + f'STUDIO_HOME="{studio_home}"\n' + + f'VENV_DIR="$STUDIO_HOME/unsloth_studio"\n' + + f'_STUDIO_HOME_REDIRECT="{redirect}"\n' + + block + + "echo RESULT=ok\n" + ) + + +def _run_install_guard( + studio_home: Path, + redirect: str, + create_share_conf: bool = False, + create_bin_shim: bool = False, + create_venv_marker: bool = False, +) -> subprocess.CompletedProcess: + venv_dir = studio_home / "unsloth_studio" + (venv_dir / "bin").mkdir(parents = True, exist_ok = True) + py = venv_dir / "bin" / "python" + py.write_text("#!/bin/sh\nexit 0\n") + py.chmod(0o755) + if create_share_conf: + (studio_home / "share").mkdir(parents = True, exist_ok = True) + (studio_home / "share" / "studio.conf").write_text("") + if create_bin_shim: + (studio_home / "bin").mkdir(parents = True, exist_ok = True) + (studio_home / "bin" / "unsloth").write_text("") + if create_venv_marker: + (venv_dir / ".unsloth-studio-owned").write_text("") + script = _build_install_guard_script(studio_home, redirect) + return subprocess.run( + ["bash", "-c", script], + env = {"PATH": "/usr/bin:/bin"}, + text = True, + capture_output = True, + ) + + +def test_env_mode_blocks_unsloth_studio_without_sentinels(tmp_path): + studio_home = tmp_path / "ws" + res = _run_install_guard(studio_home, redirect = "env") + assert res.returncode != 0, ( + "env-mode without sentinels must refuse to rm -rf $VENV_DIR; " + f"stdout={res.stdout!r} stderr={res.stderr!r}" + ) + assert "does not look like an Unsloth Studio install" in res.stderr + assert (studio_home / "unsloth_studio" / "bin" / "python").is_file() + + +def test_env_mode_passes_when_share_studio_conf_present(tmp_path): + studio_home = tmp_path / "ws" + res = _run_install_guard(studio_home, redirect = "env", create_share_conf = True) + assert res.returncode == 0, ( + f"share/studio.conf sentinel must allow cleanup;" + f" stdout={res.stdout!r} stderr={res.stderr!r}" + ) + assert "RESULT=ok" in res.stdout + assert not (studio_home / "unsloth_studio").exists() + + +def test_env_mode_passes_when_bin_unsloth_shim_present(tmp_path): + studio_home = tmp_path / "ws" + res = _run_install_guard(studio_home, redirect = "env", create_bin_shim = True) + assert res.returncode == 0, res.stderr + assert not (studio_home / "unsloth_studio").exists() + + +def test_default_mode_skips_sentinel_check(tmp_path): + studio_home = tmp_path / "ws" + res = _run_install_guard(studio_home, redirect = "default") + assert res.returncode == 0, res.stderr + assert "RESULT=ok" in res.stdout + assert not (studio_home / "unsloth_studio").exists() + + +def test_install_ps1_has_matching_env_mode_guard(): + src = INSTALL_PS1.read_text() + block_start = src.index("if (Test-Path -LiteralPath $VenvPython)") + block = src[block_start : block_start + 2000] + assert ( + "$StudioRedirectMode -eq 'env'" in block + ), "install.ps1 must gate Remove-Item $VenvDir on env-mode" + assert ( + "share\\studio.conf" in block + ), "install.ps1 guard must check share\\studio.conf sentinel" + assert ( + "bin\\unsloth.exe" in block + ), "install.ps1 guard must check bin\\unsloth.exe sentinel" + assert "Refusing to delete non-Studio venv" in block + + +def test_setup_ps1_has_writability_probe(): + src = SETUP_PS1.read_text() + idx = src.index("if (Test-Path -LiteralPath $_studioOverride -PathType Container)") + block = src[idx : idx + 2000] + assert ( + "WriteAllText" in block + ), "setup.ps1 must write-probe UNSLOTH_STUDIO_HOME like setup.sh:417" + assert ( + "is not writable" in block + ), "setup.ps1 probe failure must produce a clear writable-error message" + + +def test_env_mode_blocks_when_bin_unsloth_is_a_directory(tmp_path): + """A bare directory at $STUDIO_HOME/bin/unsloth must NOT pass the + sentinel. The previous `-e` test accepted any path type, allowing an + unrelated workspace with sibling content under unsloth_studio plus + a directory at bin/unsloth to be wiped.""" + studio_home = tmp_path / "ws" + venv = studio_home / "unsloth_studio" + (venv / "bin").mkdir(parents = True) + py = venv / "bin" / "python" + py.write_text("#!/bin/sh\nexit 0\n") + py.chmod(0o755) + (venv / "important.txt").write_text("keep me") + (studio_home / "bin" / "unsloth").mkdir(parents = True) + script = _build_install_guard_script(studio_home, "env") + res = subprocess.run( + ["bash", "-c", script], + env = {"PATH": "/usr/bin:/bin"}, + text = True, + capture_output = True, + ) + assert res.returncode != 0, ( + "directory at bin/unsloth must NOT satisfy the Studio sentinel; " + f"stdout={res.stdout!r} stderr={res.stderr!r}" + ) + assert (venv / "important.txt").is_file(), "unrelated workspace data must survive" + + +def test_env_mode_passes_when_bin_unsloth_is_a_symlink(tmp_path): + """A symlink at $STUDIO_HOME/bin/unsloth (real installer artefact) + must still satisfy the sentinel after the leaf-only tightening.""" + studio_home = tmp_path / "ws" + venv = studio_home / "unsloth_studio" + (venv / "bin").mkdir(parents = True) + py = venv / "bin" / "python" + py.write_text("#!/bin/sh\nexit 0\n") + py.chmod(0o755) + (studio_home / "bin").mkdir(parents = True) + target = studio_home / "bin" / "unsloth-real" + target.write_text("#!/bin/sh\nexit 0\n") + target.chmod(0o755) + (studio_home / "bin" / "unsloth").symlink_to(target) + script = _build_install_guard_script(studio_home, "env") + res = subprocess.run( + ["bash", "-c", script], + env = {"PATH": "/usr/bin:/bin"}, + text = True, + capture_output = True, + ) + assert res.returncode == 0, res.stderr + assert "RESULT=ok" in res.stdout + assert not venv.exists() + + +def test_install_ps1_sentinel_uses_pathtype_leaf(): + """The Test-Path checks that gate Remove-Item $VenvDir must use + -PathType Leaf so a directory at the sentinel path cannot satisfy them.""" + src = INSTALL_PS1.read_text() + block_start = src.index("if (Test-Path -LiteralPath $VenvPython)") + block = src[block_start : block_start + 2000] + assert ( + 'share\\studio.conf") -PathType Leaf' in block + ), "install.ps1 share\\studio.conf check must use -PathType Leaf" + assert ( + 'bin\\unsloth.exe") -PathType Leaf' in block + ), "install.ps1 bin\\unsloth.exe check must use -PathType Leaf" + + +def test_setup_ps1_stale_venv_has_env_mode_guard(): + """studio/setup.ps1 stale-venv rebuild branch must mirror install.ps1: + refuse to Remove-Item $VenvDir under custom-root mode unless the root + carries a Studio sentinel (in-VENV marker, share\\studio.conf, or + bin\\unsloth.exe leaf).""" + src = SETUP_PS1.read_text() + idx = src.index("Stale venv detected") + block = src[idx : idx + 1500] + assert ( + "$StudioHomeIsCustom" in block + ), "setup.ps1 stale-venv branch must gate on $StudioHomeIsCustom" + assert ( + 'share\\studio.conf") -PathType Leaf' in block + ), "setup.ps1 stale-venv guard must check share\\studio.conf with -PathType Leaf" + assert ( + 'bin\\unsloth.exe") -PathType Leaf' in block + ), "setup.ps1 stale-venv guard must check bin\\unsloth.exe with -PathType Leaf" + # The guard must fire BEFORE the destructive call. + guard_idx = block.index("$StudioHomeIsCustom") + rm_idx = block.index("Remove-Item -LiteralPath $VenvDir") + assert ( + guard_idx < rm_idx + ), "custom-root guard must precede Remove-Item -LiteralPath $VenvDir" + + +def test_setup_sh_prebuilt_llama_cpp_has_ownership_guard(): + """studio/setup.sh prebuilt llama.cpp path must call + _assert_studio_owned_or_absent before invoking install_llama_prebuilt.py + so an unrelated $UNSLOTH_STUDIO_HOME/llama.cpp is not displaced by + the helper's os.replace().""" + src = SETUP_SH.read_text() + idx = src.index("installing prebuilt llama.cpp...") + block = src[idx : idx + 2000] + assert ( + '_assert_studio_owned_or_absent "$LLAMA_CPP_DIR" "llama.cpp install"' in block + ), "setup.sh must guard the prebuilt llama.cpp path with the ownership marker" + guard_idx = block.index('_assert_studio_owned_or_absent "$LLAMA_CPP_DIR"') + # Anchor on the actual command-array entry, not the why-comment mention. + helper_idx = block.index('python "$SCRIPT_DIR/install_llama_prebuilt.py"') + assert ( + guard_idx < helper_idx + ), "ownership guard must precede the install_llama_prebuilt.py call" + + +def test_setup_ps1_prebuilt_llama_cpp_has_ownership_guard(): + """Mirror check for studio/setup.ps1: prebuilt llama.cpp path must + call Assert-StudioOwnedOrAbsent before invoking install_llama_prebuilt.py.""" + src = SETUP_PS1.read_text() + idx = src.index("installing prebuilt llama.cpp bundle (preferred path)") + block = src[idx : idx + 2000] + assert ( + 'Assert-StudioOwnedOrAbsent -Path $LlamaCppDir -Label "llama.cpp install"' + in block + ), "setup.ps1 must guard the prebuilt llama.cpp path with Assert-StudioOwnedOrAbsent" + guard_idx = block.index("Assert-StudioOwnedOrAbsent -Path $LlamaCppDir") + # Anchor on the actual command-array entry, not the why-comment mention. + helper_idx = block.index('"$PSScriptRoot\\install_llama_prebuilt.py"') + assert ( + guard_idx < helper_idx + ), "Assert-StudioOwnedOrAbsent must precede the install_llama_prebuilt.py call" + + +def test_env_mode_passes_when_venv_marker_present(tmp_path): + """install.sh env-mode guard must accept the in-VENV + .unsloth-studio-owned marker as a primary sentinel so a partial + install (uv venv created, sentinels not yet written) is recoverable + by re-running install.sh.""" + studio_home = tmp_path / "ws" + res = _run_install_guard(studio_home, redirect = "env", create_venv_marker = True) + assert res.returncode == 0, ( + f"in-VENV marker must allow cleanup; " + f"stdout={res.stdout!r} stderr={res.stderr!r}" + ) + assert "RESULT=ok" in res.stdout + assert not (studio_home / "unsloth_studio").exists() + + +def test_env_mode_blocks_when_bin_unsloth_is_symlink_to_directory(tmp_path): + """install.sh env-mode guard must NOT accept a symlink-to-directory at + bin/unsloth as a Studio sentinel. Iter1's standalone -L test let any + symlink (including symlinks to dirs and broken symlinks) bypass the + guard; iter2 dropped that test so only -f (file or symlink-to-file) + counts.""" + studio_home = tmp_path / "ws" + venv = studio_home / "unsloth_studio" + (venv / "bin").mkdir(parents = True) + py = venv / "bin" / "python" + py.write_text("#!/bin/sh\nexit 0\n") + py.chmod(0o755) + (venv / "important.txt").write_text("keep me") + (studio_home / "bin").mkdir(parents = True) + target_dir = studio_home / "bin" / "unsloth-target-dir" + target_dir.mkdir() + (studio_home / "bin" / "unsloth").symlink_to(target_dir) + script = _build_install_guard_script(studio_home, "env") + res = subprocess.run( + ["bash", "-c", script], + env = {"PATH": "/usr/bin:/bin"}, + text = True, + capture_output = True, + ) + assert res.returncode != 0, ( + "symlink-to-directory at bin/unsloth must NOT pass; " + f"stdout={res.stdout!r} stderr={res.stderr!r}" + ) + assert (venv / "important.txt").is_file(), "unrelated workspace data must survive" + + +def test_env_mode_blocks_when_bin_unsloth_is_broken_symlink(tmp_path): + """install.sh guard must reject a broken symlink at bin/unsloth.""" + studio_home = tmp_path / "ws" + venv = studio_home / "unsloth_studio" + (venv / "bin").mkdir(parents = True) + py = venv / "bin" / "python" + py.write_text("#!/bin/sh\nexit 0\n") + py.chmod(0o755) + (venv / "important.txt").write_text("keep me") + (studio_home / "bin").mkdir(parents = True) + (studio_home / "bin" / "unsloth").symlink_to(studio_home / "bin" / "does-not-exist") + script = _build_install_guard_script(studio_home, "env") + res = subprocess.run( + ["bash", "-c", script], + env = {"PATH": "/usr/bin:/bin"}, + text = True, + capture_output = True, + ) + assert res.returncode != 0, ( + "broken symlink at bin/unsloth must NOT pass; " + f"stdout={res.stdout!r} stderr={res.stderr!r}" + ) + assert (venv / "important.txt").is_file() + + +def test_install_sh_writes_venv_marker_after_uv_venv(): + """install.sh must write the .unsloth-studio-owned marker into + $VENV_DIR right after `uv venv` succeeds so the env-mode deletion + guard accepts it on the next install run.""" + src = INSTALL_SH.read_text() + create_idx = src.index('run_install_cmd "create venv" uv venv "$VENV_DIR"') + tail = src[create_idx : create_idx + 600] + assert ( + ".unsloth-studio-owned" in tail + ), "install.sh must write .unsloth-studio-owned after uv venv create" + + +def test_install_ps1_writes_venv_marker_after_uv_venv(): + """install.ps1 must write the .unsloth-studio-owned marker into + $VenvDir after `uv venv` succeeds.""" + src = INSTALL_PS1.read_text() + venv_create = src.index("uv venv $VenvDir --python") + tail = src[venv_create : venv_create + 1500] + assert ( + ".unsloth-studio-owned" in tail + ), "install.ps1 must write .unsloth-studio-owned after uv venv create" + + +def test_install_ps1_guard_accepts_venv_marker(): + """install.ps1 env-mode guard must accept the in-VENV + .unsloth-studio-owned marker as a primary sentinel.""" + src = INSTALL_PS1.read_text() + block_start = src.index("if (Test-Path -LiteralPath $VenvPython)") + block = src[block_start : block_start + 2000] + assert ( + '$VenvDir ".unsloth-studio-owned") -PathType Leaf' in block + ), "install.ps1 guard must check the in-VENV marker with -PathType Leaf" + + +def test_setup_helpers_gate_on_canonical_custom_root(): + """Both _assert_studio_owned_or_absent (setup.sh) and + Assert-StudioOwnedOrAbsent (setup.ps1) must gate on a canonical + custom-vs-legacy comparison so an explicit override that resolves + to the legacy default does not trip the guard for pre-PR T5 + sidecar venvs or llama.cpp dirs.""" + sh_src = SETUP_SH.read_text() + sh_idx = sh_src.index("_assert_studio_owned_or_absent() {") + sh_func = sh_src[sh_idx : sh_idx + 600] + assert ( + '"$_STUDIO_HOME_IS_CUSTOM" = true' in sh_func + ), "setup.sh _assert_studio_owned_or_absent must gate on _STUDIO_HOME_IS_CUSTOM" + assert ( + "_LEGACY_STUDIO_HOME=" in sh_src + and "_studio_home_canon=" in sh_src + and "_STUDIO_HOME_IS_CUSTOM=" in sh_src + ), "setup.sh must compute the canonical custom-root flag" + + ps_src = SETUP_PS1.read_text() + ps_idx = ps_src.index("function Assert-StudioOwnedOrAbsent") + ps_func = ps_src[ps_idx : ps_idx + 800] + assert ( + "$StudioHomeIsCustom -and" in ps_func + ), "setup.ps1 Assert-StudioOwnedOrAbsent must gate on $StudioHomeIsCustom" + assert ( + "$StudioOwnedMarker) -PathType Leaf" in ps_func + ), "setup.ps1 marker check must use -PathType Leaf so a directory cannot satisfy it" + + +def test_setup_ps1_inplace_git_sync_marks_studio_owned(): + """setup.ps1 in-place git-sync branch (when $LlamaCppDir/.git exists) + must call Mark-StudioOwned after a successful sync so a later prebuilt + update path's Assert-StudioOwnedOrAbsent does not exit.""" + src = SETUP_PS1.read_text() + inplace_idx = src.index('Test-Path -LiteralPath (Join-Path $LlamaCppDir ".git")') + # The in-place branch ends just before the temp-dir clone branch. + clone_idx = src.index("Cloning llama.cpp @", inplace_idx) + inplace_block = src[inplace_idx:clone_idx] + assert ( + "Mark-StudioOwned -Path $LlamaCppDir" in inplace_block + ), "in-place git-sync branch must call Mark-StudioOwned on success" + assert ( + "$StudioHomeIsCustom" in inplace_block + ), "in-place Mark-StudioOwned call should be gated on $StudioHomeIsCustom" + + +def test_setup_ps1_inplace_git_sync_asserts_studio_owned_before_mutation(): + """setup.ps1 in-place git-sync branch must call Assert-StudioOwnedOrAbsent + BEFORE any destructive git operation (remote set-url, checkout -B, clean + -fdx). Asymmetric to the prebuilt path and the temp-dir-swap path which + both guard.""" + src = SETUP_PS1.read_text() + inplace_idx = src.index('Test-Path -LiteralPath (Join-Path $LlamaCppDir ".git")') + clone_idx = src.index("Cloning llama.cpp @", inplace_idx) + inplace_block = src[inplace_idx:clone_idx] + assert ( + "Assert-StudioOwnedOrAbsent -Path $LlamaCppDir" in inplace_block + ), "in-place git-sync must Assert-StudioOwnedOrAbsent before mutating $LlamaCppDir" + guard_idx = inplace_block.index("Assert-StudioOwnedOrAbsent -Path $LlamaCppDir") + git_idx = inplace_block.index("git -C $LlamaCppDir remote set-url") + assert ( + guard_idx < git_idx + ), "Assert-StudioOwnedOrAbsent must precede the first git mutation" + + +def _extract_check_health_function() -> str: + src = INSTALL_SH.read_text() + fn_start = src.index("_check_health() {") + fn_end = src.index("\n}\n", fn_start) + 2 + return src[fn_start:fn_end] + + +def _run_check_health(expected_root_id: str, response_json: str) -> int: + fn = _extract_check_health_function() + script = ( + f"_EXPECTED_STUDIO_ROOT_ID={expected_root_id!r}\n" + "_http_get() { printf '%s' \"$1\"; }\n" + + fn.replace( + '_resp=$(_http_get "http://127.0.0.1:$_port/api/health") || return 1', + f"_resp={response_json!r}", + ) + + "\n_check_health 8888\n" + "echo rc=$?\n" + ) + res = subprocess.run( + ["bash", "-c", script], + env = {"PATH": "/usr/bin:/bin"}, + text = True, + capture_output = True, + ) + rc_lines = [l for l in res.stdout.splitlines() if l.startswith("rc=")] + return int(rc_lines[0].split("=")[1]) if rc_lines else res.returncode + + +def test_check_health_accepts_matching_studio_root_id(): + """Hex digest baked at install time matches the backend's + /api/health studio_root_id -- launcher attaches to its own backend.""" + expected_id = "a" * 64 + rc = _run_check_health( + expected_id, + f'{{"status":"healthy","service":"Unsloth UI Backend","studio_root_id":"{expected_id}"}}', + ) + assert rc == 0, f"matching studio_root_id must allow attach (rc={rc})" + + +def test_check_health_rejects_mismatched_studio_root_id(): + """Different install root → different sha256 → reject. Workspace + isolation: launcher A must not open Studio B running on the same port.""" + expected_id = "a" * 64 + other_id = "b" * 64 + rc = _run_check_health( + expected_id, + f'{{"status":"healthy","service":"Unsloth UI Backend","studio_root_id":"{other_id}"}}', + ) + assert rc != 0, "mismatched studio_root_id must reject attach (workspace isolation)" + + +def test_check_health_rejects_missing_studio_root_id_field(): + """A backend that omits studio_root_id (older or non-conforming) must + not be attached to when an expected id is baked into the launcher.""" + expected_id = "a" * 64 + rc = _run_check_health( + expected_id, + '{"status":"healthy","service":"Unsloth UI Backend"}', + ) + assert rc != 0, "missing studio_root_id field must reject attach" + + +def test_check_health_no_baked_id_accepts_any_healthy_backend(): + """If _EXPECTED_STUDIO_ROOT_ID is empty (e.g. install-time hash failed + to compute), the launcher falls back to the legacy contract and accepts + any healthy Unsloth backend.""" + rc = _run_check_health( + "", + '{"status":"healthy","service":"Unsloth UI Backend","studio_root_id":"deadbeef"}', + ) + assert rc == 0, "no baked id → accept any healthy Unsloth backend" + + +def test_check_health_rejects_non_unsloth_service(): + rc = _run_check_health( + "", + '{"status":"healthy","service":"Other UI Backend"}', + ) + assert rc != 0, "non-Unsloth service must be rejected" + + +def test_check_health_handles_arbitrary_id_token(): + """Iter3 used a raw shell match against the JSON-escaped studio_root, + which failed for paths containing `\\` or `"` (FastAPI emits `\\\\` and + `\\\"`). The per-install id token is hex-only by construction, so its + JSON form has no escapes regardless of where the install lives or what + the path contains. This test pins the round-trip on a fully arbitrary + 64-char hex token.""" + expected_id = "f0" + ("ed" * 31) # 64 hex chars, not derived from any path + rc = _run_check_health( + expected_id, + f'{{"status":"healthy","service":"Unsloth UI Backend","studio_root_id":"{expected_id}"}}', + ) + assert ( + rc == 0 + ), "arbitrary 64-hex install id must round-trip cleanly (no JSON escape issue)" + + +def test_install_ps1_test_studio_health_verifies_studio_root_id(): + """install.ps1 Test-StudioHealth must compare studio_root_id against + the install-time-baked $_ExpectedStudioRootId, not the runtime env var.""" + src = INSTALL_PS1.read_text() + fn_start = src.index("function Test-StudioHealth") + fn_end = src.index("\n}\n", fn_start) + 2 + fn = src[fn_start:fn_end] + assert ( + "studio_root_id" in fn + ), "Test-StudioHealth must inspect the studio_root_id field" + assert ( + "$_ExpectedStudioRootId" in fn + ), "Test-StudioHealth must compare against the install-time baked $_ExpectedStudioRootId" + + +def test_install_ps1_bakes_studio_root_id_into_launcher(): + """install.ps1 must persist a per-install opaque id at + $StudioHome\\share\\studio_install_id and bake the value into the + generated launcher as $_ExpectedStudioRootId so the launcher can + verify the backend belongs to THIS install. The id is generated + via a CSPRNG so /api/health does not leak the install path.""" + src = INSTALL_PS1.read_text() + assert ( + "$_studioRootId" in src + ), "install.ps1 must compute $_studioRootId for the launcher" + assert ( + '"share"' in src and "studio_install_id" in src + ), "install.ps1 must persist the id at $StudioHome\\share\\studio_install_id" + assert ( + "RandomNumberGenerator" in src + ), "install.ps1 must seed the id from a CSPRNG (RandomNumberGenerator)" + assert ( + "$_ExpectedStudioRootId" in src + ), "install.ps1 must bake $_ExpectedStudioRootId into the launcher" + + +def test_health_endpoint_exposes_studio_root_id_not_raw_path(): + """studio/backend/main.py /api/health must expose studio_root_id (a + hex digest) and NOT the raw studio_root path. Studio supports + `-H 0.0.0.0`; an unauthenticated /api/health that returns the raw + install path leaks username, home dir, workspace name, etc.""" + main_py = REPO_ROOT / "studio" / "backend" / "main.py" + src = main_py.read_text() + health_idx = src.index('@app.get("/api/health")') + health_block = src[health_idx : health_idx + 1500] + assert ( + '"studio_root_id"' in health_block + ), "/api/health must expose studio_root_id (hex digest)" + assert ( + '"studio_root":' not in health_block + ), "/api/health must NOT expose the raw studio_root path (information disclosure)" + assert ( + "_studio_root_id()" in health_block + ), "/api/health must call the _studio_root_id helper" + + +def test_install_sh_bakes_studio_root_id_into_launcher(): + """install.sh must persist a per-install opaque id at + $STUDIO_HOME/share/studio_install_id and substitute its content into + the launcher heredoc placeholder for ALL modes (env / home / default), + so the launcher's _check_health rejects sibling Studios on the same + port. The id is seeded from /dev/urandom (or python3 secrets fallback) + so /api/health does not leak the install path.""" + src = INSTALL_SH.read_text() + assert ( + "_css_studio_root_id" in src + ), "install.sh must compute _css_studio_root_id for the launcher" + assert ( + '_css_id_file="$_css_id_dir/studio_install_id"' in src + ), "install.sh must persist the id at $STUDIO_HOME/share/studio_install_id" + assert ( + "od -An -N32 -tx1 /dev/urandom" in src + ), "install.sh must seed new ids from /dev/urandom (CSPRNG)" + assert ( + "@@STUDIO_ROOT_ID@@" in src + ), "install.sh must use @@STUDIO_ROOT_ID@@ placeholder in the launcher heredoc" + assert ( + "s|@@STUDIO_ROOT_ID@@|$_css_studio_root_id|g" in src + ), "install.sh must sed-substitute @@STUDIO_ROOT_ID@@ unconditionally (not just env-mode)" + + +def test_tauri_preflight_scrubs_studio_home_env(): + """All three Tauri CLI-spawn sites that lacked the scrub must now + env_remove UNSLOTH_STUDIO_HOME and STUDIO_HOME, mirroring + process.rs / install.rs / desktop_auth.rs / update.rs.""" + preflight = ( + REPO_ROOT / "studio" / "src-tauri" / "src" / "preflight.rs" + ).read_text() + commands = (REPO_ROOT / "studio" / "src-tauri" / "src" / "commands.rs").read_text() + # Both functions in preflight.rs (run_cli_probe + probe_cli_capability) + # must scrub. Count occurrences -- expect 2 in preflight, 1 in commands. + assert ( + preflight.count('cmd.env_remove("UNSLOTH_STUDIO_HOME")') >= 2 + ), "preflight.rs must scrub UNSLOTH_STUDIO_HOME in both run_cli_probe and probe_cli_capability" + assert ( + preflight.count('cmd.env_remove("STUDIO_HOME")') >= 2 + ), "preflight.rs must scrub STUDIO_HOME in both run_cli_probe and probe_cli_capability" + assert ( + 'cmd.env_remove("UNSLOTH_STUDIO_HOME")' in commands + ), "commands.rs check_install_status must scrub UNSLOTH_STUDIO_HOME" + assert ( + 'cmd.env_remove("STUDIO_HOME")' in commands + ), "commands.rs check_install_status must scrub STUDIO_HOME" + + +def test_install_sh_shim_uses_atomic_replace(): + """install.sh shim install must use ln -sfn for atomic replace; the + older `rm -f ...; ln -s ...` left a window where the shim was missing.""" + src = INSTALL_SH.read_text() + shim_idx = src.index('_shim_path="$_LOCAL_BIN/unsloth"') + block = src[shim_idx : shim_idx + 1500] + assert ( + 'ln -sfn "$VENV_DIR/bin/unsloth" "$_shim_path"' in block + ), "install.sh must use ln -sfn for atomic shim replacement" + assert ( + 'rm -f -- "$_shim_path"' not in block + ), "the explicit rm + ln pair must be replaced by atomic ln -sfn" + + +def test_install_sh_create_shortcuts_seeds_id_from_csprng_with_python_fallback( + tmp_path, +): + """_create_shortcuts must seed new ids from /dev/urandom first (no + interpreter spawn cost on the install hot path) and fall back to + `python3 -c 'secrets.token_hex(32)'` only when urandom is unreadable. + Re-running the function with an existing id file must not regenerate + the id (otherwise re-runs would invalidate previously-baked launchers).""" + src = INSTALL_SH.read_text() + fn_start = src.index('_css_data_dir="$DATA_DIR"') + block = src[fn_start : fn_start + 3000] + urandom_idx = block.index("od -An -N32 -tx1 /dev/urandom") + py_fallback_idx = block.index("python3 -c 'import secrets;", urandom_idx) + assert ( + urandom_idx < py_fallback_idx + ), "/dev/urandom must be tried before the python3 secrets fallback" + # The id file is checked for non-empty content before we generate; this is + # what makes re-runs idempotent. + assert ( + 'if [ ! -s "$_css_id_file" ]; then' in block + ), "install.sh must skip id generation when the file already has content" + + # Behavioral check: extract the generation block and run it in isolation + # twice to confirm idempotence. + studio_home = tmp_path / "studio" + (studio_home / "share").mkdir(parents = True) + gen_script = ( + f'STUDIO_HOME="{studio_home}"\n' + '_css_id_dir="$STUDIO_HOME/share"\n' + '_css_id_file="$_css_id_dir/studio_install_id"\n' + # Replicate the generation block (kept narrowly so the test fails loud + # if install.sh changes the surrounding contract). + "gen() {\n" + ' if [ ! -s "$_css_id_file" ]; then\n' + ' _css_new_id=$(od -An -N32 -tx1 /dev/urandom 2>/dev/null | tr -d " \\n")\n' + ' printf "%s" "$_css_new_id" > "$_css_id_file.$$.tmp"\n' + ' mv "$_css_id_file.$$.tmp" "$_css_id_file"\n' + " fi\n" + ' cat "$_css_id_file"\n' + "}\n" + "a=$(gen); b=$(gen)\n" + '[ "$a" = "$b" ] || { echo MISMATCH; exit 1; }\n' + 'echo "ID=$a"\n' + 'echo "LEN=${#a}"\n' + ) + res = subprocess.run(["bash", "-c", gen_script], text = True, capture_output = True) + assert res.returncode == 0, res.stderr + out = dict( + line.split("=", 1) for line in res.stdout.strip().splitlines() if "=" in line + ) + assert ( + out.get("LEN") == "64" + ), f"id must be 64 hex chars, got LEN={out.get('LEN')!r}" + assert all( + c in "0123456789abcdef" for c in out.get("ID", "") + ), f"id must be lowercase hex, got {out.get('ID')!r}" + + +def test_install_sh_create_shortcuts_fails_fast_when_no_entropy(): + """If neither /dev/urandom nor python3 is available, _create_shortcuts + must `return 1` instead of silently baking an empty studio_root_id + (which would disable the launcher's same-install discriminator).""" + src = INSTALL_SH.read_text() + fn_start = src.index('_css_data_dir="$DATA_DIR"') + block = src[fn_start : fn_start + 3000] + assert ( + "[WARN] Cannot create launcher: no entropy source for studio_install_id" + in block + ), "install.sh must warn when neither urandom nor python3 is available" + assert ( + "[WARN] Cannot create launcher: failed to read" in block + ), "install.sh must warn when the id file read produces no content" + assert ( + block.count("return 1") >= 2 + ), "both the no-entropy branch and the empty-read branch must `return 1`" + + +def test_install_sh_bakes_installed_is_env_mode_flag_in_launcher(): + """install.sh must bake the install-time mode (env vs default/home) into + the generated launcher so PORT_FILE / namespaced LOCK_DIR cannot be + flipped on by a sourced custom-root studio.conf in the user's shell.""" + src = INSTALL_SH.read_text() + assert ( + "_INSTALLED_IS_ENV_MODE='@@INSTALLED_IS_ENV_MODE@@'" in src + ), "launcher heredoc must declare _INSTALLED_IS_ENV_MODE='@@INSTALLED_IS_ENV_MODE@@'" + assert ( + "_css_is_env_mode=false" in src + ), "install.sh must default _css_is_env_mode to false" + assert ( + '[ "$_STUDIO_HOME_REDIRECT" = "env" ] && _css_is_env_mode=true' in src + ), "install.sh must set _css_is_env_mode=true only when _STUDIO_HOME_REDIRECT=env" + assert ( + "s|@@INSTALLED_IS_ENV_MODE@@|$_css_is_env_mode|g" in src + ), "install.sh sed pipeline must substitute @@INSTALLED_IS_ENV_MODE@@" + + +def test_install_sh_launcher_gates_port_file_on_baked_flag_not_runtime_env(): + """The launcher's PORT_FILE / namespaced LOCK_DIR must be gated on the + baked $_INSTALLED_IS_ENV_MODE flag, not the runtime $UNSLOTH_STUDIO_HOME. + Sourcing a custom-root studio.conf in shell must not flip a default-mode + launcher into env-mode behavior.""" + src = INSTALL_SH.read_text() + heredoc_start = src.index("cat > \"$_css_launcher\" << 'LAUNCHER_EOF'") + heredoc_end = src.index("LAUNCHER_EOF\n", heredoc_start) + heredoc = src[heredoc_start:heredoc_end] + assert ( + 'if [ "$_INSTALLED_IS_ENV_MODE" = "true" ]; then' in heredoc + ), "launcher must gate PORT_FILE/LOCK_DIR on baked _INSTALLED_IS_ENV_MODE" + port_block_start = heredoc.index('if [ "$_INSTALLED_IS_ENV_MODE" = "true" ]; then') + port_block_end = heredoc.index("\nfi\n", port_block_start) + len("\nfi\n") + port_block = heredoc[port_block_start:port_block_end] + assert 'PORT_FILE="$DATA_DIR/studio.port"' in port_block + assert ( + 'if [ -n "${UNSLOTH_STUDIO_HOME:-}" ]; then\n if command -v cksum' + not in heredoc + ), "launcher must NOT gate PORT_FILE on runtime UNSLOTH_STUDIO_HOME" + + def _run_launcher_gate(installed_flag: str, runtime_env: dict) -> str: + # Reproduce just the LOCK_DIR/PORT_FILE init block in isolation. + script = ( + f"_INSTALLED_IS_ENV_MODE={installed_flag!r}\n" + "DATA_DIR=/tmp/test_data_dir\n" + 'LOCK_DIR="${XDG_RUNTIME_DIR:-/tmp}/unsloth-studio-launcher-$(id -u).lock"\n' + 'PORT_FILE=""\n' + port_block + '\necho "PORT_FILE=$PORT_FILE"\n' + ) + env = {"PATH": "/usr/bin:/bin"} + env.update(runtime_env) + res = subprocess.run( + ["bash", "-c", script], + text = True, + capture_output = True, + env = env, + ) + for line in res.stdout.splitlines(): + if line.startswith("PORT_FILE="): + return line[len("PORT_FILE=") :] + return "" + + # default-mode install should NEVER set PORT_FILE, even if UNSLOTH_STUDIO_HOME leaks in. + assert ( + _run_launcher_gate("false", {"UNSLOTH_STUDIO_HOME": "/tmp/leaked"}) == "" + ), "default-mode launcher must keep PORT_FILE empty even with UNSLOTH_STUDIO_HOME in env" + # env-mode install should set PORT_FILE regardless of runtime env. + assert ( + _run_launcher_gate("true", {}) == "/tmp/test_data_dir/studio.port" + ), "env-mode launcher must set PORT_FILE based on baked DATA_DIR" + + +def test_main_py_studio_root_id_caches_at_module_load(): + """_studio_root_id() is called on every /api/health poll; the id is + stable for the lifetime of the process so it must be read once at + module load and re-used (avoids a hot-path filesystem probe and + protects against transient FS errors during health polling).""" + main_py = (REPO_ROOT / "studio" / "backend" / "main.py").read_text() + assert ( + "_STUDIO_ROOT_ID_CACHE: str = _read_studio_install_id()" in main_py + ), "main.py must populate _STUDIO_ROOT_ID_CACHE from _read_studio_install_id() at module load" + fn_idx = main_py.index("def _studio_root_id() -> str:") + next_def_idx = main_py.index("\ndef ", fn_idx + 1) + fn_block = main_py[fn_idx:next_def_idx] + assert ( + "return _STUDIO_ROOT_ID_CACHE" in fn_block + ), "_studio_root_id() body must return the cached value" + assert ( + "read_text(" not in fn_block and "hashlib" not in fn_block + ), "_studio_root_id() must NOT do filesystem or hash work on every call" + + +def test_main_py_read_studio_install_id_validates_hex_and_handles_missing( + tmp_path, monkeypatch +): + """_read_studio_install_id reads $STUDIO_HOME/share/studio_install_id and + returns "" when the file is absent, empty, contains non-hex content, or + is the wrong length. "" triggers the launcher's "no baked id, accept any + healthy backend" fallback path (see test_check_health_no_baked_id_*). + Behavioral check: spin up a stub _STUDIO_ROOT_RESOLVED and exercise + _read_studio_install_id directly without importing main.py (which + pulls in heavy deps). Test the rejection rules verbatim.""" + import re + + pattern = re.compile(r"^[0-9a-f]{64}$") + + def _read(root: Path) -> str: + # Mirror the implementation; this test pins the exact contract so a + # future refactor can't silently widen what's accepted. + try: + token = (root / "share" / "studio_install_id").read_text().strip() + except (OSError, ValueError): + return "" + return token if pattern.fullmatch(token) else "" + + root = tmp_path / "studio" + (root / "share").mkdir(parents = True) + + # Missing file -> empty + assert _read(root) == "" + + id_file = root / "share" / "studio_install_id" + # Empty file -> empty + id_file.write_text("") + assert _read(root) == "" + # Non-hex content -> empty + id_file.write_text( + "not-a-hex-id-just-text-padded-to-64-chars-zzzzzzzzzzzzzzzzzzzzzz" + ) + assert _read(root) == "" + # Uppercase hex -> empty (must be lowercase) + id_file.write_text("F" * 64) + assert _read(root) == "" + # Wrong length -> empty (32 chars, not 64) + id_file.write_text("a" * 32) + assert _read(root) == "" + # Valid 64-char lowercase hex with surrounding whitespace -> stripped+accepted + valid = "0123456789abcdef" * 4 + id_file.write_text(f"\n {valid} \n") + assert _read(root) == valid + + +def test_llama_cpp_search_roots_handles_studio_root_oserror(): + """_find_llama_server_binary calls studio_root() which can raise + OSError or ValueError from Path.expanduser().resolve() (broken symlink, + null byte). The except clause must mirror sibling _kill_orphaned_servers + (which catches the same trio) so inference startup does not crash.""" + llama_cpp = ( + REPO_ROOT / "studio" / "backend" / "core" / "inference" / "llama_cpp.py" + ).read_text() + find_block_start = llama_cpp.index("_find_llama_server_binary") + find_block = llama_cpp[find_block_start : find_block_start + 4000] + assert ( + "except (ImportError, OSError, ValueError):" in find_block + ), "_find_llama_server_binary must catch (ImportError, OSError, ValueError) from studio_root()" + kill_def_idx = llama_cpp.index("def _kill_orphaned_servers") + kill_block = llama_cpp[kill_def_idx : kill_def_idx + 4000] + assert ( + "except (ImportError, OSError, ValueError):" in kill_block + ), "sibling _kill_orphaned_servers must keep its (ImportError, OSError, ValueError) handler" + + +def test_install_sh_install_id_survives_symlinked_studio_home(tmp_path): + """End-to-end behavioral check: when $STUDIO_HOME is reached via a + symlinked parent (e.g. symlinked $HOME on Linux, junctioned %USERPROFILE% + on Windows), install.sh and the backend agree on the install id BY + CONSTRUCTION because the id is read from a file whose location resolves + the same way for both. The previous sha256(canonical_path) scheme + required `cd -P/pwd -P` and Path.resolve() to produce identical strings, + which broke under symlinks/junctions and required cycles 17-27 of the + PR's review history to fully canonicalize. This is the regression test + pinning that the new design has no such drift.""" + real = tmp_path / "realhome" + real.mkdir() + link = tmp_path / "linkhome" + link.symlink_to(real) + studio_home = real / ".unsloth" / "studio" + (studio_home / "share").mkdir(parents = True) + # Write a stub install id at the canonical location. + valid_id = "ab12" * 16 + (studio_home / "share" / "studio_install_id").write_text(valid_id) + # Read it back via both the canonical and the symlinked path; both must + # see the SAME content (which is what makes install.sh's cat and the + # backend's read_text agree without any canonicalization dance). + raw_via_link = link / ".unsloth" / "studio" / "share" / "studio_install_id" + raw_direct = studio_home / "share" / "studio_install_id" + assert raw_via_link.read_text() == valid_id + assert raw_direct.read_text() == valid_id + # And install.sh's `cat` would see the same. + import subprocess as _sp + + res = _sp.run(["cat", str(raw_via_link)], capture_output = True, text = True) + assert res.returncode == 0 + assert res.stdout == valid_id + + +def test_install_sh_substitutes_root_id_before_data_dir(): + """The two-stage sed substitution must bake @@STUDIO_ROOT_ID@@ / + @@INSTALLED_IS_ENV_MODE@@ first (non-user-controlled), then @@DATA_DIR@@ + (user-controlled). A custom $DATA_DIR containing the literal text + @@STUDIO_ROOT_ID@@ must not be mutated by the global root-id sed pass.""" + src = INSTALL_SH.read_text() + root_id_idx = src.index("s|@@STUDIO_ROOT_ID@@|$_css_studio_root_id|g") + env_mode_idx = src.index("s|@@INSTALLED_IS_ENV_MODE@@|$_css_is_env_mode|g") + data_dir_idx = src.index("s|@@DATA_DIR@@|$_sed_safe|g") + assert root_id_idx < data_dir_idx, ( + "@@STUDIO_ROOT_ID@@ substitution must happen BEFORE @@DATA_DIR@@ " + "(non-user-controlled placeholders first)" + ) + assert ( + env_mode_idx < data_dir_idx + ), "@@INSTALLED_IS_ENV_MODE@@ substitution must happen BEFORE @@DATA_DIR@@" + + +def test_install_sh_root_id_pass_does_not_mutate_user_data_dir(tmp_path): + """Behavioral subprocess test: a $DATA_DIR containing the literal text + `@@STUDIO_ROOT_ID@@` must not be mutated when the placeholder pass runs + first; only the actual placeholder occurrences in the launcher template + are replaced.""" + src = INSTALL_SH.read_text() + heredoc_start = src.index("cat > \"$_css_launcher\" << 'LAUNCHER_EOF'") + heredoc_body_start = src.index("\n", heredoc_start) + 1 + heredoc_body_end = src.index("LAUNCHER_EOF\n", heredoc_start) + template = src[heredoc_body_start:heredoc_body_end] + launcher_path = tmp_path / "launch.sh" + launcher_path.write_text(template) + # Run the iter6 sed order: root-id first, then data-dir. + weird_data_dir = "/tmp/with-@@STUDIO_ROOT_ID@@/share" + root_id = "deadbeef" * 8 + is_env = "true" + script = f""" +sed -e "s|@@STUDIO_ROOT_ID@@|{root_id}|g" \\ + -e "s|@@INSTALLED_IS_ENV_MODE@@|{is_env}|g" \\ + "{launcher_path}" > "{launcher_path}.tmp" && mv "{launcher_path}.tmp" "{launcher_path}" +_sq_escaped=$(printf '%s' "{weird_data_dir}" | sed "s/'/'\\\\\\\\''/g") +_sed_safe=$(printf '%s' "$_sq_escaped" | sed 's/[\\\\&|]/\\\\&/g') +sed "s|@@DATA_DIR@@|$_sed_safe|g" "{launcher_path}" > "{launcher_path}.tmp" \\ + && mv "{launcher_path}.tmp" "{launcher_path}" +""" + subprocess.run(["bash", "-c", script], check = True) + final = launcher_path.read_text() + assert ( + f"DATA_DIR='{weird_data_dir}'" in final + ), f"DATA_DIR must be preserved verbatim (no @@STUDIO_ROOT_ID@@ mutation); got: {final[:500]}" + assert ( + f"_EXPECTED_STUDIO_ROOT_ID='{root_id}'" in final + ), "STUDIO_ROOT_ID placeholder must still be substituted in the launcher heredoc" + + +def test_install_ps1_install_id_file_layout_matches_backend_read_path(): + """install.ps1 must write the id at $StudioHome\\share\\studio_install_id + so the backend (studio/backend/main.py:_read_studio_install_id) can find + it via _STUDIO_ROOT_RESOLVED / "share" / "studio_install_id" without + mode-specific path knowledge. Persistence-across-runs is enforced by the + pre-write Test-Path check.""" + src = INSTALL_PS1.read_text() + id_idx = src.index('$_studioIdDir = Join-Path $StudioHome "share"') + context = src[id_idx : id_idx + 1500] + assert ( + '$_studioIdFile = Join-Path $_studioIdDir "studio_install_id"' in context + ), "install.ps1 must persist the id at $StudioHome\\share\\studio_install_id" + assert ( + "Test-Path -LiteralPath $_studioIdFile" in context + ), "install.ps1 must skip id generation when the file already has content (re-run idempotence)" + assert ( + "RandomNumberGenerator" in context and "GetBytes($_idBytes)" in context + ), "install.ps1 must seed new ids from a CSPRNG (RandomNumberGenerator)" + assert ( + "Move-Item -LiteralPath $_idTmp" in context + ), "install.ps1 must atomic-rename the temp file into place to avoid half-written ids" diff --git a/tests/test_studio_root_resilience.py b/tests/test_studio_root_resilience.py new file mode 100644 index 0000000000..1ce0430dc4 --- /dev/null +++ b/tests/test_studio_root_resilience.py @@ -0,0 +1,154 @@ +"""Resilience checks for Studio install-root inference under hostile +filesystem conditions: +- _infer_studio_home_from_venv must NOT propagate PermissionError / + OSError out through studio_root() (it would crash module import in + run.py / main.py / transformers_version.py / model_config.py). +- _kill_orphaned_servers must catch (ImportError, OSError, ValueError) + on the studio_root() probe so a transient resolve / sentinel failure + cannot crash server startup. +- _find_llama_server_binary must keep the custom-root in search_roots + when the inner resolve() comparison itself fails.""" + +from __future__ import annotations + +import importlib.util +import re +import sys +import textwrap +from pathlib import Path +from unittest import mock + +import pytest + + +REPO_ROOT = Path(__file__).resolve().parents[1] +STORAGE_ROOTS = ( + REPO_ROOT / "studio" / "backend" / "utils" / "paths" / "storage_roots.py" +) +LLAMA_CPP = REPO_ROOT / "studio" / "backend" / "core" / "inference" / "llama_cpp.py" + + +def _load(name: str, path: Path): + spec = importlib.util.spec_from_file_location(name, path) + assert spec is not None and spec.loader is not None + mod = importlib.util.module_from_spec(spec) + sys.modules[name] = mod + spec.loader.exec_module(mod) + return mod + + +def test_infer_studio_home_swallows_permission_error(tmp_path, monkeypatch): + candidate = tmp_path / "fake_root" + venv = candidate / "unsloth_studio" + venv.mkdir(parents = True) + monkeypatch.setattr(sys, "prefix", str(venv)) + sys.modules.pop("sr_perm", None) + mod = _load("sr_perm", STORAGE_ROOTS) + with mock.patch.object(Path, "is_file", side_effect = PermissionError("denied")): + # Must NOT raise. + assert mod._infer_studio_home_from_venv() is None + + +def test_studio_root_does_not_crash_on_permission_error(tmp_path, monkeypatch): + """studio_root() must remain callable even when the venv inference + encounters a restricted filesystem; it should fall through to the + legacy default.""" + candidate = tmp_path / "fake_root" + venv = candidate / "unsloth_studio" + venv.mkdir(parents = True) + monkeypatch.setattr(sys, "prefix", str(venv)) + monkeypatch.delenv("UNSLOTH_STUDIO_HOME", raising = False) + monkeypatch.delenv("STUDIO_HOME", raising = False) + sys.modules.pop("sr_studio_perm", None) + mod = _load("sr_studio_perm", STORAGE_ROOTS) + with mock.patch.object(Path, "is_file", side_effect = OSError("ebusy")): + result = mod.studio_root() + assert result == Path.home() / ".unsloth" / "studio" + + +def test_kill_orphan_catches_oserror_from_studio_root(): + """_kill_orphaned_servers must catch (ImportError, OSError, ValueError) + on the studio_root() probe specifically; the sister function + _find_llama_server_binary uses the same broader catch on its own probe.""" + src = LLAMA_CPP.read_text() + fn_start = src.index("def _kill_orphaned_servers") + fn_body = src[fn_start : fn_start + 4000] + # The studio_root() probe in this fn is the one that imports as `_sr` + # and assigns `_resolved_sr = _sr()`. Find the except that closes it. + probe_idx = fn_body.index("storage_roots import studio_root as _sr") + # The matching except is the next `except ...:` after the inner + # OSError/ValueError block that wraps resolve(). + after = fn_body[probe_idx:] + # Skip over the inner `except (OSError, ValueError):` that wraps resolve(). + inner_idx = after.index("except (OSError, ValueError):") + after_inner = after[inner_idx + len("except (OSError, ValueError):") :] + outer_match = re.search(r"except\s*\(?[^)]*?\)?:", after_inner) + assert outer_match, "outer except for studio_root probe missing" + clause = outer_match.group(0) + assert ( + "OSError" in clause and "ValueError" in clause + ), f"_kill_orphaned_servers studio_root probe catch too narrow: {clause!r}" + + +def _exec_search_roots_block( + home: Path, studio_root_value: Path, resolve_raises: bool +) -> list[Path]: + """Extract _find_llama_server_binary's env-mode search_roots block + and execute it with controlled inputs.""" + src = LLAMA_CPP.read_text() + block_start = src.index('legacy_llama = Path.home() / ".unsloth" / "llama.cpp"') + block_end = src.index("_seen_roots: set[str]", block_start) + raw = src[block_start:block_end] + indent = " " * 8 + block = textwrap.dedent(indent + raw) + fake_module = type(sys)("fake_storage_roots") + fake_module.studio_root = lambda: studio_root_value + sys.modules["utils.paths.storage_roots"] = fake_module + try: + original_resolve = Path.resolve + + def _resolve(self, *a, **k): + if resolve_raises: + raise OSError("ebusy") + return original_resolve(self, *a, **k) + + with ( + mock.patch.object(Path, "home", classmethod(lambda cls: home)), + mock.patch.object(Path, "resolve", _resolve), + ): + ns: dict = {"Path": Path} + exec(block, ns) # noqa: S102 + return ns["search_roots"] + finally: + sys.modules.pop("utils.paths.storage_roots", None) + + +def test_search_roots_keeps_custom_when_resolve_fails(tmp_path): + home = tmp_path / "home" + home.mkdir() + custom = tmp_path / "custom_studio" + custom.mkdir() + roots = _exec_search_roots_block( + home = home, studio_root_value = custom, resolve_raises = True + ) + # On resolve() failure, the inner except falls back to direct equality; + # custom != legacy_studio so the custom root must remain in search_roots. + assert ( + custom / "llama.cpp" in roots + ), f"custom root dropped on resolve() failure: {roots}" + # custom-mode discovery excludes the legacy tree to match _kill_orphaned_servers. + assert ( + (home / ".unsloth" / "llama.cpp") not in roots + ), f"legacy llama path must not appear in custom-mode search_roots: {roots}" + + +def test_search_roots_default_mode_uses_legacy_only(tmp_path): + home = tmp_path / "home" + home.mkdir() + legacy = home / ".unsloth" / "studio" + legacy.mkdir(parents = True) + roots = _exec_search_roots_block( + home = home, studio_root_value = legacy, resolve_raises = False + ) + # Default mode: only legacy_llama. + assert roots == [home / ".unsloth" / "llama.cpp"] diff --git a/unsloth_cli/commands/studio.py b/unsloth_cli/commands/studio.py index 140940209f..76aac3dc15 100644 --- a/unsloth_cli/commands/studio.py +++ b/unsloth_cli/commands/studio.py @@ -20,7 +20,73 @@ import typer studio_app = typer.Typer(help = "Unsloth Studio commands.") -STUDIO_HOME = Path.home() / ".unsloth" / "studio" + +# Resolve install root: UNSLOTH_STUDIO_HOME, then STUDIO_HOME alias, then +# sys.prefix inference (so a direct call to /bin/unsloth resolves after +# the installer's env var has expired), then legacy ~/.unsloth/studio. +# UNSLOTH_STUDIO_HOME wins when both env vars are set. +def _looks_like_installer_managed_studio_home(candidate: Path) -> bool: + """Sentinel check (studio.conf or bin shim) so a dev venv named + unsloth_studio is not misidentified as a custom Studio root. + """ + shim_name = "unsloth.exe" if platform.system() == "Windows" else "unsloth" + return (candidate / "share" / "studio.conf").is_file() or ( + candidate / "bin" / shim_name + ).is_file() + + +def _resolve_studio_home() -> tuple[Path, bool]: + override = (os.environ.get("UNSLOTH_STUDIO_HOME") or "").strip() + if not override: + override = (os.environ.get("STUDIO_HOME") or "").strip() + if override: + try: + return Path(override).expanduser().resolve(), True + except (OSError, ValueError): + return Path(override).expanduser(), True + try: + prefix = Path(sys.prefix).resolve() + if prefix.name == "unsloth_studio": + inferred = prefix.parent + legacy = (Path.home() / ".unsloth" / "studio").resolve() + if inferred != legacy and _looks_like_installer_managed_studio_home( + inferred + ): + return inferred, True + except (OSError, ValueError): + pass + return Path.home() / ".unsloth" / "studio", False + + +STUDIO_HOME, _STUDIO_HOME_IS_CUSTOM = _resolve_studio_home() + + +def _ensure_studio_env_exported() -> None: + """Re-export UNSLOTH_STUDIO_HOME / UNSLOTH_LLAMA_CPP_PATH only for real + custom roots so subprocesses inherit the right install. Called from each + studio subcommand entry rather than at import time, to avoid leaking env + state into unrelated importers (tests, --help, CLI introspection). + """ + if not _STUDIO_HOME_IS_CUSTOM: + return + # Truthy-check (not setdefault) so a blank UNSLOTH_STUDIO_HOME= does not + # suppress the inferred custom root. + if not os.environ.get("UNSLOTH_STUDIO_HOME"): + os.environ["UNSLOTH_STUDIO_HOME"] = str(STUDIO_HOME) + # When override == legacy default, llama.cpp stays at ~/.unsloth/llama.cpp. + try: + _legacy_studio = (Path.home() / ".unsloth" / "studio").resolve() + _is_legacy = STUDIO_HOME.resolve() == _legacy_studio + except (OSError, ValueError): + _is_legacy = STUDIO_HOME == (Path.home() / ".unsloth" / "studio") + if _is_legacy: + _llama_dir = Path.home() / ".unsloth" / "llama.cpp" + else: + _llama_dir = STUDIO_HOME / "llama.cpp" + if not os.environ.get("UNSLOTH_LLAMA_CPP_PATH"): + os.environ["UNSLOTH_LLAMA_CPP_PATH"] = str(_llama_dir) + + BOOTSTRAP_PASSWORD_FILE = ".bootstrap_password" DESKTOP_SECRET_FILE = ".desktop_secret" DEFAULT_ADMIN_USERNAME = "unsloth" @@ -427,6 +493,8 @@ def studio_default( ), ): """Launch the Unsloth Studio server.""" + # Runs before any subcommand; covers run/setup/update/etc in one place. + _ensure_studio_env_exported() if ctx.invoked_subcommand is not None: return From 7de1f4c51350f9dc88ad837059ddf2a6f20e7f61 Mon Sep 17 00:00:00 2001 From: Daniel Han Date: Tue, 5 May 2026 23:22:22 -0700 Subject: [PATCH 003/159] Route CPU-only Linux x86_64 to ggml-org/llama.cpp prebuilts (#5302) * Route CPU-only Linux x86_64 to ggml-org/llama.cpp prebuilts setup.sh hard-coded _HELPER_RELEASE_REPO=unslothai/llama.cpp for every non-Darwin host. unslothai/llama.cpp only publishes Linux CUDA bundles (app-*-linux-x64-cuda*.tar.gz), so a CPU-only Linux host walked ~30 releases looking for a non-existent app-*-linux-x64-cpu asset, exited the prebuilt planner with "no compatible Linux prebuilt asset was found", and fell through to a source build. Free CI runners (ubuntu-latest with no GPU) hit this on every install, and anyone running Studio on a Linux laptop without an NVIDIA GPU paid the ~3 minute cmake+make cost on first install. ggml-org publishes llama--bin-ubuntu-x64.tar.gz on every release and install_llama_prebuilt.py already knows how to fetch it: when called with --published-repo ggml-org/llama.cpp, the Linux x86_64 + not has_usable_nvidia branch in direct_upstream_release_plan picks up that asset directly. The fix is purely on the routing side. Tighten the gate so a Linux host routes to ggml-org only when it is x86_64 and has no GPU detection tool installed (nvidia-smi, rocminfo, amd-smi, hipconfig, hipinfo). Everything else stays on the current path: - macOS: already on ggml-org, unchanged - Windows: already on ggml-org via setup.ps1, unchanged - Linux CUDA: nvidia-smi present -> unslothai/llama.cpp, unchanged - Linux ROCm: rocminfo / amd-smi / hipconfig / hipinfo present -> unslothai/llama.cpp -> source build with HIP, unchanged - Linux Intel / Vulkan / SYCL: no NVIDIA / AMD tools, hits the new ggml-org route, gets upstream CPU asset (same as today's source-build CPU output, ~3 min faster) - Linux arm64 / s390x: not x86_64 -> unslothai/llama.cpp -> source build, unchanged * Tighten routing comment in studio/setup.sh --- studio/setup.sh | 19 +++++++++++++++++++ 1 file changed, 19 insertions(+) diff --git a/studio/setup.sh b/studio/setup.sh index ff93d2d41d..c5beb7ebd3 100755 --- a/studio/setup.sh +++ b/studio/setup.sh @@ -651,11 +651,30 @@ _LLAMA_CPP_DEGRADED=false _LLAMA_FORCE_COMPILE="${UNSLOTH_LLAMA_FORCE_COMPILE:-0}" _REQUESTED_LLAMA_TAG="${UNSLOTH_LLAMA_TAG:-${_DEFAULT_LLAMA_TAG}}" _HOST_SYSTEM="$(uname -s 2>/dev/null || true)" +_HOST_MACHINE="$(uname -m 2>/dev/null || true)" + +# Pick the release repo install_llama_prebuilt.py plans against. +# unslothai/llama.cpp ships only Linux CUDA bundles, so CPU-only Linux +# x86_64 routes to ggml-org for bin-ubuntu-x64.tar.gz. Anything with a +# GPU tool installed stays on unslothai (CUDA bundle / ROCm source build). +_LINUX_HAS_GPU=false +for _GPU_TOOL in nvidia-smi rocminfo amd-smi hipconfig hipinfo; do + if command -v "$_GPU_TOOL" >/dev/null 2>&1; then + _LINUX_HAS_GPU=true + break + fi +done + if [ "$_HOST_SYSTEM" = "Darwin" ]; then _HELPER_RELEASE_REPO="ggml-org/llama.cpp" +elif [ "$_HOST_SYSTEM" = "Linux" ] \ + && [ "$_HOST_MACHINE" = "x86_64" ] \ + && [ "$_LINUX_HAS_GPU" = false ]; then + _HELPER_RELEASE_REPO="ggml-org/llama.cpp" else _HELPER_RELEASE_REPO="unslothai/llama.cpp" fi +unset _GPU_TOOL _LLAMA_PR="${UNSLOTH_LLAMA_PR:-}" _SKIP_PREBUILT_INSTALL=false _LLAMA_PR_FORCE="${UNSLOTH_LLAMA_PR_FORCE:-${_DEFAULT_LLAMA_PR_FORCE}}" From d65149795b4d4b83e69201bb6d372f9272b34731 Mon Sep 17 00:00:00 2001 From: Manan Shah <52329525+Manan17@users.noreply.github.com> Date: Tue, 5 May 2026 23:54:58 -0700 Subject: [PATCH 004/159] feat(studio): MLX training tab on Apple Silicon (LoRA / full FT, VLM, export) (#5265) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit * Add Apple Silicon MLX routing Rewrite __init__.py: detect MLX on macOS arm64 before any torch imports Extract original GPU init to _gpu_init.py (unchanged) MLX path imports FastMLXModel from unsloth_zoo, skips all GPU code GPU path unchanged: from ._gpu_init import * * Add Apple Silicon MLX routing - Rewrite __init__.py: detect MLX on macOS arm64 before any torch imports - Extract original GPU init to _gpu_init.py (unchanged) - MLX path imports FastMLXModel from unsloth_zoo, skips all GPU code - GPU path unchanged: from ._gpu_init import * * mlx with studio * mlx with studio * updating temporary install.sh * updating temporary install.sh * adding t_v5 path * adding t_v5 path * fixing vision training * fixing vision training * adding chat * adding chat * minor * minor * Adding export and fixing training issues, inference with lora adaptors * Adding export and fixing training issues, inference with lora adaptors * fix: MLX worker pass load_in_4bit, override is_vlm based on dataset, streaming for VLM * fix: MLX worker pass load_in_4bit, override is_vlm based on dataset, streaming for VLM * Merge mlx-apple-silicon into main * update install.sh to point to main branch * update install.sh to point to main branch * fix: export returns 3 values (success, message, output_path) matching upstream worker * fix: export returns 3 values (success, message, output_path) matching upstream worker * fix(mlx): show training-process peak memory in Studio UI, not system-wide Studio UI was showing ~95 GB during MLX training because get_gpu_utilization read "In use system memory" from IORegistry's AGXAccelerator — system-wide GPU memory across all processes (training + backend + browser + Display). Now the trainer's mx.get_peak_memory value is forwarded through the progress event and surfaced via /api/train/hardware while training is active. Falls back to the system-wide reading when training is not running. * fix(mlx): show training-process peak memory in Studio UI, not system-wide Studio UI was showing ~95 GB during MLX training because get_gpu_utilization read "In use system memory" from IORegistry's AGXAccelerator — system-wide GPU memory across all processes (training + backend + browser + Display). Now the trainer's mx.get_peak_memory() value is forwarded through the progress event and surfaced via /api/train/hardware while training is active. Falls back to the system-wide reading when training is not running. * fix(mlx): make is_bfloat16_supported detect M1/M2 (no native bf16) M1 and M2 chips emulate bf16 in software on the GPU, causing 40-70% slower prefill compared to native fp16. M3+ have native bf16 (macOS Sonoma+ MPSGraph). Replaces the always-True stub with chip-aware detection via mx.device_info. * fix(mlx): make is_bfloat16_supported() detect M1/M2 (no native bf16) M1 and M2 chips emulate bf16 in software on the GPU, causing 40-70% slower prefill compared to native fp16. M3+ have native bf16 (macOS Sonoma+ MPSGraph). Replaces the always-True stub with chip-aware detection via mx.device_info(). * feat(mlx): wire training_type="Full Finetuning" through MLX worker Compute use_lora from the UI's training_type before loading the model, pass full_finetuning=not use_lora to FastMLXModel.from_pretrained, and let the existing 'if use_lora' branch skip get_peft_model. Matches the GPU worker's flow. * feat(mlx): wire training_type="Full Finetuning" through MLX worker Compute use_lora from the UI's training_type before loading the model, pass full_finetuning=not use_lora to FastMLXModel.from_pretrained, and let the existing 'if use_lora' branch skip get_peft_model. Matches the GPU worker's flow. * fix(mlx): pass save_method='merged_16bit' from Studio's export page Previously the MLX path called save_pretrained_merged with no save_method, which fell through to a no-op that didn't actually fuse LoRA into the base. Now Studio's "Merged Model" export properly fuses LoRA + dequantizes any 4-bit base to bf16, matching the GPU behavior for the same UI option. * fix(mlx): pass save_method='merged_16bit' from Studio's export page Previously the MLX path called save_pretrained_merged() with no save_method, which fell through to a no-op that didn't actually fuse LoRA into the base. Now Studio's "Merged Model" export properly fuses LoRA + dequantizes any 4-bit base to bf16, matching the GPU behavior for the same UI option. * fix(studio): pass private to MLX push, return 3-tuples consistently MLX push_to_hub branch now forwards private=private (matches GPU) Existing 2-tuple early-returns ('repo_id+token required', 'PEFT model needed') were tripping the route's 3-tuple unpack. Added a None output_path so the unpack always succeeds. * fix(studio): pass private to MLX push, return 3-tuples consistently - MLX push_to_hub branch now forwards private=private (matches GPU) - Existing 2-tuple early-returns ('repo_id+token required', 'PEFT model needed') were tripping the route's 3-tuple unpack. Added a None output_path so the unpack always succeeds. * studio wirings * studio wirings * Merge pull request #5 from Manan17/feat/quant_config studio wirings * fix(mlx): wire train_on_completions for VLM via per-template lookup Mirror the GPU worker: stop excluding VLMs and stop hardcoding template detection. Look up the model in MODEL_TO_TEMPLATE_MAPPER and fetch the per-template instruction/response markers from TEMPLATE_TO_RESPONSES_MAPPER. The frontend already force-disables train_on_completions for vision+image and audio cases, so backend just trusts the flag. * fix(mlx): wire train_on_completions for VLM via per-template lookup Mirror the GPU worker: stop excluding VLMs and stop hardcoding template detection. Look up the model in MODEL_TO_TEMPLATE_MAPPER and fetch the per-template instruction/response markers from TEMPLATE_TO_RESPONSES_MAPPER. The frontend already force-disables train_on_completions for vision+image and audio cases, so backend just trusts the flag. * wire in lora rslora, init lora weights, random_state * wire in lora rslora, init lora weights, random_state * loftq studio error message fix * loftq studio error message fix * handle unknown optim and lr scheduler * handle unknown optim and lr scheduler * Merge pull request #6 from Manan17/update/peftkwargs Update/peftkwargs * feat(mlx): pass finetune_language/attention/mlp/vision flags to FastMLXModel Studio's four UI checkboxes now actually flow through to MLX get_peft_model (which was just updated in unsloth-zoo to honor them). Also drops the incorrect train_projector wiring that tied projector LoRA to the attn/mlp flags — those are language-side toggles, not projector toggles. Co-Authored-By: Manan17 * feat(mlx): pass finetune_language/attention/mlp/vision flags to FastMLXModel Studio's four UI checkboxes now actually flow through to MLX get_peft_model (which was just updated in unsloth-zoo to honor them). Also drops the incorrect train_projector wiring that tied projector LoRA to the attn/mlp flags — those are language-side toggles, not projector toggles. Co-Authored-By: Manan17 * feat(mlx,ux): auto-imply finetune_language_layers when user picks attn/mlp UI guardrail. The four checkboxes (vision/language/attention/MLP) carry "scope × module-type" semantics that aren't obvious — picking just "Attention modules" + "MLP modules" without "Language layers" naturally reads as "fine-tune attn/mlp" but our backend reads it as "fine-tune attn/mlp modules in *no* tower" → empty target_modules → zero trainable params → crash inside value_and_grad. If user selected attn or mlp module types but no layer scope, default to language scope. Power users can still explicitly choose language=False, vision=True if they want vision-only fine-tuning of attn/mlp. Co-Authored-By: Manan17 * feat(mlx,ux): auto-imply finetune_language_layers when user picks attn/mlp UI guardrail. The four checkboxes (vision/language/attention/MLP) carry "scope × module-type" semantics that aren't obvious — picking just "Attention modules" + "MLP modules" without "Language layers" naturally reads as "fine-tune attn/mlp" but our backend reads it as "fine-tune attn/mlp modules in *no* tower" → empty target_modules → zero trainable params → crash inside value_and_grad. If user selected attn or mlp module types but no layer scope, default to language scope. Power users can still explicitly choose language=False, vision=True if they want vision-only fine-tuning of attn/mlp. Co-Authored-By: Manan17 * fix(mlx): wire top_k, repetition_penalty, and VLM top_p through to mlx-lm/mlx-vlm Inference UI sliders for top_k and repetition_penalty had no effect on MLX, and VLM top_p was also silently dropped. Plus a latent pre-existing bug: mlx_vlm.generate_step expects temperature= (long form), but we were passing temp= which silently fell into **kwargs — every VLM chat was effectively greedy regardless of the temperature slider. Text path (_generate_text): make_sampler now receives top_k in addition to temp/top_p make_logits_processors built and forwarded when repetition_penalty is non-trivial (skip when 0.0/1.0 to avoid pointless overhead) VLM path (_generate_vlm): Pass top_p, top_k, repetition_penalty as kwargs (mlx_vlm.stream_generate forwards them to generate_step's sampler/logits_processor builders) Rename temp= → temperature= so it's actually consumed Verified end-to-end with a smoke test on Qwen2.5-0.5B-Instruct (text) and Qwen2.5-VL-3B-Instruct (VLM): each of {greedy, top_p=0.5, top_k=10, rep_pen=1.5} now produces a distinct output, proving the parameters reach the sampler. Co-Authored-By: Manan17 * fix(mlx): wire top_k, repetition_penalty, and VLM top_p through to mlx-lm/mlx-vlm Inference UI sliders for top_k and repetition_penalty had no effect on MLX, and VLM top_p was also silently dropped. Plus a latent pre-existing bug: mlx_vlm.generate_step expects temperature= (long form), but we were passing temp= which silently fell into **kwargs — every VLM chat was effectively greedy regardless of the temperature slider. Text path (_generate_text): - make_sampler now receives top_k in addition to temp/top_p - make_logits_processors built and forwarded when repetition_penalty is non-trivial (skip when 0.0/1.0 to avoid pointless overhead) VLM path (_generate_vlm): - Pass top_p, top_k, repetition_penalty as kwargs (mlx_vlm.stream_generate forwards them to generate_step's sampler/logits_processor builders) - Rename temp= → temperature= so it's actually consumed Verified end-to-end with a smoke test on Qwen2.5-0.5B-Instruct (text) and Qwen2.5-VL-3B-Instruct (VLM): each of {greedy, top_p=0.5, top_k=10, rep_pen=1.5} now produces a distinct output, proving the parameters reach the sampler. Co-Authored-By: Manan17 * feat(mlx): map format_type to MLX save_method, reuse local save dir for hub push export_merged_model: format_type="4-bit (FP4)" → save_method="merged_4bit" (was hardcoded merged_16bit, ignoring the UI choice). Both export_merged_model and export_base_model now pass save_directory= to push_to_hub_merged so it reuses the just-written local folder instead of re-saving under a relative "username/model" directory. Co-Authored-By: Manan17 * feat(mlx): map format_type to MLX save_method, reuse local save dir for hub push - export_merged_model: format_type="4-bit (FP4)" → save_method="merged_4bit" (was hardcoded merged_16bit, ignoring the UI choice). - Both export_merged_model and export_base_model now pass save_directory= to push_to_hub_merged so it reuses the just-written local folder instead of re-saving under a relative "username/model" directory. Co-Authored-By: Manan17 * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * restore install * restore install * fix(mlx): restore FastVisionModel as a distinct class unsloth/__init__.py was assigning `FastVisionModel = FastLanguageModel` right after defining `class FastVisionModel(FastLanguageModel)` with a `for_training` static method. The alias erased the class binding, so the documented `FastVisionModel.for_training(model)` call from upstream Unsloth's VLM notebooks raised `AttributeError` on MLX. Remove the offending alias. `FastVisionModel` is now a real subclass of `FastLanguageModel` again — inherits `from_pretrained` / `get_peft_model` / `for_inference`, exposes `for_training` as a no-op pass-through (no-op because MLX doesn't have a train/eval mode flag; the call exists purely for GPU/MLX notebook parity). Verified end-to-end: Qwen3-VL-2B + LaTeX_OCR LoRA + vision LoRA via FastVisionModel.from_pretrained → get_peft_model → for_training → MLXTrainer.train runs 10 steps cleanly (loss 1.10 → 0.12, no NaNs, peak 5.89 GB). Studio's path (FastLanguageModel.from_pretrained for any repo, auto-detect VLM in the loader) is unaffected. Tier-1 review finding #8. * fix(mlx): restore FastVisionModel as a distinct class unsloth/__init__.py was assigning `FastVisionModel = FastLanguageModel` right after defining `class FastVisionModel(FastLanguageModel)` with a `for_training` static method. The alias erased the class binding, so the documented `FastVisionModel.for_training(model)` call from upstream Unsloth's VLM notebooks raised `AttributeError` on MLX. Remove the offending alias. `FastVisionModel` is now a real subclass of `FastLanguageModel` again — inherits `from_pretrained` / `get_peft_model` / `for_inference`, exposes `for_training` as a no-op pass-through (no-op because MLX doesn't have a train/eval mode flag; the call exists purely for GPU/MLX notebook parity). Verified end-to-end: Qwen3-VL-2B + LaTeX_OCR LoRA + vision LoRA via FastVisionModel.from_pretrained → get_peft_model → for_training → MLXTrainer.train() runs 10 steps cleanly (loss 1.10 → 0.12, no NaNs, peak 5.89 GB). Studio's path (FastLanguageModel.from_pretrained for any repo, auto-detect VLM in the loader) is unaffected. Tier-1 review finding #8. * Studio: harden MLX training and export, restore GPU init guards Studio export Restore Tuple[bool, str, Optional[str]] contract on export_merged_model, export_base_model, export_gguf, and export_lora_adapter, populating output_path on successful local saves so routes/worker/CLI/frontend details.output_path is non-empty again. Lift the GPU save_method assignment out of the local-save branch so Hub-only merged exports (save_directory='', push_to_hub=True) no longer hit UnboundLocalError on the push branch. For MLX merged and base hub-only export, stage to a tempfile.TemporaryDirectory before push_to_hub_merged instead of passing save_directory=''. Source _IS_MLX from unsloth instead of recomputing the platform check (single source of truth, also enforces mlx-package availability). Studio MLX training/inference Pass token=hf_token into FastMLXModel.from_pretrained for gated/private models, matching the inference path. Strip hf_token and wandb_token from wandb.init(config=...) so secrets do not leak into the W&B run config. Replace load_from_disk(local_datasets[0]) with the existing UnslothTrainer._resolve_local_files / _loader_for_files helpers so uploaded JSON/JSONL/CSV/Parquet files train through the normal datasets loader (load_from_disk still used for HF save_to_disk directories). Make the dataset slice helper inclusive at the end and treat 0 as a real index instead of "unset", matching the GPU and embedding paths. Add a status_message -> message alias inside _send so the existing parent pump (training.py) renders MLX status updates instead of blanks. Forward min_p through generate_chat_response into _generate_text / _generate_vlm and into make_sampler / vlm_kwargs so the sampling control is no longer a no-op on MLX. Wrap unsloth_zoo.mlx_loader / mlx_trainer imports with a clearer ImportError pointing users at install.sh for Apple Silicon. Exit the MLX stop-polling thread on EOFError/OSError instead of busy-looping when the queue/pipe is permanently closed (one-line why-safe rationale inline). Studio frontend ParamsSection subscribes to platform deviceType via the Zustand hook so the gradient checkpointing dropdown re-renders after the async device fetch completes. Studio hardware get_gpu_utilization MLX branch now reads _read_apple_gpu_stats once and derives VRAM totals from psutil, removing the second ioreg subprocess per utilization poll. Unsloth core Restore the os.geteuid == 0 guard around the CUDA ldconfig recovery that was lost when GPU initialization moved into _gpu_init.py, plus the non-root manual-fix warning branch. Non-root CUDA users no longer shell out to ldconfig at import time. Load dataprep/raw_text via importlib so the MLX import path no longer pulls torch in through dataprep/__init__.py -> synthetic.py. FastVisionModel.from_pretrained overrides the inherited delegator only to inject text_only=False; this is an extension, not a duplication, and is needed so VLM checkpoint loads keep the vision tower. Wrap the MLX-branch unsloth_zoo import with a clearer ImportError. * Studio: regression tests for MLX training/export and GPU init ldconfig guard tests/python/test_gpu_init_ldconfig_guard.py asserts the geteuid root check still wraps the ldconfig recovery and the non-root branch warns bnb users; AST + source-text inspection so the test runs without torch. tests/studio/test_export_output_path_contract.py covers the Tuple[bool, str, Optional[str]] return contract on every export method, the output_path assignment after successful local save, the Hub-only GPU save_method binding fix, the MLX hub-only TemporaryDirectory staging, and the single-source `_IS_MLX` import from unsloth. tests/studio/test_mlx_training_worker_behaviors.py covers token forwarding to FastMLXModel.from_pretrained, wandb config secret stripping, file-aware local dataset loading, status_message -> message aliasing, inclusive slice semantics, EOFError/OSError stop thread exit, and the friendly mlx_loader / mlx_trainer ImportError. * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * fix(mlx): cap inference memory + release wired on unload + tame worker pre-pin Three memory-hardening fixes for Studio's MLX path: 1. Inference applies the same Metal caps as the trainer. load_model previously only called set_wired_limit(100% of recommended) with no upper memory_limit, leaving large VLM checkpoints unbounded during the loader allocation. Add _configure_memory_limits() that sets memory_limit to 85% of recommended and wired_limit to min(recommended, memory_limit) — matching MLXTrainer's defaults so behavior is the same whether the user trains or just runs inference. 2. unload_model releases pinned memory back to the OS — but only when the cache is empty. Without this, pinned wired bytes stayed allocated to MLX after the model was gone, starving other apps. The release is guarded on `not self.models` so unloading one of several cached models doesn't un-pin weights still in use. 3. Worker pre-cap is conservative instead of aggressive. The previous pre-pin set_wired_limit(100% of recommended) competed with MLXTrainer's later more conservative cap. Replace with the same 85%-memory / min(rec, memory) pair that the trainer applies later (idempotent re-apply). Bounds the model load + LoRA setup window without over-pinning. * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * tests/studio: regression tests for the _IS_MLX dispatch gate Two gates drive every MLX-vs-CUDA dispatch decision in Studio: 1. unsloth._IS_MLX in unsloth/__init__.py — evaluated once at import time, read by Studio worker code to choose the GPU vs MLX trainer and inference paths. Defined as Darwin AND arm64 AND find_spec("mlx") is not None. 2. utils.hardware.detect_hardware() — runtime probe with priority CUDA > XPU > MLX > CPU. The MLX branch is reached only when both CUDA and XPU are unavailable and the host is Apple Silicon and mlx is importable. Neither gate had a direct test. Adds tests/studio/test_is_mlx_dispatch_gate.py with six tests: test_is_mlx_gate_uses_three_required_predicates AST-walks unsloth/__init__.py and asserts the _IS_MLX assignment is a BoolOp(And) of platform.system()=="Darwin", platform.machine()=="arm64", and find_spec("mlx") is not None. Catches accidental rewrites that drop a predicate. test_is_mlx_gate_true_on_apple_silicon_with_mlx_present Spoofs platform to Darwin/arm64, injects a fake mlx module so find_spec returns a real ModuleSpec, re-evaluates the gate expression. Verifies it flips True under the exact conditions Studio expects. test_is_mlx_gate_false_when_mlx_missing Spoofs Apple Silicon but with mlx absent. Verifies the gate stays False (so a Mac without mlx installed does not pretend to have MLX support). test_is_mlx_gate_false_on_non_apple_silicon Canary on the actual Linux+CUDA / AMD / Intel test host: the gate must remain False regardless of whether mlx happens to be importable. Protects existing GPU users from accidental MLX hijack when MLX support evolves. test_detect_hardware_picks_mlx_when_only_apple_silicon_available Forces torch.cuda and torch.xpu off, spoofs Apple Silicon, injects fake mlx and mlx.core. detect_hardware() must return DeviceType.MLX. test_detect_hardware_picks_cuda_on_real_host Canary: on a real CUDA host detect_hardware() must return DeviceType.CUDA. Protects against the MLX branch shadowing CUDA dispatch on NVIDIA / AMD ROCm hosts. Uses the same monkeypatch.setitem(sys.modules, ...) fake-mlx pattern as the existing test_mlx_inference_backend.py — no new test infrastructure, no real mlx install required. * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * Add AGPL-3.0 SPDX header to Studio MLX regression tests Four Studio MLX test files shipped without an SPDX-License-Identifier: studio/backend/tests/test_mlx_training_worker_config.py tests/studio/test_mlx_training_worker_behaviors.py tests/studio/test_export_output_path_contract.py tests/studio/test_is_mlx_dispatch_gate.py They sit in or alongside studio/backend/, which is governed by studio/LICENSE.AGPL-3.0, and exercise AGPL Studio code. Add the same "# SPDX-License-Identifier: AGPL-3.0-only" header that's already on test_mlx_inference_backend.py so the license declaration matches the code under test rather than defaulting to the repo-root Apache-2.0. * Wrap MLX submodule imports with friendly install hint The _IS_MLX block at the top of unsloth/__init__.py already catches the missing-package case with a friendly install hint, but the follow-up "from unsloth_zoo.mlx_trainer import ..." and "from unsloth_zoo.mlx_loader import ..." lines run unguarded. An Apple Silicon user who has unsloth-zoo installed but on an older version (e.g. the current PyPI release, before the MLX modules ship) sees a raw ImportError on the submodule rather than the hint that points at install.sh. Wrap the two submodule imports in the same try/except shape so the friendly install message fires whether the package is missing entirely or just predates the MLX submodules. No-op once both packages release together; smooths the transitional window where unsloth/main has merged but unsloth-zoo on PyPI has not. --------- Co-authored-by: DoubleMathew Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com> Co-authored-by: Lee Jackson <130007945+Imagineer99@users.noreply.github.com> Co-authored-by: Daniel Han --- .gitignore | 1 + install.sh | 6 + pyproject.toml | 4 +- studio/backend/core/export/export.py | 262 +++++--- .../backend/core/inference/mlx_inference.py | 395 +++++++++++ studio/backend/core/inference/worker.py | 92 +++ studio/backend/core/training/training.py | 43 +- studio/backend/core/training/worker.py | 628 ++++++++++++++++++ .../tests/test_mlx_inference_backend.py | 157 +++++ .../tests/test_mlx_training_worker_config.py | 83 +++ studio/backend/utils/hardware/hardware.py | 99 ++- .../frontend/src/components/app-sidebar.tsx | 15 +- studio/frontend/src/config/env.ts | 2 +- .../src/features/settings/settings-dialog.tsx | 8 +- .../features/settings/tabs/api-keys-tab.tsx | 2 +- .../studio/sections/params-section.tsx | 8 +- .../features/training/lib/model-defaults.ts | 9 +- .../frontend/src/hooks/use-hf-model-search.ts | 24 +- studio/frontend/src/types/training.ts | 2 +- tests/python/test_gpu_init_ldconfig_guard.py | 46 ++ .../test_export_output_path_contract.py | 122 ++++ tests/studio/test_is_mlx_dispatch_gate.py | 213 ++++++ .../test_mlx_training_worker_behaviors.py | 90 +++ unsloth/__init__.py | 429 +++--------- unsloth/_gpu_init.py | 346 ++++++++++ 25 files changed, 2623 insertions(+), 463 deletions(-) create mode 100644 studio/backend/core/inference/mlx_inference.py create mode 100644 studio/backend/tests/test_mlx_inference_backend.py create mode 100644 studio/backend/tests/test_mlx_training_worker_config.py create mode 100644 tests/python/test_gpu_init_ldconfig_guard.py create mode 100644 tests/studio/test_export_output_path_contract.py create mode 100644 tests/studio/test_is_mlx_dispatch_gate.py create mode 100644 tests/studio/test_mlx_training_worker_behaviors.py create mode 100644 unsloth/_gpu_init.py diff --git a/.gitignore b/.gitignore index b6786ee655..2f23f18d65 100644 --- a/.gitignore +++ b/.gitignore @@ -228,3 +228,4 @@ setup_leo.sh server.pid *.log package-lock.json +llama.cpp/ diff --git a/install.sh b/install.sh index ec47d016eb..9046a9bdf6 100755 --- a/install.sh +++ b/install.sh @@ -2017,6 +2017,12 @@ else fi fi +# ── Install mlx-vlm on Apple Silicon (optional, for VLM training) ── +if [ "$OS" = "macos" ] && [ "$_ARCH" = "arm64" ]; then + substep "installing mlx-vlm (VLM training support)..." + run_install_cmd "install mlx-vlm" uv pip install --python "$_VENV_PY" mlx-vlm +fi + # ── Run studio setup ── tauri_log "STEP" "Running Studio setup" # When --local, use the repo's own setup.sh directly. diff --git a/pyproject.toml b/pyproject.toml index c2f884e192..5687ea12f8 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -89,7 +89,7 @@ huggingfacenotorch = [ ] huggingface = [ "unsloth[huggingfacenotorch]", - "unsloth_zoo>=2026.5.1", + "unsloth_zoo>=2026.4.8", "torchvision", "unsloth[triton]", ] @@ -579,7 +579,7 @@ colab-ampere-torch220 = [ "flash-attn>=2.6.3 ; ('linux' in sys_platform)", ] colab-new = [ - "unsloth_zoo>=2026.5.1", + "unsloth_zoo>=2026.4.8", "packaging", "tyro", "transformers>=4.51.3,!=4.52.0,!=4.52.1,!=4.52.2,!=4.52.3,!=4.53.0,!=4.54.0,!=4.55.0,!=4.55.1,!=4.57.0,!=4.57.4,!=4.57.5,!=5.0.0,!=5.1.0,<=5.5.0", diff --git a/studio/backend/core/export/export.py b/studio/backend/core/export/export.py index 6fee5a38f7..4ab95d896f 100644 --- a/studio/backend/core/export/export.py +++ b/studio/backend/core/export/export.py @@ -9,16 +9,14 @@ Export backend - handles model exporting in various formats import glob import json import structlog +import tempfile from loggers import get_logger import os import shutil from pathlib import Path from typing import Optional, Tuple, List -from peft import PeftModel, PeftModelForCausalLM -from unsloth import FastLanguageModel, FastVisionModel +from unsloth import FastLanguageModel, FastVisionModel, _IS_MLX from huggingface_hub import HfApi, ModelCard -from transformers.modeling_utils import PushToHubMixin -import torch from utils.hardware import clear_gpu_cache from utils.models import is_vision_model, get_base_model_from_lora @@ -26,6 +24,12 @@ from utils.models.model_config import detect_audio_type from utils.paths import ensure_dir, outputs_root, resolve_export_dir, resolve_output_dir from core.inference import get_inference_backend +# GPU-only imports — guarded for Apple Silicon where these aren't needed +if not _IS_MLX: + from peft import PeftModel, PeftModelForCausalLM + from transformers.modeling_utils import PushToHubMixin + import torch + logger = get_logger(__name__) _LLAMA_CPP_SCRIPTS_WARNING_EMITTED = False @@ -225,7 +229,7 @@ class ExportBackend: model, tokenizer = FastModel.from_pretrained( model_name = checkpoint_path, max_seq_length = max_seq_length, - dtype = torch.float32, + dtype = None if _IS_MLX else torch.float32, load_in_4bit = False, trust_remote_code = trust_remote_code, ) @@ -262,8 +266,12 @@ class ExportBackend: trust_remote_code = trust_remote_code, ) - # Check if PEFT model - self.is_peft = isinstance(model, (PeftModel, PeftModelForCausalLM)) + # Check if PEFT / LoRA model + if _IS_MLX: + # MLX doesn't use PeftModel — detect LoRA via adapter_config.json + self.is_peft = adapter_config.exists() + else: + self.is_peft = isinstance(model, (PeftModel, PeftModelForCausalLM)) # Store loaded model self.current_model = model @@ -325,9 +333,7 @@ class ExportBackend: private: Whether to make the repo private Returns: - Tuple of (success, message, output_path). output_path is the - resolved absolute on-disk directory of the saved model when - ``save_directory`` was set, else None. + Tuple of (success: bool, message: str, output_path: Optional[str]) """ if not self.current_model or not self.current_tokenizer: return False, "No model loaded. Please select a checkpoint first.", None @@ -341,14 +347,17 @@ class ExportBackend: output_path: Optional[str] = None try: - # Determine save method - if format_type == "4-bit (FP4)": - save_method = "merged_4bit_forced" - elif self._audio_type == "whisper": - # Whisper uses save_method=None for local 16-bit merged save - save_method = None - else: # 16-bit (FP16) - save_method = "merged_16bit" + if _IS_MLX: + mlx_save_method = ( + "merged_4bit" if format_type == "4-bit (FP4)" else "merged_16bit" + ) + else: + if format_type == "4-bit (FP4)": + save_method = "merged_4bit_forced" + elif self._audio_type == "whisper": + save_method = None + else: + save_method = "merged_16bit" # Save locally if requested if save_directory: @@ -356,11 +365,17 @@ class ExportBackend: logger.info(f"Saving merged model locally to: {save_directory}") ensure_dir(Path(save_directory)) - self.current_model.save_pretrained_merged( - save_directory, self.current_tokenizer, save_method = save_method - ) + if _IS_MLX: + self.current_model.save_pretrained_merged( + save_directory, + self.current_tokenizer, + save_method = mlx_save_method, + ) + else: + self.current_model.save_pretrained_merged( + save_directory, self.current_tokenizer, save_method = save_method + ) - # Write export metadata so the Chat page can identify the base model self._write_export_metadata(save_directory) logger.info(f"Model saved successfully to {save_directory}") output_path = str(Path(save_directory).resolve()) @@ -376,17 +391,40 @@ class ExportBackend: logger.info(f"Pushing merged model to Hub: {repo_id}") - # Whisper uses save_method=None for local but "merged_16bit" for hub push - hub_save_method = ( - save_method if save_method is not None else "merged_16bit" - ) - self.current_model.push_to_hub_merged( - repo_id, - self.current_tokenizer, - save_method = hub_save_method, - token = hf_token, - private = private, - ) + if _IS_MLX: + if save_directory: + self.current_model.push_to_hub_merged( + repo_id, + self.current_tokenizer, + save_directory = save_directory, + token = hf_token, + private = private, + ) + else: + with tempfile.TemporaryDirectory() as tmp_dir: + self.current_model.save_pretrained_merged( + tmp_dir, + self.current_tokenizer, + save_method = mlx_save_method, + ) + self.current_model.push_to_hub_merged( + repo_id, + self.current_tokenizer, + save_directory = tmp_dir, + token = hf_token, + private = private, + ) + else: + hub_save_method = ( + save_method if save_method is not None else "merged_16bit" + ) + self.current_model.push_to_hub_merged( + repo_id, + self.current_tokenizer, + save_method = hub_save_method, + token = hf_token, + private = private, + ) logger.info(f"Model pushed successfully to {repo_id}") return True, "Model exported successfully", output_path @@ -411,9 +449,7 @@ class ExportBackend: Export base model (for non-PEFT models). Returns: - Tuple of (success, message, output_path). output_path is the - resolved absolute on-disk directory of the saved model when - ``save_directory`` was set, else None. + Tuple of (success: bool, message: str, output_path: Optional[str]) """ if not self.current_model or not self.current_tokenizer: return False, "No model loaded. Please select a checkpoint first.", None @@ -433,8 +469,16 @@ class ExportBackend: logger.info(f"Saving base model locally to: {save_directory}") ensure_dir(Path(save_directory)) - self.current_model.save_pretrained(save_directory) - self.current_tokenizer.save_pretrained(save_directory) + if _IS_MLX: + # MLX: save_pretrained_merged handles non-LoRA models too + # (fuse() is a no-op when there are no LoRA layers) + self.current_model.save_pretrained_merged( + save_directory, + self.current_tokenizer, + ) + else: + self.current_model.save_pretrained(save_directory) + self.current_tokenizer.save_pretrained(save_directory) # Write export metadata so the Chat page can identify the base model self._write_export_metadata(save_directory) @@ -452,44 +496,73 @@ class ExportBackend: logger.info(f"Pushing base model to Hub: {repo_id}") - # Get base model name from request or model config - base_model = ( - base_model_id - or self.current_model.config._name_or_path - or "unknown" - ) - - # Create repo - hf_api = HfApi(token = hf_token) - repo_id = PushToHubMixin._create_repo( - PushToHubMixin, - repo_id = repo_id, - private = private, - token = hf_token, - ) - username = repo_id.split("/")[0] - - # Create and push model card - content = MODEL_CARD.format( - username = username, - base_model = base_model, - model_type = self.current_model.config.model_type, - method = "", - extra = "unsloth", - ) - card = ModelCard(content) - card.push_to_hub( - repo_id, token = hf_token, commit_message = "Unsloth Model Card" - ) - - # Upload model files - if save_directory: - hf_api.upload_folder( - folder_path = save_directory, repo_id = repo_id, repo_type = "model" - ) - logger.info(f"Model pushed successfully to {repo_id}") + if _IS_MLX: + if save_directory: + self.current_model.push_to_hub_merged( + repo_id, + self.current_tokenizer, + save_directory = save_directory, + token = hf_token, + private = private, + ) + else: + with tempfile.TemporaryDirectory() as tmp_dir: + self.current_model.save_pretrained_merged( + tmp_dir, + self.current_tokenizer, + ) + self.current_model.push_to_hub_merged( + repo_id, + self.current_tokenizer, + save_directory = tmp_dir, + token = hf_token, + private = private, + ) else: - return False, "Local save directory required for Hub upload", None + # Get base model name from request or model config + base_model = ( + base_model_id + or self.current_model.config._name_or_path + or "unknown" + ) + + # Create repo + hf_api = HfApi(token = hf_token) + repo_id = PushToHubMixin._create_repo( + PushToHubMixin, + repo_id = repo_id, + private = private, + token = hf_token, + ) + username = repo_id.split("/")[0] + + # Create and push model card + content = MODEL_CARD.format( + username = username, + base_model = base_model, + model_type = self.current_model.config.model_type, + method = "", + extra = "unsloth", + ) + card = ModelCard(content) + card.push_to_hub( + repo_id, token = hf_token, commit_message = "Unsloth Model Card" + ) + + # Upload model files + if save_directory: + hf_api.upload_folder( + folder_path = save_directory, + repo_id = repo_id, + repo_type = "model", + ) + logger.info(f"Model pushed successfully to {repo_id}") + else: + return ( + False, + "Local save directory required for Hub upload", + None, + ) return True, "Model exported successfully", output_path @@ -519,9 +592,7 @@ class ExportBackend: hf_token: Hugging Face token Returns: - Tuple of (success, message, output_path). output_path is the - resolved absolute on-disk directory containing the .gguf - files when ``save_directory`` was set, else None. + Tuple of (success: bool, message: str, output_path: Optional[str]) """ if not self.current_model or not self.current_tokenizer: return False, "No model loaded. Please select a checkpoint first.", None @@ -692,9 +763,7 @@ class ExportBackend: Export LoRA adapter only (not merged). Returns: - Tuple of (success, message, output_path). output_path is the - resolved absolute on-disk directory of the saved adapter - when ``save_directory`` was set, else None. + Tuple of (success: bool, message: str, output_path: Optional[str]) """ if not self.current_model or not self.current_tokenizer: return False, "No model loaded. Please select a checkpoint first.", None @@ -710,8 +779,13 @@ class ExportBackend: logger.info(f"Saving LoRA adapter locally to: {save_directory}") ensure_dir(Path(save_directory)) - self.current_model.save_pretrained(save_directory) - self.current_tokenizer.save_pretrained(save_directory) + if _IS_MLX: + # MLX: save adapters.safetensors + tokenizer files + self.current_model.save_lora_adapters(save_directory) + self.current_tokenizer.save_pretrained(save_directory) + else: + self.current_model.save_pretrained(save_directory) + self.current_tokenizer.save_pretrained(save_directory) logger.info(f"Adapter saved successfully to {save_directory}") output_path = str(Path(save_directory).resolve()) @@ -726,10 +800,24 @@ class ExportBackend: logger.info(f"Pushing LoRA adapter to Hub: {repo_id}") - self.current_model.push_to_hub(repo_id, token = hf_token, private = private) - self.current_tokenizer.push_to_hub( - repo_id, token = hf_token, private = private - ) + if _IS_MLX: + with tempfile.TemporaryDirectory() as tmp_dir: + self.current_model.save_lora_adapters(tmp_dir) + self.current_tokenizer.save_pretrained(tmp_dir) + hf_api = HfApi(token = hf_token) + hf_api.create_repo(repo_id, private = private, exist_ok = True) + hf_api.upload_folder( + folder_path = tmp_dir, + repo_id = repo_id, + repo_type = "model", + ) + else: + self.current_model.push_to_hub( + repo_id, token = hf_token, private = private + ) + self.current_tokenizer.push_to_hub( + repo_id, token = hf_token, private = private + ) logger.info(f"Adapter pushed successfully to {repo_id}") return True, "LoRA adapter exported successfully", output_path diff --git a/studio/backend/core/inference/mlx_inference.py b/studio/backend/core/inference/mlx_inference.py new file mode 100644 index 0000000000..1d2b03ecb9 --- /dev/null +++ b/studio/backend/core/inference/mlx_inference.py @@ -0,0 +1,395 @@ +# SPDX-License-Identifier: AGPL-3.0-only +"""MLX inference backend for Apple Silicon. + +Drop-in replacement for InferenceBackend — same interface, uses mlx-lm/mlx-vlm +instead of torch/transformers for model loading and generation. +""" + +import threading +from typing import Optional, Generator +from loggers import get_logger + +logger = get_logger(__name__) + + +class MLXInferenceBackend: + def __init__(self): + self.models = {} + self.active_model_name = None + self.loading_models = set() + self.loaded_local_models = [] + self.device = "mlx" + self._generation_lock = threading.Lock() + + # MLX state + self._model = None + self._tokenizer = None + self._processor = None + self._is_vlm = False + self._config = {} + + # Recorded for unload to release pinned memory back to the OS. + self._memory_limits_applied = {} + + def _configure_memory_limits(self): + """Apply Metal memory caps before loading a model. + + Mirrors MLXTrainer._configure_memory_limits's defaults: + memory_limit = 85% of recommended working-set, + wired_limit = min(recommended, memory_limit). Recorded so unload + can lower wired_limit back to release pinned RAM. + """ + import mlx.core as mx + + if not mx.metal.is_available(): + return + info = mx.device_info() + rec_bytes = info.get("max_recommended_working_set_size") + if not rec_bytes or rec_bytes <= 0: + return + rec_gb = rec_bytes / 1e9 + memory_limit_gb = rec_gb * 0.85 + wired_limit_gb = min(rec_gb, memory_limit_gb) + mx.set_memory_limit(int(memory_limit_gb * 1e9)) + mx.set_wired_limit(int(wired_limit_gb * 1e9)) + self._memory_limits_applied = { + "memory_limit_gb": memory_limit_gb, + "wired_limit_gb": wired_limit_gb, + "recommended_gb": rec_gb, + } + logger.info( + "MLX memory caps: memory_limit=%.2f GB, wired_limit=%.2f GB", + memory_limit_gb, + wired_limit_gb, + ) + + def load_model( + self, + config, + max_seq_length = 2048, + load_in_4bit = True, + hf_token = None, + trust_remote_code = False, + gpu_ids = None, + dtype = None, + ) -> bool: + import mlx.core as mx + + model_name = config.identifier if hasattr(config, "identifier") else str(config) + is_vision = getattr(config, "is_vision", False) + + if hf_token: + import os + + os.environ["HF_TOKEN"] = hf_token + self._configure_memory_limits() + + is_lora = getattr(config, "is_lora", False) + + logger.info( + "Loading %s via %s (is_lora=%s)", + model_name, + "mlx-vlm" if is_vision else "mlx-lm", + is_lora, + ) + + try: + from unsloth_zoo.mlx_loader import FastMLXModel + except ImportError as e: + raise ImportError( + "Unsloth: MLX inference requires unsloth-zoo with the MLX modules " + "(unsloth_zoo.mlx_loader). Reinstall via install.sh on Apple Silicon." + ) from e + + model, tokenizer_or_processor = FastMLXModel.from_pretrained( + model_name, + max_seq_length = max_seq_length, + dtype = dtype, + load_in_4bit = load_in_4bit, + token = hf_token, + trust_remote_code = trust_remote_code, + text_only = False if is_vision else True, + ) + + if is_vision: + processor = tokenizer_or_processor + self._model = model + self._processor = processor + self._tokenizer = getattr(processor, "tokenizer", processor) + self._is_vlm = True + else: + tokenizer = tokenizer_or_processor + self._model = model + self._tokenizer = tokenizer + self._processor = None + self._is_vlm = False + + self.active_model_name = model_name + self.models[model_name] = { + "model": self._model, + "tokenizer": self._tokenizer, + "processor": self._processor, + "is_vision": is_vision, + "is_lora": getattr(config, "is_lora", False), + "is_audio": False, + "audio_type": None, + "has_audio_input": False, + } + + logger.info("Model %s loaded successfully", model_name) + return True + + def unload_model(self, model_name: str) -> bool: + import mlx.core as mx + import gc + + if model_name in self.models: + del self.models[model_name] + self._model = None + self._tokenizer = None + self._processor = None + if self.active_model_name == model_name: + self.active_model_name = None + gc.collect() + mx.clear_cache() + + if mx.metal.is_available() and self._memory_limits_applied and not self.models: + try: + mx.set_wired_limit(0) + logger.info("MLX wired_limit released back to OS on unload") + except Exception as e: + logger.warning("Failed to release wired_limit: %s", e) + self._memory_limits_applied = {} + logger.info("Model %s unloaded", model_name) + return True + + def generate_chat_response( + self, + messages, + system_prompt = "", + image = None, + temperature = 0.7, + top_p = 0.9, + top_k = 40, + min_p = 0.0, + max_new_tokens = 256, + repetition_penalty = 1.0, + cancel_event = None, + ) -> Generator[str, None, None]: + if self._model is None: + raise RuntimeError("No model loaded") + + # Build messages with system prompt + full_messages = [] + if system_prompt: + full_messages.append({"role": "system", "content": system_prompt}) + full_messages.extend(messages) + + # Inject image into the last user message for VLM + if self._is_vlm and image is not None: + for msg in reversed(full_messages): + if msg.get("role") == "user": + content = msg.get("content", "") + if isinstance(content, str): + msg["content"] = [ + {"type": "image"}, + {"type": "text", "text": content}, + ] + elif isinstance(content, list): + # Prepend image if not already there + has_image = any( + p.get("type") == "image" + for p in content + if isinstance(p, dict) + ) + if not has_image: + content.insert(0, {"type": "image"}) + break + + if self._is_vlm: + yield from self._generate_vlm( + full_messages, + image, + temperature, + top_p, + top_k, + min_p, + max_new_tokens, + repetition_penalty, + cancel_event, + ) + else: + yield from self._generate_text( + full_messages, + temperature, + top_p, + top_k, + min_p, + max_new_tokens, + repetition_penalty, + cancel_event, + ) + + def _generate_text( + self, + messages, + temperature, + top_p, + top_k, + min_p, + max_new_tokens, + repetition_penalty, + cancel_event, + ): + from mlx_lm import stream_generate + from mlx_lm.sample_utils import make_sampler, make_logits_processors + + prompt = self._tokenizer.apply_chat_template( + messages, + tokenize = False, + add_generation_prompt = True, + ) + if prompt is None: + raise RuntimeError( + "apply_chat_template returned None — tokenizer may be incompatible" + ) + + sampler = make_sampler( + temp = temperature, + top_p = top_p, + top_k = int(top_k or 0), + min_p = float(min_p or 0.0), + min_tokens_to_keep = 1, + ) + # Only build a logits processor when we actually have a non-trivial + # repetition penalty (1.0 is the no-op value). + logits_processors = None + if repetition_penalty is not None and float(repetition_penalty) not in ( + 0.0, + 1.0, + ): + logits_processors = make_logits_processors( + repetition_penalty = float(repetition_penalty), + ) + + token_ids = [] + logger.info( + "Generating: prompt_len=%d, max_tokens=%d, model=%s, tokenizer=%s", + len(prompt), + max_new_tokens, + type(self._model).__name__, + type(self._tokenizer).__name__, + ) + with self._generation_lock: + try: + gen_kwargs = dict( + prompt = prompt, + max_tokens = max_new_tokens, + sampler = sampler, + ) + if logits_processors is not None: + gen_kwargs["logits_processors"] = logits_processors + for response in stream_generate( + self._model, + self._tokenizer, + **gen_kwargs, + ): + token_ids.append(response.token) + # Decode full sequence with skip_special_tokens — same as GPU + cumulative = self._tokenizer.decode( + token_ids, + skip_special_tokens = True, + ) + yield cumulative + + if cancel_event and cancel_event.is_set(): + break + except Exception as e: + import traceback + + logger.error("stream_generate failed:\n%s", traceback.format_exc()) + raise + + def _generate_vlm( + self, + messages, + image, + temperature, + top_p, + top_k, + min_p, + max_new_tokens, + repetition_penalty, + cancel_event, + ): + from mlx_vlm import stream_generate as vlm_stream + + # Apply chat template + chat_fn = getattr(self._processor, "apply_chat_template", None) + if ( + chat_fn is None + or not hasattr(self._processor, "chat_template") + or self._processor.chat_template is None + ): + tok = getattr(self._processor, "tokenizer", self._processor) + chat_fn = tok.apply_chat_template + + prompt = chat_fn(messages, tokenize = False, add_generation_prompt = True) + + # For VLM: always use mlx_vlm's stream_generate which handles + # pixel_values properly (passes None for text-only, image for VLM) + images = [image] if image is not None else None + + cumulative = "" + logger.info( + "VLM generating: prompt_len=%d, has_image=%s", + len(prompt), + image is not None, + ) + # mlx_vlm.stream_generate forwards **kwargs into generate_step, which + # accepts temp/top_p/top_k/repetition_penalty (and builds the sampler + # + logits_processors internally). Pass them through. + # NOTE: mlx_vlm.generate_step expects ``temperature=`` (long form) — + # passing ``temp=`` silently falls into **kwargs and is ignored, + # leaving generation stuck at the default 0.0 (greedy). + vlm_kwargs = dict( + max_tokens = max_new_tokens, + temperature = temperature, + top_p = top_p, + top_k = int(top_k or 0), + min_p = float(min_p or 0.0), + ) + if repetition_penalty is not None and float(repetition_penalty) not in ( + 0.0, + 1.0, + ): + vlm_kwargs["repetition_penalty"] = float(repetition_penalty) + + with self._generation_lock: + for response in vlm_stream( + self._model, + self._processor, + prompt, + images, + **vlm_kwargs, + ): + token_text = ( + response.text if hasattr(response, "text") else str(response) + ) + cumulative += token_text + yield cumulative + if cancel_event and cancel_event.is_set(): + break + + def generate_with_adapter_control( + self, use_adapter = None, cancel_event = None, **gen_kwargs + ) -> Generator[str, None, None]: + # MLX LoRA adapter toggling not yet supported — generate normally + yield from self.generate_chat_response(cancel_event = cancel_event, **gen_kwargs) + + def reset_generation_state(self): + import mlx.core as mx + import gc + + gc.collect() + mx.clear_cache() diff --git a/studio/backend/core/inference/worker.py b/studio/backend/core/inference/worker.py index fbcce276ba..085a1ab899 100644 --- a/studio/backend/core/inference/worker.py +++ b/studio/backend/core/inference/worker.py @@ -663,6 +663,98 @@ def run_inference_process( model_name = config["model_name"] + # ── 0. MLX fast-path — skip torch/transformers entirely ── + backend_path = str(Path(__file__).resolve().parent.parent.parent) + if backend_path not in sys.path: + sys.path.insert(0, backend_path) + + from utils.hardware import hardware as _hw + + _hw.detect_hardware() + if _hw.DEVICE == _hw.DeviceType.MLX: + try: + _activate_transformers_version(model_name) + except Exception: + pass + try: + from core.inference.mlx_inference import MLXInferenceBackend + + backend = MLXInferenceBackend() + _send_response( + resp_queue, + {"type": "status", "message": "Loading model...", "ts": time.time()}, + ) + _handle_load(backend, config, resp_queue) + except Exception as exc: + _send_response( + resp_queue, + { + "type": "error", + "error": f"MLX inference init failed: {exc}", + "stack": traceback.format_exc(limit = 20), + "ts": time.time(), + }, + ) + return + + # Enter same command loop as GPU path + logger.info("MLX inference subprocess ready, entering command loop") + while True: + try: + cmd = cmd_queue.get(timeout = 1.0) + except _queue.Empty: + continue + except (EOFError, OSError): + return + if cmd is None: + continue + cmd_type = cmd.get("type", "") + try: + if cmd_type == "generate": + cancel_event.clear() + _handle_generate(backend, cmd, resp_queue, cancel_event) + elif cmd_type == "load": + if backend.active_model_name: + backend.unload_model(backend.active_model_name) + _handle_load(backend, cmd, resp_queue) + elif cmd_type == "unload": + _handle_unload(backend, cmd, resp_queue) + elif cmd_type == "cancel": + cancel_event.set() + elif cmd_type == "reset": + cancel_event.set() + backend.reset_generation_state() + _send_response(resp_queue, {"type": "reset_ack", "ts": time.time()}) + elif cmd_type == "status": + _send_response( + resp_queue, + { + "type": "status_response", + "active_model": backend.active_model_name, + "models": { + k: {kk: vv for kk, vv in v.items() if kk != "model"} + for k, v in backend.models.items() + }, + "loading": list(backend.loading_models), + "ts": time.time(), + }, + ) + elif cmd_type == "shutdown": + return + except Exception as exc: + logger.error("MLX command error (%s): %s", cmd_type, exc) + _send_response( + resp_queue, + { + "type": "gen_error" if cmd_type == "generate" else "error", + "request_id": cmd.get("request_id"), + "error": str(exc), + "stack": traceback.format_exc(limit = 20), + "ts": time.time(), + }, + ) + return + # ── 1. Activate correct transformers version BEFORE any ML imports ── try: _activate_transformers_version(model_name) diff --git a/studio/backend/core/training/training.py b/studio/backend/core/training/training.py index 5642faa189..a04ad5ef49 100644 --- a/studio/backend/core/training/training.py +++ b/studio/backend/core/training/training.py @@ -62,6 +62,7 @@ class TrainingProgress: grad_norm: Optional[float] = None num_tokens: Optional[int] = None eval_loss: Optional[float] = None + peak_memory_gb: Optional[float] = None class TrainingBackend: @@ -199,21 +200,27 @@ class TrainingBackend: config["load_in_4bit"] = False # Spawn subprocess — use locals so state is untouched on failure - resolved_gpu_ids, gpu_selection = prepare_gpu_selection( - kwargs.get("gpu_ids"), - model_name = config["model_name"], - hf_token = config["hf_token"] or None, - training_type = config["training_type"], - load_in_4bit = config["load_in_4bit"], - batch_size = config.get("batch_size", 4), - max_seq_length = config.get("max_seq_length", 2048), - lora_rank = config.get("lora_r", 16), - target_modules = config.get("target_modules"), - gradient_checkpointing = config.get("gradient_checkpointing", "unsloth"), - optimizer = config.get("optim", "adamw_8bit"), - ) - config["resolved_gpu_ids"] = resolved_gpu_ids - config["gpu_selection"] = gpu_selection + from utils.hardware import hardware as _hw + + if _hw.DEVICE == _hw.DeviceType.MLX: + config["resolved_gpu_ids"] = None + config["gpu_selection"] = None + else: + resolved_gpu_ids, gpu_selection = prepare_gpu_selection( + kwargs.get("gpu_ids"), + model_name = config["model_name"], + hf_token = config["hf_token"] or None, + training_type = config["training_type"], + load_in_4bit = config["load_in_4bit"], + batch_size = config.get("batch_size", 4), + max_seq_length = config.get("max_seq_length", 2048), + lora_rank = config.get("lora_r", 16), + target_modules = config.get("target_modules"), + gradient_checkpointing = config.get("gradient_checkpointing", "unsloth"), + optimizer = config.get("optim", "adamw_8bit"), + ) + config["resolved_gpu_ids"] = resolved_gpu_ids + config["gpu_selection"] = gpu_selection from .worker import run_training_process @@ -512,6 +519,12 @@ class TrainingBackend: self._progress.grad_norm = event.get("grad_norm") self._progress.num_tokens = event.get("num_tokens") self._progress.eval_loss = event.get("eval_loss") + _peak = event.get("peak_memory_gb") + if _peak is not None: + try: + self._progress.peak_memory_gb = float(_peak) + except (TypeError, ValueError): + pass self._progress.is_training = True status = event.get("status_message", "") if status: diff --git a/studio/backend/core/training/worker.py b/studio/backend/core/training/worker.py index 60b9e994ab..9c017db6de 100644 --- a/studio/backend/core/training/worker.py +++ b/studio/backend/core/training/worker.py @@ -338,6 +338,594 @@ def _activate_transformers_version(model_name: str) -> None: activate_transformers_for_subprocess(model_name) +def _adapt_for_mlx_vlm(items): + """Adapt GPU-path VLM dataset output for mlx-vlm consumption. + + The GPU path embeds PIL images inside messages content as + {"type": "image", "image": PIL_Image}. mlx-vlm's prepare_inputs + needs images at top-level to produce pixel_values — regardless of + model type. Extract them and leave bare {"type": "image"} placeholders. + """ + adapted = [] + for item in items: + images = [] + messages = [] + for msg in item.get("messages", []): + content = msg.get("content", "") + if isinstance(content, list): + new_content = [] + for part in content: + if isinstance(part, dict) and part.get("type") == "image": + img = part.get("image") + if img is not None: + images.append(img) + new_content.append({"type": "image"}) + else: + new_content.append(part) + messages.append({"role": msg["role"], "content": new_content}) + else: + messages.append(msg) + out = {"messages": messages} + if images: + out["image"] = images[0] if len(images) == 1 else images + elif "image" in item: + out["image"] = item["image"] + elif "images" in item: + out["images"] = item["images"] + adapted.append(out) + return adapted + + +_MLX_STUDIO_OPTIM_MAP = { + "adamw_8bit": "adamw", + "paged_adamw_8bit": "adamw", + "adamw_bnb_8bit": "adamw", + "paged_adamw_32bit": "adamw", + "adamw_torch": "adamw", + "adamw_torch_fused": "adamw", + "adamw": "adamw", + "adafactor": "adafactor", + "sgd": "sgd", + "adam": "adam", + "muon": "muon", + "lion": "lion", +} +_MLX_STUDIO_LR_SCHEDULERS = {"linear", "cosine", "constant"} + + +def _normalize_mlx_studio_optimizer(value): + raw = str(value or "adamw_8bit").strip().lower() + try: + return _MLX_STUDIO_OPTIM_MAP[raw] + except KeyError: + supported = ", ".join(sorted(_MLX_STUDIO_OPTIM_MAP)) + raise ValueError( + f"Unsupported optimizer for MLX training: {value!r}. " + f"Supported values: {supported}." + ) + + +def _normalize_mlx_studio_scheduler(value): + raw = str(value or "linear").strip().lower() + if raw not in _MLX_STUDIO_LR_SCHEDULERS: + supported = ", ".join(sorted(_MLX_STUDIO_LR_SCHEDULERS)) + raise ValueError( + f"Unsupported LR scheduler for MLX training: {value!r}. " + f"Supported values: {supported}." + ) + return raw + + +def _run_mlx_training(event_queue, stop_queue, config): + """Self-contained MLX training path for Apple Silicon. + + Uses MLXTrainer from unsloth_zoo directly -- no torch/SFTTrainer needed. + Mirrors the event_queue protocol so the parent process pump works unchanged. + """ + import time + import gc + import math + import threading + import queue as _queue + from pathlib import Path + + def _send(event_type, **kwargs): + if event_type == "status" and "message" not in kwargs: + sm = kwargs.get("status_message") + if sm is not None: + kwargs["message"] = sm + event_queue.put({"type": event_type, "ts": time.time(), **kwargs}) + + _send("status", status_message = "Loading MLX libraries...") + + import mlx.core as mx + + try: + from unsloth_zoo.mlx_loader import FastMLXModel + from unsloth_zoo.mlx_trainer import ( + MLXTrainer, + MLXTrainingConfig, + train_on_responses_only, + ) + except ImportError as e: + raise ImportError( + "Unsloth: MLX training requires unsloth-zoo with the MLX modules " + "(unsloth_zoo.mlx_loader / unsloth_zoo.mlx_trainer). Reinstall via " + "install.sh on Apple Silicon." + ) from e + from datasets import load_dataset + + if mx.metal.is_available(): + info = mx.device_info() + rec_bytes = info.get("max_recommended_working_set_size", 0) or 0 + if rec_bytes > 0: + memory_cap = int(rec_bytes * 0.85) + wired_cap = min(int(rec_bytes), memory_cap) + mx.set_memory_limit(memory_cap) + mx.set_wired_limit(wired_cap) + + model_name = config["model_name"] + hf_token = config.get("hf_token") or None + if hf_token: + os.environ["HF_TOKEN"] = hf_token + + if config.get("use_loftq"): + message = "LoftQ is not supported for MLX training yet." + _send("error", error = message) + raise NotImplementedError(message) + + optim_name = _normalize_mlx_studio_optimizer(config.get("optim", "adamw_8bit")) + lr_scheduler_type = _normalize_mlx_studio_scheduler( + config.get("lr_scheduler_type", "linear") + ) + + # ── 1. Load model ── + # Force text-only if the dataset is not an image dataset, even if the model + # has vision capabilities (e.g. Qwen3.5-VL trained on plain alpaca text). + _send("status", status_message = f"Loading {model_name}...") + is_dataset_image = bool(config.get("is_dataset_image", False)) + training_type = config.get("training_type", "LoRA/QLoRA") + use_lora = training_type == "LoRA/QLoRA" + model, tokenizer = FastMLXModel.from_pretrained( + model_name, + load_in_4bit = config.get("load_in_4bit", True), + full_finetuning = not use_lora, + text_only = None if is_dataset_image else True, + token = hf_token, + trust_remote_code = bool(config.get("trust_remote_code", False)), + random_state = config.get("random_seed", 3407), + ) + + is_vlm = bool(is_dataset_image and getattr(model, "_is_vlm_model", False)) + model._is_vlm_model = is_vlm + + # ── 2. Apply LoRA / full FT ── + # Pass gradient_checkpointing as string ("mlx"/"unsloth"/"none"/etc.) + # get_peft_model and MLXTrainer both accept strings and handle them. + gc_setting = config.get("gradient_checkpointing", "mlx") + if isinstance(gc_setting, str): + use_grad_checkpoint = ( + gc_setting if gc_setting.lower() not in ("false", "") else False + ) + else: + use_grad_checkpoint = gc_setting + + if use_lora: + _send("status", status_message = "Configuring LoRA adapters...") + peft_kwargs = dict( + r = config.get("lora_r", 16), + lora_alpha = config.get("lora_alpha", 16), + lora_dropout = config.get("lora_dropout", 0.0), + use_rslora = config.get("use_rslora", False), + init_lora_weights = config.get("init_lora_weights", True), + random_state = config.get("random_seed", 3407), + target_modules = config.get("target_modules") + or [ + "q_proj", + "k_proj", + "v_proj", + "o_proj", + "gate_proj", + "up_proj", + "down_proj", + ], + use_gradient_checkpointing = use_grad_checkpoint, + ) + finetune_language = config.get("finetune_language_layers", True) + finetune_attention = config.get("finetune_attention_modules", True) + finetune_mlp = config.get("finetune_mlp_modules", True) + finetune_vision = ( + config.get("finetune_vision_layers", False) if is_vlm else False + ) + + if ( + (finetune_attention or finetune_mlp) + and not finetune_language + and not finetune_vision + ): + finetune_language = True + + peft_kwargs["finetune_language_layers"] = finetune_language + peft_kwargs["finetune_attention_modules"] = finetune_attention + peft_kwargs["finetune_mlp_modules"] = finetune_mlp + if is_vlm: + peft_kwargs["finetune_vision_layers"] = finetune_vision + model = FastMLXModel.get_peft_model(model, **peft_kwargs) + + # ── 3. Load dataset ── + _send("status", status_message = "Loading dataset...") + hf_dataset = config.get("hf_dataset", "") + subset = config.get("subset") + train_split = config.get("train_split", "train") or "train" + eval_split = config.get("eval_split") + slice_start = config.get("dataset_slice_start") + slice_end = config.get("dataset_slice_end") + + def _slice(ds): + if slice_start is not None or slice_end is not None: + start = slice_start if slice_start is not None else 0 + end = slice_end if slice_end is not None else len(ds) - 1 + if end < start: + return ds.select([]) + ds = ds.select(range(start, min(end + 1, len(ds)))) + return ds + + def _load_local(file_paths): + from core.training.trainer import UnslothTrainer + from datasets import load_from_disk + + if len(file_paths) == 1: + p = Path(file_paths[0]) + if p.is_dir() and ( + (p / "dataset_info.json").exists() or (p / "state.json").exists() + ): + return load_from_disk(str(p)) + all_files = UnslothTrainer._resolve_local_files(file_paths) + if not all_files: + raise ValueError("No local dataset files found") + loader = UnslothTrainer._loader_for_files(all_files) + return load_dataset(loader, data_files = all_files, split = "train") + + if hf_dataset: + load_kwargs = {"split": train_split, "token": hf_token} + if subset: + load_kwargs["name"] = subset + dataset = load_dataset(hf_dataset, **load_kwargs) + dataset = _slice(dataset) + elif config.get("local_datasets"): + dataset = _load_local(config["local_datasets"]) + dataset = _slice(dataset) + else: + raise ValueError("No dataset specified") + + # Eval dataset (separate split or local file) + eval_dataset = None + if eval_split and hf_dataset: + eval_kwargs = {"split": eval_split, "token": hf_token} + if subset: + eval_kwargs["name"] = subset + try: + eval_dataset = load_dataset(hf_dataset, **eval_kwargs) + except Exception as e: + _send("status", status_message = f"Eval split load failed: {e}") + eval_dataset = None + elif config.get("local_eval_datasets"): + eval_dataset = _load_local(config["local_eval_datasets"]) + + # ── 3b. Format dataset (VLM or text) ── + # Reuse the GPU path's format pipeline for both VLM (auto-detects OCR/caption/ + # llava/sharegpt+images) and text (alpaca/sharegpt/chatml → "text" column). + format_type = config.get("format_type", "") + try: + from utils.datasets import format_and_template_dataset + + def _fmt_progress(status_message = "", **_kw): + _send("status", status_message = status_message) + + if is_vlm: + _send("status", status_message = "Formatting VLM dataset...") + vlm_info = format_and_template_dataset( + dataset, + model_name = model_name, + tokenizer = tokenizer, + is_vlm = True, + dataset_name = hf_dataset or "local", + progress_callback = _fmt_progress, + ) + if vlm_info.get("success"): + dataset = _adapt_for_mlx_vlm(vlm_info["dataset"]) + else: + errors = vlm_info.get("errors", []) + raise ValueError( + f"VLM dataset format conversion failed: {'; '.join(errors)}" + ) + if eval_dataset is not None: + ev_info = format_and_template_dataset( + eval_dataset, + model_name = model_name, + tokenizer = tokenizer, + is_vlm = True, + dataset_name = hf_dataset or "local", + ) + if ev_info.get("success"): + eval_dataset = _adapt_for_mlx_vlm(ev_info["dataset"]) + + elif format_type: + _send("status", status_message = f"Formatting dataset ({format_type})...") + info = format_and_template_dataset( + dataset, + model_name = model_name, + tokenizer = tokenizer, + is_vlm = False, + format_type = format_type, + dataset_name = hf_dataset or "local", + ) + if info.get("success", True): + dataset = info.get("dataset", dataset) + if eval_dataset is not None: + ev = format_and_template_dataset( + eval_dataset, + model_name = model_name, + tokenizer = tokenizer, + is_vlm = False, + format_type = format_type, + dataset_name = hf_dataset or "local", + ) + if ev.get("success", True): + eval_dataset = ev.get("dataset", eval_dataset) + except ImportError: + _send("status", status_message = "Format helper unavailable, using raw dataset") + + # ── 4. Resolve training steps ── + max_steps = config.get("max_steps", 0) or 0 + num_epochs = config.get("num_epochs", 3) + max_seq_length = config.get("max_seq_length", 2048) + batch_size = config.get("batch_size", 4) + grad_accum = config.get("gradient_accumulation_steps", 4) + + if max_steps <= 0: + max_steps = max( + 1, + math.ceil(len(dataset) / batch_size / grad_accum) * num_epochs, + ) + + lr_value = float(config.get("learning_rate", "2e-4")) + + # Warmup: prefer warmup_steps; fall back to warmup_ratio + warmup_steps = config.get("warmup_steps") + warmup_ratio = config.get("warmup_ratio") + if warmup_steps is None and warmup_ratio is not None: + warmup_steps = int(round(warmup_ratio * max_steps)) + if warmup_steps is None: + warmup_steps = 5 + + # ── 5. Build output dir ── + output_dir = config.get("output_dir", "") + if not output_dir: + output_dir = f"{model_name.replace('/', '_')}_{int(time.time())}" + # Resolve to ~/.unsloth/studio/outputs/ so the export page can find it + from utils.paths import resolve_output_dir, ensure_dir + + output_dir = str(resolve_output_dir(output_dir)) + ensure_dir(Path(output_dir)) + + # ── 6. Create trainer ── + eval_steps_val = config.get("eval_steps", 0) or 0 + if isinstance(eval_steps_val, float) and 0 < eval_steps_val < 1: + # Studio sometimes sends fraction-of-total-steps + eval_steps_val = max(1, int(eval_steps_val * max_steps)) + else: + eval_steps_val = int(eval_steps_val) + + trainer = MLXTrainer( + model = model, + tokenizer = tokenizer, + train_dataset = dataset, + eval_dataset = eval_dataset, + args = MLXTrainingConfig( + per_device_train_batch_size = batch_size, + gradient_accumulation_steps = grad_accum, + max_steps = max_steps, + learning_rate = lr_value, + warmup_steps = warmup_steps, + lr_scheduler_type = lr_scheduler_type, + optim = optim_name, + weight_decay = float(config.get("weight_decay", 0.001) or 0.001), + logging_steps = 1, + max_seq_length = max_seq_length, + seed = config.get("random_seed", 3407), + use_cce = True, + compile = True, + gradient_checkpointing = use_grad_checkpoint, + streaming = is_vlm, + packing = bool(config.get("packing", False)), + output_dir = output_dir, + save_steps = int(config.get("save_steps", 0) or 0), + eval_steps = eval_steps_val, + ), + ) + + # Tell the parent that eval is configured so the frontend shows the eval chart + if eval_dataset is not None and eval_steps_val > 0: + _send("eval_configured") + + # ── 7. Apply train_on_responses_only if requested ── + if config.get("train_on_completions", False): + _send("status", status_message = "Configuring response-only training...") + try: + from utils.datasets import ( + MODEL_TO_TEMPLATE_MAPPER, + TEMPLATE_TO_RESPONSES_MAPPER, + ) + + template_name = MODEL_TO_TEMPLATE_MAPPER.get(model_name.lower()) + markers = ( + TEMPLATE_TO_RESPONSES_MAPPER.get(template_name) + if template_name + else None + ) + if markers: + trainer = train_on_responses_only( + trainer, + instruction_part = markers["instruction"], + response_part = markers["response"], + ) + else: + _send( + "status", + status_message = f"train_on_completions skipped (no template for {model_name})", + ) + except Exception as e: + _send("status", status_message = f"train_on_completions failed: {e}") + + # ── 8. Setup wandb / tensorboard ── + wandb_run = None + tb_writer = None + if config.get("enable_wandb", False): + try: + import wandb as _wandb + + wandb_token = config.get("wandb_token") + if wandb_token: + os.environ["WANDB_API_KEY"] = wandb_token + _wandb_sensitive = {"hf_token", "wandb_token"} + wandb_run = _wandb.init( + project = config.get("wandb_project") or "unsloth-mlx", + config = {k: v for k, v in config.items() if k not in _wandb_sensitive}, + reinit = True, + ) + except Exception as e: + _send("status", status_message = f"wandb init failed: {e}") + if config.get("enable_tensorboard", False): + try: + from tensorboardX import SummaryWriter + except ImportError: + try: + from torch.utils.tensorboard import SummaryWriter + except ImportError: + SummaryWriter = None + if SummaryWriter is not None: + try: + tb_dir = config.get("tensorboard_dir") or f"{output_dir}/runs" + tb_writer = SummaryWriter(log_dir = tb_dir) + except Exception as e: + _send("status", status_message = f"tensorboard init failed: {e}") + else: + _send( + "status", + status_message = "tensorboard unavailable (install tensorboardX)", + ) + + # ── 9. Real-time progress callback ── + _send("status", status_message = f"Training {model_name}...") + + def _on_step(step, total, loss, lr, tok_s, peak_gb, elapsed, num_tokens): + eta = (elapsed / step * (total - step)) if step > 0 else 0 + _send( + "progress", + step = step, + epoch = round(step / total * num_epochs, 2) if total > 0 else 0, + loss = loss, + learning_rate = lr, + total_steps = total, + elapsed_seconds = elapsed, + eta_seconds = max(0, eta), + grad_norm = None, + num_tokens = num_tokens, + eval_loss = None, + status_message = None, + peak_memory_gb = peak_gb, + ) + if wandb_run is not None: + try: + wandb_run.log( + { + "train/loss": loss, + "train/learning_rate": lr, + "train/tokens_per_sec": tok_s, + "train/peak_gb": peak_gb, + "train/num_tokens": num_tokens, + }, + step = step, + ) + except Exception: + pass + if tb_writer is not None: + try: + tb_writer.add_scalar("train/loss", loss, step) + tb_writer.add_scalar("train/learning_rate", lr, step) + tb_writer.add_scalar("train/tokens_per_sec", tok_s, step) + tb_writer.add_scalar("train/peak_gb", peak_gb, step) + except Exception: + pass + + trainer.add_step_callback(_on_step) + + def _on_eval(step, eval_loss, perplexity): + _send("progress", step = step, eval_loss = eval_loss) + if wandb_run is not None: + try: + wandb_run.log( + {"eval/loss": eval_loss, "eval/perplexity": perplexity}, step = step + ) + except Exception: + pass + if tb_writer is not None: + try: + tb_writer.add_scalar("eval/loss", eval_loss, step) + tb_writer.add_scalar("eval/perplexity", perplexity, step) + except Exception: + pass + + trainer.add_eval_callback(_on_eval) + + # ── 10. Stop signal polling ── + _stop_save = [True] # mutable so thread can update; [save_flag] + + def _poll_stop(): + while True: + try: + msg = stop_queue.get(timeout = 1.0) + if msg and msg.get("type") == "stop": + _stop_save[0] = msg.get("save", True) + trainer.stop_requested = True + return + except _queue.Empty: + continue + except (EOFError, OSError): + # why safe: pipe permanently broken, no further messages can arrive + return + + stop_thread = threading.Thread(target = _poll_stop, daemon = True) + stop_thread.start() + + # ── 11. Run training ── + gc.collect() + mx.synchronize() + trainer.train() + + # ── 12. Save and finalize ── + if trainer.stop_requested and not _stop_save[0]: + # User clicked "Cancel" (save=False) — skip saving + _send("complete", output_dir = None, status_message = "Training cancelled") + else: + _send("status", status_message = "Saving model...") + mx.synchronize() + trainer.save_model(output_dir) + _send("complete", output_dir = output_dir, status_message = "Training completed") + + if tb_writer is not None: + try: + tb_writer.close() + except Exception: + pass + if wandb_run is not None: + try: + wandb_run.finish() + except Exception: + pass + + def run_training_process( *, event_queue: Any, @@ -371,6 +959,46 @@ def run_training_process( model_name = config["model_name"] + # ── 0. MLX FAST-PATH (must run before any torch/transformers imports) ── + # Apple Silicon uses MLXTrainer directly -- skip transformers version + # activation, causal-conv1d install, and torch imports entirely. + backend_path = str(Path(__file__).resolve().parent.parent.parent) + if backend_path not in sys.path: + sys.path.insert(0, backend_path) + + from utils.hardware import hardware as _hw + + _hw.detect_hardware() + if _hw.DEVICE == _hw.DeviceType.MLX: + if config.get("is_dataset_audio"): + event_queue.put( + { + "type": "error", + "error": "Audio dataset training is not yet supported on Apple Silicon.", + "stack": "", + "ts": time.time(), + } + ) + return + # Activate correct transformers version (Gemma-4 needs 5.5.0, etc.) + # Must happen before any transformers/mlx-lm imports in _run_mlx_training. + try: + _activate_transformers_version(model_name) + except Exception: + pass # Non-fatal: fall through with whatever version is installed + try: + _run_mlx_training(event_queue, stop_queue, config) + except Exception as exc: + event_queue.put( + { + "type": "error", + "error": str(exc), + "stack": traceback.format_exc(limit = 20), + "ts": time.time(), + } + ) + return + # ── 1. Activate correct transformers version BEFORE any ML imports ── try: _activate_transformers_version(model_name) diff --git a/studio/backend/tests/test_mlx_inference_backend.py b/studio/backend/tests/test_mlx_inference_backend.py new file mode 100644 index 0000000000..868e537372 --- /dev/null +++ b/studio/backend/tests/test_mlx_inference_backend.py @@ -0,0 +1,157 @@ +# SPDX-License-Identifier: AGPL-3.0-only + +import sys +import types +from types import SimpleNamespace + + +class _DummyMetal: + @staticmethod + def is_available(): + return False + + +class _DummyMX: + metal = _DummyMetal() + + @staticmethod + def set_wired_limit(_limit): + return None + + @staticmethod + def device_info(): + return {"max_recommended_working_set_size": 1024} + + +class _DummyTokenizer: + pass + + +class _DummyProcessor: + tokenizer = _DummyTokenizer() + + +class _DummyModel: + pass + + +def _install_fake_mlx(monkeypatch): + mlx_pkg = types.ModuleType("mlx") + mlx_core = types.ModuleType("mlx.core") + mlx_core.metal = _DummyMetal() + mlx_core.set_wired_limit = _DummyMX.set_wired_limit + mlx_core.device_info = _DummyMX.device_info + mlx_pkg.core = mlx_core + monkeypatch.setitem(sys.modules, "mlx", mlx_pkg) + monkeypatch.setitem(sys.modules, "mlx.core", mlx_core) + + +def _install_fake_fast_mlx(monkeypatch, calls): + class _FastMLXModel: + @staticmethod + def from_pretrained(*args, **kwargs): + calls.append((args, kwargs)) + if kwargs["text_only"] is False: + return _DummyModel(), _DummyProcessor() + return _DummyModel(), _DummyTokenizer() + + unsloth_zoo_pkg = types.ModuleType("unsloth_zoo") + mlx_loader = types.ModuleType("unsloth_zoo.mlx_loader") + mlx_loader.FastMLXModel = _FastMLXModel + unsloth_zoo_pkg.mlx_loader = mlx_loader + monkeypatch.setitem(sys.modules, "unsloth_zoo", unsloth_zoo_pkg) + monkeypatch.setitem(sys.modules, "unsloth_zoo.mlx_loader", mlx_loader) + + +def test_mlx_inference_text_load_forwards_studio_settings(monkeypatch): + _install_fake_mlx(monkeypatch) + calls = [] + _install_fake_fast_mlx(monkeypatch, calls) + + from core.inference.mlx_inference import MLXInferenceBackend + + backend = MLXInferenceBackend() + config = SimpleNamespace(identifier = "fake/text", is_vision = False, is_lora = False) + + assert backend.load_model( + config, + max_seq_length = 4096, + load_in_4bit = False, + hf_token = "hf-token", + trust_remote_code = True, + dtype = "float16", + ) + + assert calls == [ + ( + ("fake/text",), + { + "max_seq_length": 4096, + "dtype": "float16", + "load_in_4bit": False, + "token": "hf-token", + "trust_remote_code": True, + "text_only": True, + }, + ) + ] + assert backend._is_vlm is False + assert isinstance(backend._tokenizer, _DummyTokenizer) + + +def test_mlx_inference_vlm_lora_uses_unsloth_loader_without_native_adapter_rewrite( + monkeypatch, + tmp_path, +): + _install_fake_mlx(monkeypatch) + calls = [] + _install_fake_fast_mlx(monkeypatch, calls) + + def _native_vlm_load(*_args, **_kwargs): + raise AssertionError("Studio MLX VLM inference must use FastMLXModel") + + mlx_vlm = types.ModuleType("mlx_vlm") + mlx_vlm.load = _native_vlm_load + monkeypatch.setitem(sys.modules, "mlx_vlm", mlx_vlm) + + adapter_dir = tmp_path / "adapter" + adapter_dir.mkdir() + cfg_path = adapter_dir / "adapter_config.json" + original_cfg = '{"base_model_name_or_path": "fake/base", "rank": 8}\n' + cfg_path.write_text(original_cfg) + + from core.inference.mlx_inference import MLXInferenceBackend + + backend = MLXInferenceBackend() + config = SimpleNamespace( + identifier = str(adapter_dir), + is_vision = True, + is_lora = True, + base_model = "fake/base", + ) + + assert backend.load_model( + config, + max_seq_length = 8192, + load_in_4bit = True, + hf_token = "hf-token", + trust_remote_code = True, + ) + + assert calls == [ + ( + (str(adapter_dir),), + { + "max_seq_length": 8192, + "dtype": None, + "load_in_4bit": True, + "token": "hf-token", + "trust_remote_code": True, + "text_only": False, + }, + ) + ] + assert cfg_path.read_text() == original_cfg + assert backend._is_vlm is True + assert isinstance(backend._processor, _DummyProcessor) + assert isinstance(backend._tokenizer, _DummyTokenizer) diff --git a/studio/backend/tests/test_mlx_training_worker_config.py b/studio/backend/tests/test_mlx_training_worker_config.py new file mode 100644 index 0000000000..5900af4e3d --- /dev/null +++ b/studio/backend/tests/test_mlx_training_worker_config.py @@ -0,0 +1,83 @@ +# SPDX-License-Identifier: AGPL-3.0-only + +import importlib.util +import sys +import types +from pathlib import Path + +import pytest + + +def _load_worker_module(): + stub_names = ( + "structlog", + "loggers", + "utils", + "utils.hardware", + "utils.wheel_utils", + ) + previous_modules = {name: sys.modules.get(name) for name in stub_names} + + try: + sys.modules["structlog"] = types.ModuleType("structlog") + + loggers = types.ModuleType("loggers") + loggers.get_logger = lambda *_args, **_kwargs: None + sys.modules["loggers"] = loggers + + utils = types.ModuleType("utils") + utils.__path__ = [] + sys.modules["utils"] = utils + + hardware = types.ModuleType("utils.hardware") + hardware.apply_gpu_ids = lambda *_args, **_kwargs: None + sys.modules["utils.hardware"] = hardware + + wheel_utils = types.ModuleType("utils.wheel_utils") + for name in ( + "direct_wheel_url", + "flash_attn_wheel_url", + "install_wheel", + "probe_torch_wheel_env", + "url_exists", + ): + setattr(wheel_utils, name, lambda *_args, **_kwargs: None) + sys.modules["utils.wheel_utils"] = wheel_utils + + worker_path = ( + Path(__file__).resolve().parents[1] / "core" / "training" / "worker.py" + ) + spec = importlib.util.spec_from_file_location( + "mlx_training_worker_under_test", worker_path + ) + module = importlib.util.module_from_spec(spec) + assert spec.loader is not None + spec.loader.exec_module(module) + return module + finally: + for name, module in previous_modules.items(): + if module is None: + sys.modules.pop(name, None) + else: + sys.modules[name] = module + + +_worker = _load_worker_module() +_normalize_mlx_studio_optimizer = _worker._normalize_mlx_studio_optimizer +_normalize_mlx_studio_scheduler = _worker._normalize_mlx_studio_scheduler + + +def test_mlx_studio_optimizer_aliases_are_explicit(): + assert _normalize_mlx_studio_optimizer("adamw_8bit") == "adamw" + assert _normalize_mlx_studio_optimizer("paged_adamw_8bit") == "adamw" + assert _normalize_mlx_studio_optimizer("adafactor") == "adafactor" + + +def test_mlx_studio_rejects_unknown_optimizer(): + with pytest.raises(ValueError, match = "Unsupported optimizer for MLX training"): + _normalize_mlx_studio_optimizer("adamw_typo") + + +def test_mlx_studio_rejects_unknown_scheduler(): + with pytest.raises(ValueError, match = "Unsupported LR scheduler for MLX training"): + _normalize_mlx_studio_scheduler("linear_typo") diff --git a/studio/backend/utils/hardware/hardware.py b/studio/backend/utils/hardware/hardware.py index c218b7b4b9..3764e38272 100644 --- a/studio/backend/utils/hardware/hardware.py +++ b/studio/backend/utils/hardware/hardware.py @@ -143,6 +143,7 @@ def detect_hardware() -> DeviceType: # --- MLX: Apple Silicon --- if is_apple_silicon() and _has_mlx(): DEVICE = DeviceType.MLX + CHAT_ONLY = False chip = platform.processor() or platform.machine() print(f"Hardware detected: MLX — Apple Silicon ({chip})") return DEVICE @@ -270,19 +271,30 @@ def get_gpu_memory_info() -> Dict[str, Any]: import mlx.core as mx import psutil - # MLX uses unified memory — report system memory as the pool + # MLX uses unified memory. Total = system RAM. GPU memory used + # comes from IORegistry's AGXAccelerator (system-wide, no sudo). total = psutil.virtual_memory().total - # MLX doesn't expose per-process GPU allocation; report 0 as allocated - allocated = 0 + agx = _read_apple_gpu_stats() + allocated = agx.get("vram_used_bytes", 0) if agx else 0 + + try: + info = mx.device_info() + gpu_name = ( + info.get("device_name") + or platform.processor() + or platform.machine() + ) + except Exception: + gpu_name = platform.processor() or platform.machine() return { "available": True, "backend": _backend_label(device), "device": 0, - "device_name": f"Apple Silicon ({platform.processor() or platform.machine()})", + "device_name": f"Apple Silicon ({gpu_name})", "total_gb": total / (1024**3), "allocated_gb": allocated / (1024**3), - "reserved_gb": 0, + "reserved_gb": allocated / (1024**3), "free_gb": (total - allocated) / (1024**3), "utilization_pct": (allocated / total) * 100 if total else 0, } @@ -460,6 +472,39 @@ def _smi_query(func_name: str, *args, **kwargs) -> Optional[Dict[str, Any]]: return None +def _read_apple_gpu_stats() -> Dict[str, Any]: + """Query macOS IORegistry for AGX (Apple GPU) live stats. No sudo needed. + + Returns dict with utilization_pct, vram_used_bytes (system-wide GPU memory). + Returns empty dict on failure. + """ + import subprocess + import re + + try: + result = subprocess.run( + ["ioreg", "-r", "-c", "AGXAccelerator"], + capture_output = True, + timeout = 2, + ) + text = result.stdout.decode("utf-8", errors = "replace") + except Exception: + return {} + + # PerformanceStatistics block has GPU utilization and in-use memory + m = re.search(r'"PerformanceStatistics" = \{([^}]+)\}', text) + if not m: + return {} + stats_str = m.group(1) + pairs = re.findall(r'"([^"]+)"=(\d+)', stats_str) + stats = {k: int(v) for k, v in pairs} + + return { + "utilization_pct": stats.get("Device Utilization %", 0), + "vram_used_bytes": stats.get("In use system memory", 0), + } + + def get_gpu_utilization() -> Dict[str, Any]: """Return a live snapshot of device utilization information.""" device = get_device() @@ -470,6 +515,50 @@ def get_gpu_utilization() -> Dict[str, Any]: result["backend"] = _backend_label(device) return result + # MLX path: single _read_apple_gpu_stats() call carries both VRAM-used + # bytes and GPU utilization %. psutil for unified-memory total is cheap. + if device == DeviceType.MLX: + try: + import psutil + + agx = _read_apple_gpu_stats() + total_bytes = psutil.virtual_memory().total + except Exception as e: + logger.error(f"Error getting MLX GPU utilization: {e}") + return {"available": False, "backend": device.value, "error": str(e)} + if not agx: + return {"available": False, "backend": device.value} + allocated_bytes = agx.get("vram_used_bytes", 0) or 0 + vram_used_gb = allocated_bytes / (1024**3) + total_gb = total_bytes / (1024**3) + + try: + from core.training import get_training_backend + + tb = get_training_backend() + tb_progress = getattr(tb, "_progress", None) + if tb_progress is not None and getattr(tb_progress, "is_training", False): + tb_peak = getattr(tb_progress, "peak_memory_gb", None) + if tb_peak is not None and tb_peak > 0: + vram_used_gb = float(tb_peak) + except Exception: + pass + + return { + "available": True, + "backend": device.value, + "gpu_utilization_pct": agx.get("utilization_pct") if agx else None, + "temperature_c": None, + "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_limit_w": None, + "power_utilization_pct": None, + } + mem = get_gpu_memory_info() if device != DeviceType.CPU and mem.get("available"): return { diff --git a/studio/frontend/src/components/app-sidebar.tsx b/studio/frontend/src/components/app-sidebar.tsx index edcd5120eb..171b4eb92d 100644 --- a/studio/frontend/src/components/app-sidebar.tsx +++ b/studio/frontend/src/components/app-sidebar.tsx @@ -37,13 +37,12 @@ import { Delete02Icon, Download03Icon, GemIcon, - Globe02Icon, Search01Icon, PowerIcon, PencilEdit02Icon, LayoutAlignLeftIcon, - HelpCircleIcon, Settings02Icon, + SourceCodeSquareIcon, ZapIcon, } from "@hugeicons/core-free-icons"; import { @@ -528,7 +527,7 @@ export function AppSidebar() {
{displayTitle} - Unsloth + Studio
@@ -549,8 +548,8 @@ export function AppSidebar() { useSettingsDialogStore.getState().openDialog("api-keys")} > - - API + + Developer New @@ -579,12 +578,6 @@ export function AppSidebar() { - useSettingsDialogStore.getState().openDialog("about")} - > - - Help - setShutdownOpen(true)}> Shutdown diff --git a/studio/frontend/src/config/env.ts b/studio/frontend/src/config/env.ts index 72bb3fa815..3839706d25 100644 --- a/studio/frontend/src/config/env.ts +++ b/studio/frontend/src/config/env.ts @@ -50,7 +50,7 @@ export async function fetchDeviceType(): Promise { if (res.ok) { const data = (await res.json()) as { device_type?: string; chat_only?: boolean }; const deviceType = data.device_type ?? detectLocalPlatform(); - const chatOnly = data.chat_only ?? deviceType === "mac"; + const chatOnly = data.chat_only ?? false; usePlatformStore.setState({ deviceType, chatOnly, fetched: true }); return deviceType; } diff --git a/studio/frontend/src/features/settings/settings-dialog.tsx b/studio/frontend/src/features/settings/settings-dialog.tsx index 63c0a9d388..376af06e9d 100644 --- a/studio/frontend/src/features/settings/settings-dialog.tsx +++ b/studio/frontend/src/features/settings/settings-dialog.tsx @@ -10,11 +10,11 @@ import { import { cn } from "@/lib/utils"; import { Cancel01Icon, - Globe02Icon, - HelpCircleIcon, Message01Icon, PaintBrush02Icon, Settings02Icon, + SourceCodeSquareIcon, + SparklesIcon, UserIcon, } from "@hugeicons/core-free-icons"; import { HugeiconsIcon } from "@hugeicons/react"; @@ -40,8 +40,8 @@ const TABS: TabDef[] = [ { id: "profile", label: "Profile", icon: UserIcon }, { id: "appearance", label: "Appearance", icon: PaintBrush02Icon }, { id: "chat", label: "Chat", icon: Message01Icon }, - { id: "api-keys", label: "API", icon: Globe02Icon, badge: "New" }, - { id: "about", label: "Help", icon: HelpCircleIcon }, + { id: "api-keys", label: "Developer", icon: SourceCodeSquareIcon, badge: "New" }, + { id: "about", label: "Help", icon: SparklesIcon }, ]; function renderTab(tab: SettingsTab) { diff --git a/studio/frontend/src/features/settings/tabs/api-keys-tab.tsx b/studio/frontend/src/features/settings/tabs/api-keys-tab.tsx index ac9ec40543..64c6f520c0 100644 --- a/studio/frontend/src/features/settings/tabs/api-keys-tab.tsx +++ b/studio/frontend/src/features/settings/tabs/api-keys-tab.tsx @@ -63,7 +63,7 @@ export function ApiKeysTab() { return (
-

API

+

Developer

Access Unsloth programmatically via the OpenAI-compatible API.{" "} s.deviceType); const isLora = store.trainingMethod !== "full"; const showVisionLora = store.isVisionModel && store.isDatasetImage === true; const [loraOpen, setLoraOpen] = useState(false); @@ -883,7 +885,11 @@ export function ParamsSection(): ReactElement { None Standard - Unsloth + {platformDeviceType === "mac" ? ( + MLX + ) : ( + Unsloth + )} diff --git a/studio/frontend/src/features/training/lib/model-defaults.ts b/studio/frontend/src/features/training/lib/model-defaults.ts index c40a1e2282..8bc9c4e064 100644 --- a/studio/frontend/src/features/training/lib/model-defaults.ts +++ b/studio/frontend/src/features/training/lib/model-defaults.ts @@ -3,6 +3,7 @@ import type { BackendModelConfig } from "../api/models-api"; import type { TrainingConfigState } from "../types/config"; +import { usePlatformStore } from "@/config/env"; type ModelDefaultsPatch = Partial< Pick< @@ -69,7 +70,13 @@ function toStringArray(value: unknown): string[] | undefined { function toGradientCheckpointing( value: unknown, ): TrainingConfigState["gradientCheckpointing"] | undefined { - if (value === "none" || value === "true" || value === "unsloth") return value; + if (value === "none" || value === "true" || value === "unsloth" || value === "mlx") { + // On Mac, map "unsloth" → "mlx" since Unsloth GC is GPU-only + if (usePlatformStore.getState().deviceType === "mac" && value === "unsloth") { + return "mlx"; + } + return value; + } return undefined; } diff --git a/studio/frontend/src/hooks/use-hf-model-search.ts b/studio/frontend/src/hooks/use-hf-model-search.ts index 77214f38d9..efe4d726be 100644 --- a/studio/frontend/src/hooks/use-hf-model-search.ts +++ b/studio/frontend/src/hooks/use-hf-model-search.ts @@ -6,6 +6,7 @@ import { listModels } from "@huggingface/hub"; import { type CachedResult, cachedModelInfo, primeCacheFromListing } from "@/lib/hf-cache"; import { useCallback, useMemo } from "react"; import { useHfPaginatedSearch } from "./use-hf-paginated-search"; +import { usePlatformStore } from "@/config/env"; export interface HfModelResult { id: string; @@ -16,7 +17,8 @@ export interface HfModelResult { isGguf: boolean; } -const EXCLUDED_TAGS = new Set([ +/** Tags to exclude on GPU (CUDA/ROCm) — MLX models won't load on GPU. */ +const EXCLUDED_TAGS_GPU = new Set([ "gptq", "awq", "exl2", @@ -28,6 +30,18 @@ const EXCLUDED_TAGS = new Set([ "ctranslate2", ]); +/** Tags to exclude on MLX (Mac) — GPU-only quant formats won't load on MLX. */ +const EXCLUDED_TAGS_MLX = new Set([ + "gptq", + "awq", + "exl2", + "onnx", + "openvino", + "coreml", + "tflite", + "ctranslate2", +]); + // Embedding / sentence-transformer models ship with onnx/openvino as additional // export formats — they should not be excluded by the tag check above. const EMBEDDING_TAGS = new Set([ @@ -77,7 +91,7 @@ function estimateSizeFromDtypes( return total > 0 ? total : undefined; } -function makeMapModel(excludeGguf: boolean) { +function makeMapModel(excludeGguf: boolean, excludedTags: Set) { return (raw: unknown): HfModelResult | null => { const m = raw as { name: string; @@ -87,7 +101,7 @@ function makeMapModel(excludeGguf: boolean) { tags?: string[]; }; const isEmbedding = m.tags?.some((t) => EMBEDDING_TAGS.has(t)); - if (!isEmbedding && m.tags?.some((t) => EXCLUDED_TAGS.has(t))) { + if (!isEmbedding && m.tags?.some((t) => excludedTags.has(t))) { return null; } const isGguf = @@ -314,7 +328,9 @@ export function useHfModelSearch( [trimmed, searchQuery, pinnedId, task, accessToken, priorityIds], ); - const mapModel = useMemo(() => makeMapModel(excludeGguf), [excludeGguf]); + const deviceType = usePlatformStore((s) => s.deviceType); + const excludedTags = deviceType === "mac" ? EXCLUDED_TAGS_MLX : EXCLUDED_TAGS_GPU; + const mapModel = useMemo(() => makeMapModel(excludeGguf, excludedTags), [excludeGguf, excludedTags]); const search = useHfPaginatedSearch(createIter, mapModel); // Secondary sort guarantee: unsloth models always float to the top. diff --git a/studio/frontend/src/types/training.ts b/studio/frontend/src/types/training.ts index d65d14fb83..187f54a13b 100644 --- a/studio/frontend/src/types/training.ts +++ b/studio/frontend/src/types/training.ts @@ -10,7 +10,7 @@ export function isAdapterMethod(method: TrainingMethod): boolean { export type StepNumber = 1 | 2 | 3 | 4 | 5; export type DatasetSource = "huggingface" | "upload"; export type DatasetFormat = "auto" | "alpaca" | "chatml" | "sharegpt"; -export type GradientCheckpointing = "none" | "true" | "unsloth"; +export type GradientCheckpointing = "none" | "true" | "unsloth" | "mlx"; export interface WizardState { currentStep: StepNumber; diff --git a/tests/python/test_gpu_init_ldconfig_guard.py b/tests/python/test_gpu_init_ldconfig_guard.py new file mode 100644 index 0000000000..081a6132b4 --- /dev/null +++ b/tests/python/test_gpu_init_ldconfig_guard.py @@ -0,0 +1,46 @@ +import ast +from pathlib import Path + +REPO_ROOT = Path(__file__).resolve().parents[2] +GPU_INIT = REPO_ROOT / "unsloth" / "_gpu_init.py" + + +def _find_geteuid_guard(tree: ast.AST): + for node in ast.walk(tree): + if not isinstance(node, ast.If): + continue + for sub in ast.walk(node.test): + if isinstance(sub, ast.Call) and isinstance(sub.func, ast.Attribute): + if sub.func.attr == "geteuid": + return node + return None + + +def test_gpu_init_has_geteuid_guard(): + tree = ast.parse(GPU_INIT.read_text()) + guard = _find_geteuid_guard(tree) + assert ( + guard is not None + ), "_gpu_init.py must guard ldconfig recovery on os.geteuid()" + + +def test_ldconfig_calls_only_inside_geteuid_guard(): + src = GPU_INIT.read_text() + tree = ast.parse(src) + guard = _find_geteuid_guard(tree) + assert guard is not None + guard_src = ast.get_source_segment(src, guard) or "" + ldconfig_lines = [ + line for line in src.splitlines() if "ldconfig" in line and "os.system" in line + ] + for line in ldconfig_lines: + assert line.strip() in guard_src, ( + "os.system('ldconfig ...') must live inside the geteuid guard, " + f"but found unguarded: {line!r}" + ) + + +def test_non_root_branch_warns_when_bnb_present(): + src = GPU_INIT.read_text() + assert "elif bnb is not None" in src + assert "sudo ldconfig" in src diff --git a/tests/studio/test_export_output_path_contract.py b/tests/studio/test_export_output_path_contract.py new file mode 100644 index 0000000000..e99fc42091 --- /dev/null +++ b/tests/studio/test_export_output_path_contract.py @@ -0,0 +1,122 @@ +# SPDX-License-Identifier: AGPL-3.0-only + +import ast +from pathlib import Path + +REPO_ROOT = Path(__file__).resolve().parents[2] +EXPORT = REPO_ROOT / "studio" / "backend" / "core" / "export" / "export.py" + +EXPORT_FNS = ( + "export_merged_model", + "export_base_model", + "export_gguf", + "export_lora_adapter", +) + + +def _find_method(tree, cls_name, method_name): + for cls in ast.walk(tree): + if isinstance(cls, ast.ClassDef) and cls.name == cls_name: + for item in cls.body: + if isinstance(item, ast.FunctionDef) and item.name == method_name: + return item + return None + + +def _return_tuple_arity(fn): + arities = [] + for node in ast.walk(fn): + if isinstance(node, ast.Return) and isinstance(node.value, ast.Tuple): + arities.append(len(node.value.elts)) + return arities + + +def test_export_methods_return_three_tuple_annotation(): + tree = ast.parse(EXPORT.read_text()) + for fn_name in EXPORT_FNS: + fn = _find_method(tree, "ExportBackend", fn_name) + assert fn is not None, f"missing ExportBackend.{fn_name}" + ret = fn.returns + assert isinstance(ret, ast.Subscript), f"{fn_name} return must be Tuple[...]" + slc = ret.slice + elts = slc.elts if isinstance(slc, ast.Tuple) else None + assert ( + elts is not None and len(elts) == 3 + ), f"{fn_name} return annotation must be a 3-tuple, got {ast.dump(ret)}" + + +def test_export_methods_return_three_element_tuples(): + tree = ast.parse(EXPORT.read_text()) + for fn_name in EXPORT_FNS: + fn = _find_method(tree, "ExportBackend", fn_name) + assert fn is not None + arities = _return_tuple_arity(fn) + assert arities, f"{fn_name} has no tuple-return statements" + for arity in arities: + assert arity == 3, f"{fn_name} return tuple arity {arity}, expected 3" + + +def test_local_save_assigns_output_path(): + tree = ast.parse(EXPORT.read_text()) + for fn_name in EXPORT_FNS: + fn = _find_method(tree, "ExportBackend", fn_name) + assert fn is not None + assigns = [] + for node in ast.walk(fn): + if isinstance(node, ast.Assign): + for tgt in node.targets: + if isinstance(tgt, ast.Name) and tgt.id == "output_path": + assigns.append(node) + non_none = [ + a + for a in assigns + if not (isinstance(a.value, ast.Constant) and a.value.value is None) + ] + assert non_none, f"{fn_name} never assigns a non-None output_path" + + +def test_gpu_save_method_bound_for_hub_only(): + tree = ast.parse(EXPORT.read_text()) + fn = _find_method(tree, "ExportBackend", "export_merged_model") + assert fn is not None + found_pre_save_method = False + for node in ast.walk(fn): + if isinstance(node, ast.Try): + for stmt in node.body: + if isinstance(stmt, ast.If): + test = stmt.test + if isinstance(test, ast.Name) and test.id == "_IS_MLX": + for sub in ast.walk( + ast.Module(body = stmt.orelse, type_ignores = []) + ): + if isinstance(sub, ast.Assign) and any( + isinstance(t, ast.Name) and t.id == "save_method" + for t in sub.targets + ): + found_pre_save_method = True + break + if found_pre_save_method: + break + if found_pre_save_method: + break + assert found_pre_save_method, ( + "GPU save_method must be assigned at the top of the try block, " + "before the `if save_directory:` guard, so Hub-only export does not " + "raise UnboundLocalError." + ) + + +def test_mlx_hub_only_uses_temp_directory(): + src = EXPORT.read_text() + assert ( + src.count("tempfile.TemporaryDirectory") >= 3 + ), "expected TemporaryDirectory in merged, base, and lora hub-push paths" + assert "import tempfile" in src.split("class ExportBackend")[0] + + +def test_is_mlx_imported_from_unsloth(): + src = EXPORT.read_text() + assert "from unsloth import" in src + head = src.split("class ExportBackend")[0] + assert "_IS_MLX" in head + assert "_IS_MLX = platform.system()" not in src diff --git a/tests/studio/test_is_mlx_dispatch_gate.py b/tests/studio/test_is_mlx_dispatch_gate.py new file mode 100644 index 0000000000..fc07a497e7 --- /dev/null +++ b/tests/studio/test_is_mlx_dispatch_gate.py @@ -0,0 +1,213 @@ +# SPDX-License-Identifier: AGPL-3.0-only + +""" +Regression tests for the CUDA-vs-MLX dispatch gates Studio relies on. + +Two gates drive every dispatch decision in Studio's MLX path: + + 1. ``unsloth._IS_MLX`` at the top of ``unsloth/__init__.py`` -- evaluated + once at import time and read by Studio worker code to choose between + the GPU and MLX trainer / inference / export paths. Defined as + ``Darwin AND arm64 AND find_spec("mlx") is not None``. + + 2. ``utils.hardware.detect_hardware()`` -- runtime probe in the Studio + backend. Priority order: CUDA -> XPU -> MLX -> CPU. The MLX branch is + reached only when both CUDA and XPU are unavailable AND the host is + Apple Silicon AND mlx is importable. + +These gates are the canaries for "MLX support accidentally hijacks +CUDA/AMD/Intel users". The tests here: + + * verify the source-level structure of the ``_IS_MLX`` expression so an + accidental rewrite (e.g. dropping the ``arm64`` check) is caught, + * exercise the runtime gate logic under a spoofed Darwin+arm64 platform + with a fake ``mlx`` module in ``sys.modules`` to confirm both gates + flip True together, + * confirm that on the actual Linux+CUDA test host both gates remain in + their CUDA-side state. + +No real MLX install is required; uses the same ``monkeypatch.setitem`` +fake-mlx pattern as ``test_mlx_inference_backend.py``. +""" + +import ast +import importlib +import sys +import types +from pathlib import Path + + +REPO_ROOT = Path(__file__).resolve().parents[2] +UNSLOTH_INIT = REPO_ROOT / "unsloth" / "__init__.py" + + +# --------------------------------------------------------------------------- +# 1. Source-level structure check on _IS_MLX (no platform dependencies). +# --------------------------------------------------------------------------- + + +def test_is_mlx_gate_uses_three_required_predicates(): + """The _IS_MLX assignment must AND together exactly the three checks + that Studio depends on: Darwin OS, arm64 machine, and an importable + mlx package. Dropping any one of them silently breaks dispatch. + """ + tree = ast.parse(UNSLOTH_INIT.read_text()) + + target = None + for node in ast.walk(tree): + if ( + isinstance(node, ast.Assign) + and len(node.targets) == 1 + and isinstance(node.targets[0], ast.Name) + and node.targets[0].id == "_IS_MLX" + ): + target = node.value + break + assert target is not None, "_IS_MLX assignment not found in unsloth/__init__.py" + assert isinstance(target, ast.BoolOp) and isinstance( + target.op, ast.And + ), "_IS_MLX must be a BoolOp(And) of platform + mlx checks" + + expr_src = ast.unparse(target) + assert ( + "platform.system()" in expr_src and "Darwin" in expr_src + ), "_IS_MLX must check platform.system() == 'Darwin'" + assert ( + "platform.machine()" in expr_src and "arm64" in expr_src + ), "_IS_MLX must check platform.machine() == 'arm64'" + assert ( + "find_spec" in expr_src and "'mlx'" in expr_src + ), "_IS_MLX must check importlib.util.find_spec('mlx')" + + +# --------------------------------------------------------------------------- +# 2. Runtime gate behavior with the platform spoofed to Apple Silicon and a +# fake mlx module in sys.modules. Re-evaluates the same expression +# rather than reloading unsloth (which would cascade-reload torch). +# --------------------------------------------------------------------------- + + +def _evaluate_is_mlx_gate(platform_module, importlib_util): + """Re-evaluate the _IS_MLX expression using injected dependencies. + + Mirrors the assignment in unsloth/__init__.py exactly. + """ + return ( + platform_module.system() == "Darwin" + and platform_module.machine() == "arm64" + and importlib_util.find_spec("mlx") is not None + ) + + +def test_is_mlx_gate_true_on_apple_silicon_with_mlx_present(monkeypatch): + import platform + import importlib.util + + # Inject a fake mlx package so find_spec returns a non-None ModuleSpec. + fake_mlx = types.ModuleType("mlx") + fake_mlx.__spec__ = importlib.machinery.ModuleSpec("mlx", loader = None) + fake_mlx.__path__ = [] + monkeypatch.setitem(sys.modules, "mlx", fake_mlx) + + monkeypatch.setattr(platform, "system", lambda: "Darwin") + monkeypatch.setattr(platform, "machine", lambda: "arm64") + + assert _evaluate_is_mlx_gate(platform, importlib.util) is True + + +def test_is_mlx_gate_false_when_mlx_missing(monkeypatch): + import platform + import importlib.util + + # Apple Silicon platform but no mlx package -> gate must be False. + monkeypatch.delitem(sys.modules, "mlx", raising = False) + monkeypatch.setattr(platform, "system", lambda: "Darwin") + monkeypatch.setattr(platform, "machine", lambda: "arm64") + + real_find_spec = importlib.util.find_spec + + def _no_mlx(name, *args, **kwargs): + if name == "mlx": + return None + return real_find_spec(name, *args, **kwargs) + + monkeypatch.setattr(importlib.util, "find_spec", _no_mlx) + + assert _evaluate_is_mlx_gate(platform, importlib.util) is False + + +def test_is_mlx_gate_false_on_non_apple_silicon(): + """On the real Linux+CUDA / AMD / Intel test host, the gate stays False.""" + import platform + import importlib.util + + if platform.system() == "Darwin" and platform.machine() == "arm64": + # On a Mac CI runner this assertion would not apply; skip there. + import pytest + + pytest.skip("Test host is Apple Silicon; CUDA-side canary doesn't apply.") + + assert _evaluate_is_mlx_gate(platform, importlib.util) is False + + +# --------------------------------------------------------------------------- +# 3. Studio's runtime detect_hardware() picks MLX only when CUDA + XPU are +# both unavailable AND the host is Apple Silicon AND mlx is importable. +# --------------------------------------------------------------------------- + + +def _import_studio_hardware(): + """Lazy import for the Studio hardware module, with the bare-imports + convention that Studio uses (studio/backend on sys.path). + """ + studio_backend = REPO_ROOT / "studio" / "backend" + if str(studio_backend) not in sys.path: + sys.path.insert(0, str(studio_backend)) + from utils.hardware import hardware as hw # type: ignore + + return hw + + +def test_detect_hardware_picks_mlx_when_only_apple_silicon_available(monkeypatch): + hw = _import_studio_hardware() + + # Force CUDA + XPU paths off so detect_hardware falls through to MLX. + import torch + + monkeypatch.setattr(torch.cuda, "is_available", lambda: False) + if hasattr(torch, "xpu"): + monkeypatch.setattr(torch.xpu, "is_available", lambda: False) + + # Spoof Apple Silicon and provide an importable mlx.core for _has_mlx(). + import platform + + monkeypatch.setattr(platform, "system", lambda: "Darwin") + monkeypatch.setattr(platform, "machine", lambda: "arm64") + + fake_mlx = types.ModuleType("mlx") + fake_mlx_core = types.ModuleType("mlx.core") + fake_mlx.core = fake_mlx_core + monkeypatch.setitem(sys.modules, "mlx", fake_mlx) + monkeypatch.setitem(sys.modules, "mlx.core", fake_mlx_core) + + detected = hw.detect_hardware() + assert detected == hw.DeviceType.MLX, f"expected MLX, got {detected!r}" + + +def test_detect_hardware_picks_cuda_on_real_host(): + """Canary: on a real CUDA host the MLX branch must NOT be taken even + if mlx happens to be importable. Protects CUDA/AMD/Intel users from + accidental MLX dispatch when MLX support is added. + """ + import torch + + if not torch.cuda.is_available(): + import pytest + + pytest.skip("No CUDA available on this host; canary not applicable.") + + hw = _import_studio_hardware() + detected = hw.detect_hardware() + assert ( + detected == hw.DeviceType.CUDA + ), f"CUDA host must dispatch to CUDA, got {detected!r}" diff --git a/tests/studio/test_mlx_training_worker_behaviors.py b/tests/studio/test_mlx_training_worker_behaviors.py new file mode 100644 index 0000000000..6c067ea00b --- /dev/null +++ b/tests/studio/test_mlx_training_worker_behaviors.py @@ -0,0 +1,90 @@ +# SPDX-License-Identifier: AGPL-3.0-only + +import ast +from pathlib import Path + +REPO_ROOT = Path(__file__).resolve().parents[2] +WORKER = REPO_ROOT / "studio" / "backend" / "core" / "training" / "worker.py" + + +def _find_func(tree, name): + for node in ast.walk(tree): + if isinstance(node, ast.FunctionDef) and node.name == name: + return node + return None + + +def test_run_mlx_training_passes_token_to_from_pretrained(): + tree = ast.parse(WORKER.read_text()) + fn = _find_func(tree, "_run_mlx_training") + assert fn is not None + found = False + for node in ast.walk(fn): + if ( + isinstance(node, ast.Call) + and isinstance(node.func, ast.Attribute) + and node.func.attr == "from_pretrained" + and isinstance(node.func.value, ast.Name) + and node.func.value.id == "FastMLXModel" + ): + kwarg_names = {kw.arg for kw in node.keywords if kw.arg} + assert ( + "token" in kwarg_names + ), f"FastMLXModel.from_pretrained must forward token=hf_token; got {kwarg_names!r}" + found = True + assert found, "FastMLXModel.from_pretrained call not found in _run_mlx_training" + + +def test_wandb_init_strips_secret_keys(): + src = WORKER.read_text() + assert "_wandb_sensitive" in src, "expected a sensitive-key set near wandb.init" + assert '"hf_token"' in src and '"wandb_token"' in src + assert ( + "config = dict(config)" not in src + ), "wandb.init received raw config dict; secrets would leak" + + +def test_local_dataset_loader_uses_load_dataset_path(): + src = WORKER.read_text() + assert "_resolve_local_files" in src + assert "_loader_for_files" in src + assert "data_files = all_files" in src or "data_files=all_files" in src + + +def test_send_aliases_status_message_to_message(): + src = WORKER.read_text() + assert 'kwargs["message"] = sm' in src or 'kwargs["message"]=sm' in src + + +def test_slice_uses_inclusive_end_and_handles_zero(): + src = WORKER.read_text() + assert "min(end + 1, len(ds))" in src or "min(end+1, len(ds))" in src + assert "slice_start if slice_start is not None else 0" in src + assert "slice_end if slice_end is not None else len(ds) - 1" in src + + +def test_poll_stop_returns_on_broken_pipe(): + src = WORKER.read_text() + assert "except (EOFError, OSError)" in src + lines = src.splitlines() + for i, line in enumerate(lines): + if "except (EOFError, OSError)" in line: + for j in range(i + 1, min(i + 6, len(lines))): + stripped = lines[j].strip() + if not stripped or stripped.startswith("#"): + continue + assert stripped.startswith( + "return" + ), f"expected return after EOFError/OSError, got {stripped!r}" + break + break + else: + raise AssertionError("EOFError/OSError handler not found in worker.py") + + +def test_unsloth_zoo_mlx_imports_have_friendly_error(): + src = WORKER.read_text() + assert "from unsloth_zoo.mlx_loader import FastMLXModel" in src + assert "from unsloth_zoo.mlx_trainer import" in src + assert "raise ImportError" in src + assert "install.sh" in src diff --git a/unsloth/__init__.py b/unsloth/__init__.py index 9db9ae0a32..9b620a5c76 100644 --- a/unsloth/__init__.py +++ b/unsloth/__init__.py @@ -12,348 +12,117 @@ # See the License for the specific language governing permissions and # limitations under the License. -import warnings, importlib, sys -from packaging.version import Version -import os, re, subprocess, inspect, functools -import numpy as np +import os, platform, importlib.util -# Log Unsloth is being used os.environ["UNSLOTH_IS_PRESENT"] = "1" -# Check if modules that need patching are already imported -critical_modules = ["trl", "transformers", "peft"] -already_imported = [mod for mod in critical_modules if mod in sys.modules] - -# Fix some issues before importing other packages -from .import_fixes import ( - fix_message_factory_issue, - check_fbgemm_gpu_version, - disable_broken_causal_conv1d, - disable_broken_vllm, - configure_amdgpu_asic_id_table_path, - torchvision_compatibility_check, - fix_diffusers_warnings, - fix_huggingface_hub, +# Detect Apple Silicon + MLX before any torch/numpy imports +_IS_MLX = ( + platform.system() == "Darwin" + and platform.machine() == "arm64" + and importlib.util.find_spec("mlx") is not None ) -# Configure libdrm ids table path early so ROCm can resolve AMD GPU names. -configure_amdgpu_asic_id_table_path() -disable_broken_causal_conv1d() -disable_broken_vllm() -fix_message_factory_issue() -check_fbgemm_gpu_version() -torchvision_compatibility_check() -fix_diffusers_warnings() -fix_huggingface_hub() -del configure_amdgpu_asic_id_table_path -del disable_broken_causal_conv1d -del disable_broken_vllm -del fix_message_factory_issue -del check_fbgemm_gpu_version -del torchvision_compatibility_check -del fix_diffusers_warnings -del fix_huggingface_hub - -# This check is critical because Unsloth optimizes these libraries by modifying -# their code at import time. If they're imported first, the original (slower, -# more memory-intensive) implementations will be used instead of Unsloth's -# optimized versions, potentially causing OOM errors or slower training. -if already_imported: - # stacklevel=2 makes warning point to user's import line rather than this library code, - # showing them exactly where to fix the import order in their script - warnings.warn( - f"WARNING: Unsloth should be imported before [{', '.join(already_imported)}] " - f"to ensure all optimizations are applied. Your code may run slower or encounter " - f"memory issues without these optimizations.\n\n" - f"Please restructure your imports with 'import unsloth' at the top of your file.", - stacklevel = 2, - ) -del already_imported, critical_modules - -# Unsloth currently does not work on multi GPU setups - sadly we are a 2 brother team so -# enabling it will require much more work, so we have to prioritize. Please understand! -# We do have a beta version, which you can contact us about! -# Thank you for your understanding and we appreciate it immensely! - -# Fixes https://github.com/unslothai/unsloth/issues/1266 -os.environ["PROTOCOL_BUFFERS_PYTHON_IMPLEMENTATION"] = "python" - -# [TODO] Check why some GPUs don't work -# "pinned_use_cuda_host_register:True,"\ -# "pinned_num_register_threads:8" - - -from importlib.metadata import version as importlib_version -from importlib.metadata import PackageNotFoundError - -# Check for unsloth_zoo -try: - unsloth_zoo_version = importlib_version("unsloth_zoo") - if Version(unsloth_zoo_version) < Version("2026.3.4"): - print( - "Unsloth: Please update Unsloth and Unsloth-Zoo to the latest version!\n" - "Do this via `pip install --upgrade --force-reinstall --no-cache-dir --no-deps unsloth unsloth_zoo`" - ) - # if os.environ.get("UNSLOTH_DISABLE_AUTO_UPDATES", "0") == "0": - # try: - # os.system("pip install --upgrade --no-cache-dir --no-deps unsloth_zoo") - # except: - # try: - # os.system("pip install --upgrade --no-cache-dir --no-deps --user unsloth_zoo") - # except: - # raise ImportError("Unsloth: Please update unsloth_zoo via `pip install --upgrade --no-cache-dir --no-deps unsloth_zoo`") - import unsloth_zoo -except PackageNotFoundError: - raise ImportError( - f"Unsloth: Please install unsloth_zoo via `pip install unsloth_zoo` then retry!" - ) -except: - raise -del PackageNotFoundError, importlib_version - -# Try importing PyTorch and check version -try: - import torch -except ModuleNotFoundError: - raise ImportError( - "Unsloth: Pytorch is not installed. Go to https://pytorch.org/.\n" - "We have some installation instructions on our Github page." - ) -except: - raise - -from unsloth_zoo.device_type import ( - is_hip, - get_device_type, - DEVICE_TYPE, - DEVICE_TYPE_TORCH, - DEVICE_COUNT, - ALLOW_PREQUANTIZED_MODELS, -) - -# Fix other issues -from .import_fixes import ( - fix_xformers_performance_issue, - fix_vllm_aimv2_issue, - check_vllm_torch_sm100_compatibility, - fix_vllm_guided_decoding_params, - fix_trl_vllm_ascend, - fix_vllm_pdl_blackwell, - fix_triton_compiled_kernel_missing_attrs, - patch_trunc_normal_precision_issue, - ignore_logger_messages, - patch_ipykernel_hf_xet, - patch_trackio, - patch_datasets, - patch_enable_input_require_grads, - fix_openenv_no_vllm, - patch_openspiel_env_async, - fix_executorch, - patch_vllm_for_notebooks, - patch_torchcodec_audio_decoder, - disable_torchcodec_if_broken, - disable_broken_wandb, - patch_peft_weight_converter_compatibility, -) - -fix_xformers_performance_issue() -fix_vllm_aimv2_issue() -# Check vLLM + torch < 2.9.0 + SM100 compatibility BEFORE importing vLLM -check_vllm_torch_sm100_compatibility() -fix_vllm_guided_decoding_params() -fix_trl_vllm_ascend() -fix_vllm_pdl_blackwell() -fix_triton_compiled_kernel_missing_attrs() -patch_trunc_normal_precision_issue() -ignore_logger_messages() -patch_ipykernel_hf_xet() -patch_trackio() -patch_datasets() -patch_enable_input_require_grads() -fix_openenv_no_vllm() -patch_openspiel_env_async() -fix_executorch() -patch_vllm_for_notebooks() -patch_torchcodec_audio_decoder() -disable_torchcodec_if_broken() -disable_broken_wandb() -patch_peft_weight_converter_compatibility() - -del fix_xformers_performance_issue -del fix_vllm_aimv2_issue -del check_vllm_torch_sm100_compatibility -del fix_vllm_guided_decoding_params -del fix_trl_vllm_ascend -del fix_vllm_pdl_blackwell -del fix_triton_compiled_kernel_missing_attrs -del patch_trunc_normal_precision_issue -del ignore_logger_messages -del patch_ipykernel_hf_xet -del patch_trackio -del patch_datasets -del patch_enable_input_require_grads -del fix_openenv_no_vllm -del patch_openspiel_env_async -del fix_executorch -del patch_vllm_for_notebooks -del patch_torchcodec_audio_decoder -del disable_torchcodec_if_broken -del disable_broken_wandb -del patch_peft_weight_converter_compatibility - -# Torch 2.4 has including_emulation -if DEVICE_TYPE == "cuda": - major_version, minor_version = torch.cuda.get_device_capability() - SUPPORTS_BFLOAT16 = major_version >= 8 - - old_is_bf16_supported = torch.cuda.is_bf16_supported - if "including_emulation" in str(inspect.signature(old_is_bf16_supported)): - - def is_bf16_supported(including_emulation = False): - return old_is_bf16_supported(including_emulation) - - torch.cuda.is_bf16_supported = is_bf16_supported - else: - - def is_bf16_supported(): - return SUPPORTS_BFLOAT16 - - torch.cuda.is_bf16_supported = is_bf16_supported - del major_version, minor_version -elif DEVICE_TYPE == "hip": - SUPPORTS_BFLOAT16 = torch.cuda.is_bf16_supported() -elif DEVICE_TYPE == "xpu": - # torch.xpu.is_bf16_supported() does not have including_emulation - # set SUPPORTS_BFLOAT16 as torch.xpu.is_bf16_supported() - SUPPORTS_BFLOAT16 = torch.xpu.is_bf16_supported() - -# For Gradio HF Spaces? -# if "SPACE_AUTHOR_NAME" not in os.environ and "SPACE_REPO_NAME" not in os.environ: -import triton - -if DEVICE_TYPE == "cuda": - libcuda_dirs = lambda: None - if Version(triton.__version__) >= Version("3.0.0"): - try: - from triton.backends.nvidia.driver import libcuda_dirs - except: - pass - else: - from triton.common.build import libcuda_dirs - - # Try loading bitsandbytes and triton +if _IS_MLX: try: - import bitsandbytes as bnb - except: - print( - "Unsloth: `bitsandbytes` is not installed - 4bit QLoRA unallowed, but 16bit and full finetuning works!" - ) - bnb = None + import unsloth_zoo + except ImportError as _e: + raise ImportError( + "Unsloth: MLX support requires `unsloth-zoo` with MLX modules. " + "Reinstall with `pip install unsloth-zoo` or rerun install.sh." + ) from _e + # The mlx_trainer / mlx_loader submodules ship with unsloth-zoo's MLX + # support. An older installed unsloth-zoo (e.g. from PyPI before the + # MLX release lands) will satisfy `import unsloth_zoo` but be missing + # these submodules. Surface the same friendly install hint instead of + # a raw ImportError on the submodule path. try: - cdequantize_blockwise_fp32 = bnb.functional.lib.cdequantize_blockwise_fp32 - libcuda_dirs() - except: - # Only run the ldconfig recovery when we can actually run - # ldconfig (root). On non-root environments (shared HPC, - # locked-down containers, CI runners, etc.) the recovery would - # shell out to `ldconfig` and fail with "Permission denied", - # which is especially noisy for users who don't even have - # bitsandbytes installed and are just doing 16bit/full - # finetuning. libcuda_dirs() is used by both triton and bnb, - # so we still run the recovery whenever we're root, regardless - # of whether bnb is installed. - if hasattr(os, "geteuid") and os.geteuid() == 0: - warnings.warn("Unsloth: Running `ldconfig /usr/lib64-nvidia` to link CUDA.") + from unsloth_zoo.mlx_trainer import MLXTrainer, MLXTrainingConfig + from unsloth_zoo.mlx_loader import FastMLXModel + except ImportError as _e: + raise ImportError( + "Unsloth: MLX support requires an unsloth-zoo build that includes " + "`unsloth_zoo.mlx_trainer` and `unsloth_zoo.mlx_loader`. Upgrade with " + "`pip install -U unsloth-zoo` or rerun install.sh." + ) from _e - if os.path.exists("/usr/lib64-nvidia"): - os.system("ldconfig /usr/lib64-nvidia") - elif os.path.exists("/usr/local"): - # Sometimes bitsandbytes cannot be linked properly in Runpod for example - possible_cudas = ( - subprocess.check_output(["ls", "-al", "/usr/local"]) - .decode("utf-8") - .split("\n") - ) - find_cuda = re.compile(r"[\s](cuda\-[\d\.]{2,})$") - possible_cudas = [find_cuda.search(x) for x in possible_cudas] - possible_cudas = [x.group(1) for x in possible_cudas if x is not None] + # Load raw_text helpers without executing dataprep/__init__.py, which + # imports synthetic.py -> torch and would defeat the torch-free MLX path. + from pathlib import Path as _Path - # Try linking cuda folder, or everything in local - if len(possible_cudas) == 0: - os.system("ldconfig /usr/local/") - else: - find_number = re.compile(r"([\d\.]{2,})") - latest_cuda = np.argsort( - [float(find_number.search(x).group(1)) for x in possible_cudas] - )[::-1][0] - latest_cuda = possible_cudas[latest_cuda] - os.system(f"ldconfig /usr/local/{latest_cuda}") - del find_number, latest_cuda - del possible_cudas, find_cuda + _raw_text_path = _Path(__file__).resolve().parent / "dataprep" / "raw_text.py" + _raw_text_spec = importlib.util.spec_from_file_location( + "unsloth._mlx_raw_text", _raw_text_path + ) + if _raw_text_spec is None or _raw_text_spec.loader is None: + raise ImportError("Unsloth: could not load MLX raw_text dataprep helpers.") + _raw_text = importlib.util.module_from_spec(_raw_text_spec) + _raw_text_spec.loader.exec_module(_raw_text) + RawTextDataLoader = _raw_text.RawTextDataLoader + TextPreprocessor = _raw_text.TextPreprocessor + del _raw_text, _raw_text_spec, _raw_text_path, _Path - if bnb is not None: - importlib.reload(bnb) - importlib.reload(triton) - try: - libcuda_dirs = lambda: None - if Version(triton.__version__) >= Version("3.0.0"): - try: - from triton.backends.nvidia.driver import libcuda_dirs - except: - pass - else: - from triton.common.build import libcuda_dirs - cdequantize_blockwise_fp32 = ( - bnb.functional.lib.cdequantize_blockwise_fp32 - ) - libcuda_dirs() - except: - warnings.warn( - "Unsloth: CUDA is not linked properly.\n" - "Try running `python -m bitsandbytes` then `python -m xformers.info`\n" - "We tried running `ldconfig /usr/lib64-nvidia` ourselves, but it didn't work.\n" - "You need to run in your terminal `sudo ldconfig /usr/lib64-nvidia` yourself, then import Unsloth.\n" - "Also try `sudo ldconfig /usr/local/cuda-xx.x` - find the latest cuda version.\n" - "Unsloth will still run for now, but maybe it might crash - let's hope it works!" - ) - elif bnb is not None: - # Non-root + bnb installed: we can't run ldconfig ourselves, - # but bnb is going to crash later when the user actually uses - # 4bit quantization - tell them how to fix it manually so - # they're not surprised by an opaque error down the road. - warnings.warn( - "Unsloth: CUDA is not linked properly.\n" - "You need to run in your terminal `sudo ldconfig /usr/lib64-nvidia` yourself, then import Unsloth.\n" - "Also try `sudo ldconfig /usr/local/cuda-xx.x` - find the latest cuda version.\n" - "Unsloth will still run for now, but maybe it might crash - let's hope it works!" + __version__ = unsloth_zoo.__version__ + DEVICE_TYPE = "mlx" + + class FastLanguageModel: + @staticmethod + def from_pretrained(*args, **kwargs): + return FastMLXModel.from_pretrained(*args, **kwargs) + + @staticmethod + def get_peft_model(*args, **kwargs): + return FastMLXModel.get_peft_model(*args, **kwargs) + + @staticmethod + def for_inference(*args, **kwargs): + return args[0] if args else None + + class FastVisionModel(FastLanguageModel): + @staticmethod + def from_pretrained(*args, **kwargs): + kwargs.setdefault("text_only", False) + return FastMLXModel.from_pretrained(*args, **kwargs) + + @staticmethod + def for_training(*args, **kwargs): + return args[0] if args else None + + FastTextModel = FastLanguageModel + FastModel = FastLanguageModel + + class FastSentenceTransformer: + @staticmethod + def from_pretrained(*args, **kwargs): + raise NotImplementedError( + "Unsloth: FastSentenceTransformer is not yet supported on MLX." ) - del libcuda_dirs -elif DEVICE_TYPE == "hip": - # NO-OP for rocm device - pass -elif DEVICE_TYPE == "xpu": - import bitsandbytes as bnb - # TODO: check triton for intel installed properly. - pass + @staticmethod + def get_peft_model(*args, **kwargs): + raise NotImplementedError( + "Unsloth: FastSentenceTransformer is not yet supported on MLX." + ) -from .models import * -from .models import __version__ -from .save import * -from .chat_templates import * -from .tokenizer_utils import * -from .trainer import * + def is_bfloat16_supported(): + try: + import mlx.core as mx -# Export dataprep utilities for CLI and downstream users -from .dataprep.raw_text import RawTextDataLoader, TextPreprocessor -from unsloth_zoo.rl_environments import ( - check_python_modules, - create_locked_down_function, - execute_with_time_limit, - Benchmarker, - is_port_open, - launch_openenv, -) + name = mx.device_info().get("device_name", "") or "" + return not name.startswith(("Apple M1", "Apple M2")) + except Exception: + return True -# Patch TRL trainers for backwards compatibility -_patch_trl_trainer() + is_bf16_supported = is_bfloat16_supported + + class UnslothVisionDataCollator: + def __init__(self, *args, **kwargs): + raise NotImplementedError( + "Unsloth: UnslothVisionDataCollator is not used on MLX. " + "Use the MLX trainer/data path instead." + ) + +else: + # GPU path: load everything from _gpu_init + from ._gpu_init import * + from ._gpu_init import __version__ diff --git a/unsloth/_gpu_init.py b/unsloth/_gpu_init.py new file mode 100644 index 0000000000..2fc4bfde3c --- /dev/null +++ b/unsloth/_gpu_init.py @@ -0,0 +1,346 @@ +# Copyright 2023-present Daniel Han-Chen & the Unsloth team. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +import warnings, importlib, sys +from packaging.version import Version +import os, re, subprocess, inspect, functools +import numpy as np + +# Log Unsloth is being used +os.environ["UNSLOTH_IS_PRESENT"] = "1" + +# Check if modules that need patching are already imported +critical_modules = ["trl", "transformers", "peft"] +already_imported = [mod for mod in critical_modules if mod in sys.modules] + +# Fix some issues before importing other packages +from .import_fixes import ( + fix_message_factory_issue, + check_fbgemm_gpu_version, + disable_broken_causal_conv1d, + disable_broken_vllm, + configure_amdgpu_asic_id_table_path, + torchvision_compatibility_check, + fix_diffusers_warnings, + fix_huggingface_hub, +) + +# Configure libdrm ids table path early so ROCm can resolve AMD GPU names. +configure_amdgpu_asic_id_table_path() +disable_broken_causal_conv1d() +disable_broken_vllm() +fix_message_factory_issue() +check_fbgemm_gpu_version() +torchvision_compatibility_check() +fix_diffusers_warnings() +fix_huggingface_hub() +del configure_amdgpu_asic_id_table_path +del disable_broken_causal_conv1d +del disable_broken_vllm +del fix_message_factory_issue +del check_fbgemm_gpu_version +del torchvision_compatibility_check +del fix_diffusers_warnings +del fix_huggingface_hub + +# This check is critical because Unsloth optimizes these libraries by modifying +# their code at import time. If they're imported first, the original (slower, +# more memory-intensive) implementations will be used instead of Unsloth's +# optimized versions, potentially causing OOM errors or slower training. +if already_imported: + # stacklevel=2 makes warning point to user's import line rather than this library code, + # showing them exactly where to fix the import order in their script + warnings.warn( + f"WARNING: Unsloth should be imported before [{', '.join(already_imported)}] " + f"to ensure all optimizations are applied. Your code may run slower or encounter " + f"memory issues without these optimizations.\n\n" + f"Please restructure your imports with 'import unsloth' at the top of your file.", + stacklevel = 2, + ) +del already_imported, critical_modules + +# Unsloth currently does not work on multi GPU setups - sadly we are a 2 brother team so +# enabling it will require much more work, so we have to prioritize. Please understand! +# We do have a beta version, which you can contact us about! +# Thank you for your understanding and we appreciate it immensely! + +# Fixes https://github.com/unslothai/unsloth/issues/1266 +os.environ["PROTOCOL_BUFFERS_PYTHON_IMPLEMENTATION"] = "python" + +# [TODO] Check why some GPUs don't work +# "pinned_use_cuda_host_register:True,"\ +# "pinned_num_register_threads:8" + + +from importlib.metadata import version as importlib_version +from importlib.metadata import PackageNotFoundError + +# Check for unsloth_zoo +try: + unsloth_zoo_version = importlib_version("unsloth_zoo") + if Version(unsloth_zoo_version) < Version("2026.3.4"): + print( + "Unsloth: Please update Unsloth and Unsloth-Zoo to the latest version!\n" + "Do this via `pip install --upgrade --force-reinstall --no-cache-dir --no-deps unsloth unsloth_zoo`" + ) + # if os.environ.get("UNSLOTH_DISABLE_AUTO_UPDATES", "0") == "0": + # try: + # os.system("pip install --upgrade --no-cache-dir --no-deps unsloth_zoo") + # except: + # try: + # os.system("pip install --upgrade --no-cache-dir --no-deps --user unsloth_zoo") + # except: + # raise ImportError("Unsloth: Please update unsloth_zoo via `pip install --upgrade --no-cache-dir --no-deps unsloth_zoo`") + import unsloth_zoo +except PackageNotFoundError: + raise ImportError( + f"Unsloth: Please install unsloth_zoo via `pip install unsloth_zoo` then retry!" + ) +except: + raise +del PackageNotFoundError, importlib_version + +# Try importing PyTorch and check version +try: + import torch +except ModuleNotFoundError: + raise ImportError( + "Unsloth: Pytorch is not installed. Go to https://pytorch.org/.\n" + "We have some installation instructions on our Github page." + ) +except: + raise + +from unsloth_zoo.device_type import ( + is_hip, + get_device_type, + DEVICE_TYPE, + DEVICE_TYPE_TORCH, + DEVICE_COUNT, + ALLOW_PREQUANTIZED_MODELS, +) + +# Fix other issues +from .import_fixes import ( + fix_xformers_performance_issue, + fix_vllm_aimv2_issue, + check_vllm_torch_sm100_compatibility, + fix_vllm_guided_decoding_params, + fix_vllm_pdl_blackwell, + fix_triton_compiled_kernel_missing_attrs, + patch_trunc_normal_precision_issue, + ignore_logger_messages, + patch_ipykernel_hf_xet, + patch_trackio, + patch_datasets, + patch_enable_input_require_grads, + fix_openenv_no_vllm, + patch_openspiel_env_async, + fix_executorch, + patch_vllm_for_notebooks, + patch_torchcodec_audio_decoder, + disable_torchcodec_if_broken, + disable_broken_wandb, + fix_trl_vllm_ascend, + patch_peft_weight_converter_compatibility, +) + +fix_xformers_performance_issue() +fix_vllm_aimv2_issue() +# Check vLLM + torch < 2.9.0 + SM100 compatibility BEFORE importing vLLM +check_vllm_torch_sm100_compatibility() +fix_vllm_guided_decoding_params() +fix_trl_vllm_ascend() +fix_vllm_pdl_blackwell() +fix_triton_compiled_kernel_missing_attrs() +patch_trunc_normal_precision_issue() +ignore_logger_messages() +patch_ipykernel_hf_xet() +patch_trackio() +patch_datasets() +patch_enable_input_require_grads() +fix_openenv_no_vllm() +patch_openspiel_env_async() +fix_executorch() +patch_vllm_for_notebooks() +patch_torchcodec_audio_decoder() +disable_torchcodec_if_broken() +disable_broken_wandb() +patch_peft_weight_converter_compatibility() + +del fix_xformers_performance_issue +del fix_vllm_aimv2_issue +del check_vllm_torch_sm100_compatibility +del fix_vllm_guided_decoding_params +del fix_trl_vllm_ascend +del fix_vllm_pdl_blackwell +del fix_triton_compiled_kernel_missing_attrs +del patch_trunc_normal_precision_issue +del ignore_logger_messages +del patch_ipykernel_hf_xet +del patch_trackio +del patch_datasets +del patch_enable_input_require_grads +del fix_openenv_no_vllm +del patch_openspiel_env_async +del fix_executorch +del patch_vllm_for_notebooks +del patch_torchcodec_audio_decoder +del disable_torchcodec_if_broken +del disable_broken_wandb +del patch_peft_weight_converter_compatibility + +# Torch 2.4 has including_emulation +if DEVICE_TYPE == "cuda": + major_version, minor_version = torch.cuda.get_device_capability() + SUPPORTS_BFLOAT16 = major_version >= 8 + + old_is_bf16_supported = torch.cuda.is_bf16_supported + if "including_emulation" in str(inspect.signature(old_is_bf16_supported)): + + def is_bf16_supported(including_emulation = False): + return old_is_bf16_supported(including_emulation) + + torch.cuda.is_bf16_supported = is_bf16_supported + else: + + def is_bf16_supported(): + return SUPPORTS_BFLOAT16 + + torch.cuda.is_bf16_supported = is_bf16_supported + del major_version, minor_version +elif DEVICE_TYPE == "hip": + SUPPORTS_BFLOAT16 = torch.cuda.is_bf16_supported() +elif DEVICE_TYPE == "xpu": + # torch.xpu.is_bf16_supported() does not have including_emulation + # set SUPPORTS_BFLOAT16 as torch.xpu.is_bf16_supported() + SUPPORTS_BFLOAT16 = torch.xpu.is_bf16_supported() + +# For Gradio HF Spaces? +# if "SPACE_AUTHOR_NAME" not in os.environ and "SPACE_REPO_NAME" not in os.environ: +import triton + +if DEVICE_TYPE == "cuda": + libcuda_dirs = lambda: None + if Version(triton.__version__) >= Version("3.0.0"): + try: + from triton.backends.nvidia.driver import libcuda_dirs + except: + pass + else: + from triton.common.build import libcuda_dirs + + # Try loading bitsandbytes and triton + try: + import bitsandbytes as bnb + except: + print( + "Unsloth: `bitsandbytes` is not installed - 4bit QLoRA unallowed, but 16bit and full finetuning works!" + ) + bnb = None + try: + cdequantize_blockwise_fp32 = bnb.functional.lib.cdequantize_blockwise_fp32 + libcuda_dirs() + except: + if hasattr(os, "geteuid") and os.geteuid() == 0: + warnings.warn("Unsloth: Running `ldconfig /usr/lib64-nvidia` to link CUDA.") + + if os.path.exists("/usr/lib64-nvidia"): + os.system("ldconfig /usr/lib64-nvidia") + elif os.path.exists("/usr/local"): + # Sometimes bitsandbytes cannot be linked properly in Runpod for example + possible_cudas = ( + subprocess.check_output(["ls", "-al", "/usr/local"]) + .decode("utf-8") + .split("\n") + ) + find_cuda = re.compile(r"[\s](cuda\-[\d\.]{2,})$") + possible_cudas = [find_cuda.search(x) for x in possible_cudas] + possible_cudas = [x.group(1) for x in possible_cudas if x is not None] + + # Try linking cuda folder, or everything in local + if len(possible_cudas) == 0: + os.system("ldconfig /usr/local/") + else: + find_number = re.compile(r"([\d\.]{2,})") + latest_cuda = np.argsort( + [float(find_number.search(x).group(1)) for x in possible_cudas] + )[::-1][0] + latest_cuda = possible_cudas[latest_cuda] + os.system(f"ldconfig /usr/local/{latest_cuda}") + del find_number, latest_cuda + del possible_cudas, find_cuda + + if bnb is not None: + importlib.reload(bnb) + importlib.reload(triton) + try: + libcuda_dirs = lambda: None + if Version(triton.__version__) >= Version("3.0.0"): + try: + from triton.backends.nvidia.driver import libcuda_dirs + except: + pass + else: + from triton.common.build import libcuda_dirs + cdequantize_blockwise_fp32 = ( + bnb.functional.lib.cdequantize_blockwise_fp32 + ) + libcuda_dirs() + except: + warnings.warn( + "Unsloth: CUDA is not linked properly.\n" + "Try running `python -m bitsandbytes` then `python -m xformers.info`\n" + "We tried running `ldconfig /usr/lib64-nvidia` ourselves, but it didn't work.\n" + "You need to run in your terminal `sudo ldconfig /usr/lib64-nvidia` yourself, then import Unsloth.\n" + "Also try `sudo ldconfig /usr/local/cuda-xx.x` - find the latest cuda version.\n" + "Unsloth will still run for now, but maybe it might crash - let's hope it works!" + ) + elif bnb is not None: + warnings.warn( + "Unsloth: CUDA is not linked properly.\n" + "You need to run in your terminal `sudo ldconfig /usr/lib64-nvidia` yourself, then import Unsloth.\n" + "Also try `sudo ldconfig /usr/local/cuda-xx.x` - find the latest cuda version.\n" + "Unsloth will still run for now, but maybe it might crash - let's hope it works!" + ) + del libcuda_dirs +elif DEVICE_TYPE == "hip": + # NO-OP for rocm device + pass +elif DEVICE_TYPE == "xpu": + import bitsandbytes as bnb + + # TODO: check triton for intel installed properly. + pass + +from .models import * +from .models import __version__ +from .save import * +from .chat_templates import * +from .tokenizer_utils import * +from .trainer import * + +# Export dataprep utilities for CLI and downstream users +from .dataprep.raw_text import RawTextDataLoader, TextPreprocessor +from unsloth_zoo.rl_environments import ( + check_python_modules, + create_locked_down_function, + execute_with_time_limit, + Benchmarker, + is_port_open, + launch_openenv, +) + +# Patch TRL trainers for backwards compatibility +_patch_trl_trainer() From e594c38f3f9ae9bfb9fb69c0d59d966939fef2cd Mon Sep 17 00:00:00 2001 From: Daniel Han Date: Wed, 6 May 2026 07:11:27 +0000 Subject: [PATCH 005/159] Fix circular import in trainer.py after _gpu_init refactor The MLX-aware unsloth/__init__.py loads the GPU surface via "from ._gpu_init import *" at line 127. _gpu_init.py:332 then runs "from .trainer import *", and trainer.py:26 was doing "from . import is_bfloat16_supported". That last import asks for the symbol on the `unsloth` package namespace, but `unsloth/__init__.py` is paused at line 127 and the star export from _gpu_init has not yet propagated. Result on a fresh editable install (./install.sh --local): ImportError: cannot import name 'is_bfloat16_supported' from partially initialized module 'unsloth' (most likely due to a circular import) (.../unsloth/__init__.py) PyPI 2026.5.2 is unaffected because that wheel ships the older inline __init__.py which defines is_bfloat16_supported before "from .trainer import *" runs. Fix: import directly from the concrete _utils module, matching the pattern already in models/loader.py and models/llama.py. This bypasses the partially-initialized unsloth namespace and resolves cleanly via unsloth.models._utils, which has already been fully loaded by the preceding "from .models import *" in _gpu_init.py:327. Verified end-to-end: ./install.sh --local on Linux+CUDA imports unsloth cleanly, FastLanguageModel.from_pretrained loads, 10-step LoRA training on Llama-3.2-1B runs (loss 1.76, 5.85 samples/s on B200), post-train generation works. --- unsloth/trainer.py | 7 ++++++- 1 file changed, 6 insertions(+), 1 deletion(-) diff --git a/unsloth/trainer.py b/unsloth/trainer.py index eea985e958..d26389424c 100644 --- a/unsloth/trainer.py +++ b/unsloth/trainer.py @@ -23,7 +23,12 @@ from functools import wraps import trl import inspect from trl import SFTTrainer -from . import is_bfloat16_supported +# Import via the concrete _utils module instead of the partially-initialized +# `unsloth` namespace. The MLX-aware `unsloth/__init__.py` loads the GPU +# surface via `from ._gpu_init import *`, which transitively triggers +# `from .trainer import *` BEFORE the star export has propagated back to +# `unsloth`, so `from . import is_bfloat16_supported` raises at import time. +from .models._utils import is_bfloat16_supported from unsloth.utils import ( configure_padding_free, configure_sample_packing, From 23550a8e3814ad52edabf0c5c328b2db3ebbc0b1 Mon Sep 17 00:00:00 2001 From: Daniel Han Date: Wed, 6 May 2026 07:22:56 +0000 Subject: [PATCH 006/159] Trim trainer.py import-fix comment to one line --- unsloth/trainer.py | 7 ++----- 1 file changed, 2 insertions(+), 5 deletions(-) diff --git a/unsloth/trainer.py b/unsloth/trainer.py index d26389424c..01b8822bd5 100644 --- a/unsloth/trainer.py +++ b/unsloth/trainer.py @@ -23,11 +23,8 @@ from functools import wraps import trl import inspect from trl import SFTTrainer -# Import via the concrete _utils module instead of the partially-initialized -# `unsloth` namespace. The MLX-aware `unsloth/__init__.py` loads the GPU -# surface via `from ._gpu_init import *`, which transitively triggers -# `from .trainer import *` BEFORE the star export has propagated back to -# `unsloth`, so `from . import is_bfloat16_supported` raises at import time. + +# why: bypass partially-initialised unsloth ns during _gpu_init load from .models._utils import is_bfloat16_supported from unsloth.utils import ( configure_padding_free, From 0c803242efe7a2777999123bd6df78b92906fe7d Mon Sep 17 00:00:00 2001 From: Avaya Aggarwal <119044997+OnePunchMonk@users.noreply.github.com> Date: Wed, 6 May 2026 15:08:35 +0530 Subject: [PATCH 007/159] feat(studio): add Continued Pretraining (CPT) as a training method (#4677) * feat(studio): add Continued Pretraining (CPT) support Implements CPT as a first-class training method in Unsloth Studio, resolving feature request #4565. Changes: - frontend/src/types/training.ts: add 'cpt' to TrainingMethod union - frontend/src/lib/vram.ts: add 'cpt' to VramTrainingMethod (fp16 footprint) - frontend/src/features/export/constants.ts: add CPT to METHOD_LABELS - frontend/src/features/training/api/mappers.ts: map 'cpt' -> 'Continued Pretraining', force packing=true and train_on_completions=false for CPT payloads - frontend/src/features/studio/sections/model-section.tsx: add 'Continued Pretraining' option (purple dot) to Method selector; update tooltip - frontend/src/features/onboarding/.../model-selection-step.tsx: add CPT to onboarding wizard method dropdown - backend/models/training.py: update training_type field description - backend/core/training/worker.py: detect is_cpt flag, force packing=True, train_on_completions=False, pass is_cpt to _train_worker - backend/core/training/trainer.py: _train_worker reads is_cpt kwarg, forces packing on, skips train_on_responses_only for raw-text pretraining CPT behaviour: - Full model weights (no LoRA adapters), same as Full Finetuning - Sequence packing always enabled for GPU efficiency - Trains on every token (no chat-format masking) - VRAM estimated at fp16 (2.0 bytes/param) * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * Update mappers.ts * Add CPT raw dataset support and UI fixes * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * Add missing training methods module * Handle invalid raw-text rows and expose raw in onboarding --------- Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com> Co-authored-by: Roland Tannous <115670425+rolandtannous@users.noreply.github.com> Co-authored-by: Etherll <61019402+Etherll@users.noreply.github.com> Co-authored-by: Etherll --- .gitignore | 4 +- studio/backend/core/training/trainer.py | 187 ++++++++++++++-- studio/backend/core/training/training.py | 6 +- studio/backend/core/training/worker.py | 74 ++++++- studio/backend/models/training.py | 14 +- studio/backend/routes/training.py | 1 + .../tests/test_training_raw_support.py | 183 ++++++++++++++++ .../backend/utils/datasets/dataset_utils.py | 30 +++ studio/backend/utils/datasets/raw_text.py | 142 ++++++++++++ studio/frontend/src/config/training.ts | 10 + .../frontend/src/features/export/constants.ts | 1 + .../components/steps/dataset-step.tsx | 1 + .../components/steps/model-selection-step.tsx | 1 + .../components/steps/summary-step.tsx | 6 +- .../studio/historical-training-view.tsx | 21 +- .../sections/dataset-preview-dialog.tsx | 10 +- .../studio/sections/dataset-section.tsx | 1 + .../studio/sections/model-section.tsx | 13 +- .../studio/sections/params-section.tsx | 68 +++++- .../studio/sections/progress-section.tsx | 4 +- .../src/features/training/api/mappers.ts | 23 +- .../training/hooks/use-training-actions.ts | 6 +- .../features/training/lib/training-methods.ts | 48 ++++ .../training/stores/training-config-store.ts | 207 +++++++++++++++--- .../src/features/training/types/api.ts | 2 + .../src/features/training/types/config.ts | 2 + studio/frontend/src/lib/vram.ts | 16 +- studio/frontend/src/types/training.ts | 6 +- 28 files changed, 986 insertions(+), 101 deletions(-) create mode 100644 studio/backend/tests/test_training_raw_support.py create mode 100644 studio/backend/utils/datasets/raw_text.py create mode 100644 studio/frontend/src/features/training/lib/training-methods.ts diff --git a/.gitignore b/.gitignore index 2f23f18d65..ae6770bc07 100644 --- a/.gitignore +++ b/.gitignore @@ -24,8 +24,8 @@ dist/ downloads/ eggs/ .eggs/ -lib/ -lib64/ +/lib/ +/lib64/ parts/ sdist/ var/ diff --git a/studio/backend/core/training/trainer.py b/studio/backend/core/training/trainer.py index fe8d277ac0..a3f063694f 100644 --- a/studio/backend/core/training/trainer.py +++ b/studio/backend/core/training/trainer.py @@ -62,6 +62,7 @@ from datasets import Dataset, load_dataset from utils.models import is_vision_model, detect_audio_type from utils.datasets import format_and_template_dataset from utils.datasets import MODEL_TO_TEMPLATE_MAPPER, TEMPLATE_TO_RESPONSES_MAPPER +from utils.datasets.raw_text import prepare_raw_text_dataset from utils.paths import ( ensure_dir, resolve_dataset_path, @@ -125,6 +126,7 @@ class UnslothTrainer: self.load_in_4bit = True # Track quantization mode for metadata # Model state tracking + self.is_cpt = False # Set to True for Continued Pretraining self.is_vlm = False self.is_audio = False self.is_audio_vlm = ( @@ -925,6 +927,7 @@ class UnslothTrainer: use_gradient_checkpointing: str = "unsloth", use_rslora: bool = False, use_loftq: bool = False, + modules_to_save: list = None, ) -> bool: """ Prepare model for training (with optional LoRA). @@ -1121,11 +1124,14 @@ class UnslothTrainer: loftq_config = {"loftq_bits": 4, "loftq_iter": 1} if use_loftq else None, + modules_to_save = modules_to_save, ) else: # Text model LoRA logger.info(f"Text model LoRA configuration:") logger.info(f" - Target modules: {target_modules}\n") + if modules_to_save: + logger.info(f" - Modules to save: {modules_to_save}\n") self.model = FastLanguageModel.get_peft_model( self.model, @@ -1140,6 +1146,7 @@ class UnslothTrainer: loftq_config = {"loftq_bits": 4, "loftq_iter": 1} if use_loftq else None, + modules_to_save = modules_to_save, ) # Check if stopped during LoRA preparation @@ -2342,6 +2349,7 @@ class UnslothTrainer: eval_steps: float = 0.00, dataset_slice_start: int = None, dataset_slice_end: int = None, + is_cpt: bool = False, ) -> Optional[tuple]: """ Load and prepare dataset for training. @@ -2360,6 +2368,35 @@ class UnslothTrainer: False # True if eval comes from a separate HF split ) eval_enabled = eval_steps is not None and eval_steps > 0 + raw_text_mode = is_cpt or format_type == "raw" + + def _raw_mode_label() -> str: + return "CPT" if is_cpt else "raw text" + + def _apply_raw_text_prep(ds: Dataset, split_name: str) -> Dataset: + try: + result = prepare_raw_text_dataset( + ds, + mode_label = _raw_mode_label(), + split_name = split_name, + eos_token = getattr(self.tokenizer, "eos_token", None), + append_eos = True, + ) + except ValueError as exc: + error_msg = str(exc) + logger.error(error_msg) + self._update_progress(error = error_msg) + raise + + for notice in result.notices: + if notice.level == "warning": + logger.warning(notice.message) + if notice.update_status: + self._update_progress(status_message = notice.message) + else: + logger.info(f"{notice.message}\n") + + return result.dataset if local_datasets: # Load local datasets using load_dataset() so the result is @@ -2534,6 +2571,48 @@ class UnslothTrainer: processed = self._preprocess_dac_dataset(dataset, custom_format_mapping) return ({"dataset": processed, "final_format": "audio_dac"}, None) + # ========== RAW TEXT BYPASS ========== + if raw_text_mode: + logger.info( + f"{_raw_mode_label().capitalize()} mode: bypassing chat template, " + "using raw text\n" + ) + dataset = _apply_raw_text_prep(dataset, "train") + if has_separate_eval_source and eval_dataset is not None: + eval_dataset = _apply_raw_text_prep(eval_dataset, "eval") + + dataset_info = { + "dataset": dataset, + "detected_format": "raw_text", + "final_format": "raw_text", + "success": True, + } + + if has_separate_eval_source and eval_dataset is not None: + logger.info( + f"{_raw_mode_label().capitalize()}: eval dataset " + f"({len(eval_dataset)} rows) kept as raw text\n" + ) + elif eval_enabled and not has_separate_eval_source: + split_result = self._resolve_eval_split_from_dataset(dataset) + if split_result is not None: + train_portion, eval_dataset = split_result + dataset_info["dataset"] = train_portion + + train_dataset = dataset_info["dataset"] + n = len(train_dataset) if hasattr(train_dataset, "__len__") else None + n_display = f"{n:,}" if isinstance(n, int) else "streaming" + self._update_progress( + status_message = f"Dataset ready ({n_display} samples, raw text)" + ) + logger.info(f"Raw-text dataset ready ({n_display} samples)\n") + + if "text" not in train_dataset.column_names: + raise ValueError( + f"Raw-text dataset missing 'text' column: {train_dataset.column_names}" + ) + return (dataset_info, eval_dataset) + elif self.is_audio_vlm: formatted = self._format_audio_vlm_dataset( dataset, custom_format_mapping @@ -2676,6 +2755,7 @@ class UnslothTrainer: output_dir: str | None = None, num_epochs: int = 3, learning_rate: float = 2e-4, + embedding_learning_rate: float | None = None, batch_size: int = 2, gradient_accumulation_steps: int = 4, warmup_steps: int = None, @@ -2728,6 +2808,7 @@ class UnslothTrainer: "output_dir": output_dir, "num_epochs": num_epochs, "learning_rate": learning_rate, + "embedding_learning_rate": embedding_learning_rate, "batch_size": batch_size, "gradient_accumulation_steps": gradient_accumulation_steps, "warmup_steps": warmup_steps, @@ -2945,6 +3026,13 @@ class UnslothTrainer: logger.info("Configuring data collator...\n") + dataset_final_format = ( + str(dataset.get("final_format", "")).lower() + if isinstance(dataset, dict) + else "" + ) + raw_text_mode = dataset_final_format == "raw_text" + data_collator = None # Default to built-in data collator if is_deepseek_ocr: # Special DeepSeek OCR collator - auto-install if needed @@ -2984,7 +3072,7 @@ class UnslothTrainer: self._update_progress(error = error_msg, is_training = False) return - elif self.is_audio_vlm: + elif self.is_audio_vlm and not raw_text_mode: # Audio VLM collator (e.g. Gemma 3N with audio data) # Mirrors the collate_fn from Gemma3N_(4B)-Audio notebook logger.info("Configuring audio VLM data collator...\n") @@ -3026,7 +3114,7 @@ class UnslothTrainer: data_collator = audio_vlm_collate_fn logger.info("Audio VLM data collator configured\n") - elif self.is_vlm: + elif self.is_vlm and not raw_text_mode: # Standard VLM collator (images) logger.info("Using UnslothVisionDataCollator for vision model\n") from unsloth.trainer import UnslothVisionDataCollator @@ -3137,8 +3225,9 @@ class UnslothTrainer: optim_value = training_args.get("optim", "adamw_8bit") lr_scheduler_type_value = training_args.get("lr_scheduler_type", "linear") - if self.is_vlm or self.is_audio_vlm: + if (self.is_vlm or self.is_audio_vlm) and not raw_text_mode: # Vision / audio VLM config (both need skip_prepare_dataset + remove_unused_columns) + # Raw-text runs on VLM-capable models are routed to the text path below. label = "audio VLM" if self.is_audio_vlm else "vision" logger.info(f"Configuring {label} model training parameters\n") # Use provided values or defaults for vision models @@ -3160,7 +3249,14 @@ class UnslothTrainer: } ) else: - logger.info("Configuring text model training parameters\n") + is_cpt = training_args.get("is_cpt", False) + self.is_cpt = is_cpt + if is_cpt: + logger.info("Configuring Continued Pretraining (CPT) parameters\n") + elif raw_text_mode: + logger.info("Configuring raw-text training parameters\n") + else: + logger.info("Configuring text model training parameters\n") config_args.update( { "optim": optim_value, @@ -3189,9 +3285,10 @@ class UnslothTrainer: logger.info("Training configuration prepared\n") # ========== TRAINER INITIALIZATION ========== - if self.is_audio_vlm: + if self.is_audio_vlm and not raw_text_mode: # Audio VLM (e.g. Gemma 3N + audio): raw Dataset from _format_audio_vlm_dataset # Notebook uses processing_class=processor.tokenizer (text tokenizer only) + # Raw-text runs are routed to the text path below. train_dataset = ( dataset if isinstance(dataset, Dataset) else dataset["dataset"] ) @@ -3210,8 +3307,9 @@ class UnslothTrainer: if eval_dataset is not None: trainer_kwargs["eval_dataset"] = eval_dataset self.trainer = SFTTrainer(**trainer_kwargs) - elif self.is_vlm: + elif self.is_vlm and not raw_text_mode: # Image VLM: dataset is dict wrapper from format_and_template_dataset + # Raw-text runs are routed to the text path below. train_dataset = ( dataset["dataset"] if isinstance(dataset, dict) else dataset ) @@ -3242,16 +3340,48 @@ class UnslothTrainer: ) sft_tokenizer = self.tokenizer.tokenizer - trainer_kwargs = { - "model": self.model, - "tokenizer": sft_tokenizer, - "train_dataset": dataset["dataset"], - "data_collator": data_collator, - "args": SFTConfig(**config_args), - } - if eval_dataset is not None: - trainer_kwargs["eval_dataset"] = eval_dataset - self.trainer = SFTTrainer(**trainer_kwargs) + if is_cpt: + try: + from unsloth import ( + UnslothTrainer as _UnslothCPTTrainer, + UnslothTrainingArguments as _UnslothTrainingArguments, + ) + except ImportError as exc: + raise RuntimeError( + "CPT requires a newer Unsloth install that exports " + "`UnslothTrainer` and `UnslothTrainingArguments` " + "(for embedding_learning_rate support). " + "Upgrade with: `pip install -U unsloth unsloth_zoo`." + ) from exc + + embedding_lr = training_args.get("embedding_learning_rate") + logger.info( + f"CPT: using UnslothTrainer with embedding_learning_rate={embedding_lr}\n" + ) + trainer_kwargs = { + "model": self.model, + "tokenizer": sft_tokenizer, + "train_dataset": dataset["dataset"], + "data_collator": data_collator, + "args": _UnslothTrainingArguments( + embedding_learning_rate = embedding_lr, + **config_args, + ), + } + if eval_dataset is not None: + trainer_kwargs["eval_dataset"] = eval_dataset + self.trainer = _UnslothCPTTrainer(**trainer_kwargs) + else: + trainer_kwargs = { + "model": self.model, + "tokenizer": sft_tokenizer, + "train_dataset": dataset["dataset"], + "data_collator": data_collator, + "args": SFTConfig(**config_args), + } + if eval_dataset is not None: + trainer_kwargs["eval_dataset"] = eval_dataset + self.trainer = SFTTrainer(**trainer_kwargs) # Restore the full processor as processing_class so checkpoint # saves include preprocessor_config.json (needed for GGUF export). if sft_tokenizer is not self.tokenizer: @@ -3260,19 +3390,32 @@ class UnslothTrainer: # ========== TRAIN ON RESPONSES ONLY ========== # Determine if we should train on responses only + # Raw-text datasets always train on all tokens. instruction_part = None response_part = None - train_on_responses_enabled = training_args.get( - "train_on_completions", False + is_cpt = training_args.get("is_cpt", False) + train_on_responses_enabled = ( + False + if (is_cpt or raw_text_mode) + else training_args.get("train_on_completions", False) ) + if is_cpt: + logger.info( + "CPT mode: skipping train_on_responses_only — training on all tokens\n" + ) + elif raw_text_mode: + logger.info( + "Raw-text mode: skipping train_on_responses_only — training on all tokens\n" + ) + # DeepSeek OCR handles this internally in its collator, so skip # Audio VLM handles label masking in its collator, so skip if ( train_on_responses_enabled and not self.is_audio_vlm and not self.is_audio - and not (is_deepseek_ocr or dataset["final_format"].lower() == "alpaca") + and not (is_deepseek_ocr or dataset_final_format == "alpaca") ): try: logger.info("Configuring train on responses only...\n") @@ -3318,7 +3461,7 @@ class UnslothTrainer: and response_part and not self.is_audio_vlm and not self.is_audio - and not (is_deepseek_ocr or dataset["final_format"].lower() == "alpaca") + and not (is_deepseek_ocr or dataset_final_format == "alpaca") ): try: from unsloth.chat_templates import train_on_responses_only @@ -3451,7 +3594,9 @@ class UnslothTrainer: config = json.load(f) # Determine the training method - if self.load_in_4bit: + if self.is_cpt: + method = "CPT" + elif self.load_in_4bit: method = "qlora" else: method = "lora" diff --git a/studio/backend/core/training/training.py b/studio/backend/core/training/training.py index a04ad5ef49..72b13c3225 100644 --- a/studio/backend/core/training/training.py +++ b/studio/backend/core/training/training.py @@ -159,6 +159,7 @@ class TrainingBackend: "is_embedding": kwargs.get("is_embedding", False), "num_epochs": kwargs.get("num_epochs", 3), "learning_rate": kwargs.get("learning_rate", "2e-4"), + "embedding_learning_rate": kwargs.get("embedding_learning_rate"), "batch_size": kwargs.get("batch_size", 2), "gradient_accumulation_steps": kwargs.get("gradient_accumulation_steps", 4), "warmup_steps": kwargs.get("warmup_steps"), @@ -195,8 +196,9 @@ class TrainingBackend: "gpu_ids": kwargs.get("gpu_ids"), } - # Derive load_in_4bit from training_type - if config["training_type"] != "LoRA/QLoRA": + # Full finetuning always runs in 16-bit. LoRA/QLoRA and CPT preserve the + # explicit request so 4-bit adapter/raw-text runs remain possible. + if config["training_type"] == "Full Finetuning": config["load_in_4bit"] = False # Spawn subprocess — use locals so state is untouched on failure diff --git a/studio/backend/core/training/worker.py b/studio/backend/core/training/worker.py index 9c017db6de..ef5cafb175 100644 --- a/studio/backend/core/training/worker.py +++ b/studio/backend/core/training/worker.py @@ -15,6 +15,7 @@ from __future__ import annotations import structlog from loggers import get_logger +import math import os import shutil import sys @@ -1208,6 +1209,8 @@ def run_training_process( # ── 4b. Load and format dataset (LLM helper may use VRAM briefly) ── _send_status(event_queue, "Loading and formatting dataset...") hf_dataset = config.get("hf_dataset", "") + training_type = config.get("training_type", "LoRA/QLoRA") + _is_cpt_for_dataset = training_type == "Continued Pretraining" dataset_result = trainer.load_and_format_dataset( dataset_source = hf_dataset if hf_dataset and hf_dataset.strip() else None, format_type = config.get("format_type", ""), @@ -1220,6 +1223,7 @@ def run_training_process( eval_steps = config.get("eval_steps", 0.00), dataset_slice_start = config.get("dataset_slice_start"), dataset_slice_end = config.get("dataset_slice_end"), + is_cpt = _is_cpt_for_dataset, ) if isinstance(dataset_result, tuple): @@ -1305,7 +1309,9 @@ def run_training_process( _tqdm_thread.start() training_type = config.get("training_type", "LoRA/QLoRA") - use_lora = training_type == "LoRA/QLoRA" + is_cpt = training_type == "Continued Pretraining" + use_lora = training_type in ("LoRA/QLoRA", "Continued Pretraining") + cpt_trains_embeddings = False # ── 4c. Load training model (uses VRAM — dataset already formatted) ── _send_status(event_queue, "Loading model...") @@ -1337,8 +1343,41 @@ def run_training_process( ) return - # ── 4d. Prepare model (LoRA or full finetuning) ── - if use_lora: + # ── 4d. Prepare model (LoRA, full finetuning, or CPT) ── + if is_cpt: + _send_status(event_queue, "Configuring LoRA for continued pretraining...") + # embed_tokens (if the user included it) goes to modules_to_save — + # trained full-precision at embedding_learning_rate. lm_head stays as + # a LoRA target for merge compatibility (see unsloth PR #4106). + _user_modules = config.get("target_modules") or [] + wants_embed = "embed_tokens" in _user_modules + cpt_trains_embeddings = wants_embed + cpt_target_modules = [m for m in _user_modules if m != "embed_tokens"] + if not cpt_target_modules: + cpt_target_modules = [ + "q_proj", + "k_proj", + "v_proj", + "o_proj", + "gate_proj", + "up_proj", + "down_proj", + "lm_head", + ] + success = trainer.prepare_model_for_training( + use_lora = True, + target_modules = cpt_target_modules, + modules_to_save = ["embed_tokens"] if wants_embed else None, + lora_r = config.get("lora_r", 128), + lora_alpha = config.get("lora_alpha", 32), + lora_dropout = config.get("lora_dropout", 0.0), + use_gradient_checkpointing = config.get( + "gradient_checkpointing", "unsloth" + ), + use_rslora = config.get("use_rslora", False), + use_loftq = config.get("use_loftq", False), + ) + elif use_lora: _send_status(event_queue, "Configuring LoRA adapters...") success = trainer.prepare_model_for_training( use_lora = True, @@ -1379,9 +1418,9 @@ def run_training_process( ) return - # Convert learning rate + lr_default = "5e-5" if is_cpt else "2e-4" try: - lr_value = float(config.get("learning_rate", "2e-4")) + lr_value = float(config.get("learning_rate", lr_default)) except ValueError: event_queue.put( { @@ -1393,6 +1432,25 @@ def run_training_process( ) return + # embedding_learning_rate is validated by the Pydantic model (Optional[float], + # gt=0, lt=1.0); if present it is already a finite float in range. + embedding_lr_value = config.get("embedding_learning_rate") + if is_cpt: + if cpt_trains_embeddings: + if embedding_lr_value is None: + # Default embedding_learning_rate = lr/10 per Unsloth's CPT notebook. + embedding_lr_value = lr_value / 10.0 + logger.info( + f"CPT: using default embedding_learning_rate={embedding_lr_value:.1e} " + f"(lr/10). Set explicitly to override.\n" + ) + elif embedding_lr_value is not None: + logger.warning( + "CPT: embedding_learning_rate was provided but embed_tokens is " + "not being trained; ignoring the override.\n" + ) + embedding_lr_value = None + # Generate output dir resume_from_checkpoint = config.get("resume_from_checkpoint") output_dir = config.get("output_dir") or _output_dir_from_resume_checkpoint( @@ -1425,6 +1483,7 @@ def run_training_process( output_dir = output_dir, num_epochs = config.get("num_epochs", 3), learning_rate = lr_value, + embedding_learning_rate = embedding_lr_value, batch_size = config.get("batch_size", 2), gradient_accumulation_steps = config.get("gradient_accumulation_steps", 4), warmup_steps = config.get("warmup_steps"), @@ -1434,7 +1493,9 @@ def run_training_process( weight_decay = config.get("weight_decay", 0.001), random_seed = config.get("random_seed", 3407), packing = config.get("packing", False), - train_on_completions = config.get("train_on_completions", False), + train_on_completions = False + if is_cpt + else config.get("train_on_completions", False), enable_wandb = config.get("enable_wandb", False), wandb_project = config.get("wandb_project", "unsloth-training"), wandb_token = config.get("wandb_token"), @@ -1445,6 +1506,7 @@ def run_training_process( max_seq_length = config.get("max_seq_length", 2048), optim = config.get("optim", "adamw_8bit"), lr_scheduler_type = config.get("lr_scheduler_type", "linear"), + is_cpt = is_cpt, resume_from_checkpoint = resume_from_checkpoint, ) diff --git a/studio/backend/models/training.py b/studio/backend/models/training.py index a9f4caa1bb..8127af1ee6 100644 --- a/studio/backend/models/training.py +++ b/studio/backend/models/training.py @@ -16,8 +16,11 @@ class TrainingStartRequest(BaseModel): model_name: str = Field( ..., description = "Model identifier (e.g., 'unsloth/llama-3-8b-bnb-4bit')" ) - training_type: str = Field( - ..., description = "Training type: 'LoRA/QLoRA' or 'Full Finetuning'" + training_type: Literal["LoRA/QLoRA", "Full Finetuning", "Continued Pretraining"] = ( + Field( + ..., + description = "Training type: 'LoRA/QLoRA', 'Full Finetuning', or 'Continued Pretraining'", + ) ) hf_token: Optional[str] = Field(None, description = "HuggingFace token") load_in_4bit: bool = Field(True, description = "Load model in 4-bit quantization") @@ -86,6 +89,13 @@ class TrainingStartRequest(BaseModel): packing: bool = Field(False, description = "Enable sequence packing") optim: str = Field("adamw_8bit", description = "Optimizer") lr_scheduler_type: str = Field("linear", description = "Learning rate scheduler type") + embedding_learning_rate: Optional[float] = Field( + None, + gt = 0, + lt = 1.0, + description = "Separate learning rate for embedding matrices (CPT). " + "Must be in (0, 1). Should be 2-10x smaller than the main learning rate.", + ) # LoRA parameters use_lora: bool = Field(True, description = "Use LoRA (derived from training_type)") diff --git a/studio/backend/routes/training.py b/studio/backend/routes/training.py index e5195bb337..19202f3883 100644 --- a/studio/backend/routes/training.py +++ b/studio/backend/routes/training.py @@ -207,6 +207,7 @@ async def start_training( "custom_format_mapping": request.custom_format_mapping, "num_epochs": request.num_epochs, "learning_rate": request.learning_rate, + "embedding_learning_rate": request.embedding_learning_rate, "batch_size": request.batch_size, "gradient_accumulation_steps": request.gradient_accumulation_steps, "warmup_steps": request.warmup_steps, diff --git a/studio/backend/tests/test_training_raw_support.py b/studio/backend/tests/test_training_raw_support.py new file mode 100644 index 0000000000..876ee34686 --- /dev/null +++ b/studio/backend/tests/test_training_raw_support.py @@ -0,0 +1,183 @@ +# 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 asyncio +import importlib.util +import unittest +from pathlib import Path +from unittest.mock import patch + +from datasets import Dataset + +from core.training.training import TrainingBackend +from models.training import TrainingStartRequest +from utils.datasets import format_dataset, format_and_template_dataset +from utils.datasets.raw_text import prepare_raw_text_dataset + +_BACKEND_ROOT = Path(__file__).resolve().parent.parent + + +def _load_route_module(name: str, relative_path: str): + spec = importlib.util.spec_from_file_location(name, _BACKEND_ROOT / relative_path) + module = importlib.util.module_from_spec(spec) + spec.loader.exec_module(module) + return module + + +class TestTrainingRawSupport(unittest.TestCase): + def test_training_backend_preserves_cpt_4bit_and_embedding_lr(self): + backend = TrainingBackend() + + class DummyProcess: + pid = 12345 + + def start(self): + return None + + class DummyThread: + def start(self): + return None + + dummy_queue = object() + + with ( + patch( + "core.training.training.prepare_gpu_selection", + return_value = ([0], {"selection_mode": "auto"}), + ), + patch( + "core.training.training._CTX.Queue", + side_effect = [dummy_queue, dummy_queue], + ), + patch( + "core.training.training._CTX.Process", return_value = DummyProcess() + ) as mock_process, + patch( + "core.training.training.threading.Thread", + return_value = DummyThread(), + ), + ): + backend.start_training( + job_id = "test-cpt-raw", + model_name = "unsloth/test-bnb-4bit", + training_type = "Continued Pretraining", + format_type = "raw", + load_in_4bit = True, + embedding_learning_rate = 1e-5, + ) + + config = mock_process.call_args.kwargs["kwargs"]["config"] + self.assertTrue(config["load_in_4bit"]) + self.assertEqual(config["embedding_learning_rate"], 1e-5) + + def test_training_route_forwards_embedding_learning_rate(self): + training_route = _load_route_module( + "training_route_module_raw_support", + "routes/training.py", + ) + captured: dict = {} + + class DummyBackend: + current_job_id = None + + def is_training_active(self): + return False + + def start_training(self, **kwargs): + captured.update(kwargs) + return True + + request = TrainingStartRequest( + model_name = "unsloth/test-bnb-4bit", + training_type = "Continued Pretraining", + format_type = "raw", + load_in_4bit = True, + embedding_learning_rate = 1e-5, + ) + + with ( + patch.object( + training_route, + "get_training_backend", + return_value = DummyBackend(), + ), + patch.object(training_route, "load_model_defaults", return_value = {}), + patch( + "core.inference.get_inference_backend", + return_value = type( + "InferenceBackend", + (), + {"active_model_name": None}, + )(), + ), + patch( + "core.export.get_export_backend", + return_value = type( + "ExportBackend", + (), + {"current_checkpoint": None}, + )(), + ), + ): + response = asyncio.run( + training_route.start_training(request, current_subject = "test-user") + ) + + self.assertEqual(response.status, "queued") + self.assertEqual(captured["embedding_learning_rate"], 1e-5) + self.assertTrue(captured["load_in_4bit"]) + + def test_format_dataset_supports_raw_text(self): + dataset = Dataset.from_dict( + { + "body": ["hello", "world"], + "title": ["a", "b"], + "id": [1, 2], + } + ) + + result = format_dataset(dataset, format_type = "raw") + + self.assertEqual(result["final_format"], "raw_text") + self.assertIn("text", result["dataset"].column_names) + self.assertEqual(result["dataset"][0]["text"], "hello") + self.assertFalse(result["requires_manual_mapping"]) + + def test_format_and_template_dataset_supports_raw_text_without_template(self): + dataset = Dataset.from_dict({"body": ["hello raw world"]}) + + result = format_and_template_dataset( + dataset, + model_name = "unsloth/test", + tokenizer = None, + format_type = "raw", + ) + + self.assertTrue(result["success"]) + self.assertEqual(result["final_format"], "raw_text") + self.assertEqual(result["dataset"][0]["text"], "hello raw world") + + def test_prepare_raw_text_dataset_drops_null_rows_before_appending_eos(self): + dataset = Dataset.from_dict({"text": ["hello", None, "world"]}) + + result = prepare_raw_text_dataset( + dataset, + mode_label = "CPT", + split_name = "train", + eos_token = "", + append_eos = True, + ) + + self.assertEqual(len(result.dataset), 2) + self.assertEqual(result.dataset[0]["text"], "hello") + self.assertEqual(result.dataset[1]["text"], "world") + self.assertTrue( + any( + "null or non-string 'text' values" in notice.message + for notice in result.notices + ) + ) + + +if __name__ == "__main__": + unittest.main() diff --git a/studio/backend/utils/datasets/dataset_utils.py b/studio/backend/utils/datasets/dataset_utils.py index fac8c3d295..26378d64ee 100644 --- a/studio/backend/utils/datasets/dataset_utils.py +++ b/studio/backend/utils/datasets/dataset_utils.py @@ -41,6 +41,7 @@ from .chat_templates import ( get_tokenizer_chat_template, DEFAULT_ALPACA_TEMPLATE, ) +from .raw_text import prepare_raw_text_dataset from .vlm_processing import generate_smart_vlm_instruction from .data_collators import DeepSeekOCRDataCollator, VLMDataCollator from .model_mappings import TEMPLATE_TO_MODEL_MAPPER @@ -437,6 +438,20 @@ def format_dataset( # Detect multimodal first (needed for all flows) multimodal_info = detect_multimodal_dataset(dataset) + if format_type == "raw": + raw_result = prepare_raw_text_dataset(dataset) + return { + "dataset": raw_result.dataset, + "detected_format": "raw_text", + "final_format": "raw_text", + "chat_column": "text", + "is_standardized": True, + "requires_manual_mapping": False, + "is_image": multimodal_info["is_image"], + "multimodal_info": multimodal_info, + "warnings": [notice.message for notice in raw_result.notices], + } + # If user provided explicit mapping, skip detection and apply in the requested format if custom_format_mapping: try: @@ -1105,6 +1120,21 @@ def format_and_template_dataset( num_proc = num_proc, ) + if dataset_info["final_format"] == "raw_text": + summary = get_dataset_info_summary(dataset_info) + return { + "dataset": dataset_info["dataset"], + "detected_format": dataset_info["detected_format"], + "final_format": dataset_info["final_format"], + "chat_column": dataset_info.get("chat_column"), + "is_vlm": False, + "success": True, + "requires_manual_mapping": False, + "warnings": dataset_info.get("warnings", []), + "errors": [], + "summary": summary, + } + # Step 2: Apply chat template detected = dataset_info.get("detected_format", "unknown") if progress_callback and n_rows: diff --git a/studio/backend/utils/datasets/raw_text.py b/studio/backend/utils/datasets/raw_text.py new file mode 100644 index 0000000000..353145fd5a --- /dev/null +++ b/studio/backend/utils/datasets/raw_text.py @@ -0,0 +1,142 @@ +# SPDX-License-Identifier: AGPL-3.0-only +# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0 + +""" +Shared helpers for raw-text dataset preparation. +""" + +from dataclasses import dataclass +from typing import Literal + +from datasets import Dataset + + +@dataclass(frozen = True) +class RawTextNotice: + message: str + level: Literal["info", "warning"] + update_status: bool = False + + +@dataclass(frozen = True) +class RawTextPreparationResult: + dataset: Dataset + notices: list[RawTextNotice] + + +def _string_columns(dataset: Dataset) -> list[str]: + feature_map = getattr(dataset, "features", {}) or {} + string_cols: list[str] = [] + for col in dataset.column_names: + feature = feature_map.get(col) + dtype = str(getattr(feature, "dtype", "")) + if dtype in {"string", "large_string"}: + string_cols.append(col) + return string_cols + + +def _split_scope(split_name: str | None) -> str: + return f"the {split_name} split" if split_name else "this dataset" + + +def _drop_invalid_text_rows( + dataset: Dataset, + *, + mode_title: str, + split_scope: str, +) -> tuple[Dataset, list[RawTextNotice]]: + filtered_dataset = dataset.filter(lambda ex: isinstance(ex["text"], str)) + dropped_rows = len(dataset) - len(filtered_dataset) + if not dropped_rows: + return filtered_dataset, [] + + if len(filtered_dataset) == 0: + raise ValueError( + f"{mode_title} training requires at least one string 'text' value " + f"in {split_scope}; all {dropped_rows} rows were null or non-string." + ) + + return filtered_dataset, [ + RawTextNotice( + message = ( + f"{mode_title}: dropped {dropped_rows:,} row(s) with null or " + f"non-string 'text' values from {split_scope}" + ), + level = "warning", + update_status = True, + ) + ] + + +def prepare_raw_text_dataset( + dataset: Dataset, + *, + mode_label: str = "raw text", + split_name: str | None = None, + eos_token: str | None = None, + append_eos: bool = False, +) -> RawTextPreparationResult: + notices: list[RawTextNotice] = [] + mode_title = mode_label.capitalize() + split_scope = _split_scope(split_name) + + if "text" not in dataset.column_names: + string_cols = _string_columns(dataset) + if not string_cols: + raise ValueError( + f"{mode_title} training requires a string 'text' column but none " + f"was found in {split_scope} (columns: {dataset.column_names})." + ) + + renamed_col = string_cols[0] + if len(string_cols) > 1: + notices.append( + RawTextNotice( + message = ( + f"{mode_title}: dataset has {len(string_cols)} string " + f"columns ({string_cols}); auto-selecting '{renamed_col}' " + "as the training text. Rename the intended column to " + "'text' to override." + ), + level = "warning", + update_status = True, + ) + ) + notices.append( + RawTextNotice( + message = ( + f"{mode_title}: renaming column '{renamed_col}' -> 'text' " + f"for {split_scope}" + ), + level = "info", + ) + ) + dataset = dataset.rename_column(renamed_col, "text") + + dataset, invalid_row_notices = _drop_invalid_text_rows( + dataset, + mode_title = mode_title, + split_scope = split_scope, + ) + notices.extend(invalid_row_notices) + + if append_eos: + if not eos_token: + notices.append( + RawTextNotice( + message = ( + f"{mode_title}: tokenizer has no eos_token; skipping EOS " + "append. Model will not learn document boundaries." + ), + level = "warning", + ) + ) + else: + + def _append_eos(ex, _eos = eos_token): + text = ex["text"] + return {"text": text if text.endswith(_eos) else text + _eos} + + dataset = dataset.map(_append_eos) + + return RawTextPreparationResult(dataset = dataset, notices = notices) diff --git a/studio/frontend/src/config/training.ts b/studio/frontend/src/config/training.ts index 913d612838..e9fe1d679c 100644 --- a/studio/frontend/src/config/training.ts +++ b/studio/frontend/src/config/training.ts @@ -76,6 +76,13 @@ export const TARGET_MODULES = [ "down_proj", ]; +/** CPT requires embed_tokens and lm_head in addition to standard LoRA modules. */ +export const CPT_TARGET_MODULES = [ + ...TARGET_MODULES, + "embed_tokens", + "lm_head", +]; + export const OPTIMIZER_OPTIONS: ReadonlyArray<{ value: string; label: string }> = [ { value: "adamw_8bit", label: "AdamW 8-bit" }, { value: "paged_adamw_8bit", label: "Paged AdamW 8-bit" }, @@ -96,11 +103,14 @@ export const LR_SCHEDULER_OPTIONS: ReadonlyArray<{ value: string; label: string */ export const LR_DEFAULT_LORA = 2e-4; export const LR_DEFAULT_FULL = 2e-5; +export const LR_DEFAULT_CPT = 5e-5; export const DEFAULT_HYPERPARAMS = { epochs: 3, contextLength: 2048, learningRate: LR_DEFAULT_LORA, + // null = let backend auto-compute (lr/10 per Unsloth CPT recipe). Only used by CPT. + embeddingLearningRate: null as number | null, optimizerType: "adamw_8bit", lrSchedulerType: "linear", loraRank: 16, diff --git a/studio/frontend/src/features/export/constants.ts b/studio/frontend/src/features/export/constants.ts index e9c3b8c95b..c97d9c1f6e 100644 --- a/studio/frontend/src/features/export/constants.ts +++ b/studio/frontend/src/features/export/constants.ts @@ -74,6 +74,7 @@ export const METHOD_LABELS: Record = { qlora: "QLoRA", lora: "LoRA", full: "Full Fine-tune", + cpt: "Continued Pretraining", }; export const GUIDE_STEPS = [ diff --git a/studio/frontend/src/features/onboarding/components/steps/dataset-step.tsx b/studio/frontend/src/features/onboarding/components/steps/dataset-step.tsx index 9b0456f67a..2945654d68 100644 --- a/studio/frontend/src/features/onboarding/components/steps/dataset-step.tsx +++ b/studio/frontend/src/features/onboarding/components/steps/dataset-step.tsx @@ -63,6 +63,7 @@ const FORMAT_OPTIONS: { value: DatasetFormat; label: string }[] = [ { value: "alpaca", label: "Alpaca" }, { value: "chatml", label: "ChatML" }, { value: "sharegpt", label: "ShareGPT" }, + { value: "raw", label: "Raw Text" }, ]; export function DatasetStep() { diff --git a/studio/frontend/src/features/onboarding/components/steps/model-selection-step.tsx b/studio/frontend/src/features/onboarding/components/steps/model-selection-step.tsx index 1ff23cb0fc..52f61700e2 100644 --- a/studio/frontend/src/features/onboarding/components/steps/model-selection-step.tsx +++ b/studio/frontend/src/features/onboarding/components/steps/model-selection-step.tsx @@ -366,6 +366,7 @@ export function ModelSelectionStep() { QLoRA (4-bit) LoRA (16-bit) Full Fine-tune + Continued Pretraining

diff --git a/studio/frontend/src/features/onboarding/components/steps/summary-step.tsx b/studio/frontend/src/features/onboarding/components/steps/summary-step.tsx index 8840983574..1988cfe970 100644 --- a/studio/frontend/src/features/onboarding/components/steps/summary-step.tsx +++ b/studio/frontend/src/features/onboarding/components/steps/summary-step.tsx @@ -5,6 +5,7 @@ import { Badge } from "@/components/ui/badge"; import { Card, CardContent, CardHeader, CardTitle } from "@/components/ui/card"; import { Separator } from "@/components/ui/separator"; import { useTrainingConfigStore } from "@/features/training"; +import { getTrainingMethodLabel } from "@/features/training/lib/training-methods"; import { useHardwareInfo } from "@/hooks"; import { isAdapterMethod } from "@/types/training"; import { ChipIcon, Database02Icon, GpuIcon, Settings04Icon } from "@hugeicons/core-free-icons"; @@ -102,6 +103,7 @@ export function SummaryStep() { const showLoraParams = isAdapterMethod(trainingMethod); const datasetName = datasetSource === "upload" ? uploadedFile : dataset; + const trainingMethodLabel = getTrainingMethodLabel(trainingMethod); return (
@@ -150,7 +152,7 @@ export function SummaryStep() {
- +
@@ -199,7 +201,7 @@ export function SummaryStep() {
Training - {trainingMethod === "qlora" ? "QLoRA" : trainingMethod === "lora" ? "LoRA" : "Full"} + {trainingMethodLabel}
diff --git a/studio/frontend/src/features/studio/historical-training-view.tsx b/studio/frontend/src/features/studio/historical-training-view.tsx index e461fc5a90..d13ad12727 100644 --- a/studio/frontend/src/features/studio/historical-training-view.tsx +++ b/studio/frontend/src/features/studio/historical-training-view.tsx @@ -4,6 +4,7 @@ import type { TrainingViewData } from "@/features/training"; import { getTrainingRun } from "@/features/training"; import type { TrainingRunDetailResponse } from "@/features/training"; +import { parseBackendTrainingMethod } from "@/features/training/lib/training-methods"; import { type ReactElement, useEffect, useState } from "react"; import { ChartsSection } from "./sections/charts-section"; import { ProgressSection } from "./sections/progress-section"; @@ -12,15 +13,6 @@ interface HistoricalTrainingViewProps { runId: string; } -function normalizeTrainingMethod(config: Record): string { - const type = config?.training_type as string | undefined; - if (!type || type === "Full Finetuning") return "full"; - if (type === "LoRA/QLoRA") { - return config?.load_in_4bit ? "qlora" : "lora"; - } - return "full"; -} - function mapToViewData(detail: TrainingRunDetailResponse): TrainingViewData { const { run, metrics } = detail; @@ -79,7 +71,10 @@ function mapToViewData(detail: TrainingRunDetailResponse): TrainingViewData { error: run.status === "error" ? run.error_message : null, isTrainingRunning: false, modelName: run.model_name, - trainingMethod: normalizeTrainingMethod(detail.config), + trainingMethod: parseBackendTrainingMethod( + detail.config?.training_type, + detail.config?.load_in_4bit, + ), lossHistory, lrHistory, gradNormHistory, @@ -143,7 +138,11 @@ export function HistoricalTrainingView({ loraRank: detail.config.lora_r as number | undefined, loraAlpha: detail.config.lora_alpha as number | undefined, loraDropout: detail.config.lora_dropout as number | undefined, - loraVariant: detail.config.use_rslora ? "rsLoRA" : undefined, + loraVariant: detail.config.use_rslora + ? "rslora" + : detail.config.use_loftq + ? "loftq" + : "lora", } : undefined; diff --git a/studio/frontend/src/features/studio/sections/dataset-preview-dialog.tsx b/studio/frontend/src/features/studio/sections/dataset-preview-dialog.tsx index 5ad2d582c4..d05a8cf242 100644 --- a/studio/frontend/src/features/studio/sections/dataset-preview-dialog.tsx +++ b/studio/frontend/src/features/studio/sections/dataset-preview-dialog.tsx @@ -15,6 +15,7 @@ import { Badge } from "@/components/ui/badge"; import { Spinner } from "@/components/ui/spinner"; import { useTrainingActions, useTrainingConfigStore } from "@/features/training"; import { checkDatasetFormat } from "@/features/training/api/datasets-api"; +import { isRawTextDatasetFormat } from "@/features/training/lib/training-methods"; import type { CheckFormatResponse } from "@/features/training/types/datasets"; import { Database02Icon, AlertCircleIcon } from "@hugeicons/core-free-icons"; import { HugeiconsIcon } from "@hugeicons/react"; @@ -90,10 +91,11 @@ export function DatasetPreviewDialog({ const effectiveIsAudio = !!data?.is_audio; const effectiveIsVlm = isVlm || !!data?.is_image; + const isRawFormat = isRawTextDatasetFormat(datasetFormat); const hasHeuristicMapping = !data?.requires_manual_mapping && !!data?.suggested_mapping; - const mappingEnabled = !!data?.requires_manual_mapping || hasHeuristicMapping; + const mappingEnabled = !isRawFormat && (!!data?.requires_manual_mapping || hasHeuristicMapping); const showMappingFooter = mode === "mapping" && mappingEnabled; - const mappingOk = isMappingComplete(manualMapping, effectiveIsVlm, datasetFormat, effectiveIsAudio); + const mappingOk = isRawFormat || isMappingComplete(manualMapping, effectiveIsVlm, datasetFormat, effectiveIsAudio); const availableRoles = getAvailableRoles(effectiveIsVlm, datasetFormat, effectiveIsAudio); const isHfDataset = datasetSource === "huggingface"; @@ -413,7 +415,7 @@ export function DatasetPreviewDialog({ - {data.warning && ( + {data.warning && !isRawFormat && (
{data.warning} diff --git a/studio/frontend/src/features/studio/sections/dataset-section.tsx b/studio/frontend/src/features/studio/sections/dataset-section.tsx index 11c2321863..80f0a06c3f 100644 --- a/studio/frontend/src/features/studio/sections/dataset-section.tsx +++ b/studio/frontend/src/features/studio/sections/dataset-section.tsx @@ -913,6 +913,7 @@ export function DatasetSection() { Alpaca ChatML ShareGPT + Raw Text
diff --git a/studio/frontend/src/features/studio/sections/model-section.tsx b/studio/frontend/src/features/studio/sections/model-section.tsx index 775073eb64..a3c737dae5 100644 --- a/studio/frontend/src/features/studio/sections/model-section.tsx +++ b/studio/frontend/src/features/studio/sections/model-section.tsx @@ -67,6 +67,7 @@ const METHOD_DOTS: Record = { qlora: "bg-emerald-400", lora: "bg-blue-400", full: "bg-amber-400", + cpt: "bg-purple-400", }; const DARK_TRIGGER = @@ -570,7 +571,9 @@ export function ModelSection() { QLoRA uses 4-bit quantization for lowest VRAM. LoRA uses - 16-bit. Full updates all weights.{" "} + 16-bit. Full updates all weights. CPT (Continued Pretraining) + trains on raw text to adapt the model to a new domain without + chat formatting.{" "} + + + + Continued Pretraining + + diff --git a/studio/frontend/src/features/studio/sections/params-section.tsx b/studio/frontend/src/features/studio/sections/params-section.tsx index a6626a5ef8..19b75db140 100644 --- a/studio/frontend/src/features/studio/sections/params-section.tsx +++ b/studio/frontend/src/features/studio/sections/params-section.tsx @@ -34,11 +34,14 @@ import { } from "@/components/ui/tooltip"; import { CONTEXT_LENGTHS, + CPT_TARGET_MODULES, LR_SCHEDULER_OPTIONS, OPTIMIZER_OPTIONS, TARGET_MODULES, } from "@/config/training"; import { useMaxStepsEpochsToggle, useTrainingConfigStore } from "@/features/training"; +import { isRawTextDatasetFormat } from "@/features/training/lib/training-methods"; +import { isAdapterMethod } from "@/types/training"; import type { GradientCheckpointing } from "@/types/training"; import { ArrowDown01Icon, @@ -126,10 +129,13 @@ function SliderRow({ export function ParamsSection(): ReactElement { const store = useTrainingConfigStore(); const platformDeviceType = usePlatformStore((s) => s.deviceType); - const isLora = store.trainingMethod !== "full"; + const isLora = isAdapterMethod(store.trainingMethod); + const isCpt = store.trainingMethod === "cpt"; + const isRawText = isRawTextDatasetFormat(store.datasetFormat); const showVisionLora = store.isVisionModel && store.isDatasetImage === true; const [loraOpen, setLoraOpen] = useState(false); const [hyperOpen, setHyperOpen] = useState(false); + const needsExpandedHeight = isCpt || (isLora && loraOpen) || hyperOpen; const [ctxInput, setCtxInput] = useState(String(store.contextLength)); const ctxAnchorRef = useRef(null); const ctxItems = CONTEXT_LENGTHS.map(String); @@ -168,7 +174,7 @@ export function ParamsSection(): ReactElement { title="Parameters" description="Configure training hyperparameters" accent="orange" - className={`${(isLora && loraOpen) || hyperOpen + className={`${needsExpandedHeight ? "min-h-studio-config-column" : "h-studio-config-column"} duration-150`} > @@ -378,10 +384,62 @@ export function ParamsSection(): ReactElement { className="w-full font-mono" />

- Recommended: 2e-4 for LoRA, 2e-5 for full fine-tune + Recommended: 2e-4 for LoRA, 5e-5 for CPT, 2e-5 for full fine-tune

+ {/* Embedding Learning Rate (CPT only) */} + {isCpt && ( +
+ + Embedding Learning Rate + + + + + + Only used when CPT is training embed_tokens. + Embeddings are easier to destabilize than LoRA weights, so + they usually need a smaller LR. Leave blank to use + lr/10; typical working range is 2x-10x smaller + than the main LR. Increase it only if vocabulary or + domain-token adaptation is too slow. + + + + { + const raw = e.target.value; + if (raw === "") { + store.setEmbeddingLearningRate(null); + return; + } + const n = Number(raw); + store.setEmbeddingLearningRate(Number.isFinite(n) ? n : null); + }} + className="w-full font-mono" + /> +

+ Leave blank to use lr/10 (recommended). Typical range is + 2x-10x smaller than the main learning rate. +

+
+ )} + {/* LoRA Settings */} {isLora && ( @@ -516,7 +574,7 @@ export function ParamsSection(): ReactElement { Target Modules
- {TARGET_MODULES.map((mod) => { + {(isCpt ? CPT_TARGET_MODULES : TARGET_MODULES).map((mod) => { const active = store.targetModules.includes(mod); return (
)} - {!store.isEmbeddingModel && ( + {!store.isEmbeddingModel && !isCpt && !isRawText && (
({ @@ -272,7 +274,7 @@ export function ProgressSection({ {data.modelName || "--"} - {data.trainingMethod === "qlora" ? "QLoRA" : data.trainingMethod === "lora" ? "LoRA" : "Full"} + {trainingMethodLabel}
diff --git a/studio/frontend/src/features/training/api/mappers.ts b/studio/frontend/src/features/training/api/mappers.ts index 561dbe1408..5e68ccd72c 100644 --- a/studio/frontend/src/features/training/api/mappers.ts +++ b/studio/frontend/src/features/training/api/mappers.ts @@ -3,9 +3,10 @@ import type { TrainingConfigState } from "../types/config"; import type { TrainingStartRequest } from "../types/api"; - -const BACKEND_LORA_TYPE = "LoRA/QLoRA"; -const BACKEND_FULL_TYPE = "Full Finetuning"; +import { + isRawTextDatasetFormat, + toBackendTrainingType, +} from "../lib/training-methods"; function parseSliceValue(value: string | null): number | null { if (value == null) return null; @@ -16,16 +17,15 @@ function parseSliceValue(value: string | null): number | null { return num; } -export function toBackendTrainingType(trainingMethod: string): string { - return trainingMethod === "full" ? BACKEND_FULL_TYPE : BACKEND_LORA_TYPE; -} - export function buildTrainingStartPayload( config: TrainingConfigState, ): TrainingStartRequest { + const isCpt = config.trainingMethod === "cpt"; const adapterMethod = config.trainingMethod !== "full"; const isQloraMethod = config.trainingMethod === "qlora"; + const isFourBitModel = (config.selectedModel ?? "").toLowerCase().includes("4bit"); const isEmbedding = config.isEmbeddingModel; + const isRawText = isRawTextDatasetFormat(config.datasetFormat); const hfDataset = config.datasetSource === "huggingface" ? config.dataset : null; const localDatasets = config.datasetSource === "upload" && config.uploadedFile @@ -53,7 +53,7 @@ export function buildTrainingStartPayload( model_name: config.selectedModel ?? "", training_type: toBackendTrainingType(config.trainingMethod), hf_token: config.hfToken.trim() || null, - load_in_4bit: adapterMethod ? isQloraMethod : false, + load_in_4bit: (adapterMethod && isQloraMethod) || (isCpt && isFourBitModel), max_seq_length: config.contextLength, trust_remote_code: config.trustRemoteCode ?? false, hf_dataset: hfDataset, @@ -71,6 +71,10 @@ export function buildTrainingStartPayload( custom_format_mapping: customFormatMapping, num_epochs: config.epochs, learning_rate: String(config.learningRate), + embedding_learning_rate: + isCpt && config.embeddingLearningRate != null + ? config.embeddingLearningRate + : null, batch_size: config.batchSize, gradient_accumulation_steps: config.gradientAccumulation, warmup_steps: isEmbedding ? null : config.warmupSteps, @@ -91,7 +95,8 @@ export function buildTrainingStartPayload( gradient_checkpointing: config.gradientCheckpointing, use_rslora: config.loraVariant === "rslora", use_loftq: config.loraVariant === "loftq", - train_on_completions: isEmbedding ? false : config.trainOnCompletions, + // CPT always trains on full sequences (no chat format masking) + train_on_completions: (isEmbedding || isCpt || isRawText) ? false : config.trainOnCompletions, finetune_vision_layers: config.finetuneVisionLayers, finetune_language_layers: config.finetuneLanguageLayers, finetune_attention_modules: config.finetuneAttentionModules, diff --git a/studio/frontend/src/features/training/hooks/use-training-actions.ts b/studio/frontend/src/features/training/hooks/use-training-actions.ts index deaec6c6c2..cff4e929ae 100644 --- a/studio/frontend/src/features/training/hooks/use-training-actions.ts +++ b/studio/frontend/src/features/training/hooks/use-training-actions.ts @@ -8,6 +8,7 @@ import { checkDatasetFormat } from "../api/datasets-api"; import { getTrainingRun } from "../api/history-api"; import { buildTrainingStartPayload } from "../api/mappers"; import { resetTraining, startTraining, stopTraining } from "../api/train-api"; +import { isRawTextDatasetFormat } from "../lib/training-methods"; import { syncTrainingRuntimeFromBackend } from "../lib/sync-runtime"; import { validateTrainingConfig } from "../lib/validation"; import { useDatasetPreviewDialogStore } from "../stores/dataset-preview-dialog-store"; @@ -88,7 +89,10 @@ export function useTrainingActions() { }); } - const needsReview = check.requires_manual_mapping || check.detected_format === "custom_heuristic"; + const isRawFormat = isRawTextDatasetFormat(config.datasetFormat); + const needsReview = + !isRawFormat && + (check.requires_manual_mapping || check.detected_format === "custom_heuristic"); if (needsReview && !hasManualMapping(config, isVlm, isAudio)) { // Pre-fill from suggested_mapping or VLM detected columns const hint: Record = {}; diff --git a/studio/frontend/src/features/training/lib/training-methods.ts b/studio/frontend/src/features/training/lib/training-methods.ts new file mode 100644 index 0000000000..9070f2adcf --- /dev/null +++ b/studio/frontend/src/features/training/lib/training-methods.ts @@ -0,0 +1,48 @@ +// 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 { DatasetFormat, TrainingMethod } from "@/types/training"; + +const BACKEND_TRAINING_TYPE: Record = { + qlora: "LoRA/QLoRA", + lora: "LoRA/QLoRA", + full: "Full Finetuning", + cpt: "Continued Pretraining", +}; + +const TRAINING_METHOD_LABELS: Record = { + qlora: "QLoRA", + lora: "LoRA", + full: "Full", + cpt: "CPT", +}; + +export function toBackendTrainingType(trainingMethod: TrainingMethod): string { + return BACKEND_TRAINING_TYPE[trainingMethod]; +} + +export function getTrainingMethodLabel( + trainingMethod: TrainingMethod | string, +): string { + if (Object.prototype.hasOwnProperty.call(TRAINING_METHOD_LABELS, trainingMethod)) { + return TRAINING_METHOD_LABELS[trainingMethod as TrainingMethod]; + } + return TRAINING_METHOD_LABELS.full; +} + +export function parseBackendTrainingMethod( + trainingType: unknown, + loadIn4Bit: unknown, +): TrainingMethod { + if (trainingType === "Continued Pretraining") return "cpt"; + if (trainingType === "LoRA/QLoRA") { + return loadIn4Bit ? "qlora" : "lora"; + } + return "full"; +} + +export function isRawTextDatasetFormat( + datasetFormat: DatasetFormat, +): boolean { + return datasetFormat === "raw"; +} diff --git a/studio/frontend/src/features/training/stores/training-config-store.ts b/studio/frontend/src/features/training/stores/training-config-store.ts index 8214b0eb2a..ef16f641f5 100644 --- a/studio/frontend/src/features/training/stores/training-config-store.ts +++ b/studio/frontend/src/features/training/stores/training-config-store.ts @@ -1,15 +1,17 @@ // 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 { DEFAULT_HYPERPARAMS, LR_DEFAULT_FULL, LR_DEFAULT_LORA, STEPS } from "@/config/training"; +import { CPT_TARGET_MODULES, DEFAULT_HYPERPARAMS, LR_DEFAULT_CPT, LR_DEFAULT_FULL, LR_DEFAULT_LORA, STEPS, TARGET_MODULES } from "@/config/training"; import { authFetch } from "@/features/auth"; import { isAdapterMethod } from "@/types/training"; +import type { DatasetFormat } from "@/types/training"; import type { ModelType, StepNumber, TrainingMethod } from "@/types/training"; import { create } from "zustand"; import { persist } from "zustand/middleware"; import { checkDatasetFormat } from "../api/datasets-api"; import { checkVisionModel, getModelConfig } from "../api/models-api"; import { mapBackendModelConfigToTrainingPatch } from "../lib/model-defaults"; +import { isRawTextDatasetFormat } from "../lib/training-methods"; import type { BackendModelConfig } from "../api/models-api"; import type { TrainingConfigState, TrainingConfigStore } from "../types/config"; @@ -108,6 +110,11 @@ let _learningRateManuallySet = false; // setTrainingMethod can restore it when switching back from full to adapter. let _yamlLearningRate: number | undefined = undefined; +// Track whether entering CPT auto-forced datasetFormat="raw" so that +// leaving CPT can restore the prior user-visible format. +let _datasetFormatBeforeCpt: DatasetFormat | null = null; +let _datasetFormatAutoForcedByCpt = false; + const NON_PERSISTED_STATE_KEYS: ReadonlySet = new Set([ "modelType", "isCheckingVision", @@ -156,6 +163,123 @@ function canProceedForStep(state: TrainingConfigState): boolean { } } +type TrainingMethodStatePatch = Partial< + Pick< + TrainingConfigState, + | "trainingMethod" + | "learningRate" + | "loraRank" + | "loraAlpha" + | "loraVariant" + | "targetModules" + | "datasetFormat" + | "trainOnCompletions" + > +>; + +function getCptTrainingPatch(): TrainingMethodStatePatch { + return { + loraRank: 128, + loraAlpha: 32, + loraVariant: "rslora", + targetModules: CPT_TARGET_MODULES, + datasetFormat: "raw", + trainOnCompletions: false, + }; +} + +function getCptModelDefaultsPatch(): TrainingMethodStatePatch { + return { + ...getCptTrainingPatch(), + learningRate: LR_DEFAULT_CPT, + }; +} + +function getRestoreFromCptPatch(): TrainingMethodStatePatch { + return { + loraRank: DEFAULT_HYPERPARAMS.loraRank, + loraAlpha: DEFAULT_HYPERPARAMS.loraAlpha, + loraVariant: DEFAULT_HYPERPARAMS.loraVariant, + targetModules: TARGET_MODULES, + }; +} + +function clearCptDatasetFormatTracking(): void { + _datasetFormatBeforeCpt = null; + _datasetFormatAutoForcedByCpt = false; +} + +function recordCptDatasetFormatOverride(currentDatasetFormat: DatasetFormat): void { + if (isRawTextDatasetFormat(currentDatasetFormat)) { + clearCptDatasetFormatTracking(); + return; + } + _datasetFormatBeforeCpt = currentDatasetFormat; + _datasetFormatAutoForcedByCpt = true; +} + +function getRestoreDatasetFormatFromCptPatch(): TrainingMethodStatePatch { + if (!_datasetFormatAutoForcedByCpt || _datasetFormatBeforeCpt == null) { + clearCptDatasetFormatTracking(); + return {}; + } + + const previousDatasetFormat = _datasetFormatBeforeCpt; + clearCptDatasetFormatTracking(); + return { datasetFormat: previousDatasetFormat }; +} + +function resolveTrainingMethodLearningRate( + prevMethod: TrainingMethod, + nextMethod: TrainingMethod, +): number | undefined { + if (_learningRateManuallySet) { + return undefined; + } + + const wasCpt = prevMethod === "cpt"; + const wasAdapter = isAdapterMethod(prevMethod); + const nowAdapter = isAdapterMethod(nextMethod); + + if (nextMethod === "cpt") { + return LR_DEFAULT_CPT; + } + if (wasCpt && nowAdapter) { + return _yamlLearningRate ?? LR_DEFAULT_LORA; + } + if (wasAdapter && nowAdapter) { + return undefined; + } + return nowAdapter ? _yamlLearningRate ?? LR_DEFAULT_LORA : LR_DEFAULT_FULL; +} + +function buildTrainingMethodPatch( + prevMethod: TrainingMethod, + nextMethod: TrainingMethod, + currentDatasetFormat: DatasetFormat, +): TrainingMethodStatePatch { + const patch: TrainingMethodStatePatch = { trainingMethod: nextMethod }; + + if (prevMethod !== "cpt" && nextMethod === "cpt") { + recordCptDatasetFormatOverride(currentDatasetFormat); + Object.assign(patch, getCptTrainingPatch()); + } + if (prevMethod === "cpt" && nextMethod !== "cpt") { + Object.assign( + patch, + getRestoreFromCptPatch(), + getRestoreDatasetFormatFromCptPatch(), + ); + } + + const learningRate = resolveTrainingMethodLearningRate(prevMethod, nextMethod); + if (learningRate !== undefined) { + patch.learningRate = learningRate; + } + + return patch; +} + export const useTrainingConfigStore = create()( persist( (set, get) => { @@ -216,11 +340,14 @@ export const useTrainingConfigStore = create()( // Auto-select training method based on model size vs GPU memory. // If model_size * 1.5 * context_scale fits in free VRAM, use LoRA 16-bit. // Otherwise use QLoRA 4-bit. + // Auto-select LoRA vs QLoRA based on GPU memory. + // Skip if user has manually chosen CPT -- don't override it. const modelSizeBytes = modelDetails.model_size_bytes; - if (modelSizeBytes && modelSizeBytes > 0) { + if (modelSizeBytes && modelSizeBytes > 0 && get().trainingMethod !== "cpt") { void autoSelectTrainingMethod(modelSizeBytes, patch.contextLength ?? get().contextLength) .then((method) => { if (get().selectedModel !== modelName) return; + if (get().trainingMethod === "cpt") return; if (method) { const lrPatch = !_learningRateManuallySet && !modelConfigHasLR ? { learningRate: method === "full" ? LR_DEFAULT_FULL : LR_DEFAULT_LORA } @@ -230,8 +357,16 @@ export const useTrainingConfigStore = create()( }); } + // Preserve CPT hyperparams: YAML adapter defaults (r/alpha/targets/LR) + // are tuned for standard LoRA and would otherwise clobber CPT settings. + const cptOverrides = + get().trainingMethod === "cpt" + ? getCptModelDefaultsPatch() + : {}; + set({ ...patch, + ...cptOverrides, modelType: inferredModelType, isVisionModel: modelDetails.is_vision, isEmbeddingModel: isEmbedding, @@ -396,29 +531,14 @@ export const useTrainingConfigStore = create()( void loadAndApplyModelDefaults(state.selectedModel); }, setTrainingMethod: (trainingMethod) => { - if (_learningRateManuallySet) { - set({ trainingMethod }); - return; - } - - const prev = get().trainingMethod; - const wasAdapter = isAdapterMethod(prev); - const nowAdapter = isAdapterMethod(trainingMethod); - - // qlora <-> lora: same LR range, don't touch learning rate - if (wasAdapter && nowAdapter) { - set({ trainingMethod }); - return; - } - - // Category changed (adapter <-> full) - if (nowAdapter) { - // Switching TO adapter: restore YAML LR if available - set({ trainingMethod, learningRate: _yamlLearningRate ?? LR_DEFAULT_LORA }); - } else { - // Switching TO full: no YAML full-LR exists, use constant - set({ trainingMethod, learningRate: LR_DEFAULT_FULL }); - } + const state = get(); + set( + buildTrainingMethodPatch( + state.trainingMethod, + trainingMethod, + state.datasetFormat, + ), + ); }, setHfToken: (hfToken) => set({ hfToken: hfToken.trim().replace(/^["']+|["']+$/g, "") }), @@ -448,7 +568,26 @@ export const useTrainingConfigStore = create()( runDatasetCheck(uploadedFile, "train"); } }, - setDatasetFormat: (datasetFormat) => set({ datasetFormat }), + setDatasetFormat: (datasetFormat) => + set((state) => { + if (state.trainingMethod === "cpt") { + if (isRawTextDatasetFormat(datasetFormat)) { + clearCptDatasetFormatTracking(); + } + return { + datasetFormat: "raw", + trainOnCompletions: false, + }; + } + + return { + datasetFormat, + trainOnCompletions: + isRawTextDatasetFormat(datasetFormat) + ? false + : state.trainOnCompletions, + }; + }), setDataset: (dataset) => { _datasetCheckController?.abort(); _datasetCheckController = null; @@ -566,6 +705,8 @@ export const useTrainingConfigStore = create()( _learningRateManuallySet = true; set({ learningRate }); }, + setEmbeddingLearningRate: (embeddingLearningRate) => + set({ embeddingLearningRate }), setOptimizerType: (optimizerType) => set({ optimizerType }), setLrSchedulerType: (lrSchedulerType) => set({ lrSchedulerType }), setLoraRank: (loraRank) => set({ loraRank }), @@ -608,6 +749,7 @@ export const useTrainingConfigStore = create()( _trainOnCompletionsManuallySet = false; _learningRateManuallySet = false; _yamlLearningRate = undefined; + clearCptDatasetFormatTracking(); set(initialState); }, resetToModelDefaults: () => { @@ -629,7 +771,7 @@ export const useTrainingConfigStore = create()( }, { name: "unsloth_training_config_v1", - version: 9, + version: 10, migrate: (persisted, version) => { const s = persisted as Record; if (version < 2 && s.datasetSubset == null && s.datasetConfig != null) { @@ -665,6 +807,17 @@ export const useTrainingConfigStore = create()( s.weightDecay = DEFAULT_HYPERPARAMS.weightDecay; } } + if (version < 10 && s.trainingMethod === "cpt") { + // Backfill CPT defaults for state persisted before they existed. + s.loraRank = 128; + s.loraAlpha = 32; + s.loraVariant = "rslora"; + s.targetModules = CPT_TARGET_MODULES; + s.datasetFormat = "raw"; + if (s.learningRate == null || s.learningRate === LR_DEFAULT_LORA) { + s.learningRate = LR_DEFAULT_CPT; + } + } return s as unknown as TrainingConfigStore; }, partialize: partializePersistedState, diff --git a/studio/frontend/src/features/training/types/api.ts b/studio/frontend/src/features/training/types/api.ts index ae65d1a53c..fb8a2f899e 100644 --- a/studio/frontend/src/features/training/types/api.ts +++ b/studio/frontend/src/features/training/types/api.ts @@ -21,6 +21,8 @@ export interface TrainingStartRequest { custom_format_mapping?: Record | null; num_epochs: number; learning_rate: string; + /** Optional CPT embedding LR. If omitted, backend uses lr/10; typical range is 2x-10x smaller than main LR. */ + embedding_learning_rate?: number | null; batch_size: number; gradient_accumulation_steps: number; warmup_steps: number | null; diff --git a/studio/frontend/src/features/training/types/config.ts b/studio/frontend/src/features/training/types/config.ts index 2d19dea874..5b156316ca 100644 --- a/studio/frontend/src/features/training/types/config.ts +++ b/studio/frontend/src/features/training/types/config.ts @@ -41,6 +41,7 @@ export interface TrainingConfigState { epochs: number; contextLength: number; learningRate: number; + embeddingLearningRate: number | null; optimizerType: string; lrSchedulerType: string; loraRank: number; @@ -115,6 +116,7 @@ export interface TrainingConfigActions { setEpochs: (epochs: number) => void; setContextLength: (length: number) => void; setLearningRate: (rate: number) => void; + setEmbeddingLearningRate: (rate: number | null) => void; setOptimizerType: (value: string) => void; setLrSchedulerType: (value: string) => void; setLoraRank: (rank: number) => void; diff --git a/studio/frontend/src/lib/vram.ts b/studio/frontend/src/lib/vram.ts index a5abebb043..fad2c438bd 100644 --- a/studio/frontend/src/lib/vram.ts +++ b/studio/frontend/src/lib/vram.ts @@ -57,7 +57,12 @@ export type VramFitStatus = "fits" | "tight" | "exceeds"; */ export const FP16_LOADING_BYTES = 2.0; -export type TrainingMethod = "qlora" | "lora" | "full"; +export type TrainingMethod = "qlora" | "lora" | "full" | "cpt"; + +function usesQuantizedLoading(method: TrainingMethod, modelId?: string): boolean { + if (method === "qlora") return true; + return method === "cpt" && (modelId ?? "").toLowerCase().includes("4bit"); +} /** * Estimate VRAM (GB) needed to load a model with Unsloth. @@ -66,15 +71,18 @@ export type TrainingMethod = "qlora" | "lora" | "full"; * - QLoRA : 4-bit quantized via bnb -> 0.90 bytes/param (calibrated) * - LoRA : fp16 -> 2.0 bytes/param (theoretical) * - Full : fp16 -> 2.0 bytes/param (theoretical) + * - CPT : fp16 LoRA (16-bit base) -> 2.0 bytes/param (theoretical) * * Formula: totalParams * bytesPerParam + 1.4 GB overhead */ export function estimateLoadingVram( totalParams: number, method: TrainingMethod = "qlora", + modelId?: string, ): number { - const bytesPerParam = - method === "qlora" ? BNB_4BIT_LOADING_BYTES : FP16_LOADING_BYTES; + const bytesPerParam = usesQuantizedLoading(method, modelId) + ? BNB_4BIT_LOADING_BYTES + : FP16_LOADING_BYTES; const gb = (totalParams / 1e9) * bytesPerParam + LOADING_OVERHEAD_GB; return Math.round(gb * 10) / 10; } @@ -119,7 +127,7 @@ export function buildModelVramMap( continue; } - const est = estimateLoadingVram(model.totalParams, method); + const est = estimateLoadingVram(model.totalParams, method, model.id); const status = gpu.available ? checkVramFit(est, gpu.memoryTotalGb) : null; map.set(model.id, { est, status }); } diff --git a/studio/frontend/src/types/training.ts b/studio/frontend/src/types/training.ts index 187f54a13b..feca65ebda 100644 --- a/studio/frontend/src/types/training.ts +++ b/studio/frontend/src/types/training.ts @@ -2,14 +2,14 @@ // Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0 export type ModelType = "vision" | "audio" | "embeddings" | "text"; -export type TrainingMethod = "qlora" | "lora" | "full"; +export type TrainingMethod = "qlora" | "lora" | "full" | "cpt"; export function isAdapterMethod(method: TrainingMethod): boolean { - return method === "lora" || method === "qlora"; + return method === "lora" || method === "qlora" || method === "cpt"; } export type StepNumber = 1 | 2 | 3 | 4 | 5; export type DatasetSource = "huggingface" | "upload"; -export type DatasetFormat = "auto" | "alpaca" | "chatml" | "sharegpt"; +export type DatasetFormat = "auto" | "alpaca" | "chatml" | "sharegpt" | "raw"; export type GradientCheckpointing = "none" | "true" | "unsloth" | "mlx"; export interface WizardState { From 94811ba75d313eebf2b4ea1b40ca79418bc8b18d Mon Sep 17 00:00:00 2001 From: Daniel Han Date: Wed, 6 May 2026 03:31:41 -0700 Subject: [PATCH 008/159] Fix 14 stale tests under tests/studio/install/ that drifted from code (#5305) * Fix 14 stale tests under tests/studio/install/ that drifted from code All 14 failures audited locally and tracked back to test-side drift (no production-code regressions). After these test updates the entire tests/studio/install/ directory now passes: 346 passed, 1 skipped. Per failure: tests/studio/install/test_install_llama_prebuilt_logic.py (5 fails): * test_existing_install_matches_plan_with_fingerprint_linux * test_install_prebuilt_skips_download_when_existing_install_matches * test_install_prebuilt_skips_when_older_release_fallback_matches_existing_install * test_install_prebuilt_skips_same_release_fallback_attempt_when_installed * test_existing_install_matches_choice_fails_when_install_tree_incomplete All five build a fake Linux install tree via write_linux_install_shape and call existing_install_matches_choice. The matcher returns False because runtime_payload_is_healthy now requires a libllama-common.so* library in build/bin/ (added by PR #5135), and the fixture never wrote it. Add the missing library to write_linux_install_shape; matcher passes for all five tests. tests/studio/install/test_rocm_support.py (8 fails after the partial audit, one collection-tier flake): * TestEnsureRocmTorch::test_cpu_torch_gets_rocm_reinstall and TestEnsureRocmTorch::test_probe_timeout_triggers_reinstall _ensure_rocm_torch was refactored to call pip_install for the torch reinstall and pip_install_try (not pip_install) for the follow-up bitsandbytes install. The tests still asserted mock_pip.call_count == 2. Add a second @patch.object on pip_install_try and split the assertions across the two mocks. * TestInstallShStructure::test_cuda_precedence Asserted file-position-of-string ordering: looked for `if [ -z "$_smi" ]` before the first `amd-smi` literal in install.sh. The installer now defines top-level helpers `_has_amd_rocm_gpu` (uses `amd-smi`) and `_has_usable_nvidia_gpu` (uses `nvidia-smi`) before either is called from `get_torch_index_url`, so file-position ordering carries no semantic meaning. Rewrite the test to extract the `get_torch_index_url` body via a small brace-matched helper and assert the runtime ordering: NVIDIA call sits before the `if [ -z "$_smi" ]` branch and the AMD call sits inside it. * TestLiveRegression::test_get_torch_index_url_returns_cuda_on_nvidia Sed-extracted only get_torch_index_url and eval'd it -- but the function calls _has_amd_rocm_gpu and _has_usable_nvidia_gpu, so the eval'd body crashed and fell through to the CPU URL on a fully-loaded NVIDIA host. Extract the helpers alongside the function. Also pre-skip when nvidia-smi is on PATH but does not list a GPU (containers occasionally ship the binary without a driver). * TestWorkerRocmMambaSsm::test_probe_script_has_getattr_hip and TestWorkerRocmMambaSsm::test_probe_returns_hip_version_field The wheel-resolver probe subprocess (the only place where `getattr(torch.version, 'hip', None)` is emitted) was hoisted out of worker.py into studio/backend/utils/wheel_utils.py during the wheel-resolver refactor. Point the file-content assertions at wheel_utils.py and assert worker.py still consumes the `hip_version` field. * TestHardwareAmdBranching::test_hardware_branches_on_is_rocm_for_utilization TestHardwareAmdBranching::test_hardware_branches_on_is_rocm_for_visible TestHardwareAmdBranching::test_hardware_branches_on_is_rocm_for_physical_count hardware.py refactored: the IS_ROCM branch and direct `from . import amd` were hoisted out of get_gpu_utilization / get_visible_gpu_utilization into the shared `_smi_query` dispatcher. Update the first two tests to assert the dispatcher call shape (`_smi_query("get_primary_gpu_utilization", ...)` etc.) plus IS_ROCM + amd-import in `_smi_query` itself. Update the physical-count test to assert IS_ROCM + the literal `from . import amd` as that function still imports amd directly rather than going through `_smi_query`. No production-code changes; tests-only. * [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> --- .../test_install_llama_prebuilt_logic.py | 3 + tests/studio/install/test_rocm_support.py | 149 ++++++++++++++---- 2 files changed, 120 insertions(+), 32 deletions(-) diff --git a/tests/studio/install/test_install_llama_prebuilt_logic.py b/tests/studio/install/test_install_llama_prebuilt_logic.py index 79dab30129..128b90cfe2 100644 --- a/tests/studio/install/test_install_llama_prebuilt_logic.py +++ b/tests/studio/install/test_install_llama_prebuilt_logic.py @@ -754,6 +754,9 @@ def write_linux_install_shape(install_dir: Path) -> None: (install_dir / "llama-quantize").write_text("#!/bin/sh\n", encoding = "utf-8") (runtime_dir / "llama-server").write_text("#!/bin/sh\n", encoding = "utf-8") (runtime_dir / "llama-quantize").write_text("#!/bin/sh\n", encoding = "utf-8") + # Mirror the runtime payload health groups in install_llama_prebuilt.py: + # libllama-common.so* was added by PR #5135 and is required. + (runtime_dir / "libllama-common.so.0").write_bytes(b"DLL") (runtime_dir / "libllama.so.0").write_bytes(b"DLL") (runtime_dir / "libggml.so.0").write_bytes(b"DLL") (runtime_dir / "libggml-base.so.0").write_bytes(b"DLL") diff --git a/tests/studio/install/test_rocm_support.py b/tests/studio/install/test_rocm_support.py index 48831fd57b..553756cd4f 100644 --- a/tests/studio/install/test_rocm_support.py +++ b/tests/studio/install/test_rocm_support.py @@ -53,6 +53,32 @@ _has_usable_nvidia_gpu = stack_mod._has_usable_nvidia_gpu _ROCM_TORCH_INDEX = stack_mod._ROCM_TORCH_INDEX +def _extract_sh_function_body(source: str, name: str) -> str: + """Return the body of a shell function from `source` by brace matching. + + Used by structural tests that need to assert ordering of helper + calls inside a specific function rather than across the whole + install.sh file. + """ + needle = f"{name}() {{" + start = source.find(needle) + if start < 0: + return "" + depth = 0 + i = start + len(needle) - 1 # land on the opening brace + n = len(source) + while i < n: + ch = source[i] + if ch == "{": + depth += 1 + elif ch == "}": + depth -= 1 + if depth == 0: + return source[start : i + 1] + i += 1 + return source[start:] + + # ── Helper: build HostInfo for different scenarios ────────────────────────── @@ -561,12 +587,13 @@ class TestEnsureRocmTorch: _ensure_rocm_torch() mock_pip.assert_not_called() + @patch.object(stack_mod, "pip_install_try", return_value = True) @patch.object(stack_mod, "pip_install") @patch.object(stack_mod, "_has_usable_nvidia_gpu", return_value = False) @patch.object(stack_mod, "_has_rocm_gpu", return_value = True) @patch.object(stack_mod, "_detect_rocm_version", return_value = (7, 1)) def test_cpu_torch_gets_rocm_reinstall( - self, mock_ver, mock_gpu, mock_nvidia, mock_pip + self, mock_ver, mock_gpu, mock_nvidia, mock_pip, mock_pip_try ): """CPU-only torch on ROCm host should trigger reinstall.""" mock_probe = MagicMock() @@ -575,12 +602,11 @@ class TestEnsureRocmTorch: with patch("os.path.isdir", return_value = True): with patch("subprocess.run", return_value = mock_probe): _ensure_rocm_torch() - # Should call pip_install twice: once for torch, once for bitsandbytes - assert mock_pip.call_count == 2 - torch_call = mock_pip.call_args_list[0] - assert "rocm7.1" in str(torch_call) - bnb_call = mock_pip.call_args_list[1] - assert "bitsandbytes" in str(bnb_call) + # Should install torch via pip_install and bitsandbytes via pip_install_try. + assert mock_pip.call_count == 1 + assert "rocm7.1" in str(mock_pip.call_args_list[0]) + assert mock_pip_try.call_count >= 1 + assert "bitsandbytes" in str(mock_pip_try.call_args_list[0]) @patch.object(stack_mod, "pip_install") @patch.object(stack_mod, "_has_usable_nvidia_gpu", return_value = False) @@ -642,12 +668,13 @@ class TestEnsureRocmTorch: torch_call = mock_pip.call_args_list[0] assert "rocm7.1" in str(torch_call) + @patch.object(stack_mod, "pip_install_try", return_value = True) @patch.object(stack_mod, "pip_install") @patch.object(stack_mod, "_has_usable_nvidia_gpu", return_value = False) @patch.object(stack_mod, "_has_rocm_gpu", return_value = True) @patch.object(stack_mod, "_detect_rocm_version", return_value = (7, 1)) def test_probe_timeout_triggers_reinstall( - self, mock_ver, mock_gpu, mock_nvidia, mock_pip + self, mock_ver, mock_gpu, mock_nvidia, mock_pip, mock_pip_try ): """Probe subprocess timeout should not crash; should proceed to reinstall.""" with patch("os.path.isdir", return_value = True): @@ -656,8 +683,10 @@ class TestEnsureRocmTorch: ): _ensure_rocm_torch() # If probe times out, the function should treat torch as unusable and reinstall - assert mock_pip.call_count == 2 + # both torch (via pip_install) and bitsandbytes (via pip_install_try). + assert mock_pip.call_count == 1 assert "rocm7.1" in str(mock_pip.call_args_list[0]) + assert mock_pip_try.call_count >= 1 @patch.object(stack_mod, "pip_install") @patch.object(stack_mod, "_has_usable_nvidia_gpu", return_value = False) @@ -857,15 +886,33 @@ class TestInstallShStructure: assert "rocm" in source.lower() def test_cuda_precedence(self): - """ROCm detection should only run when nvidia-smi is absent.""" + """ROCm detection should only run when nvidia-smi is absent. + + install.sh defines _has_amd_rocm_gpu and _has_usable_nvidia_gpu + helpers near each other (file-position order has no semantic + meaning), so check the runtime ordering inside + get_torch_index_url instead: NVIDIA branch runs first and the + AMD/ROCm branch only fires inside the `if [ -z "$_smi" ]` + block. + """ sh_path = PACKAGE_ROOT / "install.sh" source = sh_path.read_text() - # The ROCm block should be inside the "if [ -z "$_smi" ]" branch - smi_block_start = source.find('if [ -z "$_smi" ]') - rocm_block_start = source.find("amd-smi") + body = _extract_sh_function_body(source, "get_torch_index_url") + nvidia_call = body.find("_has_usable_nvidia_gpu") + no_nvidia_branch = body.find('if [ -z "$_smi" ]') + rocm_call = body.find("_has_amd_rocm_gpu") assert ( - smi_block_start < rocm_block_start - ), "ROCm detection should be inside the 'no nvidia-smi' branch" + nvidia_call >= 0 + ), "get_torch_index_url should call _has_usable_nvidia_gpu" + assert ( + no_nvidia_branch >= 0 + ), "get_torch_index_url should gate ROCm on no-nvidia-smi" + assert ( + rocm_call > no_nvidia_branch + ), "ROCm detection should sit inside the 'no nvidia-smi' branch" + assert ( + nvidia_call < no_nvidia_branch + ), "NVIDIA detection should run before the no-nvidia-smi branch" def test_bitsandbytes_amd_install(self): """install.sh should install bitsandbytes for AMD when ROCm detected.""" @@ -963,16 +1010,32 @@ class TestLiveRegression: if not shutil.which("nvidia-smi"): pytest.skip("No nvidia-smi available") - sh_path = PACKAGE_ROOT / "install.sh" - # Extract just the function (don't source the whole installer) - result = subprocess.run( + # Skip if nvidia-smi exists but does not actually list a GPU on this + # host (containers occasionally ship the binary without a driver). + check = subprocess.run( [ "bash", "-c", - f"eval \"$(sed -n '/^get_torch_index_url()/,/^}}/p' '{sh_path}')\"; " - "get_torch_index_url", + "nvidia-smi -L 2>/dev/null | " + "awk '/^GPU[[:space:]]+[0-9]+:/{f=1} END{exit !f}'", ], capture_output = True, + ) + if check.returncode != 0: + pytest.skip("nvidia-smi is on PATH but no GPU is listed") + + sh_path = PACKAGE_ROOT / "install.sh" + # get_torch_index_url calls _has_usable_nvidia_gpu and + # _has_amd_rocm_gpu, so all three function definitions must be + # in scope when we eval the extract. + extract_cmd = ( + f"sed -n '/^_has_amd_rocm_gpu()/,/^}}$/p; " + f"/^_has_usable_nvidia_gpu()/,/^}}$/p; " + f"/^get_torch_index_url()/,/^}}$/p' '{sh_path}'" + ) + result = subprocess.run( + ["bash", "-c", f'eval "$({extract_cmd})"; get_torch_index_url'], + capture_output = True, text = True, timeout = 30, ) @@ -988,19 +1051,23 @@ class TestLiveRegression: # Load worker.py module _WORKER_PATH = PACKAGE_ROOT / "studio" / "backend" / "core" / "training" / "worker.py" +# The wheel-probe subprocess was hoisted out of worker.py into wheel_utils +# during the wheel-resolver refactor; the probe script literal lives there. +_WHEEL_UTILS_PATH = PACKAGE_ROOT / "studio" / "backend" / "utils" / "wheel_utils.py" class TestWorkerRocmMambaSsm: """Verify worker.py Mamba/SSM install logic on ROCm.""" def test_probe_returns_hip_version_field(self): - """_probe_causal_conv1d_env probe script should include hip_version.""" - source = _WORKER_PATH.read_text() - assert "hip_version" in source + """The wheel probe should include hip_version, and worker.py should + consume it.""" + assert "hip_version" in _WHEEL_UTILS_PATH.read_text() + assert "hip_version" in _WORKER_PATH.read_text() def test_probe_script_has_getattr_hip(self): """Probe script should use getattr for torch.version.hip (safe on CUDA).""" - source = _WORKER_PATH.read_text() + source = _WHEEL_UTILS_PATH.read_text() assert "getattr(torch.version, 'hip', None)" in source def test_direct_wheel_url_returns_none_without_cuda_major(self): @@ -1216,27 +1283,45 @@ class TestHardwareAmdBranching: assert "from . import amd" in source def test_hardware_branches_on_is_rocm_for_utilization(self): - """get_gpu_utilization should check IS_ROCM before choosing backend.""" + """get_gpu_utilization should dispatch to amd.py via _smi_query + when IS_ROCM, and the dispatcher itself must check IS_ROCM and + import the amd backend.""" hw_path = ( PACKAGE_ROOT / "studio" / "backend" / "utils" / "hardware" / "hardware.py" ) source = hw_path.read_text() - # Find the get_gpu_utilization function func_start = source.find("def get_gpu_utilization") func_body = source[func_start : source.find("\ndef ", func_start + 1)] - assert "IS_ROCM" in func_body - assert "amd.get_primary_gpu_utilization" in func_body + assert '_smi_query("get_primary_gpu_utilization"' in func_body + smi = source[ + source.find("def _smi_query") : source.find( + "\ndef ", source.find("def _smi_query") + 1 + ) + ] + assert "IS_ROCM" in smi + assert "from . import amd" in smi def test_hardware_branches_on_is_rocm_for_visible(self): - """get_visible_gpu_utilization should check IS_ROCM.""" + """get_visible_gpu_utilization should dispatch to amd.py via + _smi_query when IS_ROCM.""" hw_path = ( PACKAGE_ROOT / "studio" / "backend" / "utils" / "hardware" / "hardware.py" ) source = hw_path.read_text() func_start = source.find("def get_visible_gpu_utilization") func_body = source[func_start : source.find("\ndef ", func_start + 1)] - assert "IS_ROCM" in func_body - assert "amd.get_visible_gpu_utilization" in func_body + # The dispatcher call may wrap onto multiple lines; allow whitespace + # between the open paren and the literal func name argument. + import re as _re + + assert _re.search(r'_smi_query\(\s*"get_visible_gpu_utilization"', func_body) + smi = source[ + source.find("def _smi_query") : source.find( + "\ndef ", source.find("def _smi_query") + 1 + ) + ] + assert "IS_ROCM" in smi + assert "from . import amd" in smi def test_hardware_branches_on_is_rocm_for_physical_count(self): """get_physical_gpu_count should try amd.py when IS_ROCM.""" @@ -1247,7 +1332,7 @@ class TestHardwareAmdBranching: func_start = source.find("def get_physical_gpu_count") func_body = source[func_start : source.find("\ndef ", func_start + 1)] assert "IS_ROCM" in func_body - assert "amd.get_physical_gpu_count" in func_body + assert "from . import amd" in func_body # ============================================================================= From 1942e58171aa7f69381bc866a1886b35934a7151 Mon Sep 17 00:00:00 2001 From: Daniel Han Date: Wed, 6 May 2026 10:31:51 +0000 Subject: [PATCH 009/159] tests/studio: comprehensive hardware dispatch matrix Drives every supported hardware profile from a single test host by spoofing platform, torch.cuda, torch.xpu, torch.version.hip, and sys.modules['mlx'] so we can exercise the CUDA, ROCm, XPU, MLX, and CPU dispatch paths deterministically without owning the actual hardware. Profiles covered (parametrized; add a row to PROFILES to extend): nvidia_cuda Linux x86_64 + cuda available, hip=None amd_rocm Linux x86_64 + cuda available, hip="6.1" (PyTorch ROCm aliases torch.cuda over HIP) intel_xpu Linux x86_64 + cuda off, xpu available apple_silicon_mlx Darwin arm64 + cuda/xpu off + mlx in sys.modules apple_silicon_no_mlx Darwin arm64 + everything off (Mac chat-only fallback) linux_arm64_with_mlx Linux arm64 + mlx in sys.modules -- canary that the system check still guards against accidental hijack cpu_only Linux x86_64 + nothing -- pure CPU fallback For each profile the suite asserts: 1. unsloth._IS_MLX (re-evaluated under the spoof) matches expectation. 2. utils.hardware.detect_hardware() returns the right DeviceType and IS_ROCM flag. 3. utils.hardware.is_apple_silicon() agrees with the platform spoof. Plus two negative-space canaries: test_cuda_takes_priority_over_mlx_when_both_available With CUDA AND MLX both present, dispatch must pick CUDA. Protects existing GPU users from a future refactor that reorders the dispatch. test_xpu_takes_priority_over_mlx_when_both_available Same canary for Intel/XPU vs MLX. All 23 tests pass on Linux+CUDA in 1.8s with no real hardware required. Future regressions in either the unsloth _IS_MLX gate or Studio's detect_hardware priority order will fail loudly here. --- tests/studio/test_hardware_dispatch_matrix.py | 381 ++++++++++++++++++ 1 file changed, 381 insertions(+) create mode 100644 tests/studio/test_hardware_dispatch_matrix.py diff --git a/tests/studio/test_hardware_dispatch_matrix.py b/tests/studio/test_hardware_dispatch_matrix.py new file mode 100644 index 0000000000..276c57ca5a --- /dev/null +++ b/tests/studio/test_hardware_dispatch_matrix.py @@ -0,0 +1,381 @@ +# SPDX-License-Identifier: AGPL-3.0-only +""" +Comprehensive hardware dispatch matrix for Studio. + +Drives every supported hardware profile from a single test host by +spoofing platform / torch.cuda / torch.xpu / sys.modules['mlx'] so we +can exercise the CUDA, ROCm, XPU, MLX, and CPU dispatch paths +deterministically without real hardware. + +Profiles checked: + + nvidia_cuda Linux x86_64 + torch.cuda.is_available()=True, + torch.version.hip=None + amd_rocm Linux x86_64 + torch.cuda.is_available()=True, + torch.version.hip="6.1" (PyTorch ROCm aliases + torch.cuda.* over HIP) + intel_xpu Linux x86_64 + torch.cuda off, torch.xpu.is_available()=True + apple_silicon_mlx Darwin arm64 + cuda off + xpu off + mlx importable + apple_silicon_no_mlx Darwin arm64 + everything off (no mlx pkg) + linux_arm64_with_mlx Linux arm64 + mlx importable -- gate must NOT activate + (canary against accidental Linux-arm64 hijack) + cpu_only Linux x86_64 + nothing -- pure CPU fallback + +For each profile we assert three contracts: + + 1. ``unsloth._IS_MLX`` (re-evaluated under the spoof). + 2. ``utils.hardware.detect_hardware()`` ``DeviceType`` and ``IS_ROCM``. + 3. ``utils.hardware.is_apple_silicon()``. + +Add a row to ``PROFILES`` to extend coverage; tests parametrize over it +automatically. No real hardware required. +""" + +from __future__ import annotations + +import importlib +import importlib.machinery +import importlib.util +import sys +import types +from dataclasses import dataclass, field +from pathlib import Path +from typing import Optional + +import pytest + + +REPO_ROOT = Path(__file__).resolve().parents[2] +STUDIO_BACKEND = REPO_ROOT / "studio" / "backend" + + +# --------------------------------------------------------------------------- +# Profile definition +# --------------------------------------------------------------------------- + + +@dataclass +class HardwareProfile: + name: str + system: str # platform.system() value + machine: str # platform.machine() value + cuda_available: bool # torch.cuda.is_available() value + hip_version: Optional[ + str + ] # torch.version.hip; None for NVIDIA, "6.1" etc. for ROCm + xpu_available: bool # torch.xpu.is_available() value + has_mlx: bool # whether to inject a fake mlx into sys.modules + mps_available: bool # torch.backends.mps.is_available() value + + expect_is_mlx: bool # unsloth._IS_MLX + expect_device_type: ( + str # Studio DeviceType (uppercased name: "CUDA"/"XPU"/"MLX"/"CPU") + ) + expect_is_rocm: bool # Studio IS_ROCM + expect_apple_silicon: bool # Studio is_apple_silicon() + extra_notes: str = "" + + +PROFILES = [ + HardwareProfile( + name = "nvidia_cuda", + system = "Linux", + machine = "x86_64", + cuda_available = True, + hip_version = None, + xpu_available = False, + has_mlx = False, + mps_available = False, + expect_is_mlx = False, + expect_device_type = "CUDA", + expect_is_rocm = False, + expect_apple_silicon = False, + ), + HardwareProfile( + name = "amd_rocm", + system = "Linux", + machine = "x86_64", + cuda_available = True, + hip_version = "6.1", + xpu_available = False, + has_mlx = False, + mps_available = False, + expect_is_mlx = False, + expect_device_type = "CUDA", + expect_is_rocm = True, + expect_apple_silicon = False, + extra_notes = "PyTorch ROCm reuses torch.cuda.* over HIP; " + "Studio still uses DeviceType.CUDA but flips IS_ROCM=True.", + ), + HardwareProfile( + name = "intel_xpu", + system = "Linux", + machine = "x86_64", + cuda_available = False, + hip_version = None, + xpu_available = True, + has_mlx = False, + mps_available = False, + expect_is_mlx = False, + expect_device_type = "XPU", + expect_is_rocm = False, + expect_apple_silicon = False, + ), + HardwareProfile( + name = "apple_silicon_mlx", + system = "Darwin", + machine = "arm64", + cuda_available = False, + hip_version = None, + xpu_available = False, + has_mlx = True, + mps_available = True, + expect_is_mlx = True, + expect_device_type = "MLX", + expect_is_rocm = False, + expect_apple_silicon = True, + ), + HardwareProfile( + name = "apple_silicon_no_mlx", + system = "Darwin", + machine = "arm64", + cuda_available = False, + hip_version = None, + xpu_available = False, + has_mlx = False, + mps_available = True, + expect_is_mlx = False, + expect_device_type = "CPU", + expect_is_rocm = False, + expect_apple_silicon = True, + extra_notes = "Mac without mlx falls through to CPU (chat-only).", + ), + HardwareProfile( + name = "linux_arm64_with_mlx", + system = "Linux", + machine = "arm64", + cuda_available = False, + hip_version = None, + xpu_available = False, + has_mlx = True, + mps_available = False, + expect_is_mlx = False, + expect_device_type = "CPU", + expect_is_rocm = False, + expect_apple_silicon = False, + extra_notes = "Canary: Linux ARM64 with mlx package installed must NOT " + "trigger MLX dispatch; the system check is what guards it.", + ), + HardwareProfile( + name = "cpu_only", + system = "Linux", + machine = "x86_64", + cuda_available = False, + hip_version = None, + xpu_available = False, + has_mlx = False, + mps_available = False, + expect_is_mlx = False, + expect_device_type = "CPU", + expect_is_rocm = False, + expect_apple_silicon = False, + ), +] + +PROFILE_IDS = [p.name for p in PROFILES] + + +# --------------------------------------------------------------------------- +# Spoofing helpers +# --------------------------------------------------------------------------- + + +@pytest.fixture +def spoof_hardware(monkeypatch): + """Return a function that applies a HardwareProfile to the live process. + + Idempotent: each call re-applies the profile. Cleanup happens + automatically when the test exits via monkeypatch. + """ + + def _apply(profile: HardwareProfile) -> None: + import platform + import torch + + # platform spoof (used by both the unsloth gate and Studio's helpers) + monkeypatch.setattr(platform, "system", lambda: profile.system) + monkeypatch.setattr(platform, "machine", lambda: profile.machine) + + # torch.cuda.is_available + monkeypatch.setattr(torch.cuda, "is_available", lambda: profile.cuda_available) + + # torch.version.hip — None on NVIDIA, "6.1" etc. on ROCm + torch_version = torch.version + monkeypatch.setattr(torch_version, "hip", profile.hip_version, raising = False) + + # torch.xpu.is_available + get_device_name -- detect_hardware reads both. + # Real torch.xpu.get_device_name requires the XPU-compiled torch build, + # so always stub it under the spoof to keep tests hardware-agnostic. + if hasattr(torch, "xpu"): + monkeypatch.setattr( + torch.xpu, "is_available", lambda: profile.xpu_available + ) + monkeypatch.setattr( + torch.xpu, + "get_device_name", + lambda i = 0: "Intel XPU (stub)", + raising = False, + ) + elif profile.xpu_available: + xpu_stub = types.SimpleNamespace( + is_available = lambda: True, + get_device_name = lambda i = 0: "Intel XPU (stub)", + ) + monkeypatch.setattr(torch, "xpu", xpu_stub, raising = False) + + # torch.backends.mps.is_available + if hasattr(torch.backends, "mps"): + monkeypatch.setattr( + torch.backends.mps, "is_available", lambda: profile.mps_available + ) + + # mlx + mlx.core in sys.modules + if profile.has_mlx: + fake_mlx = types.ModuleType("mlx") + fake_mlx.__spec__ = importlib.machinery.ModuleSpec("mlx", loader = None) + fake_mlx.__path__ = [] + fake_mlx_core = types.ModuleType("mlx.core") + fake_mlx.core = fake_mlx_core + monkeypatch.setitem(sys.modules, "mlx", fake_mlx) + monkeypatch.setitem(sys.modules, "mlx.core", fake_mlx_core) + else: + monkeypatch.delitem(sys.modules, "mlx", raising = False) + monkeypatch.delitem(sys.modules, "mlx.core", raising = False) + real_find_spec = importlib.util.find_spec + + def _no_mlx(name, *args, **kwargs): + if name == "mlx": + return None + return real_find_spec(name, *args, **kwargs) + + monkeypatch.setattr(importlib.util, "find_spec", _no_mlx) + + return _apply + + +def _evaluate_unsloth_is_mlx_gate() -> bool: + """Re-evaluate the exact expression from unsloth/__init__.py:20-24.""" + import importlib.util + import platform + + return ( + platform.system() == "Darwin" + and platform.machine() == "arm64" + and importlib.util.find_spec("mlx") is not None + ) + + +def _import_studio_hardware_module(): + """Lazy-load Studio's hardware module under the bare-imports layout.""" + if str(STUDIO_BACKEND) not in sys.path: + sys.path.insert(0, str(STUDIO_BACKEND)) + # Force a fresh import so detect_hardware re-runs under the current spoofs. + sys.modules.pop("utils.hardware.hardware", None) + sys.modules.pop("utils.hardware", None) + from utils.hardware import hardware as hw # type: ignore + + return hw + + +# --------------------------------------------------------------------------- +# Tests +# --------------------------------------------------------------------------- + + +@pytest.mark.parametrize("profile", PROFILES, ids = PROFILE_IDS) +def test_unsloth_is_mlx_gate_matches_profile(profile, spoof_hardware): + """The _IS_MLX expression in unsloth/__init__.py flips correctly per profile.""" + spoof_hardware(profile) + actual = _evaluate_unsloth_is_mlx_gate() + assert actual is profile.expect_is_mlx, ( + f"profile {profile.name}: expected _IS_MLX={profile.expect_is_mlx}, " + f"got {actual}. {profile.extra_notes}" + ) + + +@pytest.mark.parametrize("profile", PROFILES, ids = PROFILE_IDS) +def test_studio_detect_hardware_matches_profile(profile, spoof_hardware): + """Studio's detect_hardware() routes to the right DeviceType per profile.""" + spoof_hardware(profile) + hw = _import_studio_hardware_module() + detected = hw.detect_hardware() + expected = getattr(hw.DeviceType, profile.expect_device_type) + assert detected == expected, ( + f"profile {profile.name}: expected {profile.expect_device_type}, " + f"got {detected!r}. {profile.extra_notes}" + ) + assert hw.IS_ROCM is profile.expect_is_rocm, ( + f"profile {profile.name}: expected IS_ROCM={profile.expect_is_rocm}, " + f"got {hw.IS_ROCM}" + ) + + +@pytest.mark.parametrize("profile", PROFILES, ids = PROFILE_IDS) +def test_studio_is_apple_silicon_matches_profile(profile, spoof_hardware): + """Studio's is_apple_silicon() helper agrees with platform spoof.""" + spoof_hardware(profile) + hw = _import_studio_hardware_module() + assert hw.is_apple_silicon() is profile.expect_apple_silicon, ( + f"profile {profile.name}: expected is_apple_silicon={profile.expect_apple_silicon}, " + f"got {hw.is_apple_silicon()}" + ) + + +# --------------------------------------------------------------------------- +# Negative-space tests: catch regressions where the dispatch order changes. +# --------------------------------------------------------------------------- + + +def test_cuda_takes_priority_over_mlx_when_both_available(spoof_hardware): + """If both CUDA and MLX are available, Studio MUST pick CUDA. This is the + canary that protects every existing GPU user from being silently routed + to MLX after future refactors. + """ + profile = HardwareProfile( + name = "cuda_plus_mlx", + system = "Darwin", + machine = "arm64", + cuda_available = True, + hip_version = None, + xpu_available = False, + has_mlx = True, + mps_available = True, + expect_is_mlx = True, + expect_device_type = "CUDA", + expect_is_rocm = False, + expect_apple_silicon = True, + ) + spoof_hardware(profile) + hw = _import_studio_hardware_module() + assert hw.detect_hardware() == hw.DeviceType.CUDA + + +def test_xpu_takes_priority_over_mlx_when_both_available(spoof_hardware): + """XPU is selected over MLX in the dispatch order.""" + profile = HardwareProfile( + name = "xpu_plus_mlx", + system = "Darwin", + machine = "arm64", + cuda_available = False, + hip_version = None, + xpu_available = True, + has_mlx = True, + mps_available = True, + expect_is_mlx = True, + expect_device_type = "XPU", + expect_is_rocm = False, + expect_apple_silicon = True, + ) + spoof_hardware(profile) + hw = _import_studio_hardware_module() + assert hw.detect_hardware() == hw.DeviceType.XPU From a56c959233863b0cdad90f076c8be7f9a6a3a4c5 Mon Sep 17 00:00:00 2001 From: Daniel Han Date: Wed, 6 May 2026 04:41:57 -0700 Subject: [PATCH 010/159] Add Studio PR-time CI: pin enforcement, frontend, backend, wheel smoke (#5298) * Add Studio PR-time CI: pin enforcement, frontend, backend, wheel smoke The repo currently has no PR-time CI; only release-desktop.yml (manual) and stale.yml (issue pinger). studio/backend/tests/ has 35 test files (~860 tests collected) that never run automatically. Frontend lint/typecheck/build scripts exist in package.json but are not gated on PRs either. This is the gap that let 2026.5.1 ship with the broken Studio chat-history bundle. Adds four ubuntu-latest workflows, all CPU-only and free for public repos: studio-pin-enforce.yml Greps studio/frontend/package.json for caret/tilde ranges on the @assistant-ui surface (and assistant-stream). Blocks the exact regression vector that produced 2026.5.1 (^0.12.19 resolving to a breaking 0.12.28). studio-frontend-ci.yml npm ci (strict lockfile), tree-clean check after, typecheck, vite build, bundle grep for the Studio unstable_Provider call site (<= 3 hits = OK, >= 4 = the 2026.5.1 regression), 75 MB dist budget, biome non-blocking. Uploads dist on failure. studio-backend-ci.yml Runs the existing studio/backend/tests/ suite on Python 3.10/3.11/3.12. Excludes test_studio_api.py (live model + GGUF download) and llama_cpp_load_progress_live (spawns a real llama.cpp). Local run on this branch: 861 pass, 4 skipped, 5 deselected. ruff non-blocking. wheel-smoke.yml python -m build, then verifies the produced wheel: - ships studio/frontend/package-lock.json - ships studio/frontend/dist/index.html - does NOT ship studio/frontend/node_modules/ - does NOT ship studio/frontend/bun.lock - main JS bundle has < 4 unstable_Provider hits Then installs the wheel into a fresh venv with a lightweight dep set and imports studio.backend.main. Locally validated against the wheel built from this branch. Each workflow has concurrency cancellation on the same ref. biome and ruff are gated as non-blocking until the existing accumulated drift is cleared (~470 biome errors today); remove the bypass in a follow-up. Notes verified locally: - pin enforcement: PASS (carets dropped on this branch) - frontend npm ci -> typecheck -> build -> grep -> budget: PASS - bundle: 48 MB, hits=1 - backend pytest: 861 pass, 1 GPU-pollution failure not reproducible on GPU-less runners (won't reproduce on ubuntu-latest) - wheel build: 13s, produces unsloth-2026.5.2-py3-none-any.whl - wheel content sanity: all five checks PASS * CI: install full backend dep set + refine pytest filter for CPU runners First CI run on PR #5298 surfaced two real gaps: 1. pytest collection failed at `import yaml` in utils/models/model_config. Locally my workspace venv had pyyaml from a transitive; CI's clean Python 3.10/3.11/3.12 didn't, so collection hit ModuleNotFoundError on the very first test module. Same blew up the wheel-smoke `from studio.backend.main import app` step. 2. Once the import chain was complete, ~9 tests still failed because they exercise GPU-only paths or live transformers introspection that can't run on a GPU-less `ubuntu-latest` runner regardless of code correctness: - TestGpuAutoSelection - TestPreSpawnGpuResolution - TestPerGpuFitGuardAllCounts - TestTransformersIntrospection - test_returns_cuda_when_cuda_available - test_calls_cuda_cache_when_cuda Fix: - Backend CI installs `studio/backend/requirements/studio.txt` (the declared backend dep set) + the extras the import chain needs but studio.txt omits (python-multipart, sqlalchemy, cryptography, pyyaml, jinja2, mammoth, unpdf, requests, etc.) + torch CPU wheel + transformers. - Refine the pytest -k filter to deselect the GPU/introspection-bound classes by name. Deselections are commented inline with the reason. - wheel-smoke uses the same dep set so the import smoke matches. Locally validated against the freshly-built unsloth-2026.5.2 wheel: 831 passed, 5 skipped, 35 deselected, 0 failed in 47s Studio backend imports cleanly in a fresh venv after the wheel install. * CI: collapse multiline pytest -k expression to a single line YAML's | block-scalar fed the newlines verbatim into the -k argument and pytest rejected it as 'Wrong expression passed to -k'. Same logical filter on one line. * CI: rename jobs so the GitHub UI shows what each check actually does Adds a per-job 'name:' to all four workflows so the PR check list reads: Studio pin enforcement / @assistant-ui must be pinned exactly Studio frontend CI / Frontend build + bundle sanity Studio backend CI / Backend pytest (Python 3.10|3.11|3.12) Studio backend CI / Backend ruff lint (non-blocking) Wheel build + smoke / Wheel build + content sanity + import smoke Instead of the default ' / ' which was opaque ('check', 'build', 'pytest (3.10)', 'ruff', 'wheel'). * CI: add Python 3.13 to backend pytest matrix Verified locally: 831 backend tests pass under Python 3.13 with the same filter set used for 3.10 / 3.11 / 3.12. * CI: add Studio inference smoke + Tauri build smoke Two new workflows. Both CPU-only, both free on `ubuntu-latest`. studio-inference-smoke.yml The only workflow we have that proves "Studio actually works", as opposed to "the bundle parses" or "the imports succeed": - runs install.sh --local --no-torch (lean Studio install) - downloads unsloth/gemma-4-E2B-it-GGUF UD-IQ3_XXS into actions/cache - boots Studio in api-only mode - logs in with the bootstrap password, changes it, re-logs - POST /api/inference/load on the GGUF - POST /api/inference/chat/completions and asserts a non-empty assistant response Validated end-to-end locally on a fresh main install: model loaded, chat completion returned `Hello!` against the same GGUF the workflow uses. studio-tauri-smoke.yml PR-time variant of release-desktop.yml. Linux-only debug build (`tauri build --debug --no-bundle`) on ubuntu-22.04. Catches src-tauri Cargo.toml / Rust source breakage, tauri.conf.json drift, and frontend-distDir wiring. Pinned to the same Tauri CLI version (2.10.1) as release-desktop.yml so CLI bumps surface in CI before they break the release pipeline. Mac and Windows desktop builds stay manual via release-desktop.yml because they need code-signing secrets. * CI: use 'hf download' instead of deprecated 'huggingface-cli download' huggingface_hub 1.13.0 dropped the huggingface-cli entrypoint. The replacement is the 'hf' CLI shipped with the same package. Same args, just s/huggingface-cli/hf/. * CI: assert llama.cpp prebuilt path was used on ubuntu-latest The inference-smoke job runs on ubuntu-latest (CPU-only, x86_64), which is exactly the host shape that should pick up ggml-org/llama.cpp's bin-ubuntu-x64.tar.gz prebuilt directly. If install.sh ever falls back to a source build on this runner, the studio/setup.sh routing has regressed and every CPU-only Linux user is paying a 3 minute compile cost again. Tee install.sh output to logs/install.log, then fail the job if the log contains "falling back to source build" or is missing the success marker "prebuilt installed and validated" / "prebuilt up to date and validated". Also include logs/install.log in the failure artifact so the prebuilt diagnostics are uploaded alongside studio.log when the job fails. * Tighten prebuilt-assertion comment in studio-inference-smoke * CI: switch inference-smoke model to Qwen3.5-2B UD-IQ3_XXS Drops the Gemma 4 E2B GGUF (~2.3 GB) for unsloth/Qwen3.5-2B-GGUF (UD-IQ3_XXS, ~890 MiB). Cache-miss download is roughly a third of what it was, and CPU inference on ubuntu-latest finishes well inside the 25 minute job budget. Verified locally: load via /api/inference/load returns status=loaded, is_gguf=true, supports_reasoning=true, supports_tools=true; chat completion returns a non-empty assistant message ("Hello!"). * CI: add workflow_dispatch to inference-smoke for manual cache pre-warm * CI: fold pin-enforce grep into studio-frontend-ci, drop standalone workflow The "@assistant-ui must be pinned exactly" check was its own ~7 second workflow, doing a single grep on studio/frontend/package.json. Move it into studio-frontend-ci.yml as a pre-install step (right after checkout, before any node setup so a violation fails fast). One fewer top-level check row on every PR, same coverage. Add a FIXME so this step is dropped once @assistant-ui/* and assistant-stream leave 0.x: on 1.x, caret ranges are conventional and this becomes overzealous. * CI: add Repo tests (CPU) job, mirroring unsloth-zoo PR #624 conftest The top-level tests/ tree was previously not run anywhere. 23 of its files are CPU-friendly with the right harness: pure-Python helpers, ast walks, installer logic, and CLI shape tests. Locally validated: 302 passed, 9 skipped, 12 deselected in ~7 seconds on Python 3.12. Three pieces: 1. tests/conftest.py -- GPU-free harness, mirrors the conftest landed in unslothai/unsloth-zoo PR #624. Pre-loads unsloth_zoo.device_type and unsloth.device_type under a temporarily-mocked torch.cuda.is_available() so each module's @cache permanently captures "cuda" and the import chain succeeds on a CPU runner. Also stubs torch.cuda.get_device_capability / is_bf16_supported / mem_get_info, which unsloth/__init__.py and unsloth_zoo.temporary_patches probe at import time when DEVICE_TYPE == "cuda". On a real accelerator the harness is skipped and detection runs normally. 2. Two existing tests were leaking sys.modules state across the session because they injected stubs without an __spec__ and without restoration: - tests/test_raw_text.py shoved a "datasets" stub into sys.modules. transformers' import_utils later did importlib.util.find_spec("datasets") and got ValueError: datasets.__spec__ is None. - tests/python/test_fast_sentence_transformer_redirect_lifecycle.py shoved "transformers", "sentence_transformers", and "sentence_transformers.models" stubs in. Subsequent tests that did `import transformers` got the non-package stub. Fix: set __spec__ on stubs, plus an autouse fixture in the sentence-transformer test file that restores the three keys after each test. 3. .github/workflows/studio-backend-ci.yml gains a third job, `Repo tests (CPU)`, that installs the same dep set as the backend-pytest matrix (Python 3.12 only -- the tests are version-independent), exports PYTHONPATH=studio so tests/python/* can import install_python_stack, and runs the 23-file subset above with `-m 'not server and not e2e'`. * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * CI: install unsloth_zoo for Repo CPU tests, harden conftest fallback The CPU job at run 25422050018 broke at conftest collection: the preload of unsloth.device_type pulled in `from unsloth_zoo.utils import Version` and ubuntu-latest didn't have unsloth_zoo on the path because it is an optional dep of unsloth. Two fixes: 1. Install unsloth_zoo>=2026.5.1 alongside the other deps in the Repo tests (CPU) job (it's also what unsloth's optional `huggingface` extra pins). 2. Wrap the body of _preload_device_type in conftest.py in a try/except so any import failure (missing prereq, broken module, etc.) cleanly returns False instead of aborting the entire collection. The caller already falls back to the stub device_type module on False, so the net behavior is "best effort: real device_type if possible, stub otherwise" instead of "abort the test session". * kernels.utils: guard CUDA_STREAMS / XPU_STREAMS init for DEVICE_COUNT==0 When DEVICE_COUNT is 0 (CPU host: no visible NVIDIA / AMD / Intel GPU) the dict comprehension {... for i in range(0)} was empty and the subsequent max(_CUDA_STREAMS.keys()) raised ValueError: max() iterable argument is empty during module import. That made unsloth.kernels.utils unimportable on any CPU runner, which in turn blocked all of tests/saving/**, three top-level tests/test_*.py, and tests/qlora/test_unsloth_qlora_train_and_merge.py from even collecting on CPU CI. Wrap the per-device-index dict comprehension and max() machinery in a DEVICE_COUNT > 0 guard. When DEVICE_COUNT is 0 fall back to empty containers (CUDA_STREAMS = (), WEIGHT_BUFFERS = [], ABSMAX_BUFFERS = []). The consumer functions further down in this module index these arrays by device_index but only during real GPU work, so the empty fallbacks never get touched on a CPU host. GPU-safety verified locally: with 8 visible CUDA devices, CUDA_STREAMS has 8 entries (identical to before this PR). With CUDA_VISIBLE_DEVICES="" the module imports cleanly, CUDA_STREAMS is (), and the previously blocked tests now collect (test_get_model_name passes 38 subtests, test_resolve_model_class passes 9, test_model_registry collects all 8 parametrizations). Same shape applied to the DEVICE_TYPE == "xpu" branch for symmetry. * CI: switch Repo tests (CPU) to auto-discovery + isolate flakes Three changes, locally validated end-to-end (779 passed, 11 skipped, 23 deselected, 0 failed across all three steps): 1. Repo tests (CPU, auto-discovered): replace the explicit 23-file list with `pytest tests/` plus a small set of `--ignore` and `--deselect` flags. New tests under tests/python, tests/studio (excluding the two state-sensitive files), and top-level tests/test_*.py are picked up automatically with no workflow edit. --ignore covers: - tests/qlora and tests/saving: GPU-bound by design - tests/utils: helpers folder, not tests - tests/sh: shell suite handled in its own step - two state-polluting hardware-spoof files (next step) -m 'not server and not e2e': honours markers already declared in tests/python/conftest.py --deselect: test_model_registration / test_all_model_registration hit huggingface_hub live; they belong on a network job 2. Hardware-spoof tests (state-sensitive, run in isolation): tests/studio/test_hardware_dispatch_matrix.py and tests/studio/test_is_mlx_dispatch_gate.py mutate module globals in studio.backend.utils.hardware.hardware (IS_ROCM, DEVICE) via their spoof fixtures, and the leak crosses file boundaries. Running them in their own pytest invocation avoids polluting the main sweep. Both pass cleanly in isolation: 28 passed, 1 skipped. 3. Shell installer tests: explicitly enumerated subset that does not depend on install.ps1 layout (test_install_host_defaults.sh has drifted; that's a separate followup). Test fixes folded in to keep the run green: - tests/studio/install/test_rocm_support.py::TestAmdGpuMonitoring ::test_amd_primary_gpu_with_mock now clears HIP/ROCR/CUDA_VISIBLE_DEVICES via monkeypatch so _first_visible_amd_gpu_id() does not short-circuit when the runner sets CUDA_VISIBLE_DEVICES="" to suppress CUDA. - tests/studio/test_hardware_dispatch_matrix.py::spoof_hardware fixture now stubs torch.cuda.get_device_properties when cuda_available is True so detect_hardware()'s device_name probe does not call into _cuda_init() on a CPU runner. * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * CI: install torchvision (CPU) so unsloth_zoo.vision_utils can import Run 25430652224 collected three test modules that import unsloth and crashed at unsloth_zoo/vision_utils.py:68 with ModuleNotFoundError: No module named 'torchvision' unsloth_zoo.vision_utils unconditionally imports torchvision at module scope, and unsloth.models._utils pulls vision_utils in. The Repo tests (CPU) job installed torch from the CPU index but not torchvision, so any test that imports unsloth.models.* failed at collection. Add torchvision<0.26 to the same pip install --index-url https://download.pytorch.org/whl/cpu line. * CI: install bitsandbytes (CPU build) for unsloth.models._utils import Run 25430982243 collected three test modules that import unsloth and crashed at unsloth/models/_utils.py:1166 with ModuleNotFoundError: No module named 'bitsandbytes' The bnb import there is unconditional. Recent bnb versions (>=0.45) ship a CPU build so the wheel installs on a free Linux runner and the import resolves; the kernels still raise on use but the module collects, which is enough for these CPU tests. Add 'bitsandbytes>=0.45' to the Repo tests (CPU) deps. * CI: rename workflows + guard kernels.utils CPU-torch binding Workflow renames (top-level `name:` keys; affects PR check rows): Studio backend CI -> Backend CI Studio frontend CI -> Frontend CI Studio inference smoke -> Studio GGUF CI Studio Tauri smoke -> Studio Tauri CI Wheel build + smoke -> Wheel CI Backend CI's matrix job goes from "Backend pytest (Python 3.10)" to just "(Python 3.10)" so the GitHub UI row reads "Backend CI / (Python 3.10)" rather than the old verbose form. Production guard for CPU torch (run 25431126138): unsloth/kernels/utils.py:165 was an unconditional _gpu_getCurrentRawStream = torch._C._cuda_getCurrentRawStream which raised AttributeError on a CPU-only torch wheel because the compiled CUDA backend is absent. Three test modules (test_get_model_name, test_model_registry, test_resolve_model_class) crashed at collection because their import chain reaches this line. Add a hasattr probe: when torch is built without CUDA, fall through to a no-op binding that returns 0. _get_tensor_stream is only invoked during real GPU work, so the no-op is never executed on a CPU host. GPU-safety verified locally: with 8 visible CUDA devices the binding still resolves to the real torch._C._cuda_getCurrentRawStream (behaviour identical to before this PR). The XPU branch is untouched. * [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> --- .github/workflows/studio-backend-ci.yml | 200 ++++++++++++++++++ .github/workflows/studio-frontend-ci.yml | 108 ++++++++++ .github/workflows/studio-inference-smoke.yml | 185 ++++++++++++++++ .github/workflows/studio-tauri-smoke.yml | 105 +++++++++ .github/workflows/wheel-smoke.yml | 124 +++++++++++ tests/conftest.py | 141 ++++++++++++ ...sentence_transformer_redirect_lifecycle.py | 40 +++- tests/studio/install/test_rocm_support.py | 13 +- tests/studio/test_hardware_dispatch_matrix.py | 14 ++ tests/test_raw_text.py | 5 +- unsloth/kernels/utils.py | 75 ++++--- 11 files changed, 977 insertions(+), 33 deletions(-) create mode 100644 .github/workflows/studio-backend-ci.yml create mode 100644 .github/workflows/studio-frontend-ci.yml create mode 100644 .github/workflows/studio-inference-smoke.yml create mode 100644 .github/workflows/studio-tauri-smoke.yml create mode 100644 .github/workflows/wheel-smoke.yml create mode 100644 tests/conftest.py diff --git a/.github/workflows/studio-backend-ci.yml b/.github/workflows/studio-backend-ci.yml new file mode 100644 index 0000000000..5a858888e7 --- /dev/null +++ b/.github/workflows/studio-backend-ci.yml @@ -0,0 +1,200 @@ +# SPDX-License-Identifier: AGPL-3.0-only +# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. + +# Runs the existing studio/backend/tests/ suite (~860 tests, all CPU-friendly) +# on every PR that touches the backend or unsloth library. Until this lands, +# none of those tests run automatically. Verified locally on Python 3.13 with +# the surgical exclusions below: 861 pass, 4 skipped. +# +# Exclusions: +# - tests/test_studio_api.py: end-to-end against a live model + GGUF download, +# too heavy for free runners. Run separately when GPU CI is available. +# - -k 'not llama_cpp_load_progress_live': spawns a real llama.cpp process, +# not appropriate for CPU-only runners. +# +# ruff is non-blocking initially; remove `|| true` once the backend lints clean. + +name: Backend CI + +on: + pull_request: + paths: + - 'studio/**' + - 'unsloth/**' + - 'unsloth_cli/**' + - 'tests/**' + - 'pyproject.toml' + - '.github/workflows/studio-backend-ci.yml' + push: + branches: [main, pip] + +concurrency: + group: ${{ github.workflow }}-${{ github.ref }} + cancel-in-progress: true + +jobs: + pytest: + name: (Python ${{ matrix.python }}) + runs-on: ubuntu-latest + timeout-minutes: 15 + strategy: + fail-fast: false + matrix: + python: ['3.10', '3.11', '3.12', '3.13'] + steps: + - uses: actions/checkout@v4 + + - uses: actions/setup-python@v5 + with: + python-version: '${{ matrix.python }}' + cache: 'pip' + + - name: Install backend test dependencies (CPU only) + run: | + python -m pip install --upgrade pip + # Studio's declared backend deps: + pip install -r studio/backend/requirements/studio.txt + # Extras that studio.txt does not list but the import chain needs + # (python-multipart for FastAPI form/file uploads, sqlalchemy/cryptography + # for the auth DB, yaml/jinja2 for utils.models.model_config, etc.): + pip install \ + python-multipart aiofiles sqlalchemy cryptography \ + pyyaml jinja2 mammoth unpdf requests \ + 'numpy<3' pytest pytest-asyncio httpx + # Torch CPU + transformers are required by a chunk of the backend test + # suite (gpu_selection, kv_cache_estimation, utils). CPU-only torch + # keeps the install ~250 MB / ~1 min on a clean runner. + pip install --index-url https://download.pytorch.org/whl/cpu 'torch>=2.4,<2.11' + pip install 'transformers>=4.51,<5.5' + + - name: Backend tests + working-directory: studio/backend + # Locally validated against this dep set: 831 passed, 5 skipped, 35 deselected. + # Deselections (all environment-specific, would never pass on a GPU-less + # `ubuntu-latest` runner regardless of code correctness): + # - llama_cpp_load_progress_live: spawns a real llama.cpp process + # - TestGpuAutoSelection / TestPreSpawnGpuResolution / TestPerGpuFitGuardAllCounts: + # require live transformers config introspection on real GPUs + # - TestTransformersIntrospection: same + # - test_returns_cuda_when_cuda_available / test_calls_cuda_cache_when_cuda: + # assume CUDA-capable GPU + run: | + python -m pytest tests/ -q --tb=short \ + --ignore=tests/test_studio_api.py \ + -k 'not llama_cpp_load_progress_live and not TestGpuAutoSelection and not TestPreSpawnGpuResolution and not TestPerGpuFitGuardAllCounts and not TestTransformersIntrospection and not test_returns_cuda_when_cuda_available and not test_calls_cuda_cache_when_cuda' + + repo-cpu-tests: + # Auto-discover everything under tests/ that is not GPU-bound by + # design. New tests added in covered directories are picked up + # without a workflow edit. Locally validated: 779 passed, 11 + # skipped, 23 deselected. tests/conftest.py (mirroring unsloth-zoo + # PR #624) pre-loads unsloth_zoo.device_type and unsloth.device_type + # under a mocked torch.cuda.is_available so the unsloth import + # chain succeeds on CPU. + name: Repo tests (CPU) + runs-on: ubuntu-latest + timeout-minutes: 10 + steps: + - uses: actions/checkout@v4 + + - uses: actions/setup-python@v5 + with: + python-version: '3.12' + cache: 'pip' + + - name: Install deps (shared shape with backend pytest job) + run: | + python -m pip install --upgrade pip + pip install -r studio/backend/requirements/studio.txt + pip install \ + python-multipart aiofiles sqlalchemy cryptography \ + pyyaml jinja2 mammoth unpdf requests typer \ + 'numpy<3' pytest pytest-asyncio httpx + # torchvision is needed because unsloth_zoo.vision_utils imports + # it at module scope and is reached via unsloth.models._utils. + pip install --index-url https://download.pytorch.org/whl/cpu \ + 'torch>=2.4,<2.11' 'torchvision<0.26' + pip install 'transformers>=4.51,<5.5' + # bitsandbytes is a hard import in unsloth/models/_utils.py. + # Recent versions ship a CPU build so it installs on a free + # Linux runner; the kernels still raise on use, but import + # succeeds and the package collects. + pip install 'bitsandbytes>=0.45' + # unsloth.device_type imports unsloth_zoo.utils.Version at module + # scope, so the conftest harness needs unsloth_zoo on the path + # even though it is an optional dep of unsloth. + pip install 'unsloth_zoo>=2026.5.1' + pip install -e . --no-deps + + - name: Repo tests (CPU, auto-discovered) + env: + # tests/python/* import install_python_stack from studio/. + PYTHONPATH: ${{ github.workspace }}/studio + # Skip lazy compilation work the unsloth import chain wants to + # do at import time on a real GPU. + UNSLOTH_COMPILE_DISABLE: '1' + # --ignore: GPU-bound directories (qlora and saving need real + # weights / GPU; tests/sh is a shell suite the next step + # handles; tests/utils is a helpers folder, not tests). + # State-sensitive hardware-spoofing files are pulled out and run + # in isolation in the next step because they mutate + # hardware.py module globals (IS_ROCM / DEVICE) and pollute + # downstream tests. + # -m: honour markers already declared in tests/python/conftest.py + # (`server` = needs studio venv, `e2e` = needs network). + # --deselect: two registry tests that hit huggingface_hub for + # live model existence checks; they belong on a network job. + run: | + python -m pytest tests/ -q --tb=short \ + --ignore=tests/qlora \ + --ignore=tests/saving \ + --ignore=tests/utils \ + --ignore=tests/sh \ + --ignore=tests/studio/test_hardware_dispatch_matrix.py \ + --ignore=tests/studio/test_is_mlx_dispatch_gate.py \ + -m 'not server and not e2e' \ + --deselect tests/test_model_registry.py::test_model_registration \ + --deselect tests/test_model_registry.py::test_all_model_registration + + - name: Hardware-spoof tests (state-sensitive, run in isolation) + env: + PYTHONPATH: ${{ github.workspace }}/studio + UNSLOTH_COMPILE_DISABLE: '1' + # These two files mutate hardware.py module globals at runtime + # via the spoof fixtures, which leaks state into any other test + # that imports hardware. Run them in their own pytest invocation + # so the leak does not cross file boundaries. + run: | + python -m pytest -q --tb=short \ + tests/studio/test_hardware_dispatch_matrix.py \ + tests/studio/test_is_mlx_dispatch_gate.py + + - name: Shell installer tests + # Subset that does not depend on a writable / pristine install.sh + # tree; test_install_host_defaults.sh checks install.ps1 layout + # which has drifted (separate followup). + run: | + set -e + for s in \ + tests/sh/test_get_torch_index_url.sh \ + tests/sh/test_mac_intel_compat.sh \ + tests/sh/test_tauri_install_exit_order.sh \ + tests/sh/test_torch_constraint.sh; do + echo "::group::$s" + bash "$s" + echo "::endgroup::" + done + + ruff: + name: Backend ruff lint (non-blocking) + runs-on: ubuntu-latest + timeout-minutes: 5 + steps: + - uses: actions/checkout@v4 + - uses: actions/setup-python@v5 + with: + python-version: '3.12' + cache: 'pip' + - run: pip install ruff + - name: ruff check (non-blocking until accumulated drift is cleared) + run: ruff check studio/backend || true diff --git a/.github/workflows/studio-frontend-ci.yml b/.github/workflows/studio-frontend-ci.yml new file mode 100644 index 0000000000..039bd5dd08 --- /dev/null +++ b/.github/workflows/studio-frontend-ci.yml @@ -0,0 +1,108 @@ +# SPDX-License-Identifier: AGPL-3.0-only +# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. + +# Frontend PR gate: lockfile freshness, typecheck, build, and a bundle grep +# that catches the 2026.5.1 chat-history regression at the JS level. +# +# biome runs as non-blocking for now: the codebase currently has accumulated +# ~470 errors and ~1650 warnings against the existing biome config. Surfacing +# the count in CI lets us drive it down without forcing a fleet-wide cleanup +# in the same PR. Drop `continue-on-error` once that number is zero. + +name: Frontend CI + +on: + pull_request: + paths: + - 'studio/frontend/**' + - '.github/workflows/studio-frontend-ci.yml' + push: + branches: [main, pip] + +concurrency: + group: ${{ github.workflow }}-${{ github.ref }} + cancel-in-progress: true + +jobs: + build: + name: Frontend build + bundle sanity + runs-on: ubuntu-latest + timeout-minutes: 10 + defaults: + run: + working-directory: studio/frontend + steps: + - uses: actions/checkout@v4 + + # FIXME: drop this step once @assistant-ui/* and assistant-stream + # leave 0.x -- on 1.x, caret ranges are conventional. Until then, + # every 0.minor on this surface is a SemVer-major (this is exactly + # how 2026.5.1 shipped a broken chat runtime: ^0.12.19 quietly + # resolved to 0.12.28). + - name: '@assistant-ui must be pinned exactly (no caret/tilde)' + working-directory: ${{ github.workspace }} + run: | + set -e + if grep -nE '"(@assistant-ui/[a-z-]+|assistant-stream)":[[:space:]]*"[\^~]' studio/frontend/package.json; then + echo "::error file=studio/frontend/package.json::These packages must be pinned to exact versions until they leave 0.x. Drop the leading ^ or ~." + exit 1 + fi + echo "All assistant-ui packages are pinned exactly." + + - uses: actions/setup-node@v4 + with: + node-version: '22' + cache: 'npm' + cache-dependency-path: studio/frontend/package-lock.json + + - name: Lockfile must agree with package.json (npm ci is strict) + run: npm ci --no-fund --no-audit + + - name: npm ci must not have modified the working tree + working-directory: ${{ github.workspace }} + run: | + if ! git diff --quiet -- studio/frontend; then + echo "::error::npm ci modified files; commit the updated lockfile" + git status -- studio/frontend + exit 1 + fi + + - name: Typecheck + run: npm run typecheck + + - name: Build + run: npm run build + + - name: Built bundle must not contain Studio's unstable_Provider call site + run: | + set -e + JS=$(ls dist/assets/index-*.js | head -1) + HITS=$(grep -c 'unstable_Provider:' "$JS" || echo 0) + echo "main bundle: $JS" + echo "unstable_Provider: hits=$HITS (assistant-ui internals contribute up to 3)" + if [ "$HITS" -gt 3 ]; then + echo "::error file=studio/frontend/src/features/chat/runtime-provider.tsx::Studio bundle still passes unstable_Provider through useRemoteThreadListRuntime; this is the 2026.5.1 chat-history regression. Pass adapters directly into useLocalRuntime instead." + exit 1 + fi + + - name: Bundle size budget (75 MB) + run: | + SIZE=$(du -sb dist | cut -f1) + BUDGET=$((75 * 1024 * 1024)) + echo "dist size: $SIZE bytes ($((SIZE/1024/1024)) MB), budget: $BUDGET bytes (75 MB)" + if [ "$SIZE" -gt "$BUDGET" ]; then + echo "::error::studio/frontend/dist/ exceeded the 75 MB budget. Drop dead deps (e.g. the unused next dep) or split chunks." + exit 1 + fi + + - name: Biome (non-blocking until accumulated drift is cleared) + continue-on-error: true + run: npm run biome:check + + - name: Upload built dist on failure + if: failure() + uses: actions/upload-artifact@v4 + with: + name: studio-frontend-dist + path: studio/frontend/dist + retention-days: 3 diff --git a/.github/workflows/studio-inference-smoke.yml b/.github/workflows/studio-inference-smoke.yml new file mode 100644 index 0000000000..8efe072d28 --- /dev/null +++ b/.github/workflows/studio-inference-smoke.yml @@ -0,0 +1,185 @@ +# SPDX-License-Identifier: AGPL-3.0-only +# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. + +# End-to-end smoke: install Studio via install.sh --local --no-torch, download +# a tiny GGUF, boot Studio, log in, change password, load the model, send a +# chat completion, assert a non-empty response. Only workflow that tests "the +# app actually works". +# +# Model: Qwen3.5-2B UD-IQ3_XXS (~890 MiB) -- small enough that the cache miss +# is cheap and inference fits in the 25 min CPU-runner budget. GGUF is cached +# across runs via actions/cache. + +name: Studio GGUF CI + +on: + pull_request: + paths: + - 'studio/**' + - 'unsloth/**' + - 'unsloth_cli/**' + - 'install.sh' + - 'pyproject.toml' + - '.github/workflows/studio-inference-smoke.yml' + push: + branches: [main, pip] + # Manual trigger for pre-warming the GGUF cache on main, or re-running + # against an arbitrary branch without pushing a no-op commit. + workflow_dispatch: + +concurrency: + group: ${{ github.workflow }}-${{ github.ref }} + cancel-in-progress: true + +env: + GGUF_REPO: unsloth/Qwen3.5-2B-GGUF + GGUF_FILE: Qwen3.5-2B-UD-IQ3_XXS.gguf + STUDIO_PORT: '18888' + +jobs: + inference: + name: Studio boots, loads a GGUF, answers a chat completion + runs-on: ubuntu-latest + timeout-minutes: 25 + steps: + - uses: actions/checkout@v4 + + - name: Linux dependencies for llama.cpp prebuilt + run: | + sudo apt-get update + sudo apt-get install -y --no-install-recommends \ + libcurl4-openssl-dev libssl-dev jq + + - uses: actions/setup-node@v4 + with: + node-version: '22' + cache: 'npm' + cache-dependency-path: studio/frontend/package-lock.json + + - uses: actions/setup-python@v5 + with: + python-version: '3.12' + cache: 'pip' + + - name: Cache GGUF model file + id: cache-gguf + uses: actions/cache@v4 + with: + path: gguf-cache + key: ${{ runner.os }}-gguf-${{ env.GGUF_REPO }}-${{ env.GGUF_FILE }}-v1 + + - name: Download GGUF if cache miss + if: steps.cache-gguf.outputs.cache-hit != 'true' + run: | + # huggingface-cli was deprecated in huggingface_hub 1.13; the new CLI is `hf`. + python -m pip install --upgrade huggingface_hub hf_transfer + mkdir -p gguf-cache + HF_HUB_ENABLE_HF_TRANSFER=1 \ + hf download "$GGUF_REPO" "$GGUF_FILE" --local-dir gguf-cache + + - name: Install Studio (--local, --no-torch keeps the install lean) + run: | + mkdir -p logs + set -o pipefail + bash install.sh --local --no-torch 2>&1 | tee logs/install.log + + - name: Assert llama.cpp prebuilt was installed (no source-build fallback) + # ubuntu-latest is CPU-only x86_64, so studio/setup.sh should route + # to ggml-org/llama.cpp and grab bin-ubuntu-x64.tar.gz. A source + # build here means the routing regressed. + run: | + if grep -q "falling back to source build" logs/install.log; then + echo "::error::llama.cpp prebuilt path failed on ubuntu-latest. studio/setup.sh routing regressed; CPU-only Linux x86_64 should hit ggml-org/llama.cpp's bin-ubuntu-x64.tar.gz." + grep -E "llama-prebuilt|llama.cpp" logs/install.log | tail -60 + exit 1 + fi + if ! grep -qE "prebuilt installed and validated|prebuilt up to date and validated" logs/install.log; then + echo "::error::install.log does not contain the success marker for the llama.cpp prebuilt path. Did setup.sh skip the prebuilt install?" + grep -E "llama-prebuilt|llama.cpp" logs/install.log | tail -60 + exit 1 + fi + echo "llama.cpp prebuilt path used successfully" + + - name: Reset auth + start Studio in the background + run: | + unsloth studio reset-password + mkdir -p logs + UNSLOTH_API_ONLY=1 unsloth studio -H 127.0.0.1 -p "$STUDIO_PORT" \ + > logs/studio.log 2>&1 & + echo "STUDIO_PID=$!" >> "$GITHUB_ENV" + + - name: Wait for /api/health + run: | + for i in $(seq 1 60); do + if curl -fs "http://127.0.0.1:${STUDIO_PORT}/api/health" > /tmp/health.json; then + echo "ready after ${i}s" + cat /tmp/health.json + jq -e '.status == "healthy"' /tmp/health.json + exit 0 + fi + sleep 1 + done + echo "Studio did not become healthy in 60s" + tail -200 logs/studio.log + exit 1 + + - name: Login + change bootstrap password + run: | + PW=$(cat ~/.unsloth/studio/auth/.bootstrap_password) + NEW="CIPasswordSmoke12345!" + TOKEN=$(curl -fs -X POST "http://127.0.0.1:${STUDIO_PORT}/api/auth/login" \ + -H 'content-type: application/json' \ + -d "{\"username\":\"unsloth\",\"password\":\"$PW\"}" | jq -r .access_token) + curl -fs -X POST "http://127.0.0.1:${STUDIO_PORT}/api/auth/change-password" \ + -H "Authorization: Bearer $TOKEN" -H 'content-type: application/json' \ + -d "{\"current_password\":\"$PW\",\"new_password\":\"$NEW\"}" > /dev/null + # Re-login to clear must_change_password flag. + NEW_TOKEN=$(curl -fs -X POST "http://127.0.0.1:${STUDIO_PORT}/api/auth/login" \ + -H 'content-type: application/json' \ + -d "{\"username\":\"unsloth\",\"password\":\"$NEW\"}" | jq -r .access_token) + echo "TOKEN=$NEW_TOKEN" >> "$GITHUB_ENV" + + - name: Load the GGUF into Studio + run: | + GGUF_PATH="$GITHUB_WORKSPACE/gguf-cache/${GGUF_FILE}" + ls -lh "$GGUF_PATH" + curl -fs -X POST "http://127.0.0.1:${STUDIO_PORT}/api/inference/load" \ + -H "Authorization: Bearer $TOKEN" -H 'content-type: application/json' \ + --max-time 600 \ + -d "{\"model_path\":\"$GGUF_PATH\",\"is_lora\":false,\"max_seq_length\":2048}" \ + | jq '{status, display_name, is_gguf, context_length}' + + - name: Send a chat completion + assert non-empty response + run: | + RESP=$(curl -fs -X POST "http://127.0.0.1:${STUDIO_PORT}/api/inference/chat/completions" \ + -H "Authorization: Bearer $TOKEN" -H 'content-type: application/json' \ + --max-time 900 \ + -d '{ + "messages":[{"role":"user","content":"Say hello in one short sentence."}], + "max_tokens":40, + "stream":false + }') + echo "raw response: $RESP" + CONTENT=$(echo "$RESP" | jq -r '.choices[0].message.content // empty') + echo "model response: $CONTENT" + if [ -z "$CONTENT" ]; then + echo "::error::Empty assistant response from Studio" + exit 1 + fi + + - name: Stop Studio + if: always() + run: | + kill "${STUDIO_PID}" || true + sleep 2 + ss -tln | grep ":${STUDIO_PORT}" || true + + - name: Upload Studio + install logs on failure + if: failure() + uses: actions/upload-artifact@v4 + with: + name: studio-inference-log + path: | + logs/studio.log + logs/install.log + retention-days: 7 diff --git a/.github/workflows/studio-tauri-smoke.yml b/.github/workflows/studio-tauri-smoke.yml new file mode 100644 index 0000000000..fcc9c8d963 --- /dev/null +++ b/.github/workflows/studio-tauri-smoke.yml @@ -0,0 +1,105 @@ +# SPDX-License-Identifier: AGPL-3.0-only +# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. + +# PR-time smoke for the Tauri desktop wrapper. Builds the frontend and the +# Tauri Linux debug binary, with no codesigning. Catches: +# - tauri.conf.json drift +# - src-tauri Cargo.toml or rust source breakage +# - Tauri CLI version drift (we pin 2.10.1, matching release-desktop.yml) +# - frontend output not picked up by Tauri's distDir +# +# Linux-only on a free `ubuntu-latest` runner. Mac and Windows desktop builds +# stay in release-desktop.yml (manual `workflow_dispatch`) because they need +# code-signing secrets and ~30 min of runner time each. + +name: Studio Tauri CI + +on: + pull_request: + paths: + - 'studio/frontend/**' + - 'studio/src-tauri/**' + - '.github/workflows/studio-tauri-smoke.yml' + push: + branches: [main, pip] + +concurrency: + group: ${{ github.workflow }}-${{ github.ref }} + cancel-in-progress: true + +jobs: + linux-debug-build: + name: Tauri Linux debug build (no codesign) + runs-on: ubuntu-22.04 + timeout-minutes: 25 + steps: + - uses: actions/checkout@v4 + + - name: Linux native deps for Tauri / WebKit2GTK + run: | + sudo apt-get update + sudo apt-get install -y \ + libwebkit2gtk-4.1-dev libayatana-appindicator3-dev \ + librsvg2-dev libxdo-dev libssl-dev patchelf + + - uses: actions/setup-node@v4 + with: + node-version: '24' + cache: 'npm' + cache-dependency-path: studio/frontend/package-lock.json + + - uses: dtolnay/rust-toolchain@stable + + - uses: swatinem/rust-cache@v2 + with: + workspaces: studio/src-tauri -> target + + - name: Install pinned Tauri CLI (matches release-desktop.yml) + run: npm install --save-dev --prefix studio @tauri-apps/cli@2.10.1 + + - name: Verify pinned Tauri CLI version + run: | + out="$(npx --prefix studio tauri --version)" + echo "$out" + [ "$out" = "tauri-cli 2.10.1" ] || { echo "::error::expected tauri-cli 2.10.1, got $out"; exit 1; } + + - name: Frontend build (npm ci, vite) + working-directory: studio/frontend + run: | + npm ci --no-fund --no-audit + npm run build + test -f dist/index.html + + - name: Tauri debug build (Linux, no bundle, no codesign) + # `--debug` + `--no-bundle` keeps this lean: compiles the Rust crate, + # confirms the frontend dist is wired into Tauri, but skips the AppImage + # / .deb production. Code signing is irrelevant because we never produce + # a distributable artifact. + env: + TAURI_SIGNING_PRIVATE_KEY: '' + TAURI_SIGNING_PRIVATE_KEY_PASSWORD: '' + run: npx --prefix studio tauri build --debug --no-bundle + + - name: Inspect produced binary + run: | + BIN=$(find studio/src-tauri/target/debug -maxdepth 1 -type f -executable 2>/dev/null \ + | grep -Ev '\.(d|so|dylib|dll)$' \ + | grep -Ev '/(deps|build|examples)$' \ + | head -1) + echo "binary: $BIN" + if [ -z "$BIN" ]; then + echo "::error::Tauri debug binary not produced" + ls -la studio/src-tauri/target/debug/ || true + exit 1 + fi + file "$BIN" + du -h "$BIN" + + - uses: actions/upload-artifact@v4 + if: failure() + with: + name: tauri-debug-build + path: | + studio/src-tauri/target/debug + studio/frontend/dist + retention-days: 3 diff --git a/.github/workflows/wheel-smoke.yml b/.github/workflows/wheel-smoke.yml new file mode 100644 index 0000000000..080a6bb261 --- /dev/null +++ b/.github/workflows/wheel-smoke.yml @@ -0,0 +1,124 @@ +# SPDX-License-Identifier: AGPL-3.0-only +# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. + +# Builds the PyPI wheel from the PR branch, then verifies the built wheel +# actually contains what we expect to ship and does NOT contain the broken +# Studio bundle that 2026.5.1 published. This is the single workflow that +# would have blocked the 2026.5.1 release before twine upload. +# +# Verified locally end-to-end against this branch: +# - python -m build produces unsloth--py3-none-any.whl in 13s +# - wheel content sanity passes: +# lockfile shipped, frontend dist shipped, +# no node_modules in wheel, no bun.lock in wheel, +# main bundle has unstable_Provider hits=1 (assistant-ui internals only). +# - Studio backend imports cleanly from the installed wheel with the +# lightweight dep set below. + +name: Wheel CI + +on: + pull_request: + paths: + - 'pyproject.toml' + - 'studio/**' + - 'unsloth/**' + - 'unsloth_cli/**' + - '.github/workflows/wheel-smoke.yml' + push: + branches: [main, pip] + +concurrency: + group: ${{ github.workflow }}-${{ github.ref }} + cancel-in-progress: true + +jobs: + wheel: + name: Wheel build + content sanity + import smoke + runs-on: ubuntu-latest + timeout-minutes: 15 + steps: + - uses: actions/checkout@v4 + + - uses: actions/setup-node@v4 + with: + node-version: '22' + cache: 'npm' + cache-dependency-path: studio/frontend/package-lock.json + + - uses: actions/setup-python@v5 + with: + python-version: '3.12' + + - name: Build frontend + run: | + cd studio/frontend + npm ci --no-fund --no-audit + npm run build + + - name: Build wheel + sdist + run: | + python -m pip install --upgrade pip build + rm -rf dist build ./*.egg-info + python -m build + + - name: Wheel content sanity + run: | + python - <<'PY' + import zipfile, glob, sys + w = glob.glob("dist/unsloth-*.whl") + if not w: + print("FAIL: no wheel produced"); sys.exit(2) + w = w[0] + print(f"wheel: {w}") + with zipfile.ZipFile(w) as z: + n = z.namelist() + checks = { + "lockfile shipped": any(s.endswith("studio/frontend/package-lock.json") for s in n), + "frontend dist shipped": any(s.endswith("studio/frontend/dist/index.html") for s in n), + "no node_modules": not any("studio/frontend/node_modules/" in s for s in n), + "no bun.lock": not any(s.endswith("studio/frontend/bun.lock") for s in n), + } + js = [s for s in n + if "studio/frontend/dist/assets/" in s + and s.endswith(".js") + and "/index-" in s] + if not js: + print("FAIL: no main bundle index-*.js in wheel"); sys.exit(2) + data = z.read(js[0]).decode("utf-8", "replace") + hits = data.count("unstable_Provider:") + print(f"main bundle: {js[0]}") + print(f"unstable_Provider hits: {hits} (>=4 indicates 2026.5.1 regression)") + checks["bundle has no Studio unstable_Provider call site"] = (hits < 4) + + print() + for k, v in checks.items(): + print(f" [{'PASS' if v else 'FAIL'}] {k}") + sys.exit(0 if all(checks.values()) else 1) + PY + + - name: Studio backend import smoke + # Imports `studio.backend.main:app` from the freshly-installed wheel in + # a clean venv. This catches the class of bug that 2026.5.1 shipped with: + # frontend dist missing, package-lock.json missing, or the wheel's Python + # source tree broken in a way that surfaces only at app construction time. + run: | + python -m venv /tmp/v + /tmp/v/bin/pip install --upgrade pip + /tmp/v/bin/pip install -r studio/backend/requirements/studio.txt + /tmp/v/bin/pip install \ + python-multipart aiofiles sqlalchemy cryptography \ + pyyaml jinja2 mammoth unpdf requests \ + 'numpy<3' + /tmp/v/bin/pip install --no-deps dist/unsloth-*.whl + # Run from /tmp so Python imports the installed package, not the source tree. + cd /tmp + /tmp/v/bin/python -c "from studio.backend.main import app; print('Studio backend OK:', app.title)" + + - name: Upload wheel on failure + if: failure() + uses: actions/upload-artifact@v4 + with: + name: unsloth-wheel + path: dist/ + retention-days: 7 diff --git a/tests/conftest.py b/tests/conftest.py new file mode 100644 index 0000000000..de41c50fc1 --- /dev/null +++ b/tests/conftest.py @@ -0,0 +1,141 @@ +# SPDX-License-Identifier: AGPL-3.0-only +# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. + +"""GPU-free test harness. + +unsloth's import chain hits unsloth_zoo.device_type, which calls +get_device_type() at import time and raises NotImplementedError on CI +runners with no CUDA / XPU / HIP visible. Pre-load the real +unsloth_zoo.device_type under a temporarily-mocked +torch.cuda.is_available() so its @cache permanently captures "cuda". +On a real accelerator the pre-load is skipped and detection runs +normally. + +Mirrors the conftest harness in unslothai/unsloth-zoo PR #624. +""" + +from __future__ import annotations + +import importlib.util +import os +import sys +import types + + +def _has_real_accelerator() -> bool: + try: + import torch + except Exception: + return False + for probe in ( + lambda: hasattr(torch, "cuda") and torch.cuda.is_available(), + lambda: hasattr(torch, "xpu") and torch.xpu.is_available(), + lambda: hasattr(torch, "accelerator") and torch.accelerator.is_available(), + ): + try: + if probe(): + return True + except Exception: + pass + return False + + +def _preload_device_type(package: str, prereqs: tuple[str, ...] = ()) -> bool: + """Pre-load .device_type under a mocked + torch.cuda.is_available() == True so its @cache permanently + captures "cuda". prereqs lists submodule names of that + must be loaded first (e.g. 'utils' for unsloth_zoo). Returns False + if the package or any prerequisite cannot be imported, in which + case the caller falls back to a stub.""" + target = f"{package}.device_type" + if target in sys.modules: + return True + pkg_spec = importlib.util.find_spec(package) + if pkg_spec is None or not pkg_spec.submodule_search_locations: + return False + pkg_path = pkg_spec.submodule_search_locations[0] + + skeleton_already = package in sys.modules + if not skeleton_already: + skel = types.ModuleType(package) + skel.__path__ = [pkg_path] + skel.__spec__ = pkg_spec + skel.__package__ = package + sys.modules[package] = skel + + try: + for prereq in prereqs: + full = f"{package}.{prereq}" + if full in sys.modules: + continue + prereq_path = os.path.join(pkg_path, f"{prereq}.py") + prereq_spec = importlib.util.spec_from_file_location(full, prereq_path) + prereq_mod = importlib.util.module_from_spec(prereq_spec) + sys.modules[full] = prereq_mod + prereq_spec.loader.exec_module(prereq_mod) + + device_type_path = os.path.join(pkg_path, "device_type.py") + dt_spec = importlib.util.spec_from_file_location(target, device_type_path) + dt_mod = importlib.util.module_from_spec(dt_spec) + sys.modules[target] = dt_mod + + import torch + + _orig_is_avail = torch.cuda.is_available + torch.cuda.is_available = lambda: True # type: ignore[assignment] + try: + dt_spec.loader.exec_module(dt_mod) + finally: + torch.cuda.is_available = _orig_is_avail + except Exception: + sys.modules.pop(target, None) + return False + finally: + if not skeleton_already: + sys.modules.pop(package, None) + + return True + + +def _patch_torch_cuda_for_import() -> None: + """Stub torch.cuda.* probes that fire at IMPORT time of unsloth / + unsloth_zoo when DEVICE_TYPE was forced to "cuda" above. These are + queries, not real GPU work, so returning plausible Ampere values + lets the import chain finish; tests that touch real tensors run on + CPU like normal.""" + try: + import torch.cuda.memory as _cuda_memory # type: ignore + + _cuda_memory.mem_get_info = lambda *a, **k: (0, 80 * 1024**3) + except Exception: + pass + try: + import torch + + torch.cuda.get_device_capability = lambda *a, **k: (8, 0) + torch.cuda.is_bf16_supported = lambda *a, **k: True + except Exception: + pass + + +def _install_device_type_stub(name: str) -> None: + stub = types.ModuleType(name) + stub.DEVICE_TYPE = "cuda" + stub.DEVICE_TYPE_TORCH = "cuda" + stub.DEVICE_COUNT = 1 + stub.ALLOW_PREQUANTIZED_MODELS = False + stub.is_hip = lambda: False + stub.get_device_type = lambda: "cuda" + stub.get_device_count = lambda: 1 + stub.device_synchronize = lambda *a, **k: None + stub.device_empty_cache = lambda *a, **k: None + stub.device_is_bf16_supported = lambda *a, **k: False + sys.modules[name] = stub + + +if not _has_real_accelerator(): + if not _preload_device_type("unsloth_zoo", prereqs = ("utils",)): + _install_device_type_stub("unsloth_zoo.device_type") + if not _preload_device_type("unsloth"): + _install_device_type_stub("unsloth.device_type") + _patch_torch_cuda_for_import() diff --git a/tests/python/test_fast_sentence_transformer_redirect_lifecycle.py b/tests/python/test_fast_sentence_transformer_redirect_lifecycle.py index ff9b91ec23..31d86b09a4 100644 --- a/tests/python/test_fast_sentence_transformer_redirect_lifecycle.py +++ b/tests/python/test_fast_sentence_transformer_redirect_lifecycle.py @@ -8,11 +8,45 @@ from __future__ import annotations +import importlib.util import os import pathlib import sys import types +import pytest + + +def _stub_module(name: str) -> types.ModuleType: + # __spec__ must be set so importlib.util.find_spec(name) does not raise + # ValueError if a downstream test imports the real package. + mod = types.ModuleType(name) + mod.__spec__ = importlib.util.spec_from_loader(name, loader = None) + return mod + + +_STUB_KEYS = ( + "transformers", + "sentence_transformers", + "sentence_transformers.models", +) + + +@pytest.fixture(autouse = True) +def _restore_sys_modules(): + """Snapshot the entries we shadow with stubs and restore them after each + test so a downstream test that does `import transformers` for real does + not pick up our non-package stub.""" + saved = {k: sys.modules.get(k) for k in _STUB_KEYS} + try: + yield + finally: + for k, v in saved.items(): + if v is None: + sys.modules.pop(k, None) + else: + sys.modules[k] = v + class _FakeAuto: def __init__(self, name): @@ -45,14 +79,14 @@ class _RaisingTransformer: def _build_driver(transformer_class): - transformers_mod = types.ModuleType("transformers") + transformers_mod = _stub_module("transformers") transformers_mod.AutoModel = _FakeAuto("AutoModel") transformers_mod.AutoProcessor = _FakeAuto("AutoProcessor") transformers_mod.AutoTokenizer = _FakeAuto("AutoTokenizer") sys.modules["transformers"] = transformers_mod - st_root = types.ModuleType("sentence_transformers") - st_models = types.ModuleType("sentence_transformers.models") + st_root = _stub_module("sentence_transformers") + st_models = _stub_module("sentence_transformers.models") st_models.Transformer = transformer_class sys.modules["sentence_transformers"] = st_root sys.modules["sentence_transformers.models"] = st_models diff --git a/tests/studio/install/test_rocm_support.py b/tests/studio/install/test_rocm_support.py index 553756cd4f..13df049c2d 100644 --- a/tests/studio/install/test_rocm_support.py +++ b/tests/studio/install/test_rocm_support.py @@ -1188,7 +1188,7 @@ class TestAmdGpuMonitoring: assert metrics["vram_utilization_pct"] is not None assert metrics["power_utilization_pct"] is not None - def test_amd_primary_gpu_with_mock(self): + def test_amd_primary_gpu_with_mock(self, monkeypatch): """get_primary_gpu_utilization returns correct dict with mocked amd-smi.""" amd_path = PACKAGE_ROOT / "studio" / "backend" / "utils" / "hardware" / "amd.py" _amd_spec = importlib.util.spec_from_file_location("test_amd2", amd_path) @@ -1203,6 +1203,17 @@ class TestAmdGpuMonitoring: except Exception: pytest.skip("Could not load amd module") + # _first_visible_amd_gpu_id() short-circuits to None when any of + # HIP / ROCR / CUDA_VISIBLE_DEVICES is set to "" or "-1". CI runners + # often unset CUDA at the env level by setting CUDA_VISIBLE_DEVICES + # to "" so the test must not inherit that. + for var in ( + "HIP_VISIBLE_DEVICES", + "ROCR_VISIBLE_DEVICES", + "CUDA_VISIBLE_DEVICES", + ): + monkeypatch.delenv(var, raising = False) + mock_json = json.dumps( [ { diff --git a/tests/studio/test_hardware_dispatch_matrix.py b/tests/studio/test_hardware_dispatch_matrix.py index 276c57ca5a..c7a6841936 100644 --- a/tests/studio/test_hardware_dispatch_matrix.py +++ b/tests/studio/test_hardware_dispatch_matrix.py @@ -208,6 +208,20 @@ def spoof_hardware(monkeypatch): # torch.cuda.is_available monkeypatch.setattr(torch.cuda, "is_available", lambda: profile.cuda_available) + # detect_hardware reads torch.cuda.get_device_properties(0).name when + # cuda_available is True. On a CPU CI runner that triggers _cuda_init + # and crashes with "No CUDA GPUs are available". Stub it so the + # dispatch path under test runs end-to-end. + if profile.cuda_available: + stub_props = types.SimpleNamespace( + name = "Stub GPU" if not profile.hip_version else "Stub AMD GPU", + ) + monkeypatch.setattr( + torch.cuda, + "get_device_properties", + lambda i = 0: stub_props, + raising = False, + ) # torch.version.hip — None on NVIDIA, "6.1" etc. on ROCm torch_version = torch.version diff --git a/tests/test_raw_text.py b/tests/test_raw_text.py index d8289fed20..056ea8660f 100644 --- a/tests/test_raw_text.py +++ b/tests/test_raw_text.py @@ -35,8 +35,11 @@ class MockDataset: return cls(data_dict) -# Mock datasets module +# Mock datasets module. __spec__ must be set so importlib.util.find_spec +# does not raise ValueError when transformers' import_utils probes for +# the real `datasets` package later in the test session. datasets_mock = type(sys)("datasets") +datasets_mock.__spec__ = importlib.util.spec_from_loader("datasets", loader = None) datasets_mock.Dataset = MockDataset sys.modules["datasets"] = datasets_mock diff --git a/unsloth/kernels/utils.py b/unsloth/kernels/utils.py index 09b03a597b..dd5a9cbf0e 100644 --- a/unsloth/kernels/utils.py +++ b/unsloth/kernels/utils.py @@ -161,8 +161,14 @@ else: if DEVICE_TYPE == "xpu": _gpu_getCurrentRawStream = torch._C._xpu_getCurrentRawStream # NVIDIA GPU Default Logic -else: +elif hasattr(torch._C, "_cuda_getCurrentRawStream"): _gpu_getCurrentRawStream = torch._C._cuda_getCurrentRawStream +else: + # CPU-only torch wheel (no compiled CUDA backend). _get_tensor_stream + # is only invoked during real GPU work, so a no-op binding is safe. + def _gpu_getCurrentRawStream(_index = 0): + return 0 + c_void_p = ctypes.c_void_p @@ -177,36 +183,49 @@ global XPU_STREAMS global WEIGHT_BUFFERS global ABSMAX_BUFFERS -# INTEL GPU Specific Logic +# DEVICE_COUNT == 0 = no visible accelerator (e.g. CPU-only CI runner). +# The consumer functions below only index these arrays during real GPU +# work, so empty containers are safe -- they just need to be defined so +# the module imports cleanly. if DEVICE_TYPE == "xpu": - _XPU_STREAMS = { - (index := torch.xpu.device(i).idx): ctypes.c_void_p( - torch._C._xpu_getCurrentRawStream(index) - ) - for i in range(DEVICE_COUNT) - } - XPU_STREAMS = [None] * (max(_XPU_STREAMS.keys()) + 1) - WEIGHT_BUFFERS = [None] * (max(_XPU_STREAMS.keys()) + 1) - ABSMAX_BUFFERS = [None] * (max(_XPU_STREAMS.keys()) + 1) - for k, v in _XPU_STREAMS.items(): - XPU_STREAMS[k] = v - XPU_STREAMS = tuple(XPU_STREAMS) - del _XPU_STREAMS + if DEVICE_COUNT > 0: + _XPU_STREAMS = { + (index := torch.xpu.device(i).idx): ctypes.c_void_p( + torch._C._xpu_getCurrentRawStream(index) + ) + for i in range(DEVICE_COUNT) + } + XPU_STREAMS = [None] * (max(_XPU_STREAMS.keys()) + 1) + WEIGHT_BUFFERS = [None] * (max(_XPU_STREAMS.keys()) + 1) + ABSMAX_BUFFERS = [None] * (max(_XPU_STREAMS.keys()) + 1) + for k, v in _XPU_STREAMS.items(): + XPU_STREAMS[k] = v + XPU_STREAMS = tuple(XPU_STREAMS) + del _XPU_STREAMS + else: + XPU_STREAMS = () + WEIGHT_BUFFERS = [] + ABSMAX_BUFFERS = [] else: # NVIDIA GPU Default Logic - _CUDA_STREAMS = { - (index := torch.cuda.device(i).idx): ctypes.c_void_p( - torch._C._cuda_getCurrentRawStream(index) - ) - for i in range(DEVICE_COUNT) - } - CUDA_STREAMS = [None] * (max(_CUDA_STREAMS.keys()) + 1) - WEIGHT_BUFFERS = [None] * (max(_CUDA_STREAMS.keys()) + 1) - ABSMAX_BUFFERS = [None] * (max(_CUDA_STREAMS.keys()) + 1) - for k, v in _CUDA_STREAMS.items(): - CUDA_STREAMS[k] = v - CUDA_STREAMS = tuple(CUDA_STREAMS) - del _CUDA_STREAMS + if DEVICE_COUNT > 0: + _CUDA_STREAMS = { + (index := torch.cuda.device(i).idx): ctypes.c_void_p( + torch._C._cuda_getCurrentRawStream(index) + ) + for i in range(DEVICE_COUNT) + } + CUDA_STREAMS = [None] * (max(_CUDA_STREAMS.keys()) + 1) + WEIGHT_BUFFERS = [None] * (max(_CUDA_STREAMS.keys()) + 1) + ABSMAX_BUFFERS = [None] * (max(_CUDA_STREAMS.keys()) + 1) + for k, v in _CUDA_STREAMS.items(): + CUDA_STREAMS[k] = v + CUDA_STREAMS = tuple(CUDA_STREAMS) + del _CUDA_STREAMS + else: + CUDA_STREAMS = () + WEIGHT_BUFFERS = [] + ABSMAX_BUFFERS = [] # Bitsandbytes operations ctypes_c_int = ctypes.c_int From fac2dc09b0fdf0c38a81f5bad889d58f6706d672 Mon Sep 17 00:00:00 2001 From: Lee Jackson <130007945+Imagineer99@users.noreply.github.com> Date: Wed, 6 May 2026 12:55:37 +0100 Subject: [PATCH 011/159] fix: restore API and Help menu labels (#5310) --- studio/frontend/src/components/app-sidebar.tsx | 15 +++++++++++---- .../src/features/settings/settings-dialog.tsx | 8 ++++---- .../src/features/settings/tabs/api-keys-tab.tsx | 2 +- 3 files changed, 16 insertions(+), 9 deletions(-) diff --git a/studio/frontend/src/components/app-sidebar.tsx b/studio/frontend/src/components/app-sidebar.tsx index 171b4eb92d..13b8adfa48 100644 --- a/studio/frontend/src/components/app-sidebar.tsx +++ b/studio/frontend/src/components/app-sidebar.tsx @@ -37,12 +37,13 @@ import { Delete02Icon, Download03Icon, GemIcon, + Globe02Icon, + HelpCircleIcon, Search01Icon, PowerIcon, PencilEdit02Icon, LayoutAlignLeftIcon, Settings02Icon, - SourceCodeSquareIcon, ZapIcon, } from "@hugeicons/core-free-icons"; import { @@ -527,7 +528,7 @@ export function AppSidebar() {
{displayTitle} - Studio + Unsloth
@@ -548,8 +549,8 @@ export function AppSidebar() { useSettingsDialogStore.getState().openDialog("api-keys")} > - - Developer + + API New @@ -578,6 +579,12 @@ export function AppSidebar() { + useSettingsDialogStore.getState().openDialog("about")} + > + + Help + setShutdownOpen(true)}> Shutdown diff --git a/studio/frontend/src/features/settings/settings-dialog.tsx b/studio/frontend/src/features/settings/settings-dialog.tsx index 376af06e9d..63c0a9d388 100644 --- a/studio/frontend/src/features/settings/settings-dialog.tsx +++ b/studio/frontend/src/features/settings/settings-dialog.tsx @@ -10,11 +10,11 @@ import { import { cn } from "@/lib/utils"; import { Cancel01Icon, + Globe02Icon, + HelpCircleIcon, Message01Icon, PaintBrush02Icon, Settings02Icon, - SourceCodeSquareIcon, - SparklesIcon, UserIcon, } from "@hugeicons/core-free-icons"; import { HugeiconsIcon } from "@hugeicons/react"; @@ -40,8 +40,8 @@ const TABS: TabDef[] = [ { id: "profile", label: "Profile", icon: UserIcon }, { id: "appearance", label: "Appearance", icon: PaintBrush02Icon }, { id: "chat", label: "Chat", icon: Message01Icon }, - { id: "api-keys", label: "Developer", icon: SourceCodeSquareIcon, badge: "New" }, - { id: "about", label: "Help", icon: SparklesIcon }, + { id: "api-keys", label: "API", icon: Globe02Icon, badge: "New" }, + { id: "about", label: "Help", icon: HelpCircleIcon }, ]; function renderTab(tab: SettingsTab) { diff --git a/studio/frontend/src/features/settings/tabs/api-keys-tab.tsx b/studio/frontend/src/features/settings/tabs/api-keys-tab.tsx index 64c6f520c0..ac9ec40543 100644 --- a/studio/frontend/src/features/settings/tabs/api-keys-tab.tsx +++ b/studio/frontend/src/features/settings/tabs/api-keys-tab.tsx @@ -63,7 +63,7 @@ export function ApiKeysTab() { return (
-

Developer

+

API

Access Unsloth programmatically via the OpenAI-compatible API.{" "} Date: Wed, 6 May 2026 22:16:20 +0530 Subject: [PATCH 012/159] [studio]: Fix tool reasoning trace in UI (#5314) * fix thought for 1 second issue * gemini suggesion --- .../src/components/assistant-ui/reasoning.tsx | 23 +++++++++++++++---- 1 file changed, 18 insertions(+), 5 deletions(-) diff --git a/studio/frontend/src/components/assistant-ui/reasoning.tsx b/studio/frontend/src/components/assistant-ui/reasoning.tsx index 387f8cd458..fe913baf2a 100644 --- a/studio/frontend/src/components/assistant-ui/reasoning.tsx +++ b/studio/frontend/src/components/assistant-ui/reasoning.tsx @@ -316,15 +316,28 @@ const ReasoningGroupImpl: ReasoningGroupComponent = ({ if (message.status?.type !== "running") { return false; } - const lastIndex = message.parts.length - 1; - if (lastIndex < 0) { + const parts = message.parts; + const len = parts.length; + if (len === 0) { return false; } - const lastType = message.parts[lastIndex]?.type; - if (lastType !== "reasoning") { + + let groupHasReasoning = false; + for (let i = startIndex; i <= endIndex && i < len; i += 1) { + if (parts[i]?.type === "reasoning") { + groupHasReasoning = true; + break; + } + } + if (!groupHasReasoning) { return false; } - return lastIndex >= startIndex && lastIndex <= endIndex; + for (let i = endIndex + 1; i < len; i += 1) { + if (parts[i]?.type !== "tool-call") { + return false; + } + } + return true; }); const persistedDuration = useAuiState(({ message }) => { From 948ce43584c272be016f6c7210d99a6a0bddc26a Mon Sep 17 00:00:00 2001 From: Daniel Han Date: Thu, 7 May 2026 00:12:09 -0700 Subject: [PATCH 013/159] =?UTF-8?q?fix:=203=20patch=5F*=20helpers=20?= =?UTF-8?q?=E2=80=94=20fast=5Flora=20import,=20sft=5Ftrainer=20Union,=20op?= =?UTF-8?q?enenv=20OSError=20(#5319)?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit * fix: import fast_lora_forward inside patch_fast_lora patch_fast_lora has referenced an unbound `fast_lora_forward` since ddf118a8f (2024-11-21). The function is defined at unsloth/kernels/fast_lora.py:652 and re-exported through unsloth/kernels/__init__.py:45, but it was never imported into unsloth/models/_utils.py, so calling patch_fast_lora() raises NameError: name 'fast_lora_forward' is not defined. The bug went unnoticed because no production code path calls patch_fast_lora() unconditionally. Surfaced by a new CPU-CI check that invokes every zero-arg patch_* helper across unsloth + unsloth_zoo (consolidated-tests-ci.yml on PR #5312). Importing inside the function (rather than at module top) keeps the import surface narrow and avoids a circular-import risk if unsloth.kernels.fast_lora ever needs to import from unsloth.models._utils. * fix: inject typing imports into patch_sft_trainer_tokenizer's exec namespace patch_sft_trainer_tokenizer rewrites the source of TRL's SFTTrainer methods (_prepare_non_packed_dataloader, _prepare_dataset) and re-execs them. With TRL 1.x, those methods carry `Union[...]` type hints in their signatures. The current rewrite only injects identifiers found by `dir(trl.trainer.sft_trainer)` into the exec namespace, which does not include `Union`, so exec(function, ...) raises NameError: name 'Union' is not defined. Fix: import Union, Optional, List, Any, Callable, Tuple, Dict, Iterator inside the function. exec receives `locals()` as its globals dict, so those names are visible to the executed source body. Same pattern as unsloth/models/_utils.py:patch_linear_scaling, which already injects `from typing import Union, Optional, List, Any, Callable, Tuple` into its own exec_code. Surfaced by the consolidated CPU-CI runtime patch_* check on PR #5312 in the matrix cell `transformers>=5,<6 + trl>=1,<2`. * fix: guard openenv_vllm_reload_weights against OSError from inspect.getsource TRL 0.29.1 and the 1.x line ship some openenv helpers as compiled bytecode without accessible source on disk. inspect.getsource(patch_target) raises OSError("could not get source code") in that case, which surfaces as a hard failure in patch_trl_openenv() and aborts the rest of the RL_ADDITIONAL_FUNCTIONS["openenv"] iteration. Wrap the getsource call in a try/except OSError and log a warning instead. The wake_up(tags=...) rewrite is the only thing skipped; the core weight-reload patch path stays functional. Surfaced by the consolidated CPU-CI runtime patch_* check on PR #5312 in matrix cells running TRL 0.29.1 (latest <1.0.0) and TRL 1.3.0 (latest 1.x). The pyproject pin (TRL 0.18.2-0.24.0) still gets source for this function so the original code path runs unchanged there. --- unsloth/models/_utils.py | 1 + unsloth/models/rl_replacements.py | 15 ++++++++++++++- unsloth/tokenizer_utils.py | 6 ++++++ 3 files changed, 21 insertions(+), 1 deletion(-) diff --git a/unsloth/models/_utils.py b/unsloth/models/_utils.py index 34fec53504..d3eee03325 100644 --- a/unsloth/models/_utils.py +++ b/unsloth/models/_utils.py @@ -2483,6 +2483,7 @@ def patch_tokenizer(model, tokenizer): def patch_fast_lora(): import peft.tuners.lora.bnb + from ..kernels.fast_lora import fast_lora_forward peft.tuners.lora.bnb.Linear4bit.forward = fast_lora_forward diff --git a/unsloth/models/rl_replacements.py b/unsloth/models/rl_replacements.py index 0f10847282..c2be1bf74a 100755 --- a/unsloth/models/rl_replacements.py +++ b/unsloth/models/rl_replacements.py @@ -1780,7 +1780,20 @@ def openenv_vllm_reload_weights(): patch_target_name = "generate_rollout_completions" patch_target = getattr(openenv_utils, patch_target_name) - src = inspect.getsource(patch_target) + # TRL 0.29.1+ ships some openenv helpers as compiled bytecode without + # accessible source on disk; inspect.getsource raises OSError("could + # not get source code") in that case. Skip the source-rewrite patch + # rather than crashing -- the core unsloth weight-reload path stays + # functional, only the wake_up tag rewrite is skipped. + try: + src = inspect.getsource(patch_target) + except OSError as e: + logger.warning( + f"Unsloth: Could not retrieve source for trl openenv " + f"{patch_target_name} ({e}); skipping rewrite. " + f"Weight reload still functional." + ) + return src = textwrap.dedent(src) original_src = src diff --git a/unsloth/tokenizer_utils.py b/unsloth/tokenizer_utils.py index 130894e385..67edc41d52 100644 --- a/unsloth/tokenizer_utils.py +++ b/unsloth/tokenizer_utils.py @@ -1580,6 +1580,12 @@ def patch_sft_trainer_tokenizer(): except: return all_imports = dir(trl.trainer.sft_trainer) + # Make typing names available to the exec'd source bodies. TRL >= 1.x + # type-hints _prepare_dataset / _prepare_non_packed_dataloader with + # `Union[...]` and friends; without these imports in the exec namespace + # those become NameErrors at exec time. Mirrors the pattern used in + # unsloth/models/_utils.py:patch_linear_scaling. + from typing import Union, Optional, List, Any, Callable, Tuple, Dict, Iterator # noqa: F401 for ( function_name, From 4ab096970d81520d42a5452017fec77e1a8154f8 Mon Sep 17 00:00:00 2001 From: Lee Jackson <130007945+Imagineer99@users.noreply.github.com> Date: Thu, 7 May 2026 10:17:23 +0100 Subject: [PATCH 014/159] Studio: API settings overflow with long Colab URLs (#5286) * fix: API settings overflow with long Colab URLs * fix: gentle wrapping for API usage snippets --------- Co-authored-by: Roland Tannous <115670425+rolandtannous@users.noreply.github.com> --- .../features/settings/components/usage-examples.tsx | 10 +++++----- .../frontend/src/features/settings/settings-dialog.tsx | 2 +- .../src/features/settings/tabs/api-keys-tab.tsx | 8 ++++---- 3 files changed, 10 insertions(+), 10 deletions(-) diff --git a/studio/frontend/src/features/settings/components/usage-examples.tsx b/studio/frontend/src/features/settings/components/usage-examples.tsx index a8f8ce31a6..0e68237857 100644 --- a/studio/frontend/src/features/settings/components/usage-examples.tsx +++ b/studio/frontend/src/features/settings/components/usage-examples.tsx @@ -96,11 +96,11 @@ export function UsageExamples() { }; return ( -

+

Usage examples

-
-
-
+
+
+
{TABS.map((t) => { const active = lang === t.id; return ( @@ -134,7 +134,7 @@ export function UsageExamples() { {copied ? "Copied" : "Copy"}
- 
+        
           {snippets[lang]}
         

         

diff --git a/studio/frontend/src/features/settings/settings-dialog.tsx b/studio/frontend/src/features/settings/settings-dialog.tsx
index 63c0a9d388..38002e2f14 100644
--- a/studio/frontend/src/features/settings/settings-dialog.tsx
+++ b/studio/frontend/src/features/settings/settings-dialog.tsx
@@ -163,7 +163,7 @@ export function SettingsDialog() {
             >
               
             
-            

+            

               {renderTab(activeTab)}
             

           
diff --git a/studio/frontend/src/features/settings/tabs/api-keys-tab.tsx b/studio/frontend/src/features/settings/tabs/api-keys-tab.tsx
index ac9ec40543..6cbd28f14d 100644
--- a/studio/frontend/src/features/settings/tabs/api-keys-tab.tsx
+++ b/studio/frontend/src/features/settings/tabs/api-keys-tab.tsx
@@ -61,8 +61,8 @@ export function ApiKeysTab() {
   };
 
   return (
-    

-      

+    

+      

         
API

         

           Access Unsloth programmatically via the OpenAI-compatible API.{" "}
@@ -111,7 +111,7 @@ export function ApiKeysTab() {
         )}
       
 
-      

+      

         
Access tokens

         {error ? (
           

@@ -131,7 +131,7 @@ export function ApiKeysTab() {
             No API access yet.
           

         ) : (
-          

+          

             {keys.map((k) => (
               
             ))}

From 7af8cac0148d209a0ca7b607756d6c1eea917af7 Mon Sep 17 00:00:00 2001
From: Daniel Han 
Date: Thu, 7 May 2026 02:26:58 -0700
Subject: [PATCH 015/159] tests/studio/install: parallel UNSLOTH_STUDIO_HOME
 smoke test (#5306)

* tests/studio/install: parallel UNSLOTH_STUDIO_HOME smoke test

Adds tests/studio/install/smoke_test_parallel_studio_home.py to lock in
the install-time and runtime isolation guarantees added by #5190.

The runner spawns N concurrent install.sh --local --no-torch jobs, each
with its own UNSLOTH_STUDIO_HOME and a redirected HOME, then launches N
backends on dynamically allocated ports and cross-checks every install
against its running process. Asserts:

  install-time
    - all N installs exit 0
    - per-install bin / share / llama.cpp / unsloth_studio venv tree
    - shim symlink resolves into its own venv, no cross-resolution
    - share/studio_install_id is unique across the N installs
    - share/studio.conf exports UNSLOTH_EXE / UNSLOTH_STUDIO_HOME /
      UNSLOTH_LLAMA_CPP_PATH all pointing inside the install
    - share/launch-studio.sh has @@DATA_DIR@@ substituted to its own
      share/ at install time
    - the redirected HOME stays clean: no rc-file append, no
      .desktop file, no Studio.app stub, no shared marker

  runtime
    - /api/health returns 200 with status healthy and chat_only true
    - /api/health.studio_root_id matches share/studio_install_id
      (runtime resolver agrees with install-time write)
    - studio_root_id values are pairwise distinct
    - GET / and GET /api/chat return 200 on each backend
    - /proc/PID/exe is the install's own venv python

Standalone smoke runner, not pytest collected. Default --n 4 finishes
in about 60 seconds on a warm uv cache; artifacts are removed on PASS
unless --keep is passed and kept on FAIL or ERROR for inspection.

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

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

* tests/studio/install: portability + log fd cleanup in parallel smoke

Two cleanups on the parallel UNSLOTH_STUDIO_HOME smoke runner:

- Skip the /proc/PID/exe runtime cross-resolution check on platforms
  without /proc (macOS, BSD, Windows). install.sh supports macOS, so
  the smoke should not hard-error there. The install-time symlink,
  studio.conf and launch-studio.sh assertions already pin the venv
  python target statically; the proc check stays as a Linux-only
  redundant cross-resolution catch and now returns None cleanly on
  other platforms instead of raising.

- Wrap the per-backend log file in a with-statement so its parent fd
  is released deterministically at function return. The child still
  holds its own dup'd fd via Popen, so logging continues unchanged.
  The prior code relied on local-scope GC and was fine in CPython,
  but the with form makes the intent explicit.

Smoke still passes locally: 4 parallel installs in 42s, 4 backends
healthy in 5s, all install + runtime invariants hold.

* tests/studio/install: pin UNSLOTH_STUDIO_HOME on backend launch

The launch step copied os.environ unchanged except for HOME. If the
parent shell already exports UNSLOTH_STUDIO_HOME or STUDIO_HOME (for
example, when the developer is sourcing studio.conf from an existing
install), every backend inherits it and the Studio resolver prioritises
those env vars over the per-label sys.prefix inference. The runtime
invariant block then reports the caller's install_id on every port
instead of the per-label one, and the test fails spuriously rather
than testing the right roots.

Pin UNSLOTH_STUDIO_HOME to the per-label studio_home and pop the
STUDIO_HOME alias for each launch, mirroring what _run_one_install
already does for the install step.

Verified by running the smoke with UNSLOTH_STUDIO_HOME=/nonexistent
and STUDIO_HOME=/also-bogus exported in the parent env: PASS, all four
backends report their own install_id rather than the parent value.

---------

Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com>
---
 .../smoke_test_parallel_studio_home.py        | 421 ++++++++++++++++++
 1 file changed, 421 insertions(+)
 create mode 100644 tests/studio/install/smoke_test_parallel_studio_home.py

diff --git a/tests/studio/install/smoke_test_parallel_studio_home.py b/tests/studio/install/smoke_test_parallel_studio_home.py
new file mode 100644
index 0000000000..133591fb33
--- /dev/null
+++ b/tests/studio/install/smoke_test_parallel_studio_home.py
@@ -0,0 +1,421 @@
+#!/usr/bin/env python3
+"""Smoke test: N parallel install.sh runs with distinct UNSLOTH_STUDIO_HOME
+values must produce N fully isolated installs whose backends can run
+side by side without clashing.
+
+Covers the env-override path added in #5190:
+
+    install-time
+        * N concurrent ``install.sh --local --no-torch`` runs against
+          this checkout, each pinned to its own UNSLOTH_STUDIO_HOME and
+          a redirected HOME, all exit 0.
+        * Each STUDIO_HOME contains its own bin/, share/, llama.cpp/
+          and unsloth_studio/ venv, with no cross-install absolute
+          paths.
+        * share/studio_install_id is unique across the N installs.
+        * share/studio.conf exports UNSLOTH_EXE, UNSLOTH_STUDIO_HOME
+          and UNSLOTH_LLAMA_CPP_PATH, all pointing inside this install.
+        * share/launch-studio.sh has @@DATA_DIR@@ substituted to its
+          own share/ at install time.
+        * bin/unsloth is a symlink that resolves into its own venv.
+        * The redirected HOME is left clean: no shell-rc append, no
+          .desktop file, no Studio.app stub, no shared marker.
+
+    runtime
+        * N concurrent ``bin/unsloth studio`` launches each bind their
+          own dynamically allocated free port and stay healthy.
+        * /api/health is 200, status is healthy, chat_only is true
+          under --no-torch.
+        * The studio_root_id reported by /api/health on each backend
+          equals that install's share/studio_install_id, so the
+          runtime resolver agrees with the install-time write.
+        * studio_root_id values are pairwise distinct.
+        * GET / and GET /api/chat are 200 on every backend.
+        * The Python interpreter behind each PID is the install's own
+          venv python (the bin/unsloth shim does not cross-resolve).
+
+This is an integration smoke runner, not a pytest unit test. It does
+real installs (~1 minute end to end on a warm uv cache) and is meant
+to be invoked explicitly:
+
+    python tests/studio/install/smoke_test_parallel_studio_home.py
+    python tests/studio/install/smoke_test_parallel_studio_home.py --n 6 --keep
+
+Exits 0 on PASS, 1 on FAIL, 2 on infrastructure error. Artifacts land
+under a temporary directory and are removed on PASS unless --keep is
+set; on FAIL or ERROR they are kept regardless so logs can be
+inspected.
+"""
+
+from __future__ import annotations
+
+import argparse
+import json
+import os
+import shutil
+import socket
+import subprocess
+import sys
+import tempfile
+import time
+import urllib.error
+import urllib.request
+from concurrent.futures import ThreadPoolExecutor, as_completed
+from datetime import datetime
+from pathlib import Path
+
+PACKAGE_ROOT = Path(__file__).resolve().parents[3]
+INSTALL_TIMEOUT_S = 600
+HEALTH_TIMEOUT_S = 120
+HEALTH_POLL_INTERVAL_S = 1.0
+
+
+class TestFailure(AssertionError):
+    pass
+
+
+def _log(msg: str) -> None:
+    ts = datetime.now().strftime("%H:%M:%S")
+    print(f"[smoke {ts}] {msg}", flush = True)
+
+
+def _free_port() -> int:
+    with socket.socket(socket.AF_INET, socket.SOCK_STREAM) as s:
+        s.bind(("127.0.0.1", 0))
+        return s.getsockname()[1]
+
+
+def _run_one_install(
+    label: str,
+    repo: Path,
+    studio_home: Path,
+    fake_home: Path,
+    uv_cache: Path,
+    log_path: Path,
+) -> tuple[str, int]:
+    studio_home.mkdir(parents = True, exist_ok = True)
+    fake_home.mkdir(parents = True, exist_ok = True)
+    uv_cache.mkdir(parents = True, exist_ok = True)
+    log_path.parent.mkdir(parents = True, exist_ok = True)
+    env = os.environ.copy()
+    env["HOME"] = str(fake_home)
+    env["UNSLOTH_STUDIO_HOME"] = str(studio_home)
+    env["UV_CACHE_DIR"] = str(uv_cache)
+    env["NO_COLOR"] = "1"
+    with log_path.open("w") as fh:
+        proc = subprocess.run(
+            ["bash", "install.sh", "--local", "--no-torch"],
+            cwd = str(repo),
+            env = env,
+            stdout = fh,
+            stderr = subprocess.STDOUT,
+            timeout = INSTALL_TIMEOUT_S,
+        )
+    return label, proc.returncode
+
+
+def _launch_backend(
+    studio_home: Path, fake_home: Path, port: int, log_path: Path
+) -> subprocess.Popen:
+    log_path.parent.mkdir(parents = True, exist_ok = True)
+    env = os.environ.copy()
+    env["HOME"] = str(fake_home)
+    # Pin UNSLOTH_STUDIO_HOME (and clear the alias) so the child cannot
+    # inherit a Studio root from the caller's shell. Without this, a shell
+    # that already exports either var would override the per-label sys.prefix
+    # inference and every backend would resolve to the caller's install.
+    env["UNSLOTH_STUDIO_HOME"] = str(studio_home)
+    env.pop("STUDIO_HOME", None)
+    # The child process inherits a dup of stdout via Popen, so closing the
+    # parent's handle when this function returns is safe and avoids relying
+    # on GC timing to release the fd.
+    with log_path.open("w") as fh:
+        return subprocess.Popen(
+            [
+                str(studio_home / "bin" / "unsloth"),
+                "studio",
+                "-H",
+                "127.0.0.1",
+                "-p",
+                str(port),
+                "--silent",
+            ],
+            env = env,
+            stdout = fh,
+            stderr = subprocess.STDOUT,
+            start_new_session = True,
+        )
+
+
+def _wait_for_health(port: int, timeout: float) -> dict:
+    deadline = time.time() + timeout
+    last_err: Exception | None = None
+    url = f"http://127.0.0.1:{port}/api/health"
+    while time.time() < deadline:
+        try:
+            with urllib.request.urlopen(url, timeout = 2) as r:
+                if r.status == 200:
+                    return json.loads(r.read().decode())
+        except (urllib.error.URLError, ConnectionError, OSError) as e:
+            last_err = e
+        time.sleep(HEALTH_POLL_INTERVAL_S)
+    raise TestFailure(
+        f"port {port}: /api/health never returned 200 (last_err={last_err})"
+    )
+
+
+def _http_status(port: int, path: str, timeout: float = 5.0) -> int:
+    url = f"http://127.0.0.1:{port}{path}"
+    try:
+        with urllib.request.urlopen(url, timeout = timeout) as r:
+            return r.status
+    except urllib.error.HTTPError as e:
+        return e.code
+
+
+def _check_install_layout(label: str, studio_home: Path) -> dict:
+    for d in ("bin", "share", "llama.cpp", "unsloth_studio"):
+        if not (studio_home / d).is_dir():
+            raise TestFailure(f"[{label}] missing {studio_home / d}")
+
+    shim = studio_home / "bin" / "unsloth"
+    if not shim.is_symlink():
+        raise TestFailure(f"[{label}] {shim} is not a symlink")
+    expected_target = (studio_home / "unsloth_studio" / "bin" / "unsloth").resolve()
+    if shim.resolve() != expected_target:
+        raise TestFailure(
+            f"[{label}] shim resolves to {shim.resolve()}, expected {expected_target}"
+        )
+
+    install_id_path = studio_home / "share" / "studio_install_id"
+    if not install_id_path.is_file():
+        raise TestFailure(f"[{label}] missing {install_id_path}")
+    install_id = install_id_path.read_text().strip()
+    if len(install_id) < 32:
+        raise TestFailure(f"[{label}] studio_install_id too short: {install_id!r}")
+
+    conf = (studio_home / "share" / "studio.conf").read_text()
+    must_contain = [
+        f"UNSLOTH_EXE='{studio_home}/unsloth_studio/bin/unsloth'",
+        f"export UNSLOTH_STUDIO_HOME='{studio_home}'",
+        f"export UNSLOTH_LLAMA_CPP_PATH='{studio_home}/llama.cpp'",
+    ]
+    for needle in must_contain:
+        if needle not in conf:
+            raise TestFailure(
+                f"[{label}] studio.conf missing line:\n  {needle}\n" f"actual:\n{conf}"
+            )
+
+    launcher = (studio_home / "share" / "launch-studio.sh").read_text()
+    if "@@DATA_DIR@@" in launcher:
+        raise TestFailure(f"[{label}] launch-studio.sh kept @@DATA_DIR@@ placeholder")
+    expected_data_dir_line = f"DATA_DIR='{studio_home}/share'"
+    if expected_data_dir_line not in launcher:
+        raise TestFailure(
+            f"[{label}] launch-studio.sh missing {expected_data_dir_line!r}"
+        )
+
+    return {"label": label, "studio_home": str(studio_home), "install_id": install_id}
+
+
+def _check_fake_home_clean(fake_home: Path) -> None:
+    forbidden = [
+        ".bashrc",
+        ".zshrc",
+        ".profile",
+        ".unsloth",
+        Path(".local") / "share" / "applications" / "unsloth-studio.desktop",
+        Path("Desktop") / "unsloth-studio.desktop",
+        Path("Applications") / "Unsloth Studio.app",
+    ]
+    leaked = [str(p) for p in forbidden if (fake_home / p).exists()]
+    if leaked:
+        raise TestFailure(
+            f"redirected HOME picked up persistent install pollution: {leaked}"
+        )
+
+
+def _backend_pid_python(pid: int) -> Path | None:
+    """Resolve the binary backing a running PID. Linux exposes this at
+    /proc/PID/exe; on platforms without /proc (macOS, BSD, Windows) we
+    skip this check and rely on the install-time symlink + studio.conf
+    invariants to catch cross-resolution. Returns None when /proc is
+    unavailable so the caller can skip cleanly."""
+    if sys.platform != "linux":
+        return None
+    proc_exe = Path(f"/proc/{pid}/exe")
+    if not proc_exe.exists():
+        return None
+    return proc_exe.resolve()
+
+
+def run(n_installs: int, keep: bool) -> int:
+    if n_installs < 2:
+        raise TestFailure("--n must be >= 2 to test for clashes")
+    labels = [chr(ord("a") + i) for i in range(n_installs)]
+
+    repo = PACKAGE_ROOT
+    if not (repo / "install.sh").is_file():
+        raise TestFailure(
+            f"install.sh not found at {repo}; " "run from a clone of unslothai/unsloth"
+        )
+
+    test_root = Path(tempfile.mkdtemp(prefix = "unsloth_studio_clash_"))
+    _log(f"test root: {test_root}")
+    _log(f"repo: {repo}")
+
+    backends: list[tuple[str, Path, Path, int, subprocess.Popen]] = []
+    failed = False
+    try:
+        # ---- parallel installs --------------------------------------------
+        _log(f"launching {n_installs} parallel installs (--local --no-torch)")
+        with ThreadPoolExecutor(max_workers = n_installs) as pool:
+            futures = []
+            for label in labels:
+                futures.append(
+                    pool.submit(
+                        _run_one_install,
+                        label,
+                        repo,
+                        test_root / "installs" / label,
+                        test_root / "fake_homes" / label,
+                        test_root / "uv_caches" / label,
+                        test_root / "logs" / f"install_{label}.log",
+                    )
+                )
+            for fut in as_completed(futures):
+                label, rc = fut.result()
+                _log(f"  install {label}: exit {rc}")
+                if rc != 0:
+                    raise TestFailure(
+                        f"install {label} failed (rc={rc}); see "
+                        f"{test_root / 'logs' / f'install_{label}.log'}"
+                    )
+
+        # ---- install-layout invariants ------------------------------------
+        _log("verifying install-time invariants")
+        observed = []
+        for label in labels:
+            studio_home = test_root / "installs" / label
+            obs = _check_install_layout(label, studio_home)
+            observed.append(obs)
+            _check_fake_home_clean(test_root / "fake_homes" / label)
+        ids = [o["install_id"] for o in observed]
+        if len(set(ids)) != len(ids):
+            raise TestFailure(f"studio_install_id collision: {ids}")
+        _log(f"  {len(ids)} unique studio_install_ids, all redirected HOMEs clean")
+
+        # ---- parallel backend launches ------------------------------------
+        _log(f"launching {n_installs} backends in parallel")
+        for label in labels:
+            port = _free_port()
+            studio_home = test_root / "installs" / label
+            fake_home = test_root / "fake_homes" / label
+            log_path = test_root / "logs" / f"run_{label}.log"
+            proc = _launch_backend(studio_home, fake_home, port, log_path)
+            backends.append((label, studio_home, fake_home, port, proc))
+            _log(f"  {label} -> port {port} (pid {proc.pid})")
+
+        # ---- wait for health ----------------------------------------------
+        _log("waiting for /api/health on each backend")
+        health_payloads: dict[str, dict] = {}
+        with ThreadPoolExecutor(max_workers = n_installs) as pool:
+            fut_to_label = {
+                pool.submit(_wait_for_health, port, HEALTH_TIMEOUT_S): label
+                for (label, _sh, _fh, port, _p) in backends
+            }
+            for fut in as_completed(fut_to_label):
+                label = fut_to_label[fut]
+                health_payloads[label] = fut.result()
+                _log(f"  {label}: healthy")
+
+        # ---- runtime invariants -------------------------------------------
+        _log("checking runtime invariants")
+        seen_root_ids: set[str] = set()
+        for (label, studio_home, _fh, port, proc), obs in zip(backends, observed):
+            health = health_payloads[label]
+            if health.get("status") != "healthy":
+                raise TestFailure(f"[{label}] health status != healthy: {health}")
+            if health.get("studio_root_id") != obs["install_id"]:
+                raise TestFailure(
+                    f"[{label}] runtime studio_root_id "
+                    f"{health.get('studio_root_id')!r} != install_id "
+                    f"{obs['install_id']!r}"
+                )
+            if not health.get("chat_only"):
+                raise TestFailure(f"[{label}] chat_only is not true under --no-torch")
+            if health["studio_root_id"] in seen_root_ids:
+                raise TestFailure(
+                    f"[{label}] studio_root_id collision at runtime: "
+                    f"{health['studio_root_id']}"
+                )
+            seen_root_ids.add(health["studio_root_id"])
+
+            for path in ("/", "/api/chat"):
+                code = _http_status(port, path)
+                if code != 200:
+                    raise TestFailure(f"[{label}] GET {path} -> {code}")
+
+            exe = _backend_pid_python(proc.pid)
+            if exe is not None:
+                expected_python = (
+                    studio_home / "unsloth_studio" / "bin" / "python"
+                ).resolve()
+                if exe != expected_python:
+                    raise TestFailure(
+                        f"[{label}] PID {proc.pid} exe={exe}, expected {expected_python}"
+                    )
+
+        versions = {h.get("version") for h in health_payloads.values()}
+        if len(versions) != 1:
+            raise TestFailure(f"version mismatch across installs: {versions}")
+
+        _log(
+            f"PASS: all install + runtime invariants hold "
+            f"(version={next(iter(versions))})"
+        )
+        return 0
+
+    except TestFailure as e:
+        _log(f"FAIL: {e}")
+        failed = True
+        return 1
+    except Exception as e:
+        _log(f"ERROR: {type(e).__name__}: {e}")
+        failed = True
+        return 2
+    finally:
+        for _lbl, _sh, _fh, _port, proc in backends:
+            if proc.poll() is None:
+                try:
+                    proc.terminate()
+                    proc.wait(timeout = 10)
+                except Exception:
+                    proc.kill()
+
+        if keep or failed:
+            _log(f"artifacts kept at {test_root}")
+        else:
+            shutil.rmtree(test_root, ignore_errors = True)
+            _log(f"cleaned up {test_root}")
+
+
+def main() -> int:
+    ap = argparse.ArgumentParser(description = __doc__)
+    ap.add_argument(
+        "--n",
+        type = int,
+        default = 4,
+        help = "number of parallel installs (default 4, must be >= 2)",
+    )
+    ap.add_argument(
+        "--keep",
+        action = "store_true",
+        help = "leave the temp test root on disk even on PASS",
+    )
+    args = ap.parse_args()
+    return run(args.n, args.keep)
+
+
+if __name__ == "__main__":
+    raise SystemExit(main())

From b65a7450ca6091c230e0d8c721182b2ff7d20718 Mon Sep 17 00:00:00 2001
From: Lee Jackson <130007945+Imagineer99@users.noreply.github.com>
Date: Thu, 7 May 2026 11:33:31 +0100
Subject: [PATCH 016/159] Studio: Dark theme refactor, right sidebar redesign,
 and chat UI polish (#5150)
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

* Dark theme refactor, right sidebar redesign, and chat UI polish

- Dark theme refactor
- Redesign right sidebar
- Further left sidebar adjustments
- Wider chat and content area; layout tweaks for chat content
- Rounded corners across elements for consistency
- Show chat message menu icons on menu-area hover, not only on message hover
- Assistant message menu icons now always visible; user messages keep on-hover
- Redesigned copy icon used consistently across chat blocks and messages
- Redesigned trash icon, applied consistently
- Unified icon sizing and style with the sidebar
- Adjusted icon colors across chat
- Fix on-hover background design for chat icons
- Fix tooltip from 'more' button staying visible after clicking elsewhere
- Adjust position and design of generation speed info text below messages
- Adjust design of token speed info popup
- Adjust sidebar scrollbar to cover recent chats only

* Recents sidebar rename, UI/theme refactor, layout and chat polish

UI & Theme:
- Dark theme refactor
- Consistent rounded corners across elements
- CSS polish and cleanup
- Remove unused logo image assets

Recents sidebar:
- Add 'more' button for options menu
- Support renaming conversations and training runs
- Confirmation dialog before deleting chats
- Add optional display_name column to training_runs (idempotent ALTER TABLE) so renaming doesn't lose model_name/dataset_name from the run config
- New PATCH /api/train/runs/{run_id} endpoint accepts { display_name: string | null }; empty/whitespace clears the override
- Sidebar shows display_name ?? model_name and exposes Rename in the row's More menu, mirroring the chat rename flow
- Cache last list response in localStorage and hydrate from it on mount, so recents paint instantly on F5 / route revisit; cached items are shape-validated and dropped if malformed
- Optimistic updates on rename and delete (apply locally + cache before background refresh)
- Visible toast on rename/delete failure instead of swallowed errors

Layout:
- Redesigned right sidebar
- Further left sidebar adjustments
- Updated chat content layout; chat and content area slightly widened
- Sidebar scrollbar covers recent chats only

Icons:
- Redesigned copy icon, unified across chat blocks and messages
- Redesigned trash icon to match
- Consistent icon sizing and style across chat and sidebar
- Adjusted icon colors across chat
- Fix icon on-hover background design

Chat messages:
- Menu icons now appear on hover over the menu area, not just the message
- Assistant message menu icons always visible; user messages keep on-hover (next/previous response stays visible for edited prompts)
- Repositioned and restyled generation speed info text below messages
- Restyled token generation speed popup

Tooltips:
- Removed tooltip on hover for previous/next assistant response icons
- Unified tooltip design across sidebars and chat
- Removed tooltip animations (also fixes related lag)

Model & Chat Template config:
- Merged Chat Template config into Model Configuration section
- Added revert-to-original for chat template
- Fix Chat Template config disappearing on page refresh until model reload

Performance & scroll:
- Removed chatbox movement animations across pages/navigation (fixes related UI lag)
- Fix scroll flicker at end of streaming when a code block is the final element
- Additional chat scroll improvements

Bug fixes:
- Fix 'more' button tooltip remaining visible after clicking elsewhere

* Remove sidebar localStorage cache and optimistic updates

Drops the localStorage hydration and optimistic rename/delete logic from the recents sidebar; reverts to fetching fresh on mount.

* Fix missing cn import in shared-composer (regression from merge)

* chore(sidebar): import sidebar deps from feature indexes

Re-export deleteChatItem / renameChatItem / useChatSidebarItems / SidebarItem / useChatSearchStore / ChatSearchDialog from @/features/chat, and removeTrainingUnloadGuard from @/features/training. Switch app-sidebar.tsx to consume them via the public feature indexes instead of deep paths, clearing the no-restricted-imports eslint errors. No behavior or UX change.

* fix(studio/frontend): reload training Recents sidebar after F5 refresh

The Recents sidebar showed empty after a hard refresh. The hook's inFlightRef dedup guard collided with React StrictMode's double-mount in dev: the second mount's fetch returned silently with no error, no retry, and no toast — leaving the sidebar empty until navigation.

Replace skip-if-busy dedup with abort-previous via a hook-level AbortController. This also fixes a latent race where a slow poll could resurrect a just-deleted row by clobbering the optimistic update.

Changes (all in use-training-history-sidebar.ts):
- fetchRuns aborts any in-flight request before starting a new one; post-await signal.aborted check drops stale responses.
- Optimistic helpers (applyRunUpdate, removeRun) abort in-flight fetches so they don't depend on caller discipline to invalidate stale data.
- Initial load gets bounded retry-with-backoff (500ms / 1.5s / 3.5s) and surfaces a sonner toast with a Retry action on final failure.
- Failure toast auto-dismisses on any successful load (initial retry, Retry click, or polling recovery).
- Polling pauses while the tab is hidden and catches up on visible, avoiding wasted requests during long training runs.
- Both effects own their teardown explicitly (abort + clear timer).

* Apply unified tooltip design and behavior across remaining pages for consistency

* UI polish: spacing, tooltip on source icons, letter spacing, smaller icons, consistent edit icon

- Adjust tiny spacing between elements around the UI for subtle polish
- Redesign tooltip on source icons for web search / tool use, consistent with the new design
- Adjust chat text letter spacing
- Smaller icon sizes
- Replace 'edit message' icon in chat with the new Rename icon used in Recents for consistency

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

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

* Adjust CSS for right sidebar

* Fix scrollbar UI compatibility across browsers

* fix: preserve chat preset settings on model load

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

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

* fix(studio): remove duplicate chat template status field

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

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

* chore: remove creative preset assumption

* fix(studio): align speculative decoding default

* fix(studio/chat): snap numeric param inputs to step grid

- Type a value in any param input (Temperature, Top K, Max Tokens, etc.)
  now clamps to [min, max] and snaps to the slider's step grid, killing
  off-grid values like 1.051234 and FP residue from slider drags.
- Branch picker chevrons share the action bar's 32px height + 10px radius
  via a new .aui-branch-chevron-btn utility; hover area aligns visually
  while staying narrower than the sibling icon buttons.

* fix(studio/chat): keep training-run polls converging and drop dead preset code

- Keep training-run polls converging when responses outrun the 5s interval
  (don't unconditionally abort prior in-flight; skip if one is still pending,
  mutation race still guarded).
- Drop dead Creative/Precise preset code paths (remove 'builtin-fixed' source
  variant + unreachable branches).

* fix(studio): training-run cards show custom name + model + dataset

- Training-run cards now display custom display_name + model + dataset,
  with cross-view sync on rename/delete.
- Enhance clarity of borders and colors in dark theme on export etc.

* fix(studio): match active state green to unsloth brand color

* fix(studio): preserve can_resume on training rename

* fix(studio): keep GGUF chat template override distinct

* fix(studio): treat audio input models as multimodal

* fix(studio): cancel numeric draft on Escape

* fix(studio): use default speculative mode on toggle

* fix(studio): detect GGUF audio VLM input models

* fix(studio): address final PR review findings

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

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

* fix(studio): refresh sidebar/history when a new training run starts so it appears without a manual reload

* fix: API and svg

* fix(studio/sidebar): align run rename dirty check with displayed baseline

* fix(studio/sidebar): use leading-tight on account block to prevent descender clipping with truncate

---------

Co-authored-by: sneakr 
Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com>
Co-authored-by: Roland Tannous <115670425+rolandtannous@users.noreply.github.com>
Co-authored-by: shine1i 
---
 studio/backend/core/inference/llama_cpp.py    |   12 +-
 studio/backend/models/__init__.py             |    2 +
 studio/backend/models/inference.py            |   34 +-
 studio/backend/models/training.py             |   11 +-
 studio/backend/routes/inference.py            |   32 +-
 studio/backend/routes/training_history.py     |   37 +-
 studio/backend/storage/studio_db.py           |   22 +-
 .../tests/test_inference_model_validation.py  |   36 +
 .../tests/test_training_history_update.py     |  100 ++
 studio/frontend/public/blacklogo-c.png        |  Bin 141545 -> 0 bytes
 studio/frontend/public/blacklogo.png          |  Bin 166049 -> 0 bytes
 studio/frontend/public/sidebar-logo-black.png |  Bin 9170 -> 0 bytes
 studio/frontend/public/sidebar-logo-white.png |  Bin 9057 -> 0 bytes
 studio/frontend/public/unsloth-beta-black.png |  Bin 160502 -> 0 bytes
 studio/frontend/public/unsloth-beta-white.png |  Bin 156421 -> 0 bytes
 studio/frontend/public/whitelogo-c.png        |  Bin 139810 -> 0 bytes
 studio/frontend/public/whitelogo.png          |  Bin 162842 -> 0 bytes
 .../frontend/src/components/app-sidebar.tsx   |  491 ++++--
 .../components/assistant-ui/attachment.tsx    |    2 +-
 .../components/assistant-ui/markdown-text.tsx |   20 +-
 .../assistant-ui/message-timing.tsx           |    5 +-
 .../assistant-ui/model-selector.tsx           |    8 +-
 .../assistant-ui/model-selector/pickers.tsx   |   10 +-
 .../src/components/assistant-ui/sources.tsx   |   21 +-
 .../src/components/assistant-ui/thread.tsx    |  142 +-
 .../assistant-ui/tooltip-icon-button.tsx      |    4 +-
 .../use-intent-aware-autoscroll.tsx           |  153 +-
 studio/frontend/src/components/ui/button.tsx  |  132 +-
 studio/frontend/src/components/ui/select.tsx  |  495 +++---
 studio/frontend/src/components/ui/sidebar.tsx | 1538 +++++++++--------
 studio/frontend/src/components/ui/tooltip.tsx |   22 +-
 .../src/features/chat/api/chat-adapter.ts     |    7 +
 .../frontend/src/features/chat/chat-page.tsx  |   24 +-
 .../src/features/chat/chat-settings-sheet.tsx | 1083 +++++++-----
 .../chat/components/context-usage-bar.tsx     |    7 +-
 .../chat/hooks/use-chat-model-runtime.ts      |   71 +-
 .../chat/hooks/use-chat-sidebar-items.ts      |   23 +
 studio/frontend/src/features/chat/index.ts    |    8 +
 .../features/chat/presets/preset-policy.ts    |   93 +-
 .../src/features/chat/shared-composer.tsx     |   34 +-
 .../chat/stores/chat-runtime-store.ts         |    6 +
 .../frontend/src/features/chat/types/api.ts   |   21 +
 .../src/features/export/export-page.tsx       |    6 +-
 .../src/features/settings/settings-dialog.tsx |   12 +-
 .../studio/historical-training-view.tsx       |   12 +-
 .../src/features/studio/history-card-grid.tsx |   51 +-
 .../src/features/training/api/history-api.ts  |   18 +
 .../frontend/src/features/training/events.ts  |   51 +
 .../training/hooks/use-training-actions.ts    |    3 +
 .../hooks/use-training-history-sidebar.ts     |  174 +-
 .../frontend/src/features/training/index.ts   |   16 +-
 .../src/features/training/types/history.ts    |    1 +
 studio/frontend/src/index.css                 |  619 ++++++-
 .../test_chat_preset_builtin_invariants.py    |   56 +-
 54 files changed, 3752 insertions(+), 1973 deletions(-)
 create mode 100644 studio/backend/tests/test_inference_model_validation.py
 create mode 100644 studio/backend/tests/test_training_history_update.py
 delete mode 100644 studio/frontend/public/blacklogo-c.png
 delete mode 100644 studio/frontend/public/blacklogo.png
 delete mode 100644 studio/frontend/public/sidebar-logo-black.png
 delete mode 100644 studio/frontend/public/sidebar-logo-white.png
 delete mode 100644 studio/frontend/public/unsloth-beta-black.png
 delete mode 100644 studio/frontend/public/unsloth-beta-white.png
 delete mode 100644 studio/frontend/public/whitelogo-c.png
 delete mode 100644 studio/frontend/public/whitelogo.png
 create mode 100644 studio/frontend/src/features/training/events.ts

diff --git a/studio/backend/core/inference/llama_cpp.py b/studio/backend/core/inference/llama_cpp.py
index 8da836de38..38c0261f5a 100644
--- a/studio/backend/core/inference/llama_cpp.py
+++ b/studio/backend/core/inference/llama_cpp.py
@@ -437,6 +437,7 @@ class LlamaCppBackend:
         self._effective_context_length: Optional[int] = None
         self._max_context_length: Optional[int] = None
         self._chat_template: Optional[str] = None
+        self._chat_template_override: Optional[str] = None
         self._supports_reasoning: bool = False
         self._reasoning_always_on: bool = False
         self._reasoning_style: str = "enable_thinking"
@@ -621,6 +622,10 @@ class LlamaCppBackend:
     def chat_template(self) -> Optional[str]:
         return self._chat_template
 
+    @property
+    def chat_template_override(self) -> Optional[str]:
+        return self._chat_template_override
+
     @property
     def supports_reasoning(self) -> bool:
         return self._supports_reasoning
@@ -2221,12 +2226,12 @@ class LlamaCppBackend:
                 self._speculative_type = None
 
             # Apply custom chat template override if provided
+            self._chat_template_override = chat_template_override
             if chat_template_override:
                 import tempfile
 
-                self._chat_template = chat_template_override
                 flags = detect_reasoning_flags(
-                    self._chat_template,
+                    chat_template_override,
                     self._model_identifier,
                     log_source = "GGUF chat template override",
                 )
@@ -2525,6 +2530,7 @@ class LlamaCppBackend:
             self._effective_context_length = None
             self._max_context_length = None
             self._chat_template = None
+            self._chat_template_override = None
             self._supports_reasoning = False
             self._reasoning_always_on = False
             self._reasoning_style = "enable_thinking"
@@ -4211,6 +4217,8 @@ class LlamaCppBackend:
                     return "csm"
                 if len(_tok("<|startoftranscript|>")) == 1:
                     return "whisper"
+                if len(_tok("")) == 1:
+                    return "audio_vlm"
                 if (
                     len(_tok("<|bicodec_semantic_0|>")) == 1
                     and len(_tok("<|bicodec_global_0|>")) == 1
diff --git a/studio/backend/models/__init__.py b/studio/backend/models/__init__.py
index a4fbbbe6ee..7addca02ca 100644
--- a/studio/backend/models/__init__.py
+++ b/studio/backend/models/__init__.py
@@ -15,6 +15,7 @@ from .training import (
     TrainingRunMetrics,
     TrainingRunDetailResponse,
     TrainingRunDeleteResponse,
+    TrainingRunUpdateRequest,
 )
 from .models import (
     CheckpointInfo,
@@ -81,6 +82,7 @@ __all__ = [
     "TrainingRunMetrics",
     "TrainingRunDetailResponse",
     "TrainingRunDeleteResponse",
+    "TrainingRunUpdateRequest",
     # Model management schemas
     "ModelDetails",
     "LocalModelInfo",
diff --git a/studio/backend/models/inference.py b/studio/backend/models/inference.py
index 43087cc5bf..7a4c7d0b3c 100644
--- a/studio/backend/models/inference.py
+++ b/studio/backend/models/inference.py
@@ -11,7 +11,14 @@ import time
 import uuid
 from typing import Annotated, Any, Dict, Literal, Optional, List, Union
 
-from pydantic import BaseModel, Discriminator, Field, Tag, model_validator
+from pydantic import (
+    BaseModel,
+    Discriminator,
+    Field,
+    Tag,
+    field_validator,
+    model_validator,
+)
 
 
 class LoadRequest(BaseModel):
@@ -43,6 +50,16 @@ class LoadRequest(BaseModel):
         None,
         description = "Custom Jinja2 chat template to use instead of the model's default",
     )
+
+    @field_validator("chat_template_override")
+    @classmethod
+    def normalize_blank_chat_template_override(
+        cls, value: Optional[str]
+    ) -> Optional[str]:
+        if value is not None and value.strip() == "":
+            return None
+        return value
+
     cache_type_kv: Optional[str] = Field(
         None,
         description = "KV cache data type for both K and V (e.g. 'f16', 'bf16', 'q8_0', 'q4_1', 'q5_1')",
@@ -299,10 +316,6 @@ class InferenceStatusResponse(BaseModel):
     supports_tools: bool = Field(
         False, description = "Whether the active model supports tool calling"
     )
-    chat_template: Optional[str] = Field(
-        None,
-        description = "Jinja2 chat template string for the active model",
-    )
     context_length: Optional[int] = Field(
         None, description = "Context length of the active model"
     )
@@ -314,6 +327,17 @@ class InferenceStatusResponse(BaseModel):
         None,
         description = "Model's native context length from GGUF metadata (not capped by VRAM)",
     )
+    cache_type_kv: Optional[str] = Field(
+        None,
+        description = "KV cache quantization dtype (e.g. 'q8_0'), or None for default",
+    )
+    chat_template: Optional[str] = Field(
+        None, description = "Model's default chat template (Jinja2 source), if any"
+    )
+    chat_template_override: Optional[str] = Field(
+        None,
+        description = "Active chat template override applied at load time, or None if model is using its default",
+    )
     speculative_type: Optional[str] = Field(
         None,
         description = "Active speculative decoding mode (e.g. 'ngram-simple', 'ngram-mod'), or None if disabled",
diff --git a/studio/backend/models/training.py b/studio/backend/models/training.py
index 8127af1ee6..0c5825c54e 100644
--- a/studio/backend/models/training.py
+++ b/studio/backend/models/training.py
@@ -5,7 +5,7 @@
 Pydantic schemas for Training API
 """
 
-from pydantic import BaseModel, Field, model_validator
+from pydantic import BaseModel, ConfigDict, Field, model_validator
 from typing import Any, Optional, List, Dict, Literal
 
 
@@ -224,6 +224,7 @@ class TrainingRunSummary(BaseModel):
     status: Literal["running", "completed", "stopped", "error"]
     model_name: str
     dataset_name: str
+    display_name: Optional[str] = None
     started_at: str
     ended_at: Optional[str] = None
     total_steps: Optional[int] = None
@@ -237,6 +238,14 @@ class TrainingRunSummary(BaseModel):
     resumed_later: bool = False
 
 
+class TrainingRunUpdateRequest(BaseModel):
+    """Mutable fields on a training run."""
+
+    model_config = ConfigDict(extra = "forbid")
+
+    display_name: Optional[str] = Field(None, max_length = 120)
+
+
 class TrainingRunListResponse(BaseModel):
     """Response for listing training runs."""
 
diff --git a/studio/backend/routes/inference.py b/studio/backend/routes/inference.py
index a6b00360af..6b559b9c45 100644
--- a/studio/backend/routes/inference.py
+++ b/studio/backend/routes/inference.py
@@ -474,7 +474,6 @@ async def load_model(
                     f"Model already loaded (GGUF): {model_log_label} variant={request.gguf_variant}, skipping reload"
                 )
                 inference_config = load_inference_config(llama_backend.model_identifier)
-                from utils.models import is_audio_input_type
 
                 _gguf_audio = (
                     llama_backend._audio_type
@@ -495,9 +494,7 @@ async def load_model(
                     is_gguf = True,
                     is_audio = _gguf_is_audio,
                     audio_type = _gguf_audio,
-                    has_audio_input = is_audio_input_type(_gguf_audio)
-                    if _gguf_audio
-                    else False,
+                    has_audio_input = False,
                     inference = inference_config,
                     requires_trust_remote_code = bool(
                         inference_config.get("trust_remote_code", False)
@@ -658,9 +655,10 @@ async def load_model(
                 f"Loaded GGUF model via llama-server: {model_log_label if native_grant_backed else config.identifier}"
             )
 
-            # Detect TTS audio by probing the loaded model's vocabulary
-            from utils.models import is_audio_input_type
-
+            # Detect TTS/audio marker tokens by probing the loaded model's vocabulary.
+            # GGUF audio input is not wired through the chat path yet, so do not
+            # advertise has_audio_input for GGUF models until uploaded audio is
+            # actually forwarded to llama-server.
             _gguf_audio = llama_backend.detect_audio_type()
             _gguf_is_audio = _gguf_audio in ("snac", "bicodec", "dac")
             llama_backend._is_audio = _gguf_is_audio
@@ -686,7 +684,7 @@ async def load_model(
                 is_gguf = True,
                 is_audio = _gguf_is_audio,
                 audio_type = _gguf_audio,
-                has_audio_input = is_audio_input_type(_gguf_audio),
+                has_audio_input = False,
                 inference = inference_config,
                 requires_trust_remote_code = bool(
                     inference_config.get("trust_remote_code", False)
@@ -1156,13 +1154,15 @@ async def get_status(
             ):
                 _display_model_id = os.path.basename(_model_id)
             _inference_cfg = load_inference_config(_model_id) if _model_id else None
+            _audio_type = getattr(llama_backend, "_audio_type", None)
             return InferenceStatusResponse(
                 active_model = _display_model_id,
                 is_vision = llama_backend.is_vision,
                 is_gguf = True,
                 gguf_variant = llama_backend.hf_variant,
                 is_audio = getattr(llama_backend, "_is_audio", False),
-                audio_type = getattr(llama_backend, "_audio_type", None),
+                audio_type = _audio_type,
+                has_audio_input = False,
                 loading = [],
                 loaded = [_display_model_id] if _display_model_id else [],
                 inference = _inference_cfg,
@@ -1178,6 +1178,8 @@ async def get_status(
                 context_length = llama_backend.context_length,
                 max_context_length = llama_backend.max_context_length,
                 native_context_length = llama_backend.native_context_length,
+                cache_type_kv = llama_backend.cache_type_kv,
+                chat_template_override = llama_backend.chat_template_override,
                 speculative_type = llama_backend.speculative_type,
             )
 
@@ -1669,6 +1671,12 @@ async def openai_chat_completions(
         and not _effective_enable_tools(payload)
         and (_tools_passthrough or _has_response_format)
     ):
+        if payload.audio_base64:
+            raise HTTPException(
+                status_code = 400,
+                detail = "Audio input is not supported for GGUF chat models yet.",
+            )
+
         # Preserve the vision guard that would otherwise run in the
         # non-passthrough path below: text-only tool-capable GGUFs
         # should return a clear 400 here rather than forwarding the
@@ -1716,6 +1724,12 @@ async def openai_chat_completions(
 
     # ── GGUF path: proxy to llama-server /v1/chat/completions ──
     if using_gguf:
+        if payload.audio_base64:
+            raise HTTPException(
+                status_code = 400,
+                detail = "Audio input is not supported for GGUF chat models yet.",
+            )
+
         # Reject images if this GGUF model doesn't support vision
         image_b64 = extracted_image_b64 or payload.image_base64
         if image_b64 and not llama_backend.is_vision:
diff --git a/studio/backend/routes/training_history.py b/studio/backend/routes/training_history.py
index 6f34321959..771d9f1e35 100644
--- a/studio/backend/routes/training_history.py
+++ b/studio/backend/routes/training_history.py
@@ -18,8 +18,15 @@ from models import (
     TrainingRunListResponse,
     TrainingRunMetrics,
     TrainingRunSummary,
+    TrainingRunUpdateRequest,
+)
+from storage.studio_db import (
+    delete_run,
+    get_run,
+    get_run_metrics,
+    list_runs,
+    update_run_display_name,
 )
-from storage.studio_db import delete_run, get_run, get_run_metrics, list_runs
 
 logger = get_logger(__name__)
 
@@ -73,6 +80,34 @@ async def get_training_run_detail(
     )
 
 
+@router.patch("/runs/{run_id}", response_model = TrainingRunSummary)
+async def update_training_run(
+    run_id: str,
+    payload: TrainingRunUpdateRequest,
+    current_subject: str = Depends(get_current_subject),
+):
+    """Update mutable fields on a training run (currently only display_name)."""
+    run = get_run(run_id)
+    if run is None:
+        raise HTTPException(status_code = 404, detail = f"Run {run_id} not found")
+
+    if "display_name" in payload.model_fields_set:
+        next_display = payload.display_name
+        if next_display is not None:
+            next_display = next_display.strip() or None
+        update_run_display_name(run_id, next_display)
+
+    refreshed = get_run(run_id)
+    if refreshed is None:
+        raise HTTPException(status_code = 404, detail = f"Run {run_id} not found")
+    return TrainingRunSummary(
+        **{
+            **{k: v for k, v in refreshed.items() if k != "config_json"},
+            "can_resume": can_resume_run(refreshed),
+        }
+    )
+
+
 @router.delete("/runs/{run_id}", response_model = TrainingRunDeleteResponse)
 async def delete_training_run(
     run_id: str,
diff --git a/studio/backend/storage/studio_db.py b/studio/backend/storage/studio_db.py
index 29e787c196..8dc29a9f24 100644
--- a/studio/backend/storage/studio_db.py
+++ b/studio/backend/storage/studio_db.py
@@ -75,10 +75,16 @@ def _ensure_schema(conn: sqlite3.Connection) -> None:
             output_dir TEXT,
             error_message TEXT,
             duration_seconds REAL,
-            loss_sparkline TEXT
+            loss_sparkline TEXT,
+            display_name TEXT
         )
         """
     )
+    existing_cols = {
+        row[1] for row in conn.execute("PRAGMA table_info(training_runs)").fetchall()
+    }
+    if "display_name" not in existing_cols:
+        conn.execute("ALTER TABLE training_runs ADD COLUMN display_name TEXT")
     conn.execute(
         """
         CREATE TABLE IF NOT EXISTS training_metrics (
@@ -261,6 +267,18 @@ def insert_metrics_batch(run_id: str, metrics: list[dict]) -> None:
         conn.close()
 
 
+def update_run_display_name(id: str, display_name: Optional[str]) -> None:
+    conn = get_connection()
+    try:
+        conn.execute(
+            "UPDATE training_runs SET display_name = ? WHERE id = ?",
+            (display_name, id),
+        )
+        conn.commit()
+    finally:
+        conn.close()
+
+
 def list_runs(limit: int = 50, offset: int = 0) -> dict:
     conn = get_connection()
     try:
@@ -270,7 +288,7 @@ def list_runs(limit: int = 50, offset: int = 0) -> dict:
             SELECT r.id, r.status, r.model_name, r.dataset_name, r.started_at,
                    r.ended_at, r.total_steps, r.final_step, r.final_loss,
                    r.output_dir, r.duration_seconds, r.error_message,
-                   r.loss_sparkline,
+                   r.loss_sparkline, r.display_name,
                    CASE
                        WHEN r.status = 'stopped'
                             AND r.output_dir IS NOT NULL
diff --git a/studio/backend/tests/test_inference_model_validation.py b/studio/backend/tests/test_inference_model_validation.py
new file mode 100644
index 0000000000..219affade3
--- /dev/null
+++ b/studio/backend/tests/test_inference_model_validation.py
@@ -0,0 +1,36 @@
+# SPDX-License-Identifier: AGPL-3.0-only
+# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved.
+
+import os
+import sys
+
+_backend = os.path.join(os.path.dirname(__file__), "..")
+sys.path.insert(0, _backend)
+
+from models.inference import LoadRequest
+
+
+def _base_load_request(**overrides):
+    data = {
+        "model_path": "unsloth/test-model-GGUF",
+        "hf_token": None,
+        "max_seq_length": 4096,
+        "load_in_4bit": True,
+        "is_lora": False,
+        "gguf_variant": "Q4_K_M",
+    }
+    data.update(overrides)
+    return LoadRequest.model_validate(data)
+
+
+def test_blank_chat_template_override_normalizes_to_none():
+    req = _base_load_request(chat_template_override = "   \n\t")
+
+    assert req.chat_template_override is None
+
+
+def test_nonblank_chat_template_override_is_preserved_verbatim():
+    template = "  {{ messages }}  "
+    req = _base_load_request(chat_template_override = template)
+
+    assert req.chat_template_override == template
diff --git a/studio/backend/tests/test_training_history_update.py b/studio/backend/tests/test_training_history_update.py
new file mode 100644
index 0000000000..d8a0c93622
--- /dev/null
+++ b/studio/backend/tests/test_training_history_update.py
@@ -0,0 +1,100 @@
+# SPDX-License-Identifier: AGPL-3.0-only
+# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved.
+
+import asyncio
+import os
+import sys
+
+import pytest
+from pydantic import ValidationError
+
+_backend = os.path.join(os.path.dirname(__file__), "..")
+sys.path.insert(0, _backend)
+
+from models.training import TrainingRunUpdateRequest
+from routes import training_history
+
+
+BASE_RUN = {
+    "id": "run-1",
+    "status": "stopped",
+    "model_name": "unsloth/test-model",
+    "dataset_name": "test-dataset",
+    "display_name": "Existing name",
+    "started_at": "2026-01-01T00:00:00Z",
+    "ended_at": "2026-01-01T00:01:00Z",
+    "total_steps": 10,
+    "final_step": 5,
+    "output_dir": "/tmp/run-1",
+    "resumed_later": False,
+}
+
+
+def _patch_run(monkeypatch: pytest.MonkeyPatch, payload: TrainingRunUpdateRequest):
+    stored = dict(BASE_RUN)
+    calls: list[str | None] = []
+
+    def fake_get_run(run_id: str):
+        assert run_id == "run-1"
+        return dict(stored)
+
+    def fake_update_run_display_name(run_id: str, display_name: str | None):
+        assert run_id == "run-1"
+        calls.append(display_name)
+        stored["display_name"] = display_name
+
+    monkeypatch.setattr(training_history, "get_run", fake_get_run)
+    monkeypatch.setattr(
+        training_history,
+        "update_run_display_name",
+        fake_update_run_display_name,
+    )
+    monkeypatch.setattr(training_history, "can_resume_run", lambda run: True)
+
+    result = asyncio.run(
+        training_history.update_training_run(
+            "run-1",
+            payload,
+            current_subject = "test-user",
+        )
+    )
+    return result, calls
+
+
+def test_update_run_omitted_display_name_is_noop(monkeypatch: pytest.MonkeyPatch):
+    result, calls = _patch_run(monkeypatch, TrainingRunUpdateRequest.model_validate({}))
+
+    assert calls == []
+    assert result.display_name == "Existing name"
+    assert result.can_resume is True
+
+
+def test_update_run_explicit_null_clears_display_name(monkeypatch: pytest.MonkeyPatch):
+    result, calls = _patch_run(
+        monkeypatch,
+        TrainingRunUpdateRequest.model_validate({"display_name": None}),
+    )
+
+    assert calls == [None]
+    assert result.display_name is None
+    assert result.can_resume is True
+
+
+def test_update_run_whitespace_clears_display_name(monkeypatch: pytest.MonkeyPatch):
+    result, calls = _patch_run(
+        monkeypatch,
+        TrainingRunUpdateRequest.model_validate({"display_name": "   "}),
+    )
+
+    assert calls == [None]
+    assert result.display_name is None
+
+
+def test_update_run_rejects_unknown_fields():
+    with pytest.raises(ValidationError):
+        TrainingRunUpdateRequest.model_validate({"unknown": "value"})
+
+
+def test_update_run_rejects_overlong_display_name():
+    with pytest.raises(ValidationError):
+        TrainingRunUpdateRequest.model_validate({"display_name": "x" * 121})
diff --git a/studio/frontend/public/blacklogo-c.png b/studio/frontend/public/blacklogo-c.png
deleted file mode 100644
index 7ab9959536d21dc48ed49f7782d579033014b9f9..0000000000000000000000000000000000000000
GIT binary patch
literal 0
HcmV?d00001

literal 141545
zcmZ6y1yq#V_dh%{bc=K&TuG@b-AXBlbjZ*hGjw-|q7tHjN`seXC}DdvG;kU{Zx&Nn2{I+0+Bs_q^t`95x9at5HUi0;0Y^o
z!VdV4-R6<5CI}SB1p&OD%z_->=QwClC
z{*&EQ@E&+VcJq9y4<5+%
zO9g&B=&|wxegBD#DZgOmr5gF2mEN+zF9J}{?WWmy5@ehA6OpFBZr#=;#C>=EqeYQDg7S{^d=0SFm=dl2_c)ecRBjL1ue9VSKcIqZekesUFyT+Vkkk
z{q5e3@4BUGE6h6T5iy?K<1TFb5I#2I0QDl-h{acrSaAb5-chhqD9XnkzX9umbcwU|
zm5D)%IK_mDcfudWP(LUFzCd{NyF_xjNkB{BS)#HMaJcR*(7bm@Ws5Zrj}2u7yrN)7{6XTpRB$!lBa-aYxXoe%UQvFVEPNh=f^dA0fOavN`f0K*T%DJV
z`a*p$-@8e3^*e#tvD0Y)LP+4`Mds$%`~10SMG^Yp
ze}nScP@aU}!MwNG*;GML;s;JUmzM5tAPbtr>^$L=m=~=4qW#MkelwdI57Q&yu?tD3
z%+lfL!8F}sGDVVX6PiItwkB90_ZD#g6U`uDA$m4cQ@sd3Pn)6;KTjXW+n~U^$$Yg~
z=a|45OFlMQ*1V{O*p)rz_{plzA8Pd%$NxUL=qvq`0Fndi<6dJ)-e%r9z#U}jLI!6y@fqf3=zXokbEkk!B~-{4
z=d#)m!OV{I;=B2N*waiJYyK8Q|oR^8$+w}
zCUXhA6}+d)y+;`x;V%25;wL7@Y_1_I3Sk{W&z1jtk_<>)UrDz+JaRs>&6n+=yFFmV&0PuPM$jR~4XxGy^-?MGP!HX`
zxnBoX?ixv44y7U0go+GFe5#7cj(7oitmP=#!O01zY^F115&ocBb}{&XCb9XJKF&LW
zPM$PPZP^_7UIe02O%|T`nE*MLpx?=bm3Hmsr@1R>sBJw}NX?it5~{eVM|@5*5(*0X
zx<0keD>%9B<|;+)=;|~Wa;1u^uasSb5E&f#S#kA_4W%3;{xCxVM>5Omu_SY1^ZHl^7bxe_V)!d+%rs6smfTAZP4rokCJ;#zjyI`8Z
z1RgUMozR&!#gc99{r<~yiS$dQ?$@SBBd$M-$n_Vgz3b_|HlencokUmvT}e@hMsuAz
zm6T4CiR-Nu$KTYcq?Fkvu?*FzAVzj`5aTU!7-GmQ5C_^SC0n!kei_;GiU65^$-(-h
zNx-wz_eZ~rGyxK+25DQH1c6v0Rr4z7)I)N@;h^S`y{0U@iQp^=H*847Qfs1pQ{5$?Yd6>1;
zS#-4t#1f9P4Nd(O_hf_cFnRnnP|@+wUo&po81F-t(1H~;Wzl^n`uQ6x@o^({Ey7B#
z#CNQSGBmB!
zI3F~qmPreUnYxyysj7XJqS>%Wnz~R$F`Tt*-4oqaOn%Z;k)eJCUVSsxXhxr)rG*#B
ziot5QUkCT~P)1L4eD|Lr
zgms90xlL}#x$ZEvj8Txum7*VI4dxscz+UBhgesDw?BHk&(}FZ=U{
zm;)+JCfd?c7J;{IAUvVMT-w`ZA82g3EAXoBccWJu>twpo^<_3t+hZxDM>O&k
zf~u?sZx&~;G)cNJX`(EtnI`op5!L9O_7mqkSwVZj4JI{NyZR+ntND5z-Lo>r+hZ4_
z@u#qj917Rw39dg+rbvd+ZBXT6EXg!*PfuNi7wqo3+*9L&SUX
zdJ(NhXMXbIW3{LI%VuFFl!Pn{q;6k*oh<*e6ocPOSvV1+Sb3w{lru=6$OguLl^;i)
zD3aPk7()sk)O)g9`bg3>uC!`%%?q|rRN*6|HUYuT66un``qoX_jcDp}JH-APq8%6G
zLE3*kTK8z4!BHPm&`}Q|ZgViPY7<{v+SHN12G8WA;>2yc$z7akZB_Pm(opv88lK=~
z+&G+JYj=%|>Gn|yr`$m%VjD3{E+yOC{BMMy5Zs@7%7G9bF{6<^~)S+d}#4(9kq}S&a$z_Z*QrmaZb$Z!T9u%6VT@VhWP=qCg
zcL=|16%p$I0iuB%=Hh?;J~aj2a7l1=CNu>#=ZHAtsuZA+r)a65`d=zj`uq)Kke;+_
zpAst%=}34)SC}24D7-q~>NfiVF3%UFw0Mn3d|@G8EV5^=Rs5phu&IBBR*KUL7RC0@
zPg<$wR&>1aEAu8*Fc$?!5@Q+-!hCd7-W!?g;)_kIqW#Y>ITof}x!8omQTR-ci{JSxUHUHXj}j
z?IR*kf>;cVp5?lHH?3tdydKMeJsvQHF8VKH`0S?MJ%TO{UIF{mHAE{lRBX((5&F`datlAZZmU8&THc$oD|Rpn%X~$P?m~q6Ie>OF;XzTL#%gaf|N3r
zCeez~HXMY*&V$!J=aTCM>o`YlGwiaiU{V#ncQRv@(VMkWNY@C&C1h%Xqx*ihc(a#u
zbf@5P5jFM5onIKTa^w#^^D~CVs>d~F(@s|>KPVJhHkQu@%i5%Ba}2+edc69RIi93Z
z(ClkfgYE00Rknua!KIt|jt$E+YKP+SH0qouSaMg+q
zR?dw;Y{GbZD>aJ+Jeeclj6lV$lYbdv(0QeJ&Rqq$P%qLP9L4`O0A(w5nVdWu?>CF+
zN=Y?Ss6*w1_$>$}D4RlXZV-~(dq4(1Ls^bS;fS?C6qn!h;8HTME4oD_y)1E9Ab}++
zLld(BpRdMoh3o0#Ntyw1kbC$KVdJjk0Z7OfQXQ7vXPa4WcIQr0rrfK%u{nn`YRg
zof1QB5@5JUWOptu@+gD~_J-Nd(q-zMJJxF7uweH7Ui#x|mH2IWfJuH|O7=5+E`uN<#{%8Y9-|#^njdw_=}bSAkIy0=go+#rkBnp(^C_s=OZmRZ6G7&6ua;n$ed&k
z><{O1d7x#IVm^+f`mpW0m?t1%V|G?ACL1W!w!gAwGcgg_(?;E5pgHCq7mQ4
zr@UZTRGw8l!p`Nj0n!zo$eG2xo~kqP=w%I`89buBTQ}*NvN4kaqP8|&7I1P3_t9$1
z4?j9Z&3h9sh+e2Zym3Z+_^2fDeM1FpcIiM8TT)SGAr2B1+a!=2-aKA=1WYnk%xJ2<
zuaXQ-B+o_3cE2L7J^>DW4No+#t?Zi>uhstRX&F
zM#q9<1R^LJacDkOoQYU50u2_2cMOi&-}N^Th1Y~#&%xnYThgAhT?C9^jqo2kDSt`~
zPajZcO?siHt-_R_cw4E5Y#U>qP+C3;TDQS;V1RJZh(i#H4oi*GN=-}vl=|p$p|mRQ
zKH<|bW~@J-`4mKd%*T4t%|FZk;XLlv{Rm`FB_9lH^jc!kplls4d!*nNEK=bPG#E&N6-40ClL>eKH3lDhs6Wg
zy&sezQz3Tr$jH(qKufB8i9Q8i>n#9(nBU|@U+l=s(Tw?Sg^vlh;Zjjev4na}$)H|v$$UvARaBADwW{1w8PcCoL
znP`w*DbKMS{2zQK!rJx4B1aLAM`_?kUNsrZ(}ie0Yi7CBQyenzQ0?PiAls27AMW|RbO4nzk%D&;`m
zXM20XjE(6-?{{Ho$Io7#uRN&c!#Z}qxca$z7Nl()9eHth7Ln=)ojzOqk}AXIh95qA
z0reuWLu%bXyVn4;)TN@L0~Tx$9(&r1Ft}Ud{wRT^j<|xIl}aa
zG_#m270F^g@Z(mgW5D4inVRdBhJYoBo6^gQkiY*B(L@}4lAaS7VHGEj0RG_?S@jm8I{0&g!q3?P>M1W9Q
z7QVl=-&H6Z2ZTCMw(5_b3{=JiXHHIVa8qF>%%eo{kjRy%puCch!BMELVdK-9YQ%0x
z{eF~^(d7{cB-z_ycPQ{@T{^e`J9V92BY2%Htj^9I)ANbrqo86yEr)m=vM1eU-S^2w
z1Y${UDfsf6{ABz9s8lE_bCC$&ej8lbTS?YcKzZMJA=8Fw5+&1Lxk(JHuYdkG
zQ8}Bh4&9?Er(Bmz7xMgS{QnBVZ&870NBfq0;36yegE=$7V-;fGh@mfoVvG^s`iV%{
z?FJK_w)HEPb0dk%%^57y^-&-@{j@X4TQYb;xb^P;@HcnsbsJnl0yxAdJWuNI(^M{Q
z@JF7vjf2ckvZoR7Bq`+62*gK-*>5>KPXcc^2MT(z*#Z%44miJmaqljG{!;1vFhC!-
zq5~toRjc@t9u5!EN*%3~bQruSm=Nd4p4Lsh&}hu`i7+FOm-Rmt@&m)Np8Nql>Q}%T
ze0tXiWOLn<@TFKdNR)+5!tbc6z?AV2#yCb~j#|c_Yo4l}Bb6sO{DACxdfrKeF*$xy
z@LgQf3zlg4)~PO~69?+`kwMGNX0e)oB&2R<@i#UR$J9$^!v;429?lQ;`}
zg;aTJlZ;sCC*Dyi8UQy6olb}0`LYDptbR+zUO;{EgW19u&ZgcDQUW}G?nYSoaFe<=
z)7Wq-y>XMOp5pi2A@
zlaYTpw1QE^A4aZw5yCS18P-ox&%b`R3tGI=0gkVXs~_2OBqIjb(z0X<@Gc?#wcW8E
zz;HS(R*c+T$Zu(^ZF2g`*HUzbzNo%`^YC1wfsOkc;7)kC!kV~mlx@M`9*|Yee8VnP
zPYz^Cce;&Bn7)df;Atzm*_UgA!@(Y~QnXOCU!Sgfk1#_E*1bpo8DIA@6F9Yap^7-!
z7VQ4^QQh*&uT2pP&@xcQ3r
zZ7y0qO&q|$7`2BK
zC^nbr?|_!|vyQgiH5{Kn@*f+p@XHw_u@6_}@3YkXMxwMAeAwp@=#MKv#Z8*yxa7Hx
ztjO{D
z)(Ru@d_~xQJFzWX%C^Q3BaMf2FKmw^g(;o;&PXA<)QC_jhJx$0ELjK4&|`K?U4}{u+TCzRET5T0?gA?pRBvJJKx
z1fLSN>J?WmWeE(Fomb37IHWz9Lp+phtx1o+NyB{bfE@jZSHi{!DjZf>ycf|+`@oLJ
zUhIYH#0rKHj4)Q3oF&iji5HkqB+R3!+X0P&h7r8MQ^SO6dMLHC)lv-bU!K1EeN
zq&7CyB{fY2xi4*zOFt)R$1ffQw-xxPb!U&-v2$apZz`IGCqj2y&mC<6q>f{>jn|
zswnJJ*j4IFQ7~6b`)aDCh_XFcwb+UR1~46yJ8Ee*Z~BL6rGOHLQTi%dW&`*s#gaZy
z2{SKP+-0u*cvoe7)wUx=DAsL!b5yfgz;p47os`4K;0S9<7tjtkX1s6gQwVJS;pZ2K
z-5Qnn%hS~12>3^kc44$`$=t3i8>V?XIsdK@9RBRooL%zV@YYq5|EfnqWna3Rp)eksZc6E}dK
zZM>Urr_+ZoEcEdkHci#RpbJm?k~}9
zaq4_kq>AqZ!6w!krP99y6bfQM^S>j532fW(Rp8xr(SnUDn9(3T(v76tBIjdCL`e#86YC#=a-S+Z*%e&a1BBp$JO
z#thgOUed09!SOamRM@dN+IoC8KQd@TA2ZhW1k^~Rd6nbLWC1g;JbT~OeLZ@QgB8U3
zoZn4g4zX*x%f$MJNG=KDqcS=xOW4+gS}Qp`Q(~AoTiMTP+hnj*xYxU;srIf{c!kCtnmTI8L?ofAbPvi
zDfe3dXEm9?wBMumJc5nx2}@>2%LT~(EF
zw>KL*nXg&isk}f|LY}OgYc>jaa?l`$AR~njJd)V;zShp{$XAv`87Q|G=58>V=zo>`
zKe8>csQGt=PY`(q>lau8;{}(l)w*|bh9QY}4LR#(!06h!(QwoN-1+Q$O`P%@g}f`d
z8`drPCTm+_#79Z3)rwKj5eaMHejr>V1)sf~;7C(4DnRxi`6iFx%kC1vNr!SSmTCRO
zQ1<3R9_|uY@(-X}+f8Me%Ic2D_o+{nj2zer-7-KH)u{M7PNk;+6x0AIIG`5GhwTZ9
zHM3hc)l9GB!(zITEp(g}KF#Xu{bs@yag75Bcl;ScuV>PXN<&dj<0XB?MWZAOKhhvj
z?mry{op>Afv?$u6tIOlHh_jVivzSYEVSUWK-FcP1`rRKEzue2
zx(#>sUi-^wAHX-9&=#0E&9`c(}=KP(}|y(*NT4Kw_J!
zN6&8tix&u&gcuy(xBaJG{NzD(MPF4|>M|zvSkpTm0kw1?zrP^w23X0+70d*`HBEMB
zfK2CU&FUaA^mt`*_j5HKva)|<=_GrH7&nY;L#C?U?rU*I-?i~DF+0Ugsl@X#;7bUY
znyd95Krdp_@3anXF=Oi4tU>@X+f{VeeKC}Ly9jgVnQb!Mlm>~cp~|vxnGy0P(t!Zv
z>XwY6j~RE@<%{E2W}3`ecPj&$1}RlZNgO%ab_^n$+?B$pBKZmUe5am-{GS#W
zb%eThufGn)H|;5XT{c
zHTeB(*qt_^3R3q!wgwN>*dq-GILZ&6;CQ;|eX+#%cYfh~oZx7{{cSwtKPmqNy`*8?)DjNCWuY6CD?|#H6$Qp
z!A3|T=FtQ9gr4tE!MAH{iEmVdNn*_Gux0?$qQcZ?bO^MPhTIZ>ec5JnNQS3BaF^YT{M}7D
zX#~WSBNTv68vB}8J>hsF9Ammw>b6T9?yj>1=4(Gz_9t-jZ*#IU~O
z`3uLSR`ii+m$V+rFh98Yy~_R~UD!hPyVqH!XSup-Z(|nQ9Md*Vcm+AyZ`z^KxA|Sa
zEOv@%ym;WA1G0E&2?I)d1`c$8Hv?#OO>mF5qW3aAi54gO-`fD=
zb5H+%grAGK#(e_h=IAo9_H7TYm6d7FtnbZJho=VU%^ck)G^QE^+5A(xROro`_~6T2
zgK_C=Ye)%f_{js?n#BM%C+PLie%!Ibm+tFR)`(q_Cpn)`rHwZPL#dT}8Xk9Zn?L50tanZ+-Yaf{8#J$^#pF)XgzQnxRt^bFJD){dOIBs&1PFpu5tNinv5wSJDoxi8Z#?pZz}QlSYm
zTdEe@B8INh()h5f>;v7JM8EP#?pOl@W}sQs`%fGpBOWo|{0G9M*B8M)6rocPLl#BF
z)xqRQBV2W`&z0dXZJxu$&mx2P<@q^1@9AVwoirF}Ilsh#>d1$RbtrvgF!kr};)Vd1
z_(}?+;s7O`(Wg3m3a#B*1+?{v&JPkM6RYQ_aJg2z)|RtR
zKuIcCridt~X@IW9y)~d{EeROnY`KM(N_U!C&?Hv1#NbJ0vY925Bf&zfzKZS!e0K@n
zLx#rb0t3AtGpSb>y=i>|2p62UY{8<`(~z4LsxS%=y&!IcnJ)3HF<6kM$B9E1Z0j2T
z8(n(h022MqJsIKotb3E00kFypf4i9(QFVX?%)%dK2FTm`$JHw283TI(#e>tc#=T=M
z=6d5}!sR6dSz#?@D1tAV5#mXKBp}e&Oj&)JBt-(@%2JBuBgW^979`(g@jjaB#NF1C
zZ-$kryyB+t3M7ASv_Jd#42Ry#Zhl&Y73SCf_ByU&4%GPQ$T-@vq8c6BiT|i1zKWzmG-QbxzjI#ZT6iD_K12m0&J@G
zwp$ebgXZN2MG|bB>VdgN=9KxzzY+qP@Hs8^v+tFIlwR}AQj#sT2}{tNLzyZIYFM|&
zHoIHMgX0GnDHbeUIGY`dRZ7Z`s|aPCs@Pgd0z#kO&VU2~huF85#7B#r^ycUm{sQJqi&-HCMDJ
z-Yg*`>LQ3-Hq~T9Yu#rw_;kmL+ZP;9?v+2x^dB6GKQY;PFeP6qUCoju#?qBvAL3u&
zOK|}g2=9)ny10GN%TKCs9p4uCksPmu7?tQt215wlnLmd%ESI>r
zD$2MT@CAp{9eqOmheUAFPull}eBPgV(p9TctDISDPucxIQT(&Z*t{jjKJ$MMGn+qQqs+_mAx-f!YBYk9LD!4lpE
z$4{?1N_-6Jr%Bw$`}h$9IGg~HsCP?-kNoaP&KJEPn^N`!99zRyBBYAAQIjfuW6l7c
z*A6$}9UC1=Px{FUWZTI4Sx?Z8A7QR^iA*4&q9;(DZ8w!KTWWPEkRIRJJm!1O?@ScQ
zOa5pwM3FI?gXEG$?uSh7?3NU|xGK*4RHt@igb0pea~ysUx&e824+lHR34%C_$xcw6
zWhx1wFP+C7t`eGYW4}We{21D=@g}Z|cbRzqXsW#UHaKHO)e$9FScUyV%xo~IqZFQq
zpm{%@X+zRUN3&#WSz?Ozkdcp;LW%&W&w%06tED?-d|i_{>OqETMWdai$RY%ok0jij
zC{m&}-}{M<=Ag$SL!<5JO-v)MqU}BoyWFxrW`6!mjv%|_G`9#}VdCO;=E24$?-!Z2^MFun58Y+M|65+aM*AS64)85;Xpyf7Z8RwfHFw_^w<48N&?@+hOQqLMySpBbs+DchDHQY4CO0JZadV!w>ZFBhnK&5+GU
zQ&(OWk9KsucT3~Pg1-rK+5iIs<^>9yrCHF@op@Z#>^i|Hszw2m%gsU0L^cVK@f_3U
zs+r&6N;aBHDo(7m-qH6qVnegF)-k2bzvpZau2oYw;pB4zgZT7)TzW5Q&O4G^Vq2=Q
z*-k$WqH456FRs2(V8w%}!)x{oRT@QZ@|S-^YBh(7)>&!5w2E`^1lLJ)i~l>i
z1Mp^Sz~A(k)Q(t4zga
zE=F3MU!%{k94m`}HYcbP?-sAtd_P-)$VJRCVIe`rLkmTa_8xaVvt
zWh>2PbbvR$>EhD)mM#njJlK#e%8@R1MH+pUlk#o#BIqfDb)zG8ZYWdt<@yJ)$-|K#
zgd1$HgV2p89mf=OlN;~=kn!GoqFzld*H%9Khe`(nJjtj9
zCRDZFu1X!(lHfDsKsgOjsEkjWuPbSOKi-OA!VMExZ>
zURIeBRfN{qP{q5$Wi+!BH>FY~(#-)pbkcbw1*)VzEj7k067j^Ll!zr#F@
zKRYNsE*U%H_=nK&h~uM(OscmkWkLJO?um-DB!kluhJP^++uHypx9O%cmb(w0NL5}T
zw&`#t40xiO=Unne;(x!R8G_LupB_^KHR_bUKne-0Zw_fca*mzV1MAZz<+?o}+7qKF
z-qR4Ns-{~$W{jM>D>uZvlQX#Rb!R?P#6oOM5)lZBGCyG(IfV{2WK)A*4!{hje4Vk`
z-`=meUs92Uh1%<_yKJ9;{M&A_YTeG;jkxzrjz1d34jx%Yq(^8wVjnMKAX_hhcbk8f^xytAfSuf%e?d#1jAEsqLwtgOC@F7pDQYr)*9Gshvu;uFGT8GI
z$N^G_M5zFs^nXXC0o`OG`I;Iqp4l+KWBlw_z%=>aUNEJL!J%0*L|$bH2uGA9t*8ew;D-v@^k0ndd5WMx7!P
z`m|yV`Cf00Gv-2RRQ&48ZiK^(TfnIRTkLE~@=bk3Z8+IAcW{{S3gnfu$j>236}e{n
zzZq&SVBP>}_SO8pXBo`4mH(E1r@7raJ+2y2j1SMqnTgJV%Y!#tY5*k%lG4)JiN^lm
zSHgG`zLmR<6h4hX_`H#x9!|>)h1po8&$=OOrkN$Qg$-CXrpATN+1YvaFmUepT*t`h
zWl9PyH8pj1Ma8Yt)6>GjLdD0AiPY8Ag}f$xF|*YnD)Fu_cFyZ
zdaS2iJmN=+C=fVoA3h;_R)N_>b$s-
zu<=MV+VCV;eq*xk{#?!e9&qL0p_SDgd@{!8ca&aO9&LhW{E
zJ_%aj9iO2^YtU#ND+W?zFwOX7Udq$<`+DEALl7Duv(3?VW#=c!nXRvr
z-X~2!G*j?=S%a&tj?@%(!@kH#I0zu1$Jv&s@?4Sd-Ap#;3=
z?#V*J0>&mm(&!rJ^AG*k->n@E;M`)XVTG>Ac%$5Qp
zpezPYk+)Zq9ThZ`o-jo_v~rDeOwb`kJJWcOajC(h3KDcyUAn0xn#AsiZwk
zB~W_9FWlmJt69XK5vaQqui@SCEXK{jQklVfk4vE_LF|Y3;=ukZYiN)VlaNr!BMm6B
zRVZ2)bTe3-#}Kb@-#6~eJfOLV0lPYU29fRGnO1BQp@KqWj;>2q6F)EkHzYTsM+5yG
zaBh11*6vKz{vEmz?ZK4k6g{w30J4k-VL&?&=x~t3SjIeL)Bta^xm@Gi<|EMJRgsKt
zng7{k(&n$3tXsH`vUbzpBo?KxKW2>KfAOR(@_P=?$jKNdt5EX^M`q}mlk6@qq0S2;THd
zFUN#dxQ6`LUx$JhKOK!wmNk!pZ=&%O%ljgB@UqgPCx{!6^s!RNHVR#!@XK=y{X4Ei
zY}f};HFg2~0+$&2vFd_6ux6+NK5g=X8fQWqKs15Pe$?ww)8OLb(!re-9O;l1LMzmA
z@wa)%;V8^PS2Z~$CAYNnUXjgA@i6WAy+hCccG?nlvfi|IS3$Wc}
zumwa#T#G=}ZY6%ib8}Ie>9N}w$RI9e>}R-i7*kjBQQH`SR@-BPjv?-rzm~sDJQEWY
zZCyXi1>lu{g0-_hL*@<#z#JYbDd8M#VhOSaUbeSC^p|B76%)$=kPLue9$wxDp`r4r
zX=xrivkx|ATG+~)&-e_>Okd^a=fh@9UP`W{du2PB=ysRMH?rK)`bqO;(?;)3*FDqc
z4^08jI-rY0pV&-6Z{1ACD6p>%e&{EA)ExGs9~2va%4rgntbY1uFfn(l084;!a5EOC
z1ZI>ZNfdbp5C3oENRBlgB~TTzj5v(_ZY2khO^D)Ore-A*V^!uY_z_yCmE|>*V$HO%
zysL5BWNFdcC%T=l+03|xM>n@U2RbWbNi6Vcd!j4;lPVW7eB{i~medzc%I-;o=rM0z_PVt1*5s?>nJ~7)
z!$VAipU`h;W0v*5MMg%3sj?&ajUDsz)+3{%cIWS8nsyhu7nVL)r~**;URpled-D8i
z+6i`xV)v{}vqz{{F`@CHzp|2aKKezo{Hiu#1w;Zx1TqT&_1vDxAGjXcNo#SX*pSI>
z>zT(9ZpUh47YCXC6S_iE;L?awZYAP^_BQ4~A7
z-E0dziJboXeG-)c`}U&Z;`3hy^A1L2^=UczlWbVQ^Y2~c-^1Qcu3Yb+ooDf9i0C9oKRuY
zP2PBQwr<=KD(AaB9habRC5inoibzgMnw4ld)%_=Xx=KHaN+?0?8L+32=jd
z>5SmqoxVsaamRyYr7xuB><}kiJsSRe9hm2|%I^c(qwJlcBzIE+3xfVrwj(4C?@H5A
z_TssN*}Li4^0n0T$L(zaIveHPTN%d+kXs$f?$j(c%~)bb4XCZ1LX?^)Gt8y)LthPG
z)d*i$o-p^8Gc)gh+HeFiz5|Vb?zWQ%%`PM;be$cpKhsQ8idVSuzu3OM65`TK9l5iT
zHuPB`WuTD|9td(FckA{#m;YLhDpS_fVBvccl&!e0sP%V&9!Qi>l=38M6!4C{VzLP%
zPqim<+P^Gs9lK4{Rc`iaK|r#?_S0*G;*8fCd+*vNH6*
zqd@)MJzFnd_)W!uBzUulMSZTR5yO5^;Y!V)TTXN9(h55phFK7sE1ck6Ekv#Cn)f{g
zMS|}(W%J9$ms9jV410Ou^zpNXlYMW)sPZ+d6@>7`+u7)PLU%T(S4(
zQb!^Gx_H}W9n*(HyhRgmGfuXjGvR{0kb+8#SE7#gp
z#k+_92)H35rj${Ht^-rGjojg0xF>hQiHgcT{GR#ZJZwc*XRUU4_f^(xGX=%Nfs19@
zP@@FpjRkGHb;<``tuM8hC~Wkd?q0ccnx+rCnf5yTJ^Z_f(wX~%ON`39X0EUNHn1om
z><49$-V7N!extJI-3!3FCHHt)jGE3x4L42+U#5|NZmi}0-iY^jphawX%J;JTFe9EY
zyCjUa8*%wRszFC$NFWlsFD|IS)*1%>n7n3^0putU@xZt}a6egeB!OCsrk1;2A_JCnl;d-iNn87
z)M(GECa}hl2-bqAxSw{`dMpTAwd49X)k*9H`Dw_7UY|-|ti2Dz2?UMiCGX%~%Xeyk
zaI5MnL2nKHO8z^x}7ppOc4
zO4|V{9aa`pM@Ugit8J@VU(CJUEivKWFSFfpEuR^3^sLml5&}@#-U%P=Gin+dtNG4X
zALW9?m;2IQIt}Mro$SmFiV9#qptnY&fIdf>d20hfi=vek1K<9eg|cQlVCpD#2l+sD
zBBYt3#lLVXMGg6jz@>?>$y;#W_5W#UA2W_U;qV_{fVZH}Pdekk(g9cD>Y`@`V%`ns
zMg2Y!Hj2+bkE}tDlN|tI)S*L!+qV1+sS34KydU}cD1>tb^P<>81sBUm(ZSDZD&fr9
zEN07CmDAkJD%5mv`?p3*<~N#ejlg+UN4`dngHPI8?KJGZQQ1BJ(I&p_IT4k%VG7gz
zq_E3X#*5>8I*Wm6mGA`l#)$fBR+oU#sSux*;9#ck`Pt84_vwFK
zN*z=?n>Q7uj#2BjgYvi4sNTK?{2@D8#whzkx*Q%|bCxb;%NFvrp
zbfpe1fFWc`MQdDSxXvr=tlL+II5Elk70{&Al`_zwia&+rV}Do9gJG?OL^am7hV5qT
z+)BX)le9NazE)QJLc9$5Z|LQEn*MQ9qdmwcXt#SnVs;=4ftRff6n;R7kAaLpCsGUI
z#|$il^3QesB+89za(4#?nDO5dENF19O#M4y;=_6IK!TFe2k#ynNai8?)w(9*T@~kP7
zCP==W1!_3oR3GAaSYxf|bk{o%dL>T;Vj&JS#QM@oi^A`Qp98a5i(Udis~qx{!ll9X
zMRQ+XCl;8J14$sh%BHapa5Kux81f-Nw1Qd`QQb1xD
zQaT0c5~V~M3F#a{7)nBrMg|yCq@}z6XZ+sxPd`4;wOB5lIoEa0-uqSO8lDcW7;dOc
z>7(jyQA6Mw7H7}ylb@i&9J+Z*9v|06zvhc{?Myh2bG}`2WfpvRT|dbC?n1}4vH3ft
zP7dEK@1}w=Lu=CS^89E?<}$3Kqk~-Po(uo^>0u{NUP^%GQz;rSlL&e4Eyq^FJ?6Ju
z4fdyn-Tf}}mJJ8#@RJFl!!wn#@<~s`zguz45ae6DZdAIMJUYEO1NX8lB%km3y!Ov`
z_@8$uyk-44w>MQl1Y|FIDX9by663rWz=o#6_>_!tFmEllUOC+}I$JBZY#?|Ao>pz-!YIorN=h+oGR)n~fV=}|`h&L>e($C0@`pe=Aj8@kPbPII%&Sp)q8Or$A3
z2jjIzKO^q=U7Qsew52z>Q|X`%@6N&%we)Bj*q<^Z?m?sk3{8K;kF|!o)VWIiAW)Fq4e!coId(aTo^CZJT8czqJdli1^USR7iR7U3e
z_E;P_>SESg9zV;BIq-=;XAtbHVI3^1$$^GUB$wxVQroe(NuJOpw0mn0&#l+Ybg*3>
zDc#nqz>`idmFs&v)#(q#b=diKbN*+9K*!^osUy?;4cXXNN@f~Rn$s8LI}yCkv~-D^
zSOtT#c-zO2!h+HkI0M^NHZ~q*D|2+CK)8v#;!pSgqs8kb|7UqYJ;l@-1h$4P)FRms%=SJigZuBmjD5vxp7utBa4S2v4Z5G0927
z{%bBA!xf@tX?U{J+d(*rK>1ut=K0~HHv5YcBMiVqq3$N$pUHx1G6%FNSi1kO@y&Oh
z|7Xdq#!>!vKM;~X$5sa9_D1)rv_3mn6^1o)FhcGqH%)kPC{2fL$12_ZZBx?OwNm~Zhe_KkIn?nZtT&2f99B+zkZqhY`@e`!N8$H9Nz-5V
z{#Co{=22iDO9LCVTCm9t|FqZ`ORO*@ly!_C`)Lj^i2H+Z2LrOT$yPCalQmt7!dA;j
z)-m4Ui;hlKN~;!L*nFKcxviX5oig;}DtU&EFC}Y|6RjPaAI=GF9M_rWO`C7;IFp1e
z$?a%q9WMQY3sfzSv>?h9
zcFM2b#u1d**>+zV%cULbu+m|!Y!y9=t+p9O#c0;rpu&JZ*}@POU+km0tG?*3L-m>~
zd5}?@eykw${vG2|*@0ZbvYhR!eXI8w`};3=)e0CX`+GjVtC{vBu}pQL|
zcTf0Rezr$!|E}zJ?ed7(5Ql8CTuOO+?cw|`oC;a=mpO7{+%lLW%SYr12>=Y|+X^A!m3YXcC$KyIt~(F~sY=wmlC
z%t8-S_oIR{JetTBO0VT{m7^@mDk3t}B_2m|pscAHYAlML)uB@B=ryCuU2ut6-1NJ=
zA1!@iHR*f0ZUu+;936*h>`%Z;=gxNfJJF~Eks_l;iTR~n*{vm>i0vr#U58Ju_sFY7
zA5ifnk>MCoTD5M8F3~q#5r))D+_7KHrwVWKwGn|BZw+>ZV-6sRTK|j$sZa9A?4aKX
z9r*pBQ#ir^61XOTR6yc4_=l>85#lfN1Cu_Q%x*6t$u{^wHlNd1Z8AmpuV?$k!%{pj-Qu!3
zUpBzu%x?m%nVwJQ*V78F1yIjAe~(u>ZZ;mi+>@05;0nOK$ydXwNg{CCjk=|@$n?$R
zxytRu*3B2*F&AT6evoAP8S;mIhmE=8HF3C+rR9vTdvVaei8sFirI2XsD~>w(
zQThUdL@?Z^P**zr_5P{D{tun`ogRVC`5=sX@xBln)-x
z-odYrv*?HzU{YpDyk|WYhM(8>8$Sboj`mwDl)v!*T8TfBVu@Hax=CJ>XraZ41H{ab
z_0kxQ-H)=&kD+eS#f`n?QbeL`a74Ck32(wPbNGi+c%%N>E|?k3GRQW6thWIY&k_A;
zcxXrei^paE-X!a7Ur&<(=coj+j0sDM1SVXa`*;@!fyvCH0$wp98R1LORJT
z4umR-FBN!MoJDK_kBr!T$x-s{k1LtQUv@&~#_tCe&a-PYEqw_RzY&P9WZ5z8*~5LR
z%Oa{@)1G>|KbOs
zP6l~b{!HZe{GFjGanbtB{=Zxj*EPpf0*V*(Yi-=vG?K)x*HFE$WWVqko
zYZlX){?af&)bEz7C~Ntf8p!Nx!G>P^4!U(4qQ1Ro)~d4ETAmz;{@f(*hyOVxm7oAEuBzV++J|hOe-M(k@Z1WX8{B1}CM#w1|2>RT8rfKL0
z&9Fr~_zphiDEWq<@9V%hL;?>aORK=&x5912UoR-)Wg^5X(50^!G5xQtF>5kz_z%pr
zAoscMsvc1W+uQqKDKRgq1dxWXkniQyrF5Rg&&8^=C~xV)cwsRV)k-B@x0*GD1zG`D
z8`hZNSPu|VNAR@G&?Sc0k6Gxx`Z-S!K6?AMVf|f>H3}d+Ilbx2*15+L*tmb18Q_?H
zOeao1y$H!sY|dtZy}BrWzsnWF{bu`AOwIZK4^RA2L@t!2klSSnLf=+Gzpq4s03@J+
zJZu)N;H&_e!jSX8`_}9Ujgw#xdC0fX^(qYrH(l-TJ>dn4mnR@))pa*iKUmKCe0q_VHjpLPS8j|B{xLEagpOh@0$iq&-un>l&$jWBoYyiCdmI&NDEI0#Us28+u}(J!feb*30CE8UT;4{(
zISt}olEXlJfpI|jJ0-livC$Q82U0|hDCRC;;K*!gD*fAd~r84UZ-Bwr>e#&?@b;e_nINxx$k><||
zOMKi4VlcO4x1N7V7v{M4f1YjpYJDH>-%>ucYy1sallCPVZ!tF+XMVd8|0WA#=9XCbkf#
zDMTwq=|a~pxl8)8KMe@6=Gm-`zdYJ~ab7;t{y31!?CF#dUIy4U5YTd&k|#+K`16Ik=dE?628`CXT3CBsJwrTjjNewJETGIR3nmSJ@o*6
z(kXLs2eln#pHWvQz7D5yQNngJIf+K$0DP86^}Lxn^gFH~
zca_vMaEv0m`jjv;B#l(!oxnbdd5&G+$Z6W~8)XV&Wg?~7JYjl+aB>80eRRwk#m~sq
zJMJ$!f|lcr<9%}M&n|Lrk6<}#*
zjGuGMTf>iNX^$zyLby*U?4Ce7-47eUuVM8)TVJW3ege>9Mg3BmmqQCaT}DX>|2kaY
zgD~iOia&Yar09pUxkY(!W#rKpqap)ddU77fepGt!+|mTRv2+#wF{Z8DwQ40
zY*282GP#+pp8cfH$cnX}2RKC|0DdRV0lh|9hpBbsz|PPOC#!|l#$Mm5oE?TTzKwaA
z6Uu**x^kD--qbo@dye872v^Y?<|&Jk>J;5!{0_|R
z{RiCI(4SY6bTNV<0?oG=ks>UwxdN2(Qa>Ij*WFBW$BCfjYXv>)hf$X>Zs|;$LyMk-
z9G`W+J`{E3NmN@}@{~VcgUAHe6R$6>-D26l>geP@mh2a%7&uIh5JRB47*3_WFaMad
zZuf-E9{eWj*~uJvNGZzK%$e$})xRVx<(VB$Wc~q3?{WeSFxww;mX{&jsaw{``2YUOtY#X(p!B
zd*Rj*XbIy-T}VVIMp?W#`&fl=y=)F2sx%!@+I&p8iK9wC?aEC?^Wkn~
z(gur6_55ozQDz1cyKH#0Ex*452ZqG89zZ1w$vT7Z-<3oq%fW
zttxpe|4mbQu70hgt>5>cvZ}Q4#V9BJXts>&J?5(C>mjIU57YCYYu(rJAsTq9FRdDT
zV)!K;J=s7LWjQ&a|IZn#?rJWD<*bFslvRVk_K)w-4w0qiH!bR6#&mH2th~6hP1}@Y
zI2?)Wy%S)9(GCeWI3$K8HjmCVB!5CRDYd}c?!?i+7J
z2=g=HZtw_Z6dn*y9JoA$x#S-lzpMAq>hX)b3?oX$_4OWq9_V#5^|f%JQr^Br<8}v=
zu*RU0{K~G(<*rOE*j&T)i=QfZ&pa0Ktq@|Na>Y?wHfaZ)MezXR5bm>Tc}p1dSiirb
zj?-A&x>j5722m6n82JO~C;qxcEgowJVk+<5z6Eu-jn@e+Xyq%z+PHs$O;0fn$dc$#
zAoB1~1)s1lM!sYXWNY$NX}4fSDOZOJ61*r3Q?F;dvMs>GY?7sh*^_@r7)9^vG@a1c
z)h#4om-K_U`Zk8x&du+-SfLoVW|NcSV>31b#-iS<%*n@VyuAcMlM;5~B1-vp(Lvsr6uZ
zqKT#1qmNJSpS^UuOW(okAUH>gW~7`r$z(dM3qZK
z&{k&%yto52Seh_%%~{5U1_>1z>zE`nrn+3hfPO}DGi3$iGt;50Dpc;h8E;_>2bGHk
z(2mHIG#=C^A&qMkO}$DId?
zVlwx1F!}*7?f)Zh{&c4<2=-NHbjWXP-{BKQ~=-@KJt%>j11^68?m_K=Zo><_@=?xgYJ6*OKK&
z%30lKp#~lc@z#dBgVsw~^$irRBCG*Qzh|5z#RY8|dq{kQ2h3CAsQGW}idn?PQLhU6
zg1lobC+;E(tG|Nak(P=bs$C$q^UN4|1?hxzJ+%WGM^f8)n+TBehjpccJkcj6h3toJ
z))CAMXYAb)ne}3PN4qWyuY6j(@cp%={1K`mN9$iT1?xWvAxf7T52qsl{@G~_iwUh*
zx;WkJ6hRxjNA8D9dAR>ve&P=_^2e~-yLNJ%=^g;9nvaw|)*3KR7_C8b_h;}zXG69x
zahY$~LH}nVKapQrAO9qKg0HMfU|xWu41!MK;|Y&UNUDHvxitNZ83yzY!+NPW7>6Ll
z9*EI05=+CY1j!F)BLdY7AMDLb&vbF)dDjvBX%6+;4s=9&S
z`Mx#G-nX>ns6DPd-(C!_3GhI{H7pNxaB~FxG`8_O$%lEB@9J7kRagl9KR1>-t;PJ!
z1zc)};A<-Pq*c755&7Ms6U?yAW>JuHnIN)$qo__)N3Voa=&guxfX`yl@$XpRz3OKo
zcgOb+1*O#a)S}etsCMppbSY!>$4-6~3`yOKc4F^&rkO_=Px-j}*e?I<1l^>D@w^vp
zB4k0r0GL-%YQX&Nx2I(5UYm1P4
z;~SKtD|&ub)ff#QEo!Mj?eiUZqekdBvQqYDF3)FbK~pq@oSq-ivyTYdKHTN{{@6zgiP
z^OeqbviU{s_T8rlpC%Cy2FQg&4qa_8kgdsuzZMc2V=G&qMfXHO3jtxhb(!R7|1WS*b~
zp5Y4pQELD)#Una<%e4}uKx$#MZ71pO>YSz=g&o|UXH}?DCu4wq%hc})73&9A{*HtL
zwVt(FpHtV{Mh)~pgYyLYDKXyCY^c+z&7ak1Bk%*N?vvTNx?OHA6Qf=a!xK3He4R%_{
zMZ#1N))6D3GoA|&2B)b=z)o33*rGPOcwx?MR`?cpEbDV+yN$$my`Q^5&@ss&`0;W>
zJB*@t@6LwbDrrvOgAG-%Qo4`>9sg0c4k`dhRB|mEhYJ*cj()SuEV+Q-P{wB9{K1m`
z1B3?ameo!7PSUiG_e`yBMC#S3z2?Jk{vF?3-bRzyxscLz*@DbP5N!CI3{L#mOskEN
z?UOx7z<`-=O?X&ZrvRE2;2bbYE%Y>}Mjjle3TgHyu0f
z0mJ3zqYrCr{W`?L8P+5;g3V(#C7cSUxmLR``A%`R`^y!ulC=qTrIa7ERa|#EN9*{U
z#AaO#2|MZ}*Y}Z&%9C6BkN{7=TQ=u>x)SF&q_=sW_P~$8AT|pabPxSe_N>13am(~}
z3csVas-hoFWu#1WMr5T{Vgo}E)uN^(5G_?^;j~CUBNlA4(0-FTMG#SB>E1K|luHaX
zMfHDs^6rdC!W!kzC}lz1W@cF$1%k{2pfIwcFvrlT)bcsG_gaBj9dS`4jt9WZ-_+|M
zhd0XLW(7MM(9%|hEFI^b0Wvqy92-Bp`6*4F%s)8u{_dDGoo%dd)C5pN+?n$Yx8Vv}
zN*A&wuw5KLx07R09NR9ny%k9^FyW4e@NCu?9hcsuwf&GnN$mefTr$l}%Ewcsi|O7c
zBH8OlR@@CPB>yCTbN5QH*Rt>il=~OpFi()e4x@p$m#++@Xo`=63=#C&1TBMO@i6Mg
zD45A-wG@+XvkzP!^{P%G8NfTrsv{urqUiLCO8J<8F8TcV3fjnvkDN_9C1F1_f9=
zsFt^kM2YuTC~1(f+m6o7F&v2AqlF`?;hhPd90!2F%>#}cdoqsf8vytfh-2;bM?K9a
z>gWfZA>`E~n1>Ss4WNa28Gu8LryR8txhzw@TvQiNPQ`=UNoolct3~XbdsaX}jT`s?Q~-yZYe@v48G<
zAX*q!SVX^(Dp&-WvzJr*f&8Cs;!jFYvrO2c%e}W5Dvq>sw&uXaxMSEkUKHyBMQwrY
zy0a~)CCKTp0ND%>eADrC1fOt%QT_1GqxA#j4`Q+ulk1^Q%zgj7usca)|
z5#}7L2rmgsu{;cF1+g7zOP;Uc8L?o<;pL
z$DtKq$xPc>TiiW%_{N!Na=0svBd9t@syVXjfuB8>f{v?xs8GBR%jb`GGgfQ@x
z2ypQ|^l5tlQ@sh}vP{wl`WYMRVMa*NpDtYw3UHrR&0@?y7EM*5Ym=)IiacGfAlZN0rCv=n}
z9$vEDurc``ly&csm%|sUDYGC5S41#_KO=p0>!0?Jf(rCFrI8~1E&8}c?AmT0!Bmn4
znnJIAyP{LFg-%704^FjWK30P+JqH8mwl0R+
zIId(-<%YCx$JvBr-1`;a&ApnQIe>|Xx=MB!kLo7tx!nUkPfJ7qvIe*t
z6$zSi;wPH!<)lEGUpp7fq+t|N0VJeGN#f32V8?5uIJ90_Mwg}|()ZWUX2ZFv*9vCM5E3PL(G+kDiwh22r1^`^$JJL6*3@g9D+uDqvRyA1M^b*#KX59X1V+G7
zS6`6mtR(%C?3r0A+FK3n!rW(xni)usU8?M}izKQH{80<+2Y2clh-c{UinQo-j>qd5
zeG@FDy1(-2y~FoAowJk9{q3I~zt7q@J;2y!qfpYAL={6_A0H{yezE-=_}dSRjEt1w
z=bmSelj@h`fcg1ZzFw0d!>%1ZD-?&+Ui0O;Po&Ma-RSS>LKRPYBADQD{cHvGA|c2G
z8yMy5@U!55vH-6@l?PxAU07CT!6kqGxt11I3n&H*
zA9nzC`~BlRfenA#^@-|nl=mtDmXa_&`$>moa5eQU!5HoB{MA5m>)RUrfSG^JDy<-j
z-!s5iSneA;XU9bf4!Ki9;awcW=A8)8X!(g
zZ)zLL(R^2GcJx9FBrERv+F^sqkI?Fr_Z7(!<^dJzp+u}Gf-tXzOHTt%)K<*D-KX<~
zZ|RNCjk*x5lo~sO@|4Pb(9+cnU0t`RTvj+P)!P8f^qPflhm(yQdc
z)W}VJWpSH)4Nd^6%9p*$)~WY9Q$s-(U7%bs*|OF%oTnJZ;sVUL>=1^>MvX?pR~g18T49uw|IV$P0#hyPk0?Lj?_a)aLOMRzF;1UsK%
zjd%K^S$o@_jY-lz5*<-YA;3603rPQtGSfC@%TcZuK+1Up$p16OmmrELSMC
zCTh;FcOkKd{u_6e^B(<=mVI$w|~qjF8l@79_6)g5I7I_yb2ZrHqyOC`sQa;1br@Ioe#
zKvt<1;bdj~@3YF`Mzj=boVbTZ;yjy1@f}%$rn;^36F?DH&Op{9Q|Ix|D`76&?$`Z=
ze6MXBc|QQbBb3+(26!OX&12#q#CQe%<-YM{{Eh@QEQVF;8F=pk0&pT(?$ZK`l_D^a
zDMr%2{vd3b)zH8UTrUb*T3U*w*}z5y>#J#8&WhL?PQRcA+WD_oykxy+CBd$Kg$p|;
zQgNv~vQFjH>~Bd*^b46la>Xssi2LU3g}x+l;TgXUgUVK0rQN|N$s@DJoVCBa98G;a
z>$Ip>!{{%I+>HHy!(sldM^7BQx|yC7(3nSuM&duC2e*A-Z{QHsjxUBqyy)Cjd_FEw
z^6??eeAmx~+X`|wblqV6lt-4z1MNoB@8)YD_Yeo?UHtdfJJeNIe3iQ>L>OU}rM)$`!@hxoOr1`lCioPM{S
zS1flvODNNPd9itk0|AZ`yk;v3BV*(A^(SlEXTCBe^{c!^26aL?E{bt`UrRv>!R6$U
zPTT>r3A~%h=0nk#&^Y_b(t|74^WSjIO9d8k(ekI%=b5hc`3R=HL_Ybe8|wCdr8RjF
z*rft7qmJAeOi2PkIYkw!z&QhGSP;eD>+Bje9>c%J3y4h1s1TqtCPK-0Yj+Eb0BcM%
z_GRNK9i9Shg4y2kVA35kf`Y=rJHBV`Dk)-=tTBY_0pNv|9j+RA8g~P@3s$w)0{~Xt
zb%;8YB4xZwgk|aUVX#V{u&M3>twQYBGKZPDa8+JAb2zez4#3|C@Qw5=Q-iQx&%AVBi8c(m
zYa7|U+3&*rDju%8z-QTB?IjaTedzCG6>%HVkkvAy3ggP-DHA~?NN-tjR832HK7No^
zqVJQLDrdil{Pu4{to)$6W&NA|?Qc5DQP&w@n7HGms<_8r?~$N6FXVXR@c{0
zUVoyjr>aT@%0kL|&+=7p#J2+E>QZ59=?#@yb9Rh$cn|*j8Sx?lYORHgvaac{o&F*J58652hFrZu0j66{?lkolas^6-W^Vm
z3}XWY!AB;Tsfdq%Ti7A4oavIAr9Se*@zP`g+0@VQnu)j!AHpm=NKlsJGdrQ5;~lLG
z;bn{y_eZkCMkUr^sw7{6kjgxrZ*@PQtAoCFZ+myYYRN>fFFFG$u}EoCO*Pcni>#H;
zEavHPknc)P;#Q&YIq!l?eO5(}Z56
zP6fYfFX5-kMAUj~L+htW@IB@u`V9rUhiek)<`gMQD(yJannM4+sOGU6?Zp8`x7nSq
z3*izg*v02{I~LYOVOfv&HW5X$C^ais9J0&K**Udj>iFb49G-9*BRPF5z$YqXtGH!Y
z!y_yJUFfh@T+_OFnAf1KBjmni;=K+xZaZE|kNNigrsb)V&`Y3z8v}*_(+05DEy-J^zsqTHWaSV1Q&-pNLFIt>7X5E`D{T{P?daDrHK^v5J
zI=c@ku^^CEi3Soo{c~I88#t3{DFTqr5_T;g_UX7#M35tC@6q
zl1Pq#n)D>_5#iyCqaq`dC0uw*8g{zh#C-V~e#_x-%FXh0=Sw_jpyyGi8(=7ehx7+D
z)seIJ&2&7}s$q@gPrRNS@-27mQP@F#2EE6GE)O5$=XhEg2vKw61k&fqzg3ZwD`!?4
zeMdGdlBjr|IYqxsGtadnGiHJZd&R)MAC)K#5wLCcvtUM%Yw|s>
z@OjMg6Ku)Vx4XCXtBBAB&Sc=Yy?@4P56a=*>W^EW&42~lX3_<12&~(yTWcF7*u??X
zY?&Mpl(;)D0Wbu&B-H2M$d#@C4jEO!L&`>qv^mgYRihkT_uv16@F~Tg
zNYqL>_8Jpsj?hE2VsDr~0yMlFw4nH!=nEEky_?KH9Im1Vzse+47L^}RtwnfWXJ3DF
zm~|`|$FkQcEmTn4{p>Z8ONG0z^&Y{~$H(|65PRbIt1C7Rz4Bci4@{IIqkyqbzu3!r
zOz52&+rRWjndKP`q=rA!yf*&W6=lE!2Th%4$n-qW@d4f0&$!Xrg9$j+ec`x>AUp;J
z28VvB!%tPf07Y?W|CPB)1P$d6O3h4bG`ykk{Wmj4jP?1|s*{snY$Dzc(?>Vjfg9N>
z6QSci@%I>#cs>Mr92-2m_tek(4-o5@E;g3G!56`>w^doXbHcYsbh8J$>MBN&UAbB+
zTBUn*`%
zPG}PZjLBWy*69x?ohg8a`@<0942uwM{K5~h0SyY(Nf;!Xs^-bY7B4KXt##)6iB%(S
z2Q6Rnl+I4W+{WeHJ-!{Qf2)4Du|~^!nEjc8=Ck>XZ-Wre@xrHb&sg9T`cN_LF?J!@
zAP*L}zYRcym;;P=>#*`B+~@FICJQ_M?@JN3?xdR63o{7iV7euI)T@lAB>qyx`}7(U
zOz0T>w>uGDz;tc4Io$|40|u?<&u7t17a~Sp>ms)eY6U^wgb4U0<}fR(IEn(Batrk2
zW{clHtYZKK5;8F{LH?x)*z|G?5O7x@kDd*pGm-`RnESt3>kF}YGTlIVc`n?n?i~mZ2n*oWv(1??nx?i<IHMfqO>(p1Ru8=Hx>z
zQWS61mCYsMcJzNL!v|i%3SLB*ZXbo)o1{Makvdiw%b_a
zD*dT<2(QVxwaM}&1(N;9
zK6Ta)9;X4CiI@%E&$HS*d(>oYna_ii1)Uy*M!{LYU1{a1N~G{$!?
z+tygQPwdt5)t#N4+1;F_84NI&wBk}o`oC9`__n>hxVT78Vt4D-Ev)>~(tO=4TAToD
zzU-5^^v3$uS)631um=)wO^1P08i5=aUm+XmlF3@;E&<0WrR$1IfwnWRs6Ql|eVnE$
zXwi=uYMGupC&BR|7s2TYY+M!H%Y|ecR%nEqf~(wsh=>O&$k0p=;ebYxU$|Y3WRuOS
z_r2wB2yw+3BQs(0e+uy?ou$G$V$*R~^#f$QnW(&(h)o3jZ$Joe@`JYn>l*H+R}wTE
zLA)Tp1GnogIQ{zr?bCU3^efLNud46Jf4bK?b1S6n=tKz8fljy}y3CMaS+WV<-Z$6T
zT7LpbhaOQk=P7??8TXYA4V9px{0tE!>p*X-3!)NY0=7jY7fNT?o#;C*FELk42qR0!
z(;-Jh%}ZyTP2V4j)k^v9x3S`jTw_0>efBk>pxeeIcs^
z35ff0KUvA!?lHc&Rx{)46*8&mR`qH{;|x1CeGQ{}Z#2KRI=Dq0m!*D~ty}2_qk_C^
zbIC~MC==l}q@c~V98)6pw%PK>R
zruhhZ=#aQJPBzO20T=dNF$vUIbwH-yTxFIwhxo6w)O?0;wGwGcoE%>D!kg((tQ$=g
z^fjFeoW*cK*m!do?;q)haTDKJtciNP?P9^@|5y-(eC{%9L4D)ou+Q^NQ*;CFyw@@z
zDM4r`gmTd2?8U01;6<-6mExn1j~fH*)nsb-Bf69zqisiS{U0IITFS?j=LRnLR|sn)
z_FhABA$}>GF6b2r}YaiewZi@IgB9{tW+}C1v}@eg#$sQN>rH-_7vy!N;?9dEeJLY_!}Bh{
z>0$;%`L$;|J>@`4NPgemZM%yjW)@)K9RPiOXY?npX!rh%QYEV_N;d-Tw0+6)Dwh{V
zk^_k$jjzbcZU`)ti3>OjEU50q^{OsDE(%|N=ln4=g<_ttB(?pC@@I-lAIhny&@&}u
z>5T_5JUJ3CedXTC|1yWnV{P&7f*5DRK-2X)2*3Zf+_H-HNbdQL=~?iZIpp$6hN`>~
zK^ac-u=3~VW0|2iD)~=D&5s3ruevctCUhEx5zMb)+#x!6H6meMy!0(SYmjO4w@>1X
z*oCN@*&%s2xoax1oT!CD&ZAA(8smB
zV>@3ePir(*Q*Dj$#7Tmym!P#*JJ2YDv@_bzc+D#<_aI|D?0$q>ZuIGHXP3T?%G|YD
z5I6hJ>zJA@07xW4*C}KF!@Vy#S
zV0;0AhX3t2jg0_TXUUfhW$UQEhr@TVqa^l8RQm9U_pa^#1m4x=iE22Nv@6IWs0Fpf
z5zM0yo+WW8VGtsJ>vi*)@J
z#}Qz7t&hNgGt#(W1^XZ4nx1`MVT6PznDh9D#B12UJZ*I{e0aI6Rc~}V;KHBrr5z2k
zo&}toSMEn%X0fsYL0vbeY_;|{Fs-g*;)v7wU-=Giuiq$_18D}er~#ZOjgM!CK(53(I6-^`)g1u$&KT4__-5M$~C(8z9s2}7tVYFG!g
zLqp0UaW~V9|I2S7EdG``KnrW6X)3W%{XIjrP7E^)~-B(rZp
z!0(s@X*ZQT#8?xpM>HsDrP;^gIXoaDX&sxBzapXW?*8o`9|x}sDHj&y51s5?eS_in
z2;;(($PNwGfQ#(w$tW4P*lTddptJT6HAVCl1c$X?fdCtBnU4GhO>7uX$z$<{@-M@?
z8X*VGM+PXODS;{%4PuUF6jO5|`wlAfOhf+V?&ZB}3OCtt6@E!jdQfyoju~foy+NuG
zKx=6gpY)}&@4zTOq^(D=vxh5f%}BaI#;%x%`_B)$Jj6U|gm_alD^ecI%KKhtw-`s|
z8^r566h4>Xq|f(ff07jD5ITXfKlN*)$)5{Jl^Z-L&Yc{kJ!_?&P&8sIdSG4jc8qq6
zSNDZB%%spn4>wmzZi&RuZ8@C*h;TAb_d*G#&~HHYZqD9A-W&(uJIt+8w8X}Gf)&07
z)X^_iBV4F{{PZmLYka
zq+7(1u%@%bGlPho1462iOWUJ(9mLS_nZhHRM=O5BRY_=%wX$Aq^X5+kt9^nr#=i!x
zjM^0Tj>!y(6ip#!{O{*QE1IqTDP>DyuSn96B|ZOg`h4iRMHS3IS$2P*)0}13JshmjS
zOGB1NFg3H3t8UBpJMDr6cD;$*{A@rKB~Gk8E=4jcMUVI-jOoa67NcwyhD_7NRx->k
zN>i@zvDdT`X>Ej1Dr)>BfF9k5+*%N(JPrzdRW05u??@TTOIZi)K*!%S{ocNK7&yiS
zttDyZr}O|9qgtMV$QkyFUgs-E_GwP^O%A)CLnUvsb5!~~sAJ!l)z&E>4gj;Epucw&
zAy#ddUCkZ5&;-;v2I3f0DWGl!!^b%Pic|$LA
zI6Alt?^vH#lkIQ#wN!pjLs&!Wsse*HP*bTMKntc
z*$nS$QZT2yG*p&niMm`2`-VH~1YwKlN>})=`k=~C7N_`K4l(b23`91^%Ux;WV1E*`dML1Z>qiKH3kmOpzNw?8qS~l1qV}EO#n%Qf{s;Regj&@0J(hn+&dR?
z6cXC-AN%KO+pMaqvB0HrMbX5IMeK{{dpt#&6CrxUz8wuEf!@eMXxDSfVx
zFS7JAi5_HDkN(C+JNvrfqfyTQcgoSTgW9mTu
zLqKRoc`+qbw+6-cgxh-WqWgnV?s^|NNCNG|@=Oy*+zk3D2%sB~hbU~e#^!gnJ&sb)
z!Oqp#&>Y}fE)$Z(Pxn3HNdpW++^AU7jF^G6H#xZ|a|LG*amPJrvd|;t^ak>DGZOO~
zWiDcwTyp=-+}BuB*XG6d=ynR-CuSJCMyY=PjyFkjCj>=;Bz)5GSa9OzrICDhrO%xy
z^{w5bj}U46EoZ^9b*o0x=c4v>v7GEeA{aXPI9^HywxSMCre=@Xrfi`VZpag;?X_4E
z-Pcd#S33N5KNS&mTy|bOlCQUzWzlgceAP(w><6U4zw0V(312wNQnz9iohxwq=h(_T
zqu6A!v-th|M?}}|2T^&XLa=qQ1Mv2U_hwd|*B7t7DMU=+Vg@=dX32o!GG&P52t*Sw
zfq2pv%J5$W7w4z*SQ9-^WPpy3IxkZ$g#9X!x%oN}VVVJ7_c*FO3EMvaaObiac~e8aGF)f(gO*(>>R7d!hn&%Lpx;t~MV22C0p;IX1h4r$`_7wCP>9-J
zeLC`WDt`XKO5Yft`aZbC%toD57uQJa%?A*VliD`Ce52fFs;j^#L{=OGM;C#+@75^~
zdHjK@Z90wtFjD=%qL=dey`$;n6`n?a3hATkAURW0pF?K~L=}Nt0d1r};I&WIA6D)P
z`<%Q2)ZraW4P9aw)y<#S6e%n^9JfT^rnH7G?S$Dk#t$*>dXXu{R7_hDRgJTfxIKDg
zWb+zf35*Mta@L0Nk3Ae+roWnUmigX~Bt?LUo{8k54^YyY={{kOe5kWr!LoIdt3oAZ|t7}T5W|0C(TqpAM?zY$8ZvPx3O
zs1W*4X6jm5g_~@$3fIcDcSL3=5!v%vWn6niRzeXM*CsNqnZ16``~CgloKEM~Io|j6
ze!ia1$MZ1`-L%EkN!Ln4rTMt>kJ92pYXpyz_%eG!L{7G2?M#%F-jqzUh~7Ixc_WVz
zWHyh9r$|A=3bv5Q(wTPKQj{Asc)}-s|>goc3
z_v$`<{j0V9=cKLrQUy8*1k$$}s^35K9iI%$&)m1E|1R@+Zdzj#wbyc*E5n}m&Z`Xy
zk`wZtg|4hFGaVZB-tAx)1ujvFCPflCN`tXKHIMuc+!)LbgubhAIbV+0PLybRPNh;k
zb(>`MO^3tXF#k}-uYQ{@p~fVa>5iLiK5%L?idD&2{W-GydF!G4Ir2w~z!18a(Z)JN
z<_N5SyAl$f5Zrwi>;J&CAeqcy2b3o`J0mggqc4HWKRLaUmK-%sobs+K6Er0-M=
z{vGSw_|xSC)6}~`HXWVyli}LcUF`Rz=$TaRF1*w7vOLSxdr!{=+-e@r4C1}#|5-nc
zNuN6SSpT6w`8`#GhXR=t`XgG3!D$cv(-
zj?GoiuKX78IXKb(HOOY}d4b_BO#$36|=XFC&ro8kKL)_+73iV{6kl)o>lCm<3Ur
zDiYP8z#>I~A$yNPl=?O9?(Zm*o8zKhcfJ`jzp0lssC5;Lmf3!K-`pZYC63n#F#Xyn
zGigSKhpKUYU%S98?FXQ2MsjsQM7k^yFirNDODxW`2+mghumKKme7D<%GhHc6E~{S*
zbN4H>7Aq2TX2M)~H`9x6(|vOLdt>+K($moyn-Tc6Pp+cFZ04{pIt*Py&ez)rX3#t$
z8Aa^cxFof{Cjq4LGWl_6cXG#*y(WbwRxanLddmNbLqjzYLR@uG1cu_xgajJW%k(jP
zh;L!bZ3&^vGSOmE^g`nrFFAOSp7!rA^^sn0_(k2)EP?pAKJp?Sm4#pM;J|MShSeqw
z;4Md22EU*NquZ%^3RppRShSn?Nqmm;&yC)CkGQiJoQSrWfr(}=fWxx
z0egOuNf|lw?U&c*;s#8Y&y;UNV06?N(}=2q!R!x#A>Lc<*ctx7UWvfNLo*NZ{rKJQ
zVN$)H3B{;Vk+y_OJX{o9&QgxtB0}BfTGJM$8n$@x^SOgJb|*-4YI>f5!_{RYo8Y0hvLkEVxn^La=5Li#CpHH1)CYxri(QL4>`>Ir?Y
zZ`uR-t_)^7{LX`wQ5KI)tN{N|smFo6{(0MWM!T$vmzA?CcM2+xX6(EWgj`a$F-gy_
zSaq~2w*m8xnO5tq-ih*a?BZ+Syvut$)_Y6tWUm)k^FD^{x6X{9D#tpyR98whsW@ET
zE`_|m?cw2(>^ky48}APQa|mlzm@5-CFeEaHTl@TY_=V_@+g&
zt-q84Pp;&w%fWr*gRf+71DgD4i~NSQk_n@X)i0&Y(^OSww+|{E!qp_qUHD^MZQhJt
z``|j9xg_%6cY+aiLR7nIiBPk5(`6oQC&fVqnvBO|k6QEbE1DCB8BZIPjD94Fg}}~p
zVDmWfIk-|9q6d&?BGWOUr@@EwrU5p>^p^51HVZ(?3-v2(He1O-RFn>2LTYcN2a7d`
z*f9?+Ew7PN(OOJ41^ipWMZ*BFzq6_dK4=g1N#--a!CDNaW?5*~vLYm0m!ED=`1#@J
zZY%)l=qqtZfS8P`<9`>#4U!(9o
z4M|Lc`A=K0c&~oLHZ`j+u3ZV#p>*k~B_e_|k
zP)$XZ+0{l2zv$|rz;_?GR#ljzpVl5GX|xU%ms!CJ%*~XG@8#L9U?D3{3Pj;?1s`oL
z;NAEe+5MCaC$f4#9~cm{-HRY@mOFcUSDkA2_dORE7f(~Z@`jxPkUZbDwJE}!_7w;r
z4}l$%{wNd>hce=N!B-9@U%R2tS69@sv$KK36biGP*t2h^+afNsL`ki5fH~(7(5D5i
zZy0;`T7qw(C6rNR{V7)n1xU~0ajQL}o+uj3;W`um9ty#JM@uHao~n&{_n0`D5V`Z7
z#~w*Kcau53hpg$1=~J>G8g`rc?svRAy0vLNl#^kBHn
zd>iOv>H`h}9>a&O4QKC7Ylv
z&U-}pi#pF>Wmroh@(B&9aS$~xhGno?#L5LKUd(?`6yKrnlu)kT3t!C6_M>E$d*$X=
zB}UxSe%oVaB=DT~IsMEcCQ^#Q;`BR+8VDXIYX4x&O$B;sSqlm}JnI1r^3*GN*^|SU
zvWIIAE?&I23!rs^oS!VHtjf#FKR;^?d-n%ZQeqAG4_GEyY@*bR=5hY~d%h4^qsup1c6c8J+Y+4nom{w=ff5*ose2^F0zej%Z`P4(&>bG*~P
z$#2y9V%9xbU?Pz&BB}x`OT1m#IcO;7fx0&b6ejbQ*XLl>@Dbco9zh0AQ;vG#`f+|k
zLkbWT#QcuE;XacEx*L%iOkBr>aJaJ7E|-Qcg66Yk^Mr5p1JDzo%{>+iMx8?ou{?)@G+YxNZWo!|eT`Ten
z=6RW&l$JmB#>VSE=a%I1Olujs*e4FNVgv;#xhqt$!`L>qIC4bRDR)Y#fEYn)|8|Ah
zB#dM`9BVTL1M+0rWqz*q{KXH&EeXrqT){#3!qh*K;VK130dD$!ZbQze+GhEorl7Y;
zQl1-K=e5Q2Dj8K?%hHef{651V{ojJ50lMoZ9*gaeVMZZqS**FcOUI^xTBN^X53$x(t%!^A7LZ7uW;X1`Bq)FXNhR}W+1b~{{g
zYXp_*E~S{&k4aeJhyidZ3=EY^z~{FD=Kclk^2n?&Ut)m%7?656%XbO=`>`Pye$)Zq
zr`Y`NPNai3cRE@Hf63*id<3W5_*RQb0nnM(LgWqus4lx-5$H7&atAYD#(f#bzO_$7
z)<5qIqDI)9zFXtQ&l*T4GcA0(-LKjSr2E-x!qEKCOC(
zfX1vF&0~!S0?Q(yDJNRlT~n}6Xr8gignj>;55}Hv%6dw7d{sXPSq>KJor!cZJ$X$3
zJZzu+slRj#fB(JORI{Wr6b7-z#gW{;?veQWa2!?5zLJ<&dYI*!B};R|4X3rT)5L_k
z8o~D`26!yW^@xy$IPW}E
z(`EGOsfljAJ+Hf-Q%@GcI~;?3Ij9RQuoHmhS^T4*ct!li_7&Ptl&uqx&Y#Sv3l8cVcPB_^r-Z^M2mckYKu
zVbVl8dAon(TYp8N?$>_0ZbCk;Ce}Y90#O^P*p!*GM*?d*J6nPRHn>=>zJ2VIrvKU!
zZt`*U2c%Z@m@)Q6JFdI>FI^D|?O2?YqpBk0#l4KC;u1U}Q$kIh3
z|L0ZbwOTJxX5H(`!AAh+>5i)*m!>^Vj<)f=`V^HPMJ=g`PK@wYzT{|4gnJ_j8d>*o
zKkwmt*TeQ7aoR=IM3OXgg9*@&Qnc1*^>Nbk(ZbiJceZxke~5qRYkVO0J|?UZRT)>L
z+i<~L&EXoC6xoUMrxu>pr_HFiv*#}hgok#xHa)nPvUHPo#&Ay3zmt}0Lw!;&t%Ud&
z$@1YWRf
z;FoWCp
z$*}XpI3u_>cIr-=nzc0WHr9*mnZ5pWC1`e^3jAd}SFXAZxOd6Ravg0_>geq5zn>?)
zTEl{>HMuziv#R3x7_Lc%wBKl!xZ4Lm!>v-nlm*uKr(A!vmG4Agw#T=k)PCsWM17Li
zK2EhEoSD|@LZ6yVB;2%E^>FS(am6CLE+SH#Zs-SSQf9u}GW?7Z)K+{EKDy29lAMSc
z`2xOorucqItVqQj`_!BY%xFW^#R}5`2ImSHnWP&%7t>A85>zWmnY?8+E~gCrAhvXp
zZzf#h?AgWbCavWTF3hL;$wJ=T8@hyR8ZuitFq!05?cmKqIEWN!qxN29RwEaNUfF}2
zz~!r7;(xINsuJGv6WJrhz}Fe0?8VnylhOF~aeG(QIE&sd>dy|>&iZ?&iI$YQ+?&!c
z4@>XrXA=yB0dCpp+|L8!pxFj7h1r|9t9`TLrb@W0YKr^SOC-FT6)dRLHml{N6avK$
zmKbW{iEx4;zq*|cb`!4*soo&JH5Y}^Xj|lEWPiabtav)SyHGUN{n@gsQLCB=m3PHZ
zV+pjCmIw+R$z}G1b#lAs_7>g~>B~md@8=!*z_6(syqcyx{NnBsNui)ExqpS;k>eT)
z>^w%QcJ8j$ZFY2(?a2SS`03w_F}Ij)+l;-uS%rE9KUYD(2ke3sb#;<
zu^n3H@Zcr67t&rKp!eW}j<^nhW1KjDq;8c6X#nuPX9AB1{jQw&Emtp>4t=1Oglf1q
zzSA%->Drcb$BB#hoonZ)Il(iMJ>E43HYViZYGT;hqI^Ogjg8*hK|lDaWCeI1JO>_K
zHHxgQ%>V4OcaK^Qxv9?m;FK!#UrYVI9a0fkB+_P4lJsR_s=AQgT^aG6@v;E&rD4#V>LN|BzFVLRI_&Oy(Dr{z*L-a&hOO~q>*+
zvi6f@QYqTb%O-quZ+jQ62>yq*?HQ={$CAApC<7|i(-pTcLkAOcrnfpkC8WvEv?2Yv
zQ~Ct@Y_L;*>8XyNwJolGz3Q{-CDW531LI;fR6~KAH(faH-e6*3*?=+Ld4Hkae46N-
z=1{l(dj4qhTl?cz*sDM51Qx}o(Dq_x@u=6Vb}s~nT9uV&ocZiZm1k~m-5jjP3aH9Y+Y$SRC
zeSs-K3uH*4k92ep5F|;(>3s|4G&Z*h9|@m>5rxoo2$P%xA|_D=0dI*Q(MusnKXP@*
znm*|MF4ylfQfg;VjXkLNN(Bv-ui}0FQ#l=fV&~bpt_2@?`>383pSbAJhMGUQ2t!$0
z{U&0TA9poY1wl(5_3~4TMb?$I8}RPWF~<1#Q~7)2^wb8EXr#FSqk#85=ffuX2}{f~
zoV+LXZ~)~4h1vOY4_t29EnFGRP8*&o(5OTtnBjf#%3^l-i9P_bQuNtubN!7*ZoA!e
zoo)YV_U-h`Ao1R;?3XdvO`6%yn0Q_+K`fM3p!H(sK~67
zonAv_bDDm`=loNi3he8>VL*GBD(k~0=PNza9^JlLi7Pf{W@aY#7V97u5OWxp#yM#%
zLOhfeDC=lAWk>susc6}{{Ws+8>|SIK>_ZRDC{7kr-KEZ%4OH^TlA6LTO5
zk5VO^;BV{CSxd{P9`vL$olM}zYROLtEw~8F%%zIW@#5j)L+g=yzuv!{EMP}{TdcGy**kT|{yMbg$@4zyu0&P#9uXXAS568b
zr*ei?!4YyFU&3Q41@9DtU*g2TQDWyU&3d~`0S40hMyC}&=s5nHbo%|@pljpzBA&*#
z7mZ%CqOJ{Q5atM7-&RY~ZcS(&C^B=*40~3>@2xg>QE?ZcD6}gybzZE^WsuB}04

       {children}
@@ -181,6 +211,17 @@ export function AppSidebar() {
   useEffect(() => { if (isChatRoute) setChatOpen(true); }, [isChatRoute]);
   useEffect(() => { if (isStudioRoute) setRunsOpen(true); }, [isStudioRoute]);
 
+  const scrollRef = useRef(null);
+  const [scrolled, setScrolled] = useState(false);
+  useEffect(() => {
+    const el = scrollRef.current;
+    if (!el) return;
+    const handler = () => setScrolled(el.scrollTop > 0);
+    handler();
+    el.addEventListener("scroll", handler, { passive: true });
+    return () => el.removeEventListener("scroll", handler);
+  }, []);
+
   const isRecipesRoute = pathname.startsWith("/data-recipes");
   const { displayTitle, avatarDataUrl } = useEffectiveProfile();
 
@@ -195,7 +236,7 @@ export function AppSidebar() {
     : undefined;
 
   // Training runs
-  const { items: runItems, refresh: refreshRuns } = useTrainingHistorySidebarItems(
+  const { items: runItems } = useTrainingHistorySidebarItems(
     !chatOnly && isStudioRoute,
   );
   const activeJobId = useTrainingRuntimeStore((s) => s.jobId);
@@ -213,6 +254,93 @@ export function AppSidebar() {
     });
   }
 
+  type RenameTarget =
+    | { kind: "chat"; item: SidebarItem; current: string }
+    | { kind: "run"; run: TrainingRunSummary; current: string };
+  const [renamingTarget, setRenamingTarget] = useState(
+    null,
+  );
+  const [renameDraft, setRenameDraft] = useState("");
+  const renameTrimmed = renameDraft.trim();
+  const nextRunDisplayName = renameTrimmed.length > 0 ? renameTrimmed : null;
+  const renameDirty =
+    renamingTarget !== null &&
+    (renamingTarget.kind === "chat"
+      ? renameTrimmed.length > 0 && renameTrimmed !== renamingTarget.current
+      : renameTrimmed.length > 0
+        ? renameTrimmed !== renamingTarget.current
+        : renamingTarget.run.display_name != null);
+
+  function openRenameChat(item: SidebarItem) {
+    setRenameDraft(item.title);
+    setRenamingTarget({ kind: "chat", item, current: item.title });
+  }
+  function openRenameRun(run: TrainingRunSummary) {
+    const current = run.display_name ?? run.model_name;
+    setRenameDraft(current);
+    setRenamingTarget({ kind: "run", run, current });
+  }
+  async function commitRename() {
+    const target = renamingTarget;
+    if (!target || !renameDirty) return;
+    setRenamingTarget(null);
+    if (target.kind === "chat") {
+      try {
+        await renameChatItem(target.item, renameTrimmed);
+      } catch (err) {
+        toast.error("Failed to rename chat", {
+          description: err instanceof Error ? err.message : undefined,
+        });
+      }
+      return;
+    }
+    try {
+      const updated = await renameTrainingRun(target.run.id, nextRunDisplayName);
+      emitTrainingRunUpdated(updated);
+    } catch (err) {
+      toast.error("Failed to rename run", {
+        description: err instanceof Error ? err.message : undefined,
+      });
+    }
+  }
+
+  type DeleteTarget =
+    | { kind: "chat"; item: SidebarItem }
+    | { kind: "run"; run: TrainingRunSummary };
+  const [confirmingDelete, setConfirmingDelete] =
+    useState(null);
+
+  async function commitDelete() {
+    const target = confirmingDelete;
+    if (!target) return;
+    setConfirmingDelete(null);
+    if (target.kind === "chat") {
+      try {
+        await handleDeleteThread(target.item);
+      } catch (err) {
+        toast.error("Failed to delete chat", {
+          description: err instanceof Error ? err.message : undefined,
+        });
+      }
+      return;
+    }
+    if (target.run.status === "running") {
+      toast.error("Cannot delete a running training run");
+      return;
+    }
+    try {
+      await deleteTrainingRun(target.run.id);
+      if (selectedHistoryRunId === target.run.id) {
+        setSelectedHistoryRunId(null);
+      }
+      emitTrainingRunDeleted(target.run.id);
+    } catch (err) {
+      toast.error("Failed to delete run", {
+        description: err instanceof Error ? err.message : undefined,
+      });
+    }
+  }
+
   return (
     <>
     
-      
+      
         {/* Expanded: compact logo + close toggle */}
         

           
               unsloth
             
-            
+            
               BETA
             
           
@@ -259,13 +384,17 @@ export function AppSidebar() {
                 
               
-              
+              
                 Close sidebar
               
             
@@ -274,19 +403,23 @@ export function AppSidebar() {
 
         {/* Collapsed: panel icon doubles as expand trigger */}
         {!isMobile && (
-          

+          

             
               
                 
               
-              
+              
                 Open sidebar
               
             
@@ -294,7 +427,7 @@ export function AppSidebar() {
         )}
       
 
-      
+      
         
           
             
       
 
-      
-        {/* Navigate (no header) */}
-        
-          
-            
-               {
-                  if (chatOnly) return;
-                  navigate({ to: "/studio" });
-                  closeMobileIfOpen();
-                }}
-              />
+      
+        
+          
+             {
+                if (chatOnly) return;
+                navigate({ to: "/studio" });
+                closeMobileIfOpen();
+              }}
+            />
 
-               {
-                  navigate({ to: "/data-recipes" });
-                  closeMobileIfOpen();
-                }}
-              />
+             {
+                navigate({ to: "/data-recipes" });
+                closeMobileIfOpen();
+              }}
+            />
 
-               {
-                  if (chatOnly) return;
-                  navigate({ to: "/export" });
-                  closeMobileIfOpen();
-                }}
-              />
-            
-          
-        
+             {
+                if (chatOnly) return;
+                navigate({ to: "/export" });
+                closeMobileIfOpen();
+              }}
+            />
+          
+        
+      
 
+      
         {/* Recent Chats — hide on Studio only (Eyera fac13); chatOpen = ec695 clickability */}
         {!isStudioRoute && chatItems.length > 0 && (
           
-          
-            
+          
+            
               
                 Recents
                 
               
             
             
-            
+            
               
                 {chatItems.map((item) => (
                   
                      {
                         navigate({
                           to: "/chat",
@@ -410,17 +542,38 @@ export function AppSidebar() {
                     >
                       {item.title}
                     
-                    
+                    
+                      
+                        
+                      
+                      
+                         openRenameChat(item)}>
+                          
+                          Rename
+                        
+                         setConfirmingDelete({ kind: "chat", item })}
+                        >
+                          
+                          Delete
+                        
+                      
+                    
                   
                 ))}
               
@@ -433,15 +586,15 @@ export function AppSidebar() {
         {/* Recent Runs */}
         {isStudioRoute && runItems.length > 0 && !chatOnly && (
           
-          
-            
+          
+            
               
                 Recents
                 
               
             
             
-              
+              
                 
                   {runItems.map((run) => {
                     const isActiveRun =
@@ -453,7 +606,7 @@ export function AppSidebar() {
                       >
                          {
                             setSelectedHistoryRunId(run.id);
                             closeMobileIfOpen();
@@ -468,7 +621,7 @@ export function AppSidebar() {
                               aria-hidden
                             />
                             
-                              {run.model_name}
+                              {run.display_name ?? run.model_name}
                             
                             
                               {formatRelativeShort(run.started_at)}
@@ -478,25 +631,41 @@ export function AppSidebar() {
                             {run.dataset_name}
                           
                         
-                        
+                          
+                          
+                             openRenameRun(run)}>
+                              
+                              Rename
+                            
+                            
+                                setConfirmingDelete({ kind: "run", run })
                               }
-                              await refreshRuns();
-                            } catch {
-                              // ignore — next refresh will reconcile
-                            }
-                          }}
-                          title="Delete"
-                          className="absolute right-1 top-1/2 -translate-y-1/2 flex size-5 scale-90 items-center justify-center rounded-[10px] text-sidebar-foreground/55 opacity-0 transition-all duration-150 hover:bg-destructive/12 hover:text-destructive group-hover/run-item:scale-100 group-hover/run-item:opacity-100"
-                        >
-                          
-                        
+                            >
+                              
+                              Delete
+                            
+                          
+                        
                       
                     );
                   })}
@@ -516,7 +685,7 @@ export function AppSidebar() {
                 
                   

                     
                   

                   

-                    {displayTitle}
-                    Unsloth
+                    {displayTitle}
+                    Unsloth
                   

                   
                 
@@ -536,13 +705,13 @@ export function AppSidebar() {
               
                 
                    useSettingsDialogStore.getState().openDialog()}
                   >
-                    
+                    
                     Settings
                     ⌘,
                   
@@ -559,7 +728,7 @@ export function AppSidebar() {
                     ref={anchorRef as React.Ref}
                     onSelect={(e) => { e.preventDefault(); toggleTheme(); }}
                   >
-                    {isDark ?  : }
+                    {isDark ?  : }
                     {isDark ? "Light Mode" : "Dark Mode"}
                   
                   
-                    
+                    
                     Guided Tour
                   
                 
@@ -582,11 +751,11 @@ export function AppSidebar() {
                  useSettingsDialogStore.getState().openDialog("about")}
                 >
-                  
+                  
                   Help
                 
                  setShutdownOpen(true)}>
-                  
+                  
                   Shutdown
                 
               
@@ -601,6 +770,96 @@ export function AppSidebar() {
       onOpenChange={setShutdownOpen}
       onAfterShutdown={removeTrainingUnloadGuard}
     />
+     {
+        if (!open) setConfirmingDelete(null);
+      }}
+    >
+      
+        
+          
+            {confirmingDelete?.kind === "run"
+              ? "Delete training run"
+              : "Delete chat"}
+          
+          
+            {confirmingDelete?.kind === "run" ? (
+              <>
+                Are you sure you want to delete this run{" "}
+                {confirmingDelete.run.display_name ?? confirmingDelete.run.model_name}?
+              
+            ) : confirmingDelete?.kind === "chat" ? (
+              <>
+                Are you sure you want to delete this chat{" "}
+                {confirmingDelete.item.title}?
+              
+            ) : null}
+          
+        
+        
+          
+          
+        
+      
+    
+     {
+        if (!open) setRenamingTarget(null);
+      }}
+    >
+      
+        
+          
+            {renamingTarget?.kind === "run" ? "Rename run" : "Rename chat"}
+          
+        
+         setRenameDraft(event.target.value)}
+          onKeyDown={(event) => {
+            if (event.key === "Enter") {
+              event.preventDefault();
+              void commitRename();
+            }
+          }}
+          autoFocus
+          maxLength={120}
+          placeholder={renamingTarget?.kind === "run" ? "Run name" : "Chat title"}
+          aria-label={renamingTarget?.kind === "run" ? "Run name" : "Chat title"}
+          className="focus-visible:border-input focus-visible:ring-0"
+        />
+        
+          
+          
+        
+      
+    
     
   );
 }
diff --git a/studio/frontend/src/components/assistant-ui/attachment.tsx b/studio/frontend/src/components/assistant-ui/attachment.tsx
index 074dba5320..b5b2810008 100644
--- a/studio/frontend/src/components/assistant-ui/attachment.tsx
+++ b/studio/frontend/src/components/assistant-ui/attachment.tsx
@@ -184,7 +184,7 @@ const AttachmentUI: FC = () => {
         
         {isComposer && }
       
-      
+      
         
       
     
diff --git a/studio/frontend/src/components/assistant-ui/markdown-text.tsx b/studio/frontend/src/components/assistant-ui/markdown-text.tsx
index 7eb4b21ba7..d2c6208fda 100644
--- a/studio/frontend/src/components/assistant-ui/markdown-text.tsx
+++ b/studio/frontend/src/components/assistant-ui/markdown-text.tsx
@@ -7,7 +7,7 @@ import { copyToClipboard } from "@/lib/copy-to-clipboard";
 import { preprocessLaTeX } from "@/lib/latex";
 import { openLink } from "@/lib/open-link";
 import { INTERNAL, useMessagePartText } from "@assistant-ui/react";
-import { Copy02Icon, Tick02Icon } from "@hugeicons/core-free-icons";
+import { Copy01Icon, Tick02Icon } from "@hugeicons/core-free-icons";
 import { HugeiconsIcon } from "@hugeicons/react";
 import { createCodePlugin } from "./code-plugin";
 import { createMathPlugin } from "@streamdown/math";
@@ -50,9 +50,9 @@ const COPY_RESET_MS = 2000;
 const MERMAID_SOURCE_RE = /```mermaid\s*([\s\S]*?)```/i;
 const CODE_FENCE_RE = /^```([^\r\n`]*)\r?\n([\s\S]*?)\r?\n?```$/;
 const ACTION_PANEL_CLASS =
-  "pointer-events-auto flex shrink-0 items-center gap-2 rounded-md border border-sidebar bg-sidebar/80 px-1.5 py-1 supports-[backdrop-filter]:bg-sidebar/70 supports-[backdrop-filter]:backdrop-blur dark:border-white/10 dark:bg-code-block dark:supports-[backdrop-filter]:bg-code-block";
+  "pointer-events-auto flex shrink-0 items-center gap-1";
 const ACTION_BUTTON_CLASS =
-  "cursor-pointer p-1 text-muted-foreground transition-all hover:text-foreground disabled:cursor-not-allowed disabled:opacity-50";
+  "flex size-8 cursor-pointer items-center justify-center rounded-[10px] text-chat-icon-fg transition-all hover:bg-chat-icon-bg-hover hover:text-chat-icon-fg-hover disabled:cursor-not-allowed disabled:opacity-50";
 
 type CodeFence = {
   language: string | null;
@@ -289,8 +289,9 @@ function MermaidCopyButton({ source }: { source: string }) {
       }}
     >
       
     
   );
@@ -308,7 +309,7 @@ function CodeBlockActions({
   const { copied, showCopied } = useCopiedState();
 
   return (
-    

+    

       

         
         
       

     

diff --git a/studio/frontend/src/components/assistant-ui/message-timing.tsx b/studio/frontend/src/components/assistant-ui/message-timing.tsx
index df233812b4..5ad1bdabed 100644
--- a/studio/frontend/src/components/assistant-ui/message-timing.tsx
+++ b/studio/frontend/src/components/assistant-ui/message-timing.tsx
@@ -51,7 +51,7 @@ export const MessageTiming: FC<{
           data-slot="message-timing-trigger"
           aria-label="Message timing"
           className={cn(
-            "flex items-center rounded-md p-1 font-mono text-muted-foreground text-xs tabular-nums transition-colors hover:bg-accent hover:text-accent-foreground",
+            "flex items-center rounded-[10px] p-1 font-mono text-chat-icon-fg text-[13px] tabular-nums transition-colors hover:bg-chat-icon-bg-hover hover:text-chat-icon-fg-hover",
             className,
           )}
         >
@@ -62,7 +62,8 @@ export const MessageTiming: FC<{
         side={side}
         sideOffset={8}
         data-slot="message-timing-popover"
-        className="[&_span>svg]:hidden! rounded-lg border bg-popover px-3 py-2 text-popover-foreground shadow-md"
+        variant="rich"
+        className="[&_span>svg]:hidden!"
       >
         

           {st ? (
diff --git a/studio/frontend/src/components/assistant-ui/model-selector.tsx b/studio/frontend/src/components/assistant-ui/model-selector.tsx
index 795bcb6d08..22bd7412ab 100644
--- a/studio/frontend/src/components/assistant-ui/model-selector.tsx
+++ b/studio/frontend/src/components/assistant-ui/model-selector.tsx
@@ -78,9 +78,9 @@ function ModelSelectorTrigger({
         className={cn(
           "flex min-w-0 items-center gap-2 transition-colors",
           variant === "outline" &&
-          "rounded-[8px] border border-border/60 hover:bg-[#ececec] dark:hover:bg-[#2e3035]",
-          variant === "ghost" && "rounded-[8px] hover:bg-[#ececec] dark:hover:bg-[#2e3035]",
-          variant === "muted" && "rounded-[8px] bg-muted hover:bg-muted/80",
+          "rounded-[10px] border border-border/60 hover:bg-[#ececec] dark:hover:bg-[#2d2e32]",
+          variant === "ghost" && "rounded-[10px] hover:bg-[#ececec] dark:hover:bg-[#2d2e32]",
+          variant === "muted" && "rounded-[10px] bg-muted hover:bg-muted/80",
           size === "sm" && "h-8 px-3 text-xs",
           size === "default" && "h-9 px-3.5 text-sm",
           size === "lg" && "h-10 px-4 text-sm",
@@ -145,7 +145,7 @@ function ModelSelectorContent({
       align="start"
       data-tour={dataTour}
       className={cn(
-        "w-[min(440px,calc(100vw-1rem))] max-w-[calc(100vw-1rem)] min-w-0 gap-0 p-2",
+        "menu-soft-surface ring-0 w-[min(440px,calc(100vw-1rem))] max-w-[calc(100vw-1rem)] min-w-0 gap-0 p-2",
         className,
       )}
     >
diff --git a/studio/frontend/src/components/assistant-ui/model-selector/pickers.tsx b/studio/frontend/src/components/assistant-ui/model-selector/pickers.tsx
index fae3c22caf..a0f97967ef 100644
--- a/studio/frontend/src/components/assistant-ui/model-selector/pickers.tsx
+++ b/studio/frontend/src/components/assistant-ui/model-selector/pickers.tsx
@@ -175,7 +175,10 @@ function ModelRow({
     return (
       
         {content}
-        
+        
           {label}
           {vramTooltipText}
         
@@ -187,7 +190,10 @@ function ModelRow({
     return (
       
         {content}
-        
+        
           {tooltipText}
         
       
diff --git a/studio/frontend/src/components/assistant-ui/sources.tsx b/studio/frontend/src/components/assistant-ui/sources.tsx
index 81c8b0c213..3a55c3fa78 100644
--- a/studio/frontend/src/components/assistant-ui/sources.tsx
+++ b/studio/frontend/src/components/assistant-ui/sources.tsx
@@ -104,7 +104,7 @@ function Source({
       variant={variant}
       size={size}
       className={cn(
-        "cursor-pointer outline-none focus-visible:border-ring focus-visible:ring-[3px] focus-visible:ring-ring/50",
+        "rounded-full cursor-pointer outline-none hover:bg-chat-icon-bg-hover! hover:text-chat-icon-fg-hover! focus-visible:border-ring focus-visible:ring-[3px] focus-visible:ring-ring/50",
         className,
       )}
     >
@@ -137,7 +137,7 @@ const SourceBadge: FC<{ source: SourceData }> = ({ source }) => {
   const displayTitle = source.title || domain;
 
   return (
-    
+    
       
         
           
@@ -146,16 +146,21 @@ const SourceBadge: FC<{ source: SourceData }> = ({ source }) => {
           
         
       
-      
+      
         

           
           

             

               {source.title || domain}
             

-            
{domain}

+            
{domain}

             {source.description && (
-              

+              

                 {source.description}
               

             )}
@@ -245,7 +250,7 @@ const SourcesGroup: FC = () => {
   const hiddenCount = sources.length - (visibleCount ?? sources.length);
 
   return (
-    

+    

       {/* Hidden measurement container — renders all badges to measure row positions */}
       
 {
             onClick={() => setExpanded(true)}
             className={cn(
               badgeVariants({ variant: "outline", size: "default" }),
-              "cursor-pointer text-muted-foreground hover:text-foreground",
+              "rounded-full cursor-pointer text-muted-foreground hover:bg-chat-icon-bg-hover! hover:text-chat-icon-fg-hover!",
             )}
           >
             +{hiddenCount} more
@@ -285,7 +290,7 @@ const SourcesGroup: FC = () => {
             onClick={() => setExpanded(false)}
             className={cn(
               badgeVariants({ variant: "outline", size: "default" }),
-              "cursor-pointer text-muted-foreground hover:text-foreground",
+              "rounded-full cursor-pointer text-muted-foreground hover:bg-chat-icon-bg-hover! hover:text-chat-icon-fg-hover!",
             )}
           >
             Show less
diff --git a/studio/frontend/src/components/assistant-ui/thread.tsx b/studio/frontend/src/components/assistant-ui/thread.tsx
index 0d6cd3bbf9..31a1fb21e0 100644
--- a/studio/frontend/src/components/assistant-ui/thread.tsx
+++ b/studio/frontend/src/components/assistant-ui/thread.tsx
@@ -54,10 +54,8 @@ import {
 import {
   ArrowDownIcon,
   ArrowUpIcon,
-  CheckIcon,
   ChevronLeftIcon,
   ChevronRightIcon,
-  CopyIcon,
   DownloadIcon,
   GlobeIcon,
   HeadphonesIcon,
@@ -66,14 +64,13 @@ import {
   LoaderIcon,
   MicIcon,
   MoreHorizontalIcon,
-  PencilIcon,
   RefreshCwIcon,
   SquareIcon,
   TerminalIcon,
-  Trash2Icon,
   XIcon,
 } from "lucide-react";
-import { motion } from "motion/react";
+import { Copy01Icon, Delete02Icon, Edit03Icon, Tick02Icon } from "@hugeicons/core-free-icons";
+import { HugeiconsIcon } from "@hugeicons/react";
 import {
   type FC,
   type FormEvent,
@@ -108,9 +105,9 @@ export const Thread: FC<{
     
       
@@ -121,7 +118,7 @@ export const Thread: FC<{
           scrollToBottomOnInitialize={false}
           scrollToBottomOnThreadSwitch={false}
           className={cn(
-            "aui-thread-viewport relative flex min-h-0 min-w-0 flex-1 basis-0 flex-col overflow-x-auto overflow-y-auto scroll-smooth px-5",
+            "aui-thread-viewport aui-stream-viewport relative flex min-h-0 min-w-0 flex-1 basis-0 flex-col overflow-x-auto overflow-y-auto scroll-smooth px-5",
             hideComposer ? "pt-4" : "pt-[48px]",
           )}
         >
@@ -164,7 +161,7 @@ export const Thread: FC<{
 
         {!hideComposer && (
            hideWelcome || !thread.isEmpty}>
-            

+            

               

                   
                 

-                

-                  LLMs can make mistakes. Double-check all responses.
+                

+                  LLMs can make mistakes. Double-check responses.
                 

               

             

@@ -204,7 +201,7 @@ const ThreadScrollToBottom: FC = () => {
         isAtBottom && "invisible pointer-events-none",
       )}
     >
-      
+      
     
   );
 };
@@ -253,14 +250,9 @@ const GeneratingSpinner: FC = () => {
 const ComposerAnimated: FC<{ disabled?: boolean }> = ({ disabled }) => {
   return (
     

-      
+      

         
-      
+      

     

   );
 };
@@ -306,7 +298,7 @@ const Composer: FC<{ disabled?: boolean }> = ({ disabled }) => {
       
        = ({ disabled }) => {
       {isTauri ? (
         // Phase 1 native model drops own Tauri local-path drops. Restore browser
         // attachment drops in Tauri when Phase 1d adds attachment-token bridging.
-        

+        

           {composerContent}
         

       ) : (
-        
+        
           {composerContent}
         
       )}
@@ -455,14 +447,8 @@ const ReasoningToggle: FC = () => {
         setReasoningEnabled(next);
         applyQwenThinkingParams(next);
       }}
-      className={cn(
-        "flex items-center gap-1.5 rounded-full px-2.5 py-1 text-xs font-medium transition-colors",
-        disabled
-          ? "cursor-not-allowed opacity-40"
-          : reasoningEnabled
-            ? "bg-primary/10 text-primary hover:bg-primary/20"
-            : "bg-muted text-muted-foreground hover:bg-muted-foreground/15",
-      )}
+      className="composer-pill-btn"
+      data-active={reasoningEnabled && !disabled ? "true" : "false"}
       aria-label={reasoningEnabled ? "Disable thinking" : "Enable thinking"}
     >
       {reasoningEnabled && !disabled ? (
@@ -527,14 +513,8 @@ const WebSearchToggle: FC = () => {
       type="button"
       disabled={disabled}
       onClick={() => setToolsEnabled(!toolsEnabled)}
-      className={cn(
-        "flex items-center gap-1.5 rounded-full px-2.5 py-1 text-xs font-medium transition-colors",
-        disabled
-          ? "cursor-not-allowed opacity-40"
-          : toolsEnabled
-            ? "bg-primary/10 text-primary hover:bg-primary/20"
-            : "bg-muted text-muted-foreground hover:bg-muted-foreground/15",
-      )}
+      className="composer-pill-btn"
+      data-active={toolsEnabled && !disabled ? "true" : "false"}
       aria-label={toolsEnabled ? "Disable web search" : "Enable web search"}
     >
       
@@ -557,14 +537,8 @@ const CodeToolsToggle: FC = () => {
       type="button"
       disabled={disabled}
       onClick={() => setCodeToolsEnabled(!codeToolsEnabled)}
-      className={cn(
-        "flex items-center gap-1.5 rounded-full px-2.5 py-1 text-xs font-medium transition-colors",
-        disabled
-          ? "cursor-not-allowed opacity-40"
-          : codeToolsEnabled
-            ? "bg-primary/10 text-primary hover:bg-primary/20"
-            : "bg-muted text-muted-foreground hover:bg-muted-foreground/15",
-      )}
+      className="composer-pill-btn"
+      data-active={codeToolsEnabled && !disabled ? "true" : "false"}
       aria-label={
         codeToolsEnabled ? "Disable code execution" : "Enable code execution"
       }
@@ -635,7 +609,7 @@ const ToolStatusDisplay: FC = () => {
 
 const ComposerAction: FC<{ disabled?: boolean }> = ({ disabled }) => {
   return (
-    

+    

       

         
         
@@ -725,10 +699,10 @@ const GeneratingIndicator: FC = () => {
 const AssistantMessage: FC = () => {
   return (
     
-      

+      

         
          {
         
       

 
-      

-        
+      

+        
         
       

     
@@ -789,9 +763,13 @@ const DeleteMessageButton: FC = () => {
       tooltip="Delete message"
       disabled={isRunning}
       onClick={handleDelete}
-      className="text-muted-foreground hover:text-destructive"
+      className="text-chat-icon-fg hover:text-destructive"
     >
-      
+      
     
   );
 };
@@ -817,7 +795,11 @@ const CopyButton: FC = () => {
 
   return (
     
-      {copied ?  : }
+      
     
   );
 };
@@ -826,40 +808,39 @@ const AssistantActionBar: FC = () => {
   return (
     
       
       
         
-          
+          
         
       
       
-      
       
         
           
-            
+            
           
         
          e.preventDefault()}
           className="aui-action-bar-more-content z-50 min-w-32 overflow-hidden rounded-md border bg-popover p-1 text-popover-foreground shadow-md"
         >
           
             
-              
+              
               Export as Markdown
             
           
         
       
+      
     
   );
 };
@@ -884,22 +865,21 @@ const UserMessageAudio: FC = () => {
 const UserMessage: FC = () => {
   return (
     
       
       
 
       

-        

+        

           
         

-        

+        

           
+          
         

       

-
-      
     
   );
 };
@@ -908,12 +888,12 @@ const UserActionBar: FC = () => {
   return (
     
       
       
         
-          
+          
         
       
       
@@ -981,23 +961,31 @@ const BranchPicker: FC = ({
     
       
-        
-          
-        
+        
       
-      
-         / 
+      
+        /
       
       
-        
-          
-        
+        
       
     
   );
diff --git a/studio/frontend/src/components/assistant-ui/tooltip-icon-button.tsx b/studio/frontend/src/components/assistant-ui/tooltip-icon-button.tsx
index e498999068..4d72285101 100644
--- a/studio/frontend/src/components/assistant-ui/tooltip-icon-button.tsx
+++ b/studio/frontend/src/components/assistant-ui/tooltip-icon-button.tsx
@@ -37,7 +37,9 @@ export const TooltipIconButton = forwardRef<
           {tooltip}
         
       
-      {tooltip}
+      
+        {tooltip}
+      
     
   );
 });
diff --git a/studio/frontend/src/components/assistant-ui/use-intent-aware-autoscroll.tsx b/studio/frontend/src/components/assistant-ui/use-intent-aware-autoscroll.tsx
index ad01afcdaf..d5927f77c1 100644
--- a/studio/frontend/src/components/assistant-ui/use-intent-aware-autoscroll.tsx
+++ b/studio/frontend/src/components/assistant-ui/use-intent-aware-autoscroll.tsx
@@ -60,6 +60,16 @@ const UPWARD_DETACH_THRESHOLD_PX = 2;
 // keeps the viewport pinned as long as content keeps arriving; settles
 // this long after the last change.
 const FOLLOW_SETTLE_MS = 600;
+// Maximum stabilizer compensation. The stabilizer is meant to absorb
+// sub-frame transients (~5-15px shiki re-renders, ~8px action-bar
+// reservation drift). Anything larger is almost certainly an intentional
+// content removal — message delete, regenerate's old-content clear,
+// reasoning-panel collapse — and should *not* be silently padded over,
+// which would leave persistent empty space below the last message.
+// Above this threshold we release the stabilizer immediately and let
+// the autoscroll re-pin to the new content height, which is the natural
+// behavior the user expects for those actions.
+const STABILIZER_MAX_PX = 64;
 
 export type ScrollToBottom = (behavior?: ScrollBehavior) => void;
 
@@ -202,6 +212,21 @@ export function useIntentAwareAutoScroll(): {
         return false;
       };
 
+      // Stabilizer state — see `stabilize` below for the full
+      // explanation. Lives in this closure so it resets naturally
+      // whenever the viewport remounts (Compare-pane swap, thread
+      // switch with remount, etc.).
+      let stabilizerPx = 0;
+      let maxContentHeight = 0;
+
+      const releaseStabilizer = (): void => {
+        if (stabilizerPx === 0) {
+          return;
+        }
+        stabilizerPx = 0;
+        el.style.removeProperty("--aui-scroll-stabilizer");
+      };
+
       const extendFollow = (): void => {
         if (userDetachedRef.current) {
           return;
@@ -212,6 +237,13 @@ export function useIntentAwareAutoScroll(): {
       const detach = (): void => {
         userDetachedRef.current = true;
         followUntilRef.current = 0;
+        // The stabilizer is only meaningful while we're actively
+        // pinning to the bottom. Once the user scrolls up, drop any
+        // residual padding so the bottom stays flush whenever they
+        // come back. Safe here because the user is mid-content —
+        // shrinking scrollHeight cannot cap their scrollTop.
+        releaseStabilizer();
+        maxContentHeight = el.scrollHeight;
       };
 
       const requestTick = (): void => {
@@ -334,21 +366,116 @@ export function useIntentAwareAutoScroll(): {
         requestTick();
       };
 
-      const resizeObserver = new ResizeObserver(() => {
-        extendFollow();
-        requestTick();
-      });
+      // Scroll stabilizer.
+      //
+      // Problem: when a trailing code block finalizes at stream end
+      // (Streamdown flips `isAnimating` → false, shiki re-renders the
+      // 
 with highlight spans), the block's rendered height
+      // briefly dips and then recovers a frame later. That dip shrinks
+      // `scrollHeight`, which the browser handles by *synchronously*
+      // capping `scrollTop` to the new (smaller) `scrollHeight −
+      // clientHeight`. The cap is visible as a one-frame upward jump;
+      // the recovery a frame or two later is the "snap back" the user
+      // perceives as a flicker. No amount of programmatic re-scrolling
+      // can prevent this — once `scrollHeight` drops, the cap has
+      // already happened and `scrollTop` cannot be pushed past the new
+      // max.
+      //
+      // Fix: keep `scrollHeight` monotonic across the follow window.
+      // We track the maximum *content* height (scrollHeight minus our
+      // own padding contribution) seen during follow, and compensate
+      // for any shortfall by writing the deficit into a CSS custom
+      // property `--aui-scroll-stabilizer`, which the viewport's
+      // `padding-bottom` reads. A 5px content shrink instantly grows
+      // the padding by 5px, so the browser sees no scrollHeight change
+      // and never caps scrollTop. As content naturally grows past its
+      // prior high-water mark (e.g. the next message streams in), the
+      // padding shrinks back toward zero.
+      //
+      // Self-contained: lives entirely on the viewport element via a
+      // CSS variable. Doesn't touch the composer, the action bar, the
+      // message footer, the spacer, or any other UI.
+      //
+      // Returns the post-adjustment scrollHeight so a single layout
+      // read per observer callback can feed both stabilization and
+      // pinning, avoiding a redundant flush.
+      const stabilize = (): number => {
+        const sh = el.scrollHeight;
+        const currentContent = sh - stabilizerPx;
+        const followActive =
+          !userDetachedRef.current &&
+          performance.now() < followUntilRef.current;
+        if (!followActive) {
+          // Outside the follow window we stop adjusting, but we keep
+          // `maxContentHeight` aligned with reality so the next follow
+          // session starts from the current content size, not stale.
+          maxContentHeight = currentContent;
+          return sh;
+        }
+        if (currentContent > maxContentHeight) {
+          maxContentHeight = currentContent;
+        }
+        const shrink = maxContentHeight - currentContent;
+        // Large shrinks (over STABILIZER_MAX_PX) are intentional content
+        // removals — message delete, regenerate clearing the old
+        // assistant turn, reasoning-panel collapse. Compensating for
+        // those would leave persistent empty space at the bottom of the
+        // viewport, which the user reads as "weird empty gap." Release
+        // the stabilizer instead and rebase the high-water mark; the
+        // pinIfFollowing call right after will smoothly re-anchor to
+        // the new (smaller) bottom.
+        if (shrink > STABILIZER_MAX_PX) {
+          maxContentHeight = currentContent;
+          if (stabilizerPx !== 0) {
+            stabilizerPx = 0;
+            el.style.removeProperty("--aui-scroll-stabilizer");
+          }
+          return currentContent;
+        }
+        const needed = Math.max(0, shrink);
+        if (needed !== stabilizerPx) {
+          stabilizerPx = needed;
+          el.style.setProperty(
+            "--aui-scroll-stabilizer",
+            `${stabilizerPx}px`,
+          );
+        }
+        return currentContent + stabilizerPx;
+      };
 
-      const mutationObserver = new MutationObserver(() => {
-        extendFollow();
-        requestTick();
-      });
+      // Synchronous pin-to-bottom. Observer callbacks run in the event-
+      // loop's "update the rendering" step (after layout, before paint),
+      // so the scrollTo here is composited in the same frame as the
+      // mutation that triggered the observer.
+      const pinIfFollowing = (scrollHeight: number): void => {
+        if (userDetachedRef.current) {
+          return;
+        }
+        if (performance.now() >= followUntilRef.current) {
+          return;
+        }
+        if (scrollHeight <= el.clientHeight) {
+          return;
+        }
+        el.scrollTo({ top: scrollHeight, behavior: "instant" });
+      };
 
-      const onViewportResize = () => {
+      // All three layout-change signals fan in here so there's a
+      // single place to understand "what runs when the viewport's
+      // content shape changes". Order matters: extend first so the
+      // stabilizer sees the follow window as active; stabilize before
+      // pinning so we scroll to the post-adjustment scrollHeight.
+      const onLayoutChange = (): void => {
         extendFollow();
+        const scrollHeight = stabilize();
+        pinIfFollowing(scrollHeight);
         requestTick();
       };
 
+      const resizeObserver = new ResizeObserver(onLayoutChange);
+      const mutationObserver = new MutationObserver(onLayoutChange);
+      const onViewportResize = onLayoutChange;
+
       // Fresh attach always starts pinned. `userDetachedRef` survives
       // ref rebinds (it's hook-scoped), so if the viewport element is
       // ever unmounted and remounted without an AUI lifecycle event
@@ -366,7 +493,13 @@ export function useIntentAwareAutoScroll(): {
       setIsAtBottom(true);
       requestTick();
 
-      resizeObserver.observe(el);
+      // Observe the border box, not the content box. The stabilizer
+      // writes `padding-bottom`, which shrinks the content box; if we
+      // observed that, every stabilizer adjustment would echo back as
+      // a resize and re-enter onLayoutChange. Border-box stays put
+      // through padding changes but still tracks parent-driven
+      // resizes (window, sidebar toggle) — which is all we need.
+      resizeObserver.observe(el, { box: "border-box" });
       mutationObserver.observe(el, {
         childList: true,
         subtree: true,
diff --git a/studio/frontend/src/components/ui/button.tsx b/studio/frontend/src/components/ui/button.tsx
index e95ab8faa7..9e27446989 100644
--- a/studio/frontend/src/components/ui/button.tsx
+++ b/studio/frontend/src/components/ui/button.tsx
@@ -1,69 +1,69 @@
 // 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 buttonVariants = cva(
-  "focus-visible:border-ring focus-visible:ring-ring/50 aria-invalid:ring-destructive/20 dark:aria-invalid:ring-destructive/40 aria-invalid:border-destructive dark:aria-invalid:border-destructive/50 rounded-4xl border border-transparent bg-clip-padding text-sm font-medium focus-visible:ring-[3px] aria-invalid:ring-[3px] [&_svg:not([class*='size-'])]:size-4 inline-flex items-center justify-center whitespace-nowrap transition-all disabled:pointer-events-none disabled:opacity-50 [&_svg]:pointer-events-none shrink-0 [&_svg]:shrink-0 outline-none group/button select-none cursor-pointer",
-  {
-    variants: {
-      variant: {
-        default: "bg-primary text-primary-foreground hover:bg-primary/80",
-        dark: "bg-foreground text-background hover:bg-foreground/85 dark:bg-foreground dark:text-background",
-        outline:
-          "border-border bg-input/30 hover:bg-input/50 hover:text-foreground aria-expanded:bg-muted aria-expanded:text-foreground",
-        secondary:
-          "bg-secondary text-secondary-foreground hover:bg-secondary/80 aria-expanded:bg-secondary aria-expanded:text-secondary-foreground",
-        ghost:
-          "hover:bg-muted hover:text-foreground dark:hover:bg-muted/50 aria-expanded:bg-muted aria-expanded:text-foreground",
-        destructive:
-          "bg-destructive/10 hover:bg-destructive/20 focus-visible:ring-destructive/20 dark:focus-visible:ring-destructive/40 dark:bg-destructive/20 text-destructive focus-visible:border-destructive/40 dark:hover:bg-destructive/30",
-        link: "text-primary underline-offset-4 hover:underline",
-      },
-      size: {
-        default:
-          "h-9 gap-1.5 px-3 has-data-[icon=inline-end]:pr-2.5 has-data-[icon=inline-start]:pl-2.5",
-        xs: "h-6 gap-1 px-2.5 text-xs has-data-[icon=inline-end]:pr-2 has-data-[icon=inline-start]:pl-2 [&_svg:not([class*='size-'])]:size-3",
-        sm: "h-8 gap-1 px-3 has-data-[icon=inline-end]:pr-2 has-data-[icon=inline-start]:pl-2",
-        lg: "h-10 gap-1.5 px-4 has-data-[icon=inline-end]:pr-3 has-data-[icon=inline-start]:pl-3",
-        icon: "size-9",
-        "icon-xs": "size-6 [&_svg:not([class*='size-'])]:size-3",
-        "icon-sm": "size-8",
-        "icon-lg": "size-10",
-      },
-    },
-    defaultVariants: {
-      variant: "default",
-      size: "default",
-    },
-  },
-);
-
-export function Button({
-  className,
-  variant = "default",
-  size = "default",
-  asChild = false,
-  ...props
-}: React.ComponentProps<"button"> &
-  VariantProps & {
-    asChild?: boolean;
-  }): React.ReactElement {
-  const Comp = asChild ? Slot.Root : "button";
-
-  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 buttonVariants = cva(
+  "focus-visible:border-ring focus-visible:ring-ring/50 aria-invalid:ring-destructive/20 dark:aria-invalid:ring-destructive/40 aria-invalid:border-destructive dark:aria-invalid:border-destructive/50 rounded-4xl border border-transparent text-sm font-medium focus-visible:ring-[3px] aria-invalid:ring-[3px] [&_svg:not([class*='size-'])]:size-4 inline-flex items-center justify-center whitespace-nowrap transition-all disabled:pointer-events-none disabled:opacity-50 [&_svg]:pointer-events-none shrink-0 [&_svg]:shrink-0 outline-none group/button select-none cursor-pointer",
+  {
+    variants: {
+      variant: {
+        default: "bg-primary text-primary-foreground hover:bg-primary/80",
+        dark: "bg-foreground text-background hover:bg-foreground/85 dark:bg-foreground dark:text-background",
+        outline:
+          "border-border bg-input/30 hover:bg-input/50 hover:text-foreground aria-expanded:bg-muted aria-expanded:text-foreground",
+        secondary:
+          "bg-secondary text-secondary-foreground hover:bg-secondary/80 aria-expanded:bg-secondary aria-expanded:text-secondary-foreground",
+        ghost:
+          "hover:bg-muted hover:text-foreground dark:hover:bg-muted/50 aria-expanded:bg-muted aria-expanded:text-foreground",
+        destructive:
+          "bg-destructive/10 hover:bg-destructive/20 focus-visible:ring-destructive/20 dark:focus-visible:ring-destructive/40 dark:bg-destructive/20 text-destructive focus-visible:border-destructive/40 dark:hover:bg-destructive/30",
+        link: "text-primary underline-offset-4 hover:underline",
+      },
+      size: {
+        default:
+          "h-9 gap-1.5 px-3 has-data-[icon=inline-end]:pr-2.5 has-data-[icon=inline-start]:pl-2.5",
+        xs: "h-6 gap-1 px-2.5 text-xs has-data-[icon=inline-end]:pr-2 has-data-[icon=inline-start]:pl-2 [&_svg:not([class*='size-'])]:size-3",
+        sm: "h-8 gap-1 px-3 has-data-[icon=inline-end]:pr-2 has-data-[icon=inline-start]:pl-2",
+        lg: "h-10 gap-1.5 px-4 has-data-[icon=inline-end]:pr-3 has-data-[icon=inline-start]:pl-3",
+        icon: "size-9",
+        "icon-xs": "size-6 [&_svg:not([class*='size-'])]:size-3",
+        "icon-sm": "size-8",
+        "icon-lg": "size-10",
+      },
+    },
+    defaultVariants: {
+      variant: "default",
+      size: "default",
+    },
+  },
+);
+
+export function Button({
+  className,
+  variant = "default",
+  size = "default",
+  asChild = false,
+  ...props
+}: React.ComponentProps<"button"> &
+  VariantProps & {
+    asChild?: boolean;
+  }): React.ReactElement {
+  const Comp = asChild ? Slot.Root : "button";
+
+  return (
+    
+  );
+}
diff --git a/studio/frontend/src/components/ui/select.tsx b/studio/frontend/src/components/ui/select.tsx
index f65d7c3676..4044c164e5 100644
--- a/studio/frontend/src/components/ui/select.tsx
+++ b/studio/frontend/src/components/ui/select.tsx
@@ -1,244 +1,257 @@
 // 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 { Select as SelectPrimitive } from "radix-ui";
-import type * as React from "react";
-import { createContext, useContext, useState } from "react";
-
-import { cn } from "@/lib/utils";
-import { useDialogPortalContainer } from "@/components/ui/dialog";
-import {
-  ArrowDown01Icon,
-  ArrowUp01Icon,
-  Tick02Icon,
-  UnfoldMoreIcon,
-} from "@hugeicons/core-free-icons";
-import { HugeiconsIcon } from "@hugeicons/react";
-
-const SelectOpenContext = createContext(false);
-
-function Select({
-  onOpenChange,
-  ...props
-}: React.ComponentProps) {
-  const [isOpen, setIsOpen] = useState(false);
-  return (
-    
-       {
-          setIsOpen(open);
-          onOpenChange?.(open);
-        }}
-        {...props}
-      />
-    
-  );
-}
-
-function SelectGroup({
-  className,
-  ...props
-}: React.ComponentProps) {
-  return (
-    
-  );
-}
-
-function SelectValue({
-  ...props
-}: React.ComponentProps) {
-  return ;
-}
-
-function SelectTrigger({
-  className,
-  size = "default",
-  children,
-  ...props
-}: React.ComponentProps & {
-  size?: "sm" | "default";
-}) {
-  const isOpen = useContext(SelectOpenContext);
-
-  return (
-    
-      {children}
-      
-        
-      
-    
-  );
-}
-
-function SelectContent({
-  className,
-  children,
-  position = "item-aligned",
-  align = "center",
-  container,
-  ...props
-}: React.ComponentProps & {
-  container?: HTMLElement | null;
-}) {
-  const dialogContainer = useDialogPortalContainer();
-  return (
-    
-      
-        
-        
-          {children}
-        
-        
-      
-    
-  );
-}
-
-function SelectLabel({
-  className,
-  ...props
-}: React.ComponentProps) {
-  return (
-    
-  );
-}
-
-function SelectItem({
-  className,
-  children,
-  ...props
-}: React.ComponentProps) {
-  return (
-    
-      
-        
-          
-        
-      
-      {children}
-    
-  );
-}
-
-function SelectSeparator({
-  className,
-  ...props
-}: React.ComponentProps) {
-  return (
-    
-  );
-}
-
-function SelectScrollUpButton({
-  className,
-  ...props
-}: React.ComponentProps) {
-  return (
-    
-      
-    
-  );
-}
-
-function SelectScrollDownButton({
-  className,
-  ...props
-}: React.ComponentProps) {
-  return (
-    
-      
-    
-  );
-}
-
-export {
-  Select,
-  SelectContent,
-  SelectGroup,
-  SelectItem,
-  SelectLabel,
-  SelectScrollDownButton,
-  SelectScrollUpButton,
-  SelectSeparator,
-  SelectTrigger,
-  SelectValue,
-};
+"use client";
+
+import { Select as SelectPrimitive } from "radix-ui";
+import type * as React from "react";
+import { createContext, useContext, useState } from "react";
+
+import { cn } from "@/lib/utils";
+import { useDialogPortalContainer } from "@/components/ui/dialog";
+import {
+  ArrowDown01Icon,
+  ArrowUp01Icon,
+  Tick02Icon,
+  UnfoldMoreIcon,
+} from "@hugeicons/core-free-icons";
+import { HugeiconsIcon } from "@hugeicons/react";
+
+const SelectOpenContext = createContext(false);
+
+function Select({
+  onOpenChange,
+  ...props
+}: React.ComponentProps) {
+  const [isOpen, setIsOpen] = useState(false);
+  return (
+    
+       {
+          setIsOpen(open);
+          onOpenChange?.(open);
+        }}
+        {...props}
+      />
+    
+  );
+}
+
+function SelectGroup({
+  className,
+  ...props
+}: React.ComponentProps) {
+  return (
+    
+  );
+}
+
+function SelectValue({
+  ...props
+}: React.ComponentProps) {
+  return ;
+}
+
+function SelectTrigger({
+  className,
+  size = "default",
+  children,
+  icon,
+  iconClassName,
+  animateRadius = true,
+  ...props
+}: React.ComponentProps & {
+  size?: "sm" | "default";
+  icon?: typeof UnfoldMoreIcon;
+  iconClassName?: string;
+  animateRadius?: boolean;
+}) {
+  const isOpen = useContext(SelectOpenContext);
+
+  return (
+    
+      {children}
+      
+        
+      
+    
+  );
+}
+
+function SelectContent({
+  className,
+  children,
+  position = "item-aligned",
+  align = "center",
+  container,
+  ...props
+}: React.ComponentProps & {
+  container?: HTMLElement | null;
+}) {
+  const dialogContainer = useDialogPortalContainer();
+  return (
+    
+      
+        
+        
+          {children}
+        
+        
+      
+    
+  );
+}
+
+function SelectLabel({
+  className,
+  ...props
+}: React.ComponentProps) {
+  return (
+    
+  );
+}
+
+function SelectItem({
+  className,
+  children,
+  ...props
+}: React.ComponentProps) {
+  return (
+    
+      
+        
+          
+        
+      
+      {children}
+    
+  );
+}
+
+function SelectSeparator({
+  className,
+  ...props
+}: React.ComponentProps) {
+  return (
+    
+  );
+}
+
+function SelectScrollUpButton({
+  className,
+  ...props
+}: React.ComponentProps) {
+  return (
+    
+      
+    
+  );
+}
+
+function SelectScrollDownButton({
+  className,
+  ...props
+}: React.ComponentProps) {
+  return (
+    
+      
+    
+  );
+}
+
+export {
+  Select,
+  SelectContent,
+  SelectGroup,
+  SelectItem,
+  SelectLabel,
+  SelectScrollDownButton,
+  SelectScrollUpButton,
+  SelectSeparator,
+  SelectTrigger,
+  SelectValue,
+};
diff --git a/studio/frontend/src/components/ui/sidebar.tsx b/studio/frontend/src/components/ui/sidebar.tsx
index 8eb8c51491..6be77d01b9 100644
--- a/studio/frontend/src/components/ui/sidebar.tsx
+++ b/studio/frontend/src/components/ui/sidebar.tsx
@@ -1,768 +1,770 @@
-// 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 * as React from "react"
-import { cva, type VariantProps } from "class-variance-authority"
-import { Slot } from "radix-ui"
-
-import { cn } from "@/lib/utils"
-import { Button } from "@/components/ui/button"
-import { Input } from "@/components/ui/input"
-import { Separator } from "@/components/ui/separator"
-import {
-  Sheet,
-  SheetContent,
-  SheetDescription,
-  SheetHeader,
-  SheetTitle,
-} from "@/components/ui/sheet"
-import { Skeleton } from "@/components/ui/skeleton"
-import {
-  Tooltip,
-  TooltipContent,
-  TooltipTrigger,
-} from "@/components/ui/tooltip"
-import { useIsMobile } from "@/hooks/use-mobile"
-import { HugeiconsIcon } from "@hugeicons/react"
-import { LayoutAlignLeftIcon } from "@hugeicons/core-free-icons"
-
-const noop = () => {}
-
-const SIDEBAR_WIDTH = "16rem"
-const SIDEBAR_WIDTH_ICON = "3rem"
-const SIDEBAR_KEYBOARD_SHORTCUT = "b"
-
-type SidebarContextProps = {
-  state: "expanded" | "collapsed"
-  open: boolean
-  setOpen: (open: boolean) => void
-  openMobile: boolean
-  setOpenMobile: (open: boolean) => void
-  isMobile: boolean
-  toggleSidebar: () => void
-  hasPinMode: boolean
-  pinned: boolean
-  setPinned: (value: boolean) => void
-  togglePinned: () => void
-}
-
-const SidebarContext = React.createContext(null)
-
-function useSidebar() {
-  const context = React.useContext(SidebarContext)
-  if (!context) {
-    throw new Error("useSidebar must be used within a SidebarProvider.")
-  }
-
-  return context
-}
-
-function SidebarProvider({
-  defaultOpen = true,
-  open: openProp,
-  onOpenChange: setOpenProp,
-  pinned: pinnedProp,
-  setPinned: setPinnedProp,
-  togglePinned: togglePinnedProp,
-  className,
-  style,
-  children,
-  ...props
-}: React.ComponentProps<"div"> & {
-  defaultOpen?: boolean
-  open?: boolean
-  onOpenChange?: (open: boolean) => void
-  pinned?: boolean
-  setPinned?: (value: boolean) => void
-  togglePinned?: () => void
-}) {
-  const isMobile = useIsMobile()
-  const [openMobile, setOpenMobile] = React.useState(false)
-
-  const prevIsMobileRef = React.useRef(isMobile)
-  React.useEffect(() => {
-    if (prevIsMobileRef.current && !isMobile) {
-      setOpenMobile(false)
-    }
-    prevIsMobileRef.current = isMobile
-  }, [isMobile])
-
-  // Whether pin mode is active (caller provides pinned + setPinned + togglePinned).
-  const hasPinMode = pinnedProp !== undefined && setPinnedProp !== undefined && togglePinnedProp !== undefined
-
-  // This is the internal state of the sidebar.
-  // We use openProp and setOpenProp for control from outside the component.
-  const [_open, _setOpen] = React.useState(defaultOpen)
-
-  // When pin mode is active, open is driven entirely by `pinned` (explicit
-  // user toggle). Otherwise fall back to the controlled/uncontrolled pattern.
-  const open = hasPinMode ? !!pinnedProp : (openProp ?? _open)
-
-  const setOpen = React.useCallback(
-    (value: boolean | ((value: boolean) => boolean)) => {
-      const openState = typeof value === "function" ? value(open) : value
-
-      if (hasPinMode) {
-        // In pin mode, setOpen controls pinned state.
-        setPinnedProp?.(openState)
-        return
-      }
-
-      if (setOpenProp) {
-        setOpenProp(openState)
-      } else {
-        _setOpen(openState)
-      }
-    },
-    [setOpenProp, open, hasPinMode, setPinnedProp]
-  )
-
-  // Helper to toggle the sidebar.
-  const toggleSidebar = React.useCallback(() => {
-    if (isMobile) return setOpenMobile((open) => !open)
-    if (hasPinMode && togglePinnedProp) return togglePinnedProp()
-    return setOpen((open) => !open)
-  }, [isMobile, setOpen, setOpenMobile, hasPinMode, togglePinnedProp])
-
-  // Adds a keyboard shortcut to toggle the sidebar.
-  React.useEffect(() => {
-    const handleKeyDown = (event: KeyboardEvent) => {
-      if (
-        event.key === SIDEBAR_KEYBOARD_SHORTCUT &&
-        (event.metaKey || event.ctrlKey)
-      ) {
-        event.preventDefault()
-        toggleSidebar()
-      }
-    }
-
-    window.addEventListener("keydown", handleKeyDown)
-    return () => window.removeEventListener("keydown", handleKeyDown)
-  }, [toggleSidebar])
-
-  // We add a state so that we can do data-state="expanded" or "collapsed".
-  // This makes it easier to style the sidebar with Tailwind classes.
-  const state = open ? "expanded" : "collapsed"
-
-  const pinned = pinnedProp ?? false
-  const setPinned = setPinnedProp ?? noop
-  const togglePinned = togglePinnedProp ?? noop
-
-  const contextValue = React.useMemo(
-    () => ({
-      state,
-      open,
-      setOpen,
-      isMobile,
-      openMobile,
-      setOpenMobile,
-      toggleSidebar,
-      hasPinMode,
-      pinned,
-      setPinned,
-      togglePinned,
-    }),
-    [state, open, setOpen, isMobile, openMobile, setOpenMobile, toggleSidebar, hasPinMode, pinned, setPinned, togglePinned]
-  )
-
-  return (
-    
-      

-        {children}
-      

-    
-  )
-}
-
-function Sidebar({
-  side = "left",
-  variant = "sidebar",
-  collapsible = "offcanvas",
-  className,
-  children,
-  dir,
-  ...props
-}: React.ComponentProps<"div"> & {
-  side?: "left" | "right"
-  variant?: "sidebar" | "floating" | "inset"
-  collapsible?: "offcanvas" | "icon" | "none"
-}) {
-  const { isMobile, state, openMobile, setOpenMobile, hasPinMode, pinned } = useSidebar()
-
-  if (collapsible === "none") {
-    return (
-      

-        {children}
-      

-    )
-  }
-
-  if (isMobile) {
-    return (
-      
-        
-          
-            Sidebar
-            Displays the mobile sidebar.
-          
-          
{children}

-        
-      
-    )
-  }
-
-  return (
-    

-      {/* This is what handles the sidebar gap on desktop */}
-      

-      

-        

-          {children}
-        

-      

-    

-  )
-}
-
-function SidebarTrigger({
-  className,
-  onClick,
-  ...props
-}: React.ComponentProps) {
-  const { toggleSidebar } = useSidebar()
-
-  return (
-    
-  )
-}
-
-function SidebarRail({ className, ...props }: React.ComponentProps<"button">) {
-  const { toggleSidebar } = useSidebar()
-
-  return (
-    
+  )
+}
+
+function SidebarRail({ className, ...props }: React.ComponentProps<"button">) {
+  const { toggleSidebar } = useSidebar()
+
+  return (
+    
                 
-                
+                
                   Open configuration
                 
               
diff --git a/studio/frontend/src/features/chat/chat-settings-sheet.tsx b/studio/frontend/src/features/chat/chat-settings-sheet.tsx
index 08c5ef4080..f20d621d08 100644
--- a/studio/frontend/src/features/chat/chat-settings-sheet.tsx
+++ b/studio/frontend/src/features/chat/chat-settings-sheet.tsx
@@ -22,11 +22,9 @@ import {
   DropdownMenuSeparator,
   DropdownMenuTrigger,
 } from "@/components/ui/dropdown-menu";
-import { Input } from "@/components/ui/input";
 import {
   InputGroup,
   InputGroupAddon,
-  InputGroupButton,
   InputGroupInput,
 } from "@/components/ui/input-group";
 import {
@@ -50,23 +48,20 @@ import { useIsMobile } from "@/hooks/use-mobile";
 import { cn } from "@/lib/utils";
 import {
   ArrowDown01Icon,
-  CodeIcon,
-  Delete02Icon,
-  FloppyDiskIcon,
-  Settings02Icon,
-  Settings05Icon,
-  SlidersHorizontalIcon,
-  Wrench01Icon,
+  ArrowTurnBackwardIcon,
+  InformationCircleIcon,
+  LayoutAlignRightIcon,
 } from "@hugeicons/core-free-icons";
 import { HugeiconsIcon } from "@hugeicons/react";
 import {
   Tooltip,
   TooltipContent,
+  TooltipTrigger,
 } from "@/components/ui/tooltip";
 import { Tooltip as TooltipPrimitive } from "radix-ui";
-import { AnimatePresence, motion } from "motion/react";
+import { ChevronDown } from "lucide-react";
 import { Fragment, type ReactNode } from "react";
-import { useEffect, useLayoutEffect, useMemo, useRef, useState } from "react";
+import { useEffect, useMemo, useRef, useState } from "react";
 import { toast } from "sonner";
 import { useChatRuntimeStore } from "./stores/chat-runtime-store";
 import {
@@ -174,7 +169,10 @@ function migrateLegacySystemPromptTemplates(presets: Preset[]): Preset[] {
       localStorage.setItem(LEGACY_CHAT_SYSTEM_PROMPTS_MIGRATED_KEY, raw);
       return presets;
     }
-    const mergedPresets = normalizeCustomPresets([...presets, ...importedPresets]);
+    const mergedPresets = normalizeCustomPresets([
+      ...presets,
+      ...importedPresets,
+    ]);
     saveCustomPresets(mergedPresets);
     try {
       localStorage.setItem(LEGACY_CHAT_SYSTEM_PROMPTS_MIGRATED_KEY, raw);
@@ -232,6 +230,139 @@ function loadSavedActivePreset(): string {
   }
 }
 
+function InfoHint({ children }: { children: ReactNode }) {
+  return (
+    
+      
+        
+      
+      
+        {children}
+      
+    
+  );
+}
+
+/**
+ * Editable numeric value display.
+ *
+ * Renders as a single  that *looks* like text by default —
+ * transparent background, no border, no ring — and only shows a faint
+ * surface tint on hover/focus to signal editability. When unfocused,
+ * the input shows the formatted display string (`displayValue ?? value`,
+ * so labels like "Off" / "Max" still render); on focus, it switches to
+ * the raw numeric value, selects it, and accepts free text input.
+ * Commit happens on blur or Enter; Escape reverts. The clamp-to-range
+ * happens on commit so users can type intermediate values without the
+ * input fighting them mid-keystroke. Single component shared by every
+ * slider value and the Context Length input so the click-to-edit
+ * affordance is consistent across the panel.
+ */
+function snapToStep(
+  value: number,
+  step: number,
+  min?: number,
+  max?: number,
+): number {
+  const lo = min ?? Number.NEGATIVE_INFINITY;
+  const hi = max ?? Number.POSITIVE_INFINITY;
+  const clamped = Math.min(Math.max(value, lo), hi);
+  const stepStr = String(step);
+  const decimals = stepStr.includes(".") ? stepStr.split(".")[1].length : 0;
+  const base = Number.isFinite(lo) ? lo : 0;
+  const snapped = base + Math.round((clamped - base) / step) * step;
+  const reclamped = Math.min(Math.max(snapped, lo), hi);
+  return Number(reclamped.toFixed(decimals));
+}
+
+function NumericValueInput({
+  value,
+  min,
+  max,
+  step,
+  onChange,
+  displayValue,
+  className,
+  ariaLabel,
+  size: sizeAttr,
+}: {
+  value: number;
+  min?: number;
+  max?: number;
+  step: number;
+  onChange: (v: number) => void;
+  displayValue?: string;
+  className?: string;
+  ariaLabel?: string;
+  size?: number;
+}) {
+  const [focused, setFocused] = useState(false);
+  const [draft, setDraft] = useState("");
+  const cancelBlurCommitRef = useRef(false);
+
+  const commit = (raw: string) => {
+    const parsed = Number.parseFloat(raw);
+    if (!Number.isFinite(parsed)) {
+      return;
+    }
+    const final = snapToStep(parsed, step, min, max);
+    if (final !== value) {
+      onChange(final);
+    }
+  };
+
+  return (
+     {
+        cancelBlurCommitRef.current = false;
+        setDraft(String(value));
+        setFocused(true);
+        // Defer the select() so it runs after the value swap above.
+        const target = e.currentTarget;
+        requestAnimationFrame(() => target.select());
+      }}
+      onBlur={() => {
+        if (cancelBlurCommitRef.current) {
+          cancelBlurCommitRef.current = false;
+        } else {
+          commit(draft);
+        }
+        setFocused(false);
+      }}
+      onChange={(e) => setDraft(e.target.value)}
+      onKeyDown={(e) => {
+        if (e.key === "Enter") {
+          e.currentTarget.blur();
+        } else if (e.key === "Escape") {
+          cancelBlurCommitRef.current = true;
+          setDraft(String(value));
+          e.currentTarget.blur();
+        }
+      }}
+      className={cn("panel-number-input", className)}
+    />
+  );
+}
+
 function ParamSlider({
   label,
   value,
@@ -240,6 +371,8 @@ function ParamSlider({
   step,
   onChange,
   displayValue,
+  info,
+  valueSize,
 }: {
   label: string;
   value: number;
@@ -248,21 +381,36 @@ function ParamSlider({
   step: number;
   onChange: (v: number) => void;
   displayValue?: string;
+  info?: ReactNode;
+  valueSize?: number;
 }) {
   return (
-    

-      

-        {label}
-        
-          {displayValue ?? value}
-        
+    

+      

+        

+          
+            {label}
+          
+          {info && {info}}
+        

+        
       

        onChange(v)}
+        onValueChange={([v]) => onChange(snapToStep(v, step, min, max))}
+        className="panel-slider"
       />
     

   );
@@ -306,15 +454,15 @@ function saveCollapsibleOpen(label: string, open: boolean) {
 }
 
 function CollapsibleSection({
-  icon,
   label,
   children,
   defaultOpen = false,
+  first = false,
 }: {
-  icon: Parameters[0]["icon"];
   label: string;
   children?: ReactNode;
   defaultOpen?: boolean;
+  first?: boolean;
 }) {
   const [open, setOpen] = useState(() => {
     const saved = loadCollapsibleState();
@@ -322,7 +470,12 @@ function CollapsibleSection({
   });
 
   return (
-    

+    

       
-      
-        {open && (
-          
-            
{children}

-          
+        className={cn(
+          "flex w-full cursor-pointer items-center justify-between text-[12px] font-medium normal-case tracking-[0.04em] text-nav-fg-muted transition-colors hover:text-nav-fg focus-visible:outline-none focus-visible:ring-0",
+          first ? "pt-4 pb-5" : "py-5",
         )}
-      
+      >
+        {label}
+        
+          
+        
+      
+      {open && 
{children}
}
     

   );
 }
@@ -378,18 +517,26 @@ export function ChatSettingsPanel({
 }: ChatSettingsPanelProps) {
   const isMobile = useIsMobile();
   const isGguf = useChatRuntimeStore((s) => s.activeGgufVariant) != null;
+  const hasModelContent = isGguf || Boolean(params.checkpoint);
   const speculativeType = useChatRuntimeStore((s) => s.speculativeType);
   const setSpeculativeType = useChatRuntimeStore((s) => s.setSpeculativeType);
   const loadedSpeculativeType = useChatRuntimeStore(
     (s) => s.loadedSpeculativeType,
   );
-  const currentModels = useChatRuntimeStore((s) => s.models);
   const modelRequiresTrustRemoteCode = useChatRuntimeStore(
     (s) => s.modelRequiresTrustRemoteCode,
   );
   const currentCheckpoint = params.checkpoint;
-  const currentModelIsVision =
-    currentModels.find((m) => m.id === currentCheckpoint)?.isVision ?? false;
+  const currentModelIsMultimodal = useChatRuntimeStore((s) => {
+    if (s.loadedIsMultimodal) return true;
+    const m = s.models.find((m) => m.id === currentCheckpoint);
+    return (
+      Boolean(m?.isVision) ||
+      Boolean(m?.isAudio) ||
+      Boolean(m?.hasAudioInput) ||
+      m?.audioType === "audio_vlm"
+    );
+  });
   const ggufContextLength = useChatRuntimeStore((s) => s.ggufContextLength);
   const ggufMaxContextLength = useChatRuntimeStore(
     (s) => s.ggufMaxContextLength,
@@ -415,6 +562,16 @@ export function ChatSettingsPanel({
   const ctxDirty = customContextLength !== null;
   const specDirty = speculativeType !== loadedSpeculativeType;
   const modelSettingsDirty = kvDirty || ctxDirty || specDirty;
+  const chatTemplateOverride = useChatRuntimeStore(
+    (s) => s.chatTemplateOverride,
+  );
+  const loadedChatTemplateOverride = useChatRuntimeStore(
+    (s) => s.loadedChatTemplateOverride,
+  );
+  const setChatTemplateOverride = useChatRuntimeStore(
+    (s) => s.setChatTemplateOverride,
+  );
+  const templateDirty = chatTemplateOverride !== loadedChatTemplateOverride;
   const [customPresets, setCustomPresets] = useState(() =>
     loadSavedCustomPresets(),
   );
@@ -424,10 +581,6 @@ export function ChatSettingsPanel({
   const [presetNameInput, setPresetNameInput] = useState(() =>
     loadSavedActivePreset(),
   );
-  const presetControlRowRef = useRef(null);
-  const [presetMenuWidthPx, setPresetMenuWidthPx] = useState<
-    number | undefined
-  >(undefined);
   const [systemPromptEditorOpen, setSystemPromptEditorOpen] = useState(false);
   const [systemPromptDraft, setSystemPromptDraft] = useState("");
   const [activePresetBaseline, setActivePresetBaseline] = useState(params);
@@ -442,19 +595,18 @@ export function ChatSettingsPanel({
     () => customPresets.find((preset) => preset.name === activePreset) ?? null,
     [activePreset, customPresets],
   );
+  const activeBuiltinPreset = useMemo(
+    () =>
+      BUILTIN_PRESETS.find((preset) => preset.name === activePreset) ?? null,
+    [activePreset],
+  );
   const hasUnsavedPresetChanges = useMemo(
     () => {
       if (activePresetDefinition == null) {
         return false;
       }
-      if (BUILTIN_PRESET_NAMES.has(activePresetDefinition.name)) {
-        if (activePresetDefinition.name === "Default") {
-          return activePresetSource === "modified";
-        }
-        return (
-          activePresetSource === "modified" ||
-          !isSamePresetConfig(activePresetDefinition.params, params)
-        );
+      if (activePresetDefinition.name === "Default") {
+        return activePresetSource === "modified";
       }
       return !isSamePresetConfig(activePresetDefinition.params, params);
     },
@@ -520,7 +672,10 @@ export function ChatSettingsPanel({
       : trimmed;
     setCustomPresets((prev) => {
       const next = prev.filter((p) => p.name !== saveName);
-      const merged = [...next, { name: saveName, params: toPresetParams(params) }];
+      const merged = [
+        ...next,
+        { name: saveName, params: toPresetParams(params) },
+      ];
       saveCustomPresets(merged);
       return merged;
     });
@@ -544,7 +699,8 @@ export function ChatSettingsPanel({
       return;
     }
     const fallbackPreset =
-      BUILTIN_PRESETS.find((preset) => preset.name === "Default") ?? null;
+      BUILTIN_PRESETS.find((preset) => preset.name === "Default") ??
+      null;
     setCustomPresets((prev) => {
       const next = prev.filter((preset) => preset.name !== name);
       saveCustomPresets(next);
@@ -587,28 +743,6 @@ export function ChatSettingsPanel({
   useEffect(() => {
     if (presets.some((preset) => preset.name === activePreset)) {
       const expectedSource = getPresetSource(activePreset);
-      if (activePresetDefinition != null) {
-        if (BUILTIN_PRESET_NAMES.has(activePresetDefinition.name)) {
-          if (activePresetDefinition.name === "Default") {
-            if (
-              activePresetSource !== "modified" &&
-              activePresetSource !== expectedSource
-            ) {
-              setActivePresetSource(expectedSource);
-            }
-            return;
-          }
-          const matchesActivePreset = isSamePresetConfig(
-            activePresetDefinition.params,
-            params,
-          );
-          const nextSource = matchesActivePreset ? expectedSource : "modified";
-          if (activePresetSource !== nextSource) {
-            setActivePresetSource(nextSource);
-          }
-          return;
-        }
-      }
       if (
         activePresetSource !== "modified" &&
         activePresetSource !== expectedSource
@@ -628,9 +762,7 @@ export function ChatSettingsPanel({
     }
   }, [
     activePreset,
-    activePresetDefinition,
     activePresetSource,
-    params,
     presets,
     setActivePresetSource,
   ]);
@@ -645,307 +777,302 @@ export function ChatSettingsPanel({
     }
   }, [open]);
 
-  useLayoutEffect(() => {
-    const el = presetControlRowRef.current;
-    if (!el || !open) return;
-    const measure = () => {
-      setPresetMenuWidthPx(el.getBoundingClientRect().width);
-    };
-    measure();
-    const ro = new ResizeObserver(measure);
-    ro.observe(el);
-    return () => ro.disconnect();
-  }, [open]);
-
-  const modelSection = (
-    
-      

-        {isGguf && (
-          <>
-            

-              

-                Context Length
-                 {
-                    const raw = e.target.value;
-                    if (raw === "") {
-                      setCustomContextLength(null);
-                      return;
-                    }
-                    const v = Number.parseInt(raw, 10);
-                    if (!Number.isNaN(v) && v >= 0) {
-                      const maxCtx = ctxMaxValue ?? Number.POSITIVE_INFINITY;
-                      const clamped = Math.min(v, maxCtx);
-                      setCustomContextLength(
-                        clamped === (ggufContextLength ?? 0) ? null : clamped,
-                      );
-                    }
-                  }}
-                />
-              

-               {
-                  setCustomContextLength(
-                    v === (ggufContextLength ?? 0) ? null : v,
-                  );
-                }}
-              />
-              {ggufMaxContextLength != null &&
-                typeof ctxDisplayValue === "number" &&
-                ctxDisplayValue > ggufMaxContextLength && (
-                  

-                    Exceeds estimated VRAM capacity (
-                    {ggufMaxContextLength.toLocaleString()} tokens). The model
-                    may use system RAM.
-                  

-                )}
-            

-            

-              

-                
KV Cache Dtype

-                

-                  Quantize KV cache to reduce VRAM.
-                

-              

-              

-                
-              

-            

-            {!currentModelIsVision && (
-              

-                

-                  

-                    Speculative Decoding
-                  

-                  

-                    Speed up generation with no VRAM cost.
-                  

-                

-                

-                  
-                

-              

-            )}
-            {modelSettingsDirty && (
-              

-                
-                
-              

-            )}
-          
-        )}
-        {!isGguf && params.checkpoint && (
-          <>
-            

-              

-                
Enable custom code

-                

-                  Allow models with custom code (e.g. Nemotron). Only enable if
-                  sure.
-                

-              

-              
-            

-            {trustRemoteCodeMissing && (
-              
-                
-                  Keep custom code enabled for this model
-                
-                
-                  This model requires custom code to load. You can edit the
-                  toggle, but loading will stay blocked until it is turned back
-                  on.
-                
-              
-            )}
-          
-        )}
-      

-    
-  );
-
   const settingsContent = (
     <>
       

-      

+      

         {isMobile ? (
-          
+          
             Configuration
           
         ) : (
           <>
+            
+              Configuration
+            
             
               
                 
               
-              
+              
                 Close configuration
               
             
-            
-              Configuration
-            
           
         )}
       

 
-      

-        {/* mt-4 matches the Playground sidebar gap (SidebarHeader py-3 + SidebarGroup pt-1) */}
-        

-          

-            

-              
-                
-                   setPresetNameInput(e.target.value)}
-                    onKeyDown={(e) => {
-                      if (e.key === "Enter" && presetSaveState.canSubmit) {
-                        e.preventDefault();
-                        savePresetWithName(presetNameInput);
+      

+        {hasModelContent && (
+        
+          

+            {isGguf && (
+              <>
+                

+                  

+                    
+                      Context Length
+                    
+                     {
+                        setCustomContextLength(
+                          v === (ggufContextLength ?? 0) ? null : v,
+                        );
+                      }}
+                      ariaLabel="Context Length"
+                      size={8}
+                    />
+                  

+                   {
+                      const snapped = Math.round(v);
+                      setCustomContextLength(
+                        snapped === (ggufContextLength ?? 0) ? null : snapped,
+                      );
                     }}
-                    placeholder="Preset name"
-                    maxLength={80}
-                    autoComplete="off"
-                    className={cn(
-                      "!h-8 min-h-0 min-w-0 self-stretch !pl-2.5 !pr-2 pt-1 pb-1 text-sm leading-10 md:text-sm",
-                      presetSaveState.isSaveReady &&
-                        "text-foreground placeholder:text-primary/45",
-                    )}
-                    aria-label="Inference preset name"
+                    className="panel-slider"
                   />
-                  
-                    
-                       ggufMaxContextLength && (
+                      

+                        Exceeds estimated VRAM capacity (
+                        {ggufMaxContextLength.toLocaleString()} tokens). The
+                        model may use system RAM.
+                      

+                    )}
+                

+                

+                  

+                    
+                      KV Cache Dtype
+                    
+                    
+                      Lower KV cache precision to save VRAM at the cost of some
+                      quality. f16/bf16 are full precision; q8_0/q5_1/q4_1 are
+                      quantized.
+                    
+                  

+                  

+                    
+                  

+                

+                {!currentModelIsMultimodal && (
+                  

+                    

+                      
+                        Speculative Decoding
+                      
+                      
+                        N-gram speculation; faster generation with negligible
+                        VRAM overhead. Text-only models.
+                      
+                    

+                     {
+                        setSpeculativeType(checked ? "default" : null);
+                      }}
+                    />
+                  

+                )}
+              
+            )}
+            {!isGguf && params.checkpoint && (
+              <>
+                

+                  

+                    
+                      Enable custom code
+                    
+                    
+                      Run custom Python from the model repo (e.g. Nemotron).
+                      Only enable for trusted sources.
+                    
+                  

+                  
+                

+                {trustRemoteCodeMissing && (
+                  
+                    
+                      Keep custom code enabled for this model
+                    
+                    
+                      This model requires custom code to load. You can edit the
+                      toggle, but loading will stay blocked until it is turned
+                      back on.
+                    
+                  
+                )}
+              
+            )}
+            
+            {(modelSettingsDirty || templateDirty) && (
+              

+                
+                
+              

+            )}
+          

+        
+        )}
+
+        
+          

+            
+              
+                

+                  
+                     setPresetNameInput(e.target.value)}
+                      onPointerDown={(e) => e.stopPropagation()}
+                      onClick={(e) => e.stopPropagation()}
+                      onKeyDown={(e) => {
+                        if (e.key === "Enter" && presetSaveState.canSubmit) {
+                          e.preventDefault();
+                          savePresetWithName(presetNameInput);
+                        }
+                        e.stopPropagation();
+                      }}
+                      placeholder="Preset name"
+                      maxLength={80}
+                      autoComplete="off"
+                      className={cn(
+                        "!h-9 min-h-0 min-w-0 self-stretch !pl-3.5 !pr-2 py-0 text-[13px] font-medium leading-9 text-nav-fg md:text-[13px]",
+                        presetSaveState.isSaveReady &&
+                          "placeholder:text-primary/50",
+                      )}
+                      aria-label="Inference preset name"
+                    />
+                    
+                      
-                
-                
-                  {presets.map((p, index) => (
-                    
-                       applyPreset(p.name)}>
-                        {p.name}
-                      
-                      {index === BUILTIN_PRESETS.length - 1 &&
-                        presets.length > BUILTIN_PRESETS.length && (
-                          
-                        )}
-                    
-                  ))}
-                
-              
-            

-            

+                      
+                    
+                  
+                

+              
+              
+                {presets.map((p, index) => (
+                  
+                     applyPreset(p.name)}
+                      className="flex min-h-9 items-center px-3 py-0 text-[13px] font-medium leading-[1.4] tracking-nav"
+                    >
+                      {p.name}
+                    
+                    {index === BUILTIN_PRESETS.length - 1 &&
+                      presets.length > BUILTIN_PRESETS.length && (
+                        
+                      )}
+                  
+                ))}
+              
+            
+            

               
               
             

           

-        

+        
 
-        

-          

-            
-            
-          

-