diff --git a/.git-blame-ignore-revs b/.git-blame-ignore-revs new file mode 100644 index 0000000000..17d96cd0f5 --- /dev/null +++ b/.git-blame-ignore-revs @@ -0,0 +1,8 @@ +# Commits listed here are skipped by `git blame` so that bulk, whitespace-only +# changes don't obscure the real authorship of a line. +# +# GitHub honors this file automatically. To use it locally, run once: +# git config blame.ignoreRevsFile .git-blame-ignore-revs + +# chore(studio/frontend): normalize line endings to LF +c50b8ab910f5aa56dd7ae0022d2c7b96bfe3384a diff --git a/.gitattributes b/.gitattributes index 75fba5d6ab..5f04b5e9d1 100644 --- a/.gitattributes +++ b/.gitattributes @@ -5,3 +5,9 @@ # clone (core.autocrlf=true) rewrites them to CRLF, and the trailing \r breaks # them when run in WSL/Linux (e.g. `set -e` -> "set: Illegal option -"). *.sh text eol=lf + +# Normalize Studio frontend sources to LF. Scoped to the frontend tree (rather +# than repo-wide *.ts/*.tsx/... rules) so the policy can't force LF on files +# elsewhere. text=auto lets Git detect and leave binary assets (logos, fonts) +# untouched while text files (.ts/.tsx/.json/.html/.svg/...) are stored as LF. +studio/frontend/** text=auto eol=lf diff --git a/.github/workflows/cross-platform-parity-ci.yml b/.github/workflows/cross-platform-parity-ci.yml new file mode 100644 index 0000000000..4632794587 --- /dev/null +++ b/.github/workflows/cross-platform-parity-ci.yml @@ -0,0 +1,61 @@ +# SPDX-License-Identifier: AGPL-3.0-only +# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. + +# Runs tests/python/test_cross_platform_parity.py on Windows and macOS. +# +# Why: that test is the guard that install.sh and install.ps1 stay in +# sync, but today it only runs on ubuntu-latest (auto-discovered by +# studio-backend-ci.yml's "Repo tests (CPU)" job). The test reads both +# installer scripts, and on Windows Path.read_text() defaults to the +# cp1252 locale encoding, so a non-cp1252 byte in install.sh (it already +# contains a U+274C) raises UnicodeDecodeError there even though Linux and +# macOS default to UTF-8. The reads were pinned to encoding="utf-8" in +# #6166; this job keeps that from silently regressing by exercising the +# test on the platforms it claims parity for. Pure pytest, no GPU, +# sub-second, so the matrix is cheap. + +name: Cross-platform parity + +on: + pull_request: + paths: + - 'install.sh' + - 'install.ps1' + - 'tests/python/test_cross_platform_parity.py' + - '.github/workflows/cross-platform-parity-ci.yml' + push: + branches: [main] + paths: + - 'install.sh' + - 'install.ps1' + - 'tests/python/test_cross_platform_parity.py' + - '.github/workflows/cross-platform-parity-ci.yml' + workflow_dispatch: + +concurrency: + group: ${{ github.workflow }}-${{ github.ref }} + cancel-in-progress: true + +permissions: + contents: read + +jobs: + parity: + name: parity (${{ matrix.os }}) + strategy: + fail-fast: false + matrix: + os: [windows-latest, macos-latest] + runs-on: ${{ matrix.os }} + timeout-minutes: 10 + steps: + - uses: actions/checkout@de0fac2e4500dabe0009e67214ff5f5447ce83dd # v6.0.2 + with: + persist-credentials: false + - uses: actions/setup-python@a309ff8b426b58ec0e2a45f0f869d46889d02405 # v6.2.0 + with: + python-version: '3.12' + cache: 'pip' + - run: python -m pip install -U pip pytest + - name: Cross-platform parity test + run: python -m pytest tests/python/test_cross_platform_parity.py -q diff --git a/CONTRIBUTING.md b/CONTRIBUTING.md index eb60a5a201..6eb8d1bc6e 100644 --- a/CONTRIBUTING.md +++ b/CONTRIBUTING.md @@ -27,3 +27,9 @@ Your support extends beyond code: Finally, please be mindful of our [Code of Conduct](https://github.com/unslothai/unsloth/blob/main/CODE_OF_CONDUCT.md) to ensure a welcoming and inclusive environment for everyone. Thank you so much for reading and we hope you have lots of fun using Unsloth! 🦥 + + +## Pull Request Guidelines +- Keep PRs focused on a single change +- Include a concise description and motivation +- Link related issues when applicable diff --git a/install.ps1 b/install.ps1 index a6601f36e3..b64128c730 100644 --- a/install.ps1 +++ b/install.ps1 @@ -1177,14 +1177,38 @@ shell.Run cmd, 0, False if ($SkipTorch) { $InitialGpuBranch = "no_torch" } Write-TauriDiag -GpuBranch $InitialGpuBranch -TorchIndexFamily "none" -PythonVersionForDiag $DiagPythonVersion - # ── Install uv if not present ── + # ── Install uv ── Write-TauriLog "STEP" "Installing uv package manager" - if (-not (Get-Command uv -ErrorAction SilentlyContinue)) { - substep "installing uv package manager..." + $UvMinVersion = "0.7.22" + function Test-UvVersionOk { + $cmd = Get-Command uv -ErrorAction SilentlyContinue + if (-not $cmd) { return $false } + try { + $raw = (& uv --version 2>$null | Select-Object -First 1) + } catch { + return $false + } + if ($raw -notmatch 'uv\s+([0-9]+(?:\.[0-9]+)+)') { return $false } + try { + return ([version]$Matches[1] -ge [version]$UvMinVersion) + } catch { + return $false + } + } + + if (-not (Test-UvVersionOk)) { + if (Get-Command uv -ErrorAction SilentlyContinue) { + substep "updating uv package manager..." + } else { + substep "installing uv package manager..." + } if ($script:WingetAvailable) { $prevEAP = $ErrorActionPreference $ErrorActionPreference = "Continue" - try { winget install --id=astral-sh.uv -e --source winget --accept-package-agreements --accept-source-agreements } catch {} + try { winget upgrade --id=astral-sh.uv -e --source winget --accept-package-agreements --accept-source-agreements } catch {} + if (-not (Test-UvVersionOk)) { + try { winget install --id=astral-sh.uv -e --source winget --accept-package-agreements --accept-source-agreements } catch {} + } $ErrorActionPreference = $prevEAP Refresh-SessionPath } @@ -1192,19 +1216,40 @@ shell.Run cmd, 0, False # use Astral's official PowerShell installer. This is the only # supported path on hosts without winget (Windows ARM64 runners, # corporate machines without the Store, etc.). - if (-not (Get-Command uv -ErrorAction SilentlyContinue)) { + if (-not (Test-UvVersionOk)) { substep "installing uv via https://astral.sh/uv/install.ps1..." "Yellow" Invoke-Expression (Invoke-RestMethod -Uri "https://astral.sh/uv/install.ps1") Refresh-SessionPath } } - if (-not (Get-Command uv -ErrorAction SilentlyContinue)) { + # A freshly installed uv can sit later on PATH than an older one (active + # venv, Scoop/pipx shim). Prefer a just-installed uv from a known location. + if (-not (Test-UvVersionOk)) { + $origPath = $env:PATH + foreach ($d in @($env:UV_INSTALL_DIR, $env:XDG_BIN_HOME, + (Join-Path $env:USERPROFILE ".local\bin"), + (Join-Path $env:LOCALAPPDATA "Microsoft\WinGet\Links"))) { + if ($d -and (Test-Path $d)) { + $env:PATH = "$d;$origPath" + if (Test-UvVersionOk) { break } + $env:PATH = $origPath + } + } + } + + if (-not (Test-UvVersionOk)) { step "uv" "could not be installed" "Red" substep "Install it from https://docs.astral.sh/uv/" "Yellow" return (Exit-InstallFailure "uv could not be installed") } + # When bytecode compilation is enabled, large installs can exceed uv's 60s + # default on slow machines. Default to 180s, preserving overrides ("0" disables). + if (-not $env:UV_COMPILE_BYTECODE_TIMEOUT) { + $env:UV_COMPILE_BYTECODE_TIMEOUT = "180" + } + # ── Create venv (migrate old layout if possible, otherwise fresh) ── # Pass the resolved executable path to uv so it does not re-resolve # a version string back to a conda interpreter. @@ -1529,7 +1574,7 @@ shell.Run cmd, 0, False # popping a UAC/DiskPart prompt RunAsInvoker can't suppress (manifest is # asInvoker). So only probe when a HIP SDK is present (hipinfo found -> # un-elevated) or the user opts in; else fall through to WMI name inference - # (enough to pick ROCm wheels + lemonade llama.cpp). + # (enough to pick ROCm wheels + the ROCm llama.cpp prebuilt). # An explicit opt-out (UNSLOTH_ENABLE_AMD_SMI=0/false/no/off) wins over the # HIP-SDK heuristic: a HIP SDK binary with a broken runtime can still pop the # prompt, so $HipSdkInstalled must NOT silently re-enable it. @@ -1580,7 +1625,7 @@ shell.Run cmd, 0, False # ── Arch resolution: env-var override → name inference ────────────── # Runs even when the hipinfo/amd-smi probe could NOT confirm a runtime # ($HasROCm false): the gfx arch inferred from the WMI GPU name lets the - # studio setup forward --rocm-gfx and pull a GPU-accelerated (lemonade) + # studio setup forward --rocm-gfx and pull a GPU-accelerated ROCm # llama.cpp, which bundles its own ROCm runtime. PyTorch's ROCm wheels # still require a confirmed HIP SDK -- they stay gated on $HasROCm below. if (-not $ROCmGfxArch) { @@ -1591,7 +1636,7 @@ shell.Run cmd, 0, False substep "gfx arch from UNSLOTH_ROCM_GFX_ARCH env override: $ROCmGfxArch" "Cyan" } # 2. Best-effort name → arch lookup from marketing name (amd-smi / WMI). - # Targets only arches the lemonade-sdk ROCm prebuilts cover + # Targets only arches the ROCm prebuilts cover # (gfx120X/110X/1151/1150/103X); unknown names fall back to CPU. elseif ($ROCmGpuLabel) { $nameArchTable = @( @@ -1602,9 +1647,9 @@ shell.Run cmd, 0, False @{ P = "RX 7900|RX 7800|RX 7700(?!S)|PRO W7900|PRO W7800|PRO W7700"; A = "gfx1100" } # RDNA 3 desktop/workstation (Navi 31) @{ P = "RX 7600|RX 7700S|RX 7650|PRO W7600|PRO W7500|PRO V710"; A = "gfx1102" } # RDNA 3 (Navi 33) @{ P = "780M|760M|740M|Phoenix|Hawk Point|Z1 Extreme|Z2 Extreme"; A = "gfx1103" } # RDNA 3 iGPU (Phoenix / Hawk Point) - @{ P = "RX 6900|RX 6800|RX 6750|RX 6700|PRO W6800|PRO W6900"; A = "gfx1030" } # RDNA 2 (Navi 21) -- lemonade gfx103X - @{ P = "RX 6650|RX 6600|PRO W6600|PRO W6650"; A = "gfx1032" } # RDNA 2 (Navi 23) -- lemonade gfx103X - @{ P = "RX 6500|RX 6400|RX 6300|PRO W6400|PRO W6500"; A = "gfx1034" } # RDNA 2 (Navi 24) -- lemonade gfx103X + @{ P = "RX 6900|RX 6800|RX 6750|RX 6700|PRO W6800|PRO W6900"; A = "gfx1030" } # RDNA 2 (Navi 21) -- gfx103X family + @{ P = "RX 6650|RX 6600|PRO W6600|PRO W6650"; A = "gfx1032" } # RDNA 2 (Navi 23) -- gfx103X family + @{ P = "RX 6500|RX 6400|RX 6300|PRO W6400|PRO W6500"; A = "gfx1034" } # RDNA 2 (Navi 24) -- gfx103X family ) foreach ($row in $nameArchTable) { if ($ROCmGpuLabel -match $row.P) { @@ -1881,7 +1926,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.6.3" unsloth-zoo } + $baseInstallExit = Invoke-InstallCommand { uv pip install --python $VenvPython --no-deps --reinstall-package unsloth --reinstall-package unsloth-zoo "unsloth>=2026.6.5" unsloth-zoo } if ($baseInstallExit -eq 0) { # Resolve pydantic WITH deps so pip pins pydantic-core # to the matching version (no-torch-runtime.txt below @@ -1895,7 +1940,7 @@ shell.Run cmd, 0, False } } } else { - $baseInstallExit = Invoke-InstallCommand { uv pip install --python $VenvPython --reinstall-package unsloth --reinstall-package unsloth-zoo "unsloth>=2026.6.3" unsloth-zoo } + $baseInstallExit = Invoke-InstallCommand { uv pip install --python $VenvPython --reinstall-package unsloth --reinstall-package unsloth-zoo "unsloth>=2026.6.5" unsloth-zoo } } if ($baseInstallExit -ne 0) { Write-Host "[ERROR] Failed to install unsloth (exit code $baseInstallExit)" -ForegroundColor Red @@ -1942,7 +1987,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.6.3" unsloth-zoo } + $baseInstallExit = Invoke-InstallCommand { uv pip install --python $VenvPython --no-deps --upgrade-package unsloth --upgrade-package unsloth-zoo "unsloth>=2026.6.5" unsloth-zoo } if ($baseInstallExit -eq 0) { # Same pydantic-with-deps trick as the migrated branch. $baseInstallExit = Invoke-InstallCommand { uv pip install --python $VenvPython pydantic } @@ -1954,7 +1999,7 @@ shell.Run cmd, 0, False } } } elseif ($StudioLocalInstall) { - $baseInstallExit = Invoke-InstallCommand { uv pip install --python $VenvPython --upgrade-package unsloth "unsloth>=2026.6.3" unsloth-zoo } + $baseInstallExit = Invoke-InstallCommand { uv pip install --python $VenvPython --upgrade-package unsloth "unsloth>=2026.6.5" unsloth-zoo } } else { $baseInstallExit = Invoke-InstallCommand { uv pip install --python $VenvPython --upgrade-package unsloth -- "$PackageName" } } @@ -1982,7 +2027,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.6.3" --torch-backend=auto } + $baseInstallExit = Invoke-InstallCommand { uv pip install --python $VenvPython unsloth-zoo "unsloth>=2026.6.5" --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 c7be27942e..9ec46f2013 100755 --- a/install.sh +++ b/install.sh @@ -1456,7 +1456,11 @@ fi # ── Install uv ── tauri_log "STEP" "Installing uv package manager" -UV_MIN_VERSION="0.7.14" +UV_MIN_VERSION="0.7.22" + +# When bytecode compilation is enabled, large installs can exceed uv's 60s default on slow machines. Default to 180s, preserving overrides ("0" disables). +: "${UV_COMPILE_BYTECODE_TIMEOUT:=180}" +export UV_COMPILE_BYTECODE_TIMEOUT version_ge() { # returns 0 if $1 >= $2 @@ -2053,6 +2057,37 @@ _pick_radeon_wheel() { # CPU, non-Strix WSL) skips it and normal detection runs unchanged. NEVER aborts # the installer -- always returns 0. Runs the idempotent helper (ROCm 7.2 + # librocdxg), then sources the env it persisted so detection finds the GPU. +# Export the ROCm-on-WSL env into this process and persist it to /etc/profile.d +# so non-login Studio/llama launches inherit it. Idempotent (writes only when +# the drop-in is missing); no-op without librocdxg, so never fires off WSL. +# /etc/profile.d is root-owned -- sudo-tee when not root, else ROCm vanishes +# after this shell on a non-root reinstall. Best-effort either way. +_persist_rocm_wsl_dropin() { + [ -e /opt/rocm/lib/librocdxg.so ] || [ -e /opt/rocm/lib64/librocdxg.so ] || return 0 + _rw_rocm=/opt/rocm + export HSA_ENABLE_DXG_DETECTION=1 + export TORCH_ROCM_AOTRITON_ENABLE_EXPERIMENTAL=1 + case ":${PATH}:" in + *":${_rw_rocm}/bin:"*) ;; + *) export PATH="${_rw_rocm}/bin:${PATH}" ;; + esac + export LD_LIBRARY_PATH="${_rw_rocm}/lib:${LD_LIBRARY_PATH:-}" + [ -r /etc/profile.d/unsloth-rocm-wsl.sh ] && return 0 + _rw_dropin="$( + printf '# >>> Unsloth ROCm-on-WSL (gfx1151) >>>\n' + printf 'export HSA_ENABLE_DXG_DETECTION=1\n' + printf 'export TORCH_ROCM_AOTRITON_ENABLE_EXPERIMENTAL=1\n' + printf 'export PATH="%s/bin:${PATH}"\n' "${_rw_rocm}" + printf 'export LD_LIBRARY_PATH="%s/lib:${LD_LIBRARY_PATH:-}"\n' "${_rw_rocm}" + printf '# <<< Unsloth ROCm-on-WSL (gfx1151) <<<\n' + )" + if [ "$(id -u)" = "0" ]; then + printf '%s\n' "$_rw_dropin" > /etc/profile.d/unsloth-rocm-wsl.sh 2>/dev/null || true + elif command -v sudo >/dev/null 2>&1; then + printf '%s\n' "$_rw_dropin" | sudo tee /etc/profile.d/unsloth-rocm-wsl.sh >/dev/null 2>&1 || true + fi +} + _maybe_bootstrap_rocm_wsl() { [ "${OS:-}" = "wsl" ] || return 0 [ "${SKIP_TORCH:-false}" = "false" ] || return 0 @@ -2066,6 +2101,11 @@ _maybe_bootstrap_rocm_wsl() { _ensure_rocm_probe_env if command -v rocminfo >/dev/null 2>&1 && \ rocminfo 2>/dev/null | awk '/Name:[[:space:]]*gfx1151/{found=1} END{exit !found}'; then + # rocminfo may work only via the transient env _ensure_rocm_probe_env + # just set, which dies with the installer. Persist the drop-in so login + # shells (Studio, llama.cpp) inherit it -- else a reinstall over an + # existing /opt/rocm (uninstall keeps ROCm but drops it) loses the GPU. + _persist_rocm_wsl_dropin return 0 fi # WSL GPU passthrough device must exist (present on any WSL2 GPU host). @@ -2083,31 +2123,8 @@ _maybe_bootstrap_rocm_wsl() { . /etc/profile.d/unsloth-rocm-wsl.sh || true else # librocdxg present but the env drop-in is gone (e.g. a Studio - # uninstall removed it while keeping shared ROCm). Restore the FULL - # env inline (so rocminfo is on PATH) and recreate the drop-in. - _rw_rocm=/opt/rocm - export HSA_ENABLE_DXG_DETECTION=1 - export TORCH_ROCM_AOTRITON_ENABLE_EXPERIMENTAL=1 - export PATH="${_rw_rocm}/bin:${PATH}" - export LD_LIBRARY_PATH="${_rw_rocm}/lib:${LD_LIBRARY_PATH:-}" - # Persist the drop-in so later non-login Studio launches get the env - # too. /etc/profile.d is root-owned: a plain redirect fails for a - # non-root reinstall (ROCm would silently disappear after this shell), - # so tee through sudo when not root. Best-effort -- the current shell - # already has the env, so the install proceeds either way. - _rw_dropin="$( - printf '# >>> Unsloth ROCm-on-WSL (gfx1151) >>>\n' - printf 'export HSA_ENABLE_DXG_DETECTION=1\n' - printf 'export TORCH_ROCM_AOTRITON_ENABLE_EXPERIMENTAL=1\n' - printf 'export PATH="%s/bin:${PATH}"\n' "${_rw_rocm}" - printf 'export LD_LIBRARY_PATH="%s/lib:${LD_LIBRARY_PATH:-}"\n' "${_rw_rocm}" - printf '# <<< Unsloth ROCm-on-WSL (gfx1151) <<<\n' - )" - if [ "$(id -u)" = "0" ]; then - printf '%s\n' "$_rw_dropin" > /etc/profile.d/unsloth-rocm-wsl.sh 2>/dev/null || true - elif command -v sudo >/dev/null 2>&1; then - printf '%s\n' "$_rw_dropin" | sudo tee /etc/profile.d/unsloth-rocm-wsl.sh >/dev/null 2>&1 || true - fi + # uninstall removed it while keeping shared ROCm). Restore the env. + _persist_rocm_wsl_dropin fi return 0 fi @@ -2415,7 +2432,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.6.3" unsloth-zoo + "unsloth>=2026.6.5" unsloth-zoo # Resolve pydantic WITH deps so pip pins pydantic-core to the # matching version (no-torch-runtime.txt below is --no-deps). # All transitive deps are torch-free. @@ -2428,7 +2445,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.6.3" unsloth-zoo + "unsloth>=2026.6.5" unsloth-zoo fi if [ "$STUDIO_LOCAL_INSTALL" = true ]; then substep "overlaying local repo (editable)..." @@ -2632,7 +2649,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.6.3" unsloth-zoo + "unsloth>=2026.6.5" unsloth-zoo # Same pydantic-with-deps trick as the migrated branch. run_install_cmd "install pydantic (with deps for compatible core)" \ uv pip install --python "$_VENV_PY" pydantic @@ -2650,7 +2667,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.6.3" unsloth-zoo + --upgrade-package unsloth "unsloth>=2026.6.5" 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..." @@ -2682,7 +2699,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.6.3" --torch-backend=auto + run_install_cmd "install unsloth (auto torch backend)" uv pip install --python "$_VENV_PY" unsloth-zoo "unsloth>=2026.6.5" --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..." @@ -2893,7 +2910,10 @@ if [ -t 1 ]; then case "${_reply:-y}" in [Yy]*|"") step "launch" "starting Unsloth Studio..." - "$VENV_DIR/bin/unsloth" studio -p 8888 + # Detach stdin from the `curl | sh` pipe: as a foreground server the + # studio would otherwise drain the rest of this piped script, leaving + # the shell to die parsing the now-truncated tail (`unexpected fi`). + "$VENV_DIR/bin/unsloth" studio -p 8888 =2026.6.3", + "unsloth_zoo>=2026.6.4", "wheel>=0.42.0", "packaging", "numpy", @@ -91,7 +92,7 @@ huggingfacenotorch = [ ] huggingface = [ "unsloth[huggingfacenotorch]", - "unsloth_zoo>=2026.6.3", + "unsloth_zoo>=2026.6.4", "torchvision", "unsloth[triton]", ] @@ -581,7 +582,7 @@ colab-ampere-torch220 = [ "flash-attn>=2.6.3 ; ('linux' in sys_platform)", ] colab-new = [ - "unsloth_zoo>=2026.6.3", + "unsloth_zoo>=2026.6.4", "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/assets/chat_templates/gemma-4-edge.jinja b/studio/backend/assets/chat_templates/gemma-4-edge.jinja new file mode 100644 index 0000000000..0266127233 --- /dev/null +++ b/studio/backend/assets/chat_templates/gemma-4-edge.jinja @@ -0,0 +1,397 @@ +{#- + Gemma 4 chat template (E2B / E4B edge variant), vendored for Unsloth Studio. + Source: google/gemma-4-31B-it HF discussion/PR #118 (adds the preserve_thinking + flag plus null-rendering, string-arguments validation, balanced turn tags, empty + messages handling, and OpenAI image_url/input_audio aliases). + Studio-local changes vs PR #118: + 1. preserve_thinking defaults to false (see SETUP block below). + 2. The empty "<|channel>thought\n" block on enable_thinking=false is + NOT emitted. Google ships a distinct template for E2B/E4B (google/gemma-4-E2B-it, + google/gemma-4-E4B-it) that omits it; only the 12b/26B-A4B/31B family emits it. + This file matches the E2B/E4B behavior; gemma-4.jinja keeps the larger-model one. + Applied to unsloth/gemma-4-E2B-it-GGUF and unsloth/gemma-4-E4B-it-GGUF so the + embedded GGUF template does not need re-downloading. +-#} +{%- macro format_parameters(properties, required, filter_keys=false) -%} + {%- set standard_keys = ['description', 'type', 'properties', 'required', 'nullable'] -%} + {%- set ns = namespace(found_first=false) -%} + {%- for key, value in properties | dictsort -%} + {%- set add_comma = false -%} + {%- if not filter_keys or key not in standard_keys -%} + {%- if ns.found_first %},{% endif -%} + {%- set ns.found_first = true -%} + {{ key }}:{ + {%- if value['description'] -%} + description:<|"|>{{ value['description'] }}<|"|> + {%- set add_comma = true -%} + {%- endif -%} + {%- if value['type'] | upper == 'STRING' -%} + {%- if value['enum'] -%} + {%- if add_comma %},{%- else -%} {%- set add_comma = true -%} {% endif -%} + enum:{{ format_argument(value['enum']) }} + {%- endif -%} + {%- elif value['type'] | upper == 'ARRAY' -%} + {%- if value['items'] is mapping and value['items'] -%} + {%- if add_comma %},{%- else -%} {%- set add_comma = true -%} {% endif -%} + items:{ + {%- set ns_items = namespace(found_first=false) -%} + {%- for item_key, item_value in value['items'] | dictsort -%} + {%- if item_value is not none -%} + {%- if ns_items.found_first %},{% endif -%} + {%- set ns_items.found_first = true -%} + {%- if item_key == 'properties' -%} + properties:{ + {%- if item_value is mapping -%} + {{- format_parameters(item_value, value['items']['required'] | default([])) -}} + {%- endif -%} + } + {%- elif item_key == 'required' -%} + required:[ + {%- for req_item in item_value -%} + <|"|>{{- req_item -}}<|"|> + {%- if not loop.last %},{% endif -%} + {%- endfor -%} + ] + {%- elif item_key == 'type' -%} + {%- if item_value is string -%} + type:{{ format_argument(item_value | upper) }} + {%- else -%} + type:{{ format_argument(item_value | map('upper') | list) }} + {%- endif -%} + {%- else -%} + {{ item_key }}:{{ format_argument(item_value) }} + {%- endif -%} + {%- endif -%} + {%- endfor -%} + } + {%- endif -%} + {%- endif -%} + {%- if value['nullable'] %} + {%- if add_comma %},{%- else -%} {%- set add_comma = true -%} {% endif -%} + nullable:true + {%- endif -%} + {%- if value['type'] | upper == 'OBJECT' -%} + {%- if value['properties'] is defined and value['properties'] is mapping -%} + {%- if add_comma %},{%- else -%} {%- set add_comma = true -%} {% endif -%} + properties:{ + {{- format_parameters(value['properties'], value['required'] | default([])) -}} + } + {%- elif value is mapping -%} + {%- if add_comma %},{%- else -%} {%- set add_comma = true -%} {% endif -%} + properties:{ + {{- format_parameters(value, value['required'] | default([]), filter_keys=true) -}} + } + {%- endif -%} + {%- if value['required'] -%} + {%- if add_comma %},{%- else -%} {%- set add_comma = true -%} {% endif -%} + required:[ + {%- for item in value['required'] | default([]) -%} + <|"|>{{- item -}}<|"|> + {%- if not loop.last %},{% endif -%} + {%- endfor -%} + ] + {%- endif -%} + {%- endif -%} + {%- if add_comma %},{%- else -%} {%- set add_comma = true -%} {% endif -%} + type:<|"|>{{ value['type'] | upper }}<|"|>} + {%- endif -%} + {%- endfor -%} +{%- endmacro -%} +{%- macro format_function_declaration(tool_data) -%} + declaration:{{- tool_data['function']['name'] -}}{description:<|"|>{{- tool_data['function']['description'] -}}<|"|> + {%- set params = tool_data['function']['parameters'] -%} + {%- if params -%} + ,parameters:{ + {%- if params['properties'] -%} + properties:{ {{- format_parameters(params['properties'], params['required']) -}} }, + {%- endif -%} + {%- if params['required'] -%} + required:[ + {%- for item in params['required'] -%} + <|"|>{{- item -}}<|"|> + {{- ',' if not loop.last -}} + {%- endfor -%} + ], + {%- endif -%} + {%- if params['type'] -%} + type:<|"|>{{- params['type'] | upper -}}<|"|>} + {%- endif -%} + {%- endif -%} + {%- if 'response' in tool_data['function'] -%} + {%- set response_declaration = tool_data['function']['response'] -%} + ,response:{ + {%- if response_declaration['description'] -%} + description:<|"|>{{- response_declaration['description'] -}}<|"|>, + {%- endif -%} + {%- if response_declaration['type'] | upper == 'OBJECT' -%} + type:<|"|>{{- response_declaration['type'] | upper -}}<|"|>} + {%- endif -%} + {%- endif -%} + } +{%- endmacro -%} +{%- macro format_argument(argument, escape_keys=True) -%} + {%- if argument is none -%} + {{- 'null' -}} + {%- elif argument is string -%} + {{- '<|"|>' + argument + '<|"|>' -}} + {%- elif argument is boolean -%} + {{- 'true' if argument else 'false' -}} + {%- elif argument is mapping -%} + {{- '{' -}} + {%- set ns = namespace(found_first=false) -%} + {%- for key, value in argument | dictsort -%} + {%- if ns.found_first %},{% endif -%} + {%- set ns.found_first = true -%} + {%- if escape_keys -%} + {{- '<|"|>' + key + '<|"|>' -}} + {%- else -%} + {{- key -}} + {%- endif -%} + :{{- format_argument(value, escape_keys=escape_keys) -}} + {%- endfor -%} + {{- '}' -}} + {%- elif argument is sequence -%} + {{- '[' -}} + {%- for item in argument -%} + {{- format_argument(item, escape_keys=escape_keys) -}} + {%- if not loop.last %},{% endif -%} + {%- endfor -%} + {{- ']' -}} + {%- else -%} + {{- argument -}} + {%- endif -%} +{%- endmacro -%} +{%- macro strip_thinking(text) -%} + {%- set ns = namespace(result='') -%} + {%- for part in text.split('') -%} + {%- if '<|channel>' in part -%} + {%- set ns.result = ns.result + part.split('<|channel>')[0] -%} + {%- else -%} + {%- set ns.result = ns.result + part -%} + {%- endif -%} + {%- endfor -%} + {{- ns.result | trim -}} +{%- endmacro -%} + +{%- macro format_tool_response_block(tool_name, response) -%} + {{- '<|tool_response>' -}} + {%- if response is mapping -%} + {{- 'response:' + tool_name + '{' -}} + {%- for key, value in response | dictsort -%} + {{- key -}}:{{- format_argument(value, escape_keys=False) -}} + {%- if not loop.last %},{% endif -%} + {%- endfor -%} + {{- '}' -}} + {%- else -%} + {{- 'response:' + tool_name + '{value:' + format_argument(response, escape_keys=False) + '}' -}} + {%- endif -%} + {{- '' -}} +{%- endmacro -%} + +{#- ===== SETUP ===== -#} +{%- set ns = namespace(prev_message_type=None, prev_non_tool_role=None) -%} +{%- set loop_messages = messages -%} +{%- set enable_thinking = enable_thinking | default(false) -%} +{#- Unsloth Studio: preserve_thinking defaults OFF (upstream PR #118 defaults true). -#} +{%- set preserve_thinking = preserve_thinking | default(false) -%} +{{- bos_token -}} +{#- Handle System/Tool Definitions Block -#} +{%- if enable_thinking or tools or (messages and messages[0]['role'] in ['system', 'developer']) -%} + {{- '<|turn>system\n' -}} + {#- Inject Thinking token at the very top of the FIRST system turn -#} + {%- if enable_thinking -%} + {{- '<|think|>\n' -}} + {%- set ns.prev_message_type = 'think' -%} + {%- endif -%} + {%- if messages and messages[0]['role'] in ['system', 'developer'] -%} + {%- if messages[0]['content'] is string -%} + {{- messages[0]['content'] | trim -}} + {%- elif messages[0]['content'] is sequence -%} + {%- for item in messages[0]['content'] -%} + {{- item['text'] | trim + ' '-}} + {%- endfor -%} + {%- endif -%} + {%- set loop_messages = messages[1:] -%} + {%- endif -%} + {%- if tools -%} + {%- for tool in tools %} + {{- '<|tool>' -}} + {{- format_function_declaration(tool) | trim -}} + {{- '' -}} + {%- endfor %} + {%- set ns.prev_message_type = 'tool' -%} + {%- endif -%} + {{- '\n' -}} +{%- endif %} + +{#- Pre-scan: find last user message index for reasoning guard -#} +{%- set ns_turn = namespace(last_user_idx=-1) -%} +{%- for i in range(loop_messages | length) -%} + {%- if loop_messages[i]['role'] == 'user' -%} + {%- set ns_turn.last_user_idx = i -%} + {%- endif -%} +{%- endfor -%} + +{#- Loop through messages -#} +{%- for message in loop_messages -%} + {%- if message['role'] != 'tool' -%} + {%- set ns.prev_message_type = None -%} + {%- set role = 'model' if message['role'] == 'assistant' else message['role'] -%} + {#- Detect continuation using tracked state - O(1) instead of O(n) backward scan -#} + {%- set continue_same_model_turn = (role == 'model' and ns.prev_non_tool_role == 'assistant') -%} + {%- if not continue_same_model_turn -%} + {{- '<|turn>' + role + '\n' }} + {%- endif -%} + + {#- Render reasoning/reasoning_content as thinking channel (tool-call turns only) -#} + {%- set thinking_text = message.get('reasoning') or message.get('reasoning_content') -%} + {%- set thinking_gate = (loop.index0 > ns_turn.last_user_idx) or preserve_thinking -%} + {%- if thinking_text and thinking_gate and message.get('tool_calls') -%} + {{- '<|channel>thought\n' + thinking_text + '\n' -}} + {%- endif -%} + + {%- if message.get('tool_calls') -%} + {%- for tool_call in message.get('tool_calls') -%} + {%- set function = tool_call['function'] -%} + {{- '<|tool_call>call:' + function['name'] + '{' -}} + {%- if function['arguments'] is mapping -%} + {%- set ns_args = namespace(found_first=false) -%} + {%- for key, value in function['arguments'] | dictsort -%} + {%- if ns_args.found_first %},{% endif -%} + {%- set ns_args.found_first = true -%} + {{- key -}}:{{- format_argument(value, escape_keys=False) -}} + {%- endfor -%} + {%- elif function['arguments'] is none -%} + {%- else -%} + {{- raise_exception( + "chat_template: tool_calls[].function.arguments must be a " + "JSON object (mapping), not a string. Deserialize arguments " + "before passing to the template." + ) -}} + {%- endif -%} + {{- '}' -}} + {%- endfor -%} + {%- set ns.prev_message_type = 'tool_call' -%} + {%- endif -%} + + {%- set ns_tr_out = namespace(flag=false) -%} + {%- if message.get('tool_responses') -%} + {#- Legacy: tool_responses embedded on the assistant message (Google/Gemma native) -#} + {%- for tool_response in message.get('tool_responses') -%} + {{- format_tool_response_block(tool_response['name'] | default('unknown', true), tool_response['response']) -}} + {%- set ns_tr_out.flag = true -%} + {%- set ns.prev_message_type = 'tool_response' -%} + {%- endfor -%} + {%- elif message.get('tool_calls') -%} + {#- OpenAI Chat Completions: forward-scan consecutive role:tool messages -#} + {%- set ns_tool_scan = namespace(stopped=false) -%} + {%- for k in range(loop.index0 + 1, loop_messages | length) -%} + {%- if ns_tool_scan.stopped -%} + {%- elif loop_messages[k]['role'] != 'tool' -%} + {%- set ns_tool_scan.stopped = true -%} + {%- else -%} + {%- set follow = loop_messages[k] -%} + {#- Resolve tool_call_id to function name -#} + {%- set ns_tname = namespace(name=follow.get('name') or 'unknown') -%} + {%- for tc in message.get('tool_calls') -%} + {%- if tc.get('id') == follow.get('tool_call_id') -%} + {%- set ns_tname.name = tc['function']['name'] -%} + {%- endif -%} + {%- endfor -%} + {#- Handle content as string or content-parts array -#} + {%- set tool_body = follow.get('content') -%} + {%- if tool_body is string -%} + {{- format_tool_response_block(ns_tname.name, tool_body) -}} + {%- elif tool_body is sequence and tool_body is not string -%} + {%- set ns_txt = namespace(s='') -%} + {%- for part in tool_body -%} + {%- if part.get('type') == 'text' -%} + {%- set ns_txt.s = ns_txt.s + (part.get('text') | default('')) -%} + {%- endif -%} + {%- endfor -%} + {{- format_tool_response_block(ns_tname.name, ns_txt.s) -}} + {%- for part in tool_body -%} + {%- if part.get('type') in ['image', 'image_url'] -%} + {{- '<|image|>' -}} + {%- elif part.get('type') in ['audio', 'input_audio'] -%} + {{- '<|audio|>' -}} + {%- elif part.get('type') == 'video' -%} + {{- '<|video|>' -}} + {%- endif -%} + {%- endfor -%} + {%- else -%} + {{- format_tool_response_block(ns_tname.name, tool_body) -}} + {%- endif -%} + {%- set ns_tr_out.flag = true -%} + {%- set ns.prev_message_type = 'tool_response' -%} + {%- endif -%} + {%- endfor -%} + {%- endif -%} + + {%- set captured_content -%} + {%- if message.get('content') is string -%} + {%- if role == 'model' -%} + {{- strip_thinking(message['content']) -}} + {%- else -%} + {{- message['content'] | trim -}} + {%- endif -%} + {%- elif message.get('content') is sequence -%} + {%- for item in message['content'] -%} + {%- if item.get('type') == 'text' -%} + {%- if role == 'model' -%} + {{- strip_thinking(item['text']) -}} + {%- else -%} + {{- item['text'] | trim -}} + {%- endif -%} + {%- elif item.get('type') in ['image', 'image_url'] -%} + {{- '<|image|>' -}} + {%- elif item.get('type') in ['audio', 'input_audio'] -%} + {{- '<|audio|>' -}} + {%- elif item.get('type') == 'video' -%} + {{- '<|video|>' -}} + {%- endif -%} + {%- endfor -%} + {%- endif -%} + {%- endset -%} + + {{- captured_content -}} + {%- set has_content = captured_content | trim | length > 0 -%} + + {#- Forward-scan: find next non-tool message role for continuation detection -#} + {%- set next_nt = namespace(role=None, found=false) -%} + {%- for j in range(loop.index0 + 1, loop_messages | length) -%} + {%- if not next_nt.found -%} + {%- if loop_messages[j]['role'] != 'tool' -%} + {%- set next_nt.role = loop_messages[j]['role'] -%} + {%- set next_nt.found = true -%} + {%- endif -%} + {%- endif -%} + {%- endfor -%} + + {%- set continues_into_next = ( + role == 'model' + and next_nt.role == 'assistant' + and not message.get('tool_calls') + and not ns_tr_out.flag + ) -%} + + {%- if ns.prev_message_type == 'tool_call' and not ns_tr_out.flag -%} + {{- '<|tool_response>' -}} + {%- elif continues_into_next -%} + {{- '\n' -}} + {%- elif not (ns_tr_out.flag and not has_content) -%} + {{- '\n' -}} + {%- endif -%} + + {#- Track previous non-tool role for next iteration (avoids O(n) backward scan) -#} + {%- set ns.prev_non_tool_role = message['role'] -%} + {%- endif -%} +{%- endfor -%} + +{%- if add_generation_prompt -%} + {%- if ns.prev_message_type != 'tool_response' and ns.prev_message_type != 'tool_call' -%} + {{- '<|turn>model\n' -}} + {%- endif -%} + {#- E2B/E4B do NOT emit an empty thought block when enable_thinking is false + (unlike the 12b/26B-A4B/31B family); see header. -#} +{%- endif -%} diff --git a/studio/backend/assets/chat_templates/gemma-4.jinja b/studio/backend/assets/chat_templates/gemma-4.jinja new file mode 100644 index 0000000000..65ab39df57 --- /dev/null +++ b/studio/backend/assets/chat_templates/gemma-4.jinja @@ -0,0 +1,397 @@ +{#- + Gemma 4 chat template, vendored for Unsloth Studio. + Source: google/gemma-4-31B-it HF discussion/PR #118 (adds the preserve_thinking + flag plus null-rendering, string-arguments validation, balanced turn tags, empty + messages handling, and OpenAI image_url/input_audio aliases). + Studio-local change: preserve_thinking defaults to false (see SETUP block below). + Applied to unsloth/gemma-4-*-GGUF models so the embedded GGUF template does not + need re-downloading. Keep in sync with upstream if PR #118 changes. +-#} +{%- macro format_parameters(properties, required, filter_keys=false) -%} + {%- set standard_keys = ['description', 'type', 'properties', 'required', 'nullable'] -%} + {%- set ns = namespace(found_first=false) -%} + {%- for key, value in properties | dictsort -%} + {%- set add_comma = false -%} + {%- if not filter_keys or key not in standard_keys -%} + {%- if ns.found_first %},{% endif -%} + {%- set ns.found_first = true -%} + {{ key }}:{ + {%- if value['description'] -%} + description:<|"|>{{ value['description'] }}<|"|> + {%- set add_comma = true -%} + {%- endif -%} + {%- if value['type'] | upper == 'STRING' -%} + {%- if value['enum'] -%} + {%- if add_comma %},{%- else -%} {%- set add_comma = true -%} {% endif -%} + enum:{{ format_argument(value['enum']) }} + {%- endif -%} + {%- elif value['type'] | upper == 'ARRAY' -%} + {%- if value['items'] is mapping and value['items'] -%} + {%- if add_comma %},{%- else -%} {%- set add_comma = true -%} {% endif -%} + items:{ + {%- set ns_items = namespace(found_first=false) -%} + {%- for item_key, item_value in value['items'] | dictsort -%} + {%- if item_value is not none -%} + {%- if ns_items.found_first %},{% endif -%} + {%- set ns_items.found_first = true -%} + {%- if item_key == 'properties' -%} + properties:{ + {%- if item_value is mapping -%} + {{- format_parameters(item_value, value['items']['required'] | default([])) -}} + {%- endif -%} + } + {%- elif item_key == 'required' -%} + required:[ + {%- for req_item in item_value -%} + <|"|>{{- req_item -}}<|"|> + {%- if not loop.last %},{% endif -%} + {%- endfor -%} + ] + {%- elif item_key == 'type' -%} + {%- if item_value is string -%} + type:{{ format_argument(item_value | upper) }} + {%- else -%} + type:{{ format_argument(item_value | map('upper') | list) }} + {%- endif -%} + {%- else -%} + {{ item_key }}:{{ format_argument(item_value) }} + {%- endif -%} + {%- endif -%} + {%- endfor -%} + } + {%- endif -%} + {%- endif -%} + {%- if value['nullable'] %} + {%- if add_comma %},{%- else -%} {%- set add_comma = true -%} {% endif -%} + nullable:true + {%- endif -%} + {%- if value['type'] | upper == 'OBJECT' -%} + {%- if value['properties'] is defined and value['properties'] is mapping -%} + {%- if add_comma %},{%- else -%} {%- set add_comma = true -%} {% endif -%} + properties:{ + {{- format_parameters(value['properties'], value['required'] | default([])) -}} + } + {%- elif value is mapping -%} + {%- if add_comma %},{%- else -%} {%- set add_comma = true -%} {% endif -%} + properties:{ + {{- format_parameters(value, value['required'] | default([]), filter_keys=true) -}} + } + {%- endif -%} + {%- if value['required'] -%} + {%- if add_comma %},{%- else -%} {%- set add_comma = true -%} {% endif -%} + required:[ + {%- for item in value['required'] | default([]) -%} + <|"|>{{- item -}}<|"|> + {%- if not loop.last %},{% endif -%} + {%- endfor -%} + ] + {%- endif -%} + {%- endif -%} + {%- if add_comma %},{%- else -%} {%- set add_comma = true -%} {% endif -%} + type:<|"|>{{ value['type'] | upper }}<|"|>} + {%- endif -%} + {%- endfor -%} +{%- endmacro -%} +{%- macro format_function_declaration(tool_data) -%} + declaration:{{- tool_data['function']['name'] -}}{description:<|"|>{{- tool_data['function']['description'] -}}<|"|> + {%- set params = tool_data['function']['parameters'] -%} + {%- if params -%} + ,parameters:{ + {%- if params['properties'] -%} + properties:{ {{- format_parameters(params['properties'], params['required']) -}} }, + {%- endif -%} + {%- if params['required'] -%} + required:[ + {%- for item in params['required'] -%} + <|"|>{{- item -}}<|"|> + {{- ',' if not loop.last -}} + {%- endfor -%} + ], + {%- endif -%} + {%- if params['type'] -%} + type:<|"|>{{- params['type'] | upper -}}<|"|>} + {%- endif -%} + {%- endif -%} + {%- if 'response' in tool_data['function'] -%} + {%- set response_declaration = tool_data['function']['response'] -%} + ,response:{ + {%- if response_declaration['description'] -%} + description:<|"|>{{- response_declaration['description'] -}}<|"|>, + {%- endif -%} + {%- if response_declaration['type'] | upper == 'OBJECT' -%} + type:<|"|>{{- response_declaration['type'] | upper -}}<|"|>} + {%- endif -%} + {%- endif -%} + } +{%- endmacro -%} +{%- macro format_argument(argument, escape_keys=True) -%} + {%- if argument is none -%} + {{- 'null' -}} + {%- elif argument is string -%} + {{- '<|"|>' + argument + '<|"|>' -}} + {%- elif argument is boolean -%} + {{- 'true' if argument else 'false' -}} + {%- elif argument is mapping -%} + {{- '{' -}} + {%- set ns = namespace(found_first=false) -%} + {%- for key, value in argument | dictsort -%} + {%- if ns.found_first %},{% endif -%} + {%- set ns.found_first = true -%} + {%- if escape_keys -%} + {{- '<|"|>' + key + '<|"|>' -}} + {%- else -%} + {{- key -}} + {%- endif -%} + :{{- format_argument(value, escape_keys=escape_keys) -}} + {%- endfor -%} + {{- '}' -}} + {%- elif argument is sequence -%} + {{- '[' -}} + {%- for item in argument -%} + {{- format_argument(item, escape_keys=escape_keys) -}} + {%- if not loop.last %},{% endif -%} + {%- endfor -%} + {{- ']' -}} + {%- else -%} + {{- argument -}} + {%- endif -%} +{%- endmacro -%} +{%- macro strip_thinking(text) -%} + {%- set ns = namespace(result='') -%} + {%- for part in text.split('') -%} + {%- if '<|channel>' in part -%} + {%- set ns.result = ns.result + part.split('<|channel>')[0] -%} + {%- else -%} + {%- set ns.result = ns.result + part -%} + {%- endif -%} + {%- endfor -%} + {{- ns.result | trim -}} +{%- endmacro -%} + +{%- macro format_tool_response_block(tool_name, response) -%} + {{- '<|tool_response>' -}} + {%- if response is mapping -%} + {{- 'response:' + tool_name + '{' -}} + {%- for key, value in response | dictsort -%} + {{- key -}}:{{- format_argument(value, escape_keys=False) -}} + {%- if not loop.last %},{% endif -%} + {%- endfor -%} + {{- '}' -}} + {%- else -%} + {{- 'response:' + tool_name + '{value:' + format_argument(response, escape_keys=False) + '}' -}} + {%- endif -%} + {{- '' -}} +{%- endmacro -%} + +{#- ===== SETUP ===== -#} +{%- set ns = namespace(prev_message_type=None, prev_non_tool_role=None) -%} +{%- set loop_messages = messages -%} +{%- set enable_thinking = enable_thinking | default(false) -%} +{#- Unsloth Studio: preserve_thinking defaults OFF (upstream PR #118 defaults true). -#} +{%- set preserve_thinking = preserve_thinking | default(false) -%} +{{- bos_token -}} +{#- Handle System/Tool Definitions Block -#} +{%- if enable_thinking or tools or (messages and messages[0]['role'] in ['system', 'developer']) -%} + {{- '<|turn>system\n' -}} + {#- Inject Thinking token at the very top of the FIRST system turn -#} + {%- if enable_thinking -%} + {{- '<|think|>\n' -}} + {%- set ns.prev_message_type = 'think' -%} + {%- endif -%} + {%- if messages and messages[0]['role'] in ['system', 'developer'] -%} + {%- if messages[0]['content'] is string -%} + {{- messages[0]['content'] | trim -}} + {%- elif messages[0]['content'] is sequence -%} + {%- for item in messages[0]['content'] -%} + {{- item['text'] | trim + ' '-}} + {%- endfor -%} + {%- endif -%} + {%- set loop_messages = messages[1:] -%} + {%- endif -%} + {%- if tools -%} + {%- for tool in tools %} + {{- '<|tool>' -}} + {{- format_function_declaration(tool) | trim -}} + {{- '' -}} + {%- endfor %} + {%- set ns.prev_message_type = 'tool' -%} + {%- endif -%} + {{- '\n' -}} +{%- endif %} + +{#- Pre-scan: find last user message index for reasoning guard -#} +{%- set ns_turn = namespace(last_user_idx=-1) -%} +{%- for i in range(loop_messages | length) -%} + {%- if loop_messages[i]['role'] == 'user' -%} + {%- set ns_turn.last_user_idx = i -%} + {%- endif -%} +{%- endfor -%} + +{#- Loop through messages -#} +{%- for message in loop_messages -%} + {%- if message['role'] != 'tool' -%} + {%- set ns.prev_message_type = None -%} + {%- set role = 'model' if message['role'] == 'assistant' else message['role'] -%} + {#- Detect continuation using tracked state - O(1) instead of O(n) backward scan -#} + {%- set continue_same_model_turn = (role == 'model' and ns.prev_non_tool_role == 'assistant') -%} + {%- if not continue_same_model_turn -%} + {{- '<|turn>' + role + '\n' }} + {%- endif -%} + + {#- Render reasoning/reasoning_content as thinking channel (tool-call turns only) -#} + {%- set thinking_text = message.get('reasoning') or message.get('reasoning_content') -%} + {%- set thinking_gate = (loop.index0 > ns_turn.last_user_idx) or preserve_thinking -%} + {%- if thinking_text and thinking_gate and message.get('tool_calls') -%} + {{- '<|channel>thought\n' + thinking_text + '\n' -}} + {%- endif -%} + + {%- if message.get('tool_calls') -%} + {%- for tool_call in message.get('tool_calls') -%} + {%- set function = tool_call['function'] -%} + {{- '<|tool_call>call:' + function['name'] + '{' -}} + {%- if function['arguments'] is mapping -%} + {%- set ns_args = namespace(found_first=false) -%} + {%- for key, value in function['arguments'] | dictsort -%} + {%- if ns_args.found_first %},{% endif -%} + {%- set ns_args.found_first = true -%} + {{- key -}}:{{- format_argument(value, escape_keys=False) -}} + {%- endfor -%} + {%- elif function['arguments'] is none -%} + {%- else -%} + {{- raise_exception( + "chat_template: tool_calls[].function.arguments must be a " + "JSON object (mapping), not a string. Deserialize arguments " + "before passing to the template." + ) -}} + {%- endif -%} + {{- '}' -}} + {%- endfor -%} + {%- set ns.prev_message_type = 'tool_call' -%} + {%- endif -%} + + {%- set ns_tr_out = namespace(flag=false) -%} + {%- if message.get('tool_responses') -%} + {#- Legacy: tool_responses embedded on the assistant message (Google/Gemma native) -#} + {%- for tool_response in message.get('tool_responses') -%} + {{- format_tool_response_block(tool_response['name'] | default('unknown', true), tool_response['response']) -}} + {%- set ns_tr_out.flag = true -%} + {%- set ns.prev_message_type = 'tool_response' -%} + {%- endfor -%} + {%- elif message.get('tool_calls') -%} + {#- OpenAI Chat Completions: forward-scan consecutive role:tool messages -#} + {%- set ns_tool_scan = namespace(stopped=false) -%} + {%- for k in range(loop.index0 + 1, loop_messages | length) -%} + {%- if ns_tool_scan.stopped -%} + {%- elif loop_messages[k]['role'] != 'tool' -%} + {%- set ns_tool_scan.stopped = true -%} + {%- else -%} + {%- set follow = loop_messages[k] -%} + {#- Resolve tool_call_id to function name -#} + {%- set ns_tname = namespace(name=follow.get('name') or 'unknown') -%} + {%- for tc in message.get('tool_calls') -%} + {%- if tc.get('id') == follow.get('tool_call_id') -%} + {%- set ns_tname.name = tc['function']['name'] -%} + {%- endif -%} + {%- endfor -%} + {#- Handle content as string or content-parts array -#} + {%- set tool_body = follow.get('content') -%} + {%- if tool_body is string -%} + {{- format_tool_response_block(ns_tname.name, tool_body) -}} + {%- elif tool_body is sequence and tool_body is not string -%} + {%- set ns_txt = namespace(s='') -%} + {%- for part in tool_body -%} + {%- if part.get('type') == 'text' -%} + {%- set ns_txt.s = ns_txt.s + (part.get('text') | default('')) -%} + {%- endif -%} + {%- endfor -%} + {{- format_tool_response_block(ns_tname.name, ns_txt.s) -}} + {%- for part in tool_body -%} + {%- if part.get('type') in ['image', 'image_url'] -%} + {{- '<|image|>' -}} + {%- elif part.get('type') in ['audio', 'input_audio'] -%} + {{- '<|audio|>' -}} + {%- elif part.get('type') == 'video' -%} + {{- '<|video|>' -}} + {%- endif -%} + {%- endfor -%} + {%- else -%} + {{- format_tool_response_block(ns_tname.name, tool_body) -}} + {%- endif -%} + {%- set ns_tr_out.flag = true -%} + {%- set ns.prev_message_type = 'tool_response' -%} + {%- endif -%} + {%- endfor -%} + {%- endif -%} + + {%- set captured_content -%} + {%- if message.get('content') is string -%} + {%- if role == 'model' -%} + {{- strip_thinking(message['content']) -}} + {%- else -%} + {{- message['content'] | trim -}} + {%- endif -%} + {%- elif message.get('content') is sequence -%} + {%- for item in message['content'] -%} + {%- if item.get('type') == 'text' -%} + {%- if role == 'model' -%} + {{- strip_thinking(item['text']) -}} + {%- else -%} + {{- item['text'] | trim -}} + {%- endif -%} + {%- elif item.get('type') in ['image', 'image_url'] -%} + {{- '<|image|>' -}} + {%- elif item.get('type') in ['audio', 'input_audio'] -%} + {{- '<|audio|>' -}} + {%- elif item.get('type') == 'video' -%} + {{- '<|video|>' -}} + {%- endif -%} + {%- endfor -%} + {%- endif -%} + {%- endset -%} + + {{- captured_content -}} + {%- set has_content = captured_content | trim | length > 0 -%} + + {#- Forward-scan: find next non-tool message role for continuation detection -#} + {%- set next_nt = namespace(role=None, found=false) -%} + {%- for j in range(loop.index0 + 1, loop_messages | length) -%} + {%- if not next_nt.found -%} + {%- if loop_messages[j]['role'] != 'tool' -%} + {%- set next_nt.role = loop_messages[j]['role'] -%} + {%- set next_nt.found = true -%} + {%- endif -%} + {%- endif -%} + {%- endfor -%} + + {%- set continues_into_next = ( + role == 'model' + and next_nt.role == 'assistant' + and not message.get('tool_calls') + and not ns_tr_out.flag + ) -%} + + {%- if ns.prev_message_type == 'tool_call' and not ns_tr_out.flag -%} + {{- '<|tool_response>' -}} + {%- elif continues_into_next -%} + {{- '\n' -}} + {%- elif not (ns_tr_out.flag and not has_content) -%} + {{- '\n' -}} + {%- endif -%} + + {#- Track previous non-tool role for next iteration (avoids O(n) backward scan) -#} + {%- set ns.prev_non_tool_role = message['role'] -%} + {%- endif -%} +{%- endfor -%} + +{%- if add_generation_prompt -%} + {%- if ns.prev_message_type != 'tool_response' and ns.prev_message_type != 'tool_call' -%} + {{- '<|turn>model\n' -}} + {%- endif -%} + + {%- if not enable_thinking -%} + {#- Suppress thinking - but not when awaiting tool responses -#} + {%- if ns.prev_message_type != 'tool_call' -%} + {{- '<|channel>thought\n' -}} + {%- endif -%} + {%- endif -%} +{%- endif -%} diff --git a/studio/backend/cloudflare_tunnel.py b/studio/backend/cloudflare_tunnel.py index 7c088b6e1e..e5dba69452 100644 --- a/studio/backend/cloudflare_tunnel.py +++ b/studio/backend/cloudflare_tunnel.py @@ -24,12 +24,20 @@ from pathlib import Path from typing import Optional, Tuple # cloudflared logs the quick-tunnel URL; match only the URL so we do not depend -# on the surrounding wording, which Cloudflare may change. -_URL_RE = re.compile(r"https://[A-Za-z0-9-]+\.trycloudflare\.com") +# on the surrounding wording, which Cloudflare may change. The negative lookahead +# drops cloudflared's own API host, which appears in failure lines such as +# failed to request quick Tunnel: Post "https://api.trycloudflare.com/tunnel" +# and must never be mistaken for a usable tunnel URL. +_URL_RE = re.compile(r"https://(?!api\.)[A-Za-z0-9-]+\.trycloudflare\.com") + +# cloudflared logs this once per edge connection it establishes. Until at least +# one appears the quick-tunnel URL returns Cloudflare error 1033 (HTTP 530), so +# we wait for it before advertising the URL. +_REGISTERED_MARKER = "Registered tunnel connection" _RELEASE_BASE = "https://github.com/cloudflare/cloudflared/releases/latest/download" -_URL_TIMEOUT = 15.0 # seconds to wait for the public URL before giving up +_READY_TIMEOUT = 15.0 # seconds to wait for the URL + a registered edge connection _DOWNLOAD_TIMEOUT = 60 # urlopen timeout for the one-time binary download @@ -180,13 +188,24 @@ class CloudflareTunnel: upstream stays local-only. """ - def __init__(self, port: int, binary: str): + def __init__( + self, + port: int, + binary: str, + protocol: Optional[str] = None, + ): self.port = port self.binary = binary + # None lets cloudflared pick its default (quic, with its own http2 + # fallback); set to "http2" to force it when quic is blocked. + self.protocol = protocol self._proc: Optional[subprocess.Popen] = None self._lock = threading.Lock() + self._stopped = False self._url_event = threading.Event() + self._ready_event = threading.Event() self.url: Optional[str] = None + self.ready = False self.error: Optional[str] = None def start(self) -> None: @@ -197,25 +216,33 @@ class CloudflareTunnel: f"http://localhost:{self.port}", "--no-autoupdate", ] - proc = subprocess.Popen( - cmd, - stdout = subprocess.PIPE, - stderr = subprocess.STDOUT, - stdin = subprocess.DEVNULL, - text = True, - errors = "replace", - bufsize = 1, - **_windows_hidden_kwargs(), - ) + if self.protocol: + cmd += ["--protocol", self.protocol] with self._lock: + # A stop() that landed before us (e.g. a shutdown in the caller's + # register->start window) marks the tunnel stopped; spawning now would + # orphan a process nobody owns, so refuse. + if self._stopped: + return + proc = subprocess.Popen( + cmd, + stdout = subprocess.PIPE, + stderr = subprocess.STDOUT, + stdin = subprocess.DEVNULL, + text = True, + errors = "replace", + bufsize = 1, + **_windows_hidden_kwargs(), + ) self._proc = proc threading.Thread( target = self._reader, args = (proc,), name = "cloudflared-reader", daemon = True ).start() def _reader(self, proc: subprocess.Popen) -> None: - # Drain cloudflared's output, capture the first trycloudflare URL, and - # keep draining so it never blocks on a full pipe. + # Drain cloudflared's output: capture the first trycloudflare URL and the + # first edge-connection registration, and keep draining so it never + # blocks on a full pipe. try: if proc.stdout is not None: for line in proc.stdout: @@ -224,20 +251,35 @@ class CloudflareTunnel: if match: self.url = match.group(0) self._url_event.set() + if not self.ready and _REGISTERED_MARKER in line: + self.ready = True + self._ready_event.set() except Exception: pass finally: + # stdout closed -> cloudflared has exited. Record why, and unblock any + # waiters at once instead of letting them wait out the full timeout. if self.url is None: self.error = "cloudflared exited before emitting a tunnel URL" - self._url_event.set() + elif not self.ready: + self.error = "cloudflared exited before the tunnel connection registered" + self._url_event.set() + self._ready_event.set() - def wait_for_url(self, timeout: float = _URL_TIMEOUT) -> Optional[str]: - self._url_event.wait(timeout) - return self.url + def wait_for_ready(self, timeout: float = _READY_TIMEOUT) -> Optional[str]: + """Block until the tunnel is actually serving -- the URL has been minted + *and* at least one edge connection has registered -- or until timeout. + + Returns the URL only when ready, so callers never advertise a URL that + would return Cloudflare error 1033 (HTTP 530).""" + self._ready_event.wait(timeout) + return self.url if self.ready else None def stop(self) -> None: """Terminate the tunnel. Idempotent and safe to call from a signal handler.""" with self._lock: + # Mark stopped so a start() racing behind us refuses to spawn. + self._stopped = True proc, self._proc = self._proc, None if proc is None: return @@ -260,43 +302,74 @@ class CloudflareTunnel: # enough; the lock guards the start/stop/shutdown races. _active_tunnel: Optional[CloudflareTunnel] = None _active_lock = threading.Lock() +# Latched by stop_studio_tunnel so a shutdown landing *between* a start's retry +# attempts aborts the loop instead of starting a tunnel nobody will ever stop. +_shutdown_requested = False -def start_studio_tunnel(port: int, timeout: float = _URL_TIMEOUT) -> Optional[str]: - """Start a quick tunnel and return its public URL, or None (best-effort). +def start_studio_tunnel(port: int, timeout: float = _READY_TIMEOUT) -> Optional[str]: + """Start a quick tunnel and return its public URL once it is actually + serving, or None (best-effort). - On any failure (no binary, no URL within timeout, early crash) the tunnel is - stopped and None is returned, so the caller prints a hint and continues. + Waits for cloudflared to both mint the URL and register an edge connection + before returning, so the caller never advertises a URL that yields Cloudflare + error 1033 (HTTP 530). If a URL is minted but no connection registers within + the window (e.g. quic is blocked on this network), retries once forcing the + http2 protocol. On any failure the tunnel is stopped and None is returned. """ - global _active_tunnel + global _active_tunnel, _shutdown_requested binary = ensure_cloudflared() if not binary: return None - tunnel = CloudflareTunnel(port, binary) - # Register before start/wait so a shutdown during the URL wait can stop it. with _active_lock: - prior, _active_tunnel = _active_tunnel, tunnel - if prior is not None: - prior.stop() - try: - tunnel.start() - url = tunnel.wait_for_url(timeout) - except Exception: - url = None - if url: - return url - # No URL (or crash): drop it unless a concurrent shutdown already replaced it. - with _active_lock: - if _active_tunnel is tunnel: - _active_tunnel = None - tunnel.stop() + _shutdown_requested = False # fresh session + # Default protocol first (quic, with cloudflared's own http2 fallback); if a + # URL appears but no connection registers, quic is likely blocked -> retry + # once forcing http2. + for protocol in (None, "http2"): + # Create + register under the lock, and bail if a stop already landed + # (e.g. between this and the previous attempt) so we never start a tunnel + # after shutdown has run. + with _active_lock: + if _shutdown_requested: + _active_tunnel = None + return None + tunnel = CloudflareTunnel(port, binary, protocol = protocol) + prior, _active_tunnel = _active_tunnel, tunnel + if prior is not None: + prior.stop() + try: + tunnel.start() + url = tunnel.wait_for_ready(timeout) + except Exception: + url = None + if url: + return url + saw_url = tunnel.url is not None + # Not ready: drop it, but only if we are still the active tunnel. + with _active_lock: + was_active = _active_tunnel is tunnel + if was_active: + _active_tunnel = None + tunnel.stop() + # A concurrent shutdown or start took over while we waited; retrying would + # spawn a tunnel nobody owns (orphaned after shutdown), so bail instead. + if not was_active: + return None + # No URL at all is an API/network failure, not a protocol one; forcing + # http2 will not help, so do not burn another window on it. + if not saw_url: + return None return None def stop_studio_tunnel() -> None: """Terminate the active tunnel, if any. Idempotent.""" - global _active_tunnel + global _active_tunnel, _shutdown_requested with _active_lock: + # Latch so an in-flight start_studio_tunnel won't start a fresh tunnel + # (e.g. its http2 retry) after we have already torn down. + _shutdown_requested = True tunnel, _active_tunnel = _active_tunnel, None if tunnel is not None: tunnel.stop() diff --git a/studio/backend/core/export/export.py b/studio/backend/core/export/export.py index d5a3de06df..b28b61f088 100644 --- a/studio/backend/core/export/export.py +++ b/studio/backend/core/export/export.py @@ -18,7 +18,12 @@ from utils.hardware import clear_gpu_cache from utils.models import is_vision_model, get_base_model_from_lora from utils.models.model_config import detect_audio_type -from utils.paths import ensure_dir, outputs_root, resolve_export_dir, resolve_output_dir +from utils.paths import ( + ensure_dir, + outputs_root, + resolve_export_write_dir, + resolve_output_dir, +) from core.inference import get_inference_backend # GPU-only imports — guarded for Apple Silicon where these aren't needed @@ -336,7 +341,7 @@ class ExportBackend: save_method = "merged_16bit" if save_directory: - save_directory = str(resolve_export_dir(save_directory)) + save_directory = str(resolve_export_write_dir(save_directory)) logger.info(f"Saving merged model locally to: {save_directory}") ensure_dir(Path(save_directory)) @@ -436,7 +441,7 @@ class ExportBackend: output_path: Optional[str] = None try: if save_directory: - save_directory = str(resolve_export_dir(save_directory)) + save_directory = str(resolve_export_write_dir(save_directory)) logger.info(f"Saving base model locally to: {save_directory}") ensure_dir(Path(save_directory)) @@ -563,6 +568,7 @@ class ExportBackend: return False, "No model loaded. Please select a checkpoint first.", None output_path: Optional[str] = None + model_tmp_to_cleanup: Optional[str] = None try: # unsloth expects lowercase quant method quant_method = quantization_method.lower() @@ -588,9 +594,8 @@ class ExportBackend: _LLAMA_CPP_SCRIPTS_WARNING_EMITTED = True if save_directory: - save_directory = str(resolve_export_dir(save_directory)) - # Absolute path so unsloth's relative-path internals resolve - # against the repo root cwd, not the export directory. + save_directory = str(resolve_export_write_dir(save_directory)) + # Keep unsloth relative-path internals anchored to the repo cwd. abs_save_dir = os.path.abspath(save_directory) logger.info(f"Saving GGUF model locally to: {abs_save_dir}") @@ -604,9 +609,15 @@ class ExportBackend: cwd = os.getcwd() pre_existing_ggufs = set(glob.glob(os.path.join(cwd, "*.gguf"))) - model_save_path = os.path.join(abs_save_dir, "model") + pre_existing_subs = {d.name for d in Path(abs_save_dir).iterdir() if d.is_dir()} + + # Avoid clobbering an existing user-owned model/ directory. + import uuid + + _model_tmp = os.path.join(abs_save_dir, f"_tmp_model_{uuid.uuid4().hex[:8]}") + model_tmp_to_cleanup = _model_tmp self.current_model.save_pretrained_gguf( - model_save_path, + _model_tmp, self.current_tokenizer, quantization_method = quant_method, ) @@ -618,10 +629,12 @@ class ExportBackend: shutil.move(src, dest) logger.info(f"Relocated GGUF: {os.path.basename(src)} → {abs_save_dir}/") - # Flatten any .gguf from subdirs (e.g. model_gguf/) into abs_save_dir. + # Flatten GGUF files from subdirs created during this export. for sub in list(Path(abs_save_dir).iterdir()): if not sub.is_dir(): continue + if sub.name in pre_existing_subs: + continue for src in sub.glob("*.gguf"): dest = os.path.join(abs_save_dir, src.name) shutil.move(str(src), dest) @@ -634,7 +647,7 @@ class ExportBackend: if self.current_checkpoint: ckpt = Path(self.current_checkpoint) gguf_dir = ckpt.parent / f"{ckpt.name}_gguf" - if gguf_dir.is_dir(): + if gguf_dir.is_dir() and gguf_dir.resolve() != Path(abs_save_dir).resolve(): for src in gguf_dir.glob("*.gguf"): dest = os.path.join(abs_save_dir, src.name) shutil.move(str(src), dest) @@ -683,6 +696,8 @@ class ExportBackend: ) except Exception as e: + if model_tmp_to_cleanup: + shutil.rmtree(model_tmp_to_cleanup, ignore_errors = True) logger.error(f"Error exporting GGUF model: {e}") import traceback @@ -712,7 +727,7 @@ class ExportBackend: output_path: Optional[str] = None try: if save_directory: - save_directory = str(resolve_export_dir(save_directory)) + save_directory = str(resolve_export_write_dir(save_directory)) logger.info(f"Saving LoRA adapter locally to: {save_directory}") ensure_dir(Path(save_directory)) diff --git a/studio/backend/core/inference/chat_templates.py b/studio/backend/core/inference/chat_templates.py new file mode 100644 index 0000000000..58f63ff61b --- /dev/null +++ b/studio/backend/core/inference/chat_templates.py @@ -0,0 +1,109 @@ +# SPDX-License-Identifier: AGPL-3.0-only +# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0 + +"""Bundled chat-template selection for GGUF inference. + +Some shipped GGUF quants embed an older chat template. Rather than re-cutting and +asking users to re-download every quant, Studio can override the embedded template +at llama-server launch time with a bundled, up-to-date Jinja template for known +model families. The override is wired through the existing ``chat_template_override`` +-> ``--chat-template-file`` path in ``LlamaCppBackend.load_model``. + +Currently this covers ``unsloth/gemma-4-*-GGUF``, which gains the upstream PR #118 +``preserve_thinking`` flag (defaulted OFF here) so the Studio "Preserve thinking" +toggle appears while staying disabled by default. +""" + +import re +from functools import lru_cache +from pathlib import Path +from typing import Optional + +# assets live at /assets/chat_templates/. This module is at +# /core/inference/chat_templates.py, so walk up three parents to +# (mirrors utils/inference/inference_config.py). +_ASSETS_DIR = Path(__file__).parent.parent.parent / "assets" / "chat_templates" + +# unsloth/gemma-4--GGUF (case-insensitive). The "-GGUF" suffix is retained +# on ModelConfig.identifier for HF GGUF repos, so this matches E2B / E4B / 31B / +# 26B-A4B and any future unsloth/gemma-4-*-GGUF, while excluding gemma-3, +# non-Unsloth, and non-GGUF identifiers (e.g. the bf16 "unsloth/gemma-4-E2B-it"). +_GEMMA4_GGUF_RE = re.compile(r"^unsloth/gemma-4-.+-gguf$", re.IGNORECASE) + +# Google ships two distinct gemma-4 chat templates: E2B/E4B omit the empty +# "<|channel>thought" block on enable_thinking=false, while the +# 12b/26B-A4B/31B family emits it. Route the two GGUF families to the matching +# bundled template so each keeps its model's intended behavior. +_GEMMA4_EDGE_GGUF_RE = re.compile(r"^unsloth/gemma-4-e[24]b-it-gguf$", re.IGNORECASE) + +_GEMMA4_TEMPLATE_FILE = "gemma-4.jinja" # 12b / 26B-A4B / 31B +_GEMMA4_EDGE_TEMPLATE_FILE = "gemma-4-edge.jinja" # E2B / E4B + + +def _canonical_repo_id(model_identifier: str) -> str: + """Mirror ``ModelConfig.from_identifier``: a bare HF shorthand with no owner + (e.g. ``gemma-4-E2B-it-GGUF``) defaults to the ``unsloth/`` org. The resolver + runs on the raw ``request.model_path`` (before that canonicalization), so apply + the same rule here, otherwise shorthand loads would skip the override. + """ + mid = model_identifier.strip() + if mid and "/" not in mid: + mid = f"unsloth/{mid}" + return mid + + +def is_unsloth_gemma4_gguf(model_identifier: Optional[str]) -> bool: + """True for canonical ``unsloth/gemma-4-*-GGUF`` repo identifiers (and the + owner-less shorthand that resolves to the same Unsloth repo).""" + if not model_identifier: + return False + return bool(_GEMMA4_GGUF_RE.match(_canonical_repo_id(model_identifier))) + + +def is_unsloth_gemma4_edge_gguf(model_identifier: Optional[str]) -> bool: + """True for the E2B / E4B GGUF repos, which use the edge-variant template.""" + if not model_identifier: + return False + return bool(_GEMMA4_EDGE_GGUF_RE.match(_canonical_repo_id(model_identifier))) + + +def _gemma4_template_file(model_identifier: Optional[str]) -> Optional[str]: + """Return the bundled template filename for a gemma-4 GGUF id, else None.""" + if is_unsloth_gemma4_edge_gguf(model_identifier): + return _GEMMA4_EDGE_TEMPLATE_FILE + if is_unsloth_gemma4_gguf(model_identifier): + return _GEMMA4_TEMPLATE_FILE + return None + + +@lru_cache(maxsize=8) +def load_bundled_chat_template(name: str) -> str: + """Read a bundled chat-template asset by filename (cached for the process).""" + return (_ASSETS_DIR / name).read_text(encoding="utf-8") + + +def resolve_effective_chat_template_override( + *, + model_identifier: Optional[str], + user_override: Optional[str], +) -> Optional[str]: + """Resolve which chat-template text to launch llama-server with. + + Precedence: + 1. An explicit, non-empty user override always wins (advanced users). + 2. For ``unsloth/gemma-4-*-GGUF``, return the bundled gemma-4 template + (adds ``preserve_thinking``, default off) so the embedded GGUF template + is overridden without re-downloading quants. E2B/E4B get the edge + variant; 12b/26B-A4B/31B get the standard one. + 3. Otherwise ``None`` -> llama-server renders the GGUF's embedded template. + + The result is fed to ``LlamaCppBackend.load_model(chat_template_override=...)`` + and must be computed before the route-level reload-dedup check so the live + backend state and the incoming request compare consistently. + """ + if user_override and user_override.strip(): + return user_override + template_file = _gemma4_template_file(model_identifier) + if template_file is not None: + return load_bundled_chat_template(template_file) + return None diff --git a/studio/backend/core/inference/llama_cpp.py b/studio/backend/core/inference/llama_cpp.py index 8cf37ed9ec..80b9739cc1 100644 --- a/studio/backend/core/inference/llama_cpp.py +++ b/studio/backend/core/inference/llama_cpp.py @@ -31,8 +31,12 @@ from core.inference.llama_server_args import ( extra_args_disable_mmproj, parse_cache_override, parse_ctx_override, + parse_split_mode_override, resolve_cache_type_kv, resolve_requested_ctx, + resolve_tensor_parallel, + strip_shadowing_flags, + strip_split_mode_only, ) from core.tool_healing import ( _TC_END_TAG_RE, @@ -71,6 +75,33 @@ from state.tool_approvals import ( logger = get_logger(__name__) +def _wsl_system_rocm_lib_dirs() -> "list[str]": + """System ROCm lib dir(s) to load before a prebuilt's bundled HIP, on WSL. + + The bundled bare-metal HIP can't drive WSL's /dev/dxg and segfaults on the + first GPU call; the system ROCm libs (libamdhip64 + librocdxg) can, while + the bundle still supplies libggml-hip / librocblas (gfx1151 kernels). + Mirrors install_llama_prebuilt._wsl_system_rocm_lib_dirs so a prebuilt that + passed install validation runs the same at serve time. No-op off a ROCDXG + WSL host (needs /dev/dxg, "microsoft" /proc/version, librocdxg in /opt/rocm). + """ + try: + if not os.path.exists("/dev/dxg"): + return [] + with open("/proc/version", encoding = "utf-8", errors = "replace") as fh: + if "microsoft" not in fh.read().lower(): + return [] + except OSError: + return [] + out: "list[str]" = [] + for d in ("/opt/rocm/lib", "/opt/rocm/lib64"): + if os.path.exists(os.path.join(d, "librocdxg.so")) or os.path.exists( + os.path.join(d, "librocdxg.so.1") + ): + out.append(d) + return out + + # ── Pre-compiled patterns for plan-without-action re-prompt ── # Forward-looking intent signals: the model is describing what it *will* # do rather than giving a final answer. @@ -91,15 +122,11 @@ _INTENT_SIGNAL = re.compile( ) _MAX_REPROMPTS = 1 -# Without max_tokens, llama-server defaults n_predict = n_ctx (up to 262144 for -# Qwen3.5), causing many-minute zombie decodes when cancel fails. -# t_max_predict_ms is a wall-clock backstop but per the llama.cpp README only -# fires after a newline, so we keep a token cap as the front-line limiter. -# The cap is the effective context length when known, else this floor. 4096 was -# too low: Qwen3 / gpt-oss reasoning traces and max_tokens-omitting OpenAI-API -# callers (langchain, llama-index, curl) got truncated mid-sentence. +# Default max_tokens to the effective context when known. The floor is high +# enough for reasoning-heavy GGUFs and max_tokens-omitting API clients. _DEFAULT_MAX_TOKENS_FLOOR = 32768 -_DEFAULT_T_MAX_PREDICT_MS = 600_000 # 10 min +_DEFAULT_FIRST_TOKEN_TIMEOUT_S = 1200.0 # 20 min +_DEFAULT_STREAM_STALL_TIMEOUT_S = 120.0 # 2 min _REPROMPT_MAX_CHARS = 2000 _FORCED_REPEAT_PLAN_SIGNAL = re.compile( r"\b(?:i\s+will|i'll|let\s+me|going\s+to|need\s+to|call|use|run|search|fetch|render)\b", @@ -691,6 +718,11 @@ class LlamaCppBackend: self._spec_fallback_reason: Optional[str] = None self._hf_variant: Optional[str] = None self._is_vision: bool = False + # Block-diffusion model (e.g. DiffusionGemma): served by the diffusion + # runner, not llama-server. Set from the GGUF architecture at load. + self._architecture: Optional[str] = None + self._is_diffusion: bool = False + self._diffusion_visual_bin: Optional[str] = None self._healthy = False # Set by _classify_gpu_offload after _wait_for_health. self._gpu_offload_active: Optional[bool] = None @@ -705,6 +737,8 @@ class LlamaCppBackend: self._supports_preserve_thinking: bool = False self._supports_tools: bool = False self._cache_type_kv: Optional[str] = None + # Whether --split-mode tensor was applied on the active load. + self._tensor_parallel: bool = False self._reasoning_default: bool = True self._speculative_type: Optional[str] = None # Canonical UI-facing mode the user requested @@ -795,6 +829,11 @@ class LlamaCppBackend: def is_vision(self) -> bool: return self._is_vision + @property + def is_diffusion(self) -> bool: + """True when the loaded GGUF is a block-diffusion model (DiffusionGemma).""" + return self._is_diffusion + @property def hf_variant(self) -> Optional[str]: return self._hf_variant @@ -984,8 +1023,10 @@ class LlamaCppBackend: # enable_thinking / reasoning_effort -- skip. if self._supports_reasoning and not self._reasoning_always_on: if self._reasoning_style == "reasoning_effort": - if reasoning_effort in ("low", "medium", "high"): + if reasoning_effort in ("none", "low", "medium", "high"): kwargs["reasoning_effort"] = reasoning_effort + elif reasoning_effort == "minimal": + kwargs["reasoning_effort"] = "low" elif enable_thinking is not None: kwargs["reasoning_effort"] = "high" if enable_thinking else "low" else: @@ -997,12 +1038,21 @@ class LlamaCppBackend: @property def supports_tools(self) -> bool: + # DiffusionGemma serves via the visual runner, whose live per-step canvas + # frames are dropped by the agentic tool loop; never route it through tools. + if self._is_diffusion: + return False return self._supports_tools @property def cache_type_kv(self) -> Optional[str]: return self._cache_type_kv + @property + def tensor_parallel(self) -> bool: + """Whether --split-mode tensor is active on the loaded server.""" + return self._tensor_parallel + @property def speculative_type(self) -> Optional[str]: return self._speculative_type @@ -1329,6 +1379,105 @@ class LlamaCppBackend: return False return False + # Datacenter / professional NVIDIA parts that benefit from the llama.cpp + # FP32-accum / P2P tunings. Whole-word (\b) so short markers don't match + # workstation parts as substrings: "a100" must not fire on "RTX A1000". + _DATACENTER_GPU_RE = re.compile( + r"\b(?:a100|a30|h100|h200|h800|gh200|b200|b100|b300|gb200|gb300|" + r"l40s?|l4|rtx pro 6000|rtx 6000 ada)\b" + ) + + @staticmethod + def _is_datacenter_gpu(gpu_indices = None) -> bool: + """True iff every selected NVIDIA GPU is a datacenter/professional part. + NVIDIA-only, fails open to False (consumer GeForce, ROCm, CPU and errors + are left untouched); a mixed DC+consumer selection counts as non-DC. + + gpu_indices are PHYSICAL ids (see _get_gpu_free_memory), but + get_device_properties wants mask-relative ordinals, so we rebuild the + ordinal->physical map from CUDA_VISIBLE_DEVICES and key names by physical + id. Otherwise a masked host (CUDA_VISIBLE_DEVICES=4,5,6,7, selection [4,5]) + would drop the tuning or probe the wrong GPU.""" + try: + import torch + + if getattr(torch.version, "hip", None) is not None: + return False # ROCm reuses torch.cuda.*; not a CUDA part + if not (hasattr(torch, "cuda") and torch.cuda.is_available()): + return False + count = torch.cuda.device_count() + + # Mirror _get_gpu_free_memory: map visible ordinal -> physical id via + # CUDA_VISIBLE_DEVICES; unset/unparsable leaves physical id == ordinal. + physical_ids: Optional[list[int]] = None + cvd = os.environ.get("CUDA_VISIBLE_DEVICES") + if cvd is not None: + try: + physical_ids = [int(x.strip()) for x in cvd.split(",") if x.strip()] + except ValueError: + physical_ids = None + + pattern = LlamaCppBackend._DATACENTER_GPU_RE + names_by_id: dict[int, str] = {} + for ordinal in range(count): + try: + name = (torch.cuda.get_device_properties(ordinal).name or "").lower() + except Exception: + continue + pid = ( + physical_ids[ordinal] + if physical_ids is not None and ordinal < len(physical_ids) + else ordinal + ) + names_by_id[pid] = name + + indices = list(gpu_indices) if gpu_indices else list(names_by_id) + saw = False + for _i in indices: + name = names_by_id.get(_i) + if name is None: + continue # not visible -> skip (fail conservative) + saw = True + if not pattern.search(name): + return False + return saw + except Exception: + return False + + @staticmethod + def _effective_gpu_count(gpu_indices = None) -> int: + """GPUs llama-server will use: len(selection), else the visible CUDA + device count (None = every visible GPU). 0 on error so multi-GPU tuning + stays off when the count is unknown.""" + if gpu_indices is not None: + return len(gpu_indices) + try: + import torch + if hasattr(torch, "cuda") and torch.cuda.is_available(): + return torch.cuda.device_count() + except Exception: + return 0 + return 0 + + @staticmethod + def _apply_datacenter_env(env: dict, gpu_indices = None) -> bool: + """Inject DC llama.cpp tuning into env in place via setdefault (user + values win); return whether the box qualified. Opt out with + UNSLOTH_DISABLE_DC_TUNING=1; only datacenter NVIDIA parts qualify + (consumer/ROCm/CPU/error are a no-op). Sets GGML_CUDA_FORCE_CUBLAS_COMPUTE_32F + for any qualifying GPU (FP32 accum: ~0% cost on B200, real cost on GeForce), + plus GGML_CUDA_P2P + CUDA_SCALE_LAUNCH_QUEUES=4x for multi-GPU (+33-51% pp + tensor-split, +8-16% pipeline split on B200).""" + if os.environ.get("UNSLOTH_DISABLE_DC_TUNING") == "1": + return False + if not LlamaCppBackend._is_datacenter_gpu(gpu_indices): + return False + env.setdefault("GGML_CUDA_FORCE_CUBLAS_COMPUTE_32F", "1") + if LlamaCppBackend._effective_gpu_count(gpu_indices) > 1: + env.setdefault("GGML_CUDA_P2P", "1") + env.setdefault("CUDA_SCALE_LAUNCH_QUEUES", "4x") + return True + @staticmethod def _get_gpu_free_memory() -> list[tuple[int, int]]: """Query free memory per GPU. @@ -1507,6 +1656,23 @@ class LlamaCppBackend: # buffers; 0.90 dropped 91-94% fits to CPU offload (#5106). _GPU_PIN_VRAM_FRACTION = 0.95 + # Per-GPU compute-graph buffer to reserve in tensor mode (MiB). This is the + # logits buffer (n_batch x vocab) + activation scratch that llama.cpp sizes + # via graph_reserve -- it is roughly EQUAL on every device (not proportional + # to the tensor split) and independent of context. Measured ~2.3 GB + # (gemma-3-27B) to ~3.8 GB (gemma-4-31B) on a 256k-vocab model; we reserve a + # conservative headroom above that. It is (a) subtracted from each GPU's free + # VRAM before computing --tensor-split, so the roomier GPU absorbs more + # weight and the smallest GPU keeps room for KV, and (b) reserved per device + # when capping context. The auto-fallback to layer split covers any + # underestimate. NOTE: scales with the model's vocab / batch size; tune if a + # large-vocab model OOMs at load. + _TENSOR_PARALLEL_BUFFER_RESERVE_MIB = 5120 + + # KV cache types llama.cpp accepts in tensor mode. A quantized KV cache + # aborts a --split-mode tensor load, so it's dropped for the tensor attempt. + _TENSOR_PARALLEL_KV_TYPES = frozenset({"f16", "bf16", "f32"}) + @staticmethod def _windows_pip_nvidia_dll_dirs(prefix: str) -> list[str]: """Return DLL dirs from pip-installed CUDA wheels under @@ -1809,6 +1975,7 @@ class LlamaCppBackend: ctx_checkpoints: int = 0, kv_on_gpu: bool = True, mtp_engaged: bool = False, + budget_frac: Optional[float] = None, ) -> int: """Return the largest context length that fits in GPU VRAM. @@ -1842,8 +2009,11 @@ class LlamaCppBackend: ctx_checkpoints = ctx_checkpoints, ) - # MTP engaged: carve the drafter's reserve out of the fit budget. - budget_frac = _CTX_FIT_VRAM_FRACTION - (_MTP_VRAM_RESERVE_FRAC if mtp_engaged else 0.0) + # MTP engaged: carve the drafter's reserve out of the fit budget. Callers + # can override outright (tensor-parallel mode passes a fatter margin), so + # only compute a default when none was supplied. + if budget_frac is None: + budget_frac = _CTX_FIT_VRAM_FRACTION - (_MTP_VRAM_RESERVE_FRAC if mtp_engaged else 0.0) budget_bytes = available_mib * 1024 * 1024 * budget_frac model_footprint = model_size_bytes @@ -2058,11 +2228,16 @@ class LlamaCppBackend: self._ssm_state_size = None self._shared_kv_layers = None self._nextn_predict_layers = None + self._architecture = None + self._is_diffusion = False try: + canvas_seen = False WANTED = { "general.architecture", "tokenizer.chat_template", + # Block-diffusion marker (DiffusionGemma); routes to the diffusion runner. + "diffusion.canvas_length", # Source-repo hints for the SWA resolver's HF fallback. "general.source.huggingface.repository", "general.source.url", @@ -2117,6 +2292,7 @@ class LlamaCppBackend: general[key] = val_s if key == "general.architecture": arch = val_s + self._architecture = val_s arch_keys = { f"{arch}.context_length": "context_length", f"{arch}.block_count": "n_layers", @@ -2145,6 +2321,8 @@ class LlamaCppBackend: if vtype == 4 else struct.unpack(" Optional[tuple[list, str, Optional[str]]]: + """Resolve how to launch the DiffusionGemma runner: (shim argv prefix, + visual-server binary, optional extra PYTHONPATH dir for the file override). + + Shim: UNSLOTH_DG_SHIM (a .py file) first, else the installed + unsloth_zoo.diffusion_studio.shim. Binary: DG_VISUAL_BIN first, else + alongside llama-server. Returns None if neither can be found. + """ + import importlib.util + import os + import sys + + # Visual-server binary: env override, else next to llama-server or in the + # install's build/bin (where the prebuilt/installer puts it). .exe on Windows. + visual_bin = os.environ.get("DG_VISUAL_BIN") + if not visual_bin: + name = "llama-diffusion-gemma-visual-server" + (".exe" if os.name == "nt" else "") + base = self._find_llama_server_binary() + if base: + base_dir = Path(base).parent + for cand in ( + base_dir / name, + base_dir / "build" / "bin" / name, + base_dir / "build" / "bin" / "Release" / name, + ): + if cand.is_file(): + visual_bin = str(cand) + break + if not (visual_bin and Path(visual_bin).is_file()): + return None + + # Shim: a file override (its dir goes on PYTHONPATH), else the zoo package via -m. + shim_file = os.environ.get("UNSLOTH_DG_SHIM") + if shim_file and Path(shim_file).is_file(): + return ([sys.executable, shim_file], visual_bin, str(Path(shim_file).parent)) + + # Find the installed shim without importing the heavy unsloth_zoo package + # (find_spec on the top-level package does not run its __init__). + try: + spec = importlib.util.find_spec("unsloth_zoo") + except Exception: + spec = None + if spec is not None and spec.submodule_search_locations: + pkg_dir = Path(list(spec.submodule_search_locations)[0]) + if (pkg_dir / "diffusion_studio" / "shim.py").is_file(): + return ( + [sys.executable, "-m", "unsloth_zoo.diffusion_studio.shim"], + visual_bin, + None, + ) + + return None + + def _start_diffusion_server( + self, + *, + model_path: str, + gguf_path: Optional[str], + hf_repo: Optional[str], + hf_variant: Optional[str], + model_identifier: str, + n_ctx: int, + extra_args: Optional[List[str]], + ) -> bool: + """Launch the OpenAI-compat diffusion shim (which drives the on-device + visual decoder) and wait for health. Presents the same /v1 + /health + interface as llama-server, so the rest of Studio is unchanged. + """ + import os + + assets = self._find_diffusion_assets() + if assets is None: + raise RuntimeError( + "DiffusionGemma runner not found. Install unsloth_zoo (which ships " + "unsloth_zoo.diffusion_studio.shim) or set UNSLOTH_DG_SHIM to a shim " + "file, and provide the visual-server binary via DG_VISUAL_BIN or next " + "to llama-server in the install tree." + ) + shim_cmd, visual_bin, extra_pythonpath = assets + self._diffusion_visual_bin = visual_bin + + self._kill_process() + self._port = self._find_free_port() + # Auto-size (0): the visual server probes the largest context that fits this GPU's VRAM + # (capped at the training context). An explicit in-range n_ctx overrides it. + maxtok = n_ctx if (n_ctx and 0 < n_ctx <= 65536) else 0 + gpu = os.environ.get("DG_GPU", "0") + + cmd = list(shim_cmd) + [ + "--gguf", + model_path, + "--host", + "127.0.0.1", + "--port", + str(self._port), + "--gpu", + gpu, + "--maxtok", + str(maxtok), + ] + + env = child_env_without_native_path_secret() + # `python -m unsloth_zoo.diffusion_studio.shim` imports unsloth_zoo, which + # refuses to load unless UNSLOTH_IS_PRESENT is set (normally by `import + # unsloth`). The shim never imports unsloth, so set it here as unsloth does. + env["UNSLOTH_IS_PRESENT"] = "1" + env["DG_VISUAL_BIN"] = visual_bin + env["DG_GPU"] = gpu + # The file-override shim imports its sibling visual_engine; put its dir on PYTHONPATH. + # (The zoo-package shim is an installed module and needs no PYTHONPATH change.) + if extra_pythonpath: + existing = env.get("PYTHONPATH") + env["PYTHONPATH"] = ( + (extra_pythonpath + os.pathsep + existing) if existing else extra_pythonpath + ) + + logger.info(f"Starting DiffusionGemma runner: {' '.join(cmd)}") + self._stdout_lines = [] + self._llama_log_fh = None + self._llama_log_path = None + try: + log_dir = _swa_cache_path().parent / "logs" / "diffusion-server" + log_dir.mkdir(parents = True, exist_ok = True) + self._llama_log_path = log_dir / f"diffusion-{int(time.time())}-port-{self._port}.log" + self._llama_log_fh = open(self._llama_log_path, "w", encoding = "utf-8", buffering = 1) + logger.info(f"diffusion runner stdout/stderr -> {self._llama_log_path}") + except OSError as e: + logger.debug(f"Could not open diffusion runner log file: {e}") + + # PR_SET_PDEATHSIG: the shim (and its visual server) die with this backend + # process, so a Studio crash/restart never orphans a GPU process. + popen_kwargs = dict(_windows_hidden_subprocess_kwargs()) + if sys.platform.startswith("linux"): # prctl/libc.so.6 are Linux-only + + def _pdeathsig(): + try: + import ctypes + import signal as _signal + ctypes.CDLL("libc.so.6", use_errno = True).prctl(1, _signal.SIGTERM) + except Exception: + pass + + popen_kwargs["preexec_fn"] = _pdeathsig + + self._process = subprocess.Popen( + cmd, + stdout = subprocess.PIPE, + stderr = subprocess.STDOUT, + text = True, + env = env, + **popen_kwargs, + ) + self._stdout_thread = threading.Thread( + target = self._drain_stdout, daemon = True, name = "diffusion-stdout" + ) + self._stdout_thread.start() + + # Publish state before the health wait (mirrors the llama-server path). + self._gguf_path = model_path + self._hf_repo = hf_repo + self._is_vision = False + self._is_audio = False # clear any prior TTS/audio model's routing flag + self._model_identifier = model_identifier + self._cache_type_kv = None + self._gpu_offload_active = True + if hf_variant: + self._hf_variant = hf_variant + elif gguf_path: + try: + from utils.models.model_config import _extract_quant_label + self._hf_variant = _extract_quant_label(gguf_path) + except Exception: + self._hf_variant = None + else: + self._hf_variant = None + # Provisional until the server reports the budget it resolved (auto-size picks it from VRAM). + self._effective_context_length = maxtok or self._context_length + self._max_context_length = self._context_length or maxtok or None + + healthy = self._wait_for_health(timeout = 600.0) + if healthy: + self._healthy = True + self._gpu_offload_active = True + if extra_args is not None: + self._extra_args = list(extra_args) + self._extra_args_source = (model_identifier, hf_variant) + # The visual server logs "MAXTOK=" with the context budget it actually resolved + # (auto-sized to VRAM). Read it back so the UI context bar shows the real budget. + chosen = maxtok + try: + import re as _re + for _ln in reversed(self._stdout_lines): + _m = _re.search(r"MAXTOK=(\d+)", _ln) + if _m: + chosen = int(_m.group(1)) + break + except Exception: + pass + if chosen and chosen > 0: + self._effective_context_length = chosen + self._max_context_length = chosen + self._requested_n_ctx = int(n_ctx) + else: + self._healthy = False + logger.error("DiffusionGemma runner failed to become healthy") + return healthy + # ── HF download (no lock held) ─────────────────────────────── def _download_gguf( @@ -2598,6 +2996,16 @@ class LlamaCppBackend: return str(mmproj) + def _mmproj_vram_bytes(self, launch_mmproj_path: Optional[str]) -> int: + """Return resolved mmproj VRAM bytes, or 0 when absent/unreadable.""" + if not launch_mmproj_path: + return 0 + try: + return self._get_gguf_size_bytes(launch_mmproj_path) + except OSError as e: + logger.debug(f"Could not size mmproj {launch_mmproj_path}: {e}") + return 0 + def _resolve_launch_mtp_path(self, *, mtp_draft_path: Optional[str]) -> Optional[str]: """Return mtp_draft_path iff it exists on disk, else None. @@ -2652,6 +3060,16 @@ class LlamaCppBackend: """ lowered = (output or "").lower() + # Tensor parallelism (--split-mode tensor) is arch-gated in llama.cpp; + # unsupported architectures abort the load with this marker. Point the + # user at the toggle instead of a generic invalid-GGUF/OOM message. + if "split_mode_tensor not implemented" in lowered: + return ( + "Tensor parallelism is not supported for this model's " + "architecture. Turn off Tensor Parallelism in the model " + "settings and reload." + ) + # Detect Ollama source up front so the arch branch can keep the # Ollama hint instead of the generic "unsupported arch" message. gguf = gguf_path or "" @@ -2708,6 +3126,97 @@ class LlamaCppBackend: "Check that the GGUF file is valid and you have enough memory." ) + def _plan_tensor_parallel( + self, + gpus: list[tuple[int, int]], + model_size: int, + target_ctx: int, + cache_type_kv: Optional[str] = None, + n_parallel: int = 1, + mtp_engaged: bool = False, + max_target_ctx: Optional[int] = None, + ) -> tuple[int, int, list[int], Optional[list[int]]]: + """Plan a ``--split-mode tensor`` load. Pure: no model or GPU needed. + + ``gpus`` is a list of ``(gpu_index, free_mib)``; ``model_size`` is the + weight size in bytes; ``target_ctx`` is the context to fit (the explicit + request, or the model's native length for auto). ``max_target_ctx`` is + the native/hardware ceiling used only for the UI bound (defaults to + ``target_ctx``). Returns + ``(effective_ctx, max_available_ctx, gpu_indices, tensor_split)``. + + Policy (assumes >= 2 GPUs; the caller drops the toggle below that): + - Cap context to the KV that fits the pooled VRAM after the weights and + one per-device compute-graph buffer (``_TENSOR_PARALLEL_BUFFER_RESERVE_MIB``). + llama.cpp's ``--fit`` is a no-op in tensor mode, so this is the only + cap, honored even for an explicit ``-c``. It is more accurate than the + 0.80 whole-pool heuristic, which over-reserves and leaves VRAM unused. + - ``tensor_split`` is None (llama.cpp's even default, safe for every arch + incl. Gemma 3n which GGML_ASSERTs on a weighted split) when an even + share fits the smallest GPU; otherwise it is weighted by + ``(free - buffer)`` so the roomier GPU absorbs more weight and the + smallest GPU keeps room for KV. + """ + # Drop GPUs that can't hold the per-device compute-graph buffer; they'd + # OOM in tensor mode. load_model already filters before calling, so this + # is defense-in-depth that also keeps the pure function self-contained + # (and unit-testable without a GPU). + reserve_mib = self._TENSOR_PARALLEL_BUFFER_RESERVE_MIB + usable_gpus = [g for g in gpus if g[1] >= reserve_mib] + gpu_indices = sorted(idx for idx, _ in usable_gpus) + if len(gpu_indices) < 2: + # Tensor parallelism is meaningless on <2 GPUs (the caller drops the + # toggle before this); be defensive and never emit a split here. + return ( + target_ctx if target_ctx > 0 else 4096, + target_ctx if target_ctx > 0 else 4096, + gpu_indices, + None, + ) + free_by_idx = {idx: free for idx, free in usable_gpus} + pool_mib = sum(free_by_idx.values()) + kv_budget_b = (pool_mib - len(gpu_indices) * reserve_mib) * 1024 * 1024 - model_size + if mtp_engaged: + # MTP keeps a draft model + its own KV cache on GPU. + kv_budget_b -= 2 * 1024**3 + + def _fit_ctx(ctx: int) -> int: + # Largest context whose KV fits the pooled budget. Floors small, but + # never raises an explicit ctx above what was asked. + if self._can_estimate_kv() and ctx > 0: + ctx_floor = min(2048, ctx) + if kv_budget_b <= 0: + # Weights + buffers exceed the pool -> floor; the load then + # falls back to layer split. + return ctx_floor + kv_at = self._estimate_kv_cache_bytes(ctx, cache_type_kv, n_parallel = n_parallel) + if kv_at <= kv_budget_b: + return ctx + return max(ctx_floor, int(ctx * kv_budget_b / kv_at)) + # KV size unknown -> can't prove a safe cap; floor. + return min(4096, ctx) if ctx > 0 else 4096 + + # max_available_ctx is the hardware ceiling for the UI bound, sized from + # the native context independent of an explicit small -c (which only + # caps effective_ctx). + max_ctx_target = max_target_ctx if (max_target_ctx and max_target_ctx > 0) else target_ctx + max_available_ctx = _fit_ctx(max_ctx_target) + effective_ctx = min(_fit_ctx(target_ctx), max_available_ctx) + + min_free_mib = min(free_by_idx.values()) + kv_bytes = ( + self._estimate_kv_cache_bytes(effective_ctx, cache_type_kv, n_parallel = n_parallel) + if (self._can_estimate_kv() and effective_ctx > 0) + else 0 + ) + even_share_mib = (model_size + kv_bytes) / len(gpu_indices) / (1024 * 1024) + tensor_split: Optional[list[int]] = None + if even_share_mib > (min_free_mib - reserve_mib): + adj = [max(0, int(free_by_idx[i] - reserve_mib)) for i in gpu_indices] + if sum(adj) > 0: + tensor_split = adj + return effective_ctx, max_available_ctx, gpu_indices, tensor_split + @staticmethod def _is_projector_incompatibility(output: str) -> bool: """True when llama-server aborted because it cannot load the model's @@ -2830,6 +3339,7 @@ class LlamaCppBackend: cache_type_kv: Optional[str] = None, speculative_type: Optional[str] = None, spec_draft_n_max: Optional[int] = None, + tensor_parallel: bool = False, n_threads: Optional[int] = None, n_gpu_layers: Optional[int] = None, # caller compat, unused n_parallel: int = 1, @@ -2861,6 +3371,7 @@ class LlamaCppBackend: cache_type_kv = cache_type_kv, speculative_type = speculative_type, spec_draft_n_max = spec_draft_n_max, + tensor_parallel = tensor_parallel, chat_template_override = chat_template_override, extra_args = extra_args, is_vision = is_vision, @@ -2917,13 +3428,9 @@ class LlamaCppBackend: with self._lock: self._kill_process() + # Resolve llama-server now but defer a not-found error: a block-diffusion + # GGUF uses the diffusion runner, and its arch is only known after the header. binary = self._find_llama_server_binary() - if not binary: - raise RuntimeError( - "llama-server binary not found. " - "Run setup.sh to build it, install llama.cpp, " - "or set LLAMA_SERVER_PATH environment variable." - ) # ── Phase 2: download (NO lock held, so cancel can proceed) ── # mtp_draft_path arrives set for local Gemma loads (detected @@ -2978,6 +3485,30 @@ class LlamaCppBackend: logger.info("Load cancelled after download phase") return False + # Block-diffusion GGUFs (DiffusionGemma) cannot run on llama-server; + # serve them with the diffusion runner (same OpenAI-compat interface). + if self._is_diffusion: + with self._lock: + if self._cancel_event.is_set(): + logger.info("Load cancelled before diffusion server start") + return False + return self._start_diffusion_server( + model_path = model_path, + gguf_path = gguf_path, + hf_repo = hf_repo, + hf_variant = hf_variant, + model_identifier = model_identifier, + n_ctx = n_ctx, + extra_args = extra_args, + ) + + if not binary: + raise RuntimeError( + "llama-server binary not found. " + "Run setup.sh to build it, install llama.cpp, " + "or set LLAMA_SERVER_PATH environment variable." + ) + # Outside ``self._lock`` so /unload, /cancel, /status aren't # blocked. ``unload_model`` also records the kill, so the # frontend /unload+/load Apply path engages the wait here even @@ -3000,15 +3531,69 @@ class LlamaCppBackend: requested_ctx = resolve_requested_ctx(extra_args, n_ctx) cache_override = parse_cache_override(extra_args) cache_type_kv = resolve_cache_type_kv(extra_args, cache_type_kv) + # A user --split-mode in extras last-wins-overrides the + # toggle, so reconcile it back into tensor_parallel state. + split_mode_override = parse_split_mode_override(extra_args) + tensor_parallel = resolve_tensor_parallel(extra_args, tensor_parallel) + # Tensor mode aborts on a quantized KV cache, so drop it for the + # tensor attempt (and strip any inherited/explicit --cache-type + # that would re-impose it when appended last). The layer-split + # fallback re-runs with tensor_parallel False and keeps the type. + if ( + tensor_parallel + and cache_type_kv + and cache_type_kv.strip().lower() not in self._TENSOR_PARALLEL_KV_TYPES + ): + logger.info( + "Tensor parallelism requires a non-quantized KV cache; " + "ignoring cache type %s for the tensor attempt.", + cache_type_kv, + ) + cache_type_kv = None + if extra_args: + extra_args = strip_shadowing_flags( + extra_args, + strip_context = False, + strip_cache = True, + strip_spec = False, + strip_template = False, + strip_split_mode = False, + ) if ctx_override is not None and ctx_override > 0: logger.info(f"User --ctx-size {ctx_override} honored; skipping auto-reduce") if cache_override is not None: logger.info(f"User --cache-type-k/-v {cache_override} honored for KV estimate") + if split_mode_override is not None: + logger.info( + f"User --split-mode {split_mode_override} honored; " + "reconciled into tensor_parallel state" + ) effective_ctx = requested_ctx if requested_ctx > 0 else (self._context_length or 0) max_available_ctx = self._context_length or effective_ctx gpus: list[tuple[int, int]] = [] + # Keep fit-budget and launch-flag mmproj resolution in sync. + launch_mmproj_path = None + if not extra_args_disable_mmproj(extra_args): + launch_mmproj_path = self._resolve_launch_mmproj_path( + model_path = model_path, + mmproj_path = mmproj_path, + ) + # Need both a resolved mmproj AND the config vision flag; a stray + # mmproj passing the family-name heuristic must not flip a non-VLM + # GGUF into vision mode. + effective_is_vision = bool(launch_mmproj_path) and bool(is_vision) + if is_vision and not effective_is_vision: + logger.warning( + "Vision-capable GGUF loaded without a usable mmproj; " + "image input will be disabled for this session" + ) try: - model_size = self._get_gguf_size_bytes(model_path) + gguf_size = self._get_gguf_size_bytes(model_path) + # Include GPU-loaded mmproj in the fit budget (#5825). + mmproj_size = ( + self._mmproj_vram_bytes(launch_mmproj_path) if effective_is_vision else 0 + ) + model_size = gguf_size + mmproj_size gpus = self._get_gpu_free_memory() # Resolve effective context: 0 means let llama-server use @@ -3065,6 +3650,8 @@ class LlamaCppBackend: # Auto n_ctx=0 (native): prefer fewer GPUs with reduced # context, since multi-GPU is slower. gpu_indices, use_fit = None, True + # Per-GPU weight proportions for tensor mode (None = even). + tp_tensor_split: Optional[list[int]] = None explicit_ctx = requested_ctx > 0 # MTP draft model lives outside the main estimates; carve # its reserve out of every fit budget and pin threshold so @@ -3072,10 +3659,67 @@ class LlamaCppBackend: _mtp_reserve = _MTP_VRAM_RESERVE_FRAC if _mtp_will_engage else 0.0 _pin_fraction = self._GPU_PIN_VRAM_FRACTION - _mtp_reserve - if gpus and self._can_estimate_kv() and effective_ctx > 0: - # Largest hardware-aware cap from the native context - # across all usable GPU subsets (for UI bounds), - # independent of the requested context. + # Tensor mode allocates a compute-graph buffer on every + # participating GPU, so a GPU with less free VRAM than that + # reserve can't host it and would OOM at load. Drop those + # from the tensor-parallel set up front (gpu_indices below + # becomes the CUDA_VISIBLE_DEVICES mask, so they're excluded + # from llama-server entirely, not just given zero weight). + tp_gpus = gpus + if tensor_parallel: + reserve_mib = self._TENSOR_PARALLEL_BUFFER_RESERVE_MIB + tp_gpus = [g for g in gpus if g[1] >= reserve_mib] + + if tensor_parallel and len(tp_gpus) < 2: + # Tensor parallelism needs >= 2 usable GPUs. On a single + # GPU --split-mode tensor is a no-op; with 0 GPUs (CPU-only + # or probe failed) it must not reach llama-server; and a + # GPU below the buffer reserve can't participate. Drop the + # flag and fall through to normal layer/CPU allocation. + logger.info( + "Tensor parallelism requested but only %d of %d GPU(s) " + "have enough free VRAM for the compute buffer; " + "ignoring (needs >= 2).", + len(tp_gpus), + len(gpus), + ) + tensor_parallel = False + # A user --split-mode tensor in extras is appended after + # Studio's flags, so it would still reach llama-server and + # fail here; strip it so the downgrade actually applies. + extra_args = strip_split_mode_only(extra_args) + + if tensor_parallel and tp_gpus: + # Tensor-parallel allocation: use all usable GPUs, weight + # the split by (free - buffer), and cap context to the + # pooled VRAM after weights + per-device compute-graph + # buffers. See _plan_tensor_parallel for the policy. + target_ctx = ( + effective_ctx + if explicit_ctx + else (self._context_length or effective_ctx) + ) + ( + effective_ctx, + max_available_ctx, + gpu_indices, + tp_tensor_split, + ) = self._plan_tensor_parallel( + tp_gpus, + model_size, + target_ctx, + cache_type_kv = cache_type_kv, + n_parallel = n_parallel, + mtp_engaged = _mtp_will_engage, + # Report the UI ceiling from native ctx, not the + # explicit small request. + max_target_ctx = self._context_length or target_ctx, + ) + use_fit = False + elif gpus and self._can_estimate_kv() and effective_ctx > 0: + # Compute the largest hardware-aware cap from the model's + # native context across all usable GPU subsets (for UI + # bounds), independent of the currently requested context. native_ctx_for_cap = self._context_length or effective_ctx if native_ctx_for_cap > 0: ranked_for_cap = sorted(gpus, key = lambda g: g[1], reverse = True) @@ -3189,8 +3833,12 @@ class LlamaCppBackend: kv_cache_bytes = self._estimate_kv_cache_bytes( effective_ctx, cache_type_kv, n_parallel = n_parallel ) + mmproj_note = ( + f"mmproj: {mmproj_size / (1024**3):.1f} GB, " if mmproj_size else "" + ) logger.info( - f"GGUF size: {model_size / (1024**3):.1f} GB, " + f"GGUF size: {gguf_size / (1024**3):.1f} GB, " + f"{mmproj_note}" f"est. KV cache: {kv_cache_bytes / (1024**3):.1f} GB, " f"context: {effective_ctx}, " f"GPUs free: {gpus}, selected: {gpu_indices}, fit: {use_fit}" @@ -3198,24 +3846,9 @@ class LlamaCppBackend: except Exception as e: logger.warning(f"GPU selection failed ({e}), using --fit on") gpu_indices, use_fit = None, True + tp_tensor_split = None effective_ctx = requested_ctx # fall back to original - launch_mmproj_path = None - if not extra_args_disable_mmproj(extra_args): - launch_mmproj_path = self._resolve_launch_mmproj_path( - model_path = model_path, - mmproj_path = mmproj_path, - ) - # Need both a resolved mmproj AND the config vision flag; a stray - # mmproj passing the family-name heuristic must not flip a non-VLM - # GGUF into vision mode. - effective_is_vision = bool(launch_mmproj_path) and bool(is_vision) - if is_vision and not effective_is_vision: - logger.warning( - "Vision-capable GGUF loaded without a usable mmproj; " - "image input will be disabled for this session" - ) - # Audio input straight from the mmproj (clip.has_audio_encoder), # independent of token names. self._mmproj_has_audio = False @@ -3297,6 +3930,27 @@ class LlamaCppBackend: else: self._cache_type_kv = None + # Tensor parallelism: split the model across GPUs by tensor + # rather than by layer. Multi-GPU only -- a no-op on a single + # GPU. Default (layer split) is left implicit by omitting the + # flag. See llama.cpp --split-mode. + if tensor_parallel: + cmd.extend(["--split-mode", "tensor"]) + if tp_tensor_split and len(tp_tensor_split) > 1: + cmd.extend( + [ + "--tensor-split", + ",".join(str(int(x)) for x in tp_tensor_split), + ] + ) + self._tensor_parallel = True + logger.info( + "Tensor parallelism: --split-mode tensor, --tensor-split %s", + tp_tensor_split, + ) + else: + self._tensor_parallel = False + # Speculative decoding. See _build_speculative_flags for the # mode resolution, benchmarks, and llama.cpp references. launch_mtp_draft_path = self._resolve_launch_mtp_path( @@ -3335,6 +3989,7 @@ class LlamaCppBackend: self._chat_template_file = tempfile.NamedTemporaryFile( mode = "w", + encoding = "utf-8", suffix = ".jinja", delete = False, prefix = "unsloth_chat_template_", @@ -3356,6 +4011,13 @@ class LlamaCppBackend: thinking_default = False self._reasoning_default = thinking_default reasoning_kw = self._reasoning_kwargs(thinking_default) + # preserve_thinking is an independent kwarg. Default it OFF + # at launch so direct OpenAI-compatible callers that omit the + # field match the UI's default-off behavior (the bundled + # gemma-4 template also defaults it false; the frontend sends + # preserve_thinking per request once toggled on). + if self._supports_preserve_thinking: + reasoning_kw["preserve_thinking"] = False cmd.extend( [ "--chat-template-kwargs", @@ -3406,6 +4068,14 @@ class LlamaCppBackend: env.setdefault("GGML_CUDA_ENABLE_UNIFIED_MEMORY", "1") logger.info("AMD unified-memory APU: set GGML_CUDA_ENABLE_UNIFIED_MEMORY=1") + # DC NVIDIA GPUs: FP32 accum (+ P2P / launch queues for multi-GPU). + # See _apply_datacenter_env; opt out with UNSLOTH_DISABLE_DC_TUNING=1. + if self._apply_datacenter_env(env, gpu_indices): + multi_gpu = self._effective_gpu_count(gpu_indices) > 1 + logger.info( + f"Data-center GPU detected: applied DC llama.cpp env tuning (multi_gpu={multi_gpu})" + ) + if sys.platform == "win32": # Ordering: see _build_windows_path_dirs. #5106. path_dirs = self._build_windows_path_dirs( @@ -3431,7 +4101,15 @@ class LlamaCppBackend: # plus CUDA runtime libs (libcudart, libcublas, etc.) import platform - lib_dirs = [binary_dir] + lib_dirs = [] + # WSL: system HIP before the bundle's (which segfaults on + # /dev/dxg). Mirror install_llama_prebuilt.binary_env, which + # validates the prebuilt with this same ordering. + for _wsl_rocm in _wsl_system_rocm_lib_dirs(): + lib_dirs.append(_wsl_rocm) + if lib_dirs: + env.setdefault("HSA_ENABLE_DXG_DETECTION", "1") + lib_dirs.append(binary_dir) _arch = platform.machine() # x86_64, aarch64, etc. # Pip-installed nvidia CUDA runtime libs. The prebuilt @@ -4065,6 +4743,7 @@ class LlamaCppBackend: is_vision: bool, gguf_path: Optional[str] = None, spec_draft_n_max: Optional[int] = None, + tensor_parallel: bool = False, mtp_draft_path: Optional[str] = None, ) -> bool: """True iff the live server already satisfies these load kwargs. @@ -4101,6 +4780,12 @@ class LlamaCppBackend: if _norm(self._cache_type_kv) != _norm(cache_type_kv): return False + # Reconcile a user --split-mode in extras (load_model does the same), so + # an extras-driven tensor load isn't seen as a mismatch that needlessly + # kills/reloads a healthy server. + if self._tensor_parallel != resolve_tensor_parallel(extra_args, tensor_parallel): + return False + # Compare on the canonical requested mode. With --spec-type in # extra_args the backend stores None; mirror that here. if _extra_args_set_spec_type(extra_args): @@ -4172,6 +4857,12 @@ class LlamaCppBackend: return None return saw_gpu_buffer + def load_cancelled(self) -> bool: + """True if a load was cancelled (e.g. via unload/_cancel_event) and not + yet consumed by the next load_model. Lets the tensor->layer fallback + avoid restarting a load the user just cancelled.""" + return self._cancel_event.is_set() + def unload_model(self) -> bool: """Terminate the subprocess and cancel any in-flight download.""" self._cancel_event.set() @@ -4204,6 +4895,7 @@ class LlamaCppBackend: self._supports_preserve_thinking = False self._supports_tools = False self._cache_type_kv = None + self._tensor_parallel = False self._speculative_type = None self._requested_spec_mode = None self._spec_draft_n_max = None @@ -4612,28 +5304,84 @@ class LlamaCppBackend: @staticmethod def _iter_text_cancellable( - response: "httpx.Response", cancel_event: Optional[threading.Event] = None + response: "httpx.Response", + cancel_event: Optional[threading.Event] = None, + stall_timeout_s: float = _DEFAULT_STREAM_STALL_TIMEOUT_S, + first_token_deadline: Optional[float] = None, + post_first_chunk_read_timeout_s: Optional[float] = _DEFAULT_STREAM_STALL_TIMEOUT_S, ) -> Generator[str, None, None]: - """Iterate an httpx streaming response with cancel support. - - Checks cancel_event between chunks and on ReadTimeout; the - _stream_with_retry watcher also closes the response on cancel. - """ + """Iterate a stream while polling cancel and stall timeouts.""" text_iter = response.iter_text() + if first_token_deadline is None: + first_token_deadline = time.monotonic() + _DEFAULT_FIRST_TOKEN_TIMEOUT_S + last_chunk_at: Optional[float] = None while True: if cancel_event is not None and cancel_event.is_set(): response.close() return try: + if last_chunk_at is None: + remaining_s = first_token_deadline - time.monotonic() + if remaining_s <= 0: + raise httpx.ReadTimeout("The model did not produce a first token in time.") + LlamaCppBackend._set_stream_read_timeout(response, remaining_s) chunk = next(text_iter) + if chunk: + if last_chunk_at is None and post_first_chunk_read_timeout_s is not None: + LlamaCppBackend._set_stream_read_timeout( + response, + post_first_chunk_read_timeout_s, + ) + last_chunk_at = time.monotonic() yield chunk except StopIteration: return except httpx.ReadTimeout: - # No data within the timeout window -- loop back and re-check - # cancel_event. + now = time.monotonic() + if last_chunk_at is None: + if now >= first_token_deadline: + raise + elif now - last_chunk_at >= stall_timeout_s: + raise httpx.ReadTimeout("The model stopped producing tokens mid-response.") continue + @staticmethod + def _set_stream_read_timeout(response: "httpx.Response", read_timeout_s: float) -> None: + """Lower only post-header stream reads; keep prefill timeout long.""" + try: + timeout_ext = response.request.extensions.get("timeout") + if isinstance(timeout_ext, dict): + timeout_ext["read"] = read_timeout_s + except Exception: + logger.debug("Could not lower response read timeout", exc_info = True) + + @staticmethod + def _shutdown_active_httpx_sockets(client: "httpx.Client") -> None: + """Best-effort interrupt for a sync httpx request blocked before headers.""" + try: + pool = getattr(getattr(client, "_transport", None), "_pool", None) + connections = list(getattr(pool, "_connections", []) or []) + for connection in connections: + inner = getattr(connection, "_connection", None) + stream = getattr(inner, "_network_stream", None) + sock = getattr(stream, "_sock", None) + if sock is None: + continue + try: + sock.shutdown(socket.SHUT_RDWR) + except OSError: + pass + try: + sock.close() + except OSError: + pass + except Exception: + logger.debug("Could not shutdown active httpx socket", exc_info = True) + try: + client.close() + except Exception: + logger.debug("Could not close httpx client", exc_info = True) + @staticmethod @contextlib.contextmanager def _stream_with_retry( @@ -4642,38 +5390,28 @@ class LlamaCppBackend: payload: dict, cancel_event: Optional[threading.Event] = None, headers: Optional[dict] = None, + first_token_deadline: Optional[float] = None, ): - """Open an httpx streaming POST with cancel support. - - Sends once with a long read timeout (120 s) so prefill finishes without - a retry storm (the old 0.5 s timeout caused duplicate POSTs every half - second). A watcher thread cancels by closing the response. httpx can't - interrupt a blocked read before the response exists, so cancel during - the header wait (1-5 s prefill) is deferred until headers arrive. - """ + """Open one streaming POST and let cancel interrupt prefill or reads.""" if cancel_event is not None and cancel_event.is_set(): raise GeneratorExit - # Background watcher: close the response if cancel is requested. - # Only effective after response headers arrive (httpx limitation). _cancel_closed = threading.Event() _response_ref: list = [None] def _cancel_watcher(): while not _cancel_closed.is_set(): if cancel_event.wait(timeout = 0.3): - # Cancel requested. Poll until the response object exists - # so we can close it, or until the main thread finishes - # (_cancel_closed set in finally). while not _cancel_closed.is_set(): r = _response_ref[0] - if r is not None: - try: + try: + if r is not None: r.close() - return - except Exception as e: - logger.debug(f"Error closing response in cancel watcher: {e}") - # Response not created yet -- wait briefly and retry + else: + LlamaCppBackend._shutdown_active_httpx_sockets(client) + return + except Exception as e: + logger.debug(f"Error closing request in cancel watcher: {e}") _cancel_closed.wait(timeout = 0.1) return @@ -4683,12 +5421,12 @@ class LlamaCppBackend: watcher.start() try: - # Long read timeout so prefill can finish without a retry storm. - # Cancel during prefill and streaming is handled by the watcher - # thread closing the response, unblocking any httpx read. + if first_token_deadline is None: + first_token_deadline = time.monotonic() + _DEFAULT_FIRST_TOKEN_TIMEOUT_S + prefill_read_timeout = max(0.1, first_token_deadline - time.monotonic()) prefill_timeout = httpx.Timeout( connect = 30, - read = 120.0, + read = prefill_read_timeout, write = 10, pool = 10, ) @@ -4704,7 +5442,7 @@ class LlamaCppBackend: raise GeneratorExit yield response return - except (httpx.ReadError, httpx.RemoteProtocolError, httpx.CloseError): + except (httpx.RequestError, RuntimeError): # Response was closed by the cancel watcher if cancel_event is not None and cancel_event.is_set(): raise GeneratorExit @@ -4759,14 +5497,12 @@ class LlamaCppBackend: ) if _reasoning_kw is not None: payload["chat_template_kwargs"] = _reasoning_kw - # Cap to the effective context length when known, else the floor. - # The wall-clock backstop below stops a stuck model regardless. + # Default cap to the model context when known. payload["max_tokens"] = ( max_tokens if max_tokens is not None else (self._effective_context_length or _DEFAULT_MAX_TOKENS_FLOOR) ) - payload["t_max_predict_ms"] = _DEFAULT_T_MAX_PREDICT_MS if stop: payload["stop"] = stop if seed is not None: @@ -4782,20 +5518,20 @@ class LlamaCppBackend: _metadata_finish_reason = None try: - # _stream_with_retry uses a 120 s read timeout so prefill can - # finish. Cancel during streaming is handled by the watcher - # thread (closes the response on cancel_event). + # Prefill can use the long first-token timeout; body reads are lowered after headers. stream_timeout = httpx.Timeout(connect = 10, read = 0.5, write = 10, pool = 10) _auth_headers = {"Authorization": f"Bearer {self._api_key}"} if self._api_key else None with httpx.Client( timeout = stream_timeout, limits = httpx.Limits(max_keepalive_connections = 0) ) as client: + first_token_deadline = time.monotonic() + _DEFAULT_FIRST_TOKEN_TIMEOUT_S with self._stream_with_retry( client, url, payload, cancel_event, headers = _auth_headers, + first_token_deadline = first_token_deadline, ) as response: if response.status_code != 200: error_body = response.read().decode() @@ -4806,7 +5542,11 @@ class LlamaCppBackend: buffer = "" has_content_tokens = False reasoning_text = "" - for raw_chunk in self._iter_text_cancellable(response, cancel_event): + for raw_chunk in self._iter_text_cancellable( + response, + cancel_event, + first_token_deadline = first_token_deadline, + ): buffer += raw_chunk while "\n" in buffer: line, buffer = buffer.split("\n", 1) @@ -4834,6 +5574,12 @@ class LlamaCppBackend: try: data = json.loads(line[6:]) + # Diffusion frame (per-step canvas) from the shim: forward untouched so + # the frontend renders it in place. No assistant text, so it never enters + # the cumulative content. + if data.get("type") == "diffusion_frame": + yield data + continue # Capture server timings/usage from final chunks. _chunk_timings = data.get("timings") if _chunk_timings: @@ -5030,7 +5776,6 @@ class LlamaCppBackend: if max_tokens is not None else (self._effective_context_length or _DEFAULT_MAX_TOKENS_FLOOR) ) - payload["t_max_predict_ms"] = _DEFAULT_T_MAX_PREDICT_MS if stop: payload["stop"] = stop if seed is not None: @@ -5076,12 +5821,14 @@ class LlamaCppBackend: timeout = stream_timeout, limits = httpx.Limits(max_keepalive_connections = 0), ) as client: + first_token_deadline = time.monotonic() + _DEFAULT_FIRST_TOKEN_TIMEOUT_S with self._stream_with_retry( client, url, payload, cancel_event, headers = _auth_headers, + first_token_deadline = first_token_deadline, ) as response: if response.status_code != 200: error_body = response.read().decode() @@ -5093,6 +5840,7 @@ class LlamaCppBackend: for raw_chunk in self._iter_text_cancellable( response, cancel_event, + first_token_deadline = first_token_deadline, ): raw_buf += raw_chunk while "\n" in raw_buf: @@ -5742,7 +6490,6 @@ class LlamaCppBackend: if max_tokens is not None else (self._effective_context_length or _DEFAULT_MAX_TOKENS_FLOOR) ) - stream_payload["t_max_predict_ms"] = _DEFAULT_T_MAX_PREDICT_MS if stop: stream_payload["stop"] = stop if seed is not None: @@ -5765,12 +6512,14 @@ class LlamaCppBackend: with httpx.Client( timeout = stream_timeout, limits = httpx.Limits(max_keepalive_connections = 0) ) as client: + first_token_deadline = time.monotonic() + _DEFAULT_FIRST_TOKEN_TIMEOUT_S with self._stream_with_retry( client, url, stream_payload, cancel_event, headers = _auth_headers, + first_token_deadline = first_token_deadline, ) as response: if response.status_code != 200: error_body = response.read().decode() @@ -5779,7 +6528,11 @@ class LlamaCppBackend: ) buffer = "" - for raw_chunk in self._iter_text_cancellable(response, cancel_event): + for raw_chunk in self._iter_text_cancellable( + response, + cancel_event, + first_token_deadline = first_token_deadline, + ): buffer += raw_chunk while "\n" in buffer: line, buffer = buffer.split("\n", 1) diff --git a/studio/backend/core/inference/llama_server_args.py b/studio/backend/core/inference/llama_server_args.py index 00f8c66d5c..69a86fa3ba 100644 --- a/studio/backend/core/inference/llama_server_args.py +++ b/studio/backend/core/inference/llama_server_args.py @@ -109,6 +109,7 @@ def validate_extra_args(args: Optional[Iterable[str]]) -> list[str]: out.append(token) parse_ctx_override(out) parse_cache_override(out) + parse_split_mode_override(out) return out @@ -157,8 +158,20 @@ _TEMPLATE_FLAGS: frozenset[str] = frozenset( "--no-jinja", } ) +# Multi-GPU split mode shadows the Tensor Parallelism toggle +# (--split-mode tensor). Pass-through stays allowed so users keep the +# row/none/layer modes the toggle doesn't expose, but it's stripped on +# inherit and reconciled into the round-tripped tensor_parallel state. +# --tensor-split is coupled to the split mode and is stripped with it: Studio +# owns the tensor-mode split ratios, so an inherited/stale --tensor-split must +# not last-wins-override Studio's computed asymmetric split. +_SPLIT_MODE_FLAGS: frozenset[str] = frozenset({"-sm", "--split-mode"}) +_TENSOR_SPLIT_FLAGS: frozenset[str] = frozenset({"-ts", "--tensor-split"}) +_SPLIT_SHADOWING_FLAGS: frozenset[str] = _SPLIT_MODE_FLAGS | _TENSOR_SPLIT_FLAGS -_SHADOWING_FLAGS: frozenset[str] = _CONTEXT_FLAGS | _CACHE_FLAGS | _SPEC_FLAGS | _TEMPLATE_FLAGS +_SHADOWING_FLAGS: frozenset[str] = ( + _CONTEXT_FLAGS | _CACHE_FLAGS | _SPEC_FLAGS | _TEMPLATE_FLAGS | _SPLIT_SHADOWING_FLAGS +) # Shadowing flags that take no value -- strip the flag only, not the next token. _BOOLEAN_SHADOWING_FLAGS: frozenset[str] = frozenset({"--spec-default", "--jinja", "--no-jinja"}) @@ -213,11 +226,12 @@ def resolve_requested_ctx(args: Optional[Iterable[str]], fallback_n_ctx: int) -> return override if override is not None else fallback_n_ctx -def parse_cache_override(args: Optional[Iterable[str]]) -> Optional[str]: - """Return the last-wins cache type if extras pass cache flags. +def _last_flag_value(args: Optional[Iterable[str]], flags: frozenset[str]) -> Optional[str]: + """Return the last-wins string value among ``flags`` in extras, or None. - Recognises -ctk (key) and -ctv (value); treats both as one setting, - since Studio's KV estimate has a single cache_type_kv knob. + Handles both ``--flag=value`` and ``--flag value`` forms and raises if a + matched flag has no (or an empty) value. Shared by the single-knob + last-wins parsers (cache type, split mode). """ if not args: return None @@ -228,7 +242,7 @@ def parse_cache_override(args: Optional[Iterable[str]]) -> Optional[str]: while i < n: tok = tokens[i] flag = _flag_name(tok) - if flag is None or flag not in _CACHE_FLAGS: + if flag is None or flag not in flags: i += 1 continue @@ -249,6 +263,17 @@ def parse_cache_override(args: Optional[Iterable[str]]) -> Optional[str]: return override +def parse_cache_override(args: Optional[Iterable[str]]) -> Optional[str]: + """Return the last-wins cache type if extras pass cache flags. + + Mirrors parse_ctx_override but for cache type. Recognises both -ctk + (key) and -ctv (value). When both flags appear, returns the last-wins + value, treating key and value cache flags as the same setting because + Studio's KV estimate has a single cache_type_kv knob. + """ + return _last_flag_value(args, _CACHE_FLAGS) + + def resolve_cache_type_kv( args: Optional[Iterable[str]], fallback_cache_type_kv: Optional[str] ) -> Optional[str]: @@ -260,6 +285,30 @@ def resolve_cache_type_kv( return override if override is not None else fallback_cache_type_kv +def parse_split_mode_override(args: Optional[Iterable[str]]) -> Optional[str]: + """Return the last-wins ``--split-mode`` / ``-sm`` value from extras. + + Mirrors parse_cache_override for the multi-GPU split mode. Returns the + raw mode string (e.g. ``tensor`` / ``row`` / ``none`` / ``layer``), or + None when extras don't set it. + """ + return _last_flag_value(args, _SPLIT_MODE_FLAGS) + + +def resolve_tensor_parallel(args: Optional[Iterable[str]], fallback_tensor_parallel: bool) -> bool: + """Return the tensor-parallel state load_model should treat as requested. + + A user-supplied ``--split-mode`` in extras last-wins-overrides the + toggle, so reconcile it back into the boolean: any explicit split mode + means tensor-parallel is on iff that mode is ``tensor``. Falls back to + the toggle value when extras don't set it. + """ + override = parse_split_mode_override(args) + if override is None: + return fallback_tensor_parallel + return override.strip().lower() == "tensor" + + _MMPROJ_DISABLE_FLAGS: frozenset[str] = frozenset({"--no-mmproj", "--no-mmproj-auto"}) _MMPROJ_ENABLE_FLAGS: frozenset[str] = frozenset({"--mmproj-auto"}) @@ -289,12 +338,15 @@ def strip_shadowing_flags( strip_cache: bool = True, strip_spec: bool = True, strip_template: bool = True, + strip_split_mode: bool = True, ) -> list[str]: """Strip flags that shadow first-class Studio settings. Used when inheriting a previous load's ``llama_extra_args`` so an - inherited `-c 4096` can't override the current `max_seq_length` (same for - cache / spec / template). Each ``strip_*`` toggle controls one group. + inherited `-c 4096` can't override the current `max_seq_length` + (same for cache / spec / template / split-mode). Each ``strip_*`` + toggle controls one group; the route only strips groups whose + first-class field the caller actually supplied. """ shadowing: set[str] = set() if strip_context: @@ -305,6 +357,8 @@ def strip_shadowing_flags( shadowing |= _SPEC_FLAGS if strip_template: shadowing |= _TEMPLATE_FLAGS + if strip_split_mode: + shadowing |= _SPLIT_SHADOWING_FLAGS tokens = [str(a) for a in (args or [])] out: list[str] = [] @@ -325,3 +379,20 @@ def strip_shadowing_flags( else: i += 1 return out + + +def strip_split_mode_only(args: Optional[Iterable[str]]) -> Optional[list[str]]: + """Remove the split-mode group (``--split-mode`` / ``-sm`` and the coupled + ``--tensor-split`` / ``-ts``) from ``args``, keeping every other shadow flag. + Preserves a None/empty input so the inherit-vs-explicit-empty distinction + survives. Used where tensor mode is being forced off (downgrade / fallback).""" + if not args: + return args + return strip_shadowing_flags( + args, + strip_context = False, + strip_cache = False, + strip_spec = False, + strip_template = False, + strip_split_mode = True, + ) diff --git a/studio/backend/core/inference/mcp_client.py b/studio/backend/core/inference/mcp_client.py index a0d79bbdf4..5a36d90c5d 100644 --- a/studio/backend/core/inference/mcp_client.py +++ b/studio/backend/core/inference/mcp_client.py @@ -8,6 +8,7 @@ import json import os import shlex import sys +import time from typing import Any, Optional from loggers import get_logger @@ -16,6 +17,14 @@ logger = get_logger(__name__) MCP_TOOL_PREFIX = "mcp__" +# A failed probe isn't cached (a recovered server must come back), but it's +# recorded so a down server isn't re-probed -- and the chat send re-hung for +# the full timeout -- on every message. Cool off for this long after a failure; +# much longer for OAuth, whose probe can hang up to _OAUTH_PROBE_TIMEOUT, +# so that hang doesn't recur every minute. +FAILED_PROBE_COOLOFF_SECONDS = 60.0 +OAUTH_FAILED_PROBE_COOLOFF_SECONDS = 300.0 + _oauth_token_store = None @@ -227,6 +236,53 @@ async def list_tools_async( return await asyncio.wait_for(_fetch(), timeout = timeout) +# Discovered-tool cache, keyed by MCP server id. get_enabled_mcp_tools() +# probes a server only on a cache miss, keeping MCP discovery off the chat +# send's critical path -- tool schemas are stable within a session. The +# /refresh route warms it; a URL/header/OAuth change or a delete evicts it. +# Successful probes are cached indefinitely. +_tool_cache: dict[str, list[dict]] = {} + +# server_id -> monotonic time before which a failed server must not be +# re-probed (see record_probe_failure). Cleared on a successful probe or +# eviction. +_probe_cooloff_until: dict[str, float] = {} + +# MCP server fields whose change invalidates a server's discovered tools: the +# endpoint/auth used to probe it (url, headers, oauth) or whether it's used at +# all (is_enabled). A rename does not. The update route's eviction and +# get_enabled_mcp_tools' mid-probe guard both key off this so they can't drift. +TOOL_CACHE_INVALIDATING_FIELDS = frozenset({"url", "headers_json", "use_oauth", "is_enabled"}) + + +def get_cached_tools(server_id: str) -> Optional[list[dict]]: + return _tool_cache.get(server_id) + + +def cache_tools(server_id: str, tools: list[dict]) -> None: + _tool_cache[server_id] = tools + _probe_cooloff_until.pop(server_id, None) + + +def record_probe_failure(server_id: str, use_oauth: bool = False) -> None: + cooloff = OAUTH_FAILED_PROBE_COOLOFF_SECONDS if use_oauth else FAILED_PROBE_COOLOFF_SECONDS + _probe_cooloff_until[server_id] = time.monotonic() + cooloff + + +def in_failure_cooloff(server_id: str) -> bool: + return _probe_cooloff_until.get(server_id, 0.0) > time.monotonic() + + +def invalidate_tool_cache(server_id: Optional[str] = None) -> None: + """Evict one server's cached tools, or every entry when server_id is None.""" + if server_id is None: + _tool_cache.clear() + _probe_cooloff_until.clear() + else: + _tool_cache.pop(server_id, None) + _probe_cooloff_until.pop(server_id, None) + + def _flatten_result(result: Any) -> str: parts = [] for block in getattr(result, "content", None) or []: diff --git a/studio/backend/core/inference/providers.py b/studio/backend/core/inference/providers.py index f7a2da2833..5b72373c03 100644 --- a/studio/backend/core/inference/providers.py +++ b/studio/backend/core/inference/providers.py @@ -249,6 +249,23 @@ PROVIDER_REGISTRY: dict[str, dict[str, Any]] = { # Surfaced via the frontend's CUSTOM_PROVIDER_PRESETS, not the dropdown. "hidden": True, }, + "custom": { + "display_name": "Custom", + # User-supplied via provider_base_url. + "base_url": "", + "default_models": [], + "supports_streaming": True, + "supports_vision": True, + "supports_tool_calling": True, + "auth_header": "Authorization", + "auth_prefix": "Bearer ", + "notes": ( + "User-supplied OpenAI-compatible server. Routed to " + "/v1/chat/completions; /models is optional." + ), + # Surfaced by the frontend's generic Custom option, not the dropdown. + "hidden": True, + }, "ollama": { "display_name": "Ollama", "base_url": "http://localhost:11434/v1", diff --git a/studio/backend/core/inference/tensor_fallback.py b/studio/backend/core/inference/tensor_fallback.py new file mode 100644 index 0000000000..73687165b8 --- /dev/null +++ b/studio/backend/core/inference/tensor_fallback.py @@ -0,0 +1,70 @@ +# SPDX-License-Identifier: AGPL-3.0-only +# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0 + +"""Tensor-parallel -> layer-split auto-fallback for GGUF loads. + +Kept in its own module (no FastAPI / httpx deps) so the orchestration can be +unit-tested with a fake loader, without a GPU or a running llama-server. +""" + +from __future__ import annotations + +import logging +from typing import Awaitable, Callable, Optional + +from core.inference.llama_server_args import ( + resolve_tensor_parallel, + strip_split_mode_only, +) + +logger = logging.getLogger(__name__) + + +async def load_with_tensor_fallback( + attempt_load: Callable[[bool, Optional[list[str]]], Awaitable[bool]], + *, + requested_tensor: bool, + extra_args: Optional[list[str]], + label: str = "", + cancelled: Optional[Callable[[], bool]] = None, +) -> bool: + """Run a GGUF load with the tensor-parallel -> layer-split auto-fallback. + + ``attempt_load(tensor_parallel, extra_args)`` performs one load and returns + True on success; it *raises* on a hard crash (llama-server aborts on some + archs / older builds), which is treated the same as a False return. + + Tensor mode can be requested by the toggle or by a ``--split-mode tensor`` + in ``extra_args`` (an allowed shadow flag), so the retry is keyed on whether + tensor mode is actually engaged, and it strips ``--split-mode`` from the + extras so the layer retry can't relaunch the same failing tensor load. A + non-tensor load keeps its original contract and propagates exceptions. + + ``cancelled()`` distinguishes a real tensor-start failure from a user + cancellation: ``attempt_load`` also returns False when the load was + cancelled, so without this the helper would restart a load the user just + cancelled. + """ + tensor_requested = resolve_tensor_parallel(extra_args, requested_tensor) + try: + success = await attempt_load(requested_tensor, extra_args) + except Exception as exc: + if not tensor_requested: + raise + logger.warning("Tensor-parallel load raised for '%s': %s", label, exc) + success = False + + if success or not tensor_requested: + return success + + # The first attempt returned False because the user cancelled, not because + # tensor mode is unsupported -- do not relaunch the cancelled load. + if cancelled is not None and cancelled(): + return success + + logger.warning( + "Tensor-parallel load failed for '%s'; retrying with layer split " + "(this model may not support tensor parallelism)", + label, + ) + return await attempt_load(False, strip_split_mode_only(extra_args)) diff --git a/studio/backend/core/inference/tools.py b/studio/backend/core/inference/tools.py index b27fa6ff73..43c9610282 100644 --- a/studio/backend/core/inference/tools.py +++ b/studio/backend/core/inference/tools.py @@ -24,11 +24,16 @@ import urllib.request from core.inference.mcp_client import ( MCP_TOOL_PREFIX, + TOOL_CACHE_INVALIDATING_FIELDS, + cache_tools, call_tool_sync, + get_cached_tools, + in_failure_cooloff, is_stdio, list_tools_async, parse_server_headers, probe_timeout, + record_probe_failure, stdio_mcp_enabled, ) from storage import mcp_servers_db @@ -634,28 +639,56 @@ async def get_enabled_mcp_tools() -> list[dict]: if not servers: return [] - results = await asyncio.gather( - *( - list_tools_async( - url = s["url"], - headers = parse_server_headers(s), - timeout = probe_timeout(s["url"], bool(s.get("use_oauth"))), - use_oauth = bool(s.get("use_oauth")), - ) - for s in servers - ), - return_exceptions = True, - ) + # Skip servers still in their post-failure cool-off, otherwise a down + # server gets re-probed -- and blocks the send for the full timeout -- on + # every message. + uncached = [ + s for s in servers if get_cached_tools(s["id"]) is None and not in_failure_cooloff(s["id"]) + ] + if uncached: + results = await asyncio.gather( + *( + list_tools_async( + url = s["url"], + headers = parse_server_headers(s), + timeout = probe_timeout(s["url"], bool(s.get("use_oauth"))), + use_oauth = bool(s.get("use_oauth")), + ) + for s in uncached + ), + return_exceptions = True, + ) + # An edit/delete can land while we await a probe (up to 305 s for + # OAuth); its cache eviction is a no-op against an entry we haven't + # written yet. Re-read and drop a result whose server changed or + # was removed mid-probe, else a stale tool list caches indefinitely. + current = {s["id"]: s for s in mcp_servers_db.list_servers()} + for server, payload in zip(uncached, results): + # Guard the failure branch too: a stale failure must not park a + # cool-off on the fresh config, or the server the user just fixed + # is skipped for the whole window. + fresh = current.get(server["id"]) + if fresh is None or any( + fresh.get(k) != server.get(k) for k in TOOL_CACHE_INVALIDATING_FIELDS + ): + continue + if isinstance(payload, BaseException): + logger.warning( + "MCP server '%s' (%s) discovery failed: %s", + server.get("display_name") or server["id"], + server.get("url"), + payload, + ) + # Failures aren't cached, but record one so a down server + # isn't re-probed every send during the cool-off. + record_probe_failure(server["id"], bool(fresh.get("use_oauth"))) + continue + cache_tools(server["id"], payload) specs: list[dict] = [] - for server, payload in zip(servers, results): - if isinstance(payload, BaseException): - logger.warning( - "MCP server '%s' (%s) discovery failed: %s", - server.get("display_name") or server["id"], - server.get("url"), - payload, - ) + for server in servers: + payload = get_cached_tools(server["id"]) + if payload is None: continue specs.extend(_mcp_specs_for_server(server, payload)) return specs @@ -773,6 +806,7 @@ def _search_knowledge_base(arguments: dict, rag_scope: dict | None) -> str: query = str(query), scope_kb_id = scope.get("kb_id"), scope_thread_id = scope.get("thread_id"), + scope_project_id = scope.get("project_id"), top_k = top_k, **_scope_retrieval_kwargs(scope), ) @@ -880,6 +914,7 @@ def build_rag_autoinject(conversation: list[dict], rag_scope: dict | None) -> di query = query, scope_kb_id = rag_scope.get("kb_id"), scope_thread_id = rag_scope.get("thread_id"), + scope_project_id = rag_scope.get("project_id"), top_k = top_k, min_dense_score = floor, **_scope_retrieval_kwargs(rag_scope), diff --git a/studio/backend/core/rag/ingestion.py b/studio/backend/core/rag/ingestion.py index 5ac0639f11..c0c9a9f656 100644 --- a/studio/backend/core/rag/ingestion.py +++ b/studio/backend/core/rag/ingestion.py @@ -35,6 +35,22 @@ def _sha256_file(path: str) -> str: return h.hexdigest() +def _remove_upload(stored_path: str | None, *, keep_path: str | None = None) -> None: + if not stored_path: + return + try: + target = os.path.realpath(stored_path) + if keep_path is not None and target == os.path.realpath(keep_path): + return + from utils.paths import rag_uploads_root + + uploads = os.path.realpath(str(rag_uploads_root())) + if os.path.isfile(target) and os.path.commonpath([uploads, target]) == uploads: + os.remove(target) + except Exception: # noqa: BLE001 - upload cleanup must not block ingestion. + logger.warning("failed to remove RAG upload %s", stored_path, exc_info = True) + + def _emit(job_id: str, event: dict) -> None: with _jobs_lock: q = _jobs.get(job_id) @@ -152,6 +168,7 @@ def start_ingestion( filename: str, stored_path: str, *, + project_id: str | None = None, model_name: str | None = None, ) -> tuple[str, str]: """Create the document + job rows and spawn the worker, returning @@ -167,11 +184,15 @@ def start_ingestion( existing = store.document_by_hash(conn, scope, sha) if existing is not None: job_id = _new_job(conn, existing, scope, status = "completed", progress = 1.0) + _remove_upload(stored_path) with _jobs_lock: _jobs[job_id] = queue.Queue() _emit(job_id, {"type": "complete", "num_chunks": 0, "deduped": True}) _emit(job_id, None) return existing, job_id + for failed in store.failed_documents_by_hash(conn, scope, sha): + store.delete_document(conn, failed["id"]) + _remove_upload(failed.get("stored_path"), keep_path = stored_path) document_id = store.create_document( conn, @@ -180,6 +201,7 @@ def start_ingestion( sha256 = sha, kb_id = kb_id, thread_id = thread_id, + project_id = project_id, status = "pending", stored_path = stored_path, ) diff --git a/studio/backend/core/rag/retrieval.py b/studio/backend/core/rag/retrieval.py index b16e8b1c94..fe6a033a52 100644 --- a/studio/backend/core/rag/retrieval.py +++ b/studio/backend/core/rag/retrieval.py @@ -22,7 +22,7 @@ class Hit: def retrieve_lexical( conn: sqlite3.Connection, - scope: str, + scope: str | list[str], query: str, k: int | None = None, ) -> list[Hit]: @@ -32,7 +32,7 @@ def retrieve_lexical( def retrieve_dense( conn: sqlite3.Connection, - scope: str, + scope: str | list[str], query: str, k: int | None = None, *, @@ -69,7 +69,7 @@ def _rrf(rankings: list[list[Hit]], rrf_k: int, top_k: int) -> list[Hit]: def retrieve_hybrid( conn: sqlite3.Connection, - scope: str, + scope: str | list[str], query: str, *, k: int | None = None, diff --git a/studio/backend/core/rag/store.py b/studio/backend/core/rag/store.py index a77b9af7f8..7d58931e53 100644 --- a/studio/backend/core/rag/store.py +++ b/studio/backend/core/rag/store.py @@ -29,6 +29,15 @@ def thread_scope(thread_id: str) -> str: return f"thread_{thread_id}" +def project_scope(project_id: str) -> str: + return f"project_{project_id}" + + +def _scopes(scope) -> list[str]: + """Search helpers accept one scope or several (e.g. project + thread).""" + return [scope] if isinstance(scope, str) else list(scope) + + def _f32(vector) -> bytes: """Pack a vector into float32 bytes for vec0.""" return struct.pack(f"{len(vector)}f", *(float(x) for x in vector)) @@ -96,19 +105,21 @@ def create_document( sha256: str, kb_id: str | None = None, thread_id: str | None = None, + project_id: str | None = None, status: str = "pending", stored_path: str | None = None, document_id: str | None = None, ) -> str: document_id = document_id or str(uuid.uuid4()) conn.execute( - "INSERT INTO documents(id, scope, kb_id, thread_id, filename, sha256, status, " - "stored_path, created_at) VALUES(?,?,?,?,?,?,?,?,?)", + "INSERT INTO documents(id, scope, kb_id, thread_id, project_id, filename, sha256, " + "status, stored_path, created_at) VALUES(?,?,?,?,?,?,?,?,?,?)", ( document_id, scope, kb_id, thread_id, + project_id, filename, sha256, status, @@ -137,7 +148,8 @@ def set_document_status( def list_documents(conn: sqlite3.Connection, scope: str) -> list[dict]: rows = conn.execute( - "SELECT id, scope, kb_id, thread_id, filename, sha256, status, error, num_chunks, created_at " + "SELECT id, scope, kb_id, thread_id, project_id, filename, sha256, status, error, " + "num_chunks, created_at " "FROM documents WHERE scope=? ORDER BY created_at DESC", (scope,), ).fetchall() @@ -151,11 +163,21 @@ def get_document(conn: sqlite3.Connection, document_id: str) -> dict | None: def document_by_hash(conn: sqlite3.Connection, scope: str, sha256: str) -> str | None: row = conn.execute( - "SELECT id FROM documents WHERE scope=? AND sha256=?", (scope, sha256) + "SELECT id FROM documents WHERE scope=? AND sha256=? AND status!='failed' " + "ORDER BY created_at DESC LIMIT 1", + (scope, sha256), ).fetchone() return row["id"] if row else None +def failed_documents_by_hash(conn: sqlite3.Connection, scope: str, sha256: str) -> list[dict]: + rows = conn.execute( + "SELECT id, stored_path FROM documents WHERE scope=? AND sha256=? AND status='failed'", + (scope, sha256), + ).fetchall() + return [dict(r) for r in rows] + + def add_chunks( conn: sqlite3.Connection, scope: str, @@ -220,30 +242,41 @@ def delete_document(conn: sqlite3.Connection, document_id: str) -> None: conn.commit() -def search_lexical(conn: sqlite3.Connection, scope: str, query: str, k: int): - """BM25 lexical search. Returns [(chunk_id, score)], higher = better.""" +def search_lexical(conn: sqlite3.Connection, scope, query: str, k: int): + """BM25 lexical search over one scope or several. Returns + [(chunk_id, score)], higher = better.""" mq = _match_query(query) if not mq: return [] + scopes = _scopes(scope) + if not scopes: + return [] + placeholders = ",".join("?" * len(scopes)) rows = conn.execute( - "SELECT chunk_id, bm25(chunks_fts) AS s FROM chunks_fts " - "WHERE chunks_fts MATCH ? AND scope=? ORDER BY s LIMIT ?", - (mq, scope, k), + f"SELECT chunk_id, bm25(chunks_fts) AS s FROM chunks_fts " + f"WHERE chunks_fts MATCH ? AND scope IN ({placeholders}) ORDER BY s LIMIT ?", + (mq, *scopes, k), ).fetchall() # bm25() is negative (more negative = better); flip to higher-is-better. return [(r["chunk_id"], -r["s"]) for r in rows] -def search_dense(conn: sqlite3.Connection, scope: str, vector, k: int): - """Cosine KNN over vec0. Returns [(chunk_id, 1 - distance)].""" +def search_dense(conn: sqlite3.Connection, scope, vector, k: int): + """Cosine KNN over vec0 for one scope or several. Returns + [(chunk_id, 1 - distance)]. vec0 KNN constrains its partition key by + equality, so multi-scope runs one query per scope and merges by score.""" if not rag_db.vec_table_exists(conn): return [] - rows = conn.execute( - "SELECT chunk_id, distance FROM chunks_vec " - "WHERE scope=? AND embedding MATCH ? ORDER BY distance LIMIT ?", - (scope, _f32(vector), k), - ).fetchall() - return [(r["chunk_id"], 1.0 - r["distance"]) for r in rows] + out: list[tuple[str, float]] = [] + for s in _scopes(scope): + rows = conn.execute( + "SELECT chunk_id, distance FROM chunks_vec " + "WHERE scope=? AND embedding MATCH ? ORDER BY distance LIMIT ?", + (s, _f32(vector), k), + ).fetchall() + out.extend((r["chunk_id"], 1.0 - r["distance"]) for r in rows) + out.sort(key = lambda t: t[1], reverse = True) + return out[:k] def chunks_by_id(conn: sqlite3.Connection, ids) -> dict: diff --git a/studio/backend/core/rag/tool.py b/studio/backend/core/rag/tool.py index 1b0d87590a..ccb1b47e63 100644 --- a/studio/backend/core/rag/tool.py +++ b/studio/backend/core/rag/tool.py @@ -3,7 +3,8 @@ """``search_knowledge_base`` LLM tool: scope resolution + hit formatting. -KB scope wins over thread scope. Hits render as ```` blocks for the model, +KB scope wins; otherwise project and thread scopes combine so project chats also +see their own attachments. Hits render as ```` blocks for the model, plus a parallel citation source-map for clickable sources. Each call opens and closes its own ``rag_db`` connection. """ @@ -15,7 +16,7 @@ from xml.sax.saxutils import quoteattr from storage import rag_db from . import config, retrieval -from .store import kb_scope, thread_scope +from .store import kb_scope, project_scope, thread_scope SEARCH_KNOWLEDGE_BASE_TOOL = { "type": "function", @@ -42,12 +43,23 @@ SEARCH_KNOWLEDGE_BASE_TOOL = { } -def _resolve_scope(scope_kb_id: str | None, scope_thread_id: str | None) -> str | None: +def _resolve_scope( + scope_kb_id: str | None, + scope_thread_id: str | None, + scope_project_id: str | None = None, +) -> str | list[str] | None: + """KB (an explicit pick) is exclusive; project and thread scopes combine so a + project chat also retrieves from its own attached documents.""" if scope_kb_id: return kb_scope(scope_kb_id) + scopes = [] + if scope_project_id: + scopes.append(project_scope(scope_project_id)) if scope_thread_id: - return thread_scope(scope_thread_id) - return None + scopes.append(thread_scope(scope_thread_id)) + if not scopes: + return None + return scopes[0] if len(scopes) == 1 else scopes def _format(rows, hits) -> tuple[str, list[dict]]: @@ -83,6 +95,7 @@ def search_knowledge_base_with_sources( query: str, scope_kb_id: str | None = None, scope_thread_id: str | None = None, + scope_project_id: str | None = None, top_k: int | None = None, min_score: float = 0.0, model_name: str | None = None, @@ -92,7 +105,7 @@ def search_knowledge_base_with_sources( rendered ```` block's ``id``.""" if not query or not query.strip(): return "Error: query is empty.", [] - scope = _resolve_scope(scope_kb_id, scope_thread_id) + scope = _resolve_scope(scope_kb_id, scope_thread_id, scope_project_id) if scope is None: return "No documents are attached to this chat.", [] @@ -124,6 +137,7 @@ def search_for_autoinject( query: str, scope_kb_id: str | None = None, scope_thread_id: str | None = None, + scope_project_id: str | None = None, top_k: int | None = None, min_dense_score: float = 0.70, model_name: str | None = None, @@ -138,7 +152,7 @@ def search_for_autoinject( """ if not query or not query.strip(): return None - scope = _resolve_scope(scope_kb_id, scope_thread_id) + scope = _resolve_scope(scope_kb_id, scope_thread_id, scope_project_id) if scope is None: return None k = top_k or config.TOP_K_HYBRID @@ -177,6 +191,7 @@ def search_knowledge_base( query: str, scope_kb_id: str | None = None, scope_thread_id: str | None = None, + scope_project_id: str | None = None, top_k: int | None = None, min_score: float = 0.0, model_name: str | None = None, @@ -186,6 +201,7 @@ def search_knowledge_base( query = query, scope_kb_id = scope_kb_id, scope_thread_id = scope_thread_id, + scope_project_id = scope_project_id, top_k = top_k, min_score = min_score, model_name = model_name, diff --git a/studio/backend/core/training/resume.py b/studio/backend/core/training/resume.py index 165c1c2cf1..2a4a198610 100644 --- a/studio/backend/core/training/resume.py +++ b/studio/backend/core/training/resume.py @@ -3,6 +3,7 @@ """Helpers for validating resumable training outputs.""" +import json from pathlib import Path from typing import Optional @@ -56,9 +57,29 @@ def normalize_resume_output_dir(path_value: str) -> str: return str(path) +def _run_config(run: dict) -> dict: + raw_config = run.get("config_json") + if isinstance(raw_config, dict): + return raw_config + if not isinstance(raw_config, str) or not raw_config.strip(): + return {} + try: + parsed = json.loads(raw_config) + except (json.JSONDecodeError, TypeError): + return {} + return parsed if isinstance(parsed, dict) else {} + + +def _uses_s3_dataset(run: dict) -> bool: + config = _run_config(run) + return config.get("dataset_source") == "s3" or "s3_dataset" in config + + def can_resume_run(run: dict) -> bool: if run.get("resumed_later"): return False + if _uses_s3_dataset(run): + return False final_step = run.get("final_step") total_steps = run.get("total_steps") diff --git a/studio/backend/core/training/s3_dataset.py b/studio/backend/core/training/s3_dataset.py new file mode 100644 index 0000000000..3d05d19c75 --- /dev/null +++ b/studio/backend/core/training/s3_dataset.py @@ -0,0 +1,228 @@ +# SPDX-License-Identifier: AGPL-3.0-only +# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0 + +""" +S3 dataset loader. + +Downloads dataset files (parquet / json / jsonl / csv) from an AWS S3 bucket +to a local temp directory so the existing local-file dataset path can consume +them. boto3 is an optional dependency and is imported lazily — callers should +gate on :func:`boto3_available` before invoking the loader. + +The S3 config dict mirrors ``models.training.S3Config.model_dump()`` (snake_case +keys): bucket, region, prefix, access_key_id, secret_access_key, use_iam_role. +Credentials are read once to build the client and never logged or persisted. +""" + +from __future__ import annotations + +import logging +import os +import shutil +import tempfile +from importlib.util import find_spec +from typing import Callable, Optional + +logger = logging.getLogger(__name__) + +# Extensions the local-file loader (UnslothTrainer._loader_for_files) understands. +SUPPORTED_EXTENSIONS = (".parquet", ".json", ".jsonl", ".csv") +_JSON_EXTENSIONS = (".json", ".jsonl") +_IGNORED_METADATA_FILENAMES = { + "dataset_info.json", + "metadata.json", + "schema.json", + "state.json", +} + + +class S3DownloadCancelled(RuntimeError): + """Raised when the caller cancels an S3 dataset download.""" + + +class S3DatasetDownload: + def __init__( + self, + files: list[str], + temp_dir: Optional[str] = None, + ): + self.files = files + self.temp_dir = temp_dir + + def cleanup(self) -> None: + if not self.temp_dir: + return + shutil.rmtree(self.temp_dir, ignore_errors = True) + self.temp_dir = None + + +def boto3_available() -> bool: + """True if boto3 can be imported (without importing it).""" + return find_spec("boto3") is not None + + +def _build_s3_client(s3_config: dict): + """Create a boto3 S3 client from the config dict. + + Uses explicit access keys when provided, otherwise falls back to the + default credential chain (IAM role / instance profile / env / shared creds). + """ + import boto3 # lazy: optional dependency + + region = s3_config.get("region") or "us-east-1" + use_iam_role = bool(s3_config.get("use_iam_role")) + access_key_id = s3_config.get("access_key_id") + secret_access_key = s3_config.get("secret_access_key") + + if not use_iam_role and access_key_id and secret_access_key: + return boto3.client( + "s3", + region_name = region, + aws_access_key_id = access_key_id, + aws_secret_access_key = secret_access_key, + ) + # IAM role / instance profile / ambient credentials + return boto3.client("s3", region_name = region) + + +def _list_dataset_keys(client, bucket: str, prefix: Optional[str]) -> list[str]: + """List object keys under ``prefix`` that have a supported data extension.""" + paginator = client.get_paginator("list_objects_v2") + list_kwargs = {"Bucket": bucket} + if prefix: + list_kwargs["Prefix"] = prefix + + keys: list[str] = [] + for page in paginator.paginate(**list_kwargs): + for obj in page.get("Contents", []): + key = obj["Key"] + if key.endswith("/"): + continue # directory placeholder + if os.path.basename(key).lower() in _IGNORED_METADATA_FILENAMES: + continue + if key.lower().endswith(SUPPORTED_EXTENSIONS): + keys.append(key) + return keys + + +def _extension_family(key: str) -> str: + ext = os.path.splitext(key)[1].lower() + if ext in _JSON_EXTENSIONS: + return "json" + return ext.lstrip(".") + + +def _validate_single_extension_family(keys: list[str]) -> None: + families: list[str] = [] + for key in keys: + family = _extension_family(key) + if family not in families: + families.append(family) + + if len(families) <= 1: + return + + raise ValueError( + "S3 prefix contains mixed dataset formats " + f"({', '.join(families)}). Keep one dataset format under the selected prefix." + ) + + +def _unique_local_path(target_dir: str, filename: str, used_paths: set[str]) -> str: + """Return an unused flattened path for an S3 object basename.""" + stem, ext = os.path.splitext(filename) + candidate = os.path.join(target_dir, filename) + suffix = 1 + while candidate in used_paths or os.path.exists(candidate): + candidate = os.path.join(target_dir, f"{stem}_{suffix}{ext}") + suffix += 1 + used_paths.add(candidate) + return candidate + + +def _raise_if_cancelled(cancel_callback: Optional[Callable[[], bool]]) -> None: + if cancel_callback is not None and cancel_callback(): + raise S3DownloadCancelled("S3 dataset download cancelled") + + +def prepare_s3_dataset_download( + s3_config: dict, + dest_dir: Optional[str] = None, + cancel_callback: Optional[Callable[[], bool]] = None, +) -> S3DatasetDownload: + """Download supported dataset files from S3 to a local directory. + + Returns the local files plus the owned temporary directory, when one was + created. Call ``cleanup()`` after the dataset loader has materialized data. + + Raises ``RuntimeError`` if boto3 is missing, and ``ValueError`` if the + bucket/prefix contains no supported dataset files. + """ + if not boto3_available(): + raise RuntimeError("S3 dataset loading requires boto3. Install it with: pip install boto3") + + bucket = s3_config.get("bucket") + if not bucket: + raise ValueError("s3_config.bucket is required") + prefix = s3_config.get("prefix") + + _raise_if_cancelled(cancel_callback) + client = _build_s3_client(s3_config) + + keys = _list_dataset_keys(client, bucket, prefix) + _raise_if_cancelled(cancel_callback) + if not keys: + where = f"s3://{bucket}/{prefix}" if prefix else f"s3://{bucket}" + raise ValueError( + f"No supported dataset files ({', '.join(SUPPORTED_EXTENSIONS)}) " + f"found under {where}" + ) + + _validate_single_extension_family(keys) + + owns_temp_dir = dest_dir is None + target_dir = dest_dir or tempfile.mkdtemp(prefix = "unsloth_s3_dataset_") + try: + os.makedirs(target_dir, exist_ok = True) + + local_files: list[str] = [] + used_paths: set[str] = set() + for key in keys: + _raise_if_cancelled(cancel_callback) + filename = os.path.basename(key) + local_path = _unique_local_path(target_dir, filename, used_paths) + download_kwargs = {} + if cancel_callback is not None: + download_kwargs["Callback"] = lambda _bytes: _raise_if_cancelled(cancel_callback) + client.download_file(bucket, key, local_path, **download_kwargs) + _raise_if_cancelled(cancel_callback) + local_files.append(local_path) + except Exception: + if owns_temp_dir: + shutil.rmtree(target_dir, ignore_errors = True) + raise + + logger.info( + "Downloaded %d dataset file(s) from s3://%s/%s to %s", + len(local_files), + bucket, + prefix or "", + target_dir, + ) + return S3DatasetDownload( + files = local_files, + temp_dir = target_dir if owns_temp_dir else None, + ) + + +def download_s3_dataset( + s3_config: dict, + dest_dir: Optional[str] = None, + cancel_callback: Optional[Callable[[], bool]] = None, +) -> list[str]: + download = prepare_s3_dataset_download( + s3_config, + dest_dir = dest_dir, + cancel_callback = cancel_callback, + ) + return download.files diff --git a/studio/backend/core/training/trainer.py b/studio/backend/core/training/trainer.py index 57342b2453..018b66ea92 100644 --- a/studio/backend/core/training/trainer.py +++ b/studio/backend/core/training/trainer.py @@ -2227,6 +2227,7 @@ class UnslothTrainer: dataset_slice_start: int = None, dataset_slice_end: int = None, is_cpt: bool = False, + s3_config: dict = None, ) -> Optional[tuple]: """ Load and prepare a dataset for training. @@ -2237,6 +2238,9 @@ class UnslothTrainer: Returns (dataset_info, eval_dataset) or None on error; eval_dataset may be None if no eval split is available. """ + from core.training.s3_dataset import S3DownloadCancelled + + s3_download = None try: dataset = None eval_dataset = None @@ -2272,6 +2276,22 @@ class UnslothTrainer: return result.dataset + # S3 datasets are downloaded to a local temp dir and then consumed + # through the same local-file path below. + if s3_config and not local_datasets: + from core.training.s3_dataset import prepare_s3_dataset_download + + self._update_progress(status_message = "Downloading dataset from S3...") + s3_download = prepare_s3_dataset_download( + s3_config, + cancel_callback = lambda: self.should_stop, + ) + local_datasets = s3_download.files + if self.should_stop: + logger.info("Stopped during S3 download\n") + return None + logger.info(f"Downloaded {len(local_datasets)} file(s) from S3\n") + if local_datasets: # Use load_dataset() for an Arrow-backed result; in-memory # Dataset.from_list() has no cache and forces num_proc=1 during @@ -2539,10 +2559,16 @@ class UnslothTrainer: return (dataset_info, eval_dataset) + except S3DownloadCancelled: + logger.info("Stopped during S3 download\n") + return None except Exception as e: logger.error(f"Error loading dataset: {e}") self._update_progress(error = str(e)) return None + finally: + if s3_download is not None: + s3_download.cleanup() def _auto_detect_eval_split_from_hf( self, dataset_source: str, subset: str @@ -3367,7 +3393,19 @@ class UnslothTrainer: # ========== PROGRESS TRACKING ========== self.trainer.add_callback(self._create_progress_callback()) - num_samples = len(dataset["dataset"] if isinstance(dataset, dict) else dataset) + num_samples = None + if hasattr(self.trainer, "train_dataset") and self.trainer.train_dataset is not None: + try: + num_samples = len(self.trainer.train_dataset) + except TypeError: + logger.debug( + "train_dataset does not support len(); falling back to " + "raw dataset size for step estimation." + ) + + if num_samples is None: + num_samples = len(dataset["dataset"] if isinstance(dataset, dict) else dataset) + batch_size = training_args.get("batch_size", 2) total_steps = self._calculate_total_steps( num_samples, @@ -3376,10 +3414,8 @@ class UnslothTrainer: training_args.get("num_epochs", 3), training_args.get("max_steps", 0), ) - self._update_progress(total_steps = total_steps) - # ========== START TRAINING ========== - self._update_progress(status_message = "Starting training...") + self._update_progress(total_steps = total_steps, status_message = "Starting training...") logger.info("Starting training...\n") self.trainer.train(resume_from_checkpoint = training_args.get("resume_from_checkpoint")) diff --git a/studio/backend/core/training/training.py b/studio/backend/core/training/training.py index 2764101d93..6dd42976c7 100644 --- a/studio/backend/core/training/training.py +++ b/studio/backend/core/training/training.py @@ -40,6 +40,34 @@ logger = get_logger(__name__) _HF_TMP_CHECKPOINT_RE = re.compile(r"^tmp-checkpoint-\d+$") +def _sanitize_db_config(config: dict[str, Any]) -> dict[str, Any]: + db_config = { + k: v for k, v in config.items() if k not in {"hf_token", "wandb_token", "s3_config"} + } + s3_config = config.get("s3_config") + if hasattr(s3_config, "model_dump"): + s3_config = s3_config.model_dump() + if isinstance(s3_config, dict) and s3_config: + db_config["dataset_source"] = "s3" + db_config["s3_dataset"] = { + "bucket": s3_config.get("bucket"), + "region": s3_config.get("region"), + "prefix": s3_config.get("prefix"), + "use_iam_role": bool(s3_config.get("use_iam_role")), + } + return db_config + + +def _s3_dataset_name(s3_dataset: Any) -> Optional[str]: + if not isinstance(s3_dataset, dict): + return None + bucket = s3_dataset.get("bucket") + if not bucket: + return None + prefix = s3_dataset.get("prefix") + return f"s3://{bucket}/{prefix}" if prefix else f"s3://{bucket}" + + def _cleanup_cancelled_checkpoints(output_dir: str | os.PathLike) -> None: """Remove only HF Trainer ``tmp-checkpoint-/`` partials after a cancel. @@ -236,6 +264,7 @@ class TrainingBackend: "resume_from_checkpoint": kwargs.get("resume_from_checkpoint"), "trust_remote_code": kwargs.get("trust_remote_code", False), "gpu_ids": kwargs.get("gpu_ids"), + "s3_config": kwargs.get("s3_config"), } # Full finetuning always runs in 16-bit; LoRA/QLoRA/CPT keep the request. @@ -309,7 +338,7 @@ class TrainingBackend: self._run_finalized = False self._db_run_created = False self._db_total_steps_set = False - self._db_config = {k: v for k, v in config.items() if k not in {"hf_token", "wandb_token"}} + self._db_config = _sanitize_db_config(config) self._db_started_at = datetime.now(timezone.utc).isoformat() # Assign subprocess handles after state reset. @@ -732,8 +761,11 @@ class TrainingBackend: try: from storage.studio_db import create_run - dataset_name = self._db_config.get("hf_dataset") or next( - iter(self._db_config.get("local_datasets") or []), "unknown" + dataset_name = ( + self._db_config.get("hf_dataset") + or next(iter(self._db_config.get("local_datasets") or []), None) + or _s3_dataset_name(self._db_config.get("s3_dataset")) + or "unknown" ) create_run( id = self.current_job_id, diff --git a/studio/backend/core/training/worker.py b/studio/backend/core/training/worker.py index 18b25cb4fe..1120744a2d 100644 --- a/studio/backend/core/training/worker.py +++ b/studio/backend/core/training/worker.py @@ -1336,6 +1336,32 @@ def _run_mlx_training(event_queue, stop_queue, config): kwargs["message"] = sm event_queue.put({"type": event_type, "ts": time.time(), **kwargs}) + _stop_save = [True] + _stop_requested = [False] + _trainer_ref = [None] + + def _is_stop_requested(): + return _stop_requested[0] + + 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) + _stop_requested[0] = True + trainer = _trainer_ref[0] + if trainer is not None: + trainer.stop_requested = True + return + except _queue.Empty: + continue + except (EOFError, OSError): + return + + stop_thread = threading.Thread(target = _poll_stop, daemon = True) + stop_thread.start() + _send("status", status_message = "Loading MLX libraries...") import mlx.core as mx @@ -1515,6 +1541,26 @@ def _run_mlx_training(event_queue, stop_queue, config): elif config.get("local_datasets"): dataset = _load_local(config["local_datasets"]) dataset = _slice(dataset) + elif config.get("s3_config"): + from core.training.s3_dataset import ( + S3DownloadCancelled, + prepare_s3_dataset_download, + ) + + _send("status", status_message = "Downloading dataset from S3...") + try: + s3_download = prepare_s3_dataset_download( + config["s3_config"], + cancel_callback = _is_stop_requested, + ) + try: + dataset = _load_local(s3_download.files) + finally: + s3_download.cleanup() + except S3DownloadCancelled: + _send("complete", output_dir = None, status_message = "Training cancelled") + return + dataset = _slice(dataset) else: raise ValueError("No dataset specified") @@ -1693,6 +1739,9 @@ def _run_mlx_training(event_queue, stop_queue, config): eval_steps = eval_steps_val, ), ) + _trainer_ref[0] = trainer + if _stop_requested[0]: + trainer.stop_requested = True # Tell the parent eval is configured so the frontend shows the eval chart if eval_dataset is not None and eval_steps_val > 0: @@ -1733,7 +1782,7 @@ def _run_mlx_training(event_queue, stop_queue, config): wandb_token = config.get("wandb_token") if wandb_token: os.environ["WANDB_API_KEY"] = wandb_token - _wandb_sensitive = {"hf_token", "wandb_token"} + _wandb_sensitive = {"hf_token", "wandb_token", "s3_config"} 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}, @@ -1835,26 +1884,6 @@ def _run_mlx_training(event_queue, stop_queue, config): 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): - # Safe: pipe permanently broken, no more messages can arrive. - return - - stop_thread = threading.Thread(target = _poll_stop, daemon = True) - stop_thread.start() - # ── 11. Run training ── gc.collect() mx.synchronize() @@ -2462,7 +2491,7 @@ def run_training_process(*, event_queue: Any, stop_queue: Any, config: dict) -> def _on_progress(progress: TrainingProgress): has_train_loss = progress.step > 0 and progress.loss is not None has_eval_loss = progress.eval_loss is not None - if has_train_loss or has_eval_loss: + if (progress.step == 0 and progress.total_steps > 0) or has_train_loss or has_eval_loss: event_queue.put( { "type": "progress", @@ -2546,6 +2575,7 @@ def run_training_process(*, event_queue: Any, stop_queue: Any, config: dict) -> dataset_slice_start = config.get("dataset_slice_start"), dataset_slice_end = config.get("dataset_slice_end"), is_cpt = _is_cpt_for_dataset, + s3_config = config.get("s3_config"), ) if isinstance(dataset_result, tuple): @@ -3010,20 +3040,9 @@ def _run_embedding_training(event_queue: Any, stop_queue: Any, config: dict) -> subset = config.get("subset") or None train_split = config.get("train_split", "train") or "train" - if hf_dataset and hf_dataset.strip(): - hf_token = config.get("hf_token", "") - hf_token = hf_token if hf_token and hf_token.strip() else None - dataset = load_dataset( - hf_dataset.strip(), - subset, - split = train_split, - token = hf_token, - ) - elif local_datasets: - # Load local file(s) — mirrors the non-embedding pipeline's directory - # handling so recipe outputs (parquet-files/) work. + def _load_local_embedding_dataset(dataset_paths: list[str]): all_files: list[str] = [] - for dataset_file in local_datasets: + for dataset_file in dataset_paths: file_path = ( dataset_file if os.path.isabs(dataset_file) @@ -3053,17 +3072,58 @@ def _run_embedding_training(event_queue: Any, stop_queue: Any, config: dict) -> else: all_files.append(file_path) - if all_files: - first_ext = Path(all_files[0]).suffix.lower() - if first_ext in (".json", ".jsonl"): - loader = "json" - elif first_ext == ".csv": - loader = "csv" - elif first_ext == ".parquet": - loader = "parquet" - else: - raise ValueError(f"Unsupported local dataset format: {all_files[0]}") - dataset = load_dataset(loader, data_files = all_files, split = "train") + if not all_files: + raise ValueError("No local dataset files found") + + first_ext = Path(all_files[0]).suffix.lower() + if first_ext in (".json", ".jsonl"): + loader = "json" + elif first_ext == ".csv": + loader = "csv" + elif first_ext == ".parquet": + loader = "parquet" + else: + raise ValueError(f"Unsupported local dataset format: {all_files[0]}") + return load_dataset(loader, data_files = all_files, split = "train") + + if hf_dataset and hf_dataset.strip(): + hf_token = config.get("hf_token", "") + hf_token = hf_token if hf_token and hf_token.strip() else None + dataset = load_dataset( + hf_dataset.strip(), + subset, + split = train_split, + token = hf_token, + ) + elif local_datasets: + dataset = _load_local_embedding_dataset(local_datasets) + elif config.get("s3_config"): + from core.training.s3_dataset import ( + S3DownloadCancelled, + prepare_s3_dataset_download, + ) + + _send_status(event_queue, "Downloading dataset from S3...") + s3_download = None + try: + s3_download = prepare_s3_dataset_download( + config["s3_config"], + cancel_callback = lambda: _should_stop, + ) + dataset = _load_local_embedding_dataset(s3_download.files) + except S3DownloadCancelled: + event_queue.put( + { + "type": "complete", + "output_dir": None, + "status_message": "Training cancelled", + "ts": time.time(), + } + ) + return + finally: + if s3_download is not None: + s3_download.cleanup() else: event_queue.put( { diff --git a/studio/backend/hub/utils/inventory_scan.py b/studio/backend/hub/utils/inventory_scan.py index fe3aadfd38..0f7ce6fe34 100644 --- a/studio/backend/hub/utils/inventory_scan.py +++ b/studio/backend/hub/utils/inventory_scan.py @@ -142,13 +142,18 @@ def _compute_all_hf_cache_scans() -> list: logger.warning("Could not scan active HF cache: %s", exc) for extra_fn in (legacy_hf_cache_dir, hf_default_cache_dir): - extra = extra_fn() - if extra.is_dir() and str(extra.resolve()) not in seen: - seen.add(str(extra.resolve())) - try: - scans.append(scan_cache_dir(cache_dir = str(extra))) - except Exception as exc: - logger.warning("Could not scan HF cache %s: %s", extra, exc) + try: + extra = extra_fn() + # is_dir()/resolve() can raise on an inaccessible path; skip it. + if not extra.is_dir(): + continue + resolved = str(extra.resolve()) + if resolved in seen: + continue + seen.add(resolved) + scans.append(scan_cache_dir(cache_dir = str(extra))) + except Exception as exc: + logger.warning("Could not scan HF cache %s: %s", extra_fn.__name__, exc) return scans diff --git a/studio/backend/main.py b/studio/backend/main.py index 2b34b11fef..064e261753 100644 --- a/studio/backend/main.py +++ b/studio/backend/main.py @@ -193,6 +193,11 @@ if _STUDIO_ROOT_RESOLVED != _LEGACY_STUDIO_ROOT: if not os.environ.get("UNSLOTH_LLAMA_CPP_PATH"): os.environ["UNSLOTH_LLAMA_CPP_PATH"] = str(_STUDIO_ROOT_RESOLVED / "llama.cpp") +# The studio bundles unsloth_zoo; declare unsloth present (as `import unsloth` +# does) so its lazy submodule imports (export, hardware, mlx) and the +# DiffusionGemma runner never trip the install guard on a clean install. +os.environ.setdefault("UNSLOTH_IS_PRESENT", "1") + import hashlib import mimetypes import re as _re diff --git a/studio/backend/models/export.py b/studio/backend/models/export.py index 0d54b435a6..584f82dea0 100644 --- a/studio/backend/models/export.py +++ b/studio/backend/models/export.py @@ -3,14 +3,14 @@ """Pydantic schemas for Export API.""" -from pathlib import Path +from pathlib import Path, PureWindowsPath from pydantic import BaseModel, Field, field_validator from typing import List, Optional, Literal, Dict, Any def _validate_save_directory(value: str) -> str: - """Reject save_directory values that escape the export root.""" + """Validate save_directory — allows absolute paths (user may want a different drive).""" if value is None: raise ValueError("save_directory is required") raw = str(value).strip() @@ -20,15 +20,15 @@ def _validate_save_directory(value: str) -> str: raise ValueError("save_directory may not contain null bytes") if any(ch in raw for ch in ("\r", "\n")): raise ValueError("save_directory may not contain control characters") - if len(raw) > 255: - raise ValueError("save_directory must be <= 255 characters") path = Path(raw).expanduser() - if path.is_absolute(): - raise ValueError( - "save_directory must be a name or relative path under the " - "export root; absolute paths are rejected" - ) - if ".." in path.parts: + path_parts = (*path.parts, *PureWindowsPath(raw).parts, *raw.replace("\\", "/").split("/")) + if any(len(part) > 255 for part in path_parts if part not in ("", ".", "/", "\\")): + raise ValueError("save_directory path components must be <= 255 characters") + if ( + ".." in path.parts + or ".." in PureWindowsPath(raw).parts + or ".." in raw.replace("\\", "/").split("/") + ): raise ValueError("save_directory may not contain '..' segments") return raw diff --git a/studio/backend/models/inference.py b/studio/backend/models/inference.py index 039a5d75dd..04c63c2247 100644 --- a/studio/backend/models/inference.py +++ b/studio/backend/models/inference.py @@ -87,6 +87,15 @@ class LoadRequest(BaseModel): "'mtp' or 'mtp+ngram'." ), ) + tensor_parallel: bool = Field( + False, + description = ( + "Split the model across GPUs by tensor (--split-mode tensor) " + "instead of by layer for GGUF models. Only affects multi-GPU " + "setups, where it can make generation significantly faster. " + "No effect on a single GPU. Ignored for non-GGUF models." + ), + ) llama_extra_args: Optional[List[str]] = Field( None, description = ( @@ -160,6 +169,9 @@ class LoadResponse(BaseModel): is_vision: bool = Field(False, description = "Whether model is a vision model") is_lora: bool = Field(False, description = "Whether model is a LoRA adapter") is_gguf: bool = Field(False, description = "Whether model is a GGUF model (llama.cpp)") + is_diffusion: bool = Field( + False, description = "Whether model is a block-diffusion model (DiffusionGemma)" + ) is_audio: bool = Field(False, description = "Whether model is a TTS audio model") audio_type: Optional[str] = Field(None, description = "Audio codec type: snac, csm, bicodec, dac") has_audio_input: bool = Field(False, description = "Whether model accepts audio input (ASR)") @@ -224,6 +236,10 @@ class LoadResponse(BaseModel): "None when the platform default is in effect." ), ) + tensor_parallel: bool = Field( + False, + description = "Whether tensor-parallel split (--split-mode tensor) is active.", + ) class UnloadResponse(BaseModel): @@ -273,6 +289,9 @@ class InferenceStatusResponse(BaseModel): ) is_vision: bool = Field(False, description = "Whether the active model is a vision model") is_gguf: bool = Field(False, description = "Whether the active model is a GGUF model (llama.cpp)") + is_diffusion: bool = Field( + False, description = "Whether the active model is a block-diffusion model (DiffusionGemma)" + ) gguf_variant: Optional[str] = Field(None, description = "GGUF quantization variant (e.g. Q4_K_M)") is_audio: bool = Field(False, description = "Whether the active model is a TTS audio model") audio_type: Optional[str] = Field(None, description = "Audio codec type: snac, csm, bicodec, dac") @@ -339,6 +358,10 @@ class InferenceStatusResponse(BaseModel): "None when the platform default is in effect." ), ) + tensor_parallel: bool = Field( + False, + description = "Whether tensor-parallel split (--split-mode tensor) is active.", + ) llama_cpp_supports_mtp: bool = Field( True, description = ( @@ -1314,6 +1337,23 @@ class ResponsesOutputMessage(BaseModel): content: list[ResponsesOutputTextContent] = Field(default_factory = list) +class ResponsesOutputReasoningContent(BaseModel): + """A reasoning text content block inside a reasoning output item.""" + + type: Literal["reasoning_text"] = "reasoning_text" + text: str + + +class ResponsesOutputReasoning(BaseModel): + """A top-level reasoning output item in the Responses API response.""" + + type: Literal["reasoning"] = "reasoning" + id: str = Field(default_factory = lambda: f"rs_{uuid.uuid4().hex[:12]}") + status: Literal["completed", "in_progress", "incomplete"] = "completed" + summary: list = Field(default_factory = list) + content: Optional[list[ResponsesOutputReasoningContent]] = None + + class ResponsesOutputFunctionCall(BaseModel): """A function-call output item in the Responses API response. @@ -1328,7 +1368,11 @@ class ResponsesOutputFunctionCall(BaseModel): status: Literal["completed", "in_progress", "incomplete"] = "completed" -ResponsesOutputItem = Union[ResponsesOutputMessage, ResponsesOutputFunctionCall] +ResponsesOutputItem = Union[ + ResponsesOutputMessage, + ResponsesOutputReasoning, + ResponsesOutputFunctionCall, +] class ResponsesUsage(BaseModel): diff --git a/studio/backend/models/providers.py b/studio/backend/models/providers.py index 8f90e19df6..5a75246c07 100644 --- a/studio/backend/models/providers.py +++ b/studio/backend/models/providers.py @@ -108,6 +108,9 @@ class ProviderTestRequest(BaseModel): base_url: Optional[str] = Field( None, description = "Custom base URL (overrides registry default)" ) + model_id: Optional[str] = Field( + None, description = "Model ID for providers that need a chat probe" + ) class ProviderTestResult(BaseModel): diff --git a/studio/backend/models/training.py b/studio/backend/models/training.py index 4a3303c734..ca04591178 100644 --- a/studio/backend/models/training.py +++ b/studio/backend/models/training.py @@ -29,6 +29,43 @@ _MIN_VISION_IMAGE_SIZE = 256 _MAX_VISION_IMAGE_SIZE = 2048 +class S3Config(BaseModel): + """S3 bucket configuration for loading datasets from AWS S3""" + + # Accept both snake_case and the frontend's camelCase field names. + model_config = ConfigDict(populate_by_name = True) + + bucket: str = Field(..., description = "S3 bucket name") + region: str = Field("us-east-1", description = "AWS region") + prefix: Optional[str] = Field(None, description = "Optional path prefix within bucket") + access_key_id: Optional[str] = Field( + None, + alias = "accessKeyId", + description = "AWS access key ID (optional if using IAM role)", + ) + secret_access_key: Optional[str] = Field( + None, + alias = "secretAccessKey", + description = "AWS secret access key (optional if using IAM role)", + ) + use_iam_role: bool = Field( + False, + alias = "useIamRole", + description = "Use IAM role credentials instead of access keys", + ) + + @model_validator(mode = "after") + def _check_credentials(self) -> "S3Config": + # Require either IAM role auth or a full key pair so credentials are + # never half-configured. + if not self.use_iam_role and not (self.access_key_id and self.secret_access_key): + raise ValueError( + "s3_config requires either use_iam_role=True or both " + "access_key_id and secret_access_key" + ) + return self + + def _parse_lr(v: Any) -> float: """Parse learning_rate as a positive float strictly below _MAX_LR_VALUE.""" if v is None: @@ -338,6 +375,12 @@ class TrainingStartRequest(BaseModel): description = "Physical GPU indices to use, for example [0, 1]. Omit or pass [] to use automatic selection. Explicit gpu_ids are unsupported when the parent CUDA_VISIBLE_DEVICES uses UUID/MIG entries.", ) + # S3 dataset source configuration + s3_config: Optional[S3Config] = Field( + None, + description = "S3 bucket configuration for loading datasets from AWS S3. Requires boto3 to be installed.", + ) + @model_validator(mode = "after") def _check_steps_or_epochs(self) -> "TrainingStartRequest": # Each accepts 0 as "use the other"; both 0 means nothing to train. diff --git a/studio/backend/requirements/no-torch-runtime.txt b/studio/backend/requirements/no-torch-runtime.txt index 85294114b1..6efe91d448 100644 --- a/studio/backend/requirements/no-torch-runtime.txt +++ b/studio/backend/requirements/no-torch-runtime.txt @@ -68,3 +68,9 @@ trl>=0.18.2,!=0.19.0,<=0.24.0 sentence-transformers cut_cross_entropy pillow + +# RAG store + document parsing, mirroring studio.txt. Pinned here because +# this file installs --no-deps; without them Studio runs with RAG disabled. +sqlite-vec==0.1.9 +pymupdf==1.27.2.3 +python-docx==1.2.0 diff --git a/studio/backend/requirements/studio.txt b/studio/backend/requirements/studio.txt index 9bffb81549..96fef60471 100644 --- a/studio/backend/requirements/studio.txt +++ b/studio/backend/requirements/studio.txt @@ -17,6 +17,7 @@ structlog>=24.1.0 diceware ddgs cryptography>=42.0.0 +boto3>=1.34.0 # optional: S3 dataset loading httpx>=0.27.0 fastmcp>=3.0.2 # RAG (knowledge bases, hybrid retrieval). sentence-transformers lives in diff --git a/studio/backend/routes/chat_history.py b/studio/backend/routes/chat_history.py index 782bba3be0..2ea572e30e 100644 --- a/studio/backend/routes/chat_history.py +++ b/studio/backend/routes/chat_history.py @@ -332,6 +332,35 @@ async def delete_project( status_code = 404, detail = f"Project {project_id} not found", ) + # Best-effort: drop the project's RAG sources (lazy import keeps RAG optional). + try: + import os + + from storage import rag_db + if rag_db.RAG_AVAILABLE: + from core.rag import store as rag_store + from utils.paths import rag_uploads_root + + uploads = os.path.realpath(str(rag_uploads_root())) + conn = rag_db.get_connection() + try: + scope = rag_store.project_scope(project_id) + for doc in rag_store.list_documents(conn, scope): + full = rag_store.get_document(conn, doc["id"]) or {} + rag_store.delete_document(conn, doc["id"]) + stored = full.get("stored_path") + # Also remove the uploaded file; confined to the uploads root. + if stored: + target = os.path.realpath(stored) + if ( + os.path.isfile(target) + and os.path.commonpath([uploads, target]) == uploads + ): + os.remove(target) + finally: + conn.close() + except Exception: # noqa: BLE001 - source cleanup must not block project deletion + logger.warning("failed to delete RAG sources for project %s", project_id, exc_info = True) return ChatProject(**project) diff --git a/studio/backend/routes/export.py b/studio/backend/routes/export.py index ef37daa158..f7b3a56a71 100644 --- a/studio/backend/routes/export.py +++ b/studio/backend/routes/export.py @@ -154,19 +154,41 @@ async def get_export_status(current_subject: str = Depends(get_current_subject)) ) +def _try_register_external_export(path: Path) -> tuple[bool, Optional[str]]: + """Best-effort registration so absolute exports show up in local scans.""" + try: + from storage.studio_db import add_scan_folder + folder = add_scan_folder(str(path)) + return True, str(folder.get("path") or path) + except Exception as exc: + logger.warning("Could not register export scan folder %s: %s", path, exc) + return False, None + + def _export_details(output_path: Optional[str]) -> Optional[Dict[str, Any]]: - """Return the export path relative to exports_root, hiding the install path.""" + """Return relative export paths, keeping external absolute paths visible.""" if not output_path: return None try: from utils.paths.storage_roots import exports_root + path = Path(output_path) + # If it's outside exports_root, return the full absolute path + # so users can find their files on a different drive. + if path.is_absolute(): + try: + path.resolve().relative_to(exports_root().resolve()) + except ValueError: + registered, registered_path = _try_register_external_export(path) + return { + "output_path": str(path), + "scan_folder_registered": registered, + "scan_folder_path": registered_path, + } rel = os.path.relpath(output_path, exports_root()) - if rel.startswith(".."): - rel = os.path.basename(output_path) return {"output_path": rel} except Exception: - return {"output_path": os.path.basename(output_path)} + return {"output_path": output_path} @router.post("/export/merged", response_model = ExportOperationResponse) diff --git a/studio/backend/routes/inference.py b/studio/backend/routes/inference.py index 9ff30cd6db..e9353cd803 100644 --- a/studio/backend/routes/inference.py +++ b/studio/backend/routes/inference.py @@ -121,6 +121,19 @@ def _template_raise_message(error_text: str, chat_template: Optional[str]) -> Op def _friendly_error(exc: Exception) -> str: """Extract a user-friendly message from known llama-server errors.""" + if isinstance(exc, httpx.ReadTimeout): + if "stopped producing tokens" in str(exc).lower(): + return ( + "The model stopped producing tokens before the response " + "completed. Try stopping and retrying, or reduce max tokens." + ) + return ( + "The model is still processing the prompt but did not produce a " + "first token within 20 minutes. Try reducing context length, " + "using more GPU offload, or loading a smaller model." + ) + if isinstance(exc, httpx.TimeoutException): + return "Timed out communicating with the model server. Try again shortly." # httpx transport failures from the async pass-through helpers. Any # RequestError subclass (ConnectError, ReadError, RemoteProtocolError, # WriteError, PoolTimeout, ...) means the llama-server subprocess is @@ -224,7 +237,11 @@ def _openai_stream_error_chunk(exc) -> dict: (a code-less error hides it).""" _cls = _classify_llama_generation_error(exc) if _cls: - return openai_error_body(_friendly_error(exc), status = 400, code = "context_length_exceeded") + return openai_error_body( + _friendly_error(exc), + status = 400, + code = "context_length_exceeded", + ) if _cls is False: return openai_error_body(_friendly_error(exc), status = 400) return openai_error_body(_friendly_error(exc), status = 500) @@ -415,17 +432,15 @@ def _apply_overflow_truncation(body: dict, err_text: str) -> bool: return True -def _anthropic_stream_error_event(exc): - """Anthropic in-band SSE ``error`` event for a mid-stream failure, or ``None`` - to fall through to a normal message_delta finish. Returns an event only for a - classifiable upstream client error (context overflow / 4xx) so a streaming - over-context request surfaces a real error instead of a silent empty - end_turn message.""" - if _classify_llama_generation_error(exc) is None: +def _anthropic_stream_error_event(exc, *, force: bool = False): + """Return an Anthropic in-band stream error event when one is useful.""" + _cls = _classify_llama_generation_error(exc) + if _cls is None and not force: return None + status = 400 if _cls is not None else 500 return build_anthropic_sse_event( "error", - anthropic_error_body(_friendly_error(exc), status = 400), + anthropic_error_body(_friendly_error(exc), status = status), ) @@ -588,17 +603,20 @@ try: from core.inference import get_inference_backend from core.inference.llama_cpp import ( LlamaCppBackend, + _DEFAULT_FIRST_TOKEN_TIMEOUT_S, _DEFAULT_MAX_TOKENS_FLOOR, - _DEFAULT_T_MAX_PREDICT_MS, + _DEFAULT_STREAM_STALL_TIMEOUT_S, _canonicalize_spec_mode, _extra_args_set_spec_type, _hf_offline_if_dns_dead, detect_reasoning_flags, ) from core.inference.llama_server_args import ( + resolve_tensor_parallel, strip_shadowing_flags, validate_extra_args, ) + from core.inference.tensor_fallback import load_with_tensor_fallback from utils.models import ModelConfig from utils.inference import load_inference_config from utils.models.model_config import ( @@ -619,17 +637,20 @@ except ImportError: from core.inference import get_inference_backend from core.inference.llama_cpp import ( LlamaCppBackend, + _DEFAULT_FIRST_TOKEN_TIMEOUT_S, _DEFAULT_MAX_TOKENS_FLOOR, - _DEFAULT_T_MAX_PREDICT_MS, + _DEFAULT_STREAM_STALL_TIMEOUT_S, _canonicalize_spec_mode, _extra_args_set_spec_type, _hf_offline_if_dns_dead, detect_reasoning_flags, ) from core.inference.llama_server_args import ( + resolve_tensor_parallel, strip_shadowing_flags, validate_extra_args, ) + from core.inference.tensor_fallback import load_with_tensor_fallback from utils.models import ModelConfig from utils.inference import load_inference_config from utils.models.model_config import ( @@ -644,6 +665,152 @@ except ImportError: verify_native_path_lease, ) + +def _llama_non_streaming_generation_timeout() -> httpx.Timeout: + return httpx.Timeout( + connect = _DEFAULT_FIRST_TOKEN_TIMEOUT_S, + read = _DEFAULT_FIRST_TOKEN_TIMEOUT_S, + write = _DEFAULT_FIRST_TOKEN_TIMEOUT_S, + pool = _DEFAULT_FIRST_TOKEN_TIMEOUT_S, + ) + + +def _llama_streaming_generation_timeout() -> httpx.Timeout: + return httpx.Timeout( + connect = _DEFAULT_FIRST_TOKEN_TIMEOUT_S, + read = _DEFAULT_FIRST_TOKEN_TIMEOUT_S, + write = _DEFAULT_FIRST_TOKEN_TIMEOUT_S, + pool = _DEFAULT_FIRST_TOKEN_TIMEOUT_S, + ) + + +def _set_stream_response_read_timeout( + response: httpx.Response, read_timeout_s: float = _DEFAULT_STREAM_STALL_TIMEOUT_S +) -> None: + try: + timeout_ext = response.request.extensions.get("timeout") + if isinstance(timeout_ext, dict): + timeout_ext["read"] = read_timeout_s + except Exception: + pass + + +async def _preheader_cancelled(cancel_event = None, request: Optional[Request] = None) -> bool: + if cancel_event is not None and cancel_event.is_set(): + return True + if request is not None and await request.is_disconnected(): + if cancel_event is not None: + cancel_event.set() + return True + return False + + +async def _wait_preheader_cancel(cancel_event = None, request: Optional[Request] = None) -> None: + while not await _preheader_cancelled(cancel_event, request): + await asyncio.sleep(0.05) + + +async def _send_stream_with_preheader_cancel( + client: httpx.AsyncClient, + req: httpx.Request, + cancel_event = None, + request: Optional[Request] = None, +) -> Optional[httpx.Response]: + if cancel_event is None and request is None: + return await client.send(req, stream = True) + if await _preheader_cancelled(cancel_event, request): + return None + + send_task = asyncio.create_task(client.send(req, stream = True)) + cancel_task = asyncio.create_task(_wait_preheader_cancel(cancel_event, request)) + + async def _stop_send_task() -> None: + try: + await client.aclose() + except Exception: + pass + send_task.cancel() + try: + await send_task + except (asyncio.CancelledError, Exception): + pass + + try: + done, _pending = await asyncio.wait( + {send_task, cancel_task}, + return_when = asyncio.FIRST_COMPLETED, + ) + if send_task in done: + return await send_task + + await _stop_send_task() + return None + except asyncio.CancelledError: + if cancel_event is not None: + cancel_event.set() + await _stop_send_task() + raise + finally: + cancel_task.cancel() + try: + await cancel_task + except (asyncio.CancelledError, Exception): + pass + + +async def _aiter_llama_stream_items( + async_iter, + *, + cancel_event = None, + request: Optional[Request] = None, + first_token_deadline: Optional[float] = None, + response: Optional[httpx.Response] = None, + post_first_item_read_timeout_s: Optional[float] = _DEFAULT_STREAM_STALL_TIMEOUT_S, +): + if first_token_deadline is None: + first_token_deadline = time.monotonic() + _DEFAULT_FIRST_TOKEN_TIMEOUT_S + last_item_at: Optional[float] = None + while True: + if cancel_event is not None and cancel_event.is_set(): + return + if request is not None and await request.is_disconnected(): + if cancel_event is not None: + cancel_event.set() + return + waiting_first_item = last_item_at is None + try: + if waiting_first_item: + remaining_s = first_token_deadline - time.monotonic() + if remaining_s <= 0: + raise httpx.ReadTimeout("The model did not produce a first token in time.") + if response is not None: + _set_stream_response_read_timeout(response, remaining_s) + item = await asyncio.wait_for(async_iter.__anext__(), timeout = remaining_s) + else: + item = await async_iter.__anext__() + except asyncio.TimeoutError as exc: + if waiting_first_item: + raise httpx.ReadTimeout("The model did not produce a first token in time.") from exc + raise + except StopAsyncIteration: + return + except httpx.ReadTimeout: + now = time.monotonic() + if last_item_at is None: + if now >= first_token_deadline: + raise + continue + raise httpx.ReadTimeout("The model stopped producing tokens mid-response.") + if ( + last_item_at is None + and response is not None + and post_first_item_read_timeout_s is not None + ): + _set_stream_response_read_timeout(response, post_first_item_read_timeout_s) + last_item_at = time.monotonic() + yield item + + from models.inference import ( LoadRequest, UnloadRequest, @@ -678,6 +845,8 @@ from models.inference import ( ResponsesFunctionCallOutputInputItem, ResponsesOutputTextContent, ResponsesOutputMessage, + ResponsesOutputReasoning, + ResponsesOutputReasoningContent, ResponsesOutputFunctionCall, ResponsesUsage, ResponsesResponse, @@ -708,6 +877,7 @@ from state.tool_approvals import resolve_tool_decision from core.inference.key_exchange import decrypt_api_key from core.inference.providers import get_provider_info, get_base_url from core.inference.external_provider import ExternalProviderClient +from core.inference.chat_templates import resolve_effective_chat_template_override from storage import providers_db from utils.utils import safe_error_detail, log_and_http_error @@ -1159,9 +1329,34 @@ def _normalise_settings_str(value: Optional[str]) -> Optional[str]: return value -def _request_matches_loaded_settings(request: LoadRequest, llama_backend: LlamaCppBackend) -> bool: +def _should_strip_split_mode(request: LoadRequest, backend_extra: Optional[list[str]]) -> bool: + """Whether an inherited --split-mode should be stripped on reload. + + The binary Tensor Parallelism toggle can't carry --split-mode's row/none/ + layer modes, so only strip when the toggle overrides it: tensor being turned + on, or the inherited mode is tensor (toggle turning it off). Non-tensor modes + survive. Shared by the inheritance strip and the already-loaded stale check + so they agree on what reload would do. + """ + fields_set = getattr(request, "model_fields_set", set()) + return "tensor_parallel" in fields_set and ( + request.tensor_parallel or resolve_tensor_parallel(backend_extra, False) + ) + + +def _request_matches_loaded_settings( + request: LoadRequest, + llama_backend: LlamaCppBackend, + effective_chat_template_override: Optional[str] = None, +) -> bool: """True iff every runtime setting on the request matches the loaded server. - Caller has already checked model+variant+is_loaded. See #5401.""" + Caller has already checked model+variant+is_loaded. See #5401. + + ``effective_chat_template_override`` is the resolved template that will be + launched (user override, else a bundled family template such as the + gemma-4 override), so the dedup compares against what the backend actually + holds rather than the raw request field. Defaults to the request field for + callers that do not resolve a bundled override.""" # Compare requested n_ctx (not effective) so VRAM-cap doesn't mask an # Auto-vs-explicit slider flip. if request.max_seq_length != llama_backend.requested_n_ctx: @@ -1170,6 +1365,24 @@ def _request_matches_loaded_settings(request: LoadRequest, llama_backend: LlamaC llama_backend.cache_type_kv ): return False + # Reconcile a user --split-mode in extras into the effective tensor state. + # When the request omits llama_extra_args ("inherit"), compare using the + # stored extras stripped the way the reload strips them, so an extras-driven + # tensor load isn't seen as a mismatch that needlessly reloads the server. + backend_extra = list(llama_backend.extra_args) if llama_backend.extra_args else [] + effective_extra = ( + request.llama_extra_args + if request.llama_extra_args is not None + else strip_shadowing_flags( + backend_extra, + strip_split_mode = _should_strip_split_mode(request, backend_extra), + ) + ) + if ( + resolve_tensor_parallel(effective_extra, request.tensor_parallel) + != llama_backend.tensor_parallel + ): + return False # Spec decoding works on vision models too (MTP is mmproj-compatible, # llama.cpp #22673; the old ``not is_vision`` gate is gone), so compare # the real requested mode -- coercing vision to ``off`` here used to @@ -1183,15 +1396,29 @@ def _request_matches_loaded_settings(request: LoadRequest, llama_backend: LlamaC if backend_mode in ("mtp", "mtp+ngram") and request.spec_draft_n_max is not None: if int(request.spec_draft_n_max) != (llama_backend.spec_draft_n_max or 0): return False - if (request.chat_template_override or None) != (llama_backend.chat_template_override or None): + _effective_cto = ( + effective_chat_template_override + if effective_chat_template_override is not None + else request.chat_template_override + ) + if (_effective_cto or None) != (llama_backend.chat_template_override or None): return False # llama_extra_args=None means "inherit"; only an explicit differing list # forces a reload. On the inherit path, refuse to match if stored extras # contain any shadow flag, so the reload path strips them rather than - # leaving a stale override in effect. - backend_extra = list(llama_backend.extra_args) if llama_backend.extra_args else [] + # leaving a stale override in effect. (backend_extra computed above.) if request.llama_extra_args is None: - if backend_extra and strip_shadowing_flags(backend_extra) != backend_extra: + # Mirror the reload's conditional split-mode strip, so a preserved + # non-tensor mode (row/none/layer) isn't seen as stale and doesn't + # trigger a needless reload of a healthy server. + if ( + backend_extra + and strip_shadowing_flags( + backend_extra, + strip_split_mode = _should_strip_split_mode(request, backend_extra), + ) + != backend_extra + ): return False else: if list(request.llama_extra_args) != backend_extra: @@ -1300,6 +1527,17 @@ async def load_model( # Version switching is handled by the subprocess-based inference # backend -- no ensure_transformers_version() needed here. + # Resolve the effective chat-template override once, up front: an + # explicit user override, else a bundled family template (e.g. the + # gemma-4 override that ships preserve_thinking without re-downloading + # quants), else None. Used for both the reload-dedup check below and the + # load_model calls, so the live backend state and the incoming request + # compare against the same template text. + effective_chat_template_override = resolve_effective_chat_template_override( + model_identifier = model_identifier, + user_override = request.chat_template_override, + ) + # ── Already-loaded check: skip reload if the exact model is active ── backend = get_inference_backend() llama_backend = get_llama_cpp_backend() @@ -1317,7 +1555,9 @@ async def load_model( and llama_backend.model_identifier and llama_backend.model_identifier.lower() == model_identifier.lower() # Match runtime settings so Apply isn't dropped (#5401). - and _request_matches_loaded_settings(request, llama_backend) + and _request_matches_loaded_settings( + request, llama_backend, effective_chat_template_override + ) # Skip if a prior audio probe failed -- let load_model retry. and getattr(llama_backend, "_audio_probed", True) ): @@ -1342,6 +1582,7 @@ async def load_model( is_vision = llama_backend._is_vision, is_lora = False, is_gguf = True, + is_diffusion = llama_backend.is_diffusion, is_audio = _gguf_is_audio, audio_type = _gguf_audio, has_audio_input = getattr(llama_backend, "_has_audio_input", False), @@ -1360,6 +1601,7 @@ async def load_model( chat_template = llama_backend.chat_template, speculative_type = llama_backend.requested_spec_mode, spec_draft_n_max = llama_backend.spec_draft_n_max, + tensor_parallel = llama_backend.tensor_parallel, ) else: if ( @@ -1480,7 +1722,13 @@ async def load_model( else: # Strip only the groups whose first-class field was set by # the caller, so an inherited --chat-template-file survives - # an Apply that omits chat_template_override. + # an Apply that omits chat_template_override. A bundled family + # template (e.g. the gemma-4 override) is an effective + # first-class template setting even when the raw request + # omits chat_template_override, so strip the inherited + # --chat-template-file in that case too -- otherwise the stale + # extra arg (appended last) shadows the bundled template while + # Studio reports the bundled template's capabilities. fields_set = getattr(request, "model_fields_set", set()) stripped = strip_shadowing_flags( llama_backend.extra_args, @@ -1489,7 +1737,13 @@ async def load_model( strip_spec = ( "speculative_type" in fields_set or "spec_draft_n_max" in fields_set ), - strip_template = "chat_template_override" in fields_set, + strip_template = ( + "chat_template_override" in fields_set + or effective_chat_template_override is not None + ), + strip_split_mode = _should_strip_split_mode( + request, llama_backend.extra_args + ), ) try: extra_llama_args = validate_extra_args(stripped) @@ -1514,22 +1768,25 @@ async def load_model( # during the (potentially long) GGUF download + llama-server start. _n_parallel = getattr(fastapi_request.app.state, "llama_parallel_slots", 1) + # Load kwargs common to HF and local modes; the two differ only by + # the model-source args (hf_repo/-token vs gguf_path/mmproj). + _common_load_kwargs = dict( + model_identifier = config.identifier, + is_vision = config.is_vision, + n_ctx = request.max_seq_length, + chat_template_override = effective_chat_template_override, + cache_type_kv = request.cache_type_kv, + speculative_type = request.speculative_type, + spec_draft_n_max = request.spec_draft_n_max, + n_parallel = _n_parallel, + extra_args = extra_llama_args, + ) if config.gguf_hf_repo: # HF mode: download via huggingface_hub then start llama-server - success = await asyncio.to_thread( - llama_backend.load_model, + _source_load_kwargs = dict( hf_repo = config.gguf_hf_repo, hf_variant = config.gguf_variant, hf_token = request.hf_token, - model_identifier = config.identifier, - is_vision = config.is_vision, - n_ctx = request.max_seq_length, - chat_template_override = request.chat_template_override, - cache_type_kv = request.cache_type_kv, - speculative_type = request.speculative_type, - spec_draft_n_max = request.spec_draft_n_max, - n_parallel = _n_parallel, - extra_args = extra_llama_args, ) else: # Local mode: llama-server loads via -m @@ -1548,8 +1805,7 @@ async def load_model( except HTTPException as exc: logger.warning("Dropping MTP drafter for native load: %s", exc.detail) config.gguf_mtp_file = None - success = await asyncio.to_thread( - llama_backend.load_model, + _source_load_kwargs = dict( gguf_path = config.gguf_file, mmproj_path = config.gguf_mmproj_file, mtp_draft_path = config.gguf_mtp_file, @@ -1557,17 +1813,36 @@ async def load_model( # the same string the inheritance check at the top of /load # uses (#5401 followup). hf_variant = config.gguf_variant, - model_identifier = config.identifier, - is_vision = config.is_vision, - n_ctx = request.max_seq_length, - chat_template_override = request.chat_template_override, - cache_type_kv = request.cache_type_kv, - speculative_type = request.speculative_type, - spec_draft_n_max = request.spec_draft_n_max, - n_parallel = _n_parallel, - extra_args = extra_llama_args, ) + # Run a single load attempt with the given tensor flag + extras. + async def _attempt_gguf_load( + tensor_parallel: bool, attempt_extra_args: Optional[list[str]] + ) -> bool: + attempt_kwargs = { + **_common_load_kwargs, + "extra_args": attempt_extra_args, + } + return await asyncio.to_thread( + llama_backend.load_model, + **_source_load_kwargs, + **attempt_kwargs, + tensor_parallel = tensor_parallel, + ) + + # Tensor parallelism is arch-gated in llama.cpp and crashes some loads + # outright (e.g. Gemma 3n aborts with a GGML_ASSERT). The helper auto- + # falls back to layer split so the checkbox never blocks a model from + # loading; the response reports the backend's actual tensor_parallel + # state so the UI toggle reflects the fallback. + success = await load_with_tensor_fallback( + _attempt_gguf_load, + requested_tensor = request.tensor_parallel, + extra_args = extra_llama_args, + label = config.identifier, + cancelled = llama_backend.load_cancelled, + ) + if not success: raise HTTPException( status_code = 500, @@ -1595,6 +1870,7 @@ async def load_model( is_vision = llama_backend.is_vision, is_lora = False, is_gguf = True, + is_diffusion = llama_backend.is_diffusion, is_audio = _gguf_is_audio, audio_type = _gguf_audio, has_audio_input = llama_backend._has_audio_input, @@ -1612,6 +1888,7 @@ async def load_model( chat_template = llama_backend.chat_template, speculative_type = llama_backend.requested_spec_mode, spec_draft_n_max = llama_backend.spec_draft_n_max, + tensor_parallel = llama_backend.tensor_parallel, ) # ── Standard path: load via Unsloth/transformers ────────── @@ -2072,11 +2349,27 @@ async def get_status(current_subject: str = Depends(get_current_subject)): _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) + # Don't surface Studio's auto-applied bundled family template (e.g. the + # gemma-4 override) as a user-authored override: the frontend adopts + # status.chat_template_override as editable state and would otherwise + # re-send it as an explicit override for a later, unrelated model. Only + # expose a genuine user override. + _reported_chat_template_override = llama_backend.chat_template_override + _auto_chat_template_override = resolve_effective_chat_template_override( + model_identifier = _model_id, + user_override = None, + ) + if ( + _auto_chat_template_override is not None + and _reported_chat_template_override == _auto_chat_template_override + ): + _reported_chat_template_override = None return InferenceStatusResponse( active_model = _display_model_id, model_identifier = None if _native_grant_backed else _model_id, is_vision = llama_backend.is_vision, is_gguf = True, + is_diffusion = llama_backend.is_diffusion, gguf_variant = llama_backend.hf_variant, is_audio = getattr(llama_backend, "_is_audio", False), audio_type = _audio_type, @@ -2097,9 +2390,10 @@ async def get_status(current_subject: str = Depends(get_current_subject)): 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, + chat_template_override = _reported_chat_template_override, speculative_type = llama_backend.requested_spec_mode, spec_draft_n_max = llama_backend.spec_draft_n_max, + tensor_parallel = llama_backend.tensor_parallel, llama_cpp_supports_mtp = _supports_mtp, spec_fallback_reason = llama_backend.spec_fallback_reason, llama_cpp_prebuilt_stale = _stale, @@ -3899,6 +4193,10 @@ async def openai_chat_completions( _stream_usage = cumulative.get("usage") _stream_timings = cumulative.get("timings") _stream_finish = cumulative.get("finish_reason") + elif cumulative.get("type") == "diffusion_frame": + # Diffusion frame (per-step canvas): pass through as a raw SSE line on the + # tool_status channel. No assistant text, so it never enters the cumulative diff. + yield f"data: {json.dumps(cumulative)}\n\n" else: logger.warning( "gguf_stream_chunks: unexpected dict event: %s", @@ -4800,6 +5098,8 @@ async def openai_completions(request: Request, current_subject: str = Depends(ge ) body = await request.json() + if body.get("max_tokens") is None: + body["max_tokens"] = llama_backend.context_length or _DEFAULT_MAX_TOKENS_FLOOR target_url = f"{llama_backend.base_url}/v1/completions" is_stream = body.get("stream", False) @@ -4817,15 +5117,27 @@ async def openai_completions(request: Request, current_subject: str = Depends(ge # honor stream_options.include_usage per event, while keeping SSE # framing and token bytes intact. _include_usage = bool((body.get("stream_options") or {}).get("include_usage")) - client = httpx.AsyncClient(timeout = 600) + client = httpx.AsyncClient(timeout = _llama_streaming_generation_timeout()) resp = None bytes_iter = None try: req = client.build_request("POST", target_url, json = body) - resp = await client.send(req, stream = True) + first_token_deadline = time.monotonic() + _DEFAULT_FIRST_TOKEN_TIMEOUT_S + resp = await _send_stream_with_preheader_cancel(client, req, request = request) + if resp is None: + return + if resp.status_code != 200: + err_bytes = await resp.aread() + err_text = err_bytes.decode("utf-8", errors = "replace") + raise RuntimeError(f"llama-server returned {resp.status_code}: {err_text}") bytes_iter = resp.aiter_bytes() buffer = b"" - async for chunk in bytes_iter: + async for chunk in _aiter_llama_stream_items( + bytes_iter, + request = request, + first_token_deadline = first_token_deadline, + response = resp, + ): buffer += chunk while b"\n\n" in buffer: event, buffer = buffer.split(b"\n\n", 1) @@ -4841,6 +5153,9 @@ async def openai_completions(request: Request, current_subject: str = Depends(ge yield out + b"\n\n" except Exception as e: logger.error("openai_completions stream error: %s", e) + error_chunk = _openai_stream_error_chunk(e) + yield f"data: {json.dumps(error_chunk)}\n\n".encode("utf-8") + return finally: if bytes_iter is not None: try: @@ -4860,7 +5175,11 @@ async def openai_completions(request: Request, current_subject: str = Depends(ge return StreamingResponse(_stream(), media_type = "text/event-stream") else: async with httpx.AsyncClient() as client: - resp = await client.post(target_url, json = body, timeout = 600) + resp = await client.post( + target_url, + json = body, + timeout = _llama_non_streaming_generation_timeout(), + ) if resp.status_code != 200: raise _openai_passthrough_error(resp.status_code, resp.text) @@ -4898,7 +5217,7 @@ async def openai_embeddings(request: Request, current_subject: str = Depends(get target_url = f"{llama_backend.base_url}/v1/embeddings" async with httpx.AsyncClient() as client: - resp = await client.post(target_url, json = body, timeout = 600) + resp = await client.post(target_url, json = body, timeout = _DEFAULT_FIRST_TOKEN_TIMEOUT_S) return Response( content = resp.content, status_code = resp.status_code, @@ -5001,6 +5320,149 @@ def _responses_tool_output_text(output: Union[str, list]) -> str: return "(no output)" +_RESPONSES_THINK_OPEN = "" +_RESPONSES_THINK_CLOSE = "" +_RESPONSES_REASONING_EFFORTS = {"none", "minimal", "low", "medium", "high", "max", "xhigh"} + + +def _coerce_responses_reasoning_text(value: Any) -> str: + if value is None: + return "" + if isinstance(value, str): + return value + if isinstance(value, list): + return "".join(_coerce_responses_reasoning_text(part) for part in value) + if isinstance(value, dict): + for key in ("text", "reasoning_text", "content"): + text = _coerce_responses_reasoning_text(value.get(key)) + if text: + return text + return "" + return json.dumps(value) + + +def _responses_marker_holdback(text: str, markers: tuple[str, ...]) -> int: + """Number of trailing chars to retain because they may start a marker.""" + for size in range(min(len(text), max(len(m) for m in markers) - 1), 0, -1): + suffix = text[-size:] + if any(marker.startswith(suffix) for marker in markers): + return size + return 0 + + +class _ResponsesReasoningExtractor: + """Split local markup into Responses reasoning and visible text.""" + + def __init__(self, *, parse_think_markers: bool = False) -> None: + self._buffer = "" + self._in_reasoning = False + self._parse_think_markers = parse_think_markers + + def feed( + self, + text: str = "", + reasoning_content: Any = None, + ) -> tuple[str, str]: + reasoning_parts: list[str] = [] + visible_parts: list[str] = [] + structured_reasoning = _coerce_responses_reasoning_text(reasoning_content) + if structured_reasoning: + reasoning_parts.append(structured_reasoning) + if text: + self._buffer += text + if not self._parse_think_markers: + visible_parts.append(self._buffer) + self._buffer = "" + return "".join(reasoning_parts), "".join(visible_parts) + + while self._buffer: + if self._in_reasoning: + close_idx = self._buffer.find(_RESPONSES_THINK_CLOSE) + if close_idx != -1: + reasoning_parts.append(self._buffer[:close_idx]) + self._buffer = self._buffer[close_idx + len(_RESPONSES_THINK_CLOSE) :] + self._in_reasoning = False + continue + keep = _responses_marker_holdback(self._buffer, (_RESPONSES_THINK_CLOSE,)) + if keep == len(self._buffer): + break + reasoning_parts.append(self._buffer[:-keep] if keep else self._buffer) + self._buffer = self._buffer[-keep:] if keep else "" + break + + open_idx = self._buffer.find(_RESPONSES_THINK_OPEN) + close_idx = self._buffer.find(_RESPONSES_THINK_CLOSE) + if close_idx != -1 and (open_idx == -1 or close_idx < open_idx): + visible_parts.append(self._buffer[:close_idx]) + self._buffer = self._buffer[close_idx + len(_RESPONSES_THINK_CLOSE) :] + continue + if open_idx != -1: + visible_parts.append(self._buffer[:open_idx]) + self._buffer = self._buffer[open_idx + len(_RESPONSES_THINK_OPEN) :] + self._in_reasoning = True + continue + + keep = _responses_marker_holdback( + self._buffer, + (_RESPONSES_THINK_OPEN, _RESPONSES_THINK_CLOSE), + ) + if keep == len(self._buffer): + break + visible_parts.append(self._buffer[:-keep] if keep else self._buffer) + self._buffer = self._buffer[-keep:] if keep else "" + break + + return "".join(reasoning_parts), "".join(visible_parts) + + def finish(self) -> tuple[str, str]: + if not self._buffer: + return "", "" + remaining = self._buffer + self._buffer = "" + if not self._parse_think_markers: + return "", remaining + if self._in_reasoning: + self._in_reasoning = False + return remaining, "" + return "", remaining.replace(_RESPONSES_THINK_CLOSE, "") + + +def _extract_responses_reasoning( + text: str = "", + reasoning_content: Any = None, + *, + parse_think_markers: bool = False, +) -> tuple[str, str]: + extractor = _ResponsesReasoningExtractor(parse_think_markers = parse_think_markers) + reasoning, visible = extractor.feed(text, reasoning_content) + final_reasoning, final_visible = extractor.finish() + return reasoning + final_reasoning, visible + final_visible + + +def _responses_should_parse_think_markers( + chat_req: ChatCompletionRequest, llama_backend: Any = None +) -> bool: + if llama_backend is not None and getattr(llama_backend, "is_loaded", False): + if getattr(llama_backend, "reasoning_always_on", False): + return True + if not getattr(llama_backend, "supports_reasoning", False): + return False + if chat_req.enable_thinking is True: + return True + return chat_req.enable_thinking is None and chat_req.reasoning_effort not in (None, "none") + + +def _responses_reasoning_output_item(reasoning_text: str, item_id: Optional[str] = None) -> dict: + kwargs: dict[str, Any] = { + "status": "completed", + "summary": [], + "content": [ResponsesOutputReasoningContent(text = reasoning_text)], + } + if item_id is not None: + kwargs["id"] = item_id + return ResponsesOutputReasoning(**kwargs).model_dump() + + def _normalise_responses_input(payload: ResponsesRequest) -> list[ChatMessage]: """Convert a ResponsesRequest's ``input`` into a Chat-format ``ChatMessage`` list. @@ -5159,6 +5621,33 @@ def _build_chat_request( if payload.parallel_tool_calls is not None: chat_kwargs["parallel_tool_calls"] = payload.parallel_tool_calls + # ``chat_template_kwargs`` (e.g. ``{"enable_thinking": true}``) arrives via + # the Responses extra-body: ResponsesRequest has ``extra="allow"``, so the + # OpenAI SDK's ``extra_body`` spread lands the dict in ``model_extra``. The + # downstream Chat Completions paths consume the typed ``enable_thinking`` + # field -- the non-streaming path lifts it in ``openai_chat_completions`` + # only when it is still ``None``, and the streaming pass-through reads + # ``payload.enable_thinking`` directly -- so lift it here, mirroring that + # handler, to cover both Responses paths. + explicit_enable_thinking = False + _extra = getattr(payload, "model_extra", None) + if isinstance(_extra, dict): + _tpl_kw = _extra.get("chat_template_kwargs") + if isinstance(_tpl_kw, dict) and "enable_thinking" in _tpl_kw: + chat_kwargs["enable_thinking"] = bool(_tpl_kw["enable_thinking"]) + explicit_enable_thinking = True + + if isinstance(payload.reasoning, dict): + effort = payload.reasoning.get("effort") + if isinstance(effort, str) and effort in _RESPONSES_REASONING_EFFORTS: + if not explicit_enable_thinking: + chat_kwargs["reasoning_effort"] = effort + chat_kwargs["enable_thinking"] = effort != "none" + elif chat_kwargs.get("enable_thinking") is False: + chat_kwargs["reasoning_effort"] = "none" + elif effort != "none": + chat_kwargs["reasoning_effort"] = effort + return ChatCompletionRequest(**chat_kwargs) @@ -5202,10 +5691,18 @@ async def _responses_non_streaming( choices = body.get("choices", []) text = "" + reasoning_text = "" tool_calls: list[dict] = [] if choices: msg = choices[0].get("message", {}) or {} - text = msg.get("content", "") or "" + raw_content = msg.get("content", "") or "" + raw_text = raw_content if isinstance(raw_content, str) else json.dumps(raw_content) + llama_backend = get_llama_cpp_backend() + reasoning_text, text = _extract_responses_reasoning( + raw_text, + msg.get("reasoning_content"), + parse_think_markers = _responses_should_parse_think_markers(chat_req, llama_backend), + ) tool_calls = msg.get("tool_calls") or [] usage_data = body.get("usage", {}) @@ -5219,6 +5716,10 @@ async def _responses_non_streaming( # the model produced content, so clients expecting a pure tool-call turn # (finish_reason="tool_calls") don't see a spurious empty message item. output_items: list[dict] = [] + if reasoning_text and not text and not tool_calls: + text = reasoning_text + if reasoning_text: + output_items.append(_responses_reasoning_output_item(reasoning_text)) if text: msg_id = f"msg_{uuid.uuid4().hex[:12]}" output_items.append( @@ -5266,16 +5767,15 @@ async def _responses_stream( avoids that. Non-GGUF falls back to the wrapper (which doesn't use httpx, so the issue doesn't apply). - Text deltas arrive as ``response.output_text.delta`` on a single - ``message`` output item at ``output_index=0``. Each tool call from + Output items are allocated as upstream deltas appear. Reasoning/text deltas + open top-level ``reasoning`` / ``message`` items; each tool call from ``delta.tool_calls[]`` is promoted to its own top-level ``function_call`` - output item (one per distinct ``tool_calls[].index``) and relayed as + item (one per distinct ``tool_calls[].index``) and relayed as ``response.function_call_arguments.delta`` / ``.done`` events so clients (Codex, OpenAI Python SDK) can reconstruct the call incrementally and reply with a ``function_call_output`` item next turn. """ resp_id = f"resp_{uuid.uuid4().hex[:12]}" - msg_id = f"msg_{uuid.uuid4().hex[:12]}" created_at = int(time.time()) chat_req = _build_chat_request(payload, messages, stream = True) @@ -5315,61 +5815,188 @@ async def _responses_stream( async def event_generator(): full_text = "" + full_reasoning = "" input_tokens = 0 output_tokens = 0 + extractor = _ResponsesReasoningExtractor( + parse_think_markers = _responses_should_parse_think_markers(chat_req, llama_backend) + ) + reasoning_state: dict[str, Any] = {"output_index": None, "item_id": None, "opened": False} + message_state: dict[str, Any] = {"output_index": None, "item_id": None, "opened": False} # Per-tool-call state keyed by Chat Completions `tool_calls[].index`, # stable across chunks for the same call. Values: # {output_index, item_id, call_id, name, arguments, opened} tool_call_state: dict[int, dict] = {} - # Text message lives at output_index 0; tool calls claim 1, 2, ... - next_output_index = 1 + next_output_index = 0 + + def _sse(event_name: str, payload: dict) -> str: + return f"event: {event_name}\ndata: {json.dumps(payload)}\n\n" + + def _claim_output_index() -> int: + nonlocal next_output_index + output_index = next_output_index + next_output_index += 1 + return output_index + + def _ensure_reasoning_open() -> list[str]: + if reasoning_state["opened"]: + return [] + reasoning_state["output_index"] = _claim_output_index() + reasoning_state["item_id"] = f"rs_{uuid.uuid4().hex[:12]}" + reasoning_state["opened"] = True + output_index = reasoning_state["output_index"] + item_id = reasoning_state["item_id"] + return [ + _sse( + "response.output_item.added", + { + "type": "response.output_item.added", + "output_index": output_index, + "item": { + "type": "reasoning", + "id": item_id, + "status": "in_progress", + "summary": [], + "content": [], + }, + }, + ), + _sse( + "response.content_part.added", + { + "type": "response.content_part.added", + "item_id": item_id, + "output_index": output_index, + "content_index": 0, + "part": {"type": "reasoning_text", "text": ""}, + }, + ), + ] + + def _ensure_message_open() -> list[str]: + if message_state["opened"]: + return [] + message_state["output_index"] = _claim_output_index() + message_state["item_id"] = f"msg_{uuid.uuid4().hex[:12]}" + message_state["opened"] = True + output_index = message_state["output_index"] + item_id = message_state["item_id"] + return [ + _sse( + "response.output_item.added", + { + "type": "response.output_item.added", + "output_index": output_index, + "item": { + "type": "message", + "id": item_id, + "status": "in_progress", + "role": "assistant", + "content": [], + }, + }, + ), + _sse( + "response.content_part.added", + { + "type": "response.content_part.added", + "item_id": item_id, + "output_index": output_index, + "content_index": 0, + "part": {"type": "output_text", "text": "", "annotations": []}, + }, + ), + ] def _snapshot_output() -> list[dict]: """Snapshot of all completed output items for response.completed.""" - items: list[dict] = [ - { - "type": "message", - "id": msg_id, - "status": "completed", - "role": "assistant", - "content": [ + indexed_items: list[tuple[int, dict]] = [] + if reasoning_state["opened"]: + indexed_items.append( + ( + reasoning_state["output_index"], { - "type": "output_text", - "text": full_text, - "annotations": [], - } - ], - } - ] - for st in sorted(tool_call_state.values(), key = lambda s: s["output_index"]): - items.append( - { - "type": "function_call", - "id": st["item_id"], - "status": "completed", - "call_id": st["call_id"], - "name": st["name"], - "arguments": st["arguments"], - } + "type": "reasoning", + "id": reasoning_state["item_id"], + "status": "completed", + "summary": [], + "content": [{"type": "reasoning_text", "text": full_reasoning}], + }, + ) ) - return items + if message_state["opened"]: + indexed_items.append( + ( + message_state["output_index"], + { + "type": "message", + "id": message_state["item_id"], + "status": "completed", + "role": "assistant", + "content": [ + { + "type": "output_text", + "text": full_text, + "annotations": [], + } + ], + }, + ) + ) + for st in tool_call_state.values(): + indexed_items.append( + ( + st["output_index"], + { + "type": "function_call", + "id": st["item_id"], + "status": "completed", + "call_id": st["call_id"], + "name": st["name"], + "arguments": st["arguments"], + }, + ) + ) + return [item for _, item in sorted(indexed_items, key = lambda pair: pair[0])] + + def _failed_response_payload(exc: Exception, status_code: int) -> dict: + return { + "type": "response.failed", + "response": { + "id": resp_id, + "object": "response", + "created_at": created_at, + "status": "failed", + "model": payload.model, + "output": _snapshot_output(), + "usage": { + "input_tokens": input_tokens, + "output_tokens": output_tokens, + "total_tokens": input_tokens + output_tokens, + }, + "error": { + "code": status_code, + "message": _friendly_error(exc), + }, + }, + } # ── Preamble events ── - yield f"event: response.created\ndata: {json.dumps({'type': 'response.created', 'response': {'id': resp_id, 'object': 'response', 'created_at': created_at, 'status': 'in_progress', 'model': payload.model, 'output': [], 'usage': {'input_tokens': 0, 'output_tokens': 0, 'total_tokens': 0}}})}\n\n" - - # output_item.added (text message at output_index 0) - output_item = { - "type": "message", - "id": msg_id, - "status": "in_progress", - "role": "assistant", - "content": [], - } - yield f"event: response.output_item.added\ndata: {json.dumps({'type': 'response.output_item.added', 'output_index': 0, 'item': output_item})}\n\n" - - # content_part.added - content_part = {"type": "output_text", "text": "", "annotations": []} - yield f"event: response.content_part.added\ndata: {json.dumps({'type': 'response.content_part.added', 'item_id': msg_id, 'output_index': 0, 'content_index': 0, 'part': content_part})}\n\n" + yield _sse( + "response.created", + { + "type": "response.created", + "response": { + "id": resp_id, + "object": "response", + "created_at": created_at, + "status": "in_progress", + "model": payload.model, + "output": [], + "usage": {"input_tokens": 0, "output_tokens": 0, "total_tokens": 0}, + }, + }, + ) # ── Direct httpx lifecycle to llama-server ── # Full same-task open + close, same pattern as @@ -5377,16 +6004,33 @@ async def _responses_stream( # `async with`, explicit aclose of lines_iter BEFORE resp / client so # the innermost httpcore byte stream is finalised in this task (not via # the asyncgen GC in a sibling task). - client = httpx.AsyncClient(timeout = 600) + client = httpx.AsyncClient(timeout = _llama_streaming_generation_timeout()) resp = None lines_iter = None try: req = client.build_request("POST", target_url, json = body) + first_token_deadline = time.monotonic() + _DEFAULT_FIRST_TOKEN_TIMEOUT_S try: - resp = await client.send(req, stream = True) + resp = await _send_stream_with_preheader_cancel(client, req, request = request) + if resp is None: + return except httpx.RequestError as e: logger.error("responses stream: upstream unreachable: %s", e) - yield f"event: response.failed\ndata: {json.dumps({'type': 'response.failed', 'response': {'id': resp_id, 'object': 'response', 'created_at': created_at, 'status': 'failed', 'model': payload.model, 'output': [], 'error': {'code': 502, 'message': _friendly_error(e)}}})}\n\n" + yield _sse( + "response.failed", + { + "type": "response.failed", + "response": { + "id": resp_id, + "object": "response", + "created_at": created_at, + "status": "failed", + "model": payload.model, + "output": [], + "error": {"code": 502, "message": _friendly_error(e)}, + }, + }, + ) return if resp.status_code != 200: @@ -5397,13 +6041,33 @@ async def _responses_stream( resp.status_code, err_text[:500], ) - yield f"event: response.failed\ndata: {json.dumps({'type': 'response.failed', 'response': {'id': resp_id, 'object': 'response', 'created_at': created_at, 'status': 'failed', 'model': payload.model, 'output': [], 'error': {'code': resp.status_code, 'message': f'llama-server error: {err_text[:500]}'}}})}\n\n" + yield _sse( + "response.failed", + { + "type": "response.failed", + "response": { + "id": resp_id, + "object": "response", + "created_at": created_at, + "status": "failed", + "model": payload.model, + "output": [], + "error": { + "code": resp.status_code, + "message": f"llama-server error: {err_text[:500]}", + }, + }, + }, + ) return lines_iter = resp.aiter_lines() - async for raw_line in lines_iter: - if await request.is_disconnected(): - break + async for raw_line in _aiter_llama_stream_items( + lines_iter, + request = request, + first_token_deadline = first_token_deadline, + response = resp, + ): if not raw_line: continue if not raw_line.startswith("data: "): @@ -5427,17 +6091,38 @@ async def _responses_stream( continue delta = choices[0].get("delta", {}) or {} - content = delta.get("content") - if content: - full_text += content - delta_event = { - "type": "response.output_text.delta", - "item_id": msg_id, - "output_index": 0, - "content_index": 0, - "delta": content, - } - yield f"event: response.output_text.delta\ndata: {json.dumps(delta_event)}\n\n" + reasoning_delta, visible_delta = extractor.feed( + delta.get("content") or "", + delta.get("reasoning_content"), + ) + if reasoning_delta: + for event in _ensure_reasoning_open(): + yield event + full_reasoning += reasoning_delta + yield _sse( + "response.reasoning_text.delta", + { + "type": "response.reasoning_text.delta", + "item_id": reasoning_state["item_id"], + "output_index": reasoning_state["output_index"], + "content_index": 0, + "delta": reasoning_delta, + }, + ) + if visible_delta: + for event in _ensure_message_open(): + yield event + full_text += visible_delta + yield _sse( + "response.output_text.delta", + { + "type": "response.output_text.delta", + "item_id": message_state["item_id"], + "output_index": message_state["output_index"], + "content_index": 0, + "delta": visible_delta, + }, + ) for tc in delta.get("tool_calls") or []: idx = tc.get("index", 0) @@ -5447,14 +6132,13 @@ async def _responses_stream( # First chunk for this tool call -- allocate an # output_index and emit output_item.added. st = { - "output_index": next_output_index, + "output_index": _claim_output_index(), "item_id": f"fc_{uuid.uuid4().hex[:12]}", "call_id": tc.get("id") or "", "name": fn.get("name") or "", "arguments": "", "opened": False, } - next_output_index += 1 tool_call_state[idx] = st else: # Later chunks sometimes carry id/name only once; merge @@ -5477,7 +6161,7 @@ async def _responses_stream( "arguments": "", }, } - yield f"event: response.output_item.added\ndata: {json.dumps(item_added)}\n\n" + yield _sse("response.output_item.added", item_added) st["opened"] = True arg_delta = fn.get("arguments") or "" @@ -5489,7 +6173,7 @@ async def _responses_stream( "output_index": st["output_index"], "delta": arg_delta, } - yield f"event: response.function_call_arguments.delta\ndata: {json.dumps(args_delta_event)}\n\n" + yield _sse("response.function_call_arguments.delta", args_delta_event) elif arg_delta: # Buffer args until we can open the item (some models # send id/name in the same chunk as the first arg delta; @@ -5502,6 +6186,12 @@ async def _responses_stream( output_tokens = usage.get("completion_tokens", output_tokens) except Exception as e: logger.error("responses stream error: %s", e) + status_code = 400 if _classify_llama_generation_error(e) is not None else 500 + yield _sse( + "response.failed", + _failed_response_payload(e, status_code), + ) + return finally: if lines_iter is not None: try: @@ -5518,8 +6208,134 @@ async def _responses_stream( except Exception: pass - # ── Closing events for tool calls ── - for st in sorted(tool_call_state.values(), key = lambda s: s["output_index"]): + final_reasoning, final_visible = extractor.finish() + if final_reasoning: + for event in _ensure_reasoning_open(): + yield event + full_reasoning += final_reasoning + yield _sse( + "response.reasoning_text.delta", + { + "type": "response.reasoning_text.delta", + "item_id": reasoning_state["item_id"], + "output_index": reasoning_state["output_index"], + "content_index": 0, + "delta": final_reasoning, + }, + ) + if final_visible: + for event in _ensure_message_open(): + yield event + full_text += final_visible + yield _sse( + "response.output_text.delta", + { + "type": "response.output_text.delta", + "item_id": message_state["item_id"], + "output_index": message_state["output_index"], + "content_index": 0, + "delta": final_visible, + }, + ) + if full_reasoning and not full_text and not tool_call_state: + for event in _ensure_message_open(): + yield event + full_text = full_reasoning + yield _sse( + "response.output_text.delta", + { + "type": "response.output_text.delta", + "item_id": message_state["item_id"], + "output_index": message_state["output_index"], + "content_index": 0, + "delta": full_text, + }, + ) + + close_items: list[tuple[int, str, dict[str, Any]]] = [] + if reasoning_state["opened"]: + close_items.append((reasoning_state["output_index"], "reasoning", reasoning_state)) + if message_state["opened"]: + close_items.append((message_state["output_index"], "message", message_state)) + close_items.extend((st["output_index"], "tool", st) for st in tool_call_state.values()) + + for _, kind, st in sorted(close_items, key = lambda item: item[0]): + if kind == "reasoning": + yield _sse( + "response.reasoning_text.done", + { + "type": "response.reasoning_text.done", + "item_id": st["item_id"], + "output_index": st["output_index"], + "content_index": 0, + "text": full_reasoning, + }, + ) + yield _sse( + "response.content_part.done", + { + "type": "response.content_part.done", + "item_id": st["item_id"], + "output_index": st["output_index"], + "content_index": 0, + "part": {"type": "reasoning_text", "text": full_reasoning}, + }, + ) + yield _sse( + "response.output_item.done", + { + "type": "response.output_item.done", + "output_index": st["output_index"], + "item": { + "type": "reasoning", + "id": st["item_id"], + "status": "completed", + "summary": [], + "content": [{"type": "reasoning_text", "text": full_reasoning}], + }, + }, + ) + continue + + if kind == "message": + yield _sse( + "response.output_text.done", + { + "type": "response.output_text.done", + "item_id": st["item_id"], + "output_index": st["output_index"], + "content_index": 0, + "text": full_text, + }, + ) + yield _sse( + "response.content_part.done", + { + "type": "response.content_part.done", + "item_id": st["item_id"], + "output_index": st["output_index"], + "content_index": 0, + "part": {"type": "output_text", "text": full_text, "annotations": []}, + }, + ) + yield _sse( + "response.output_item.done", + { + "type": "response.output_item.done", + "output_index": st["output_index"], + "item": { + "type": "message", + "id": st["item_id"], + "status": "completed", + "role": "assistant", + "content": [ + {"type": "output_text", "text": full_text, "annotations": []} + ], + }, + }, + ) + continue + # If id/name never arrived (malformed upstream), synthesise so the # client still sees a coherent frame sequence. if not st["opened"]: @@ -5537,20 +6353,16 @@ async def _responses_stream( "arguments": "", }, } - yield f"event: response.output_item.added\ndata: {json.dumps(item_added)}\n\n" + yield _sse("response.output_item.added", item_added) if st["arguments"]: - yield ( - "event: response.function_call_arguments.delta\n" - "data: " - + json.dumps( - { - "type": "response.function_call_arguments.delta", - "item_id": st["item_id"], - "output_index": st["output_index"], - "delta": st["arguments"], - } - ) - + "\n\n" + yield _sse( + "response.function_call_arguments.delta", + { + "type": "response.function_call_arguments.delta", + "item_id": st["item_id"], + "output_index": st["output_index"], + "delta": st["arguments"], + }, ) st["opened"] = True @@ -5561,7 +6373,7 @@ async def _responses_stream( "name": st["name"], "arguments": st["arguments"], } - yield f"event: response.function_call_arguments.done\ndata: {json.dumps(args_done)}\n\n" + yield _sse("response.function_call_arguments.done", args_done) item_done = { "type": "response.output_item.done", @@ -5575,14 +6387,7 @@ async def _responses_stream( "arguments": st["arguments"], }, } - yield f"event: response.output_item.done\ndata: {json.dumps(item_done)}\n\n" - - # ── Closing events for text message ── - yield f"event: response.output_text.done\ndata: {json.dumps({'type': 'response.output_text.done', 'item_id': msg_id, 'output_index': 0, 'content_index': 0, 'text': full_text})}\n\n" - - yield f"event: response.content_part.done\ndata: {json.dumps({'type': 'response.content_part.done', 'item_id': msg_id, 'output_index': 0, 'content_index': 0, 'part': {'type': 'output_text', 'text': full_text, 'annotations': []}})}\n\n" - - yield f"event: response.output_item.done\ndata: {json.dumps({'type': 'response.output_item.done', 'output_index': 0, 'item': {'type': 'message', 'id': msg_id, 'status': 'completed', 'role': 'assistant', 'content': [{'type': 'output_text', 'text': full_text, 'annotations': []}]}})}\n\n" + yield _sse("response.output_item.done", item_done) # response.completed total_tokens = input_tokens + output_tokens @@ -5602,7 +6407,7 @@ async def _responses_stream( }, }, } - yield f"event: response.completed\ndata: {json.dumps(completed_response)}\n\n" + yield _sse("response.completed", completed_response) return StreamingResponse( event_generator(), @@ -6496,7 +7301,6 @@ def _build_passthrough_payload( body["max_tokens"] = ( max_tokens if max_tokens is not None else (backend_ctx or _DEFAULT_MAX_TOKENS_FLOOR) ) - body["t_max_predict_ms"] = _DEFAULT_T_MAX_PREDICT_MS # Normalize stop the same way the non-passthrough path does (the passthrough # was previously the one path that forwarded an empty stop string verbatim). _stop = _normalize_stop_sequences(stop) @@ -6606,7 +7410,7 @@ async def _anthropic_passthrough_stream( # `try: ... except Exception: pass` so nested anyio cleanup noise can't # bubble out. client = httpx.AsyncClient( - timeout = 600, + timeout = _llama_streaming_generation_timeout(), limits = httpx.Limits(max_keepalive_connections = 0), ) resp = None @@ -6614,7 +7418,12 @@ async def _anthropic_passthrough_stream( cancel_watcher = None try: req = client.build_request("POST", target_url, json = body) - resp = await client.send(req, stream = True) + first_token_deadline = time.monotonic() + _DEFAULT_FIRST_TOKEN_TIMEOUT_S + resp = await _send_stream_with_preheader_cancel( + client, req, cancel_event, request = request + ) + if resp is None: + return # Upstream client error (e.g. over-context 400) arrives before any # SSE. The 200 stream headers are already flushed, so surface it as @@ -6643,12 +7452,13 @@ async def _anthropic_passthrough_stream( # The watcher closes `resp` on cancel, raising in aiter_lines. cancel_watcher = asyncio.create_task(_await_cancel_then_close(cancel_event, resp)) lines_iter = resp.aiter_lines() - async for raw_line in lines_iter: - if cancel_event.is_set(): - break - if await request.is_disconnected(): - cancel_event.set() - break + async for raw_line in _aiter_llama_stream_items( + lines_iter, + cancel_event = cancel_event, + request = request, + first_token_deadline = first_token_deadline, + response = resp, + ): if not raw_line or not raw_line.startswith("data: "): continue data_str = raw_line[6:] @@ -6662,11 +7472,26 @@ async def _anthropic_passthrough_stream( _drop_parallel_tool_call_deltas(chunk) for line in emitter.feed_chunk(chunk): yield line - except (httpx.RemoteProtocolError, httpx.ReadError, httpx.CloseError): + except (httpx.RemoteProtocolError, httpx.ReadError, httpx.CloseError) as e: if not cancel_event.is_set(): - raise + logger.error("anthropic_messages passthrough stream error: %s", e) + event = _anthropic_stream_error_event( + e, + force = True, + ) + if event is not None: + yield event + return except Exception as e: - logger.error("anthropic_messages passthrough stream error: %s", e) + if not cancel_event.is_set(): + logger.error("anthropic_messages passthrough stream error: %s", e) + event = _anthropic_stream_error_event( + e, + force = True, + ) + if event is not None: + yield event + return finally: if cancel_watcher is not None: cancel_watcher.cancel() @@ -6740,7 +7565,11 @@ async def _anthropic_passthrough_non_streaming( ) async with httpx.AsyncClient() as client: - resp = await client.post(target_url, json = body, timeout = 600) + resp = await client.post( + target_url, + json = body, + timeout = _llama_non_streaming_generation_timeout(), + ) if resp.status_code != 200: raise HTTPException( @@ -7094,15 +7923,13 @@ async def _openai_passthrough_stream( _tracker = _TrackedCancel(cancel_event, *_cancel_keys) _tracker.__enter__() - # Outer guard: asyncio.CancelledError at `await client.send(...)` is a - # BaseException that bypasses `except httpx.RequestError`; without this the - # tracker leaks. The generator's finally only runs once iteration starts. + # Keep tracker cleanup paired if pre-header dispatch is cancelled. try: # Dispatch BEFORE returning StreamingResponse so transport errors and # non-200 upstream statuses surface as real HTTP errors -- OpenAI SDKs # rely on status codes to raise APIError/BadRequestError. client = httpx.AsyncClient( - timeout = 600, + timeout = _llama_streaming_generation_timeout(), limits = httpx.Limits(max_keepalive_connections = 0), ) resp = None @@ -7112,7 +7939,10 @@ async def _openai_passthrough_stream( while True: try: req = client.build_request("POST", target_url, json = body) - resp = await client.send(req, stream = True) + first_token_deadline = time.monotonic() + _DEFAULT_FIRST_TOKEN_TIMEOUT_S + resp = await _send_stream_with_preheader_cancel( + client, req, cancel_event, request = request + ) except httpx.RequestError as e: # llama-server subprocess crashed / starting / unreachable. logger.error("openai passthrough stream: upstream unreachable: %s", e) @@ -7129,6 +7959,21 @@ async def _openai_passthrough_stream( status_code = 502, detail = _friendly_error(e), ) + if resp is None: + try: + await client.aclose() + except Exception: + pass + _tracker.__exit__(None, None, None) + return StreamingResponse( + iter(()), + media_type = "text/event-stream", + headers = { + "Cache-Control": "no-cache", + "Connection": "keep-alive", + "X-Accel-Buffering": "no", + }, + ) if resp.status_code == 200: break @@ -7172,12 +8017,13 @@ async def _openai_passthrough_stream( cancel_watcher = asyncio.create_task(_await_cancel_then_close(cancel_event, resp)) try: lines_iter = resp.aiter_lines() - async for raw_line in lines_iter: - if cancel_event.is_set(): - break - if await request.is_disconnected(): - cancel_event.set() - break + async for raw_line in _aiter_llama_stream_items( + lines_iter, + cancel_event = cancel_event, + request = request, + first_token_deadline = first_token_deadline, + response = resp, + ): if not raw_line: continue if not raw_line.startswith("data: "): @@ -7256,7 +8102,11 @@ async def _openai_passthrough_non_streaming(llama_backend, payload, model_name): while True: try: async with httpx.AsyncClient() as client: - resp = await client.post(target_url, json = body, timeout = 600) + resp = await client.post( + target_url, + json = body, + timeout = _llama_non_streaming_generation_timeout(), + ) except httpx.RequestError as e: # llama-server subprocess crashed / starting / unreachable. Surface the # same friendly message the sync chat path emits so operators don't see diff --git a/studio/backend/routes/mcp_servers.py b/studio/backend/routes/mcp_servers.py index 3001c6b7c9..37d99a222e 100644 --- a/studio/backend/routes/mcp_servers.py +++ b/studio/backend/routes/mcp_servers.py @@ -10,12 +10,16 @@ from fastapi import APIRouter, Depends, HTTPException from auth.authentication import get_current_subject from core.inference.mcp_client import ( + TOOL_CACHE_INVALIDATING_FIELDS, + cache_tools, clear_oauth_tokens_async, + invalidate_tool_cache, is_stdio, list_tools_async, parse_server_headers, parse_stdio_command, probe_timeout, + record_probe_failure, stdio_mcp_enabled, ) from core.inference.mcp_config_import import parse_mcp_config @@ -198,6 +202,11 @@ async def update_mcp_server( ): await clear_oauth_tokens_async(old["url"]) mcp_servers_db.update_server(server_id, changes) + # A new endpoint/auth makes cached tools wrong and disabling makes them + # unreachable, so drop them and let the next send re-probe; a rename + # leaves them valid. + if changes.keys() & TOOL_CACHE_INVALIDATING_FIELDS: + invalidate_tool_cache(server_id) return _row_to_response(mcp_servers_db.get_server(server_id)) @@ -209,6 +218,7 @@ async def delete_mcp_server(server_id: str, current_subject: str = Depends(get_c if old.get("use_oauth"): await clear_oauth_tokens_async(old["url"]) mcp_servers_db.delete_server(server_id) + invalidate_tool_cache(server_id) @router.post("/{server_id}/refresh", response_model = McpServerProbeResult) @@ -238,8 +248,23 @@ async def refresh_mcp_server_tools( error = str(exc), exc_info = True, ) + current = mcp_servers_db.get_server(server_id) + if current is not None and not any( + current.get(k) != server.get(k) for k in TOOL_CACHE_INVALIDATING_FIELDS + ): + # Start the cool-off so the next chat send doesn't immediately re-hang + # on this server's timeout. If the row changed while the probe was + # awaiting, the failure belongs to the old config and must not park + # the newly edited server. + record_probe_failure(server_id, use_oauth) return McpServerProbeResult(ok = False, error = safe_curated_detail(exc)) + # Warm the chat-path cache so the next send skips re-probing. + current = mcp_servers_db.get_server(server_id) + if current is not None and not any( + current.get(k) != server.get(k) for k in TOOL_CACHE_INVALIDATING_FIELDS + ): + cache_tools(server_id, tools) return McpServerProbeResult(ok = True, tool_count = len(tools)) diff --git a/studio/backend/routes/models.py b/studio/backend/routes/models.py index 7151af33f3..a2f2eca81b 100644 --- a/studio/backend/routes/models.py +++ b/studio/backend/routes/models.py @@ -51,6 +51,14 @@ def _is_hidden_model(*values: str | None) -> bool: return any(v and any(n in v.lower() for n in needles) for v in values) +def _safe_resolve(path: Path) -> Optional[str]: + """resolve() to a string, or None when the path is inaccessible.""" + try: + return str(path.resolve()) + except OSError: + return None + + backend_path = Path(__file__).parent.parent.parent if str(backend_path) not in sys.path: sys.path.insert(0, str(backend_path)) @@ -676,9 +684,9 @@ async def list_local_models( # trusted Path objects are used for FS access; the user string is # used for matching only, never for path construction. allowed_roots: list[Path] = [Path("./models").resolve(), hf_cache_dir] - if legacy_hf.is_dir(): + if _safe_is_dir(legacy_hf): allowed_roots.append(legacy_hf) - if hf_default.is_dir(): + if _safe_is_dir(hf_default): allowed_roots.append(hf_default) try: from utils.paths import studio_root, outputs_root @@ -702,15 +710,20 @@ async def list_local_models( try: local_models = _scan_models_dir(models_root) + _scan_hf_cache(hf_cache_dir) + # Resolve once; an inaccessible aux cache must skip that scan, not 500. + hf_cache_real = _safe_resolve(hf_cache_dir) + legacy_real = _safe_resolve(legacy_hf) + default_real = _safe_resolve(hf_default) + # Scan legacy Unsloth HF cache for backward compatibility. - if legacy_hf.is_dir() and legacy_hf.resolve() != hf_cache_dir.resolve(): + if _safe_is_dir(legacy_hf) and legacy_real != hf_cache_real: local_models += _scan_hf_cache(legacy_hf) # Scan HF system default cache (may differ under env overrides). if ( - hf_default.is_dir() - and hf_default.resolve() != hf_cache_dir.resolve() - and hf_default.resolve() != legacy_hf.resolve() + _safe_is_dir(hf_default) + and default_real != hf_cache_real + and default_real != legacy_real ): local_models += _scan_hf_cache(hf_default) @@ -2069,13 +2082,10 @@ async def get_gguf_variants( best = _pick_best_gguf(filenames) default_variant = _extract_quant_label(best) if best else None - # Which variants are fully downloaded in the HF cache. For split - # GGUFs ALL shards must be present, so sum cached bytes per variant - # vs. the expected total. Cache dir casing may differ from the - # canonical repo_id, so match case-insensitively. - cached_bytes_by_quant: dict[str, int] = {} + # Per-snapshot so a split GGUF's shards must all sit in one snapshot; + # mmproj adapters are excluded so they can't inflate a quant's bytes. + cached_bytes_by_quant_per_snapshot: list[dict[str, int]] = [] try: - import re as _re from huggingface_hub import constants as hf_constants if not _is_valid_repo_id(repo_id): @@ -2088,21 +2098,31 @@ async def get_gguf_variants( snapshots = entry / "snapshots" if snapshots.is_dir(): for snap in snapshots.iterdir(): + by_quant: dict[str, int] = {} for f in _iter_gguf_paths(snap): - q = _extract_quant_label(f.name) - cached_bytes_by_quant[q] = ( - cached_bytes_by_quant.get(q, 0) + f.stat().st_size - ) + if _is_mmproj_filename(f.name): + continue + try: + size = f.stat().st_size + except OSError: + continue # broken symlink / unreadable: skip + q = _extract_quant_label(f.name).lower() + by_quant[q] = by_quant.get(q, 0) + size + if by_quant: + cached_bytes_by_quant_per_snapshot.append(by_quant) break except Exception: pass def _is_fully_downloaded(variant) -> bool: - cached = cached_bytes_by_quant.get(variant.quant, 0) - if cached == 0 or variant.size_bytes == 0: + if variant.size_bytes == 0: return False - # Rounding tolerance (symlinks vs real sizes). - return cached >= variant.size_bytes * 0.99 + # Complete within one snapshot (tolerance for symlink size jitter). + quant = variant.quant.lower() + return any( + by_quant.get(quant, 0) >= variant.size_bytes * 0.99 + for by_quant in cached_bytes_by_quant_per_snapshot + ) return GgufVariantsResponse( repo_id = repo_id, @@ -2157,16 +2177,26 @@ async def get_gguf_download_progress( for entry in cache_dir.iterdir(): if entry.name.lower() == target: # Completed .gguf files for this variant in snapshots. + # Exclude mmproj so a vision adapter can't satisfy a same-label + # main variant (e.g. mmproj-F16 vs an F16 weight). for f in _iter_gguf_paths(entry): + if _is_mmproj_filename(f.name): + continue fname = f.name.lower().replace("-", "").replace("_", "") if not variant_lower or variant_lower in fname: - downloaded_bytes += f.stat().st_size + try: + downloaded_bytes += f.stat().st_size + except OSError: + continue # broken symlink / unreadable: skip # In-progress (.incomplete) downloads in blobs. blobs_dir = entry / "blobs" if blobs_dir.is_dir(): for f in blobs_dir.iterdir(): if f.is_file() and f.name.endswith(".incomplete"): - in_progress_bytes += f.stat().st_size + try: + in_progress_bytes += f.stat().st_size + except OSError: + continue break total_progress_bytes = downloaded_bytes + in_progress_bytes @@ -2300,11 +2330,22 @@ def _get_repo_size_cached(repo_id: str) -> int: def _all_hf_cache_scans(): - """scan_cache_dir results for the active, legacy, and default HF caches.""" + """scan_cache_dir for the active, legacy, and default HF caches. + + Each probe is isolated: an unreadable auxiliary cache (permission denied, + broken symlink, OS-redirected ~/.cache) is skipped, not fatal, so the + Downloaded list never blanks out and downloads never leak into Recommended. + """ from huggingface_hub import scan_cache_dir from utils.paths import legacy_hf_cache_dir, hf_default_cache_dir - scans = [scan_cache_dir()] + scans = [] + # Guard the active cache too: degrade to "no downloads" instead of raising. + try: + scans.append(scan_cache_dir()) + except Exception as exc: + logger.warning("Could not scan active HF cache: %s", exc) + seen: set[str] = set() try: # Resolve the active cache dir for dedup. @@ -2314,13 +2355,18 @@ def _all_hf_cache_scans(): pass for extra_fn in (legacy_hf_cache_dir, hf_default_cache_dir): - extra = extra_fn() - if extra.is_dir() and str(extra.resolve()) not in seen: - seen.add(str(extra.resolve())) - try: - scans.append(scan_cache_dir(cache_dir = str(extra))) - except Exception as exc: - logger.warning("Could not scan HF cache %s: %s", extra, exc) + try: + extra = extra_fn() + # is_dir()/resolve() can raise on an inaccessible path; skip it. + if not extra.is_dir(): + continue + resolved = str(extra.resolve()) + if resolved in seen: + continue + seen.add(resolved) + scans.append(scan_cache_dir(cache_dir = str(extra))) + except Exception as exc: + logger.warning("Could not scan HF cache %s: %s", extra_fn.__name__, exc) return scans @@ -2379,6 +2425,38 @@ def _repo_has_gguf_files(repo_info) -> bool: return _repo_gguf_size_bytes(repo_info) > 0 +def _blob_mtime(f) -> float: + """Blob modification time in epoch seconds (0.0 if unknown). + + Prefers HF metadata ``blob_last_modified``, falls back to stat(); uses + only mtimes (portable across Windows, macOS, Linux), never path parsing. + """ + ts = getattr(f, "blob_last_modified", None) + if isinstance(ts, (int, float)) and ts > 0: + return float(ts) + blob_path = getattr(f, "blob_path", None) + if blob_path: + try: + return float(Path(blob_path).stat().st_mtime) + except OSError: + pass + return 0.0 + + +def _repo_gguf_last_modified(repo_info) -> float: + """Newest mtime among a repo's primary (non-mmproj) GGUF blobs. + + Drives the Downloaded list's "last downloaded" ordering and groups a + multi-quant repo by its most recently downloaded quant. + """ + latest = 0.0 + for revision in repo_info.revisions: + for f in revision.files: + if _is_main_gguf_filename(f.file_name): + latest = max(latest, _blob_mtime(f)) + return latest + + @router.get("/cached-gguf") async def list_cached_gguf(current_subject: str = Depends(get_current_subject)): """List GGUF repos downloaded to HF cache, legacy Unsloth cache, and HF default cache.""" @@ -2399,17 +2477,30 @@ async def list_cached_gguf(current_subject: str = Depends(get_current_subject)): continue key = repo_id.lower() existing = seen_lower.get(key) + last_modified = _repo_gguf_last_modified(repo_info) if existing is None or total_size > existing["size_bytes"]: - seen_lower[key] = { + row = { "repo_id": repo_id, "size_bytes": total_size, "cache_path": str(repo_info.repo_path), } + # Keep the newest timestamp across duplicate caches; + # attach only when known so absent rows sort as oldest. + lm = max(last_modified, (existing or {}).get("last_modified", 0.0)) + if lm > 0: + row["last_modified"] = lm + seen_lower[key] = row + elif last_modified > existing.get("last_modified", 0.0): + existing["last_modified"] = last_modified except Exception as e: repo_label = getattr(repo_info, "repo_id", "") logger.warning(f"Skipping cached GGUF repo {repo_label}: {e}") continue - cached = sorted(seen_lower.values(), key = lambda c: c["repo_id"]) + # Newest download first; stable repo_id tie-break for equal/missing mtimes. + cached = sorted( + seen_lower.values(), + key = lambda c: (-(c.get("last_modified") or 0.0), c["repo_id"].lower()), + ) return {"cached": cached} except Exception as e: logger.error(f"Error listing cached GGUF repos: {e}", exc_info = True) @@ -2447,18 +2538,39 @@ async def list_cached_models(current_subject: str = Depends(get_current_subject) ) if not has_weights: continue + last_modified = max( + ( + _blob_mtime(f) + for rev in repo_info.revisions + for f in rev.files + if f.file_name.endswith(_WEIGHT_EXTENSIONS) + ), + default = 0.0, + ) key = repo_id.lower() existing = seen_lower.get(key) if existing is None or total_size > existing["size_bytes"]: - seen_lower[key] = { + row = { "repo_id": repo_id, "size_bytes": total_size, } + # Keep the newest timestamp across duplicate caches; + # attach only when known so absent rows sort as oldest. + lm = max(last_modified, (existing or {}).get("last_modified", 0.0)) + if lm > 0: + row["last_modified"] = lm + seen_lower[key] = row + elif last_modified > existing.get("last_modified", 0.0): + existing["last_modified"] = last_modified except Exception as e: repo_label = getattr(repo_info, "repo_id", "") logger.warning(f"Skipping cached model repo {repo_label}: {e}") continue - cached = sorted(seen_lower.values(), key = lambda c: c["repo_id"]) + # Newest download first; stable repo_id tie-break for equal/missing mtimes. + cached = sorted( + seen_lower.values(), + key = lambda c: (-(c.get("last_modified") or 0.0), c["repo_id"].lower()), + ) return {"cached": cached} except Exception as e: logger.error(f"Error listing cached models: {e}", exc_info = True) diff --git a/studio/backend/routes/providers.py b/studio/backend/routes/providers.py index 64d1c3eab3..5a55c9b0bb 100644 --- a/studio/backend/routes/providers.py +++ b/studio/backend/routes/providers.py @@ -190,7 +190,8 @@ async def test_provider( """ Test connectivity to an external provider. - Makes a lightweight GET /models call to verify the API key works. + Makes a lightweight GET /models call to verify the API key works. Generic + custom endpoints use a chat-completions probe because /models is optional. encrypted_api_key is decrypted server-side and never stored. """ info = get_provider_info(payload.provider_type) @@ -212,6 +213,14 @@ async def test_provider( ) base_url = payload.base_url or info["base_url"] + if payload.provider_type == "custom": + if not base_url: + return ProviderTestResult( + success = False, + message = "Connection failed: Base URL is required for custom providers.", + models_count = None, + ) + client = ExternalProviderClient( provider_type = payload.provider_type, base_url = base_url, @@ -220,6 +229,26 @@ async def test_provider( ) try: + if payload.provider_type == "custom": + model_id = (payload.model_id or "").strip() + if not model_id: + return ProviderTestResult( + success = False, + message = "Connection failed: add a model ID to test custom providers.", + models_count = None, + ) + await client.chat_completion( + messages = [{"role": "user", "content": "ping"}], + model = model_id, + temperature = 0.0, + top_p = 1.0, + max_tokens = 1, + ) + return ProviderTestResult( + success = True, + message = "Connected successfully. Chat completions endpoint responded.", + models_count = None, + ) if info.get("model_list_mode") == "curated": await client.verify_models_endpoint_lightweight() return ProviderTestResult( diff --git a/studio/backend/routes/rag.py b/studio/backend/routes/rag.py index 8aba73caac..8d23240fd5 100644 --- a/studio/backend/routes/rag.py +++ b/studio/backend/routes/rag.py @@ -75,6 +75,19 @@ def _save_upload(file: UploadFile) -> tuple[str, str]: return stored_path, filename +def _remove_stored_upload(stored_path: str | None) -> None: + """Best-effort cleanup for files saved by _save_upload.""" + if not stored_path: + return + try: + uploads = os.path.realpath(str(rag_uploads_root())) + target = os.path.realpath(stored_path) + if os.path.isfile(target) and os.path.commonpath([uploads, target]) == uploads: + os.remove(target) + except Exception: # noqa: BLE001 - DB/index deletion has already succeeded. + logger.warning("failed to remove RAG upload %s", stored_path, exc_info = True) + + def _doc_view(row: dict) -> dict: return { "id": row["id"], @@ -84,6 +97,7 @@ def _doc_view(row: dict) -> dict: "numChunks": row.get("num_chunks") or 0, "kbId": row.get("kb_id"), "threadId": row.get("thread_id"), + "projectId": row.get("project_id"), "createdAt": row.get("created_at"), } @@ -102,6 +116,7 @@ class SearchRequest(BaseModel): query: str kb_id: str | None = None thread_id: str | None = None + project_id: str | None = None top_k: int = Field(default = config.TOP_K_HYBRID, ge = 1, le = 50) min_score: float = 0.0 mode: str = "hybrid" # hybrid | lexical | dense @@ -244,14 +259,50 @@ def list_thread_documents(thread_id: str, subject: str = Depends(get_current_sub conn.close() +@router.post("/projects/{project_id}/documents") +async def upload_project_document( + project_id: str, + file: UploadFile = File(...), + subject: str = Depends(get_current_subject), +) -> dict: + _require_rag() + from storage.studio_db import get_chat_project + + if get_chat_project(project_id) is None: + raise HTTPException(status_code = 404, detail = "Project not found") + stored_path, filename = _save_upload(file) + document_id, job_id = ingestion.start_ingestion( + store.project_scope(project_id), + None, + None, + filename, + stored_path, + project_id = project_id, + ) + return {"documentId": document_id, "jobId": job_id, "filename": filename} + + +@router.get("/projects/{project_id}/documents") +def list_project_documents(project_id: str, subject: str = Depends(get_current_subject)) -> dict: + _require_rag() + conn = rag_db.get_connection() + try: + docs = store.list_documents(conn, store.project_scope(project_id)) + return {"documents": [_doc_view(d) for d in docs]} + finally: + conn.close() + + @router.delete("/documents/{document_id}") def delete_document(document_id: str, subject: str = Depends(get_current_subject)) -> dict: _require_rag() conn = rag_db.get_connection() try: - if store.get_document(conn, document_id) is None: + doc = store.get_document(conn, document_id) + if doc is None: raise HTTPException(status_code = 404, detail = "Document not found") store.delete_document(conn, document_id) + _remove_stored_upload(doc.get("stored_path")) return {"ok": True} finally: conn.close() @@ -297,10 +348,15 @@ def search(payload: SearchRequest, subject: str = Depends(get_current_subject)) _require_rag() if payload.kb_id: scope = store.kb_scope(payload.kb_id) - elif payload.thread_id: - scope = store.thread_scope(payload.thread_id) else: - raise HTTPException(status_code = 400, detail = "Provide kb_id or thread_id") + scopes = [] + if payload.project_id: + scopes.append(store.project_scope(payload.project_id)) + if payload.thread_id: + scopes.append(store.thread_scope(payload.thread_id)) + if not scopes: + raise HTTPException(status_code = 400, detail = "Provide kb_id, project_id, or thread_id") + scope = scopes[0] if len(scopes) == 1 else scopes conn = rag_db.get_connection() try: diff --git a/studio/backend/routes/training.py b/studio/backend/routes/training.py index d7687ffdee..281f03bcaf 100644 --- a/studio/backend/routes/training.py +++ b/studio/backend/routes/training.py @@ -125,6 +125,17 @@ async def start_training( backend = get_training_backend() + # S3 dataset loading needs the optional boto3 dependency. Reject early + # with a clear message so credentials are never accepted and then + # silently dropped on a host without boto3 installed. + if request.s3_config is not None: + from core.training.s3_dataset import boto3_available + if not boto3_available(): + raise HTTPException( + status_code = 501, + detail = "S3 dataset loading requires boto3. Install it with: pip install boto3", + ) + # Check before mutating state. if backend.is_training_active(): existing_job_id: Optional[str] = getattr(backend, "current_job_id", "") @@ -235,6 +246,7 @@ async def start_training( "resume_from_checkpoint": request.resume_from_checkpoint, "trust_remote_code": request.trust_remote_code, "gpu_ids": request.gpu_ids, + "s3_config": request.s3_config.model_dump() if request.s3_config else None, } # Training page has no trust_remote_code toggle; as a safety net consult @@ -293,6 +305,10 @@ async def start_training( error = None, ) + except HTTPException: + # Deliberate rejections (S3 not implemented, resume validation) must + # reach the client with their original status, not a generic 500. + raise except ValueError as e: logger.warning("Rejected training GPU selection: %s", e) # Deliberate user-facing GPU-selection validation message. diff --git a/studio/backend/run.py b/studio/backend/run.py index a883154c4a..84991a71bb 100644 --- a/studio/backend/run.py +++ b/studio/backend/run.py @@ -545,6 +545,11 @@ if _STUDIO_ROOT_RESOLVED != _LEGACY_STUDIO_ROOT: if not os.environ.get("UNSLOTH_LLAMA_CPP_PATH"): os.environ["UNSLOTH_LLAMA_CPP_PATH"] = str(_STUDIO_ROOT_RESOLVED / "llama.cpp") +# The studio bundles unsloth_zoo; declare unsloth present (as `import unsloth` +# does) so its lazy submodule imports (export, hardware, mlx) and the +# DiffusionGemma runner never trip the install guard on a clean install. +os.environ.setdefault("UNSLOTH_IS_PRESENT", "1") + def _write_pid_file(): """Write the current process PID to the studio PID file.""" @@ -1027,9 +1032,13 @@ def run_server( _cloudflare_enabled = cloudflare and host == "0.0.0.0" and not api_only and not _IS_COLAB if _cloudflare_enabled: try: # best-effort: any failure must not block startup - from cloudflare_tunnel import start_studio_tunnel + from cloudflare_tunnel import start_studio_tunnel, stop_studio_tunnel + _cloudflare_url = start_studio_tunnel(port) app.state.cloudflare_url = _cloudflare_url + # Backstop: tear the tunnel down even on an abnormal exit that bypasses + # _graceful_shutdown (e.g. an exception after startup -> sys.exit). Idempotent. + atexit.register(stop_studio_tunnel) except Exception as e: logger.debug("Cloudflare tunnel skipped: %s", e) diff --git a/studio/backend/storage/rag_db.py b/studio/backend/storage/rag_db.py index 601fb73d0a..564e3284f8 100644 --- a/studio/backend/storage/rag_db.py +++ b/studio/backend/storage/rag_db.py @@ -57,6 +57,7 @@ def _ensure_schema(conn: sqlite3.Connection) -> None: scope TEXT NOT NULL, kb_id TEXT, thread_id TEXT, + project_id TEXT, filename TEXT NOT NULL, sha256 TEXT NOT NULL, status TEXT NOT NULL DEFAULT 'pending', @@ -102,6 +103,10 @@ def _ensure_schema(conn: sqlite3.Connection) -> None: ); """ ) + # Lazy upgrade for databases created before project sources existed. + cols = {r[1] for r in conn.execute("PRAGMA table_info(documents)").fetchall()} + if "project_id" not in cols: + conn.execute("ALTER TABLE documents ADD COLUMN project_id TEXT") def get_connection() -> sqlite3.Connection: diff --git a/studio/backend/storage/studio_db.py b/studio/backend/storage/studio_db.py index da8d9b5e66..85cfacbc27 100644 --- a/studio/backend/storage/studio_db.py +++ b/studio/backend/storage/studio_db.py @@ -652,7 +652,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.display_name, + r.loss_sparkline, r.display_name, r.config_json, CASE WHEN r.status = 'stopped' AND r.output_dir IS NOT NULL diff --git a/studio/backend/tests/test_cached_gguf_routes.py b/studio/backend/tests/test_cached_gguf_routes.py index 20aa37a2b4..5a4ca68ab4 100644 --- a/studio/backend/tests/test_cached_gguf_routes.py +++ b/studio/backend/tests/test_cached_gguf_routes.py @@ -371,3 +371,156 @@ def test_list_cached_gguf_includes_vision_repo_with_main_gguf_and_mmproj(monkeyp "cache_path": str(vision_repo.repo_path), } ] + + +def _gfile(name: str, size: int, mtime: float) -> SimpleNamespace: + """A cached file carrying a Hugging Face ``blob_last_modified`` timestamp.""" + return SimpleNamespace( + file_name = name, + size_on_disk = size, + blob_path = None, + blob_last_modified = mtime, + ) + + +def test_all_hf_cache_scans_survives_inaccessible_aux_cache(monkeypatch, tmp_path): + """An unreadable auxiliary cache (e.g. an inaccessible + ``~/.cache/huggingface/hub``) must be skipped, not abort the scan. + Regression guard for ``extra.is_dir()`` raising and wiping the response. + """ + import huggingface_hub + import utils.paths as paths_mod + + active = SimpleNamespace( + repos = [_repo("Org/Active", [_file("Q4_K_M.gguf", 5_000)], tmp_path / "active")] + ) + + def _fake_scan(cache_dir = None): + if cache_dir is None: + return active + raise AssertionError("auxiliary scan should have been skipped") + + class _Boom: + def is_dir(self): + raise PermissionError(13, "Permission denied") + + def resolve(self): + raise PermissionError(13, "Permission denied") + + monkeypatch.setattr(huggingface_hub, "scan_cache_dir", _fake_scan) + monkeypatch.setattr(paths_mod, "legacy_hf_cache_dir", lambda: _Boom()) + monkeypatch.setattr(paths_mod, "hf_default_cache_dir", lambda: _Boom()) + + scans = models_route._all_hf_cache_scans() + assert scans == [active] + + # End-to-end: the endpoint still returns the active cache's repo. + monkeypatch.setattr(models_route, "_all_hf_cache_scans", lambda: [active]) + result = asyncio.run(models_route.list_cached_gguf(current_subject = "test-user")) + assert result["cached"] == [ + { + "repo_id": "Org/Active", + "size_bytes": 5_000, + "cache_path": str(tmp_path / "active"), + } + ] + + +def test_list_cached_gguf_sorts_newest_first_grouping_by_latest_quant(monkeypatch, tmp_path): + """Downloaded is ordered newest-first, and a multi-quant repo is placed by + its most recently downloaded quant (``last_modified`` = newest quant).""" + older = _repo( + "Org/Older", + [_gfile("Older-Q4_K_M.gguf", 5_000, 1_000.0)], + tmp_path / "models--Org--Older", + ) + newer = _repo( + "Org/Newer", + [ + _gfile("Newer-Q4_K_M.gguf", 5_000, 2_000.0), + _gfile("Newer-Q8_0.gguf", 9_000, 3_000.0), # newest quant in the repo + ], + tmp_path / "models--Org--Newer", + ) + + monkeypatch.setattr( + models_route, + "_all_hf_cache_scans", + lambda: [SimpleNamespace(repos = [older, newer])], + ) + + result = asyncio.run(models_route.list_cached_gguf(current_subject = "test-user")) + + assert [c["repo_id"] for c in result["cached"]] == ["Org/Newer", "Org/Older"] + assert result["cached"][0]["last_modified"] == 3_000.0 + assert result["cached"][1]["last_modified"] == 1_000.0 + + +def test_list_cached_gguf_dedupe_keeps_newest_timestamp(monkeypatch, tmp_path): + """Same repo in two caches with equal size keeps the newest last_modified, + regardless of scan order.""" + older = _repo("org/dupe", [_gfile("dupe-Q4_K_M.gguf", 5_000, 1_000.0)], tmp_path / "a") + newer = _repo("org/dupe", [_gfile("dupe-Q4_K_M.gguf", 5_000, 9_000.0)], tmp_path / "b") + for scans in ([older, newer], [newer, older]): # both orders + monkeypatch.setattr( + models_route, + "_all_hf_cache_scans", + lambda s = scans: [SimpleNamespace(repos = [s[0]]), SimpleNamespace(repos = [s[1]])], + ) + result = asyncio.run(models_route.list_cached_gguf(current_subject = "t")) + assert len(result["cached"]) == 1 + assert result["cached"][0]["last_modified"] == 9_000.0 + + +def test_gguf_variants_mmproj_does_not_mark_quant_downloaded(monkeypatch, tmp_path): + """The per-quant 'downloaded' flag is driven by the real weight file in a + single snapshot; an mmproj vision adapter (matching a quant label) must + not make that quant appear downloaded.""" + import huggingface_hub.constants as hf_constants + + variants = [ + SimpleNamespace(filename = "model-Q4_K_M.gguf", quant = "Q4_K_M", size_bytes = 10_000), + SimpleNamespace(filename = "model-F16.gguf", quant = "F16", size_bytes = 20_000), + ] + monkeypatch.setattr( + models_route, "list_gguf_variants", lambda repo_id, hf_token = None: (variants, True) + ) + monkeypatch.setattr(hf_constants, "HF_HUB_CACHE", str(tmp_path)) + + snap = tmp_path / "models--org--repo" / "snapshots" / "rev" + snap.mkdir(parents = True) + (snap / "model-Q4_K_M.gguf").write_bytes(b"x" * 10_000) # real weight, fully present + (snap / "mmproj-F16.gguf").write_bytes(b"y" * 20_000) # mmproj adapter, label "F16" + + result = asyncio.run( + models_route.get_gguf_variants( + repo_id = "org/repo", hf_token = None, current_subject = "test-user" + ) + ) + + flags = {v.quant: v.downloaded for v in result.variants} + assert flags["Q4_K_M"] is True + assert flags["F16"] is False + + +def test_gguf_download_progress_excludes_mmproj(monkeypatch, tmp_path): + """A cached mmproj adapter must not count toward a same-label main + variant's download progress (mmproj-F16 vs an F16 weight).""" + import huggingface_hub.constants as hf_constants + + monkeypatch.setattr(hf_constants, "HF_HUB_CACHE", str(tmp_path)) + snap = tmp_path / "models--org--repo" / "snapshots" / "rev" + snap.mkdir(parents = True) + (snap / "mmproj-F16.gguf").write_bytes(b"y" * 20_000) # only the adapter on disk + + result = asyncio.run( + models_route.get_gguf_download_progress( + repo_id = "org/repo", + variant = "F16", + expected_bytes = 20_000, + current_subject = "test-user", + ) + ) + + assert result["downloaded_bytes"] == 0 + assert result["progress"] == 0 diff --git a/studio/backend/tests/test_cloudflare_tunnel.py b/studio/backend/tests/test_cloudflare_tunnel.py index 873547631d..8042240b64 100644 --- a/studio/backend/tests/test_cloudflare_tunnel.py +++ b/studio/backend/tests/test_cloudflare_tunnel.py @@ -52,6 +52,26 @@ def test_url_regex_no_match_on_unrelated(): assert ct._URL_RE.search("INF connecting to https://api.cloudflare.com/v4") is None +def test_url_regex_ignores_api_endpoint(): + # cloudflared's failure line names its own API host; it must never be taken + # as the tunnel URL (it returns a 404 and is not a quick tunnel). + line = ( + 'failed to request quick Tunnel: Post "https://api.trycloudflare.com/tunnel": ' + "context deadline exceeded" + ) + assert ct._URL_RE.search(line) is None + + +def test_url_regex_skips_api_host_but_matches_real_url(): + blob = ( + 'ERR failed to request quick Tunnel: Post "https://api.trycloudflare.com/tunnel"\n' + "INF | https://brave-mountain-river-clouds.trycloudflare.com |\n" + ) + m = ct._URL_RE.search(blob) + assert m is not None + assert m.group(0) == "https://brave-mountain-river-clouds.trycloudflare.com" + + # ── asset mapping ──────────────────────────────────────────────────── @@ -295,9 +315,88 @@ def test_stop_terminates_process(): t.stop() -def test_wait_for_url_times_out_without_blocking(): +def test_start_after_stop_does_not_spawn(monkeypatch): + # If stop() lands before start() (a concurrent shutdown in the caller's + # register->start window), start() must NOT spawn a cloudflared process -- + # nobody would own it and it would be orphaned. t = ct.CloudflareTunnel(8080, "/bin/cloudflared") - assert t.wait_for_url(timeout = 0.05) is None + spawned = [] + + class _FakeProc: + stdout = None + + def poll(self): + return 0 + + monkeypatch.setattr(ct.subprocess, "Popen", lambda *a, **k: (spawned.append(a), _FakeProc())[1]) + t.stop() # proc is None -> no-op terminate, but marks the tunnel stopped + t.start() # must short-circuit before Popen + assert spawned == [] + assert t._proc is None + + +def test_wait_for_ready_times_out_without_blocking(): + t = ct.CloudflareTunnel(8080, "/bin/cloudflared") + assert t.wait_for_ready(timeout = 0.05) is None + + +def _fake_proc(text): + return types.SimpleNamespace(stdout = io.StringIO(text)) + + +def test_reader_captures_url_and_registration(): + t = ct.CloudflareTunnel(8080, "/bin/cloudflared") + t._reader( + _fake_proc( + "INF Requesting new quick Tunnel on trycloudflare.com...\n" + "INF | https://words-here-abc.trycloudflare.com |\n" + "INF Registered tunnel connection connIndex=0 protocol=http2\n" + ) + ) + assert t.url == "https://words-here-abc.trycloudflare.com" + assert t.ready is True + assert t.wait_for_ready(0) == t.url + assert t.error is None # a fully-registered tunnel records no error + + +def test_reader_url_without_registration_is_not_ready(): + # A URL but no "Registered tunnel connection" (e.g. quic control stream + # fails) must not be advertised -- it returns Cloudflare error 1033. + t = ct.CloudflareTunnel(8080, "/bin/cloudflared") + t._reader( + _fake_proc( + "INF | https://words-here-abc.trycloudflare.com |\n" + 'ERR failed to serve tunnel connection error="control stream failure"\n' + ) + ) + assert t.url == "https://words-here-abc.trycloudflare.com" + assert t.ready is False + assert t.wait_for_ready(0) is None + assert t.error == "cloudflared exited before the tunnel connection registered" + + +def test_reader_handles_none_stdout(): + # Popen.stdout can be None; _reader must not crash and must leave the tunnel + # un-ready so wait_for_ready returns None. + t = ct.CloudflareTunnel(8080, "/bin/cloudflared") + t._reader(types.SimpleNamespace(stdout = None)) + assert t.url is None + assert t.ready is False + assert t.wait_for_ready(0) is None + assert t.error == "cloudflared exited before emitting a tunnel URL" + + +def test_reader_ignores_api_endpoint_failure_line(): + t = ct.CloudflareTunnel(8080, "/bin/cloudflared") + t._reader( + _fake_proc( + "ERR failed to request quick Tunnel: Post " + '"https://api.trycloudflare.com/tunnel": context deadline exceeded\n' + ) + ) + assert t.url is None + assert t.wait_for_ready(0) is None + assert t.error == "cloudflared exited before emitting a tunnel URL" def test_start_studio_tunnel_no_binary(monkeypatch): @@ -306,18 +405,23 @@ def test_start_studio_tunnel_no_binary(monkeypatch): def test_start_studio_tunnel_registers_before_wait(monkeypatch): - # The tunnel must be visible to stop_studio_tunnel() during the URL wait, - # else a shutdown in that window orphans cloudflared. + # The tunnel must be visible to stop_studio_tunnel() during the readiness + # wait, else a shutdown in that window orphans cloudflared. seen = {} class _Stub: - def __init__(self, port, binary): + def __init__( + self, + port, + binary, + protocol = None, + ): self.url = None def start(self): pass - def wait_for_url(self, timeout): + def wait_for_ready(self, timeout): seen["active_during_wait"] = ct._active_tunnel is self self.url = "https://x.trycloudflare.com" return self.url @@ -338,13 +442,18 @@ def test_start_studio_tunnel_clears_and_stops_on_no_url(monkeypatch): seen = {} class _Stub: - def __init__(self, port, binary): + def __init__( + self, + port, + binary, + protocol = None, + ): self.url = None def start(self): pass - def wait_for_url(self, timeout): + def wait_for_ready(self, timeout): return None def stop(self): @@ -359,13 +468,18 @@ def test_start_studio_tunnel_clears_and_stops_on_no_url(monkeypatch): def test_start_studio_tunnel_returns_url(monkeypatch): class _StubTunnel: - def __init__(self, port, binary): + def __init__( + self, + port, + binary, + protocol = None, + ): self.url = None def start(self): self.url = "https://stub-xyz.trycloudflare.com" - def wait_for_url(self, timeout): + def wait_for_ready(self, timeout): return self.url def stop(self): @@ -379,6 +493,169 @@ def test_start_studio_tunnel_returns_url(monkeypatch): ct.stop_studio_tunnel() +def test_start_studio_tunnel_falls_back_to_http2(monkeypatch): + # First attempt mints a URL but never registers (quic blocked); the http2 + # retry registers and wins. + attempts = [] + + class _Stub: + def __init__( + self, + port, + binary, + protocol = None, + ): + self.protocol = protocol + self.url = None + attempts.append(protocol) + + def start(self): + self.url = "https://words.trycloudflare.com" # URL always minted + + def wait_for_ready(self, timeout): + return self.url if self.protocol == "http2" else None + + def stop(self): + pass + + monkeypatch.setattr(ct, "ensure_cloudflared", lambda: "/bin/cloudflared") + monkeypatch.setattr(ct, "CloudflareTunnel", _Stub) + try: + assert ct.start_studio_tunnel(8080) == "https://words.trycloudflare.com" + assert attempts == [None, "http2"] # default first, then forced http2 + finally: + ct.stop_studio_tunnel() + + +def test_start_studio_tunnel_no_retry_when_shutdown_between_attempts(monkeypatch): + # A stop() landing in the gap AFTER the failed first attempt is cleaned up but + # BEFORE the http2 retry registers must abort the loop -- not start a second + # tunnel that nobody will ever stop (Codex review). Simulated by having the + # first attempt's stop() (called during cleanup) trigger the shutdown. + attempts = [] + + class _Stub: + def __init__( + self, + port, + binary, + protocol = None, + ): + self.url = None + attempts.append(protocol) + + def start(self): + self.url = "https://words.trycloudflare.com" # URL minted, never ready + + def wait_for_ready(self, timeout): + return None + + def stop(self): + ct.stop_studio_tunnel() # a concurrent shutdown lands in the gap + + monkeypatch.setattr(ct, "ensure_cloudflared", lambda: "/bin/cloudflared") + monkeypatch.setattr(ct, "CloudflareTunnel", _Stub) + assert ct.start_studio_tunnel(8080) is None + assert attempts == [None] # http2 retry aborted after shutdown + assert ct._active_tunnel is None + + +def test_start_studio_tunnel_no_http2_retry_when_no_url(monkeypatch): + # No URL at all is an API/network failure; the http2 fallback would not help, + # so it must be skipped (don't burn a second timeout window). + attempts = [] + + class _Stub: + def __init__( + self, + port, + binary, + protocol = None, + ): + self.url = None + attempts.append(protocol) + + def start(self): + pass # never mints a URL + + def wait_for_ready(self, timeout): + return None + + def stop(self): + pass + + monkeypatch.setattr(ct, "ensure_cloudflared", lambda: "/bin/cloudflared") + monkeypatch.setattr(ct, "CloudflareTunnel", _Stub) + assert ct.start_studio_tunnel(8080) is None + assert attempts == [None] + + +def test_start_studio_tunnel_both_protocols_fail_registration(monkeypatch): + # Both quic and http2 mint a URL but neither registers -> both attempts are + # exhausted and None is returned (no dead URL advertised). + attempts = [] + + class _Stub: + def __init__( + self, + port, + binary, + protocol = None, + ): + self.url = None + attempts.append(protocol) + + def start(self): + self.url = "https://words.trycloudflare.com" # URL minted, never ready + + def wait_for_ready(self, timeout): + return None + + def stop(self): + pass + + monkeypatch.setattr(ct, "ensure_cloudflared", lambda: "/bin/cloudflared") + monkeypatch.setattr(ct, "CloudflareTunnel", _Stub) + assert ct.start_studio_tunnel(8080) is None + assert attempts == [None, "http2"] + assert ct._active_tunnel is None + + +def test_start_studio_tunnel_aborts_retry_on_concurrent_shutdown(monkeypatch): + # If a concurrent stop_studio_tunnel() clears _active_tunnel while we wait, + # the retry loop must NOT start a second (http2) tunnel: shutdown is already + # done, so nothing would ever stop it and it would be orphaned. + attempts = [] + + class _Stub: + def __init__( + self, + port, + binary, + protocol = None, + ): + self.url = None + attempts.append(protocol) + + def start(self): + self.url = "https://words.trycloudflare.com" # URL minted (saw_url True) + + def wait_for_ready(self, timeout): + # Simulate stop_studio_tunnel() landing during the wait. + with ct._active_lock: + ct._active_tunnel = None + return None # never registered + + def stop(self): + pass + + monkeypatch.setattr(ct, "ensure_cloudflared", lambda: "/bin/cloudflared") + monkeypatch.setattr(ct, "CloudflareTunnel", _Stub) + assert ct.start_studio_tunnel(8080) is None + assert attempts == [None] # no http2 retry -> no orphaned second tunnel + assert ct._active_tunnel is None + + # ── run.py source-level pins (AST / source, no heavy import) ───────── @@ -419,6 +696,13 @@ def test_argparse_cloudflare_default_true(): assert _argparse_default(_RUN_PY.read_text(), "--cloudflare") is True +def test_run_server_registers_tunnel_atexit_backstop(): + # An abnormal exit (exception after startup -> sys.exit) bypasses + # _graceful_shutdown; an atexit backstop must still stop the tunnel. + src = _RUN_PY.read_text() + assert "atexit.register(stop_studio_tunnel)" in src + + def test_run_server_gates_tunnel_on_wildcard(): # Guard against accidentally widening the trigger beyond 0.0.0.0. source = _RUN_PY.read_text() diff --git a/studio/backend/tests/test_datacenter_gpu_tuning.py b/studio/backend/tests/test_datacenter_gpu_tuning.py new file mode 100644 index 0000000000..fd9b291e8a --- /dev/null +++ b/studio/backend/tests/test_datacenter_gpu_tuning.py @@ -0,0 +1,278 @@ +# SPDX-License-Identifier: AGPL-3.0-only +# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0 + +"""Data-center llama.cpp env tuning: FP32 accum (+ P2P / launch queues for +multi-GPU) must apply only to datacenter NVIDIA parts, never consumer GeForce, +AMD/ROCm, CPU or macOS. User values win; UNSLOTH_DISABLE_DC_TUNING=1 disables. +""" + +from __future__ import annotations + +import sys +import types + +import pytest + +from core.inference.llama_cpp import LlamaCppBackend + + +def _fake_torch( + names, + *, + hip = None, + cuda_ok = True, +): + """torch stub: version.hip, cuda.is_available/device_count, get_device_properties(i).name.""" + t = types.ModuleType("torch") + t.version = types.SimpleNamespace(hip = hip) + t.cuda = types.SimpleNamespace( + is_available = lambda: cuda_ok, + device_count = lambda: len(names), + get_device_properties = lambda i: types.SimpleNamespace(name = names[i]), + ) + return t + + +@pytest.fixture(autouse = True) +def _clear_cuda_visible_devices(monkeypatch): + """Detection reads CUDA_VISIBLE_DEVICES, so clear it by default (run unmasked, + physical id == ordinal) regardless of host; masked tests set it explicitly.""" + monkeypatch.delenv("CUDA_VISIBLE_DEVICES", raising = False) + + +# --------------------------------------------------------------------------- +# _is_datacenter_gpu +# --------------------------------------------------------------------------- + + +@pytest.mark.parametrize( + "names,expected", + [ + # Datacenter / professional parts. + (["NVIDIA A100-SXM4-80GB"], True), + (["NVIDIA A30"], True), + (["NVIDIA H100 80GB HBM3"], True), + (["NVIDIA H200"], True), + (["NVIDIA H800"], True), + (["NVIDIA GH200 480GB"], True), + (["NVIDIA B200"], True), + (["NVIDIA GB200"], True), + (["NVIDIA L40S"], True), + (["NVIDIA L4"], True), + (["NVIDIA RTX PRO 6000 Blackwell Server Edition"], True), + (["NVIDIA RTX 6000 Ada Generation"], True), + # Consumer GeForce: never. + (["NVIDIA GeForce RTX 4090"], False), + (["NVIDIA GeForce RTX 5090"], False), + (["NVIDIA GeForce RTX 3090"], False), + (["NVIDIA GeForce RTX 2080 Ti"], False), + (["NVIDIA GeForce GTX 1080"], False), + # Workstation/laptop: short markers must not match as substrings + # ("a100" in "A1000", "a30" in "A3000"). + (["NVIDIA RTX A1000 Laptop GPU"], False), + (["NVIDIA RTX A1000 6GB Laptop GPU"], False), + (["NVIDIA RTX A3000 Laptop GPU"], False), + # Homogeneous multi-DC: all must match. + (["NVIDIA B200", "NVIDIA B200"], True), + (["NVIDIA H100 80GB HBM3", "NVIDIA H100 80GB HBM3"], True), + # Mixed DC + consumer: non-DC, so tuning never lands on the GeForce. + (["NVIDIA B200", "NVIDIA GeForce RTX 4090"], False), + (["NVIDIA GeForce RTX 4090", "NVIDIA B200"], False), + ], +) +def test_is_datacenter_gpu(monkeypatch, names, expected): + monkeypatch.setitem(sys.modules, "torch", _fake_torch(names)) + assert LlamaCppBackend._is_datacenter_gpu() is expected + + +def test_is_datacenter_gpu_respects_selection(monkeypatch): + # A mixed box where only the DC GPU is selected -> True; only consumer -> False. + monkeypatch.setitem( + sys.modules, + "torch", + _fake_torch(["NVIDIA B200", "NVIDIA GeForce RTX 4090"]), + ) + assert LlamaCppBackend._is_datacenter_gpu([0]) is True + assert LlamaCppBackend._is_datacenter_gpu([1]) is False + assert LlamaCppBackend._is_datacenter_gpu([0, 1]) is False + + +def test_is_datacenter_gpu_out_of_range_indices_skipped(monkeypatch): + monkeypatch.setitem(sys.modules, "torch", _fake_torch(["NVIDIA B200"])) + # Out-of-range / negative indices are skipped; the one valid DC GPU still wins. + assert LlamaCppBackend._is_datacenter_gpu([0, 5, -1]) is True + # Only invalid indices -> nothing seen -> False (fail closed for the flag). + assert LlamaCppBackend._is_datacenter_gpu([5, 9]) is False + + +def test_is_datacenter_gpu_masked_host_physical_ids(monkeypatch): + # Mask 4,5,6,7 -> ordinals 0..3 == physical 4..7. PHYSICAL selection [4,5] + # must resolve, not index out of range (the pre-fix bug: 4 >= device_count). + monkeypatch.setenv("CUDA_VISIBLE_DEVICES", "4,5,6,7") + monkeypatch.setitem(sys.modules, "torch", _fake_torch(["NVIDIA B200"] * 4)) + assert LlamaCppBackend._is_datacenter_gpu([4, 5]) is True + assert LlamaCppBackend._is_datacenter_gpu([4, 5, 6, 7]) is True + assert LlamaCppBackend._is_datacenter_gpu(None) is True + assert LlamaCppBackend._is_datacenter_gpu([0, 1]) is False # not visible -> skip + + +def test_is_datacenter_gpu_masked_host_reordered(monkeypatch): + # Reordered mask preserves order: ordinal 0 -> physical 7, 1 -> 4, ... + monkeypatch.setenv("CUDA_VISIBLE_DEVICES", "7,4,5,6") + monkeypatch.setitem(sys.modules, "torch", _fake_torch(["NVIDIA H100 80GB HBM3"] * 4)) + assert LlamaCppBackend._is_datacenter_gpu([7, 4]) is True + + +def test_is_datacenter_gpu_masked_host_mixed_class(monkeypatch): + # Mask 4,5: physical 4 = GeForce, physical 5 = B200. Detection must follow the + # selected physical GPU, not a same-numbered ordinal. + monkeypatch.setenv("CUDA_VISIBLE_DEVICES", "4,5") + monkeypatch.setitem( + sys.modules, + "torch", + _fake_torch(["NVIDIA GeForce RTX 4090", "NVIDIA B200"]), + ) + assert LlamaCppBackend._is_datacenter_gpu([4]) is False + assert LlamaCppBackend._is_datacenter_gpu([5]) is True + assert LlamaCppBackend._is_datacenter_gpu([4, 5]) is False + + +def test_is_datacenter_gpu_unparsable_mask_falls_back(monkeypatch): + # Unparsable (UUID) mask falls back to physical id == ordinal (mirrors + # _get_gpu_free_memory), so ordinal lookup still classifies the device. + monkeypatch.setenv("CUDA_VISIBLE_DEVICES", "GPU-abcdef12") + monkeypatch.setitem(sys.modules, "torch", _fake_torch(["NVIDIA B200"])) + assert LlamaCppBackend._is_datacenter_gpu([0]) is True + + +def test_is_datacenter_gpu_rocm_is_false(monkeypatch): + # ROCm reuses torch.cuda.*; an MI300X must not qualify. + monkeypatch.setitem( + sys.modules, + "torch", + _fake_torch(["AMD Instinct MI300X"], hip = "6.2.0"), + ) + assert LlamaCppBackend._is_datacenter_gpu() is False + + +def test_is_datacenter_gpu_no_cuda_is_false(monkeypatch): + monkeypatch.setitem(sys.modules, "torch", _fake_torch([], cuda_ok = False)) + assert LlamaCppBackend._is_datacenter_gpu() is False + + +def test_is_datacenter_gpu_missing_torch_is_false(monkeypatch): + monkeypatch.setitem(sys.modules, "torch", None) + assert LlamaCppBackend._is_datacenter_gpu() is False + + +# --------------------------------------------------------------------------- +# _effective_gpu_count +# --------------------------------------------------------------------------- + + +def test_effective_gpu_count_explicit_selection(monkeypatch): + monkeypatch.setitem(sys.modules, "torch", _fake_torch(["NVIDIA B200"] * 4)) + assert LlamaCppBackend._effective_gpu_count([0]) == 1 + assert LlamaCppBackend._effective_gpu_count([0, 1, 2]) == 3 + + +def test_effective_gpu_count_none_uses_visible(monkeypatch): + # None -> visible device count. + monkeypatch.setitem(sys.modules, "torch", _fake_torch(["NVIDIA B200"] * 4)) + assert LlamaCppBackend._effective_gpu_count(None) == 4 + + +def test_effective_gpu_count_no_cuda_is_zero(monkeypatch): + monkeypatch.setitem(sys.modules, "torch", _fake_torch([], cuda_ok = False)) + assert LlamaCppBackend._effective_gpu_count(None) == 0 + + +def test_effective_gpu_count_missing_torch_is_zero(monkeypatch): + monkeypatch.setitem(sys.modules, "torch", None) + assert LlamaCppBackend._effective_gpu_count(None) == 0 + + +# --------------------------------------------------------------------------- +# _apply_datacenter_env (the env-injection decision) +# --------------------------------------------------------------------------- + + +def test_apply_env_single_dc_gpu_sets_only_fp32(monkeypatch): + monkeypatch.delenv("UNSLOTH_DISABLE_DC_TUNING", raising = False) + monkeypatch.setitem(sys.modules, "torch", _fake_torch(["NVIDIA B200"])) + env: dict = {} + assert LlamaCppBackend._apply_datacenter_env(env, [0]) is True + assert env == {"GGML_CUDA_FORCE_CUBLAS_COMPUTE_32F": "1"} + assert "GGML_CUDA_P2P" not in env # no multi-GPU flags on one GPU + assert "CUDA_SCALE_LAUNCH_QUEUES" not in env + + +def test_apply_env_multi_dc_gpu_sets_all(monkeypatch): + monkeypatch.delenv("UNSLOTH_DISABLE_DC_TUNING", raising = False) + monkeypatch.setitem(sys.modules, "torch", _fake_torch(["NVIDIA B200"] * 4)) + env: dict = {} + assert LlamaCppBackend._apply_datacenter_env(env, [0, 1]) is True + assert env["GGML_CUDA_FORCE_CUBLAS_COMPUTE_32F"] == "1" + assert env["GGML_CUDA_P2P"] == "1" + assert env["CUDA_SCALE_LAUNCH_QUEUES"] == "4x" + + +def test_apply_env_none_indices_uses_visible_count(monkeypatch): + # None on a 2x DC box -> multi-GPU flags applied. + monkeypatch.delenv("UNSLOTH_DISABLE_DC_TUNING", raising = False) + monkeypatch.setitem(sys.modules, "torch", _fake_torch(["NVIDIA H100", "NVIDIA H100"])) + env: dict = {} + assert LlamaCppBackend._apply_datacenter_env(env, None) is True + assert env["GGML_CUDA_P2P"] == "1" + assert env["CUDA_SCALE_LAUNCH_QUEUES"] == "4x" + + +def test_apply_env_consumer_gpu_is_noop(monkeypatch): + monkeypatch.delenv("UNSLOTH_DISABLE_DC_TUNING", raising = False) + monkeypatch.setitem(sys.modules, "torch", _fake_torch(["NVIDIA GeForce RTX 4090"] * 2)) + env: dict = {} + assert LlamaCppBackend._apply_datacenter_env(env, [0, 1]) is False + assert env == {} + + +def test_apply_env_user_value_wins(monkeypatch): + monkeypatch.delenv("UNSLOTH_DISABLE_DC_TUNING", raising = False) + monkeypatch.setitem(sys.modules, "torch", _fake_torch(["NVIDIA B200"] * 2)) + env = { + "GGML_CUDA_FORCE_CUBLAS_COMPUTE_32F": "0", # user explicitly disabled + "CUDA_SCALE_LAUNCH_QUEUES": "8x", # user override + } + assert LlamaCppBackend._apply_datacenter_env(env, [0, 1]) is True + # setdefault must not clobber user values; the unset one still defaults. + assert env["GGML_CUDA_FORCE_CUBLAS_COMPUTE_32F"] == "0" + assert env["CUDA_SCALE_LAUNCH_QUEUES"] == "8x" + assert env["GGML_CUDA_P2P"] == "1" + + +def test_apply_env_disable_flag_respected(monkeypatch): + monkeypatch.setenv("UNSLOTH_DISABLE_DC_TUNING", "1") + monkeypatch.setitem(sys.modules, "torch", _fake_torch(["NVIDIA B200"] * 2)) + env: dict = {} + assert LlamaCppBackend._apply_datacenter_env(env, [0, 1]) is False + assert env == {} + + +def test_apply_env_fail_open_on_detection_error(monkeypatch): + monkeypatch.delenv("UNSLOTH_DISABLE_DC_TUNING", raising = False) + monkeypatch.setitem(sys.modules, "torch", None) # detection raises -> False + env: dict = {} + assert LlamaCppBackend._apply_datacenter_env(env, [0]) is False + assert env == {} + + +def test_apply_env_masked_host_multi_dc(monkeypatch): + # End-to-end masked host (mask 4,5,6,7, physical selection [4,5]): pre-fix + # applied no tuning; now all three multi-GPU flags must be set. + monkeypatch.delenv("UNSLOTH_DISABLE_DC_TUNING", raising = False) + monkeypatch.setenv("CUDA_VISIBLE_DEVICES", "4,5,6,7") + monkeypatch.setitem(sys.modules, "torch", _fake_torch(["NVIDIA B200"] * 4)) + env: dict = {} + assert LlamaCppBackend._apply_datacenter_env(env, [4, 5]) is True + assert env["GGML_CUDA_FORCE_CUBLAS_COMPUTE_32F"] == "1" + assert env["GGML_CUDA_P2P"] == "1" + assert env["CUDA_SCALE_LAUNCH_QUEUES"] == "4x" diff --git a/studio/backend/tests/test_export_absolute_paths.py b/studio/backend/tests/test_export_absolute_paths.py new file mode 100644 index 0000000000..f333a590c2 --- /dev/null +++ b/studio/backend/tests/test_export_absolute_paths.py @@ -0,0 +1,472 @@ +# 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 importlib.machinery +import importlib.util +import sys +import types +from pathlib import Path + +import pytest + + +_BACKEND_DIR = Path(__file__).resolve().parent.parent + + +def _load_module( + module_name: str, + relative_path: str, + monkeypatch = None, +): + path = _BACKEND_DIR / relative_path + spec = importlib.util.spec_from_file_location(module_name, path) + assert spec is not None + assert spec.loader is not None + module = importlib.util.module_from_spec(spec) + if monkeypatch is None: + sys.modules[module_name] = module + else: + monkeypatch.setitem(sys.modules, module_name, module) + spec.loader.exec_module(module) + return module + + +class _DummyLogger: + def __getattr__(self, _name): + return lambda *args, **kwargs: None + + +class _Router: + def get(self, *args, **kwargs): + return lambda fn: fn + + def post(self, *args, **kwargs): + return lambda fn: fn + + def delete(self, *args, **kwargs): + return lambda fn: fn + + +class _HTTPException(Exception): + def __init__( + self, + status_code: int, + detail: str | None = None, + ): + super().__init__(detail) + self.status_code = status_code + self.detail = detail + + +class _LocalModelInfo: + def __init__(self, **kwargs): + self.__dict__.update(kwargs) + + +def _identity_decorator(*_args, **_kwargs): + return lambda fn: fn + + +def _install_lightweight_backend_stubs(monkeypatch): + fastapi = types.ModuleType("fastapi") + fastapi.APIRouter = lambda: _Router() + fastapi.Body = lambda default = None, **_kwargs: default + fastapi.Depends = lambda dependency = None, **_kwargs: dependency + fastapi.HTTPException = _HTTPException + fastapi.Query = lambda default = None, **_kwargs: default + fastapi.Request = object + monkeypatch.setitem(sys.modules, "fastapi", fastapi) + + fastapi_responses = types.ModuleType("fastapi.responses") + fastapi_responses.StreamingResponse = object + monkeypatch.setitem(sys.modules, "fastapi.responses", fastapi_responses) + + monkeypatch.setitem( + sys.modules, + "structlog", + types.SimpleNamespace( + BoundLogger = _DummyLogger, + get_logger = lambda *args, **kwargs: _DummyLogger(), + ), + ) + loggers = types.ModuleType("loggers") + loggers.get_logger = lambda *args, **kwargs: _DummyLogger() + monkeypatch.setitem(sys.modules, "loggers", loggers) + + auth_pkg = types.ModuleType("auth") + auth_mod = types.ModuleType("auth.authentication") + auth_mod.get_current_subject = lambda: None + monkeypatch.setitem(sys.modules, "auth", auth_pkg) + monkeypatch.setitem(sys.modules, "auth.authentication", auth_mod) + + core_pkg = types.ModuleType("core") + core_export = types.ModuleType("core.export") + core_export.get_export_backend = lambda: None + core_inference = types.ModuleType("core.inference") + core_inference.get_inference_backend = lambda: None + monkeypatch.setitem(sys.modules, "core", core_pkg) + monkeypatch.setitem(sys.modules, "core.export", core_export) + monkeypatch.setitem(sys.modules, "core.inference", core_inference) + + utils_pkg = types.ModuleType("utils") + utils_pkg.__path__ = [] + utils_paths = types.ModuleType("utils.paths") + storage_roots = _load_module( + "utils.paths.storage_roots", + "utils/paths/storage_roots.py", + monkeypatch, + ) + utils_pkg.paths = utils_paths + utils_paths.storage_roots = storage_roots + utils_paths.is_local_path = lambda value: Path(str(value)).is_absolute() + utils_paths.outputs_root = lambda: Path("outputs") + utils_paths.exports_root = storage_roots.exports_root + utils_paths.resolve_cached_repo_id_case = lambda value: value + utils_paths.resolve_output_dir = lambda value = None: Path(value or "outputs") + utils_paths.resolve_export_dir = storage_roots.resolve_export_dir + monkeypatch.setitem(sys.modules, "utils", utils_pkg) + monkeypatch.setitem(sys.modules, "utils.paths", utils_paths) + + utils_utils = types.ModuleType("utils.utils") + utils_utils.log_and_http_error = lambda *args, **kwargs: (_ for _ in ()).throw( + _HTTPException(kwargs.get("status_code", 500), kwargs.get("detail")) + ) + utils_utils.safe_error_detail = lambda value: str(value) + monkeypatch.setitem(sys.modules, "utils.utils", utils_utils) + + utils_models = types.ModuleType("utils.models") + for name in ( + "scan_trained_models", + "scan_exported_models", + "scan_checkpoints", + "list_gguf_variants", + ): + setattr(utils_models, name, lambda *args, **kwargs: []) + for name in ( + "get_base_model_from_checkpoint", + "get_base_model_from_lora", + "load_model_defaults", + ): + setattr(utils_models, name, lambda *args, **kwargs: None) + utils_models.is_vision_model = lambda *args, **kwargs: False + utils_models.is_embedding_model = lambda *args, **kwargs: False + utils_models.ModelConfig = object + monkeypatch.setitem(sys.modules, "utils.models", utils_models) + + utils_model_config = types.ModuleType("utils.models.model_config") + utils_model_config._pick_best_gguf = lambda variants: variants[0] if variants else None + utils_model_config._extract_quant_label = lambda value: value + utils_model_config.is_audio_input_type = lambda *args, **kwargs: None + monkeypatch.setitem( + sys.modules, + "utils.models.model_config", + utils_model_config, + ) + + models_pkg = types.ModuleType("models") + models_pkg.__path__ = [] + for name in ( + "CheckpointInfo", + "CheckpointListResponse", + "LocalModelListResponse", + "ModelCheckpoints", + "ModelDetails", + "LoRAScanResponse", + "LoRAInfo", + "ModelListResponse", + "LoadCheckpointRequest", + "ExportStatusResponse", + "ExportOperationResponse", + "ExportMergedModelRequest", + "ExportBaseModelRequest", + "ExportGGUFRequest", + "ExportLoRAAdapterRequest", + ): + setattr(models_pkg, name, object) + models_pkg.LocalModelInfo = _LocalModelInfo + monkeypatch.setitem(sys.modules, "models", models_pkg) + + models_models = types.ModuleType("models.models") + for name in ( + "BrowseEntry", + "BrowseFoldersResponse", + "GgufVariantDetail", + "GgufVariantsResponse", + "ScanFolderInfo", + "AddScanFolderRequest", + ): + setattr(models_models, name, object) + models_models.ModelType = str + monkeypatch.setitem(sys.modules, "models.models", models_models) + + models_responses = types.ModuleType("models.responses") + for name in ( + "LoRABaseModelResponse", + "VisionCheckResponse", + "EmbeddingCheckResponse", + ): + setattr(models_responses, name, object) + monkeypatch.setitem(sys.modules, "models.responses", models_responses) + + +def _install_pydantic_stub(monkeypatch): + pydantic = types.ModuleType("pydantic") + pydantic.BaseModel = object + pydantic.Field = lambda default = None, **_kwargs: default + pydantic.field_validator = _identity_decorator + monkeypatch.setitem(sys.modules, "pydantic", pydantic) + + +def _install_export_backend_stubs(monkeypatch): + _install_lightweight_backend_stubs(monkeypatch) + + unsloth = types.ModuleType("unsloth") + unsloth.FastLanguageModel = object + unsloth.FastVisionModel = object + unsloth._IS_MLX = True + unsloth.__spec__ = importlib.machinery.ModuleSpec("unsloth", loader = None) + monkeypatch.setitem(sys.modules, "unsloth", unsloth) + + unsloth_zoo = types.ModuleType("unsloth_zoo") + unsloth_zoo.__path__ = [] + unsloth_zoo.__spec__ = importlib.machinery.ModuleSpec( + "unsloth_zoo", + loader = None, + is_package = True, + ) + llama_cpp = types.ModuleType("unsloth_zoo.llama_cpp") + llama_cpp.LLAMA_CPP_DEFAULT_DIR = str(Path("/tmp/llama.cpp")) + llama_cpp._resolve_local_convert_script = lambda *args, **kwargs: None + llama_cpp.__spec__ = importlib.machinery.ModuleSpec( + "unsloth_zoo.llama_cpp", + loader = None, + ) + monkeypatch.setitem(sys.modules, "unsloth_zoo", unsloth_zoo) + monkeypatch.setitem(sys.modules, "unsloth_zoo.llama_cpp", llama_cpp) + + huggingface_hub = types.ModuleType("huggingface_hub") + huggingface_hub.HfApi = object + huggingface_hub.ModelCard = object + monkeypatch.setitem(sys.modules, "huggingface_hub", huggingface_hub) + + utils_hardware = types.ModuleType("utils.hardware") + utils_hardware.clear_gpu_cache = lambda: None + monkeypatch.setitem(sys.modules, "utils.hardware", utils_hardware) + + utils_models = sys.modules["utils.models"] + utils_models.get_base_model_from_lora = lambda *args, **kwargs: None + utils_models.is_vision_model = lambda *args, **kwargs: False + + utils_model_config = sys.modules["utils.models.model_config"] + utils_model_config.detect_audio_type = lambda *args, **kwargs: None + + utils_paths = sys.modules["utils.paths"] + utils_paths.ensure_dir = lambda path: Path(path).mkdir(parents = True, exist_ok = True) + utils_paths.resolve_export_write_dir = lambda value = None: Path(value or "exports") + utils_paths.resolve_output_dir = lambda value = None: Path(value or "outputs") + + +def test_gguf_export_cleans_temp_dir_when_post_processing_fails(tmp_path, monkeypatch): + _install_export_backend_stubs(monkeypatch) + export_mod = _load_module("test_core_export_backend", "core/export/export.py", monkeypatch) + + cwd = tmp_path / "cwd" + save_dir = tmp_path / "export" + cwd.mkdir() + monkeypatch.chdir(cwd) + monkeypatch.setattr(export_mod, "resolve_export_write_dir", lambda _value: save_dir) + monkeypatch.setattr( + export_mod.shutil, + "move", + lambda *args, **kwargs: (_ for _ in ()).throw(OSError("move failed")), + ) + + class _Model: + def save_pretrained_gguf(self, model_save_path, tokenizer, quantization_method): + Path(model_save_path).mkdir(parents = True) + (Path(model_save_path) / "model.safetensors").write_bytes(b"weights") + (cwd / "converted.gguf").write_bytes(b"gguf") + + backend = export_mod.ExportBackend.__new__(export_mod.ExportBackend) + backend.current_model = _Model() + backend.current_tokenizer = object() + backend.current_checkpoint = None + + success, message, output_path = backend.export_gguf(str(save_dir), "Q4_K_M") + + assert success is False + assert "move failed" in message + assert output_path is None + assert list(save_dir.glob("_tmp_model_*")) == [] + + +def test_save_directory_validator_rejects_windows_parent_segments(monkeypatch): + _install_pydantic_stub(monkeypatch) + export_models = _load_module("test_models_export", "models/export.py", monkeypatch) + + with pytest.raises(ValueError, match = r"\.\."): + export_models._validate_save_directory(r"E:\AI\..\secret") + + +def test_save_directory_validator_allows_deep_absolute_paths(monkeypatch, tmp_path): + _install_pydantic_stub(monkeypatch) + export_models = _load_module("test_models_export_deep_path", "models/export.py", monkeypatch) + + deep_path = tmp_path + for index in range(40): + deep_path /= f"segment-{index:02d}" + raw = str(deep_path) + + assert len(raw) > 255 + assert export_models._validate_save_directory(raw) == raw + + +def test_save_directory_validator_rejects_long_path_component(monkeypatch, tmp_path): + _install_pydantic_stub(monkeypatch) + export_models = _load_module( + "test_models_export_long_component", "models/export.py", monkeypatch + ) + + with pytest.raises(ValueError, match = "path components"): + export_models._validate_save_directory(str(tmp_path / ("a" * 256))) + + +def test_export_write_dir_accepts_external_absolute_but_read_dir_rejects(tmp_path, monkeypatch): + storage_roots = _load_module( + "test_storage_roots_accept_external", + "utils/paths/storage_roots.py", + ) + + export_root = tmp_path / "exports" + external = tmp_path / "external" + export_root.mkdir() + external.mkdir() + monkeypatch.setattr(storage_roots, "exports_root", lambda: export_root) + + assert storage_roots.resolve_export_write_dir(str(external)) == external + + with pytest.raises(ValueError, match = "path escapes root"): + storage_roots.resolve_export_dir(str(external)) + + +def test_export_write_dir_accepts_expanded_home_path(tmp_path, monkeypatch): + storage_roots = _load_module( + "test_storage_roots_accept_home_path", + "utils/paths/storage_roots.py", + ) + + export_root = tmp_path / "exports" + home = tmp_path / "home" + export_root.mkdir() + home.mkdir() + monkeypatch.setattr(storage_roots, "exports_root", lambda: export_root) + if storage_roots.os.name == "nt": + monkeypatch.setenv("USERPROFILE", str(home)) + else: + monkeypatch.setenv("HOME", str(home)) + + assert storage_roots.resolve_export_write_dir("~/exports/model") == home / "exports" / "model" + + +def test_resolve_export_write_dir_rejects_backslash_parent_segment(): + storage_roots = _load_module( + "test_storage_roots_reject_parent", + "utils/paths/storage_roots.py", + ) + + with pytest.raises(ValueError, match = r"\.\."): + storage_roots.resolve_export_write_dir(r"exports\..\outside") + + +def test_export_write_dir_handles_non_native_windows_absolute_as_relative(tmp_path, monkeypatch): + storage_roots = _load_module( + "test_storage_roots_non_native_windows_path", + "utils/paths/storage_roots.py", + ) + + export_root = tmp_path / "exports" + export_root.mkdir() + monkeypatch.setattr(storage_roots, "exports_root", lambda: export_root) + + if storage_roots.os.name == "nt": + pytest.skip("Windows drive paths are native on Windows") + + assert ( + storage_roots.resolve_export_write_dir(r"C:\exports\model") + == export_root / r"C:\exports\model" + ) + + +def test_export_details_registers_external_absolute_output(tmp_path, monkeypatch): + _install_lightweight_backend_stubs(monkeypatch) + export_route = _load_module( + "test_routes_export_external", + "routes/export.py", + monkeypatch, + ) + + output = tmp_path / "Gemma4_26B_gguf" + output.mkdir() + export_root = tmp_path / "studio" / "exports" + export_root.mkdir(parents = True) + registered = [] + + monkeypatch.setattr( + export_route, + "_try_register_external_export", + lambda path: (registered.append(path) is None, str(path)), + ) + monkeypatch.setattr( + "utils.paths.storage_roots.exports_root", + lambda: export_root, + ) + + details = export_route._export_details(str(output)) + + assert details == { + "output_path": str(output), + "scan_folder_registered": True, + "scan_folder_path": str(output), + } + assert registered == [output] + + +def test_export_details_does_not_register_contained_exports(tmp_path, monkeypatch): + _install_lightweight_backend_stubs(monkeypatch) + export_route = _load_module( + "test_routes_export_contained", + "routes/export.py", + monkeypatch, + ) + + export_root = tmp_path / "exports" + output = export_root / "model-gguf" + output.mkdir(parents = True) + + monkeypatch.setattr( + export_route, + "_try_register_external_export", + lambda path: pytest.fail(f"unexpected registration: {path}"), + ) + monkeypatch.setattr( + "utils.paths.storage_roots.exports_root", + lambda: export_root, + ) + + assert export_route._export_details(str(output)) == {"output_path": "model-gguf"} + + +def test_registered_absolute_export_folder_is_discoverable(tmp_path, monkeypatch): + _install_lightweight_backend_stubs(monkeypatch) + models_route = _load_module("test_routes_models", "routes/models.py", monkeypatch) + + export_dir = tmp_path / "Gemma4_26B_gguf" + export_dir.mkdir() + gguf_file = export_dir / "Gemma4_26B.BF16-00001-of-00002.gguf" + gguf_file.write_bytes(b"gguf") + + found = models_route._scan_models_dir(export_dir) + + assert len(found) == 1 + assert found[0].path == str(gguf_file) + assert found[0].source == "models_dir" diff --git a/studio/backend/tests/test_external_provider_usage_chunk.py b/studio/backend/tests/test_external_provider_usage_chunk.py index ebfd6a8f50..3ec9133718 100644 --- a/studio/backend/tests/test_external_provider_usage_chunk.py +++ b/studio/backend/tests/test_external_provider_usage_chunk.py @@ -144,6 +144,14 @@ def _make_openai_client() -> ExternalProviderClient: ) +def _make_custom_client() -> ExternalProviderClient: + return ExternalProviderClient( + provider_type = "custom", + base_url = "http://custom.example/v1", + api_key = "", + ) + + def _anthropic_sse(events: list[dict]) -> bytes: chunks: list[str] = [] for event in events: @@ -180,6 +188,137 @@ def _usage_chunks(lines: list[str]) -> list[dict]: return out +def test_custom_provider_registry_is_hidden(): + from core.inference.providers import get_provider_info, list_available_providers + + info = get_provider_info("custom") + assert info is not None + assert info["hidden"] is True + assert "custom" not in {p["provider_type"] for p in list_available_providers()} + + +def test_custom_provider_uses_chat_completions_without_auth_key(monkeypatch): + captured: dict = {} + + def handler(request: httpx.Request) -> httpx.Response: + captured["url"] = str(request.url) + captured["headers"] = dict(request.headers) + captured["body"] = json.loads(request.content.decode("utf-8")) + return httpx.Response( + 200, + content = b'data: {"choices":[{"delta":{"content":"ok"}}]}\n\ndata: [DONE]\n\n', + headers = {"content-type": "text/event-stream"}, + ) + + _mock_http_client(monkeypatch, handler) + + async def run(): + client = _make_custom_client() + lines = await _collect( + client.stream_chat_completion( + messages = [{"role": "user", "content": "ping"}], + model = "Qwen/Qwen3-0.6B", + temperature = 0.7, + top_p = 0.95, + max_tokens = 64, + ) + ) + await client.close() + return lines + + lines = _drive(run()) + assert captured["url"] == "http://custom.example/v1/chat/completions" + assert "authorization" not in {k.lower() for k in captured["headers"]} + assert captured["body"]["model"] == "Qwen/Qwen3-0.6B" + assert any("ok" in line for line in lines) + + +def test_custom_provider_test_endpoint_probes_chat_completion(monkeypatch): + import importlib.util + import sys + from pathlib import Path + + module_path = Path(__file__).resolve().parents[1] / "routes" / "providers.py" + spec = importlib.util.spec_from_file_location("_providers_route_under_test", module_path) + assert spec is not None + assert spec.loader is not None + providers_route = importlib.util.module_from_spec(spec) + sys.modules[spec.name] = providers_route + spec.loader.exec_module(providers_route) + + captured: dict = {} + + class _FakeClient: + def __init__(self, **kwargs): + captured["init"] = kwargs + + async def chat_completion(self, **kwargs): + captured["chat_completion"] = kwargs + return {"choices": [{"message": {"content": "ok"}}]} + + async def list_models(self): + raise AssertionError("custom provider test must not call /models") + + async def close(self): + captured["closed"] = True + + monkeypatch.setattr(providers_route, "ExternalProviderClient", _FakeClient) + + async def run(): + return await providers_route.test_provider( + providers_route.ProviderTestRequest( + provider_type = "custom", + base_url = "http://custom.example/v1", + model_id = "Qwen/Qwen3-0.6B", + ), + current_subject = "unsloth", + ) + + result = _drive(run()) + assert result.success is True + assert result.models_count is None + assert captured["init"]["provider_type"] == "custom" + assert captured["chat_completion"]["model"] == "Qwen/Qwen3-0.6B" + assert captured["chat_completion"]["max_tokens"] == 1 + assert captured["closed"] is True + + +def test_custom_provider_test_endpoint_requires_model_id(monkeypatch): + import importlib.util + import sys + from pathlib import Path + + module_path = Path(__file__).resolve().parents[1] / "routes" / "providers.py" + spec = importlib.util.spec_from_file_location("_providers_route_under_test", module_path) + assert spec is not None + assert spec.loader is not None + providers_route = importlib.util.module_from_spec(spec) + sys.modules[spec.name] = providers_route + spec.loader.exec_module(providers_route) + + class _FakeClient: + def __init__(self, **kwargs): + pass + + async def close(self): + pass + + monkeypatch.setattr(providers_route, "ExternalProviderClient", _FakeClient) + + async def run(): + return await providers_route.test_provider( + providers_route.ProviderTestRequest( + provider_type = "custom", + base_url = "http://custom.example/v1", + ), + current_subject = "unsloth", + ) + + result = _drive(run()) + assert result.success is False + assert "model ID" in result.message + + def test_anthropic_stream_emits_usage_chunk_before_done(monkeypatch): sse_events = [ { diff --git a/studio/backend/tests/test_gemma4_chat_template_override.py b/studio/backend/tests/test_gemma4_chat_template_override.py new file mode 100644 index 0000000000..f726741aa5 --- /dev/null +++ b/studio/backend/tests/test_gemma4_chat_template_override.py @@ -0,0 +1,345 @@ +# SPDX-License-Identifier: AGPL-3.0-only +# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0 + +"""Auto-override of the chat template for ``unsloth/gemma-4-*-GGUF``. + +Studio ships a bundled ``gemma-4.jinja`` (PR #118 based, ``preserve_thinking`` +defaulted off) and applies it to gemma-4 GGUF loads via the existing +``chat_template_override`` -> ``--chat-template-file`` path, so users do not need +to re-download quants. Pins the family matcher, the resolver precedence, the +bundled asset's reasoning/tool capabilities (which drive the "Preserve thinking" +UI toggle), the Jinja gate behaviour, and the reload-dedup interaction. +""" + +from __future__ import annotations + +import importlib.util +import sys +import types as _types +from pathlib import Path + +_BACKEND_DIR = str(Path(__file__).resolve().parent.parent) +if _BACKEND_DIR not in sys.path: + sys.path.insert(0, _BACKEND_DIR) + +import pytest + +# ── chat_templates is dependency-light: load it directly so the pure-logic +# tests run without the studio venv / core.inference package side effects. ── +_CT_PATH = Path(_BACKEND_DIR) / "core" / "inference" / "chat_templates.py" +_ct_spec = importlib.util.spec_from_file_location("_gemma4_ct_test", _CT_PATH) +chat_templates = importlib.util.module_from_spec(_ct_spec) +_ct_spec.loader.exec_module(chat_templates) + +is_unsloth_gemma4_gguf = chat_templates.is_unsloth_gemma4_gguf +resolve_effective_chat_template_override = chat_templates.resolve_effective_chat_template_override +load_bundled_chat_template = chat_templates.load_bundled_chat_template +is_unsloth_gemma4_edge_gguf = chat_templates.is_unsloth_gemma4_edge_gguf + +BUNDLED = load_bundled_chat_template("gemma-4.jinja") # 12b / 26B-A4B / 31B +EDGE = load_bundled_chat_template("gemma-4-edge.jinja") # E2B / E4B + + +# ── Stubs so core.inference.llama_cpp imports without the full studio venv ── +def _stub_modules_ctx(): + """patch.dict context that stubs the heavy deps llama_cpp pulls in at import, + but only those NOT already importable (real httpx / structlog are kept when + present, e.g. in CI), and removes the stubs on exit so other tests are not + polluted.""" + from unittest.mock import patch + + _loggers_stub = _types.ModuleType("loggers") + _loggers_stub.get_logger = lambda name: __import__("logging").getLogger(name) + _structlog_stub = _types.ModuleType("structlog") + _structlog_stub.get_logger = lambda *a, **k: __import__("logging").getLogger("stub") + _httpx_stub = _types.ModuleType("httpx") + for _exc in ( + "ConnectError", + "TimeoutException", + "ReadTimeout", + "ReadError", + "RemoteProtocolError", + "CloseError", + ): + setattr(_httpx_stub, _exc, type(_exc, (Exception,), {})) + _httpx_stub.Timeout = type("T", (), {"__init__": lambda s, *a, **k: None}) + _httpx_stub.Client = type( + "C", + (), + { + "__init__": lambda s, **kw: None, + "__enter__": lambda s: s, + "__exit__": lambda s, *a: None, + }, + ) + overrides = { + name: stub + for name, stub in ( + ("loggers", _loggers_stub), + ("structlog", _structlog_stub), + ("httpx", _httpx_stub), + ) + if name not in sys.modules + } + return patch.dict(sys.modules, overrides) + + +def _detect_reasoning_flags(): + with _stub_modules_ctx(): + from core.inference.llama_cpp import detect_reasoning_flags + return detect_reasoning_flags + + +# ── Family matcher ─────────────────────────────────────────────────── + + +@pytest.mark.parametrize( + "model_id,expected", + [ + ("unsloth/gemma-4-E2B-it-GGUF", True), + ("unsloth/gemma-4-E4B-it-GGUF", True), + ("unsloth/gemma-4-31B-it-GGUF", True), + ("unsloth/gemma-4-26B-A4B-it-GGUF", True), + ("UNSLOTH/GEMMA-4-E2B-IT-GGUF", True), # case-insensitive + ("gemma-4-E2B-it-GGUF", True), # owner-less shorthand -> unsloth/ + ("gemma-4-31B-it-GGUF", True), # owner-less shorthand -> unsloth/ + ("unsloth/gemma-4-E2B-it", False), # bf16, not GGUF + ("unsloth/gemma-3-4b-it-GGUF", False), # gemma 3 + ("google/gemma-4-31B-it-GGUF", False), # not unsloth + ("unsloth/Qwen3.5-9B-MTP-GGUF", False), + ("/home/user/models/gemma-4-E2B.Q4_K_M.gguf", False), # local path + ("", False), + (None, False), + ], +) +def test_is_unsloth_gemma4_gguf(model_id, expected): + assert is_unsloth_gemma4_gguf(model_id) is expected + + +# ── Resolver precedence ────────────────────────────────────────────── + + +@pytest.mark.parametrize( + "model_id,expected_edge", + [ + ("unsloth/gemma-4-E2B-it-GGUF", True), + ("unsloth/gemma-4-E4B-it-GGUF", True), + ("UNSLOTH/GEMMA-4-E4B-IT-GGUF", True), + ("unsloth/gemma-4-12b-it-GGUF", False), + ("unsloth/gemma-4-26B-A4B-it-GGUF", False), + ("unsloth/gemma-4-31B-it-GGUF", False), + ("unsloth/gemma-3-4b-it-GGUF", False), + ], +) +def test_is_unsloth_gemma4_edge_gguf(model_id, expected_edge): + assert is_unsloth_gemma4_edge_gguf(model_id) is expected_edge + + +def test_resolver_returns_edge_template_for_e2b_e4b(): + for mid in ("unsloth/gemma-4-E2B-it-GGUF", "unsloth/gemma-4-E4B-it-GGUF"): + out = resolve_effective_chat_template_override(model_identifier = mid, user_override = None) + assert out == EDGE + assert out != BUNDLED + + +def test_resolver_handles_owner_less_shorthand(): + # ModelConfig.from_identifier prefixes unsloth/ for bare ids; the resolver + # runs before that, so it must apply the same normalization. + assert ( + resolve_effective_chat_template_override( + model_identifier = "gemma-4-E2B-it-GGUF", user_override = None + ) + == EDGE + ) + assert ( + resolve_effective_chat_template_override( + model_identifier = "gemma-4-31B-it-GGUF", user_override = None + ) + == BUNDLED + ) + + +def test_resolver_returns_standard_template_for_larger_models(): + for mid in ( + "unsloth/gemma-4-12b-it-GGUF", + "unsloth/gemma-4-26B-A4B-it-GGUF", + "unsloth/gemma-4-31B-it-GGUF", + ): + out = resolve_effective_chat_template_override(model_identifier = mid, user_override = None) + assert out == BUNDLED + + +def test_resolver_user_override_wins(): + out = resolve_effective_chat_template_override( + model_identifier = "unsloth/gemma-4-E2B-it-GGUF", user_override = "MY TEMPLATE" + ) + assert out == "MY TEMPLATE" + + +def test_resolver_blank_override_falls_back_to_bundled(): + out = resolve_effective_chat_template_override( + model_identifier = "unsloth/gemma-4-31B-it-GGUF", user_override = " " + ) + assert out == BUNDLED + + +def test_resolver_none_for_non_gemma(): + assert ( + resolve_effective_chat_template_override( + model_identifier = "unsloth/Llama-3.2-1B-Instruct-GGUF", user_override = None + ) + is None + ) + + +# ── Bundled asset content + capability classification ──────────────── + + +@pytest.mark.parametrize("tpl", [BUNDLED, EDGE]) +def test_bundled_template_has_preserve_thinking_defaulted_off(tpl): + assert "preserve_thinking" in tpl + assert "preserve_thinking | default(false)" in tpl + + +@pytest.mark.parametrize("name", ["gemma-4.jinja", "gemma-4-edge.jinja"]) +def test_bundled_templates_are_ascii(name): + # The temp file written for --chat-template-file must encode on any locale. + # Keeping the bundled templates ASCII avoids UnicodeEncodeError on non-UTF-8 + # Windows locales (cp932/cp1252) regardless of the writer's encoding. + text = load_bundled_chat_template(name) + non_ascii = sorted({c for c in text if ord(c) > 127}) + assert not non_ascii, f"{name} has non-ASCII chars: {non_ascii}" + + +@pytest.mark.parametrize("tpl", [BUNDLED, EDGE]) +def test_detect_reasoning_flags_on_bundled_template(tpl): + detect_reasoning_flags = _detect_reasoning_flags() + flags = detect_reasoning_flags(tpl, "unsloth/gemma-4-E2B-it-GGUF") + assert flags["supports_reasoning"] is True + assert flags["reasoning_style"] == "enable_thinking" + assert flags["reasoning_always_on"] is False + # This is what makes the "Preserve thinking" toggle appear in the UI. + assert flags["supports_preserve_thinking"] is True + assert flags["supports_tools"] is True + + +def test_edge_template_omits_empty_thought_block_on_thinking_off(): + """E2B/E4B must NOT emit the empty <|channel>thought block when + thinking is disabled; the larger-model template must. This is the only + intended difference between the two bundled templates.""" + EMPTY = "<|channel>thought\n" + msgs = [{"role": "user", "content": "hi"}] + edge_off = _render_with(EDGE, msgs, enable_thinking = False) + std_off = _render_with(BUNDLED, msgs, enable_thinking = False) + assert EMPTY not in edge_off, "edge (E2B/E4B) should not emit empty thought block" + assert EMPTY in std_off, "standard (12b/26B/31B) should emit empty thought block" + # With thinking ON neither appends the empty block at the prompt tail. + assert EMPTY not in _render_with(EDGE, msgs, enable_thinking = True) + + +# ── Jinja gate behaviour (off = omit prior reasoning, on = keep) ───── + + +def _render_with(tpl, messages, **kw): + pytest.importorskip("jinja2") # transitive via transformers; skip in minimal envs + from jinja2 import Environment, BaseLoader + + def raise_exception(msg): + raise RuntimeError(msg) + + env = Environment(loader = BaseLoader()) + return env.from_string(tpl).render( + messages = messages, + bos_token = "", + raise_exception = raise_exception, + add_generation_prompt = True, + **kw, + ) + + +def _render(messages, **kw): + return _render_with(BUNDLED, messages, **kw) + + +def _convo_with_prior_tool_reasoning(): + # Assistant tool-call turn with reasoning, BEFORE the last user message. + return [ + {"role": "user", "content": "q1"}, + { + "role": "assistant", + "reasoning_content": "SECRET_THOUGHT", + "tool_calls": [{"id": "c1", "function": {"name": "f", "arguments": {"x": 1}}}], + }, + {"role": "tool", "tool_call_id": "c1", "content": "42"}, + {"role": "user", "content": "q2"}, + ] + + +def test_preserve_thinking_off_omits_prior_reasoning(): + # default(false): kwarg unset -> prior reasoning dropped before last user turn. + assert "SECRET_THOUGHT" not in _render(_convo_with_prior_tool_reasoning()) + + +def test_preserve_thinking_on_keeps_prior_reasoning(): + assert "SECRET_THOUGHT" in _render(_convo_with_prior_tool_reasoning(), preserve_thinking = True) + + +def test_enable_thinking_gates_think_token(): + assert "<|think|>" in _render([{"role": "user", "content": "hi"}], enable_thinking = True) + assert "<|think|>" not in _render([{"role": "user", "content": "hi"}], enable_thinking = False) + + +# ── Reload dedup interaction (why the route resolves the effective override) ── + + +def test_already_in_target_state_consistent_with_bundled_override(): + """The backend dedup compares the incoming override against the live one. + + The route resolves the bundled template up front so a re-load that omits + ``chat_template_override`` still matches (no spurious reload), while a raw + ``None`` would not. + """ + LlamaCppBackend = _import_backend() + + class _FakeProcess: + def terminate(self): ... + def wait(self, timeout = None): + return 0 + + def kill(self): ... + def poll(self): + return 0 + + backend = LlamaCppBackend() + backend._process = _FakeProcess() + backend._healthy = True + backend._model_identifier = "unsloth/gemma-4-E2B-it-GGUF" + backend._hf_variant = "Q4_K_M" + backend._requested_n_ctx = 8192 + backend._cache_type_kv = None + backend._speculative_type = None + backend._requested_spec_mode = "auto" + backend._chat_template_override = BUNDLED # live server launched with the bundle + backend._is_vision = False + backend._extra_args = None + backend._gguf_path = None + + common = dict( + model_identifier = "unsloth/gemma-4-E2B-it-GGUF", + hf_variant = "Q4_K_M", + n_ctx = 8192, + cache_type_kv = None, + speculative_type = None, + extra_args = None, + is_vision = False, + ) + # Effective (resolved bundled) override -> already loaded, no reload. + assert backend._already_in_target_state(chat_template_override = BUNDLED, **common) is True + # Raw None (unresolved) -> false match, would force a needless reload. + assert backend._already_in_target_state(chat_template_override = None, **common) is False + + +def _import_backend(): + with _stub_modules_ctx(): + from core.inference.llama_cpp import LlamaCppBackend + return LlamaCppBackend diff --git a/studio/backend/tests/test_gguf_route_cursor_reset.py b/studio/backend/tests/test_gguf_route_cursor_reset.py index fdbf31f9de..dc397b2d08 100644 --- a/studio/backend/tests/test_gguf_route_cursor_reset.py +++ b/studio/backend/tests/test_gguf_route_cursor_reset.py @@ -59,11 +59,16 @@ def _make_backend(monkeypatch, streams: list[list[str]], payloads: list[dict]): payload, _cancel_event, headers = None, + first_token_deadline = None, ): payloads.append(copy.deepcopy(payload)) yield type("FakeResponse", (), {"status_code": 200, "chunks": streams.pop(0)})() - def fake_iter_text_cancellable(response, _cancel_event): + def fake_iter_text_cancellable( + response, + _cancel_event, + first_token_deadline = None, + ): yield from response.chunks monkeypatch.setattr(backend, "_stream_with_retry", fake_stream_with_retry) diff --git a/studio/backend/tests/test_lemonade_llamacpp_rocm_bins_mock.py b/studio/backend/tests/test_lemonade_llamacpp_rocm_bins_mock.py deleted file mode 100644 index 7c6c514b8f..0000000000 --- a/studio/backend/tests/test_lemonade_llamacpp_rocm_bins_mock.py +++ /dev/null @@ -1,523 +0,0 @@ -# SPDX-License-Identifier: AGPL-3.0-only -# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0 - -"""Validates that the installer correctly resolves lemonade ROCm prebuilt assets. - -Uses a faked HostInfo so no AMD GPU is needed. Network calls to the lemonade -GitHub API are stubbed out so the suite runs without internet access and is -not subject to rate limits. -""" - -from __future__ import annotations - -import importlib -import sys -from pathlib import Path -from unittest.mock import patch - -import pytest - -_studio = Path(__file__).resolve().parent.parent.parent -if str(_studio) not in sys.path: - sys.path.insert(0, str(_studio)) - -_mod = importlib.import_module("install_llama_prebuilt") -HostInfo = _mod.HostInfo -resolve_lemonade_rocm_choice = getattr(_mod, "resolve_lemonade_rocm_choice", None) -_LEMONADE_GFX_FAMILIES = getattr(_mod, "_LEMONADE_GFX_FAMILIES", None) - -if resolve_lemonade_rocm_choice is None or _LEMONADE_GFX_FAMILIES is None: - pytest.skip("PR symbols not present - check branch", allow_module_level = True) - - -@pytest.fixture(autouse = True) -def _clear_lemonade_release_cache(): - """Prevent cross-test pollution of the lemonade release lru_cache and - selection-log dedup set when tests vary the fetch_json mock return value.""" - _cache = getattr(_mod, "_fetch_lemonade_release_cached", None) - _logged: set | None = getattr(_mod, "_lemonade_selection_logged", None) - if _cache is not None and hasattr(_cache, "cache_clear"): - _cache.cache_clear() - if _logged is not None: - _logged.clear() - yield - if _cache is not None and hasattr(_cache, "cache_clear"): - _cache.cache_clear() - if _logged is not None: - _logged.clear() - - -_STUB_TAG = "b1262" -_STUB_OS_PREFIXES = ("ubuntu", "windows") -_STUB_FAMILIES = ("gfx1151", "gfx1150", "gfx120X", "gfx110X", "gfx103X") - - -def _stub_lemonade_release() -> dict: - """Minimal lemonade release payload covering all supported GPU/OS combinations.""" - assets = [ - { - "name": f"llama-{_STUB_TAG}-{prefix}-rocm-{family}-x64.zip", - "browser_download_url": ( - f"https://github.com/lemonade-sdk/llamacpp-rocm/releases/download/" - f"{_STUB_TAG}/llama-{_STUB_TAG}-{prefix}-rocm-{family}-x64.zip" - ), - } - for prefix in _STUB_OS_PREFIXES - for family in _STUB_FAMILIES - ] - return {"tag_name": _STUB_TAG, "assets": assets} - - -def _make_rocm_host(gfx_target: str, *, windows: bool = False) -> HostInfo: - return HostInfo( - system = "Windows" if windows else "Linux", - machine = "amd64" if windows else "x86_64", - is_windows = windows, - is_linux = not windows, - is_macos = False, - is_x86_64 = True, - is_arm64 = False, - nvidia_smi = None, - driver_cuda_version = None, - compute_caps = [], - visible_cuda_devices = None, - has_physical_nvidia = False, - has_usable_nvidia = False, - has_rocm = True, - rocm_gfx_target = gfx_target, - ) - - -def _lookup_family(gfx: str) -> str | None: - for prefix, family in _LEMONADE_GFX_FAMILIES: - if gfx.startswith(prefix): - return family - return None - - -# --------------------------------------------------------------------------- -# GPU family mapping -# --------------------------------------------------------------------------- - - -@pytest.mark.parametrize( - "gfx,expected_family", - [ - ("gfx1151", "gfx1151"), - ("gfx1150", "gfx1150"), - ("gfx1201", "gfx120X"), - ("gfx1200", "gfx120X"), - ("gfx1100", "gfx110X"), - ("gfx1030", "gfx103X"), - ], -) -def test_gpu_family_mapping(gfx, expected_family): - assert _lookup_family(gfx) == expected_family - - -def test_unknown_gpu_not_in_families(): - assert _lookup_family("gfx999") is None - - -# --------------------------------------------------------------------------- -# Asset resolution - hits real lemonade GitHub API -# --------------------------------------------------------------------------- - - -@pytest.mark.parametrize( - "gfx,os_prefix,windows", - [ - ("gfx1151", "ubuntu", False), - ("gfx1150", "ubuntu", False), - ("gfx1201", "ubuntu", False), - ("gfx1100", "ubuntu", False), - ("gfx1030", "ubuntu", False), - ("gfx1151", "windows", True), - ("gfx1100", "windows", True), - ], -) -def test_asset_resolves_for_known_gpu(gfx, os_prefix, windows): - host = _make_rocm_host(gfx, windows = windows) - with patch.object(_mod, "fetch_json", return_value = _stub_lemonade_release()): - result = resolve_lemonade_rocm_choice(host, os_prefix, "default", llama_tag = "latest") - assert result is not None, f"Installer will NOT fetch lemonade binary for {gfx} ({os_prefix})" - assert _lookup_family(gfx) in result.name - assert result.url.startswith("https://github.com/lemonade-sdk/llamacpp-rocm") - - -def test_unknown_gpu_falls_through_to_upstream(): - host = _make_rocm_host("gfx999") - result = resolve_lemonade_rocm_choice(host, "ubuntu", "default", llama_tag = "latest") - assert result is None - - -# --------------------------------------------------------------------------- -# The Linux attempt builder must plan a lemonade ROCm attempt for AMD-only hosts. -# This is the path setup.sh actually invokes (fork hosts now select from the -# manifest), so the lemonade integration is useless if it isn't wired in here. -# --------------------------------------------------------------------------- - -_linux_published_attempts = getattr(_mod, "_linux_published_attempts", None) -direct_upstream_release_plan = getattr(_mod, "direct_upstream_release_plan", None) - -PublishedLlamaArtifact = _mod.PublishedLlamaArtifact -PublishedReleaseBundle = _mod.PublishedReleaseBundle - - -def _rocm_bundle(gfx_family: str, mapped_targets: list[str]) -> "PublishedReleaseBundle": - """A fork manifest bundle exposing a per-gfx linux-rocm artifact, so - published_rocm_choice_for_host can match the host before the lemonade - fallback is appended.""" - asset_name = f"app-b9457-linux-x64-rocm-{gfx_family}.tar.gz" - artifact = PublishedLlamaArtifact( - asset_name = asset_name, - install_kind = "linux-rocm", - runtime_line = None, - coverage_class = None, - supported_sms = [], - min_sm = None, - max_sm = None, - bundle_profile = None, - rank = 1000, - gfx_target = gfx_family, - mapped_targets = mapped_targets, - ) - return PublishedReleaseBundle( - repo = "unslothai/llama.cpp", - release_tag = "v1.0", - upstream_tag = "b9457", - assets = {asset_name: f"https://example.invalid/{asset_name}"}, - artifacts = [artifact], - ) - - -@pytest.mark.skipif( - _linux_published_attempts is None, - reason = "Linux attempt builder not present on this branch", -) -def test_linux_attempts_include_fork_rocm_and_lemonade_for_rocm_host(): - host = _make_rocm_host("gfx1151") - bundle = _rocm_bundle("gfx1151", ["gfx1151"]) - with patch.object(_mod, "fetch_json", return_value = _stub_lemonade_release()): - attempts = _linux_published_attempts(host, bundle, "latest") - kinds = [a.install_kind for a in attempts] - assert "linux-rocm" in kinds, f"builder did not include any linux-rocm attempt; got {kinds}" - sources = {a.source_label for a in attempts if a.install_kind == "linux-rocm"} - # The fork's own per-gfx bundle is preferred, with the lemonade prebuilt as - # the fallback -- both must be present for a covered ROCm host. - assert "published" in sources, f"fork ROCm bundle missing; got {sources}" - assert "lemonade" in sources, f"lemonade ROCm fallback missing; got {sources}" - lemonade_attempt = next(a for a in attempts if a.source_label == "lemonade") - assert "gfx1151" in lemonade_attempt.name - - -@pytest.mark.skipif( - direct_upstream_release_plan is None, - reason = "direct release planners not present on this branch", -) -def test_direct_upstream_plan_includes_lemonade_for_windows_hip_host(): - host = _make_rocm_host("gfx1151", windows = True) - release = { - "tag_name": "b9022", - "name": "b9022", - "assets": [], - } - with patch.object(_mod, "fetch_json", return_value = _stub_lemonade_release()): - plan = direct_upstream_release_plan(release, host, "ggml-org/llama.cpp", "latest") - assert plan is not None, "Windows ROCm host should plan a lemonade HIP attempt" - kinds = [a.install_kind for a in plan.attempts] - assert "windows-hip" in kinds, f"planner did not include a lemonade HIP attempt; got {kinds}" - - -@pytest.mark.skipif( - direct_upstream_release_plan is None, - reason = "direct release planners not present on this branch", -) -def test_windows_hip_falls_back_to_upstream_when_lemonade_unavailable(): - """If lemonade returns None (e.g. gfx999 or transient API failure), the planner - must still include the upstream HIP asset rather than silently downgrading to CPU.""" - host = _make_rocm_host("gfx999", windows = True) - hip_asset = "llama-b9022-bin-win-hip-radeon-x64.zip" - release = { - "tag_name": "b9022", - "name": "b9022", - "assets": [ - { - "name": hip_asset, - "browser_download_url": f"https://example.invalid/{hip_asset}", - }, - ], - } - plan = direct_upstream_release_plan(release, host, "ggml-org/llama.cpp", "latest") - assert plan is not None - kinds = [a.install_kind for a in plan.attempts] - assert "windows-hip" in kinds, f"upstream HIP asset not included as fallback; got {kinds}" - hip_attempt = next(a for a in plan.attempts if a.install_kind == "windows-hip") - assert hip_attempt.source_label == "upstream" - - -# ── Follow-up: pinned-tag URL helper, URL trust pinning, opt-out env, autouse cache clear ── - - -def test_lemonade_release_api_url_pinned_tag(): - """A pinned llama_tag must produce the /releases/tags/ URL.""" - assert _mod._lemonade_release_api_for("b1262").endswith("/releases/tags/b1262") - assert _mod._lemonade_release_api_for("latest").endswith("/releases/latest") - assert _mod._lemonade_release_api_for("").endswith("/releases/latest") - - -def test_lemonade_release_api_url_encodes_tag(): - """Unexpected slashes / hashes in the tag must be URL-encoded so the URL - cannot be reshaped (defence in depth -- tags should already be sanitised - upstream).""" - url = _mod._lemonade_release_api_for("b1260/../latest") - assert "/releases/tags/b1260%2F..%2Flatest" in url - assert "//latest" not in url.split("/releases/tags/", 1)[1] - - -def test_lemonade_resolver_skipped_by_opt_out_env(monkeypatch): - """UNSLOTH_DISABLE_LEMONADE_ROCM=1 must short-circuit the resolver.""" - monkeypatch.setenv("UNSLOTH_DISABLE_LEMONADE_ROCM", "1") - host = _make_rocm_host("gfx1151") - res = resolve_lemonade_rocm_choice(host, "ubuntu", "linux-rocm", llama_tag = "latest") - assert res is None - - -def test_lemonade_resolver_rejects_non_github_url(monkeypatch): - """If the GitHub API response somehow contained an off-host download URL, - the resolver must refuse to use it (lemonade assets are not in the - approved-hash manifest).""" - bad_release = { - "tag_name": _STUB_TAG, - "assets": [ - { - "name": f"llama-{_STUB_TAG}-ubuntu-rocm-gfx1151-x64.zip", - "browser_download_url": "https://attacker.invalid/llama.zip", - }, - ], - } - host = _make_rocm_host("gfx1151") - with patch.object(_mod, "fetch_json", return_value = bad_release): - res = resolve_lemonade_rocm_choice(host, "ubuntu", "linux-rocm", llama_tag = "latest") - assert res is None - - -def test_lemonade_resolver_rejects_http_scheme(): - assert not _mod._is_trusted_github_release_url( - "http://github.com/lemonade-sdk/llamacpp-rocm/releases/download/x/y.zip", - "lemonade-sdk/llamacpp-rocm", - ) - - -def test_lemonade_resolver_accepts_github_cdn(): - # Real GitHub release CDN URLs carry the /github-production-release-asset- prefix. - assert _mod._is_trusted_github_release_url( - "https://objects.githubusercontent.com/github-production-release-asset-abc123/456/789?token=x", - "lemonade-sdk/llamacpp-rocm", - ) - - -def test_lemonade_resolver_rejects_arbitrary_cdn_path(): - # A CDN URL without the release-asset path prefix must be rejected. - assert not _mod._is_trusted_github_release_url( - "https://objects.githubusercontent.com/abc/def", - "lemonade-sdk/llamacpp-rocm", - ) - - -def test_lemonade_resolver_accepts_release_path(): - url = "https://github.com/lemonade-sdk/llamacpp-rocm/releases/download/b1262/llama-b1262-ubuntu-rocm-gfx1151-x64.zip" - assert _mod._is_trusted_github_release_url(url, "lemonade-sdk/llamacpp-rocm") - - -def test_lemonade_resolver_rejects_wrong_repo(): - """A github.com release URL for a different repo must be rejected.""" - assert not _mod._is_trusted_github_release_url( - "https://github.com/attacker/llamacpp-rocm/releases/download/x/y.zip", - "lemonade-sdk/llamacpp-rocm", - ) - - -def test_lemonade_resolver_rejects_empty_browser_download_url(): - """An asset entry with an empty browser_download_url must fall through.""" - release = { - "tag_name": _STUB_TAG, - "assets": [ - { - "name": f"llama-{_STUB_TAG}-ubuntu-rocm-gfx1151-x64.zip", - "browser_download_url": "", - }, - ], - } - host = _make_rocm_host("gfx1151") - with patch.object(_mod, "fetch_json", return_value = release): - res = resolve_lemonade_rocm_choice(host, "ubuntu", "linux-rocm", llama_tag = "latest") - assert res is None - - -def test_lemonade_runtime_patterns_include_hip_runtime(): - """linux-rocm overlay must use a broad lib glob to catch all bundled .so files. - - Lemonade ZIPs carry transitive deps (libamd_comgr, libLLVM, libclang-cpp, - ...) whose names change across ROCm releases. A broad ``lib*.so*`` glob - avoids having to enumerate every transitive dependency by name. - """ - from install_llama_prebuilt import runtime_patterns_for_choice, AssetChoice - - choice = AssetChoice( - repo = "lemonade-sdk/llamacpp-rocm", - tag = "b1262", - name = "llama-b1262-ubuntu-rocm-gfx1151-x64.zip", - url = "https://github.com/lemonade-sdk/llamacpp-rocm/releases/download/b1262/x.zip", - source_label = "lemonade", - install_kind = "linux-rocm", - ) - pats = runtime_patterns_for_choice(choice) - # The broad glob must be present so every .so in the lemonade bundle - # (including transitive deps added in future ROCm releases) gets overlaid. - assert "lib*.so*" in pats, f"'lib*.so*' missing from linux-rocm patterns: {pats}" - - -_pick_rocm_gfx_target = getattr(_mod, "_pick_rocm_gfx_target", None) - - -@pytest.mark.skipif( - _pick_rocm_gfx_target is None, - reason = "_pick_rocm_gfx_target not present on this branch", -) -def test_pick_rocm_gfx_target_honors_cuda_visible_devices(monkeypatch): - """AMD HIP honours CUDA_VISIBLE_DEVICES identically to HIP_VISIBLE_DEVICES; - on a gfx1151 + gfx1100 mixed host, CUDA_VISIBLE_DEVICES=1 must select gfx1100.""" - # Two GPUs; rocminfo reports each token twice (as in the real tool output). - probe_out = "gfx1151\ngfx1151\ngfx1100\ngfx1100" - monkeypatch.delenv("HIP_VISIBLE_DEVICES", raising = False) - monkeypatch.delenv("ROCR_VISIBLE_DEVICES", raising = False) - monkeypatch.setenv("CUDA_VISIBLE_DEVICES", "1") - assert _pick_rocm_gfx_target(probe_out) == "gfx1100" - - -@pytest.mark.skipif( - _pick_rocm_gfx_target is None, - reason = "_pick_rocm_gfx_target not present on this branch", -) -def test_pick_rocm_gfx_target_cuda_visible_devices_minus_one_returns_none(monkeypatch): - """CUDA_VISIBLE_DEVICES=-1 means no GPU visible; resolver must return None.""" - probe_out = "gfx1151\ngfx1100" - monkeypatch.delenv("HIP_VISIBLE_DEVICES", raising = False) - monkeypatch.delenv("ROCR_VISIBLE_DEVICES", raising = False) - monkeypatch.setenv("CUDA_VISIBLE_DEVICES", "-1") - assert _pick_rocm_gfx_target(probe_out) is None - - -@pytest.mark.skipif( - _pick_rocm_gfx_target is None, - reason = "_pick_rocm_gfx_target not present on this branch", -) -def test_pick_rocm_gfx_target_same_arch_multi_gpu(monkeypatch): - """Regression: [gfx1100, gfx1100, gfx1151] with HIP_VISIBLE_DEVICES=2 must - return gfx1151, not fall back to GPU 0 due to dict.fromkeys collapsing the - two gfx1100 entries into one and making index 2 out of range.""" - # Simulate rocminfo output for 3 GPUs (2x gfx1100 dGPU + 1x gfx1151 APU). - # Each GPU gets its own Agent section with a few token mentions. - probe_out = ( - "***\nAgent 1\n***\n gfx1100 some info\n gfx1100\n" - "***\nAgent 2\n***\n gfx1100 some info\n gfx1100\n" - "***\nAgent 3\n***\n gfx1151 some info\n gfx1151\n" - ) - monkeypatch.delenv("ROCR_VISIBLE_DEVICES", raising = False) - monkeypatch.delenv("CUDA_VISIBLE_DEVICES", raising = False) - monkeypatch.setenv("HIP_VISIBLE_DEVICES", "2") - assert _pick_rocm_gfx_target(probe_out) == "gfx1151" - - -# --------------------------------------------------------------------------- -# Fork release scan: Windows ROCm resolves lemonade by the requested tag -# --------------------------------------------------------------------------- - -_resolve_release_asset_choice = getattr(_mod, "resolve_release_asset_choice", None) -_ApprovedReleaseChecksums = getattr(_mod, "ApprovedReleaseChecksums", None) - - -@pytest.mark.skipif( - _resolve_release_asset_choice is None or _ApprovedReleaseChecksums is None, - reason = "fork release planner not present on this branch", -) -def test_fork_scan_windows_rocm_resolves_lemonade_by_requested_tag(): - """The fork release scan pins llama_tag to per-release upstream tags - (b9457, ...) that lemonade's own tag series never contains, so the - lemonade lookup must use the requested tag ("latest") instead. Pinning - lemonade to the per-release tag 404s on every scanned release and a - Windows ROCm host ends in a rate-limited fatal instead of the lemonade - prebuilt.""" - host = _make_rocm_host("gfx1151", windows = True) - # No windows-rocm artifact in the bundle, matching current fork releases. - bundle = _rocm_bundle("gfx1151", ["gfx1151"]) - checksums = _ApprovedReleaseChecksums( - repo = "unslothai/llama.cpp", - release_tag = "v1.0", - upstream_tag = "b9457", - artifacts = {}, - ) - seen_urls: list[str] = [] - - def _fake_fetch(api_url, *args, **kwargs): - seen_urls.append(api_url) - if "lemonade-sdk" in api_url: - if api_url.endswith("/releases/latest"): - return _stub_lemonade_release() - raise RuntimeError(f"unexpected pinned lemonade fetch: {api_url}") - # ggml-org asset listing for the upstream HIP/CPU filename fallbacks. - return {"tag_name": "b9457", "assets": []} - - with patch.object(_mod, "fetch_json", side_effect = _fake_fetch): - attempts = _resolve_release_asset_choice( - host, - "b9457", # concrete per-release upstream tag from the scan loop - bundle, - checksums, - requested_tag = "latest", - ) - - lemonade = [a for a in attempts if a.source_label == "lemonade"] - assert lemonade, f"lemonade attempt missing for Windows ROCm host; got {attempts}" - assert "gfx1151" in lemonade[0].name - assert any( - u.endswith("/releases/latest") for u in seen_urls - ), f"lemonade was never resolved via /releases/latest; fetches: {seen_urls}" - assert not any( - "lemonade-sdk" in u and "/releases/tags/" in u for u in seen_urls - ), f"lemonade lookup was pinned to the fork release tag: {seen_urls}" - - -@pytest.mark.skipif( - direct_upstream_release_plan is None, - reason = "direct release planners not present on this branch", -) -def test_direct_upstream_plan_includes_lemonade_for_linux_rocm_host(): - """A Linux ROCm host on the ggml-org direct path (e.g. a --published-repo - override) must plan lemonade before the CPU tarball, mirroring the Windows - branch. The lemonade planning previously lived in the removed - --simple-policy dispatcher, so without this leg such hosts silently - install the CPU build.""" - host = _make_rocm_host("gfx1151") - release = { - "tag_name": "b9022", - "name": "b9022", - "assets": [ - { - "name": "llama-b9022-bin-ubuntu-x64.tar.gz", - "browser_download_url": ( - "https://github.com/ggml-org/llama.cpp/releases/download/" - "b9022/llama-b9022-bin-ubuntu-x64.tar.gz" - ), - } - ], - } - with patch.object(_mod, "fetch_json", return_value = _stub_lemonade_release()): - plan = direct_upstream_release_plan(release, host, "ggml-org/llama.cpp", "latest") - assert plan is not None, "Linux ROCm host should produce a direct plan" - kinds = [a.install_kind for a in plan.attempts] - sources = [a.source_label for a in plan.attempts] - assert "linux-rocm" in kinds, f"lemonade ROCm attempt missing; got {kinds}" - assert sources[0] == "lemonade", f"lemonade must be the first attempt; got {sources}" - assert "gfx1151" in plan.attempts[0].name diff --git a/studio/backend/tests/test_llama_cpp_freshness.py b/studio/backend/tests/test_llama_cpp_freshness.py index f8e4619ded..f90c4ba0e7 100644 --- a/studio/backend/tests/test_llama_cpp_freshness.py +++ b/studio/backend/tests/test_llama_cpp_freshness.py @@ -433,3 +433,90 @@ def test_fetch_latest_release_tag_uses_publish_time(monkeypatch): ] monkeypatch.setattr(urllib.request, "urlopen", lambda req, timeout = 5.0: _Resp(payload)) assert fr._fetch_latest_release_tag("unslothai/llama.cpp") == "b9596-mix-e6f2453" + + +# reset_caches(drop_disk=...) -- post-update stale same-base mix disk cache. + + +def _seed_disk_cache(tmp_path: Path, latest_tag: str) -> Path: + # Matches _cache_path_for under the fixture's stubbed _cache_dir. + cache_dir = tmp_path / ".freshness" + cache_dir.mkdir(exist_ok = True) + cache_file = cache_dir / "unslothai__llama.cpp.json" + cache_file.write_text(json.dumps({"fetched_at": time.time(), "latest_tag": latest_tag})) + return cache_file + + +def test_reset_caches_drop_disk_removes_disk_cache(tmp_path): + cache_file = _seed_disk_cache(tmp_path, "b9596-mix-aaa") + assert cache_file.exists() + fr.reset_caches(drop_disk = True) + assert not cache_file.exists() + + +def test_reset_caches_default_keeps_disk_cache(tmp_path): + # The no-arg form is in-memory only (its existing test-only contract); it + # must not delete the on-disk cache. + cache_file = _seed_disk_cache(tmp_path, "b9596-mix-aaa") + fr.reset_caches() + assert cache_file.exists() + + +def test_reset_caches_drop_disk_on_missing_dir_is_noop(tmp_path): + # Fresh machine, no cache dir yet: drop_disk must be a quiet no-op. + assert not (tmp_path / ".freshness").exists() + fr.reset_caches(drop_disk = True) # must not raise + + +def test_drop_disk_lets_banner_fail_open_after_same_base_mix_swap(monkeypatch, tmp_path): + # P2 #2: the disk cache holds a still-fresh same-base mix (b9596-mix-aaa) + # from before an update to a *different* same-base mix (b9596-mix-bbb). + # The post-install path drops the disk cache; if the forced refresh is then + # offline, latest reads as None and the banner fails open -- instead of + # replaying the stale b9596-mix-aaa and falsely reading "behind". + _seed_disk_cache(tmp_path, "b9596-mix-aaa") + install_dir = tmp_path / "llama.cpp" + _write_marker( + install_dir, + tag = "b9596", + release_tag = "b9596-mix-bbb", + installed_at_utc = (datetime.now(tz = timezone.utc) - timedelta(days = 5)) + .isoformat() + .replace("+00:00", "Z"), + ) + bin_path = _fake_binary(install_dir, layout = "root") + # GitHub unreachable for the rest of the test (the offline post-install + # refresh, and the later status check). + monkeypatch.setattr(fr, "_fetch_latest_release_tag", lambda repo, timeout = 5.0: None) + + fr.reset_caches(drop_disk = True) # exactly what the apply path now does + info = fr.check_prebuilt_freshness(str(bin_path)) + assert info["latest_tag"] is None + assert info["behind"] is False + assert info["stale"] is False + + +def test_in_memory_only_reset_replays_stale_same_base_mix(monkeypatch, tmp_path): + # Contrast/guard for the case above: an in-memory-only reset leaves the + # stale same-base mix on disk, so an offline check replays it and falsely + # reads behind/stale. This is exactly the failure drop_disk removes; if a + # future change makes the no-arg reset also clear disk, the apply-path call + # and this guard should be revisited together. + _seed_disk_cache(tmp_path, "b9596-mix-aaa") + install_dir = tmp_path / "llama.cpp" + _write_marker( + install_dir, + tag = "b9596", + release_tag = "b9596-mix-bbb", + installed_at_utc = (datetime.now(tz = timezone.utc) - timedelta(days = 5)) + .isoformat() + .replace("+00:00", "Z"), + ) + bin_path = _fake_binary(install_dir, layout = "root") + monkeypatch.setattr(fr, "_fetch_latest_release_tag", lambda repo, timeout = 5.0: None) + + fr.reset_caches() # in-memory only -> stale disk value survives + info = fr.check_prebuilt_freshness(str(bin_path)) + assert info["latest_tag"] == "b9596-mix-aaa" + assert info["behind"] is True + assert info["stale"] is True diff --git a/studio/backend/tests/test_llama_cpp_tool_loop.py b/studio/backend/tests/test_llama_cpp_tool_loop.py index 3c121c281d..aacb029ff6 100644 --- a/studio/backend/tests/test_llama_cpp_tool_loop.py +++ b/studio/backend/tests/test_llama_cpp_tool_loop.py @@ -52,11 +52,16 @@ def _make_backend(monkeypatch, streams: list[list[str]], payloads: list[dict]): payload, _cancel_event, headers = None, + first_token_deadline = None, ): payloads.append(copy.deepcopy(payload)) yield type("FakeResponse", (), {"status_code": 200, "chunks": streams.pop(0)})() - def fake_iter_text_cancellable(response, _cancel_event): + def fake_iter_text_cancellable( + response, + _cancel_event, + first_token_deadline = None, + ): yield from response.chunks monkeypatch.setattr(backend, "_stream_with_retry", fake_stream_with_retry) diff --git a/studio/backend/tests/test_llama_cpp_update.py b/studio/backend/tests/test_llama_cpp_update.py index 3b8f511a7c..e207306c91 100644 --- a/studio/backend/tests/test_llama_cpp_update.py +++ b/studio/backend/tests/test_llama_cpp_update.py @@ -420,8 +420,8 @@ def test_start_update_installer_failure_reports_error(monkeypatch, tmp_path): # --- installer-argument construction (mirrors the post-#5963 setup scripts) --- -def test_rocm_install_args_lemonade_gfx(): - # Lemonade HIP app bundle: gfx family lives in the asset name. +def test_rocm_install_args_gfx_family(): + # Per-gfx ROCm bundle: gfx family lives in the asset name. assert upd._rocm_install_args("app-b9585-linux-x64-rocm-gfx110X.tar.gz") == [ "--rocm-gfx", "gfx110x", diff --git a/studio/backend/tests/test_llama_route_timeouts.py b/studio/backend/tests/test_llama_route_timeouts.py new file mode 100644 index 0000000000..5aee6198ba --- /dev/null +++ b/studio/backend/tests/test_llama_route_timeouts.py @@ -0,0 +1,87 @@ +# 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 time +from types import SimpleNamespace + +_backend = os.path.join(os.path.dirname(__file__), "..") +sys.path.insert(0, _backend) + +import routes.inference as inf_mod # noqa: E402 + + +def test_non_streaming_generation_timeout_has_read_deadline(): + timeout = inf_mod._llama_non_streaming_generation_timeout() + assert timeout.read == inf_mod._DEFAULT_FIRST_TOKEN_TIMEOUT_S + + +def test_stream_first_item_deadline_after_headers(): + async def _run(): + class _Never: + async def __anext__(self): + await asyncio.Future() + + started = time.monotonic() + try: + async for _ in inf_mod._aiter_llama_stream_items( + _Never(), + first_token_deadline = started + 0.02, + ): + pass + except inf_mod.httpx.ReadTimeout: + pass + else: + raise AssertionError("first item deadline did not fire") + assert time.monotonic() - started < 0.5 + + asyncio.run(_run()) + + +def test_preheader_send_cleanup_on_disconnect_and_cancel(): + async def _run(cancel_parent): + state = SimpleNamespace(disconnected = False, closed = False, cancelled = False) + started = asyncio.Event() + + class _Client: + async def send( + self, + req, + stream = False, + ): + started.set() + try: + await asyncio.Future() + except asyncio.CancelledError: + state.cancelled = True + raise + + async def aclose(self): + state.closed = True + + class _Request: + async def is_disconnected(self): + return state.disconnected + + task = asyncio.create_task( + inf_mod._send_stream_with_preheader_cancel(_Client(), object(), request = _Request()) + ) + await started.wait() + if cancel_parent: + task.cancel() + try: + await task + except asyncio.CancelledError: + pass + else: + raise AssertionError("helper cancellation did not propagate") + else: + state.disconnected = True + assert await task is None + assert state.closed + assert state.cancelled + + asyncio.run(_run(False)) + asyncio.run(_run(True)) diff --git a/studio/backend/tests/test_llama_server_args.py b/studio/backend/tests/test_llama_server_args.py index 6ae9d21e47..f3a3ea1ec4 100644 --- a/studio/backend/tests/test_llama_server_args.py +++ b/studio/backend/tests/test_llama_server_args.py @@ -25,8 +25,11 @@ _spec.loader.exec_module(_lsa) is_managed_flag = _lsa.is_managed_flag parse_cache_override = _lsa.parse_cache_override parse_ctx_override = _lsa.parse_ctx_override +parse_split_mode_override = _lsa.parse_split_mode_override resolve_cache_type_kv = _lsa.resolve_cache_type_kv +resolve_tensor_parallel = _lsa.resolve_tensor_parallel strip_shadowing_flags = _lsa.strip_shadowing_flags +strip_split_mode_only = _lsa.strip_split_mode_only extra_args_disable_mmproj = _lsa.extra_args_disable_mmproj validate_extra_args = _lsa.validate_extra_args @@ -510,6 +513,112 @@ def test_strip_shadowing_flags_defaults_strip_everything(): assert out == [] +# ── --split-mode (Tensor Parallelism toggle) ───────────────────────── +# Soft-shadowed exactly like --cache-type-*: pass-through allowed (keeps +# the row/none/layer modes the boolean toggle doesn't expose), stripped +# on inherit, and reconciled back into the round-tripped tensor_parallel +# state. + + +@pytest.mark.parametrize( + "args", + [ + ["--split-mode", "tensor"], + ["--split-mode", "row"], + ["--split-mode", "none"], + ["--split-mode", "layer"], + ["-sm", "tensor"], + ["--split-mode=row"], + ["-sm=tensor"], + ], +) +def test_split_mode_passes_through(args): + # Not denylisted -- a user keeps row/none/layer via extras. + assert validate_extra_args(args) == args + + +def test_split_mode_is_not_managed(): + assert is_managed_flag("--split-mode") is False + assert is_managed_flag("-sm") is False + + +@pytest.mark.parametrize( + "args,expected", + [ + (None, None), + ([], None), + (["--top-k", "20"], None), + (["--split-mode", "tensor"], "tensor"), + (["--split-mode", "row"], "row"), + (["-sm", "none"], "none"), + (["--split-mode=layer"], "layer"), + (["-sm=tensor"], "tensor"), + # last-wins when supplied twice + (["-sm", "row", "--split-mode", "tensor"], "tensor"), + ], +) +def test_parse_split_mode_override(args, expected): + assert parse_split_mode_override(args) == expected + + +@pytest.mark.parametrize( + "args", + [ + ["--split-mode"], + ["-sm"], + ["--split-mode", "-c", "4096"], # next token is a flag, not a value + ], +) +def test_parse_split_mode_override_rejects_malformed_values(args): + with pytest.raises(ValueError, match = "split-mode|'-sm'"): + parse_split_mode_override(args) + + +def test_validate_extra_args_rejects_malformed_split_mode(): + # Validation catches a value-less --split-mode at the boundary, + # mirroring the early --ctx-size / --cache-type checks. + with pytest.raises(ValueError, match = "split-mode"): + validate_extra_args(["--split-mode"]) + + +@pytest.mark.parametrize( + "args,fallback,expected", + [ + # No override -> fall back to the toggle value, both directions. + (["--top-k", "20"], True, True), + (["--top-k", "20"], False, False), + (None, True, True), + ([], False, False), + # Explicit override wins: tensor -> on, anything else -> off, + # regardless of the toggle fallback. + (["--split-mode", "tensor"], False, True), + (["-sm", "tensor"], False, True), + (["--split-mode", "row"], True, False), + (["--split-mode", "none"], True, False), + (["--split-mode", "layer"], True, False), + (["--split-mode=tensor"], False, True), + # Case-insensitive on the mode string. + (["--split-mode", "TENSOR"], False, True), + # last-wins across multiple --split-mode flags. + (["-sm", "tensor", "--split-mode", "row"], True, False), + ], +) +def test_resolve_tensor_parallel(args, fallback, expected): + assert resolve_tensor_parallel(args, fallback) is expected + + +def test_strip_shadowing_flags_drops_split_mode_when_requested(): + out = strip_shadowing_flags( + ["--split-mode", "row", "--top-k", "20"], + strip_context = False, + strip_cache = False, + strip_spec = False, + strip_template = False, + strip_split_mode = True, + ) + assert out == ["--top-k", "20"] + + def test_extra_args_disable_mmproj_detects_flag(): assert extra_args_disable_mmproj(["--no-mmproj"]) is True assert extra_args_disable_mmproj(["--threads", "12", "--no-mmproj"]) is True @@ -540,6 +649,86 @@ def test_strip_shadowing_flags_drops_model_draft_with_spec(): assert out == ["--top-k", "20"] +def test_strip_shadowing_flags_keeps_split_mode_when_not_requested(): + # No tensor_parallel field supplied on the Apply -> an inherited + # --split-mode survives (mirrors the chat-template keep behavior). + out = strip_shadowing_flags( + ["--split-mode", "row", "--top-k", "20"], + strip_context = True, + strip_cache = True, + strip_spec = True, + strip_template = True, + strip_split_mode = False, + ) + assert out == ["--split-mode", "row", "--top-k", "20"] + + +def test_strip_shadowing_flags_drops_split_mode_short_alias_and_equals(): + assert strip_shadowing_flags(["-sm", "tensor", "--top-k", "20"], strip_split_mode = True) == [ + "--top-k", + "20", + ] + assert strip_shadowing_flags(["--split-mode=row", "--seed", "-1"], strip_split_mode = True) == [ + "--seed", + "-1", + ] + + +def test_strip_shadowing_flags_defaults_strip_split_mode_too(): + # The route's already-loaded comparator (no kwargs) must see a stored + # --split-mode as a shadowing flag so it forces a reload. + assert strip_shadowing_flags(["--split-mode", "tensor"]) == [] + + +@pytest.mark.parametrize( + "args", + [ + ["--split-mode", "tensor", "-c", "4096"], + ["-sm", "tensor", "-c", "4096"], + ["--split-mode=tensor", "-c", "4096"], + ["-sm=tensor", "-c", "4096"], + ], +) +def test_strip_split_mode_only_keeps_other_shadow_flags(args): + # Every --split-mode form (long/short, space/=) is dropped; -c survives. + assert strip_split_mode_only(args) == ["-c", "4096"] + + +def test_strip_split_mode_only_preserves_none_and_empty(): + # None means "inherit"; [] means "explicit empty" -- both must round-trip. + assert strip_split_mode_only(None) is None + assert strip_split_mode_only([]) == [] + + +def test_strip_shadowing_flags_drops_tensor_split_with_split_mode(): + # --tensor-split is coupled to the split mode: stripped together so a stale + # ratio can't override Studio's computed tensor split. Other flags survive. + out = strip_shadowing_flags( + ["--split-mode", "row", "--tensor-split", "1,1", "--top-k", "20"], + strip_context = False, + strip_cache = False, + strip_spec = False, + strip_template = False, + strip_split_mode = True, + ) + assert out == ["--top-k", "20"] + + +def test_strip_shadowing_flags_keeps_tensor_split_when_not_requested(): + # strip_split_mode=False keeps the whole split group (mode + ratios). + assert strip_shadowing_flags( + ["--tensor-split", "1,1", "--top-k", "20"], strip_split_mode = False + ) == ["--tensor-split", "1,1", "--top-k", "20"] + + +def test_strip_split_mode_only_drops_tensor_split_too(): + # Downgrade / layer fallback must drop the coupled --tensor-split (all forms). + assert strip_split_mode_only( + ["--split-mode", "tensor", "--tensor-split", "1,1", "-c", "4096"] + ) == ["-c", "4096"] + assert strip_split_mode_only(["-sm=tensor", "-ts=3,1"]) == [] + + def test_strip_shadowing_flags_keeps_model_draft_without_spec(): out = strip_shadowing_flags( ["--model-draft", "/custom/mtp.gguf"], diff --git a/studio/backend/tests/test_mcp_servers.py b/studio/backend/tests/test_mcp_servers.py index ede3cf15d4..90b1ade03c 100644 --- a/studio/backend/tests/test_mcp_servers.py +++ b/studio/backend/tests/test_mcp_servers.py @@ -598,3 +598,614 @@ def test_safetensors_agentic_empty_allowlist_still_means_allow_all(): ) # Empty allow-list = run anything (preserved contract). assert calls == [("python", {"code": "1"})] or len(calls) >= 1 + + +# ── discovery cache ───────────────────────────────────────────────── + + +def _one_tool(name = "echo"): + return [{"name": name, "inputSchema": {"type": "object", "properties": {}}}] + + +def test_get_enabled_mcp_tools_caches_discovery(tmp_path, monkeypatch): + """A second send must serve tools from cache instead of re-probing.""" + import asyncio + + _reset_db(tmp_path, monkeypatch) + from core.inference import mcp_client + from core.inference import tools as tools_mod + + monkeypatch.setattr(mcp_client, "_tool_cache", {}) + mcp_servers_db.create_server(id = "s1", display_name = "A", url = "https://x/mcp", is_enabled = True) + + calls: list[str] = [] + + async def fake( + url, + headers = None, + timeout = None, + use_oauth = False, + ): + calls.append(url) + return _one_tool() + + monkeypatch.setattr(tools_mod, "list_tools_async", fake) + + first = asyncio.run(tools_mod.get_enabled_mcp_tools()) + second = asyncio.run(tools_mod.get_enabled_mcp_tools()) + + assert len(calls) == 1 # probed once, cache hit on the second send + assert [t["function"]["name"] for t in first] == ["mcp__s1__echo"] + assert first == second + + +def test_get_enabled_mcp_tools_does_not_cache_failures(tmp_path, monkeypatch): + """A failed probe isn't cached: once the cool-off elapses, it's retried.""" + import asyncio + + _reset_db(tmp_path, monkeypatch) + from core.inference import mcp_client + from core.inference import tools as tools_mod + + monkeypatch.setattr(mcp_client, "_tool_cache", {}) + monkeypatch.setattr(mcp_client, "_probe_cooloff_until", {}) + mcp_servers_db.create_server(id = "s1", display_name = "A", url = "https://x/mcp", is_enabled = True) + + attempts = {"n": 0} + + async def fake( + url, + headers = None, + timeout = None, + use_oauth = False, + ): + attempts["n"] += 1 + if attempts["n"] == 1: + raise RuntimeError("server down") + return _one_tool() + + monkeypatch.setattr(tools_mod, "list_tools_async", fake) + + assert asyncio.run(tools_mod.get_enabled_mcp_tools()) == [] # failure -> empty + # Expire the cool-off (an until-time in the past) so the server is retried. + mcp_client._probe_cooloff_until["s1"] = 0.0 + second = asyncio.run(tools_mod.get_enabled_mcp_tools()) + assert attempts["n"] == 2 # retried after the cool-off, not cached + assert [t["function"]["name"] for t in second] == ["mcp__s1__echo"] + + +def test_refresh_warms_tool_cache(tmp_path, monkeypatch): + """Clicking Refresh must populate the cache the chat path reads.""" + import asyncio + + _reset_db(tmp_path, monkeypatch) + from core.inference import mcp_client + from core.inference import tools as tools_mod + import routes.mcp_servers as routes_mcp + + monkeypatch.setattr(mcp_client, "_tool_cache", {}) + mcp_servers_db.create_server(id = "s1", display_name = "A", url = "https://x/mcp", is_enabled = True) + + async def fake_refresh( + url, + headers = None, + timeout = None, + use_oauth = False, + ): + return _one_tool() + + monkeypatch.setattr(routes_mcp, "list_tools_async", fake_refresh) + res = asyncio.run(routes_mcp.refresh_mcp_server_tools("s1", current_subject = "u")) + assert res.ok and res.tool_count == 1 + + def boom(*a, **k): + raise AssertionError("chat path re-probed despite a warm cache") + + monkeypatch.setattr(tools_mod, "list_tools_async", boom) + specs = asyncio.run(tools_mod.get_enabled_mcp_tools()) + assert [t["function"]["name"] for t in specs] == ["mcp__s1__echo"] + + +def test_update_url_evicts_tool_cache(tmp_path, monkeypatch): + """Re-pointing the URL must drop the old endpoint's cached tools.""" + import asyncio + + _reset_db(tmp_path, monkeypatch) + from core.inference import mcp_client + from models.mcp_servers import McpServerUpdate + import routes.mcp_servers as routes_mcp + + monkeypatch.setattr(mcp_client, "_tool_cache", {"s1": _one_tool("stale")}) + mcp_servers_db.create_server(id = "s1", display_name = "A", url = "https://old/mcp", is_enabled = True) + + asyncio.run( + routes_mcp.update_mcp_server( + "s1", McpServerUpdate(url = "https://new/mcp"), current_subject = "u" + ) + ) + assert mcp_client.get_cached_tools("s1") is None + + +def test_update_display_name_keeps_tool_cache(tmp_path, monkeypatch): + """A rename touches no endpoint, so the cache must survive it.""" + import asyncio + + _reset_db(tmp_path, monkeypatch) + from core.inference import mcp_client + from models.mcp_servers import McpServerUpdate + import routes.mcp_servers as routes_mcp + + cached = _one_tool() + monkeypatch.setattr(mcp_client, "_tool_cache", {"s1": cached}) + mcp_servers_db.create_server(id = "s1", display_name = "A", url = "https://x/mcp", is_enabled = True) + + asyncio.run( + routes_mcp.update_mcp_server("s1", McpServerUpdate(display_name = "B"), current_subject = "u") + ) + assert mcp_client.get_cached_tools("s1") == cached + + +def test_update_disable_evicts_tool_cache(tmp_path, monkeypatch): + """Disabling a server must drop its cached tools, not leave them unread.""" + import asyncio + + _reset_db(tmp_path, monkeypatch) + from core.inference import mcp_client + from models.mcp_servers import McpServerUpdate + import routes.mcp_servers as routes_mcp + + monkeypatch.setattr(mcp_client, "_tool_cache", {"s1": _one_tool()}) + mcp_servers_db.create_server(id = "s1", display_name = "A", url = "https://x/mcp", is_enabled = True) + + asyncio.run( + routes_mcp.update_mcp_server("s1", McpServerUpdate(is_enabled = False), current_subject = "u") + ) + assert mcp_client.get_cached_tools("s1") is None + + +def test_delete_evicts_tool_cache(tmp_path, monkeypatch): + """Deleting a server must not leave its tools cached.""" + import asyncio + + _reset_db(tmp_path, monkeypatch) + from core.inference import mcp_client + import routes.mcp_servers as routes_mcp + + monkeypatch.setattr(mcp_client, "_tool_cache", {"s1": _one_tool()}) + mcp_servers_db.create_server(id = "s1", display_name = "A", url = "https://x/mcp", is_enabled = True) + asyncio.run(routes_mcp.delete_mcp_server("s1", current_subject = "u")) + assert mcp_client.get_cached_tools("s1") is None + + +def test_invalidate_tool_cache_clears_all(monkeypatch): + from core.inference import mcp_client + + monkeypatch.setattr(mcp_client, "_tool_cache", {"a": _one_tool(), "b": _one_tool()}) + mcp_client.invalidate_tool_cache() + assert mcp_client.get_cached_tools("a") is None + assert mcp_client.get_cached_tools("b") is None + + +def test_get_enabled_mcp_tools_probes_only_uncached(tmp_path, monkeypatch): + """An already-cached server must not be re-probed alongside a cold one.""" + import asyncio + + _reset_db(tmp_path, monkeypatch) + from core.inference import mcp_client + from core.inference import tools as tools_mod + + monkeypatch.setattr(mcp_client, "_tool_cache", {"s1": _one_tool("cached")}) + mcp_servers_db.create_server(id = "s1", display_name = "A", url = "https://a/mcp", is_enabled = True) + mcp_servers_db.create_server(id = "s2", display_name = "B", url = "https://b/mcp", is_enabled = True) + + probed: list[str] = [] + + async def fake( + url, + headers = None, + timeout = None, + use_oauth = False, + ): + probed.append(url) + return _one_tool("fresh") + + monkeypatch.setattr(tools_mod, "list_tools_async", fake) + + specs = asyncio.run(tools_mod.get_enabled_mcp_tools()) + assert probed == ["https://b/mcp"] # only the uncached server is probed + assert sorted(t["function"]["name"] for t in specs) == ["mcp__s1__cached", "mcp__s2__fresh"] + + +def test_get_enabled_mcp_tools_partial_failure_caches_healthy(tmp_path, monkeypatch): + """One server failing must not stop the others from being cached/served.""" + import asyncio + + _reset_db(tmp_path, monkeypatch) + from core.inference import mcp_client + from core.inference import tools as tools_mod + + monkeypatch.setattr(mcp_client, "_tool_cache", {}) + monkeypatch.setattr(mcp_client, "_probe_cooloff_until", {}) + mcp_servers_db.create_server(id = "s1", display_name = "A", url = "https://bad/mcp", is_enabled = True) + mcp_servers_db.create_server(id = "s2", display_name = "B", url = "https://good/mcp", is_enabled = True) + + async def fake( + url, + headers = None, + timeout = None, + use_oauth = False, + ): + if "bad" in url: + raise RuntimeError("down") + return _one_tool("ok") + + monkeypatch.setattr(tools_mod, "list_tools_async", fake) + + specs = asyncio.run(tools_mod.get_enabled_mcp_tools()) + assert [t["function"]["name"] for t in specs] == ["mcp__s2__ok"] + assert mcp_client.get_cached_tools("s1") is None # failure not cached + assert mcp_client.get_cached_tools("s2") == _one_tool("ok") # healthy cached + + +def test_get_enabled_mcp_tools_caches_empty_tool_list(tmp_path, monkeypatch): + """A server exposing zero tools is cached as [] (a hit), not re-probed.""" + import asyncio + + _reset_db(tmp_path, monkeypatch) + from core.inference import mcp_client + from core.inference import tools as tools_mod + + monkeypatch.setattr(mcp_client, "_tool_cache", {}) + mcp_servers_db.create_server(id = "s1", display_name = "A", url = "https://x/mcp", is_enabled = True) + + calls: list[str] = [] + + async def fake( + url, + headers = None, + timeout = None, + use_oauth = False, + ): + calls.append(url) + return [] + + monkeypatch.setattr(tools_mod, "list_tools_async", fake) + + assert asyncio.run(tools_mod.get_enabled_mcp_tools()) == [] + assert asyncio.run(tools_mod.get_enabled_mcp_tools()) == [] + assert len(calls) == 1 # [] is a cache hit, not re-probed every send + assert mcp_client.get_cached_tools("s1") == [] + + +def test_update_headers_evicts_tool_cache(tmp_path, monkeypatch): + """Changing auth headers must drop tools discovered under the old headers.""" + import asyncio + + _reset_db(tmp_path, monkeypatch) + from core.inference import mcp_client + from models.mcp_servers import McpServerUpdate + import routes.mcp_servers as routes_mcp + + monkeypatch.setattr(mcp_client, "_tool_cache", {"s1": _one_tool()}) + mcp_servers_db.create_server(id = "s1", display_name = "A", url = "https://x/mcp", is_enabled = True) + + asyncio.run( + routes_mcp.update_mcp_server( + "s1", + McpServerUpdate(headers = {"Authorization": "Bearer new"}), + current_subject = "u", + ) + ) + assert mcp_client.get_cached_tools("s1") is None + + +def test_get_enabled_mcp_tools_skips_cache_when_config_changes_mid_probe(tmp_path, monkeypatch): + """A config edit landing during an in-flight probe must not be clobbered + by the now-stale probe result (TOCTOU on the cache write).""" + import asyncio + + _reset_db(tmp_path, monkeypatch) + from core.inference import mcp_client + from core.inference import tools as tools_mod + + monkeypatch.setattr(mcp_client, "_tool_cache", {}) + mcp_servers_db.create_server(id = "s1", display_name = "A", url = "https://old/mcp", is_enabled = True) + + async def fake( + url, + headers = None, + timeout = None, + use_oauth = False, + ): + # Simulate a PUT landing while we are awaiting the probe. + mcp_servers_db.update_server("s1", {"url": "https://new/mcp"}) + return _one_tool() + + monkeypatch.setattr(tools_mod, "list_tools_async", fake) + + specs = asyncio.run(tools_mod.get_enabled_mcp_tools()) + assert specs == [] # stale result is neither served... + assert mcp_client.get_cached_tools("s1") is None # ...nor cached + + +def test_get_enabled_mcp_tools_no_cooloff_when_config_changes_mid_failed_probe( + tmp_path, monkeypatch +): + """An edit landing while a probe of the OLD config is failing must not park + a cool-off on the now-fresh config -- else the re-pointed server the user + just fixed is needlessly skipped for the whole cool-off window.""" + import asyncio + + _reset_db(tmp_path, monkeypatch) + from core.inference import mcp_client + from core.inference import tools as tools_mod + + monkeypatch.setattr(mcp_client, "_tool_cache", {}) + monkeypatch.setattr(mcp_client, "_probe_cooloff_until", {}) + mcp_servers_db.create_server(id = "s1", display_name = "A", url = "https://old/mcp", is_enabled = True) + + async def fake( + url, + headers = None, + timeout = None, + use_oauth = False, + ): + # The user re-points the server while the old endpoint's probe fails. + mcp_servers_db.update_server("s1", {"url": "https://new/mcp"}) + raise RuntimeError("old endpoint down") + + monkeypatch.setattr(tools_mod, "list_tools_async", fake) + + assert asyncio.run(tools_mod.get_enabled_mcp_tools()) == [] + # The failure was for the OLD config, so the new one must stay re-probable. + assert not mcp_client.in_failure_cooloff("s1") + + +def test_get_enabled_mcp_tools_no_cooloff_when_server_deleted_mid_failed_probe( + tmp_path, monkeypatch +): + """A delete landing while a probe fails must not leave an orphan cool-off + entry keyed by the since-removed server id.""" + import asyncio + + _reset_db(tmp_path, monkeypatch) + from core.inference import mcp_client + from core.inference import tools as tools_mod + + monkeypatch.setattr(mcp_client, "_tool_cache", {}) + monkeypatch.setattr(mcp_client, "_probe_cooloff_until", {}) + mcp_servers_db.create_server(id = "s1", display_name = "A", url = "https://x/mcp", is_enabled = True) + + async def fake( + url, + headers = None, + timeout = None, + use_oauth = False, + ): + mcp_servers_db.delete_server("s1") + raise RuntimeError("down") + + monkeypatch.setattr(tools_mod, "list_tools_async", fake) + + assert asyncio.run(tools_mod.get_enabled_mcp_tools()) == [] + assert "s1" not in mcp_client._probe_cooloff_until # no orphan cool-off + + +def test_get_enabled_mcp_tools_skips_failed_server_during_cooloff(tmp_path, monkeypatch): + """A down server is probed once, then skipped during the cool-off instead + of being re-probed (and re-hung) on every send.""" + import asyncio + + _reset_db(tmp_path, monkeypatch) + from core.inference import mcp_client + from core.inference import tools as tools_mod + + monkeypatch.setattr(mcp_client, "_tool_cache", {}) + monkeypatch.setattr(mcp_client, "_probe_cooloff_until", {}) + mcp_servers_db.create_server(id = "s1", display_name = "A", url = "https://x/mcp", is_enabled = True) + + attempts = {"n": 0} + + async def fake( + url, + headers = None, + timeout = None, + use_oauth = False, + ): + attempts["n"] += 1 + raise RuntimeError("down") + + monkeypatch.setattr(tools_mod, "list_tools_async", fake) + + assert asyncio.run(tools_mod.get_enabled_mcp_tools()) == [] # probes, fails + assert asyncio.run(tools_mod.get_enabled_mcp_tools()) == [] # within cool-off + assert asyncio.run(tools_mod.get_enabled_mcp_tools()) == [] # still skipped + assert attempts["n"] == 1 # only the first send probed + + +def test_cache_tools_clears_failure_cooloff(monkeypatch): + """A successful probe lifts a server's failure cool-off.""" + from core.inference import mcp_client + + monkeypatch.setattr(mcp_client, "_tool_cache", {}) + monkeypatch.setattr(mcp_client, "_probe_cooloff_until", {}) + mcp_client.record_probe_failure("s1") + assert mcp_client.in_failure_cooloff("s1") + mcp_client.cache_tools("s1", _one_tool()) + assert not mcp_client.in_failure_cooloff("s1") + + +def test_oauth_failure_cools_off_longer_than_plain(monkeypatch): + """An OAuth server's failure cools off longer than a plain server's, so its + multi-minute probe hang doesn't recur every minute.""" + from core.inference import mcp_client + + monkeypatch.setattr(mcp_client, "_probe_cooloff_until", {}) + mcp_client.record_probe_failure("plain", use_oauth = False) + mcp_client.record_probe_failure("oauth", use_oauth = True) + assert mcp_client._probe_cooloff_until["oauth"] > mcp_client._probe_cooloff_until["plain"] + + +def test_invalidate_clears_failure_cooloff(monkeypatch): + """Eviction drops the failure cool-off so an edited server re-probes at once.""" + from core.inference import mcp_client + + monkeypatch.setattr(mcp_client, "_tool_cache", {}) + monkeypatch.setattr(mcp_client, "_probe_cooloff_until", {"s1": 1.0, "s2": 2.0}) + mcp_client.invalidate_tool_cache("s1") + assert "s1" not in mcp_client._probe_cooloff_until + assert "s2" in mcp_client._probe_cooloff_until + mcp_client.invalidate_tool_cache() + assert mcp_client._probe_cooloff_until == {} + + +def test_refresh_failure_records_cooloff(tmp_path, monkeypatch): + """A failed manual refresh starts the cool-off so the next chat send does + not immediately hang on the down server.""" + import asyncio + + _reset_db(tmp_path, monkeypatch) + from core.inference import mcp_client + import routes.mcp_servers as routes_mcp + + monkeypatch.setattr(mcp_client, "_tool_cache", {}) + monkeypatch.setattr(mcp_client, "_probe_cooloff_until", {}) + mcp_servers_db.create_server(id = "s1", display_name = "A", url = "https://x/mcp", is_enabled = True) + + async def boom( + url, + headers = None, + timeout = None, + use_oauth = False, + ): + raise RuntimeError("down") + + monkeypatch.setattr(routes_mcp, "list_tools_async", boom) + res = asyncio.run(routes_mcp.refresh_mcp_server_tools("s1", current_subject = "u")) + assert res.ok is False + assert mcp_client.in_failure_cooloff("s1") + + +def test_refresh_drops_result_when_config_changes_mid_probe(tmp_path, monkeypatch): + """A manual refresh must not warm the chat cache with tools discovered + under an old config if the server is edited while the probe is in flight.""" + import asyncio + + _reset_db(tmp_path, monkeypatch) + from core.inference import mcp_client + import routes.mcp_servers as routes_mcp + + monkeypatch.setattr(mcp_client, "_tool_cache", {}) + mcp_servers_db.create_server(id = "s1", display_name = "A", url = "https://old/mcp", is_enabled = True) + + async def fake_refresh( + url, + headers = None, + timeout = None, + use_oauth = False, + ): + mcp_servers_db.update_server("s1", {"url": "https://new/mcp"}) + return _one_tool("stale") + + monkeypatch.setattr(routes_mcp, "list_tools_async", fake_refresh) + res = asyncio.run(routes_mcp.refresh_mcp_server_tools("s1", current_subject = "u")) + assert res.ok and res.tool_count == 1 + assert mcp_client.get_cached_tools("s1") is None + + +def test_refresh_failure_no_cooloff_when_config_changes_mid_probe(tmp_path, monkeypatch): + """A manual refresh failure for an old config must not cool off the freshly + edited server.""" + import asyncio + + _reset_db(tmp_path, monkeypatch) + from core.inference import mcp_client + import routes.mcp_servers as routes_mcp + + monkeypatch.setattr(mcp_client, "_tool_cache", {}) + monkeypatch.setattr(mcp_client, "_probe_cooloff_until", {}) + mcp_servers_db.create_server(id = "s1", display_name = "A", url = "https://old/mcp", is_enabled = True) + + async def boom( + url, + headers = None, + timeout = None, + use_oauth = False, + ): + mcp_servers_db.update_server("s1", {"url": "https://new/mcp"}) + raise RuntimeError("old endpoint down") + + monkeypatch.setattr(routes_mcp, "list_tools_async", boom) + res = asyncio.run(routes_mcp.refresh_mcp_server_tools("s1", current_subject = "u")) + assert res.ok is False + assert not mcp_client.in_failure_cooloff("s1") + + +def test_get_enabled_mcp_tools_drops_result_when_server_deleted_mid_probe(tmp_path, monkeypatch): + """A delete landing while a probe is in flight must drop the now-orphan + result -- the `fresh is None` arm of the mid-probe TOCTOU guard. The + result is neither served nor cached under the since-removed id.""" + import asyncio + + _reset_db(tmp_path, monkeypatch) + from core.inference import mcp_client + from core.inference import tools as tools_mod + + monkeypatch.setattr(mcp_client, "_tool_cache", {}) + monkeypatch.setattr(mcp_client, "_probe_cooloff_until", {}) + mcp_servers_db.create_server(id = "s1", display_name = "A", url = "https://x/mcp", is_enabled = True) + + async def fake( + url, + headers = None, + timeout = None, + use_oauth = False, + ): + # Simulate a DELETE landing while we await the probe. + mcp_servers_db.delete_server("s1") + return _one_tool() + + monkeypatch.setattr(tools_mod, "list_tools_async", fake) + + specs = asyncio.run(tools_mod.get_enabled_mcp_tools()) + assert specs == [] # orphan result not served + assert mcp_client.get_cached_tools("s1") is None # nor cached under a gone id + + +def test_oauth_probe_failure_in_chat_path_uses_long_cooloff(tmp_path, monkeypatch): + """When an OAuth server fails discovery during a send, the chat path must + record the OAuth (long) cool-off, not the plain one -- otherwise its + multi-minute browser hang recurs every minute.""" + import asyncio + import time + + _reset_db(tmp_path, monkeypatch) + from core.inference import mcp_client + from core.inference import tools as tools_mod + + monkeypatch.setattr(mcp_client, "_tool_cache", {}) + monkeypatch.setattr(mcp_client, "_probe_cooloff_until", {}) + mcp_servers_db.create_server( + id = "s1", + display_name = "A", + url = "https://x/mcp", + is_enabled = True, + use_oauth = True, + ) + + async def boom( + url, + headers = None, + timeout = None, + use_oauth = False, + ): + raise RuntimeError("oauth down") + + monkeypatch.setattr(tools_mod, "list_tools_async", boom) + + assert asyncio.run(tools_mod.get_enabled_mcp_tools()) == [] + assert mcp_client.in_failure_cooloff("s1") + # The recorded window must exceed the plain cool-off, proving the OAuth + # branch (use_oauth=True) fired -- not the 60 s default. + remaining = mcp_client._probe_cooloff_until["s1"] - time.monotonic() + assert remaining > mcp_client.FAILED_PROBE_COOLOFF_SECONDS diff --git a/studio/backend/tests/test_mcp_stdio_pr5863.py b/studio/backend/tests/test_mcp_stdio_pr5863.py index c6a4898d5d..9a3e8d6882 100644 --- a/studio/backend/tests/test_mcp_stdio_pr5863.py +++ b/studio/backend/tests/test_mcp_stdio_pr5863.py @@ -19,6 +19,10 @@ from storage import mcp_servers_db def _reset_db(tmp_path, monkeypatch): monkeypatch.setenv("UNSLOTH_STUDIO_HOME", str(tmp_path)) monkeypatch.setattr(mcp_servers_db, "_schema_ready", False) + # The discovered-tool cache is process-global and keyed by server id; tests + # reuse "stdio1", so clear it (and the failure cool-off) for isolation — + # otherwise a prior test's warm cache makes discovery skip its probe. + mcp_client.invalidate_tool_cache() def _enable(monkeypatch): diff --git a/studio/backend/tests/test_mmproj_vram_accounting.py b/studio/backend/tests/test_mmproj_vram_accounting.py new file mode 100644 index 0000000000..bee289f183 --- /dev/null +++ b/studio/backend/tests/test_mmproj_vram_accounting.py @@ -0,0 +1,40 @@ +# SPDX-License-Identifier: AGPL-3.0-only +# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0 + +"""Tests for mmproj VRAM accounting in GGUF fit budgeting (#5825).""" + +from __future__ import annotations + +from pathlib import Path + +from core.inference.llama_cpp import LlamaCppBackend + + +def _write(path: Path, n_bytes: int) -> Path: + path.parent.mkdir(parents = True, exist_ok = True) + path.write_bytes(b"\x00" * n_bytes) + return path + + +def _backend() -> LlamaCppBackend: + return LlamaCppBackend.__new__(LlamaCppBackend) + + +def test_counts_resolved_projector_size(tmp_path: Path): + mmproj = _write(tmp_path / "Qwen3.5-9B-BF16-mmproj.gguf", 1024) + + got = _backend()._mmproj_vram_bytes(str(mmproj)) + + assert got == 1024 + + +def test_zero_when_no_projector_resolved(tmp_path: Path): + assert _backend()._mmproj_vram_bytes(None) == 0 + + +def test_zero_when_projector_missing_on_disk(tmp_path: Path): + missing = tmp_path / "Qwen3.5-9B-BF16-mmproj.gguf" # never created + + got = _backend()._mmproj_vram_bytes(str(missing)) + + assert got == 0 diff --git a/studio/backend/tests/test_openai_tool_passthrough.py b/studio/backend/tests/test_openai_tool_passthrough.py index 79e8c9ae49..b54f4c130d 100644 --- a/studio/backend/tests/test_openai_tool_passthrough.py +++ b/studio/backend/tests/test_openai_tool_passthrough.py @@ -1,13 +1,7 @@ # SPDX-License-Identifier: AGPL-3.0-only # Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. -"""Tests for the OpenAI /v1/chat/completions client-side tool pass-through. - -Covers ChatMessage tool/assistant roles, ChatCompletionRequest tool fields and -extra="allow", anthropic_tool_choice_to_openai, _build_passthrough_payload -tool_choice propagation, and _friendly_error's httpx-to-"Lost connection" -mapping. No server or GPU required. -""" +"""Tests for the OpenAI /v1/chat/completions client-side tool pass-through.""" import os import sys @@ -28,11 +22,13 @@ from models.inference import ( ChatMessage, CompletionChoice, CompletionMessage, + ResponsesRequest, ) from core.inference.anthropic_compat import ( anthropic_tool_choice_to_openai, ) from routes.inference import ( + _build_chat_request, _build_openai_passthrough_body, _build_passthrough_payload, _clamp_finish_reason, @@ -786,6 +782,22 @@ class TestPassthroughReasoningKwargs: ) assert body["chat_template_kwargs"] == {"reasoning_effort": "high"} + def test_reasoning_effort_none_forwarded_for_effort_style_models(self): + body = _build_openai_passthrough_body( + self._payload(enable_thinking = False, reasoning_effort = "none"), + backend_ctx = 4096, + llama_backend = _reasoning_backend(reasoning_style = "reasoning_effort"), + ) + assert body["chat_template_kwargs"] == {"reasoning_effort": "none"} + + def test_reasoning_effort_minimal_maps_to_low_for_effort_style_models(self): + body = _build_openai_passthrough_body( + self._payload(enable_thinking = True, reasoning_effort = "minimal"), + backend_ctx = 4096, + llama_backend = _reasoning_backend(reasoning_style = "reasoning_effort"), + ) + assert body["chat_template_kwargs"] == {"reasoning_effort": "low"} + def test_enable_thinking_maps_to_effort_for_effort_style_models(self): body = _build_openai_passthrough_body( self._payload(enable_thinking = False), @@ -896,12 +908,6 @@ class TestOpenAICompatibilityHelpers: class TestFriendlyErrorHttpx: - """When llama-server is down, httpx RequestError strings lack the - "Lost connection to llama-server" substring the sync path keys off, so the - old substring-only `_friendly_error` returned a useless generic message. - These tests pin the new isinstance-based mapping. - """ - def _req(self): return httpx.Request("POST", "http://127.0.0.1:65535/v1/chat/completions") @@ -919,7 +925,7 @@ class TestFriendlyErrorHttpx: def test_read_timeout_mapped(self): exc = httpx.ReadTimeout("timed out", request = self._req()) - assert "Lost connection" in _friendly_error(exc) + assert "first token within 20 minutes" in _friendly_error(exc) def test_non_httpx_unchanged(self): # Non-httpx exceptions still fall through to the substring heuristics @@ -1396,3 +1402,47 @@ class TestGgufVisionToolRouting: assert seen_seeds == expected assert [choice["index"] for choice in body["choices"]] == [0, 1, 2] + + +# ===================================================================== +# Responses API -> Chat Completions translation: chat_template_kwargs +# (e.g. {"enable_thinking": true}) sent via the Responses extra-body must +# reach the built ChatCompletionRequest's typed ``enable_thinking`` field, +# otherwise /v1/responses silently ignores reasoning control (issue #6198). +# ===================================================================== + + +class TestResponsesChatTemplateKwargs: + _messages = [ChatMessage(role = "user", content = "What is 100 - 67?")] + + def test_enable_thinking_lifted_from_extra_body(self): + payload = ResponsesRequest( + model = "qwen-local", + input = "What is 100 - 67?", + chat_template_kwargs = {"enable_thinking": True}, + ) + chat_req = _build_chat_request(payload, self._messages, stream = False) + assert chat_req.enable_thinking is True + + def test_enable_thinking_false_lifted_from_extra_body(self): + payload = ResponsesRequest( + model = "qwen-local", + input = "hi", + chat_template_kwargs = {"enable_thinking": False}, + ) + chat_req = _build_chat_request(payload, self._messages, stream = True) + assert chat_req.enable_thinking is False + + def test_no_chat_template_kwargs_leaves_enable_thinking_unset(self): + payload = ResponsesRequest(model = "qwen-local", input = "hi") + chat_req = _build_chat_request(payload, self._messages, stream = False) + assert chat_req.enable_thinking is None + + def test_chat_template_kwargs_without_enable_thinking_is_ignored(self): + payload = ResponsesRequest( + model = "qwen-local", + input = "hi", + chat_template_kwargs = {"some_other_flag": True}, + ) + chat_req = _build_chat_request(payload, self._messages, stream = False) + assert chat_req.enable_thinking is None diff --git a/studio/backend/tests/test_rag_ingestion.py b/studio/backend/tests/test_rag_ingestion.py index c5fe15876d..f0b71bc23b 100644 --- a/studio/backend/tests/test_rag_ingestion.py +++ b/studio/backend/tests/test_rag_ingestion.py @@ -39,7 +39,7 @@ def test_ingestion_lifecycle_pending_to_completed(rag_home, stub_embeddings, tmp conn = rag_db.get_connection() try: - assert store.get_document(conn, doc_id)["status"] == "pending" + assert store.get_document(conn, doc_id)["status"] in {"pending", "running", "completed"} finally: conn.close() @@ -83,6 +83,133 @@ def test_ingestion_dedupe_by_hash(rag_home, stub_embeddings, tmp_path): conn.close() +def test_ingestion_dedupe_removes_duplicate_upload(rag_home, stub_embeddings): + from utils.paths import ensure_dir, rag_uploads_root + + uploads = ensure_dir(rag_uploads_root()) + first_path = uploads / "doc.txt" + duplicate_path = uploads / "copy.txt" + first_path.write_text("alpha bravo charlie", encoding = "utf-8") + duplicate_path.write_text("alpha bravo charlie", encoding = "utf-8") + scope = store.project_scope("P1") + + doc_id, job_id = ingestion.start_ingestion( + scope, + None, + None, + "doc.txt", + str(first_path), + project_id = "P1", + ) + _drain(job_id) + _wait_completed(job_id) + + doc_id2, job_id2 = ingestion.start_ingestion( + scope, + None, + None, + "copy.txt", + str(duplicate_path), + project_id = "P1", + ) + events = _drain(job_id2) + assert doc_id2 == doc_id + assert any(e.get("deduped") for e in events) + assert first_path.exists() + assert not duplicate_path.exists() + + +def test_ingestion_retry_replaces_failed_hash(rag_home, stub_embeddings): + from utils.paths import ensure_dir, rag_uploads_root + + uploads = ensure_dir(rag_uploads_root()) + old_path = uploads / "failed.txt" + retry_path = uploads / "retry.txt" + old_path.write_text("alpha bravo charlie", encoding = "utf-8") + retry_path.write_text("alpha bravo charlie", encoding = "utf-8") + scope = store.project_scope("P1") + sha = ingestion._sha256_file(str(old_path)) + + conn = rag_db.get_connection() + try: + failed_id = store.create_document( + conn, + scope = scope, + filename = "failed.txt", + sha256 = sha, + project_id = "P1", + status = "failed", + stored_path = str(old_path), + ) + finally: + conn.close() + + doc_id, job_id = ingestion.start_ingestion( + scope, + None, + None, + "retry.txt", + str(retry_path), + project_id = "P1", + ) + events = _drain(job_id) + assert doc_id != failed_id + assert not any(e.get("deduped") for e in events) + assert not old_path.exists() + assert retry_path.exists() + + status = _wait_completed(job_id) + assert status["status"] == "completed" + conn = rag_db.get_connection() + try: + assert store.get_document(conn, failed_id) is None + assert store.get_document(conn, doc_id)["status"] == "completed" + finally: + conn.close() + + +def test_delete_document_route_removes_stored_upload(rag_home): + from fastapi import FastAPI + from fastapi.testclient import TestClient + + from auth.authentication import get_current_subject + from routes.rag import router + from utils.paths import ensure_dir, rag_uploads_root + + upload = ensure_dir(rag_uploads_root()) / "delete-me.txt" + upload.write_text("alpha bravo", encoding = "utf-8") + scope = store.project_scope("P1") + + conn = rag_db.get_connection() + try: + doc_id = store.create_document( + conn, + scope = scope, + filename = "delete-me.txt", + sha256 = "delete-route-sha", + project_id = "P1", + status = "completed", + stored_path = str(upload), + ) + finally: + conn.close() + + app = FastAPI() + app.include_router(router, prefix = "/api/rag") + app.dependency_overrides[get_current_subject] = lambda: "tester" + client = TestClient(app) + + res = client.delete(f"/api/rag/documents/{doc_id}") + assert res.status_code == 200 + assert not upload.exists() + + conn = rag_db.get_connection() + try: + assert store.get_document(conn, doc_id) is None + finally: + conn.close() + + def test_ingestion_delete_removes_all_rows(rag_home, stub_embeddings, tmp_path): path = _write(tmp_path, "doc.txt", "alpha bravo charlie delta") scope = store.kb_scope("K1") diff --git a/studio/backend/tests/test_responses_tool_passthrough.py b/studio/backend/tests/test_responses_tool_passthrough.py index 69fb0a78c2..ae7ff729bd 100644 --- a/studio/backend/tests/test_responses_tool_passthrough.py +++ b/studio/backend/tests/test_responses_tool_passthrough.py @@ -36,6 +36,7 @@ import json import httpx import pytest +from fastapi.responses import JSONResponse from pydantic import ValidationError from models.inference import ( @@ -46,6 +47,7 @@ from models.inference import ( ResponsesInputMessage, ResponsesOutputFunctionCall, ResponsesOutputMessage, + ResponsesOutputReasoning, ResponsesOutputTextContent, ResponsesOutputTextPart, ResponsesRequest, @@ -59,6 +61,7 @@ from routes.inference import ( _chat_tool_calls_to_responses_output, _normalise_responses_input, _responses_tool_output_text, + _responses_non_streaming, _responses_stream, _translate_responses_tool_choice_to_chat, _translate_responses_tools_to_chat, @@ -284,6 +287,59 @@ class TestBuildChatRequest: assert chat_req.parallel_tool_calls is False + def test_chat_template_kwargs_enable_thinking_true_is_lifted(self): + payload = ResponsesRequest( + input = "hi", + chat_template_kwargs = {"enable_thinking": True}, + ) + messages = [ChatMessage(role = "user", content = "hi")] + + chat_req = _build_chat_request(payload, messages, stream = False) + + assert chat_req.enable_thinking is True + + def test_chat_template_kwargs_enable_thinking_false_is_lifted(self): + payload = ResponsesRequest( + input = "hi", + chat_template_kwargs = {"enable_thinking": False}, + ) + messages = [ChatMessage(role = "user", content = "hi")] + + chat_req = _build_chat_request(payload, messages, stream = False) + + assert chat_req.enable_thinking is False + + def test_reasoning_effort_high_enables_local_thinking(self): + payload = ResponsesRequest(input = "hi", reasoning = {"effort": "high"}) + messages = [ChatMessage(role = "user", content = "hi")] + + chat_req = _build_chat_request(payload, messages, stream = False) + + assert chat_req.reasoning_effort == "high" + assert chat_req.enable_thinking is True + + def test_reasoning_effort_none_disables_local_thinking(self): + payload = ResponsesRequest(input = "hi", reasoning = {"effort": "none"}) + messages = [ChatMessage(role = "user", content = "hi")] + + chat_req = _build_chat_request(payload, messages, stream = False) + + assert chat_req.reasoning_effort == "none" + assert chat_req.enable_thinking is False + + def test_explicit_enable_thinking_false_disables_reasoning_effort(self): + payload = ResponsesRequest( + input = "hi", + reasoning = {"effort": "high"}, + chat_template_kwargs = {"enable_thinking": False}, + ) + messages = [ChatMessage(role = "user", content = "hi")] + + chat_req = _build_chat_request(payload, messages, stream = False) + + assert chat_req.reasoning_effort == "none" + assert chat_req.enable_thinking is False + # ===================================================================== # _normalise_responses_input — multi-turn tool mapping @@ -544,6 +600,119 @@ class TestChatToolCallsToResponsesOutput: assert items[0]["arguments"] == "" +# ===================================================================== +# Non-streaming Responses adapter +# ===================================================================== + + +class TestResponsesNonStreamingAdapter: + class _Request: + pass + + @staticmethod + def _run_with_message( + monkeypatch, + message, + payload = None, + llama_backend = None, + ): + import routes.inference as inf_mod + + async def fake_chat_completions(chat_req, request): + return JSONResponse( + content = { + "model": "test-model", + "choices": [{"message": message}], + "usage": {"prompt_tokens": 2, "completion_tokens": 3}, + } + ) + + monkeypatch.setattr(inf_mod, "openai_chat_completions", fake_chat_completions) + if llama_backend is not None: + monkeypatch.setattr(inf_mod, "get_llama_cpp_backend", lambda: llama_backend) + payload = payload or ResponsesRequest(input = "hi") + messages = [ChatMessage(role = "user", content = "hi")] + + async def run(): + response = await _responses_non_streaming( + payload, messages, TestResponsesNonStreamingAdapter._Request() + ) + return json.loads(response.body.decode()) + + return asyncio.run(run()) + + def test_think_block_becomes_reasoning_item_before_message(self, monkeypatch): + payload = ResponsesRequest(input = "hi", reasoning = {"effort": "high"}) + body = self._run_with_message( + monkeypatch, + {"content": "plan33"}, + payload = payload, + ) + + assert [item["type"] for item in body["output"]] == ["reasoning", "message"] + assert body["output"][0]["content"] == [{"type": "reasoning_text", "text": "plan"}] + assert body["output"][0]["summary"] == [] + assert body["output"][1]["content"][0]["text"] == "33" + assert "" not in body["output"][1]["content"][0]["text"] + assert "" not in body["output"][1]["content"][0]["text"] + + def test_literal_think_tags_remain_visible_without_reasoning_request(self, monkeypatch): + body = self._run_with_message(monkeypatch, {"content": "show x tags"}) + + assert [item["type"] for item in body["output"]] == ["message"] + assert body["output"][0]["content"][0]["text"] == "show x tags" + + def test_non_reasoning_gguf_keeps_literal_think_tags_visible(self, monkeypatch): + payload = ResponsesRequest(input = "hi", reasoning = {"effort": "high"}) + body = self._run_with_message( + monkeypatch, + {"content": "show x tags"}, + payload = payload, + llama_backend = SimpleNamespace( + is_loaded = True, + reasoning_always_on = False, + supports_reasoning = False, + ), + ) + + assert [item["type"] for item in body["output"]] == ["message"] + assert body["output"][0]["content"][0]["text"] == "show x tags" + + def test_structured_reasoning_content_extracts_text_parts(self, monkeypatch): + body = self._run_with_message( + monkeypatch, + { + "content": "33", + "reasoning_content": [ + {"type": "reasoning_text", "text": "plan"}, + {"type": "reasoning_text", "text": " next"}, + ], + }, + ) + + assert [item["type"] for item in body["output"]] == ["reasoning", "message"] + assert body["output"][0]["content"] == [{"type": "reasoning_text", "text": "plan next"}] + assert body["output"][1]["content"][0]["text"] == "33" + + def test_plain_content_remains_message_only(self, monkeypatch): + body = self._run_with_message(monkeypatch, {"content": "33"}) + + assert [item["type"] for item in body["output"]] == ["message"] + assert body["output"][0]["content"][0]["text"] == "33" + + def test_reasoning_only_is_also_visible_message_text(self, monkeypatch): + payload = ResponsesRequest(input = "hi", reasoning = {"effort": "high"}) + body = self._run_with_message( + monkeypatch, + {"content": "plan"}, + payload = payload, + ) + + assert [item["type"] for item in body["output"]] == ["reasoning", "message"] + assert body["output"][0]["content"][0]["text"] == "plan" + assert body["output"][1]["content"][0]["text"] == "plan" + + # ===================================================================== # Streaming Responses adapter # ===================================================================== @@ -570,6 +739,262 @@ class TestResponsesStreamAdapter: if line.startswith(prefix) ] + @staticmethod + def _install_stream_mock( + monkeypatch, + chunks, + *, + supports_reasoning = True, + reasoning_always_on = False, + ): + import routes.inference as inf_mod + + def handler(request: httpx.Request) -> httpx.Response: + content = "".join(f"data: {json.dumps(chunk)}\n\n" for chunk in chunks) + content += "data: [DONE]\n\n" + return httpx.Response( + 200, + content = content.encode(), + headers = {"content-type": "text/event-stream"}, + ) + + transport = httpx.MockTransport(handler) + real_async_client = httpx.AsyncClient + + def _client(*args, **kwargs): + return real_async_client( + transport = transport, + timeout = kwargs.get("timeout", 600), + ) + + monkeypatch.setattr(inf_mod.httpx, "AsyncClient", _client) + monkeypatch.setattr( + inf_mod, + "get_llama_cpp_backend", + lambda: SimpleNamespace( + is_loaded = True, + is_vision = False, + context_length = 4096, + base_url = "http://llama.test", + supports_reasoning = supports_reasoning, + reasoning_always_on = reasoning_always_on, + _request_reasoning_kwargs = ( + lambda enable_thinking = None, reasoning_effort = None, preserve_thinking = None: None + ), + ), + ) + + def test_split_think_markers_stream_as_reasoning_and_visible_text(self, monkeypatch): + chunks = [ + {"choices": [{"delta": {"content": "pla"}}]}, + {"choices": [{"delta": {"content": "n33"}}]}, + {"choices": [], "usage": {"prompt_tokens": 2, "completion_tokens": 3}}, + ] + self._install_stream_mock(monkeypatch, chunks) + payload = ResponsesRequest(input = "hi", stream = True, reasoning = {"effort": "high"}) + messages = [ChatMessage(role = "user", content = "hi")] + + async def run(): + response = await _responses_stream(payload, messages, self._Request()) + return await self._collect(response) + + lines = asyncio.run(run()) + + reasoning_deltas = self._payloads(lines, "response.reasoning_text.delta") + text_deltas = self._payloads(lines, "response.output_text.delta") + assert "".join(event["delta"] for event in reasoning_deltas) == "plan" + assert "".join(event["delta"] for event in text_deltas) == "33" + completed = self._payloads(lines, "response.completed")[0] + assert [item["type"] for item in completed["response"]["output"]] == [ + "reasoning", + "message", + ] + assert completed["response"]["output"][0]["content"][0]["text"] == "plan" + assert completed["response"]["output"][1]["content"][0]["text"] == "33" + + def test_literal_think_tags_stream_as_visible_text_without_reasoning_request(self, monkeypatch): + chunks = [ + {"choices": [{"delta": {"content": "show x tags"}}]}, + {"choices": [], "usage": {"prompt_tokens": 2, "completion_tokens": 3}}, + ] + self._install_stream_mock(monkeypatch, chunks) + payload = ResponsesRequest(input = "hi", stream = True) + messages = [ChatMessage(role = "user", content = "hi")] + + async def run(): + response = await _responses_stream(payload, messages, self._Request()) + return await self._collect(response) + + lines = asyncio.run(run()) + + reasoning_deltas = self._payloads(lines, "response.reasoning_text.delta") + text_deltas = self._payloads(lines, "response.output_text.delta") + assert reasoning_deltas == [] + assert "".join(event["delta"] for event in text_deltas) == "show x tags" + completed = self._payloads(lines, "response.completed")[0] + assert [item["type"] for item in completed["response"]["output"]] == ["message"] + assert completed["response"]["output"][0]["content"][0]["text"] == ( + "show x tags" + ) + + def test_non_reasoning_gguf_stream_keeps_literal_think_tags_visible(self, monkeypatch): + chunks = [ + {"choices": [{"delta": {"content": "show x tags"}}]}, + {"choices": [], "usage": {"prompt_tokens": 2, "completion_tokens": 3}}, + ] + self._install_stream_mock(monkeypatch, chunks, supports_reasoning = False) + payload = ResponsesRequest(input = "hi", stream = True, reasoning = {"effort": "high"}) + messages = [ChatMessage(role = "user", content = "hi")] + + async def run(): + response = await _responses_stream(payload, messages, self._Request()) + return await self._collect(response) + + lines = asyncio.run(run()) + + reasoning_deltas = self._payloads(lines, "response.reasoning_text.delta") + text_deltas = self._payloads(lines, "response.output_text.delta") + assert reasoning_deltas == [] + assert "".join(event["delta"] for event in text_deltas) == "show x tags" + completed = self._payloads(lines, "response.completed")[0] + assert [item["type"] for item in completed["response"]["output"]] == ["message"] + assert completed["response"]["output"][0]["content"][0]["text"] == ( + "show x tags" + ) + + def test_reasoning_only_streams_as_visible_message_text(self, monkeypatch): + chunks = [ + {"choices": [{"delta": {"content": "plan"}}]}, + {"choices": [], "usage": {"prompt_tokens": 2, "completion_tokens": 3}}, + ] + self._install_stream_mock(monkeypatch, chunks) + payload = ResponsesRequest(input = "hi", stream = True, reasoning = {"effort": "high"}) + messages = [ChatMessage(role = "user", content = "hi")] + + async def run(): + response = await _responses_stream(payload, messages, self._Request()) + return await self._collect(response) + + lines = asyncio.run(run()) + + reasoning_deltas = self._payloads(lines, "response.reasoning_text.delta") + text_deltas = self._payloads(lines, "response.output_text.delta") + assert "".join(event["delta"] for event in reasoning_deltas) == "plan" + assert "".join(event["delta"] for event in text_deltas) == "plan" + completed = self._payloads(lines, "response.completed")[0] + assert [item["type"] for item in completed["response"]["output"]] == [ + "reasoning", + "message", + ] + assert completed["response"]["output"][0]["content"][0]["text"] == "plan" + assert completed["response"]["output"][1]["content"][0]["text"] == "plan" + + def test_structured_reasoning_content_streams_as_reasoning(self, monkeypatch): + chunks = [ + {"choices": [{"delta": {"reasoning_content": "plan"}}]}, + {"choices": [{"delta": {"content": "33"}}]}, + {"choices": [], "usage": {"prompt_tokens": 2, "completion_tokens": 3}}, + ] + self._install_stream_mock(monkeypatch, chunks) + payload = ResponsesRequest(input = "hi", stream = True) + messages = [ChatMessage(role = "user", content = "hi")] + + async def run(): + response = await _responses_stream(payload, messages, self._Request()) + return await self._collect(response) + + lines = asyncio.run(run()) + + reasoning_deltas = self._payloads(lines, "response.reasoning_text.delta") + text_deltas = self._payloads(lines, "response.output_text.delta") + assert "".join(event["delta"] for event in reasoning_deltas) == "plan" + assert "".join(event["delta"] for event in text_deltas) == "33" + completed = self._payloads(lines, "response.completed")[0] + assert completed["response"]["output"][0]["type"] == "reasoning" + assert completed["response"]["output"][1]["type"] == "message" + + def test_structured_reasoning_content_parts_stream_as_reasoning(self, monkeypatch): + chunks = [ + { + "choices": [ + { + "delta": { + "reasoning_content": { + "content": [ + {"type": "reasoning_text", "text": "plan"}, + {"type": "reasoning_text", "text": " next"}, + ] + } + } + } + ] + }, + {"choices": [{"delta": {"content": "33"}}]}, + {"choices": [], "usage": {"prompt_tokens": 2, "completion_tokens": 3}}, + ] + self._install_stream_mock(monkeypatch, chunks) + payload = ResponsesRequest(input = "hi", stream = True) + messages = [ChatMessage(role = "user", content = "hi")] + + async def run(): + response = await _responses_stream(payload, messages, self._Request()) + return await self._collect(response) + + lines = asyncio.run(run()) + + reasoning_deltas = self._payloads(lines, "response.reasoning_text.delta") + text_deltas = self._payloads(lines, "response.output_text.delta") + assert "".join(event["delta"] for event in reasoning_deltas) == "plan next" + assert "".join(event["delta"] for event in text_deltas) == "33" + assert "reasoning_text" not in "".join(event["delta"] for event in reasoning_deltas) + completed = self._payloads(lines, "response.completed")[0] + assert completed["response"]["output"][0]["content"][0]["text"] == "plan next" + assert completed["response"]["output"][1]["content"][0]["text"] == "33" + + def test_tool_first_stream_closes_items_in_output_index_order(self, monkeypatch): + chunks = [ + { + "choices": [ + { + "delta": { + "tool_calls": [ + { + "index": 0, + "id": "call_0", + "type": "function", + "function": {"name": "lookup", "arguments": "{}"}, + } + ] + } + } + ] + }, + {"choices": [{"delta": {"content": "done"}}]}, + {"choices": [], "usage": {"prompt_tokens": 2, "completion_tokens": 3}}, + ] + self._install_stream_mock(monkeypatch, chunks) + payload = ResponsesRequest(input = "hi", stream = True) + messages = [ChatMessage(role = "user", content = "hi")] + + async def run(): + response = await _responses_stream(payload, messages, self._Request()) + return await self._collect(response) + + lines = asyncio.run(run()) + + done_events = self._payloads(lines, "response.output_item.done") + assert [event["output_index"] for event in done_events] == [0, 1] + assert [event["item"]["type"] for event in done_events] == ["function_call", "message"] + completed = self._payloads(lines, "response.completed")[0] + assert [item["type"] for item in completed["response"]["output"]] == [ + "function_call", + "message", + ] + def test_requests_usage_and_caps_parallel_tool_calls(self, monkeypatch): import routes.inference as inf_mod @@ -678,6 +1103,15 @@ class TestResponsesStreamAdapter: class TestResponsesOutputFunctionCall: + def test_reasoning_output_item_serialises_full_reasoning_content(self): + item = ResponsesOutputReasoning(content = [{"type": "reasoning_text", "text": "plan"}]) + d = item.model_dump() + assert d["type"] == "reasoning" + assert d["id"].startswith("rs_") + assert d["status"] == "completed" + assert d["summary"] == [] + assert d["content"] == [{"type": "reasoning_text", "text": "plan"}] + def test_direct_construction(self): fc = ResponsesOutputFunctionCall( call_id = "call_1", @@ -778,6 +1212,26 @@ class TestCodexStyleRequestShapes: assert len(req.input) == 3 assert isinstance(req.input[1], ResponsesUnknownInputItem) + def test_emitted_reasoning_item_replay_is_dropped_for_local_chat(self): + payload = ResponsesRequest( + input = [ + {"role": "user", "content": "Hi"}, + { + "type": "reasoning", + "id": "rs_1", + "summary": [], + "content": [{"type": "reasoning_text", "text": "plan"}], + }, + {"role": "assistant", "content": "33"}, + {"role": "user", "content": "Continue"}, + ], + ) + + msgs = _normalise_responses_input(payload) + + assert [m.role for m in msgs] == ["user", "assistant", "user"] + assert all("plan" not in (m.content or "") for m in msgs if isinstance(m.content, str)) + def test_unknown_content_part_type_accepted(self): """Unknown content-part types (e.g. future input_audio) validate as ResponsesUnknownContentPart so the request doesn't 422.""" diff --git a/studio/backend/tests/test_s3_dataset.py b/studio/backend/tests/test_s3_dataset.py new file mode 100644 index 0000000000..f47db565ff --- /dev/null +++ b/studio/backend/tests/test_s3_dataset.py @@ -0,0 +1,274 @@ +# SPDX-License-Identifier: AGPL-3.0-only +# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0 + +"""Tests for the S3 dataset loader (core.training.s3_dataset). + +boto3 is optional and may be absent in CI, so the S3 client is mocked: a fake +client provides a paginator over a synthetic bucket listing and writes files on +download_file. No network or real AWS credentials are involved. +""" + +import importlib.util +import os +from pathlib import Path + +import pytest + +# Load the modules under test directly by path. Importing them through their +# packages (core.training / models) would execute heavy package __init__ chains +# (structlog, torch, …) that aren't needed for these unit tests. +_BACKEND = Path(__file__).resolve().parents[1] + + +def _load(mod_name, rel_path): + spec = importlib.util.spec_from_file_location(mod_name, _BACKEND / rel_path) + module = importlib.util.module_from_spec(spec) + spec.loader.exec_module(module) + return module + + +s3_dataset = _load("s3_dataset", "core/training/s3_dataset.py") +S3Config = _load("models_training_s3", "models/training.py").S3Config + + +class _FakePaginator: + def __init__(self, keys): + self._keys = keys + + def paginate(self, **kwargs): + prefix = kwargs.get("Prefix") + contents = [{"Key": k} for k in self._keys if prefix is None or k.startswith(prefix)] + # Emit in two pages to exercise pagination handling. + mid = len(contents) // 2 + yield {"Contents": contents[:mid]} + yield {"Contents": contents[mid:]} + + +class _FakeS3Client: + def __init__(self, keys): + self._keys = keys + self.downloaded = [] + + def get_paginator(self, name): + assert name == "list_objects_v2" + return _FakePaginator(self._keys) + + def download_file(self, bucket, key, local_path, **kwargs): + self.downloaded.append((bucket, key, local_path)) + callback = kwargs.get("Callback") + if callback is not None: + callback(1) + with open(local_path, "w", encoding = "utf-8") as f: + f.write(f"content-of:{key}") + + +@pytest.fixture +def fake_client(monkeypatch): + """Force boto3_available True and stub the client builder.""" + keys = [ + "datasets/train.parquet", + "datasets/extra.parquet", + "datasets/notes.txt", # filtered out (unsupported) + "datasets/subdir/", # directory placeholder, skipped + "other/ignore.parquet", # filtered out by prefix + ] + client = _FakeS3Client(keys) + monkeypatch.setattr(s3_dataset, "boto3_available", lambda: True) + monkeypatch.setattr(s3_dataset, "_build_s3_client", lambda cfg: client) + return client + + +def _cfg(**overrides): + base = { + "bucket": "my-bucket", + "region": "us-east-1", + "prefix": "datasets/", + "access_key_id": "AKIA_TEST", + "secret_access_key": "secret", + "use_iam_role": False, + } + base.update(overrides) + return base + + +def test_downloads_only_supported_files_under_prefix(fake_client, tmp_path): + files = s3_dataset.download_s3_dataset(_cfg(), dest_dir = str(tmp_path)) + names = sorted(os.path.basename(f) for f in files) + # txt is unsupported, the directory placeholder is skipped, and the + # "other/" key is excluded by the prefix filter. + assert names == ["extra.parquet", "train.parquet"] + for f in files: + assert os.path.exists(f) + + +def test_allows_json_and_jsonl_family(monkeypatch, tmp_path): + client = _FakeS3Client(["datasets/train.json", "datasets/extra.jsonl"]) + monkeypatch.setattr(s3_dataset, "boto3_available", lambda: True) + monkeypatch.setattr(s3_dataset, "_build_s3_client", lambda cfg: client) + + files = s3_dataset.download_s3_dataset(_cfg(), dest_dir = str(tmp_path)) + + assert sorted(os.path.basename(f) for f in files) == ["extra.jsonl", "train.json"] + + +def test_ignores_common_json_metadata_files(monkeypatch, tmp_path): + client = _FakeS3Client( + [ + "datasets/train.parquet", + "datasets/schema.json", + "datasets/metadata.json", + "datasets/dataset_info.json", + ] + ) + monkeypatch.setattr(s3_dataset, "boto3_available", lambda: True) + monkeypatch.setattr(s3_dataset, "_build_s3_client", lambda cfg: client) + + files = s3_dataset.download_s3_dataset(_cfg(), dest_dir = str(tmp_path)) + + assert [os.path.basename(f) for f in files] == ["train.parquet"] + + +def test_raises_when_prefix_contains_mixed_formats(monkeypatch, tmp_path): + client = _FakeS3Client(["datasets/train.parquet", "datasets/stray.csv"]) + monkeypatch.setattr(s3_dataset, "boto3_available", lambda: True) + monkeypatch.setattr(s3_dataset, "_build_s3_client", lambda cfg: client) + + with pytest.raises(ValueError, match = "mixed dataset formats"): + s3_dataset.download_s3_dataset(_cfg(), dest_dir = str(tmp_path)) + + assert client.downloaded == [] + + +def test_raises_when_no_supported_files(monkeypatch, tmp_path): + client = _FakeS3Client(["datasets/readme.txt"]) + monkeypatch.setattr(s3_dataset, "boto3_available", lambda: True) + monkeypatch.setattr(s3_dataset, "_build_s3_client", lambda cfg: client) + with pytest.raises(ValueError, match = "No supported dataset files"): + s3_dataset.download_s3_dataset(_cfg(), dest_dir = str(tmp_path)) + + +def test_raises_when_boto3_missing(monkeypatch, tmp_path): + monkeypatch.setattr(s3_dataset, "boto3_available", lambda: False) + with pytest.raises(RuntimeError, match = "requires boto3"): + s3_dataset.download_s3_dataset(_cfg(), dest_dir = str(tmp_path)) + + +def test_basename_collisions_are_disambiguated(monkeypatch, tmp_path): + # Two keys share a basename under different sub-prefixes. + client = _FakeS3Client(["datasets/a/train.parquet", "datasets/b/train.parquet"]) + monkeypatch.setattr(s3_dataset, "boto3_available", lambda: True) + monkeypatch.setattr(s3_dataset, "_build_s3_client", lambda cfg: client) + files = s3_dataset.download_s3_dataset(_cfg(), dest_dir = str(tmp_path)) + assert len(files) == 2 + assert len(set(files)) == 2 # no overwrite + + +def test_basename_collision_skips_existing_generated_suffix(monkeypatch, tmp_path): + client = _FakeS3Client( + [ + "datasets/a/train.parquet", + "datasets/b/train_1.parquet", + "datasets/c/train.parquet", + ] + ) + monkeypatch.setattr(s3_dataset, "boto3_available", lambda: True) + monkeypatch.setattr(s3_dataset, "_build_s3_client", lambda cfg: client) + + files = s3_dataset.download_s3_dataset(_cfg(), dest_dir = str(tmp_path)) + + assert [os.path.basename(f) for f in files] == [ + "train.parquet", + "train_1.parquet", + "train_2.parquet", + ] + assert len(set(files)) == 3 + assert (tmp_path / "train_1.parquet").read_text(encoding = "utf-8") == ( + "content-of:datasets/b/train_1.parquet" + ) + assert (tmp_path / "train_2.parquet").read_text(encoding = "utf-8") == ( + "content-of:datasets/c/train.parquet" + ) + + +def test_download_handle_cleans_owned_temp_dir(monkeypatch, tmp_path): + target_dir = tmp_path / "owned-download" + client = _FakeS3Client(["datasets/train.parquet"]) + monkeypatch.setattr(s3_dataset, "boto3_available", lambda: True) + monkeypatch.setattr(s3_dataset, "_build_s3_client", lambda cfg: client) + monkeypatch.setattr(s3_dataset.tempfile, "mkdtemp", lambda prefix: str(target_dir)) + + download = s3_dataset.prepare_s3_dataset_download(_cfg()) + + assert target_dir.exists() + assert download.files == [str(target_dir / "train.parquet")] + download.cleanup() + assert not target_dir.exists() + + +def test_dest_dir_is_not_removed_by_cleanup(monkeypatch, tmp_path): + client = _FakeS3Client(["datasets/train.parquet"]) + monkeypatch.setattr(s3_dataset, "boto3_available", lambda: True) + monkeypatch.setattr(s3_dataset, "_build_s3_client", lambda cfg: client) + + download = s3_dataset.prepare_s3_dataset_download(_cfg(), dest_dir = str(tmp_path)) + + download.cleanup() + assert tmp_path.exists() + assert (tmp_path / "train.parquet").exists() + + +def test_cancel_callback_aborts_and_removes_temp_dir(monkeypatch, tmp_path): + target_dir = tmp_path / "cancelled-download" + client = _FakeS3Client(["datasets/train.parquet"]) + monkeypatch.setattr(s3_dataset, "boto3_available", lambda: True) + monkeypatch.setattr(s3_dataset, "_build_s3_client", lambda cfg: client) + monkeypatch.setattr(s3_dataset.tempfile, "mkdtemp", lambda prefix: str(target_dir)) + calls = 0 + + def cancel_after_download_starts(): + nonlocal calls + calls += 1 + return calls >= 4 + + with pytest.raises(s3_dataset.S3DownloadCancelled): + s3_dataset.prepare_s3_dataset_download( + _cfg(), + cancel_callback = cancel_after_download_starts, + ) + + assert not target_dir.exists() + + +# ── S3Config model (camelCase aliases + credential validation) ── + + +def test_s3config_accepts_camelcase_aliases(): + cfg = S3Config.model_validate( + { + "bucket": "b", + "region": "eu-west-1", + "accessKeyId": "AKIA", + "secretAccessKey": "shh", + } + ) + assert cfg.access_key_id == "AKIA" + assert cfg.secret_access_key == "shh" + # model_dump() yields snake_case for the loader. + assert cfg.model_dump()["access_key_id"] == "AKIA" + + +def test_s3config_accepts_snake_case(): + cfg = S3Config.model_validate( + {"bucket": "b", "access_key_id": "AKIA", "secret_access_key": "shh"} + ) + assert cfg.access_key_id == "AKIA" + + +def test_s3config_requires_credentials_or_iam(): + with pytest.raises(ValueError): + S3Config.model_validate({"bucket": "b"}) + + +def test_s3config_iam_role_needs_no_keys(): + cfg = S3Config.model_validate({"bucket": "b", "useIamRole": True}) + assert cfg.use_iam_role is True diff --git a/studio/backend/tests/test_studio_api.py b/studio/backend/tests/test_studio_api.py index 27f695b744..928b636e3e 100644 --- a/studio/backend/tests/test_studio_api.py +++ b/studio/backend/tests/test_studio_api.py @@ -149,6 +149,7 @@ def test_help_output(): "--host", "--frontend", "--silent", + "--tensor-parallel", ]: assert flag in out, f"Missing flag {flag!r} in --help output" print(" PASS --help shows all flags") diff --git a/studio/backend/tests/test_tensor_parallel.py b/studio/backend/tests/test_tensor_parallel.py new file mode 100644 index 0000000000..30bfb91a08 --- /dev/null +++ b/studio/backend/tests/test_tensor_parallel.py @@ -0,0 +1,578 @@ +# SPDX-License-Identifier: AGPL-3.0-only +# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0 + +"""Backend contract for the Tensor Parallelism toggle. + +The toggle threads a single ``tensor_parallel`` bool from the chat UI +through the load request to a ``--split-mode tensor`` llama-server flag, +and round-trips it back via the load/status responses so the switch +reflects what is actually running. These tests pin: + + * the pydantic request/response/status contract (snake_case key, + default False), + * the backend ``tensor_parallel`` property and its reset on unload, + * the ``_already_in_target_state`` reload-detection branch, and + * that ``--split-mode tensor`` is emitted only behind the toggle. +""" + +from __future__ import annotations + +import asyncio +import inspect +import sys +import types as _types +from pathlib import Path + +import pytest + +_BACKEND_DIR = str(Path(__file__).resolve().parent.parent) +if _BACKEND_DIR not in sys.path: + sys.path.insert(0, _BACKEND_DIR) + +# Same external-dep stubs as the other llama_cpp unit tests so importing +# the backend doesn't drag in structlog / httpx / loggers. +_loggers_stub = _types.ModuleType("loggers") +_loggers_stub.get_logger = lambda name: __import__("logging").getLogger(name) +sys.modules.setdefault("loggers", _loggers_stub) + +_structlog_stub = _types.ModuleType("structlog") +_structlog_stub.get_logger = lambda *a, **k: __import__("logging").getLogger("stub") +sys.modules.setdefault("structlog", _structlog_stub) + +_httpx_stub = _types.ModuleType("httpx") +for _exc in ( + "ConnectError", + "TimeoutException", + "ReadTimeout", + "ReadError", + "RemoteProtocolError", + "CloseError", +): + setattr(_httpx_stub, _exc, type(_exc, (Exception,), {})) +_httpx_stub.Timeout = type("T", (), {"__init__": lambda s, *a, **k: None}) +_httpx_stub.Client = type( + "C", + (), + { + "__init__": lambda s, **kw: None, + "__enter__": lambda s: s, + "__exit__": lambda s, *a: None, + }, +) +sys.modules.setdefault("httpx", _httpx_stub) + +from core.inference import llama_cpp as llama_cpp_module +from core.inference.llama_cpp import LlamaCppBackend +from core.inference.llama_server_args import resolve_tensor_parallel +from core.inference.tensor_fallback import load_with_tensor_fallback +from models.inference import ( + InferenceStatusResponse, + LoadRequest, + LoadResponse, +) + + +# ── Pydantic contract (snake_case key, default False) ──────────────── + + +def test_load_request_defaults_tensor_parallel_false(): + req = LoadRequest(model_path = "owner/repo") + assert req.tensor_parallel is False + + +def test_load_request_accepts_tensor_parallel(): + req = LoadRequest(model_path = "owner/repo", tensor_parallel = True) + assert req.tensor_parallel is True + + +def test_load_request_round_trips_json_key(): + # The frontend sends the snake_case key verbatim. + req = LoadRequest.model_validate({"model_path": "owner/repo", "tensor_parallel": True}) + assert req.tensor_parallel is True + assert req.model_dump()["tensor_parallel"] is True + + +@pytest.mark.parametrize("model_cls", [LoadResponse, InferenceStatusResponse]) +def test_response_models_emit_tensor_parallel(model_cls): + # Default False, and the key is always present in the JSON body. + if model_cls is LoadResponse: + default = model_cls( + status = "loaded", + model = "owner/repo", + display_name = "repo", + inference = {}, + ) + on = model_cls( + status = "loaded", + model = "owner/repo", + display_name = "repo", + inference = {}, + tensor_parallel = True, + ) + else: + default = model_cls() + on = model_cls(tensor_parallel = True) + assert default.model_dump()["tensor_parallel"] is False + assert on.model_dump()["tensor_parallel"] is True + + +# ── Backend property + reset ───────────────────────────────────────── + + +class _FakeProcess: + """Stand-in for subprocess.Popen so _kill_process is a no-op.""" + + def terminate(self): + pass + + def wait(self, timeout = None): + return 0 + + def kill(self): + pass + + def poll(self): + return 0 + + +def test_tensor_parallel_property_defaults_false(): + assert LlamaCppBackend().tensor_parallel is False + + +def test_tensor_parallel_property_reflects_field(): + backend = LlamaCppBackend() + backend._tensor_parallel = True + assert backend.tensor_parallel is True + + +def test_unload_resets_tensor_parallel(): + backend = LlamaCppBackend() + backend._process = _FakeProcess() + backend._tensor_parallel = True + backend.unload_model() + assert backend.tensor_parallel is False + + +# ── _already_in_target_state reload-detection branch ───────────────── + + +def _loaded_backend(tensor_parallel: bool) -> LlamaCppBackend: + backend = LlamaCppBackend() + backend._process = _FakeProcess() # is_loaded only checks "is not None" + backend._healthy = True + backend._model_identifier = "owner/repo" + backend._hf_variant = "Q4_K_M" + backend._requested_n_ctx = 8192 + backend._cache_type_kv = None + backend._requested_spec_mode = "auto" + backend._chat_template_override = None + backend._is_vision = False + backend._extra_args = None + backend._gguf_path = None + backend._tensor_parallel = tensor_parallel + return backend + + +def _target_state(backend: LlamaCppBackend, tensor_parallel: bool) -> bool: + return backend._already_in_target_state( + gguf_path = None, + model_identifier = "owner/repo", + hf_variant = "Q4_K_M", + n_ctx = 8192, + cache_type_kv = None, + speculative_type = "auto", + chat_template_override = None, + extra_args = None, + is_vision = False, + tensor_parallel = tensor_parallel, + ) + + +@pytest.mark.parametrize("flag", [True, False]) +def test_already_in_target_state_matches_same_tensor_parallel(flag): + assert _target_state(_loaded_backend(flag), flag) is True + + +@pytest.mark.parametrize( + "loaded,requested", + [(False, True), (True, False)], +) +def test_already_in_target_state_reloads_on_tensor_parallel_change(loaded, requested): + # Flipping the toggle either direction must force a reload so the + # command is rebuilt with/without --split-mode tensor. + assert _target_state(_loaded_backend(loaded), requested) is False + + +def test_already_in_target_state_reconciles_split_mode_extras(): + # Tensor engaged via --split-mode in extras (boolean omitted/default False) + # must match a server already running tensor mode -- no spurious reload. + backend = _loaded_backend(tensor_parallel = True) + backend._extra_args = ["--split-mode", "tensor"] + assert ( + backend._already_in_target_state( + gguf_path = None, + model_identifier = "owner/repo", + hf_variant = "Q4_K_M", + n_ctx = 8192, + cache_type_kv = None, + speculative_type = "auto", + chat_template_override = None, + extra_args = ["--split-mode", "tensor"], + is_vision = False, + tensor_parallel = False, + ) + is True + ) + + +# ── --split-mode tensor is emitted only behind the toggle ──────────── + + +def _load_model_source() -> str: + return inspect.getsource(llama_cpp_module.LlamaCppBackend.load_model) + + +def test_split_mode_tensor_is_gated_on_the_toggle(): + src = _load_model_source() + assert ( + 'cmd.extend(["--split-mode", "tensor"])' in src + ), "the tensor-parallel flag emission must be present in load_model" + # The emission lives behind `if tensor_parallel:` -- it must never be + # part of the unconditional base cmd list. + base_start = src.find("cmd = [") + base_end = src.find("\n ]", base_start) + base_block = src[base_start:base_end] if base_end > base_start else "" + assert ( + "--split-mode" not in base_block + ), "--split-mode must be conditional, not in the base cmd list" + gate = src.find("if tensor_parallel:") + emit = src.find('cmd.extend(["--split-mode", "tensor"])') + assert 0 <= gate < emit, "emission must sit under `if tensor_parallel:`" + + +def test_proportional_tensor_split_is_emitted_in_tensor_mode(): + # Asymmetric GPUs (e.g. 48 GB + 24 GB) OOM the smaller card under the + # even default; the allocator weights --tensor-split by free VRAM. Pin + # that the flag is emitted from inside the tensor-parallel block. + src = _load_model_source() + assert '"--tensor-split"' in src + gate = src.find("if tensor_parallel:") + ts = src.find('"--tensor-split"') + nxt_else = src.find("self._tensor_parallel = False") + assert 0 <= gate < ts < nxt_else, "--tensor-split must be emitted under `if tensor_parallel:`" + + +# ── tensor-mode allocation: conservative VRAM budget ───────────────── + + +def _kv_seeded_backend() -> LlamaCppBackend: + # Minimal GGUF metadata so _can_estimate_kv() is True (legacy KV path). + backend = LlamaCppBackend() + backend._n_layers = 32 + backend._embedding_length = 4096 + backend._n_heads = 32 + backend._n_kv_heads = 8 + backend._context_length = 131072 + return backend + + +def test_fit_context_budget_frac_override_is_tighter(): + backend = _kv_seeded_backend() + model_size = 8 * 1024**3 + pool_mib = 24 * 1024 # tight enough that KV capping bites + + fit_default = backend._fit_context_to_vram(131072, pool_mib, model_size, "f16") + fit_tp = backend._fit_context_to_vram(131072, pool_mib, model_size, "f16", budget_frac = 0.80) + assert fit_tp < 131072, "expected the context to be capped at this VRAM tier" + assert fit_tp <= fit_default, "a tighter budget must not allow MORE context" + # Omitting the override must reproduce the default budget exactly. + assert backend._fit_context_to_vram(131072, pool_mib, model_size, "f16") == fit_default + + +# ── unsupported-arch load failure -> clean message ─────────────────── + + +def test_split_mode_tensor_arch_failure_message(): + msg = LlamaCppBackend._classify_llama_start_failure( + "llama_model_create: LLAMA_SPLIT_MODE_TENSOR not implemented for " + "architecture 'deepseek2'", + None, + "unsloth/DeepSeek-V3-GGUF", + ) + assert "Tensor parallelism is not supported" in msg + + +def test_unrelated_arch_failure_not_hijacked_by_tensor_message(): + msg = LlamaCppBackend._classify_llama_start_failure( + "unknown model architecture: 'flux'", "/models/flux.gguf", None + ) + assert "Tensor parallelism" not in msg + + +# ── _plan_tensor_parallel: the allocation math (pure, no model/GPU) ─── +# Seeded full-attention KV (~128 KiB/token) via _kv_seeded_backend, so the +# context cap + split are deterministic. Asserts relationships rather than +# magic numbers so the KV estimate can evolve without breaking these. + +_GB = 1024**3 +_ASYM = [(0, 48000), (1, 24000)] # asymmetric pool, 72000 MiB +_SYM = [(0, 24000), (1, 24000)] # symmetric pool + + +def _plan( + model_gb, + target = 131072, + gpus = _ASYM, + mtp = False, +): + b = _kv_seeded_backend() + return b, b._plan_tensor_parallel(gpus, int(model_gb * _GB), target, mtp_engaged = mtp) + + +def _kv_budget_b(model_gb, gpus = _ASYM): + reserve = LlamaCppBackend._TENSOR_PARALLEL_BUFFER_RESERVE_MIB + return (sum(f for _, f in gpus) - len(gpus) * reserve) * 1024 * 1024 - int(model_gb * _GB) + + +def test_tp_plan_weighted_split_on_asymmetric_big_model(): + b, (ec, mac, gi, ts) = _plan(50) + reserve = b._TENSOR_PARALLEL_BUFFER_RESERVE_MIB + assert gi == [0, 1] + # split weighted by (free - buffer), not raw free + assert ts == [48000 - reserve, 24000 - reserve] + assert ec < 131072 # capped below native + + +def test_tp_plan_even_split_when_model_fits(): + # A small model whose even share fits the smallest GPU -> llama.cpp's even + # default (None), which is safe for archs that crash on a weighted split. + _, (ec, mac, gi, ts) = _plan(4) + assert ts is None + + +def test_tp_plan_symmetric_gpus_use_even_split(): + _, (ec, mac, gi, ts) = _plan(8, gpus = _SYM) + assert ts is None + + +def test_tp_plan_context_fits_pool_budget_no_oom(): + b, (ec, mac, gi, ts) = _plan(50) + # the chosen context's KV must fit the pooled budget (weights + buffers) + assert b._estimate_kv_cache_bytes(ec) <= _kv_budget_b(50) + + +def test_tp_plan_uses_available_vram_not_wasteful(): + # when the cap engages, the chosen context nearly fills the budget + b, (ec, mac, gi, ts) = _plan(50) + assert b._estimate_kv_cache_bytes(ec) >= 0.9 * _kv_budget_b(50) + + +def test_tp_plan_weights_exceed_pool_floors_context(): + # 70 GB > pool minus per-GPU reserves -> floor (triggers layer fallback) + _, (ec, mac, gi, ts) = _plan(70) + assert ec == 2048 + + +def test_tp_plan_floor_never_exceeds_explicit_small_context(): + # An explicit context below the 2048 floor must not be raised: a caller + # asking for 1024 should not have KV sized for 2048 (avoidable OOM). + _, (ec, mac, gi, ts) = _plan(70, target = 1024) # weights exceed pool -> floor path + assert ec == 1024 + _, (ec2, *_rest) = _plan(50, target = 1024) # cap path with a tiny budget + assert ec2 <= 1024 + + +def test_tp_plan_explicit_context_honored_when_it_fits(): + _, (ec, mac, gi, ts) = _plan(50, target = 8192) + assert ec == 8192 + + +def test_tp_plan_explicit_context_capped_when_too_large(): + _, (ec, mac, gi, ts) = _plan(50, target = 131072) + assert 2048 <= ec < 131072 + + +def test_tp_plan_max_available_ctx_reports_native_not_explicit_ctx(): + # An explicit small ctx caps effective_ctx but the UI ceiling + # (max_available_ctx) must reflect the native/hardware cap, not the request. + b = _kv_seeded_backend() + ec, mac, _gi, _ts = b._plan_tensor_parallel(_ASYM, int(50 * _GB), 8192, max_target_ctx = 131072) + _, native_mac, *_ = b._plan_tensor_parallel(_ASYM, int(50 * _GB), 131072) + assert ec == 8192 # explicit request honored for the load + assert mac == native_mac > ec # ceiling reflects the hardware cap + + +def test_tp_plan_mtp_reserves_extra_and_shrinks_context(): + _, (ec_no, *_rest) = _plan(50) + _, (ec_mtp, *_rest) = _plan(50, mtp = True) + assert ec_mtp < ec_no + + +def test_tp_plan_no_kv_metadata_floors_context(): + b = LlamaCppBackend() # no KV metadata -> can't size safely + ec, mac, gi, ts = b._plan_tensor_parallel(_ASYM, int(50 * _GB), 131072) + assert ec <= 4096 + + +def test_tp_plan_single_gpu_never_splits(): + # The toggle is a no-op without >= 2 GPUs (most dev/CI machines). Even if + # the planner is reached, it must not emit a tensor split. + b = _kv_seeded_backend() + ec, mac, gi, ts = b._plan_tensor_parallel([(0, 24000)], int(8 * _GB), 8192) + assert ts is None + assert gi == [0] + + +def test_tp_plan_zero_gpus_never_splits(): + b = _kv_seeded_backend() + ec, mac, gi, ts = b._plan_tensor_parallel([], int(8 * _GB), 8192) + assert ts is None + assert gi == [] + + +def test_tp_plan_drops_gpu_below_buffer_reserve(): + # A GPU with less free VRAM than the per-device compute-buffer reserve + # can't host tensor mode; it's excluded, which here leaves <2 usable -> no + # split (and gpu_indices reflects only the usable device). + b = _kv_seeded_backend() + reserve = LlamaCppBackend._TENSOR_PARALLEL_BUFFER_RESERVE_MIB + ec, mac, gi, ts = b._plan_tensor_parallel([(0, 48000), (1, reserve - 1)], int(8 * _GB), 8192) + assert gi == [0] + assert ts is None + + +# ── route auto-fallback survives a *raised* tensor-load crash ───────── +# A tensor-incompatible model makes load_model RAISE (Gemma 3n aborts) rather +# than return False. The /load fallback helper must catch that and retry with +# layer split -- stripping any --split-mode from the extras so the retry can't +# relaunch tensor -- while a non-tensor load propagates its exception. These +# exercise the real helper with a fake loader (no GPU, no llama-server). + + +class _RecordingLoader: + """Fake ``attempt_load``: crashes whenever tensor mode is effectively + engaged (via the bool or a ``--split-mode`` in extras), like a real + tensor-incompatible model; succeeds on layer split.""" + + def __init__(self): + self.calls: list[tuple] = [] + + async def __call__(self, tensor_parallel, extra_args): + self.calls.append((tensor_parallel, list(extra_args) if extra_args else extra_args)) + if resolve_tensor_parallel(extra_args, tensor_parallel): + raise RuntimeError("llama-server failed to start") + return True + + +def test_tensor_fallback_retries_layer_on_crash(): + loader = _RecordingLoader() + ok = asyncio.run( + load_with_tensor_fallback(loader, requested_tensor = True, extra_args = None, label = "m") + ) + assert ok is True + # tensor first (crashes), then layer split. + assert [c[0] for c in loader.calls] == [True, False] + + +def test_tensor_fallback_no_retry_on_success(): + calls: list[bool] = [] + + async def _ok(tensor_parallel, extra_args): + calls.append(tensor_parallel) + return True + + ok = asyncio.run( + load_with_tensor_fallback(_ok, requested_tensor = True, extra_args = None, label = "m") + ) + assert ok is True + assert calls == [True] # no fallback when the tensor load succeeds + + +def test_tensor_fallback_retries_when_tensor_returns_false(): + # load_model can signal failure by *returning False* (not only by raising); + # that must trigger the layer-split retry just like a crash does. + calls: list[bool] = [] + + async def _false_on_tensor(tensor_parallel, extra_args): + calls.append(tensor_parallel) + return not resolve_tensor_parallel(extra_args, tensor_parallel) + + ok = asyncio.run( + load_with_tensor_fallback( + _false_on_tensor, requested_tensor = True, extra_args = None, label = "m" + ) + ) + assert ok is True + assert calls == [True, False] + + +def test_tensor_fallback_returns_false_when_both_attempts_fail(): + # Tensor fails and the layer retry also fails -> the helper returns False so + # the route raises its own HTTP 500 (it does not crash mid-flight). + calls: list[bool] = [] + + async def _always_false(tensor_parallel, extra_args): + calls.append(tensor_parallel) + return False + + ok = asyncio.run( + load_with_tensor_fallback(_always_false, requested_tensor = True, extra_args = None, label = "m") + ) + assert ok is False + assert calls == [True, False] # tried tensor, then layer split + + +def test_tensor_fallback_skips_layer_retry_when_cancelled(): + # load_model returns False on a user cancellation too. When cancelled() is + # True, the helper must NOT relaunch the load the user just cancelled. + calls: list[bool] = [] + + async def _false_on_tensor(tensor_parallel, extra_args): + calls.append(tensor_parallel) + return False + + ok = asyncio.run( + load_with_tensor_fallback( + _false_on_tensor, + requested_tensor = True, + extra_args = None, + label = "m", + cancelled = lambda: True, + ) + ) + assert ok is False + assert calls == [True] # no layer-split retry after cancellation + + +@pytest.mark.parametrize( + "extras", + [ + ["--split-mode", "tensor", "-c", "4096"], + ["-sm", "tensor", "-c", "4096"], + ["--split-mode=tensor", "-c", "4096"], + ["-sm=tensor", "-c", "4096"], + ], +) +def test_tensor_fallback_strips_split_mode_from_extras_on_retry(extras): + # Tensor engaged via extras (boolean False); the retry must drop every + # --split-mode form (long/short, space/=) but keep the user's other flags, + # else resolve_tensor_parallel re-enables tensor and relaunches the crash. + loader = _RecordingLoader() + ok = asyncio.run( + load_with_tensor_fallback(loader, requested_tensor = False, extra_args = extras, label = "m") + ) + assert ok is True + assert len(loader.calls) == 2 + assert loader.calls[1][1] == ["-c", "4096"] # split-mode stripped, -c kept + + +def test_tensor_fallback_propagates_non_tensor_crash(): + async def _always_raise(tensor_parallel, extra_args): + raise RuntimeError("bad model") + + with pytest.raises(RuntimeError, match = "bad model"): + asyncio.run( + load_with_tensor_fallback( + _always_raise, requested_tensor = False, extra_args = None, label = "m" + ) + ) diff --git a/studio/backend/tests/test_training_resume.py b/studio/backend/tests/test_training_resume.py new file mode 100644 index 0000000000..91fdac9961 --- /dev/null +++ b/studio/backend/tests/test_training_resume.py @@ -0,0 +1,93 @@ +# SPDX-License-Identifier: AGPL-3.0-only +# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0 + +"""Regression tests for resumable training run eligibility.""" + +import importlib.util +import json +from pathlib import Path + + +_BACKEND = Path(__file__).resolve().parents[1] + + +def _load_resume_module(): + spec = importlib.util.spec_from_file_location( + "training_resume_under_test", + _BACKEND / "core" / "training" / "resume.py", + ) + module = importlib.util.module_from_spec(spec) + assert spec.loader is not None + spec.loader.exec_module(module) + return module + + +resume = _load_resume_module() + + +def _stopped_run(**overrides): + run = { + "status": "stopped", + "final_step": 5, + "total_steps": 10, + "output_dir": "/tmp/unsloth-output", + "resumed_later": False, + "config_json": json.dumps({"hf_dataset": "org/dataset"}), + } + run.update(overrides) + return run + + +def test_can_resume_run_allows_checkpointed_non_s3_run(monkeypatch): + monkeypatch.setattr(resume, "has_resume_state", lambda _path: True) + + assert resume.can_resume_run(_stopped_run()) is True + + +def test_can_resume_run_rejects_s3_dataset_source(monkeypatch): + monkeypatch.setattr(resume, "has_resume_state", lambda _path: True) + + run = _stopped_run( + config_json = json.dumps( + { + "dataset_source": "s3", + "s3_dataset": { + "bucket": "training-data", + "prefix": "datasets/", + "region": "us-east-1", + "use_iam_role": True, + }, + } + ) + ) + + assert resume.can_resume_run(run) is False + + +def test_can_resume_run_rejects_s3_metadata_marker(monkeypatch): + monkeypatch.setattr(resume, "has_resume_state", lambda _path: True) + + run = _stopped_run(config_json = json.dumps({"s3_dataset": {"bucket": "training-data"}})) + + assert resume.can_resume_run(run) is False + + +def test_list_runs_includes_config_json_for_resume_policy(monkeypatch, tmp_path): + from storage import studio_db + + monkeypatch.setenv("UNSLOTH_STUDIO_HOME", str(tmp_path)) + monkeypatch.setattr(studio_db, "_schema_ready", False) + config_json = json.dumps({"dataset_source": "s3", "s3_dataset": {"bucket": "training-data"}}) + + studio_db.create_run( + id = "run-s3", + model_name = "unsloth/test-model", + dataset_name = "s3://training-data", + config_json = config_json, + started_at = "2026-01-01T00:00:00Z", + total_steps = 10, + ) + + result = studio_db.list_runs() + + assert result["runs"][0]["config_json"] == config_json diff --git a/studio/backend/tests/test_vision_cache.py b/studio/backend/tests/test_vision_cache.py index 2fee50d842..d1bdec8449 100644 --- a/studio/backend/tests/test_vision_cache.py +++ b/studio/backend/tests/test_vision_cache.py @@ -30,6 +30,7 @@ _loggers_stub.get_logger = lambda name: __import__("logging").getLogger(name) sys.modules.setdefault("loggers", _loggers_stub) from utils.models.model_config import ( + ModelConfig, is_vision_model, _is_vision_model_uncached, _vision_detection_cache, @@ -120,6 +121,99 @@ class TestVisionCacheSubprocessPath: mock_raw_config.assert_called_once_with("unsloth/gemma-4-E4B-it", hf_token = None) +# --------------------------------------------------------------------------- +# Local GGUF capability path +# --------------------------------------------------------------------------- + + +class TestLocalGgufVisionDetection: + @patch( + "utils.models.model_config._is_vision_model_subprocess", + side_effect = AssertionError("GGUF must not use Transformers vision detection"), + ) + def test_qwen36_gguf_with_mmproj_skips_transformers(self, mock_subprocess, tmp_path): + model = tmp_path / "Qwen3.6-27B-UD-Q4_K_XL-MTP.gguf" + model.write_bytes(b"") + (tmp_path / "mmproj-F32.gguf").write_bytes(b"") + + assert is_vision_model(str(model)) is True + mock_subprocess.assert_not_called() + + @patch( + "utils.models.model_config._is_vision_model_subprocess", + side_effect = AssertionError("GGUF must not use Transformers vision detection"), + ) + def test_direct_gguf_in_variant_subdir_finds_snapshot_mmproj(self, mock_subprocess, tmp_path): + variant_dir = tmp_path / "BF16" + variant_dir.mkdir() + model = variant_dir / "Qwen3.6-27B-UD-Q4_K_XL-MTP.gguf" + model.write_bytes(b"") + (tmp_path / "mmproj-F32.gguf").write_bytes(b"") + + assert is_vision_model(str(model)) is True + mock_subprocess.assert_not_called() + + @patch( + "utils.models.model_config._is_vision_model_subprocess", + side_effect = AssertionError("GGUF must not use Transformers vision detection"), + ) + def test_qwen36_gguf_without_mmproj_skips_transformers(self, mock_subprocess, tmp_path): + model = tmp_path / "Qwen3.6-27B-UD-Q4_K_XL-MTP.gguf" + model.write_bytes(b"") + + assert is_vision_model(str(model)) is False + mock_subprocess.assert_not_called() + + def test_local_gguf_check_observes_mmproj_added_later(self, tmp_path): + model = tmp_path / "Qwen3.6-27B-UD-Q4_K_XL-MTP.gguf" + model.write_bytes(b"") + + assert is_vision_model(str(model)) is False + (tmp_path / "mmproj-F32.gguf").write_bytes(b"") + assert is_vision_model(str(model)) is True + + @patch( + "utils.models.model_config._is_vision_model_subprocess", + side_effect = AssertionError("GGUF must not use Transformers vision detection"), + ) + def test_ui_selection_returns_local_gguf_config(self, mock_subprocess, tmp_path): + model = tmp_path / "Qwen3.6-27B-UD-Q4_K_XL-MTP.gguf" + model.write_bytes(b"") + mmproj = tmp_path / "mmproj-F32.gguf" + mmproj.write_bytes(b"") + + config = ModelConfig.from_ui_selection(str(model), None) + + assert config is not None + assert config.is_gguf is True + assert config.is_vision is True + assert config.gguf_mmproj_file == str(mmproj.resolve()) + mock_subprocess.assert_not_called() + + @patch( + "utils.models.model_config._is_vision_model_subprocess", + side_effect = AssertionError("GGUF must not use Transformers vision detection"), + ) + def test_ui_selection_direct_gguf_in_variant_subdir_keeps_mmproj( + self, mock_subprocess, tmp_path + ): + variant_dir = tmp_path / "BF16" + variant_dir.mkdir() + model = variant_dir / "Qwen3.6-27B-UD-Q4_K_XL-MTP.gguf" + model.write_bytes(b"") + mmproj = tmp_path / "mmproj-F32.gguf" + mmproj.write_bytes(b"") + + config = ModelConfig.from_ui_selection(str(model), None) + + assert config is not None + assert config.is_gguf is True + assert config.is_vision is True + assert config.gguf_mmproj_file == str(mmproj.resolve()) + mock_subprocess.assert_not_called() + + +# --------------------------------------------------------------------------- # Exception handling — cache the False fallback diff --git a/studio/backend/utils/llama_cpp_freshness.py b/studio/backend/utils/llama_cpp_freshness.py index f5fd745334..87d0d2ec01 100644 --- a/studio/backend/utils/llama_cpp_freshness.py +++ b/studio/backend/utils/llama_cpp_freshness.py @@ -301,7 +301,22 @@ def format_stale_warning(info: dict) -> str: ) -def reset_caches() -> None: - """Test-only: drop all in-memory caches.""" +def reset_caches(*, drop_disk: bool = False) -> None: + """Drop the in-memory freshness caches. The no-arg form is test-only. + + With ``drop_disk = True`` also delete the on-disk 24h release cache. Used by + the post-install/update path: in-memory clearing alone leaves the stale + same-base value on disk, so if the post-install GitHub refresh can't reach + the network, ``latest_published_release`` would replay that stale disk value + (see its last-good fallback) and the banner could linger. Dropping the disk + cache makes latest read as None in that offline case, so the banner fails + open (off) instead of pointing at the just-replaced build.""" _marker_cache.clear() _release_memo.clear() + if drop_disk: + import shutil + + # _cache_dir() is a dedicated freshness-only subdir; it is re-created on + # the next _save_disk_cache. ignore_errors so a missing/locked dir is a + # no-op rather than breaking an otherwise successful install. + shutil.rmtree(_cache_dir(), ignore_errors = True) diff --git a/studio/backend/utils/llama_cpp_update.py b/studio/backend/utils/llama_cpp_update.py index 654ade6cd4..8b90d36bc5 100644 --- a/studio/backend/utils/llama_cpp_update.py +++ b/studio/backend/utils/llama_cpp_update.py @@ -310,8 +310,9 @@ def get_update_status(*, force_refresh: bool = False) -> dict: def _rocm_install_args(asset: Optional[str]) -> list[str]: """Forward --rocm-gfx/--has-rocm from the marker asset, mirroring setup.sh. - The installer probe can miss the gfx arch on amd-smi-only hosts; lemonade - bundles carry the family in the name (rocm-gfx110X), fork bundles only rocm/hip.""" + The installer probe can miss the gfx arch on amd-smi-only hosts; per-gfx + ROCm bundles carry the family in the name (rocm-gfx110X), version-tagged + bundles only rocm/hip.""" if not asset: return [] low = asset.lower() @@ -406,10 +407,13 @@ def _run_update(install_dir: Path, repo: str, asset: Optional[str], script: Path tail = "".join(tail_lines).strip()[-1500:] raise RuntimeError(f"installer exited {returncode}: {tail or 'no output'}") - # New UNSLOTH_PREBUILT_INFO.json is on disk; drop in-memory caches and - # re-prime the 24h disk freshness cache with the true newest, so the - # banner can't linger on a stale same-base value after the swap. - reset_caches() + # New UNSLOTH_PREBUILT_INFO.json is on disk; drop the in-memory AND the + # on-disk freshness caches, then re-prime the 24h disk cache with the + # true newest, so the banner can't linger on a stale same-base value + # after the swap. drop_disk matters when the refresh below can't reach + # GitHub: without it, latest_published_release would replay the stale + # disk value; with it, latest reads as None and the banner fails open. + reset_caches(drop_disk = True) try: latest_published_release(repo, force_refresh = True) except Exception as exc: # pragma: no cover - network defensive diff --git a/studio/backend/utils/models/model_config.py b/studio/backend/utils/models/model_config.py index f61c210cf6..f91b8fb8c9 100644 --- a/studio/backend/utils/models/model_config.py +++ b/studio/backend/utils/models/model_config.py @@ -721,6 +721,25 @@ def is_vision_model(model_name: str, hf_token: Optional[str] = None) -> bool: model_name: Model identifier (HF repo or local path) hf_token: Optional HF token for gated/private models """ + # Local GGUF models are served by llama-server. Their multimodal + # capability comes from a companion mmproj, not a Transformers config. + # Do not cache this lookup: a projector may be added beside an existing + # weight file after it was first inspected. + if is_local_path(model_name): + local_path = normalize_path(model_name) + gguf_file = detect_gguf_model(local_path) + if gguf_file: + companion_root = _local_gguf_companion_search_root(local_path, gguf_file) + mmproj_file = detect_mmproj_file(gguf_file, search_root = companion_root) + is_vision = mmproj_file is not None + logger.debug( + "Local GGUF vision check for '%s': mmproj=%s, is_vision=%s", + gguf_file, + mmproj_file, + is_vision, + ) + return is_vision + # Normalize model name so different casings of the same repo share a key try: if is_local_path(model_name): @@ -1324,6 +1343,7 @@ def _extract_quant_label(filename: str) -> str: "model-UD-IQ1_S.gguf" → "UD-IQ1_S" "model-UD-TQ1_0.gguf" → "UD-TQ1_0" "MXFP4_MOE/model-MXFP4_MOE-0001.gguf"→ "MXFP4_MOE" + "Qwen3.6-IQ4_XS-3.53bpw.gguf" → "IQ4_XS-3.53bpw" """ import re @@ -1339,6 +1359,10 @@ def _extract_quant_label(filename: str) -> str: r"|Q[0-9]+_[0-9]+" # Standard: Q8_0, Q5_1 r"|Q[0-9]+_K" # Short K-quant: Q6_K r"|BF16|F16|F32)" # Full precision + # Optional bits-per-weight modifier so repos that ship multiple + # files at the same base quant (e.g. byteshape's IQ4_XS at 3.53, + # 3.97, 4.19 bpw) don't collapse into a single merged variant. + r"(-[0-9]+(?:\.[0-9]+)?bpw)?" ) match = re.search(quant_re, stem, re.IGNORECASE) # Subdir layouts like ``BF16/foo.gguf`` keep the quant in the directory, @@ -1353,11 +1377,41 @@ def _extract_quant_label(filename: str) -> str: break if match: prefix = match.group(1) or "" - return f"{prefix}{match.group(2)}" + bpw = match.group(3) or "" + return f"{prefix}{match.group(2)}{bpw}" # Fallback: last hyphen-separated segment return stem.split("-")[-1] +def _local_gguf_companion_search_root(selected_path: str, gguf_file: str) -> str: + """Directory to scan upward from for local GGUF companion files.""" + import re + + selected = Path(selected_path) + gguf_path = Path(gguf_file) + if selected.suffix.lower() != ".gguf": + return selected_path + + gguf_dir = gguf_path.parent + if not gguf_dir.name: + return str(gguf_dir) + + quant_dir_re = ( + r"(UD-)?(" + r"MXFP[0-9]+(?:_[A-Z0-9]+)*" + r"|IQ[0-9]+_[A-Z]+(?:_[A-Z0-9]+)?" + r"|TQ[0-9]+_[0-9]+" + r"|Q[0-9]+_K_[A-Z]+" + r"|Q[0-9]+_[0-9]+" + r"|Q[0-9]+_K" + r"|BF16|F16|F32" + r")" + ) + if re.fullmatch(quant_dir_re, gguf_dir.name, re.IGNORECASE): + return str(gguf_dir.parent) + return str(gguf_dir) + + def _iter_hf_cache_snapshots(repo_id: str): """Yield HF cache snapshot dirs for *repo_id*, newest first. @@ -1371,12 +1425,11 @@ def _iter_hf_cache_snapshots(repo_id: str): return cache_dir = Path(hf_constants.HF_HUB_CACHE) - if not cache_dir.is_dir(): - return - target = f"models--{repo_id.replace('/', '--')}".lower() repo_dir: Optional[Path] = None try: + if not cache_dir.is_dir(): + return for entry in cache_dir.iterdir(): if entry.is_dir() and entry.name.lower() == target: repo_dir = entry @@ -1387,10 +1440,9 @@ def _iter_hf_cache_snapshots(repo_id: str): return snapshots = repo_dir / "snapshots" - if not snapshots.is_dir(): - return - try: + if not snapshots.is_dir(): + return snap_dirs = [s for s in snapshots.iterdir() if s.is_dir()] except OSError: return @@ -2277,10 +2329,10 @@ class ModelConfig: except Exception as e: logger.debug(f"Could not read export metadata: {e}") - # Pass search_root=path so detect_mmproj_file walks up to the - # snapshot root: the weight may sit in a quant subdir while - # mmproj-*.gguf lives at the root. - mmproj_file = detect_mmproj_file(gguf_file, search_root = path) + # Direct file selections may point into a quant subdir while + # mmproj-*.gguf lives at the snapshot root. + companion_root = _local_gguf_companion_search_root(path, gguf_file) + mmproj_file = detect_mmproj_file(gguf_file, search_root = companion_root) if mmproj_file: gguf_is_vision = True logger.info(f"Detected mmproj for vision: {mmproj_file}") @@ -2288,7 +2340,7 @@ class ModelConfig: logger.warning(f"Base model is vision but no mmproj file found in {gguf_dir}") # Separate MTP drafter sibling (Gemma 4), mirroring mmproj. - mtp_file = detect_mtp_file(gguf_file, search_root = path) + mtp_file = detect_mtp_file(gguf_file, search_root = companion_root) if mtp_file: logger.info(f"Detected MTP drafter: {mtp_file}") @@ -2478,6 +2530,15 @@ class ModelConfig: identifier = resolved_identifier path = resolved_identifier + # Keep existing local GGUF selections on the llama-server path. This + # constructor is still used by older inference helpers and must not + # describe a .gguf weight file as loadable by FastVisionModel. + if is_local and not is_lora and detect_gguf_model(path): + gguf_config = cls.from_identifier(path, hf_token = hf_token) + if gguf_config is not None: + gguf_config.display_name = display_name + return gguf_config + # --- Base Model and Vision Detection --- base_model = None is_vision = False diff --git a/studio/backend/utils/paths/__init__.py b/studio/backend/utils/paths/__init__.py index 6913a3be73..eff9b64678 100644 --- a/studio/backend/utils/paths/__init__.py +++ b/studio/backend/utils/paths/__init__.py @@ -43,13 +43,14 @@ from .storage_roots import ( resolve_under_root, resolve_output_dir, resolve_export_dir, + resolve_export_write_dir, resolve_tensorboard_dir, resolve_dataset_path, ) # Re-export shim: mark project-path helpers as used so the import-hoist # safety net does not flag them as unused. -_REEXPORTED = (documents_root, project_workspaces_root) +_REEXPORTED = (documents_root, project_workspaces_root, resolve_export_write_dir) __all__ = [ "normalize_path", @@ -89,6 +90,7 @@ __all__ = [ "resolve_under_root", "resolve_output_dir", "resolve_export_dir", + "resolve_export_write_dir", "resolve_tensorboard_dir", "resolve_dataset_path", ] diff --git a/studio/backend/utils/paths/storage_roots.py b/studio/backend/utils/paths/storage_roots.py index 2d4f5ac243..d336bc2e71 100644 --- a/studio/backend/utils/paths/storage_roots.py +++ b/studio/backend/utils/paths/storage_roots.py @@ -6,7 +6,7 @@ from __future__ import annotations import json import os import sys -from pathlib import Path +from pathlib import Path, PurePosixPath, PureWindowsPath import tempfile @@ -317,6 +317,26 @@ def _clean_relative_path(path_value: str, *, strip_prefixes: tuple[str, ...] = ( return Path(*parts) if parts else Path() +def _has_parent_segment(raw: str, path: Path) -> bool: + """Return true when a user path contains a parent-directory segment. + + On POSIX, ``Path("E:\\foo\\..\\bar")`` treats backslashes as normal + characters, so check both the host parser and Windows-style parsing. + """ + if ".." in path.parts: + return True + if ".." in PureWindowsPath(raw).parts: + return True + return ".." in raw.replace("\\", "/").split("/") + + +def _is_absolute_user_path(path: Path) -> bool: + expanded = str(path) + if os.name == "nt": + return path.is_absolute() and PureWindowsPath(expanded).is_absolute() + return path.is_absolute() and PurePosixPath(expanded).is_absolute() + + def _assert_contained(resolved: Path, root: Path) -> None: """Raise ValueError if ``resolved`` realpaths outside ``root``.""" try: @@ -351,10 +371,10 @@ def resolve_under_root( raise ValueError("path may not contain null bytes") path = Path(raw).expanduser() - if ".." in path.parts: + if _has_parent_segment(raw, path): raise ValueError(f"path may not contain '..' segments: {raw!r}") - if path.is_absolute(): + if _is_absolute_user_path(path): _assert_contained(path, root) return path @@ -373,6 +393,36 @@ def resolve_output_dir(path_value: str | None = None) -> Path: def resolve_export_dir(path_value: str | None = None) -> Path: + """Resolve an export directory — contained under exports_root(). + + Used by scan/read endpoints. Use :func:`resolve_export_write_dir` + for the export write path where absolute paths are accepted. + """ + return resolve_under_root( + path_value, + root = exports_root(), + strip_prefixes = ("exports",), + ) + + +def resolve_export_write_dir(path_value: str | None = None) -> Path: + """Resolve an export save directory — accepts absolute paths. + + Unlike :func:`resolve_export_dir`, this function passes absolute + paths through as-is so users can target a different drive when + their Studio install lives on a constrained system volume + (see :gh-issue:`6082`). Used only by the export write path. + """ + if not path_value or not str(path_value).strip(): + return exports_root() + raw = str(path_value).strip() + if "\x00" in raw: + raise ValueError("path may not contain null bytes") + path = Path(raw).expanduser() + if _has_parent_segment(raw, path): + raise ValueError(f"path may not contain '..' segments: {raw!r}") + if _is_absolute_user_path(path): + return path return resolve_under_root( path_value, root = exports_root(), diff --git a/studio/frontend/index.html b/studio/frontend/index.html index 4f81ffd4ff..0fbb4eaeeb 100644 --- a/studio/frontend/index.html +++ b/studio/frontend/index.html @@ -1,16 +1,16 @@ - + - - - - - - Unsloth Studio - - -
- - - + + + + + + Unsloth Studio + + +
+ + + diff --git a/studio/frontend/public/hub/profile/logo/meta.svg b/studio/frontend/public/hub/profile/logo/meta.svg index 9fa656bd6b..fe3709aeea 100644 --- a/studio/frontend/public/hub/profile/logo/meta.svg +++ b/studio/frontend/public/hub/profile/logo/meta.svg @@ -1,19 +1,19 @@ - - -Logo of Meta Platforms -- Graphic created by Detmar Owen - - - - - - - - - - - - - - - + + +Logo of Meta Platforms -- Graphic created by Detmar Owen + + + + + + + + + + + + + + + \ No newline at end of file diff --git a/studio/frontend/public/provider-logos/misc/meta.svg b/studio/frontend/public/provider-logos/misc/meta.svg index 9fa656bd6b..fe3709aeea 100644 --- a/studio/frontend/public/provider-logos/misc/meta.svg +++ b/studio/frontend/public/provider-logos/misc/meta.svg @@ -1,19 +1,19 @@ - - -Logo of Meta Platforms -- Graphic created by Detmar Owen - - - - - - - - - - - - - - - + + +Logo of Meta Platforms -- Graphic created by Detmar Owen + + + + + + + + + + + + + + + \ No newline at end of file diff --git a/studio/frontend/src/app/provider.tsx b/studio/frontend/src/app/provider.tsx index 91d99e238e..39e04843a9 100644 --- a/studio/frontend/src/app/provider.tsx +++ b/studio/frontend/src/app/provider.tsx @@ -259,8 +259,16 @@ function TauriWrapper({ children }: { children: ReactNode }) { <> {children} - - +
+ + +
); } diff --git a/studio/frontend/src/app/routes/__root.tsx b/studio/frontend/src/app/routes/__root.tsx index c05b0e6451..f5179e68f7 100644 --- a/studio/frontend/src/app/routes/__root.tsx +++ b/studio/frontend/src/app/routes/__root.tsx @@ -6,7 +6,7 @@ import { Navbar } from "@/components/navbar"; import { fetchDeviceType, usePlatformStore } from "@/config/env"; import { SidebarInset, SidebarProvider } from "@/components/ui/sidebar"; import { SettingsDialog, useSettingsDialogStore } from "@/features/settings"; -import { useChatRuntimeStore } from "@/features/chat"; +import { clearNewChatDraft, useChatRuntimeStore } from "@/features/chat"; import { useTrainingUnloadGuard } from "@/features/training"; import { useSidebarPin } from "@/hooks/use-sidebar-pin"; import { useT, type TranslationKey } from "@/i18n"; @@ -15,6 +15,7 @@ import { createRootRoute, redirect, useMatches, + useNavigate, useRouterState, } from "@tanstack/react-router"; import { AnimatePresence, motion } from "motion/react"; @@ -78,6 +79,7 @@ function RootLayout() { const hideNavbar = HIDDEN_NAVBAR_ROUTES.includes(pathname); const isChatRoute = pathname.startsWith("/chat"); const { pinned, setPinned, togglePinned } = useSidebarPin(); + const navigate = useNavigate(); useTrainingUnloadGuard(); @@ -107,11 +109,24 @@ function RootLayout() { if ((e.metaKey || e.ctrlKey) && e.key === ",") { e.preventDefault(); useSettingsDialogStore.getState().openDialog(); + return; + } + // Cmd/Ctrl+Shift+O opens a new chat. + if ((e.metaKey || e.ctrlKey) && e.shiftKey && e.code === "KeyO") { + e.preventDefault(); + clearNewChatDraft(); // fresh chat starts empty, no bleed from the last one + const chatRuntime = useChatRuntimeStore.getState(); + chatRuntime.setActiveThreadId(null); + chatRuntime.setActiveProjectId(null); + void navigate({ + to: "/chat", + search: { new: crypto.randomUUID() }, + }); } }; window.addEventListener("keydown", handler); return () => window.removeEventListener("keydown", handler); - }, []); + }, [navigate]); useEffect(() => { if (isChatRoute) return; diff --git a/studio/frontend/src/components/app-sidebar.tsx b/studio/frontend/src/components/app-sidebar.tsx index f1a200b09b..23fe5e188b 100644 --- a/studio/frontend/src/components/app-sidebar.tsx +++ b/studio/frontend/src/components/app-sidebar.tsx @@ -45,6 +45,8 @@ import { Switch } from "@/components/ui/switch"; import { useAnimatedThemeToggle } from "@/components/ui/animated-theme-toggler"; import { cn } from "@/lib/utils"; import { + Archive01Icon, + ArchiveRestoreIcon, ChefHatIcon, CursorInfo02Icon, DashboardCircleIcon, @@ -64,7 +66,8 @@ import { PowerIcon, PencilEdit02Icon, LayoutAlignLeftIcon, - Settings02Icon, + Setting07Icon, + Sun03Icon, TestTube01Icon, ZapIcon, } from "@hugeicons/core-free-icons"; @@ -80,16 +83,19 @@ import { } from "@/components/ui/tooltip"; import { Tooltip as TooltipPrimitive } from "radix-ui"; import { HugeiconsIcon } from "@hugeicons/react"; -import { ChevronDown, ChevronsUpDown, MoreHorizontalIcon, Moon, Sun } from "lucide-react"; +import { ChevronDown, MoreHorizontalIcon, Moon } from "lucide-react"; import { Link, useNavigate, useRouterState } from "@tanstack/react-router"; import { + archiveChatItem, ChatSearchDialog, + clearNewChatDraft, createChatProject, deleteChatProject, deleteChatItem, moveChatItemToProject, renameChatItem, renameChatProject, + unarchiveChatItem, useChatRuntimeStore, useChatProjects, useChatSearchStore, @@ -291,15 +297,16 @@ export function AppSidebar() { const activeProjectId = isChatRoute ? ((search.project as string | undefined) ?? null) : null; - const { items: allChatItems } = useChatSidebarItems({ - enabled: !isStudioRoute, - requireMessages: false, - }); + const { items: allChatItems, archivedItems: archivedChatItems } = + useChatSidebarItems({ + enabled: !isStudioRoute, + requireMessages: false, + }); const recentChatItems = useMemo( () => allChatItems.filter((item) => !item.projectId), [allChatItems], ); - const chatItems = allChatItems; + const [archivedOpen, setArchivedOpen] = useState(false); const storeThreadId = useChatRuntimeStore((s) => s.activeThreadId); const setActiveThreadId = useChatRuntimeStore((s) => s.setActiveThreadId); const activeThreadId = isChatRoute @@ -349,6 +356,7 @@ export function AppSidebar() { function openNewChat(projectId = activeProjectId) { if (chatDisabled) return; + clearNewChatDraft(); setActiveThreadId(null); useChatRuntimeStore.getState().setActiveProjectId(projectId); navigate({ to: "/chat", search: chatSearchForProject(projectId) }); @@ -374,6 +382,33 @@ export function AppSidebar() { }); } + async function handleArchiveThread(item: SidebarItem) { + try { + await archiveChatItem(item, activeThreadId, (view) => { + navigate({ + to: "/chat", + search: item.projectId + ? { project: item.projectId } + : { new: view.newThreadNonce }, + }); + }); + } catch (err) { + toast.error("Failed to archive chat", { + description: err instanceof Error ? err.message : undefined, + }); + } + } + + async function handleUnarchiveThread(item: SidebarItem) { + try { + await unarchiveChatItem(item); + } catch (err) { + toast.error("Failed to unarchive chat", { + description: err instanceof Error ? err.message : undefined, + }); + } + } + type RenameTarget = | { kind: "chat"; item: SidebarItem; current: string } | { kind: "project"; project: ProjectRecord; current: string } @@ -610,7 +645,7 @@ export function AppSidebar() { side="bottom" align="start" sideOffset={0} - className="unsloth-plus-menu menu-flat-destructive w-52" + className="unsloth-plus-menu menu-flat-destructive w-56" > openRenameChat(item)}> @@ -692,6 +727,10 @@ export function AppSidebar() { + void handleArchiveThread(item)}> + + Archive + setConfirmingDelete({ kind: "chat", item })} @@ -949,7 +988,7 @@ export function AppSidebar() { - + {recentChatItems.map((item) => renderChatSidebarItem(item, "recent"), @@ -961,6 +1000,85 @@ export function AppSidebar() { )} + {/* Archived chats — hidden on Studio + when nothing is archived */} + {!isStudioRoute && archivedChatItems.length > 0 && ( + + + + + Archived + + + + + + + {archivedChatItems.map((item) => ( + + { + navigate({ + to: "/chat", + search: + item.type === "single" + ? { thread: item.id } + : { compare: item.id }, + }); + closeMobileIfOpen(); + }} + > + {item.title} + + + + + + + openRenameChat(item)}> + + Rename + + void handleUnarchiveThread(item)}> + + Unarchive + + setConfirmingDelete({ kind: "chat", item })} + > + + Delete + + + + + ))} + + + + + + )} + {isStudioRoute && runItems.length > 0 && !chatOnly && ( @@ -971,7 +1089,7 @@ export function AppSidebar() { - + {runItems.map((run) => { // Explicit selection wins. Otherwise highlight the active @@ -1092,20 +1210,25 @@ export function AppSidebar() { {displayTitle} Unsloth - + {/* settings cog (replaces the up/down chevron) */} + useSettingsDialogStore.getState().openDialog()} > - + {t("shell.navigation.settings")} ⌘, @@ -1114,7 +1237,7 @@ export function AppSidebar() { > {t("shell.navigation.api")} - + {t("common.new")} @@ -1122,7 +1245,7 @@ export function AppSidebar() { ref={anchorRef as React.Ref} onSelect={(e) => { e.preventDefault(); toggleTheme(); }} > - {isDark ? : } + {isDark ? : } {isDark ? t("shell.navigation.lightMode") diff --git a/studio/frontend/src/components/assistant-ui/markdown-text.tsx b/studio/frontend/src/components/assistant-ui/markdown-text.tsx index ef9b33fced..8f91b419d2 100644 --- a/studio/frontend/src/components/assistant-ui/markdown-text.tsx +++ b/studio/frontend/src/components/assistant-ui/markdown-text.tsx @@ -8,7 +8,8 @@ import { copyToClipboard } from "@/lib/copy-to-clipboard"; import { preprocessLaTeX } from "@/lib/latex"; import { openLink } from "@/lib/open-link"; import { INTERNAL, useAuiState, useMessagePartText } from "@assistant-ui/react"; -import { Copy01Icon, Download01Icon, Tick02Icon } from "@hugeicons/core-free-icons"; +import { Tick02Icon } from "@/lib/tick-icon"; +import { Copy01Icon, Download01Icon } from "@hugeicons/core-free-icons"; import { HugeiconsIcon } from "@hugeicons/react"; import { createMathPlugin } from "@streamdown/math"; import { mermaid } from "@streamdown/mermaid"; @@ -284,9 +285,15 @@ function CodeBlockActions({ ); } +// DiffusionGemma renders its denoising live in the bubble (see DiffusionCanvas in +// thread.tsx) and has the artifacts canvas on by default, so a full-HTML answer +// (e.g. a playable game) renders as an interactive card without the global toggle. function StreamdownBlock(props: BlockProps) { const shouldCollapseHtmlArtifacts = useChatRuntimeStore( - (state) => state.artifactsEnabled || state.collapseHtmlArtifacts, + (state) => + state.artifactsEnabled || + state.collapseHtmlArtifacts || + state.loadedIsDiffusion, ); const messageHasRenderableRenderHtmlTool = useAuiState(({ message }) => message.parts.some(isRenderableRenderHtmlToolPart), diff --git a/studio/frontend/src/components/assistant-ui/message-timing.tsx b/studio/frontend/src/components/assistant-ui/message-timing.tsx index 447fd763e3..8d40f587ad 100644 --- a/studio/frontend/src/components/assistant-ui/message-timing.tsx +++ b/studio/frontend/src/components/assistant-ui/message-timing.tsx @@ -19,6 +19,11 @@ const formatNumber = (n: number): string => { return n.toLocaleString(); }; +const formatRate = (r: number | undefined): string => { + if (r === undefined || !Number.isFinite(r)) return "—"; + return `${Math.round(r).toLocaleString()} tok/s`; +}; + /** * Shows streaming stats as a badge with hover tooltip. * When server timings are available (GGUF), shows prompt eval, generation, @@ -51,6 +56,10 @@ export const MessageTiming: FC<{ st?.cache_n ?? custom?.contextUsage?.cachedTokens ?? 0; // Anthropic-only cache-write count. const cacheWrites = custom?.contextUsage?.cacheWriteTokens ?? 0; + // DiffusionGemma reports separately-labelled throughput (no prefill, so no "prompt + // speed"), matching the CLI: in-step parallel, effective (canvas*blocks/wall), and + // output (answer tokens/wall). + const isDiffusion = (st as { diffusion?: boolean } | undefined)?.diffusion === true; // Guard unphysical tok/s: llama.cpp emits predicted_ms=0 on no-op turns, // blowing the rate up to Infinity. Require >=1 token, a non-zero decode @@ -93,6 +102,99 @@ export const MessageTiming: FC<{ >
{st ? ( + isDiffusion ? ( + <> + {/* DiffusionGemma: honest throughput (no autoregressive prompt speed) */} + {timing.firstTokenTime !== undefined && ( +
+ First token + + {formatTimingMs(timing.firstTokenTime)} + +
+ )} + {st?.diffusion_parallel_tok_s != null && ( +
+ Speed (in-step) + + {formatRate(st.diffusion_parallel_tok_s)} + +
+ )} + {st?.diffusion_effective_tok_s != null && ( +
+ Effective + + {formatRate(st.diffusion_effective_tok_s)} + +
+ )} + {st?.diffusion_output_tok_s != null && ( +
+ Output + + {formatRate(st.diffusion_output_tok_s)} + +
+ )} + {st?.diffusion_steps != null && ( +
+ Denoising + + {formatNumber(st.diffusion_steps)} steps + {st?.diffusion_blocks != null + ? `, ${formatNumber(st.diffusion_blocks)} block${st.diffusion_blocks === 1 ? "" : "s"}` + : ""} + +
+ )} + {st?.diffusion_canvas != null && ( +
+ Canvas + + {formatNumber(st.diffusion_canvas)} tokens + +
+ )} + {(st?.diffusion_wall_ms ?? st?.predicted_ms) != null && ( +
+ Generation + + {formatTimingMs(st.diffusion_wall_ms ?? st.predicted_ms)} + +
+ )} + {timing.tokenCount !== undefined && ( +
+ Answer tokens + + {formatNumber(timing.tokenCount)} + +
+ )} + {(st?.diffusion_prompt_n ?? st?.prompt_n) != null && ( +
+ Prompt + + {formatNumber(st.diffusion_prompt_n ?? st.prompt_n)} tokens + +
+ )} +
+
+ Total + + {formatTimingMs(timing.totalStreamTime)} + +
+
+ Chunks + + {timing.totalChunks} + +
+ + ) : ( <> {/* Server-side metrics (GGUF) */} {st?.prompt_ms != null && ( @@ -135,6 +237,30 @@ export const MessageTiming: FC<{
)} + {timing.firstTokenTime !== undefined && ( +
+ First token + + {formatTimingMs(timing.firstTokenTime)} + +
+ )} + {st?.diffusion_steps != null && ( +
+ Denoising steps + + {formatNumber(st.diffusion_steps)} + +
+ )} + {st?.diffusion_blocks != null && ( +
+ Blocks + + {formatNumber(st.diffusion_blocks)} + +
+ )} {cacheHits > 0 && (
Cache hits @@ -165,6 +291,7 @@ export const MessageTiming: FC<{
+ ) ) : ( <> {/* Client-side metrics (safetensors + external provider fallback) */} diff --git a/studio/frontend/src/components/assistant-ui/model-selector.tsx b/studio/frontend/src/components/assistant-ui/model-selector.tsx index b8d1bc0f00..b97c34152f 100644 --- a/studio/frontend/src/components/assistant-ui/model-selector.tsx +++ b/studio/frontend/src/components/assistant-ui/model-selector.tsx @@ -21,7 +21,9 @@ import { Search01Icon, } from "@hugeicons/core-free-icons"; import { HugeiconsIcon } from "@hugeicons/react"; -import { useMemo, useState } from "react"; +import { type KeyboardEvent, useMemo, useState } from "react"; +import { Input } from "../ui/input"; +import { HubModelPicker, LoraModelPicker } from "./model-selector/pickers"; import type { DeletedModelRef, ExternalModelOption, @@ -29,8 +31,6 @@ import type { ModelOption, ModelSelectorChangeMeta, } from "./model-selector/types"; -import { HubModelPicker, LoraModelPicker } from "./model-selector/pickers"; -import { Input } from "../ui/input"; const PROVIDER_LOGO_EXT: Record = { openai: "svg", @@ -189,7 +189,7 @@ function ModelSelectorTrigger({ @@ -241,13 +241,57 @@ function ModelSelectorContent({ return "hub"; }, [externalModels, loraModels, value]); + function focusActiveModelOption(root: HTMLElement): boolean { + const option = + root.querySelector( + '[role="tabpanel"]:not([hidden]) [data-model-picker-active-option="true"]', + ) ?? + root.querySelector( + '[data-model-picker-active-option="true"]', + ) ?? + root.querySelector( + '[role="tabpanel"]:not([hidden]) [data-model-picker-option]', + ) ?? + root.querySelector( + "[data-model-picker-option]", + ); + if (!option) { + return false; + } + option.focus(); + return true; + } + + function handlePickerEntryKeyDown(event: KeyboardEvent) { + if (event.key !== "ArrowDown") { + return; + } + + const target = event.target; + if (!(target instanceof HTMLElement)) { + return; + } + const isPickerSearchInput = target.matches( + "[data-model-picker-search-input]", + ); + const isTabTrigger = Boolean(target.closest('[role="tab"]')); + if (!isPickerSearchInput && !isTabTrigger) { + return; + } + + if (focusActiveModelOption(event.currentTarget)) { + event.preventDefault(); + } + } + return ( @@ -307,7 +351,7 @@ function ModelSelectorContent({ )} {onPickLocalModel ? ( -
+
-
+
- {!cachedReady && !showHfSection ? ( + {/* First-load spinner only when nothing cached is shown yet. */} + {!cachedReady && + !showHfSection && + visibleCachedGguf.length === 0 && + visibleCachedModelRows.length === 0 ? (
Loading models…
- ) : !showHfSection && - (cachedGguf.length > 0 || - (!chatOnly && cachedModels.length > 0)) ? ( + ) : null} + + {/* Downloaded stays visible (filtered) while searching. */} + {visibleCachedGguf.length > 0 || visibleCachedModelRows.length > 0 ? ( <> } collapsed={downloadedCollapsed} onToggle={() => setDownloadedCollapsed((v) => !v)} >Downloaded - {!downloadedCollapsed && cachedGguf.map((c) => ( -
- - setExpandedGguf((prev) => - prev === c.repo_id ? null : c.repo_id, - ) - } - vramStatus={null} - /> - {expandedGguf === c.repo_id && ( - { - await deleteCachedModel(c.repo_id, quant); - refreshCachedLists(); - }} - /> - )} -
- ))} - {!downloadedCollapsed && !chatOnly && - cachedModels.map((c) => ( -
-
+ {!downloadedCollapsed && + visibleCachedGguf.map((c) => { + const optionKey = makeModelOptionKey("downloaded-gguf", c.repo_id); + return ( +
- onSelect(c.repo_id, { - source: "hub", - isLora: false, - isDownloaded: true, - }) + setExpandedGguf((prev) => + prev === c.repo_id ? null : c.repo_id, + ) + } + onArrowDownIntoChildren={ + expandedGguf === c.repo_id + ? () => { + const focused = focusFirstChildOption(optionKey); + return focused; + } + : undefined } vramStatus={null} /> + {expandedGguf === c.repo_id && ( + + hubModelList.focusOption(optionKey) + } + onNavigatePastEnd={() => + hubModelList.moveFocus(optionKey, "next") + } + gpuGb={gpu.available ? gpu.memoryTotalGb : undefined} + systemRamGb={ + gpu.available ? gpu.systemRamAvailableGb : undefined + } + onDeleteVariant={async (quant) => { + await deleteCachedModel(c.repo_id, quant); + refreshCachedLists(); + }} + /> + )}
- - This will remove{" "} - - {c.repo_id} - {" "} - from disk. You can re-download it later. - - } - successMessage={`Deleted ${c.repo_id}`} - onConfirm={() => deleteCachedModel(c.repo_id)} - onDeleted={refreshCachedLists} - /> -
- ))} + ); + })} + {!downloadedCollapsed && + visibleCachedModelRows.map((c) => { + const optionKey = makeModelOptionKey("downloaded-model", c.repo_id); + return ( +
+
+ + onSelect(c.repo_id, { + source: "hub", + isLora: false, + isDownloaded: true, + }) + } + vramStatus={null} + /> +
+ + This will remove{" "} + + {c.repo_id} + {" "} + from disk. You can re-download it later. + + } + successMessage={`Deleted ${c.repo_id}`} + onConfirm={() => deleteCachedModel(c.repo_id)} + onDeleted={refreshCachedLists} + /> +
+ ); + })} ) : null} @@ -986,6 +1358,7 @@ export function HubModelPicker({ LM Studio {lmStudioModels.map((m) => { const isGguf = isGgufRepo(m.id) || isGgufRepo(m.display_name); + const optionKey = makeModelOptionKey("lm-studio", m.id); return (
{ if (isGguf) { setExpandedGguf((prev) => @@ -1007,12 +1384,27 @@ export function HubModelPicker({ }); } }} + onArrowDownIntoChildren={ + expandedGguf === m.id + ? () => { + const focused = focusFirstChildOption(optionKey); + return focused; + } + : undefined + } vramStatus={null} /> {expandedGguf === m.id && ( + hubModelList.focusOption(optionKey) + } + onNavigatePastEnd={() => + hubModelList.moveFocus(optionKey, "next") + } gpuGb={gpu.available ? gpu.memoryTotalGb : undefined} systemRamGb={ gpu.available ? gpu.systemRamAvailableGb : undefined @@ -1185,12 +1577,17 @@ export function HubModelPicker({ // Single .gguf files (e.g. Ollama blobs) load directly; // GGUF repos/directories expand to pick a variant. const isDirectGguf = isGgufFile; + const optionKey = makeModelOptionKey("custom-folder", m.id); return (
{ if (isDirectGguf) { onSelect(m.id, { @@ -1210,12 +1607,27 @@ export function HubModelPicker({ }); } }} + onArrowDownIntoChildren={ + expandedGguf === m.id + ? () => { + const focused = focusFirstChildOption(optionKey); + return focused; + } + : undefined + } vramStatus={null} /> {expandedGguf === m.id && ( + hubModelList.focusOption(optionKey) + } + onNavigatePastEnd={() => + hubModelList.moveFocus(optionKey, "next") + } gpuGb={gpu.available ? gpu.memoryTotalGb : undefined} systemRamGb={ gpu.available ? gpu.systemRamAvailableGb : undefined @@ -1231,7 +1643,7 @@ export function HubModelPicker({ {!showHfSection && cachedReady ? ( <> } + icon={} collapsed={recommendedCollapsed} onToggle={() => setRecommendedCollapsed((v) => !v)} >Recommended @@ -1242,6 +1654,7 @@ export function HubModelPicker({ ) : ( visibleRecommendedIds.map((id) => { const vram = recommendedVramMap.get(id); + const optionKey = makeModelOptionKey("recommended", id); return (
{ if (isKnownGgufRepo(id)) { setExpandedGguf((prev) => (prev === id ? null : id)); @@ -1264,11 +1681,26 @@ export function HubModelPicker({ } vramEst={isKnownGgufRepo(id) ? undefined : vram?.est} gpuGb={gpu.available ? gpu.memoryTotalGb : undefined} + onArrowDownIntoChildren={ + expandedGguf === id + ? () => { + const focused = focusFirstChildOption(optionKey); + return focused; + } + : undefined + } /> {expandedGguf === id && ( + hubModelList.focusOption(optionKey) + } + onNavigatePastEnd={() => + hubModelList.moveFocus(optionKey, "next") + } gpuGb={gpu.available ? gpu.memoryTotalGb : undefined} systemRamGb={ gpu.available ? gpu.systemRamAvailableGb : undefined @@ -1292,9 +1724,10 @@ export function HubModelPicker({ {showHfSection && filteredRecommendedIds.length > 0 ? ( <> - }>Recommended + }>Recommended {filteredRecommendedIds.map((id) => { const vram = recommendedVramMap.get(id); + const optionKey = makeModelOptionKey("search-recommended", id); return (
{ if (isKnownGgufRepo(id)) { setExpandedGguf((prev) => (prev === id ? null : id)); @@ -1317,11 +1754,26 @@ export function HubModelPicker({ } vramEst={isKnownGgufRepo(id) ? undefined : vram?.est} gpuGb={gpu.available ? gpu.memoryTotalGb : undefined} + onArrowDownIntoChildren={ + expandedGguf === id + ? () => { + const focused = focusFirstChildOption(optionKey); + return focused; + } + : undefined + } /> {expandedGguf === id && ( + hubModelList.focusOption(optionKey) + } + onNavigatePastEnd={() => + hubModelList.moveFocus(optionKey, "next") + } gpuGb={gpu.available ? gpu.memoryTotalGb : undefined} systemRamGb={ gpu.available ? gpu.systemRamAvailableGb : undefined @@ -1340,7 +1792,9 @@ export function HubModelPicker({ Hugging Face )} {hfIds.length === 0 && !isLoading ? ( - filteredRecommendedIds.length === 0 ? ( + filteredRecommendedIds.length === 0 && + visibleCachedGguf.length === 0 && + visibleCachedModelRows.length === 0 ? (
No matching models.
@@ -1349,6 +1803,7 @@ export function HubModelPicker({ hfIds.map((id) => { const vram = vramMap.get(id); const isSearchGguf = isKnownGgufRepo(id); + const optionKey = makeModelOptionKey("search-hf", id); return (
{ if (isSearchGguf) { setExpandedGguf((prev) => (prev === id ? null : id)); @@ -1371,11 +1830,26 @@ export function HubModelPicker({ } vramEst={isSearchGguf ? undefined : vram?.est} gpuGb={gpu.available ? gpu.memoryTotalGb : undefined} + onArrowDownIntoChildren={ + expandedGguf === id + ? () => { + const focused = focusFirstChildOption(optionKey); + return focused; + } + : undefined + } /> {expandedGguf === id && ( + hubModelList.focusOption(optionKey) + } + onNavigatePastEnd={() => + hubModelList.moveFocus(optionKey, "next") + } gpuGb={gpu.available ? gpu.memoryTotalGb : undefined} systemRamGb={ gpu.available ? gpu.systemRamAvailableGb : undefined @@ -1467,6 +1941,26 @@ export function LoraModelPicker({ }); }, [normalized, query]); + const loraOptionKeys = useMemo( + () => + grouped.flatMap(([, adapters]) => + adapters.map((adapter) => makeModelOptionKey("lora", adapter.id)), + ), + [grouped], + ); + const selectedLoraOptionKey = useMemo( + () => + value + ? loraOptionKeys.find((optionKey) => optionKey.endsWith(`::${value}`)) + : undefined, + [loraOptionKeys, value], + ); + const loraModelList = useRovingModelList({ + label: "Fine-tuned models", + optionKeys: loraOptionKeys, + selectedOptionKey: selectedLoraOptionKey, + }); + return (
@@ -1478,11 +1972,12 @@ export function LoraModelPicker({ value={query} onChange={(event) => setQuery(event.target.value)} placeholder="Search trained models" + data-model-picker-search-input={true} className="h-9 border-[#f2f2f2] dark:border-input pl-8" />
-
+
{grouped.length === 0 ? (
@@ -1500,11 +1995,12 @@ export function LoraModelPicker({ const isMerged = adapter.exportType === "merged"; const isGguf = adapter.exportType === "gguf"; const isExportedGguf = isExported && isGguf; - const canDelete = (isTraining || isExported) && !isExportedGguf; + const canDelete = canDeleteLoraModel(adapter); const isTrainingFull = isTraining && isMerged; const isLocalGgufDir = isLocal && (isGgufRepo(adapter.id) || isGgufRepo(adapter.name)); + const optionKey = makeModelOptionKey("lora", adapter.id); const tag = isLocal ? isLocalGgufDir ? "GGUF" @@ -1535,6 +2031,10 @@ export function LoraModelPicker({ label={adapter.name} meta={meta} selected={value === adapter.id} + optionProps={loraModelList.getOptionProps( + optionKey, + value === adapter.id, + )} onClick={() => { if (isLocalGgufDir || isExportedGguf) { setExpandedGguf((prev) => @@ -1562,6 +2062,14 @@ export function LoraModelPicker({ } + onArrowDownIntoChildren={ + expandedGguf === adapter.id + ? () => { + const focused = focusFirstChildOption(optionKey); + return focused; + } + : undefined + } />
{canDelete && ( @@ -1596,6 +2104,13 @@ export function LoraModelPicker({ + loraModelList.focusOption(optionKey) + } + onNavigatePastEnd={() => + loraModelList.moveFocus(optionKey, "next") + } gpuGb={gpu.available ? gpu.memoryTotalGb : undefined} systemRamGb={ gpu.available ? gpu.systemRamAvailableGb : undefined diff --git a/studio/frontend/src/components/assistant-ui/thread.tsx b/studio/frontend/src/components/assistant-ui/thread.tsx index 100143ed33..b63915fa35 100644 --- a/studio/frontend/src/components/assistant-ui/thread.tsx +++ b/studio/frontend/src/components/assistant-ui/thread.tsx @@ -69,6 +69,14 @@ import { getExternalReasoningCapabilities } from "@/features/chat/provider-capab import { useRagToolDisabled } from "@/features/chat/hooks/use-rag-tool-disabled"; import { useChatRuntimeStore } from "@/features/chat/stores/chat-runtime-store"; import { useExternalProvidersStore } from "@/features/chat/stores/external-providers-store"; +import { + PLUS_MENU_ORDER, + composerDraftKey, + readComposerDraft, + type PlusMenuItemId, + usePlusMenuPrefsStore, + writeComposerDraft, +} from "@/features/chat"; import { deleteThreadMessage } from "@/features/chat/utils/delete-thread-message"; import { ThreadDocumentsBar } from "@/features/rag/components/thread-documents-bar"; import { KnowledgeBaseComposerButton } from "@/features/rag/components/knowledge-base-composer-button"; @@ -78,6 +86,7 @@ import { applyQwenThinkingParams } from "@/features/chat/utils/qwen-params"; import { isTauri } from "@/lib/api-base"; import { copyToClipboard } from "@/lib/copy-to-clipboard"; import { toast } from "@/lib/toast"; +import { Tick02Icon } from "@/lib/tick-icon"; import { cn } from "@/lib/utils"; import { ActionBarMorePrimitive, @@ -107,7 +116,6 @@ import { Image03Icon, McpServerIcon, PencilRulerIcon, - Tick02Icon, } from "@hugeicons/core-free-icons"; import { HugeiconsIcon } from "@hugeicons/react"; import { useNavigate } from "@tanstack/react-router"; @@ -134,6 +142,7 @@ import { type KeyboardEvent, type DragEvent as ReactDragEvent, type ReactNode, + Fragment, createContext, useCallback, useContext, @@ -937,6 +946,31 @@ const Composer: FC<{ const referenceThreadId = threadId ?? activeThreadId ?? null; const hasSendableContent = composerText.trim().length > 0 || hasAttachments || hasPendingAudio; + + // Per-thread draft autosave: restore on mount, then mirror composer text + // into localStorage (debounced) so a half-typed message survives a + // navigation or reload. Cleared once empty (i.e. after a send). Setting the + // text even when no draft exists keeps a thread from inheriting the + // previous thread's composer contents. + const draftKey = composerDraftKey(activeThreadId); + const lastDraftKeyRef = useRef(draftKey); + useEffect(() => { + const draft = readComposerDraft(draftKey) ?? ""; + const composer = aui.composer(); + if (composer.getState().isEditing) { + composer.setText(draft); + } + }, [draftKey, aui]); + useEffect(() => { + // After a thread switch composerText can still hold the previous + // thread's text; skip that cycle so it isn't saved under the new key. + if (lastDraftKeyRef.current !== draftKey) { + lastDraftKeyRef.current = draftKey; + return; + } + const t = setTimeout(() => writeComposerDraft(draftKey, composerText), 300); + return () => clearTimeout(t); + }, [composerText, draftKey]); // Two-row layout shows once the input wraps or a tool is on. Tools can // pre-select before a model loads, so an active toggle expands it either way. const composerExpanded = @@ -2093,17 +2127,179 @@ const ComposerToolsMenu: FC<{ side?: "top" | "bottom" }> = ({ const messageCount = useAuiState(({ thread }) => thread.messages.length); const { startQueue } = useContext(PromptQueueContext); + const plusPins = usePlusMenuPrefsStore((s) => s.pins); + const [recentPrompts, setRecentPrompts] = useState([]); const refreshRecentPrompts = useCallback(async () => { try { const rows = await listPromptEntries(); - setRecentPrompts( - [...rows].sort((a, b) => b.updatedAt - a.updatedAt).slice(0, 3), - ); + const byRecent = [...rows].sort((a, b) => b.updatedAt - a.updatedAt); + // Pinned prompts take over the submenu; fall back to the 3 most recent + // when nothing is pinned. + const pinnedIds = usePlusMenuPrefsStore.getState().pinnedPromptIds; + const pinned = byRecent.filter((p) => pinnedIds.includes(p.id)); + setRecentPrompts(pinned.length > 0 ? pinned : byRecent.slice(0, 3)); } catch { } }, []); + // Adjustable "+" menu items, keyed by id. Pinned ones render at the top + // level; the rest fall into the "More" overflow submenu. The core items + // (photos, web search, code) and "More" itself are always shown and live + // outside this map. + const plusMenuNodes: Record = { + chatWithFiles: ( + setRagEnabled(!ragEnabled)} + > + + Chat with Files + {ragEnabled && !ragDisabled ? ( + + ) : null} + + ), + mcp: ( + setMcpEnabledForChat(!mcpEnabledForChat)} + > + + MCP + {mcpEnabledForChat && !mcpDisabled ? ( + + ) : null} + + ), + savedPrompts: ( + + + + Saved prompts + + + {recentPrompts.map((p) => ( + aui.composer().setText(p.text)} + > + {p.name} + + ))} + {recentPrompts.length > 0 ? : null} + setPromptStorageOpen(true)}> + All saved prompts… + + + + ), + compareChat: ( + startCompare()}> + + Compare chat + + ), + exportChat: ( + + + + Export chat + + + { + if (!activeThreadId) return; + exportConversationRawJsonl(activeThreadId).catch(() => + toast.error("Export failed."), + ); + }} + > + Raw JSONL + + { + if (!activeThreadId) return; + exportConversationCsv(activeThreadId).catch(() => + toast.error("Export failed."), + ); + }} + > + CSV + + { + if (!activeThreadId) return; + exportConversationShareGPT(activeThreadId).catch(() => + toast.error("Export failed."), + ); + }} + > + ShareGPT JSONL + + + + ), + canvas: ( + setArtifactsEnabled(!artifactsEnabled)} + > + + Canvas + {artifactsEnabled ? ( + + ) : null} + + ), + projects: ( + + + + Projects + + + setNewProjectOpen(true)}> + + New project + + Recents + {recentProjects.length > 0 ? ( + recentProjects.map((project) => ( + openProject(project.id)} + > + + {project.name} + + )) + ) : ( + + No recent projects + + )} + + + ), + }; + const pinnedPlusItems = PLUS_MENU_ORDER.filter((id) => plusPins[id]); + const overflowPlusItems = PLUS_MENU_ORDER.filter((id) => !plusPins[id]); + return ( <> = ({ align="start" sideOffset={0} avoidCollisions={true} - className="unsloth-plus-menu w-[212px]" + className="unsloth-plus-menu w-[244px]" // Don't refocus the + on close; restored focus showed a stray ring. onCloseAutoFocus={(event) => event.preventDefault()} > @@ -2220,165 +2416,22 @@ const ComposerToolsMenu: FC<{ side?: "top" | "bottom" }> = ({ )} - setRagEnabled(!ragEnabled)} - > - - Chat with Files - {ragEnabled && !ragDisabled ? ( - - ) : null} - - setMcpEnabledForChat(!mcpEnabledForChat)} - > - - MCP - {mcpEnabledForChat && !mcpDisabled ? ( - - ) : null} - - - - - More - - - - - - Saved prompts - - - {recentPrompts.map((p) => ( - aui.composer().setText(p.text)} - > - {p.name} - - ))} - {recentPrompts.length > 0 ? : null} - setPromptStorageOpen(true)}> - All saved prompts… - - - - startCompare()}> - - Compare chat - - - - - Export chat - - - { - if (!activeThreadId) return; - exportConversationRawJsonl(activeThreadId).catch(() => - toast.error("Export failed."), - ); - }} - > - Raw JSONL - - { - if (!activeThreadId) return; - exportConversationCsv(activeThreadId).catch(() => - toast.error("Export failed."), - ); - }} - > - CSV - - { - if (!activeThreadId) return; - exportConversationShareGPT(activeThreadId).catch(() => - toast.error("Export failed."), - ); - }} - > - ShareGPT JSONL - - - - setArtifactsEnabled(!artifactsEnabled)} - > - - Canvas - {artifactsEnabled ? ( - - ) : null} - - - - - - - - Projects - - - setNewProjectOpen(true)}> - - New project - - Recents - {recentProjects.length > 0 ? ( - recentProjects.map((project) => ( - openProject(project.id)} - > - - {project.name} - - )) - ) : ( - - No recent projects - - )} - - + {pinnedPlusItems.map((id) => ( + {plusMenuNodes[id]} + ))} + {overflowPlusItems.length > 0 ? ( + + + + More + + + {overflowPlusItems.map((id) => ( + {plusMenuNodes[id]} + ))} + + + ) : null} { + const isRunning = useAuiState( + ({ message }) => message.status?.type === "running", + ); + // A non-null canvas is set only by diffusion_frame events (diffusion models only), + // so it is a sufficient gate; loadedIsDiffusion can lag the first frame on a fresh load. + const canvas = useChatRuntimeStore((s) => s.activeDiffusionCanvas); + if (!isRunning || !canvas) { + return null; + } + const stepLabel = + canvas.total > 0 ? `step ${canvas.step + 1}/${canvas.total}` : "denoising"; + return ( +
+
+ + Denoising + + block {canvas.block + 1} - {stepLabel} + +
+ 
+        {canvas.text}
+
+
+ ); +}; + const AssistantMessage: FC = () => { return ( {
+ { side="bottom" align="start" onCloseAutoFocus={(e) => e.preventDefault()} - className="aui-action-bar-more-content z-50 min-w-32 overflow-hidden rounded-md [--radius:1.1rem] bg-popover p-1 text-popover-foreground shadow-[0_2px_8px_-2px_rgba(0,0,0,0.16)] dark:shadow-none" + className="aui-action-bar-more-content z-50 min-w-32 overflow-hidden rounded-full bg-popover p-1 text-popover-foreground shadow-[0_2px_8px_-2px_rgba(0,0,0,0.16)] dark:shadow-none" > - + Export as Markdown diff --git a/studio/frontend/src/components/llama-update-banner.tsx b/studio/frontend/src/components/llama-update-banner.tsx index 83320f8e15..8f8bc5ca7f 100644 --- a/studio/frontend/src/components/llama-update-banner.tsx +++ b/studio/frontend/src/components/llama-update-banner.tsx @@ -5,13 +5,81 @@ import { Button } from "@/components/ui/button"; import { useLlamaUpdateCheck } from "@/hooks/use-llama-update-check"; import { useShowLlamaUpdateBanner } from "@/hooks/use-llama-update-pref"; import { toast } from "@/lib/toast"; +import { cn } from "@/lib/utils"; +import { Download } from "lucide-react"; import { AnimatePresence, motion } from "motion/react"; -import type { ReactElement } from "react"; +import { type ReactElement, useEffect, useRef, useState } from "react"; const EASE_OUT_QUART: [number, number, number, number] = [0.165, 0.84, 0.44, 1]; +// Backend progress is coarse (5% steps, ~0.9 max) and the extract tail emits no +// signal. Creep toward this cap so the bar keeps moving rather than freezing. +const RUNNING_CAP = 0.95; + +// Smoothed 0..1 bar progress: eases toward real `progress`, trickles toward a +// ceiling when idle, animates to 100% when `done`. Resets to 0 on each start. +function useSmoothedProgress( + active: boolean, + progress: number | null, + done: boolean, +): number { + const [display, setDisplay] = useState(0); + const displayRef = useRef(0); + const progressRef = useRef(progress); + const doneRef = useRef(done); + progressRef.current = progress; + doneRef.current = done; + + useEffect(() => { + if (!active) { + displayRef.current = 0; + setDisplay(0); + return; + } + let raf = 0; + let last = performance.now(); + const tick = (now: number) => { + // rAF timestamps can predate the performance.now() captured above, so + // clamp dt at 0 to keep the first frame from stepping backwards. + const dt = Math.max(0, Math.min((now - last) / 1000, 0.1)); + last = now; + const current = displayRef.current; + const real = progressRef.current ?? 0; + let target: number; + let speed: number; // approach rate (fraction of remaining gap per second) + if (doneRef.current) { + target = 1; + speed = 5; + } else if (real > current) { + target = real; // catch up to a freshly observed milestone + speed = 4; + } else { + target = RUNNING_CAP; // no signal: creep toward the cap, never frozen + speed = 0.3; + } + const cap = doneRef.current ? 1 : RUNNING_CAP; + const next = Math.min( + current + (target - current) * Math.min(speed * dt, 1), + cap, + ); + displayRef.current = next; + setDisplay(next); + if (doneRef.current && next > 0.999) { + return; + } + raf = requestAnimationFrame(tick); + }; + raf = requestAnimationFrame(tick); + return () => cancelAnimationFrame(raf); + }, [active]); + + return display; +} interface LlamaUpdateBannerProps { enabled?: boolean; + // false: fill the parent instead of self-anchoring, so banners can stack in a + // shared container. true (default) keeps standalone desktop mounts working. + positioned?: boolean; } /** @@ -23,6 +91,7 @@ interface LlamaUpdateBannerProps { */ export function LlamaUpdateBanner({ enabled = true, + positioned = true, }: LlamaUpdateBannerProps): ReactElement | null { const showBannerPref = useShowLlamaUpdateBanner(); const { status, visible, applying, apply, dismiss, snooze } = @@ -46,6 +115,13 @@ export function LlamaUpdateBanner({ const show = visible && status != null && (status.update_available || applying); const updateProgress = status?.job.progress ?? null; + const jobSucceeded = status?.job.state === "success"; + // Drives the bar so it animates continuously; aria reports the real value. + const displayProgress = useSmoothedProgress( + applying, + updateProgress, + jobSucceeded, + ); return ( @@ -55,10 +131,14 @@ export function LlamaUpdateBanner({ animate={{ opacity: 1, y: 0, scale: 1 }} exit={{ opacity: 0, y: 8, scale: 0.97 }} transition={{ duration: 0.35, ease: EASE_OUT_QUART }} - className="fixed bottom-4 right-4 z-[9998] w-[calc(100vw-2rem)] max-w-[340px]" + className={cn( + positioned + ? "fixed bottom-4 right-4 z-[9998] w-[calc(100vw-2rem)] max-w-[400px]" + : "pointer-events-auto w-full", + )} data-testid="llama-update-banner" > -
+
{applying ? null : ( +
+
)}
diff --git a/studio/frontend/src/components/ui/accordion.tsx b/studio/frontend/src/components/ui/accordion.tsx index 7754c78a11..35de233858 100644 --- a/studio/frontend/src/components/ui/accordion.tsx +++ b/studio/frontend/src/components/ui/accordion.tsx @@ -1,98 +1,98 @@ // SPDX-License-Identifier: AGPL-3.0-only // Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0 -"use client"; - -import { Accordion as AccordionPrimitive } from "radix-ui"; -import type * as React from "react"; - -import { cn } from "@/lib/utils"; -import { ArrowDown01Icon, ArrowUp01Icon } from "@hugeicons/core-free-icons"; -import { HugeiconsIcon } from "@hugeicons/react"; - -function Accordion({ - className, - ...props -}: React.ComponentProps) { - return ( - - ); -} - -function AccordionItem({ - className, - ...props -}: React.ComponentProps) { - return ( - - ); -} - -function AccordionTrigger({ - className, - children, - ...props -}: React.ComponentProps) { - return ( - - - {children} - - - - - ); -} - -function AccordionContent({ - className, - children, - ...props -}: React.ComponentProps) { - return ( - -
- {children} -
- - ); -} - -export { Accordion, AccordionItem, AccordionTrigger, AccordionContent }; +"use client"; + +import { Accordion as AccordionPrimitive } from "radix-ui"; +import type * as React from "react"; + +import { cn } from "@/lib/utils"; +import { ArrowDown01Icon, ArrowUp01Icon } from "@hugeicons/core-free-icons"; +import { HugeiconsIcon } from "@hugeicons/react"; + +function Accordion({ + className, + ...props +}: React.ComponentProps) { + return ( + + ); +} + +function AccordionItem({ + className, + ...props +}: React.ComponentProps) { + return ( + + ); +} + +function AccordionTrigger({ + className, + children, + ...props +}: React.ComponentProps) { + return ( + + + {children} + + + + + ); +} + +function AccordionContent({ + className, + children, + ...props +}: React.ComponentProps) { + return ( + +
+ {children} +
+ + ); +} + +export { Accordion, AccordionItem, AccordionTrigger, AccordionContent }; diff --git a/studio/frontend/src/components/ui/alert-dialog.tsx b/studio/frontend/src/components/ui/alert-dialog.tsx index f5c1dbacca..97f4be7f44 100644 --- a/studio/frontend/src/components/ui/alert-dialog.tsx +++ b/studio/frontend/src/components/ui/alert-dialog.tsx @@ -1,50 +1,50 @@ // SPDX-License-Identifier: AGPL-3.0-only // Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0 -import { AlertDialog as AlertDialogPrimitive } from "radix-ui"; -import type * as React from "react"; - -import { Button } from "@/components/ui/button"; -import { cn } from "@/lib/utils"; - -function AlertDialog({ - ...props -}: React.ComponentProps) { - return ; -} - -function AlertDialogTrigger({ - ...props -}: React.ComponentProps) { - return ( - - ); -} - -function AlertDialogPortal({ - ...props -}: React.ComponentProps) { - return ( - - ); -} - -function AlertDialogOverlay({ - className, - ...props -}: React.ComponentProps) { - return ( - - ); -} - +import { AlertDialog as AlertDialogPrimitive } from "radix-ui"; +import type * as React from "react"; + +import { Button } from "@/components/ui/button"; +import { cn } from "@/lib/utils"; + +function AlertDialog({ + ...props +}: React.ComponentProps) { + return ; +} + +function AlertDialogTrigger({ + ...props +}: React.ComponentProps) { + return ( + + ); +} + +function AlertDialogPortal({ + ...props +}: React.ComponentProps) { + return ( + + ); +} + +function AlertDialogOverlay({ + className, + ...props +}: React.ComponentProps) { + return ( + + ); +} + function AlertDialogContent({ className, size = "default", @@ -60,143 +60,143 @@ function AlertDialogContent({ - - ); -} - -function AlertDialogHeader({ - className, - ...props -}: React.ComponentProps<"div">) { - return ( -
- ); -} - -function AlertDialogFooter({ - className, - ...props -}: React.ComponentProps<"div">) { - return ( -
- ); -} - -function AlertDialogMedia({ - className, - ...props -}: React.ComponentProps<"div">) { - return ( -
- ); -} - -function AlertDialogTitle({ - className, - ...props -}: React.ComponentProps) { - return ( - - ); -} - -function AlertDialogDescription({ - className, - ...props -}: React.ComponentProps) { - return ( - - ); -} - -function AlertDialogAction({ - className, - variant = "default", - size = "default", - ...props -}: React.ComponentProps & - Pick, "variant" | "size">) { - return ( - - ); -} - -function AlertDialogCancel({ - className, - variant = "outline", - size = "default", - ...props -}: React.ComponentProps & - Pick, "variant" | "size">) { - return ( - - ); -} - -export { - AlertDialog, - AlertDialogAction, - AlertDialogCancel, - AlertDialogContent, - AlertDialogDescription, - AlertDialogFooter, - AlertDialogHeader, - AlertDialogMedia, - AlertDialogOverlay, - AlertDialogPortal, - AlertDialogTitle, - AlertDialogTrigger, -}; + className={cn( + "data-open:animate-in data-closed:animate-out data-closed:fade-out-0 data-open:fade-in-0 data-closed:zoom-out-95 data-open:zoom-in-95 bg-background ring-foreground/5 gap-6 rounded-4xl p-6 ring-1 duration-100 data-[size=default]:max-w-xs data-[size=sm]:max-w-xs data-[size=default]:sm:max-w-md group/alert-dialog-content fixed top-1/2 left-1/2 z-50 grid w-full -translate-x-1/2 -translate-y-1/2 outline-none", + className, + )} + {...props} + /> + + ); +} + +function AlertDialogHeader({ + className, + ...props +}: React.ComponentProps<"div">) { + return ( +
+ ); +} + +function AlertDialogFooter({ + className, + ...props +}: React.ComponentProps<"div">) { + return ( +
+ ); +} + +function AlertDialogMedia({ + className, + ...props +}: React.ComponentProps<"div">) { + return ( +
+ ); +} + +function AlertDialogTitle({ + className, + ...props +}: React.ComponentProps) { + return ( + + ); +} + +function AlertDialogDescription({ + className, + ...props +}: React.ComponentProps) { + return ( + + ); +} + +function AlertDialogAction({ + className, + variant = "default", + size = "default", + ...props +}: React.ComponentProps & + Pick, "variant" | "size">) { + return ( + + ); +} + +function AlertDialogCancel({ + className, + variant = "outline", + size = "default", + ...props +}: React.ComponentProps & + Pick, "variant" | "size">) { + return ( + + ); +} + +export { + AlertDialog, + AlertDialogAction, + AlertDialogCancel, + AlertDialogContent, + AlertDialogDescription, + AlertDialogFooter, + AlertDialogHeader, + AlertDialogMedia, + AlertDialogOverlay, + AlertDialogPortal, + AlertDialogTitle, + AlertDialogTrigger, +}; diff --git a/studio/frontend/src/components/ui/alert.tsx b/studio/frontend/src/components/ui/alert.tsx index a4a5f4c4b7..094b607d9a 100644 --- a/studio/frontend/src/components/ui/alert.tsx +++ b/studio/frontend/src/components/ui/alert.tsx @@ -1,79 +1,79 @@ // SPDX-License-Identifier: AGPL-3.0-only // Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0 -import { type VariantProps, cva } from "class-variance-authority"; -import type * as React from "react"; - -import { cn } from "@/lib/utils"; - -const alertVariants = cva( - "grid gap-0.5 rounded-lg border px-4 py-3 text-left text-sm has-data-[slot=alert-action]:relative has-data-[slot=alert-action]:pr-18 has-[>svg]:grid-cols-[auto_1fr] has-[>svg]:gap-x-2.5 *:[svg]:row-span-2 *:[svg]:translate-y-0.5 *:[svg]:text-current *:[svg:not([class*='size-'])]:size-4 w-full relative group/alert", - { - variants: { - variant: { - default: "bg-card text-card-foreground", - destructive: - "text-destructive bg-card *:data-[slot=alert-description]:text-destructive/90 *:[svg]:text-current", - }, - }, - defaultVariants: { - variant: "default", - }, - }, -); - -function Alert({ - className, - variant, - ...props -}: React.ComponentProps<"div"> & VariantProps) { - return ( -
- ); -} - -function AlertTitle({ className, ...props }: React.ComponentProps<"div">) { - return ( -
svg]/alert:col-start-2 [&_a]:hover:text-foreground [&_a]:underline [&_a]:underline-offset-3", - className, - )} - {...props} - /> - ); -} - -function AlertDescription({ - className, - ...props -}: React.ComponentProps<"div">) { - return ( -
- ); -} - -function AlertAction({ className, ...props }: React.ComponentProps<"div">) { - return ( -
- ); -} - -export { Alert, AlertTitle, AlertDescription, AlertAction }; +import { type VariantProps, cva } from "class-variance-authority"; +import type * as React from "react"; + +import { cn } from "@/lib/utils"; + +const alertVariants = cva( + "grid gap-0.5 rounded-lg border px-4 py-3 text-left text-sm has-data-[slot=alert-action]:relative has-data-[slot=alert-action]:pr-18 has-[>svg]:grid-cols-[auto_1fr] has-[>svg]:gap-x-2.5 *:[svg]:row-span-2 *:[svg]:translate-y-0.5 *:[svg]:text-current *:[svg:not([class*='size-'])]:size-4 w-full relative group/alert", + { + variants: { + variant: { + default: "bg-card text-card-foreground", + destructive: + "text-destructive bg-card *:data-[slot=alert-description]:text-destructive/90 *:[svg]:text-current", + }, + }, + defaultVariants: { + variant: "default", + }, + }, +); + +function Alert({ + className, + variant, + ...props +}: React.ComponentProps<"div"> & VariantProps) { + return ( +
+ ); +} + +function AlertTitle({ className, ...props }: React.ComponentProps<"div">) { + return ( +
svg]/alert:col-start-2 [&_a]:hover:text-foreground [&_a]:underline [&_a]:underline-offset-3", + className, + )} + {...props} + /> + ); +} + +function AlertDescription({ + className, + ...props +}: React.ComponentProps<"div">) { + return ( +
+ ); +} + +function AlertAction({ className, ...props }: React.ComponentProps<"div">) { + return ( +
+ ); +} + +export { Alert, AlertTitle, AlertDescription, AlertAction }; diff --git a/studio/frontend/src/components/ui/animated-shiny-text.tsx b/studio/frontend/src/components/ui/animated-shiny-text.tsx index 4c650f1003..8d366ca3d6 100644 --- a/studio/frontend/src/components/ui/animated-shiny-text.tsx +++ b/studio/frontend/src/components/ui/animated-shiny-text.tsx @@ -1,41 +1,41 @@ // SPDX-License-Identifier: AGPL-3.0-only // Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0 -import type { ComponentPropsWithoutRef, CSSProperties, FC } from "react" - -import { cn } from "@/lib/utils" - -export interface AnimatedShinyTextProps extends ComponentPropsWithoutRef<"span"> { - shimmerWidth?: number -} - -export const AnimatedShinyText: FC = ({ - children, - className, - shimmerWidth = 100, - ...props -}) => { - return ( - - {children} - - ) -} +import type { ComponentPropsWithoutRef, CSSProperties, FC } from "react" + +import { cn } from "@/lib/utils" + +export interface AnimatedShinyTextProps extends ComponentPropsWithoutRef<"span"> { + shimmerWidth?: number +} + +export const AnimatedShinyText: FC = ({ + children, + className, + shimmerWidth = 100, + ...props +}) => { + return ( + + {children} + + ) +} diff --git a/studio/frontend/src/components/ui/aspect-ratio.tsx b/studio/frontend/src/components/ui/aspect-ratio.tsx index cb605f01eb..2471f4333d 100644 --- a/studio/frontend/src/components/ui/aspect-ratio.tsx +++ b/studio/frontend/src/components/ui/aspect-ratio.tsx @@ -1,12 +1,12 @@ // SPDX-License-Identifier: AGPL-3.0-only // Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0 -import { AspectRatio as AspectRatioPrimitive } from "radix-ui"; - -function AspectRatio({ - ...props -}: React.ComponentProps) { - return ; -} - -export { AspectRatio }; +import { AspectRatio as AspectRatioPrimitive } from "radix-ui"; + +function AspectRatio({ + ...props +}: React.ComponentProps) { + return ; +} + +export { AspectRatio }; diff --git a/studio/frontend/src/components/ui/avatar.tsx b/studio/frontend/src/components/ui/avatar.tsx index 2250bb849a..31262b32f7 100644 --- a/studio/frontend/src/components/ui/avatar.tsx +++ b/studio/frontend/src/components/ui/avatar.tsx @@ -1,113 +1,113 @@ // SPDX-License-Identifier: AGPL-3.0-only // Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0 -import { Avatar as AvatarPrimitive } from "radix-ui"; -import type * as React from "react"; - -import { cn } from "@/lib/utils"; - -function Avatar({ - className, - size = "default", - ...props -}: React.ComponentProps & { - size?: "default" | "sm" | "lg"; -}) { - return ( - - ); -} - -function AvatarImage({ - className, - ...props -}: React.ComponentProps) { - return ( - - ); -} - -function AvatarFallback({ - className, - ...props -}: React.ComponentProps) { - return ( - - ); -} - -function AvatarBadge({ className, ...props }: React.ComponentProps<"span">) { - return ( - svg]:hidden", - "group-data-[size=default]/avatar:size-2.5 group-data-[size=default]/avatar:[&>svg]:size-2", - "group-data-[size=lg]/avatar:size-3 group-data-[size=lg]/avatar:[&>svg]:size-2", - className, - )} - {...props} - /> - ); -} - -function AvatarGroup({ className, ...props }: React.ComponentProps<"div">) { - return ( -
- ); -} - -function AvatarGroupCount({ - className, - ...props -}: React.ComponentProps<"div">) { - return ( -
svg]:size-4 group-has-data-[size=lg]/avatar-group:[&>svg]:size-5 group-has-data-[size=sm]/avatar-group:[&>svg]:size-3 ring-background relative flex shrink-0 items-center justify-center ring-2", - className, - )} - {...props} - /> - ); -} - -export { - Avatar, - AvatarImage, - AvatarFallback, - AvatarGroup, - AvatarGroupCount, - AvatarBadge, -}; +import { Avatar as AvatarPrimitive } from "radix-ui"; +import type * as React from "react"; + +import { cn } from "@/lib/utils"; + +function Avatar({ + className, + size = "default", + ...props +}: React.ComponentProps & { + size?: "default" | "sm" | "lg"; +}) { + return ( + + ); +} + +function AvatarImage({ + className, + ...props +}: React.ComponentProps) { + return ( + + ); +} + +function AvatarFallback({ + className, + ...props +}: React.ComponentProps) { + return ( + + ); +} + +function AvatarBadge({ className, ...props }: React.ComponentProps<"span">) { + return ( + svg]:hidden", + "group-data-[size=default]/avatar:size-2.5 group-data-[size=default]/avatar:[&>svg]:size-2", + "group-data-[size=lg]/avatar:size-3 group-data-[size=lg]/avatar:[&>svg]:size-2", + className, + )} + {...props} + /> + ); +} + +function AvatarGroup({ className, ...props }: React.ComponentProps<"div">) { + return ( +
+ ); +} + +function AvatarGroupCount({ + className, + ...props +}: React.ComponentProps<"div">) { + return ( +
svg]:size-4 group-has-data-[size=lg]/avatar-group:[&>svg]:size-5 group-has-data-[size=sm]/avatar-group:[&>svg]:size-3 ring-background relative flex shrink-0 items-center justify-center ring-2", + className, + )} + {...props} + /> + ); +} + +export { + Avatar, + AvatarImage, + AvatarFallback, + AvatarGroup, + AvatarGroupCount, + AvatarBadge, +}; diff --git a/studio/frontend/src/components/ui/badge.tsx b/studio/frontend/src/components/ui/badge.tsx index 3951ae9de0..0f2f334986 100644 --- a/studio/frontend/src/components/ui/badge.tsx +++ b/studio/frontend/src/components/ui/badge.tsx @@ -1,54 +1,54 @@ // SPDX-License-Identifier: AGPL-3.0-only // Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0 -/* eslint-disable react-refresh/only-export-components */ - -import { type VariantProps, cva } from "class-variance-authority"; -import { Slot } from "radix-ui"; -import type * as React from "react"; - -import { cn } from "@/lib/utils"; - -export const badgeVariants = cva( - "h-5 gap-1 rounded-4xl border border-transparent px-2 py-0.5 text-xs font-medium transition-all has-data-[icon=inline-end]:pr-1.5 has-data-[icon=inline-start]:pl-1.5 [&>svg]:size-3! inline-flex items-center justify-center w-fit whitespace-nowrap shrink-0 [&>svg]:pointer-events-none focus-visible:border-ring focus-visible:ring-ring/50 focus-visible:ring-[3px] aria-invalid:ring-destructive/20 dark:aria-invalid:ring-destructive/40 aria-invalid:border-destructive overflow-hidden group/badge", - { - variants: { - variant: { - default: "bg-primary text-primary-foreground [a]:hover:bg-primary/80", - secondary: - "bg-secondary text-secondary-foreground [a]:hover:bg-secondary/80", - destructive: - "bg-destructive/10 [a]:hover:bg-destructive/20 focus-visible:ring-destructive/20 dark:focus-visible:ring-destructive/40 text-destructive dark:bg-destructive/20", - outline: - "border-border text-foreground [a]:hover:bg-muted [a]:hover:text-muted-foreground bg-input/30", - ghost: - "hover:bg-muted hover:text-muted-foreground dark:hover:bg-muted/50", - link: "text-primary underline-offset-4 hover:underline", - }, - }, - defaultVariants: { - variant: "default", - }, - }, -); - -export function Badge({ - className, - variant = "default", - asChild = false, - ...props -}: React.ComponentProps<"span"> & - VariantProps & { - asChild?: boolean; - }): React.ReactElement { - const Comp = asChild ? Slot.Root : "span"; - - return ( - - ); -} +/* eslint-disable react-refresh/only-export-components */ + +import { type VariantProps, cva } from "class-variance-authority"; +import { Slot } from "radix-ui"; +import type * as React from "react"; + +import { cn } from "@/lib/utils"; + +export const badgeVariants = cva( + "h-5 gap-1 rounded-4xl border border-transparent px-2 py-0.5 text-xs font-medium transition-all has-data-[icon=inline-end]:pr-1.5 has-data-[icon=inline-start]:pl-1.5 [&>svg]:size-3! inline-flex items-center justify-center w-fit whitespace-nowrap shrink-0 [&>svg]:pointer-events-none focus-visible:border-ring focus-visible:ring-ring/50 focus-visible:ring-[3px] aria-invalid:ring-destructive/20 dark:aria-invalid:ring-destructive/40 aria-invalid:border-destructive overflow-hidden group/badge", + { + variants: { + variant: { + default: "bg-primary text-primary-foreground [a]:hover:bg-primary/80", + secondary: + "bg-secondary text-secondary-foreground [a]:hover:bg-secondary/80", + destructive: + "bg-destructive/10 [a]:hover:bg-destructive/20 focus-visible:ring-destructive/20 dark:focus-visible:ring-destructive/40 text-destructive dark:bg-destructive/20", + outline: + "border-border text-foreground [a]:hover:bg-muted [a]:hover:text-muted-foreground bg-input/30", + ghost: + "hover:bg-muted hover:text-muted-foreground dark:hover:bg-muted/50", + link: "text-primary underline-offset-4 hover:underline", + }, + }, + defaultVariants: { + variant: "default", + }, + }, +); + +export function Badge({ + className, + variant = "default", + asChild = false, + ...props +}: React.ComponentProps<"span"> & + VariantProps & { + asChild?: boolean; + }): React.ReactElement { + const Comp = asChild ? Slot.Root : "span"; + + return ( + + ); +} diff --git a/studio/frontend/src/components/ui/breadcrumb.tsx b/studio/frontend/src/components/ui/breadcrumb.tsx index dc026994ce..a2dad8783f 100644 --- a/studio/frontend/src/components/ui/breadcrumb.tsx +++ b/studio/frontend/src/components/ui/breadcrumb.tsx @@ -1,126 +1,126 @@ // SPDX-License-Identifier: AGPL-3.0-only // Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0 -import { Slot } from "radix-ui"; -import type * as React from "react"; - -import { cn } from "@/lib/utils"; -import { - ArrowRight01Icon, - MoreHorizontalCircle01Icon, -} from "@hugeicons/core-free-icons"; -import { HugeiconsIcon } from "@hugeicons/react"; - -function Breadcrumb({ className, ...props }: React.ComponentProps<"nav">) { - return ( -