diff --git a/.github/dependabot.yml b/.github/dependabot.yml index a06cb1d114..4a0bfa70f1 100644 --- a/.github/dependabot.yml +++ b/.github/dependabot.yml @@ -9,19 +9,6 @@ updates: actions: patterns: ["*"] - - package-ecosystem: "pip" - directories: - - "/" - - "/studio/backend/plugins/data-designer-unstructured-seed" - - "/studio/backend/requirements" - - "/unsloth/kernels/moe" - schedule: - interval: "weekly" - open-pull-requests-limit: 10 - groups: - pip: - patterns: ["*"] - - package-ecosystem: "bun" directory: "/studio/frontend" schedule: diff --git a/.pre-commit-config.yaml b/.pre-commit-config.yaml index f3b5987a9c..309c7d9aca 100644 --- a/.pre-commit-config.yaml +++ b/.pre-commit-config.yaml @@ -1,6 +1,6 @@ repos: - repo: https://github.com/astral-sh/ruff-pre-commit - rev: v0.15.8 + rev: v0.15.10 hooks: - id: ruff args: diff --git a/install.ps1 b/install.ps1 index 5ddb42ea7e..61f31d85b4 100644 --- a/install.ps1 +++ b/install.ps1 @@ -754,7 +754,7 @@ shell.Run cmd, 0, False # ── Choose the correct PyTorch index URL based on driver CUDA version ── # Mirrors Get-PytorchCudaTag in setup.ps1. function Get-TorchIndexUrl { - $baseUrl = "https://download.pytorch.org/whl" + $baseUrl = if ($env:UNSLOTH_PYTORCH_MIRROR) { $env:UNSLOTH_PYTORCH_MIRROR.TrimEnd('/') } else { "https://download.pytorch.org/whl" } if (-not $NvidiaSmiExe) { return "$baseUrl/cpu" } try { $output = & $NvidiaSmiExe 2>&1 | Out-String @@ -819,7 +819,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.4.2" unsloth-zoo } + $baseInstallExit = Invoke-InstallCommand { uv pip install --python $VenvPython --no-deps --reinstall-package unsloth --reinstall-package unsloth-zoo "unsloth>=2026.4.4" unsloth-zoo } if ($baseInstallExit -eq 0) { $NoTorchReq = Find-NoTorchRuntimeFile if ($NoTorchReq) { @@ -827,7 +827,7 @@ shell.Run cmd, 0, False } } } else { - $baseInstallExit = Invoke-InstallCommand { uv pip install --python $VenvPython --reinstall-package unsloth --reinstall-package unsloth-zoo "unsloth>=2026.4.2" unsloth-zoo } + $baseInstallExit = Invoke-InstallCommand { uv pip install --python $VenvPython --reinstall-package unsloth --reinstall-package unsloth-zoo "unsloth>=2026.4.4" unsloth-zoo } } if ($baseInstallExit -ne 0) { Write-Host "[ERROR] Failed to install unsloth (exit code $baseInstallExit)" -ForegroundColor Red @@ -857,7 +857,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.4.2" unsloth-zoo } + $baseInstallExit = Invoke-InstallCommand { uv pip install --python $VenvPython --no-deps --upgrade-package unsloth --upgrade-package unsloth-zoo "unsloth>=2026.4.4" unsloth-zoo } if ($baseInstallExit -eq 0) { $NoTorchReq = Find-NoTorchRuntimeFile if ($NoTorchReq) { @@ -865,7 +865,7 @@ shell.Run cmd, 0, False } } } elseif ($StudioLocalInstall) { - $baseInstallExit = Invoke-InstallCommand { uv pip install --python $VenvPython --upgrade-package unsloth "unsloth>=2026.4.2" unsloth-zoo } + $baseInstallExit = Invoke-InstallCommand { uv pip install --python $VenvPython --upgrade-package unsloth "unsloth>=2026.4.4" unsloth-zoo } } else { $baseInstallExit = Invoke-InstallCommand { uv pip install --python $VenvPython --upgrade-package unsloth "$PackageName" } } @@ -886,7 +886,7 @@ shell.Run cmd, 0, False # Fallback: GPU detection failed to produce a URL -- let uv resolve torch substep "installing unsloth (this may take a few minutes)..." if ($StudioLocalInstall) { - $baseInstallExit = Invoke-InstallCommand { uv pip install --python $VenvPython unsloth-zoo "unsloth>=2026.4.2" --torch-backend=auto } + $baseInstallExit = Invoke-InstallCommand { uv pip install --python $VenvPython unsloth-zoo "unsloth>=2026.4.4" --torch-backend=auto } if ($baseInstallExit -ne 0) { Write-Host "[ERROR] Failed to install unsloth (exit code $baseInstallExit)" -ForegroundColor Red return diff --git a/install.sh b/install.sh index 053f334d2b..0dbcdf380e 100755 --- a/install.sh +++ b/install.sh @@ -94,6 +94,45 @@ run_install_cmd() { return $_rc } +# Install bitsandbytes on AMD ROCm hosts. Uses the continuous-release_main +# wheel for the ROCm 4-bit GEMV fix (bnb PR #1887, post-0.49.2); bnb <= 0.49.2 +# NaNs at decode shape on every AMD GPU. Falls back to PyPI >=0.49.1 if the +# pre-release URL is unreachable. Drop the pin once bnb 0.50+ ships on PyPI. +_install_bnb_rocm() { + _label="$1" + _venv_py="$2" + case "$_ARCH" in + x86_64|amd64) + _bnb_whl_url="https://github.com/bitsandbytes-foundation/bitsandbytes/releases/download/continuous-release_main/bitsandbytes-1.33.7.preview-py3-none-manylinux_2_24_x86_64.whl" + ;; + aarch64|arm64) + _bnb_whl_url="https://github.com/bitsandbytes-foundation/bitsandbytes/releases/download/continuous-release_main/bitsandbytes-1.33.7.preview-py3-none-manylinux_2_24_aarch64.whl" + ;; + *) + _bnb_whl_url="" + ;; + esac + # uv rejects the continuous-release_main bitsandbytes wheel because the + # filename version (1.33.7rc0) does not match the embedded metadata version + # (0.50.0.dev0). pip accepts the mismatch, so bootstrap pip and use it. + if ! "$_venv_py" -m pip --version >/dev/null 2>&1; then + if ! run_maybe_quiet "$_venv_py" -m ensurepip --upgrade; then + run_maybe_quiet uv pip install --python "$_venv_py" pip || \ + substep "[WARN] could not bootstrap pip; bitsandbytes install will likely fail" "$C_WARN" + fi + fi + if [ -n "$_bnb_whl_url" ]; then + substep "installing bitsandbytes for AMD ROCm (pre-release, PR #1887)..." + if run_install_cmd "$_label (pre-release)" "$_venv_py" -m pip install \ + --force-reinstall --no-cache-dir --no-deps "$_bnb_whl_url"; then + return 0 + fi + substep "[WARN] bnb pre-release install failed; falling back to PyPI (4-bit decode broken on ROCm)" "$C_WARN" + fi + run_install_cmd "$_label (pypi fallback)" "$_venv_py" -m pip install \ + --force-reinstall --no-cache-dir --no-deps "bitsandbytes>=0.49.1" +} + if [ "$_next_is_package" = true ]; then echo "❌ ERROR: --package requires an argument." >&2 exit 1 @@ -978,22 +1017,122 @@ _find_no_torch_runtime() { fi } +# ── AMD ROCm GPU detection helper ── +# Returns 0 (true) if an actual AMD GPU is present, 1 (false) otherwise. +# Checks rocminfo for gfx[1-9]* (excludes gfx000 CPU agent) and +# amd-smi list for GPU data rows (excludes header-only output). +_has_amd_rocm_gpu() { + if command -v rocminfo >/dev/null 2>&1 && \ + rocminfo 2>/dev/null | awk '/Name:[[:space:]]*gfx[0-9]/ && !/Name:[[:space:]]*gfx000/{found=1} END{exit !found}'; then + return 0 + elif command -v amd-smi >/dev/null 2>&1 && \ + amd-smi list 2>/dev/null | awk '/^GPU[[:space:]]*[:\[][[:space:]]*[0-9]/{ found=1 } END{ exit !found }'; then + return 0 + fi + return 1 +} + +# ── NVIDIA usable-GPU helper ── +# Returns 0 (true) only if nvidia-smi is present AND actually lists a GPU. +# Prevents AMD-only hosts with a stale nvidia-smi on PATH from being routed +# into the CUDA branch. +_has_usable_nvidia_gpu() { + _nvsmi="" + if command -v nvidia-smi >/dev/null 2>&1; then + _nvsmi="nvidia-smi" + elif [ -x "/usr/bin/nvidia-smi" ]; then + _nvsmi="/usr/bin/nvidia-smi" + else + return 1 + fi + "$_nvsmi" -L 2>/dev/null | awk '/^GPU[[:space:]]+[0-9]+:/{found=1} END{exit !found}' +} + # ── Detect GPU and choose PyTorch index URL ── # Mirrors Get-TorchIndexUrl in install.ps1. # On CPU-only machines this returns the cpu index, avoiding the solver # dead-end where --torch-backend=auto resolves to unsloth==2024.8. get_torch_index_url() { - _base="https://download.pytorch.org/whl" + _base="${UNSLOTH_PYTORCH_MIRROR:-https://download.pytorch.org/whl}" + _base="${_base%/}" # macOS: always CPU (no CUDA support) case "$(uname -s)" in Darwin) echo "$_base/cpu"; return ;; esac - # Try nvidia-smi + # Try nvidia-smi -- require the binary to actually list a usable GPU. + # Presence of the binary alone (container leftovers, stale driver + # packages) is not sufficient: otherwise an AMD-only host would + # silently install CUDA wheels. _smi="" - if command -v nvidia-smi >/dev/null 2>&1; then - _smi="nvidia-smi" - elif [ -x "/usr/bin/nvidia-smi" ]; then - _smi="/usr/bin/nvidia-smi" + if _has_usable_nvidia_gpu; then + if command -v nvidia-smi >/dev/null 2>&1; then + _smi="nvidia-smi" + elif [ -x "/usr/bin/nvidia-smi" ]; then + _smi="/usr/bin/nvidia-smi" + fi + fi + if [ -z "$_smi" ]; then + # No NVIDIA GPU -- check for AMD ROCm GPU. + # PyTorch only publishes ROCm wheels for linux-x86_64; skip the + # ROCm branch entirely on aarch64 / arm64 / other architectures + # so non-x86_64 Linux hosts fall back cleanly to CPU wheels. + case "$(uname -m)" in + x86_64|amd64) : ;; + *) echo "$_base/cpu"; return ;; + esac + if ! _has_amd_rocm_gpu; then + echo "$_base/cpu"; return + fi + # AMD GPU confirmed -- detect ROCm version + _rocm_tag="" + _rocm_tag=$({ command -v amd-smi >/dev/null 2>&1 && \ + amd-smi version 2>/dev/null | awk -F'ROCm version: ' \ + 'NF>1{gsub(/[^0-9.]/, "", $2); split($2,a,"."); print "rocm"a[1]"."a[2]; ok=1; exit} END{exit !ok}'; } || \ + { [ -r /opt/rocm/.info/version ] && \ + awk -F. '{print "rocm"$1"."$2; exit}' /opt/rocm/.info/version; } || \ + { command -v hipconfig >/dev/null 2>&1 && \ + hipconfig --version 2>/dev/null | awk 'NR==1 && /^[0-9]/{split($1,a,"."); if(a[1]+0>0){print "rocm"a[1]"."a[2]; found=1}} END{exit !found}'; } || \ + { command -v dpkg-query >/dev/null 2>&1 && \ + ver="$(dpkg-query -W -f='${Version}\n' rocm-core 2>/dev/null)" && \ + [ -n "$ver" ] && \ + printf '%s\n' "$ver" | sed 's/^[0-9]*://' | awk -F'[.-]' '{print "rocm"$1"."$2; exit}'; } || \ + { command -v rpm >/dev/null 2>&1 && \ + ver="$(rpm -q --qf '%{VERSION}\n' rocm-core 2>/dev/null)" && \ + [ -n "$ver" ] && \ + printf '%s\n' "$ver" | awk -F'[.-]' '{print "rocm"$1"."$2; exit}'; }) 2>/dev/null + # Validate _rocm_tag: must match "rocmX.Y" with major >= 1 + case "$_rocm_tag" in + rocm[1-9]*.[0-9]*) : ;; # valid (major >= 1) + *) _rocm_tag="" ;; # reject malformed (empty, garbled, or major=0) + esac + if [ -n "$_rocm_tag" ]; then + # Minimum supported: ROCm 6.0 (no PyTorch wheels exist for older) + case "$_rocm_tag" in + rocm[1-5].*) echo "$_base/cpu"; return ;; + esac + # ROCm 7.2 only has torch 2.11.0 which exceeds current bounds + # (<2.11.0). Fall back to rocm7.1 index which has torch 2.10.0. + # Enumerate explicit versions rather than matching rocm6.* so + # a host on ROCm 6.5 or 6.6 (no PyTorch wheels published) is + # clipped down to the last supported 6.x (rocm6.4) instead of + # constructing https://download.pytorch.org/whl/rocm6.5 which + # returns HTTP 403. PyTorch only ships: rocm5.7, 6.0, 6.1, 6.2, + # 6.3, 6.4, 7.0, 7.1, 7.2 (and 5.7 is below our minimum). + # TODO: uncomment rocm7.2 when the torch upper bound is bumped + # to >=2.11.0. + case "$_rocm_tag" in + rocm6.0|rocm6.0.*|rocm6.1|rocm6.1.*|rocm6.2|rocm6.2.*|rocm6.3|rocm6.3.*|rocm6.4|rocm6.4.*|rocm7.0|rocm7.0.*|rocm7.1|rocm7.1.*) + echo "$_base/$_rocm_tag" ;; + rocm6.*) + # ROCm 6.5+ (no published PyTorch wheels): clip down + # to the last supported 6.x wheel set. + echo "$_base/rocm6.4" ;; + *) + # ROCm 7.2+ (including future 10.x+): cap to rocm7.1 + echo "$_base/rocm7.1" ;; + esac + return + fi + echo "$_base/cpu"; return fi - if [ -z "$_smi" ]; then echo "$_base/cpu"; return; fi # Parse CUDA version from nvidia-smi output (POSIX-safe, no grep -P) _cuda_ver=$(LC_ALL=C $_smi 2>/dev/null \ | sed -n 's/.*CUDA Version:[[:space:]]*\([0-9][0-9]*\.[0-9][0-9]*\).*/\1/p' \ @@ -1011,20 +1150,157 @@ get_torch_index_url() { elif [ "$_major" -ge 11 ]; then echo "$_base/cu118" else echo "$_base/cpu"; fi } + +get_radeon_wheel_url() { + # Only meaningful on Linux. Picks a repo.radeon.com base URL whose listing + # contains torch wheels. Tries paths like rocm-rel-7.2.1/, rocm-rel-7.2/, + # rocm-rel-7.1.1/, rocm-rel-7.1/ (AMD publishes both M.m and M.m.p dirs). + # Accepts both X.Y and X.Y.Z host versions since /opt/rocm/.info/version + # and hipconfig --version can return either shape. + case "$(uname -s)" in Linux) ;; *) echo ""; return ;; esac + + # Detect ROCm version (X.Y or X.Y.Z) -- try amd-smi, then + # /opt/rocm/.info/version, then hipconfig. + _full_ver="" + _full_ver=$({ command -v amd-smi >/dev/null 2>&1 && \ + amd-smi version 2>/dev/null | awk -F'ROCm version: ' \ + 'NF>1{if(match($2,/[0-9]+\.[0-9]+(\.[0-9]+)?/)){print substr($2,RSTART,RLENGTH); ok=1; exit}} END{exit !ok}'; } || \ + { [ -r /opt/rocm/.info/version ] && \ + awk 'match($0,/[0-9]+\.[0-9]+(\.[0-9]+)?/){print substr($0,RSTART,RLENGTH); found=1; exit} END{exit !found}' /opt/rocm/.info/version; } || \ + { command -v hipconfig >/dev/null 2>&1 && \ + hipconfig --version 2>/dev/null | awk 'NR==1 && match($0,/[0-9]+\.[0-9]+(\.[0-9]+)?/){print substr($0,RSTART,RLENGTH); found=1} END{exit !found}'; }) 2>/dev/null + + # Validate: must be X.Y or X.Y.Z with X >= 1 + case "$_full_ver" in + [1-9]*.[0-9]*.[0-9]*) : ;; # X.Y.Z + [1-9]*.[0-9]*) : ;; # X.Y + *) echo ""; return ;; + esac + echo "https://repo.radeon.com/rocm/manylinux/rocm-rel-${_full_ver}/" +} + +# ── Radeon repo wheel selection helpers ────────────────────────────────────── +# Fetches the Radeon repo directory listing once into _RADEON_LISTING (global). +# _RADEON_PYTAG holds the CPython tag for the running interpreter (e.g. cp312). +# _RADEON_BASE_URL holds the base URL for relative-href resolution. +_RADEON_LISTING="" +_RADEON_PYTAG="" +_RADEON_BASE_URL="" + +_radeon_fetch_listing() { + # Usage: _radeon_fetch_listing BASE_URL + # Populates _RADEON_LISTING, _RADEON_PYTAG, _RADEON_BASE_URL. + _RADEON_BASE_URL="$1" + _RADEON_PYTAG=$("$_VENV_PY" -c " +import sys +print('cp{}{}'.format(sys.version_info.major, sys.version_info.minor)) +" 2>/dev/null) || return 1 + if command -v curl >/dev/null 2>&1; then + _RADEON_LISTING=$(curl -fsSL --max-time 20 "$_RADEON_BASE_URL" 2>/dev/null) + elif command -v wget >/dev/null 2>&1; then + _RADEON_LISTING=$(wget -qO- --timeout=20 "$_RADEON_BASE_URL" 2>/dev/null) + fi + [ -n "$_RADEON_LISTING" ] || return 1 +} + +_pick_radeon_wheel() { + # Usage: _pick_radeon_wheel PACKAGE_NAME + # Scans $_RADEON_LISTING for the newest wheel whose filename starts exactly + # with PACKAGE_NAME- and matches _RADEON_PYTAG + linux_x86_64. + # Prints the full URL (resolving relative hrefs against _RADEON_BASE_URL). + # + # POSIX-compliant pipeline: all href parsing, filtering, and version + # selection is done inside a single awk script rather than reaching + # for GNU extensions (grep -o, sort -V) that would break under BSD + # or BusyBox coreutils. + _pkg="$1" + [ -n "$_RADEON_LISTING" ] || return 1 + [ -n "$_RADEON_PYTAG" ] || return 1 + _tag="$_RADEON_PYTAG" + _href=$(printf '%s\n' "$_RADEON_LISTING" \ + | awk -v pkg="$_pkg" -v tag="$_tag" ' + BEGIN { max_pad = ""; max_url = "" } + { + line = $0 + while (match(line, /href="[^"]*"/)) { + # Strip the leading href=" (6 chars) and trailing " (1 char) + url = substr(line, RSTART + 6, RLENGTH - 7) + line = substr(line, RSTART + RLENGTH) + + # Extract basename, strip query / fragment + n = split(url, p, "/") + base = p[n] + sub(/[?#].*/, "", base) + + prefix = pkg "-" + # Match cpXY-cpXY or cpXY-abi3 with any linux x86_64 + # platform tag (linux_x86_64, manylinux_2_28_x86_64, + # manylinux2014_x86_64, etc.) + if (substr(base, 1, length(prefix)) == prefix && + index(base, "-" tag "-") > 0 && + match(base, /x86_64\.whl$/)) { + # Extract the version component (first + # dotted-number run) and pad each piece so a + # plain lexical comparison gives us the newest. + if (match(base, /[0-9]+\.[0-9]+(\.[0-9]+)?/)) { + ver = substr(base, RSTART, RLENGTH) + m = split(ver, v, ".") + pad = "" + for (i = 1; i <= m; i++) + pad = pad sprintf("%08d", v[i]) + if (pad > max_pad) { + max_pad = pad + max_url = url + } + } + } + } + } + END { if (max_url != "") print max_url }') + [ -z "$_href" ] && return 1 + case "$_href" in + http*) printf '%s\n' "$_href" ;; + *) printf '%s\n' "${_RADEON_BASE_URL%/}/${_href#/}" ;; + esac +} + TORCH_INDEX_URL=$(get_torch_index_url) +# Auto-detect GPU for AMD ROCm based +# get_torch_index_url must have chosen */rocm* +# (gfx in rocminfo or amd-smi list). Then require rocminfo "Marketing Name:.*Radeon". +_amd_gpu_radeon=false +case "$TORCH_INDEX_URL" in + */rocm*) + if _has_amd_rocm_gpu && command -v rocminfo >/dev/null 2>&1 && \ + rocminfo 2>/dev/null | grep -q 'Marketing Name:.*Radeon'; then + _amd_gpu_radeon=true + fi + ;; +esac + # ── Print CPU-only hint when no GPU detected ── case "$TORCH_INDEX_URL" in */cpu) if [ "$SKIP_TORCH" = false ] && [ "$OS" != "macos" ]; then echo "" - echo " NOTE: No NVIDIA GPU detected (nvidia-smi not found)." + echo " NOTE: No GPU detected (nvidia-smi and ROCm not found)." echo " Installing CPU-only PyTorch. If you only need GGUF chat/inference," echo " re-run with --no-torch for a faster, lighter install:" echo " curl -fsSL https://unsloth.ai/install.sh | sh -s -- --no-torch" + echo " AMD ROCm users: see https://docs.unsloth.ai/get-started/install-and-update/amd" echo "" fi ;; + */rocm*) + echo "" + if [ "$_amd_gpu_radeon" = true ]; then + echo " AMD Radeon + ROCm detected -- installing PyTorch wheels from repo.radeon.com" + else + echo " AMD ROCm detected -- installing ROCm-enabled PyTorch ($TORCH_INDEX_URL)" + fi + echo "" + ;; esac # ── Install unsloth directly into the venv (no activation needed) ── @@ -1040,7 +1316,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.4.2" unsloth-zoo + "unsloth>=2026.4.4" unsloth-zoo _NO_TORCH_RT="$(_find_no_torch_runtime)" if [ -n "$_NO_TORCH_RT" ]; then run_install_cmd "install no-torch runtime deps" uv pip install --python "$_VENV_PY" --no-deps -r "$_NO_TORCH_RT" @@ -1048,21 +1324,162 @@ 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.4.2" unsloth-zoo + "unsloth>=2026.4.4" unsloth-zoo fi if [ "$STUDIO_LOCAL_INSTALL" = true ]; then substep "overlaying local repo (editable)..." run_install_cmd "overlay local repo" uv pip install --python "$_VENV_PY" -e "$_REPO_ROOT" --no-deps fi + # AMD ROCm: install bitsandbytes even in migrated environments so + # existing ROCm installs gain the AMD bitsandbytes build without a + # fresh reinstall. + if [ "$SKIP_TORCH" = false ]; then + case "$TORCH_INDEX_URL" in + */rocm*) + _install_bnb_rocm "install bitsandbytes (AMD)" "$_VENV_PY" + # Repair ROCm torch if overwritten during migrated install + _has_hip=$("$_VENV_PY" -c "import torch; print(getattr(torch.version,'hip','') or '')" 2>/dev/null || true) + if [ -z "$_has_hip" ]; then + substep "repairing ROCm torch (overwritten by dependency resolution)..." + run_install_cmd "repair ROCm torch" uv pip install --python "$_VENV_PY" \ + "$TORCH_CONSTRAINT" torchvision torchaudio \ + --index-url "$TORCH_INDEX_URL" \ + --force-reinstall + fi + ;; + esac + fi elif [ -n "$TORCH_INDEX_URL" ]; then # Fresh: Step 1 - install torch from explicit index (skip when --no-torch or Intel Mac) if [ "$SKIP_TORCH" = true ]; then substep "skipping PyTorch (--no-torch or Intel Mac x86_64)." "$C_WARN" + elif [ "$_amd_gpu_radeon" = true ]; then + _radeon_url=$(get_radeon_wheel_url) + if [ -n "$_radeon_url" ]; then + _radeon_listing_ok=false + if _radeon_fetch_listing "$_radeon_url" 2>/dev/null; then + _radeon_listing_ok=true + else + # Try shorter X.Y path (AMD publishes both X.Y.Z and X.Y dirs) + _radeon_url_short=$(printf '%s\n' "$_radeon_url" \ + | sed 's|rocm-rel-\([0-9]*\)\.\([0-9]*\)\.[0-9]*/|rocm-rel-\1.\2/|') + if [ "$_radeon_url_short" != "$_radeon_url" ] && \ + _radeon_fetch_listing "$_radeon_url_short" 2>/dev/null; then + _radeon_listing_ok=true + fi + fi + + if [ "$_radeon_listing_ok" = true ]; then + # Require torch, torchvision, torchaudio wheels to all resolve + # from the Radeon listing. If any is missing for this Python + # tag, fall through to the standard ROCm index instead of + # silently mixing Radeon wheels with PyPI defaults. + _torch_whl=$(_pick_radeon_wheel "torch" 2>/dev/null) || _torch_whl="" + _tv_whl=$(_pick_radeon_wheel "torchvision" 2>/dev/null) || _tv_whl="" + _ta_whl=$(_pick_radeon_wheel "torchaudio" 2>/dev/null) || _ta_whl="" + _tri_whl=$(_pick_radeon_wheel "triton" 2>/dev/null) || _tri_whl="" + # Sanity-check torch / torchvision / torchaudio are a + # matching release. The Radeon repo publishes multiple + # generations simultaneously, so picking the highest-version + # wheel for each package independently can assemble a + # mismatched trio (e.g. torch 2.9.1 + torchvision 0.23.0 + + # torchaudio 2.9.0 from the current rocm-rel-7.2.1 index). + # Check that torch and torchaudio share the same X.Y public + # version prefix, and that torchvision's minor correctly + # pairs with torch's minor (torchvision = torch.minor - 5 + # since torch 2.4 -> torchvision 0.19 -> torch 2.9 -> + # torchvision 0.24). + # URL-decode each wheel name so %2B -> + before version + # extraction. Real Radeon wheel hrefs are percent-encoded + # (torch-2.10.0%2Brocm7.2.0...), so a plain [+-] terminator + # in the sed regex below would never match and + # _radeon_versions_match would stay false for every real + # listing, silently forcing a fallback to the generic + # ROCm index. + _torch_ver="" + _tv_ver="" + _ta_ver="" + if [ -n "$_torch_whl" ]; then + _torch_name=$(printf '%s' "${_torch_whl##*/}" | sed 's/%2[Bb]/+/g') + _torch_ver=$(printf '%s\n' "$_torch_name" | sed -n 's|^torch-\([0-9][0-9]*\.[0-9][0-9]*\)\(\.[0-9][0-9]*\)\{0,1\}[+-].*|\1|p') + fi + if [ -n "$_tv_whl" ]; then + _tv_name=$(printf '%s' "${_tv_whl##*/}" | sed 's/%2[Bb]/+/g') + _tv_ver=$(printf '%s\n' "$_tv_name" | sed -n 's|^torchvision-\([0-9][0-9]*\.[0-9][0-9]*\)\(\.[0-9][0-9]*\)\{0,1\}[+-].*|\1|p') + fi + if [ -n "$_ta_whl" ]; then + _ta_name=$(printf '%s' "${_ta_whl##*/}" | sed 's/%2[Bb]/+/g') + _ta_ver=$(printf '%s\n' "$_ta_name" | sed -n 's|^torchaudio-\([0-9][0-9]*\.[0-9][0-9]*\)\(\.[0-9][0-9]*\)\{0,1\}[+-].*|\1|p') + fi + _radeon_versions_match=false + if [ -n "$_torch_ver" ] && [ -n "$_tv_ver" ] && [ -n "$_ta_ver" ]; then + _torch_major=${_torch_ver%%.*} + _torch_minor=${_torch_ver#*.} + _ta_major=${_ta_ver%%.*} + _ta_minor=${_ta_ver#*.} + _tv_major=${_tv_ver%%.*} + _tv_minor=${_tv_ver#*.} + # torchvision expected minor (e.g. torch 2.9 -> 0.24) + _expected_tv_minor=$((_torch_minor + 15)) + if [ "$_torch_major" = "$_ta_major" ] && \ + [ "$_torch_minor" = "$_ta_minor" ] && \ + [ "$_tv_major" = "0" ] && \ + [ "$_tv_minor" = "$_expected_tv_minor" ]; then + _radeon_versions_match=true + fi + fi + if [ -z "$_torch_whl" ] || [ -z "$_tv_whl" ] || [ -z "$_ta_whl" ] || \ + [ "$_radeon_versions_match" != true ]; then + substep "[WARN] Radeon repo lacks a compatible wheel set for this Python; falling back to ROCm index ($TORCH_INDEX_URL)" "$C_WARN" + run_install_cmd "install PyTorch" uv pip install --python "$_VENV_PY" \ + "$TORCH_CONSTRAINT" torchvision torchaudio \ + --index-url "$TORCH_INDEX_URL" + else + substep "installing PyTorch from Radeon repo (${_RADEON_BASE_URL})..." + # Pass explicit wheel URLs so the matched trio is + # installed together. --find-links lets uv discover + # the Radeon listing for any local lookup, and PyPI + # (not disabled) provides transitive deps like + # filelock / sympy / networkx which are not in the + # Radeon listing. + if [ -n "$_tri_whl" ]; then + run_install_cmd "install triton + PyTorch" uv pip install --python "$_VENV_PY" \ + --find-links "$_RADEON_BASE_URL" \ + "$_tri_whl" "$_torch_whl" "$_tv_whl" "$_ta_whl" + else + run_install_cmd "install PyTorch" uv pip install --python "$_VENV_PY" \ + --find-links "$_RADEON_BASE_URL" \ + "$_torch_whl" "$_tv_whl" "$_ta_whl" + fi + fi + else + substep "[WARN] Radeon repo unavailable; falling back to ROCm index ($TORCH_INDEX_URL)" "$C_WARN" + run_install_cmd "install PyTorch" uv pip install --python "$_VENV_PY" \ + "$TORCH_CONSTRAINT" torchvision torchaudio \ + --index-url "$TORCH_INDEX_URL" + fi + else + substep "[WARN] Radeon GPU detected but could not detect full ROCm version; falling back to ROCm index" "$C_WARN" + run_install_cmd "install PyTorch" uv pip install --python "$_VENV_PY" \ + "$TORCH_CONSTRAINT" torchvision torchaudio \ + --index-url "$TORCH_INDEX_URL" + fi else substep "installing PyTorch ($TORCH_INDEX_URL)..." run_install_cmd "install PyTorch" uv pip install --python "$_VENV_PY" "$TORCH_CONSTRAINT" torchvision torchaudio \ --index-url "$TORCH_INDEX_URL" fi + # AMD ROCm: install bitsandbytes (once, after torch, for all ROCm paths). + # Gate on SKIP_TORCH=false so a user running with --no-torch on a ROCm + # host stays in GGUF-only mode rather than pulling in bitsandbytes, + # which is only useful once torch is present for training. + if [ "$SKIP_TORCH" = false ]; then + case "$TORCH_INDEX_URL" in + */rocm*) + _install_bnb_rocm "install bitsandbytes (AMD)" "$_VENV_PY" + ;; + esac + fi # Fresh: Step 2 - install unsloth, preserving pre-installed torch substep "installing unsloth (this may take a few minutes)..." if [ "$SKIP_TORCH" = true ]; then @@ -1070,7 +1487,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.4.2" unsloth-zoo + "unsloth>=2026.4.4" unsloth-zoo _NO_TORCH_RT="$(_find_no_torch_runtime)" if [ -n "$_NO_TORCH_RT" ]; then run_install_cmd "install no-torch runtime deps" uv pip install --python "$_VENV_PY" --no-deps -r "$_NO_TORCH_RT" @@ -1081,18 +1498,34 @@ 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.4.2" unsloth-zoo + --upgrade-package unsloth "unsloth>=2026.4.4" unsloth-zoo substep "overlaying local repo (editable)..." run_install_cmd "overlay local repo" uv pip install --python "$_VENV_PY" -e "$_REPO_ROOT" --no-deps else run_install_cmd "install unsloth" uv pip install --python "$_VENV_PY" \ --upgrade-package unsloth "$PACKAGE_NAME" fi + # AMD ROCm: repair torch if the unsloth/unsloth-zoo install pulled in + # CUDA torch from PyPI, overwriting the ROCm wheels installed in Step 1. + if [ "$SKIP_TORCH" = false ]; then + case "$TORCH_INDEX_URL" in + */rocm*) + _has_hip=$("$_VENV_PY" -c "import torch; print(getattr(torch.version,'hip','') or '')" 2>/dev/null || true) + if [ -z "$_has_hip" ]; then + substep "repairing ROCm torch (overwritten by dependency resolution)..." + run_install_cmd "repair ROCm torch" uv pip install --python "$_VENV_PY" \ + "$TORCH_CONSTRAINT" torchvision torchaudio \ + --index-url "$TORCH_INDEX_URL" \ + --force-reinstall + fi + ;; + esac + fi else # Fallback: GPU detection failed to produce a URL -- let uv resolve torch 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.4.2" --torch-backend=auto + run_install_cmd "install unsloth (auto torch backend)" uv pip install --python "$_VENV_PY" unsloth-zoo "unsloth>=2026.4.4" --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 else diff --git a/install_gemma4_mlx.sh b/scripts/install_gemma4_mlx.sh similarity index 66% rename from install_gemma4_mlx.sh rename to scripts/install_gemma4_mlx.sh index b653af9154..26415735b8 100755 --- a/install_gemma4_mlx.sh +++ b/scripts/install_gemma4_mlx.sh @@ -9,9 +9,7 @@ set -e # # This script: # 1. Creates a Python virtual environment -# 2. Installs uv, mlx, mlx-lm, transformers -# 3. Downloads gemma4.py and gemma4_text.py from unsloth repo -# 4. Installs them into mlx-lm's models directory +# 2. Installs uv, mlx-vlm, transformers # ============================================================ # ── Output style (inspired by unsloth/install.sh) ───────────── @@ -123,13 +121,9 @@ fi _VENV_PY="$VENV_DIR/bin/python" -# ── Repo config ────────────────────────────────────────────── -BRANCH="fix/ui-fix" -REPO_URL="https://raw.githubusercontent.com/unslothai/unsloth/refs/heads/${BRANCH}" - # ── Install dependencies ────────────────────────────────────── -step "install" "installing mlx, mlx-lm..." -uv pip install --python "$_VENV_PY" -q mlx mlx-lm 2>/dev/null +step "install" "installing mlx-vlm..." +uv pip install --python "$_VENV_PY" -q mlx-vlm substep "done" step "install" "installing transformers>=5.5.0..." @@ -145,40 +139,11 @@ else fi fi -# ── Find mlx-lm models directory ───────────────────────────── -MLX_MODELS=$("$_VENV_PY" -c "import mlx_lm; print(mlx_lm.__path__[0])")/models -step "models dir" "$MLX_MODELS" - -# ── Download and install Gemma 4 model files ────────────────── - -step "download" "installing Gemma 4 model files..." - -_install_model_file() { - _fname="$1" - if curl -fsSL "${REPO_URL}/unsloth/models/${_fname}" -o "${MLX_MODELS}/${_fname}" 2>/dev/null; then - substep "downloaded ${_fname} from branch ${BRANCH}" - elif [ -f "./${_fname}" ]; then - substep "using local ./${_fname}" - cp "./${_fname}" "${MLX_MODELS}/${_fname}" - else - fail "Could not install ${_fname}. Tried: - 1) ${REPO_URL}/unsloth/models/${_fname} - 2) Local file ./${_fname} - - To fix, download the file manually and place it in the current directory, - then re-run this script." - fi -} - -_install_model_file "gemma4.py" -_install_model_file "gemma4_text.py" - -# Verify files were installed correctly -if "$_VENV_PY" -c "from mlx_lm.models.gemma4_text import ProportionalRoPE" 2>/dev/null; then - substep "model files verified" +# ── Verify installation ────────────────────────────────────── +if "$_VENV_PY" -c "import mlx_vlm"; then + substep "mlx-vlm verified" else - fail "Model files installed but verification failed (ProportionalRoPE import error). - Try manually from: https://github.com/unslothai/unsloth/tree/feature/${BRANCH}" + fail "Installation verification failed." fi # ── Done ────────────────────────────────────────────────────── @@ -186,18 +151,19 @@ echo "" printf " ${C_TITLE}%s${C_RST}\n" "Gemma 4 MLX installed!" printf " ${C_DIM}%s${C_RST}\n" "$RULE" echo "" -step "available models" "unsloth/gemma-4-E2B-it-UD-MLX-4bit (/BF16)" -substep "unsloth/gemma-4-E4B-it-UD-MLX-4bit (/BF16)" +step "available models" "unsloth/gemma-4-E2B-it-UD-MLX-4bit" +substep "unsloth/gemma-4-E4B-it-UD-MLX-4bit" +substep "unsloth/gemma-4-26b-a4b-it-UD-MLX-4bit" +substep "unsloth/gemma-4-31b-it-UD-MLX-4bit" echo "" step "venv activate" "source ${VENV_DIR}/bin/activate" echo "" -step "quick start" "python -m mlx_lm chat --model unsloth/gemma-4-E2B-it-UD-MLX-4bit --max-tokens 200" +step "text chat" "python -m mlx_vlm.chat --model unsloth/gemma-4-E2B-it-UD-MLX-4bit" echo "" -step "python API" "from mlx_lm import load, generate" -substep "model, tokenizer = load('unsloth/gemma-4-E2B-it-UD-MLX-4bit')" -substep "messages = [{'role': 'user', 'content': 'Hello!'}]" -substep "prompt = tokenizer.apply_chat_template(messages, tokenize=False, add_generation_prompt=True)" -substep "print(generate(model, tokenizer, prompt=prompt, max_tokens=200))" +step "vision chat" "python -m mlx_vlm.chat --model unsloth/gemma-4-31b-it-UD-MLX-4bit" +substep "Use /image path/to/image.jpg to load an image" +echo "" +step "gradio UI" "python -m mlx_vlm.chat_ui --model unsloth/gemma-4-31b-it-UD-MLX-4bit" echo "" printf " ${C_DIM}%s${C_RST}\n" "$RULE" echo "" diff --git a/studio/backend/auth/authentication.py b/studio/backend/auth/authentication.py index b39f915764..da59ba9a1a 100644 --- a/studio/backend/auth/authentication.py +++ b/studio/backend/auth/authentication.py @@ -10,10 +10,12 @@ from fastapi.security import HTTPAuthorizationCredentials, HTTPBearer import jwt from .storage import ( + API_KEY_PREFIX, get_jwt_secret, get_user_and_secret, load_jwt_secret, save_refresh_token, + validate_api_key, verify_refresh_token, ) @@ -137,6 +139,18 @@ async def _get_current_subject( ... """ token = credentials.credentials + + # --- API key path (sk-unsloth-...) --- + if token.startswith(API_KEY_PREFIX): + username = validate_api_key(token) + if username is None: + raise HTTPException( + status_code = status.HTTP_401_UNAUTHORIZED, + detail = "Invalid or expired API key", + ) + return username + + # --- JWT path --- subject = _decode_subject_without_verification(token) if subject is None: raise HTTPException( diff --git a/studio/backend/auth/storage.py b/studio/backend/auth/storage.py index 1395574cce..7d55a2dc59 100644 --- a/studio/backend/auth/storage.py +++ b/studio/backend/auth/storage.py @@ -72,7 +72,22 @@ def clear_bootstrap_password() -> None: def _hash_token(token: str) -> str: - """SHA-256 hash helper used for refresh token storage.""" + """SHA-256 hash helper used for refresh token storage. + + Plain SHA-256 is intentional here: refresh tokens are high-entropy + random strings from ``secrets.token_urlsafe(48)`` (384 bits of + entropy), so a slow KDF (Argon2 / bcrypt / PBKDF2) provides zero + additional security — no attacker can brute-force 2^384 regardless + of hash speed — while adding tens of ms of CPU to every refresh. + See the OWASP Password Storage Cheat Sheet on fast-vs-slow hashing + of high-entropy inputs. + + API keys use the separate ``_pbkdf2_api_key`` helper below, which + runs PBKDF2-HMAC-SHA256 with a persistent server-side salt — not + for cryptographic reasons (128-bit random tokens don't need slow + hashing), but because CodeQL's ``py/weak-sensitive-data-hashing`` + query mislabels API keys as passwords and demands a KDF. + """ return hashlib.sha256(token.encode("utf-8")).hexdigest() @@ -103,6 +118,29 @@ def get_connection() -> sqlite3.Connection: ); """ ) + conn.execute( + """ + CREATE TABLE IF NOT EXISTS api_keys ( + id INTEGER PRIMARY KEY AUTOINCREMENT, + username TEXT NOT NULL, + key_prefix TEXT NOT NULL, + key_hash TEXT NOT NULL UNIQUE, + name TEXT NOT NULL DEFAULT '', + created_at TEXT NOT NULL, + last_used_at TEXT, + expires_at TEXT, + is_active INTEGER NOT NULL DEFAULT 1 + ); + """ + ) + conn.execute( + """ + CREATE TABLE IF NOT EXISTS app_secrets ( + key TEXT PRIMARY KEY, + value TEXT NOT NULL + ); + """ + ) columns = {row["name"] for row in conn.execute("PRAGMA table_info(auth_user)")} if "must_change_password" not in columns: conn.execute( @@ -112,6 +150,89 @@ def get_connection() -> sqlite3.Connection: return conn +# ── API-key PBKDF2 salt ──────────────────────────────────────────────── +# +# Module-level cache for the persistent API-key PBKDF2 salt. Populated +# lazily on first use via ``_get_or_create_api_key_pbkdf2_salt``. Not +# protected by a lock because (a) the ``INSERT OR IGNORE`` provides +# atomicity at the SQLite layer and (b) concurrent populations converge +# on the same value, so the worst case is a harmless duplicate read on +# startup. +_api_key_pbkdf2_salt_cache: Optional[bytes] = None + + +def _get_or_create_api_key_pbkdf2_salt() -> bytes: + """Return the persistent API-key PBKDF2 salt, generating it once if missing. + + Stored as a hex-encoded 32-byte random value in the ``app_secrets`` + table under key ``"api_key_pbkdf2_salt"``. Regenerated only if the row + is missing (i.e. fresh install, or operator manually deleted the row + and accepts invalidating existing API keys). + """ + global _api_key_pbkdf2_salt_cache + if _api_key_pbkdf2_salt_cache is not None: + return _api_key_pbkdf2_salt_cache + + conn = get_connection() + try: + cur = conn.execute( + "SELECT value FROM app_secrets WHERE key = ?", + ("api_key_pbkdf2_salt",), + ) + row = cur.fetchone() + if row is None: + new_value = secrets.token_hex(32) # 32 bytes -> 64 hex chars + conn.execute( + "INSERT OR IGNORE INTO app_secrets (key, value) VALUES (?, ?)", + ("api_key_pbkdf2_salt", new_value), + ) + conn.commit() + cur = conn.execute( + "SELECT value FROM app_secrets WHERE key = ?", + ("api_key_pbkdf2_salt",), + ) + row = cur.fetchone() + salt = bytes.fromhex(row["value"]) + finally: + conn.close() + + _api_key_pbkdf2_salt_cache = salt + return salt + + +_API_KEY_PBKDF2_ITERATIONS = 100_000 + + +def _pbkdf2_api_key(raw_key: str) -> str: + """PBKDF2-HMAC-SHA256 an API key with a persistent server-side salt. + + Used for API-key storage ONLY, not refresh tokens. Matches the + PBKDF2 algorithm + iteration count used by the password hasher in + ``auth/hashing.py`` so the codebase is consistent on which KDF it + uses for credential storage. + + Notes on why a slow KDF here is *only* a CodeQL appeasement and + *not* a cryptographic requirement: API keys are cryptographically + random 128-bit tokens (via ``secrets.token_hex``), so brute force + against 2^128 is infeasible regardless of hash speed. CodeQL's + ``py/weak-sensitive-data-hashing`` query mislabels these tokens as + "password" sensitive data and then demands a KDF from its + allowlist (Argon2 / scrypt / bcrypt / PBKDF2). Per the query's + own recommendation page we use PBKDF2. The persistent salt is + still loaded from ``app_secrets`` so an attacker dumping the + ``api_keys`` table alone cannot derive hashes for candidate + tokens without also obtaining the salt row. + """ + salt = _get_or_create_api_key_pbkdf2_salt() + dk = hashlib.pbkdf2_hmac( + "sha256", + raw_key.encode("utf-8"), + salt, + _API_KEY_PBKDF2_ITERATIONS, + ) + return dk.hex() + + def is_initialized() -> bool: """Check if auth is ready for login (at least one user exists in DB).""" conn = get_connection() @@ -357,3 +478,105 @@ def revoke_user_refresh_tokens(username: str) -> None: conn.commit() finally: conn.close() + + +# --------------------------------------------------------------------------- +# API key management +# --------------------------------------------------------------------------- + +API_KEY_PREFIX = "sk-unsloth-" + + +def create_api_key( + username: str, + name: str, + expires_at: Optional[str] = None, +) -> Tuple[str, dict]: + """Create a new API key for *username*. + + Returns ``(raw_key, row_dict)`` where *raw_key* is shown to the user + exactly once. The database only stores the SHA-256 hash. + """ + raw_key = API_KEY_PREFIX + secrets.token_hex(16) + key_hash = _pbkdf2_api_key(raw_key) + key_prefix = raw_key[len(API_KEY_PREFIX) : len(API_KEY_PREFIX) + 8] + now = datetime.now(timezone.utc).isoformat() + + conn = get_connection() + try: + conn.execute( + """ + INSERT INTO api_keys (username, key_prefix, key_hash, name, created_at, expires_at) + VALUES (?, ?, ?, ?, ?, ?) + """, + (username, key_prefix, key_hash, name, now, expires_at), + ) + conn.commit() + cur = conn.execute("SELECT * FROM api_keys WHERE key_hash = ?", (key_hash,)) + row = cur.fetchone() + return raw_key, dict(row) + finally: + conn.close() + + +def list_api_keys(username: str) -> list: + """Return all API keys for *username* (never exposes ``key_hash``).""" + conn = get_connection() + try: + cur = conn.execute( + """ + SELECT id, username, key_prefix, name, created_at, last_used_at, expires_at, is_active + FROM api_keys + WHERE username = ? + ORDER BY created_at DESC + """, + (username,), + ) + return [dict(row) for row in cur.fetchall()] + finally: + conn.close() + + +def revoke_api_key(username: str, key_id: int) -> bool: + """Soft-delete an API key. Returns True if a matching row was found.""" + conn = get_connection() + try: + cursor = conn.execute( + "UPDATE api_keys SET is_active = 0 WHERE id = ? AND username = ?", + (key_id, username), + ) + conn.commit() + return cursor.rowcount > 0 + finally: + conn.close() + + +def validate_api_key(raw_key: str) -> Optional[str]: + """Validate *raw_key* and return the owning username, or ``None``. + + Also updates ``last_used_at`` on success. + """ + key_hash = _pbkdf2_api_key(raw_key) + conn = get_connection() + try: + cur = conn.execute( + "SELECT id, username, is_active, expires_at FROM api_keys WHERE key_hash = ?", + (key_hash,), + ) + row = cur.fetchone() + if row is None: + return None + if not row["is_active"]: + return None + if row["expires_at"] is not None: + expires = datetime.fromisoformat(row["expires_at"]) + if datetime.now(timezone.utc) > expires: + return None + conn.execute( + "UPDATE api_keys SET last_used_at = ? WHERE id = ?", + (datetime.now(timezone.utc).isoformat(), row["id"]), + ) + conn.commit() + return row["username"] + finally: + conn.close() diff --git a/studio/backend/core/__init__.py b/studio/backend/core/__init__.py index d8d95e2f1a..d39815c437 100644 --- a/studio/backend/core/__init__.py +++ b/studio/backend/core/__init__.py @@ -31,6 +31,7 @@ __all__ = [ # Config "ModelConfig", "is_vision_model", + "scan_trained_models", "scan_trained_loras", "load_model_defaults", "get_base_model_from_lora", @@ -72,6 +73,7 @@ def __getattr__(name): if name in ( "is_vision_model", "ModelConfig", + "scan_trained_models", "scan_trained_loras", "load_model_defaults", "get_base_model_from_lora", @@ -79,14 +81,15 @@ def __getattr__(name): from utils.models import ( is_vision_model, ModelConfig, - scan_trained_loras, + scan_trained_models, load_model_defaults, get_base_model_from_lora, ) globals()["is_vision_model"] = is_vision_model globals()["ModelConfig"] = ModelConfig - globals()["scan_trained_loras"] = scan_trained_loras + globals()["scan_trained_models"] = scan_trained_models + globals()["scan_trained_loras"] = scan_trained_models globals()["load_model_defaults"] = load_model_defaults globals()["get_base_model_from_lora"] = get_base_model_from_lora return globals()[name] diff --git a/studio/backend/core/data_recipe/oxc-validator/package.json b/studio/backend/core/data_recipe/oxc-validator/package.json index d1c765a2e1..111ae2b257 100644 --- a/studio/backend/core/data_recipe/oxc-validator/package.json +++ b/studio/backend/core/data_recipe/oxc-validator/package.json @@ -4,7 +4,7 @@ "version": "0.0.1", "type": "module", "dependencies": { - "oxc-parser": "^0.121.0", + "oxc-parser": "^0.123.0", "oxlint": "^1.51.0" } } diff --git a/studio/backend/core/export/export.py b/studio/backend/core/export/export.py index 966e045b13..d8f2e8fa37 100644 --- a/studio/backend/core/export/export.py +++ b/studio/backend/core/export/export.py @@ -310,7 +310,7 @@ class ExportBackend: repo_id: Optional[str] = None, hf_token: Optional[str] = None, private: bool = False, - ) -> Tuple[bool, str]: + ) -> Tuple[bool, str, Optional[str]]: """ Export merged model (for PEFT models). @@ -323,14 +323,21 @@ class ExportBackend: private: Whether to make the repo private Returns: - Tuple of (success: bool, message: str) + Tuple of (success, message, output_path). output_path is the + resolved absolute on-disk directory of the saved model when + ``save_directory`` was set, else None. """ if not self.current_model or not self.current_tokenizer: - return False, "No model loaded. Please select a checkpoint first." + return False, "No model loaded. Please select a checkpoint first.", None if not self.is_peft: - return False, "This is not a PEFT model. Use 'Export Base Model' instead." + return ( + False, + "This is not a PEFT model. Use 'Export Base Model' instead.", + None, + ) + output_path: Optional[str] = None try: # Determine save method if format_type == "4-bit (FP4)": @@ -354,6 +361,7 @@ class ExportBackend: # Write export metadata so the Chat page can identify the base model self._write_export_metadata(save_directory) logger.info(f"Model saved successfully to {save_directory}") + output_path = str(Path(save_directory).resolve()) # Push to hub if requested if push_to_hub: @@ -361,6 +369,7 @@ class ExportBackend: return ( False, "Repository ID and Hugging Face token required for Hub upload", + None, ) logger.info(f"Pushing merged model to Hub: {repo_id}") @@ -378,14 +387,14 @@ class ExportBackend: ) logger.info(f"Model pushed successfully to {repo_id}") - return True, "Model exported successfully" + return True, "Model exported successfully", output_path except Exception as e: logger.error(f"Error exporting merged model: {e}") import traceback logger.error(traceback.format_exc()) - return False, f"Export failed: {str(e)}" + return False, f"Export failed: {str(e)}", None def export_base_model( self, @@ -395,22 +404,26 @@ class ExportBackend: hf_token: Optional[str] = None, private: bool = False, base_model_id: Optional[str] = None, - ) -> Tuple[bool, str]: + ) -> Tuple[bool, str, Optional[str]]: """ Export base model (for non-PEFT models). Returns: - Tuple of (success: bool, message: str) + Tuple of (success, message, output_path). output_path is the + resolved absolute on-disk directory of the saved model when + ``save_directory`` was set, else None. """ if not self.current_model or not self.current_tokenizer: - return False, "No model loaded. Please select a checkpoint first." + return False, "No model loaded. Please select a checkpoint first.", None if self.is_peft: return ( False, "This is a PEFT model. Use 'Merged Model' export type instead.", + None, ) + output_path: Optional[str] = None try: # Save locally if requested if save_directory: @@ -424,6 +437,7 @@ class ExportBackend: # Write export metadata so the Chat page can identify the base model self._write_export_metadata(save_directory) logger.info(f"Model saved successfully to {save_directory}") + output_path = str(Path(save_directory).resolve()) # Push to hub if requested if push_to_hub: @@ -431,6 +445,7 @@ class ExportBackend: return ( False, "Repository ID and Hugging Face token required for Hub upload", + None, ) logger.info(f"Pushing base model to Hub: {repo_id}") @@ -472,16 +487,16 @@ class ExportBackend: ) logger.info(f"Model pushed successfully to {repo_id}") else: - return False, "Local save directory required for Hub upload" + return False, "Local save directory required for Hub upload", None - return True, "Model exported successfully" + return True, "Model exported successfully", output_path except Exception as e: logger.error(f"Error exporting base model: {e}") import traceback logger.error(traceback.format_exc()) - return False, f"Export failed: {str(e)}" + return False, f"Export failed: {str(e)}", None def export_gguf( self, @@ -490,7 +505,7 @@ class ExportBackend: push_to_hub: bool = False, repo_id: Optional[str] = None, hf_token: Optional[str] = None, - ) -> Tuple[bool, str]: + ) -> Tuple[bool, str, Optional[str]]: """ Export model in GGUF format. @@ -502,11 +517,14 @@ class ExportBackend: hf_token: Hugging Face token Returns: - Tuple of (success: bool, message: str) + Tuple of (success, message, output_path). output_path is the + resolved absolute on-disk directory containing the .gguf + files when ``save_directory`` was set, else None. """ if not self.current_model or not self.current_tokenizer: - return False, "No model loaded. Please select a checkpoint first." + return False, "No model loaded. Please select a checkpoint first.", None + output_path: Optional[str] = None try: # Convert quantization method to lowercase for unsloth quant_method = quantization_method.lower() @@ -601,6 +619,7 @@ class ExportBackend: abs_save_dir, "\n ".join(os.path.basename(f) for f in final_ggufs) or "(none)", ) + output_path = str(Path(abs_save_dir).resolve()) # Push to hub if requested if push_to_hub: @@ -608,6 +627,7 @@ class ExportBackend: return ( False, "Repository ID and Hugging Face token required for Hub upload", + None, ) logger.info(f"Pushing GGUF model to Hub: {repo_id}") @@ -620,14 +640,18 @@ class ExportBackend: ) logger.info(f"GGUF model pushed successfully to {repo_id}") - return True, f"GGUF model exported successfully ({quantization_method})" + return ( + True, + f"GGUF model exported successfully ({quantization_method})", + output_path, + ) except Exception as e: logger.error(f"Error exporting GGUF model: {e}") import traceback logger.error(traceback.format_exc()) - return False, f"GGUF export failed: {str(e)}" + return False, f"GGUF export failed: {str(e)}", None def export_lora_adapter( self, @@ -636,19 +660,22 @@ class ExportBackend: repo_id: Optional[str] = None, hf_token: Optional[str] = None, private: bool = False, - ) -> Tuple[bool, str]: + ) -> Tuple[bool, str, Optional[str]]: """ Export LoRA adapter only (not merged). Returns: - Tuple of (success: bool, message: str) + Tuple of (success, message, output_path). output_path is the + resolved absolute on-disk directory of the saved adapter + when ``save_directory`` was set, else None. """ if not self.current_model or not self.current_tokenizer: - return False, "No model loaded. Please select a checkpoint first." + return False, "No model loaded. Please select a checkpoint first.", None if not self.is_peft: - return False, "This is not a PEFT model. No adapter to export." + return False, "This is not a PEFT model. No adapter to export.", None + output_path: Optional[str] = None try: # Save locally if requested if save_directory: @@ -659,6 +686,7 @@ class ExportBackend: self.current_model.save_pretrained(save_directory) self.current_tokenizer.save_pretrained(save_directory) logger.info(f"Adapter saved successfully to {save_directory}") + output_path = str(Path(save_directory).resolve()) # Push to hub if requested if push_to_hub: @@ -666,6 +694,7 @@ class ExportBackend: return ( False, "Repository ID and Hugging Face token required for Hub upload", + None, ) logger.info(f"Pushing LoRA adapter to Hub: {repo_id}") @@ -676,14 +705,14 @@ class ExportBackend: ) logger.info(f"Adapter pushed successfully to {repo_id}") - return True, "LoRA adapter exported successfully" + return True, "LoRA adapter exported successfully", output_path except Exception as e: logger.error(f"Error exporting LoRA adapter: {e}") import traceback logger.error(traceback.format_exc()) - return False, f"Adapter export failed: {str(e)}" + return False, f"Adapter export failed: {str(e)}", None # Global export backend instance diff --git a/studio/backend/core/export/orchestrator.py b/studio/backend/core/export/orchestrator.py index 500bc9e706..206dbd6dbb 100644 --- a/studio/backend/core/export/orchestrator.py +++ b/studio/backend/core/export/orchestrator.py @@ -16,19 +16,25 @@ Pattern follows core/inference/orchestrator.py. import atexit import structlog +from collections import deque from loggers import get_logger import multiprocessing as mp import queue import threading import time from pathlib import Path -from typing import Any, List, Optional, Tuple +from typing import Any, Deque, Dict, List, Optional, Tuple from utils.paths import outputs_root logger = get_logger(__name__) _CTX = mp.get_context("spawn") +# Maximum number of captured log lines kept in memory per export +# orchestrator. Acts as scrollback for the live export log panel in the +# UI. 4000 lines is ~1 MB worst-case at 256 chars/line. +_LOG_BUFFER_MAXLEN = 4000 + class ExportOrchestrator: """ @@ -44,6 +50,9 @@ class ExportOrchestrator: self._proc: Optional[mp.Process] = None self._cmd_queue: Any = None self._resp_queue: Any = None + # Serializes export operations (load_checkpoint, export_*, + # cleanup) so concurrent HTTP requests can never interleave + # commands on the subprocess queue. Previously unused. self._lock = threading.Lock() # Local state mirrors (updated from subprocess responses) @@ -51,9 +60,103 @@ class ExportOrchestrator: self.is_vision: bool = False self.is_peft: bool = False + # ── Live log capture ───────────────────────────────────── + # Thread-safe ring buffer of log lines forwarded from the + # worker subprocess. Powers the GET /api/export/logs/stream + # SSE endpoint that the export dialog consumes. + self._log_buffer: Deque[Dict[str, Any]] = deque(maxlen = _LOG_BUFFER_MAXLEN) + self._log_lock = threading.Lock() + # Monotonically increasing sequence number. Never reset across + # operations, so SSE clients can use it as a stable cursor even + # if clear_logs() is called mid-session. + self._log_seq: int = 0 + # Snapshot of _log_seq captured at the start of the current run + # (updated by clear_logs()). The SSE endpoint defaults its + # cursor to this value so a client that connects AFTER the + # worker has already emitted its first lines still sees the + # full run. Every line appended during the current run has seq + # strictly greater than _run_start_seq, and every line from + # prior runs has seq less than or equal to it. + self._run_start_seq: int = 0 + # True while an export operation (load/export/cleanup) is + # running. The SSE endpoint ends the stream 1 second after + # this flips back to False to drain any trailing log lines. + self._export_active: bool = False + atexit.register(self._cleanup) logger.info("ExportOrchestrator initialized (subprocess mode)") + # ------------------------------------------------------------------ + # Live log capture helpers + # ------------------------------------------------------------------ + + def _append_log(self, entry: Dict[str, Any]) -> None: + """Append a log line from the worker subprocess to the buffer. + + Entries look like {"type": "log", "stream": "stdout"|"stderr", + "line": "...", "ts": ...}. Each is stamped with a monotonic + seq number before it lands in the buffer so SSE clients can + cursor through new lines. + """ + line = entry.get("line") + if not line: + return + with self._log_lock: + self._log_seq += 1 + self._log_buffer.append( + { + "seq": self._log_seq, + "stream": entry.get("stream", "stdout"), + "line": line, + "ts": entry.get("ts", time.time()), + } + ) + + def clear_logs(self) -> None: + """Drop any buffered log lines from a previous operation. + + Called at the start of each export op so the UI shows only the + output of the current run. The seq counter is NOT reset, so an + SSE client that captured the cursor before clear_logs() will + still see new lines (with strictly greater seq numbers). + + Also snapshots the current seq into ``_run_start_seq`` so the + SSE endpoint can anchor its default cursor at the start of + this run. Anything appended after this call has seq strictly + greater than the snapshot and is reachable via + ``get_logs_since(get_run_start_seq())``. + """ + with self._log_lock: + self._log_buffer.clear() + self._run_start_seq = self._log_seq + + def get_logs_since(self, cursor: int) -> Tuple[List[Dict[str, Any]], int]: + """Return log entries with seq > cursor, plus the new cursor.""" + with self._log_lock: + new_entries = [entry for entry in self._log_buffer if entry["seq"] > cursor] + if new_entries: + return new_entries, new_entries[-1]["seq"] + return [], cursor + + def get_current_log_seq(self) -> int: + """Return the current seq counter without reading any entries.""" + with self._log_lock: + return self._log_seq + + def get_run_start_seq(self) -> int: + """Return the seq value captured at the start of the current run. + + The SSE endpoint uses this as the default cursor so a client + that connects AFTER the worker has already started emitting + output still sees every line from the current run. + """ + with self._log_lock: + return self._run_start_seq + + def is_export_active(self) -> bool: + """True while an export / load / cleanup command is running.""" + return self._export_active + # ------------------------------------------------------------------ # Subprocess lifecycle # ------------------------------------------------------------------ @@ -179,8 +282,26 @@ class ExportOrchestrator: error_msg = resp.get("error", "Unknown error") raise RuntimeError(f"Subprocess error: {error_msg}") + if rtype == "log": + # Forwarded stdout/stderr line from the worker process. + self._append_log(resp) + continue + if rtype == "status": - logger.info("Export subprocess status: %s", resp.get("message", "")) + message = resp.get("message", "") + logger.info("Export subprocess status: %s", message) + # Surface status messages in the live log panel too so + # users see high level progress (e.g. "Importing + # Unsloth...", "Loading checkpoint: ...") alongside + # subprocess output. + if message: + self._append_log( + { + "stream": "status", + "line": message, + "ts": resp.get("ts", time.time()), + } + ) continue # Other response types during wait — skip @@ -231,37 +352,47 @@ class ExportOrchestrator: "hf_token": hf_token, } - # Always kill existing subprocess and spawn fresh. - if self._ensure_subprocess_alive(): - self._shutdown_subprocess() - elif self._proc is not None: - self._shutdown_subprocess(timeout = 2) + with self._lock: + # Start a fresh log buffer for this operation so the UI + # sees only the current run's output. + self.clear_logs() + self._export_active = True + try: + # Always kill existing subprocess and spawn fresh. + if self._ensure_subprocess_alive(): + self._shutdown_subprocess() + elif self._proc is not None: + self._shutdown_subprocess(timeout = 2) - logger.info("Spawning fresh export subprocess for '%s'", checkpoint_path) - self._spawn_subprocess(sub_config) + logger.info( + "Spawning fresh export subprocess for '%s'", checkpoint_path + ) + self._spawn_subprocess(sub_config) - try: - resp = self._wait_response("loaded", timeout = 300) - except RuntimeError as exc: - self._shutdown_subprocess(timeout = 5) - self.current_checkpoint = None - self.is_vision = False - self.is_peft = False - return False, str(exc) + try: + resp = self._wait_response("loaded") + except RuntimeError as exc: + self._shutdown_subprocess(timeout = 5) + self.current_checkpoint = None + self.is_vision = False + self.is_peft = False + return False, str(exc) - if resp.get("success"): - self.current_checkpoint = resp.get("checkpoint") - self.is_vision = resp.get("is_vision", False) - self.is_peft = resp.get("is_peft", False) - logger.info("Checkpoint '%s' loaded in subprocess", checkpoint_path) - return True, resp.get("message", "Loaded successfully") - else: - error = resp.get("message", "Failed to load checkpoint") - logger.error("Failed to load checkpoint: %s", error) - self.current_checkpoint = None - self.is_vision = False - self.is_peft = False - return False, error + if resp.get("success"): + self.current_checkpoint = resp.get("checkpoint") + self.is_vision = resp.get("is_vision", False) + self.is_peft = resp.get("is_peft", False) + logger.info("Checkpoint '%s' loaded in subprocess", checkpoint_path) + return True, resp.get("message", "Loaded successfully") + else: + error = resp.get("message", "Failed to load checkpoint") + logger.error("Failed to load checkpoint: %s", error) + self.current_checkpoint = None + self.is_vision = False + self.is_peft = False + return False, error + finally: + self._export_active = False def export_merged_model( self, @@ -271,7 +402,7 @@ class ExportOrchestrator: repo_id: Optional[str] = None, hf_token: Optional[str] = None, private: bool = False, - ) -> Tuple[bool, str]: + ) -> Tuple[bool, str, Optional[str]]: """Export merged PEFT model.""" return self._run_export( "merged", @@ -293,7 +424,7 @@ class ExportOrchestrator: hf_token: Optional[str] = None, private: bool = False, base_model_id: Optional[str] = None, - ) -> Tuple[bool, str]: + ) -> Tuple[bool, str, Optional[str]]: """Export base model (non-PEFT).""" return self._run_export( "base", @@ -314,7 +445,7 @@ class ExportOrchestrator: push_to_hub: bool = False, repo_id: Optional[str] = None, hf_token: Optional[str] = None, - ) -> Tuple[bool, str]: + ) -> Tuple[bool, str, Optional[str]]: """Export model in GGUF format.""" return self._run_export( "gguf", @@ -334,7 +465,7 @@ class ExportOrchestrator: repo_id: Optional[str] = None, hf_token: Optional[str] = None, private: bool = False, - ) -> Tuple[bool, str]: + ) -> Tuple[bool, str, Optional[str]]: """Export LoRA adapter only.""" return self._run_export( "lora", @@ -347,46 +478,74 @@ class ExportOrchestrator: }, ) - def _run_export(self, export_type: str, params: dict) -> Tuple[bool, str]: - """Send an export command to the subprocess and wait for result.""" - if not self._ensure_subprocess_alive(): - return False, "No export subprocess running. Load a checkpoint first." + def _run_export( + self, export_type: str, params: dict + ) -> Tuple[bool, str, Optional[str]]: + """Send an export command to the subprocess and wait for result. - cmd = {"type": "export", "export_type": export_type, **params} + Returns ``(success, message, output_path)``. ``output_path`` is the + resolved on-disk directory the worker actually wrote to (None when + the export only pushed to Hub or failed before any file was + written). Surfaced via the export route's ``details.output_path`` + so the dialog's success screen can show the user where the model + landed. + """ + with self._lock: + if not self._ensure_subprocess_alive(): + return ( + False, + "No export subprocess running. Load a checkpoint first.", + None, + ) - try: - self._send_cmd(cmd) - resp = self._wait_response( - f"export_{export_type}_done", - timeout = 3600, # GGUF for 30B+ models can take 30+ min - ) - return resp.get("success", False), resp.get("message", "") - except RuntimeError as exc: - return False, str(exc) + self.clear_logs() + self._export_active = True + try: + cmd = {"type": "export", "export_type": export_type, **params} + try: + self._send_cmd(cmd) + resp = self._wait_response( + f"export_{export_type}_done", + timeout = 3600, # GGUF for 30B+ models can take 30+ min + ) + return ( + resp.get("success", False), + resp.get("message", ""), + resp.get("output_path"), + ) + except RuntimeError as exc: + return False, str(exc), None + finally: + self._export_active = False def cleanup_memory(self) -> bool: """Cleanup export-related models from memory.""" - if not self._ensure_subprocess_alive(): - # No subprocess — just clear local state - self.current_checkpoint = None - self.is_vision = False - self.is_peft = False - return True + with self._lock: + if not self._ensure_subprocess_alive(): + # No subprocess — just clear local state + self.current_checkpoint = None + self.is_vision = False + self.is_peft = False + return True - try: - self._send_cmd({"type": "cleanup"}) - resp = self._wait_response("cleanup_done", timeout = 30) - success = resp.get("success", False) - except RuntimeError: - success = False + self._export_active = True + try: + try: + self._send_cmd({"type": "cleanup"}) + resp = self._wait_response("cleanup_done", timeout = 30) + success = resp.get("success", False) + except RuntimeError: + success = False - # Shut down subprocess after cleanup — no model loaded - self._shutdown_subprocess() + # Shut down subprocess after cleanup — no model loaded + self._shutdown_subprocess() - self.current_checkpoint = None - self.is_vision = False - self.is_peft = False - return success + self.current_checkpoint = None + self.is_vision = False + self.is_peft = False + return success + finally: + self._export_active = False def scan_checkpoints( self, outputs_dir: str = str(outputs_root()) diff --git a/studio/backend/core/export/worker.py b/studio/backend/core/export/worker.py index 6af6ff1193..f77b1966c4 100644 --- a/studio/backend/core/export/worker.py +++ b/studio/backend/core/export/worker.py @@ -17,10 +17,12 @@ Pattern follows core/inference/worker.py and core/training/worker.py. from __future__ import annotations +import errno import structlog from loggers import get_logger import os import sys +import threading import time import traceback from pathlib import Path @@ -29,38 +31,164 @@ from typing import Any logger = get_logger(__name__) -def _activate_transformers_version(model_name: str) -> None: - """Activate the correct transformers version BEFORE any ML imports. +# Gate that controls whether captured stdout/stderr lines are forwarded +# to the parent's resp_queue (and from there to the export-dialog SSE +# stream). Closed by default so the noisy bootstrap phase -- transformers +# venv activation, Unsloth/torch imports, base-model resolution, "Top +# GGUF/hub models" lists, vision detection, weight loading bars -- is +# suppressed in the UI. _handle_export() opens the gate at the start of +# the actual export work and leaves it open; the orchestrator always +# spawns a fresh subprocess for the next checkpoint load (see +# orchestrator._spawn_subprocess) which resets this state. +# +# Lines dropped while the gate is closed are still echoed to the saved +# original stdout/stderr fds so the server console / log file keeps the +# full output for debugging. +_log_forward_gate = threading.Event() - If the model needs transformers 5.x, prepend the pre-installed .venv_t5/ - directory to sys.path. Otherwise do nothing (default 4.57.x in .venv/). + +def _setup_log_capture(resp_queue: Any) -> None: + """Redirect fds 1 and 2 through pipes so every line printed by this + worker process and any child process it spawns is forwarded to the + parent process via resp_queue as {"type": "log", ...} messages. + + Must be called BEFORE LogConfig.setup_logging and BEFORE any ML + imports, otherwise library handlers may capture the original stderr + reference and bypass the pipe. + + Lines are also echoed back to the original stdout/stderr so the + server console keeps receiving the full subprocess output, even + while ``_log_forward_gate`` is closed. """ + + try: + saved_out_fd = os.dup(1) + saved_err_fd = os.dup(2) + except OSError: + # dup failed (exotic platforms) - give up quietly, export still + # works, just no live log streaming. + return + + try: + r_out, w_out = os.pipe() + r_err, w_err = os.pipe() + except OSError: + os.close(saved_out_fd) + os.close(saved_err_fd) + return + + try: + os.dup2(w_out, 1) + os.dup2(w_err, 2) + except OSError: + for fd in (saved_out_fd, saved_err_fd, r_out, w_out, r_err, w_err): + try: + os.close(fd) + except OSError: + pass + return + + # Close the write ends we just dup2'd (fds 1 and 2 are the real + # write ends now). + os.close(w_out) + os.close(w_err) + + # Replace Python's sys.stdout/sys.stderr with line-buffered writers + # bound to the (now-redirected) fds 1 and 2. + try: + sys.stdout = os.fdopen(1, "w", buffering = 1, encoding = "utf-8", errors = "replace") + sys.stderr = os.fdopen(2, "w", buffering = 1, encoding = "utf-8", errors = "replace") + except Exception: + pass + + def _reader(read_fd: int, stream_name: str, echo_fd: int) -> None: + buf = bytearray() + while True: + try: + chunk = os.read(read_fd, 4096) + except OSError as exc: + if exc.errno == errno.EBADF: + break + continue + if not chunk: + break + # Echo to the original fd so the server console still sees + # the full output. + try: + os.write(echo_fd, chunk) + except OSError: + pass + buf.extend(chunk) + # Split on \n OR \r so tqdm-style progress bars update. + while True: + nl = -1 + for i, b in enumerate(buf): + if b == 0x0A or b == 0x0D: + nl = i + break + if nl < 0: + break + line = bytes(buf[:nl]).decode("utf-8", errors = "replace") + del buf[: nl + 1] + if not line: + continue + if not _log_forward_gate.is_set(): + # Gate closed (bootstrap phase) -- already echoed to + # the saved console fd above; drop the line so the + # export dialog doesn't see import / vendoring noise. + continue + try: + resp_queue.put_nowait( + { + "type": "log", + "stream": stream_name, + "line": line, + "ts": time.time(), + } + ) + except Exception: + # Queue put failed (full, closed, etc.) - drop the + # line rather than crash the reader thread. + pass + if buf and _log_forward_gate.is_set(): + try: + resp_queue.put_nowait( + { + "type": "log", + "stream": stream_name, + "line": bytes(buf).decode("utf-8", errors = "replace"), + "ts": time.time(), + } + ) + except Exception: + pass + + t_out = threading.Thread( + target = _reader, + args = (r_out, "stdout", saved_out_fd), + daemon = True, + name = "export-log-stdout", + ) + t_err = threading.Thread( + target = _reader, + args = (r_err, "stderr", saved_err_fd), + daemon = True, + name = "export-log-stderr", + ) + t_out.start() + t_err.start() + + +def _activate_transformers_version(model_name: str) -> None: + """Activate the correct transformers version BEFORE any ML imports.""" # Ensure backend is on path for utils imports backend_path = str(Path(__file__).resolve().parent.parent.parent) if backend_path not in sys.path: sys.path.insert(0, backend_path) - from utils.transformers_version import ( - needs_transformers_5, - _resolve_base_model, - _ensure_venv_t5_exists, - _VENV_T5_DIR, - ) + from utils.transformers_version import activate_transformers_for_subprocess - resolved = _resolve_base_model(model_name) - if needs_transformers_5(resolved): - if not _ensure_venv_t5_exists(): - raise RuntimeError( - f"Cannot activate transformers 5.x: .venv_t5 missing at {_VENV_T5_DIR}" - ) - if _VENV_T5_DIR not in sys.path: - sys.path.insert(0, _VENV_T5_DIR) - logger.info("Activated transformers 5.x from %s", _VENV_T5_DIR) - # Propagate to child subprocesses (e.g. GGUF converter) - _pp = os.environ.get("PYTHONPATH", "") - os.environ["PYTHONPATH"] = _VENV_T5_DIR + (os.pathsep + _pp if _pp else "") - else: - logger.info("Using default transformers (4.57.x) for %s", model_name) + activate_transformers_for_subprocess(model_name) def _send_response(resp_queue: Any, response: dict) -> None: @@ -78,6 +206,19 @@ def _handle_load(backend, cmd: dict, resp_queue: Any) -> None: load_in_4bit = cmd.get("load_in_4bit", True) trust_remote_code = cmd.get("trust_remote_code", False) + # Auto-enable trust_remote_code for NemotronH/Nano models. + if not trust_remote_code: + _NEMOTRON_TRUST_SUBSTRINGS = ("nemotron_h", "nemotron-h", "nemotron-3-nano") + _cp_lower = checkpoint_path.lower() + if any(sub in _cp_lower for sub in _NEMOTRON_TRUST_SUBSTRINGS) and ( + _cp_lower.startswith("unsloth/") or _cp_lower.startswith("nvidia/") + ): + trust_remote_code = True + logger.info( + "Auto-enabled trust_remote_code for Nemotron model: %s", + checkpoint_path, + ) + try: _send_response( resp_queue, @@ -126,9 +267,17 @@ def _handle_export(backend, cmd: dict, resp_queue: Any) -> None: export_type = cmd["export_type"] # "merged", "base", "gguf", "lora" response_type = f"export_{export_type}_done" + # Open the log forwarding gate so the user sees the actual export + # progress (Unsloth merge bars, file copies, GGUF conversion, etc.) + # in the live log panel. The gate stays open for the rest of this + # subprocess's life; the orchestrator spawns a fresh subprocess for + # the next checkpoint load, which resets the gate to closed. + _log_forward_gate.set() + + output_path: Any = None try: if export_type == "merged": - success, message = backend.export_merged_model( + success, message, output_path = backend.export_merged_model( save_directory = cmd.get("save_directory", ""), format_type = cmd.get("format_type", "16-bit (FP16)"), push_to_hub = cmd.get("push_to_hub", False), @@ -137,7 +286,7 @@ def _handle_export(backend, cmd: dict, resp_queue: Any) -> None: private = cmd.get("private", False), ) elif export_type == "base": - success, message = backend.export_base_model( + success, message, output_path = backend.export_base_model( save_directory = cmd.get("save_directory", ""), push_to_hub = cmd.get("push_to_hub", False), repo_id = cmd.get("repo_id"), @@ -146,7 +295,7 @@ def _handle_export(backend, cmd: dict, resp_queue: Any) -> None: base_model_id = cmd.get("base_model_id"), ) elif export_type == "gguf": - success, message = backend.export_gguf( + success, message, output_path = backend.export_gguf( save_directory = cmd.get("save_directory", ""), quantization_method = cmd.get("quantization_method", "Q4_K_M"), push_to_hub = cmd.get("push_to_hub", False), @@ -154,7 +303,7 @@ def _handle_export(backend, cmd: dict, resp_queue: Any) -> None: hf_token = cmd.get("hf_token"), ) elif export_type == "lora": - success, message = backend.export_lora_adapter( + success, message, output_path = backend.export_lora_adapter( save_directory = cmd.get("save_directory", ""), push_to_hub = cmd.get("push_to_hub", False), repo_id = cmd.get("repo_id"), @@ -170,6 +319,7 @@ def _handle_export(backend, cmd: dict, resp_queue: Any) -> None: "type": response_type, "success": success, "message": message, + "output_path": output_path, "ts": time.time(), }, ) @@ -181,6 +331,7 @@ def _handle_export(backend, cmd: dict, resp_queue: Any) -> None: "type": response_type, "success": False, "message": str(exc), + "output_path": None, "stack": traceback.format_exc(limit = 20), "ts": time.time(), }, @@ -226,10 +377,26 @@ def run_export_process( """ import queue as _queue + # Install fd-level stdout/stderr capture FIRST so every subsequent + # print and every child process inherits the redirected fds. This + # is what powers the live export log stream in the UI. + _setup_log_capture(resp_queue) + os.environ["TOKENIZERS_PARALLELISM"] = "false" os.environ["PYTHONWARNINGS"] = ( "ignore" # Suppress warnings at C-level before imports ) + # Force unbuffered output from any child Python process (e.g. the + # GGUF converter) so their prints surface in the log stream as they + # happen rather than at the end. + os.environ["PYTHONUNBUFFERED"] = "1" + # tqdm defaults to a 10-second mininterval when stdout is not a tty + # (which it isn't here -- we redirected fd 1/2 to a pipe). That makes + # multi-step progress bars look frozen in the export log panel. Force + # frequent flushes so the user sees movement during merge / GGUF + # conversion. Has no effect on single-step bars (e.g. "Copying 1 + # files") which only emit start/end events regardless. + os.environ.setdefault("TQDM_MININTERVAL", "0.5") import warnings from loggers.config import LogConfig diff --git a/studio/backend/core/inference/anthropic_compat.py b/studio/backend/core/inference/anthropic_compat.py new file mode 100644 index 0000000000..e7b40a60ce --- /dev/null +++ b/studio/backend/core/inference/anthropic_compat.py @@ -0,0 +1,488 @@ +# SPDX-License-Identifier: AGPL-3.0-only +# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. + +""" +Anthropic Messages API ↔ OpenAI format translation utilities. + +Pure functions and a stateful stream emitter — no FastAPI, no I/O. +""" + +from __future__ import annotations + +import json +from typing import Any, Optional, Union + + +def anthropic_messages_to_openai( + messages: list[dict], + system: Optional[Union[str, list]] = None, +) -> list[dict]: + """Convert Anthropic messages + system to OpenAI-format message dicts.""" + result: list[dict] = [] + + # System prompt + if system: + if isinstance(system, str): + result.append({"role": "system", "content": system}) + elif isinstance(system, list): + parts = [] + for block in system: + if isinstance(block, dict) and block.get("type") == "text": + parts.append(block["text"]) + elif isinstance(block, str): + parts.append(block) + if parts: + result.append({"role": "system", "content": "\n".join(parts)}) + + for msg in messages: + role = msg["role"] if isinstance(msg, dict) else msg.role + content = msg["content"] if isinstance(msg, dict) else msg.content + + if isinstance(content, str): + result.append({"role": role, "content": content}) + continue + + # Content is a list of blocks + text_parts: list[str] = [] + tool_calls: list[dict] = [] + tool_results: list[dict] = [] + + for block in content: + b = block if isinstance(block, dict) else block.model_dump() + btype = b.get("type", "") + + if btype == "text": + text_parts.append(b["text"]) + elif btype == "tool_use": + tool_calls.append( + { + "id": b["id"], + "type": "function", + "function": { + "name": b["name"], + "arguments": json.dumps(b["input"]), + }, + } + ) + elif btype == "tool_result": + tc = b.get("content", "") + if isinstance(tc, list): + tc = " ".join( + p["text"] + for p in tc + if isinstance(p, dict) and p.get("type") == "text" + ) + tool_results.append( + { + "role": "tool", + "tool_call_id": b["tool_use_id"], + "content": str(tc), + } + ) + + if role == "assistant": + msg_dict: dict[str, Any] = {"role": "assistant"} + if text_parts: + msg_dict["content"] = "\n".join(text_parts) + if tool_calls: + msg_dict["tool_calls"] = tool_calls + result.append(msg_dict) + elif role == "user": + if text_parts: + result.append({"role": "user", "content": "\n".join(text_parts)}) + for tr in tool_results: + result.append(tr) + + return result + + +def anthropic_tools_to_openai(tools: list) -> list[dict]: + """Convert Anthropic tool definitions to OpenAI function-tool format.""" + result = [] + for t in tools: + td = t if isinstance(t, dict) else t.model_dump() + result.append( + { + "type": "function", + "function": { + "name": td["name"], + "description": td.get("description", ""), + "parameters": td.get("input_schema", {}), + }, + } + ) + return result + + +def build_anthropic_sse_event(event_type: str, data: dict) -> str: + """Format a single Anthropic SSE event.""" + return f"event: {event_type}\ndata: {json.dumps(data)}\n\n" + + +class AnthropicStreamEmitter: + """Converts generator events from generate_chat_completion_with_tools() + into Anthropic Messages SSE strings.""" + + def __init__(self) -> None: + self.block_index: int = 0 + self._text_block_open: bool = False + self._prev_text: str = "" + self._usage: dict = {} + + def start(self, message_id: str, model: str) -> list[str]: + """Emit message_start and open the first text content block.""" + events = [] + events.append( + build_anthropic_sse_event( + "message_start", + { + "type": "message_start", + "message": { + "id": message_id, + "type": "message", + "role": "assistant", + "content": [], + "model": model, + "stop_reason": None, + "stop_sequence": None, + "usage": {"input_tokens": 0, "output_tokens": 0}, + }, + }, + ) + ) + events.extend(self._open_text_block()) + return events + + def feed(self, event: dict) -> list[str]: + """Process one generator event, return SSE strings.""" + etype = event.get("type", "") + if etype == "content": + return self._handle_content(event) + elif etype == "tool_start": + return self._handle_tool_start(event) + elif etype == "tool_end": + return self._handle_tool_end(event) + elif etype == "metadata": + self._usage = event.get("usage", {}) + return [] + # status events — no Anthropic equivalent + return [] + + def finish(self, stop_reason: str = "end_turn") -> list[str]: + """Close any open block and emit message_delta + message_stop.""" + events = [] + if self._text_block_open: + events.append(self._close_block()) + events.append( + build_anthropic_sse_event( + "message_delta", + { + "type": "message_delta", + "delta": {"stop_reason": stop_reason, "stop_sequence": None}, + "usage": { + "output_tokens": self._usage.get("completion_tokens", 0), + }, + }, + ) + ) + events.append( + build_anthropic_sse_event( + "message_stop", + { + "type": "message_stop", + }, + ) + ) + return events + + def _handle_content(self, event: dict) -> list[str]: + cumulative = event.get("text", "") + new_text = cumulative[len(self._prev_text) :] + self._prev_text = cumulative + if not new_text: + return [] + if not self._text_block_open: + events = self._open_text_block() + else: + events = [] + events.append( + build_anthropic_sse_event( + "content_block_delta", + { + "type": "content_block_delta", + "index": self.block_index, + "delta": {"type": "text_delta", "text": new_text}, + }, + ) + ) + return events + + def _handle_tool_start(self, event: dict) -> list[str]: + events = [] + # Close current text block if open + if self._text_block_open: + events.append(self._close_block()) + # Open a tool_use block + self.block_index += 1 + events.append( + build_anthropic_sse_event( + "content_block_start", + { + "type": "content_block_start", + "index": self.block_index, + "content_block": { + "type": "tool_use", + "id": event.get("tool_call_id", ""), + "name": event.get("tool_name", ""), + "input": {}, + }, + }, + ) + ) + # Emit the arguments as input_json_delta + args = event.get("arguments", {}) + if args: + events.append( + build_anthropic_sse_event( + "content_block_delta", + { + "type": "content_block_delta", + "index": self.block_index, + "delta": { + "type": "input_json_delta", + "partial_json": json.dumps(args), + }, + }, + ) + ) + return events + + def _handle_tool_end(self, event: dict) -> list[str]: + events = [] + # Close the tool_use block + events.append(self._close_block()) + # Emit custom tool_result event (non-standard, ignored by SDKs) + events.append( + build_anthropic_sse_event( + "tool_result", + { + "type": "tool_result", + "tool_use_id": event.get("tool_call_id", ""), + "content": event.get("result", ""), + }, + ) + ) + # Open a new text block for the model's next response + self.block_index += 1 + events.extend(self._open_text_block()) + # Reset text tracking for the next synthesis turn + self._prev_text = "" + return events + + def _open_text_block(self) -> list[str]: + self._text_block_open = True + return [ + build_anthropic_sse_event( + "content_block_start", + { + "type": "content_block_start", + "index": self.block_index, + "content_block": {"type": "text", "text": ""}, + }, + ) + ] + + def _close_block(self) -> str: + self._text_block_open = False + return build_anthropic_sse_event( + "content_block_stop", + { + "type": "content_block_stop", + "index": self.block_index, + }, + ) + + +class AnthropicPassthroughEmitter: + """Converts llama-server's OpenAI-format streaming chunks into Anthropic SSE. + + Used for the client-side tool-use pass-through path: the client (e.g. Claude + Code) sends its own tool definitions in the ``tools`` field and expects to + execute them itself. We forward them to llama-server and translate the + streaming response back to Anthropic format without executing anything. + """ + + def __init__(self) -> None: + self.block_index: int = -1 + self._current_block_type: Optional[str] = None # "text" | "tool_use" | None + self._tool_call_states: dict = {} # delta index -> {block_index, id, name} + self._usage: dict = {} + self._stop_reason: str = "end_turn" + + def start(self, message_id: str, model: str) -> list[str]: + return [ + build_anthropic_sse_event( + "message_start", + { + "type": "message_start", + "message": { + "id": message_id, + "type": "message", + "role": "assistant", + "content": [], + "model": model, + "stop_reason": None, + "stop_sequence": None, + "usage": {"input_tokens": 0, "output_tokens": 0}, + }, + }, + ) + ] + + def feed_chunk(self, chunk: dict) -> list[str]: + """Process one OpenAI streaming chat.completion.chunk.""" + events: list[str] = [] + + # usage-only chunks carry token totals + usage = chunk.get("usage") + if usage: + self._usage = usage + + choices = chunk.get("choices") or [] + if not choices: + return events + + choice = choices[0] + delta = choice.get("delta") or {} + finish_reason = choice.get("finish_reason") + + # ── Text content ── + content = delta.get("content") + if content: + if self._current_block_type != "text": + if self._current_block_type is not None: + events.append(self._close_current_block()) + events.extend(self._open_text_block()) + events.append( + build_anthropic_sse_event( + "content_block_delta", + { + "type": "content_block_delta", + "index": self.block_index, + "delta": {"type": "text_delta", "text": content}, + }, + ) + ) + + # ── Tool calls (streaming deltas) ── + tool_calls = delta.get("tool_calls") or [] + for tc in tool_calls: + tc_idx = tc.get("index", 0) + fn = tc.get("function") or {} + if tc_idx not in self._tool_call_states: + # New tool call — close prior block, open tool_use block + if self._current_block_type is not None: + events.append(self._close_current_block()) + tc_id = tc.get("id", "") + tc_name = fn.get("name", "") + self.block_index += 1 + self._current_block_type = "tool_use" + self._tool_call_states[tc_idx] = { + "block_index": self.block_index, + "id": tc_id, + "name": tc_name, + } + events.append( + build_anthropic_sse_event( + "content_block_start", + { + "type": "content_block_start", + "index": self.block_index, + "content_block": { + "type": "tool_use", + "id": tc_id, + "name": tc_name, + "input": {}, + }, + }, + ) + ) + + args_delta = fn.get("arguments", "") + if args_delta: + events.append( + build_anthropic_sse_event( + "content_block_delta", + { + "type": "content_block_delta", + "index": self._tool_call_states[tc_idx]["block_index"], + "delta": { + "type": "input_json_delta", + "partial_json": args_delta, + }, + }, + ) + ) + + # ── Finish reason ── + if finish_reason: + if finish_reason == "tool_calls": + self._stop_reason = "tool_use" + elif finish_reason == "length": + self._stop_reason = "max_tokens" + else: + self._stop_reason = "end_turn" + + return events + + def finish(self) -> list[str]: + events: list[str] = [] + if self._current_block_type is not None: + events.append(self._close_current_block()) + events.append( + build_anthropic_sse_event( + "message_delta", + { + "type": "message_delta", + "delta": { + "stop_reason": self._stop_reason, + "stop_sequence": None, + }, + "usage": { + "output_tokens": self._usage.get("completion_tokens", 0), + }, + }, + ) + ) + events.append( + build_anthropic_sse_event( + "message_stop", + {"type": "message_stop"}, + ) + ) + return events + + def _open_text_block(self) -> list[str]: + self.block_index += 1 + self._current_block_type = "text" + return [ + build_anthropic_sse_event( + "content_block_start", + { + "type": "content_block_start", + "index": self.block_index, + "content_block": {"type": "text", "text": ""}, + }, + ) + ] + + def _close_current_block(self) -> str: + idx = self.block_index + self._current_block_type = None + return build_anthropic_sse_event( + "content_block_stop", + { + "type": "content_block_stop", + "index": idx, + }, + ) diff --git a/studio/backend/core/inference/inference.py b/studio/backend/core/inference/inference.py index 867bdefc62..4c140013a0 100644 --- a/studio/backend/core/inference/inference.py +++ b/studio/backend/core/inference/inference.py @@ -253,6 +253,10 @@ class InferenceBackend: """ Load any model: base, LoRA adapter, text, or vision. """ + # GGUF uses max_seq_length=0 as "model default"; Unsloth crashes on it. + if max_seq_length <= 0: + max_seq_length = 2048 + try: model_name = config.identifier diff --git a/studio/backend/core/inference/llama_cpp.py b/studio/backend/core/inference/llama_cpp.py index c84ac640df..b53fc513de 100644 --- a/studio/backend/core/inference/llama_cpp.py +++ b/studio/backend/core/inference/llama_cpp.py @@ -167,7 +167,17 @@ class LlamaCppBackend: @property def max_context_length(self) -> Optional[int]: - """Return the maximum context currently available on this hardware.""" + """Return the largest context that fits on this hardware at load time. + + This is the "safe zone" threshold the UI renders warnings + against. For a model whose weights fit on some GPU subset, it + is the binary-search cap from ``_fit_context_to_vram`` for that + subset. For a model whose weights exceed 90% of every GPU + subset, it is the 4096 fallback -- the spec's default when the + model will not fit. The UI slider ceiling is + ``native_context_length``; dragging above ``max_context_length`` + triggers the "might be slower" warning. + """ return self._max_context_length or self._context_length @property @@ -175,6 +185,96 @@ class LlamaCppBackend: """Return the model's native context length from GGUF metadata.""" return self._context_length + def load_progress(self) -> Optional[dict]: + """Return live model-load progress, or None if not loading. + + While llama-server is warming up, its process is typically in + kernel state D (disk sleep) mmap'ing the weight shards into + page cache before pushing layers to VRAM. During that window + ``/api/inference/status`` only reports ``loading``, which gives + the UI nothing to display besides a spinner that looks stuck + for minutes on large MoE models. + + This method samples ``/proc//status VmRSS`` against the + sum of the GGUF shard sizes so the UI can render a real bar + and compute rate / ETA. Returns ``None`` when no load is in + flight (no process, or process already healthy). + + Shape:: + + { + "phase": "mmap" | "ready", + "bytes_loaded": int, # VmRSS of the llama-server + "bytes_total": int, # sum of shard file sizes + "fraction": float, # bytes_loaded / bytes_total, 0..1 + } + + Linux-only in the current implementation. On macOS/Windows the + equivalent would be a different API; this returns ``None`` on + platforms where ``/proc//status`` is unavailable. + """ + proc = self._process + if proc is None: + return None + pid = proc.pid + if pid is None: + return None + + # Sum up shard sizes (primary + any extras sitting alongside). + bytes_total = 0 + gguf_path = self._gguf_path + if gguf_path: + primary = Path(gguf_path) + try: + if primary.is_file(): + bytes_total += primary.stat().st_size + except OSError: + pass + # Extra shards live alongside the primary with the same prefix + # before the shard index (e.g. ``-00001-of-00004.gguf``). + try: + parent = primary.parent + stem = primary.name + m = _SHARD_RE.match(stem) + prefix = m.group(1) if m else None + if prefix and parent.is_dir(): + for sibling in parent.iterdir(): + if ( + sibling.is_file() + and sibling.name.startswith(prefix) + and sibling.name != stem + and sibling.suffix == ".gguf" + ): + try: + bytes_total += sibling.stat().st_size + except OSError: + pass + except OSError: + pass + + # Read VmRSS from /proc//status. Kilobytes on Linux. + bytes_loaded = 0 + try: + with open(f"/proc/{pid}/status", "r", encoding = "utf-8") as f: + for line in f: + if line.startswith("VmRSS:"): + kb = int(line.split()[1]) + bytes_loaded = kb * 1024 + break + except (FileNotFoundError, PermissionError, ValueError, OSError): + return None + + phase = "ready" if self._healthy else "mmap" + fraction = 0.0 + if bytes_total > 0: + fraction = min(1.0, bytes_loaded / bytes_total) + return { + "phase": phase, + "bytes_loaded": bytes_loaded, + "bytes_total": bytes_total, + "fraction": round(fraction, 4), + } + @property def chat_template(self) -> Optional[str]: return self._chat_template @@ -904,10 +1004,34 @@ class LlamaCppBackend: try: import os - from huggingface_hub import get_paths_info + from huggingface_hub import get_paths_info, try_to_load_from_cache path_infos = list(get_paths_info(hf_repo, all_gguf_files, token = hf_token)) - total_download_bytes = sum((p.size or 0) for p in path_infos) + total_bytes = sum((p.size or 0) for p in path_infos) + + # Subtract bytes already present in the HF cache so we only + # preflight against what we actually have to download. Without + # this, re-loading a cached large model (e.g. MiniMax-M2.7-GGUF + # at 131 GB) fails cold whenever free disk is below the full + # weight footprint, even though nothing needs downloading. + already_cached_bytes = 0 + for p in path_infos: + if not p.size: + continue + try: + cached_path = try_to_load_from_cache(hf_repo, p.path) + except Exception: + cached_path = None + if isinstance(cached_path, str) and os.path.exists(cached_path): + try: + on_disk = os.path.getsize(cached_path) + except OSError: + on_disk = 0 + # Count as satisfied only when the full blob is present. + if on_disk >= p.size: + already_cached_bytes += p.size + + total_download_bytes = max(0, total_bytes - already_cached_bytes) if total_download_bytes > 0: cache_dir = os.environ.get( @@ -919,9 +1043,11 @@ class LlamaCppBackend: total_gb = total_download_bytes / (1024**3) free_gb = free_bytes / (1024**3) + cached_gb = already_cached_bytes / (1024**3) logger.info( - f"GGUF download: {total_gb:.1f} GB needed, " + f"GGUF download: {total_gb:.1f} GB needed " + f"({cached_gb:.1f} GB already cached), " f"{free_gb:.1f} GB free on disk" ) @@ -1063,6 +1189,7 @@ class LlamaCppBackend: speculative_type: Optional[str] = None, n_threads: Optional[int] = None, n_gpu_layers: Optional[int] = None, # Accepted for caller compat, unused + n_parallel: int = 1, ) -> bool: """ Start llama-server with a GGUF model. @@ -1192,36 +1319,28 @@ class LlamaCppBackend: best_cap = max(best_cap, capped) if best_cap > 0: max_available_ctx = best_cap + else: + # Weights exceed 90% of every GPU subset's free + # memory, so there is no fitting context. Anchor + # the UI's "safe zone" threshold at 4096 (the + # spec's default when the model cannot fit) so + # the ctx slider shows the "might be slower" + # warning as soon as the user drags above the + # fallback default instead of never. + max_available_ctx = min(4096, native_ctx_for_cap) if explicit_ctx: - # Try to honor the user's requested context exactly. + # Honor the user's requested context verbatim. If it + # fits, pin GPUs and skip --fit; if it doesn't, ship + # -c --fit on and let llama-server flex + # -ngl (CPU layer offload). The UI is expected to + # have surfaced the "might be slower" warning before + # the user submitted a ctx above the fit ceiling. requested_total = model_size + self._estimate_kv_cache_bytes( effective_ctx, cache_type_kv ) gpu_indices, use_fit = self._select_gpus(requested_total, gpus) - - # Full context doesn't fit anywhere -- cap it on the - # best GPU subset we can find (fewest GPUs first). - if use_fit: - ranked = sorted(gpus, key = lambda g: g[1], reverse = True) - for n_gpus in range(1, len(ranked) + 1): - subset = ranked[:n_gpus] - pool_mib = sum(free for _, free in subset) - capped = self._fit_context_to_vram( - effective_ctx, - pool_mib, - model_size, - cache_type_kv, - ) - kv = self._estimate_kv_cache_bytes( - capped, cache_type_kv - ) - total_mib = (model_size + kv) / (1024 * 1024) - if total_mib <= pool_mib * 0.90: - effective_ctx = capped - gpu_indices = sorted(idx for idx, _ in subset) - use_fit = False - break + # No silent shrink: effective_ctx stays == n_ctx. else: # Auto context: prefer fewer GPUs, cap context to fit. ranked = sorted(gpus, key = lambda g: g[1], reverse = True) @@ -1241,6 +1360,13 @@ class LlamaCppBackend: gpu_indices = sorted(idx for idx, _ in subset) use_fit = False break + else: + # No subset can host the weights (weights alone + # exceed 90% of every pool). Per spec, default + # the UI-visible context to 4096 and let + # --fit on flex -ngl so llama-server offloads + # layers to CPU RAM. + effective_ctx = min(4096, effective_ctx) elif gpus: # Can't estimate KV -- fall back to file-size-only check. @@ -1251,6 +1377,13 @@ class LlamaCppBackend: model_size_gb = round(model_size / (1024**3), 2), ) gpu_indices, use_fit = self._select_gpus(model_size, gpus) + if use_fit and not explicit_ctx: + # Weights don't fit on any subset. Default the UI to + # 4096 so the slider doesn't land on an unusable native + # context. --fit on will flex -ngl at runtime. + effective_ctx = ( + min(4096, effective_ctx) if effective_ctx > 0 else 4096 + ) if effective_ctx < original_ctx: kv_est = self._estimate_kv_cache_bytes(effective_ctx, cache_type_kv) @@ -1283,9 +1416,11 @@ class LlamaCppBackend: "-c", str(effective_ctx) if effective_ctx > 0 else "0", "--parallel", - "1", # Single-user studio, saves VRAM + str(n_parallel), "--flash-attn", "on", # Force flash attention for speed + # Error out at n_ctx instead of silently rotating the KV cache; frontend catches it and points the user at "Context Length". + "--no-context-shift", ] if use_fit: @@ -1531,7 +1666,12 @@ class LlamaCppBackend: ) self._stdout_thread.start() - self._gguf_path = gguf_path + # Store the resolved on-disk path, not the caller's kwarg. In + # HF mode the caller passes gguf_path=None and the real path + # (``model_path``) is what llama-server is actually mmap'ing. + # Downstream consumers (load_progress, log lines, etc.) need + # the path that exists on disk. + self._gguf_path = model_path self._hf_repo = hf_repo # For local GGUF files, extract variant from filename if not provided if hf_variant: diff --git a/studio/backend/core/inference/worker.py b/studio/backend/core/inference/worker.py index 7f7291a56d..fbcce276ba 100644 --- a/studio/backend/core/inference/worker.py +++ b/studio/backend/core/inference/worker.py @@ -34,37 +34,15 @@ from utils.hardware import apply_gpu_ids def _activate_transformers_version(model_name: str) -> None: - """Activate the correct transformers version BEFORE any ML imports. - - If the model needs transformers 5.x, prepend the pre-installed .venv_t5/ - directory to sys.path. Otherwise do nothing (default 4.57.x in .venv/). - """ + """Activate the correct transformers version BEFORE any ML imports.""" # Ensure backend is on path for utils imports backend_path = str(Path(__file__).resolve().parent.parent.parent) if backend_path not in sys.path: sys.path.insert(0, backend_path) - from utils.transformers_version import ( - needs_transformers_5, - _resolve_base_model, - _ensure_venv_t5_exists, - _VENV_T5_DIR, - ) + from utils.transformers_version import activate_transformers_for_subprocess - resolved = _resolve_base_model(model_name) - if needs_transformers_5(resolved): - if not _ensure_venv_t5_exists(): - raise RuntimeError( - f"Cannot activate transformers 5.x: .venv_t5 missing at {_VENV_T5_DIR}" - ) - if _VENV_T5_DIR not in sys.path: - sys.path.insert(0, _VENV_T5_DIR) - logger.info("Activated transformers 5.x from %s", _VENV_T5_DIR) - # Propagate to child subprocesses (e.g. GGUF converter) - _pp = os.environ.get("PYTHONPATH", "") - os.environ["PYTHONPATH"] = _VENV_T5_DIR + (os.pathsep + _pp if _pp else "") - else: - logger.info("Using default transformers (4.57.x) for %s", model_name) + activate_transformers_for_subprocess(model_name) def _decode_image(image_base64: str): @@ -309,19 +287,21 @@ def _handle_load(backend, config: dict, resp_queue: Any) -> None: except Exception as e: logger.warning("Could not read adapter_config.json: %s", e) - # Auto-enable trust_remote_code for unsloth/* transformers 5.x models - # (matches the training worker logic in core/training/worker.py) + # Auto-enable trust_remote_code for NemotronH/Nano models only. + # NemotronH has config parsing bugs requiring trust_remote_code=True. + # Other transformers 5.x models are native and do NOT need it. + # NOTE: Must NOT match Llama-Nemotron (standard Llama architecture). + _NEMOTRON_TRUST_SUBSTRINGS = ("nemotron_h", "nemotron-h", "nemotron-3-nano") trust_remote_code = config.get("trust_remote_code", False) if not trust_remote_code: - from utils.transformers_version import needs_transformers_5 - model_name = config["model_name"] - if needs_transformers_5(model_name) and model_name.lower().startswith( - "unsloth/" + _mn_lower = model_name.lower() + if any(sub in _mn_lower for sub in _NEMOTRON_TRUST_SUBSTRINGS) and ( + _mn_lower.startswith("unsloth/") or _mn_lower.startswith("nvidia/") ): trust_remote_code = True logger.info( - "Auto-enabled trust_remote_code for unsloth/* transformers 5.x model: %s", + "Auto-enabled trust_remote_code for Nemotron model: %s", model_name, ) diff --git a/studio/backend/core/training/worker.py b/studio/backend/core/training/worker.py index 0454eada89..8ab2b5b2be 100644 --- a/studio/backend/core/training/worker.py +++ b/studio/backend/core/training/worker.py @@ -16,26 +16,31 @@ from __future__ import annotations import structlog from loggers import get_logger import os -import platform import shutil import sys import time import traceback -import json import subprocess as _sp from pathlib import Path -from typing import Any -import urllib.error -import urllib.request +from typing import Any, Callable logger = get_logger(__name__) from utils.hardware import apply_gpu_ids +from utils.wheel_utils import ( + direct_wheel_url, + flash_attn_wheel_url, + install_wheel, + probe_torch_wheel_env, + url_exists, +) _CAUSAL_CONV1D_RELEASE_TAG = "v1.6.1.post4" _CAUSAL_CONV1D_PACKAGE_VERSION = "1.6.1" _MAMBA_SSM_RELEASE_TAG = "v2.3.1" _MAMBA_SSM_PACKAGE_VERSION = "2.3.1" +_FLASH_ATTN_RUNTIME_MIN_SEQ_LEN = 32768 +_FLASH_ATTN_SKIP_ENV = "UNSLOTH_STUDIO_SKIP_FLASHATTN_INSTALL" def _model_wants_causal_conv1d(model_name: str) -> bool: @@ -59,206 +64,186 @@ def _model_wants_causal_conv1d(model_name: str) -> bool: ) -def _causal_conv1d_platform_tag() -> str | None: - machine = platform.machine().lower() - if sys.platform.startswith("linux"): - if machine in {"x86_64", "amd64"}: - return "linux_x86_64" - if machine in {"aarch64", "arm64"}: - return "linux_aarch64" - return None - # No prebuilt wheels published for macOS or Windows - return None - - -def _probe_causal_conv1d_env() -> dict[str, str] | None: - try: - probe = _sp.run( - [ - sys.executable, - "-c", - ( - "import json, sys, re, torch; " - "parts = torch.__version__.split('+', 1)[0].split('.')[:2]; " - "minor = re.sub(r'[^0-9].*', '', parts[1]) if len(parts) > 1 else '0'; " - "torch_mm = parts[0] + '.' + minor; " - "print(json.dumps({" - "'python_tag': f'cp{sys.version_info.major}{sys.version_info.minor}', " - "'torch_mm': torch_mm, " - "'cuda_major': str(int(str(torch.version.cuda).split('.', 1)[0])) if torch.version.cuda else '', " - "'cxx11abi': str(torch._C._GLIBCXX_USE_CXX11_ABI).upper()" - "}))" - ), - ], - stdout = _sp.PIPE, - stderr = _sp.PIPE, - text = True, - timeout = 30, - ) - except _sp.TimeoutExpired: - logger.warning("Torch environment probe timed out after 30s") - return None - if probe.returncode != 0: - logger.warning( - "Failed to probe torch environment for causal-conv1d wheel:\n%s", - probe.stdout, - ) - return None - - try: - return json.loads(probe.stdout.strip()) - except json.JSONDecodeError: - logger.warning( - "Failed to parse torch environment probe output: %s", probe.stdout - ) - return None - - -def _direct_wheel_url( - *, - filename_prefix: str, - package_version: str, - release_tag: str, - release_base_url: str, - env: dict[str, str] | None = None, -) -> str | None: - env = env or _probe_causal_conv1d_env() - platform_tag = _causal_conv1d_platform_tag() - if env is None or platform_tag is None or not env.get("cuda_major"): - return None - - filename = ( - f"{filename_prefix}-{package_version}" - f"+cu{env['cuda_major']}torch{env['torch_mm']}" - f"cxx11abi{env['cxx11abi']}-{env['python_tag']}-{env['python_tag']}-{platform_tag}.whl" - ) - return f"{release_base_url}/{release_tag}/{filename}" - - -def _url_exists(url: str) -> bool: - try: - request = urllib.request.Request(url, method = "HEAD") - with urllib.request.urlopen(request, timeout = 10): - return True - except urllib.error.HTTPError as exc: - if exc.code == 404: - return False - logger.warning("Unexpected HTTP error while probing %s: %s", url, exc) - return False - except Exception as exc: - logger.warning("Failed to probe %s: %s", url, exc) - return False - - def _install_package_wheel_first( *, event_queue: Any, import_name: str, display_name: str, pypi_name: str, - pypi_version: str, - filename_prefix: str, - release_tag: str, - release_base_url: str, -) -> None: + pypi_version: str | None = None, + filename_prefix: str | None = None, + release_tag: str | None = None, + release_base_url: str | None = None, + wheel_url_builder: Callable[[dict[str, str] | None], str | None] | None = None, + pypi_spec: str | None = None, + pypi_status_message: str | None = None, +) -> bool: try: __import__(import_name) logger.info("%s already installed", display_name) - return + return True except ImportError: pass - env = _probe_causal_conv1d_env() - wheel_url = _direct_wheel_url( - filename_prefix = filename_prefix, - package_version = pypi_version, - release_tag = release_tag, - release_base_url = release_base_url, - env = env, - ) + env = probe_torch_wheel_env(timeout = 30) + if wheel_url_builder is not None: + wheel_url = wheel_url_builder(env) + else: + wheel_url = direct_wheel_url( + filename_prefix = filename_prefix, + package_version = pypi_version, + release_tag = release_tag, + release_base_url = release_base_url, + env = env, + ) if wheel_url is None: logger.info("No compatible %s wheel candidate", display_name) - else: - if _url_exists(wheel_url): - _send_status(event_queue, f"Installing prebuilt {display_name} wheel...") - installed = False - # Try uv first if available, then fall back to pip - if shutil.which("uv"): - uv_cmd = [ - "uv", - "pip", - "install", - "--python", - sys.executable, - "--no-deps", - wheel_url, - ] - result = _sp.run( - uv_cmd, - stdout = _sp.PIPE, - stderr = _sp.STDOUT, - text = True, - ) - if result.returncode == 0: - installed = True - else: - logger.warning( - "uv failed to install %s wheel:\n%s", - display_name, - result.stdout, - ) - if not installed: - pip_cmd = [ - sys.executable, - "-m", - "pip", - "install", - "--no-deps", - wheel_url, - ] - result = _sp.run( - pip_cmd, - stdout = _sp.PIPE, - stderr = _sp.STDOUT, - text = True, - ) - if result.returncode == 0: - installed = True - else: - logger.warning( - "pip failed to install %s wheel:\n%s", - display_name, - result.stdout, - ) - if installed: + elif url_exists(wheel_url): + _send_status(event_queue, f"Installing prebuilt {display_name} wheel...") + for installer, result in install_wheel( + wheel_url, + python_executable = sys.executable, + use_uv = bool(shutil.which("uv")), + run = _sp.run, + ): + if result.returncode == 0: logger.info("Installed prebuilt %s wheel successfully", display_name) - return + return True + logger.warning( + "%s failed to install %s wheel:\n%s", + installer, + display_name, + result.stdout, + ) + else: + logger.info("No published %s wheel found: %s", display_name, wheel_url) + + is_hip = env and env.get("hip_version") + if is_hip and not shutil.which("hipcc"): + logger.error( + "%s requires hipcc for source compilation on ROCm. " + "Install the ROCm HIP SDK: https://rocm.docs.amd.com", + display_name, + ) + _send_status( + event_queue, + f"{display_name}: hipcc not found (ROCm HIP SDK required)", + ) + return False + + if pypi_spec is None: + pypi_spec = f"{pypi_name}=={pypi_version}" + + if pypi_status_message is None: + if is_hip: + pypi_status_message = ( + f"Compiling {display_name} from source for ROCm " + "(this may take several minutes)..." + ) else: - logger.info("No published %s wheel found: %s", display_name, wheel_url) + pypi_status_message = f"Installing {display_name} from PyPI..." + + _send_status(event_queue, pypi_status_message) + + # Prefer uv for faster dependency resolution when available + plain_pypi_install = pypi_version is None + if plain_pypi_install: + if shutil.which("uv"): + pypi_cmd = [ + "uv", + "pip", + "install", + "--python", + sys.executable, + pypi_spec, + ] + else: + pypi_cmd = [sys.executable, "-m", "pip", "install", pypi_spec] + else: + if shutil.which("uv"): + pypi_cmd = [ + "uv", + "pip", + "install", + "--python", + sys.executable, + "--no-build-isolation", + "--no-deps", + ] + # Avoid stale cache artifacts from partial HIP source builds + if is_hip: + pypi_cmd.append("--no-cache") + pypi_cmd.append(pypi_spec) + else: + pypi_cmd = [ + sys.executable, + "-m", + "pip", + "install", + "--no-build-isolation", + "--no-deps", + "--no-cache-dir", + pypi_spec, + ] + + # Source compilation on ROCm can take 10-30 minutes; use a generous + # timeout. Non-HIP installs preserve the pre-existing "no timeout" + # behaviour so unrelated slow installs (e.g. causal-conv1d source + # build on Linux aarch64 or unsupported torch/CUDA combinations) + # are not aborted at 5 minutes by this PR. + _run_kwargs: dict[str, Any] = { + "stdout": _sp.PIPE, + "stderr": _sp.STDOUT, + "text": True, + } + if is_hip: + _run_kwargs["timeout"] = 1800 + + try: + result = _sp.run(pypi_cmd, **_run_kwargs) + except _sp.TimeoutExpired: + logger.error( + "%s installation timed out after %ds", + display_name, + _run_kwargs.get("timeout"), + ) + _send_status( + event_queue, + f"{display_name} installation timed out after " + f"{_run_kwargs.get('timeout')}s", + ) + return False - _send_status(event_queue, f"Installing {display_name} from PyPI...") - pypi_cmd = [ - sys.executable, - "-m", - "pip", - "install", - "--no-build-isolation", - "--no-deps", - "--no-cache-dir", - f"{pypi_name}=={pypi_version}", - ] - result = _sp.run( - pypi_cmd, - stdout = _sp.PIPE, - stderr = _sp.STDOUT, - text = True, - ) if result.returncode != 0: - logger.error("Failed to install %s from PyPI:\n%s", display_name, result.stdout) - return + if is_hip: + # Surface a clear error for ROCm source build failures + error_lines = (result.stdout or "").strip().splitlines() + snippet = "\n".join(error_lines[-5:]) if error_lines else "(no output)" + logger.error( + "Failed to compile %s for ROCm:\n%s", + display_name, + result.stdout, + ) + _send_status( + event_queue, + f"Failed to compile {display_name} for ROCm. " + "Check that hipcc and ROCm development headers are installed.\n" + f"{snippet}", + ) + else: + logger.error( + "Failed to install %s from PyPI:\n%s", + display_name, + result.stdout, + ) + return False - logger.info("Installed %s from PyPI", display_name) + if is_hip: + logger.info("Compiled and installed %s from source for ROCm", display_name) + else: + logger.info("Installed %s from PyPI", display_name) + return True def _ensure_causal_conv1d_fast_path(event_queue: Any, model_name: str) -> None: @@ -305,38 +290,41 @@ def _ensure_mamba_ssm(event_queue: Any, model_name: str) -> None: ) -def _activate_transformers_version(model_name: str) -> None: - """Activate the correct transformers version BEFORE any ML imports. +def _should_try_runtime_flash_attn_install(max_seq_length: int) -> bool: + if os.getenv(_FLASH_ATTN_SKIP_ENV) == "1": + return False + if max_seq_length < _FLASH_ATTN_RUNTIME_MIN_SEQ_LEN: + return False + return sys.platform.startswith("linux") - If the model needs transformers 5.x, prepend the pre-installed .venv_t5/ - directory to sys.path. Otherwise do nothing (default 4.57.x in .venv/). - """ + +def _ensure_flash_attn_for_long_context(event_queue: Any, max_seq_length: int) -> None: + if not _should_try_runtime_flash_attn_install(max_seq_length): + return + + installed = _install_package_wheel_first( + event_queue = event_queue, + import_name = "flash_attn", + display_name = "flash-attn", + pypi_name = "flash-attn", + wheel_url_builder = flash_attn_wheel_url, + pypi_spec = "flash-attn", + pypi_status_message = "Installing flash-attn from PyPI for long-context training...", + ) + if not installed: + _send_status(event_queue, "Continuing without flash-attn") + + +def _activate_transformers_version(model_name: str) -> None: + """Activate the correct transformers version BEFORE any ML imports.""" # Ensure backend is on path for utils imports backend_path = str(Path(__file__).resolve().parent.parent.parent) if backend_path not in sys.path: sys.path.insert(0, backend_path) - from utils.transformers_version import ( - needs_transformers_5, - _resolve_base_model, - _ensure_venv_t5_exists, - _VENV_T5_DIR, - ) + from utils.transformers_version import activate_transformers_for_subprocess - resolved = _resolve_base_model(model_name) - if needs_transformers_5(resolved): - if not _ensure_venv_t5_exists(): - raise RuntimeError( - f"Cannot activate transformers 5.x: .venv_t5 missing at {_VENV_T5_DIR}" - ) - if _VENV_T5_DIR not in sys.path: - sys.path.insert(0, _VENV_T5_DIR) - logger.info("Activated transformers 5.x from %s", _VENV_T5_DIR) - # Propagate to child subprocesses (e.g. GGUF converter) - _pp = os.environ.get("PYTHONPATH", "") - os.environ["PYTHONPATH"] = _VENV_T5_DIR + (os.pathsep + _pp if _pp else "") - else: - logger.info("Using default transformers (4.57.x) for %s", model_name) + activate_transformers_for_subprocess(model_name) def run_training_process( @@ -386,25 +374,22 @@ def run_training_process( ) return - # ── 1a. Auto-enable trust_remote_code for unsloth/* transformers 5.x models ── - # Some newer architectures (e.g. NemotronH) have config parsing bugs in - # transformers that require trust_remote_code=True as a workaround. - # Only auto-enable for unsloth/* prefixed models (trusted source). - # Exclude Gemma 4 since it is a native transformers 5.5 model and - # trust_remote_code=True would bypass the compiler (disabling fused CE). - from utils.transformers_version import needs_transformers_5 - + # ── 1a. Auto-enable trust_remote_code for NemotronH/Nano models ── + # NemotronH has config parsing bugs in transformers that require + # trust_remote_code=True as a workaround. Other transformers 5.x models + # (Qwen3.5, Gemma 4, etc.) are native and do NOT need it — enabling it + # bypasses the compiler (disabling fused CE). + # NOTE: Must NOT match Llama-Nemotron (standard Llama architecture). + _NEMOTRON_TRUST_SUBSTRINGS = ("nemotron_h", "nemotron-h", "nemotron-3-nano") _lowered = model_name.lower() - _is_native_t5 = any(x in _lowered for x in ("gemma-4", "gemma4")) if ( - needs_transformers_5(model_name) - and _lowered.startswith("unsloth/") - and not _is_native_t5 + any(sub in _lowered for sub in _NEMOTRON_TRUST_SUBSTRINGS) + and (_lowered.startswith("unsloth/") or _lowered.startswith("nvidia/")) and not config.get("trust_remote_code", False) ): config["trust_remote_code"] = True logger.info( - "Auto-enabled trust_remote_code for unsloth/* transformers 5.x model: %s", + "Auto-enabled trust_remote_code for Nemotron model: %s", model_name, ) @@ -412,6 +397,10 @@ def run_training_process( try: _ensure_causal_conv1d_fast_path(event_queue, model_name) _ensure_mamba_ssm(event_queue, model_name) + _ensure_flash_attn_for_long_context( + event_queue, + int(config.get("max_seq_length", 2048)), + ) except Exception as exc: event_queue.put( { diff --git a/studio/backend/main.py b/studio/backend/main.py index ad19ee9679..8a40791c06 100644 --- a/studio/backend/main.py +++ b/studio/backend/main.py @@ -237,6 +237,7 @@ async def get_system_info(): import platform import psutil from utils.hardware import get_device + from utils.hardware.hardware import _backend_label visibility_info = get_backend_visible_gpu_info() gpu_info = { @@ -250,7 +251,10 @@ async def get_system_info(): return { "platform": platform.platform(), "python_version": platform.python_version(), - "device_backend": get_device().value, + # Use the centralized _backend_label helper so the /api/system + # endpoint reports "rocm" on AMD hosts instead of "cuda", matching + # the /api/hardware and /api/gpu-visibility endpoints. + "device_backend": _backend_label(get_device()), "cpu_count": psutil.cpu_count(), "memory": { "total_gb": round(memory.total / 1e9, 2), @@ -349,7 +353,7 @@ def setup_frontend(app: FastAPI, build_path: Path): @app.get("/{full_path:path}") async def serve_frontend(full_path: str): - if full_path.startswith("api"): + if full_path in {"api", "v1"} or full_path.startswith(("api/", "v1/")): return {"error": "API endpoint not found"} file_path = (build_path / full_path).resolve() diff --git a/studio/backend/models/auth.py b/studio/backend/models/auth.py index 73d21130ae..c55e646508 100644 --- a/studio/backend/models/auth.py +++ b/studio/backend/models/auth.py @@ -5,6 +5,8 @@ Pydantic schemas for Authentication API """ +from typing import Optional + from pydantic import BaseModel, Field @@ -45,3 +47,44 @@ class ChangePasswordRequest(BaseModel): new_password: str = Field( ..., min_length = 8, description = "Replacement password (minimum 8 characters)" ) + + +# --------------------------------------------------------------------------- +# API key schemas +# --------------------------------------------------------------------------- + + +class CreateApiKeyRequest(BaseModel): + """Request body to create a new API key.""" + + name: str = Field(..., description = "Human-readable label for this key") + expires_in_days: Optional[int] = Field( + None, description = "Number of days until the key expires (None = never)" + ) + + +class ApiKeyResponse(BaseModel): + """Public representation of an API key (never contains the raw key).""" + + id: int + name: str + key_prefix: str = Field( + ..., description = "First 8 characters after sk-unsloth- for display" + ) + created_at: str + last_used_at: Optional[str] = None + expires_at: Optional[str] = None + is_active: bool + + +class CreateApiKeyResponse(BaseModel): + """Returned once when a key is created -- ``key`` is never shown again.""" + + key: str = Field(..., description = "Full API key (shown once)") + api_key: ApiKeyResponse + + +class ApiKeyListResponse(BaseModel): + """List of API keys for the authenticated user.""" + + api_keys: list[ApiKeyResponse] diff --git a/studio/backend/models/inference.py b/studio/backend/models/inference.py index cf08ecbc12..4917a14579 100644 --- a/studio/backend/models/inference.py +++ b/studio/backend/models/inference.py @@ -94,6 +94,10 @@ class ValidateModelResponse(BaseModel): is_gguf: bool = Field(False, description = "Whether this is a GGUF model (llama.cpp)") is_lora: bool = Field(False, description = "Whether this is a LoRA adapter") is_vision: bool = Field(False, description = "Whether this is a vision-capable model") + requires_trust_remote_code: bool = Field( + False, + description = "Whether the model defaults require trust_remote_code to be enabled for loading.", + ) class GenerateRequest(BaseModel): @@ -137,6 +141,10 @@ class LoadResponse(BaseModel): inference: dict = Field( ..., description = "Inference parameters (temperature, top_p, top_k, min_p)" ) + requires_trust_remote_code: bool = Field( + False, + description = "Whether the model defaults require trust_remote_code to be enabled for loading.", + ) context_length: Optional[int] = Field( None, description = "Model's native context length (from GGUF metadata)" ) @@ -180,6 +188,39 @@ class UnloadResponse(BaseModel): model: str = Field(..., description = "Model identifier that was unloaded") +class LoadProgressResponse(BaseModel): + """Progress of the active GGUF load, sampled on demand. + + Used by the UI to show a real progress bar during the + post-download warmup window (mmap + CUDA upload), rather than a + generic "Starting model..." spinner that freezes for minutes on + large MoE models. + """ + + phase: Optional[str] = Field( + None, + description = ( + "Load phase: 'mmap' (weights paging into RAM via mmap), " + "'ready' (llama-server reported healthy), or null when no " + "load is in flight." + ), + ) + bytes_loaded: int = Field( + 0, + description = ( + "Bytes of the model already resident in the llama-server " + "process (VmRSS on Linux)." + ), + ) + bytes_total: int = Field( + 0, + description = "Total bytes across all GGUF shards for the active model.", + ) + fraction: float = Field( + 0.0, description = "bytes_loaded / bytes_total, clamped to 0..1." + ) + + class InferenceStatusResponse(BaseModel): """Current inference backend status""" @@ -213,6 +254,10 @@ class InferenceStatusResponse(BaseModel): inference: Optional[Dict[str, Any]] = Field( None, description = "Recommended inference parameters for the active model" ) + requires_trust_remote_code: bool = Field( + False, + description = "Whether the active model requires trust_remote_code to be enabled for loading.", + ) supports_reasoning: bool = Field( False, description = "Whether the active model supports reasoning/thinking mode" ) @@ -444,3 +489,241 @@ class ChatCompletion(BaseModel): model: str = "default" choices: list[CompletionChoice] usage: CompletionUsage = Field(default_factory = CompletionUsage) + + +# ===================================================================== +# OpenAI Responses API Models (/v1/responses) +# ===================================================================== + + +# ── Request models ────────────────────────────────────────────── + + +class ResponsesInputTextPart(BaseModel): + """Text content part in a Responses API message (type=input_text).""" + + type: Literal["input_text"] + text: str + + +class ResponsesInputImagePart(BaseModel): + """Image content part in a Responses API message (type=input_image).""" + + type: Literal["input_image"] + image_url: str = Field(..., description = "data:image/png;base64,... or https://...") + detail: Optional[Literal["auto", "low", "high"]] = "auto" + + +ResponsesContentPart = Union[ResponsesInputTextPart, ResponsesInputImagePart] + + +class ResponsesInputMessage(BaseModel): + """A single message in the Responses API input array.""" + + role: Literal["system", "user", "assistant", "developer"] + content: Union[str, list[ResponsesContentPart]] + + +class ResponsesRequest(BaseModel): + """OpenAI Responses API request.""" + + model: str = Field("default", description = "Model identifier") + input: Union[str, list[ResponsesInputMessage]] = Field( + default = [], + description = "Input text or message list", + ) + instructions: Optional[str] = Field( + None, description = "System / developer instructions" + ) + temperature: Optional[float] = Field(None, ge = 0.0, le = 2.0) + top_p: Optional[float] = Field(None, ge = 0.0, le = 1.0) + max_output_tokens: Optional[int] = Field(None, ge = 1) + stream: bool = Field(False, description = "Whether to stream the response via SSE") + + # Accepted but ignored -- keeps SDK clients from failing on unsupported fields + tools: Optional[list] = None + tool_choice: Optional[Any] = None + previous_response_id: Optional[str] = None + store: Optional[bool] = None + metadata: Optional[dict] = None + truncation: Optional[Any] = None + user: Optional[str] = None + text: Optional[Any] = None + reasoning: Optional[Any] = None + + model_config = {"extra": "allow"} + + +# ── Response models ───────────────────────────────────────────── + + +class ResponsesOutputTextContent(BaseModel): + """A text content block inside an output message.""" + + type: Literal["output_text"] = "output_text" + text: str + annotations: list = Field(default_factory = list) + + +class ResponsesOutputMessage(BaseModel): + """An output message in the Responses API response.""" + + type: Literal["message"] = "message" + id: str = Field(default_factory = lambda: f"msg_{uuid.uuid4().hex[:12]}") + status: Literal["completed", "in_progress"] = "completed" + role: Literal["assistant"] = "assistant" + content: list[ResponsesOutputTextContent] = Field(default_factory = list) + + +class ResponsesUsage(BaseModel): + """Token usage for a Responses API response (input_tokens, not prompt_tokens).""" + + input_tokens: int = 0 + output_tokens: int = 0 + total_tokens: int = 0 + + +class ResponsesResponse(BaseModel): + """Top-level Responses API response object.""" + + id: str = Field(default_factory = lambda: f"resp_{uuid.uuid4().hex[:12]}") + object: Literal["response"] = "response" + created_at: int = Field(default_factory = lambda: int(time.time())) + status: Literal["completed", "in_progress", "failed"] = "completed" + model: str = "default" + output: list[ResponsesOutputMessage] = Field(default_factory = list) + usage: ResponsesUsage = Field(default_factory = ResponsesUsage) + error: Optional[Any] = None + incomplete_details: Optional[Any] = None + instructions: Optional[str] = None + metadata: dict = Field(default_factory = dict) + temperature: Optional[float] = None + top_p: Optional[float] = None + max_output_tokens: Optional[int] = None + previous_response_id: Optional[str] = None + text: Optional[Any] = None + tool_choice: Optional[Any] = None + tools: list = Field(default_factory = list) + truncation: Optional[Any] = None + + +# ===================================================================== +# Anthropic Messages API Models (/v1/messages) +# ===================================================================== + + +# ── Request models ───────────────────────────────────────────── + + +class AnthropicTextBlock(BaseModel): + type: Literal["text"] + text: str + + +class AnthropicImageSource(BaseModel): + type: Literal["base64", "url"] + media_type: Optional[str] = None + data: Optional[str] = None + url: Optional[str] = None + + +class AnthropicImageBlock(BaseModel): + type: Literal["image"] + source: AnthropicImageSource + + +class AnthropicToolUseBlock(BaseModel): + type: Literal["tool_use"] + id: str + name: str + input: dict + + +class AnthropicToolResultBlock(BaseModel): + type: Literal["tool_result"] + tool_use_id: str + content: Union[str, list] = "" + + +AnthropicContentBlock = Union[ + AnthropicTextBlock, + AnthropicImageBlock, + AnthropicToolUseBlock, + AnthropicToolResultBlock, +] + + +class AnthropicMessage(BaseModel): + role: Literal["user", "assistant"] + content: Union[str, list[AnthropicContentBlock]] + + +class AnthropicTool(BaseModel): + name: str + description: Optional[str] = None + input_schema: dict + + +class AnthropicMessagesRequest(BaseModel): + model: str = "default" + max_tokens: Optional[int] = None + messages: list[AnthropicMessage] + system: Optional[Union[str, list]] = None + tools: Optional[list[AnthropicTool]] = None + tool_choice: Optional[Any] = None + stream: bool = False + temperature: Optional[float] = None + top_p: Optional[float] = None + top_k: Optional[int] = None + stop_sequences: Optional[list[str]] = None + metadata: Optional[dict] = None + # [x-unsloth] extensions — mirror the OpenAI endpoint convenience fields + min_p: Optional[float] = Field( + None, ge = 0.0, le = 1.0, description = "[x-unsloth] Min-p sampling threshold" + ) + repetition_penalty: Optional[float] = Field( + None, ge = 1.0, le = 2.0, description = "[x-unsloth] Repetition penalty" + ) + presence_penalty: Optional[float] = Field( + None, ge = 0.0, le = 2.0, description = "[x-unsloth] Presence penalty" + ) + enable_tools: Optional[bool] = None + enabled_tools: Optional[list[str]] = None + session_id: Optional[str] = None + model_config = {"extra": "allow"} + + +# ── Response models ──────────────────────────────────────────── + + +class AnthropicUsage(BaseModel): + input_tokens: int = 0 + output_tokens: int = 0 + + +class AnthropicResponseTextBlock(BaseModel): + type: Literal["text"] = "text" + text: str + + +class AnthropicResponseToolUseBlock(BaseModel): + type: Literal["tool_use"] = "tool_use" + id: str + name: str + input: dict + + +AnthropicResponseBlock = Union[ + AnthropicResponseTextBlock, AnthropicResponseToolUseBlock +] + + +class AnthropicMessagesResponse(BaseModel): + id: str = Field(default_factory = lambda: f"msg_{uuid.uuid4().hex[:24]}") + type: Literal["message"] = "message" + role: Literal["assistant"] = "assistant" + content: list[AnthropicResponseBlock] = Field(default_factory = list) + model: str = "default" + stop_reason: Optional[str] = None + stop_sequence: Optional[str] = None + usage: AnthropicUsage = Field(default_factory = AnthropicUsage) diff --git a/studio/backend/models/models.py b/studio/backend/models/models.py index f67014a17b..46ca4e3784 100644 --- a/studio/backend/models/models.py +++ b/studio/backend/models/models.py @@ -213,3 +213,68 @@ class ScanFolderInfo(BaseModel): id: int = Field(..., description = "Database row ID") path: str = Field(..., description = "Normalized absolute path") created_at: str = Field(..., description = "ISO 8601 creation timestamp") + + +class BrowseEntry(BaseModel): + """A directory entry surfaced by the folder browser.""" + + name: str = Field(..., description = "Entry name (basename, not full path)") + has_models: bool = Field( + False, + description = ( + "Hint that the directory likely contains models " + "(*.gguf, *.safetensors, config.json, or HF-style " + "`models--*` subfolders). Used by the UI to highlight " + "promising candidates; the scanner itself is authoritative." + ), + ) + hidden: bool = Field( + False, + description = "Name starts with a dot (e.g. `.cache`)", + ) + + +class BrowseFoldersResponse(BaseModel): + """Response schema for the folder browser endpoint.""" + + current: str = Field(..., description = "Absolute path of the directory just listed") + parent: Optional[str] = Field( + None, + description = ( + "Parent directory of `current`, or null if `current` is the " + "filesystem root. The frontend uses this to render an `Up` row." + ), + ) + entries: List[BrowseEntry] = Field( + default_factory = list, + description = ( + "Subdirectories of `current`. Sorted with model-bearing " + "directories first, then alphabetically case-insensitive; " + "hidden entries come last within each group." + ), + ) + suggestions: List[str] = Field( + default_factory = list, + description = ( + "Handy starting points (home, HF cache, already-registered " + "scan folders). Rendered as quick-pick chips above the list." + ), + ) + truncated: bool = Field( + False, + description = ( + "True when the listing was capped because the directory had " + "more subfolders than the server is willing to enumerate in " + "one request. The UI should show a hint telling the user to " + "narrow their path." + ), + ) + model_files_here: int = Field( + 0, + description = ( + "Count of GGUF/safetensors files immediately inside " + "``current``. Used by the UI to surface a hint on leaf " + "model directories (which otherwise look `empty` because " + "they contain only files, no subdirectories)." + ), + ) diff --git a/studio/backend/requirements/extras-no-deps.txt b/studio/backend/requirements/extras-no-deps.txt index 9934bacd24..a16571567d 100644 --- a/studio/backend/requirements/extras-no-deps.txt +++ b/studio/backend/requirements/extras-no-deps.txt @@ -5,6 +5,10 @@ julius torchcodec snac +# peft 0.19.0 causes export subprocess shutdown issues in Studio; +# installing with --no-deps to avoid pulling in torch>=0.11.0 +peft==0.18.1 + # TRL and related packages trl==0.23.1 git+https://github.com/meta-pytorch/OpenEnv.git @@ -13,4 +17,4 @@ torch-c-dlpack-ext sentence_transformers==5.2.0 transformers==4.57.6 pytorch_tokenizers -kernels +kernels==0.12.1 diff --git a/studio/backend/routes/auth.py b/studio/backend/routes/auth.py index db37ed837d..5cd23bd450 100644 --- a/studio/backend/routes/auth.py +++ b/studio/backend/routes/auth.py @@ -7,11 +7,17 @@ Authentication API routes from fastapi import APIRouter, Depends, HTTPException, status +from datetime import datetime, timedelta, timezone + from models.auth import ( + ApiKeyListResponse, + ApiKeyResponse, AuthLoginRequest, - RefreshTokenRequest, AuthStatusResponse, ChangePasswordRequest, + CreateApiKeyRequest, + CreateApiKeyResponse, + RefreshTokenRequest, ) from models.users import Token from auth import storage, hashing @@ -131,3 +137,68 @@ async def change_password( token_type = "bearer", must_change_password = False, ) + + +# --------------------------------------------------------------------------- +# API key management +# --------------------------------------------------------------------------- + + +def _row_to_api_key_response(row: dict) -> ApiKeyResponse: + return ApiKeyResponse( + id = row["id"], + name = row["name"], + key_prefix = row["key_prefix"], + created_at = row["created_at"], + last_used_at = row.get("last_used_at"), + expires_at = row.get("expires_at"), + is_active = bool(row["is_active"]), + ) + + +@router.post("/api-keys", response_model = CreateApiKeyResponse) +async def create_api_key( + payload: CreateApiKeyRequest, + current_subject: str = Depends(get_current_subject), +) -> CreateApiKeyResponse: + """Create a new API key. The raw key is returned once and cannot be retrieved later.""" + expires_at = None + if payload.expires_in_days is not None: + expires_at = ( + datetime.now(timezone.utc) + timedelta(days = payload.expires_in_days) + ).isoformat() + + raw_key, row = storage.create_api_key( + username = current_subject, + name = payload.name, + expires_at = expires_at, + ) + return CreateApiKeyResponse( + key = raw_key, + api_key = _row_to_api_key_response(row), + ) + + +@router.get("/api-keys", response_model = ApiKeyListResponse) +async def list_api_keys( + current_subject: str = Depends(get_current_subject), +) -> ApiKeyListResponse: + """List all API keys for the authenticated user (raw keys are never exposed).""" + rows = storage.list_api_keys(current_subject) + return ApiKeyListResponse( + api_keys = [_row_to_api_key_response(r) for r in rows], + ) + + +@router.delete("/api-keys/{key_id}") +async def revoke_api_key( + key_id: int, + current_subject: str = Depends(get_current_subject), +) -> dict: + """Revoke (soft-delete) an API key.""" + if not storage.revoke_api_key(current_subject, key_id): + raise HTTPException( + status_code = status.HTTP_404_NOT_FOUND, + detail = "API key not found", + ) + return {"detail": "API key revoked"} diff --git a/studio/backend/routes/data_recipe/jobs.py b/studio/backend/routes/data_recipe/jobs.py index 1d5eceee03..00546b47a4 100644 --- a/studio/backend/routes/data_recipe/jobs.py +++ b/studio/backend/routes/data_recipe/jobs.py @@ -5,7 +5,9 @@ from __future__ import annotations +from datetime import timedelta from typing import Any +from urllib.parse import urlparse from fastapi import APIRouter, HTTPException, Query, Request from fastapi.responses import JSONResponse, StreamingResponse @@ -26,6 +28,161 @@ from models.data_recipe import ( router = APIRouter() +def _resolve_local_v1_endpoint(request: Request) -> str: + """Return the loopback /v1 URL for the actual backend listen port. + + Resolution order: + 1. ``app.state.server_port`` - explicitly published by run.py after + the uvicorn server has bound. This is the most reliable source + because it survives reverse proxies, TLS terminators and tunnels. + 2. ``request.scope["server"]`` - the real (host, port) tuple uvicorn + sets when the request is dispatched. Used when Studio is started + outside ``run_server`` (e.g. ``uvicorn studio.backend.main:app``). + 3. ``request.base_url`` parsed - last resort for test fixtures that + do not route through a live uvicorn server. + """ + port: Any = getattr(request.app.state, "server_port", None) + if not isinstance(port, int) or port <= 0: + server = request.scope.get("server") + if ( + isinstance(server, tuple) + and len(server) >= 2 + and isinstance(server[1], int) + and server[1] > 0 + ): + port = server[1] + else: + parsed = urlparse(str(request.base_url)) + port = parsed.port if parsed.port is not None else 8888 + return f"http://127.0.0.1:{int(port)}/v1" + + +def _used_llm_model_aliases(recipe: dict[str, Any]) -> set[str]: + """Return the set of model_aliases that are actually referenced by an + LLM column. Used to narrow the "Chat model loaded" gate so that orphan + model_config nodes on the canvas do not block unrelated recipe runs. + + The ``llm-`` prefix matches the existing convention in + ``core/data_recipe/service.py::_recipe_has_llm_columns`` and covers all + LLM column types emitted by the frontend (llm-text, llm-code, + llm-structured, llm-judge). + """ + aliases: set[str] = set() + for column in recipe.get("columns", []): + if not isinstance(column, dict): + continue + column_type = column.get("column_type") + if not isinstance(column_type, str) or not column_type.startswith("llm-"): + continue + alias = column.get("model_alias") + if isinstance(alias, str) and alias: + aliases.add(alias) + return aliases + + +def _inject_local_providers(recipe: dict[str, Any], request: Request) -> None: + """ + Mutate recipe dict in-place: for any provider with is_local=True, + generate a JWT and fill in the endpoint pointing at this server. + """ + providers = recipe.get("model_providers") + if not providers: + return + + # Collect local providers and pop is_local from ALL dicts unconditionally. + # Strict `is True` guard so malformed payloads (is_local: 1, + # is_local: "true") do not accidentally trigger the loopback rewrite. + local_indices: list[int] = [] + for i, provider in enumerate(providers): + if not isinstance(provider, dict): + continue + is_local = provider.pop("is_local", None) + if is_local is True: + local_indices.append(i) + + if not local_indices: + return + + endpoint = _resolve_local_v1_endpoint(request) + + # Only gate on model-loaded if a local provider is actually reachable + # from an LLM column through a model_config. Orphan model_config nodes + # that reference a local provider but that no LLM column uses should + # not block runs; the recipe would never call /v1 for them. + local_names = { + providers[i].get("name") for i in local_indices if providers[i].get("name") + } + used_aliases = _used_llm_model_aliases(recipe) + referenced_providers = { + mc.get("provider") + for mc in recipe.get("model_configs", []) + if ( + isinstance(mc, dict) + and mc.get("provider") + and mc.get("alias") in used_aliases + ) + } + + token = "" + if local_names & referenced_providers: + # Verify a model is loaded. + # NOTE: This is a point-in-time check (TOCTOU). The model could be unloaded + # or swapped after this check but before the recipe subprocess calls /v1. + # The inference endpoint returns a clear 400 in that case. + # + # Imports are deferred to avoid circular dependencies with inference modules. + from routes.inference import get_llama_cpp_backend + from core.inference import get_inference_backend + + llama = get_llama_cpp_backend() + model_loaded = llama.is_loaded + if not model_loaded: + backend = get_inference_backend() + model_loaded = bool(backend.active_model_name) + if not model_loaded: + raise ValueError( + "No model loaded in Chat. Load a model first, then run the recipe." + ) + + from auth.authentication import ( + create_access_token, + ) # deferred: avoids circular import + + # Uses the "unsloth" admin subject. If the user changes their password, + # the JWT secret rotates and this token becomes invalid mid-run. + # Acceptable for v1 - recipes typically finish well within one session. + token = create_access_token( + subject = "unsloth", + expires_delta = timedelta(hours = 24), + ) + + # Defensively strip any stale "external"-only fields the frontend may + # have left on the dict (extra_headers/extra_body/api_key_env). The UI + # hides these inputs in local mode but the payload builder still serializes + # them, so a previously external provider that flipped to local can carry + # invalid JSON or rogue auth headers into the local /v1 call. + for i in local_indices: + providers[i]["endpoint"] = endpoint + providers[i]["api_key"] = token + providers[i]["provider_type"] = "openai" + providers[i].pop("api_key_env", None) + providers[i].pop("extra_headers", None) + providers[i].pop("extra_body", None) + + # Force skip_health_check on any model_config that references a local + # provider. The local /v1/models endpoint only lists the real loaded + # model (e.g. "unsloth/llama-3.2-1b") and not the placeholder "local" + # that the recipe sends as the model id, so data_designer's pre-flight + # health check would otherwise fail before the first completion call. + # The backend route ignores the model id field in chat completions, so + # skipping the check is safe. + for mc in recipe.get("model_configs", []): + if not isinstance(mc, dict): + continue + if mc.get("provider") in local_names: + mc["skip_health_check"] = True + + def _normalize_run_name(value: Any) -> str | None: if value is None: return None @@ -40,7 +197,7 @@ def _normalize_run_name(value: Any) -> str | None: @router.post("/jobs", response_class = JSONResponse, response_model = JobCreateResponse) -def create_job(payload: RecipePayload): +def create_job(payload: RecipePayload, request: Request): recipe = payload.recipe if not recipe.get("columns"): raise HTTPException(status_code = 400, detail = "Recipe must include columns.") @@ -67,6 +224,11 @@ def create_job(payload: RecipePayload): status_code = 400, detail = f"invalid run_config: {exc}" ) from exc + try: + _inject_local_providers(recipe, request) + except ValueError as exc: + raise HTTPException(status_code = 400, detail = str(exc)) from exc + mgr = get_job_manager() try: job_id = mgr.start(recipe = recipe, run = run) diff --git a/studio/backend/routes/data_recipe/validate.py b/studio/backend/routes/data_recipe/validate.py index a793a3b172..555e3eaa06 100644 --- a/studio/backend/routes/data_recipe/validate.py +++ b/studio/backend/routes/data_recipe/validate.py @@ -68,6 +68,20 @@ def _collect_validation_errors(recipe: dict[str, Any]) -> list[ValidateError]: return errors +def _patch_local_providers(recipe: dict[str, Any]) -> None: + """Strip is_local and fill a dummy endpoint so validation doesn't choke. + + Uses a strict `is True` check to match _inject_local_providers in + jobs.py - malformed payloads with truthy but non-boolean is_local + values should not be treated as local. + """ + for provider in recipe.get("model_providers", []): + if not isinstance(provider, dict): + continue + if provider.pop("is_local", None) is True: + provider["endpoint"] = "http://127.0.0.1" + + @router.post("/validate", response_model = ValidateResponse) def validate(payload: RecipePayload) -> ValidateResponse: recipe = payload.recipe @@ -77,6 +91,8 @@ def validate(payload: RecipePayload) -> ValidateResponse: errors = [ValidateError(message = "Recipe must include columns.")], ) + _patch_local_providers(recipe) + try: validate_recipe(recipe) except RuntimeError as exc: diff --git a/studio/backend/routes/datasets.py b/studio/backend/routes/datasets.py index 8333009626..206af2a66f 100644 --- a/studio/backend/routes/datasets.py +++ b/studio/backend/routes/datasets.py @@ -11,10 +11,55 @@ import json import sys from pathlib import Path from uuid import uuid4 -from fastapi import APIRouter, Depends, HTTPException, UploadFile +from typing import Optional +from fastapi import APIRouter, Depends, HTTPException, Query, UploadFile +import re as _re import structlog from loggers import get_logger +_VALID_REPO_ID = _re.compile(r"^[A-Za-z0-9._-]+/[A-Za-z0-9._-]+$") + + +def _is_valid_repo_id(repo_id: str) -> bool: + return bool(_VALID_REPO_ID.fullmatch(repo_id)) + + +_dataset_size_cache: dict[str, int] = {} + + +def _get_dataset_size_cached(repo_id: str) -> int: + if repo_id in _dataset_size_cache: + return _dataset_size_cache[repo_id] + try: + from huggingface_hub import dataset_info as hf_dataset_info + + info = hf_dataset_info(repo_id, token = None, files_metadata = True) + total = sum(s.size for s in info.siblings if getattr(s, "size", None)) + _dataset_size_cache[repo_id] = total + return total + except Exception: + return 0 + + +def _resolve_hf_cache_realpath(repo_dir: Path) -> Optional[str]: + """Pick the most useful on-disk path for a HF cache repo dir. + + Mirrors the helper in routes/models.py: prefer the most-recent + snapshot dir, fall back to the cache repo root, return resolved + realpath. Duplicated here to keep routes/datasets.py self-contained. + """ + try: + snapshots_dir = repo_dir / "snapshots" + if snapshots_dir.is_dir(): + snaps = [s for s in snapshots_dir.iterdir() if s.is_dir()] + if snaps: + latest = max(snaps, key = lambda s: s.stat().st_mtime) + return str(latest.resolve()) + return str(repo_dir.resolve()) + except Exception: + return None + + # Add backend directory to path backend_path = Path(__file__).parent.parent.parent if str(backend_path) not in sys.path: @@ -308,6 +353,89 @@ def list_local_datasets( return LocalDatasetsResponse(datasets = _build_local_dataset_items()) +@router.get("/download-progress") +async def get_dataset_download_progress( + repo_id: str = Query( + ..., description = "HuggingFace dataset repo ID, e.g. 'unsloth/LaTeX_OCR'" + ), + current_subject: str = Depends(get_current_subject), +): + """Return download progress for a HuggingFace dataset repo. + + Mirrors ``GET /api/models/download-progress`` but scans the + ``datasets--owner--name`` cache directory under HF_HUB_CACHE. + Modern ``datasets``/``huggingface_hub`` caches both raw model and + raw dataset blobs in HF_HUB_CACHE; the ``datasets`` library writes + its processed Arrow shards elsewhere, but the in-progress *download* + bytes are observable here. Returns ``cache_path`` so the UI can + show users where the dataset blobs landed on disk. + """ + _empty = { + "downloaded_bytes": 0, + "expected_bytes": 0, + "progress": 0, + "cache_path": None, + } + try: + if not _is_valid_repo_id(repo_id): + return _empty + + from huggingface_hub import constants as hf_constants + + cache_dir = Path(hf_constants.HF_HUB_CACHE) + target = f"datasets--{repo_id.replace('/', '--')}".lower() + completed_bytes = 0 + in_progress_bytes = 0 + cache_path: Optional[str] = None + + if cache_dir.is_dir(): + for entry in cache_dir.iterdir(): + if entry.name.lower() != target: + continue + cache_path = _resolve_hf_cache_realpath(entry) + blobs_dir = entry / "blobs" + if not blobs_dir.is_dir(): + break + for f in blobs_dir.iterdir(): + if not f.is_file(): + continue + if f.name.endswith(".incomplete"): + in_progress_bytes += f.stat().st_size + else: + completed_bytes += f.stat().st_size + break + + downloaded_bytes = completed_bytes + in_progress_bytes + if downloaded_bytes == 0: + return {**_empty, "cache_path": cache_path} + + expected_bytes = _get_dataset_size_cached(repo_id) + if expected_bytes <= 0: + return { + "downloaded_bytes": downloaded_bytes, + "expected_bytes": 0, + "progress": 0, + "cache_path": cache_path, + } + + # Same 95% completion threshold as the model endpoint -- HF blob + # dedup makes completed_bytes drift slightly under expected_bytes, + # and inter-file gaps would otherwise look like "done". + if completed_bytes >= expected_bytes * 0.95: + progress = 1.0 + else: + progress = min(downloaded_bytes / expected_bytes, 0.99) + return { + "downloaded_bytes": downloaded_bytes, + "expected_bytes": expected_bytes, + "progress": round(progress, 3), + "cache_path": cache_path, + } + except Exception as e: + logger.warning(f"Error checking dataset download progress for {repo_id}: {e}") + return _empty + + @router.post("/check-format", response_model = CheckFormatResponse) def check_format( request: CheckFormatRequest, diff --git a/studio/backend/routes/export.py b/studio/backend/routes/export.py index 3e60eaaf20..798859fc87 100644 --- a/studio/backend/routes/export.py +++ b/studio/backend/routes/export.py @@ -5,9 +5,15 @@ Export API routes: checkpoint discovery and model export operations. """ +import asyncio +import json import sys +import time from pathlib import Path -from fastapi import APIRouter, Depends, HTTPException, Query +from typing import Any, AsyncGenerator, Dict, List, Optional, Tuple + +from fastapi import APIRouter, Depends, HTTPException, Query, Request +from fastapi.responses import StreamingResponse import structlog from loggers import get_logger @@ -97,7 +103,11 @@ async def load_checkpoint( logger.warning("Could not stop training: %s", e) backend = get_export_backend() - success, message = backend.load_checkpoint( + # load_checkpoint spawns and waits on a subprocess and can take + # minutes. Run it in a worker thread so the event loop stays + # free to serve the live log SSE stream concurrently. + success, message = await asyncio.to_thread( + backend.load_checkpoint, checkpoint_path = request.checkpoint_path, max_seq_length = request.max_seq_length, load_in_4bit = request.load_in_4bit, @@ -129,7 +139,7 @@ async def cleanup_export_memory( """ try: backend = get_export_backend() - success = backend.cleanup_memory() + success = await asyncio.to_thread(backend.cleanup_memory) if not success: raise HTTPException( @@ -173,6 +183,17 @@ async def get_export_status( ) +def _export_details(output_path: Optional[str]) -> Optional[Dict[str, Any]]: + """Wrap the resolved on-disk export path into the details dict the + frontend reads to populate the Export Complete screen. Returns None + when the export had no local component (Hub-only push) so the + Pydantic field stays absent rather than ``{"output_path": null}``. + """ + if not output_path: + return None + return {"output_path": output_path} + + @router.post("/export/merged", response_model = ExportOperationResponse) async def export_merged_model( request: ExportMergedModelRequest, @@ -185,7 +206,8 @@ async def export_merged_model( """ try: backend = get_export_backend() - success, message = backend.export_merged_model( + success, message, output_path = await asyncio.to_thread( + backend.export_merged_model, save_directory = request.save_directory, format_type = request.format_type, push_to_hub = request.push_to_hub, @@ -197,7 +219,11 @@ async def export_merged_model( if not success: raise HTTPException(status_code = 400, detail = message) - return ExportOperationResponse(success = True, message = message) + return ExportOperationResponse( + success = True, + message = message, + details = _export_details(output_path), + ) except HTTPException: raise except Exception as e: @@ -220,7 +246,8 @@ async def export_base_model( """ try: backend = get_export_backend() - success, message = backend.export_base_model( + success, message, output_path = await asyncio.to_thread( + backend.export_base_model, save_directory = request.save_directory, push_to_hub = request.push_to_hub, repo_id = request.repo_id, @@ -232,7 +259,11 @@ async def export_base_model( if not success: raise HTTPException(status_code = 400, detail = message) - return ExportOperationResponse(success = True, message = message) + return ExportOperationResponse( + success = True, + message = message, + details = _export_details(output_path), + ) except HTTPException: raise except Exception as e: @@ -255,7 +286,8 @@ async def export_gguf( """ try: backend = get_export_backend() - success, message = backend.export_gguf( + success, message, output_path = await asyncio.to_thread( + backend.export_gguf, save_directory = request.save_directory, quantization_method = request.quantization_method, push_to_hub = request.push_to_hub, @@ -266,7 +298,11 @@ async def export_gguf( if not success: raise HTTPException(status_code = 400, detail = message) - return ExportOperationResponse(success = True, message = message) + return ExportOperationResponse( + success = True, + message = message, + details = _export_details(output_path), + ) except HTTPException: raise except Exception as e: @@ -289,7 +325,8 @@ async def export_lora_adapter( """ try: backend = get_export_backend() - success, message = backend.export_lora_adapter( + success, message, output_path = await asyncio.to_thread( + backend.export_lora_adapter, save_directory = request.save_directory, push_to_hub = request.push_to_hub, repo_id = request.repo_id, @@ -300,7 +337,11 @@ async def export_lora_adapter( if not success: raise HTTPException(status_code = 400, detail = message) - return ExportOperationResponse(success = True, message = message) + return ExportOperationResponse( + success = True, + message = message, + details = _export_details(output_path), + ) except HTTPException: raise except Exception as e: @@ -309,3 +350,155 @@ async def export_lora_adapter( status_code = 500, detail = f"Failed to export LoRA adapter: {str(e)}", ) + + +# ───────────────────────────────────────────────────────────────────── +# Live export log stream (Server-Sent Events) +# ───────────────────────────────────────────────────────────────────── +# +# The export worker subprocess redirects its stdout/stderr into a pipe +# that a reader thread forwards to the orchestrator as log entries (see +# core/export/worker.py::_setup_log_capture and +# core/export/orchestrator.py::_append_log). This endpoint streams +# those entries to the browser so the export dialog can show a live +# terminal-style output panel while load_checkpoint / export_merged / +# export_gguf / export_lora / export_base run. +# +# Shape follows the training progress SSE endpoint +# (routes/training.py::stream_training_progress): each event carries +# `id`, `event`, and `data` fields, the stream starts with a `retry:` +# directive, and `Last-Event-ID` is honored on reconnect. + + +def _format_sse(data: str, event: str, event_id: Optional[int] = None) -> str: + """Format a single SSE message with id/event/data fields.""" + lines = [] + if event_id is not None: + lines.append(f"id: {event_id}") + lines.append(f"event: {event}") + lines.append(f"data: {data}") + lines.append("") + lines.append("") + return "\n".join(lines) + + +@router.get("/logs/stream") +async def stream_export_logs( + request: Request, + since: Optional[int] = Query( + None, + description = "Return log entries with seq strictly greater than this cursor.", + ), + current_subject: str = Depends(get_current_subject), +): + """ + Stream live stdout/stderr output from the export worker subprocess + as Server-Sent Events. + + Events: + - `log` : a single log line (data: {"stream","line","ts"}) + - `heartbeat`: periodic keepalive when no new lines are available + - `complete` : emitted once the export worker is idle and no new + lines arrived for ~1 second. Clients should close. + - `error` : unrecoverable server-side error + + The `id:` field on each event is the log entry's monotonic seq + number so the browser can resume via `Last-Event-ID` on reconnect. + """ + backend = get_export_backend() + + # Determine starting cursor. Explicit `since` wins, then + # Last-Event-ID header on reconnect, otherwise start from the + # run-start snapshot captured by clear_logs() so the client sees + # every line emitted since the current run began -- even if the + # SSE connection opened after the POST that kicked off the export. + # Using get_current_log_seq() here would lose the early bootstrap + # lines that arrive in the gap between POST and SSE connect. + last_event_id = request.headers.get("last-event-id") + if since is None and last_event_id is not None: + try: + since = int(last_event_id) + except ValueError: + pass + + if since is None: + cursor = backend.get_run_start_seq() + else: + cursor = max(0, int(since)) + + async def event_generator() -> AsyncGenerator[str, None]: + nonlocal cursor + # Tell the browser to reconnect after 3 seconds if the + # connection drops mid-export. + yield "retry: 3000\n\n" + + last_yield = time.monotonic() + idle_since: Optional[float] = None + try: + while True: + if await request.is_disconnected(): + return + + entries, new_cursor = backend.get_logs_since(cursor) + if entries: + for entry in entries: + payload = json.dumps( + { + "stream": entry.get("stream", "stdout"), + "line": entry.get("line", ""), + "ts": entry.get("ts"), + } + ) + yield _format_sse( + payload, + event = "log", + event_id = int(entry.get("seq", 0)), + ) + cursor = new_cursor + last_yield = time.monotonic() + idle_since = None + else: + now = time.monotonic() + if now - last_yield > 10.0: + yield _format_sse("{}", event = "heartbeat") + last_yield = now + if not backend.is_export_active(): + # Give the reader thread a moment to drain any + # trailing lines the worker process printed + # just before signalling done. + if idle_since is None: + idle_since = now + elif now - idle_since > 1.0: + yield _format_sse( + "{}", + event = "complete", + event_id = cursor, + ) + return + else: + idle_since = None + + await asyncio.sleep(0.1) + except asyncio.CancelledError: + # Client disconnected mid-yield. Don't re-raise, just end + # the generator cleanly so StreamingResponse finalizes. + return + except Exception as exc: + logger.error("Export log stream failed: %s", exc, exc_info = True) + try: + yield _format_sse( + json.dumps({"error": str(exc)}), + event = "error", + ) + except Exception: + pass + + return StreamingResponse( + event_generator(), + media_type = "text/event-stream", + headers = { + "Cache-Control": "no-cache", + "Connection": "keep-alive", + "X-Accel-Buffering": "no", + }, + ) diff --git a/studio/backend/routes/inference.py b/studio/backend/routes/inference.py index 30ff7da49c..4246f0056b 100644 --- a/studio/backend/routes/inference.py +++ b/studio/backend/routes/inference.py @@ -11,9 +11,10 @@ import time import uuid from pathlib import Path from fastapi import APIRouter, Depends, HTTPException, Request, status -from fastapi.responses import StreamingResponse, JSONResponse +from fastapi.responses import StreamingResponse, JSONResponse, Response from typing import Optional import json +import httpx import structlog from loggers import get_logger import asyncio @@ -71,11 +72,13 @@ from models.inference import ( UnloadRequest, GenerateRequest, LoadResponse, + LoadProgressResponse, UnloadResponse, InferenceStatusResponse, ChatCompletionRequest, ChatCompletionChunk, ChatCompletion, + ChatMessage, ChunkChoice, ChoiceDelta, CompletionChoice, @@ -83,6 +86,28 @@ from models.inference import ( CompletionUsage, ValidateModelRequest, ValidateModelResponse, + TextContentPart, + ImageContentPart, + ImageUrl, + ResponsesRequest, + ResponsesInputMessage, + ResponsesInputTextPart, + ResponsesInputImagePart, + ResponsesOutputTextContent, + ResponsesOutputMessage, + ResponsesUsage, + ResponsesResponse, + AnthropicMessagesRequest, + AnthropicMessagesResponse, + AnthropicResponseTextBlock, + AnthropicResponseToolUseBlock, + AnthropicUsage, +) +from core.inference.anthropic_compat import ( + anthropic_messages_to_openai, + anthropic_tools_to_openai, + AnthropicStreamEmitter, + AnthropicPassthroughEmitter, ) from auth.authentication import get_current_subject @@ -121,6 +146,7 @@ def get_llama_cpp_backend() -> LlamaCppBackend: @router.post("/load", response_model = LoadResponse) async def load_model( request: LoadRequest, + fastapi_request: Request, current_subject: str = Depends(get_current_subject), ): """ @@ -173,6 +199,9 @@ async def load_model( if _gguf_audio else False, inference = inference_config, + requires_trust_remote_code = bool( + inference_config.get("trust_remote_code", False) + ), context_length = llama_backend.context_length, max_context_length = llama_backend.max_context_length, native_context_length = llama_backend.native_context_length, @@ -210,6 +239,9 @@ async def load_model( audio_type = _model_info.get("audio_type"), has_audio_input = _model_info.get("has_audio_input", False), inference = inference_config, + requires_trust_remote_code = bool( + inference_config.get("trust_remote_code", False) + ), chat_template = _chat_template, ) @@ -252,6 +284,8 @@ async def load_model( # Run in a thread so the event loop stays free for progress # polling and other requests during the (potentially long) # GGUF download + llama-server startup. + _n_parallel = getattr(fastapi_request.app.state, "llama_parallel_slots", 1) + if config.gguf_hf_repo: # HF mode: download via huggingface_hub then start llama-server success = await asyncio.to_thread( @@ -265,6 +299,7 @@ async def load_model( chat_template_override = request.chat_template_override, cache_type_kv = request.cache_type_kv, speculative_type = request.speculative_type, + n_parallel = _n_parallel, ) else: # Local mode: llama-server loads via -m @@ -278,6 +313,7 @@ async def load_model( chat_template_override = request.chat_template_override, cache_type_kv = request.cache_type_kv, speculative_type = request.speculative_type, + n_parallel = _n_parallel, ) if not success: @@ -312,6 +348,9 @@ async def load_model( audio_type = _gguf_audio, has_audio_input = is_audio_input_type(_gguf_audio), inference = inference_config, + requires_trust_remote_code = bool( + inference_config.get("trust_remote_code", False) + ), context_length = llama_backend.context_length, max_context_length = llama_backend.max_context_length, native_context_length = llama_backend.native_context_length, @@ -450,6 +489,9 @@ async def load_model( audio_type = config.audio_type, has_audio_input = config.has_audio_input, inference = inference_config, + requires_trust_remote_code = bool( + inference_config.get("trust_remote_code", False) + ), chat_template = _chat_template, ) @@ -505,6 +547,9 @@ async def validate_model( is_gguf = getattr(config, "is_gguf", False), is_lora = getattr(config, "is_lora", False), is_vision = getattr(config, "is_vision", False), + requires_trust_remote_code = bool( + load_inference_config(config.identifier).get("trust_remote_code", False) + ), ) except HTTPException: @@ -650,6 +695,9 @@ async def get_status( loading = [], loaded = [_model_id], inference = _inference_cfg, + requires_trust_remote_code = bool( + (_inference_cfg or {}).get("trust_remote_code", False) + ), supports_reasoning = llama_backend.supports_reasoning, reasoning_always_on = llama_backend.reasoning_always_on, supports_tools = llama_backend.supports_tools, @@ -677,6 +725,11 @@ async def get_status( supports_reasoning = False if backend.active_model_name and hasattr(backend, "_is_gpt_oss_model"): supports_reasoning = backend._is_gpt_oss_model() + inference_config = ( + load_inference_config(backend.active_model_name) + if backend.active_model_name + else None + ) return InferenceStatusResponse( active_model = backend.active_model_name, @@ -687,6 +740,10 @@ async def get_status( has_audio_input = has_audio_input, loading = list(getattr(backend, "loading_models", set())), loaded = list(backend.models.keys()), + inference = inference_config, + requires_trust_remote_code = bool( + (inference_config or {}).get("trust_remote_code", False) + ), supports_reasoning = supports_reasoning, ) @@ -695,6 +752,34 @@ async def get_status( raise HTTPException(status_code = 500, detail = f"Failed to get status: {str(e)}") +@router.get("/load-progress", response_model = LoadProgressResponse) +async def get_load_progress( + current_subject: str = Depends(get_current_subject), +): + """ + Return the active GGUF load's mmap/upload progress. + + During the warmup window after a GGUF download -- when llama-server + is paging ~tens-to-hundreds of GB of shards into the page cache + before pushing layers to VRAM -- ``/api/inference/status`` only + shows a generic spinner. This endpoint exposes sampled progress so + the UI can render a real bar plus rate/ETA during that window. + + Returns an empty payload (``phase=null, bytes=0``) when no load is + in flight. The frontend should stop polling once ``phase`` becomes + ``ready``. + """ + try: + llama_backend = get_llama_cpp_backend() + progress = llama_backend.load_progress() + if progress is None: + return LoadProgressResponse() + return LoadProgressResponse(**progress) + except Exception as e: + logger.warning(f"Error sampling load progress: {e}") + return LoadProgressResponse() + + # ===================================================================== # Audio (TTS) Generation (/audio/generate) # ===================================================================== @@ -1816,3 +1901,1071 @@ async def openai_list_models( ) return {"object": "list", "data": models} + + +# ===================================================================== +# OpenAI-Compatible Completions Proxy (/completions → /v1/completions) +# ===================================================================== + + +@router.post("/completions") +async def openai_completions( + request: Request, + current_subject: str = Depends(get_current_subject), +): + """ + OpenAI-compatible text completions endpoint (non-chat). + + Transparently proxies to the running llama-server's ``/v1/completions``. + Only available when a GGUF model is loaded. + """ + llama_backend = get_llama_cpp_backend() + if not llama_backend.is_loaded: + raise HTTPException( + status_code = 503, + detail = "No GGUF model loaded. Load a GGUF model first.", + ) + + body = await request.json() + target_url = f"{llama_backend.base_url}/v1/completions" + is_stream = body.get("stream", False) + + if is_stream: + + async def _stream(): + # Manual httpx client/response lifecycle — see + # _anthropic_passthrough_stream for the full rationale. Briefly: + # `async with` inside an async generator causes + # "Attempted to exit cancel scope in a different task" / + # "async generator ignored GeneratorExit" on Python 3.13 + + # httpcore 1.0.x when the generator is orphaned and finalized + # by GC. Closing via a finally block that catches Exception + # (but not BaseException) suppresses the anyio cleanup noise + # while letting GeneratorExit propagate cleanly. + client = httpx.AsyncClient(timeout = 600) + resp = None + try: + req = client.build_request("POST", target_url, json = body) + resp = await client.send(req, stream = True) + async for chunk in resp.aiter_bytes(): + yield chunk + except Exception as e: + logger.error("openai_completions stream error: %s", e) + finally: + if resp is not None: + try: + await resp.aclose() + except Exception: + pass + try: + await client.aclose() + except Exception: + pass + + 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) + return Response( + content = resp.content, + status_code = resp.status_code, + media_type = "application/json", + ) + + +# ===================================================================== +# OpenAI-Compatible Embeddings Proxy (/embeddings → /v1/embeddings) +# ===================================================================== + + +@router.post("/embeddings") +async def openai_embeddings( + request: Request, + current_subject: str = Depends(get_current_subject), +): + """ + OpenAI-compatible embeddings endpoint. + + Transparently proxies to the running llama-server's ``/v1/embeddings``. + Only available when a GGUF model is loaded. + Note: the loaded model must support pooling; otherwise llama-server + will return an error (expected). + """ + llama_backend = get_llama_cpp_backend() + if not llama_backend.is_loaded: + raise HTTPException( + status_code = 503, + detail = "No GGUF model loaded. Load a GGUF model first.", + ) + + body = await request.json() + 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) + return Response( + content = resp.content, + status_code = resp.status_code, + media_type = "application/json", + ) + + +# ===================================================================== +# OpenAI Responses API (/responses → /v1/responses) +# ===================================================================== + + +def _normalise_responses_input(payload: ResponsesRequest) -> list[ChatMessage]: + """Convert a ResponsesRequest into a list of ChatMessage for the completions backend.""" + messages: list[ChatMessage] = [] + + # System / developer instructions + if payload.instructions: + messages.append(ChatMessage(role = "system", content = payload.instructions)) + + # Simple string input + if isinstance(payload.input, str): + if payload.input: + messages.append(ChatMessage(role = "user", content = payload.input)) + return messages + + # List of ResponsesInputMessage + for msg in payload.input: + role = "system" if msg.role == "developer" else msg.role + + if isinstance(msg.content, str): + messages.append(ChatMessage(role = role, content = msg.content)) + else: + # Convert Responses content parts -> Chat content parts + parts = [] + for part in msg.content: + if isinstance(part, ResponsesInputTextPart): + parts.append(TextContentPart(type = "text", text = part.text)) + elif isinstance(part, ResponsesInputImagePart): + parts.append( + ImageContentPart( + type = "image_url", + image_url = ImageUrl(url = part.image_url, detail = part.detail), + ) + ) + messages.append(ChatMessage(role = role, content = parts if parts else "")) + + return messages + + +def _build_chat_request( + payload: ResponsesRequest, messages: list[ChatMessage], stream: bool +) -> ChatCompletionRequest: + """Build a ChatCompletionRequest from a ResponsesRequest.""" + chat_kwargs = dict( + model = payload.model, + messages = messages, + stream = stream, + ) + if payload.temperature is not None: + chat_kwargs["temperature"] = payload.temperature + if payload.top_p is not None: + chat_kwargs["top_p"] = payload.top_p + if payload.max_output_tokens is not None: + chat_kwargs["max_tokens"] = payload.max_output_tokens + return ChatCompletionRequest(**chat_kwargs) + + +async def _responses_non_streaming( + payload: ResponsesRequest, + messages: list[ChatMessage], + request: Request, +) -> JSONResponse: + """Handle a non-streaming Responses API call.""" + chat_req = _build_chat_request(payload, messages, stream = False) + result = await openai_chat_completions(chat_req, request) + + # openai_chat_completions returns a JSONResponse for non-streaming + if isinstance(result, JSONResponse): + body = json.loads(result.body.decode()) + elif isinstance(result, Response): + body = json.loads(result.body.decode()) + else: + body = result + + # Extract content and usage from the Chat Completions response + choices = body.get("choices", []) + text = "" + if choices: + msg = choices[0].get("message", {}) + text = msg.get("content", "") or "" + + usage_data = body.get("usage", {}) + input_tokens = usage_data.get("prompt_tokens", 0) + output_tokens = usage_data.get("completion_tokens", 0) + + resp_id = f"resp_{uuid.uuid4().hex[:12]}" + msg_id = f"msg_{uuid.uuid4().hex[:12]}" + + response = ResponsesResponse( + id = resp_id, + created_at = int(time.time()), + status = "completed", + model = body.get("model", payload.model), + output = [ + ResponsesOutputMessage( + id = msg_id, + status = "completed", + role = "assistant", + content = [ + ResponsesOutputTextContent(text = text), + ], + ), + ], + usage = ResponsesUsage( + input_tokens = input_tokens, + output_tokens = output_tokens, + total_tokens = input_tokens + output_tokens, + ), + temperature = payload.temperature, + top_p = payload.top_p, + max_output_tokens = payload.max_output_tokens, + instructions = payload.instructions, + ) + return JSONResponse(content = response.model_dump()) + + +async def _responses_stream( + payload: ResponsesRequest, + messages: list[ChatMessage], + request: Request, +): + """Handle a streaming Responses API call, emitting named SSE events.""" + resp_id = f"resp_{uuid.uuid4().hex[:12]}" + msg_id = f"msg_{uuid.uuid4().hex[:12]}" + item_id = f"item_{uuid.uuid4().hex[:12]}" + created_at = int(time.time()) + + chat_req = _build_chat_request(payload, messages, stream = True) + result = await openai_chat_completions(chat_req, request) + + async def event_generator(): + full_text = "" + input_tokens = 0 + output_tokens = 0 + + # ── 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 + 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" + + # ── Stream delta events from the inner chat completions stream ── + if isinstance(result, StreamingResponse): + async for raw_chunk in result.body_iterator: + if isinstance(raw_chunk, bytes): + raw_chunk = raw_chunk.decode("utf-8", errors = "replace") + + for line in raw_chunk.split("\n"): + line = line.strip() + if not line.startswith("data: "): + continue + data_str = line[6:] + if data_str == "[DONE]": + continue + try: + chunk_data = json.loads(data_str) + except json.JSONDecodeError: + continue + + choices = chunk_data.get("choices", []) + if not choices: + # Check for usage in final chunk + usage = chunk_data.get("usage") + if usage: + input_tokens = usage.get("prompt_tokens", input_tokens) + output_tokens = usage.get( + "completion_tokens", output_tokens + ) + continue + + delta = choices[0].get("delta", {}) + 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" + + # Check for usage in chunk + usage = chunk_data.get("usage") + if usage: + input_tokens = usage.get("prompt_tokens", input_tokens) + output_tokens = usage.get("completion_tokens", output_tokens) + + # ── Closing events ── + # output_text.done + 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" + + # content_part.done + 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" + + # output_item.done + 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" + + # response.completed + total_tokens = input_tokens + output_tokens + completed_response = { + "type": "response.completed", + "response": { + "id": resp_id, + "object": "response", + "created_at": created_at, + "status": "completed", + "model": payload.model, + "output": [ + { + "type": "message", + "id": msg_id, + "status": "completed", + "role": "assistant", + "content": [ + { + "type": "output_text", + "text": full_text, + "annotations": [], + } + ], + } + ], + "usage": { + "input_tokens": input_tokens, + "output_tokens": output_tokens, + "total_tokens": total_tokens, + }, + }, + } + yield f"event: response.completed\ndata: {json.dumps(completed_response)}\n\n" + + return StreamingResponse( + event_generator(), + media_type = "text/event-stream", + headers = { + "Cache-Control": "no-cache", + "Connection": "keep-alive", + "X-Accel-Buffering": "no", + }, + ) + + +@router.post("/responses") +async def openai_responses( + payload: ResponsesRequest, + request: Request, + current_subject: str = Depends(get_current_subject), +): + """ + OpenAI Responses API endpoint. + + Accepts the Responses-format request, converts it to a + ChatCompletionRequest internally, and returns a response + matching the OpenAI Responses API schema (output array, + input_tokens/output_tokens, named SSE events for streaming). + """ + messages = _normalise_responses_input(payload) + if not messages: + raise HTTPException(status_code = 400, detail = "No input provided.") + + if payload.stream: + return await _responses_stream(payload, messages, request) + return await _responses_non_streaming(payload, messages, request) + + +# ===================================================================== +# Anthropic-Compatible Messages API (/messages → /v1/messages) +# ===================================================================== + + +@router.post("/messages") +async def anthropic_messages( + payload: AnthropicMessagesRequest, + request: Request, + current_subject: str = Depends(get_current_subject), +): + """ + Anthropic-compatible Messages API endpoint. + + Translates Anthropic message format to internal OpenAI format, runs + through the existing agentic tool loop when tools are provided, and + returns responses in Anthropic Messages API format (streaming SSE or + non-streaming JSON). + """ + llama_backend = get_llama_cpp_backend() + if not llama_backend.is_loaded: + raise HTTPException( + status_code = 503, + detail = "No GGUF model loaded. Load a GGUF model first.", + ) + + model_name = getattr(llama_backend, "model_identifier", None) or payload.model + message_id = f"msg_{uuid.uuid4().hex[:24]}" + + # ── Translate Anthropic → OpenAI ────────────────────────── + openai_messages = anthropic_messages_to_openai( + [m.model_dump() for m in payload.messages], + payload.system, + ) + + temperature = payload.temperature if payload.temperature is not None else 0.6 + top_p = payload.top_p if payload.top_p is not None else 0.95 + top_k = payload.top_k if payload.top_k is not None else 20 + min_p = payload.min_p if payload.min_p is not None else 0.01 + repetition_penalty = ( + payload.repetition_penalty if payload.repetition_penalty is not None else 1.0 + ) + presence_penalty = ( + payload.presence_penalty if payload.presence_penalty is not None else 0.0 + ) + stop = payload.stop_sequences or None + + # tool_choice is declared on AnthropicMessagesRequest for Anthropic SDK + # compatibility (the SDK often sets it by default), but it is not + # currently honored by Unsloth's backend. Warn once per request so the + # silent drop is visible to operators instead of looking like a model + # quality issue to clients. + if payload.tool_choice is not None: + logger.warning( + "anthropic_messages.tool_choice_ignored", + tool_choice = payload.tool_choice, + note = ( + "tool_choice is accepted for Anthropic SDK compatibility but not " + "honored by Unsloth. Use enable_tools / enabled_tools (server-side " + "built-in tools) or restrict the `tools` array (client-side) to " + "control which tools the model sees." + ), + ) + + cancel_event = threading.Event() + + # ── Tool routing ────────────────────────────────────────── + # Three paths: + # 1. enable_tools=true → server-side execution of built-in tools (Unsloth shorthand) + # 2. tools=[...] only → client-side pass-through (standard Anthropic behavior) + # 3. neither → plain chat + server_tools = payload.enable_tools and llama_backend.supports_tools + client_tools = ( + not server_tools + and payload.tools + and len(payload.tools) > 0 + and llama_backend.supports_tools + ) + + # ── Client-side pass-through path ───────────────────────── + if client_tools: + openai_tools = anthropic_tools_to_openai(payload.tools) + + if payload.stream: + return await _anthropic_passthrough_stream( + request, + cancel_event, + llama_backend, + openai_messages, + openai_tools, + temperature, + top_p, + top_k, + payload.max_tokens, + message_id, + model_name, + stop = stop, + min_p = min_p, + repetition_penalty = repetition_penalty, + presence_penalty = presence_penalty, + ) + return await _anthropic_passthrough_non_streaming( + llama_backend, + openai_messages, + openai_tools, + temperature, + top_p, + top_k, + payload.max_tokens, + message_id, + model_name, + stop = stop, + min_p = min_p, + repetition_penalty = repetition_penalty, + presence_penalty = presence_penalty, + ) + + if server_tools: + from core.inference.tools import ALL_TOOLS + + if payload.enabled_tools is not None: + openai_tools = [ + t for t in ALL_TOOLS if t["function"]["name"] in payload.enabled_tools + ] + else: + openai_tools = ALL_TOOLS + + # Build tool-use system prompt nudge (same logic as /chat/completions) + _tool_names = {t["function"]["name"] for t in openai_tools} + _has_web = "web_search" in _tool_names + _has_code = "python" in _tool_names or "terminal" in _tool_names + + _date_line = f"The current date is {_date.today().isoformat()}." + _model_size_b = _extract_model_size_b(model_name) + _is_small_model = _model_size_b is not None and _model_size_b < 9 + + if _is_small_model: + _web_tips = "Do not repeat the same search query." + else: + _web_tips = ( + "When you search and find a relevant URL in the results, " + "fetch its full content by calling web_search with the url parameter. " + "Do not repeat the same search query. If a search returns " + "no useful results, try rephrasing or fetching a result URL directly." + ) + _code_tips = ( + "Use code execution for math, calculations, data processing, " + "or to parse and analyze information from tool results." + ) + + if _has_web and _has_code: + _nudge = ( + _date_line + " " + "You have access to tools. When appropriate, prefer using " + "tools rather than answering from memory. " + + _web_tips + + " " + + _code_tips + ) + elif _has_code: + _nudge = ( + _date_line + " " + "You have access to tools. When appropriate, prefer using " + "code execution rather than answering from memory. " + _code_tips + ) + elif _has_web: + _nudge = ( + _date_line + " " + "You have access to tools. When appropriate, prefer using " + "web search for up-to-date or uncertain factual " + "information rather than answering from memory. " + _web_tips + ) + else: + _nudge = "" + + if _nudge: + _nudge += _TOOL_ACTION_NUDGE + # Inject into system prompt + if openai_messages and openai_messages[0].get("role") == "system": + openai_messages[0]["content"] = ( + openai_messages[0]["content"].rstrip() + "\n\n" + _nudge + ) + else: + openai_messages.insert(0, {"role": "system", "content": _nudge}) + + # Strip stale tool-call XML from conversation + for _msg in openai_messages: + if _msg.get("role") == "assistant" and isinstance(_msg.get("content"), str): + _msg["content"] = _TOOL_XML_RE.sub("", _msg["content"]).strip() + + def _run_tool_gen(): + return llama_backend.generate_chat_completion_with_tools( + messages = openai_messages, + tools = openai_tools, + temperature = temperature, + top_p = top_p, + top_k = top_k, + min_p = min_p, + repetition_penalty = repetition_penalty, + presence_penalty = presence_penalty, + max_tokens = payload.max_tokens, + stop = stop, + cancel_event = cancel_event, + max_tool_iterations = 25, + auto_heal_tool_calls = True, + tool_call_timeout = 300, + session_id = payload.session_id, + ) + + if payload.stream: + return await _anthropic_tool_stream( + request, + cancel_event, + _run_tool_gen, + message_id, + model_name, + ) + return await _anthropic_tool_non_streaming( + _run_tool_gen, + message_id, + model_name, + ) + + # ── No-tool path ────────────────────────────────────────── + def _run_plain_gen(): + return llama_backend.generate_chat_completion( + messages = openai_messages, + temperature = temperature, + top_p = top_p, + top_k = top_k, + min_p = min_p, + repetition_penalty = repetition_penalty, + presence_penalty = presence_penalty, + max_tokens = payload.max_tokens, + stop = stop, + cancel_event = cancel_event, + ) + + if payload.stream: + return await _anthropic_plain_stream( + request, + cancel_event, + _run_plain_gen, + message_id, + model_name, + ) + return await _anthropic_plain_non_streaming( + _run_plain_gen, + message_id, + model_name, + ) + + +async def _anthropic_tool_stream( + request, + cancel_event, + run_gen, + message_id, + model_name, +): + """Streaming response for the tool-calling path.""" + _sentinel = object() + + async def _stream(): + emitter = AnthropicStreamEmitter() + for line in emitter.start(message_id, model_name): + yield line + + gen = run_gen() + try: + while True: + if await request.is_disconnected(): + cancel_event.set() + return + event = await asyncio.to_thread(next, gen, _sentinel) + if event is _sentinel: + break + # Strip leaked tool-call XML from content events + if event.get("type") == "content": + event = dict(event) + event["text"] = _TOOL_XML_RE.sub("", event["text"]) + for line in emitter.feed(event): + yield line + except Exception as e: + logger.error("anthropic_messages stream error: %s", e) + + for line in emitter.finish("end_turn"): + yield line + + return StreamingResponse( + _stream(), + media_type = "text/event-stream", + headers = { + "Cache-Control": "no-cache", + "Connection": "keep-alive", + "X-Accel-Buffering": "no", + }, + ) + + +async def _anthropic_plain_stream( + request, + cancel_event, + run_gen, + message_id, + model_name, +): + """Streaming response for the no-tool path.""" + _sentinel = object() + + async def _stream(): + emitter = AnthropicStreamEmitter() + for line in emitter.start(message_id, model_name): + yield line + + gen = run_gen() + try: + while True: + if await request.is_disconnected(): + cancel_event.set() + return + cumulative = await asyncio.to_thread(next, gen, _sentinel) + if cumulative is _sentinel: + break + if isinstance(cumulative, dict): + if cumulative.get("type") == "metadata": + for line in emitter.feed(cumulative): + yield line + continue + # Plain generator yields cumulative text strings + for line in emitter.feed({"type": "content", "text": cumulative}): + yield line + except Exception as e: + logger.error("anthropic_messages stream error: %s", e) + + for line in emitter.finish("end_turn"): + yield line + + return StreamingResponse( + _stream(), + media_type = "text/event-stream", + headers = { + "Cache-Control": "no-cache", + "Connection": "keep-alive", + "X-Accel-Buffering": "no", + }, + ) + + +async def _anthropic_tool_non_streaming(run_gen, message_id, model_name): + """Non-streaming response for the tool-calling path. + + Builds ``content_blocks`` in generation order (text → tool_use → text → + tool_use → ...), mirroring the streaming emitter's behavior. Deltas + within a single synthesis turn are merged into the trailing text block; + tool_use blocks interrupt the text sequence and open a new text block on + the next content event. + + ``prev_text`` is reset on ``tool_end`` because + ``generate_chat_completion_with_tools`` yields cumulative content *per + turn* — the first content event of turn N+1 must diff against an empty + baseline, not against turn N's final length. + """ + content_blocks: list = [] + usage = {} + prev_text = "" + + for event in run_gen(): + etype = event.get("type", "") + if etype == "content": + # Strip leaked tool-call XML + clean = _TOOL_XML_RE.sub("", event["text"]) + new = clean[len(prev_text) :] + prev_text = clean + if new: + if content_blocks and isinstance( + content_blocks[-1], AnthropicResponseTextBlock + ): + content_blocks[-1].text += new + else: + content_blocks.append(AnthropicResponseTextBlock(text = new)) + elif etype == "tool_start": + content_blocks.append( + AnthropicResponseToolUseBlock( + id = event["tool_call_id"], + name = event["tool_name"], + input = event.get("arguments", {}), + ) + ) + elif etype == "tool_end": + prev_text = "" + elif etype == "metadata": + usage = event.get("usage", {}) + + resp = AnthropicMessagesResponse( + id = message_id, + model = model_name, + content = content_blocks, + stop_reason = "end_turn", + usage = AnthropicUsage( + input_tokens = usage.get("prompt_tokens", 0), + output_tokens = usage.get("completion_tokens", 0), + ), + ) + return JSONResponse(content = resp.model_dump()) + + +async def _anthropic_plain_non_streaming(run_gen, message_id, model_name): + """Non-streaming response for the no-tool path.""" + text_parts = [] + usage = {} + prev_text = "" + + for cumulative in run_gen(): + if isinstance(cumulative, dict): + if cumulative.get("type") == "metadata": + usage = cumulative.get("usage", {}) + continue + new = cumulative[len(prev_text) :] + prev_text = cumulative + if new: + text_parts.append(new) + + full_text = "".join(text_parts) + content_blocks = [] + if full_text: + content_blocks.append(AnthropicResponseTextBlock(text = full_text)) + + resp = AnthropicMessagesResponse( + id = message_id, + model = model_name, + content = content_blocks, + stop_reason = "end_turn", + usage = AnthropicUsage( + input_tokens = usage.get("prompt_tokens", 0), + output_tokens = usage.get("completion_tokens", 0), + ), + ) + return JSONResponse(content = resp.model_dump()) + + +# ===================================================================== +# Client-side tool pass-through (Anthropic-native tools field) +# ===================================================================== + + +def _build_passthrough_payload( + openai_messages, + openai_tools, + temperature, + top_p, + top_k, + max_tokens, + stream, + stop = None, + min_p = None, + repetition_penalty = None, + presence_penalty = None, +): + body = { + "messages": openai_messages, + "tools": openai_tools, + "tool_choice": "auto", + "temperature": temperature, + "top_p": top_p, + "top_k": top_k, + "stream": stream, + } + if stream: + body["stream_options"] = {"include_usage": True} + if max_tokens is not None: + body["max_tokens"] = max_tokens + if stop: + body["stop"] = stop + if min_p is not None: + body["min_p"] = min_p + if repetition_penalty is not None: + # llama-server's field is "repeat_penalty", not "repetition_penalty" + body["repeat_penalty"] = repetition_penalty + if presence_penalty is not None: + body["presence_penalty"] = presence_penalty + return body + + +async def _anthropic_passthrough_stream( + request, + cancel_event, + llama_backend, + openai_messages, + openai_tools, + temperature, + top_p, + top_k, + max_tokens, + message_id, + model_name, + stop = None, + min_p = None, + repetition_penalty = None, + presence_penalty = None, +): + """Streaming client-side pass-through: forward tools to llama-server and + translate its streaming response to Anthropic SSE without executing anything.""" + target_url = f"{llama_backend.base_url}/v1/chat/completions" + body = _build_passthrough_payload( + openai_messages, + openai_tools, + temperature, + top_p, + top_k, + max_tokens, + True, + stop = stop, + min_p = min_p, + repetition_penalty = repetition_penalty, + presence_penalty = presence_penalty, + ) + + async def _stream(): + emitter = AnthropicPassthroughEmitter() + for line in emitter.start(message_id, model_name): + yield line + + # Manage the httpx client and response MANUALLY — no `async with`. + # + # On Python 3.13 + httpcore 1.0.x, an orphaned async generator (e.g. + # when the client disconnects mid-stream and Starlette drops the + # StreamingResponse iterator without explicitly calling aclose()) + # is finalized by Python's asyncgen GC hook in a DIFFERENT asyncio + # task than the one that originally entered the httpx context + # managers. When `async with` exits run in the wrong task, httpcore's + # internal `HTTP11ConnectionByteStream.aclose()` hits + # `anyio.CancelScope.__exit__` with a mismatched task and raises + # RuntimeError("Attempted to exit cancel scope in a different task"), + # which escapes as "Exception ignored in:" because it happens during + # GC finalization outside any user-owned try/except. + # + # The fix: do not use `async with` for the client/response. Close + # them in a finally block wrapped in `try: ... except Exception: pass`. + # This narrowly suppresses RuntimeError / other Exception subclasses + # from the anyio cleanup noise while letting GeneratorExit (a + # BaseException, not Exception) propagate through cleanly so the + # generator terminates as Python expects. + client = httpx.AsyncClient(timeout = 600) + resp = None + try: + req = client.build_request("POST", target_url, json = body) + resp = await client.send(req, stream = True) + + async for raw_line in resp.aiter_lines(): + if await request.is_disconnected(): + cancel_event.set() + break + if not raw_line or not raw_line.startswith("data: "): + continue + data_str = raw_line[6:] + if data_str.strip() == "[DONE]": + break + try: + chunk = json.loads(data_str) + except json.JSONDecodeError: + continue + for line in emitter.feed_chunk(chunk): + yield line + except Exception as e: + logger.error("anthropic_messages passthrough stream error: %s", e) + finally: + if resp is not None: + try: + await resp.aclose() + except Exception: + pass + try: + await client.aclose() + except Exception: + pass + + for line in emitter.finish(): + yield line + + return StreamingResponse( + _stream(), + media_type = "text/event-stream", + headers = { + "Cache-Control": "no-cache", + "Connection": "keep-alive", + "X-Accel-Buffering": "no", + }, + ) + + +async def _anthropic_passthrough_non_streaming( + llama_backend, + openai_messages, + openai_tools, + temperature, + top_p, + top_k, + max_tokens, + message_id, + model_name, + stop = None, + min_p = None, + repetition_penalty = None, + presence_penalty = None, +): + """Non-streaming client-side pass-through.""" + target_url = f"{llama_backend.base_url}/v1/chat/completions" + body = _build_passthrough_payload( + openai_messages, + openai_tools, + temperature, + top_p, + top_k, + max_tokens, + False, + stop = stop, + min_p = min_p, + repetition_penalty = repetition_penalty, + presence_penalty = presence_penalty, + ) + + async with httpx.AsyncClient() as client: + resp = await client.post(target_url, json = body, timeout = 600) + + if resp.status_code != 200: + raise HTTPException( + status_code = resp.status_code, + detail = f"llama-server error: {resp.text[:500]}", + ) + + data = resp.json() + choice = (data.get("choices") or [{}])[0] + message = choice.get("message") or {} + finish_reason = choice.get("finish_reason") + + content_blocks = [] + text = message.get("content") or "" + if text: + text = _TOOL_XML_RE.sub("", text).strip() + if text: + content_blocks.append(AnthropicResponseTextBlock(text = text)) + + tool_calls = message.get("tool_calls") or [] + for tc in tool_calls: + fn = tc.get("function") or {} + try: + args = json.loads(fn.get("arguments", "{}")) + except json.JSONDecodeError: + args = {} + content_blocks.append( + AnthropicResponseToolUseBlock( + id = tc.get("id", ""), + name = fn.get("name", ""), + input = args, + ) + ) + + if tool_calls: + stop_reason = "tool_use" + elif finish_reason == "length": + stop_reason = "max_tokens" + else: + stop_reason = "end_turn" + + usage = data.get("usage") or {} + resp_obj = AnthropicMessagesResponse( + id = message_id, + model = model_name, + content = content_blocks, + stop_reason = stop_reason, + usage = AnthropicUsage( + input_tokens = usage.get("prompt_tokens", 0), + output_tokens = usage.get("completion_tokens", 0), + ), + ) + return JSONResponse(content = resp_obj.model_dump()) diff --git a/studio/backend/routes/models.py b/studio/backend/routes/models.py index 3f361ca5eb..9e7168eed6 100644 --- a/studio/backend/routes/models.py +++ b/studio/backend/routes/models.py @@ -32,8 +32,9 @@ from auth.authentication import get_current_subject # Import backend functions try: from utils.models import ( - scan_trained_loras, + scan_trained_models, scan_exported_models, + get_base_model_from_checkpoint, load_model_defaults, get_base_model_from_lora, is_vision_model, @@ -62,8 +63,9 @@ except ImportError: if str(parent_backend) not in sys.path: sys.path.insert(0, str(parent_backend)) from utils.models import ( - scan_trained_loras, + scan_trained_models, scan_exported_models, + get_base_model_from_checkpoint, load_model_defaults, get_base_model_from_lora, is_vision_model, @@ -99,6 +101,8 @@ from models import ( ModelListResponse, ) from models.models import ( + BrowseEntry, + BrowseFoldersResponse, GgufVariantDetail, GgufVariantsResponse, ModelType, @@ -571,6 +575,529 @@ async def remove_scan_folder_endpoint( return {"ok": True} +# Heuristic ceiling on how many children to stat when checking whether a +# directory "looks like" it contains models. Keeps the browser snappy +# even when a directory has thousands of unrelated entries. +_BROWSE_MODEL_HINT_PROBE = 64 +# Hard cap on how many subdirectory entries we send back. Pointing the +# browser at something like ``/usr/lib`` or ``/proc`` must not stat-storm +# the process or send tens of thousands of rows to the client. +_BROWSE_ENTRY_CAP = 2000 + + +def _count_model_files(directory: Path, cap: int = 200) -> int: + """Count GGUF/safetensors files immediately inside *directory*. + Used to surface a count-hint on the response so the UI can tell + users that a leaf directory (no subdirs, only weights) is a valid + "Use this folder" target. + + Bounded by *visited entries*, not by *match count*: in directories + with many non-model files (or many subdirectories) the scan still + stops after ``cap`` entries so a UI hint never costs more than a + bounded directory walk. + """ + n = 0 + visited = 0 + try: + for f in directory.iterdir(): + visited += 1 + if visited > cap: + break + try: + if f.is_file(): + low = f.name.lower() + if low.endswith((".gguf", ".safetensors")): + n += 1 + except OSError: + continue + except PermissionError as e: + logger.debug("browse-folders: permission denied counting %s: %s", directory, e) + return 0 + except OSError as e: + logger.debug("browse-folders: OS error counting %s: %s", directory, e) + return 0 + return n + + +def _has_direct_model_signal(directory: Path) -> bool: + """Return True if *directory* has an immediate child that signals + it holds a model: a GGUF/safetensors/config.json file, or a + `models--*` subdir (HF hub cache). Bounded by + ``_BROWSE_MODEL_HINT_PROBE`` to stay fast.""" + try: + it = directory.iterdir() + except OSError: + return False + try: + for i, child in enumerate(it): + if i >= _BROWSE_MODEL_HINT_PROBE: + break + try: + name = child.name + if child.is_file(): + low = name.lower() + if low.endswith((".gguf", ".safetensors")): + return True + if low in ("config.json", "adapter_config.json"): + return True + elif child.is_dir() and name.startswith("models--"): + return True + except OSError: + continue + except OSError: + return False + return False + + +def _looks_like_model_dir(directory: Path) -> bool: + """Bounded heuristic used by the folder browser to flag directories + worth exploring. False negatives are fine; the real scanner is + authoritative. + + Three signals, cheapest first: + + 1. Directory name itself: ``models--*`` is the HuggingFace hub cache + layout (``blobs``/``refs``/``snapshots`` children wouldn't match + the file-level probes below). + 2. An immediate child is a weight file or config (handled by + :func:`_has_direct_model_signal`). + 3. A grandchild has a direct signal -- this catches the + ``publisher/model/weights.gguf`` layout used by LM Studio and + Ollama. We probe at most the first + ``_BROWSE_MODEL_HINT_PROBE`` child directories, each of which is + checked with a bounded :func:`_has_direct_model_signal` call, + so the total cost stays O(PROBE^2) worst-case. + """ + if directory.name.startswith("models--"): + return True + if _has_direct_model_signal(directory): + return True + # Grandchild probe: LM Studio / Ollama publisher/model layout. + try: + it = directory.iterdir() + except OSError: + return False + try: + for i, child in enumerate(it): + if i >= _BROWSE_MODEL_HINT_PROBE: + break + try: + if not child.is_dir(): + continue + except OSError: + continue + # Fast name check first + if child.name.startswith("models--"): + return True + if _has_direct_model_signal(child): + return True + except OSError: + return False + return False + + +def _build_browse_allowlist() -> list[Path]: + """Return the list of root directories the folder browser is allowed + to walk. The same list is used to seed the sidebar suggestion chips, + so chip targets are always reachable. + + Roots include the current user's HOME, the resolved HF cache dirs, + Studio's own outputs/exports/studio root, registered scan folders, + and well-known third-party local-LLM dirs (LM Studio, Ollama, + `~/models`). Each is added only if it currently resolves to a real + directory, so we never produce a "dead" sandbox boundary the user + can't navigate into. + """ + from utils.paths import ( + hf_default_cache_dir, + legacy_hf_cache_dir, + well_known_model_dirs, + ) + from storage.studio_db import list_scan_folders + + candidates: list[Path] = [] + + def _add(p: Optional[Path]) -> None: + if p is None: + return + try: + resolved = p.resolve() + except OSError: + return + if resolved.is_dir(): + candidates.append(resolved) + + _add(Path.home()) + _add(_resolve_hf_cache_dir()) + try: + _add(hf_default_cache_dir()) + except Exception: # noqa: BLE001 -- best-effort + pass + try: + _add(legacy_hf_cache_dir()) + except Exception: # noqa: BLE001 -- best-effort + pass + try: + from utils.paths import ( + exports_root, + outputs_root, + studio_root, + ) + + _add(studio_root()) + _add(outputs_root()) + _add(exports_root()) + except Exception as exc: # noqa: BLE001 -- best-effort + logger.debug("browse-folders: studio roots unavailable: %s", exc) + try: + for folder in list_scan_folders(): + p = folder.get("path") + if p: + _add(Path(p)) + except Exception as exc: # noqa: BLE001 -- best-effort + logger.debug("browse-folders: could not load scan folders: %s", exc) + try: + for p in well_known_model_dirs(): + _add(p) + except Exception as exc: # noqa: BLE001 -- best-effort + logger.debug("browse-folders: well-known dirs unavailable: %s", exc) + + # Dedupe while preserving order. + seen: set[str] = set() + deduped: list[Path] = [] + for p in candidates: + key = str(p) + if key in seen: + continue + seen.add(key) + deduped.append(p) + return deduped + + +def _is_path_inside_allowlist(target: Path, allowed_roots: list[Path]) -> bool: + """Return True if *target* equals or is a descendant of any allowed + root. The comparison uses ``os.path.realpath`` so symlinks cannot be + used to escape the sandbox. + """ + try: + target_real = os.path.realpath(str(target)) + except OSError: + return False + for root in allowed_roots: + try: + root_real = os.path.realpath(str(root)) + except OSError: + continue + if target_real == root_real or target_real.startswith(root_real + os.sep): + return True + return False + + +def _normalize_browse_request_path(path: Optional[str]) -> str: + """Normalize the browse request path lexically, without touching the FS.""" + if path is None or not path.strip(): + return os.path.normpath(str(Path.home())) + + expanded = os.path.expanduser(path.strip()) + if not os.path.isabs(expanded): + expanded = os.path.join(str(Path.cwd()), expanded) + return os.path.normpath(expanded) + + +def _browse_relative_parts(requested_path: str, root: Path) -> Optional[list[str]]: + """Return validated relative path components under ``root``.""" + root_text = os.path.normpath(str(root)) + try: + rel_text = os.path.relpath(requested_path, root_text) + except ValueError: + return None + + if rel_text == ".": + return [] + if rel_text == ".." or rel_text.startswith(f"..{os.sep}"): + return None + + parts = [part for part in rel_text.split(os.sep) if part not in ("", ".")] + altsep = os.altsep + for part in parts: + if part == ".." or os.sep in part or (altsep and altsep in part): + return None + return parts + + +def _match_browse_child(current: Path, name: str) -> Optional[Path]: + """Return the immediate child named ``name`` under ``current``.""" + try: + for child in current.iterdir(): + if child.name == name: + return child + except PermissionError: + raise HTTPException( + status_code = 403, + detail = f"Permission denied reading {current}", + ) from None + except OSError as exc: + raise HTTPException( + status_code = 500, + detail = f"Could not read {current}: {exc}", + ) from exc + return None + + +def _resolve_browse_target(path: Optional[str], allowed_roots: list[Path]) -> Path: + """Resolve a requested browse path by walking from trusted allowlist roots.""" + requested_path = _normalize_browse_request_path(path) + resolved_roots: list[Path] = [] + seen_roots: set[str] = set() + for root in sorted(allowed_roots, key = lambda p: len(str(p)), reverse = True): + try: + resolved = root.resolve() + except OSError: + continue + key = str(resolved) + if key in seen_roots: + continue + seen_roots.add(key) + resolved_roots.append(resolved) + + for root in resolved_roots: + parts = _browse_relative_parts(requested_path, root) + if parts is None: + continue + + current = root + for part in parts: + child = _match_browse_child(current, part) + if child is None: + raise HTTPException( + status_code = 404, + detail = f"Path does not exist: {requested_path}", + ) + try: + resolved_child = child.resolve() + except OSError as exc: + raise HTTPException( + status_code = 400, + detail = f"Invalid path: {exc}", + ) from exc + if not _is_path_inside_allowlist(resolved_child, resolved_roots): + raise HTTPException( + status_code = 403, + detail = ( + "Path is not in the browseable allowlist. Register it via " + "POST /api/models/scan-folders first, or pick a directory " + "under your home folder." + ), + ) + current = resolved_child + + if not current.is_dir(): + raise HTTPException( + status_code = 400, + detail = f"Not a directory: {current}", + ) + return current + + raise HTTPException( + status_code = 403, + detail = ( + "Path is not in the browseable allowlist. Register it via " + "POST /api/models/scan-folders first, or pick a directory " + "under your home folder." + ), + ) + + +@router.get("/browse-folders", response_model = BrowseFoldersResponse) +async def browse_folders( + path: Optional[str] = Query( + None, + description = ( + "Directory to list. If omitted, defaults to the current user's " + "home directory. Tilde (`~`) and relative paths are expanded. " + "Must resolve inside the allowlist of browseable roots (HOME, " + "HF cache, Studio dirs, registered scan folders, well-known " + "model dirs)." + ), + ), + show_hidden: bool = Query( + False, + description = "Include entries whose name starts with a dot", + ), + current_subject: str = Depends(get_current_subject), +): + """ + List immediate subdirectories of *path* for the Custom Folders picker. + + The frontend uses this to render a modal folder browser without needing + a native OS dialog (Studio is served over HTTP, so the browser can't + reveal absolute paths on the host). The endpoint is read-only and does + not create, move, or delete anything. It simply enumerates visible + subdirectories so the user can click their way to a folder and hand + the resulting string back to POST `/api/models/scan-folders`. + + Sandbox: requests are bounded to the allowlist returned by + :func:`_build_browse_allowlist` (HOME, HF cache, Studio dirs, + registered scan folders, well-known model dirs). Paths outside the + allowlist return 403 so users cannot probe ``/etc``, ``/proc``, + ``/root`` (when not HOME), or other sensitive system locations + even if the server process can read them. Symlinks are resolved + via ``os.path.realpath`` before the check, so symlink traversal + cannot escape the sandbox either. + + Sorting: directories that look like they hold models come first, then + plain directories, then hidden entries (if `show_hidden=true`). + """ + from utils.paths import hf_default_cache_dir, well_known_model_dirs + from storage.studio_db import list_scan_folders + + # Build the allowlist once -- both the sandbox check below and the + # suggestion chips use the same set, so chips are always navigable. + allowed_roots = _build_browse_allowlist() + + try: + target = _resolve_browse_target(path, allowed_roots) + except HTTPException: + requested_path = _normalize_browse_request_path(path) + if path is not None and path.strip(): + logger.warning( + "browse-folders: rejected path %r (normalized=%s)", + path, + requested_path, + ) + raise + + # Enumerate immediate subdirectories with a bounded cap so a stray + # query against ``/usr/lib`` or ``/proc`` can't stat-storm the process. + entries: list[BrowseEntry] = [] + truncated = False + visited = 0 + try: + it = target.iterdir() + except PermissionError: + raise HTTPException( + status_code = 403, + detail = f"Permission denied reading {target}", + ) + except OSError as exc: + raise HTTPException( + status_code = 500, + detail = f"Could not read {target}: {exc}", + ) + + try: + for child in it: + # Bound by *visited entries*, not by *appended entries*: in + # directories full of files (or hidden subdirs when + # ``show_hidden=False``) the cap on ``len(entries)`` would + # never trigger and we'd still stat every child. Counting + # visits keeps the worst-case work to ``_BROWSE_ENTRY_CAP`` + # iterdir/is_dir calls regardless of how many of them + # survive the filters below. + visited += 1 + if visited > _BROWSE_ENTRY_CAP: + truncated = True + break + try: + if not child.is_dir(): + continue + except OSError: + continue + name = child.name + is_hidden = name.startswith(".") + if is_hidden and not show_hidden: + continue + entries.append( + BrowseEntry( + name = name, + has_models = _looks_like_model_dir(child), + hidden = is_hidden, + ) + ) + except PermissionError as exc: + logger.debug( + "browse-folders: permission denied during enumeration of %s: %s", + target, + exc, + ) + except OSError as exc: + # Rare: iterdir succeeded but reading a specific entry failed. + logger.warning("browse-folders: partial enumeration of %s: %s", target, exc) + + # Model-bearing dirs first, then plain, then hidden; case-insensitive + # alphabetical within each bucket. + def _sort_key(e: BrowseEntry) -> tuple[int, str]: + bucket = 0 if e.has_models else (2 if e.hidden else 1) + return (bucket, e.name.lower()) + + entries.sort(key = _sort_key) + + # Parent is None at the filesystem root (`p.parent == p`) AND when + # the parent would step outside the sandbox -- otherwise the up-row + # would 403 on click. Users can still hop to other allowed roots + # via the suggestion chips below. + parent: Optional[str] + if target.parent == target or not _is_path_inside_allowlist( + target.parent, allowed_roots + ): + parent = None + else: + parent = str(target.parent) + + # Handy starting points for the quick-pick chips. + suggestions: list[str] = [] + seen_sug: set[str] = set() + + def _add_sug(p: Optional[Path]) -> None: + if p is None: + return + try: + resolved = str(p.resolve()) + except OSError: + return + if resolved in seen_sug: + return + if Path(resolved).is_dir(): + seen_sug.add(resolved) + suggestions.append(resolved) + + # Home always comes first -- it's the safe fallback when everything + # else is cold. + _add_sug(Path.home()) + # The HF cache root the process is actually using. + try: + _add_sug(hf_default_cache_dir()) + except Exception: + pass + # Already-registered scan folders (what the user has curated). + try: + for folder in list_scan_folders(): + _add_sug(Path(folder.get("path", ""))) + except Exception as exc: + logger.debug("browse-folders: could not load scan folders: %s", exc) + # Directories commonly used by other local-LLM tools: LM Studio + # (`~/.lmstudio/models` + legacy `~/.cache/lm-studio/models` + + # user-configured downloadsFolder from LM Studio's settings.json), + # Ollama (`~/.ollama/models` + common system paths + OLLAMA_MODELS + # env var), and generic user-choice spots (`~/models`, `~/Models`). + # Each helper only returns paths that currently exist so we never + # show dead chips. + try: + for p in well_known_model_dirs(): + _add_sug(p) + except Exception as exc: + logger.debug("browse-folders: could not load well-known dirs: %s", exc) + + return BrowseFoldersResponse( + current = str(target), + parent = parent, + entries = entries, + suggestions = suggestions, + truncated = truncated, + model_files_here = _count_model_files(target), + ) + + @router.get("/list") async def list_models( current_subject: str = Depends(get_current_subject), @@ -791,15 +1318,16 @@ async def scan_loras( lora_list = [] # Scan training outputs - trained_loras = scan_trained_loras(outputs_dir = resolved_outputs_dir) - for display_name, adapter_path in trained_loras: - base_model = get_base_model_from_lora(adapter_path) + trained_models = scan_trained_models(outputs_dir = resolved_outputs_dir) + for display_name, model_path, model_type in trained_models: + base_model = get_base_model_from_checkpoint(model_path) lora_list.append( LoRAInfo( display_name = display_name, - adapter_path = adapter_path, + adapter_path = model_path, base_model = base_model, source = "training", + export_type = model_type, ) ) @@ -989,7 +1517,7 @@ async def get_gguf_variants( snapshots = entry / "snapshots" if snapshots.is_dir(): for snap in snapshots.iterdir(): - for f in snap.rglob("*.gguf"): + 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 @@ -1058,7 +1586,7 @@ async def get_gguf_download_progress( for entry in cache_dir.iterdir(): if entry.name.lower() == target: # Count completed .gguf files matching this variant in snapshots - for f in entry.rglob("*.gguf"): + for f in _iter_gguf_paths(entry): fname = f.name.lower().replace("-", "").replace("_", "") if not variant_lower or variant_lower in fname: downloaded_bytes += f.stat().st_size @@ -1088,6 +1616,25 @@ async def get_gguf_download_progress( return {"downloaded_bytes": 0, "expected_bytes": expected_bytes, "progress": 0} +def _resolve_hf_cache_realpath(repo_dir: Path) -> Optional[str]: + """Pick the most useful on-disk path for a HF cache repo. + + Prefers the most-recent snapshot dir (what `from_pretrained` actually + points at). Falls back to the cache repo root. Returns the resolved + realpath so symlinks under snapshots/ are followed back to blobs/. + """ + try: + snapshots_dir = repo_dir / "snapshots" + if snapshots_dir.is_dir(): + snaps = [s for s in snapshots_dir.iterdir() if s.is_dir()] + if snaps: + latest = max(snaps, key = lambda s: s.stat().st_mtime) + return str(latest.resolve()) + return str(repo_dir.resolve()) + except Exception: + return None + + @router.get("/download-progress") async def get_download_progress( repo_id: str = Query(..., description = "HuggingFace repo ID"), @@ -1098,8 +1645,16 @@ async def get_download_progress( Checks the local HF cache for completed blobs and in-progress (.incomplete) downloads. Uses the HF API to determine the expected total size on the first call, then caches it for subsequent polls. + Also returns ``cache_path``: the realpath of the snapshot directory + (or the cache repo root if no snapshot exists yet) so the UI can + show users where the weights actually live on disk. """ - _empty = {"downloaded_bytes": 0, "expected_bytes": 0, "progress": 0} + _empty = { + "downloaded_bytes": 0, + "expected_bytes": 0, + "progress": 0, + "cache_path": None, + } try: if not _is_valid_repo_id(repo_id): return _empty @@ -1110,10 +1665,12 @@ async def get_download_progress( target = f"models--{repo_id.replace('/', '--')}".lower() completed_bytes = 0 in_progress_bytes = 0 + cache_path: Optional[str] = None for entry in cache_dir.iterdir(): if entry.name.lower() != target: continue + cache_path = _resolve_hf_cache_realpath(entry) blobs_dir = entry / "blobs" if not blobs_dir.is_dir(): break @@ -1128,7 +1685,7 @@ async def get_download_progress( downloaded_bytes = completed_bytes + in_progress_bytes if downloaded_bytes == 0: - return _empty + return {**_empty, "cache_path": cache_path} # Get expected size from HF API (cached per repo_id) expected_bytes = _get_repo_size_cached(repo_id) @@ -1138,6 +1695,7 @@ async def get_download_progress( "downloaded_bytes": downloaded_bytes, "expected_bytes": 0, "progress": 0, + "cache_path": cache_path, } # Use 95% threshold for completion (blob deduplication can make @@ -1153,6 +1711,7 @@ async def get_download_progress( "downloaded_bytes": downloaded_bytes, "expected_bytes": expected_bytes, "progress": round(progress, 3), + "cache_path": cache_path, } except Exception as e: logger.warning(f"Error checking download progress for {repo_id}: {e}") @@ -1203,6 +1762,62 @@ def _all_hf_cache_scans(): return scans +def _is_gguf_filename(name: str) -> bool: + return name.lower().endswith(".gguf") + + +def _is_mmproj_filename(name: str) -> bool: + """Match GGUF vision-adapter (mmproj) files. Kept consistent with + ``utils.models.model_config._is_mmproj``.""" + return "mmproj" in name.lower() + + +def _is_main_gguf_filename(name: str) -> bool: + """A GGUF file that is a primary weight artifact, not an mmproj + vision adapter.""" + return _is_gguf_filename(name) and not _is_mmproj_filename(name) + + +def _iter_gguf_paths(root: Path): + for path in root.rglob("*"): + if path.is_file() and _is_gguf_filename(path.name): + yield path + + +def _repo_gguf_size_bytes(repo_info) -> int: + """Return the total on-disk size of primary GGUF weight files across + all revisions, excluding mmproj vision-adapter files. + + Hugging Face hardlinks blobs shared between revisions, so this + deduplicates by blob path (or, as a fallback, by revision commit + hash + filename) to avoid double-counting the same bytes. Files + with an unknown size (``size_on_disk is None``, e.g. a partial or + interrupted download) are treated as zero bytes. mmproj files are + excluded so that repos whose only ``.gguf`` artifact is a vision + adapter are not classified as GGUF repos: the variant selector + filters mmproj out and would otherwise show zero pickable variants. + """ + unique_blobs: dict[str, int] = {} + for revision in repo_info.revisions: + rev_id = getattr(revision, "commit_hash", None) or str(id(revision)) + for f in revision.files: + if _is_main_gguf_filename(f.file_name): + blob_path = getattr(f, "blob_path", None) + size = f.size_on_disk or 0 + if blob_path: + unique_blobs[str(blob_path)] = size + else: + unique_blobs[f"{rev_id}:{f.file_name}"] = size + return sum(unique_blobs.values()) + + +def _repo_has_gguf_files(repo_info) -> bool: + """Return True when any revision in a cached repo contains a + primary GGUF weight file. Repos whose only ``.gguf`` artifact is + an mmproj vision adapter are not treated as GGUF here.""" + return _repo_gguf_size_bytes(repo_info) > 0 + + @router.get("/cached-gguf") async def list_cached_gguf( current_subject: str = Depends(get_current_subject), @@ -1214,28 +1829,25 @@ async def list_cached_gguf( seen_lower: dict[str, dict] = {} for hf_cache in cache_scans: for repo_info in hf_cache.repos: - if repo_info.repo_type != "model": + try: + if repo_info.repo_type != "model": + continue + repo_id = repo_info.repo_id + total_size = _repo_gguf_size_bytes(repo_info) + if total_size == 0: + continue + key = repo_id.lower() + existing = seen_lower.get(key) + if existing is None or total_size > existing["size_bytes"]: + seen_lower[key] = { + "repo_id": repo_id, + "size_bytes": total_size, + "cache_path": str(repo_info.repo_path), + } + except Exception as e: + repo_label = getattr(repo_info, "repo_id", "") + logger.warning(f"Skipping cached GGUF repo {repo_label}: {e}") continue - repo_id = repo_info.repo_id - if not repo_id.upper().endswith("-GGUF"): - continue - total_size = 0 - has_gguf = False - for revision in repo_info.revisions: - for f in revision.files: - if f.file_name.endswith(".gguf"): - has_gguf = True - total_size += f.size_on_disk - if not has_gguf: - continue - key = repo_id.lower() - existing = seen_lower.get(key) - if existing is None or total_size > existing["size_bytes"]: - seen_lower[key] = { - "repo_id": repo_id, - "size_bytes": total_size, - "cache_path": str(repo_info.repo_path), - } cached = sorted(seen_lower.values(), key = lambda c: c["repo_id"]) return {"cached": cached} except Exception as e: @@ -1256,30 +1868,37 @@ async def list_cached_models( seen_lower: dict[str, dict] = {} for hf_cache in cache_scans: for repo_info in hf_cache.repos: - if repo_info.repo_type != "model": + try: + if repo_info.repo_type != "model": + continue + repo_id = repo_info.repo_id + if _repo_has_gguf_files(repo_info): + continue + total_size = sum( + (f.size_on_disk or 0) + for rev in repo_info.revisions + for f in rev.files + ) + if total_size == 0: + continue + has_weights = any( + f.file_name.endswith(_WEIGHT_EXTENSIONS) + for rev in repo_info.revisions + for f in rev.files + ) + if not has_weights: + continue + key = repo_id.lower() + existing = seen_lower.get(key) + if existing is None or total_size > existing["size_bytes"]: + seen_lower[key] = { + "repo_id": repo_id, + "size_bytes": total_size, + } + except Exception as e: + repo_label = getattr(repo_info, "repo_id", "") + logger.warning(f"Skipping cached model repo {repo_label}: {e}") continue - repo_id = repo_info.repo_id - if repo_id.upper().endswith("-GGUF"): - continue - total_size = sum( - f.size_on_disk for rev in repo_info.revisions for f in rev.files - ) - if total_size == 0: - continue - has_weights = any( - f.file_name.endswith(_WEIGHT_EXTENSIONS) - for rev in repo_info.revisions - for f in rev.files - ) - if not has_weights: - continue - key = repo_id.lower() - existing = seen_lower.get(key) - if existing is None or total_size > existing["size_bytes"]: - seen_lower[key] = { - "repo_id": repo_id, - "size_bytes": total_size, - } cached = sorted(seen_lower.values(), key = lambda c: c["repo_id"]) return {"cached": cached} except Exception as e: @@ -1356,7 +1975,7 @@ async def delete_cached_model( deleted_count = 0 for rev in target_repo.revisions: for f in rev.files: - if not f.file_name.endswith(".gguf"): + if not _is_gguf_filename(f.file_name): continue quant = _extract_quant_label(f.file_name) if quant.lower() != variant.lower(): diff --git a/studio/backend/run.py b/studio/backend/run.py index 9c3622988e..9675b9ea4c 100644 --- a/studio/backend/run.py +++ b/studio/backend/run.py @@ -248,6 +248,7 @@ def run_server( port: int = 8888, frontend_path: Path = Path(__file__).resolve().parent.parent / "frontend" / "dist", silent: bool = False, + llama_parallel_slots: int = 1, ): """ Start the FastAPI server. @@ -257,6 +258,7 @@ def run_server( port: Port to bind to (auto-increments if in use) frontend_path: Path to frontend build directory (optional) silent: Suppress startup messages + llama_parallel_slots: Number of parallel slots for llama-server Note: Signal handlers are NOT registered here so that embedders @@ -324,6 +326,15 @@ def run_server( _server = uvicorn.Server(config) _shutdown_event = Event() + # Expose the actual bound port so request-handling code can build + # loopback URLs that point at the real backend, not whatever port a + # reverse proxy or tunnel exposed in the request URL. Only publish + # an explicit value when we know the concrete port; for ephemeral + # binds (port==0) leave it unset and let request handlers fall back + # to the ASGI request scope or request.base_url. + app.state.server_port = port if port and port > 0 else None + app.state.llama_parallel_slots = llama_parallel_slots + # Run server in a daemon thread def _run(): asyncio.run(_server.serve()) diff --git a/studio/backend/tests/conftest.py b/studio/backend/tests/conftest.py index 053e9b85d9..6aa6d314c1 100644 --- a/studio/backend/tests/conftest.py +++ b/studio/backend/tests/conftest.py @@ -3,14 +3,136 @@ """ Shared pytest configuration for the backend test suite. -Ensures that the backend root is on sys.path so that -`import utils.utils` (and similar flat imports) resolve correctly. + +Responsibilities: +1. Put the backend root on sys.path so `from models.inference import ...` + (and similar flat imports) resolve in test modules — mirrors how the + app itself is launched. +2. Provide a hybrid ``studio_server`` session fixture for end-to-end tests + (see ``test_studio_api.py``). The fixture supports two invocation modes: + + a. **External server.** If ``UNSLOTH_E2E_BASE_URL`` is set, tests point + at an already-running Studio instance. ``UNSLOTH_E2E_API_KEY`` must + also be set. This is the fast-iteration mode: start the server once + with ``unsloth studio run ...``, then run pytest against it many + times with no per-run GGUF load cost. + + b. **Fixture-managed server.** Otherwise, the fixture launches a fresh + server via ``_start_server`` and tears it down at session end. This + is the one-shot mode for CI or a clean-slate verification run. + + The model / variant for mode (b) come from ``--unsloth-model`` / + ``--unsloth-gguf-variant`` pytest options, then ``UNSLOTH_E2E_MODEL`` / + ``UNSLOTH_E2E_VARIANT`` env vars, then the defaults in + ``test_studio_api.py``. """ +import os import sys from pathlib import Path +import pytest + # Add backend root to sys.path (mirrors how the app itself is launched) _backend_root = Path(__file__).resolve().parent.parent if str(_backend_root) not in sys.path: sys.path.insert(0, str(_backend_root)) + + +# ── Pytest CLI options ─────────────────────────────────────────────── + + +def pytest_addoption(parser): + group = parser.getgroup( + "unsloth-e2e", + "Unsloth Studio end-to-end test options", + ) + group.addoption( + "--unsloth-model", + action = "store", + default = None, + help = ( + "GGUF model id used when starting a server for e2e tests. " + "Ignored if UNSLOTH_E2E_BASE_URL is set. Overrides " + "UNSLOTH_E2E_MODEL env var. Defaults to test_studio_api.py's " + "DEFAULT_MODEL." + ), + ) + group.addoption( + "--unsloth-gguf-variant", + action = "store", + default = None, + help = ( + "GGUF variant used when starting a server for e2e tests. " + "Ignored if UNSLOTH_E2E_BASE_URL is set. Overrides " + "UNSLOTH_E2E_VARIANT env var. Defaults to test_studio_api.py's " + "DEFAULT_VARIANT." + ), + ) + + +# ── E2E server fixtures ────────────────────────────────────────────── + + +@pytest.fixture(scope = "session") +def studio_server(request): + """Yield ``(base_url, api_key)`` for e2e tests. + + Resolution order: + + 1. If ``UNSLOTH_E2E_BASE_URL`` is set → point at that server, + require ``UNSLOTH_E2E_API_KEY`` alongside (skip if missing). + 2. Otherwise → start a fresh ``unsloth studio run`` subprocess via + the existing ``_start_server`` helper in ``test_studio_api.py`` + and tear it down on session teardown. + + Session-scoped so the expensive GGUF load happens at most once per + pytest invocation. Lazily instantiated — tests that don't request + the fixture (e.g. the unit tests in ``test_anthropic_messages.py`` + or ``test_help_output``) do not trigger server startup. + """ + external_url = os.environ.get("UNSLOTH_E2E_BASE_URL") + if external_url: + api_key = os.environ.get("UNSLOTH_E2E_API_KEY") + if not api_key: + pytest.skip( + "UNSLOTH_E2E_BASE_URL is set but UNSLOTH_E2E_API_KEY is " + "missing — tests that require auth cannot run against an " + "external server without it.", + ) + yield external_url, api_key + return + + # Lazy import: pytest has already loaded test_studio_api into + # sys.modules by the time any test requests this fixture, so this + # is a cache hit, not a re-execution. + import test_studio_api as _e2e + + model = ( + request.config.getoption("--unsloth-model") + or os.environ.get("UNSLOTH_E2E_MODEL") + or _e2e.DEFAULT_MODEL + ) + variant = ( + request.config.getoption("--unsloth-gguf-variant") + or os.environ.get("UNSLOTH_E2E_VARIANT") + or _e2e.DEFAULT_VARIANT + ) + + proc, api_key = _e2e._start_server(model, variant) + try: + yield f"http://{_e2e.HOST}:{_e2e.PORT}", api_key + finally: + _e2e._kill_server(proc) + + +@pytest.fixture +def base_url(studio_server): + """Base URL for the e2e Studio server (from ``studio_server``).""" + return studio_server[0] + + +@pytest.fixture +def api_key(studio_server): + """API key for the e2e Studio server (from ``studio_server``).""" + return studio_server[1] diff --git a/studio/backend/tests/test_anthropic_messages.py b/studio/backend/tests/test_anthropic_messages.py new file mode 100644 index 0000000000..ec432df9e6 --- /dev/null +++ b/studio/backend/tests/test_anthropic_messages.py @@ -0,0 +1,774 @@ +# SPDX-License-Identifier: AGPL-3.0-only +# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. + +""" +Tests for the Anthropic Messages API schemas and translation layer. +No running server or GPU required. +""" + +import sys +import os +import json + +_backend = os.path.join(os.path.dirname(__file__), "..") +sys.path.insert(0, _backend) + +from models.inference import ( + AnthropicMessagesRequest, + AnthropicMessagesResponse, + AnthropicMessage, + AnthropicTextBlock, + AnthropicToolUseBlock, + AnthropicToolResultBlock, + AnthropicTool, + AnthropicUsage, + AnthropicResponseTextBlock, + AnthropicResponseToolUseBlock, +) +from core.inference.anthropic_compat import ( + anthropic_messages_to_openai, + anthropic_tools_to_openai, + build_anthropic_sse_event, + AnthropicStreamEmitter, + AnthropicPassthroughEmitter, +) + + +# ===================================================================== +# Pydantic model tests +# ===================================================================== + + +class TestAnthropicModels: + def test_minimal_request(self): + req = AnthropicMessagesRequest( + messages = [{"role": "user", "content": "Hi"}], + ) + assert req.max_tokens is None + assert req.model == "default" + assert req.stream is False + + def test_max_tokens_optional(self): + req = AnthropicMessagesRequest( + max_tokens = 100, + messages = [{"role": "user", "content": "Hi"}], + ) + assert req.max_tokens == 100 + + def test_system_as_string(self): + req = AnthropicMessagesRequest( + max_tokens = 50, + messages = [{"role": "user", "content": "Hi"}], + system = "You are helpful.", + ) + assert req.system == "You are helpful." + + def test_tools_field_parses(self): + req = AnthropicMessagesRequest( + max_tokens = 100, + messages = [{"role": "user", "content": "Hi"}], + tools = [{"name": "web_search", "input_schema": {"type": "object"}}], + ) + assert len(req.tools) == 1 + assert req.tools[0].name == "web_search" + + def test_extra_fields_accepted(self): + req = AnthropicMessagesRequest( + max_tokens = 100, + messages = [{"role": "user", "content": "Hi"}], + some_future_field = "hello", + ) + assert req.max_tokens == 100 + + def test_stream_defaults_false(self): + req = AnthropicMessagesRequest( + max_tokens = 100, + messages = [{"role": "user", "content": "Hi"}], + ) + assert req.stream is False + + def test_enable_tools_shorthand(self): + req = AnthropicMessagesRequest( + messages = [{"role": "user", "content": "Hi"}], + enable_tools = True, + enabled_tools = ["web_search", "python"], + session_id = "my-session", + ) + assert req.enable_tools is True + assert req.enabled_tools == ["web_search", "python"] + assert req.session_id == "my-session" + + def test_extension_fields_default_none(self): + req = AnthropicMessagesRequest( + messages = [{"role": "user", "content": "Hi"}], + ) + assert req.enable_tools is None + assert req.enabled_tools is None + assert req.session_id is None + + def test_response_model_defaults(self): + resp = AnthropicMessagesResponse() + assert resp.type == "message" + assert resp.role == "assistant" + assert resp.id.startswith("msg_") + assert resp.content == [] + assert resp.usage.input_tokens == 0 + + +# ===================================================================== +# Message translation tests +# ===================================================================== + + +class TestAnthropicMessagesToOpenAI: + def test_simple_user_message(self): + msgs = [{"role": "user", "content": "Hello"}] + result = anthropic_messages_to_openai(msgs) + assert result == [{"role": "user", "content": "Hello"}] + + def test_system_string_prepended(self): + msgs = [{"role": "user", "content": "Hello"}] + result = anthropic_messages_to_openai(msgs, system = "Be brief.") + assert result[0] == {"role": "system", "content": "Be brief."} + assert result[1] == {"role": "user", "content": "Hello"} + + def test_system_as_block_list(self): + system = [ + {"type": "text", "text": "Be brief."}, + {"type": "text", "text": "Be accurate."}, + ] + msgs = [{"role": "user", "content": "Hello"}] + result = anthropic_messages_to_openai(msgs, system = system) + assert result[0]["role"] == "system" + assert "Be brief." in result[0]["content"] + assert "Be accurate." in result[0]["content"] + + def test_multi_turn_conversation(self): + msgs = [ + {"role": "user", "content": "Hi"}, + {"role": "assistant", "content": "Hello!"}, + {"role": "user", "content": "How are you?"}, + ] + result = anthropic_messages_to_openai(msgs) + assert len(result) == 3 + assert result[0]["role"] == "user" + assert result[1]["role"] == "assistant" + assert result[2]["role"] == "user" + + def test_assistant_tool_use_maps_to_tool_calls(self): + msgs = [ + { + "role": "assistant", + "content": [ + {"type": "text", "text": "Let me search."}, + { + "type": "tool_use", + "id": "tu_1", + "name": "web_search", + "input": {"query": "test"}, + }, + ], + } + ] + result = anthropic_messages_to_openai(msgs) + assert len(result) == 1 + m = result[0] + assert m["role"] == "assistant" + assert m["content"] == "Let me search." + assert len(m["tool_calls"]) == 1 + tc = m["tool_calls"][0] + assert tc["id"] == "tu_1" + assert tc["function"]["name"] == "web_search" + assert json.loads(tc["function"]["arguments"]) == {"query": "test"} + + def test_tool_result_maps_to_tool_role(self): + msgs = [ + { + "role": "user", + "content": [ + { + "type": "tool_result", + "tool_use_id": "tu_1", + "content": "Result text", + }, + ], + } + ] + result = anthropic_messages_to_openai(msgs) + assert len(result) == 1 + assert result[0]["role"] == "tool" + assert result[0]["tool_call_id"] == "tu_1" + assert result[0]["content"] == "Result text" + + def test_mixed_text_and_tool_use_blocks(self): + msgs = [ + { + "role": "assistant", + "content": [ + {"type": "text", "text": "Thinking..."}, + { + "type": "tool_use", + "id": "tu_1", + "name": "python", + "input": {"code": "1+1"}, + }, + { + "type": "tool_use", + "id": "tu_2", + "name": "terminal", + "input": {"command": "ls"}, + }, + ], + } + ] + result = anthropic_messages_to_openai(msgs) + assert len(result) == 1 + m = result[0] + assert m["content"] == "Thinking..." + assert len(m["tool_calls"]) == 2 + + def test_tool_result_with_list_content(self): + msgs = [ + { + "role": "user", + "content": [ + { + "type": "tool_result", + "tool_use_id": "tu_1", + "content": [ + {"type": "text", "text": "Line 1"}, + {"type": "text", "text": "Line 2"}, + ], + }, + ], + } + ] + result = anthropic_messages_to_openai(msgs) + assert result[0]["content"] == "Line 1 Line 2" + + +# ===================================================================== +# Tool translation tests +# ===================================================================== + + +class TestAnthropicToolsToOpenAI: + def test_single_tool(self): + tools = [ + { + "name": "web_search", + "description": "Search", + "input_schema": { + "type": "object", + "properties": {"query": {"type": "string"}}, + }, + } + ] + result = anthropic_tools_to_openai(tools) + assert len(result) == 1 + assert result[0]["type"] == "function" + assert result[0]["function"]["name"] == "web_search" + assert result[0]["function"]["parameters"]["type"] == "object" + + def test_multiple_tools(self): + tools = [ + {"name": "a", "description": "Tool A", "input_schema": {}}, + {"name": "b", "description": "Tool B", "input_schema": {}}, + ] + result = anthropic_tools_to_openai(tools) + assert len(result) == 2 + assert result[0]["function"]["name"] == "a" + assert result[1]["function"]["name"] == "b" + + def test_empty_list(self): + assert anthropic_tools_to_openai([]) == [] + + def test_pydantic_model_input(self): + tool = AnthropicTool( + name = "test", description = "desc", input_schema = {"type": "object"} + ) + result = anthropic_tools_to_openai([tool]) + assert result[0]["function"]["name"] == "test" + + +# ===================================================================== +# SSE event helper tests +# ===================================================================== + + +class TestBuildAnthropicSSEEvent: + def test_basic_event(self): + result = build_anthropic_sse_event("message_start", {"type": "message_start"}) + assert result.startswith("event: message_start\n") + assert "data: " in result + assert result.endswith("\n\n") + + def test_data_is_valid_json(self): + result = build_anthropic_sse_event("test", {"key": "value"}) + data_line = result.split("\n")[1] + payload = json.loads(data_line.removeprefix("data: ")) + assert payload == {"key": "value"} + + +# ===================================================================== +# Stream emitter tests +# ===================================================================== + + +class TestAnthropicStreamEmitter: + def test_start_emits_message_start_and_content_block_start(self): + e = AnthropicStreamEmitter() + events = e.start("msg_123", "test-model") + assert len(events) == 2 + assert "message_start" in events[0] + assert "content_block_start" in events[1] + assert '"type": "text"' in events[1] + + def test_content_delta_emits_text_delta(self): + e = AnthropicStreamEmitter() + e.start("msg_1", "m") + events = e.feed({"type": "content", "text": "Hello"}) + assert len(events) == 1 + parsed = json.loads(events[0].split("data: ")[1]) + assert parsed["delta"]["type"] == "text_delta" + assert parsed["delta"]["text"] == "Hello" + + def test_cumulative_content_diffs_correctly(self): + e = AnthropicStreamEmitter() + e.start("msg_1", "m") + e.feed({"type": "content", "text": "Hel"}) + events = e.feed({"type": "content", "text": "Hello"}) + parsed = json.loads(events[0].split("data: ")[1]) + assert parsed["delta"]["text"] == "lo" + + def test_empty_content_diff_no_event(self): + e = AnthropicStreamEmitter() + e.start("msg_1", "m") + e.feed({"type": "content", "text": "Hi"}) + events = e.feed({"type": "content", "text": "Hi"}) + assert events == [] + + def test_tool_start_closes_text_opens_tool_block(self): + e = AnthropicStreamEmitter() + e.start("msg_1", "m") + e.feed({"type": "content", "text": "Thinking"}) + events = e.feed( + { + "type": "tool_start", + "tool_name": "web_search", + "tool_call_id": "tc_1", + "arguments": {"query": "test"}, + } + ) + # content_block_stop + content_block_start(tool_use) + content_block_delta(input_json) + assert len(events) == 3 + assert "content_block_stop" in events[0] + assert "tool_use" in events[1] + assert "input_json_delta" in events[2] + + def test_tool_end_closes_tool_opens_new_text_block(self): + e = AnthropicStreamEmitter() + e.start("msg_1", "m") + e.feed( + { + "type": "tool_start", + "tool_name": "t", + "tool_call_id": "tc_1", + "arguments": {}, + } + ) + events = e.feed( + { + "type": "tool_end", + "tool_name": "t", + "tool_call_id": "tc_1", + "result": "done", + } + ) + # content_block_stop (tool) + tool_result + content_block_start (new text) + assert len(events) == 3 + assert "content_block_stop" in events[0] + assert "tool_result" in events[1] + parsed = json.loads(events[1].split("data: ")[1]) + assert parsed["content"] == "done" + assert parsed["tool_use_id"] == "tc_1" + assert "content_block_start" in events[2] + assert '"type": "text"' in events[2] + + def test_finish_emits_stop_events(self): + e = AnthropicStreamEmitter() + e.start("msg_1", "m") + events = e.finish("end_turn") + # content_block_stop + message_delta + message_stop + assert len(events) == 3 + assert "content_block_stop" in events[0] + assert "message_delta" in events[1] + assert "end_turn" in events[1] + assert "message_stop" in events[2] + + def test_metadata_captured_in_finish_usage(self): + e = AnthropicStreamEmitter() + e.start("msg_1", "m") + e.feed( + { + "type": "metadata", + "usage": {"prompt_tokens": 10, "completion_tokens": 20}, + } + ) + events = e.finish("end_turn") + delta_event = [ev for ev in events if "message_delta" in ev][0] + parsed = json.loads(delta_event.split("data: ")[1]) + assert parsed["usage"]["output_tokens"] == 20 + + def test_status_events_ignored(self): + e = AnthropicStreamEmitter() + e.start("msg_1", "m") + events = e.feed({"type": "status", "text": "Searching..."}) + assert events == [] + + def test_no_tool_calls_simple_text_flow(self): + e = AnthropicStreamEmitter() + start_events = e.start("msg_1", "m") + content_events = e.feed({"type": "content", "text": "Hello world"}) + meta_events = e.feed( + {"type": "metadata", "usage": {"prompt_tokens": 5, "completion_tokens": 2}} + ) + end_events = e.finish("end_turn") + + assert len(start_events) == 2 + assert len(content_events) == 1 + assert meta_events == [] + assert len(end_events) == 3 + + def test_block_index_increments(self): + e = AnthropicStreamEmitter() + e.start("msg_1", "m") + assert e.block_index == 0 + e.feed( + { + "type": "tool_start", + "tool_name": "t", + "tool_call_id": "tc_1", + "arguments": {}, + } + ) + assert e.block_index == 1 + e.feed( + { + "type": "tool_end", + "tool_name": "t", + "tool_call_id": "tc_1", + "result": "ok", + } + ) + assert e.block_index == 2 + + def test_text_after_tool_resets_prev_text(self): + e = AnthropicStreamEmitter() + e.start("msg_1", "m") + e.feed({"type": "content", "text": "Before tool"}) + e.feed( + { + "type": "tool_start", + "tool_name": "t", + "tool_call_id": "tc_1", + "arguments": {}, + } + ) + e.feed( + { + "type": "tool_end", + "tool_name": "t", + "tool_call_id": "tc_1", + "result": "ok", + } + ) + # After tool_end, prev_text should be reset + events = e.feed({"type": "content", "text": "After tool"}) + parsed = json.loads(events[0].split("data: ")[1]) + assert parsed["delta"]["text"] == "After tool" + + +# ===================================================================== +# Pass-through emitter tests (client-side tool execution path) +# ===================================================================== + + +class TestAnthropicPassthroughEmitter: + def _parse(self, event_str): + return json.loads(event_str.split("data: ")[1]) + + def test_start_emits_message_start_only(self): + e = AnthropicPassthroughEmitter() + events = e.start("msg_1", "test-model") + assert len(events) == 1 + assert "message_start" in events[0] + parsed = self._parse(events[0]) + assert parsed["message"]["id"] == "msg_1" + assert parsed["message"]["model"] == "test-model" + + def test_text_chunk_opens_text_block_and_emits_delta(self): + e = AnthropicPassthroughEmitter() + e.start("msg_1", "m") + chunk = {"choices": [{"delta": {"content": "Hello"}}]} + events = e.feed_chunk(chunk) + # content_block_start + content_block_delta + assert len(events) == 2 + assert "content_block_start" in events[0] + assert '"type": "text"' in events[0] + delta = self._parse(events[1]) + assert delta["delta"]["type"] == "text_delta" + assert delta["delta"]["text"] == "Hello" + + def test_sequential_text_chunks_single_block(self): + e = AnthropicPassthroughEmitter() + e.start("msg_1", "m") + events1 = e.feed_chunk({"choices": [{"delta": {"content": "Hello"}}]}) + events2 = e.feed_chunk({"choices": [{"delta": {"content": " world"}}]}) + # First chunk opens the block, second only emits delta + assert len(events1) == 2 + assert len(events2) == 1 + assert self._parse(events2[0])["delta"]["text"] == " world" + + def test_tool_call_opens_tool_use_block(self): + e = AnthropicPassthroughEmitter() + e.start("msg_1", "m") + chunk = { + "choices": [ + { + "delta": { + "tool_calls": [ + { + "index": 0, + "id": "call_1", + "type": "function", + "function": {"name": "Bash", "arguments": ""}, + } + ] + } + } + ] + } + events = e.feed_chunk(chunk) + assert len(events) == 1 + parsed = self._parse(events[0]) + assert parsed["type"] == "content_block_start" + assert parsed["content_block"]["type"] == "tool_use" + assert parsed["content_block"]["id"] == "call_1" + assert parsed["content_block"]["name"] == "Bash" + + def test_tool_call_arguments_streamed_as_input_json_delta(self): + e = AnthropicPassthroughEmitter() + e.start("msg_1", "m") + # Open the tool call + e.feed_chunk( + { + "choices": [ + { + "delta": { + "tool_calls": [ + { + "index": 0, + "id": "c1", + "type": "function", + "function": {"name": "Bash", "arguments": ""}, + } + ] + } + } + ] + } + ) + # Stream argument fragments + events1 = e.feed_chunk( + { + "choices": [ + { + "delta": { + "tool_calls": [ + {"index": 0, "function": {"arguments": '{"cmd'}} + ] + } + } + ] + } + ) + events2 = e.feed_chunk( + { + "choices": [ + { + "delta": { + "tool_calls": [ + {"index": 0, "function": {"arguments": '": "ls"}'}} + ] + } + } + ] + } + ) + parsed1 = self._parse(events1[0]) + parsed2 = self._parse(events2[0]) + assert parsed1["delta"]["type"] == "input_json_delta" + assert parsed1["delta"]["partial_json"] == '{"cmd' + assert parsed2["delta"]["partial_json"] == '": "ls"}' + + def test_text_then_tool_closes_text_block(self): + e = AnthropicPassthroughEmitter() + e.start("msg_1", "m") + e.feed_chunk({"choices": [{"delta": {"content": "Let me check."}}]}) + events = e.feed_chunk( + { + "choices": [ + { + "delta": { + "tool_calls": [ + { + "index": 0, + "id": "c1", + "type": "function", + "function": {"name": "Bash", "arguments": ""}, + } + ] + } + } + ] + } + ) + # Should close text block and open tool_use block + assert "content_block_stop" in events[0] + assert "content_block_start" in events[1] + assert '"type": "tool_use"' in events[1] + + def test_finish_reason_tool_calls_sets_tool_use_stop(self): + e = AnthropicPassthroughEmitter() + e.start("msg_1", "m") + e.feed_chunk( + { + "choices": [ + { + "delta": { + "tool_calls": [ + { + "index": 0, + "id": "c1", + "type": "function", + "function": {"name": "Bash", "arguments": "{}"}, + } + ] + } + } + ] + } + ) + e.feed_chunk({"choices": [{"delta": {}, "finish_reason": "tool_calls"}]}) + events = e.finish() + delta_event = [ev for ev in events if "message_delta" in ev][0] + parsed = self._parse(delta_event) + assert parsed["delta"]["stop_reason"] == "tool_use" + + def test_finish_reason_stop_sets_end_turn(self): + e = AnthropicPassthroughEmitter() + e.start("msg_1", "m") + e.feed_chunk({"choices": [{"delta": {"content": "Hi"}}]}) + e.feed_chunk({"choices": [{"delta": {}, "finish_reason": "stop"}]}) + events = e.finish() + delta_event = [ev for ev in events if "message_delta" in ev][0] + parsed = self._parse(delta_event) + assert parsed["delta"]["stop_reason"] == "end_turn" + + def test_finish_reason_length_sets_max_tokens(self): + e = AnthropicPassthroughEmitter() + e.start("msg_1", "m") + e.feed_chunk({"choices": [{"delta": {"content": "Hi"}}]}) + e.feed_chunk({"choices": [{"delta": {}, "finish_reason": "length"}]}) + events = e.finish() + delta_event = [ev for ev in events if "message_delta" in ev][0] + parsed = self._parse(delta_event) + assert parsed["delta"]["stop_reason"] == "max_tokens" + + def test_finish_closes_current_block(self): + e = AnthropicPassthroughEmitter() + e.start("msg_1", "m") + e.feed_chunk({"choices": [{"delta": {"content": "Hi"}}]}) + events = e.finish() + assert "content_block_stop" in events[0] + assert "message_delta" in events[1] + assert "message_stop" in events[2] + + def test_usage_chunk_captured(self): + e = AnthropicPassthroughEmitter() + e.start("msg_1", "m") + e.feed_chunk({"choices": [{"delta": {"content": "Hi"}}]}) + e.feed_chunk( + { + "choices": [], + "usage": {"prompt_tokens": 10, "completion_tokens": 5}, + } + ) + events = e.finish() + delta_event = [ev for ev in events if "message_delta" in ev][0] + parsed = self._parse(delta_event) + assert parsed["usage"]["output_tokens"] == 5 + + def test_empty_chunk_returns_no_events(self): + e = AnthropicPassthroughEmitter() + e.start("msg_1", "m") + events = e.feed_chunk({"choices": []}) + assert events == [] + + def test_no_blocks_at_all_still_produces_valid_finish(self): + e = AnthropicPassthroughEmitter() + e.start("msg_1", "m") + events = e.finish() + # No content_block_stop because no block was opened + assert not any("content_block_stop" in ev for ev in events) + assert any("message_delta" in ev for ev in events) + assert any("message_stop" in ev for ev in events) + + def test_multiple_tool_calls_distinct_blocks(self): + e = AnthropicPassthroughEmitter() + e.start("msg_1", "m") + # First tool call + e.feed_chunk( + { + "choices": [ + { + "delta": { + "tool_calls": [ + { + "index": 0, + "id": "c1", + "type": "function", + "function": {"name": "Bash", "arguments": "{}"}, + } + ] + } + } + ] + } + ) + # Second tool call (different index) + events = e.feed_chunk( + { + "choices": [ + { + "delta": { + "tool_calls": [ + { + "index": 1, + "id": "c2", + "type": "function", + "function": {"name": "Read", "arguments": "{}"}, + } + ] + } + } + ] + } + ) + # Should close block 0, open block 1 + assert "content_block_stop" in events[0] + assert "content_block_start" in events[1] + parsed = self._parse(events[1]) + assert parsed["content_block"]["name"] == "Read" + assert parsed["content_block"]["id"] == "c2" diff --git a/studio/backend/tests/test_browse_folders_route.py b/studio/backend/tests/test_browse_folders_route.py new file mode 100644 index 0000000000..19a83987d3 --- /dev/null +++ b/studio/backend/tests/test_browse_folders_route.py @@ -0,0 +1,86 @@ +# 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 os +import sys +import types +from pathlib import Path + +import pytest +from fastapi import HTTPException + +# Keep this test runnable in lightweight environments where optional logging +# deps are not installed. +if "structlog" not in sys.modules: + + class _DummyLogger: + def __getattr__(self, _name): + return lambda *args, **kwargs: None + + sys.modules["structlog"] = types.SimpleNamespace( + BoundLogger = _DummyLogger, + get_logger = lambda *args, **kwargs: _DummyLogger(), + ) + +import routes.models as models_route + + +def test_resolve_browse_target_returns_allowed_directory(tmp_path): + allowed = tmp_path / "allowed" + target = allowed / "models" / "nested" + target.mkdir(parents = True) + + resolved = models_route._resolve_browse_target(str(target), [allowed]) + + assert resolved == target.resolve() + + +def test_resolve_browse_target_rejects_outside_allowlist(tmp_path): + allowed = tmp_path / "allowed" + disallowed = tmp_path / "disallowed" + allowed.mkdir() + disallowed.mkdir() + + with pytest.raises(HTTPException) as exc_info: + models_route._resolve_browse_target(str(disallowed), [allowed]) + + assert exc_info.value.status_code == 403 + + +def test_resolve_browse_target_rejects_file_path(tmp_path): + allowed = tmp_path / "allowed" + allowed.mkdir() + model_file = allowed / "model.gguf" + model_file.write_text("gguf") + + with pytest.raises(HTTPException) as exc_info: + models_route._resolve_browse_target(str(model_file), [allowed]) + + assert exc_info.value.status_code == 400 + + +def test_resolve_browse_target_allows_symlink_into_other_allowed_root(tmp_path): + home_root = tmp_path / "home" + scan_root = tmp_path / "scan" + target = scan_root / "nested" + home_root.mkdir() + target.mkdir(parents = True) + (home_root / "scan-link").symlink_to(scan_root, target_is_directory = True) + + resolved = models_route._resolve_browse_target( + str(home_root / "scan-link" / "nested"), + [home_root, scan_root], + ) + + assert resolved == target.resolve() + + +@pytest.mark.skipif(os.altsep is not None, reason = "POSIX-only path semantics") +def test_resolve_browse_target_allows_backslash_in_posix_segment(tmp_path): + allowed = tmp_path / "allowed" + target = allowed / r"dir\name" + target.mkdir(parents = True) + + resolved = models_route._resolve_browse_target(str(target), [allowed]) + + assert resolved == target.resolve() diff --git a/studio/backend/tests/test_cached_gguf_routes.py b/studio/backend/tests/test_cached_gguf_routes.py new file mode 100644 index 0000000000..05aae8fb75 --- /dev/null +++ b/studio/backend/tests/test_cached_gguf_routes.py @@ -0,0 +1,398 @@ +# SPDX-License-Identifier: AGPL-3.0-only +# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0 + +import asyncio +import sys +import types +from pathlib import Path +from types import SimpleNamespace + +# Keep this test runnable in lightweight environments where optional logging +# deps are not installed. +if "structlog" not in sys.modules: + + class _DummyLogger: + def __getattr__(self, _name): + return lambda *args, **kwargs: None + + sys.modules["structlog"] = types.SimpleNamespace( + BoundLogger = _DummyLogger, + get_logger = lambda *args, **kwargs: _DummyLogger(), + ) + +import routes.models as models_route + + +def _repo( + repo_id: str, + files: list[SimpleNamespace], + repo_path: Path, + *, + revisions: list[SimpleNamespace] | None = None, +) -> SimpleNamespace: + return SimpleNamespace( + repo_id = repo_id, + repo_type = "model", + repo_path = repo_path, + revisions = revisions or [SimpleNamespace(files = files)], + ) + + +def _file( + name: str, + size_on_disk: int, + *, + blob_path: str | None = None, +) -> SimpleNamespace: + return SimpleNamespace( + file_name = name, + size_on_disk = size_on_disk, + blob_path = blob_path, + ) + + +def test_iter_gguf_paths_matches_extension_case_insensitively(tmp_path): + nested = tmp_path / "snapshots" / "rev" + nested.mkdir(parents = True) + lower = nested / "Q4_K_M.gguf" + upper = nested / "Q8_0.GGUF" + other = nested / "README.md" + lower.write_text("a") + upper.write_text("b") + other.write_text("c") + + result = sorted(path.name for path in models_route._iter_gguf_paths(tmp_path)) + + assert result == ["Q4_K_M.gguf", "Q8_0.GGUF"] + + +def test_list_cached_gguf_includes_non_suffix_repo_when_cache_contains_gguf( + monkeypatch, tmp_path +): + repo = _repo( + "HauhauCS/Gemma-4-E4B-Uncensored-HauhauCS-Aggressive", + [_file("Q4_K_M.gguf", 5_000), _file("README.md", 10)], + tmp_path / "models--HauhauCS--Gemma", + ) + scan = SimpleNamespace(repos = [repo]) + + monkeypatch.setattr(models_route, "_all_hf_cache_scans", lambda: [scan]) + + result = asyncio.run(models_route.list_cached_gguf(current_subject = "test-user")) + + assert result["cached"] == [ + { + "repo_id": "HauhauCS/Gemma-4-E4B-Uncensored-HauhauCS-Aggressive", + "size_bytes": 5_000, + "cache_path": str(repo.repo_path), + } + ] + + +def test_list_cached_gguf_matches_extension_case_insensitively(monkeypatch, tmp_path): + repo = _repo( + "Org/Model-Without-Suffix", + [_file("Q8_0.GGUF", 7_000)], + tmp_path / "models--Org--Model-Without-Suffix", + ) + scan = SimpleNamespace(repos = [repo]) + + monkeypatch.setattr(models_route, "_all_hf_cache_scans", lambda: [scan]) + + result = asyncio.run(models_route.list_cached_gguf(current_subject = "test-user")) + + assert result["cached"] == [ + { + "repo_id": "Org/Model-Without-Suffix", + "size_bytes": 7_000, + "cache_path": str(repo.repo_path), + } + ] + + +def test_list_cached_gguf_skips_repos_without_positive_gguf_size(monkeypatch, tmp_path): + missing = _repo( + "Org/ReadmeOnly", + [_file("README.md", 10)], + tmp_path / "models--Org--ReadmeOnly", + ) + zero = _repo( + "Org/ZeroSize", + [_file("Q4_K_M.gguf", 0)], + tmp_path / "models--Org--ZeroSize", + ) + scan = SimpleNamespace(repos = [missing, zero]) + + monkeypatch.setattr(models_route, "_all_hf_cache_scans", lambda: [scan]) + + result = asyncio.run(models_route.list_cached_gguf(current_subject = "test-user")) + + assert result["cached"] == [] + + +def test_list_cached_gguf_keeps_largest_duplicate_repo_across_scans( + monkeypatch, tmp_path +): + smaller = _repo( + "Org/Dupe", + [_file("Q4_K_M.gguf", 2_000)], + tmp_path / "models--Org--Dupe-a", + ) + larger = _repo( + "org/dupe", + [_file("Q4_K_M.gguf", 5_000), _file("Q6_K.gguf", 1_000)], + tmp_path / "models--Org--Dupe-b", + ) + + monkeypatch.setattr( + models_route, + "_all_hf_cache_scans", + lambda: [ + SimpleNamespace(repos = [smaller]), + SimpleNamespace(repos = [larger]), + ], + ) + + result = asyncio.run(models_route.list_cached_gguf(current_subject = "test-user")) + + assert result["cached"] == [ + { + "repo_id": "org/dupe", + "size_bytes": 6_000, + "cache_path": str(larger.repo_path), + } + ] + + +def test_list_cached_gguf_dedupes_shared_blobs_across_revisions(monkeypatch, tmp_path): + shared = "blobs/shared-q4" + repo = _repo( + "Org/SharedBlobRepo", + [], + tmp_path / "models--Org--SharedBlobRepo", + revisions = [ + SimpleNamespace(files = [_file("Q4_K_M.gguf", 5_000, blob_path = shared)]), + SimpleNamespace(files = [_file("Q4_K_M.gguf", 5_000, blob_path = shared)]), + ], + ) + + monkeypatch.setattr( + models_route, + "_all_hf_cache_scans", + lambda: [SimpleNamespace(repos = [repo])], + ) + + result = asyncio.run(models_route.list_cached_gguf(current_subject = "test-user")) + + assert result["cached"] == [ + { + "repo_id": "Org/SharedBlobRepo", + "size_bytes": 5_000, + "cache_path": str(repo.repo_path), + } + ] + + +def test_list_cached_models_skips_non_suffix_repo_when_gguf_files_exist( + monkeypatch, tmp_path +): + mixed = _repo( + "Org/MixedRepo", + [ + _file("Q4_K_M.gguf", 5_000), + _file("model.safetensors", 10_000), + ], + tmp_path / "models--Org--MixedRepo", + ) + + monkeypatch.setattr( + models_route, + "_all_hf_cache_scans", + lambda: [SimpleNamespace(repos = [mixed])], + ) + + result = asyncio.run(models_route.list_cached_models(current_subject = "test-user")) + + assert result["cached"] == [] + + +def test_list_cached_gguf_includes_mixed_repo_with_gguf_and_safetensors( + monkeypatch, tmp_path +): + """Mirror of the _skips_ test: the mixed repo should still surface in + cached-gguf so the picker can show it as a GGUF download.""" + mixed = _repo( + "Org/MixedRepo", + [ + _file("Q4_K_M.gguf", 5_000), + _file("model.safetensors", 10_000), + ], + tmp_path / "models--Org--MixedRepo", + ) + + monkeypatch.setattr( + models_route, + "_all_hf_cache_scans", + lambda: [SimpleNamespace(repos = [mixed])], + ) + + result = asyncio.run(models_route.list_cached_gguf(current_subject = "test-user")) + + assert result["cached"] == [ + { + "repo_id": "Org/MixedRepo", + "size_bytes": 5_000, + "cache_path": str(mixed.repo_path), + } + ] + + +def test_list_cached_gguf_handles_none_size_on_disk(monkeypatch, tmp_path): + """A partial/interrupted GGUF download has ``size_on_disk = None``. The + route must treat the unknown bytes as zero instead of raising TypeError + out of ``sum()`` and wiping the entire response.""" + partial = _repo( + "Org/PartialDownload", + [_file("Q4_K_M.gguf", None), _file("Q6_K.gguf", 5_000)], + tmp_path / "models--Org--PartialDownload", + ) + + monkeypatch.setattr( + models_route, + "_all_hf_cache_scans", + lambda: [SimpleNamespace(repos = [partial])], + ) + + result = asyncio.run(models_route.list_cached_gguf(current_subject = "test-user")) + + assert result["cached"] == [ + { + "repo_id": "Org/PartialDownload", + "size_bytes": 5_000, + "cache_path": str(partial.repo_path), + } + ] + + +def test_list_cached_gguf_skips_malformed_repo_without_wiping_response( + monkeypatch, tmp_path +): + """One repo raising during classification must not poison the response + for every other repo in the scan.""" + + class _ExplodingRepo: + repo_id = "Org/Broken" + repo_type = "model" + repo_path = tmp_path / "models--Org--Broken" + + @property + def revisions(self): + raise RuntimeError("boom") + + healthy = _repo( + "Org/Healthy", + [_file("Q4_K_M.gguf", 5_000)], + tmp_path / "models--Org--Healthy", + ) + + monkeypatch.setattr( + models_route, + "_all_hf_cache_scans", + lambda: [SimpleNamespace(repos = [_ExplodingRepo(), healthy])], + ) + + result = asyncio.run(models_route.list_cached_gguf(current_subject = "test-user")) + + assert result["cached"] == [ + { + "repo_id": "Org/Healthy", + "size_bytes": 5_000, + "cache_path": str(healthy.repo_path), + } + ] + + +def test_list_cached_gguf_skips_repo_with_only_mmproj_gguf(monkeypatch, tmp_path): + """A repo whose only ``.gguf`` artifact is an mmproj vision adapter + must not be classified as a GGUF repo: the variant selector filters + mmproj out and the picker would otherwise show zero variants.""" + mmproj_only = _repo( + "Org/MmprojOnly", + [ + _file("mmproj-Q8_0.gguf", 5_000), + _file("model.safetensors", 10_000), + ], + tmp_path / "models--Org--MmprojOnly", + ) + + monkeypatch.setattr( + models_route, + "_all_hf_cache_scans", + lambda: [SimpleNamespace(repos = [mmproj_only])], + ) + + result = asyncio.run(models_route.list_cached_gguf(current_subject = "test-user")) + + assert result["cached"] == [] + + +def test_list_cached_models_includes_repo_with_only_mmproj_gguf(monkeypatch, tmp_path): + """Mirror of the cached-gguf skip: a safetensors repo with an + auxiliary mmproj vision adapter must still surface in cached-models + so the user can load it as a normal model.""" + mmproj_aux = _repo( + "Org/MmprojAux", + [ + _file("mmproj-Q8_0.gguf", 5_000), + _file("model.safetensors", 10_000), + ], + tmp_path / "models--Org--MmprojAux", + ) + + monkeypatch.setattr( + models_route, + "_all_hf_cache_scans", + lambda: [SimpleNamespace(repos = [mmproj_aux])], + ) + + result = asyncio.run(models_route.list_cached_models(current_subject = "test-user")) + + assert result["cached"] == [ + { + "repo_id": "Org/MmprojAux", + "size_bytes": 15_000, + } + ] + + +def test_list_cached_gguf_includes_vision_repo_with_main_gguf_and_mmproj( + monkeypatch, tmp_path +): + """A vision-capable GGUF repo (main weight + mmproj adapter) is still + a GGUF repo. The reported size is the main weight size; mmproj is + excluded from the GGUF-size accounting because it is filtered out at + classification time.""" + vision_repo = _repo( + "Org/VisionGguf", + [ + _file("Q4_K_M.gguf", 5_000), + _file("mmproj-Q8_0.gguf", 1_000), + ], + tmp_path / "models--Org--VisionGguf", + ) + + monkeypatch.setattr( + models_route, + "_all_hf_cache_scans", + lambda: [SimpleNamespace(repos = [vision_repo])], + ) + + result = asyncio.run(models_route.list_cached_gguf(current_subject = "test-user")) + + assert result["cached"] == [ + { + "repo_id": "Org/VisionGguf", + "size_bytes": 5_000, + "cache_path": str(vision_repo.repo_path), + } + ] diff --git a/studio/backend/tests/test_export_log_cursor.py b/studio/backend/tests/test_export_log_cursor.py new file mode 100644 index 0000000000..734ca522c9 --- /dev/null +++ b/studio/backend/tests/test_export_log_cursor.py @@ -0,0 +1,179 @@ +# 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 the export log ring-buffer cursor semantics. + +Context: the live export log SSE stream has a race where the frontend +opens the SSE connection AFTER the POST that starts the export. Any +lines the worker subprocess emits during the gap between POST and SSE +connect get buffered with seqs 1..k, and then the SSE default cursor +`get_current_log_seq()` returns k -- so lines 1..k are forever +unreachable to that client. + +Fix: `clear_logs()` snapshots the pre-run seq into `_run_start_seq` +(exposed via `get_run_start_seq()`), and `routes/export.py` defaults +the SSE cursor to that snapshot instead of the current seq. Every line +appended during the current run has seq strictly greater than the +snapshot, so the client sees the full run regardless of when it +connects. + +These tests exercise the orchestrator-side contract only (no +subprocess, no FastAPI, no frontend). The routes-level integration +with get_run_start_seq() is a one-line edit covered by manual testing +and the frontend build. +""" + +from __future__ import annotations + +import sys +import types +from pathlib import Path + +import pytest + + +# Backend root on sys.path so `from core.export.orchestrator import ...` +# and friends resolve without the studio app bootstrap. +_BACKEND_DIR = Path(__file__).resolve().parent.parent +if str(_BACKEND_DIR) not in sys.path: + sys.path.insert(0, str(_BACKEND_DIR)) + +# ExportOrchestrator imports structlog and a few heavy modules at the +# top of orchestrator.py. Stub the ones we don't need in these unit +# tests so the import succeeds on machines without the full studio +# venv. +_loggers_stub = types.ModuleType("loggers") +_loggers_stub.get_logger = lambda name: __import__("logging").getLogger(name) +sys.modules.setdefault("loggers", _loggers_stub) + +# structlog is only used for a module-level import; a bare stub is +# enough because we never call into it in these tests. +sys.modules.setdefault("structlog", types.ModuleType("structlog")) + +# utils.paths.outputs_root is only called inside scan_checkpoints which +# we don't hit in these tests. Provide a stub module so the top-level +# import in orchestrator.py resolves. +_utils_pkg = types.ModuleType("utils") +_utils_pkg.__path__ = [] # mark as package +_utils_paths_stub = types.ModuleType("utils.paths") +_utils_paths_stub.outputs_root = lambda: Path("/tmp") +sys.modules.setdefault("utils", _utils_pkg) +sys.modules.setdefault("utils.paths", _utils_paths_stub) + + +@pytest.fixture +def orchestrator(): + """Fresh ExportOrchestrator with only the log-buffer state exercised.""" + from core.export.orchestrator import ExportOrchestrator + + return ExportOrchestrator() + + +def _append(orch, line: str, stream: str = "stdout") -> None: + """Shortcut for simulating a worker log message.""" + orch._append_log({"type": "log", "stream": stream, "line": line, "ts": 0.0}) + + +# --------------------------------------------------------------------------- +# clear_logs() semantics +# --------------------------------------------------------------------------- + + +def test_run_start_seq_is_zero_before_any_logs(orchestrator) -> None: + """A brand-new orchestrator must report run_start_seq == 0 so a + first SSE connection picks up every line from seq 1 onward.""" + assert orchestrator.get_run_start_seq() == 0 + + +def test_clear_logs_snapshots_current_seq(orchestrator) -> None: + """clear_logs() must capture _log_seq BEFORE clearing the buffer, + so subsequent runs can anchor their SSE cursor at the snapshot.""" + _append(orchestrator, "old run line 1") + _append(orchestrator, "old run line 2") + _append(orchestrator, "old run line 3") + assert orchestrator.get_current_log_seq() == 3 + + orchestrator.clear_logs() + + assert orchestrator.get_run_start_seq() == 3 + assert orchestrator.get_current_log_seq() == 3 # seq counter preserved + + +# --------------------------------------------------------------------------- +# Race regression: SSE connects AFTER lines have been emitted +# --------------------------------------------------------------------------- + + +def test_sse_default_cursor_catches_all_current_run_lines(orchestrator) -> None: + """Simulate the POST-then-SSE race: worker starts emitting lines + immediately after clear_logs(), SSE connects several lines later. + Using get_run_start_seq() as the default cursor MUST return every + line emitted since clear_logs() ran. + + Pre-fix, the SSE defaulted to get_current_log_seq() at connect + time, which would return the last-seen seq and miss lines N+1..M. + """ + # Previous run leaves some buffered lines. + _append(orchestrator, "previous run line A") + _append(orchestrator, "previous run line B") + + # New run starts: orchestrator clears the buffer and snapshots seq. + orchestrator.clear_logs() + run_start = orchestrator.get_run_start_seq() + + # Worker emits early lines BEFORE the SSE connects. + _append(orchestrator, "Importing Unsloth...") + _append(orchestrator, "Loading checkpoint: /foo/bar") + _append(orchestrator, "Starting export...") + + # SSE connects now and asks "give me everything after the run + # start cursor". + entries, new_cursor = orchestrator.get_logs_since(run_start) + + # All three early lines must be present. Pre-fix this was []. + lines = [e["line"] for e in entries] + assert lines == [ + "Importing Unsloth...", + "Loading checkpoint: /foo/bar", + "Starting export...", + ] + assert new_cursor == entries[-1]["seq"] + + +def test_sse_default_cursor_excludes_previous_run(orchestrator) -> None: + """After clear_logs(), lines from the PREVIOUS run must not leak + into the new run's SSE stream. Pre-fix this worked correctly + (clear_logs cleared the deque); the fix must preserve it. + """ + _append(orchestrator, "previous run line 1") + _append(orchestrator, "previous run line 2") + _append(orchestrator, "previous run line 3") + assert orchestrator.get_current_log_seq() == 3 + + orchestrator.clear_logs() + run_start = orchestrator.get_run_start_seq() + + _append(orchestrator, "new run line") + + entries, _ = orchestrator.get_logs_since(run_start) + assert [e["line"] for e in entries] == ["new run line"] + + +def test_clear_logs_twice_advances_run_start(orchestrator) -> None: + """Back-to-back clear_logs() calls (e.g. cleanup -> load -> + export in the same dialog session) must each re-anchor run_start + at the current seq, so successive runs each start with a fresh + low-water mark.""" + _append(orchestrator, "run 1 line a") + _append(orchestrator, "run 1 line b") + + orchestrator.clear_logs() + assert orchestrator.get_run_start_seq() == 2 + + _append(orchestrator, "run 2 line a") + _append(orchestrator, "run 2 line b") + _append(orchestrator, "run 2 line c") + + orchestrator.clear_logs() + assert orchestrator.get_run_start_seq() == 5 diff --git a/studio/backend/tests/test_llama_cpp_cache_aware_disk_check.py b/studio/backend/tests/test_llama_cpp_cache_aware_disk_check.py new file mode 100644 index 0000000000..255c04a956 --- /dev/null +++ b/studio/backend/tests/test_llama_cpp_cache_aware_disk_check.py @@ -0,0 +1,243 @@ +# 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 cache-aware disk-space preflight in +``LlamaCppBackend.load_model``. + +The preflight used to compare the repo's total GGUF download size against +free disk without accounting for bytes already present in the Hugging +Face cache. That made re-loading a cached large model (e.g. +``unsloth/MiniMax-M2.7-GGUF`` at 131 GB) fail cold whenever free disk was +below the full weight footprint, even though nothing needed +downloading. + +These tests exercise the preflight arithmetic in isolation by driving +``get_paths_info`` and ``try_to_load_from_cache`` through ``mock.patch``. +No network, GPU, or subprocess use. + +Cross-platform: Linux, macOS, Windows, WSL. +""" + +from __future__ import annotations + +import sys +import tempfile +import types as _types +from pathlib import Path +from unittest.mock import patch + +import pytest + +# --------------------------------------------------------------------------- +# Stub heavy / unavailable external dependencies before importing the +# module under test. Same pattern as test_kv_cache_estimation.py. +# --------------------------------------------------------------------------- + +_BACKEND_DIR = str(Path(__file__).resolve().parent.parent) +if _BACKEND_DIR not in sys.path: + sys.path.insert(0, _BACKEND_DIR) + +# loggers +_loggers_stub = _types.ModuleType("loggers") +_loggers_stub.get_logger = lambda name: __import__("logging").getLogger(name) +sys.modules.setdefault("loggers", _loggers_stub) + +# structlog +_structlog_stub = _types.ModuleType("structlog") +sys.modules.setdefault("structlog", _structlog_stub) + +# httpx +_httpx_stub = _types.ModuleType("httpx") +for _exc_name in ( + "ConnectError", + "TimeoutException", + "ReadTimeout", + "ReadError", + "RemoteProtocolError", + "CloseError", +): + setattr(_httpx_stub, _exc_name, type(_exc_name, (Exception,), {})) + + +class _FakeTimeout: + def __init__(self, *a, **kw): + pass + + +_httpx_stub.Timeout = _FakeTimeout +_httpx_stub.Client = type( + "Client", + (), + { + "__init__": lambda self, **kw: None, + "__enter__": lambda self: self, + "__exit__": lambda self, *a: None, + }, +) +sys.modules.setdefault("httpx", _httpx_stub) + + +# --------------------------------------------------------------------------- +# Helpers +# --------------------------------------------------------------------------- + +GIB = 1024**3 + + +class _FakePathInfo: + """Mimics huggingface_hub's RepoFile-ish return type from get_paths_info.""" + + def __init__(self, path: str, size: int): + self.path = path + self.size = size + + +def _preflight( + repo_files, + cached_files, + free_bytes, + hf_repo = "unsloth/Example-GGUF", + hf_token = None, +): + """Run the preflight arithmetic as written in llama_cpp.py and return + the decision outcome as a dict. + + ``repo_files``: list of (filename, remote_bytes). + ``cached_files``: dict {filename: on_disk_bytes} for files already in cache. + ``free_bytes``: value returned by shutil.disk_usage(cache_dir).free. + """ + import os + import shutil + + path_infos = [_FakePathInfo(name, size) for name, size in repo_files] + + with tempfile.TemporaryDirectory() as tmp: + # Create SPARSE files for the cached ones so os.path.exists / + # os.path.getsize pass without actually allocating bytes on disk. + # This is critical when simulating multi-GB models. + cache_paths = {} + for name, sz in cached_files.items(): + p = Path(tmp) / name.replace("/", "_") + with open(p, "wb") as fh: + if sz > 0: + fh.truncate(sz) # sparse allocation: no data blocks written + cache_paths[name] = str(p) + + def fake_try_to_load_from_cache(repo_id, filename): + return cache_paths.get(filename) + + # Mirror the same variable names and control flow as the real code + # so behavioral drift is caught immediately. + total_bytes = sum((p.size or 0) for p in path_infos) + already_cached_bytes = 0 + for p in path_infos: + if not p.size: + continue + cached_path = fake_try_to_load_from_cache(hf_repo, p.path) + if isinstance(cached_path, str) and os.path.exists(cached_path): + try: + on_disk = os.path.getsize(cached_path) + except OSError: + on_disk = 0 + if on_disk >= p.size: + already_cached_bytes += p.size + + total_download_bytes = max(0, total_bytes - already_cached_bytes) + needed_download = total_download_bytes > free_bytes + return { + "total_bytes": total_bytes, + "already_cached_bytes": already_cached_bytes, + "total_download_bytes": total_download_bytes, + "would_raise_disk_error": (needed_download and total_download_bytes > 0), + } + + +# --------------------------------------------------------------------------- +# Tests +# --------------------------------------------------------------------------- + + +class TestCacheAwarePreflight: + def test_fully_cached_model_does_not_require_disk(self): + """The MiniMax case: 131 GB weights cached, only 36 GB free. + Preflight must not raise.""" + shards = [(f"UD-Q4_K_XL/shard-{i}.gguf", 35 * GIB) for i in range(4)] + cached = {name: size for name, size in shards} + out = _preflight( + repo_files = shards, + cached_files = cached, + free_bytes = 36 * GIB, + ) + assert out["total_download_bytes"] == 0 + assert out["already_cached_bytes"] == 140 * GIB + assert out["would_raise_disk_error"] is False + + def test_partial_cache_only_counts_remaining_bytes(self): + """Two of four shards cached: preflight against remaining 70 GB.""" + shards = [(f"UD-Q4_K_XL/shard-{i}.gguf", 35 * GIB) for i in range(4)] + cached = { + shards[0][0]: shards[0][1], + shards[1][0]: shards[1][1], + } + out = _preflight( + repo_files = shards, + cached_files = cached, + free_bytes = 80 * GIB, + ) + assert out["already_cached_bytes"] == 70 * GIB + assert out["total_download_bytes"] == 70 * GIB + assert out["would_raise_disk_error"] is False + + def test_partial_cache_insufficient_disk_for_rest_still_raises(self): + """Two of four shards cached; remaining 70 GB still bigger than + free disk -> preflight correctly wants to raise.""" + shards = [(f"UD-Q4_K_XL/shard-{i}.gguf", 35 * GIB) for i in range(4)] + cached = { + shards[0][0]: shards[0][1], + shards[1][0]: shards[1][1], + } + out = _preflight( + repo_files = shards, + cached_files = cached, + free_bytes = 50 * GIB, + ) + assert out["total_download_bytes"] == 70 * GIB + assert out["would_raise_disk_error"] is True + + def test_nothing_cached_preserves_existing_behavior(self): + """Cold-cache path still compares full download vs free disk.""" + shards = [("UD-Q4_K_XL/shard-0.gguf", 40 * GIB)] + out = _preflight( + repo_files = shards, + cached_files = {}, + free_bytes = 50 * GIB, + ) + assert out["already_cached_bytes"] == 0 + assert out["total_download_bytes"] == 40 * GIB + assert out["would_raise_disk_error"] is False + + def test_incomplete_cached_blob_is_not_credited(self): + """A partial file on disk (e.g. interrupted download) is not + counted as cached -- we still require bytes for it.""" + shards = [("UD-Q4_K_XL/shard-0.gguf", 40 * GIB)] + partial = {"UD-Q4_K_XL/shard-0.gguf": 10 * GIB} + out = _preflight( + repo_files = shards, + cached_files = partial, + free_bytes = 50 * GIB, + ) + assert out["already_cached_bytes"] == 0 + assert out["total_download_bytes"] == 40 * GIB + assert out["would_raise_disk_error"] is False + + def test_zero_size_path_infos_do_not_crash(self): + """A path_info with size=0 should not be credited or break the + arithmetic.""" + shards = [("mmproj.gguf", 0), ("UD-Q4_K_XL/shard-0.gguf", 40 * GIB)] + out = _preflight( + repo_files = shards, + cached_files = {}, + free_bytes = 50 * GIB, + ) + assert out["already_cached_bytes"] == 0 + assert out["total_bytes"] == 40 * GIB diff --git a/studio/backend/tests/test_llama_cpp_context_fit.py b/studio/backend/tests/test_llama_cpp_context_fit.py new file mode 100644 index 0000000000..f498655347 --- /dev/null +++ b/studio/backend/tests/test_llama_cpp_context_fit.py @@ -0,0 +1,389 @@ +# 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 GGUF load-time context auto-fit decision. + +Guards two regressions in ``LlamaCppBackend.load_model``: + +1. **Auto mode on weights-exceed-VRAM** (``n_ctx == 0``): when the model + weights alone exceed 90% of every GPU subset's free memory, the + auto-pick loop used to exit without matching, leaving + ``effective_ctx`` at the model's native context (e.g. 196608 for + MiniMax-M2.7). The intended default per Studio's UI spec is 4096 so + the slider lands on a usable value; the user can still drag higher + and trigger ``--fit on`` with a warning. + +2. **Explicit ctx silently shrunk when KV overflows**: with fittable + weights but a requested ctx whose KV cache pushes total memory over + 90% of VRAM, the old code binary-searched a smaller ctx and emitted + ``-c -ngl -1`` without informing the caller. The UI had + already surfaced its "might be slower" warning and expects the user's + explicit ctx to be honored with ``--fit on`` flexing ``-ngl`` instead. + +Tests avoid GPU probing, subprocess spawning, and GGUF I/O by driving the +post-metadata decision block directly against a stubbed instance. + +Requires no GPU, network, or external libraries beyond pytest. +Cross-platform: Linux, macOS, Windows, WSL. +""" + +from __future__ import annotations + +import sys +import types as _types +from pathlib import Path + +import pytest + +# --------------------------------------------------------------------------- +# Stub heavy / unavailable external dependencies before importing the +# module under test. Same pattern as test_kv_cache_estimation.py. +# --------------------------------------------------------------------------- + +_BACKEND_DIR = str(Path(__file__).resolve().parent.parent) +if _BACKEND_DIR not in sys.path: + sys.path.insert(0, _BACKEND_DIR) + +# loggers +_loggers_stub = _types.ModuleType("loggers") +_loggers_stub.get_logger = lambda name: __import__("logging").getLogger(name) +sys.modules.setdefault("loggers", _loggers_stub) + +# structlog +_structlog_stub = _types.ModuleType("structlog") +sys.modules.setdefault("structlog", _structlog_stub) + +# httpx +_httpx_stub = _types.ModuleType("httpx") +for _exc_name in ( + "ConnectError", + "TimeoutException", + "ReadTimeout", + "ReadError", + "RemoteProtocolError", + "CloseError", +): + setattr(_httpx_stub, _exc_name, type(_exc_name, (Exception,), {})) + + +class _FakeTimeout: + def __init__(self, *a, **kw): + pass + + +_httpx_stub.Timeout = _FakeTimeout +_httpx_stub.Client = type( + "Client", + (), + { + "__init__": lambda self, **kw: None, + "__enter__": lambda self: self, + "__exit__": lambda self, *a: None, + }, +) +sys.modules.setdefault("httpx", _httpx_stub) + +from core.inference.llama_cpp import LlamaCppBackend + + +# --------------------------------------------------------------------------- +# Helpers +# --------------------------------------------------------------------------- + +GIB = 1024**3 +FALLBACK_CTX = 4096 + + +def _make_backend( + native_ctx = 131072, + n_layers = 80, + n_kv_heads = 8, + n_heads = 64, + kv_key_length = 128, + kv_value_length = 128, +): + """Create a LlamaCppBackend instance with GGUF metadata fields set and + the helpers used by the decision block stubbed out.""" + inst = LlamaCppBackend.__new__(LlamaCppBackend) + inst._context_length = native_ctx + inst._n_layers = n_layers + inst._n_kv_heads = n_kv_heads + inst._n_heads = n_heads + inst._embedding_length = 8192 + inst._kv_key_length = kv_key_length + inst._kv_value_length = kv_value_length + inst._kv_lora_rank = None + inst._sliding_window = None + inst._ssm_inner_size = None + inst._full_attention_interval = None + inst._key_length_mla = None + return inst + + +def _drive( + n_ctx, + model_gib, + gpus, + native_ctx = 131072, + kv_per_token_bytes = 325_000, + can_estimate_kv = True, +): + """Drive the post-metadata portion of load_model with stubbed inputs. + + Mirrors the decision block at llama_cpp.py:1137-1296 so we can assert + the command that would be built, without subprocesses or GPU probes. + """ + inst = _make_backend(native_ctx = native_ctx) + model_size = int(model_gib * GIB) + cache_type_kv = None + + def fake_estimate(n_ctx_, _type = None): + return 0 if n_ctx_ <= 0 else n_ctx_ * kv_per_token_bytes + + inst._estimate_kv_cache_bytes = fake_estimate + inst._can_estimate_kv = lambda: can_estimate_kv + + context_length = inst._context_length + + effective_ctx = n_ctx if n_ctx > 0 else (context_length or 0) + max_available_ctx = context_length or effective_ctx + if n_ctx > 0: + effective_ctx = n_ctx + elif context_length is not None: + effective_ctx = context_length + else: + effective_ctx = 0 + original_ctx = effective_ctx + max_available_ctx = context_length or effective_ctx + + gpu_indices, use_fit = None, True + explicit_ctx = n_ctx > 0 + + if gpus and inst._can_estimate_kv() and effective_ctx > 0: + native_ctx_for_cap = context_length or effective_ctx + if native_ctx_for_cap > 0: + ranked_for_cap = sorted(gpus, key = lambda g: g[1], reverse = True) + best_cap = 0 + for n_gpus in range(1, len(ranked_for_cap) + 1): + subset = ranked_for_cap[:n_gpus] + pool_mib = sum(free for _, free in subset) + capped = inst._fit_context_to_vram( + native_ctx_for_cap, + pool_mib, + model_size, + cache_type_kv, + ) + kv = inst._estimate_kv_cache_bytes(capped, cache_type_kv) + total_mib = (model_size + kv) / (1024 * 1024) + if total_mib <= pool_mib * 0.90: + best_cap = max(best_cap, capped) + if best_cap > 0: + max_available_ctx = best_cap + + if explicit_ctx: + requested_total = model_size + inst._estimate_kv_cache_bytes( + effective_ctx, cache_type_kv + ) + gpu_indices, use_fit = inst._select_gpus(requested_total, gpus) + else: + ranked = sorted(gpus, key = lambda g: g[1], reverse = True) + matched = False + for n_gpus in range(1, len(ranked) + 1): + subset = ranked[:n_gpus] + pool_mib = sum(free for _, free in subset) + capped = inst._fit_context_to_vram( + effective_ctx, + pool_mib, + model_size, + cache_type_kv, + ) + kv = inst._estimate_kv_cache_bytes(capped, cache_type_kv) + total_mib = (model_size + kv) / (1024 * 1024) + if total_mib <= pool_mib * 0.90: + effective_ctx = capped + gpu_indices = sorted(idx for idx, _ in subset) + use_fit = False + matched = True + break + if not matched: + effective_ctx = min(FALLBACK_CTX, effective_ctx) + elif gpus: + gpu_indices, use_fit = inst._select_gpus(model_size, gpus) + if use_fit and not explicit_ctx: + effective_ctx = ( + min(FALLBACK_CTX, effective_ctx) if effective_ctx > 0 else FALLBACK_CTX + ) + + return { + "c_arg": effective_ctx if effective_ctx > 0 else 0, + "use_fit": use_fit, + "gpu_indices": gpu_indices, + "max_available_ctx": max_available_ctx, + "original_ctx": original_ctx, + } + + +# --------------------------------------------------------------------------- +# Auto mode, model weights exceed VRAM (Bug A guard) +# --------------------------------------------------------------------------- + + +class TestAutoModeWeightsExceedVRAM: + """``n_ctx == 0`` on a model whose weights don't fit anywhere.""" + + def test_minimax_like_single_gpu(self): + plan = _drive( + n_ctx = 0, + model_gib = 131, + gpus = [(0, 97_000)], + native_ctx = 196608, + ) + assert plan["c_arg"] == FALLBACK_CTX + assert plan["use_fit"] is True + assert plan["gpu_indices"] is None + # UI slider ceiling stays at native: user can still drag higher + # and get the "might be slower" path. + assert plan["max_available_ctx"] == 196608 + + def test_multi_gpu_all_subsets_fail(self): + plan = _drive( + n_ctx = 0, + model_gib = 400, + gpus = [(0, 80_000), (1, 80_000), (2, 80_000), (3, 80_000)], + native_ctx = 131072, + ) + assert plan["c_arg"] == FALLBACK_CTX + assert plan["use_fit"] is True + assert plan["gpu_indices"] is None + + def test_no_kv_metadata_auto(self): + """File-size-only fallback path also defaults to 4096.""" + plan = _drive( + n_ctx = 0, + model_gib = 131, + gpus = [(0, 97_000)], + native_ctx = 196608, + can_estimate_kv = False, + ) + assert plan["c_arg"] == FALLBACK_CTX + assert plan["use_fit"] is True + + +# --------------------------------------------------------------------------- +# Explicit ctx, KV overflows fittable weights (Bug B guard) +# --------------------------------------------------------------------------- + + +class TestExplicitCtxRespectsUser: + """``n_ctx > 0`` must never be silently shrunk.""" + + def test_fittable_weights_oversized_kv(self): + # 8 GB weights + 131k ctx KV on 24 GB VRAM. + # Budget = 21.6 GB, KV at 131k >> 13.6 GB remaining, so + # _select_gpus flips use_fit=True. + plan = _drive( + n_ctx = 131072, + model_gib = 8, + gpus = [(0, 24_000)], + native_ctx = 131072, + ) + assert plan["c_arg"] == 131072 + assert plan["use_fit"] is True + assert plan["gpu_indices"] is None + + def test_explicit_that_fits_uses_ngl(self): + plan = _drive( + n_ctx = 8192, + model_gib = 8, + gpus = [(0, 24_000)], + native_ctx = 131072, + ) + assert plan["c_arg"] == 8192 + assert plan["use_fit"] is False + assert plan["gpu_indices"] == [0] + + def test_explicit_on_weights_exceed_vram(self): + # User drags the slider to 32k on a too-big model: honored. + plan = _drive( + n_ctx = 32768, + model_gib = 131, + gpus = [(0, 97_000)], + native_ctx = 196608, + ) + assert plan["c_arg"] == 32768 + assert plan["use_fit"] is True + + def test_explicit_at_fallback_on_too_big(self): + plan = _drive( + n_ctx = FALLBACK_CTX, + model_gib = 131, + gpus = [(0, 97_000)], + native_ctx = 196608, + ) + assert plan["c_arg"] == FALLBACK_CTX + assert plan["use_fit"] is True + + def test_explicit_below_floor_honored(self): + # 2048 is below --fit-ctx default; still honored since user set it. + plan = _drive( + n_ctx = 2048, + model_gib = 8, + gpus = [(0, 24_000)], + ) + assert plan["c_arg"] == 2048 + + +# --------------------------------------------------------------------------- +# Non-regression: fittable + auto still auto-picks largest fitting ctx +# --------------------------------------------------------------------------- + + +class TestFittableAutoPickRegressions: + def test_small_model_one_gpu(self): + plan = _drive( + n_ctx = 0, + model_gib = 8, + gpus = [(0, 24_000)], + native_ctx = 131072, + kv_per_token_bytes = 8192, + ) + assert plan["use_fit"] is False + assert plan["gpu_indices"] == [0] + assert plan["c_arg"] > FALLBACK_CTX + + def test_medium_model_needs_multi_gpu(self): + plan = _drive( + n_ctx = 0, + model_gib = 60, + gpus = [(0, 40_000), (1, 40_000)], + native_ctx = 131072, + kv_per_token_bytes = 8192, + ) + assert plan["use_fit"] is False + assert plan["gpu_indices"] == [0, 1] + + def test_no_kv_metadata_fittable_auto(self): + plan = _drive( + n_ctx = 0, + model_gib = 8, + gpus = [(0, 24_000)], + native_ctx = 131072, + can_estimate_kv = False, + ) + assert plan["use_fit"] is False + assert plan["gpu_indices"] == [0] + + +# --------------------------------------------------------------------------- +# Platform-agnostic input shape +# --------------------------------------------------------------------------- + + +@pytest.mark.parametrize("platform_tag", ["linux", "windows", "mac", "rocm"]) +def test_identical_decision_across_platforms(platform_tag): + """The decision function takes ``[(gpu_idx, free_mib), ...]`` regardless + of how upstream (nvidia-smi / nvidia-smi.exe / Metal / rocm-smi) produced + it. Identical inputs must yield identical plans.""" + plan_a = _drive(n_ctx = 0, model_gib = 8, gpus = [(0, 24_000)]) + plan_b = _drive(n_ctx = 0, model_gib = 8, gpus = [(0, 24_000)]) + assert plan_a == plan_b, platform_tag diff --git a/studio/backend/tests/test_llama_cpp_load_progress.py b/studio/backend/tests/test_llama_cpp_load_progress.py new file mode 100644 index 0000000000..f46751b798 --- /dev/null +++ b/studio/backend/tests/test_llama_cpp_load_progress.py @@ -0,0 +1,258 @@ +# 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 ``LlamaCppBackend.load_progress()``. + +The chat settings flow and the training overlay both show a generic +"Starting model..." spinner during the window after a GGUF download +finishes and before llama-server reports healthy. For small models +that window is a second or two and nobody notices. For large MoE GGUFs +(MiniMax-M2.7, Qwen3.5-397B-A17B, etc.) the llama-server process spends +minutes in kernel state D, paging tens or hundreds of GB of shards +into the page cache. The UI has no way to show a real progress bar, +rate, or ETA during that window. + +``load_progress()`` samples ``/proc//status VmRSS`` (what the +kernel has actually paged in) against the total shard file size on +disk, so the frontend can render a real bar plus rate/ETA. This +module pins that contract: + + * returns ``None`` when no load is in flight + * returns ``{"phase": "mmap", ...}`` while the subprocess is alive + but ``_healthy`` is False + * returns ``{"phase": "ready", ...}`` once ``_healthy`` flips + * ``bytes_total`` is derived from the resolved on-disk path + (which the paired fix assigns to ``self._gguf_path`` on both the + local-GGUF and HF-download code paths) + * ``bytes_loaded`` is VmRSS in bytes, capped by total, rounded + * ``fraction`` is clamped to 0..1 and rounded to 4 decimal places + +Linux-only via ``/proc``. On platforms without ``/proc`` the method +returns ``None`` instead of raising. +Cross-platform test: skips cleanly on macOS / Windows if ``/proc`` is +not available. +""" + +from __future__ import annotations + +import os +import sys +import tempfile +import types as _types +from pathlib import Path +from unittest.mock import patch + +import pytest + +# --------------------------------------------------------------------------- +# Stub heavy / unavailable external dependencies before importing the +# module under test. Same pattern as test_kv_cache_estimation.py. +# --------------------------------------------------------------------------- + +_BACKEND_DIR = str(Path(__file__).resolve().parent.parent) +if _BACKEND_DIR not in sys.path: + sys.path.insert(0, _BACKEND_DIR) + +_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") +sys.modules.setdefault("structlog", _structlog_stub) + +_httpx_stub = _types.ModuleType("httpx") +for _exc_name in ( + "ConnectError", + "TimeoutException", + "ReadTimeout", + "ReadError", + "RemoteProtocolError", + "CloseError", +): + setattr(_httpx_stub, _exc_name, type(_exc_name, (Exception,), {})) + + +class _FakeTimeout: + def __init__(self, *a, **kw): + pass + + +_httpx_stub.Timeout = _FakeTimeout +_httpx_stub.Client = type( + "Client", + (), + { + "__init__": lambda self, **kw: None, + "__enter__": lambda self: self, + "__exit__": lambda self, *a: None, + }, +) +sys.modules.setdefault("httpx", _httpx_stub) + +from core.inference.llama_cpp import LlamaCppBackend + + +# --------------------------------------------------------------------------- +# Helpers +# --------------------------------------------------------------------------- + + +def _make_instance(): + inst = LlamaCppBackend.__new__(LlamaCppBackend) + inst._process = None + inst._gguf_path = None + inst._healthy = False + return inst + + +class _FakeProc: + """Minimal stand-in for subprocess.Popen that just carries a pid.""" + + def __init__(self, pid: int): + self.pid = pid + + +def _write_sparse_file(path: Path, size_bytes: int) -> None: + """Create a sparse file of the given size without allocating blocks.""" + with open(path, "wb") as fh: + if size_bytes > 0: + fh.truncate(size_bytes) + + +# --------------------------------------------------------------------------- +# Tests +# --------------------------------------------------------------------------- + + +class TestLoadProgressEmptyStates: + def test_returns_none_when_no_process(self): + inst = _make_instance() + assert inst.load_progress() is None + + def test_returns_none_when_process_has_no_pid(self): + inst = _make_instance() + inst._process = _FakeProc(pid = None) # type: ignore[arg-type] + assert inst.load_progress() is None + + +class TestLoadProgressSingleShard: + def test_mmap_phase_for_alive_but_unhealthy(self, tmp_path): + """VmRSS below total -> phase='mmap', fraction reflects progress.""" + gguf = tmp_path / "model.gguf" + _write_sparse_file(gguf, 40 * 1024**3) # 40 GB + + inst = _make_instance() + inst._process = _FakeProc(pid = os.getpid()) # use our own pid + inst._gguf_path = str(gguf) + inst._healthy = False + + # Patch /proc read to claim 10 GB RSS. + def fake_open(path, *args, **kwargs): + if str(path).startswith("/proc/"): + import io + + return io.StringIO(f"Name:\ttest\nVmRSS:\t{10 * 1024 ** 2}\tkB\n") + return open(path, *args, **kwargs) # fall through + + with patch("builtins.open", side_effect = fake_open): + out = inst.load_progress() + + assert out is not None + assert out["phase"] == "mmap" + assert out["bytes_total"] == 40 * 1024**3 + assert out["bytes_loaded"] == 10 * 1024**3 + assert 0.24 < out["fraction"] < 0.26 # ~25% + + def test_ready_phase_when_healthy(self, tmp_path): + gguf = tmp_path / "model.gguf" + _write_sparse_file(gguf, 8 * 1024**3) + + inst = _make_instance() + inst._process = _FakeProc(pid = os.getpid()) + inst._gguf_path = str(gguf) + inst._healthy = True + + def fake_open(path, *args, **kwargs): + if str(path).startswith("/proc/"): + import io + + return io.StringIO(f"VmRSS:\t{8 * 1024 ** 2}\tkB\n") + return open(path, *args, **kwargs) + + with patch("builtins.open", side_effect = fake_open): + out = inst.load_progress() + + assert out is not None + assert out["phase"] == "ready" + assert out["bytes_total"] == 8 * 1024**3 + assert out["bytes_loaded"] == 8 * 1024**3 + assert out["fraction"] == 1.0 + + +class TestLoadProgressMultiShard: + """Shard-aware total: for ``*-00001-of-00004.gguf`` primaries the + method sums sibling files with the same prefix.""" + + def test_sharded_total_aggregates_siblings(self, tmp_path): + for i in range(1, 5): + _write_sparse_file( + tmp_path / f"model-{i:05d}-of-00004.gguf", + size_bytes = 20 * 1024**3, + ) + # Drop an unrelated .gguf in the same folder -- must not be counted. + _write_sparse_file(tmp_path / "mmproj-BF16.gguf", 2 * 1024**3) + + inst = _make_instance() + inst._process = _FakeProc(pid = os.getpid()) + inst._gguf_path = str(tmp_path / "model-00001-of-00004.gguf") + inst._healthy = False + + def fake_open(path, *args, **kwargs): + if str(path).startswith("/proc/"): + import io + + return io.StringIO("VmRSS:\t0\tkB\n") + return open(path, *args, **kwargs) + + with patch("builtins.open", side_effect = fake_open): + out = inst.load_progress() + + assert out is not None + assert out["bytes_total"] == 80 * 1024**3 # 4 x 20 GB, no mmproj + + +class TestLoadProgressDegradation: + """Broken / unusual inputs never raise; they produce best-effort output.""" + + def test_missing_gguf_path_still_reports_rss(self, tmp_path): + inst = _make_instance() + inst._process = _FakeProc(pid = os.getpid()) + inst._gguf_path = None + inst._healthy = False + + def fake_open(path, *args, **kwargs): + if str(path).startswith("/proc/"): + import io + + return io.StringIO("VmRSS:\t1024\tkB\n") + return open(path, *args, **kwargs) + + with patch("builtins.open", side_effect = fake_open): + out = inst.load_progress() + + assert out is not None + assert out["phase"] == "mmap" + assert out["bytes_total"] == 0 + assert out["bytes_loaded"] == 1024 * 1024 + assert out["fraction"] == 0.0 + + def test_unreadable_proc_returns_none(self, tmp_path): + inst = _make_instance() + # Pid that doesn't exist -> /proc read fails. + inst._process = _FakeProc(pid = 999_999_999) + inst._gguf_path = str(tmp_path / "model.gguf") # doesn't need to exist + inst._healthy = False + + out = inst.load_progress() + # FileNotFoundError on /proc path -> load_progress returns None. + assert out is None diff --git a/studio/backend/tests/test_llama_cpp_load_progress_live.py b/studio/backend/tests/test_llama_cpp_load_progress_live.py new file mode 100644 index 0000000000..beed8713c1 --- /dev/null +++ b/studio/backend/tests/test_llama_cpp_load_progress_live.py @@ -0,0 +1,202 @@ +# SPDX-License-Identifier: AGPL-3.0-only +# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0 + +"""Live, no-mock integration test for ``LlamaCppBackend.load_progress()``. + +The companion files (``test_llama_cpp_load_progress.py`` and +``test_llama_cpp_load_progress_matrix.py``) patch ``builtins.open`` to +feed synthetic VmRSS values. This file is the opposite: it uses **real** +subprocesses, **real** file sizes, and the **real** ``/proc`` +interface. It is the sanity check that the contract we keep in the +mocked tests still maps to what the kernel actually returns on a live +Linux system. + +Why both: the mocked tests can be fooled by a buggy implementation that +parses ``/proc`` output in a format the kernel no longer uses, or that +makes assumptions about ``Path.stat()`` vs ``os.path.getsize``. This +file hits the real APIs so any format drift gets caught. + +Skipped cleanly on non-Linux (no ``/proc``). +""" + +from __future__ import annotations + +import os +import subprocess +import sys +import time +import types as _types +from pathlib import Path + +import pytest + +# --------------------------------------------------------------------------- +# Same stubs as the matrix file (keep self-contained so the file can be +# run standalone as well as via the full suite). +# --------------------------------------------------------------------------- + +_BACKEND_DIR = str(Path(__file__).resolve().parent.parent) +if _BACKEND_DIR not in sys.path: + sys.path.insert(0, _BACKEND_DIR) + +_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") +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("Timeout", (), {"__init__": lambda self, *a, **k: None}) +_httpx_stub.Client = type( + "Client", + (), + { + "__init__": lambda self, **kw: None, + "__enter__": lambda self: self, + "__exit__": lambda self, *a: None, + }, +) +sys.modules.setdefault("httpx", _httpx_stub) + +from core.inference.llama_cpp import LlamaCppBackend + + +pytestmark = pytest.mark.skipif( + not Path("/proc").exists(), + reason = "live /proc test is Linux-only", +) + + +def _make_backend(pid: int, gguf_path: str, healthy: bool = False): + inst = LlamaCppBackend.__new__(LlamaCppBackend) + inst._process = type("P", (), {"pid": pid})() + inst._gguf_path = gguf_path + inst._healthy = healthy + return inst + + +def test_live_rss_matches_kernel_vmrss(tmp_path): + """Spawn a real child, let it allocate real bytes, confirm + ``bytes_loaded`` tracks the kernel's VmRSS within a sane tolerance.""" + # Child that allocates ~100 MB of zero'd bytes and then idles. + script = tmp_path / "burn.py" + script.write_text( + "import time, sys\n" + "buf = bytearray(100 * 1024 * 1024)\n" # 100 MB + "# touch every page so RSS actually grows\n" + "for i in range(0, len(buf), 4096):\n" + " buf[i] = 1\n" + "sys.stdout.write('ready\\n')\n" + "sys.stdout.flush()\n" + "time.sleep(10)\n" + ) + proc = subprocess.Popen( + [sys.executable, str(script)], + stdout = subprocess.PIPE, + stderr = subprocess.PIPE, + ) + try: + # Wait for the child to finish touching pages. + ready = proc.stdout.readline() + assert ready.strip() == b"ready" + + # Create a fake 200 MB sparse gguf so bytes_total is concrete. + gguf = tmp_path / "model.gguf" + with open(gguf, "wb") as f: + f.truncate(200 * 1024 * 1024) + + inst = _make_backend(proc.pid, str(gguf), healthy = False) + out = inst.load_progress() + + assert out is not None, "load_progress returned None for live pid" + assert out["phase"] == "mmap" + assert out["bytes_total"] == 200 * 1024 * 1024 + # VmRSS for the Python child includes the interpreter + the 100MB + # buffer, so a realistic floor is 50 MB and ceiling is 200 MB. + assert ( + out["bytes_loaded"] >= 50 * 1024 * 1024 + ), f"bytes_loaded unexpectedly low: {out['bytes_loaded']}" + assert out["bytes_loaded"] <= 200 * 1024 * 1024 + assert 0.0 < out["fraction"] <= 1.0 + finally: + proc.terminate() + try: + proc.wait(timeout = 5) + except subprocess.TimeoutExpired: + proc.kill() + + +def test_live_ready_phase_when_healthy(tmp_path): + gguf = tmp_path / "m.gguf" + with open(gguf, "wb") as f: + f.truncate(1 * 1024 * 1024) + + inst = _make_backend(os.getpid(), str(gguf), healthy = True) + out = inst.load_progress() + assert out is not None + assert out["phase"] == "ready" + assert out["bytes_total"] == 1 * 1024 * 1024 + # Self-pid RSS is well above 1 MiB for CPython; fraction caps at 1. + assert out["fraction"] == 1.0 + + +def test_live_dead_pid_returns_none(tmp_path): + """A recently-dead pid may linger in /proc for ms; use a clearly + invalid id so the read reliably fails.""" + gguf = tmp_path / "m.gguf" + gguf.touch() + + inst = _make_backend(9_999_999_999, str(gguf), healthy = False) + out = inst.load_progress() + assert out is None + + +def test_live_shard_aggregation_counts_real_files(tmp_path): + """With 4 real sibling shards on disk, ``bytes_total`` equals their + summed size to the byte.""" + shard_size = 7 * 1024 * 1024 # 7 MB each + for i in range(1, 5): + f = tmp_path / f"model-{i:05d}-of-00004.gguf" + with open(f, "wb") as fh: + fh.truncate(shard_size) + # Unrelated file in same dir -- must not be counted. + with open(tmp_path / "config.json", "wb") as fh: + fh.truncate(123) + + inst = _make_backend( + os.getpid(), + str(tmp_path / "model-00001-of-00004.gguf"), + healthy = False, + ) + out = inst.load_progress() + assert out is not None + assert out["bytes_total"] == 4 * shard_size + + +def test_live_repeated_polling_stays_sane(tmp_path): + """Sampling the same backend 20 times should not raise or produce + non-numeric output, even under normal kernel RSS jitter.""" + gguf = tmp_path / "m.gguf" + with open(gguf, "wb") as f: + f.truncate(500 * 1024 * 1024) + + inst = _make_backend(os.getpid(), str(gguf), healthy = False) + seen = [] + for _ in range(20): + out = inst.load_progress() + assert out is not None + assert isinstance(out["bytes_loaded"], int) + assert isinstance(out["bytes_total"], int) + assert 0.0 <= out["fraction"] <= 1.0 + seen.append(out["bytes_loaded"]) + time.sleep(0.01) + # RSS of a healthy Python process doesn't go below ~5 MB. + assert min(seen) > 1 * 1024 * 1024 diff --git a/studio/backend/tests/test_llama_cpp_load_progress_matrix.py b/studio/backend/tests/test_llama_cpp_load_progress_matrix.py new file mode 100644 index 0000000000..a88450ec0b --- /dev/null +++ b/studio/backend/tests/test_llama_cpp_load_progress_matrix.py @@ -0,0 +1,473 @@ +# SPDX-License-Identifier: AGPL-3.0-only +# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0 + +"""Extended test matrix for ``LlamaCppBackend.load_progress()``. + +Companion to ``test_llama_cpp_load_progress.py`` (which pins the basic +contract). This file widens coverage to the edge cases that bit users +or were hypothesized to bite them on cross-platform installs: + + * Platform matrix — macOS/Windows simulation via ``/proc`` absence. + * ``VmRSS`` parsing — tab vs space delimiter, missing line, malformed + integer. + * Filesystem edges — HF-cache symlinks, broken symlinks, nonexistent + paths, relative paths. + * Shard aggregation — partial multi-shard downloads where some shards + are still ``.incomplete``, two shard series in the same dir, + ``mmproj-*.gguf`` sibling exclusion for non-sharded primaries, + single-file models. + * Lifecycle races — process set before ``_gguf_path`` is assigned, + process dead mid-sample, ``_healthy`` flipped to True. + * Concurrent sampling — 10 threads × 50 iterations against a single + backend, hitting real ``/proc`` (no mocks — see the note in + ``TestConcurrentSampling`` for why). + * Fraction bounds — capped at 1.0 when RSS exceeds total; 0.0 when + total is zero. + +All tests are Linux-only in practice (we stub ``/proc`` where needed). +The stable subset runs in well under a second. +""" + +from __future__ import annotations + +import io +import os +import sys +import threading +import types as _types +from pathlib import Path +from unittest.mock import patch + +import pytest + +# --------------------------------------------------------------------------- +# Stub heavy / unavailable external dependencies before importing the +# module under test. Same pattern as test_llama_cpp_load_progress.py. +# --------------------------------------------------------------------------- + +_BACKEND_DIR = str(Path(__file__).resolve().parent.parent) +if _BACKEND_DIR not in sys.path: + sys.path.insert(0, _BACKEND_DIR) + +_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") +sys.modules.setdefault("structlog", _structlog_stub) + +_httpx_stub = _types.ModuleType("httpx") +for _exc_name in ( + "ConnectError", + "TimeoutException", + "ReadTimeout", + "ReadError", + "RemoteProtocolError", + "CloseError", +): + setattr(_httpx_stub, _exc_name, type(_exc_name, (Exception,), {})) + + +class _FakeTimeout: + def __init__(self, *a, **kw): + pass + + +_httpx_stub.Timeout = _FakeTimeout +_httpx_stub.Client = type( + "Client", + (), + { + "__init__": lambda self, **kw: None, + "__enter__": lambda self: self, + "__exit__": lambda self, *a: None, + }, +) +sys.modules.setdefault("httpx", _httpx_stub) + +from core.inference.llama_cpp import LlamaCppBackend + + +# --------------------------------------------------------------------------- +# Helpers +# --------------------------------------------------------------------------- + + +def _make(): + inst = LlamaCppBackend.__new__(LlamaCppBackend) + inst._process = None + inst._gguf_path = None + inst._healthy = False + return inst + + +class _Proc: + def __init__(self, pid): + self.pid = pid + + +def _sparse(path, size): + with open(path, "wb") as f: + if size > 0: + f.truncate(size) + + +def _fake_proc_reader(rss_kb): + """Return an ``open()`` replacement that fakes /proc reads with a VmRSS line.""" + + def fake_open(path, *args, **kwargs): + if str(path).startswith("/proc/"): + return io.StringIO(f"VmRSS:\t{rss_kb}\tkB\n") + return open(path, *args, **kwargs) + + return fake_open + + +# --------------------------------------------------------------------------- +# A. Platform matrix +# --------------------------------------------------------------------------- + + +class TestPlatformMatrix: + """The method is Linux-first via /proc. On macOS/Windows it must + degrade to None rather than crash.""" + + def test_linux_live_proc_is_self_pid(self, tmp_path): + """Self-pid /proc read uses the real kernel interface.""" + gguf = tmp_path / "m.gguf" + _sparse(gguf, 1 * 1024**3) + inst = _make() + inst._process = _Proc(os.getpid()) + inst._gguf_path = str(gguf) + inst._healthy = False + out = inst.load_progress() + assert out is not None + assert out["phase"] == "mmap" + assert out["bytes_total"] == 1 * 1024**3 + # Our Python process has some RSS -- just sanity-check positive. + assert out["bytes_loaded"] > 0 + + def test_macos_no_proc_returns_none(self, tmp_path): + """Simulate macOS: /proc open fails with FileNotFoundError.""" + gguf = tmp_path / "m.gguf" + _sparse(gguf, 1 * 1024**3) + inst = _make() + inst._process = _Proc(pid = 12345) + inst._gguf_path = str(gguf) + + def fake_open(path, *args, **kwargs): + if str(path).startswith("/proc/"): + raise FileNotFoundError(f"No such file: {path}") + return open(path, *args, **kwargs) + + with patch("builtins.open", side_effect = fake_open): + out = inst.load_progress() + assert out is None + + def test_windows_no_proc_returns_none(self, tmp_path): + """Simulate Windows: opening /proc raises PermissionError or OSError.""" + gguf = tmp_path / "m.gguf" + _sparse(gguf, 1 * 1024**3) + inst = _make() + inst._process = _Proc(pid = 4567) + inst._gguf_path = str(gguf) + + def fake_open(path, *args, **kwargs): + if str(path).startswith("/proc/"): + raise PermissionError("access denied") + return open(path, *args, **kwargs) + + with patch("builtins.open", side_effect = fake_open): + out = inst.load_progress() + assert out is None + + +# --------------------------------------------------------------------------- +# B. VmRSS parsing edge cases +# --------------------------------------------------------------------------- + + +class TestVmRSSParsing: + def test_standard_tab_delimited(self, tmp_path): + gguf = tmp_path / "m.gguf" + _sparse(gguf, 4 * 1024**3) + inst = _make() + inst._process = _Proc(os.getpid()) + inst._gguf_path = str(gguf) + with patch("builtins.open", side_effect = _fake_proc_reader(2 * 1024**2)): + out = inst.load_progress() + assert out["bytes_loaded"] == 2 * 1024**3 + + def test_space_separated_fallback(self, tmp_path): + """Some kernels emit single-space rather than tab.""" + gguf = tmp_path / "m.gguf" + _sparse(gguf, 4 * 1024**3) + inst = _make() + inst._process = _Proc(os.getpid()) + inst._gguf_path = str(gguf) + + def fake_open(path, *a, **kw): + if str(path).startswith("/proc/"): + return io.StringIO("VmRSS: 4194304 kB\n") + return open(path, *a, **kw) + + with patch("builtins.open", side_effect = fake_open): + out = inst.load_progress() + assert out["bytes_loaded"] == 4 * 1024**3 + + def test_missing_vmrss_line(self, tmp_path): + """Kernel with VmRSS stripped (zombie / kthread) -> 0.""" + gguf = tmp_path / "m.gguf" + _sparse(gguf, 1 * 1024**3) + inst = _make() + inst._process = _Proc(os.getpid()) + inst._gguf_path = str(gguf) + + def fake_open(path, *a, **kw): + if str(path).startswith("/proc/"): + return io.StringIO("Name:\ttest\nState:\tZ (zombie)\n") + return open(path, *a, **kw) + + with patch("builtins.open", side_effect = fake_open): + out = inst.load_progress() + assert out is not None + assert out["bytes_loaded"] == 0 + assert out["fraction"] == 0.0 + + def test_malformed_vmrss_value(self, tmp_path): + """Non-integer VmRSS value should be treated as if the line were + absent (early ValueError caught).""" + gguf = tmp_path / "m.gguf" + _sparse(gguf, 1 * 1024**3) + inst = _make() + inst._process = _Proc(os.getpid()) + inst._gguf_path = str(gguf) + + def fake_open(path, *a, **kw): + if str(path).startswith("/proc/"): + return io.StringIO("VmRSS:\tXXXX\tkB\n") + return open(path, *a, **kw) + + with patch("builtins.open", side_effect = fake_open): + out = inst.load_progress() + # The implementation catches ValueError on int() and returns None. + assert out is None + + +# --------------------------------------------------------------------------- +# C. Filesystem edge cases +# --------------------------------------------------------------------------- + + +class TestFilesystemEdges: + def test_symlink_primary_follows_to_blob(self, tmp_path): + """HF cache stores blobs under blobs/ and symlinks them from + snapshots/. The method must follow the symlink.""" + blob = tmp_path / "blob" + _sparse(blob, 12 * 1024**3) + snap = tmp_path / "snap" + snap.mkdir() + link = snap / "m.gguf" + link.symlink_to(blob) + + inst = _make() + inst._process = _Proc(os.getpid()) + inst._gguf_path = str(link) + with patch("builtins.open", side_effect = _fake_proc_reader(6 * 1024**2)): + out = inst.load_progress() + assert out["bytes_total"] == 12 * 1024**3 + + def test_broken_symlink_skipped(self, tmp_path): + snap = tmp_path / "snap" + snap.mkdir() + link = snap / "m.gguf" + link.symlink_to(tmp_path / "missing-blob") + inst = _make() + inst._process = _Proc(os.getpid()) + inst._gguf_path = str(link) + with patch("builtins.open", side_effect = _fake_proc_reader(1024)): + out = inst.load_progress() + assert out["bytes_total"] == 0 + assert out["bytes_loaded"] == 1024 * 1024 + + def test_nonexistent_path_skipped(self, tmp_path): + inst = _make() + inst._process = _Proc(os.getpid()) + inst._gguf_path = str(tmp_path / "ghost.gguf") + with patch("builtins.open", side_effect = _fake_proc_reader(1024)): + out = inst.load_progress() + assert out["bytes_total"] == 0 + + def test_relative_gguf_path(self, tmp_path): + """Relative paths shouldn't crash; behaviour depends on CWD but + the method must not raise.""" + cwd = os.getcwd() + try: + os.chdir(tmp_path) + _sparse(Path("rel.gguf"), 8 * 1024**3) + inst = _make() + inst._process = _Proc(os.getpid()) + inst._gguf_path = "rel.gguf" + with patch("builtins.open", side_effect = _fake_proc_reader(0)): + out = inst.load_progress() + assert out is not None + assert out["bytes_total"] == 8 * 1024**3 + finally: + os.chdir(cwd) + + +# --------------------------------------------------------------------------- +# D. Shard aggregation +# --------------------------------------------------------------------------- + + +class TestShardAggregation: + def test_partial_multi_shard_download(self, tmp_path): + """Primary present but shards 2..N still downloading as + ``.incomplete``. Sums only the fully-arrived ``.gguf`` files.""" + _sparse(tmp_path / "m-00001-of-00004.gguf", 30 * 1024**3) + _sparse(tmp_path / "m-00002-of-00004.gguf", 30 * 1024**3) + # 3 and 4 still downloading as .incomplete + _sparse(tmp_path / "m-00003-of-00004.gguf.incomplete", 5 * 1024**3) + inst = _make() + inst._process = _Proc(os.getpid()) + inst._gguf_path = str(tmp_path / "m-00001-of-00004.gguf") + with patch("builtins.open", side_effect = _fake_proc_reader(0)): + out = inst.load_progress() + assert out["bytes_total"] == 60 * 1024**3 # only the .gguf siblings + + def test_two_shard_series_in_same_dir(self, tmp_path): + """Defensive: if two quant series share a dir, prefix filter + only sums siblings of the chosen primary.""" + for i in range(1, 3): + _sparse(tmp_path / f"m_q4-{i:05d}-of-00002.gguf", 10 * 1024**3) + _sparse(tmp_path / f"m_q8-{i:05d}-of-00002.gguf", 20 * 1024**3) + inst = _make() + inst._process = _Proc(os.getpid()) + inst._gguf_path = str(tmp_path / "m_q8-00001-of-00002.gguf") + with patch("builtins.open", side_effect = _fake_proc_reader(0)): + out = inst.load_progress() + assert out["bytes_total"] == 40 * 1024**3 # just q8 series + + def test_mmproj_sibling_not_counted(self, tmp_path): + """Vision models drop an ``mmproj-*.gguf`` alongside. For a + single-file (non-sharded) primary we only count the primary.""" + _sparse(tmp_path / "m.gguf", 8 * 1024**3) + _sparse(tmp_path / "mmproj-BF16.gguf", 2 * 1024**3) + inst = _make() + inst._process = _Proc(os.getpid()) + inst._gguf_path = str(tmp_path / "m.gguf") + with patch("builtins.open", side_effect = _fake_proc_reader(0)): + out = inst.load_progress() + # Non-sharded primary: only the primary is counted. + assert out["bytes_total"] == 8 * 1024**3 + + def test_single_file_model(self, tmp_path): + """Non-sharded model: primary only.""" + _sparse(tmp_path / "small.gguf", 4 * 1024**3) + inst = _make() + inst._process = _Proc(os.getpid()) + inst._gguf_path = str(tmp_path / "small.gguf") + with patch("builtins.open", side_effect = _fake_proc_reader(2 * 1024**2)): + out = inst.load_progress() + assert out["bytes_total"] == 4 * 1024**3 + assert out["bytes_loaded"] == 2 * 1024**3 + + +# --------------------------------------------------------------------------- +# E. Lifecycle races +# --------------------------------------------------------------------------- + + +class TestLifecycleRaces: + def test_process_set_but_gguf_path_not_yet(self, tmp_path): + """Moment between Popen and self._gguf_path=model_path.""" + inst = _make() + inst._process = _Proc(os.getpid()) + inst._gguf_path = None + with patch("builtins.open", side_effect = _fake_proc_reader(1024)): + out = inst.load_progress() + assert out is not None + assert out["phase"] == "mmap" + assert out["bytes_total"] == 0 + assert out["bytes_loaded"] == 1024 * 1024 + + def test_process_died_mid_sample(self, tmp_path): + """/proc/ disappears -> None.""" + _sparse(tmp_path / "m.gguf", 1 * 1024**3) + inst = _make() + inst._process = _Proc(pid = 999_999_999) + inst._gguf_path = str(tmp_path / "m.gguf") + assert inst.load_progress() is None + + def test_healthy_true_ready_phase(self, tmp_path): + _sparse(tmp_path / "m.gguf", 1 * 1024**3) + inst = _make() + inst._process = _Proc(os.getpid()) + inst._gguf_path = str(tmp_path / "m.gguf") + inst._healthy = True + with patch("builtins.open", side_effect = _fake_proc_reader(1024)): + out = inst.load_progress() + assert out["phase"] == "ready" + + +# --------------------------------------------------------------------------- +# F. Concurrent sampling (simulates multiple browser tabs polling) +# --------------------------------------------------------------------------- + + +class TestConcurrentSampling: + def test_parallel_invocations_never_raise(self, tmp_path): + """Many concurrent samplers hitting the same backend must not raise. + + We intentionally do NOT patch ``builtins.open`` here because + ``unittest.mock.patch`` is not thread-safe: interleaved + enter/exit across threads can leak a Mock into ``builtins.open`` + and poison every subsequent test in the session. Instead, we + let each thread hit the real ``/proc/self/status`` of the test + process, which is exactly the code path that matters in prod. + """ + _sparse(tmp_path / "m.gguf", 1 * 1024**3) + inst = _make() + inst._process = _Proc(os.getpid()) + inst._gguf_path = str(tmp_path / "m.gguf") + errors = [] + + def run(): + try: + for _ in range(50): + inst.load_progress() + except Exception as e: # pragma: no cover + errors.append(e) + + threads = [threading.Thread(target = run) for _ in range(10)] + for t in threads: + t.start() + for t in threads: + t.join() + assert not errors, errors + + +# --------------------------------------------------------------------------- +# G. Fraction bounds +# --------------------------------------------------------------------------- + + +class TestFractionBounds: + def test_fraction_capped_at_one(self, tmp_path): + _sparse(tmp_path / "m.gguf", 1 * 1024**3) + inst = _make() + inst._process = _Proc(os.getpid()) + inst._gguf_path = str(tmp_path / "m.gguf") + # RSS > total (post-paged-in + extra structures) + with patch("builtins.open", side_effect = _fake_proc_reader(2 * 1024**2)): + out = inst.load_progress() + assert 0.0 <= out["fraction"] <= 1.0 + + def test_fraction_zero_when_total_zero(self): + inst = _make() + inst._process = _Proc(os.getpid()) + inst._gguf_path = None + with patch("builtins.open", side_effect = _fake_proc_reader(1024**2)): + out = inst.load_progress() + assert out["fraction"] == 0.0 diff --git a/studio/backend/tests/test_llama_cpp_max_context_threshold.py b/studio/backend/tests/test_llama_cpp_max_context_threshold.py new file mode 100644 index 0000000000..5fd0243c9f --- /dev/null +++ b/studio/backend/tests/test_llama_cpp_max_context_threshold.py @@ -0,0 +1,244 @@ +# 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 ``max_context_length`` warning-threshold semantics. + +``/api/inference/status.max_context_length`` is what the ctx slider in +the chat settings sheet reads to decide when to render the "Exceeds +estimated VRAM capacity. The model may use system RAM." warning: + + ctxDisplayValue > ggufMaxContextLength → show warning + +For models whose weights fit on some GPU subset, the warning threshold +is the largest ctx that fits fully in VRAM (the binary-search cap from +``_fit_context_to_vram``). For models whose weights exceed 90% of every +GPU subset's free memory, the warning must fire as soon as the user +drags above the 4096 spec default (otherwise a user loading e.g. +MiniMax-M2.7 on a 97 GB GPU sees a slider up to 196608 with no +indication that any value above 4096 will trigger ``--fit on`` and +degrade performance). + +These tests pin both cases. No GPU probing, no subprocess, no GGUF I/O. +Cross-platform: Linux, macOS, Windows, WSL. +""" + +from __future__ import annotations + +import sys +import types as _types +from pathlib import Path + +import pytest + +# --------------------------------------------------------------------------- +# Stub heavy / unavailable external dependencies before importing the +# module under test. Same pattern as test_kv_cache_estimation.py. +# --------------------------------------------------------------------------- + +_BACKEND_DIR = str(Path(__file__).resolve().parent.parent) +if _BACKEND_DIR not in sys.path: + sys.path.insert(0, _BACKEND_DIR) + +# loggers +_loggers_stub = _types.ModuleType("loggers") +_loggers_stub.get_logger = lambda name: __import__("logging").getLogger(name) +sys.modules.setdefault("loggers", _loggers_stub) + +# structlog +_structlog_stub = _types.ModuleType("structlog") +sys.modules.setdefault("structlog", _structlog_stub) + +# httpx +_httpx_stub = _types.ModuleType("httpx") +for _exc_name in ( + "ConnectError", + "TimeoutException", + "ReadTimeout", + "ReadError", + "RemoteProtocolError", + "CloseError", +): + setattr(_httpx_stub, _exc_name, type(_exc_name, (Exception,), {})) + + +class _FakeTimeout: + def __init__(self, *a, **kw): + pass + + +_httpx_stub.Timeout = _FakeTimeout +_httpx_stub.Client = type( + "Client", + (), + { + "__init__": lambda self, **kw: None, + "__enter__": lambda self: self, + "__exit__": lambda self, *a: None, + }, +) +sys.modules.setdefault("httpx", _httpx_stub) + +from core.inference.llama_cpp import LlamaCppBackend + + +# --------------------------------------------------------------------------- +# Helpers +# --------------------------------------------------------------------------- + +GIB = 1024**3 + + +def _make_backend(native_ctx = 131072): + inst = LlamaCppBackend.__new__(LlamaCppBackend) + inst._context_length = native_ctx + inst._n_layers = 80 + inst._n_kv_heads = 8 + inst._n_heads = 64 + inst._embedding_length = 8192 + inst._kv_key_length = 128 + inst._kv_value_length = 128 + inst._kv_lora_rank = None + inst._sliding_window = None + inst._ssm_inner_size = None + inst._full_attention_interval = None + inst._key_length_mla = None + return inst + + +def _compute_max_available_ctx(native_ctx, model_gib, gpus, kv_per_token_bytes = 325_000): + """Run the ceiling-probe block from load_model and return the final + ``max_available_ctx`` value the backend would assign to + ``_max_context_length``. + """ + inst = _make_backend(native_ctx = native_ctx) + model_size = int(model_gib * GIB) + + inst._estimate_kv_cache_bytes = ( + lambda n, _t = None: 0 if n <= 0 else n * kv_per_token_bytes + ) + inst._can_estimate_kv = lambda: True + + context_length = inst._context_length + effective_ctx = context_length + max_available_ctx = context_length + + cache_type_kv = None + native_ctx_for_cap = context_length + + ranked_for_cap = sorted(gpus, key = lambda g: g[1], reverse = True) + best_cap = 0 + for n_gpus in range(1, len(ranked_for_cap) + 1): + subset = ranked_for_cap[:n_gpus] + pool_mib = sum(free for _, free in subset) + capped = inst._fit_context_to_vram( + native_ctx_for_cap, + pool_mib, + model_size, + cache_type_kv, + ) + kv = inst._estimate_kv_cache_bytes(capped, cache_type_kv) + total_mib = (model_size + kv) / (1024 * 1024) + if total_mib <= pool_mib * 0.90: + best_cap = max(best_cap, capped) + if best_cap > 0: + max_available_ctx = best_cap + else: + max_available_ctx = min(4096, native_ctx_for_cap) + + return max_available_ctx + + +# --------------------------------------------------------------------------- +# Weights exceed every GPU subset's VRAM (MiniMax-M2.7-like) +# --------------------------------------------------------------------------- + + +class TestMaxContextLengthForWeightsExceedVRAM: + """The UI ``max_context_length`` threshold must fall back to 4096 so + the warning fires as soon as the user drags above the spec default. + """ + + def test_minimax_like(self): + """131 GB weights, single 97 GB GPU, native ctx 196608.""" + got = _compute_max_available_ctx( + native_ctx = 196608, + model_gib = 131, + gpus = [(0, 97_000)], + ) + assert got == 4096 + + def test_multi_gpu_all_subsets_fail(self): + """400 GB weights across a 4x80 GB pool (320 GB total, still too small).""" + got = _compute_max_available_ctx( + native_ctx = 131072, + model_gib = 400, + gpus = [(0, 80_000), (1, 80_000), (2, 80_000), (3, 80_000)], + ) + assert got == 4096 + + def test_native_below_fallback_is_preserved(self): + """If the model's native ctx is itself smaller than 4096, do not + advertise a larger value than the model supports.""" + got = _compute_max_available_ctx( + native_ctx = 2048, + model_gib = 200, + gpus = [(0, 80_000)], + ) + assert got == 2048 + + +# --------------------------------------------------------------------------- +# Fittable models (regression guard) +# --------------------------------------------------------------------------- + + +class TestMaxContextLengthForFittableModels: + """The existing best-cap behaviour must be unchanged.""" + + def test_small_model_fits_easily(self): + """8 GB model on 24 GB GPU: should auto-pick a large ctx.""" + got = _compute_max_available_ctx( + native_ctx = 131072, + model_gib = 8, + gpus = [(0, 24_000)], + kv_per_token_bytes = 8192, + ) + assert got > 4096 + assert got <= 131072 + + def test_medium_model_multi_gpu(self): + """60 GB model split across 2 GPUs: picks a fitting ctx.""" + got = _compute_max_available_ctx( + native_ctx = 131072, + model_gib = 60, + gpus = [(0, 40_000), (1, 40_000)], + kv_per_token_bytes = 8192, + ) + assert got > 4096 + + def test_tiny_model_on_huge_gpu_near_native(self): + """2 GB model, 80 GB GPU, negligible KV: should approach native.""" + got = _compute_max_available_ctx( + native_ctx = 131072, + model_gib = 2, + gpus = [(0, 80_000)], + kv_per_token_bytes = 64, + ) + assert got >= 131072 - 256 # rounded to 256 boundary + + +# --------------------------------------------------------------------------- +# Property plumbing +# --------------------------------------------------------------------------- + + +class TestMaxContextLengthProperty: + def test_falls_back_to_native_when_unset(self): + inst = _make_backend(native_ctx = 131072) + inst._max_context_length = None + assert inst.max_context_length == 131072 + + def test_returns_stored_value_when_set(self): + inst = _make_backend(native_ctx = 131072) + inst._max_context_length = 4096 + assert inst.max_context_length == 4096 diff --git a/studio/backend/tests/test_llama_cpp_no_context_shift.py b/studio/backend/tests/test_llama_cpp_no_context_shift.py new file mode 100644 index 0000000000..b9f25faf88 --- /dev/null +++ b/studio/backend/tests/test_llama_cpp_no_context_shift.py @@ -0,0 +1,137 @@ +# SPDX-License-Identifier: AGPL-3.0-only +# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0 + +"""``--no-context-shift`` launch-flag contract. + +When llama-server runs with its default context-shift behavior, the UI +has no way to tell the user that the KV cache has been rotated -- +earlier turns silently vanish from the conversation. The Studio +backend always passes ``--no-context-shift`` so the server returns a +clean error instead, and the chat adapter can point the user at the +``Context Length`` input in the settings panel. + +This file is a static read of the launch command: we ask +``LlamaCppBackend`` to assemble its ``cmd`` list and assert the flag +is always present. Testing via the real subprocess would require an +actual GGUF on disk, which is out of scope for the fast test suite. +""" + +from __future__ import annotations + +import inspect +import sys +import types as _types +from pathlib import Path + +import pytest + +# --------------------------------------------------------------------------- +# Same external-dep stubs as the other llama_cpp tests. +# --------------------------------------------------------------------------- + +_BACKEND_DIR = str(Path(__file__).resolve().parent.parent) +if _BACKEND_DIR not in sys.path: + sys.path.insert(0, _BACKEND_DIR) + +_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") +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 + + +def _load_model_source() -> str: + """Return the source of ``LlamaCppBackend.load_model``. + + Using ``inspect.getsource`` instead of reading the file directly + scopes the assertions to the function that actually launches + llama-server, so neither the presence check nor the location check + can be fooled by a stray occurrence of ``"--no-context-shift"`` + elsewhere in the module. + """ + return inspect.getsource(llama_cpp_module.LlamaCppBackend.load_model) + + +def test_no_context_shift_is_in_load_model(): + """The flag is part of the static launch-command template. + + We check the source of ``load_model`` rather than mocking the whole + call chain (GPU probing, GGUF stat, etc.): the flag is written as + a literal in one place and any regression has to delete it, which + a text search will catch. + """ + assert '"--no-context-shift"' in _load_model_source(), ( + "llama-server must be launched with --no-context-shift so the " + "UI can surface a clean 'context full' error instead of silently " + "losing old turns to a KV-cache rotation." + ) + + +def test_flag_sits_inside_the_base_cmd_list(): + """Pin the flag's location so a future refactor can't accidentally + move it into a branch that only fires on some code paths. + + We slice from ``cmd = [`` to the first ``]`` at the same indent. + Using ``inspect.getsource`` means the function lives in its own + string and there are no siblings to worry about, so a plain + bracket search would also work -- anchoring on the trailing indent + just keeps the slice from wandering into a later expression if the + opening literal ever grows an in-line comment trailing it. + """ + source = _load_model_source() + start = source.find("cmd = [") + assert start >= 0, "could not find the base cmd = [...] block" + # Find the first line containing only ``]`` (possibly indented). + # Works for any indentation style the formatter picks. + rest = source[start:] + end_rel = -1 + for line_start, line in _iter_lines_with_offset(rest): + if line_start == 0: + # Skip the opening ``cmd = [`` line itself. + continue + if line.strip() == "]": + end_rel = line_start + break + assert end_rel > 0, "could not find end of cmd = [...] block" + block = rest[:end_rel] + assert '"--no-context-shift"' in block, ( + "--no-context-shift must be in the base cmd list, not in a " + "conditional branch -- otherwise some code paths would still " + "run with silent context shift enabled." + ) + # Also pin that it is next to -c / --ctx so the grouping makes sense. + assert '"-c"' in block + assert '"--flash-attn"' in block + + +def _iter_lines_with_offset(text: str): + """Yield (offset, line) pairs over ``text`` without losing offsets.""" + offset = 0 + for line in text.splitlines(keepends = True): + yield offset, line + offset += len(line) diff --git a/studio/backend/tests/test_pytorch_mirror.py b/studio/backend/tests/test_pytorch_mirror.py new file mode 100644 index 0000000000..5844f209b6 --- /dev/null +++ b/studio/backend/tests/test_pytorch_mirror.py @@ -0,0 +1,55 @@ +# 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 UNSLOTH_PYTORCH_MIRROR env var in install_python_stack.py.""" + +from __future__ import annotations + +import importlib +import os +import sys +from pathlib import Path + +import pytest + +# install_python_stack.py lives at repo_root/studio/install_python_stack.py +_INSTALL_SCRIPT = Path(__file__).resolve().parents[2] / "install_python_stack.py" + +OFFICIAL_URL = "https://download.pytorch.org/whl" + + +def _reload_whl_base(monkeypatch, mirror_value = None): + """(Re-)import install_python_stack with a controlled env and return _PYTORCH_WHL_BASE.""" + # Remove cached module so the module-level assignment re-executes + sys.modules.pop("install_python_stack", None) + + if mirror_value is None: + monkeypatch.delenv("UNSLOTH_PYTORCH_MIRROR", raising = False) + else: + monkeypatch.setenv("UNSLOTH_PYTORCH_MIRROR", mirror_value) + + # Temporarily add the script's directory to sys.path for import + script_dir = str(_INSTALL_SCRIPT.parent) + monkeypatch.syspath_prepend(script_dir) + + import install_python_stack + + return install_python_stack._PYTORCH_WHL_BASE + + +class TestPyTorchMirrorEnvVar: + """UNSLOTH_PYTORCH_MIRROR controls _PYTORCH_WHL_BASE in install_python_stack.""" + + def test_unset_uses_official_url(self, monkeypatch): + assert _reload_whl_base(monkeypatch) == OFFICIAL_URL + + def test_empty_string_falls_back_to_official(self, monkeypatch): + assert _reload_whl_base(monkeypatch, "") == OFFICIAL_URL + + def test_custom_mirror_is_used(self, monkeypatch): + mirror = "https://mirrors.nju.edu.cn/pytorch/whl" + assert _reload_whl_base(monkeypatch, mirror) == mirror + + def test_trailing_slash_stripped(self, monkeypatch): + result = _reload_whl_base(monkeypatch, "https://example.com/whl/") + assert result == "https://example.com/whl" diff --git a/studio/backend/tests/test_responses_api.py b/studio/backend/tests/test_responses_api.py new file mode 100644 index 0000000000..5b55f87259 --- /dev/null +++ b/studio/backend/tests/test_responses_api.py @@ -0,0 +1,328 @@ +# SPDX-License-Identifier: AGPL-3.0-only +# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. + +""" +Tests for the OpenAI Responses API schemas and input normalisation. +These tests do NOT require a running server or GPU -- they validate +the Pydantic models and the _normalise_responses_input helper. +""" + +import sys +import os +import json +import re + +# Ensure backend is on path +_backend = os.path.join(os.path.dirname(__file__), "..") +sys.path.insert(0, _backend) + +from models.inference import ( + ResponsesRequest, + ResponsesInputMessage, + ResponsesInputTextPart, + ResponsesInputImagePart, + ResponsesOutputTextContent, + ResponsesOutputMessage, + ResponsesUsage, + ResponsesResponse, + ChatMessage, + TextContentPart, + ImageContentPart, + ImageUrl, + ChatCompletionRequest, +) + + +# ── _normalise_responses_input: copied from routes/inference.py ── +# We cannot import routes.inference directly because routes/__init__.py +# pulls in heavy dependencies (structlog/twisted/torch). This is a +# direct copy of the function for testing purposes. + + +def _normalise_responses_input(payload: ResponsesRequest) -> list: + """Convert a ResponsesRequest into a list of ChatMessage for the completions backend.""" + messages = [] + + # System / developer instructions + if payload.instructions: + messages.append(ChatMessage(role = "system", content = payload.instructions)) + + # Simple string input + if isinstance(payload.input, str): + if payload.input: + messages.append(ChatMessage(role = "user", content = payload.input)) + return messages + + # List of ResponsesInputMessage + for msg in payload.input: + role = "system" if msg.role == "developer" else msg.role + + if isinstance(msg.content, str): + messages.append(ChatMessage(role = role, content = msg.content)) + else: + # Convert Responses content parts -> Chat content parts + parts = [] + for part in msg.content: + if isinstance(part, ResponsesInputTextPart): + parts.append(TextContentPart(type = "text", text = part.text)) + elif isinstance(part, ResponsesInputImagePart): + parts.append( + ImageContentPart( + type = "image_url", + image_url = ImageUrl(url = part.image_url, detail = part.detail), + ) + ) + messages.append(ChatMessage(role = role, content = parts if parts else "")) + + return messages + + +# ===================================================================== +# Schema validation tests +# ===================================================================== + + +class TestResponsesRequest: + """Validate ResponsesRequest accepts the shapes the OpenAI SDK sends.""" + + def test_minimal_string_input(self): + req = ResponsesRequest(input = "Hello") + assert req.input == "Hello" + assert req.stream is False + assert req.model == "default" + + def test_message_list_input(self): + req = ResponsesRequest( + input = [ + {"role": "user", "content": "Hi"}, + {"role": "assistant", "content": "Hello!"}, + ], + ) + assert len(req.input) == 2 + assert req.input[0].role == "user" + assert req.input[0].content == "Hi" + + def test_multimodal_input(self): + req = ResponsesRequest( + input = [ + { + "role": "user", + "content": [ + {"type": "input_text", "text": "What is in this image?"}, + { + "type": "input_image", + "image_url": "https://example.com/img.png", + }, + ], + }, + ], + ) + parts = req.input[0].content + assert len(parts) == 2 + assert isinstance(parts[0], ResponsesInputTextPart) + assert isinstance(parts[1], ResponsesInputImagePart) + + def test_instructions_field(self): + req = ResponsesRequest( + input = "test", + instructions = "You are a helpful assistant.", + ) + assert req.instructions == "You are a helpful assistant." + + def test_extra_fields_accepted(self): + """OpenAI SDK may send fields we don't model -- extra='allow' should pass.""" + req = ResponsesRequest( + input = "test", + tools = [{"type": "web_search_preview"}], + store = True, + metadata = {"key": "value"}, + previous_response_id = "resp_abc123", + ) + assert req.tools == [{"type": "web_search_preview"}] + assert req.store is True + + def test_stream_flag(self): + req = ResponsesRequest(input = "test", stream = True) + assert req.stream is True + + def test_temperature_and_top_p(self): + req = ResponsesRequest(input = "test", temperature = 0.8, top_p = 0.9) + assert req.temperature == 0.8 + assert req.top_p == 0.9 + + def test_max_output_tokens(self): + req = ResponsesRequest(input = "test", max_output_tokens = 512) + assert req.max_output_tokens == 512 + + def test_developer_role(self): + req = ResponsesRequest( + input = [{"role": "developer", "content": "System instructions"}], + ) + assert req.input[0].role == "developer" + + +# ===================================================================== +# Response model tests +# ===================================================================== + + +class TestResponsesResponse: + """Validate response models serialise correctly.""" + + def test_basic_response(self): + resp = ResponsesResponse( + model = "test-model", + output = [ + ResponsesOutputMessage( + content = [ResponsesOutputTextContent(text = "Hello!")] + ), + ], + usage = ResponsesUsage(input_tokens = 10, output_tokens = 5, total_tokens = 15), + ) + d = resp.model_dump() + assert d["object"] == "response" + assert d["status"] == "completed" + assert d["output"][0]["type"] == "message" + assert d["output"][0]["content"][0]["type"] == "output_text" + assert d["output"][0]["content"][0]["text"] == "Hello!" + assert d["usage"]["input_tokens"] == 10 + assert d["usage"]["output_tokens"] == 5 + assert d["usage"]["total_tokens"] == 15 + # Must NOT have prompt_tokens / completion_tokens + assert "prompt_tokens" not in d["usage"] + assert "completion_tokens" not in d["usage"] + + def test_id_format(self): + resp = ResponsesResponse() + assert resp.id.startswith("resp_") + + def test_output_message_id_format(self): + msg = ResponsesOutputMessage() + assert msg.id.startswith("msg_") + + def test_annotations_default_empty(self): + part = ResponsesOutputTextContent(text = "hi") + assert part.annotations == [] + + def test_response_json_roundtrip(self): + resp = ResponsesResponse( + model = "gpt-4", + output = [ + ResponsesOutputMessage( + content = [ResponsesOutputTextContent(text = "ok")], + ), + ], + usage = ResponsesUsage(input_tokens = 1, output_tokens = 1, total_tokens = 2), + ) + j = json.loads(resp.model_dump_json()) + assert j["object"] == "response" + assert j["output"][0]["role"] == "assistant" + assert j["output"][0]["status"] == "completed" + + +# ===================================================================== +# Input normalisation tests +# ===================================================================== + + +class TestNormaliseResponsesInput: + """Test _normalise_responses_input converts Responses input to ChatMessages.""" + + def test_string_input(self): + payload = ResponsesRequest(input = "Hello world") + msgs = _normalise_responses_input(payload) + assert len(msgs) == 1 + assert msgs[0].role == "user" + assert msgs[0].content == "Hello world" + + def test_instructions_become_system_message(self): + payload = ResponsesRequest( + input = "Hi", + instructions = "Be concise.", + ) + msgs = _normalise_responses_input(payload) + assert len(msgs) == 2 + assert msgs[0].role == "system" + assert msgs[0].content == "Be concise." + assert msgs[1].role == "user" + assert msgs[1].content == "Hi" + + def test_message_list(self): + payload = ResponsesRequest( + input = [ + {"role": "user", "content": "First"}, + {"role": "assistant", "content": "Response"}, + {"role": "user", "content": "Second"}, + ], + ) + msgs = _normalise_responses_input(payload) + assert len(msgs) == 3 + assert msgs[0].role == "user" + assert msgs[1].role == "assistant" + assert msgs[2].role == "user" + + def test_developer_role_maps_to_system(self): + payload = ResponsesRequest( + input = [{"role": "developer", "content": "Instructions"}], + ) + msgs = _normalise_responses_input(payload) + assert msgs[0].role == "system" + assert msgs[0].content == "Instructions" + + def test_multimodal_parts(self): + payload = ResponsesRequest( + input = [ + { + "role": "user", + "content": [ + {"type": "input_text", "text": "Describe this:"}, + { + "type": "input_image", + "image_url": "data:image/png;base64,abc", + }, + ], + }, + ], + ) + msgs = _normalise_responses_input(payload) + assert len(msgs) == 1 + content = msgs[0].content + assert isinstance(content, list) + assert len(content) == 2 + assert isinstance(content[0], TextContentPart) + assert content[0].text == "Describe this:" + assert isinstance(content[1], ImageContentPart) + assert content[1].image_url.url == "data:image/png;base64,abc" + + def test_empty_string_input(self): + payload = ResponsesRequest(input = "") + msgs = _normalise_responses_input(payload) + assert len(msgs) == 0 + + def test_empty_list_input(self): + payload = ResponsesRequest(input = []) + msgs = _normalise_responses_input(payload) + assert len(msgs) == 0 + + def test_instructions_only(self): + payload = ResponsesRequest(input = "", instructions = "System msg") + msgs = _normalise_responses_input(payload) + assert len(msgs) == 1 + assert msgs[0].role == "system" + + def test_instructions_plus_message_list(self): + payload = ResponsesRequest( + input = [{"role": "user", "content": "Hello"}], + instructions = "Be brief.", + ) + msgs = _normalise_responses_input(payload) + assert len(msgs) == 2 + assert msgs[0].role == "system" + assert msgs[0].content == "Be brief." + assert msgs[1].role == "user" + + +if __name__ == "__main__": + import pytest + + pytest.main([__file__, "-v"]) diff --git a/studio/backend/tests/test_studio_api.py b/studio/backend/tests/test_studio_api.py new file mode 100644 index 0000000000..9cc17c89fb --- /dev/null +++ b/studio/backend/tests/test_studio_api.py @@ -0,0 +1,643 @@ +# SPDX-License-Identifier: AGPL-3.0-only +# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0 + +""" +End-to-end tests for Unsloth Studio's HTTP API surface. + +Covers the OpenAI-compatible and Anthropic-compatible endpoints exposed +by the server that ``unsloth studio run`` boots, plus API key +authentication and the CLI's ``--help`` output: + + 1. curl -- basic chat completions (non-streaming) + 2. curl -- streaming chat completions + 3. Python OpenAI SDK -- streaming completions + 4. curl -- with tools (web_search + python) + 5. Anthropic Messages API -- basic non-streaming + 6. Anthropic Messages API -- streaming SSE + 7. Anthropic Python SDK -- non-streaming + 8. Anthropic Messages API -- streaming with tools + +Training, export, fine-tuning, and chat-UI concerns are out of scope — +see the unit suites elsewhere under ``studio/backend/tests/`` for those. + +Usage: + + # Script mode — launches its own server via ``unsloth studio run``. + python tests/test_studio_api.py + python tests/test_studio_api.py --model unsloth/... --gguf-variant ... + + # Pytest mode, external server — start a Studio server yourself, + # then point pytest at it. Fastest iteration loop. + unsloth studio run --model unsloth/Qwen3-1.7B-GGUF --gguf-variant UD-Q4_K_XL & + export UNSLOTH_E2E_BASE_URL=http://127.0.0.1:8080 + export UNSLOTH_E2E_API_KEY=sk-unsloth-... # from the server banner + pytest tests/test_studio_api.py -v + + # Pytest mode, fixture-managed server — pytest launches and tears + # down the server itself. One-shot verification, CI-friendly. + pytest tests/test_studio_api.py -v \\ + --unsloth-model unsloth/Qwen3-1.7B-GGUF \\ + --unsloth-gguf-variant UD-Q4_K_XL + +The ``base_url`` / ``api_key`` parameters on the test functions resolve +via the ``studio_server`` session fixture in ``conftest.py``. + +Requires a GPU and ~2 GB of disk for the GGUF download. +""" + +from __future__ import annotations + +import argparse +import json +import os +import re +import signal +import subprocess +import sys +import time +import urllib.error +import urllib.request +from pathlib import Path + + +# ── Configuration ──────────────────────────────────────────────────── + +DEFAULT_MODEL = "unsloth/Qwen3-1.7B-GGUF" +DEFAULT_VARIANT = "UD-Q4_K_XL" +PORT = 18222 # high port unlikely to collide +HOST = "127.0.0.1" +STARTUP_TIMEOUT = 120 # seconds to wait for banner +LOG_FILE = ( + Path(__file__).resolve().parent.parent.parent.parent + / "temp" + / "test_studio_api.log" +) + + +# ── Helpers ────────────────────────────────────────────────────────── + + +def _http( + method: str, + url: str, + *, + body: dict | None = None, + headers: dict | None = None, + timeout: int = 60, +) -> tuple[int, str]: + """Minimal stdlib HTTP helper. Returns (status_code, body_text).""" + data = json.dumps(body).encode() if body else None + req = urllib.request.Request(url, data = data, headers = headers or {}, method = method) + if body: + req.add_header("Content-Type", "application/json") + try: + with urllib.request.urlopen(req, timeout = timeout) as resp: + return resp.status, resp.read().decode() + except urllib.error.HTTPError as exc: + return exc.code, exc.read().decode(errors = "replace") + + +def _stream_http( + url: str, + *, + body: dict, + headers: dict, + timeout: int = 60, +) -> tuple[int, list[dict]]: + """POST a streaming request and collect SSE chunks.""" + data = json.dumps(body).encode() + req = urllib.request.Request(url, data = data, headers = headers, method = "POST") + req.add_header("Content-Type", "application/json") + chunks: list[dict] = [] + try: + with urllib.request.urlopen(req, timeout = timeout) as resp: + status = resp.status + for raw_line in resp: + line = raw_line.decode().strip() + if line.startswith("data: ") and line != "data: [DONE]": + try: + chunks.append(json.loads(line[6:])) + except json.JSONDecodeError: + pass + return status, chunks + except urllib.error.HTTPError as exc: + return exc.code, [] + + +# ── Test functions ─────────────────────────────────────────────────── + + +def test_help_output(): + """``unsloth studio run --help`` should show all documented options.""" + result = subprocess.run( + ["unsloth", "studio", "run", "--help"], + capture_output = True, + text = True, + timeout = 15, + ) + out = result.stdout + assert result.returncode == 0, f"--help exited with {result.returncode}" + + for flag in [ + "--model", + "--gguf-variant", + "--max-seq-length", + "--load-in-4bit", + "--api-key-name", + "--port", + "--host", + "--frontend", + "--silent", + ]: + assert flag in out, f"Missing flag {flag!r} in --help output" + print(" PASS --help shows all flags") + + +def test_curl_basic(base_url: str, api_key: str): + """Example 1: basic non-streaming chat completion via HTTP.""" + status, text = _http( + "POST", + f"{base_url}/v1/chat/completions", + body = { + "messages": [{"role": "user", "content": "Say just the word hello"}], + "stream": False, + }, + headers = {"Authorization": f"Bearer {api_key}"}, + ) + assert status == 200, f"Expected 200, got {status}: {text[:300]}" + data = json.loads(text) + assert "choices" in data, f"Missing 'choices' in response: {text[:300]}" + content = data["choices"][0]["message"]["content"] + assert len(content) > 0, "Empty assistant content" + print(f" PASS curl basic: {content[:80]!r}") + + +def _collect_streamed_content(chunks: list[dict]) -> str: + """Extract text from SSE chunks, skipping role-only and usage chunks.""" + parts = [] + for c in chunks: + choices = c.get("choices", []) + if not choices: + continue + delta = choices[0].get("delta", {}) + part = delta.get("content") + if part: + parts.append(part) + return "".join(parts) + + +def test_curl_streaming(base_url: str, api_key: str): + """Example 2: streaming chat completion via HTTP SSE.""" + status, chunks = _stream_http( + f"{base_url}/v1/chat/completions", + body = { + "messages": [{"role": "user", "content": "Count from 1 to 3"}], + "stream": True, + }, + headers = {"Authorization": f"Bearer {api_key}"}, + ) + assert status == 200, f"Expected 200, got {status}" + assert len(chunks) > 0, "No SSE chunks received" + full = _collect_streamed_content(chunks) + assert len(full) > 0, "Streamed content is empty" + print(f" PASS curl streaming: got {len(chunks)} chunks, {len(full)} chars") + + +def test_openai_sdk(base_url: str, api_key: str): + """Example 3: OpenAI Python SDK streaming completion.""" + try: + from openai import OpenAI + except ImportError: + print(" SKIP openai SDK not installed") + return + + client = OpenAI(base_url = f"{base_url}/v1", api_key = api_key) + response = client.chat.completions.create( + model = "current", + messages = [ + {"role": "user", "content": "What is 2+2? Answer with just the number."} + ], + stream = True, + ) + content_parts = [] + for chunk in response: + if not chunk.choices: + continue + delta_content = chunk.choices[0].delta.content + if delta_content: + content_parts.append(delta_content) + full = "".join(content_parts) + assert len(full) > 0, "OpenAI SDK returned empty content" + print(f" PASS OpenAI SDK streaming: {full.strip()[:80]!r}") + + +def test_curl_with_tools(base_url: str, api_key: str): + """Example 4: chat completion with tool calling enabled. + + Note: when ``enable_tools`` is set the server always returns SSE + streaming regardless of the ``stream`` flag, so we parse SSE chunks. + The model may or may not produce visible content -- tool orchestration + can intercept the response -- so we only assert the endpoint succeeds. + """ + status, chunks = _stream_http( + f"{base_url}/v1/chat/completions", + body = { + "messages": [ + { + "role": "user", + "content": "What is 123 * 456? Use code to compute it.", + } + ], + "stream": True, + "enable_tools": True, + "enabled_tools": ["python"], + "session_id": "test-session", + }, + headers = {"Authorization": f"Bearer {api_key}"}, + timeout = 120, + ) + assert status == 200, f"Expected 200, got {status}" + assert len(chunks) > 0, "No SSE chunks received for tools request" + + # Check that at least one chunk has the expected shape + has_valid_chunk = any("choices" in c or "type" in c for c in chunks) + assert has_valid_chunk, "No valid chunks in tools response" + full = _collect_streamed_content(chunks) + print(f" PASS curl with tools: {len(chunks)} chunks, {len(full)} chars content") + + +def test_invalid_key_rejected(base_url: str): + """Requests with a bad API key should be rejected.""" + status, _text = _http( + "POST", + f"{base_url}/v1/chat/completions", + body = { + "messages": [{"role": "user", "content": "Hello"}], + "stream": False, + }, + headers = {"Authorization": "Bearer sk-unsloth-boguskey123"}, + ) + assert status == 401, f"Expected 401 for invalid key, got {status}" + print(" PASS invalid API key rejected (401)") + + +def test_no_key_rejected(base_url: str): + """Requests without any auth header should be rejected.""" + status, _text = _http( + "POST", + f"{base_url}/v1/chat/completions", + body = { + "messages": [{"role": "user", "content": "Hello"}], + "stream": False, + }, + ) + assert status == 401 or status == 403, f"Expected 401/403 for no key, got {status}" + print(f" PASS no API key rejected ({status})") + + +# ── Anthropic SSE helper ───────────────────────────────────────────── + + +def _stream_anthropic_http( + url: str, + *, + body: dict, + headers: dict, + timeout: int = 60, +) -> tuple[int, list[tuple[str, dict]]]: + """POST a streaming request and collect Anthropic SSE events. + + Returns (status, [(event_type, data_dict), ...]). + """ + data = json.dumps(body).encode() + req = urllib.request.Request(url, data = data, headers = headers, method = "POST") + req.add_header("Content-Type", "application/json") + events: list[tuple[str, dict]] = [] + try: + with urllib.request.urlopen(req, timeout = timeout) as resp: + status = resp.status + current_event = None + for raw_line in resp: + line = raw_line.decode().strip() + if line.startswith("event: "): + current_event = line[7:] + elif line.startswith("data: ") and current_event: + try: + events.append((current_event, json.loads(line[6:]))) + except json.JSONDecodeError: + pass + current_event = None + return status, events + except urllib.error.HTTPError as exc: + return exc.code, [] + + +def _collect_anthropic_text(events: list[tuple[str, dict]]) -> str: + """Extract text content from Anthropic SSE events.""" + parts = [] + for etype, data in events: + if etype == "content_block_delta": + delta = data.get("delta", {}) + if delta.get("type") == "text_delta": + parts.append(delta.get("text", "")) + return "".join(parts) + + +# ── Anthropic /v1/messages test functions ──────────────────────────── + + +def test_anthropic_basic(base_url: str, api_key: str): + """Anthropic Messages API: non-streaming.""" + status, text = _http( + "POST", + f"{base_url}/v1/messages", + body = { + "model": "default", + "max_tokens": 100, + "messages": [{"role": "user", "content": "Say just the word hello"}], + }, + headers = {"Authorization": f"Bearer {api_key}"}, + ) + assert status == 200, f"Expected 200, got {status}: {text[:300]}" + data = json.loads(text) + assert data.get("type") == "message", f"Expected type 'message': {text[:300]}" + assert data.get("role") == "assistant" + content = data.get("content", []) + assert len(content) > 0, "Empty content array" + text_block = content[-1] + assert text_block.get("type") == "text", f"Expected text block: {text_block}" + assert len(text_block.get("text", "")) > 0, "Empty text in response" + print(f" PASS anthropic basic: {text_block['text'][:80]!r}") + + +def test_anthropic_streaming(base_url: str, api_key: str): + """Anthropic Messages API: streaming SSE.""" + status, events = _stream_anthropic_http( + f"{base_url}/v1/messages", + body = { + "model": "default", + "max_tokens": 100, + "messages": [{"role": "user", "content": "Count from 1 to 3"}], + "stream": True, + }, + headers = {"Authorization": f"Bearer {api_key}"}, + ) + assert status == 200, f"Expected 200, got {status}" + assert len(events) > 0, "No SSE events received" + + event_types = [e[0] for e in events] + assert "message_start" in event_types, "Missing message_start event" + assert "message_stop" in event_types, "Missing message_stop event" + + full = _collect_anthropic_text(events) + assert len(full) > 0, "Streamed text content is empty" + print(f" PASS anthropic streaming: {len(events)} events, {len(full)} chars") + + +def test_anthropic_sdk(base_url: str, api_key: str): + """Anthropic Python SDK: non-streaming.""" + try: + from anthropic import Anthropic + except ImportError: + print(" SKIP anthropic SDK not installed") + return + + client = Anthropic(base_url = f"{base_url}/v1", api_key = api_key) + message = client.messages.create( + model = "default", + max_tokens = 100, + messages = [ + {"role": "user", "content": "What is 2+2? Answer with just the number."} + ], + ) + assert message.role == "assistant" + assert len(message.content) > 0, "Empty content" + text = message.content[0].text + assert len(text) > 0, "Empty text" + print(f" PASS Anthropic SDK: {text.strip()[:80]!r}") + + +def test_anthropic_with_tools(base_url: str, api_key: str): + """Anthropic Messages API: streaming with tools.""" + status, events = _stream_anthropic_http( + f"{base_url}/v1/messages", + body = { + "model": "default", + "max_tokens": 1024, + "messages": [ + { + "role": "user", + "content": "What is 123 * 456? Use code to compute it.", + } + ], + "tools": [ + { + "name": "python", + "description": "Execute Python code in a sandbox and return stdout/stderr.", + "input_schema": { + "type": "object", + "properties": { + "code": { + "type": "string", + "description": "The Python code to run", + }, + }, + "required": ["code"], + }, + } + ], + "stream": True, + }, + headers = {"Authorization": f"Bearer {api_key}"}, + timeout = 120, + ) + assert status == 200, f"Expected 200, got {status}" + assert len(events) > 0, "No SSE events received for tools request" + + event_types = [e[0] for e in events] + assert "message_start" in event_types, "Missing message_start" + assert "message_stop" in event_types, "Missing message_stop" + + full = _collect_anthropic_text(events) + print( + f" PASS anthropic with tools: {len(events)} events, {len(full)} chars content" + ) + + +# ── Server lifecycle ───────────────────────────────────────────────── + + +def _start_server(model: str, variant: str | None) -> tuple[subprocess.Popen, str]: + """Launch ``unsloth studio run`` and parse the API key from its banner. + + Returns (process, api_key). + """ + cmd = [ + "unsloth", + "studio", + "run", + "--model", + model, + "--port", + str(PORT), + "--host", + HOST, + "--api-key-name", + "test", + ] + if variant: + cmd.extend(["--gguf-variant", variant]) + + LOG_FILE.parent.mkdir(parents = True, exist_ok = True) + log_fh = open(LOG_FILE, "w") + proc = subprocess.Popen( + cmd, + stdout = log_fh, + stderr = subprocess.STDOUT, + preexec_fn = os.setsid, + ) + + # Wait for the banner containing the API key + api_key = None + deadline = time.monotonic() + STARTUP_TIMEOUT + while time.monotonic() < deadline: + time.sleep(2) + if proc.poll() is not None: + log_fh.flush() + log_text = LOG_FILE.read_text() + raise RuntimeError( + f"Server exited early (code {proc.returncode}):\n{log_text[-2000:]}" + ) + log_text = LOG_FILE.read_text() + m = re.search(r"API Key:\s+(sk-unsloth-[a-f0-9]+)", log_text) + if m: + api_key = m.group(1) + break + + if not api_key: + log_text = LOG_FILE.read_text() + _kill_server(proc) + raise RuntimeError( + f"Timed out waiting for API key in server output:\n{log_text[-2000:]}" + ) + + # Wait a moment for the model to be fully loaded + time.sleep(2) + return proc, api_key + + +def _kill_server(proc: subprocess.Popen): + """Send SIGTERM to the process group and wait for cleanup.""" + try: + os.killpg(os.getpgid(proc.pid), signal.SIGTERM) + except (ProcessLookupError, PermissionError): + pass + try: + proc.wait(timeout = 10) + except subprocess.TimeoutExpired: + try: + os.killpg(os.getpgid(proc.pid), signal.SIGKILL) + except (ProcessLookupError, PermissionError): + pass + proc.wait(timeout = 5) + + +# ── Main ───────────────────────────────────────────────────────────── + + +def main(): + parser = argparse.ArgumentParser( + description = "End-to-end tests for unsloth studio run" + ) + parser.add_argument( + "--model", + default = DEFAULT_MODEL, + help = f"Model to test with (default: {DEFAULT_MODEL})", + ) + parser.add_argument( + "--gguf-variant", + default = DEFAULT_VARIANT, + help = f"GGUF variant (default: {DEFAULT_VARIANT})", + ) + args = parser.parse_args() + + passed = 0 + failed = 0 + skipped = 0 + + def run_test(fn, *a, **kw): + nonlocal passed, failed, skipped + try: + fn(*a, **kw) + passed += 1 + except AssertionError as exc: + failed += 1 + print(f" FAIL {fn.__name__}: {exc}") + except Exception as exc: + failed += 1 + print(f" ERROR {fn.__name__}: {type(exc).__name__}: {exc}") + + # ── 1. Test --help (no server needed) ──────────────────────────── + print("\n[1/11] Testing --help output") + run_test(test_help_output) + + # ── 2-11. Start server and run API tests ───────────────────────── + print( + f"\nStarting server: {args.model} (variant={args.gguf_variant}) on port {PORT}..." + ) + proc = None + try: + proc, api_key = _start_server(args.model, args.gguf_variant) + base_url = f"http://{HOST}:{PORT}" + print(f"Server ready. API Key: {api_key[:20]}...\n") + + print("[2/11] Testing curl basic (non-streaming)") + run_test(test_curl_basic, base_url, api_key) + + print("[3/11] Testing curl streaming") + run_test(test_curl_streaming, base_url, api_key) + + print("[4/11] Testing OpenAI Python SDK (streaming)") + run_test(test_openai_sdk, base_url, api_key) + + print("[5/11] Testing curl with tools") + run_test(test_curl_with_tools, base_url, api_key) + + print("[6/11] Testing invalid API key rejection") + run_test(test_invalid_key_rejected, base_url) + + print("[7/11] Testing no API key rejection") + run_test(test_no_key_rejected, base_url) + + print("[8/11] Testing Anthropic basic (non-streaming)") + run_test(test_anthropic_basic, base_url, api_key) + + print("[9/11] Testing Anthropic streaming") + run_test(test_anthropic_streaming, base_url, api_key) + + print("[10/11] Testing Anthropic Python SDK") + run_test(test_anthropic_sdk, base_url, api_key) + + print("[11/11] Testing Anthropic with tools") + run_test(test_anthropic_with_tools, base_url, api_key) + + except RuntimeError as exc: + print(f"\nFATAL: Server failed to start: {exc}") + failed += 11 # count remaining tests as failed + finally: + if proc: + print("\nStopping server...") + _kill_server(proc) + print("Server stopped.") + + # ── Summary ────────────────────────────────────────────────────── + total = passed + failed + print(f"\n{'=' * 40}") + print(f"Results: {passed}/{total} passed, {failed} failed") + print(f"Log: {LOG_FILE}") + print(f"{'=' * 40}") + sys.exit(1 if failed else 0) + + +if __name__ == "__main__": + main() diff --git a/studio/backend/tests/test_trained_model_scan.py b/studio/backend/tests/test_trained_model_scan.py new file mode 100644 index 0000000000..84be681fca --- /dev/null +++ b/studio/backend/tests/test_trained_model_scan.py @@ -0,0 +1,101 @@ +# 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 Studio trained-model discovery used by Chat.""" + +import json +from pathlib import Path +import sys +import types as _types +import importlib + + +_BACKEND_DIR = str(Path(__file__).resolve().parent.parent) +if _BACKEND_DIR not in sys.path: + sys.path.insert(0, _BACKEND_DIR) + +_loggers_stub = _types.ModuleType("loggers") +_loggers_stub.get_logger = lambda name: __import__("logging").getLogger(name) +sys.modules.setdefault("loggers", _loggers_stub) + +from unittest.mock import patch + +from utils.models.model_config import ( + ModelConfig, + get_base_model_from_checkpoint, + get_base_model_from_lora, + scan_trained_models, +) + + +def test_scan_trained_models_includes_lora_and_full_finetune_outputs(tmp_path: Path): + lora_dir = tmp_path / "unsloth_SmolLM-135M_1775412608" + lora_dir.mkdir() + (lora_dir / "adapter_config.json").write_text( + json.dumps({"base_model_name_or_path": "HuggingFaceTB/SmolLM-135M"}) + ) + (lora_dir / "adapter_model.safetensors").write_bytes(b"") + + full_dir = tmp_path / "unsloth_SmolLM-135M_full_1775412609" + full_dir.mkdir() + (full_dir / "config.json").write_text( + json.dumps({"_name_or_path": "HuggingFaceTB/SmolLM-135M"}) + ) + (full_dir / "model.safetensors").write_bytes(b"") + + found = { + name: (path, model_type) + for name, path, model_type in scan_trained_models(str(tmp_path)) + } + + assert found[lora_dir.name] == (str(lora_dir), "lora") + assert found[full_dir.name] == (str(full_dir), "merged") + + +def test_get_base_model_from_checkpoint_falls_back_to_full_finetune_config( + tmp_path: Path, +): + (tmp_path / "config.json").write_text( + json.dumps({"_name_or_path": "HuggingFaceTB/SmolLM-135M"}) + ) + (tmp_path / "model.safetensors").write_bytes(b"") + + assert get_base_model_from_checkpoint(str(tmp_path)) == "HuggingFaceTB/SmolLM-135M" + + +def test_get_base_model_from_lora_rejects_full_finetune_dirs(tmp_path: Path): + (tmp_path / "config.json").write_text( + json.dumps({"_name_or_path": "HuggingFaceTB/SmolLM-135M"}) + ) + (tmp_path / "model.safetensors").write_bytes(b"") + + assert get_base_model_from_lora(str(tmp_path)) is None + + +@patch("utils.models.model_config.is_audio_input_type", return_value = False) +@patch("utils.models.model_config.detect_audio_type", return_value = None) +@patch("utils.models.model_config.is_vision_model", return_value = False) +def test_model_config_full_finetune_local_path_is_not_lora( + _mock_vision, + _mock_audio_type, + _mock_audio_input, + tmp_path: Path, +): + (tmp_path / "config.json").write_text( + json.dumps({"_name_or_path": "unsloth/Qwen3-4B"}) + ) + (tmp_path / "model.safetensors").write_bytes(b"") + + config = ModelConfig.from_identifier(str(tmp_path)) + + assert config is not None + assert config.is_lora is False + assert config.base_model is None + + +def test_scan_trained_loras_aliases_scan_trained_models(): + utils_models = importlib.import_module("utils.models") + core_module = importlib.import_module("core") + + assert utils_models.scan_trained_loras is utils_models.scan_trained_models + assert core_module.scan_trained_loras is core_module.scan_trained_models diff --git a/studio/backend/tests/test_training_worker_flash_attn.py b/studio/backend/tests/test_training_worker_flash_attn.py new file mode 100644 index 0000000000..986958408e --- /dev/null +++ b/studio/backend/tests/test_training_worker_flash_attn.py @@ -0,0 +1,154 @@ +# SPDX-License-Identifier: AGPL-3.0-only +# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0 + +from __future__ import annotations + +import builtins +import subprocess +import sys +from unittest import mock + +from core.training import worker + + +def _missing_flash_attn_import(): + real_import = builtins.__import__ + + def fake_import(name, globals = None, locals = None, fromlist = (), level = 0): + if name == "flash_attn": + raise ImportError + return real_import(name, globals, locals, fromlist, level) + + return fake_import + + +def test_should_try_runtime_flash_attn_install_threshold_and_skip(monkeypatch): + monkeypatch.delenv(worker._FLASH_ATTN_SKIP_ENV, raising = False) + assert worker._should_try_runtime_flash_attn_install(32767) is False + assert worker._should_try_runtime_flash_attn_install( + 32768 + ) is sys.platform.startswith("linux") + + monkeypatch.setenv(worker._FLASH_ATTN_SKIP_ENV, "1") + assert worker._should_try_runtime_flash_attn_install(32768) is False + + +def test_runtime_flash_attn_prefers_prebuilt_wheel(monkeypatch): + statuses: list[str] = [] + + monkeypatch.delenv(worker._FLASH_ATTN_SKIP_ENV, raising = False) + monkeypatch.setattr(builtins, "__import__", _missing_flash_attn_import()) + monkeypatch.setattr( + worker, + "flash_attn_wheel_url", + lambda env: "https://example.com/fa.whl", + ) + monkeypatch.setattr(worker, "url_exists", lambda url: True) + monkeypatch.setattr( + worker, + "_send_status", + lambda queue, message: statuses.append(message), + ) + monkeypatch.setattr( + worker, + "install_wheel", + lambda *args, **kwargs: [("pip", subprocess.CompletedProcess(["pip"], 0, ""))], + ) + + worker._ensure_flash_attn_for_long_context(event_queue = [], max_seq_length = 32768) + + assert statuses == ["Installing prebuilt flash-attn wheel..."] + + +def test_runtime_flash_attn_falls_back_to_pypi(monkeypatch): + calls: list[list[str]] = [] + statuses: list[str] = [] + + monkeypatch.delenv(worker._FLASH_ATTN_SKIP_ENV, raising = False) + monkeypatch.setattr(builtins, "__import__", _missing_flash_attn_import()) + monkeypatch.setattr( + worker, + "probe_torch_wheel_env", + lambda timeout = 30: { + "python_tag": "cp313", + "torch_mm": "2.10", + "cuda_major": "13", + "cxx11abi": "TRUE", + "platform_tag": "linux_x86_64", + }, + ) + monkeypatch.setattr( + worker, + "flash_attn_wheel_url", + lambda env: "https://example.com/fa.whl", + ) + monkeypatch.setattr(worker, "url_exists", lambda url: False) + monkeypatch.setattr(worker.shutil, "which", lambda name: None) + monkeypatch.setattr( + worker, + "_send_status", + lambda queue, message: statuses.append(message), + ) + monkeypatch.setattr(worker, "install_wheel", mock.Mock()) + + def fake_run(cmd, stdout = None, stderr = None, text = None): + calls.append(list(cmd)) + return subprocess.CompletedProcess(cmd, 0, "") + + monkeypatch.setattr(worker._sp, "run", fake_run) + + worker._ensure_flash_attn_for_long_context(event_queue = [], max_seq_length = 32768) + + assert statuses == ["Installing flash-attn from PyPI for long-context training..."] + assert calls == [[sys.executable, "-m", "pip", "install", "flash-attn"]] + + +def test_runtime_flash_attn_skip_env_avoids_all_install_work(monkeypatch): + monkeypatch.setenv(worker._FLASH_ATTN_SKIP_ENV, "1") + monkeypatch.setattr(worker._sp, "run", mock.Mock()) + + worker._ensure_flash_attn_for_long_context(event_queue = [], max_seq_length = 32768) + + worker._sp.run.assert_not_called() + + +def test_causal_conv1d_fast_path_preserves_wheel_first_install_args(monkeypatch): + install_mock = mock.Mock(return_value = True) + monkeypatch.setattr(worker, "_install_package_wheel_first", install_mock) + + worker._ensure_causal_conv1d_fast_path( + event_queue = [], + model_name = "tiiuae/Falcon-H1-0.5B-Instruct", + ) + + install_mock.assert_called_once_with( + event_queue = [], + import_name = "causal_conv1d", + display_name = "causal-conv1d", + pypi_name = "causal-conv1d", + pypi_version = worker._CAUSAL_CONV1D_PACKAGE_VERSION, + filename_prefix = "causal_conv1d", + release_tag = worker._CAUSAL_CONV1D_RELEASE_TAG, + release_base_url = "https://github.com/Dao-AILab/causal-conv1d/releases/download", + ) + + +def test_mamba_ssm_path_preserves_wheel_first_install_args(monkeypatch): + install_mock = mock.Mock(return_value = True) + monkeypatch.setattr(worker, "_install_package_wheel_first", install_mock) + + worker._ensure_mamba_ssm( + event_queue = [], + model_name = "tiiuae/Falcon-H1-0.5B-Instruct", + ) + + install_mock.assert_called_once_with( + event_queue = [], + import_name = "mamba_ssm", + display_name = "mamba-ssm", + pypi_name = "mamba-ssm", + pypi_version = worker._MAMBA_SSM_PACKAGE_VERSION, + filename_prefix = "mamba_ssm", + release_tag = worker._MAMBA_SSM_RELEASE_TAG, + release_base_url = "https://github.com/state-spaces/mamba/releases/download", + ) diff --git a/studio/backend/tests/test_transformers_version.py b/studio/backend/tests/test_transformers_version.py index f3dae537c7..c031c2fea3 100644 --- a/studio/backend/tests/test_transformers_version.py +++ b/studio/backend/tests/test_transformers_version.py @@ -31,8 +31,11 @@ sys.modules.setdefault("loggers", _loggers_stub) from utils.transformers_version import ( _resolve_base_model, _check_tokenizer_config_needs_v5, + _check_config_needs_550, _tokenizer_class_cache, + _config_needs_550_cache, needs_transformers_5, + get_transformers_tier, ) @@ -188,3 +191,148 @@ class TestNeedsTransformers5: # We test the full resolution chain here: resolved = _resolve_base_model(str(tmp_path)) assert needs_transformers_5(resolved) is True + + +# --------------------------------------------------------------------------- +# _check_config_needs_550 — config.json architecture/model_type check +# --------------------------------------------------------------------------- + + +class TestCheckConfigNeeds550: + """Tests for _check_config_needs_550() local config.json checks.""" + + def setup_method(self): + _config_needs_550_cache.clear() + + def test_gemma4_architecture(self, tmp_path: Path): + """config.json with Gemma4ForConditionalGeneration should return True.""" + cfg = { + "architectures": ["Gemma4ForConditionalGeneration"], + "model_type": "gemma4", + } + (tmp_path / "config.json").write_text(json.dumps(cfg)) + + assert _check_config_needs_550(str(tmp_path)) is True + + def test_gemma4_model_type_only(self, tmp_path: Path): + """config.json with model_type=gemma4 (no architectures) should return True.""" + cfg = {"model_type": "gemma4"} + (tmp_path / "config.json").write_text(json.dumps(cfg)) + + assert _check_config_needs_550(str(tmp_path)) is True + + def test_llama_architecture(self, tmp_path: Path): + """config.json with LlamaForCausalLM should return False.""" + cfg = {"architectures": ["LlamaForCausalLM"], "model_type": "llama"} + (tmp_path / "config.json").write_text(json.dumps(cfg)) + + assert _check_config_needs_550(str(tmp_path)) is False + + def test_no_config_json(self, tmp_path: Path): + """Missing config.json should return False (fail-open).""" + # Patch network call to avoid real fetch + with patch("urllib.request.urlopen") as mock_urlopen: + mock_urlopen.side_effect = Exception("no network") + assert _check_config_needs_550(str(tmp_path)) is False + + def test_result_is_cached(self, tmp_path: Path): + """Subsequent calls should use the cache.""" + cfg = {"architectures": ["Gemma4ForConditionalGeneration"]} + (tmp_path / "config.json").write_text(json.dumps(cfg)) + + key = str(tmp_path) + _check_config_needs_550(key) + assert key in _config_needs_550_cache + assert _config_needs_550_cache[key] is True + + def test_local_file_skips_network(self, tmp_path: Path): + """When local config.json exists, no network request should be made.""" + cfg = {"architectures": ["LlamaForCausalLM"]} + (tmp_path / "config.json").write_text(json.dumps(cfg)) + + with patch("urllib.request.urlopen") as mock_urlopen: + _check_config_needs_550(str(tmp_path)) + mock_urlopen.assert_not_called() + + +# --------------------------------------------------------------------------- +# get_transformers_tier — tier detection +# --------------------------------------------------------------------------- + + +class TestGetTransformersTier: + """Tests for get_transformers_tier() tiered version detection.""" + + def setup_method(self): + _tokenizer_class_cache.clear() + _config_needs_550_cache.clear() + + def test_gemma4_substring_returns_550(self): + assert get_transformers_tier("google/gemma-4-E2B-it") == "550" + + def test_gemma4_alt_substring_returns_550(self): + assert get_transformers_tier("unsloth/gemma4-E4B-it") == "550" + + def test_gemma4_config_json_returns_550(self, tmp_path: Path): + """Local checkpoint with Gemma4 architecture → 550.""" + cfg = { + "architectures": ["Gemma4ForConditionalGeneration"], + "model_type": "gemma4", + } + (tmp_path / "config.json").write_text(json.dumps(cfg)) + + assert get_transformers_tier(str(tmp_path)) == "550" + + def test_qwen35_returns_530(self): + with patch( + "utils.transformers_version._check_config_needs_550", + return_value = False, + ): + assert get_transformers_tier("Qwen/Qwen3.5-9B") == "530" + + def test_ministral_returns_530(self): + with patch( + "utils.transformers_version._check_config_needs_550", + return_value = False, + ): + assert ( + get_transformers_tier("mistralai/Ministral-3-8B-Instruct-2512") == "530" + ) + + def test_llama_returns_default(self): + with ( + patch( + "utils.transformers_version._check_config_needs_550", + return_value = False, + ), + patch( + "utils.transformers_version._check_tokenizer_config_needs_v5", + return_value = False, + ), + ): + assert get_transformers_tier("meta-llama/Llama-3-8B") == "default" + + def test_550_checked_before_530(self): + """Ensure 5.5.0 is checked first — a model matching both should get 550.""" + # This shouldn't happen in practice, but verifies priority + assert get_transformers_tier("gemma-4-model") == "550" + + def test_needs_transformers_5_compat(self): + """needs_transformers_5 should return True for both 530 and 550 models.""" + assert needs_transformers_5("google/gemma-4-E2B-it") is True + with patch( + "utils.transformers_version._check_config_needs_550", + return_value = False, + ): + assert needs_transformers_5("Qwen/Qwen3.5-9B") is True + with ( + patch( + "utils.transformers_version._check_config_needs_550", + return_value = False, + ), + patch( + "utils.transformers_version._check_tokenizer_config_needs_v5", + return_value = False, + ), + ): + assert needs_transformers_5("meta-llama/Llama-3-8B") is False diff --git a/studio/backend/tests/test_utils.py b/studio/backend/tests/test_utils.py index 50557c6718..64c9907119 100644 --- a/studio/backend/tests/test_utils.py +++ b/studio/backend/tests/test_utils.py @@ -191,8 +191,14 @@ class TestGetGpuMemoryInfo: assert "backend" in get_gpu_memory_info() def test_backend_matches_device(self): + # The backend field uses _backend_label, which swaps "cuda" for + # "rocm" when running on an AMD host (IS_ROCM=True) so the UI + # can render the correct label. On CUDA / XPU / MLX / CPU hosts + # it is equivalent to `get_device().value`. + from utils.hardware.hardware import _backend_label + result = get_gpu_memory_info() - assert result["backend"] == get_device().value + assert result["backend"] == _backend_label(get_device()) # --- When a GPU IS available --- diff --git a/studio/backend/utils/datasets/model_mappings.py b/studio/backend/utils/datasets/model_mappings.py index 95b4791574..7fcac637c6 100644 --- a/studio/backend/utils/datasets/model_mappings.py +++ b/studio/backend/utils/datasets/model_mappings.py @@ -215,6 +215,21 @@ TEMPLATE_TO_MODEL_MAPPER = { "google/gemma-3n-E2B-it", "unsloth/gemma-3n-E2B-it-unsloth-bnb-4bit", ), + "gemma-4": ( + "unsloth/gemma-4-E2B-it", + "google/gemma-4-E2B-it", + "unsloth/gemma-4-E4B-it", + "google/gemma-4-E4B-it", + "unsloth/gemma-4-E2B-it-unsloth-bnb-4bit", + "unsloth/gemma-4-E4B-it-unsloth-bnb-4bit", + ), + "gemma-4-thinking": ( + "unsloth/gemma-4-26B-A4B-it", + "google/gemma-4-26B-A4B-it", + "unsloth/gemma-4-31B-it", + "unsloth/gemma-4-31B-it-unsloth-bnb-4bit", + "google/gemma-4-31B-it", + ), "qwen2.5": ( "unsloth/Qwen2.5-0.5B-Instruct-unsloth-bnb-4bit", "unsloth/Qwen2.5-0.5B-Instruct", @@ -399,6 +414,15 @@ TEMPLATE_TO_MODEL_MAPPER = { "THUDM/GLM-4.7-Flash", "unsloth/GLM-4.7-Flash-bnb-4bit", ), + "lfm-2": ( + "unsloth/LFM2-1.2B", + "LiquidAI/LFM2-1.2B", + "unsloth/LFM2-1.2B-unsloth-bnb-4bit", + ), + "lfm-2.5": ( + "unsloth/LFM2.5-1.2B-Instruct", + "LiquidAI/LFM2.5-1.2B-Instruct", + ), } MODEL_TO_TEMPLATE_MAPPER = {} @@ -414,6 +438,14 @@ for key, values in TEMPLATE_TO_MODEL_MAPPER.items(): TEMPLATE_TO_RESPONSES_MAPPER = { + "gemma-4-thinking": { + "instruction": "<|turn>user\n", + "response": "<|turn>model\n", + }, + "gemma-4": { + "instruction": "<|turn>user\n", + "response": "<|turn>model\n", + }, "gemma-3": { "instruction": "user\n", "response": "model\n", @@ -514,6 +546,10 @@ TEMPLATE_TO_RESPONSES_MAPPER = { "instruction": "<|im_start|>user\n", "response": "<|im_start|>assistant\n", }, + "lfm-2.5": { + "instruction": "<|im_start|>user\n", + "response": "<|im_start|>assistant\n", + }, "starling": { "instruction": "GPT4 Correct User: ", "response": "GPT4 Correct Assistant: ", diff --git a/studio/backend/utils/hardware/__init__.py b/studio/backend/utils/hardware/__init__.py index aaa0452406..400b5dd066 100644 --- a/studio/backend/utils/hardware/__init__.py +++ b/studio/backend/utils/hardware/__init__.py @@ -5,6 +5,7 @@ Hardware detection and GPU utilities """ +from . import hardware as _hardware from .hardware import ( DeviceType, DEVICE, @@ -49,6 +50,7 @@ __all__ = [ "DeviceType", "DEVICE", "CHAT_ONLY", + "IS_ROCM", "detect_hardware", "get_device", "is_apple_silicon", @@ -81,3 +83,11 @@ __all__ = [ "extract_arch_config", "estimate_training_vram", ] + + +def __getattr__(name: str): + """Resolve IS_ROCM at access time so callers always see the live value + after detect_hardware() runs (it flips the flag in hardware.py).""" + if name == "IS_ROCM": + return getattr(_hardware, "IS_ROCM") + raise AttributeError(name) diff --git a/studio/backend/utils/hardware/amd.py b/studio/backend/utils/hardware/amd.py new file mode 100644 index 0000000000..755314ca3a --- /dev/null +++ b/studio/backend/utils/hardware/amd.py @@ -0,0 +1,382 @@ +# SPDX-License-Identifier: AGPL-3.0-only +# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0 + +"""AMD GPU monitoring via amd-smi. + +Mirrors the nvidia.py module structure so hardware.py can swap backends +based on IS_ROCM. All functions return the same dict shapes as their +nvidia.py counterparts. +""" + +import json +import math +import os +import re +import subprocess +from typing import Any, Optional + +from loggers import get_logger + +logger = get_logger(__name__) + + +def _run_amd_smi(*args: str, timeout: int = 5) -> Optional[Any]: + """Run amd-smi with the given arguments and return parsed JSON, or None.""" + try: + result = subprocess.run( + ["amd-smi", *args, "--json"], + capture_output = True, + text = True, + timeout = timeout, + ) + except (OSError, subprocess.TimeoutExpired) as e: + logger.warning("amd-smi query failed: %s", e) + return None + if result.returncode != 0 or not result.stdout.strip(): + logger.warning("amd-smi returned code %d", result.returncode) + return None + try: + return json.loads(result.stdout) + except json.JSONDecodeError: + logger.warning("Failed to parse amd-smi JSON output") + return None + + +def _parse_numeric(value: Any) -> Optional[float]: + """Extract a numeric value from amd-smi output (may be str, int, float, or dict).""" + if value is None: + return None + # Newer amd-smi versions emit {"value": 10, "unit": "W"} + if isinstance(value, dict): + return _parse_numeric(value.get("value")) + if isinstance(value, (int, float)): + f = float(value) + return f if math.isfinite(f) else None + if isinstance(value, str): + # Strip units like "W", "C", "%", "MB", "MiB", "GB", "GiB" etc. + cleaned = re.sub(r"\s*[A-Za-z/%]+$", "", value.strip()) + if not cleaned or cleaned.lower() in ("n/a", "none", "unknown"): + return None + try: + return float(cleaned) + except (ValueError, TypeError): + return None + return None + + +def _parse_memory_mb(value: Any) -> Optional[float]: + """Parse a memory value from amd-smi output and return MB. + + Handles bare numbers (assumed MB -- the amd-smi convention on every + version we have seen), dict-shaped values with explicit units + (``{"value": 192, "unit": "GiB"}`` on newer releases), and plain + strings like ``"8192 MiB"``. + """ + unit = "" + raw_value = value + + if isinstance(value, dict): + unit = str(value.get("unit", "")).strip().lower() + raw_value = value.get("value") + elif isinstance(value, str): + # Extract unit suffix from strings like "192 GiB" or "8192 MB" + m = re.match(r"^\s*([\d.]+)\s*([A-Za-z]+)\s*$", value.strip()) + if m: + unit = m.group(2).lower() + + num = _parse_numeric(raw_value if isinstance(value, dict) else value) + if num is None: + return None + + # Unit conversion -- GPU tools (including amd-smi) use binary units even + # when labeling them "GB" or "MB", so treat GB/GiB and MB/MiB the same. + if "gib" in unit or "gb" in unit: + return num * 1024 + if "mib" in unit or "mb" in unit: + return num + if "kib" in unit or "kb" in unit: + return num / 1024 + if unit in ("b", "byte", "bytes"): + # Plain bytes + return num / (1024 * 1024) + + # No explicit unit -- default to MB, which is the amd-smi convention + # for bare numeric values. A previous heuristic assumed values above + # ~10M were bytes, but that misclassifies small VRAM allocations + # (e.g. 5 MB = 5,242,880 reported without a unit) as ~5 TB. Modern + # amd-smi always ships explicit units, so the heuristic branch only + # fired for legacy output where MB was already the convention. + return num + + +def _extract_gpu_metrics(gpu_data: dict) -> dict[str, Any]: + """Extract standardized metrics from a single GPU's amd-smi data.""" + # amd-smi metric output structure varies by version; try common paths + usage = gpu_data.get("usage", gpu_data.get("gpu_activity", {})) + if isinstance(usage, dict): + gpu_util = _parse_numeric( + usage.get("gfx_activity", usage.get("gpu_use_percent")) + ) + else: + gpu_util = _parse_numeric(usage) + + # Temperature -- try multiple keys in priority order. + # dict.get() returns "N/A" strings rather than falling through, + # so we must try each key and check if it parses to a real number. + temp_data = gpu_data.get("temperature", {}) + temp = None + if isinstance(temp_data, dict): + for temp_key in ("edge", "temperature_edge", "hotspot", "temperature_hotspot"): + temp = _parse_numeric(temp_data.get(temp_key)) + if temp is not None: + break + else: + temp = _parse_numeric(temp_data) + + # Power + power_data = gpu_data.get("power", {}) + if isinstance(power_data, dict): + power_draw = _parse_numeric( + power_data.get( + "current_socket_power", + power_data.get("average_socket_power", power_data.get("socket_power")), + ) + ) + power_limit = _parse_numeric( + power_data.get("power_cap", power_data.get("max_power_limit")) + ) + else: + power_draw = None + power_limit = None + + # VRAM -- unit-aware parsing to handle varying amd-smi output formats. + # Newer amd-smi versions may return {"value": 192, "unit": "GiB"}. + # Newer amd-smi uses "mem_usage" with "total_vram" / "used_vram" keys; + # older versions use "vram" or "fb_memory_usage" with "used" / "total". + vram_data = gpu_data.get( + "mem_usage", + gpu_data.get("vram", gpu_data.get("fb_memory_usage", {})), + ) + if isinstance(vram_data, dict): + vram_used_mb = _parse_memory_mb( + vram_data.get( + "used_vram", vram_data.get("vram_used", vram_data.get("used")) + ) + ) + vram_total_mb = _parse_memory_mb( + vram_data.get( + "total_vram", vram_data.get("vram_total", vram_data.get("total")) + ) + ) + else: + vram_used_mb = None + vram_total_mb = None + + # Build the standardized dict (same shape as nvidia._build_gpu_metrics) + vram_used_gb = round(vram_used_mb / 1024, 2) if vram_used_mb is not None else None + vram_total_gb = ( + round(vram_total_mb / 1024, 2) if vram_total_mb is not None else None + ) + vram_util = ( + round((vram_used_mb / vram_total_mb) * 100, 1) + if vram_used_mb is not None and vram_total_mb is not None and vram_total_mb > 0 + else None + ) + power_util = ( + round((power_draw / power_limit) * 100, 1) + if power_draw is not None and power_limit is not None and power_limit > 0 + else None + ) + + return { + "gpu_utilization_pct": gpu_util, + "temperature_c": temp, + "vram_used_gb": vram_used_gb, + "vram_total_gb": vram_total_gb, + "vram_utilization_pct": vram_util, + "power_draw_w": power_draw, + "power_limit_w": power_limit, + "power_utilization_pct": power_util, + } + + +def _has_real_metrics(metrics: dict[str, Any]) -> bool: + """Return True when ``metrics`` contains at least one non-None value. + + ``amd-smi`` can return a zero-exit JSON envelope that is missing every + expected field (error response, unsupported card, hipless container). + In that case ``_extract_gpu_metrics`` produces a dict where every value + is ``None`` -- callers must surface this as ``available: False`` rather + than ``available: True`` with empty data. + """ + return any(value is not None for value in metrics.values()) + + +def get_physical_gpu_count() -> Optional[int]: + """Return physical AMD GPU count via amd-smi, or None on failure.""" + data = _run_amd_smi("list") + if data is None: + return None + if isinstance(data, list): + return len(data) + # Some versions return a dict with a "gpu" / "gpus" key. Guard the + # .get() access with an isinstance check so a malformed scalar / + # string response from amd-smi cannot raise AttributeError. + if not isinstance(data, dict): + return None + gpus = data.get("gpu", data.get("gpus", [])) + if isinstance(gpus, list): + return len(gpus) + return None + + +def _first_visible_amd_gpu_id() -> Optional[str]: + """Return the physical AMD GPU id that should be treated as 'primary'. + + Honours HIP_VISIBLE_DEVICES / ROCR_VISIBLE_DEVICES / CUDA_VISIBLE_DEVICES + in that order (HIP respects all three). Returns ``"0"`` when none are + set, and ``None`` when the env var explicitly narrows to zero GPUs + ("" or "-1"), so callers can short-circuit to "available: False". + """ + for env_name in ( + "HIP_VISIBLE_DEVICES", + "ROCR_VISIBLE_DEVICES", + "CUDA_VISIBLE_DEVICES", + ): + raw = os.environ.get(env_name) + if raw is None: + continue + raw = raw.strip() + if raw == "" or raw == "-1": + return None + # Filter out empty tokens after splitting. This tolerates minor + # typos like ``HIP_VISIBLE_DEVICES=",1"`` (leading comma, user + # clearly meant to narrow to device 1) while still falling + # through to the next env var when every token is empty + # (e.g. ``,,,``). + tokens = [t.strip() for t in raw.split(",") if t.strip()] + if tokens: + return tokens[0] + return "0" + + +def get_primary_gpu_utilization() -> dict[str, Any]: + """Return utilization metrics for the primary visible AMD GPU.""" + gpu_idx = _first_visible_amd_gpu_id() + if gpu_idx is None: + return {"available": False} + data = _run_amd_smi("metric", "-g", gpu_idx) + if data is None: + return {"available": False} + + # amd-smi may return: + # - a list of GPU dicts (older versions) + # - a dict with a "gpu_data" key wrapping a list (newer versions) + # - a single GPU dict (rare) + if isinstance(data, dict) and "gpu_data" in data: + data = data["gpu_data"] + if isinstance(data, list): + if len(data) == 0: + return {"available": False} + gpu_data = data[0] + else: + gpu_data = data + + metrics = _extract_gpu_metrics(gpu_data) + if not _has_real_metrics(metrics): + # amd-smi returned a JSON envelope with no usable fields (error + # response or unsupported card). Surface as unavailable rather + # than available-with-empty-data so the UI does not render a + # ghost device. + return {"available": False} + metrics["available"] = True + return metrics + + +def get_visible_gpu_utilization( + parent_visible_ids: Optional[list[int]], + parent_cuda_visible_devices: Optional[str] = None, +) -> dict[str, Any]: + """Return utilization metrics for visible AMD GPUs.""" + if parent_visible_ids is None: + return { + "available": False, + "backend_cuda_visible_devices": parent_cuda_visible_devices, + "parent_visible_gpu_ids": [], + "devices": [], + "index_kind": "unresolved", + } + + data = _run_amd_smi("metric") + if data is None: + return { + "available": False, + "backend_cuda_visible_devices": parent_cuda_visible_devices, + "parent_visible_gpu_ids": parent_visible_ids or [], + "devices": [], + "index_kind": "physical", + } + + # Extract a device list from amd-smi's envelope. Newer versions return + # a JSON array directly, older versions return a dict with a "gpus" / + # "gpu" key wrapping the list. Guard non-dict / non-list envelopes + # (scalar / string fallbacks from malformed output) so the .get() + # access cannot raise AttributeError on an unexpected shape. + if isinstance(data, list): + gpu_list = data + elif isinstance(data, dict): + # Newer amd-smi wraps output in {"gpu_data": [...]} + gpu_list = data.get("gpu_data", data.get("gpus", data.get("gpu", [data]))) + else: + gpu_list = [data] + visible_set = set(parent_visible_ids) + ordinal_map = {gpu_id: ordinal for ordinal, gpu_id in enumerate(parent_visible_ids)} + + devices = [] + for fallback_idx, gpu_data in enumerate(gpu_list): + # Skip non-dict entries defensively: if amd-smi ever ships a + # scalar inside its "gpus" array (observed on some malformed + # output), _extract_gpu_metrics would raise AttributeError on + # the first .get() call. + if not isinstance(gpu_data, dict): + continue + # Use AMD-reported GPU ID when available, fall back to enumeration + # index. Newer amd-smi versions wrap scalars as ``{"value": 0, + # "unit": "none"}``, so route raw_id through ``_parse_numeric`` + # which already handles bare ints, floats, strings, and that + # dict shape uniformly. + raw_id = gpu_data.get( + "gpu", gpu_data.get("gpu_id", gpu_data.get("id", fallback_idx)) + ) + parsed_id = _parse_numeric(raw_id) + if parsed_id is None: + logger.debug( + "amd-smi GPU id %r could not be parsed; falling back to " + "enumeration index %d", + raw_id, + fallback_idx, + ) + idx = fallback_idx + else: + idx = int(parsed_id) + if idx not in visible_set: + continue + metrics = _extract_gpu_metrics(gpu_data) + if not _has_real_metrics(metrics): + # Skip ghost entries: an amd-smi response that decodes to a + # dict but contains no usable fields (error envelope, etc.) + # would otherwise show up as a device row with all-None + # numbers in the UI. + continue + metrics["index"] = idx + metrics["index_kind"] = "physical" + metrics["visible_ordinal"] = ordinal_map.get(idx, len(devices)) + devices.append(metrics) + + return { + "available": len(devices) > 0, + "backend_cuda_visible_devices": parent_cuda_visible_devices, + "parent_visible_gpu_ids": parent_visible_ids or [], + "devices": devices, + "index_kind": "physical", + } diff --git a/studio/backend/utils/hardware/hardware.py b/studio/backend/utils/hardware/hardware.py index b6d3faf6d7..be31c00a78 100644 --- a/studio/backend/utils/hardware/hardware.py +++ b/studio/backend/utils/hardware/hardware.py @@ -43,6 +43,26 @@ class DeviceType(str, Enum): DEVICE: Optional[DeviceType] = None CHAT_ONLY: bool = True # No CUDA GPU -> GGUF chat only (Mac, CPU-only, etc.) +IS_ROCM: bool = ( + False # True when running on AMD ROCm (HIP) -- routes GPU monitoring to amd.py +) + + +def _backend_label(device: DeviceType) -> str: + """Return the user-facing backend name for API responses. + + Internally we still represent ROCm hosts as ``DeviceType.CUDA`` because + ROCm torch sets ``torch.cuda.is_available() = True`` and reuses the whole + ``torch.cuda.*`` API surface, so branching on ``DeviceType`` stays + consistent with the rest of the codebase. For the JSON responses served + to the Studio frontend and other clients, however, "cuda" is misleading + on an AMD machine. This helper swaps the label to ``"rocm"`` when the + module-level ``IS_ROCM`` flag is set so the UI can render the correct + backend name without every caller having to duplicate the check. + """ + if IS_ROCM and device == DeviceType.CUDA: + return "rocm" + return device.value # ========== Detection ========== @@ -85,10 +105,11 @@ def detect_hardware() -> DeviceType: 2. MLX (Apple Silicon via MLX framework) 3. CPU (fallback) """ - global DEVICE, CHAT_ONLY - CHAT_ONLY = True # reset -- only CUDA sets it to False + global DEVICE, CHAT_ONLY, IS_ROCM + CHAT_ONLY = True # reset -- only CUDA/ROCm sets it to False + IS_ROCM = False - # --- CUDA: try PyTorch --- + # --- CUDA / ROCm: try PyTorch --- if _has_torch(): import torch @@ -96,7 +117,16 @@ def detect_hardware() -> DeviceType: DEVICE = DeviceType.CUDA CHAT_ONLY = False device_name = torch.cuda.get_device_properties(0).name - print(f"Hardware detected: CUDA — {device_name}") + + # Distinguish AMD ROCm (HIP) from NVIDIA CUDA for display purposes. + # DeviceType stays CUDA since torch.cuda.* works on ROCm via HIP. + if getattr(torch.version, "hip", None) is not None: + IS_ROCM = True + print( + f"Hardware detected: ROCm (HIP {torch.version.hip}) -- {device_name}" + ) + else: + print(f"Hardware detected: CUDA -- {device_name}") return DEVICE # --- XPU: Intel GPU --- @@ -186,7 +216,7 @@ def get_gpu_memory_info() -> Dict[str, Any]: return { "available": True, - "backend": device.value, + "backend": _backend_label(device), "device": idx, "device_name": props.name, "total_gb": total / (1024**3), @@ -197,7 +227,11 @@ def get_gpu_memory_info() -> Dict[str, Any]: } except Exception as e: logger.error(f"Error getting CUDA GPU info: {e}") - return {"available": False, "backend": device.value, "error": str(e)} + return { + "available": False, + "backend": _backend_label(device), + "error": str(e), + } # ---- XPU path (Intel GPU) ---- if device == DeviceType.XPU: @@ -213,7 +247,7 @@ def get_gpu_memory_info() -> Dict[str, Any]: return { "available": True, - "backend": device.value, + "backend": _backend_label(device), "device": idx, "device_name": props.name, "total_gb": total / (1024**3), @@ -224,7 +258,11 @@ def get_gpu_memory_info() -> Dict[str, Any]: } except Exception as e: logger.error("Error getting XPU GPU info: %s", e) - return {"available": False, "backend": device.value, "error": str(e)} + return { + "available": False, + "backend": _backend_label(device), + "error": str(e), + } # ---- MLX path (Apple Silicon) ---- if device == DeviceType.MLX: @@ -239,7 +277,7 @@ def get_gpu_memory_info() -> Dict[str, Any]: return { "available": True, - "backend": device.value, + "backend": _backend_label(device), "device": 0, "device_name": f"Apple Silicon ({platform.processor() or platform.machine()})", "total_gb": total / (1024**3), @@ -250,7 +288,11 @@ def get_gpu_memory_info() -> Dict[str, Any]: } except Exception as e: logger.error(f"Error getting MLX GPU info: {e}") - return {"available": False, "backend": device.value, "error": str(e)} + return { + "available": False, + "backend": _backend_label(device), + "error": str(e), + } # ---- CPU-only ---- return {"available": False, "backend": "cpu"} @@ -315,13 +357,15 @@ def get_package_versions() -> Dict[str, Optional[str]]: except PackageNotFoundError: versions[name] = None - # CUDA toolkit version bundled with torch + # GPU runtime version bundled with torch try: import torch versions["cuda"] = getattr(torch.version, "cuda", None) + versions["rocm"] = getattr(torch.version, "hip", None) except Exception: versions["cuda"] = None + versions["rocm"] = None return versions @@ -387,26 +431,50 @@ def _torch_get_per_device_info(device_indices: list[int]) -> list[Dict[str, Any] # ========== Live GPU Utilization ========== +def _smi_query(func_name: str, *args, **kwargs) -> Optional[Dict[str, Any]]: + """Run a query against the appropriate SMI backend (amd-smi or nvidia-smi). + + Returns the result dict if available, or None on failure/unavailability. + """ + if IS_ROCM: + backend_name = "amd-smi" + try: + from . import amd as _backend + except Exception as e: + logger.warning("%s import failed: %s", backend_name, e) + return None + else: + backend_name = "nvidia-smi" + try: + from . import nvidia as _backend + except Exception as e: + logger.warning("%s import failed: %s", backend_name, e) + return None + try: + func = getattr(_backend, func_name) + result = func(*args, **kwargs) + if result.get("available"): + return result + except Exception as e: + logger.warning("%s %s query failed: %s", backend_name, func_name, e) + return None + + def get_gpu_utilization() -> Dict[str, Any]: """Return a live snapshot of device utilization information.""" device = get_device() if device == DeviceType.CUDA: - try: - from . import nvidia - - result = nvidia.get_primary_gpu_utilization() - if result.get("available"): - result["backend"] = device.value - return result - except Exception as e: - logger.warning("nvidia-smi utilization query failed: %s", e) + result = _smi_query("get_primary_gpu_utilization") + if result is not None: + result["backend"] = _backend_label(device) + return result mem = get_gpu_memory_info() if device != DeviceType.CPU and mem.get("available"): return { "available": True, - "backend": device.value, + "backend": _backend_label(device), "gpu_utilization_pct": None, "temperature_c": None, "vram_used_gb": round(mem.get("allocated_gb", 0), 2), @@ -417,7 +485,7 @@ def get_gpu_utilization() -> Dict[str, Any]: "power_utilization_pct": None, } - return {"available": False, "backend": device.value} + return {"available": False, "backend": _backend_label(device)} def get_visible_gpu_utilization() -> Dict[str, Any]: @@ -425,18 +493,14 @@ def get_visible_gpu_utilization() -> Dict[str, Any]: if device == DeviceType.CUDA: parent_visible_spec = _get_parent_visible_gpu_spec() - try: - from . import nvidia - - result = nvidia.get_visible_gpu_utilization( - parent_visible_spec["numeric_ids"], - parent_cuda_visible_devices = parent_visible_spec["raw"], - ) - if result.get("available"): - result["backend"] = device.value - return result - except Exception as e: - logger.warning("nvidia-smi visible GPU utilization query failed: %s", e) + result = _smi_query( + "get_visible_gpu_utilization", + parent_visible_spec["numeric_ids"], + parent_cuda_visible_devices = parent_visible_spec["raw"], + ) + if result is not None: + result["backend"] = _backend_label(device) + return result # Torch-based fallback for CUDA (nvidia-smi unavailable, AMD ROCm) and XPU (Intel) if device in (DeviceType.CUDA, DeviceType.XPU): @@ -475,7 +539,7 @@ def get_visible_gpu_utilization() -> Dict[str, Any]: ) return { "available": True, - "backend": device.value, + "backend": _backend_label(device), "parent_visible_gpu_ids": parent_ids, "devices": devices, "index_kind": index_kind, @@ -486,14 +550,14 @@ def get_visible_gpu_utilization() -> Dict[str, Any]: if not mem.get("available"): return { "available": False, - "backend": device.value, + "backend": _backend_label(device), "parent_visible_gpu_ids": [], "devices": [], "index_kind": "relative", } return { "available": True, - "backend": device.value, + "backend": _backend_label(device), "parent_visible_gpu_ids": [0], "devices": [ { @@ -515,7 +579,7 @@ def get_visible_gpu_utilization() -> Dict[str, Any]: return { "available": False, - "backend": device.value, + "backend": _backend_label(device), "parent_visible_gpu_ids": [], "devices": [], "index_kind": "relative", @@ -529,7 +593,21 @@ _visible_gpu_count: Optional[int] = None def _get_parent_visible_gpu_spec() -> Dict[str, Any]: - cuda_visible = os.environ.get("CUDA_VISIBLE_DEVICES") + # ROCm uses HIP_VISIBLE_DEVICES / ROCR_VISIBLE_DEVICES in addition to + # CUDA_VISIBLE_DEVICES (which HIP also respects). Check ROCm-specific + # env vars first so multi-GPU AMD setups are handled correctly. + # Use explicit None checks (not `or`) so empty string "" is honoured + # as "no visible GPUs" rather than falling through to CUDA_VISIBLE_DEVICES. + cuda_visible = None + if IS_ROCM: + hip_vis = os.environ.get("HIP_VISIBLE_DEVICES") + rocr_vis = os.environ.get("ROCR_VISIBLE_DEVICES") + if hip_vis is not None: + cuda_visible = hip_vis + elif rocr_vis is not None: + cuda_visible = rocr_vis + if cuda_visible is None: + cuda_visible = os.environ.get("CUDA_VISIBLE_DEVICES") if cuda_visible is None: return { @@ -1109,15 +1187,17 @@ def get_physical_gpu_count() -> int: if device == DeviceType.CUDA: try: - from . import nvidia - - count = nvidia.get_physical_gpu_count() + if IS_ROCM: + from . import amd as _smi_mod + else: + from . import nvidia as _smi_mod + count = _smi_mod.get_physical_gpu_count() if count is not None: _physical_gpu_count = count return _physical_gpu_count except Exception: pass - # nvidia-smi unavailable or failed — fall back to torch + # SMI tool unavailable or failed -- fall back to torch count = _torch_get_physical_gpu_count() _physical_gpu_count = count if count is not None else 1 return _physical_gpu_count @@ -1136,12 +1216,25 @@ def get_physical_gpu_count() -> int: return _physical_gpu_count +def _backend_visible_devices_env() -> Optional[str]: + """Return the raw visibility env string that applies to this backend. + + On ROCm, HIP_VISIBLE_DEVICES / ROCR_VISIBLE_DEVICES take precedence + over CUDA_VISIBLE_DEVICES; the helper mirrors the resolution logic in + ``_get_parent_visible_gpu_spec`` so ``backend_cuda_visible_devices`` + reports the value that is actually narrowing the visible device set. + """ + if IS_ROCM: + return _get_parent_visible_gpu_spec().get("raw") + return os.environ.get("CUDA_VISIBLE_DEVICES") + + def get_backend_visible_gpu_info() -> Dict[str, Any]: device = get_device() if device in (DeviceType.CUDA, DeviceType.XPU): parent_visible_ids = get_parent_visible_gpu_ids() - # Try nvidia-smi first (NVIDIA only) - if device == DeviceType.CUDA: + # Try native SMI tool first (nvidia-smi for NVIDIA, skipped for ROCm) + if device == DeviceType.CUDA and not IS_ROCM: try: from . import nvidia @@ -1151,7 +1244,7 @@ def get_backend_visible_gpu_info() -> Dict[str, Any]: parent_visible_spec["raw"], ) if result.get("available"): - result["backend"] = device.value + result["backend"] = _backend_label(device) return result except Exception as e: logger.warning("Backend GPU visibility query failed: %s", e) @@ -1180,8 +1273,8 @@ def get_backend_visible_gpu_info() -> Dict[str, Any]: ] return { "available": True, - "backend": device.value, - "backend_cuda_visible_devices": os.environ.get("CUDA_VISIBLE_DEVICES"), + "backend": _backend_label(device), + "backend_cuda_visible_devices": _backend_visible_devices_env(), "parent_visible_gpu_ids": parent_visible_ids, "devices": devices, "index_kind": index_kind, @@ -1189,8 +1282,8 @@ def get_backend_visible_gpu_info() -> Dict[str, Any]: return { "available": False, - "backend": device.value, - "backend_cuda_visible_devices": os.environ.get("CUDA_VISIBLE_DEVICES"), + "backend": _backend_label(device), + "backend_cuda_visible_devices": _backend_visible_devices_env(), "parent_visible_gpu_ids": parent_visible_ids, "devices": [], "index_kind": "physical", @@ -1201,7 +1294,7 @@ def get_backend_visible_gpu_info() -> Dict[str, Any]: if not mem.get("available"): return { "available": False, - "backend": device.value, + "backend": _backend_label(device), "backend_cuda_visible_devices": os.environ.get("CUDA_VISIBLE_DEVICES"), "parent_visible_gpu_ids": [], "devices": [], @@ -1209,7 +1302,7 @@ def get_backend_visible_gpu_info() -> Dict[str, Any]: } return { "available": True, - "backend": device.value, + "backend": _backend_label(device), "backend_cuda_visible_devices": os.environ.get("CUDA_VISIBLE_DEVICES"), "parent_visible_gpu_ids": [0], "devices": [ @@ -1226,7 +1319,7 @@ def get_backend_visible_gpu_info() -> Dict[str, Any]: return { "available": False, - "backend": device.value, + "backend": _backend_label(device), "backend_cuda_visible_devices": os.environ.get("CUDA_VISIBLE_DEVICES"), "parent_visible_gpu_ids": [], "devices": [], @@ -1246,17 +1339,20 @@ def get_visible_gpu_count() -> int: if _visible_gpu_count is not None: return _visible_gpu_count - cuda_visible = os.environ.get("CUDA_VISIBLE_DEVICES") - if cuda_visible is not None: - # "" means zero GPUs, "0" means 1, "0,1,2" means 3 - cuda_visible = cuda_visible.strip() - if cuda_visible == "" or cuda_visible == "-1": + # Use _get_parent_visible_gpu_spec() which already handles + # HIP_VISIBLE_DEVICES / ROCR_VISIBLE_DEVICES on ROCm. + visible_spec = _get_parent_visible_gpu_spec() + if visible_spec["raw"] is not None: + raw = visible_spec["raw"].strip() + if raw == "" or raw == "-1": _visible_gpu_count = 0 + elif visible_spec["numeric_ids"] is not None: + _visible_gpu_count = len(visible_spec["numeric_ids"]) else: - _visible_gpu_count = len([x for x in cuda_visible.split(",") if x.strip()]) + _visible_gpu_count = len([x for x in raw.split(",") if x.strip()]) return _visible_gpu_count - # CUDA_VISIBLE_DEVICES not set -- try torch, fall back to physical count + # No visibility env var set -- try torch, fall back to physical count try: import torch @@ -1288,8 +1384,24 @@ def apply_gpu_ids(gpu_ids) -> None: value = str(gpu_ids) os.environ["CUDA_VISIBLE_DEVICES"] = value + # Keep ROCm visibility env vars in sync so _get_parent_visible_gpu_spec() + # picks up the narrowed set on AMD systems. Workers can call + # apply_gpu_ids() before detect_hardware() runs (so IS_ROCM is still + # its default False), so also mirror the selection whenever the + # parent process already set a ROCm visibility variable -- that + # way a downstream ROCm process inherits the narrowed mask even + # before Studio's hardware detection has classified the host. + _inherits_rocm_visibility = ( + "HIP_VISIBLE_DEVICES" in os.environ or "ROCR_VISIBLE_DEVICES" in os.environ + ) + if IS_ROCM or _inherits_rocm_visibility: + os.environ["HIP_VISIBLE_DEVICES"] = value + os.environ["ROCR_VISIBLE_DEVICES"] = value _visible_gpu_count = None - logger.info("Applied gpu_ids: CUDA_VISIBLE_DEVICES='%s'", value) + if IS_ROCM or _inherits_rocm_visibility: + logger.info("Applied gpu_ids: CUDA_VISIBLE_DEVICES='%s' (rocm)", value) + else: + logger.info("Applied gpu_ids: CUDA_VISIBLE_DEVICES='%s'", value) def get_device_map( diff --git a/studio/backend/utils/models/__init__.py b/studio/backend/utils/models/__init__.py index a81682d8b7..808e2b012e 100644 --- a/studio/backend/utils/models/__init__.py +++ b/studio/backend/utils/models/__init__.py @@ -13,8 +13,9 @@ from .model_config import ( detect_audio_type, is_audio_input_type, VALID_AUDIO_TYPES, - scan_trained_loras, + scan_trained_models, scan_exported_models, + get_base_model_from_checkpoint, load_model_defaults, get_base_model_from_lora, load_model_config, @@ -25,6 +26,8 @@ from .model_config import ( ) from .checkpoints import scan_checkpoints +scan_trained_loras = scan_trained_models + __all__ = [ "ModelConfig", "GgufVariantInfo", @@ -33,8 +36,10 @@ __all__ = [ "detect_audio_type", "is_audio_input_type", "VALID_AUDIO_TYPES", + "scan_trained_models", "scan_trained_loras", "scan_exported_models", + "get_base_model_from_checkpoint", "load_model_defaults", "get_base_model_from_lora", "load_model_config", diff --git a/studio/backend/utils/models/model_config.py b/studio/backend/utils/models/model_config.py index 61226e52cb..44754520e3 100644 --- a/studio/backend/utils/models/model_config.py +++ b/studio/backend/utils/models/model_config.py @@ -493,8 +493,9 @@ _VLM_MODEL_TYPES = { "minicpmv", } -# Pre-computed .venv_t5 path and backend dir for subprocess version switching. -_VENV_T5_DIR = str(Path.home() / ".unsloth" / "studio" / ".venv_t5") +# Pre-computed .venv_t5 paths and backend dir for subprocess version switching. +# Vision check uses 5.5.0 (newest, recognizes all architectures). +_VENV_T5_DIR = str(Path.home() / ".unsloth" / "studio" / ".venv_t5_550") _BACKEND_DIR = str(Path(__file__).resolve().parent.parent.parent) # Inline script executed in a subprocess with transformers 5.x activated. @@ -903,24 +904,85 @@ def _is_mmproj(filename: str) -> bool: return "mmproj" in filename.lower() -def detect_mmproj_file(path: str) -> Optional[str]: +def _is_gguf_filename(filename: str) -> bool: + return filename.lower().endswith(".gguf") + + +def _iter_gguf_files(directory: Path, recursive: bool = False): + if not directory.is_dir(): + return + iterator = directory.rglob("*") if recursive else directory.iterdir() + for f in iterator: + if f.is_file() and _is_gguf_filename(f.name): + yield f + + +def detect_mmproj_file(path: str, search_root: Optional[str] = None) -> Optional[str]: """ - Find the mmproj (vision projection) GGUF file in a directory. + Find the mmproj (vision projection) GGUF file for a given model. Args: - path: Directory to search — or a .gguf file (uses its parent dir). + path: Directory to search — or a .gguf file (uses its parent dir + as the starting point). + search_root: Optional outer directory that should also be scanned + (and any directory between it and ``path``). This handles + local layouts where the model weights live in a quant-named + subdir (``snapshot/BF16/foo.gguf``) but the mmproj sits at + the snapshot root (``snapshot/mmproj-BF16.gguf``). When + ``None``, only the immediate parent dir is scanned, matching + the historical behavior. Returns: Full path to the mmproj .gguf file, or None if not found. """ p = Path(path) - search_dir = p.parent if p.is_file() else p - if not search_dir.is_dir(): + start_dir = p.parent if p.is_file() else p + if not start_dir.is_dir(): return None - for f in search_dir.glob("*.gguf"): - if _is_mmproj(f.name): - return str(f.resolve()) + # Build the list of dirs to scan: immediate dir first, then walk up + # to (and including) ``search_root`` if it is an ancestor. We walk + # incrementally rather than recursing into ``search_root`` so we + # don't accidentally pick up an mmproj from a sibling subdir + # belonging to a different model variant. + seen: set[Path] = set() + scan_order: list[Path] = [] + + def _add(d: Path) -> None: + try: + resolved = d.resolve() + except OSError: + return + if resolved in seen or not resolved.is_dir(): + return + seen.add(resolved) + scan_order.append(resolved) + + _add(start_dir) + if search_root is not None: + try: + root_resolved = Path(search_root).resolve() + start_resolved = start_dir.resolve() + # Only walk if start_dir is inside (or equal to) search_root. + if root_resolved == start_resolved or ( + start_resolved.is_relative_to(root_resolved) + if hasattr(start_resolved, "is_relative_to") + else str(start_resolved).startswith(str(root_resolved) + "/") + ): + cur = start_resolved + # Walk up from start_dir to (and including) root_resolved. + while cur != root_resolved and cur.parent != cur: + cur = cur.parent + _add(cur) + if cur == root_resolved: + break + except OSError: + pass + + for d in scan_order: + for f in _iter_gguf_files(d): + if _is_mmproj(f.name): + return str(f.resolve()) return None @@ -941,7 +1003,7 @@ def detect_gguf_model(path: str) -> Optional[str]: p = Path(path) # Case 1: direct .gguf file - if p.suffix == ".gguf" and p.is_file(): + if p.suffix.lower() == ".gguf" and p.is_file(): if _is_mmproj(p.name): return None return str(p.resolve()) @@ -949,7 +1011,7 @@ def detect_gguf_model(path: str) -> Optional[str]: # Case 2: directory containing .gguf files (skip mmproj) if p.is_dir(): gguf_files = sorted( - (f for f in p.glob("*.gguf") if not _is_mmproj(f.name)), + (f for f in _iter_gguf_files(p) if not _is_mmproj(f.name)), key = lambda f: f.stat().st_size, reverse = True, ) @@ -1014,7 +1076,7 @@ def _pick_best_gguf(filenames: list[str]) -> Optional[str]: Prefers quantization levels in _GGUF_QUANT_PREFERENCE order. Falls back to the first .gguf file found. """ - gguf_files = [f for f in filenames if f.endswith(".gguf")] + gguf_files = [f for f in filenames if f.lower().endswith(".gguf")] if not gguf_files: return None @@ -1099,7 +1161,7 @@ def list_gguf_variants( for sibling in info.siblings: fname = sibling.rfilename - if not fname.endswith(".gguf"): + if not fname.lower().endswith(".gguf"): continue size = sibling.size or 0 @@ -1170,7 +1232,11 @@ def list_local_gguf_variants( quant_first_file: dict[str, str] = {} has_vision = False - for f in sorted(p.glob("*.gguf")): + # Recurse so variant-specific subdirectories (e.g. ``BF16/...gguf`` + # used by some HF GGUF repos for the largest quants) are picked up. + # Filenames in the result preserve the relative subpath so that + # ``_find_local_gguf_by_variant`` can locate the file again. + for f in sorted(_iter_gguf_files(p, recursive = True)): if _is_mmproj(f.name): has_vision = True continue @@ -1180,8 +1246,14 @@ def list_local_gguf_variants( size = 0 quant = _extract_quant_label(f.name) quant_totals[quant] = quant_totals.get(quant, 0) + size + # Only compute the (potentially expensive) relative path when this + # is the first file we've seen for this quant -- after that we'd + # discard the result anyway. Use posix-style separators so the + # filename matches what ``list_gguf_variants`` (the remote HF + # API path) returns on every platform; otherwise Windows would + # emit ``BF16\foo.gguf`` here. if quant not in quant_first_file: - quant_first_file[quant] = f.name + quant_first_file[quant] = f.relative_to(p).as_posix() variants = [ GgufVariantInfo( @@ -1207,9 +1279,11 @@ def _find_local_gguf_by_variant(directory: str, variant: str) -> Optional[str]: if p is None: return None + # Recurse into subdirectories so variants stored under a quant-named + # subdir (e.g. ``BF16/foo-BF16-00001-of-00002.gguf``) are found. matches = sorted( f - for f in p.glob("*.gguf") + for f in _iter_gguf_files(p, recursive = True) if not _is_mmproj(f.name) and _extract_quant_label(f.name) == variant ) if matches: @@ -1321,46 +1395,89 @@ def is_embedding_model(model_name: str, hf_token: Optional[str] = None) -> bool: return False -def scan_trained_loras(outputs_dir: str = str(outputs_root())) -> List[Tuple[str, str]]: +def _has_model_weight_files(model_dir: Path) -> bool: + """Return True when a directory contains loadable model weights.""" + for item in model_dir.iterdir(): + if not item.is_file(): + continue + + suffix = item.suffix.lower() + if suffix == ".safetensors": + return True + if suffix == ".gguf": + return "mmproj" not in item.name.lower() + if suffix == ".bin": + name = item.name.lower() + if ( + name.startswith("pytorch_model") + or name.startswith("model") + or name.startswith("adapter_model") + or name.startswith("consolidated") + ): + return True + return False + + +def _detect_training_output_type(model_dir: Path) -> Optional[str]: + """Classify a Studio training output as LoRA or full finetune.""" + adapter_config = model_dir / "adapter_config.json" + adapter_model = model_dir / "adapter_model.safetensors" + if adapter_config.exists() or adapter_model.exists(): + return "lora" + + config_file = model_dir / "config.json" + if config_file.exists() and _has_model_weight_files(model_dir): + return "merged" + + return None + + +def _looks_like_lora_adapter(model_dir: Path) -> bool: + return model_dir.is_dir() and ( + (model_dir / "adapter_config.json").exists() + or any(model_dir.glob("adapter_model*.safetensors")) + or any(model_dir.glob("adapter_model*.bin")) + ) + + +def scan_trained_models( + outputs_dir: str = str(outputs_root()), +) -> List[Tuple[str, str, str]]: """ - Scan outputs folder for trained LoRA adapters. + Scan outputs folder for trained Studio models. Returns: - List of tuples: [(display_name, adapter_path), ...] - - Example: - [ - ("unsloth_Meta-Llama-3.1_...", "./outputs/unsloth_Meta-Llama-3.1_.../"), - ("my_finetuned_model", "./outputs/my_finetuned_model/"), - ] + List of tuples: [(display_name, model_path, model_type), ...] + model_type is "lora" for adapter runs and "merged" for full finetunes. """ - trained_loras = [] + trained_models = [] outputs_path = resolve_output_dir(outputs_dir) if not outputs_path.exists(): logger.warning(f"Outputs directory not found: {outputs_dir}") - return trained_loras + return trained_models try: for item in outputs_path.iterdir(): if item.is_dir(): - # Check if this directory contains a LoRA adapter - adapter_config = item / "adapter_config.json" - adapter_model = item / "adapter_model.safetensors" + model_type = _detect_training_output_type(item) + if model_type is None: + continue - if adapter_config.exists() or adapter_model.exists(): - display_name = item.name - adapter_path = str(item) - trained_loras.append((display_name, adapter_path)) - logger.debug(f"Found trained LoRA: {display_name}") + display_name = item.name + model_path = str(item) + trained_models.append((display_name, model_path, model_type)) + logger.debug("Found trained model: %s (%s)", display_name, model_type) # Sort by modification time (newest first) - trained_loras.sort(key = lambda x: Path(x[1]).stat().st_mtime, reverse = True) + trained_models.sort(key = lambda x: Path(x[1]).stat().st_mtime, reverse = True) logger.info( - f"Found {len(trained_loras)} trained LoRA adapters in {outputs_dir}" + "Found %s trained models in %s", + len(trained_models), + outputs_dir, ) - return trained_loras + return trained_models except Exception as e: logger.error(f"Error scanning outputs folder: {e}") @@ -1394,7 +1511,9 @@ def scan_exported_models( # Check for flat GGUF export (e.g. exports/gemma-3-4b-it-finetune-gguf/) # Filter out mmproj (vision projection) files — they aren't loadable as main models - gguf_files = [f for f in run_dir.glob("*.gguf") if not _is_mmproj(f.name)] + gguf_files = [ + f for f in _iter_gguf_files(run_dir) if not _is_mmproj(f.name) + ] if gguf_files: base_model = None export_meta = run_dir / "export_metadata.json" @@ -1421,7 +1540,7 @@ def scan_exported_models( has_weights = any(checkpoint_dir.glob("*.safetensors")) or any( checkpoint_dir.glob("*.bin") ) - has_gguf = any(checkpoint_dir.glob("*.gguf")) + has_gguf = any(_iter_gguf_files(checkpoint_dir)) base_model = None export_type = None @@ -1444,7 +1563,7 @@ def scan_exported_models( pass elif has_gguf: export_type = "gguf" - gguf_list = list(checkpoint_dir.glob("*.gguf")) + gguf_list = list(_iter_gguf_files(checkpoint_dir)) # Check checkpoint_dir first, then fall back to parent run_dir # (export.py writes metadata to the top-level export directory) for meta_dir in (checkpoint_dir, run_dir): @@ -1493,6 +1612,68 @@ def scan_exported_models( return [] +def get_base_model_from_checkpoint(checkpoint_path: str) -> Optional[str]: + """Read the base model name from a local training or checkpoint directory.""" + try: + checkpoint_path_obj = Path(checkpoint_path) + + adapter_config_path = checkpoint_path_obj / "adapter_config.json" + if adapter_config_path.exists(): + with open(adapter_config_path, "r") as f: + config = json.load(f) + base_model = config.get("base_model_name_or_path") + if base_model: + logger.info( + "Detected base model from adapter_config.json: %s", base_model + ) + return base_model + + config_path = checkpoint_path_obj / "config.json" + if config_path.exists(): + with open(config_path, "r") as f: + config = json.load(f) + for key in ("model_name", "_name_or_path"): + base_model = config.get(key) + if base_model and str(base_model) != str(checkpoint_path_obj): + logger.info( + "Detected base model from config.json (%s): %s", + key, + base_model, + ) + return base_model + + training_args_path = checkpoint_path_obj / "training_args.bin" + if training_args_path.exists(): + try: + import torch + + training_args = torch.load(training_args_path) + if hasattr(training_args, "model_name_or_path"): + base_model = training_args.model_name_or_path + logger.info( + "Detected base model from training_args.bin: %s", base_model + ) + return base_model + except Exception as e: + logger.warning(f"Could not load training_args.bin: {e}") + + dir_name = checkpoint_path_obj.name + if dir_name.startswith("unsloth_"): + parts = dir_name.split("_") + if len(parts) >= 2: + model_parts = parts[1:-1] + base_model = "unsloth/" + "_".join(model_parts) + logger.info("Detected base model from directory name: %s", base_model) + return base_model + + logger.warning(f"Could not detect base model for checkpoint: {checkpoint_path}") + return None + + except Exception as e: + logger.error(f"Error reading base model from checkpoint config: {e}") + return None + + def get_base_model_from_lora(lora_path: str) -> Optional[str]: """ Read the base model name from a LoRA adapter's config. @@ -1501,16 +1682,14 @@ def get_base_model_from_lora(lora_path: str) -> Optional[str]: lora_path: Path to the LoRA adapter directory Returns: - Base model identifier (e.g., "unsloth/Meta-Llama-3.1-8B-Instruct-bnb-4bit") - or None if not found - - Example: - >>> get_base_model_from_lora("./outputs/unsloth_Meta-Llama-3.1_.../") - "unsloth/Meta-Llama-3.1-8B-Instruct-bnb-4bit" + Base model identifier or None if not found """ try: lora_path_obj = Path(lora_path) + if not _looks_like_lora_adapter(lora_path_obj): + return None + # Try adapter_config.json first adapter_config_path = lora_path_obj / "adapter_config.json" if adapter_config_path.exists(): @@ -1814,8 +1993,16 @@ class ModelConfig: except Exception as e: logger.debug(f"Could not read export metadata: {e}") - # If vision (or mmproj happens to exist), find the mmproj file - mmproj_file = detect_mmproj_file(gguf_file) + # If vision (or mmproj happens to exist), find the mmproj + # file. The recursive variant scan in + # ``_find_local_gguf_by_variant`` may have returned a + # weight file inside a quant-named subdir (e.g. + # ``.../BF16/foo.gguf``) while ``mmproj-*.gguf`` lives + # at the snapshot root. Pass ``search_root=path`` so + # ``detect_mmproj_file`` walks up to the snapshot root + # instead of seeing only the weight file's immediate + # parent. + mmproj_file = detect_mmproj_file(gguf_file, search_root = path) if mmproj_file: gguf_is_vision = True logger.info(f"Detected mmproj for vision: {mmproj_file}") @@ -1883,7 +2070,11 @@ class ModelConfig: # Auto-detect LoRA for local paths (check adapter_config.json on disk) if not is_lora and is_local: - detected_base = get_base_model_from_lora(path) + detected_base = ( + get_base_model_from_lora(path) + if _looks_like_lora_adapter(Path(path)) + else None + ) if detected_base: is_lora = True logger.info( diff --git a/studio/backend/utils/paths/__init__.py b/studio/backend/utils/paths/__init__.py index 11709ae56e..92191dccdd 100644 --- a/studio/backend/utils/paths/__init__.py +++ b/studio/backend/utils/paths/__init__.py @@ -34,6 +34,7 @@ from .storage_roots import ( legacy_hf_cache_dir, hf_default_cache_dir, lmstudio_model_dirs, + well_known_model_dirs, ensure_dir, ensure_studio_directories, resolve_under_root, @@ -70,6 +71,7 @@ __all__ = [ "legacy_hf_cache_dir", "hf_default_cache_dir", "lmstudio_model_dirs", + "well_known_model_dirs", "ensure_dir", "ensure_studio_directories", "resolve_under_root", diff --git a/studio/backend/utils/paths/storage_roots.py b/studio/backend/utils/paths/storage_roots.py index 4841c5d0a3..b52609b06b 100644 --- a/studio/backend/utils/paths/storage_roots.py +++ b/studio/backend/utils/paths/storage_roots.py @@ -130,6 +130,51 @@ def lmstudio_model_dirs() -> list[Path]: return dirs +def well_known_model_dirs() -> list[Path]: + """Return directories commonly used by other local LLM tools. + + Used by the folder browser to offer quick-pick chips. Returns only + paths that exist on disk, so the UI never shows dead chips. Order + reflects a rough "likelihood the user has models here" -- LM Studio + and Ollama first, then the generic fallbacks. + """ + candidates: list[Path] = [] + + # LM Studio (reuses the logic above, including settings.json override) + candidates.extend(lmstudio_model_dirs()) + + # Ollama -- both the user-level and common system-wide install paths + # (https://github.com/ollama/ollama/issues/733). + ollama_env = os.environ.get("OLLAMA_MODELS") + if ollama_env: + candidates.append(Path(ollama_env).expanduser()) + candidates.append(Path.home() / ".ollama" / "models") + candidates.append(Path("/usr/share/ollama/.ollama/models")) + candidates.append(Path("/var/lib/ollama/.ollama/models")) + + # HF hub cache root (separate from the explicit HF cache chip) + candidates.append(Path.home() / ".cache" / "huggingface" / "hub") + + # Generic "my models" spots users tend to drop things into + for name in ("models", "Models"): + candidates.append(Path.home() / name) + + # Deduplicate while preserving order; keep only extant dirs + out: list[Path] = [] + seen: set[str] = set() + for p in candidates: + try: + resolved = str(p.resolve()) + except OSError: + continue + if resolved in seen: + continue + if Path(resolved).is_dir(): + seen.add(resolved) + out.append(Path(resolved)) + return out + + def _setup_cache_env() -> None: """Set cache environment variables for HuggingFace, uv, and vLLM. diff --git a/studio/backend/utils/transformers_version.py b/studio/backend/utils/transformers_version.py index 07e4a5c000..36c3a4c22d 100644 --- a/studio/backend/utils/transformers_version.py +++ b/studio/backend/utils/transformers_version.py @@ -5,20 +5,25 @@ Automatic transformers version switching. Some newer model architectures (Ministral-3, GLM-4.7-Flash, Qwen3-30B-A3B MoE, -tiny_qwen3_moe) require transformers>=5.3.0, while everything else needs the -default 4.57.x that ships with Unsloth. +tiny_qwen3_moe) require transformers>=5.3.0, while Gemma 4 models require +transformers>=5.5.0. Everything else needs the default 4.57.x that ships +with Unsloth. + +Two separate target directories are maintained: + - .venv_t5_530/ — transformers 5.3.0 (Ministral-3, GLM, Qwen3 MoE, etc.) + - .venv_t5_550/ — transformers 5.5.0 (Gemma 4) When loading a LoRA adapter with a custom name, we resolve the base model from ``adapter_config.json`` and check *that* against the model list. Strategy: Training and inference run in subprocesses that activate the correct version - via sys.path (prepending .venv_t5/ for 5.x models). See: + via sys.path (prepending the appropriate .venv_t5_*/ directory). See: - core/training/worker.py - core/inference/worker.py For export (still in-process), ensure_transformers_version() does a lightweight - sys.path swap using the same .venv_t5/ directory pre-installed by setup.sh. + sys.path swap using the same directories pre-installed by setup.sh. """ import importlib @@ -39,7 +44,7 @@ logger = get_logger(__name__) # --------------------------------------------------------------------------- # Lowercase substrings — if ANY appears anywhere in the lowered model name, -# we need transformers 5.x. +# we need transformers 5.3.0. TRANSFORMERS_5_MODEL_SUBSTRINGS: tuple[str, ...] = ( "ministral-3-", # Ministral-3-{3,8,14}B-{Instruct,Reasoning,Base}-2512 "glm-4.7-flash", # GLM-4.7-Flash @@ -47,10 +52,24 @@ TRANSFORMERS_5_MODEL_SUBSTRINGS: tuple[str, ...] = ( "qwen3.5", # Qwen3.5 family (35B-A3B, etc.) "qwen3-next", # Qwen3-Next and variants "tiny_qwen3_moe", # imdatta0/tiny_qwen3_moe_2.8B_0.7B + "lfm2.5-vl-450m", # LiquidAI/LFM2.5-VL-450M +) + +# Lowercase substrings for models that require transformers 5.5.0 (checked first). +TRANSFORMERS_550_MODEL_SUBSTRINGS: tuple[str, ...] = ( "gemma-4", # Gemma-4 (E2B-it, E4B-it, 31B-it, 26B-A4B-it) "gemma4", # Gemma-4 alternate naming ) +# Architecture classes / model_type values that require transformers 5.5.0. +# Checked via config.json (local or HuggingFace). +_TRANSFORMERS_550_ARCHITECTURES: set[str] = { + "Gemma4ForConditionalGeneration", +} +_TRANSFORMERS_550_MODEL_TYPES: set[str] = { + "gemma4", +} + # Tokenizer classes that only exist in transformers>=5.x _TRANSFORMERS_5_TOKENIZER_CLASSES: set[str] = { "TokenizersBackend", @@ -59,12 +78,61 @@ _TRANSFORMERS_5_TOKENIZER_CLASSES: set[str] = { # Cache for dynamic tokenizer_config.json lookups to avoid repeated fetches _tokenizer_class_cache: dict[str, bool] = {} -# Versions -TRANSFORMERS_5_VERSION = "5.5.0" -TRANSFORMERS_DEFAULT_VERSION = "4.57.6" +# Cache for dynamic config.json lookups (architecture/model_type checks) +_config_needs_550_cache: dict[str, bool] = {} -# Pre-installed directory for transformers 5.x — created by setup.sh / setup.ps1 -_VENV_T5_DIR = str(Path.home() / ".unsloth" / "studio" / ".venv_t5") +# Versions +TRANSFORMERS_550_VERSION = "5.5.0" +TRANSFORMERS_530_VERSION = "5.3.0" +TRANSFORMERS_DEFAULT_VERSION = "4.57.6" +# Backwards-compat alias — points to 5.5.0 (the highest 5.x tier). +# Consumers should prefer TRANSFORMERS_530_VERSION / TRANSFORMERS_550_VERSION. +TRANSFORMERS_5_VERSION = TRANSFORMERS_550_VERSION + +# Pre-installed directories — created by setup.sh / setup.ps1 +_VENV_T5_530_DIR = str(Path.home() / ".unsloth" / "studio" / ".venv_t5_530") +_VENV_T5_550_DIR = str(Path.home() / ".unsloth" / "studio" / ".venv_t5_550") +# Backwards-compat alias +_VENV_T5_DIR = _VENV_T5_550_DIR + + +def activate_transformers_for_subprocess(model_name: str) -> None: + """Activate the correct transformers version in a subprocess worker. + + Call this BEFORE any ML imports. Resolves LoRA adapters to their base + model, determines the required tier, and prepends the appropriate + ``.venv_t5_*`` directory to ``sys.path``. Also propagates the path + via ``PYTHONPATH`` for child processes (e.g. GGUF converter). + + Used by training, inference, and export workers. + """ + resolved = _resolve_base_model(model_name) + tier = get_transformers_tier(resolved) + + if tier == "550": + if not _ensure_venv_t5_550_exists(): + raise RuntimeError( + f"Cannot activate transformers 5.5.0: " + f".venv_t5_550 missing at {_VENV_T5_550_DIR}" + ) + if _VENV_T5_550_DIR not in sys.path: + sys.path.insert(0, _VENV_T5_550_DIR) + logger.info("Activated transformers 5.5.0 from %s", _VENV_T5_550_DIR) + _pp = os.environ.get("PYTHONPATH", "") + os.environ["PYTHONPATH"] = _VENV_T5_550_DIR + (os.pathsep + _pp if _pp else "") + elif tier == "530": + if not _ensure_venv_t5_530_exists(): + raise RuntimeError( + f"Cannot activate transformers 5.3.0: " + f".venv_t5_530 missing at {_VENV_T5_530_DIR}" + ) + if _VENV_T5_530_DIR not in sys.path: + sys.path.insert(0, _VENV_T5_530_DIR) + logger.info("Activated transformers 5.3.0 from %s", _VENV_T5_530_DIR) + _pp = os.environ.get("PYTHONPATH", "") + os.environ["PYTHONPATH"] = _VENV_T5_530_DIR + (os.pathsep + _pp if _pp else "") + else: + logger.info("Using default transformers (4.57.x) for %s", model_name) def _resolve_base_model(model_name: str) -> str: @@ -192,18 +260,103 @@ def _check_tokenizer_config_needs_v5(model_name: str) -> bool: return False -def needs_transformers_5(model_name: str) -> bool: - """Return True if *model_name* belongs to an architecture that requires - ``transformers>=5.3.0``. +def _check_config_needs_550(model_name: str) -> bool: + """Check ``config.json`` for architectures or model_type that require + transformers 5.5.0 (e.g. Gemma 4). - First checks the hardcoded substring list for known models, then - dynamically fetches ``tokenizer_config.json`` from HuggingFace to check - if the tokenizer_class (e.g. ``TokenizersBackend``) requires v5. + Checks locally first, then falls back to fetching from HuggingFace. + Results are cached in ``_config_needs_550_cache``. + Returns False on any error (fail-open to lower tier). + """ + if model_name in _config_needs_550_cache: + return _config_needs_550_cache[model_name] + + def _check_cfg(cfg: dict) -> bool: + archs = cfg.get("architectures", []) + if any(a in _TRANSFORMERS_550_ARCHITECTURES for a in archs): + return True + if cfg.get("model_type") in _TRANSFORMERS_550_MODEL_TYPES: + return True + return False + + # --- Check local config.json first ------------------------------------ + local_path = Path(model_name) + local_cfg = local_path / "config.json" + if local_cfg.is_file(): + try: + with open(local_cfg) as f: + cfg = json.load(f) + result = _check_cfg(cfg) + if result: + logger.info( + "Local config.json check: %s needs transformers 5.5.0 " + "(architectures=%s, model_type=%s)", + model_name, + cfg.get("architectures", []), + cfg.get("model_type"), + ) + _config_needs_550_cache[model_name] = result + return result + except Exception as exc: + logger.debug("Could not read %s: %s", local_cfg, exc) + + # --- Fall back to fetching from HuggingFace --------------------------- + import urllib.request + + url = f"https://huggingface.co/{model_name}/raw/main/config.json" + try: + req = urllib.request.Request(url, headers = {"User-Agent": "unsloth-studio"}) + with urllib.request.urlopen(req, timeout = 10) as resp: + cfg = json.loads(resp.read().decode()) + result = _check_cfg(cfg) + if result: + logger.info( + "Dynamic config.json check: %s needs transformers 5.5.0 " + "(architectures=%s, model_type=%s)", + model_name, + cfg.get("architectures", []), + cfg.get("model_type"), + ) + _config_needs_550_cache[model_name] = result + return result + except Exception as exc: + logger.debug("Could not fetch config.json for '%s': %s", model_name, exc) + _config_needs_550_cache[model_name] = False + return False + + +def get_transformers_tier(model_name: str) -> str: + """Return the transformers tier required for *model_name*. + + Returns ``"550"`` for models needing transformers 5.5.0 (e.g. Gemma 4), + ``"530"`` for models needing transformers 5.3.0 (e.g. Ministral-3, Qwen3 MoE), + or ``"default"`` for everything else (4.57.x). + + The 5.5.0 check runs first, then 5.3.0. """ lowered = model_name.lower() + + # --- Fast substring checks (no I/O) ------------------------------------ + if any(sub in lowered for sub in TRANSFORMERS_550_MODEL_SUBSTRINGS): + return "550" if any(sub in lowered for sub in TRANSFORMERS_5_MODEL_SUBSTRINGS): - return True - return _check_tokenizer_config_needs_v5(model_name) + return "530" + + # --- Slow config fallbacks (local file first, then network) ----------- + if _check_config_needs_550(model_name): + return "550" + if _check_tokenizer_config_needs_v5(model_name): + return "530" + + return "default" + + +def needs_transformers_5(model_name: str) -> bool: + """Return True if *model_name* requires any transformers 5.x version. + + Convenience wrapper around :func:`get_transformers_tier`. + """ + return get_transformers_tier(model_name) != "default" # --------------------------------------------------------------------------- @@ -258,27 +411,36 @@ def _purge_modules() -> int: return len(to_remove) -_VENV_T5_PACKAGES = ( - f"transformers=={TRANSFORMERS_5_VERSION}", +_VENV_T5_530_PACKAGES = ( + f"transformers=={TRANSFORMERS_530_VERSION}", "huggingface_hub==1.8.0", "hf_xet==1.4.2", "tiktoken", ) +_VENV_T5_550_PACKAGES = ( + f"transformers=={TRANSFORMERS_550_VERSION}", + "huggingface_hub==1.8.0", + "hf_xet==1.4.2", + "tiktoken", +) -def _venv_t5_is_valid() -> bool: - """Return True if .venv_t5/ has all required packages at the correct versions.""" - if not os.path.isdir(_VENV_T5_DIR) or not os.listdir(_VENV_T5_DIR): +# Backwards-compat alias +_VENV_T5_PACKAGES = _VENV_T5_550_PACKAGES + + +def _venv_dir_is_valid(venv_dir: str, packages: tuple[str, ...]) -> bool: + """Return True if *venv_dir* has all *packages* at the correct versions.""" + if not os.path.isdir(venv_dir) or not os.listdir(venv_dir): return False - # Check that the key package directories exist AND match the required version - for pkg_spec in _VENV_T5_PACKAGES: + for pkg_spec in packages: parts = pkg_spec.split("==") pkg_name = parts[0] pkg_version = parts[1] if len(parts) > 1 else None pkg_name_norm = pkg_name.replace("-", "_") # Check directory exists if not any( - (Path(_VENV_T5_DIR) / d).is_dir() + (Path(venv_dir) / d).is_dir() for d in (pkg_name_norm, pkg_name_norm.replace("_", "-")) ): return False @@ -287,7 +449,7 @@ def _venv_t5_is_valid() -> bool: continue # Check version via .dist-info metadata dist_info_found = False - for di in Path(_VENV_T5_DIR).glob(f"{pkg_name_norm}-*.dist-info"): + for di in Path(venv_dir).glob(f"{pkg_name_norm}-*.dist-info"): metadata = di / "METADATA" if not metadata.is_file(): continue @@ -296,7 +458,8 @@ def _venv_t5_is_valid() -> bool: installed_ver = line.split(":", 1)[1].strip() if installed_ver != pkg_version: logger.info( - ".venv_t5 has %s==%s but need %s", + "%s has %s==%s but need %s", + venv_dir, pkg_name, installed_ver, pkg_version, @@ -311,8 +474,13 @@ def _venv_t5_is_valid() -> bool: return True -def _install_to_venv_t5(pkg: str) -> bool: - """Install a single package into .venv_t5/, preferring uv then pip.""" +def _venv_t5_is_valid() -> bool: + """Backwards-compat: check the 5.5.0 venv.""" + return _venv_dir_is_valid(_VENV_T5_550_DIR, _VENV_T5_550_PACKAGES) + + +def _install_to_dir(pkg: str, target_dir: str) -> bool: + """Install a single package into *target_dir*, preferring uv then pip.""" # Try uv first (faster) if already on PATH -- do NOT install uv at runtime if shutil.which("uv"): result = subprocess.run( @@ -323,7 +491,7 @@ def _install_to_venv_t5(pkg: str) -> bool: "--python", sys.executable, "--target", - _VENV_T5_DIR, + target_dir, "--no-deps", "--upgrade", pkg, @@ -344,7 +512,7 @@ def _install_to_venv_t5(pkg: str) -> bool: "pip", "install", "--target", - _VENV_T5_DIR, + target_dir, "--no-deps", "--upgrade", pkg, @@ -359,47 +527,62 @@ def _install_to_venv_t5(pkg: str) -> bool: return True -def _ensure_venv_t5_exists() -> bool: - """Ensure .venv_t5/ exists with all required packages. Install if missing.""" - if _venv_t5_is_valid(): +def _ensure_venv_dir(venv_dir: str, packages: tuple[str, ...], label: str) -> bool: + """Ensure *venv_dir* exists with all *packages*. Install if missing.""" + if _venv_dir_is_valid(venv_dir, packages): return True logger.warning( - ".venv_t5 not found or incomplete at %s -- installing at runtime", _VENV_T5_DIR + "%s not found or incomplete at %s -- installing at runtime", label, venv_dir ) - shutil.rmtree(_VENV_T5_DIR, ignore_errors = True) - os.makedirs(_VENV_T5_DIR, exist_ok = True) - for pkg in _VENV_T5_PACKAGES: - if not _install_to_venv_t5(pkg): + shutil.rmtree(venv_dir, ignore_errors = True) + os.makedirs(venv_dir, exist_ok = True) + for pkg in packages: + if not _install_to_dir(pkg, venv_dir): return False - logger.info("Installed transformers 5.x to %s", _VENV_T5_DIR) + logger.info("Installed %s to %s", label, venv_dir) return True -def _activate_5x() -> None: - """Prepend .venv_t5/ to sys.path, purge stale modules, reimport.""" - if not _ensure_venv_t5_exists(): - raise RuntimeError( - f"Cannot activate transformers 5.x: .venv_t5 missing at {_VENV_T5_DIR}" - ) +def _ensure_venv_t5_530_exists() -> bool: + """Ensure .venv_t5_530/ exists with transformers 5.3.0.""" + return _ensure_venv_dir( + _VENV_T5_530_DIR, _VENV_T5_530_PACKAGES, "transformers 5.3.0" + ) - if _VENV_T5_DIR not in sys.path: - sys.path.insert(0, _VENV_T5_DIR) - logger.info("Prepended %s to sys.path", _VENV_T5_DIR) + +def _ensure_venv_t5_550_exists() -> bool: + """Ensure .venv_t5_550/ exists with transformers 5.5.0.""" + return _ensure_venv_dir( + _VENV_T5_550_DIR, _VENV_T5_550_PACKAGES, "transformers 5.5.0" + ) + + +def _ensure_venv_t5_exists() -> bool: + """Backwards-compat: ensure the 5.5.0 venv exists.""" + return _ensure_venv_t5_550_exists() + + +def _activate_venv(venv_dir: str, label: str) -> None: + """Prepend *venv_dir* to sys.path, purge stale modules, reimport.""" + if venv_dir not in sys.path: + sys.path.insert(0, venv_dir) + logger.info("Prepended %s to sys.path", venv_dir) count = _purge_modules() logger.info("Purged %d cached modules", count) import transformers - logger.info("Loaded transformers %s", transformers.__version__) + logger.info("Loaded transformers %s (%s)", transformers.__version__, label) def _deactivate_5x() -> None: - """Remove .venv_t5/ from sys.path, purge stale modules, reimport.""" - while _VENV_T5_DIR in sys.path: - sys.path.remove(_VENV_T5_DIR) - logger.info("Removed %s from sys.path", _VENV_T5_DIR) + """Remove all .venv_t5_*/ dirs from sys.path, purge stale modules, reimport.""" + for d in (_VENV_T5_530_DIR, _VENV_T5_550_DIR): + while d in sys.path: + sys.path.remove(d) + logger.info("Removed venv_t5 dirs from sys.path") count = _purge_modules() logger.info("Purged %d cached modules", count) @@ -412,9 +595,10 @@ def _deactivate_5x() -> None: def ensure_transformers_version(model_name: str) -> None: """Ensure the correct ``transformers`` version is active for *model_name*. - Uses sys.path with .venv_t5/ (pre-installed by setup.sh): - • Need 5.x → prepend .venv_t5/ to sys.path, purge modules. - • Need 4.x → remove .venv_t5/ from sys.path, purge modules. + Uses sys.path with .venv_t5_530/ or .venv_t5_550/ (pre-installed by setup.sh): + • Need 5.5.0 → prepend .venv_t5_550/ to sys.path, purge modules. + • Need 5.3.0 → prepend .venv_t5_530/ to sys.path, purge modules. + • Need 4.x → remove all .venv_t5_*/ from sys.path, purge modules. For LoRA adapters with custom names, the base model is resolved from ``adapter_config.json`` before checking. @@ -424,8 +608,21 @@ def ensure_transformers_version(model_name: str) -> None: """ # Resolve LoRA adapters to their base model for accurate detection resolved = _resolve_base_model(model_name) - want_5 = needs_transformers_5(resolved) - target_version = TRANSFORMERS_5_VERSION if want_5 else TRANSFORMERS_DEFAULT_VERSION + tier = get_transformers_tier(resolved) + + if tier == "550": + target_version = TRANSFORMERS_550_VERSION + venv_dir = _VENV_T5_550_DIR + ensure_fn = _ensure_venv_t5_550_exists + elif tier == "530": + target_version = TRANSFORMERS_530_VERSION + venv_dir = _VENV_T5_530_DIR + ensure_fn = _ensure_venv_t5_530_exists + else: + target_version = TRANSFORMERS_DEFAULT_VERSION + venv_dir = None + ensure_fn = None + target_major = int(target_version.split(".")[0]) # Check what's actually loaded in memory @@ -441,8 +638,17 @@ def ensure_transformers_version(model_name: str) -> None: # --- Already correct? --------------------------------------------------- if in_memory is not None: + if in_memory == target_version: + logger.info( + "transformers %s already loaded — correct for '%s'", + in_memory, + model_name, + ) + return + # Different 5.x → need to switch (e.g. 5.3.0 loaded but need 5.5.0) in_memory_major = int(in_memory.split(".")[0]) - if in_memory_major == target_major: + if in_memory_major == target_major and venv_dir is None: + # Both are default (4.x) — close enough logger.info( "transformers %s already loaded — correct for '%s'", in_memory, @@ -451,9 +657,16 @@ def ensure_transformers_version(model_name: str) -> None: return # --- Switch version ----------------------------------------------------- - if want_5: - logger.info("Activating transformers %s via .venv_t5…", TRANSFORMERS_5_VERSION) - _activate_5x() + if venv_dir is not None: + # First remove any other 5.x venv from sys.path + _deactivate_5x() + if not ensure_fn(): + raise RuntimeError( + f"Cannot activate transformers {target_version}: " + f"venv missing at {venv_dir}" + ) + logger.info("Activating transformers %s…", target_version) + _activate_venv(venv_dir, f"transformers {target_version}") else: logger.info( "Reverting to default transformers %s…", TRANSFORMERS_DEFAULT_VERSION diff --git a/studio/backend/utils/wheel_utils.py b/studio/backend/utils/wheel_utils.py new file mode 100644 index 0000000000..00240f1e69 --- /dev/null +++ b/studio/backend/utils/wheel_utils.py @@ -0,0 +1,170 @@ +# SPDX-License-Identifier: AGPL-3.0-only +# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0 + +from __future__ import annotations + +import json +import logging +import platform +import shutil +import subprocess +import sys +import urllib.error +import urllib.request +from typing import Callable + +_logger = logging.getLogger(__name__) + +FLASH_ATTN_RELEASE_BASE_URL = ( + "https://github.com/Dao-AILab/flash-attention/releases/download" +) + + +def linux_wheel_platform_tag() -> str | None: + machine = platform.machine().lower() + if sys.platform.startswith("linux"): + if machine in {"x86_64", "amd64"}: + return "linux_x86_64" + if machine in {"aarch64", "arm64"}: + return "linux_aarch64" + # No prebuilt wheels published for macOS or Windows + return None + + +def probe_torch_wheel_env(*, timeout: int | None = None) -> dict[str, str] | None: + platform_tag = linux_wheel_platform_tag() + if platform_tag is None: + return None + + try: + probe = subprocess.run( + [ + sys.executable, + "-c", + ( + "import json, sys, re, torch; " + "parts = torch.__version__.split('+', 1)[0].split('.')[:2]; " + "minor = re.sub(r'[^0-9].*', '', parts[1]) if len(parts) > 1 else '0'; " + "torch_mm = parts[0] + '.' + minor; " + "print(json.dumps({" + "'python_tag': f'cp{sys.version_info.major}{sys.version_info.minor}', " + "'torch_mm': torch_mm, " + "'cuda_major': str(int(str(torch.version.cuda).split('.', 1)[0])) if torch.version.cuda else '', " + "'hip_version': str(torch.version.hip) if getattr(torch.version, 'hip', None) else '', " + "'cxx11abi': str(torch._C._GLIBCXX_USE_CXX11_ABI).upper()" + "}))" + ), + ], + stdout = subprocess.PIPE, + stderr = subprocess.PIPE, + text = True, + timeout = timeout, + ) + except subprocess.TimeoutExpired: + return None + + if probe.returncode != 0: + return None + + try: + env = json.loads(probe.stdout.strip()) + except json.JSONDecodeError: + return None + env["platform_tag"] = platform_tag + return env + + +def direct_wheel_url( + *, + filename_prefix: str, + package_version: str, + release_tag: str, + release_base_url: str, + env: dict[str, str] | None, +) -> str | None: + if env is None or not env.get("cuda_major"): + return None + + filename = ( + f"{filename_prefix}-{package_version}" + f"+cu{env['cuda_major']}torch{env['torch_mm']}" + f"cxx11abi{env['cxx11abi']}-{env['python_tag']}-{env['python_tag']}" + f"-{env['platform_tag']}.whl" + ) + return f"{release_base_url}/{release_tag}/{filename}" + + +def flash_attn_package_version(torch_mm: str) -> str | None: + if torch_mm == "2.10": + return "2.8.1" + try: + major, minor = (int(part) for part in torch_mm.split(".", 1)) + except ValueError: + return None + if major == 2 and 4 <= minor <= 9: + return "2.8.3" + return None + + +def flash_attn_wheel_url(env: dict[str, str] | None) -> str | None: + if env is None: + return None + package_version = flash_attn_package_version(env["torch_mm"]) + if package_version is None: + return None + return direct_wheel_url( + filename_prefix = "flash_attn", + package_version = package_version, + release_tag = f"v{package_version}", + release_base_url = FLASH_ATTN_RELEASE_BASE_URL, + env = env, + ) + + +def install_wheel( + wheel_url: str, + *, + python_executable: str, + use_uv: bool, + uv_needs_system: bool = False, + run: Callable[..., subprocess.CompletedProcess[str]] = subprocess.run, +) -> list[tuple[str, subprocess.CompletedProcess[str]]]: + attempts: list[tuple[str, subprocess.CompletedProcess[str]]] = [] + + # Try uv first if available, then fall back to pip + if use_uv and shutil.which("uv"): + uv_cmd = ["uv", "pip", "install"] + if uv_needs_system: + uv_cmd.append("--system") + uv_cmd.extend(["--python", python_executable, "--no-deps", wheel_url]) + result = run( + uv_cmd, + stdout = subprocess.PIPE, + stderr = subprocess.STDOUT, + text = True, + ) + attempts.append(("uv", result)) + if result.returncode == 0: + return attempts + + pip_cmd = [python_executable, "-m", "pip", "install", "--no-deps", wheel_url] + result = run( + pip_cmd, + stdout = subprocess.PIPE, + stderr = subprocess.STDOUT, + text = True, + ) + attempts.append(("pip", result)) + return attempts + + +def url_exists(url: str) -> bool: + try: + request = urllib.request.Request(url, method = "HEAD") + with urllib.request.urlopen(request, timeout = 10): + return True + except urllib.error.HTTPError as exc: + _logger.debug("url_exists(%s): HTTP %s", url, exc.code) + except (urllib.error.URLError, TimeoutError) as exc: + _logger.debug("url_exists(%s): %s", url, exc) + return False diff --git a/studio/frontend/package.json b/studio/frontend/package.json index d9bab4de7e..ffb3c65719 100644 --- a/studio/frontend/package.json +++ b/studio/frontend/package.json @@ -25,7 +25,7 @@ "@fontsource-variable/figtree": "^5.2.10", "@fontsource-variable/inter": "^5.2.8", "@fontsource-variable/space-grotesk": "^5.2.10", - "@hugeicons/core-free-icons": "^3.1.1", + "@hugeicons/core-free-icons": "^4.1.1", "@hugeicons/react": "^1.1.5", "@huggingface/hub": "^2.9.0", "@langchain/core": "^1.1.27", @@ -34,8 +34,8 @@ "@radix-ui/react-select": "^2.2.6", "@radix-ui/react-separator": "^1.1.8", "@radix-ui/react-slot": "^1.2.4", - "@streamdown/cjk": "1.0.2", - "@streamdown/code": "1.0.2", + "@streamdown/cjk": "1.0.3", + "@streamdown/code": "1.1.1", "@streamdown/math": "1.0.2", "@streamdown/mermaid": "1.0.2", "@tailwindcss/vite": "^4.2.2", @@ -53,7 +53,7 @@ "dexie": "^4.3.0", "js-yaml": "^4.1.1", "katex": "^0.16.28", - "lucide-react": "^0.577.0", + "lucide-react": "^1.7.0", "mammoth": "^1.11.0", "motion": "^12.34.0", "next": "^16.1.6", @@ -65,9 +65,9 @@ "react-resizable-panels": "^4.6.4", "recharts": "3.7.0", "remark-gfm": "^4.0.1", - "shadcn": "^3.8.4", + "shadcn": "^4.2.0", "sonner": "^2.0.7", - "streamdown": "2.3.0", + "streamdown": "2.5.0", "tailwind-merge": "^3.4.0", "tailwindcss": "^4.1.18", "tw-animate-css": "^1.4.0", @@ -79,14 +79,14 @@ "@biomejs/biome": "^1.9.4", "@eslint/js": "^9.39.1", "@types/js-yaml": "^4.0.9", - "@types/node": "^24.10.1", + "@types/node": "^25.5.2", "@types/react": "^19.2.5", "@types/react-dom": "^19.2.3", "@vitejs/plugin-react": "^6.0.1", "eslint": "^9.39.1", "eslint-plugin-react-hooks": "^7.0.1", - "eslint-plugin-react-refresh": "^0.4.26", - "globals": "^16.5.0", + "eslint-plugin-react-refresh": "^0.5.2", + "globals": "^17.4.0", "typescript": "~5.9.3", "typescript-eslint": "^8.55.0", "vite": "^8.0.1" diff --git a/studio/frontend/src/app/router.tsx b/studio/frontend/src/app/router.tsx index 13ff8a5cbe..d507929758 100644 --- a/studio/frontend/src/app/router.tsx +++ b/studio/frontend/src/app/router.tsx @@ -13,6 +13,7 @@ import { Route as loginRoute } from "./routes/login"; import { Route as onboardingRoute } from "./routes/onboarding"; import { Route as changePasswordRoute } from "./routes/change-password"; import { Route as studioRoute } from "./routes/studio"; +import { Route as apiKeysRoute } from "./routes/api-keys"; const routeTree = rootRoute.addChildren([ indexRoute, @@ -25,6 +26,7 @@ const routeTree = rootRoute.addChildren([ exportRoute, dataRecipesRoute, dataRecipeRoute, + apiKeysRoute, ]); export const router = createRouter({ routeTree }); diff --git a/studio/frontend/src/app/routes/__root.tsx b/studio/frontend/src/app/routes/__root.tsx index d7780c6743..e1bbdc03f7 100644 --- a/studio/frontend/src/app/routes/__root.tsx +++ b/studio/frontend/src/app/routes/__root.tsx @@ -13,7 +13,14 @@ import { AnimatePresence, motion } from "motion/react"; import { Suspense } from "react"; import { AppProvider } from "../provider"; -const CHAT_ONLY_ALLOWED = new Set(["/", "/chat", "/login", "/signup", "/change-password"]); +const CHAT_ONLY_ALLOWED = new Set([ + "/", + "/chat", + "/login", + "/signup", + "/change-password", + "/api-keys", +]); function isChatOnlyAllowed(pathname: string): boolean { if (CHAT_ONLY_ALLOWED.has(pathname)) return true; @@ -40,11 +47,12 @@ function RootLayout() { return ( {!hideNavbar && } - + diff --git a/studio/frontend/src/app/routes/api-keys.tsx b/studio/frontend/src/app/routes/api-keys.tsx new file mode 100644 index 0000000000..5846690d7b --- /dev/null +++ b/studio/frontend/src/app/routes/api-keys.tsx @@ -0,0 +1,18 @@ +// SPDX-License-Identifier: AGPL-3.0-only +// Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0 + +import { createRoute } from "@tanstack/react-router"; +import { lazy } from "react"; +import { requireAuth } from "../auth-guards"; +import { Route as rootRoute } from "./__root"; + +const ApiKeysPage = lazy(() => + import("@/features/auth/api-keys-page").then((m) => ({ default: m.ApiKeysPage })), +); + +export const Route = createRoute({ + getParentRoute: () => rootRoute, + path: "/api-keys", + beforeLoad: () => requireAuth(), + component: ApiKeysPage, +}); diff --git a/studio/frontend/src/components/assistant-ui/markdown-text.tsx b/studio/frontend/src/components/assistant-ui/markdown-text.tsx index 91ef78fcf9..c7974db365 100644 --- a/studio/frontend/src/components/assistant-ui/markdown-text.tsx +++ b/studio/frontend/src/components/assistant-ui/markdown-text.tsx @@ -41,7 +41,7 @@ const COPY_RESET_MS = 2000; const MERMAID_SOURCE_RE = /```mermaid\s*([\s\S]*?)```/i; const CODE_FENCE_RE = /^```([^\r\n`]*)\r?\n([\s\S]*?)\r?\n?```$/; const ACTION_PANEL_CLASS = - "pointer-events-auto flex shrink-0 items-center gap-2 rounded-md border border-sidebar bg-sidebar/80 px-1.5 py-1 supports-[backdrop-filter]:bg-sidebar/70 supports-[backdrop-filter]:backdrop-blur"; + "pointer-events-auto flex shrink-0 items-center gap-2 rounded-md border border-sidebar bg-sidebar/80 px-1.5 py-1 supports-[backdrop-filter]:bg-sidebar/70 supports-[backdrop-filter]:backdrop-blur dark:border-white/10 dark:bg-code-block dark:supports-[backdrop-filter]:bg-code-block"; const ACTION_BUTTON_CLASS = "cursor-pointer p-1 text-muted-foreground transition-all hover:text-foreground disabled:cursor-not-allowed disabled:opacity-50"; diff --git a/studio/frontend/src/components/assistant-ui/model-selector.tsx b/studio/frontend/src/components/assistant-ui/model-selector.tsx index 441c2b48e4..08e69bbf93 100644 --- a/studio/frontend/src/components/assistant-ui/model-selector.tsx +++ b/studio/frontend/src/components/assistant-ui/model-selector.tsx @@ -203,11 +203,22 @@ export function ModelSelector({ ? lora.name.split("/")[0].trim() : lora.name; // Show type tag instead of base model name + const isLocal = lora.source === "local"; + const isTraining = lora.source === "training"; const isExported = lora.source === "exported"; const isMerged = lora.exportType === "merged"; - const tag = isExported - ? isMerged ? "Merged · Exported" : "LoRA" - : "LoRA"; + const isGguf = lora.exportType === "gguf"; + const tag = isLocal + ? isGguf + ? "GGUF" + : "Local" + : isTraining && isMerged + ? "Full finetune" + : isExported + ? isMerged + ? "Merged · Exported" + : "LoRA · Exported" + : "LoRA"; all.set(lora.id, { ...lora, name: displayName, diff --git a/studio/frontend/src/components/assistant-ui/model-selector/folder-browser.tsx b/studio/frontend/src/components/assistant-ui/model-selector/folder-browser.tsx new file mode 100644 index 0000000000..42bd1716a1 --- /dev/null +++ b/studio/frontend/src/components/assistant-ui/model-selector/folder-browser.tsx @@ -0,0 +1,328 @@ +// 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 { + Dialog, + DialogClose, + DialogContent, + DialogFooter, + DialogHeader, + DialogTitle, +} from "@/components/ui/dialog"; +import { Spinner } from "@/components/ui/spinner"; +import { + type BrowseFoldersResponse, + browseFolders, +} from "@/features/chat/api/chat-api"; +import { cn } from "@/lib/utils"; +import { ArrowUp02Icon, Folder02Icon } from "@hugeicons/core-free-icons"; +import { HugeiconsIcon } from "@hugeicons/react"; +import { useCallback, useEffect, useMemo, useRef, useState } from "react"; + +export interface FolderBrowserProps { + open: boolean; + onOpenChange: (open: boolean) => void; + /** Called with the absolute path the user confirmed. */ + onSelect: (path: string) => void; + /** Optional initial directory. Defaults to the user's home on the server. */ + initialPath?: string; +} + +function splitBreadcrumb(path: string): { label: string; value: string }[] { + if (!path) return []; + // Distinguish path styles BEFORE normalizing separators. On POSIX + // backslashes are valid filename characters, so we cannot blindly + // rewrite ``\`` -> ``/`` -- doing so would mangle directory names + // like ``my\backup`` into ``my/backup`` and produce breadcrumb + // values that 404 on the server. Only Windows-style absolute paths + // (drive letter, or UNC ``\\server\share``) get the conversion. + const isWindowsDrive = /^[A-Za-z]:[\\/]/.test(path) || /^[A-Za-z]:$/.test(path); + const isUnc = /^\\\\/.test(path); + const isWindows = isWindowsDrive || isUnc; + const normalized = isWindows ? path.replace(/\\/g, "/") : path; + const segments = normalized.split("/"); + const parts: { label: string; value: string }[] = []; + + // POSIX absolute path: leading empty segment from split("/") + if (segments[0] === "") { + parts.push({ label: "/", value: "/" }); + let cur = ""; + for (const seg of segments.slice(1)) { + if (!seg) continue; + cur = `${cur}/${seg}`; + parts.push({ label: seg, value: cur }); + } + return parts; + } + + // Windows-ish drive path (C:, D:): first segment is the drive. Use + // ``C:/`` (drive-absolute) as the crumb value so clicking the drive + // root navigates to the root of the drive rather than the + // drive-relative current working directory on that drive (``C:`` + // alone resolves to ``CWD-on-C``, not ``C:\``). + if (/^[A-Za-z]:$/.test(segments[0])) { + const driveRoot = `${segments[0]}/`; + let cur = driveRoot; + parts.push({ label: segments[0], value: driveRoot }); + for (const seg of segments.slice(1)) { + if (!seg) continue; + cur = cur.endsWith("/") ? `${cur}${seg}` : `${cur}/${seg}`; + parts.push({ label: seg, value: cur }); + } + return parts; + } + + // Fallback: relative / UNC-ish. Render as-is as a single crumb. + return [{ label: path, value: path }]; +} + +export function FolderBrowser({ + open, + onOpenChange, + onSelect, + initialPath, +}: FolderBrowserProps) { + const [data, setData] = useState(null); + const [path, setPath] = useState(initialPath); + const [showHidden, setShowHidden] = useState(false); + const [loading, setLoading] = useState(false); + const [error, setError] = useState(null); + const abortRef = useRef(null); + + const navigate = useCallback( + ( + target: string | undefined, + hidden: boolean, + opts?: { fallbackOnError?: boolean }, + ) => { + abortRef.current?.abort(); + const ctrl = new AbortController(); + abortRef.current = ctrl; + setLoading(true); + setError(null); + // Forward the signal so cancelled navigation actually cancels the + // backend enumeration instead of just discarding the response. + browseFolders(target, hidden, ctrl.signal) + .then((res) => { + if (ctrl.signal.aborted) return; + setData(res); + setPath(res.current); + }) + .catch((err) => { + if (ctrl.signal.aborted) return; + // Surface the error, but if the very first request (typically + // a typo'd or denylisted ``initialPath``) fails AND the + // browser is empty (no ``data`` to render against), fall + // back to the user's HOME so the modal is navigable instead + // of an irrecoverable dead end. + const message = err instanceof Error ? err.message : String(err); + setError(message); + if (opts?.fallbackOnError && target !== undefined) { + // Re-issue without a target -> backend defaults to HOME. + // Don't recurse if HOME itself fails (paranoia: shouldn't + // happen since the sandbox allowlist always includes HOME). + queueMicrotask(() => navigate(undefined, hidden)); + } + }) + .finally(() => { + if (!ctrl.signal.aborted) setLoading(false); + }); + }, + [], + ); + + // Fetch when the dialog opens. Only re-run when the dialog transitions + // closed -> open; subsequent navigation is driven by `navigate()` so we + // don't want `path` in the dependency list here. + // eslint-disable-next-line react-hooks/exhaustive-deps + useEffect(() => { + if (!open) return; + // ``fallbackOnError``: if the user-supplied ``initialPath`` is bad + // (typo, denylisted, deleted) we recover into HOME instead of + // showing an empty modal with no breadcrumbs/entries. + navigate(initialPath, showHidden, { fallbackOnError: true }); + // eslint-disable-next-line react-hooks/exhaustive-deps + }, [open]); + + const handleConfirm = useCallback(() => { + if (!path) return; + onSelect(path); + onOpenChange(false); + }, [onSelect, onOpenChange, path]); + + const crumbs = useMemo( + () => (data?.current ? splitBreadcrumb(data.current) : []), + [data?.current], + ); + + return ( + + + + + Browse for folder + + + + {/* Breadcrumb */} +
+ {crumbs.length === 0 ? ( + (loading…) + ) : ( + crumbs.map((c, i) => ( + + + {i < crumbs.length - 1 && ( + / + )} + + )) + )} +
+ + {/* Suggestions (quick-pick chips) */} + {data?.suggestions && data.suggestions.length > 0 && ( +
+ {data.suggestions.map((s) => ( + + ))} +
+ )} + + {/* Entry list */} +
+ {error && ( +
{error}
+ )} + {!error && loading && ( +
+ + Loading… +
+ )} + {!error && !loading && data && ( + <> + {/* Up row */} + {data.parent !== null && ( + + )} + {data.entries.length === 0 && !(data.model_files_here && data.model_files_here > 0) && ( +
+ (empty directory) +
+ )} + {data.model_files_here !== undefined && data.model_files_here > 0 && ( +
+ {data.model_files_here} model file{data.model_files_here === 1 ? "" : "s"} in this folder. Click "Use this folder" to scan it. +
+ )} + {data.truncated === true && ( +
+ Showing first {data.entries.length} entries. Narrow the path + to see more. +
+ )} + {data.entries.map((e) => ( + + ))} + + )} +
+ + {/* Footer */} + + +
+ + + + +
+ + + + ); +} diff --git a/studio/frontend/src/components/assistant-ui/model-selector/pickers.tsx b/studio/frontend/src/components/assistant-ui/model-selector/pickers.tsx index 74ca2542d4..2f661c2e72 100644 --- a/studio/frontend/src/components/assistant-ui/model-selector/pickers.tsx +++ b/studio/frontend/src/components/assistant-ui/model-selector/pickers.tsx @@ -48,6 +48,7 @@ import type { VramFitStatus } from "@/lib/vram"; import { checkVramFit, estimateLoadingVram } from "@/lib/vram"; import { Add01Icon, Cancel01Icon, Folder02Icon, Search01Icon } from "@hugeicons/core-free-icons"; import { HugeiconsIcon } from "@hugeicons/react"; +import { FolderBrowser } from "./folder-browser"; import { Trash2Icon } from "lucide-react"; import { type ReactNode, @@ -399,10 +400,14 @@ function GgufVariantExpander({ ); } -// ── Detect GGUF repos by naming convention ──────────────────── +// ── Detect GGUF repos by naming convention or hub tag ──────────────────── -function isGgufRepo(id: string): boolean { - return id.toUpperCase().includes("-GGUF"); +function hasGgufSuffix(id: string): boolean { + return /-GGUF(?:$|-)/i.test(id); +} + +function isGgufRepo(id: string, hintedIsGguf?: boolean): boolean { + return Boolean(hintedIsGguf) || hasGgufSuffix(id); } /** Extract param count label from model name (e.g. "Qwen3-0.6B" -> "0.6B"). */ @@ -451,6 +456,33 @@ export function HubModelPicker({ const { results, isLoading, isLoadingMore, fetchMore } = useHfModelSearch(debouncedQuery); + // Sets of lowercased repo ids that the store or HF search have + // confirmed are GGUF. Absence means "no hint" and lets hasGgufSuffix + // take over as fallback, rather than conflating unknown with known- + // not-GGUF. Keys are lowercased so that store IDs and HF search IDs + // that differ only by casing still match the same hint. + const modelGgufIds = useMemo(() => { + const ids = new Set(); + for (const model of models) { + if (model.isGguf) ids.add(model.id.toLowerCase()); + } + return ids; + }, [models]); + const resultGgufIds = useMemo(() => { + const ids = new Set(); + for (const result of results) { + if (result.isGguf) ids.add(result.id.toLowerCase()); + } + return ids; + }, [results]); + const isKnownGgufRepo = useCallback( + (id: string): boolean => { + const key = id.toLowerCase(); + return isGgufRepo(id, resultGgufIds.has(key) || modelGgufIds.has(key)); + }, + [modelGgufIds, resultGgufIds], + ); + // Track which GGUF repo is expanded for variant selection const [expandedGguf, setExpandedGguf] = useState(null); @@ -481,6 +513,7 @@ export function HubModelPicker({ const [folderError, setFolderError] = useState(null); const [showFolderInput, setShowFolderInput] = useState(false); const [folderLoading, setFolderLoading] = useState(false); + const [showFolderBrowser, setShowFolderBrowser] = useState(false); const refreshLocalModelsList = useCallback(() => { listLocalModels() @@ -506,11 +539,22 @@ export function HubModelPicker({ .catch(() => {}); }, []); - const handleAddFolder = useCallback(async () => { - const trimmed = folderInput.trim(); + const handleAddFolder = useCallback(async (overridePath?: string) => { + // Accept an explicit path so the folder browser can submit the + // chosen path in the same tick it calls `setFolderInput`; reading + // `folderInput` alone would race the state update. + const raw = overridePath !== undefined ? overridePath : folderInput; + const trimmed = raw.trim(); if (!trimmed || folderLoading) return; setFolderError(null); setFolderLoading(true); + // True when the request originated from the folder browser's + // ``onSelect`` (one-click "Use this folder"). In that flow the + // typed-input panel is closed, so the inline ``folderError`` + // paragraph is invisible. Surface failures via toast instead so + // the action doesn't appear to silently no-op when the backend + // rejects (denylisted path, sandbox 403, etc.). + const fromBrowser = overridePath !== undefined; try { const created = await addScanFolder(trimmed); // Backend returns existing row for duplicates, so deduplicate @@ -526,7 +570,11 @@ export function HubModelPicker({ // Background reconciliation with the server void refreshScanFolders(); } catch (e) { - setFolderError(e instanceof Error ? e.message : "Failed to add folder"); + const message = e instanceof Error ? e.message : "Failed to add folder"; + setFolderError(message); + if (fromBrowser) { + toast.error("Couldn't add folder", { description: message }); + } } finally { setFolderLoading(false); } @@ -569,7 +617,11 @@ export function HubModelPicker({ refreshLocalModelsList(); refreshScanFolders(); - if (alreadyCached) return; + // Always refetch cached GGUF/model lists. The module-level caches give + // an instant render with stale data (no spinner flash), but newly + // downloaded repos won't appear unless we re-hit the backend on every + // mount. Initial state already has cachedReady=alreadyCached, so the + // background refresh is invisible when we already had data. let done = 0; const check = () => { if (++done >= 2) setCachedReady(true); @@ -588,7 +640,7 @@ export function HubModelPicker({ }) .catch(() => {}) .finally(check); - }, [alreadyCached, refreshLocalModelsList, refreshScanFolders]); + }, [refreshLocalModelsList, refreshScanFolders]); const handleDeleteConfirm = useCallback(async () => { if (!deleteTarget) return; @@ -625,17 +677,17 @@ export function HubModelPicker({ const recommendedIds = useMemo(() => { const all = dedupe([...models.map((model) => model.id), value ?? ""]) .filter((id) => !downloadedSet.has(id.toLowerCase())) - .filter((id) => !chatOnly || isGgufRepo(id)) + .filter((id) => !chatOnly || isKnownGgufRepo(id)) .filter((id) => !/-FP8[-.]|FP8-Dynamic/i.test(id)); // Sort: GGUFs first, then hub models const gguf: string[] = []; const hub: string[] = []; for (const id of all) { - if (isGgufRepo(id)) gguf.push(id); + if (isKnownGgufRepo(id)) gguf.push(id); else hub.push(id); } return [...gguf, ...hub]; - }, [models, value, downloadedSet, chatOnly]); + }, [models, value, downloadedSet, chatOnly, isKnownGgufRepo]); // Infinite scroll paging for the recommended section const [recommendedPage, setRecommendedPage] = useState(1); @@ -645,7 +697,7 @@ export function HubModelPicker({ }, [models, chatOnly]); const visibleRecommendedIds = useMemo(() => { - const hubStartIndex = recommendedIds.findIndex((id) => !isGgufRepo(id)); + const hubStartIndex = recommendedIds.findIndex((id) => !isKnownGgufRepo(id)); const allGguf = hubStartIndex === -1 ? recommendedIds @@ -659,7 +711,7 @@ export function HubModelPicker({ result.push(...allHub.slice(p * 4, (p + 1) * 4)); } return result; - }, [recommendedIds, recommendedPage]); + }, [recommendedIds, recommendedPage, isKnownGgufRepo]); const hasMoreRecommended = visibleRecommendedIds.length < recommendedIds.length; @@ -681,8 +733,8 @@ export function HubModelPicker({ const ids = showHfSection ? [...new Set([...visibleRecommendedIds, ...filteredRecommendedIds])] : visibleRecommendedIds; - return ids.filter((id) => !isGgufRepo(id)); - }, [visibleRecommendedIds, showHfSection, filteredRecommendedIds]); + return ids.filter((id) => !isKnownGgufRepo(id)); + }, [visibleRecommendedIds, showHfSection, filteredRecommendedIds, isKnownGgufRepo]); const { paramCountById: recommendedParamCountById } = useRecommendedModelVram(idsForVram); @@ -697,9 +749,9 @@ export function HubModelPicker({ return results .map((result) => result.id) .filter((id) => !recommendedSet.has(id)) - .filter((id) => !chatOnly || isGgufRepo(id)) + .filter((id) => !chatOnly || isKnownGgufRepo(id)) .filter((id) => !/-FP8[-.]|FP8-Dynamic/i.test(id)); - }, [recommendedSet, results, showHfSection, chatOnly]); + }, [recommendedSet, results, showHfSection, chatOnly, isKnownGgufRepo]); const metricsById = useMemo( () => @@ -800,14 +852,14 @@ export function HubModelPicker({ /** Handle clicking a model row — GGUF repos expand, others load directly. */ const handleModelClick = useCallback( (id: string) => { - if (isGgufRepo(id)) { + if (isKnownGgufRepo(id)) { // Toggle GGUF variant expander setExpandedGguf((prev) => (prev === id ? null : id)); } else { onSelect(id, { source: "hub", isLora: false }); } }, - [onSelect], + [onSelect, isKnownGgufRepo], ); return ( @@ -841,14 +893,18 @@ export function HubModelPicker({ (cachedGguf.length > 0 || (!chatOnly && cachedModels.length > 0)) ? ( <> - {"\uD83E\uDDA5"} Downloaded + Downloaded {cachedGguf.map((c) => (
handleModelClick(c.repo_id)} + onClick={() => + setExpandedGguf((prev) => + prev === c.repo_id ? null : c.repo_id, + ) + } vramStatus={null} /> {expandedGguf === c.repo_id && ( @@ -909,7 +965,7 @@ export function HubModelPicker({ { @@ -945,30 +1001,42 @@ export function HubModelPicker({ {!showHfSection ? ( <> -
+
Custom Folders - +
+ + +
{/* Folder paths */} {scanFolders.map((f) => (
handleRemoveFolder(f.id)} aria-label={`Remove folder ${f.path}`} - className="shrink-0 rounded p-0.5 text-muted-foreground/40 opacity-100 md:opacity-0 md:group-hover:opacity-100 focus-visible:opacity-100 transition-opacity hover:text-destructive" + className="shrink-0 rounded p-1 text-foreground/70 transition-colors hover:bg-destructive/10 hover:text-destructive focus-visible:bg-destructive/10 focus-visible:text-destructive" > - +
))} @@ -1007,7 +1075,17 @@ export function HubModelPicker({ /> +
)} - {/* Empty state */} - {scanFolders.length === 0 && customFolderModels.length === 0 && !showFolderInput && ( - - )} + { + setFolderInput(picked); + setFolderError(null); + // One-click UX: the "Use this folder" button submits + // the scan folder directly. Pass the path explicitly + // because `folderInput` state hasn't flushed yet. + void handleAddFolder(picked); + }} + /> + {/* Models from custom folders */} {customFolderModels.map((m) => { const isGguf = isGgufRepo(m.id) || isGgufRepo(m.display_name) || - m.path.endsWith(".gguf"); + m.path.toLowerCase().endsWith(".gguf"); return (
- {"\uD83E\uDDA5"} Recommended + Recommended {visibleRecommendedIds.length === 0 ? (
No default models. @@ -1089,16 +1171,22 @@ export function HubModelPicker({ handleModelClick(id)} + onClick={() => { + if (isKnownGgufRepo(id)) { + setExpandedGguf((prev) => (prev === id ? null : id)); + } else { + handleModelClick(id); + } + }} vramStatus={ - isGgufRepo(id) ? null : (vram?.status ?? null) + isKnownGgufRepo(id) ? null : (vram?.status ?? null) } - vramEst={isGgufRepo(id) ? undefined : vram?.est} + vramEst={isKnownGgufRepo(id) ? undefined : vram?.est} gpuGb={gpu.available ? gpu.memoryTotalGb : undefined} /> {expandedGguf === id && ( @@ -1128,7 +1216,7 @@ export function HubModelPicker({ {showHfSection && filteredRecommendedIds.length > 0 ? ( <> - {"\uD83E\uDDA5"} Recommended + Recommended {filteredRecommendedIds.map((id) => { const vram = recommendedVramMap.get(id); return ( @@ -1136,16 +1224,22 @@ export function HubModelPicker({ handleModelClick(id)} + onClick={() => { + if (isKnownGgufRepo(id)) { + setExpandedGguf((prev) => (prev === id ? null : id)); + } else { + handleModelClick(id); + } + }} vramStatus={ - isGgufRepo(id) ? null : (vram?.status ?? null) + isKnownGgufRepo(id) ? null : (vram?.status ?? null) } - vramEst={isGgufRepo(id) ? undefined : vram?.est} + vramEst={isKnownGgufRepo(id) ? undefined : vram?.est} gpuGb={gpu.available ? gpu.memoryTotalGb : undefined} /> {expandedGguf === id && ( @@ -1175,24 +1269,31 @@ export function HubModelPicker({ No matching models.
) : null - ) : ( + ) : ( hfIds.map((id) => { const vram = vramMap.get(id); + const isSearchGguf = isKnownGgufRepo(id); return (
handleModelClick(id)} + onClick={() => { + if (isSearchGguf) { + setExpandedGguf((prev) => (prev === id ? null : id)); + } else { + handleModelClick(id); + } + }} vramStatus={ - isGgufRepo(id) ? null : (vram?.status ?? null) + isSearchGguf ? null : (vram?.status ?? null) } - vramEst={isGgufRepo(id) ? undefined : vram?.est} + vramEst={isSearchGguf ? undefined : vram?.est} gpuGb={gpu.available ? gpu.memoryTotalGb : undefined} /> {expandedGguf === id && ( @@ -1330,7 +1431,7 @@ export function LoraModelPicker({ setQuery(event.target.value)} - placeholder="Search local adapters" + placeholder="Search trained models" className="h-9 pl-8" />
@@ -1339,7 +1440,7 @@ export function LoraModelPicker({
{grouped.length === 0 ? (
- No adapters found. + No trained models found.
) : ( grouped.map(([baseModel, adapters], index) => ( @@ -1348,9 +1449,11 @@ export function LoraModelPicker({ {baseModel} {adapters.map((adapter) => { const isLocal = adapter.source === "local"; + const isTraining = adapter.source === "training"; const isExported = adapter.source === "exported"; const isMerged = adapter.exportType === "merged"; const isGguf = adapter.exportType === "gguf"; + const isTrainingFull = isTraining && isMerged; const isLocalGgufDir = isLocal && (isGgufRepo(adapter.id) || isGgufRepo(adapter.name)); @@ -1360,6 +1463,8 @@ export function LoraModelPicker({ : "Local" : isGguf ? "GGUF" + : isTrainingFull + ? "Full" : isExported ? isMerged ? "Merged" @@ -1369,6 +1474,8 @@ export function LoraModelPicker({ ? isLocalGgufDir ? "GGUF" : "Local" + : isTrainingFull + ? "Full finetune" : isExported ? `${tag} · Exported` : tag; diff --git a/studio/frontend/src/components/assistant-ui/model-selector/types.ts b/studio/frontend/src/components/assistant-ui/model-selector/types.ts index f70cfc3b01..215dd2b38e 100644 --- a/studio/frontend/src/components/assistant-ui/model-selector/types.ts +++ b/studio/frontend/src/components/assistant-ui/model-selector/types.ts @@ -8,6 +8,7 @@ export interface ModelOption { name: string; description?: string; icon?: ReactNode; + isGguf?: boolean; } export interface LoraModelOption extends ModelOption { @@ -24,4 +25,3 @@ export interface ModelSelectorChangeMeta { isDownloaded?: boolean; expectedBytes?: number; } - diff --git a/studio/frontend/src/components/assistant-ui/thread.tsx b/studio/frontend/src/components/assistant-ui/thread.tsx index 4db0eda0ad..8f41987fbf 100644 --- a/studio/frontend/src/components/assistant-ui/thread.tsx +++ b/studio/frontend/src/components/assistant-ui/thread.tsx @@ -56,9 +56,12 @@ import { RefreshCwIcon, SquareIcon, TerminalIcon, + Trash2Icon, XIcon, } from "lucide-react"; import { type FC, useCallback, useEffect, useRef, useState } from "react"; +import { toast } from "sonner"; +import { deleteThreadMessage } from "@/features/chat/utils/delete-thread-message"; import { useChatRuntimeStore } from "@/features/chat/stores/chat-runtime-store"; export const Thread: FC<{ hideComposer?: boolean; hideWelcome?: boolean }> = ({ @@ -89,14 +92,27 @@ export const Thread: FC<{ hideComposer?: boolean; hideWelcome?: boolean }> = ({ }} /> - + {!hideComposer && (
)} - +
+ +
!thread.isEmpty}> {!hideComposer && } @@ -112,7 +128,7 @@ const ThreadScrollToBottom: FC = () => { @@ -622,6 +638,41 @@ const AssistantMessage: FC = () => { const COPY_RESET_MS = 2000; +const DeleteMessageButton: FC = () => { + const aui = useAui(); + const messageId = useAuiState(({ message }) => message.id); + const isRunning = useAuiState(({ thread }) => thread.isRunning); + + const handleDelete = async () => { + const remoteId = aui.threadListItem().getState().remoteId; + const thread = aui.thread(); + try { + await deleteThreadMessage({ + thread: { + export: () => thread.export(), + import: (data) => thread.import(data), + }, + messageId, + remoteId, + }); + } catch (error) { + console.error("Failed to delete message", error); + toast.error("Failed to delete message"); + } + }; + + return ( + + + + ); +}; + const CopyButton: FC = () => { const aui = useAui(); const [copied, setCopied] = useState(false); @@ -660,6 +711,7 @@ const AssistantActionBar: FC = () => { + @@ -735,6 +787,7 @@ const UserActionBar: FC = () => { + ); }; diff --git a/studio/frontend/src/components/navbar.tsx b/studio/frontend/src/components/navbar.tsx index 121c559db8..2d310a81c3 100644 --- a/studio/frontend/src/components/navbar.tsx +++ b/studio/frontend/src/components/navbar.tsx @@ -29,6 +29,7 @@ import { ChefHatIcon, Copy01Icon, CursorInfo02Icon, + Key01Icon, PackageIcon, Tick02Icon, ZapIcon, @@ -416,6 +417,20 @@ export function Navbar() {
+
+ + + API Keys + +
{tourId ? (
+ ); +} + +function RevealKeyDialog({ + open, + rawKey, + onClose, +}: { + open: boolean; + rawKey: string; + onClose: () => void; +}) { + return ( + !o && onClose()}> + + + API Key Created + + Copy this key now. It will not be shown again. + + +
+ + {rawKey} + + +
+
+ +

+ Store this key securely. You will not be able to see it again after closing this dialog. +

+
+ + + + + + ); +} + +function CreateKeyForm({ onCreated }: { onCreated: (rawKey: string) => void }) { + const [name, setName] = useState(""); + const [expiresInDays, setExpiresInDays] = useState(""); + const [loading, setLoading] = useState(false); + + const handleSubmit = async (e: React.FormEvent) => { + e.preventDefault(); + if (!name.trim()) return; + setLoading(true); + try { + const days = expiresInDays ? parseInt(expiresInDays, 10) : null; + const result = await createApiKey(name.trim(), days); + onCreated(result.key); + setName(""); + setExpiresInDays(""); + } finally { + setLoading(false); + } + }; + + return ( +
+
+ + setName(e.target.value)} + required + /> +
+
+ + setExpiresInDays(e.target.value)} + /> +
+ +
+ ); +} + +function KeysTable({ + keys, + onRevoke, +}: { + keys: ApiKey[]; + onRevoke: (id: number) => void; +}) { + if (keys.length === 0) { + return ( +

+ No API keys yet. Create one above. +

+ ); + } + + return ( +
+ + + + + + + + + + + + {keys.map((k) => ( + + + + + + + + + ))} + +
NameKeyCreatedLast usedExpires +
{k.name} + + sk-unsloth-{k.key_prefix}... + + {formatDate(k.created_at)}{formatDate(k.last_used_at)}{formatDate(k.expires_at)} + {k.is_active ? ( + + ) : ( + Revoked + )} +
+
+ ); +} + +function UsageExamples() { + const base = window.location.origin; + + const curlExample = `curl ${base}/v1/chat/completions \\ + -H "Authorization: Bearer sk-unsloth-YOUR_KEY" \\ + -H "Content-Type: application/json" \\ + -d '{ + "messages": [{"role": "user", "content": "Hello"}], + "stream": true + }'`; + + const pythonExample = `from openai import OpenAI + +client = OpenAI( + base_url="${base}/v1", + api_key="sk-unsloth-YOUR_KEY", +) + +response = client.chat.completions.create( + model="current", + messages=[{"role": "user", "content": "Hello"}], + stream=True, +) +for chunk in response: + print(chunk.choices[0].delta.content or "", end="")`; + + const toolsExample = `curl ${base}/v1/chat/completions \\ + -H "Authorization: Bearer sk-unsloth-YOUR_KEY" \\ + -H "Content-Type: application/json" \\ + -d '{ + "messages": [{"role": "user", "content": "Search for Python 3.13 features"}], + "stream": true, + "enable_tools": true, + "enabled_tools": ["web_search", "python"], + "session_id": "my-session" + }'`; + + return ( +
+

Usage examples

+
+
+

curl

+ 
+            {curlExample}
+
+
+
+

Python (OpenAI SDK)

+ 
+            {pythonExample}
+
+
+
+

With tools (web search + code execution)

+ 
+            {toolsExample}
+
+
+
+
+ ); +} + +// --------------------------------------------------------------------------- +// Page +// --------------------------------------------------------------------------- + +export function ApiKeysPage() { + const [keys, setKeys] = useState([]); + const [revealedKey, setRevealedKey] = useState(null); + const [error, setError] = useState(null); + + const loadKeys = useCallback(async () => { + try { + setError(null); + const loaded = await fetchApiKeys(); + setKeys(loaded); + } catch (e) { + setError(e instanceof Error ? e.message : "Failed to load API keys"); + } + }, []); + + useEffect(() => { + void loadKeys(); + }, [loadKeys]); + + const handleCreated = (rawKey: string) => { + setRevealedKey(rawKey); + void loadKeys(); + }; + + const handleRevoke = async (keyId: number) => { + try { + await revokeApiKey(keyId); + void loadKeys(); + } catch { + setError("Failed to revoke key"); + } + }; + + return ( + +
+
+
+ +
+
+

API Keys

+

+ Create keys to access Unsloth Studio programmatically via the OpenAI-compatible API. +

+
+
+ + {error && ( +
+ + {error} +
+ )} + + + + +
+ + setRevealedKey(null)} + /> + + ); +} diff --git a/studio/frontend/src/features/auth/index.ts b/studio/frontend/src/features/auth/index.ts index 75db92432c..9962f6d431 100644 --- a/studio/frontend/src/features/auth/index.ts +++ b/studio/frontend/src/features/auth/index.ts @@ -1,6 +1,7 @@ // SPDX-License-Identifier: AGPL-3.0-only // Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0 +export { ApiKeysPage } from "./api-keys-page"; export { LoginPage } from "./login-page"; export { ChangePasswordPage } from "./change-password-page"; export { authFetch, refreshSession } from "./api"; diff --git a/studio/frontend/src/features/chat/api/chat-adapter.ts b/studio/frontend/src/features/chat/api/chat-adapter.ts index 3d1bce2905..60904f3cb5 100644 --- a/studio/frontend/src/features/chat/api/chat-adapter.ts +++ b/studio/frontend/src/features/chat/api/chat-adapter.ts @@ -11,6 +11,7 @@ import { listGgufVariants, loadModel, streamChatCompletions, + validateModel, } from "./chat-api"; import { db } from "../db"; import { useChatRuntimeStore } from "../stores/chat-runtime-store"; @@ -46,6 +47,41 @@ type RunMessage = RunMessages[number]; /** Tracks which user messages were sent with an audio file (messageId → filename). */ export const sentAudioNames = new Map(); +/** + * Match error messages that indicate the request filled or would fill + * the KV cache, so the UI can show a dedicated toast pointing at the + * ``Context Length`` setting. + * + * Two wordings reach the client and both must hit: + * + * 1. The raw llama-server text when ``--no-context-shift`` trips -- + * "the request exceeds the available context size (N tokens)". + * 2. The rewritten friendly text emitted by + * ``backend/routes/inference.py::_friendly_error`` -- "Message too + * long: X tokens exceeds the Y-token context window. Try + * increasing the Context Length ..." This is the one most users + * see on the streaming GGUF path. + * + * We match on substrings rather than full regexes because both layers + * have drifted across versions (llama.cpp master has tweaked the + * phrasing; ``_friendly_error`` has gone through several copy edits). + */ +export function isContextLimitError(message: string): boolean { + if (!message) return false; + const m = message.toLowerCase(); + return ( + // Raw llama-server wording. + m.includes("context size") || + m.includes("context shift") || + m.includes("exceeds the available context") || + // Backend _friendly_error rewrite. + m.includes("message too long") || + m.includes("context window") || + // n_ctx mentions that carry an "exceed"/"full" signal. + (m.includes("n_ctx") && (m.includes("exceed") || m.includes("full"))) + ); +} + /** Parse "Title: ...\nURL: ...\nSnippet: ..." blocks into source content parts. */ function parseSourcesFromResult(raw: string): { type: "source"; sourceType: "url"; id: string; url: string; title: string; metadata?: { description: string } }[] { if (!raw) return []; @@ -252,13 +288,39 @@ function waitForModelReady(abortSignal?: AbortSignal): Promise { * without selecting one. Prefers GGUF (picks smallest cached variant), * falls back to smallest cached safetensors model. */ -async function autoLoadSmallestModel(): Promise { - const hfToken = useChatRuntimeStore.getState().hfToken || null; +async function autoLoadSmallestModel(): Promise<{ + loaded: boolean; + blockedByTrustRemoteCode: boolean; +}> { + const store = useChatRuntimeStore.getState(); + const hfToken = store.hfToken || null; + const trustRemoteCode = store.params.trustRemoteCode ?? false; const toastId = toast("Loading a model…", { description: "Auto-selecting the smallest downloaded model.", duration: 5000, closeButton: true, }); + let blockedByTrustRemoteCode = false; + let hadNonTrustFailure = false; + + async function canAutoLoad(payload: { + model_path: string; + max_seq_length: number; + is_lora: boolean; + gguf_variant?: string | null; + }): Promise { + const validation = await validateModel({ + ...payload, + hf_token: hfToken, + load_in_4bit: true, + trust_remote_code: trustRemoteCode, + }); + if (validation.requires_trust_remote_code && !trustRemoteCode) { + blockedByTrustRemoteCode = true; + return false; + } + return true; + } try { const [ggufRepos, modelRepos] = await Promise.all([ listCachedGguf().catch(() => []), @@ -277,6 +339,16 @@ async function autoLoadSmallestModel(): Promise { .sort((a, b) => a.size_bytes - b.size_bytes); if (downloaded.length > 0) { const variant = downloaded[0]; + if ( + !(await canAutoLoad({ + model_path: repo.repo_id, + max_seq_length: 0, + is_lora: false, + gguf_variant: variant.quant, + })) + ) { + continue; + } const loadResp = await loadModel({ model_path: repo.repo_id, hf_token: hfToken, @@ -284,10 +356,13 @@ async function autoLoadSmallestModel(): Promise { load_in_4bit: true, is_lora: false, gguf_variant: variant.quant, - trust_remote_code: false, + trust_remote_code: trustRemoteCode, }); useChatRuntimeStore.getState().setCheckpoint(repo.repo_id, variant.quant); const store = useChatRuntimeStore.getState(); + store.setModelRequiresTrustRemoteCode( + loadResp.requires_trust_remote_code ?? false, + ); store.setParams({ ...store.params, maxTokens: loadResp.context_length ?? 131072 }); // Add model to store so the selector shows the name const autoModel: ChatModelSummary = { @@ -319,9 +394,10 @@ async function autoLoadSmallestModel(): Promise { chatTemplateOverride: null, }); toast.success(`Loaded ${repo.repo_id} (${variant.quant})`, { id: toastId }); - return true; + return { loaded: true, blockedByTrustRemoteCode: false }; } } catch { + hadNonTrustFailure = true; continue; } } @@ -332,6 +408,16 @@ async function autoLoadSmallestModel(): Promise { const sorted = [...modelRepos].sort((a, b) => a.size_bytes - b.size_bytes); for (const repo of sorted) { try { + if ( + !(await canAutoLoad({ + model_path: repo.repo_id, + max_seq_length: 4096, + is_lora: false, + gguf_variant: null, + })) + ) { + continue; + } const sfLoadResp = await loadModel({ model_path: repo.repo_id, hf_token: hfToken, @@ -339,10 +425,13 @@ async function autoLoadSmallestModel(): Promise { load_in_4bit: true, is_lora: false, gguf_variant: null, - trust_remote_code: false, + trust_remote_code: trustRemoteCode, }); useChatRuntimeStore.getState().setCheckpoint(repo.repo_id); const store = useChatRuntimeStore.getState(); + store.setModelRequiresTrustRemoteCode( + sfLoadResp.requires_trust_remote_code ?? false, + ); store.setParams({ ...store.params, maxTokens: 4096 }); const sfModel: ChatModelSummary = { id: repo.repo_id, @@ -355,8 +444,9 @@ async function autoLoadSmallestModel(): Promise { store.setModels([...store.models, sfModel]); } toast.success(`Loaded ${repo.repo_id}`, { id: toastId }); - return true; + return { loaded: true, blockedByTrustRemoteCode: false }; } catch { + hadNonTrustFailure = true; continue; } } @@ -369,6 +459,17 @@ async function autoLoadSmallestModel(): Promise { duration: 30000, }); try { + if ( + !(await canAutoLoad({ + model_path: "unsloth/Qwen3.5-4B-GGUF", + max_seq_length: 0, + is_lora: false, + gguf_variant: "UD-Q4_K_XL", + })) + ) { + toast.dismiss(toastId); + return { loaded: false, blockedByTrustRemoteCode }; + } const loadResp = await loadModel({ model_path: "unsloth/Qwen3.5-4B-GGUF", hf_token: hfToken, @@ -376,10 +477,13 @@ async function autoLoadSmallestModel(): Promise { load_in_4bit: true, is_lora: false, gguf_variant: "UD-Q4_K_XL", - trust_remote_code: false, + trust_remote_code: trustRemoteCode, }); useChatRuntimeStore.getState().setCheckpoint("unsloth/Qwen3.5-4B-GGUF", "UD-Q4_K_XL"); const store = useChatRuntimeStore.getState(); + store.setModelRequiresTrustRemoteCode( + loadResp.requires_trust_remote_code ?? false, + ); store.setParams({ ...store.params, maxTokens: loadResp.context_length ?? 131072 }); const defaultModel: ChatModelSummary = { id: "unsloth/Qwen3.5-4B-GGUF", @@ -406,14 +510,24 @@ async function autoLoadSmallestModel(): Promise { chatTemplateOverride: null, }); toast.success("Loaded Qwen3.5-4B (UD-Q4_K_XL)", { id: toastId }); - return true; + return { loaded: true, blockedByTrustRemoteCode: false }; } catch { toast.dismiss(toastId); - return false; + hadNonTrustFailure = true; + return { + loaded: false, + blockedByTrustRemoteCode: + blockedByTrustRemoteCode && !hadNonTrustFailure, + }; } } catch { toast.dismiss(toastId); - return false; + hadNonTrustFailure = true; + return { + loaded: false, + blockedByTrustRemoteCode: + blockedByTrustRemoteCode && !hadNonTrustFailure, + }; } } @@ -434,11 +548,19 @@ export function createOpenAIStreamAdapter(): ChatModelAdapter { if (!useChatRuntimeStore.getState().params.checkpoint) { // Auto-load the smallest downloaded model - const loaded = await autoLoadSmallestModel(); + const { loaded, blockedByTrustRemoteCode } = + await autoLoadSmallestModel(); if (!loaded) { - toast.error("No model loaded", { - description: "Pick a model in the top bar, then retry.", - }); + toast.error( + blockedByTrustRemoteCode + ? "Enable custom code to auto-load this model" + : "No model loaded", + { + description: blockedByTrustRemoteCode + ? 'Turn on "Enable custom code" in Chat Settings, or pick another model in the top bar.' + : "Pick a model in the top bar, then retry.", + }, + ); throw new Error("Load a model first."); } } @@ -781,9 +903,24 @@ export function createOpenAIStreamAdapter(): ChatModelAdapter { } catch (err) { settleFirstTokenErr(err instanceof Error ? err : new Error("Generation failed")); if (!abortSignal.aborted) { - toast.error("Generation failed", { - description: err instanceof Error ? err.message : "Unknown error", - }); + const msg = err instanceof Error ? err.message : String(err); + if (isContextLimitError(msg)) { + // llama-server was launched with --no-context-shift, so it + // returns a hard error instead of silently dropping old + // turns from the KV cache. Point the user at the exact + // control that raises the ceiling. + toast.error("Context limit reached", { + description: + "The conversation has filled the model's context window. " + + "Increase \"Context Length\" in the chat Settings panel (⚙ in the top-right), " + + "or start a new chat.", + duration: 8000, + }); + } else { + toast.error("Generation failed", { + description: msg || "Unknown error", + }); + } } throw err; } finally { diff --git a/studio/frontend/src/features/chat/api/chat-api.ts b/studio/frontend/src/features/chat/api/chat-api.ts index 7bdd76296b..9aacfc5af4 100644 --- a/studio/frontend/src/features/chat/api/chat-api.ts +++ b/studio/frontend/src/features/chat/api/chat-api.ts @@ -117,14 +117,59 @@ export async function getGgufDownloadProgress( return parseJsonOrThrow(response); } +export interface DownloadProgressResponse { + downloaded_bytes: number; + expected_bytes: number; + progress: number; + /** + * Resolved on-disk path of the snapshot dir (or cache repo root if no + * snapshot exists yet). Null when nothing has been written to the + * cache for this repo. + */ + cache_path: string | null; +} + export async function getDownloadProgress( repoId: string, -): Promise<{ downloaded_bytes: number; expected_bytes: number; progress: number }> { +): Promise { const params = new URLSearchParams({ repo_id: repoId }); const response = await authFetch(`/api/models/download-progress?${params}`); return parseJsonOrThrow(response); } +export async function getDatasetDownloadProgress( + repoId: string, +): Promise { + const params = new URLSearchParams({ repo_id: repoId }); + const response = await authFetch(`/api/datasets/download-progress?${params}`); + return parseJsonOrThrow(response); +} + +export type ModelLoadPhase = "mmap" | "ready" | null; + +export interface LoadProgressResponse { + /** + * Load phase: ``"mmap"`` while the llama-server subprocess is paging + * weight shards into RAM, ``"ready"`` once it has reported healthy, + * or ``null`` when no load is in flight. + */ + phase: ModelLoadPhase; + bytes_loaded: number; + bytes_total: number; + fraction: number; +} + +/** + * Fetch the active GGUF load's mmap/upload progress. Complements + * ``getDownloadProgress`` / ``getGgufDownloadProgress`` for the window + * between "download complete" and "chat ready", which for large MoE + * models can be several minutes of otherwise-opaque spinning. + */ +export async function getLoadProgress(): Promise { + const response = await authFetch(`/api/inference/load-progress`); + return parseJsonOrThrow(response); +} + export interface LocalModelInfo { id: string; display_name: string; @@ -202,6 +247,42 @@ export async function removeScanFolder(id: number): Promise { await parseJsonOrThrow(response); } +export interface BrowseEntry { + name: string; + has_models: boolean; + hidden: boolean; +} + +export interface BrowseFoldersResponse { + current: string; + parent: string | null; + entries: BrowseEntry[]; + suggestions: string[]; + truncated?: boolean; + model_files_here?: number; +} + +export async function browseFolders( + path?: string, + showHidden = false, + signal?: AbortSignal, +): Promise { + const params = new URLSearchParams(); + if (path !== undefined && path !== null) params.set("path", path); + if (showHidden) params.set("show_hidden", "true"); + const qs = params.toString(); + // Forward the AbortSignal through authFetch -> fetch so that a + // navigation cancelled in the FolderBrowser (rapid breadcrumb / row / + // hidden-toggle clicks) actually cancels the in-flight HTTP request + // server-side, instead of merely dropping the response client-side + // while the backend keeps walking large directory trees. + const response = await authFetch( + `/api/models/browse-folders${qs ? `?${qs}` : ""}`, + signal ? { signal } : undefined, + ); + return parseJsonOrThrow(response); +} + export async function listGgufVariants( repoId: string, hfToken?: string, diff --git a/studio/frontend/src/features/chat/chat-page.tsx b/studio/frontend/src/features/chat/chat-page.tsx index cf1ba11d7b..c3153cdec7 100644 --- a/studio/frontend/src/features/chat/chat-page.tsx +++ b/studio/frontend/src/features/chat/chat-page.tsx @@ -71,6 +71,7 @@ type LoraCandidate = { id: string; baseModel: string; updatedAt?: number; + exportType?: "lora" | "merged" | "gguf"; }; function normalizeModelRef(value: string | null | undefined): string { @@ -81,12 +82,13 @@ function pickBestLoraForBase( loras: LoraCandidate[], baseModel: string | null, ): LoraCandidate | null { - if (loras.length === 0) return null; - const sorted = [...loras].sort( + const adapterOnly = loras.filter((lora) => lora.exportType === "lora"); + if (adapterOnly.length === 0) return null; + const sorted = [...adapterOnly].sort( (a, b) => (b.updatedAt ?? -1) - (a.updatedAt ?? -1), ); const normalizedBase = normalizeModelRef(baseModel); - if (!normalizedBase) return sorted[0]; + if (!normalizedBase) return sorted[0] ?? null; const exact = sorted.find( (lora) => normalizeModelRef(lora.baseModel) === normalizedBase, @@ -101,7 +103,7 @@ function pickBestLoraForBase( normalizedBase.includes(normalizedLoraBase) ); }); - return partial ?? sorted[0]; + return partial ?? sorted[0] ?? null; } function messageHasImage(message: MessageRecord): boolean { @@ -154,7 +156,8 @@ type CompareModelSelection = { function useIsLoraCompare(): boolean { return useChatRuntimeStore((s) => { const cp = s.params.checkpoint; - return cp ? s.loras.some((l) => l.id === cp) : false; + const selected = cp ? s.loras.find((l) => l.id === cp) : undefined; + return selected?.exportType === "lora"; }); } @@ -235,7 +238,7 @@ const LoraCompareContent = memo(function LoraCompareContent({
- Fine-tuned (LoRA) + Fine-tuned
@@ -251,7 +254,7 @@ const LoraCompareContent = memo(function LoraCompareContent({
-
+
@@ -387,7 +390,7 @@ const GeneralCompareContent = memo(function GeneralCompareContent({
-
+
): string { + const normalizedBase = baseName.trim() || "Imported Prompt"; + let nextName = normalizedBase; + let suffix = 2; + while (usedNames.has(nextName)) { + nextName = `${normalizedBase} ${suffix}`; + suffix += 1; + } + usedNames.add(nextName); + return nextName; +} + +function migrateLegacySystemPromptTemplates(presets: Preset[]): Preset[] { + if (!canUseStorage()) return presets; + try { + const raw = localStorage.getItem(LEGACY_CHAT_SYSTEM_PROMPTS_KEY); + if (!raw) return presets; + if (localStorage.getItem(LEGACY_CHAT_SYSTEM_PROMPTS_MIGRATED_KEY) === raw) { + return presets; + } + let parsed: unknown; + try { + parsed = JSON.parse(raw) as unknown; + } catch { + localStorage.removeItem(LEGACY_CHAT_SYSTEM_PROMPTS_KEY); + localStorage.setItem(LEGACY_CHAT_SYSTEM_PROMPTS_MIGRATED_KEY, raw); + return presets; + } + if (!Array.isArray(parsed)) { + localStorage.removeItem(LEGACY_CHAT_SYSTEM_PROMPTS_KEY); + localStorage.setItem(LEGACY_CHAT_SYSTEM_PROMPTS_MIGRATED_KEY, raw); + return presets; + } + const usedNames = new Set([ + ...BUILTIN_PRESETS.map((preset) => preset.name), + ...presets.map((preset) => preset.name), + ]); + const seenImportedConfigKeys = new Set( + [...BUILTIN_PRESETS, ...presets].map((preset) => + JSON.stringify({ + temperature: preset.params.temperature, + topP: preset.params.topP, + topK: preset.params.topK, + minP: preset.params.minP, + repetitionPenalty: preset.params.repetitionPenalty, + presencePenalty: preset.params.presencePenalty, + maxSeqLength: preset.params.maxSeqLength, + maxTokens: preset.params.maxTokens, + systemPrompt: preset.params.systemPrompt, + trustRemoteCode: preset.params.trustRemoteCode ?? false, + }), + ), + ); + const importedPresets = parsed + .filter((item): item is LegacySystemPromptTemplate => { + if (!item || typeof item !== "object") return false; + const maybe = item as Partial; + return ( + typeof maybe.name === "string" && typeof maybe.content === "string" + ); + }) + .map((template) => ({ + template, + importedParams: { + ...defaultInferenceParams, + systemPrompt: template.content, + }, + })) + .filter(({ importedParams }) => { + const configKey = JSON.stringify({ + temperature: importedParams.temperature, + topP: importedParams.topP, + topK: importedParams.topK, + minP: importedParams.minP, + repetitionPenalty: importedParams.repetitionPenalty, + presencePenalty: importedParams.presencePenalty, + maxSeqLength: importedParams.maxSeqLength, + maxTokens: importedParams.maxTokens, + systemPrompt: importedParams.systemPrompt, + trustRemoteCode: importedParams.trustRemoteCode ?? false, + }); + if (seenImportedConfigKeys.has(configKey)) return false; + seenImportedConfigKeys.add(configKey); + return true; + }) + .map(({ template, importedParams }) => ({ + name: getUniquePresetName(`${template.name} Prompt`, usedNames), + params: importedParams, + })); + if (importedPresets.length === 0) { + localStorage.removeItem(LEGACY_CHAT_SYSTEM_PROMPTS_KEY); + localStorage.setItem(LEGACY_CHAT_SYSTEM_PROMPTS_MIGRATED_KEY, raw); + return presets; + } + const mergedPresets = [...presets, ...importedPresets]; + localStorage.setItem(CHAT_PRESETS_KEY, JSON.stringify(mergedPresets)); + try { + localStorage.setItem(LEGACY_CHAT_SYSTEM_PROMPTS_MIGRATED_KEY, raw); + localStorage.removeItem(LEGACY_CHAT_SYSTEM_PROMPTS_KEY); + } catch { + // ignore cleanup failure after successful import write + } + return mergedPresets; + } catch { + return presets; + } +} + function loadSavedCustomPresets(): Preset[] { if (!canUseStorage()) return []; try { const raw = localStorage.getItem(CHAT_PRESETS_KEY); - if (!raw) return []; + if (!raw) { + return migrateLegacySystemPromptTemplates([]); + } const parsed = JSON.parse(raw) as unknown; - if (!Array.isArray(parsed)) return []; - return parsed + if (!Array.isArray(parsed)) { + return migrateLegacySystemPromptTemplates([]); + } + const presets = parsed .filter((item): item is Preset => { if (!item || typeof item !== "object") return false; const maybe = item as Partial; @@ -111,13 +250,10 @@ function loadSavedCustomPresets(): Preset[] { ...preset.params, }, })) - .filter( - (preset) => - preset.name.length > 0 && - !BUILTIN_PRESETS.some((builtin) => builtin.name === preset.name), - ); + .filter((preset) => preset.name.length > 0); + return migrateLegacySystemPromptTemplates(presets); } catch { - return []; + return migrateLegacySystemPromptTemplates([]); } } @@ -130,6 +266,82 @@ function loadSavedActivePreset(): string { } } +type PresetSaveMode = + | "disabled" + | "overwrite-active" + | "overwrite-other" + | "create"; + +interface PresetSaveState { + mode: PresetSaveMode; + canSubmit: boolean; + isSaveReady: boolean; + buttonLabel: string; + title: string; +} + +function isSamePresetConfig(a: InferenceParams, b: InferenceParams): boolean { + return ( + a.temperature === b.temperature && + a.topP === b.topP && + a.topK === b.topK && + a.minP === b.minP && + a.repetitionPenalty === b.repetitionPenalty && + a.presencePenalty === b.presencePenalty && + a.maxSeqLength === b.maxSeqLength && + a.maxTokens === b.maxTokens && + a.systemPrompt === b.systemPrompt && + (a.trustRemoteCode ?? false) === (b.trustRemoteCode ?? false) + ); +} + +function getPresetSaveState({ + rawName, + activePreset, + presets, + activePresetDirty, +}: { + rawName: string; + activePreset: string; + presets: Preset[]; + activePresetDirty: boolean; +}): PresetSaveState { + const trimmedName = rawName.trim(); + if (!trimmedName) { + return { + mode: "disabled", + canSubmit: false, + isSaveReady: false, + buttonLabel: "Save", + title: "Enter a preset name", + }; + } + + const matchingPreset = presets.find((preset) => preset.name === trimmedName); + if (matchingPreset) { + const isActiveMatch = matchingPreset.name === activePreset; + return { + mode: isActiveMatch ? "overwrite-active" : "overwrite-other", + canSubmit: !isActiveMatch || activePresetDirty, + isSaveReady: !isActiveMatch || activePresetDirty, + buttonLabel: isActiveMatch && !activePresetDirty ? "Saved" : "Overwrite", + title: isActiveMatch + ? activePresetDirty + ? "Save current settings to this preset" + : "No unsaved changes" + : `Overwrite preset "${trimmedName}"`, + }; + } + + return { + mode: "create", + canSubmit: true, + isSaveReady: true, + buttonLabel: "Save as New", + title: `Save current settings as "${trimmedName}"`, + }; +} + function ParamSlider({ label, value, @@ -286,6 +498,9 @@ export function ChatSettingsPanel({ (s) => s.loadedSpeculativeType, ); const currentModels = useChatRuntimeStore((s) => s.models); + const modelRequiresTrustRemoteCode = useChatRuntimeStore( + (s) => s.modelRequiresTrustRemoteCode, + ); const currentCheckpoint = params.checkpoint; const currentModelIsVision = currentModels.find((m) => m.id === currentCheckpoint)?.isVision ?? false; @@ -316,13 +531,57 @@ export function ChatSettingsPanel({ const [activePreset, setActivePreset] = useState(() => loadSavedActivePreset(), ); - const [savePresetOpen, setSavePresetOpen] = useState(false); - const [presetNameDraft, setPresetNameDraft] = useState(""); - const presets = useMemo( - () => [...BUILTIN_PRESETS, ...customPresets], - [customPresets], + const [presetNameInput, setPresetNameInput] = useState(() => + loadSavedActivePreset(), ); - const isBuiltinPreset = BUILTIN_PRESETS.some((p) => p.name === activePreset); + const presetControlRowRef = useRef(null); + const [presetMenuWidthPx, setPresetMenuWidthPx] = useState< + number | undefined + >(undefined); + const [systemPromptEditorOpen, setSystemPromptEditorOpen] = useState(false); + const [systemPromptDraft, setSystemPromptDraft] = useState(""); + const presets = useMemo(() => { + const overrides = new Set(customPresets.map((preset) => preset.name)); + return [ + ...BUILTIN_PRESETS.filter((preset) => !overrides.has(preset.name)), + ...customPresets, + ]; + }, [customPresets]); + const activePresetDefinition = useMemo( + () => presets.find((preset) => preset.name === activePreset) ?? null, + [activePreset, presets], + ); + const activeCustomPreset = useMemo( + () => customPresets.find((preset) => preset.name === activePreset) ?? null, + [activePreset, customPresets], + ); + const activeBuiltinPreset = useMemo( + () => + BUILTIN_PRESETS.find((preset) => preset.name === activePreset) ?? null, + [activePreset], + ); + const activePresetDirty = useMemo( + () => + activePresetDefinition == null + ? false + : !isSamePresetConfig(activePresetDefinition.params, params), + [activePresetDefinition, params], + ); + const presetSaveState = useMemo( + () => + getPresetSaveState({ + rawName: presetNameInput, + activePreset, + presets, + activePresetDirty, + }), + [activePreset, activePresetDirty, presetNameInput, presets], + ); + const systemPromptEditorDirty = systemPromptDraft !== params.systemPrompt; + const trustRemoteCodeMissing = + Boolean(currentCheckpoint) && + modelRequiresTrustRemoteCode && + !(params.trustRemoteCode ?? false); function set(key: K) { return (v: InferenceParams[K]) => onParamsChange({ ...params, [key]: v }); @@ -331,11 +590,19 @@ export function ChatSettingsPanel({ function applyPreset(name: string) { const p = presets.find((pr) => pr.name === name); if (p) { + if ( + modelRequiresTrustRemoteCode && + !(p.params.trustRemoteCode ?? false) + ) { + toast.warning("This configuration turns custom code off", { + description: + "The current model needs custom code enabled to load. Keep it on for this model.", + }); + return; + } onParamsChange({ ...p.params, - systemPrompt: params.systemPrompt, checkpoint: params.checkpoint, - trustRemoteCode: params.trustRemoteCode, }); setActivePreset(name); if (canUseStorage()) { @@ -348,32 +615,23 @@ export function ChatSettingsPanel({ } } - function openSavePresetDialog() { - setPresetNameDraft(activePreset === "Default" ? "" : activePreset); - setSavePresetOpen(true); - } - function savePresetWithName(rawName: string) { const trimmed = rawName.trim(); if (!trimmed) { - return; - } - if (BUILTIN_PRESETS.some((preset) => preset.name === trimmed)) { + toast.error("Enter a preset name"); return; } setCustomPresets((prev) => { - const next = [ - ...prev.filter((preset) => preset.name !== trimmed), - { name: trimmed, params: { ...params } }, - ]; + const next = prev.filter((p) => p.name !== trimmed); + const merged = [...next, { name: trimmed, params: { ...params } }]; if (canUseStorage()) { try { - localStorage.setItem(CHAT_PRESETS_KEY, JSON.stringify(next)); + localStorage.setItem(CHAT_PRESETS_KEY, JSON.stringify(merged)); } catch { // ignore } } - return next; + return merged; }); if (canUseStorage()) { try { @@ -383,11 +641,31 @@ export function ChatSettingsPanel({ } } setActivePreset(trimmed); - setSavePresetOpen(false); + setPresetNameInput(trimmed); } function deletePreset(name: string) { - if (BUILTIN_PRESETS.some((p) => p.name === name)) { + const hasCustomPreset = customPresets.some( + (preset) => preset.name === name, + ); + if (!hasCustomPreset) { + return; + } + const builtinPreset = BUILTIN_PRESETS.find((preset) => preset.name === name); + const fallbackPreset = + builtinPreset ?? + BUILTIN_PRESETS.find((preset) => preset.name === "Default") ?? + null; + if ( + activePreset === name && + fallbackPreset && + modelRequiresTrustRemoteCode && + !(fallbackPreset.params.trustRemoteCode ?? false) + ) { + toast.warning("Reset would turn custom code off", { + description: + "The current model needs custom code enabled to load. Keep it on for this model.", + }); return; } setCustomPresets((prev) => { @@ -402,17 +680,33 @@ export function ChatSettingsPanel({ return next; }); if (activePreset === name) { - setActivePreset("Default"); - if (canUseStorage()) { - try { - localStorage.setItem(CHAT_ACTIVE_PRESET_KEY, "Default"); - } catch { - // ignore + if (fallbackPreset) { + onParamsChange({ + ...fallbackPreset.params, + checkpoint: params.checkpoint, + }); + setActivePreset(fallbackPreset.name); + if (canUseStorage()) { + try { + localStorage.setItem(CHAT_ACTIVE_PRESET_KEY, fallbackPreset.name); + } catch { + // ignore + } } } } } + function openSystemPromptEditor() { + setSystemPromptDraft(params.systemPrompt); + setSystemPromptEditorOpen(true); + } + + function saveSystemPromptEditor() { + set("systemPrompt")(systemPromptDraft); + setSystemPromptEditorOpen(false); + } + useEffect(() => { if (presets.some((preset) => preset.name === activePreset)) return; setActivePreset("Default"); @@ -425,6 +719,28 @@ export function ChatSettingsPanel({ } }, [activePreset, presets]); + useEffect(() => { + setPresetNameInput(activePreset); + }, [activePreset]); + + useEffect(() => { + if (!open) { + setSystemPromptEditorOpen(false); + } + }, [open]); + + useLayoutEffect(() => { + const el = presetControlRowRef.current; + if (!el || !open) return; + const measure = () => { + setPresetMenuWidthPx(el.getBoundingClientRect().width); + }; + measure(); + const ro = new ResizeObserver(measure); + ro.observe(el); + return () => ro.disconnect(); + }, [open]); + const settingsContent = ( <>
@@ -440,52 +756,138 @@ export function ChatSettingsPanel({
{/* mt-4 matches the Playground sidebar gap (SidebarHeader py-3 + SidebarGroup pt-1) */}
-
- - - +
+
+ + + setPresetNameInput(e.target.value)} + onKeyDown={(e) => { + if (e.key === "Enter" && presetSaveState.canSubmit) { + e.preventDefault(); + savePresetWithName(presetNameInput); + } + }} + placeholder="Preset name" + maxLength={80} + autoComplete="off" + className={cn( + "!h-8 min-h-0 min-w-0 self-stretch !pl-2.5 !pr-2 pt-1 pb-1 text-sm leading-10 md:text-sm", + presetSaveState.isSaveReady && + "text-foreground placeholder:text-primary/45", + )} + aria-label="Inference preset name" + /> + + + + + + + + + + {presets.map((p) => ( + applyPreset(p.name)} + > + {p.name} + + ))} + + +
+
+ + +
- +
+ + +