diff --git a/.gitattributes b/.gitattributes index 5f04b5e9d1..0025f2a697 100644 --- a/.gitattributes +++ b/.gitattributes @@ -6,7 +6,7 @@ # them when run in WSL/Linux (e.g. `set -e` -> "set: Illegal option -"). *.sh text eol=lf -# Normalize Studio frontend sources to LF. Scoped to the frontend tree (rather +# Normalize Unsloth frontend sources to LF. Scoped to the frontend tree (rather # than repo-wide *.ts/*.tsx/... rules) so the policy can't force LF on files # elsewhere. text=auto lets Git detect and leave binary assets (logos, fonts) # untouched while text files (.ts/.tsx/.json/.html/.svg/...) are stored as LF. diff --git a/.github/scripts/agent-guides-drive.sh b/.github/scripts/agent-guides-drive.sh index f4189a159e..b63ac94b93 100755 --- a/.github/scripts/agent-guides-drive.sh +++ b/.github/scripts/agent-guides-drive.sh @@ -36,6 +36,23 @@ AGENT="${2:?usage: agent-guides-drive.sh }" # Determinism (seed/temp) is applied at the server level by # serve-unsloth-run.sh --extra; agents inherit it through the API. TIMEOUT="${AGENT_INVOKE_TIMEOUT:-180}" +# opencode is the slow outlier. Unlike the print-mode agents (claude -p, codex +# exec) it runs a full turn AND a separate small_model call to name the session, +# so one connection reply takes ~8 min on a CPU-served 4B -- right at the shared +# 600s cap, so the cell flaked when a run drifted past a ~480s success. Give it +# headroom (still well under the 40-min job budget); the fast agents keep the +# tight cap that still catches a real headless-TTY hang. +case "$AGENT" in + opencode) + # Double it, but only for a bare-integer seconds value. A GNU timeout(1) + # duration suffix (s/m/h/d, including floats like 0.5s) is left unchanged so + # the arithmetic never sees a non-number; timeout(1) parses it directly. + case "$TIMEOUT" in + *[!0-9]*) ;; + *) TIMEOUT=$(( TIMEOUT * 2 )) ;; + esac + ;; +esac # Claude refuses --dangerously-skip-permissions outside a sandbox; the CI runner # IS the sandbox, so declare it (mirrors unslothai/scripts launcher.sh). Harmless @@ -166,8 +183,8 @@ parse_connect() { echo "[$AGENT] connect --no-launch printed:"; cat_redacted "$raw" CONNECT_ENV="$(grep -E '^(export |unset )' "$raw" || true)" # The launch command is the last non-export, non-status line. start.py - # prints "Studio · model " and "Updated ..." status lines first. - CONNECT_CMD="$(grep -vE '^(export |unset |Studio |Updated |Disabled |Warning|Loading)' "$raw" \ + # prints "Unsloth · model " and "Updated ..." status lines first. + CONNECT_CMD="$(grep -vE '^(export |unset |Unsloth |Updated |Disabled |Warning|Loading)' "$raw" \ | grep -E '[^[:space:]]' | tail -1)" [ -n "$CONNECT_CMD" ] || guide_fail "could not parse a launch command from connect --no-launch output" redact "$raw" diff --git a/.github/scripts/assert-llama-loads.sh b/.github/scripts/assert-llama-loads.sh index c2ffe27469..62ef80d364 100755 --- a/.github/scripts/assert-llama-loads.sh +++ b/.github/scripts/assert-llama-loads.sh @@ -2,7 +2,7 @@ # SPDX-License-Identifier: AGPL-3.0-only # Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. # -# Assert Studio installed a llama.cpp that loads and runs on THIS macOS. Tests +# Assert Unsloth installed a llama.cpp that loads and runs on THIS macOS. Tests # the contract that matters (binaries load and their minimum-OS is <= this host) # instead of the old "did install.sh fall back to a source build?" grep, since a # source build with a correct deployment target is a valid outcome. diff --git a/.github/scripts/assert-prompt-cache.sh b/.github/scripts/assert-prompt-cache.sh index f5b6b075eb..8c28569f77 100755 --- a/.github/scripts/assert-prompt-cache.sh +++ b/.github/scripts/assert-prompt-cache.sh @@ -31,7 +31,7 @@ # (llama_cpp.py:337-340). So default: ~/.unsloth/studio/logs/llama-server/. # #

is the INTERNAL llama-server port (self._find_free_port(), -# llama_cpp.py:3489 / :4641) -- a RANDOM port, NOT the Studio port. So we must +# llama_cpp.py:3489 / :4641) -- a RANDOM port, NOT the Unsloth port. So we must # NOT filter the log glob by STUDIO_PORT (the brief's `port-` # glob would never match). We pick the newest llama-*.log instead. # diff --git a/.github/scripts/hf-download-with-retry.sh b/.github/scripts/hf-download-with-retry.sh index 013a459f46..6dec93356a 100755 --- a/.github/scripts/hf-download-with-retry.sh +++ b/.github/scripts/hf-download-with-retry.sh @@ -3,7 +3,7 @@ # Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0 # # Download a single file from a Hugging Face repo with a stall-retry -# watchdog. Used by the Studio CI workflows so a hung hf-xet transfer +# watchdog. Used by the Unsloth CI workflows so a hung hf-xet transfer # kills + retries instead of silently consuming the job's timeout. # # Usage: hf-download-with-retry.sh REPO FILE LOCAL_DIR @@ -35,7 +35,7 @@ REPO="${1:?usage: hf-download-with-retry.sh REPO FILE [LOCAL_DIR]}" FILE="${2:?usage: hf-download-with-retry.sh REPO FILE [LOCAL_DIR]}" # LOCAL_DIR is optional. If empty, hf falls back to HF_HUB_CACHE # (~/.cache/huggingface/hub) which is the desired path for callers -# that populate HF_HOME for a downstream Studio model load. +# that populate HF_HOME for a downstream Unsloth model load. LOCAL_DIR="${3:-}" # Stall threshold per attempt, in seconds. Override with diff --git a/.github/scripts/run-studio-permission-browser.sh b/.github/scripts/run-studio-permission-browser.sh new file mode 100755 index 0000000000..e5a9a4c135 --- /dev/null +++ b/.github/scripts/run-studio-permission-browser.sh @@ -0,0 +1,70 @@ +#!/usr/bin/env bash +# SPDX-License-Identifier: AGPL-3.0-only +# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. + +set -euo pipefail + +port="${1:?usage: $0 PORT BROWSER [CHANNEL]}" +browser="${2:?usage: $0 PORT BROWSER [CHANNEL]}" +channel="${3:-}" +slug="$browser${channel:+-$channel}" +artifact_dir="logs/playwright-permissions-$slug" +server_log="logs/studio-permissions-$slug.log" +studio_home="${UNSLOTH_STUDIO_HOME:-$HOME/.unsloth/studio}" +set -- +if [ -n "${STUDIO_PERMISSION_FRONTEND:-}" ]; then + set -- -f "$STUDIO_PERMISSION_FRONTEND" +fi + +mkdir -p "$artifact_dir" +# Wipe (not reset-password): the boot below must re-seed a fresh .bootstrap_password. +rm -rf "$studio_home/auth" +UNSLOTH_API_ONLY=1 unsloth studio -H 127.0.0.1 -p "$port" "$@" \ + >"$server_log" 2>&1 & +studio_pid=$! + +cleanup() { + kill "$studio_pid" 2>/dev/null || true + wait "$studio_pid" 2>/dev/null || true +} +trap cleanup EXIT + +healthy=0 +for _ in $(seq 1 180); do + if curl -fs "http://127.0.0.1:$port/api/health" >/dev/null; then + healthy=1 + break + fi + if ! kill -0 "$studio_pid" 2>/dev/null; then + tail -100 "$server_log" || true + exit 1 + fi + sleep 1 +done +if [ "$healthy" -ne 1 ]; then + tail -100 "$server_log" || true + exit 1 +fi + +old_password=$(cat "$studio_home/auth/.bootstrap_password") +new_password="CIPerm-$(python -c 'import secrets; print(secrets.token_urlsafe(16))')" +if [ "${GITHUB_ACTIONS:-}" = "true" ]; then + echo "::add-mask::$old_password" + echo "::add-mask::$new_password" +fi + +export BASE_URL="http://127.0.0.1:$port" +export STUDIO_OLD_PW="$old_password" +export STUDIO_NEW_PW="$new_password" +export STUDIO_UI_STRICT=1 +export STUDIO_UI_PERMISSION_ONLY=1 +export STUDIO_UI_WALL_TIMEOUT_S=240 +export STUDIO_PLAYWRIGHT_BROWSER="$browser" +export PW_ART_DIR="$artifact_dir" +if [ -n "$channel" ]; then + export STUDIO_PLAYWRIGHT_CHANNEL="$channel" +else + unset STUDIO_PLAYWRIGHT_CHANNEL || true +fi + +python tests/studio/playwright_chat_ui.py diff --git a/.github/workflows/consolidated-tests-ci.yml b/.github/workflows/consolidated-tests-ci.yml index fa84471d36..afad1b6c46 100644 --- a/.github/workflows/consolidated-tests-ci.yml +++ b/.github/workflows/consolidated-tests-ci.yml @@ -7,7 +7,7 @@ # # Why a separate workflow: # - studio-backend-ci.yml's "Repo tests (CPU)" job already auto-discovers -# tests/ minus tests/qlora, tests/saving, tests/utils, tests/sh. The 16 +# tests/ minus tests/qlora, tests/saving, tests/utils, tests/sh. The 17 # Bucket-A tests below live inside those --ignore dirs (CPU-runnable but # historically excluded with their GPU siblings); pulling them out into # a sibling job keeps the existing 760-passed baseline stable while we @@ -268,11 +268,13 @@ jobs: tests/saving/test_save_shell_injection.py \ tests/saving/test_patch_saving_none_tokenizer.py \ tests/saving/test_fix_sentencepiece_gguf_robustness.py \ + tests/saving/test_fix_sentencepiece_tokenizer_guard.py \ tests/saving/test_compressed_export_schemes.py \ tests/saving/test_export_api_surface.py \ tests/saving/test_export_dispatch.py \ tests/saving/test_imatrix_export.py \ tests/saving/test_gguf_single_pass_export.py \ + tests/saving/test_offline_gguf_vlm_tokenizer_7481.py \ tests/utils/test_attention_masks.py \ tests/utils/test_trunc_normal_patch.py \ tests/python/test_fast_language_model_text_only.py @@ -358,22 +360,23 @@ jobs: tests/saving/test_save_shell_injection.py \ tests/saving/test_patch_saving_none_tokenizer.py \ tests/saving/test_fix_sentencepiece_gguf_robustness.py \ + tests/saving/test_fix_sentencepiece_tokenizer_guard.py \ tests/saving/test_compressed_export_schemes.py \ tests/saving/test_export_api_surface.py \ tests/saving/test_export_dispatch.py \ tests/saving/test_imatrix_export.py \ tests/saving/test_gguf_single_pass_export.py \ + tests/saving/test_offline_gguf_vlm_tokenizer_7481.py \ tests/utils/test_attention_masks.py \ tests/utils/test_trunc_normal_patch.py \ tests/python/test_fast_language_model_text_only.py \ tests/test_bad_mappings_redirect.py \ tests/test_prefetch_snapshot_scope.py \ tests/test_gemma_2b_mapper_key.py \ - --deselect 'tests/utils/test_attention_masks.py::test_run_attention_flash_varlen_receives_window_and_softcap' - # The deselected test monkeypatches flash_attn_varlen_func, which is - # only bound on the module when `flash_attn` is importable. flash_attn - # requires CUDA + dev toolchain, which the CPU-only ubuntu-latest - # runner does not have. The other Bucket-A tests pass cleanly. + tests/test_raw_text_json_loading.py + # test_run_attention_flash_varlen_receives_window_and_softcap was deselected + # until attention_dispatch.py predefined flash_attn_varlen_func as None; it + # monkeypatches that name, so it no longer needs flash_attn on this runner. - name: unsloth_zoo @ ${{ env.UNSLOTH_ZOO_REF }} — full pytest (CPU) # 106 of 111 test_* in unsloth_zoo are CPU-only. The two CUDA-skip @@ -2127,7 +2130,7 @@ jobs: pip show unsloth_zoo echo "::endgroup::" echo "Consolidated job done. Coverage:" - echo " - 16 unsloth Bucket-A tests under tests/saving/ + tests/utils/" + echo " - 17 unsloth Bucket-A tests under tests/saving/ + tests/utils/" echo " - unsloth_zoo @ ${UNSLOTH_ZOO_REF} pytest tests/ (5 GPU cases deselected)" echo " - unsloth_zoo.compiler.test_apply_fused_lm_head" diff --git a/.github/workflows/cross-platform-parity-ci.yml b/.github/workflows/cross-platform-parity-ci.yml index bb7dcbf8e4..45ce231743 100644 --- a/.github/workflows/cross-platform-parity-ci.yml +++ b/.github/workflows/cross-platform-parity-ci.yml @@ -1,18 +1,16 @@ # SPDX-License-Identifier: AGPL-3.0-only # Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. -# Runs installer parity and autostart opt-out tests on Windows and macOS. +# Runs installer parity and autostart opt-out tests across all three platforms. # -# Why: that test is the guard that install.sh and install.ps1 stay in -# sync, but today it only runs on ubuntu-latest (auto-discovered by -# studio-backend-ci.yml's "Repo tests (CPU)" job). The test reads both -# installer scripts, and on Windows Path.read_text() defaults to the -# cp1252 locale encoding, so a non-cp1252 byte in install.sh (it already -# contains a U+274C) raises UnicodeDecodeError there even though Linux and -# macOS default to UTF-8. The reads were pinned to encoding="utf-8" in -# #6166; this job keeps that from silently regressing by exercising the -# test on the platforms it claims parity for. Pure pytest, no GPU, -# sub-second, so the matrix is cheap. +# Why: the parity test guards that install.sh and install.ps1 stay in sync. +# It originally ran only on ubuntu-latest through studio-backend-ci.yml. +# On Windows, Path.read_text() defaults to the cp1252 locale encoding, so a +# non-cp1252 byte in install.sh raises UnicodeDecodeError even though Linux +# and macOS default to UTF-8. The reads were pinned to encoding="utf-8" in +# #6166; this matrix keeps that from silently regressing. Pure pytest, no GPU, +# sub-second, so the matrix is cheap. Linux also runs the POSIX rollback test +# under dash, matching the supported curl-to-sh installer path. name: Cross-platform parity @@ -23,6 +21,8 @@ on: - 'install.ps1' - 'tests/test_installer_skip_autostart.py' - 'tests/python/test_cross_platform_parity.py' + - 'tests/sh/test_install_rollback_lifecycle.sh' + - 'tests/studio/test_install_rollback_lifecycle.ps1' - '.github/workflows/cross-platform-parity-ci.yml' push: branches: [main] @@ -31,6 +31,8 @@ on: - 'install.ps1' - 'tests/test_installer_skip_autostart.py' - 'tests/python/test_cross_platform_parity.py' + - 'tests/sh/test_install_rollback_lifecycle.sh' + - 'tests/studio/test_install_rollback_lifecycle.ps1' - '.github/workflows/cross-platform-parity-ci.yml' workflow_dispatch: @@ -47,7 +49,7 @@ jobs: strategy: fail-fast: false matrix: - os: [windows-latest, macos-latest] + os: [ubuntu-latest, windows-latest, macos-latest] runs-on: ${{ matrix.os }} timeout-minutes: 10 steps: @@ -67,3 +69,10 @@ jobs: tests/python/test_cross_platform_parity.py tests/test_installer_skip_autostart.py -q + - name: PowerShell rollback lifecycle tests + if: runner.os == 'Windows' + shell: pwsh + run: pwsh -NoProfile -File tests/studio/test_install_rollback_lifecycle.ps1 + - name: POSIX rollback lifecycle tests + if: runner.os == 'Linux' + run: sh tests/sh/test_install_rollback_lifecycle.sh diff --git a/.github/workflows/lint-ci.yml b/.github/workflows/lint-ci.yml index bd859a6e9e..e1f0afd299 100644 --- a/.github/workflows/lint-ci.yml +++ b/.github/workflows/lint-ci.yml @@ -13,10 +13,10 @@ # committed YAML / JSON config. # # TypeScript and Rust are NOT duplicated here on purpose: -# - Studio Frontend CI runs `npm run typecheck` (= `tsc --noEmit`) +# - Unsloth Frontend CI runs `npm run typecheck` (= `tsc --noEmit`) # and `npm run build` (vite/swc) on every studio/frontend/** # change, which is a full TS AST + type check. -# - Studio Tauri CI runs `tauri build --debug --no-bundle` on +# - Unsloth Tauri CI runs `tauri build --debug --no-bundle` on # every studio/src-tauri/** or studio/frontend/** change, which # compiles the Rust crate (= cargo check + cargo build). # Each is a stricter check than a parse-only step would be, so a diff --git a/.github/workflows/local-agent-guides-ci.yml b/.github/workflows/local-agent-guides-ci.yml index 25796bd5cf..0dc0cc66d7 100644 --- a/.github/workflows/local-agent-guides-ci.yml +++ b/.github/workflows/local-agent-guides-ci.yml @@ -154,7 +154,7 @@ jobs: path: gguf-cache key: ${{ runner.os }}-gguf-${{ env.GGUF_REPO }}-${{ env.GGUF_FILE }}-v1 - - name: Install Studio (--local, --no-torch) + - name: Install Unsloth (--local, --no-torch) env: GH_TOKEN: ${{ secrets.GITHUB_TOKEN }} # Gated off PR (see note above); public GGUF still downloads. @@ -167,7 +167,9 @@ jobs: # ── boot the server under test (factored helper) ────────────────── - name: Serve unsloth run --disable-tools (gemma-4-E4B) run: | - unsloth studio reset-password + # Wipe, not reset-password: since #7573 the reset rotates in place and + # prints the new passphrase, which would land unmasked in the job log. + rm -rf ~/.unsloth/studio/auth bash .github/scripts/serve-unsloth-run.sh \ --gguf-file "$GITHUB_WORKSPACE/gguf-cache/${GGUF_FILE}" \ --port "$STUDIO_PORT" --log-dir logs \ @@ -256,7 +258,7 @@ jobs: done fi - - name: Stop Studio + - name: Stop Unsloth if: always() run: | # Guard the PID: an unset/zero UNSLOTH_SERVER_PID would make @@ -359,7 +361,7 @@ jobs: path: gguf-cache key: ${{ runner.os }}-gguf-${{ env.GGUF_REPO }}-${{ env.GGUF_FILE }}-v1 - - name: Install Studio (--local, --no-torch) + - name: Install Unsloth (--local, --no-torch) env: GH_TOKEN: ${{ secrets.GITHUB_TOKEN }} # Gated off PR (see note above); public GGUF still downloads. @@ -371,7 +373,7 @@ jobs: - name: Serve unsloth run --disable-tools (gemma-4-E4B) run: | - unsloth studio reset-password + rm -rf ~/.unsloth/studio/auth bash .github/scripts/serve-unsloth-run.sh \ --gguf-file "$GITHUB_WORKSPACE/gguf-cache/${GGUF_FILE}" \ --port "$STUDIO_PORT" --log-dir logs \ @@ -448,7 +450,7 @@ jobs: done fi - - name: Stop Studio + - name: Stop Unsloth if: always() run: | # Guard the PID: an unset/zero UNSLOTH_SERVER_PID would make @@ -543,7 +545,7 @@ jobs: path: gguf-cache key: ${{ runner.os }}-gguf-${{ env.GGUF_REPO }}-${{ env.GGUF_FILE }}-v1 - - name: Install Studio (--local, --no-torch) + - name: Install Unsloth (--local, --no-torch) env: GH_TOKEN: ${{ secrets.GITHUB_TOKEN }} HF_TOKEN: ${{ secrets.HF_TOKEN }} @@ -554,7 +556,7 @@ jobs: - name: Serve unsloth run --disable-tools (gemma-4-E4B) run: | - unsloth studio reset-password + rm -rf ~/.unsloth/studio/auth bash .github/scripts/serve-unsloth-run.sh \ --gguf-file "$GITHUB_WORKSPACE/gguf-cache/${GGUF_FILE}" \ --port "$STUDIO_PORT" --log-dir logs \ @@ -620,7 +622,7 @@ jobs: done fi - - name: Stop Studio + - name: Stop Unsloth if: always() run: | if [ -n "${UNSLOTH_SERVER_PID:-}" ] && [ "${UNSLOTH_SERVER_PID}" != "0" ]; then @@ -706,7 +708,7 @@ jobs: path: hf-cache key: ${{ runner.os }}-hf-${{ env.GGUF_REPO }}-${{ env.GGUF_VARIANT }}-v2 - - name: Install Studio (--local, --no-torch) + - name: Install Unsloth (--local, --no-torch) env: GH_TOKEN: ${{ secrets.GITHUB_TOKEN }} # Gated off PR (see note above); public GGUF still downloads. @@ -718,7 +720,7 @@ jobs: - name: Serve unsloth run --disable-tools (gemma-3-270m) run: | - unsloth studio reset-password + rm -rf ~/.unsloth/studio/auth bash .github/scripts/serve-unsloth-run.sh \ --model "$GGUF_REPO" --gguf-variant "$GGUF_VARIANT" \ --port "$STUDIO_PORT" --log-dir logs \ @@ -764,7 +766,7 @@ jobs: done fi - - name: Stop Studio + - name: Stop Unsloth if: always() run: | # Guard the PID: an unset/zero UNSLOTH_SERVER_PID would make diff --git a/.github/workflows/mlx-ci.yml b/.github/workflows/mlx-ci.yml index a2f716a93c..aadf0b54e6 100644 --- a/.github/workflows/mlx-ci.yml +++ b/.github/workflows/mlx-ci.yml @@ -130,7 +130,7 @@ jobs: # MLX support landed after the most recent unsloth-zoo PyPI # release; the wheel still raises NotImplementedError on # Apple Silicon when device_type.get_device_type() runs - # unguarded. Studio's own install.sh overlays unsloth-zoo + # unguarded. Unsloth's own install.sh overlays unsloth-zoo # from git main for the same reason. Pulling deps lets pip # resolve the platform-conditional MLX-only wheels (mlx, # mlx-lm, mlx-vlm gated on darwin+arm64 in unsloth-zoo's @@ -317,13 +317,13 @@ jobs: echo done - # Validates the macOS prebuilt path Studio's setup.sh uses (#5963): install the + # Validates the macOS prebuilt path Unsloth's setup.sh uses (#5963): install the # unslothai/llama.cpp fork's latest release, download a small public GGUF, and # check llama-server /completion end to end. Split and placed last so the # untrusted binary runs only in the final smoke step, after every HF_TOKEN step, # leaving no token-bearing step or shared workspace for a tampered prebuilt to # corrupt. GH_TOKEN: releases API; HF_TOKEN (withheld on PR): probe + GGUF fetch. - - name: Studio prebuilt llama.cpp install + GGUF download (Mac M1) + - name: Unsloth prebuilt llama.cpp install + GGUF download (Mac M1) env: GH_TOKEN: ${{ secrets.GITHUB_TOKEN }} GITHUB_TOKEN: ${{ secrets.GITHUB_TOKEN }} @@ -344,12 +344,12 @@ jobs: # Final step: runs the downloaded binaries with no secrets present, and clears # the GitHub Actions command files so a tampered prebuilt cannot influence the job. - - name: Studio prebuilt llama.cpp GGUF inference smoke (Mac M1) + - name: Unsloth prebuilt llama.cpp GGUF inference smoke (Mac M1) run: | set -euo pipefail unset GITHUB_ENV GITHUB_PATH GITHUB_OUTPUT GITHUB_STEP_SUMMARY INSTALL_DIR="$HOME/.unsloth-studio-prebuilt-test/llama.cpp" - # Studio bundles only llama-server + llama-quantize (not llama-cli); + # Unsloth bundles only llama-server + llama-quantize (not llama-cli); # inference goes through llama-server's HTTP /completion endpoint. LLAMA_SERVER="$INSTALL_DIR/build/bin/llama-server" LLAMA_QUANT="$INSTALL_DIR/build/bin/llama-quantize" @@ -400,4 +400,4 @@ jobs: tail -40 /tmp/llama-server.log exit 1 fi - echo "OK: Studio prebuilt llama.cpp on Mac M1 + GGUF /completion works" + echo "OK: Unsloth prebuilt llama.cpp on Mac M1 + GGUF /completion works" diff --git a/.github/workflows/release-desktop.yml b/.github/workflows/release-desktop.yml index 188e078f90..0a8d71610d 100644 --- a/.github/workflows/release-desktop.yml +++ b/.github/workflows/release-desktop.yml @@ -4,7 +4,7 @@ on: workflow_dispatch: inputs: studio_version: - description: 'Studio version tag to release (for example, v0.1.39-beta)' + description: 'Unsloth version tag to release (for example, v0.1.39-beta)' type: string required: true pypi_version: @@ -19,6 +19,19 @@ on: permissions: contents: read +env: + DESKTOP_RELEASE_NOTES: | + Desktop app for Unsloth Studio. + + **macOS**: Download the Apple Silicon `.dmg`. + **Windows**: Download the `-setup.exe` installer. + **Linux**: Download `.deb` for Ubuntu/Debian. `.AppImage` is experimental. + + > Linux in-app updates are AppImage-oriented. Package installs should update by downloading a new package. + > Linux AppImage can show a blank window on some Tauri/WebKitGTK + Wayland/Mesa stacks; use `.deb` when available. + > Linux AppImage on Ubuntu 24.04+ may require: `sudo apt install libfuse2t64` + > First-run system dependency elevation is supported on Ubuntu/Debian. Other Linux distributions should install system packages manually. + concurrency: group: release-desktop-${{ github.repository }} cancel-in-progress: false @@ -56,7 +69,7 @@ jobs: if not studio_version: sys.exit('studio_version is required, for example v0.1.39-beta') if re.fullmatch(r'v?20\d{2}\.\d+\.\d+(?:[-+][0-9A-Za-z.-]+)?', studio_version): - sys.exit(f'studio_version must be a Studio SemVer tag, not a date-style backend version: {studio_version}') + sys.exit(f'studio_version must be an Unsloth SemVer tag, not a date-style backend version: {studio_version}') semver_tag = re.compile( r'^v(0|[1-9]\d*)\.(0|[1-9]\d*)\.(0|[1-9]\d*)' @@ -133,7 +146,7 @@ jobs: print(f'pypi_version={pypi_version}', file=output) PY - - name: Verify PyPI package and Studio stamp + - name: Verify PyPI package and Unsloth stamp shell: bash env: STUDIO_VERSION: ${{ steps.prepare.outputs.studio_version }} @@ -198,7 +211,7 @@ jobs: fi python3 scripts/stamp_studio_release.py --verify-dist "$RUNNER_TEMP/pypi-unsloth-dist" --expected "$STUDIO_VERSION" else - echo "scripts/stamp_studio_release.py not found; release-desktop requires #5308 to verify the PyPI Studio stamp." >&2 + echo "scripts/stamp_studio_release.py not found; release-desktop requires #5308 to verify the PyPI Unsloth stamp." >&2 exit 1 fi @@ -295,14 +308,6 @@ jobs: PY build: - # TODO: split into a "build (no secrets)" + "publish (secrets)" job pair - # with actions/upload-artifact handoff so the matrix build cannot - # publish a Release on its own. The current matrix runs across - # Linux/macOS/Windows in a single job, so the split needs artefact - # collection across the OS matrix and is out of scope for this - # hardening pass. - permissions: - contents: write # tauri-apps/tauri-action creates / uploads a GitHub Release strategy: fail-fast: false max-parallel: 1 @@ -311,15 +316,21 @@ jobs: - platform: macos-latest args: '--target aarch64-apple-darwin' label: macOS (Apple Silicon) + artifact: macos-aarch64 + release_arch: aarch64 # - platform: macos-latest # args: '--target x86_64-apple-darwin' # label: macOS (Intel) - platform: ubuntu-22.04 args: '' label: Linux (x64) + artifact: linux-x64 + release_arch: x64 - platform: windows-latest args: '' label: Windows (x64) + artifact: windows-x64 + release_arch: x64 name: Build ${{ matrix.label }} needs: prepare-version @@ -465,41 +476,18 @@ jobs: if (chmodIdx !== -1 && sha256Idx > chmodIdx) { throw new Error('Desktop Linux release must verify the linuxdeploy digest before chmod +x'); } - const releaseBodies = []; - for (let i = 0; i < lines.length; i += 1) { - const match = lines[i].match(/^(\s*)releaseBody:\s*\|\s*$/); - if (!match) continue; - const baseIndent = match[1].length; - const bodyLines = []; - i += 1; - for (; i < lines.length; i += 1) { - const line = lines[i]; - if (line.trim() === '') { - bodyLines.push(''); - continue; - } - const indent = line.match(/^\s*/)[0].length; - if (indent <= baseIndent) { - i -= 1; - break; - } - bodyLines.push(line.slice(baseIndent + 2)); - } - releaseBodies.push(bodyLines.join('\n')); + const releaseBody = process.env.DESKTOP_RELEASE_NOTES; + if (!releaseBody) { + throw new Error('DESKTOP_RELEASE_NOTES must not be empty'); } - if (releaseBodies.length === 0) { - throw new Error('Expected at least one desktop release body'); + if (/\brpm\b|\.rpm/i.test(releaseBody)) { + throw new Error('Desktop release body must not advertise RPM packages'); } - for (const body of releaseBodies) { - if (/\brpm\b|\.rpm/i.test(body)) { - throw new Error('Desktop release body must not advertise RPM packages'); - } - if (/AppImage.*universal|universal.*AppImage/i.test(body)) { - throw new Error('Desktop release body must not advertise AppImage as universal'); - } - if (!/AppImage.*experimental/i.test(body)) { - throw new Error('Desktop release body must mark AppImage as experimental'); - } + if (/AppImage.*universal|universal.*AppImage/i.test(releaseBody)) { + throw new Error('Desktop release body must not advertise AppImage as universal'); + } + if (!/AppImage.*experimental/i.test(releaseBody)) { + throw new Error('Desktop release body must mark AppImage as experimental'); } JS @@ -644,48 +632,33 @@ jobs: dest="$tools_dir/linuxdeploy-x86_64.AppImage" curl -fsSL "$LINUXDEPLOY_URL" -o "$dest" # Verify the digest BEFORE the binary is ever marked executable. The - # next step builds the AppImage with the Tauri signing key and a - # contents:write GITHUB_TOKEN in scope, so a substituted linuxdeploy - # that ran here could exfiltrate signing material or tamper with - # published release artifacts. Fail closed on any mismatch. + # next step builds the AppImage with the Tauri signing key, so a + # substituted linuxdeploy that ran here could exfiltrate signing + # material or tamper with release artifacts. Fail closed on any + # mismatch. echo "${LINUXDEPLOY_SHA256} ${dest}" | sha256sum -c - chmod +x "$dest" - # ── Linux: build + sign + upload ── + # ── Linux: build + sign ── - name: Build Linux app + id: build_linux if: matrix.platform == 'ubuntu-22.04' uses: tauri-apps/tauri-action@84b9d35b5fc46c1e45415bdb6144030364f7ebc5 env: - GITHUB_TOKEN: ${{ secrets.GITHUB_TOKEN }} TAURI_SIGNING_PRIVATE_KEY: ${{ secrets.TAURI_SIGNING_PRIVATE_KEY }} TAURI_SIGNING_PRIVATE_KEY_PASSWORD: ${{ secrets.TAURI_SIGNING_PRIVATE_KEY_PASSWORD }} XDG_CACHE_HOME: ${{ runner.temp }}/tauri-tools-cache with: projectPath: studio tauriScript: npx --prefix . tauri - tagName: ${{ needs.prepare-version.outputs.desktop_release_tag }} - releaseName: 'Unsloth Studio (Desktop) ${{ needs.prepare-version.outputs.studio_version }}' - releaseBody: | - Desktop app for Unsloth Studio. - - **macOS**: Download the Apple Silicon `.dmg`. - **Windows**: Download the `-setup.exe` installer. - **Linux**: Download `.deb` for Ubuntu/Debian. `.AppImage` is experimental. - - > Linux in-app updates are AppImage-oriented. Package installs should update by downloading a new package. - > Linux AppImage can show a blank window on some Tauri/WebKitGTK + Wayland/Mesa stacks; use `.deb` when available. - > Linux AppImage on Ubuntu 24.04+ may require: `sudo apt install libfuse2t64` - > First-run system dependency elevation is supported on Ubuntu/Debian. Other Linux distributions should install system packages manually. - releaseDraft: ${{ inputs.draft }} - prerelease: ${{ needs.prepare-version.outputs.prerelease }} args: -v ${{ matrix.args }} - # ── macOS: build + sign + notarize + upload ── + # ── macOS: build + sign + notarize ── - name: Build macOS app + id: build_macos if: matrix.platform == 'macos-latest' uses: tauri-apps/tauri-action@84b9d35b5fc46c1e45415bdb6144030364f7ebc5 env: - GITHUB_TOKEN: ${{ secrets.GITHUB_TOKEN }} TAURI_SIGNING_PRIVATE_KEY: ${{ secrets.TAURI_SIGNING_PRIVATE_KEY }} TAURI_SIGNING_PRIVATE_KEY_PASSWORD: ${{ secrets.TAURI_SIGNING_PRIVATE_KEY_PASSWORD }} APPLE_SIGNING_IDENTITY: ${{ secrets.APPLE_SIGNING_IDENTITY }} @@ -695,29 +668,14 @@ jobs: with: projectPath: studio tauriScript: npx --prefix . tauri - tagName: ${{ needs.prepare-version.outputs.desktop_release_tag }} - releaseName: 'Unsloth Studio (Desktop) ${{ needs.prepare-version.outputs.studio_version }}' - releaseBody: | - Desktop app for Unsloth Studio. - - **macOS**: Download the Apple Silicon `.dmg`. - **Windows**: Download the `-setup.exe` installer. - **Linux**: Download `.deb` for Ubuntu/Debian. `.AppImage` is experimental. - - > Linux in-app updates are AppImage-oriented. Package installs should update by downloading a new package. - > Linux AppImage can show a blank window on some Tauri/WebKitGTK + Wayland/Mesa stacks; use `.deb` when available. - > Linux AppImage on Ubuntu 24.04+ may require: `sudo apt install libfuse2t64` - > First-run system dependency elevation is supported on Ubuntu/Debian. Other Linux distributions should install system packages manually. - releaseDraft: ${{ inputs.draft }} - prerelease: ${{ needs.prepare-version.outputs.prerelease }} args: -v ${{ matrix.args }} - # ── Windows: build + sign + upload ── + # ── Windows: build + sign ── - name: Build Windows app + id: build_windows if: matrix.platform == 'windows-latest' uses: tauri-apps/tauri-action@84b9d35b5fc46c1e45415bdb6144030364f7ebc5 env: - GITHUB_TOKEN: ${{ secrets.GITHUB_TOKEN }} TAURI_SIGNING_PRIVATE_KEY: ${{ secrets.TAURI_SIGNING_PRIVATE_KEY }} TAURI_SIGNING_PRIVATE_KEY_PASSWORD: ${{ secrets.TAURI_SIGNING_PRIVATE_KEY_PASSWORD }} AZURE_CLIENT_ID: ${{ secrets.AZURE_CLIENT_ID }} @@ -728,44 +686,252 @@ jobs: with: projectPath: studio tauriScript: npx --prefix . tauri - tagName: ${{ needs.prepare-version.outputs.desktop_release_tag }} - releaseName: 'Unsloth Studio (Desktop) ${{ needs.prepare-version.outputs.studio_version }}' - releaseBody: | - Desktop app for Unsloth Studio. - - **macOS**: Download the Apple Silicon `.dmg`. - **Windows**: Download the `-setup.exe` installer. - **Linux**: Download `.deb` for Ubuntu/Debian. `.AppImage` is experimental. - - > Linux in-app updates are AppImage-oriented. Package installs should update by downloading a new package. - > Linux AppImage can show a blank window on some Tauri/WebKitGTK + Wayland/Mesa stacks; use `.deb` when available. - > Linux AppImage on Ubuntu 24.04+ may require: `sudo apt install libfuse2t64` - > First-run system dependency elevation is supported on Ubuntu/Debian. Other Linux distributions should install system packages manually. - releaseDraft: ${{ inputs.draft }} - prerelease: ${{ needs.prepare-version.outputs.prerelease }} args: -v ${{ matrix.args }} - # Release process note: only non-draft workflow runs advance the public - # desktop-latest updater channel. Draft builds are for private review; if a - # draft is manually published later, this channel intentionally remains - # unchanged until a narrow manual channel-publish flow is added or a public - # desktop release is created by running this workflow with draft=false. - publish-updater-channel: - name: Publish desktop updater channel + - name: Stage release assets + shell: bash + env: + ARTIFACT_PATHS: ${{ steps.build_linux.outputs.artifactPaths || steps.build_macos.outputs.artifactPaths || steps.build_windows.outputs.artifactPaths }} + RELEASE_ARCH: ${{ matrix.release_arch }} + run: | + set -euo pipefail + if command -v python3 >/dev/null 2>&1; then + PYTHON=python3 + else + PYTHON=python + fi + "$PYTHON" <<'PY' + import json + import os + import pathlib + import re + import shutil + import sys + import unicodedata + + raw_paths = os.environ.get('ARTIFACT_PATHS', '') + try: + artifact_paths = json.loads(raw_paths) + except json.JSONDecodeError as error: + sys.exit(f'Invalid tauri-action artifactPaths output: {error}') + if not isinstance(artifact_paths, list) or not artifact_paths: + sys.exit('tauri-action did not return any release artifacts') + + destination = pathlib.Path(os.environ['RUNNER_TEMP'], 'desktop-release-assets') + destination.mkdir(parents=True, exist_ok=True) + staged = [] + for raw_path in artifact_paths: + source = pathlib.Path(raw_path) + if not source.is_file(): + continue + name = source.name + for extension in ('.app.tar.gz.sig', '.app.tar.gz'): + if name.endswith(extension): + name = f'{name[:-len(extension)]}_{os.environ["RELEASE_ARCH"]}{extension}' + break + name = unicodedata.normalize('NFD', name) + name = ''.join(character for character in name if not unicodedata.combining(character)) + name = re.sub(r'[ ()\[\]{}]', '.', name) + while '..' in name: + name = name.replace('..', '.') + target = destination / name + if target.exists(): + sys.exit(f'Duplicate staged release asset name: {name}') + shutil.copy2(source, target) + staged.append(name) + + if not staged: + sys.exit('No release files were staged') + print('Staged release assets:') + print('\n'.join(sorted(staged))) + PY + + - name: Upload signed release assets + uses: actions/upload-artifact@043fb46d1a93c77aae656e7c1c64a875d1fc6a0a # v7.0.1 + with: + name: desktop-release-${{ matrix.artifact }} + path: ${{ runner.temp }}/desktop-release-assets/* + if-no-files-found: error + compression-level: 0 + retention-days: 1 + + # Only this job gets write access; builds hand off signed files via artifacts. + # Draft runs do not advance the public desktop-latest channel. + publish-release: + name: Publish desktop release needs: [prepare-version, build] - if: ${{ !inputs.draft }} runs-on: ubuntu-latest permissions: - contents: write + contents: write # create the versioned Release and replace updater-channel metadata env: GH_REPO: ${{ github.repository }} APP_VERSION: ${{ needs.prepare-version.outputs.app_version }} + PYPI_VERSION: ${{ needs.prepare-version.outputs.pypi_version }} STUDIO_VERSION: ${{ needs.prepare-version.outputs.studio_version }} DESKTOP_RELEASE_TAG: ${{ needs.prepare-version.outputs.desktop_release_tag }} DESKTOP_PRERELEASE: ${{ needs.prepare-version.outputs.prerelease }} steps: + - name: Harden runner (audit) + uses: step-security/harden-runner@a5ad31d6a139d249332a2605b85202e8c0b78450 # v2.19.1 + with: + egress-policy: audit + + - name: Download signed release assets + uses: actions/download-artifact@3e5f45b2cfb9172054b4087a40e8e0b5a5461e7c # v8.0.1 + with: + pattern: desktop-release-* + path: ${{ runner.temp }}/desktop-release-assets + merge-multiple: true + + - name: Validate release asset set + shell: bash + run: | + set -euo pipefail + python3 <<'PY' + import pathlib + import os + import sys + + asset_dir = pathlib.Path(os.environ['RUNNER_TEMP'], 'desktop-release-assets') + files = [path for path in asset_dir.iterdir() if path.is_file()] + required_suffixes = ( + '.dmg', + '.app.tar.gz', + '.app.tar.gz.sig', + '.deb', + '.AppImage', + '.AppImage.sig', + '-setup.exe', + '-setup.exe.sig', + ) + for suffix in required_suffixes: + matches = [path for path in files if path.name.endswith(suffix)] + if len(matches) != 1: + sys.exit(f'Expected exactly one {suffix} release asset, found {len(matches)}') + if any(path.name == 'latest.json' for path in files): + sys.exit('Build artifacts must not supply latest.json') + print('\n'.join(sorted(path.name for path in files))) + PY + + - name: Create or validate versioned release + shell: bash + env: + GH_TOKEN: ${{ github.token }} + RELEASE_DRAFT: ${{ inputs.draft }} + run: | + set -euo pipefail + notes_file="$RUNNER_TEMP/desktop-release-notes.md" + printf '%s\n' "$DESKTOP_RELEASE_NOTES" > "$notes_file" + + release_json="$RUNNER_TEMP/versioned-release.json" + # REST tag lookup omits drafts; `gh release view` also checks pending tags. + if gh release view "$DESKTOP_RELEASE_TAG" \ + --json tagName,isDraft,isPrerelease > "$release_json" 2>/dev/null; then + python3 <<'PY' + import json + import os + import pathlib + import sys + + release = json.loads(pathlib.Path(os.environ['RUNNER_TEMP'], 'versioned-release.json').read_text()) + expected_draft = os.environ['RELEASE_DRAFT'].lower() == 'true' + expected_prerelease = os.environ['DESKTOP_PRERELEASE'].lower() == 'true' + if release.get('tagName') != os.environ['DESKTOP_RELEASE_TAG']: + sys.exit('Existing desktop release tag does not match the requested tag') + if bool(release.get('isDraft')) != expected_draft: + sys.exit('Existing desktop release draft state does not match the workflow input') + if bool(release.get('isPrerelease')) != expected_prerelease: + sys.exit('Existing desktop release prerelease state does not match the requested version') + PY + else + release_flags=( + --title "Unsloth Studio (Desktop) ${STUDIO_VERSION}" + --notes-file "$notes_file" + --target "$GITHUB_SHA" + ) + if [ "$RELEASE_DRAFT" = "true" ]; then + release_flags+=(--draft) + fi + if [ "$DESKTOP_PRERELEASE" = "true" ]; then + release_flags+=(--prerelease) + fi + gh release create "$DESKTOP_RELEASE_TAG" "${release_flags[@]}" + fi + + - name: Publish versioned release assets + shell: bash + env: + GH_TOKEN: ${{ github.token }} + run: | + set -euo pipefail + gh release upload "$DESKTOP_RELEASE_TAG" "$RUNNER_TEMP/desktop-release-assets"/* --clobber + + - name: Generate and publish versioned updater metadata + shell: bash + env: + GH_TOKEN: ${{ github.token }} + run: | + set -euo pipefail + python3 <<'PY' + import datetime + import json + import os + import pathlib + import sys + import urllib.parse + + asset_dir = pathlib.Path(os.environ['RUNNER_TEMP'], 'desktop-release-assets') + files = [path for path in asset_dir.iterdir() if path.is_file()] + + def exactly_one(suffix: str) -> pathlib.Path: + matches = [path for path in files if path.name.endswith(suffix)] + if len(matches) != 1: + sys.exit(f'Expected exactly one {suffix} updater asset, found {len(matches)}') + return matches[0] + + def entry(signature_suffix: str) -> dict[str, str]: + signature_path = exactly_one(signature_suffix) + bundle_name = signature_path.name.removesuffix('.sig') + bundle_path = asset_dir / bundle_name + if not bundle_path.is_file(): + sys.exit(f'Missing updater bundle for {signature_path.name}: {bundle_name}') + encoded_tag = urllib.parse.quote(os.environ['DESKTOP_RELEASE_TAG'], safe='') + encoded_name = urllib.parse.quote(bundle_name, safe='') + return { + 'signature': signature_path.read_text(), + 'url': ( + f'https://github.com/{os.environ["GITHUB_REPOSITORY"]}/releases/download/' + f'{encoded_tag}/{encoded_name}' + ), + } + + darwin = entry('.app.tar.gz.sig') + linux = entry('.AppImage.sig') + windows = entry('.exe.sig') + notes = pathlib.Path(os.environ['RUNNER_TEMP'], 'desktop-release-notes.md').read_text() + metadata = { + 'version': os.environ['APP_VERSION'], + # App version is SemVer; CHANGELOG.md is keyed by the backend release. + 'pypi_version': os.environ['PYPI_VERSION'], + 'notes': notes, + 'pub_date': datetime.datetime.now(datetime.timezone.utc).isoformat(timespec='milliseconds').replace('+00:00', 'Z'), + 'platforms': { + 'darwin-aarch64': darwin, + 'darwin-aarch64-app': darwin, + 'linux-x86_64': linux, + 'linux-x86_64-appimage': linux, + 'windows-x86_64': windows, + 'windows-x86_64-nsis': windows, + }, + } + output = pathlib.Path(os.environ['RUNNER_TEMP'], 'latest.json') + output.write_text(json.dumps(metadata, indent=2) + '\n') + PY + gh release upload "$DESKTOP_RELEASE_TAG" "$RUNNER_TEMP/latest.json" --clobber + - name: Download versioned updater metadata + if: ${{ !inputs.draft }} shell: bash env: GH_TOKEN: ${{ github.token }} @@ -790,6 +956,7 @@ jobs: test -s "$RUNNER_TEMP/desktop-updater/latest.json" - name: Validate versioned updater metadata + if: ${{ !inputs.draft }} shell: bash run: | python3 <<'PY' @@ -849,6 +1016,7 @@ jobs: PY - name: Ensure desktop updater channel release + if: ${{ !inputs.draft }} shell: bash env: GH_TOKEN: ${{ github.token }} @@ -881,6 +1049,7 @@ jobs: PY - name: Prevent updater channel downgrade + if: ${{ !inputs.draft }} shell: bash env: GH_TOKEN: ${{ github.token }} @@ -971,6 +1140,7 @@ jobs: PY - name: Publish desktop updater channel metadata + if: ${{ !inputs.draft }} shell: bash env: GH_TOKEN: ${{ github.token }} diff --git a/.github/workflows/security-audit.yml b/.github/workflows/security-audit.yml index 1275d12216..27eafbedea 100644 --- a/.github/workflows/security-audit.yml +++ b/.github/workflows/security-audit.yml @@ -36,8 +36,8 @@ # - unsloth `huggingfacenotorch` extras (the canonical install path # for fine-tuning users; pulls transformers / peft / accelerate / # trl / datasets / diffusers / sentence-transformers / etc.) -# - all six Studio backend requirements files -# - Studio frontend (npm) and Tauri shell (cargo) +# - all six Unsloth backend requirements files +# - Unsloth frontend (npm) and Tauri shell (cargo) # Each Python step builds a filtered dep list from pyproject.toml + # requirements/*.txt before auditing. We do NOT install any of these # -- pip-audit resolves through PyPI metadata, scan_packages.py @@ -218,7 +218,7 @@ jobs: # on the runner). A comment line is left in place so the # skipped specs are obvious in the artifact. # The `huggingface` extra is `huggingfacenotorch` plus torch / - # torchvision / triton, deliberately skipped: Studio backend + # torchvision / triton, deliberately skipped: Unsloth backend # already pins a torch and the +cu* / +cpu local-version tags # trip up the PyPI resolver in `-r` mode. run: | @@ -253,7 +253,7 @@ jobs: # `-r requirements.txt` resolves the requirements through pip's # dependency resolver against PyPI metadata and audits the # resolved tree without ever executing setup.py / install - # hooks. Way faster than installing the full Studio runtime + # hooks. Way faster than installing the full Unsloth runtime # and -- critically -- safer: an attacker who has compromised # a transitive dep cannot run code in this job. # @@ -326,9 +326,9 @@ jobs: } >> "$GITHUB_STEP_SUMMARY" # ───────────────────────────────────────────────────────────── - # npm: Studio frontend + # npm: Unsloth frontend # ───────────────────────────────────────────────────────────── - - name: npm audit (Studio frontend) + - name: npm audit (Unsloth frontend) # `npm audit` resolves the lockfile through the npmjs.com # advisory DB. `--audit-level=high` filters the noise floor # to only HIGH and CRITICAL. We do NOT pass --omit=dev: a @@ -342,7 +342,7 @@ jobs: # Always also write the full JSON for grep-ability. npm audit --json > ../../logs-npm-audit.json || true { - echo "## npm audit (Studio frontend)" + echo "## npm audit (Unsloth frontend)" echo echo '```' tail -200 ../../logs-npm-audit.txt @@ -350,9 +350,9 @@ jobs: } >> "$GITHUB_STEP_SUMMARY" # ───────────────────────────────────────────────────────────── - # cargo: Studio Tauri shell + # cargo: Unsloth Tauri shell # ───────────────────────────────────────────────────────────── - - name: cargo audit (Studio Tauri) + - name: cargo audit (Unsloth Tauri) # `--deny warnings` would make the job fail on any advisory. # Keep non-blocking initially; drop continue-on-error after # the baseline closes. @@ -362,7 +362,7 @@ jobs: set +e cargo audit | tee ../../logs-cargo-audit.txt { - echo "## cargo audit (Studio Tauri)" + echo "## cargo audit (Unsloth Tauri)" echo echo '```' tail -200 ../../logs-cargo-audit.txt @@ -559,7 +559,7 @@ jobs: # ───────────────────────────────────────────────────────────── # CycloneDX SBOM. Lets downstream consumers audit what's - # actually shipped in unsloth wheels and the Studio backend + # actually shipped in unsloth wheels and the Unsloth backend # runtime. Generates one JSON file per requirements input plus # a combined SBOM keyed off pyproject.toml; uploads as a build # artifact (and a future step can attest it via SLSA). @@ -740,7 +740,7 @@ jobs: # `--with-deps` makes the scan transitive: every package the # declared set resolves to gets fetched and pattern-scanned, not # just the top-level pins. Resolving the full transitive closure - # of the unsloth + Studio dep tree downloads several hundred + # of the unsloth + Unsloth dep tree downloads several hundred # archives, hence the longer timeout. # # Sharded across runners for wall-clock parallelism. Each shard @@ -749,7 +749,7 @@ jobs: # composition tries to balance load: # - hf-stack: pyproject extras + no-torch-runtime # (~150 archives, transformers/peft/accelerate/...) - # - studio: FastAPI/Studio backend + overrides + extras-no-deps + # - studio: FastAPI/Unsloth backend + overrides + extras-no-deps # (~150 archives, smaller scientific stack) # - extras: the heavy openai-whisper / scikit-learn / librosa # stack (~250 archives, dominant cost) @@ -964,7 +964,7 @@ jobs: # documented at scripts/scan_npm_packages.py top-of-file. The # script is stdlib-only so adding it does not increase the # transitive supply-chain surface. - name: npm scan-packages (Studio frontend tarballs) + name: npm scan-packages (Unsloth frontend tarballs) runs-on: ubuntu-latest timeout-minutes: 30 needs: [] @@ -1173,7 +1173,7 @@ jobs: with: python-version: '3.12' - - name: Install Studio frontend deps (--ignore-scripts) + - name: Install Unsloth frontend deps (--ignore-scripts) # `npm audit signatures` requires node_modules to be populated. # `--ignore-scripts` is mandatory: this is exactly the lever the # new-install-script gate below protects against, and we must diff --git a/.github/workflows/startup-profile-ci.yml b/.github/workflows/startup-profile-ci.yml new file mode 100644 index 0000000000..fbde99836d --- /dev/null +++ b/.github/workflows/startup-profile-ci.yml @@ -0,0 +1,156 @@ +# SPDX-License-Identifier: AGPL-3.0-only +# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. + +# Measures where Studio's startup time goes, on each platform. +# +# Nothing recorded a number before: main.py logs "lifespan startup completed in X ms" +# and studio_test_kit polls /healthz, but both throw the elapsed time away. A first +# local run (Linux, warm cache, 18-core server) put `import main` at 5.7-6.6s BEFORE +# the server can bind, dominated by eager module-level imports pulled in by routes: +# torch ~1.9s self, unsloth_zoo ~0.8s, routes ~0.6s, transformers ~0.5s. +# +# Not a gate yet: --max-healthz-seconds exists, but a budget should come from +# observed numbers rather than a guess. + +name: Startup profile + +on: + pull_request: + paths: + # The measured import graph is the whole backend tree: main.py imports auth, + # core, hub, loggers, models, picker, routes and utils at module scope. + - 'studio/backend/**' + - '!studio/backend/tests/**' + # The launch phase spawns `unsloth studio --api-only`, so the CLI counts too. + - 'unsloth_cli/**' + - 'studio/src-tauri/src/preflight**' + # The profiler hardcodes the desktop argv that process.rs::backend_args builds, + # so a change there must schedule a run or the two silently diverge. + - 'studio/src-tauri/src/process.rs' + - 'scripts/profile_startup.py' + - '.github/workflows/startup-profile-ci.yml' + # The job profiles whatever `install.sh --local` built: the installers pick the + # venv's Python and the dependency specs, and pyproject's include list is what + # makes --local overlay studio.backend*. + - 'install.sh' + - 'install.ps1' + - 'pyproject.toml' + # --local also runs the checkout's setup scripts (install.sh picks + # $_REPO_ROOT/studio/setup.sh, the editable install resolves setup.ps1 to the + # repo), and both call install_python_stack.py, which picks the dependencies. + - 'studio/setup.sh' + - 'studio/setup.ps1' + - 'studio/install_python_stack.py' + workflow_dispatch: + inputs: + repeats: + description: 'launch repeats per OS (median reported)' + type: string + default: '3' + +concurrency: + group: ${{ github.workflow }}-${{ github.ref }} + cancel-in-progress: true + +permissions: + contents: read + +jobs: + profile: + name: startup ${{ matrix.os }} + runs-on: ${{ matrix.os }} + timeout-minutes: 60 + continue-on-error: true + strategy: + fail-fast: false + matrix: + os: [ubuntu-latest, macos-14, windows-latest] + + env: + UNSLOTH_STUDIO_HOME: ${{ github.workspace }}/.studio-home + # A wildcard bind calls ifconfig.me on the startup path; loopback times our code. + UNSLOTH_STUDIO_DISABLE_PUBLIC_CHECK: '1' + + steps: + - uses: actions/checkout@de0fac2e4500dabe0009e67214ff5f5447ce83dd # v6.0.2 + with: + persist-credentials: false + + - name: Install Studio + shell: bash + env: + GH_TOKEN: ${{ secrets.GITHUB_TOKEN }} + run: | + set -o pipefail + mkdir -p logs + # --local is load-bearing: it overlays the checkout, so the profiled server + # is this diff. Without it install.sh resolves unsloth from PyPI. + if [ "${{ runner.os }}" = "Windows" ]; then + pwsh -NoProfile -File ./install.ps1 --local 2>&1 | tee logs/install.log + else + bash install.sh --local 2>&1 | tee logs/install.log + fi + + - name: Profile startup + shell: bash + run: | + BIN="$UNSLOTH_STUDIO_HOME/unsloth_studio/bin/unsloth" + [ -x "$BIN" ] || BIN="$UNSLOTH_STUDIO_HOME/unsloth_studio/Scripts/unsloth.exe" + [ -x "$BIN" ] || BIN="" + # Profile imports with the INSTALLED interpreter: that venv is what launches. + PY="$UNSLOTH_STUDIO_HOME/unsloth_studio/bin/python" + [ -x "$PY" ] || PY="$UNSLOTH_STUDIO_HOME/unsloth_studio/Scripts/python.exe" + [ -x "$PY" ] || PY="$(command -v python3 || command -v python)" + python3 scripts/profile_startup.py \ + --python "$PY" \ + ${BIN:+--bin "$BIN"} \ + --repeats "${{ inputs.repeats || '3' }}" \ + --json "startup-${{ matrix.os }}.json" 2>&1 | tee logs/profile.log + + - name: Summary + if: always() + shell: bash + run: | + f="startup-${{ matrix.os }}.json" + [ -f "$f" ] || { echo "no profile produced"; exit 0; } + python3 - "$f" >> "$GITHUB_STEP_SUMMARY" <<'PY' + import json, sys + d = json.load(open(sys.argv[1])) + print(f"### {d['platform']} / {d['machine']} (py {d['python']}, {d['cpu_count']} cpu)\n") + imp = d.get("imports", {}) + # Gate on ok: a failed `import main` still leaves rows, so a total can lie. + if imp.get("ok"): + print(f"**`import main`: {imp['total_seconds']}s**\n") + print("| package | self ms |") + print("|---|---:|") + for k, v in list(imp.get("self_by_package_ms", {}).items())[:8]: + print(f"| {k} | {v} |") + print() + else: + print("**`import main` failed - no valid import profile**\n") + print("```\n" + (imp.get("error") or "")[-1500:] + "\n```\n") + lau = d.get("launch") or {} + runs = len(lau.get("runs") or []) + failed = lau.get("failed_runs") or 0 + if lau.get("healthz_median_seconds") is not None: + # The aggregates cover only the runs that reached healthz, so flag the + # failures: bare numbers would read as a normal fast startup. + note = f" _({runs - failed} of {runs} launches; {failed} never became healthy)_" if failed else "" + print(f"**time to a healthy port: {lau['healthz_median_seconds']}s median, " + f"{lau['healthz_max_seconds']}s max**{note}\n") + elif lau.get("skipped"): + print(f"_launch phase skipped: {lau['skipped']}_\n") + elif runs: + print(f"**no launch measurement: all {runs} launches failed to become healthy**\n") + PY + + - name: Upload profile + if: always() + uses: actions/upload-artifact@043fb46d1a93c77aae656e7c1c64a875d1fc6a0a # v7.0.1 + with: + name: startup-profile-${{ matrix.os }} + path: | + startup-*.json + logs/ + retention-days: 14 + if-no-files-found: warn diff --git a/.github/workflows/studio-api-smoke.yml b/.github/workflows/studio-api-smoke.yml index 15efee382e..1cfa66fea4 100644 --- a/.github/workflows/studio-api-smoke.yml +++ b/.github/workflows/studio-api-smoke.yml @@ -1,7 +1,7 @@ # SPDX-License-Identifier: AGPL-3.0-only # Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. -# Studio API & Auth Tests -- HTTP-level integration tests for the +# Unsloth API & Auth Tests -- HTTP-level integration tests for the # FastAPI surface. No Playwright, no model UI; tests/studio/test_studio_api_smoke.py # runs ~30 s and asserts: # - CORS hardening (no wildcard + credentials, no bootstrap leak) @@ -15,7 +15,7 @@ # Reuses the GGUF cache key from studio-ui-smoke.yml so the model # download is one cache-hit on the second job. -name: Studio API CI +name: Unsloth API CI on: pull_request: @@ -40,7 +40,7 @@ permissions: jobs: api-smoke: - name: Studio API & Auth Tests + name: Unsloth API & Auth Tests runs-on: ubuntu-latest timeout-minutes: 12 env: @@ -98,7 +98,7 @@ jobs: path: hf-cache key: ${{ runner.os }}-hf-${{ env.GGUF_REPO }}-${{ env.GGUF_VARIANT }}-v2 - - name: Install Studio (--local, --no-torch) + - name: Install Unsloth (--local, --no-torch) env: GH_TOKEN: ${{ secrets.GITHUB_TOKEN }} # Withheld on PR: this step runs checked-out PR code; public GGUF still downloads. @@ -111,9 +111,10 @@ jobs: - name: Install pyjwt for the JWT-expiry forge test run: pip install 'pyjwt>=2.6' - - name: Reset auth + boot Studio (API-only) + - name: Reset auth + boot Unsloth (API-only) run: | - unsloth studio reset-password + # Wipe (not reset-password): the boot below must re-seed a fresh .bootstrap_password. + rm -rf ~/.unsloth/studio/auth mkdir -p logs UNSLOTH_API_ONLY=1 unsloth studio -H 127.0.0.1 -p "$STUDIO_PORT" \ > logs/studio.log 2>&1 & @@ -144,7 +145,7 @@ jobs: echo "STUDIO_NEW_PW=$NEW" >> "$GITHUB_ENV" echo "STUDIO_NEW2_PW=$NEW2" >> "$GITHUB_ENV" - - name: Run Studio API & Auth tests + - name: Run Unsloth API & Auth tests # The script is named WITHOUT a `test_` prefix so it isn't # auto-collected by pytest in Backend CI's `tests/` walk # (which doesn't set BASE_URL and would crash at import). @@ -153,7 +154,7 @@ jobs: STUDIO_AUTH_DIR: /home/runner/.unsloth/studio/auth run: python tests/studio/studio_api_smoke.py - - name: Stop Studio + - name: Stop Unsloth if: always() run: | kill "${STUDIO_PID}" 2>/dev/null || true diff --git a/.github/workflows/studio-backend-ci.yml b/.github/workflows/studio-backend-ci.yml index 3022127a2b..dd5efbb299 100644 --- a/.github/workflows/studio-backend-ci.yml +++ b/.github/workflows/studio-backend-ci.yml @@ -30,6 +30,13 @@ on: - 'unsloth/**' - 'unsloth_cli/**' - 'tests/**' + # The root installers: tests/sh/*.sh and tests/studio/install/* assert + # against these two files, so a change here must run the suite that + # covers it. Without them an install-only edit (the shape most AMD/ROCm + # routing fixes take) skipped Backend CI entirely. + - 'install.sh' + - 'install.ps1' + - 'scripts/**' - 'pyproject.toml' - '.github/workflows/studio-backend-ci.yml' push: @@ -64,7 +71,7 @@ jobs: - name: Install backend test dependencies (CPU only) run: | python -m pip install --upgrade pip - # Studio's declared backend deps: + # Unsloth's declared backend deps: pip install -r studio/backend/requirements/studio.txt # Extras that studio.txt does not list but the import chain needs # (python-multipart for FastAPI form/file uploads, sqlalchemy/cryptography @@ -193,6 +200,7 @@ jobs: --ignore=tests/sh \ --ignore=tests/studio/test_hardware_dispatch_matrix.py \ --ignore=tests/studio/test_is_mlx_dispatch_gate.py \ + --ignore=tests/studio/test_xpu_spoof_pipeline.py \ --ignore=tests/vllm_compat \ --ignore=tests/version_compat \ -m 'not server and not e2e' \ @@ -205,36 +213,53 @@ jobs: env: PYTHONPATH: ${{ github.workspace }}/studio UNSLOTH_COMPILE_DISABLE: '1' - # These two files mutate hardware.py module globals at runtime - # via the spoof fixtures, which leaks state into any other test - # that imports hardware. Run them in their own pytest invocation - # so the leak does not cross file boundaries. + # These files mutate hardware.py module globals at runtime via the + # spoof fixtures (CUDA/ROCm/XPU/MLX/CPU), which leaks state into any + # other test that imports hardware. Run them in their own pytest + # invocation so the leak does not cross file boundaries. run: | python -m pytest -q --tb=short \ tests/studio/test_hardware_dispatch_matrix.py \ - tests/studio/test_is_mlx_dispatch_gate.py + tests/studio/test_is_mlx_dispatch_gate.py \ + tests/studio/test_xpu_spoof_pipeline.py + + - name: CLI tests (unsloth_cli) + # unsloth_cli/tests had no CI at all: `unsloth_cli/**` was only a paths + # trigger and a ruff target, so 673 tests covering the studio launcher, + # the pre-exposure gate and the auth secret writers ran nowhere, and + # four of them had been failing on main unnoticed. + # Own step, not folded into the tests/ discovery above: pyproject's + # testpaths is tests/, and this suite needs no PYTHONPATH or CUDA spoof + # (it self-bootstraps sys.path and imports neither unsloth nor torch). + run: python -m pytest unsloth_cli/tests -q --tb=short - name: Shell installer tests - # Subset that does not depend on a writable / pristine install.sh - # tree; test_install_host_defaults.sh checks install.ps1 layout - # which has drifted (separate followup). + # Auto-discovered rather than allowlisted. The old hardcoded list had + # silently fallen seven files behind tests/run_all.sh, including + # test_strixhalo_wsl_reroute.sh -- the only shell coverage of the ROCm + # WSL reroute -- so that suite never ran on a PR. Skips are explicit, + # each with a reason, and tests/studio/test_ci_shell_suite_coverage.py + # fails if this step stops discovering the directory or the skip list + # grows without one. + # + # Skipped: + # test_install_host_defaults.sh: asserts an install.ps1 layout that + # has drifted (separate followup). + # test_install_rollback_lifecycle.sh: already runs on both platforms + # in cross-platform-parity-ci.yml. run: | set -e - for s in \ - tests/sh/test_get_torch_index_url.sh \ - tests/sh/test_mac_intel_compat.sh \ - tests/sh/test_node_decision.sh \ - tests/sh/test_studio_home_node_dir.sh \ - tests/sh/test_system_node_readonly.sh \ - tests/sh/test_nvcc_meets_llama_minimum.sh \ - tests/sh/test_resolve_cuda_archs.sh \ - tests/sh/test_tauri_install_exit_order.sh \ - tests/sh/test_torch_constraint.sh \ - tests/sh/test_torch_flavor.sh \ - tests/sh/test_with_llama_cpp_dir_flag.sh \ - tests/sh/test_with_llama_cpp_dir_link_behavior.sh; do + skip="test_install_host_defaults.sh test_install_rollback_lifecycle.sh" + found=0 + for s in tests/sh/test_*.sh; do + case " $skip " in + *" $(basename "$s") "*) echo "skipping $s (see workflow comment)"; continue ;; + esac + found=$((found + 1)) echo "::group::$s" bash "$s" echo "::endgroup::" done + [ "$found" -gt 0 ] || { echo "::error::no shell tests discovered under tests/sh"; exit 1; } + echo "ran $found shell installer test files" diff --git a/.github/workflows/studio-export-capability-ci.yml b/.github/workflows/studio-export-capability-ci.yml index 1ee6489209..83df3ed476 100644 --- a/.github/workflows/studio-export-capability-ci.yml +++ b/.github/workflows/studio-export-capability-ci.yml @@ -9,7 +9,7 @@ # export is validated separately. No GPU / model / llama.cpp: the tests mock the probes and block # torch/unsloth, so the job installs only a CPU PyTorch plus import deps. -name: Studio export capability +name: Unsloth export capability on: pull_request: diff --git a/.github/workflows/studio-frontend-ci.yml b/.github/workflows/studio-frontend-ci.yml index b42086f191..773e555c8b 100644 --- a/.github/workflows/studio-frontend-ci.yml +++ b/.github/workflows/studio-frontend-ci.yml @@ -133,10 +133,13 @@ jobs: - name: Typecheck run: npm run typecheck + - name: Unit tests + run: npm test + - name: Build run: npm run build - - name: Built bundle must not contain Studio's unstable_Provider call site + - name: Built bundle must not contain Unsloth's unstable_Provider call site run: | set -e JS=$(ls dist/assets/index-*.js | head -1) @@ -144,7 +147,7 @@ jobs: echo "main bundle: $JS" echo "unstable_Provider: hits=$HITS (assistant-ui internals contribute up to 3)" if [ "$HITS" -gt 3 ]; then - echo "::error file=studio/frontend/src/features/chat/runtime-provider.tsx::Studio bundle still passes unstable_Provider through useRemoteThreadListRuntime; this is the 2026.5.1 chat-history regression. Pass adapters directly into useLocalRuntime instead." + echo "::error file=studio/frontend/src/features/chat/runtime-provider.tsx::Unsloth bundle still passes unstable_Provider through useRemoteThreadListRuntime; this is the 2026.5.1 chat-history regression. Pass adapters directly into useLocalRuntime instead." exit 1 fi diff --git a/.github/workflows/studio-inference-smoke.yml b/.github/workflows/studio-inference-smoke.yml index cf8c021e38..c37c9555bf 100644 --- a/.github/workflows/studio-inference-smoke.yml +++ b/.github/workflows/studio-inference-smoke.yml @@ -1,7 +1,7 @@ # SPDX-License-Identifier: AGPL-3.0-only # Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. -# Three end-to-end smoke jobs that boot a freshly-installed Studio and +# Three end-to-end smoke jobs that boot a freshly-installed Unsloth and # exercise the surfaces real users hit through the OpenAI / Anthropic # SDKs and curl. Each job picks the smallest model that exercises the # behaviour under test, primes HF_HOME via actions/cache, and shares @@ -27,7 +27,7 @@ # All three jobs run in parallel. Total wall time is dominated by job 3 # on a cold cache; warm cache cuts that to ~3 min. -name: Studio GGUF CI +name: Unsloth GGUF CI on: pull_request: @@ -112,7 +112,7 @@ jobs: path: hf-cache key: ${{ runner.os }}-hf-${{ env.GGUF_REPO }}-${{ env.GGUF_VARIANT }}-v2 - - name: Install Studio (--local, --no-torch) + - name: Install Unsloth (--local, --no-torch) env: GH_TOKEN: ${{ secrets.GITHUB_TOKEN }} # Withheld on PR: this step runs checked-out PR code; public GGUF still downloads. @@ -125,9 +125,10 @@ jobs: - name: Install OpenAI + Anthropic Python SDKs run: pip install 'openai>=1.50' 'anthropic>=0.40' - - name: Reset auth + boot Studio (API-only) + - name: Reset auth + boot Unsloth (API-only) run: | - unsloth studio reset-password + # Wipe (not reset-password): the boot below must re-seed a fresh .bootstrap_password. + rm -rf ~/.unsloth/studio/auth mkdir -p logs UNSLOTH_API_ONLY=1 unsloth studio -H 127.0.0.1 -p "$STUDIO_PORT" \ > logs/studio.log 2>&1 & @@ -142,7 +143,7 @@ jobs: fi sleep 1 done - echo "Studio did not become healthy in 180s" + echo "Unsloth did not become healthy in 180s" tail -200 logs/studio.log exit 1 @@ -229,11 +230,11 @@ jobs: return replies def run_anthropic(): - # Two SDK quirks vs. Studio: + # Two SDK quirks vs. Unsloth: # 1. base_url must NOT include /v1 -- the SDK appends # /v1/messages itself; otherwise the request hits # /v1/v1/messages and 405s. - # 2. The SDK sends `x-api-key` by default, but Studio's + # 2. The SDK sends `x-api-key` by default, but Unsloth's # auth layer is HTTPBearer-only. Override via # default_headers so Authorization: Bearer ... is # sent instead. @@ -276,7 +277,7 @@ jobs: print( f"[{label}] WARN non-determinism at temperature=0.0 across " f"{len(determinism_failures)} of {len(first)} turn(s); " - f"small-quant model drift, not a Studio regression. " + f"small-quant model drift, not an Unsloth regression. " f"Details: " + " | ".join(determinism_failures) ) # Sanity: turn-2 reply should mention the earlier question, and @@ -290,7 +291,7 @@ jobs: print(f"[{label}] {status_word} -- 4 turns, history grounded ('paris' present)") PY - - name: Stop Studio + - name: Stop Unsloth if: always() run: | kill "${STUDIO_PID}" 2>/dev/null || true @@ -323,7 +324,7 @@ jobs: # store xet chunks + blobs + snapshots = ~4 GiB compressed -- # 4-5x file-size inflation, dominated by xet chunks. Use main's # `--local-dir gguf-cache` pattern to cache the flat .gguf only. - # Studio's /api/inference/load accepts either a HF repo (which + # Unsloth's /api/inference/load accepts either a HF repo (which # uses HF_HOME) or an absolute file path; passing the absolute # path keeps the test off HF_HOME entirely so the cache size # tracks the GGUF file 1:1. The OpenAI/Anth and JSON+images @@ -380,7 +381,7 @@ jobs: path: gguf-cache key: ${{ runner.os }}-gguf-${{ env.GGUF_REPO }}-${{ env.GGUF_FILE }}-v1 - - name: Install Studio (--local, --no-torch) + - name: Install Unsloth (--local, --no-torch) env: GH_TOKEN: ${{ secrets.GITHUB_TOKEN }} # Withheld on PR: this step runs checked-out PR code; public GGUF still downloads. @@ -390,7 +391,7 @@ jobs: set -o pipefail bash install.sh --local --no-torch 2>&1 | tee logs/install.log - - name: Reset auth + boot Studio (API-only, default tool policy) + - name: Reset auth + boot Unsloth (API-only, default tool policy) # We deliberately use the API-only mode rather than # `unsloth studio run` because the latter calls # `set_tool_policy(...)` with a resolved bool: on loopback the @@ -400,7 +401,7 @@ jobs: # tool_policy=None so each request's `enable_tools` field is # honoured. run: | - unsloth studio reset-password + rm -rf ~/.unsloth/studio/auth mkdir -p logs UNSLOTH_API_ONLY=1 unsloth studio -H 127.0.0.1 -p "$STUDIO_PORT" \ > logs/studio.log 2>&1 & @@ -503,7 +504,7 @@ jobs: that the tool path executed. A shared CI runner can stall the stream transport (the - connection opening, or a mid-stream read) even when Studio + connection opening, or a mid-stream read) even when Unsloth is healthy, so retry a stall once with a fresh request capped at 300s. A stall means the stream did NOT complete, so partial events are normally NOT returned (an early @@ -575,11 +576,11 @@ jobs: def _tool_invoked(events): """Structural check: True iff some SSE payload is a real - tool envelope (Studio tool_start/tool_end, Anthropic + tool envelope (Unsloth tool_start/tool_end, Anthropic tool_use/tool_result, OpenAI non-empty delta.tool_calls / message.tool_calls / finish_reason='tool_calls' / role:'tool' / function_call). tool_status is NOT - evidence: Studio emits empty tool_status events on + evidence: Unsloth emits empty tool_status events on iteration boundaries even when no tool ran. """ for raw in events: @@ -698,7 +699,7 @@ jobs: attempt has structural invocation evidence. WARN (not FAIL) if invoked but no attempt produces the expected literal in tool_end.result -- small-quant Qwen3.5-2B can - emit OpenAI tool_calls deltas without Studio's GGUF + emit OpenAI tool_calls deltas without Unsloth's GGUF agentic loop intercepting them, and that GGUF-vs-OpenAI format mismatch is out of scope for #5642. """ @@ -729,6 +730,7 @@ jobs: content, events = post_sse("/v1/chat/completions", { "messages": [{"role": "user", "content": prompt}], "enable_tools": True, + "permission_mode": "full", "enabled_tools": enabled, "session_id": f"{session}-att{attempt_i}", "temperature": TOOL_PROBE_TEMP, @@ -810,7 +812,7 @@ jobs: # because (a) the search may legitimately return no results, # and (b) DuckDuckGo upstream blocks GHA IP ranges often # enough that requiring a tool_call marker would create - # red-herring failures from infra rather than from Studio. + # red-herring failures from infra rather than from Unsloth. try: # Best-effort and bounded: a single 180s attempt keeps a stall # from eating the job's timeout-minutes (it already WARNs, so a @@ -818,6 +820,7 @@ jobs: content, events = post_sse("/v1/chat/completions", { "messages": [{"role": "user", "content": "Search the web for 'unsloth ai github' and summarise."}], "enable_tools": True, + "permission_mode": "full", "enabled_tools": ["web_search"], "session_id": "ci-tool-calling-web", "temperature": 0.0, @@ -832,7 +835,7 @@ jobs: print(f"[tools] WARN web_search probe failed (non-blocking): {exc}") # ── 5. Thinking on / off ───────────────────────────────────── - # Studio strips think blocks from message.content for tools-mode + # Unsloth strips think blocks from message.content for tools-mode # responses, so we toggle plain chat (no enable_tools) and look # at the surfaced reasoning_content / message.thinking field. def thinking_call(enable): @@ -846,7 +849,7 @@ jobs: }) assert status == 200 msg = data["choices"][0]["message"] - # Studio surfaces thinking via reasoning_content (OpenAI + # Unsloth surfaces thinking via reasoning_content (OpenAI # extension). Fall back to inline markers for # robustness across template versions. raw = (msg.get("content") or "") + (msg.get("reasoning_content") or "") @@ -866,7 +869,7 @@ jobs: print(f"[tools] PASS thinking on/off (on={len(on_text)} chars, off={len(off_text)} chars)") PY - - name: Stop Studio + - name: Stop Unsloth if: always() run: | kill "${STUDIO_PID}" 2>/dev/null || true @@ -958,7 +961,7 @@ jobs: path: hf-cache key: ${{ runner.os }}-hf-${{ env.GGUF_REPO }}-${{ env.GGUF_VARIANT }}-${{ env.MMPROJ_FILE }}-v2 - - name: Install Studio (--local, --no-torch) + - name: Install Unsloth (--local, --no-torch) env: GH_TOKEN: ${{ secrets.GITHUB_TOKEN }} # Withheld on PR: this step runs checked-out PR code; public GGUF still downloads. @@ -971,12 +974,12 @@ jobs: - name: Install OpenAI + Anthropic Python SDKs run: pip install 'openai>=1.50' 'anthropic>=0.40' - - name: Reset auth + boot Studio (API-only) + - name: Reset auth + boot Unsloth (API-only) # See Job 2's comment: API-only mode keeps tool_policy=None so # response_format requests aren't routed through the agentic # tool loop. run: | - unsloth studio reset-password + rm -rf ~/.unsloth/studio/auth mkdir -p logs UNSLOTH_API_ONLY=1 unsloth studio -H 127.0.0.1 -p "$STUDIO_PORT" \ > logs/studio.log 2>&1 & @@ -1074,13 +1077,13 @@ jobs: # llama.cpp's HTTP server supports OpenAI-compatible JSON # mode: `response_format: {"type": "json_object"}` constrains # the model to emit syntactically-valid JSON. We use raw HTTP - # rather than the OpenAI SDK so that the field shape Studio + # rather than the OpenAI SDK so that the field shape Unsloth # forwards to llama-server is unambiguous (the SDK rewrites # response_format depending on which variant it recognises). # We deliberately do NOT pass a strict JSON schema -- on # small Gemma-4 quants the GBNF-from-schema path occasionally # produces empty output, and JSON mode is the surface we care - # about exposing through Studio. + # about exposing through Unsloth. status, data = post("/v1/chat/completions", { "model": "default", "messages": [ @@ -1110,7 +1113,7 @@ jobs: print(f"[json] PASS json_object -> {parsed}") # ── 2. OpenAI image_url (data URI base64) ─────────────────── - # 64x64 solid-red PNG. stb_image (used by Studio's image + # 64x64 solid-red PNG. stb_image (used by Unsloth's image # normaliser at routes/inference.py:3410) rejects 4x4 or # smaller PNGs as truncated, so we go up to 64x64 -- still # tiny in token cost. The assertion is loose: any non-empty @@ -1146,9 +1149,9 @@ jobs: print("[image/openai] PASS image_url accepted, non-empty response") # ── 3. Anthropic source/base64 image ──────────────────────── - # Two SDK quirks vs. Studio: base_url must NOT include /v1 + # Two SDK quirks vs. Unsloth: base_url must NOT include /v1 # (the SDK appends it itself; otherwise /v1/v1/messages -> 405), - # and Studio's auth is HTTPBearer-only so the SDK's default + # and Unsloth's auth is HTTPBearer-only so the SDK's default # x-api-key header is ignored -- send Authorization: Bearer # via default_headers. anthropic = Anthropic( @@ -1182,7 +1185,7 @@ jobs: print("[image/anthropic] PASS source/base64 accepted, non-empty response") PY - - name: Stop Studio + - name: Stop Unsloth if: always() run: | kill "${STUDIO_PID}" 2>/dev/null || true diff --git a/.github/workflows/studio-load-orchestrator-ci.yml b/.github/workflows/studio-load-orchestrator-ci.yml index 93d1a7742d..8710efc2bd 100644 --- a/.github/workflows/studio-load-orchestrator-ci.yml +++ b/.github/workflows/studio-load-orchestrator-ci.yml @@ -1,7 +1,7 @@ # SPDX-License-Identifier: AGPL-3.0-only # Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. # -# Event-loop regression test for the Studio model-load orchestrator. +# Event-loop regression test for the Unsloth model-load orchestrator. # Pins down issue #5642 (Win10 UI freeze on model load): the /load # route calls LlamaCppBackend.detect_audio_type synchronously, blocking # the FastAPI event loop on a chain of sync httpx.Client.post() probes. @@ -14,7 +14,7 @@ # danielhanchen/unsloth-staging-2 (Ubuntu / macOS / Windows all # green at PR time). -name: Studio load-orchestrator CI +name: Unsloth load-orchestrator CI on: pull_request: diff --git a/.github/workflows/studio-mac-api-smoke.yml b/.github/workflows/studio-mac-api-smoke.yml index 617ce189dc..c2307f17a1 100644 --- a/.github/workflows/studio-mac-api-smoke.yml +++ b/.github/workflows/studio-mac-api-smoke.yml @@ -33,7 +33,7 @@ permissions: jobs: api-smoke: - name: Studio API & Auth Tests + name: Unsloth API & Auth Tests runs-on: macos-14 timeout-minutes: 25 env: @@ -83,7 +83,7 @@ jobs: path: hf-cache key: ${{ runner.os }}-hf-${{ env.GGUF_REPO }}-${{ env.GGUF_VARIANT }}-v2 - - name: Install Studio (--local, --no-torch) + - name: Install Unsloth (--local, --no-torch) env: GH_TOKEN: ${{ secrets.GITHUB_TOKEN }} # Withheld on PR: this step runs checked-out PR code; public GGUF still downloads. @@ -99,9 +99,10 @@ jobs: - name: Install pyjwt for the JWT-expiry forge test run: pip install 'pyjwt>=2.6' - - name: Reset auth + boot Studio (API-only) + - name: Reset auth + boot Unsloth (API-only) run: | - unsloth studio reset-password + # Wipe (not reset-password): the boot below must re-seed a fresh .bootstrap_password. + rm -rf ~/.unsloth/studio/auth mkdir -p logs UNSLOTH_API_ONLY=1 unsloth studio -H 127.0.0.1 -p "$STUDIO_PORT" \ > logs/studio.log 2>&1 & @@ -129,13 +130,13 @@ jobs: echo "STUDIO_NEW_PW=$NEW" >> "$GITHUB_ENV" echo "STUDIO_NEW2_PW=$NEW2" >> "$GITHUB_ENV" - - name: Run Studio API & Auth tests + - name: Run Unsloth API & Auth tests env: BASE_URL: http://127.0.0.1:18895 STUDIO_AUTH_DIR: /Users/runner/.unsloth/studio/auth run: python tests/studio/studio_api_smoke.py - - name: Stop Studio + - name: Stop Unsloth if: always() run: | kill "${STUDIO_PID}" 2>/dev/null || true diff --git a/.github/workflows/studio-mac-inference-smoke.yml b/.github/workflows/studio-mac-inference-smoke.yml index d3d765aa84..1dbf86ae98 100644 --- a/.github/workflows/studio-mac-inference-smoke.yml +++ b/.github/workflows/studio-mac-inference-smoke.yml @@ -1,7 +1,7 @@ # SPDX-License-Identifier: AGPL-3.0-only # Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. -# Three end-to-end smoke jobs that boot a freshly-installed Studio and +# Three end-to-end smoke jobs that boot a freshly-installed Unsloth and # exercise the surfaces real users hit through the OpenAI / Anthropic # SDKs and curl. Each job picks the smallest model that exercises the # behaviour under test, primes a model cache via actions/cache, and @@ -108,7 +108,7 @@ jobs: path: hf-cache key: ${{ runner.os }}-hf-${{ env.GGUF_REPO }}-${{ env.GGUF_VARIANT }}-v2 - - name: Install Studio (--local, --no-torch) + - name: Install Unsloth (--local, --no-torch) env: GH_TOKEN: ${{ secrets.GITHUB_TOKEN }} # Withheld on PR: this step runs checked-out PR code; public GGUF still downloads. @@ -124,9 +124,10 @@ jobs: - name: Install OpenAI + Anthropic Python SDKs run: pip install 'openai>=1.50' 'anthropic>=0.40' - - name: Reset auth + boot Studio (API-only) + - name: Reset auth + boot Unsloth (API-only) run: | - unsloth studio reset-password + # Wipe (not reset-password): the boot below must re-seed a fresh .bootstrap_password. + rm -rf ~/.unsloth/studio/auth mkdir -p logs UNSLOTH_API_ONLY=1 unsloth studio -H 127.0.0.1 -p "$STUDIO_PORT" \ > logs/studio.log 2>&1 & @@ -141,7 +142,7 @@ jobs: fi sleep 1 done - echo "Studio did not become healthy in 180s" + echo "Unsloth did not become healthy in 180s" tail -200 logs/studio.log exit 1 @@ -228,11 +229,11 @@ jobs: return replies def run_anthropic(): - # Two SDK quirks vs. Studio: + # Two SDK quirks vs. Unsloth: # 1. base_url must NOT include /v1 -- the SDK appends # /v1/messages itself; otherwise the request hits # /v1/v1/messages and 405s. - # 2. The SDK sends `x-api-key` by default, but Studio's + # 2. The SDK sends `x-api-key` by default, but Unsloth's # auth layer is HTTPBearer-only. Override via # default_headers so Authorization: Bearer ... is # sent instead. @@ -283,7 +284,7 @@ jobs: print(f"[{label}] OK -- 4 turns, run1 == run2, history grounded") PY - - name: Stop Studio + - name: Stop Unsloth if: always() run: | kill "${STUDIO_PID}" 2>/dev/null || true @@ -363,7 +364,7 @@ jobs: path: gguf-cache key: ${{ runner.os }}-gguf-${{ env.GGUF_REPO }}-${{ env.GGUF_FILE }}-v1 - - name: Install Studio (--local, --no-torch) + - name: Install Unsloth (--local, --no-torch) env: GH_TOKEN: ${{ secrets.GITHUB_TOKEN }} # Withheld on PR: this step runs checked-out PR code; public GGUF still downloads. @@ -376,7 +377,7 @@ jobs: - name: Assert llama.cpp loads on this macOS run: bash .github/scripts/assert-llama-loads.sh - - name: Reset auth + boot Studio (API-only, default tool policy) + - name: Reset auth + boot Unsloth (API-only, default tool policy) # We deliberately use the API-only mode rather than # `unsloth studio run` because the latter calls # `set_tool_policy(...)` with a resolved bool: on loopback the @@ -386,7 +387,7 @@ jobs: # tool_policy=None so each request's `enable_tools` field is # honoured. run: | - unsloth studio reset-password + rm -rf ~/.unsloth/studio/auth mkdir -p logs UNSLOTH_API_ONLY=1 unsloth studio -H 127.0.0.1 -p "$STUDIO_PORT" \ > logs/studio.log 2>&1 & @@ -478,7 +479,7 @@ jobs: call with enable_tools=true must use this helper. A shared CI runner can stall the stream transport (the - connection opening, or a mid-stream read) even when Studio + connection opening, or a mid-stream read) even when Unsloth is healthy, so harden the read three ways: retry a stall once with a fresh request capped at 300s; return any text already streamed before a stall (a stall on the trailing @@ -574,11 +575,11 @@ jobs: assert status == 200, f"tool call status {status}: {data}" choice = data["choices"][0] tool_calls = (choice.get("message") or {}).get("tool_calls") or [] - # Studio's contract: when tool_choice='required', llama.cpp's + # Unsloth's contract: when tool_choice='required', llama.cpp's # grammar should force a tool_calls payload. On Mac that # contract is sometimes broken by the underlying quant; the # PASS path is "tool_calls present + correct schema", the - # WARN path documents Studio still returned 200 with a + # WARN path documents Unsloth still returned 200 with a # well-formed choices[] envelope. if tool_calls: tc = tool_calls[0] @@ -612,6 +613,7 @@ jobs: content = post_sse("/v1/chat/completions", { "messages": [{"role": "user", "content": "What is 123 * 456? Use the python tool to compute it and tell me the number."}], "enable_tools": True, + "permission_mode": "full", "enabled_tools": ["python"], "session_id": "ci-tool-calling-py", "temperature": TEMP, @@ -647,6 +649,7 @@ jobs: content = post_sse("/v1/chat/completions", { "messages": [{"role": "user", "content": "Search the web for 'unsloth ai github' and summarise."}], "enable_tools": True, + "permission_mode": "full", "enabled_tools": ["web_search"], "session_id": "ci-tool-calling-web", "temperature": TEMP, @@ -658,7 +661,7 @@ jobs: print(f"[tools] WARN web_search probe failed (non-blocking): {exc}") # ── 4. Thinking on / off ───────────────────────────────────── - # Studio strips think blocks from message.content for tools-mode + # Unsloth strips think blocks from message.content for tools-mode # responses, so we toggle plain chat (no enable_tools) and look # at the surfaced reasoning_content / message.thinking field. def thinking_call(enable): @@ -676,7 +679,7 @@ jobs: }, timeout = 180) assert status == 200 msg = data["choices"][0]["message"] - # Studio surfaces thinking via reasoning_content (OpenAI + # Unsloth surfaces thinking via reasoning_content (OpenAI # extension). Fall back to inline markers for # robustness across template versions. raw = (msg.get("content") or "") + (msg.get("reasoning_content") or "") @@ -702,7 +705,7 @@ jobs: print(f"[tools] PASS thinking on/off (on={len(on_text)} chars, off={len(off_text)} chars)") PY - - name: Stop Studio + - name: Stop Unsloth if: always() run: | kill "${STUDIO_PID}" 2>/dev/null || true @@ -808,7 +811,7 @@ jobs: path: gguf-cache key: ${{ runner.os }}-gguf-${{ env.GGUF_REPO }}-${{ env.GGUF_FILE }}-${{ env.MMPROJ_FILE }}-v2 - - name: Install Studio (--local, --no-torch) + - name: Install Unsloth (--local, --no-torch) env: GH_TOKEN: ${{ secrets.GITHUB_TOKEN }} # Withheld on PR: this step runs checked-out PR code; public GGUF still downloads. @@ -824,12 +827,12 @@ jobs: - name: Install OpenAI + Anthropic Python SDKs run: pip install 'openai>=1.50' 'anthropic>=0.40' - - name: Reset auth + boot Studio (API-only) + - name: Reset auth + boot Unsloth (API-only) # See Job 2's comment: API-only mode keeps tool_policy=None so # response_format requests aren't routed through the agentic # tool loop. run: | - unsloth studio reset-password + rm -rf ~/.unsloth/studio/auth mkdir -p logs UNSLOTH_API_ONLY=1 unsloth studio -H 127.0.0.1 -p "$STUDIO_PORT" \ > logs/studio.log 2>&1 & @@ -927,13 +930,13 @@ jobs: # llama.cpp's HTTP server supports OpenAI-compatible JSON # mode: `response_format: {"type": "json_object"}` constrains # the model to emit syntactically-valid JSON. We use raw HTTP - # rather than the OpenAI SDK so that the field shape Studio + # rather than the OpenAI SDK so that the field shape Unsloth # forwards to llama-server is unambiguous (the SDK rewrites # response_format depending on which variant it recognises). # We deliberately do NOT pass a strict JSON schema -- on # small Gemma-4 quants the GBNF-from-schema path occasionally # produces empty output, and JSON mode is the surface we care - # about exposing through Studio. + # about exposing through Unsloth. status, data = post("/v1/chat/completions", { "model": "default", "messages": [ @@ -1005,7 +1008,7 @@ jobs: ) # ── 2. OpenAI image_url (data URI base64) ─────────────────── - # 64x64 solid-red PNG. stb_image (used by Studio's image + # 64x64 solid-red PNG. stb_image (used by Unsloth's image # normaliser at routes/inference.py:3410) rejects 4x4 or # smaller PNGs as truncated, so we go up to 64x64 -- still # tiny in token cost. The assertion is loose: any non-empty @@ -1021,11 +1024,11 @@ jobs: # The Mac prebuilt llama.cpp server has a known crash when # processing image inputs alongside the gemma-4-E2B mmproj # (server disconnects mid-completion). This is upstream - # llama.cpp behaviour, not Studio. Wrap both SDK calls in + # llama.cpp behaviour, not Unsloth. Wrap both SDK calls in # try/except so an upstream crash registers as a WARN rather - # than failing the whole job. Studio's contract (OpenAI/ + # than failing the whole job. Unsloth's contract (OpenAI/ # Anthropic image fields are accepted and forwarded) is - # validated by the request body Studio constructs, not by + # validated by the request body Unsloth constructs, not by # whether llama.cpp can decode it on Mac Metal. client = OpenAI(base_url = f"{BASE}/v1", api_key = KEY) try: @@ -1051,14 +1054,14 @@ jobs: except Exception as exc: print( f"[image/openai] WARN image_url SDK call raised: {type(exc).__name__}: " - f"{exc}. Likely upstream llama.cpp Mac+vision crash, NOT a Studio " - f"regression. Studio successfully forwarded the request." + f"{exc}. Likely upstream llama.cpp Mac+vision crash, NOT an Unsloth " + f"regression. Unsloth successfully forwarded the request." ) # ── 3. Anthropic source/base64 image ──────────────────────── - # Two SDK quirks vs. Studio: base_url must NOT include /v1 + # Two SDK quirks vs. Unsloth: base_url must NOT include /v1 # (the SDK appends it itself; otherwise /v1/v1/messages -> 405), - # and Studio's auth is HTTPBearer-only so the SDK's default + # and Unsloth's auth is HTTPBearer-only so the SDK's default # x-api-key header is ignored -- send Authorization: Bearer # via default_headers. anthropic = Anthropic( @@ -1097,11 +1100,11 @@ jobs: print( f"[image/anthropic] WARN anthropic image SDK call raised: " f"{type(exc).__name__}: {exc}. Likely upstream llama.cpp Mac+vision " - f"crash, NOT a Studio regression." + f"crash, NOT an Unsloth regression." ) PY - - name: Stop Studio + - name: Stop Unsloth if: always() run: | kill "${STUDIO_PID}" 2>/dev/null || true diff --git a/.github/workflows/studio-mac-install-matrix.yml b/.github/workflows/studio-mac-install-matrix.yml index 362305cdd4..e990f752d4 100644 --- a/.github/workflows/studio-mac-install-matrix.yml +++ b/.github/workflows/studio-mac-install-matrix.yml @@ -1,7 +1,7 @@ # SPDX-License-Identifier: AGPL-3.0-only # Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. -# Proves Studio's llama.cpp install loads on every supported macOS. The heavy +# Proves Unsloth's llama.cpp install loads on every supported macOS. The heavy # app smokes stay single-OS; this matrix covers the OS-version dimension cheaply # (install.sh + binary-load assert). Regression guard for the macOS-version # selection in studio/install_llama_prebuilt.py. @@ -60,7 +60,7 @@ jobs: with: python-version: '3.12' - - name: Install Studio (--local, --no-torch) + - name: Install Unsloth (--local, --no-torch) env: GH_TOKEN: ${{ secrets.GITHUB_TOKEN }} # Withheld on PR: this step runs checked-out PR code; public GGUF still downloads. diff --git a/.github/workflows/studio-mac-ui-smoke.yml b/.github/workflows/studio-mac-ui-smoke.yml index 20ca247b9f..3bed2fcdff 100644 --- a/.github/workflows/studio-mac-ui-smoke.yml +++ b/.github/workflows/studio-mac-ui-smoke.yml @@ -19,6 +19,7 @@ on: - 'install.sh' - 'pyproject.toml' - 'tests/studio/**' + - '.github/scripts/run-studio-permission-browser.sh' - '.github/workflows/studio-mac-ui-smoke.yml' push: branches: [main, pip] @@ -83,7 +84,7 @@ jobs: path: hf-cache key: ${{ runner.os }}-hf-${{ env.GGUF_REPO }}-${{ env.GGUF_VARIANT }}-v2 - - name: Install Studio (--local, --no-torch) + - name: Install Unsloth (--local, --no-torch) env: GH_TOKEN: ${{ secrets.GITHUB_TOKEN }} # Withheld on PR: this step runs checked-out PR code; public GGUF still downloads. @@ -96,7 +97,7 @@ jobs: - name: Assert llama.cpp loads on this macOS run: bash .github/scripts/assert-llama-loads.sh - - name: Install Playwright + Chromium + - name: Install Playwright browsers # No --with-deps on Mac: that flag installs Linux apt packages. # GitHub-hosted macos-14 ships the system frameworks Chromium # needs already. @@ -112,7 +113,7 @@ jobs: # in-script retry recover from any residual flakes. run: | pip install 'playwright>=1.55,<1.58' - python -m playwright install chromium + python -m playwright install chromium webkit - name: Patch Playwright pipeTransport.js to tolerate malformed JSON # In Playwright 1.55-1.58, pipeTransport.js does @@ -143,9 +144,10 @@ jobs: print(f"pipeTransport.js: patched JSON.parse calls in {path}") PY - - name: Reset auth + boot Studio + - name: Reset auth + boot Unsloth run: | - unsloth studio reset-password + # Wipe (not reset-password): the boot below must re-seed a fresh .bootstrap_password. + rm -rf ~/.unsloth/studio/auth mkdir -p logs UNSLOTH_API_ONLY=1 unsloth studio -H 127.0.0.1 -p "$STUDIO_PORT" \ > logs/studio.log 2>&1 & @@ -188,8 +190,8 @@ jobs: # dies mid-test, (2) Chromium net::ERR_NO_BUFFER_SPACE when the # runner's kernel briefly runs out of socket buffers, and (3) a # goto 'interrupted by another navigation' when the SPA auth - # guard redirects mid-navigation. The retry FULLY resets Studio - # (kill, reset-password, reboot, wait /api/health, re-export + # guard redirects mid-navigation. The retry FULLY resets Unsloth + # (kill, wipe auth, reboot, wait /api/health, re-export # bootstrap pw) before re-running the script. A real test failure # (assertion / timeout) does NOT match any pattern so it bypasses # retry and surfaces immediately. @@ -209,10 +211,10 @@ jobs: || grep -q "ERR_NO_BUFFER_SPACE" logs/playwright_attempt_${attempt}.log \ || grep -q "interrupted by another navigation" logs/playwright_attempt_${attempt}.log; } \ && [ "$attempt" -lt "$max_attempts" ]; then - echo "::warning::Playwright flake on attempt ${attempt}; resetting Studio and retrying..." + echo "::warning::Playwright flake on attempt ${attempt}; resetting Unsloth and retrying..." kill "${STUDIO_PID}" 2>/dev/null || true sleep 2 - unsloth studio reset-password + rm -rf ~/.unsloth/studio/auth UNSLOTH_API_ONLY=1 unsloth studio -H 127.0.0.1 -p "$STUDIO_PORT" \ > "logs/studio_retry_${attempt}.log" 2>&1 & STUDIO_PID=$! @@ -238,15 +240,19 @@ jobs: exit "$rc" done - - name: Stop Studio (chat-ui ends with Shutdown click; this is belt-and-suspenders) + - name: Stop Unsloth (chat-ui ends with Shutdown click; this is belt-and-suspenders) if: always() run: | kill "${STUDIO_PID}" 2>/dev/null || true sleep 2 - - name: Reset auth + boot Studio for extra UI tests (port 18897) + - name: Cross-browser permission controls run: | - unsloth studio reset-password + bash .github/scripts/run-studio-permission-browser.sh 18895 webkit + + - name: Reset auth + boot Unsloth for extra UI tests (port 18897) + run: | + rm -rf ~/.unsloth/studio/auth mkdir -p logs UNSLOTH_API_ONLY=1 unsloth studio -H 127.0.0.1 -p 18897 \ > logs/studio_extra.log 2>&1 & @@ -271,7 +277,7 @@ jobs: echo "STUDIO_EXTRA_OLD_PW=$OLD" >> "$GITHUB_ENV" echo "STUDIO_EXTRA_NEW_PW=$NEW" >> "$GITHUB_ENV" - - name: Drive Compare/Recipes/Export/Studio/Settings with Playwright + - name: Drive Compare/Recipes/Export/Unsloth/Settings with Playwright env: BASE_URL: http://127.0.0.1:18897 STUDIO_OLD_PW: ${{ env.STUDIO_EXTRA_OLD_PW }} @@ -300,10 +306,10 @@ jobs: || grep -q "ERR_NO_BUFFER_SPACE" logs/playwright_extra_attempt_${attempt}.log \ || grep -q "interrupted by another navigation" logs/playwright_extra_attempt_${attempt}.log; } \ && [ "$attempt" -lt "$max_attempts" ]; then - echo "::warning::Playwright flake on attempt ${attempt}; resetting Studio and retrying..." + echo "::warning::Playwright flake on attempt ${attempt}; resetting Unsloth and retrying..." kill "${STUDIO_EXTRA_PID}" 2>/dev/null || true sleep 2 - unsloth studio reset-password + rm -rf ~/.unsloth/studio/auth UNSLOTH_API_ONLY=1 unsloth studio -H 127.0.0.1 -p 18897 \ > "logs/studio_extra_retry_${attempt}.log" 2>&1 & STUDIO_EXTRA_PID=$! @@ -327,7 +333,7 @@ jobs: exit "$rc" done - - name: Stop second Studio + - name: Stop second Unsloth if: always() run: | kill "${STUDIO_EXTRA_PID}" 2>/dev/null || true @@ -343,5 +349,7 @@ jobs: logs/studio_extra.log logs/install.log logs/playwright + logs/playwright-permissions-* logs/playwright_extra + logs/studio-permissions-*.log retention-days: 7 diff --git a/.github/workflows/studio-mac-update-smoke.yml b/.github/workflows/studio-mac-update-smoke.yml index d104306c7e..fe9880f3ca 100644 --- a/.github/workflows/studio-mac-update-smoke.yml +++ b/.github/workflows/studio-mac-update-smoke.yml @@ -4,15 +4,15 @@ # Mac counterpart to studio-update-smoke.yml. Verifies that on a real # Apple Silicon (macos-14, M1) runner: # -# 1. install.sh --local --no-torch installs Studio AND auto-fetches +# 1. install.sh --local --no-torch installs Unsloth AND auto-fetches # the prebuilt llama.cpp Mac binary (llama-bNNNN-bin-macos-arm64 # from ggml-org/llama.cpp). Hitting the source-build fallback is -# treated as an Unsloth bug -- Studio must always pick the +# treated as an Unsloth bug -- Unsloth must always pick the # prebuilt on Mac. # 2. unsloth studio update --local is idempotent. Two consecutive # runs both report "prebuilt up to date and validated", no # source-build fallback. -# 3. The installed Studio still boots and /api/health returns +# 3. The installed Unsloth still boots and /api/health returns # healthy after the update path. name: Mac Studio Update CI @@ -42,7 +42,7 @@ permissions: jobs: update-idempotency: - name: Studio Updating Tests + name: Unsloth Updating Tests runs-on: macos-14 timeout-minutes: 30 steps: @@ -59,7 +59,7 @@ jobs: python-version: '3.12' cache: 'pip' - - name: Install Studio (--local, --no-torch) + - name: Install Unsloth (--local, --no-torch) env: GH_TOKEN: ${{ secrets.GITHUB_TOKEN }} # Withheld on PR: this step runs checked-out PR code; public GGUF still downloads. @@ -106,7 +106,7 @@ jobs: grep -qE "prebuilt up to date and validated|prebuilt installed and validated" logs/update2.log echo "second update was clean" - - name: Boot Studio briefly to confirm the install is still usable + - name: Boot Unsloth briefly to confirm the install is still usable run: | mkdir -p logs UNSLOTH_API_ONLY=1 unsloth studio -H 127.0.0.1 -p 18891 \ @@ -123,13 +123,13 @@ jobs: sleep 1 done if [ -z "$HEALTHY" ]; then - echo "Studio failed to come up after \`update\`" + echo "Unsloth failed to come up after \`update\`" tail -200 logs/studio.log kill "$PID" 2>/dev/null || true exit 1 fi kill "$PID" 2>/dev/null || true - echo "post-update Studio /api/health OK" + echo "post-update Unsloth /api/health OK" - name: Uninstall and verify clean # Round-trip through scripts/uninstall.sh on real macOS. As a side diff --git a/.github/workflows/studio-tauri-smoke.yml b/.github/workflows/studio-tauri-smoke.yml index 018857de68..c6dad07f37 100644 --- a/.github/workflows/studio-tauri-smoke.yml +++ b/.github/workflows/studio-tauri-smoke.yml @@ -12,7 +12,7 @@ # stay in release-desktop.yml (manual `workflow_dispatch`) because they need # code-signing secrets and ~30 min of runner time each. -name: Studio Tauri CI +name: Unsloth Tauri CI on: pull_request: @@ -91,6 +91,16 @@ jobs: npm run build test -f dist/index.html + # The crate carries ~100 unit tests (native_file_dialogs, preflight, + # install, desktop_auth, ...) that nothing ran until now: this workflow + # only ever built. Run them here, where the toolchain and the WebKit dev + # packages are already installed, so a broken assertion fails the PR + # instead of sitting unnoticed. `--no-fail-fast` reports every failing + # test in one run rather than stopping at the first. + - name: Rust unit tests (studio/src-tauri) + working-directory: studio/src-tauri + run: cargo test --no-fail-fast + - name: Tauri debug build (Linux, no bundle, no codesign) # `--debug` + `--no-bundle` keeps this lean: compiles the Rust crate, # confirms the frontend dist is wired into Tauri, but skips the AppImage diff --git a/.github/workflows/studio-ui-smoke.yml b/.github/workflows/studio-ui-smoke.yml index 297a585430..3a0713f301 100644 --- a/.github/workflows/studio-ui-smoke.yml +++ b/.github/workflows/studio-ui-smoke.yml @@ -1,8 +1,8 @@ # SPDX-License-Identifier: AGPL-3.0-only # Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. -# End-to-end Studio chat UI smoke via Playwright + Chromium against a -# headless Linux runner. Boots Studio with the smallest GGUF +# End-to-end Unsloth chat UI smoke via Playwright + Chromium against a +# headless Linux runner. Boots Unsloth with the smallest GGUF # (gemma-3-270m-it UD-Q4_K_XL, ~254 MiB), drives the actual frontend # bundle, and asserts the full bootstrap-password / change-password / # send-message / persist-on-reload journey works end to end. @@ -14,7 +14,7 @@ # frontend-only CI happily pass while the actual user-visible UI is # broken (cf. the 2026.5.1 chat-history release). -name: Studio UI CI +name: Unsloth UI CI on: pull_request: @@ -27,6 +27,7 @@ on: # The Playwright test files themselves -- a PR that ONLY edits # the test must still trigger UI CI. - 'tests/studio/**' + - '.github/scripts/run-studio-permission-browser.sh' - '.github/workflows/studio-ui-smoke.yml' push: branches: [main, pip] @@ -97,7 +98,7 @@ jobs: path: hf-cache key: ${{ runner.os }}-hf-${{ env.GGUF_REPO }}-${{ env.GGUF_VARIANT }}-v2 - - name: Install Studio (--local, --no-torch) + - name: Install Unsloth (--local, --no-torch) env: GH_TOKEN: ${{ secrets.GITHUB_TOKEN }} # Withheld on PR: this step runs checked-out PR code; public GGUF still downloads. @@ -107,17 +108,15 @@ jobs: set -o pipefail bash install.sh --local --no-torch 2>&1 | tee logs/install.log - - name: Install Playwright + Chromium + - name: Install Playwright browsers run: | pip install 'playwright>=1.45' - # --with-deps installs the OS-level runtime libs Chromium - # needs (libnss3, libxkbcommon, etc.). About 30 s on a - # warm runner. - python -m playwright install --with-deps chromium + python -m playwright install --with-deps chromium firefox webkit - - name: Reset auth + boot Studio + - name: Reset auth + boot Unsloth run: | - unsloth studio reset-password + # Wipe (not reset-password): the boot below must re-seed a fresh .bootstrap_password. + rm -rf ~/.unsloth/studio/auth mkdir -p logs UNSLOTH_API_ONLY=1 unsloth studio -H 127.0.0.1 -p "$STUDIO_PORT" \ > logs/studio.log 2>&1 & @@ -147,7 +146,7 @@ jobs: # NEW + NEW2 are generated freshly per CI run via secrets.token_urlsafe # rather than hardcoded. If a workflow gets compromised, the # attacker can't replay a known-good rotated password against - # any future / parallel Studio install -- the rotated value + # any future / parallel Unsloth install -- the rotated value # only ever exists for the lifetime of this single job, masked # in the log via ::add-mask::. run: | @@ -165,31 +164,37 @@ jobs: env: BASE_URL: http://127.0.0.1:18892 # The test file lives in the repo so it can be run locally - # against a freshly-installed Studio (BASE_URL=...; STUDIO_OLD_PW= + # against a freshly-installed Unsloth (BASE_URL=...; STUDIO_OLD_PW= # $(cat ~/.unsloth/studio/auth/.bootstrap_password); python ...). PW_ART_DIR: logs/playwright # Strict mode: in CI a missing button / nav / dialog must # FAIL the test. Locally the test still runs against partial - # Studio installs without STUDIO_UI_STRICT. + # Unsloth installs without STUDIO_UI_STRICT. STUDIO_UI_STRICT: '1' run: | mkdir -p logs/playwright python tests/studio/playwright_chat_ui.py - - name: Stop Studio (chat-ui ends with Shutdown click; this is belt-and-suspenders) + - name: Stop Unsloth (chat-ui ends with Shutdown click; this is belt-and-suspenders) if: always() run: | kill "${STUDIO_PID}" 2>/dev/null || true sleep 2 + - name: Cross-browser permission controls + run: | + bash .github/scripts/run-studio-permission-browser.sh 18893 firefox + bash .github/scripts/run-studio-permission-browser.sh 18893 webkit + bash .github/scripts/run-studio-permission-browser.sh 18893 chromium chrome + # The chat UI test ends by clicking the Shutdown menuitem, which # leaves the server dead. The extra UI test (Compare / Recipes / - # Export / Studio / Settings) needs a fresh Studio, so we boot a + # Export / Unsloth / Settings) needs a fresh Unsloth, so we boot a # second one on a different port. Boot is fast (~3-5s on the # warm install we already did) so this adds little wall time. - - name: Reset auth + boot Studio for extra UI tests (port 18894) + - name: Reset auth + boot Unsloth for extra UI tests (port 18894) run: | - unsloth studio reset-password + rm -rf ~/.unsloth/studio/auth mkdir -p logs UNSLOTH_API_ONLY=1 unsloth studio -H 127.0.0.1 -p 18894 \ > logs/studio_extra.log 2>&1 & @@ -214,7 +219,7 @@ jobs: echo "STUDIO_EXTRA_OLD_PW=$OLD" >> "$GITHUB_ENV" echo "STUDIO_EXTRA_NEW_PW=$NEW" >> "$GITHUB_ENV" - - name: Drive Compare/Recipes/Export/Studio/Settings with Playwright + - name: Drive Compare/Recipes/Export/Unsloth/Settings with Playwright env: BASE_URL: http://127.0.0.1:18894 STUDIO_OLD_PW: ${{ env.STUDIO_EXTRA_OLD_PW }} @@ -227,18 +232,75 @@ jobs: mkdir -p logs/playwright_extra python tests/studio/playwright_extra_ui.py - - name: Stop second Studio + - name: UI font size scaling regression (Playwright) + env: + BASE_URL: http://127.0.0.1:18894 + STUDIO_PW: ${{ env.STUDIO_EXTRA_NEW_PW }} + PW_ART_DIR: logs/playwright_fontscale + run: | + mkdir -p logs/playwright_fontscale + python tests/studio/playwright_ui_font_scale.py + + - name: Stop second Unsloth if: always() run: | kill "${STUDIO_EXTRA_PID}" 2>/dev/null || true sleep 2 - # IME + multilingual paste regression (issue #5318 / PR #5327). - # Third Studio on its own port so a hang here cannot poison the - # earlier UI tests. No GGUF -- the bug surface is the composer. - - name: Reset auth + boot Studio for IME / i18n tests (port 18896) + # Model-picker per-model-config regression (PR #7207 re-land of #6647). + # Fourth Unsloth on its own port; loads the tiny GGUF and drives the + # picker's run-settings surface: Context Length persists across a reload, + # Reset clears the stored override (never pins it), and the infra models + # (RAG embedder + llama.cpp probe) stay hidden from the picker. + - name: Reset auth + boot Unsloth for model-config tests (port 18898) run: | - unsloth studio reset-password + rm -rf ~/.unsloth/studio/auth + mkdir -p logs + UNSLOTH_API_ONLY=1 unsloth studio -H 127.0.0.1 -p 18898 \ + > logs/studio_modelcfg.log 2>&1 & + echo "STUDIO_MODELCFG_PID=$!" >> "$GITHUB_ENV" + + - name: Wait for /api/health on 18898 + run: | + for i in $(seq 1 180); do + if curl -fs "http://127.0.0.1:18898/api/health" > /tmp/health4.json; then + jq -e '.status == "healthy"' /tmp/health4.json && break + fi + sleep 1 + done + jq -e '.status == "healthy"' /tmp/health4.json + + - name: Pass bootstrap pw for model-config test + run: | + NEW="CIModelCfg-$(python -c 'import secrets; print(secrets.token_urlsafe(16))')" + echo "::add-mask::$NEW" + echo "STUDIO_MODELCFG_NEW_PW=$NEW" >> "$GITHUB_ENV" + + - name: Drive model-picker per-model-config with Playwright + env: + BASE_URL: http://127.0.0.1:18898 + STUDIO_NEW_PW: ${{ env.STUDIO_MODELCFG_NEW_PW }} + PW_ART_DIR: logs/playwright_modelcfg + STUDIO_UI_STRICT: '1' + GGUF_REPO: ${{ env.GGUF_REPO }} + GGUF_VARIANT: ${{ env.GGUF_VARIANT }} + STUDIO_MODEL_HINT: gemma-3-270m + run: | + mkdir -p logs/playwright_modelcfg + python tests/studio/playwright_model_config.py + + - name: Stop fourth Unsloth + if: always() + run: | + kill "${STUDIO_MODELCFG_PID}" 2>/dev/null || true + sleep 2 + + # IME + multilingual paste regression (issue #5318 / PR #5327). + # Third Unsloth on its own port so a hang here cannot poison the + # earlier UI tests. No GGUF -- the bug surface is the composer. + - name: Reset auth + boot Unsloth for IME / i18n tests (port 18896) + run: | + rm -rf ~/.unsloth/studio/auth mkdir -p logs UNSLOTH_API_ONLY=1 unsloth studio -H 127.0.0.1 -p 18896 \ > logs/studio_ime.log 2>&1 & @@ -256,7 +318,7 @@ jobs: - name: Pass bootstrap pw for IME / i18n test # IME smoke does the change-password against the bootstrap that - # Studio's frontend injects into the page, so it only needs the + # Unsloth's frontend injects into the page, so it only needs the # NEW password. run: | NEW="CIIme-$(python -c 'import secrets; print(secrets.token_urlsafe(16))')" @@ -273,7 +335,7 @@ jobs: mkdir -p logs/playwright_ime python tests/studio/playwright_chat_ime_i18n.py - - name: Stop third Studio + - name: Stop third Unsloth if: always() run: | kill "${STUDIO_IME_PID}" 2>/dev/null || true @@ -293,10 +355,15 @@ jobs: path: | logs/studio.log logs/studio_extra.log + logs/studio_modelcfg.log logs/studio_ime.log logs/install.log logs/server-logs/ logs/playwright + logs/playwright-permissions-* logs/playwright_extra + logs/playwright_fontscale + logs/playwright_modelcfg logs/playwright_ime + logs/studio-permissions-*.log retention-days: 7 diff --git a/.github/workflows/studio-update-smoke.yml b/.github/workflows/studio-update-smoke.yml index 08a79afacd..047840e41c 100644 --- a/.github/workflows/studio-update-smoke.yml +++ b/.github/workflows/studio-update-smoke.yml @@ -9,7 +9,7 @@ # This catches regressions in setup.sh's update path that the existing # GGUF / wheel jobs would miss because they only invoke install.sh once. -name: Studio Update CI +name: Unsloth Update CI on: pull_request: @@ -36,7 +36,7 @@ permissions: jobs: update-idempotency: - name: Studio Updating Tests + name: Unsloth Updating Tests runs-on: ubuntu-latest timeout-minutes: 15 steps: @@ -63,7 +63,7 @@ jobs: # post-step then fatal-errors with "Cache folder path is # retrieved for pip but doesn't exist on disk". - - name: Install Studio (--local, --no-torch) + - name: Install Unsloth (--local, --no-torch) # Pass the workflow token so the llama.cpp prebuilt installer's # GitHub-API call to list releases isn't rate-limited (60/hr # unauthenticated). Without this, three consecutive install + @@ -122,7 +122,7 @@ jobs: grep -qE "prebuilt up to date and validated|prebuilt installed and validated" logs/update2.log echo "second update was clean" - - name: Boot Studio briefly to confirm the install is still usable + - name: Boot Unsloth briefly to confirm the install is still usable # If `update --local` accidentally broke the venv or wiped the # llama-server binary, the server would fail to start here. run: | @@ -138,13 +138,53 @@ jobs: sleep 1 done if ! jq -e '.status == "healthy"' /tmp/health.json 2>/dev/null; then - echo "Studio failed to come up after `update`" + echo "Unsloth failed to come up after `update`" tail -200 logs/studio.log kill "$PID" 2>/dev/null || true exit 1 fi kill "$PID" 2>/dev/null || true - echo "post-update Studio /api/health OK" + echo "post-update Unsloth /api/health OK" + + - name: A complete install reports itself complete + run: | + set -o pipefail + unsloth studio verify-install + unsloth studio desktop-capabilities --json | tee /tmp/caps.json + jq -e '.studio_install_ok == true' /tmp/caps.json + jq -e '.desktop_manageability_version >= 2' /tmp/caps.json + + - name: An incomplete install must not report itself ready + # An installer killed part-way leaves a working CLI but no studio.txt + # deps, which the old preflight called ManagedReady. The manifest is + # written last, so removing it reproduces that state. + run: | + set -o pipefail + # install.sh's default root, resolved explicitly: `python` on PATH + # here is setup-python's, not the managed venv. + MANIFEST="$HOME/.unsloth/studio/unsloth_studio/unsloth_install_manifest.json" + test -f "$MANIFEST" || { echo "::error::installer never wrote $MANIFEST"; exit 1; } + rm -f "$MANIFEST" + unsloth studio desktop-capabilities --json | tee /tmp/caps_bad.json + jq -e '.studio_install_ok == false' /tmp/caps_bad.json + if unsloth studio verify-install; then + echo "::error::verify-install passed on an install with no manifest" + exit 1 + fi + echo "incomplete install correctly reported not-ready" + + - name: Update repairs an incomplete install + # `--local` bypasses setup.sh's PyPI version compare, so this asserts + # the repair OUTCOME. The non-local fast path the desktop Repair button + # uses is covered by tests/studio/install/test_setup_fast_path_guard.py. + env: + GH_TOKEN: ${{ secrets.GITHUB_TOKEN }} + run: | + set -o pipefail + unsloth studio update --local 2>&1 | tee logs/update_repair.log + unsloth studio verify-install + unsloth studio desktop-capabilities --json | jq -e '.studio_install_ok == true' + echo "update repaired the incomplete install" - name: Uninstall and verify clean # Round-trip the installer through scripts/uninstall.sh: confirms the diff --git a/.github/workflows/studio-windows-api-smoke.yml b/.github/workflows/studio-windows-api-smoke.yml index e9abd2d669..b328939846 100644 --- a/.github/workflows/studio-windows-api-smoke.yml +++ b/.github/workflows/studio-windows-api-smoke.yml @@ -9,7 +9,7 @@ # (Section 6) is Linux-only and short-circuits on non-POSIX; the rest # is platform-portable. -name: Windows Studio API CI +name: Windows Unsloth API CI on: pull_request: @@ -34,7 +34,7 @@ permissions: jobs: api-smoke: - name: Studio API & Auth Tests + name: Unsloth API & Auth Tests runs-on: windows-latest timeout-minutes: 30 defaults: @@ -105,7 +105,7 @@ jobs: # studio-windows-update-smoke.yml for the full rationale -- # creating an empty studio/frontend/dist trips setup.ps1's # mtime-based staleness check into "frontend up to date, skip - # rebuild" and Studio boots with an empty dist directory. + # rebuild" and Unsloth boots with an empty dist directory. # Add-MpPreference accepts paths that do not yet exist. foreach ($p in @( "$env:USERPROFILE\.unsloth", @@ -121,7 +121,7 @@ jobs: } } - - name: Install Studio (--local, --no-torch) + - name: Install Unsloth (--local, --no-torch) shell: pwsh env: GH_TOKEN: ${{ secrets.GITHUB_TOKEN }} @@ -161,7 +161,7 @@ jobs: echo "install.ps1 installed the Windows prebuilt llama.cpp:" cat "$INFO" - - name: Add Studio shim to GITHUB_PATH + - name: Add Unsloth shim to GITHUB_PATH # install.ps1's User-PATH update doesn't propagate to a # running Git Bash session; export the shim dir so the # next `unsloth ...` invocation finds it. @@ -177,9 +177,10 @@ jobs: - name: Install pyjwt for the JWT-expiry forge test run: python -m pip install 'pyjwt>=2.6' - - name: Reset auth + boot Studio (API-only) + - name: Reset auth + boot Unsloth (API-only) run: | - unsloth studio reset-password + # Wipe (not reset-password): the boot below must re-seed a fresh .bootstrap_password. + rm -rf ~/.unsloth/studio/auth mkdir -p logs UNSLOTH_API_ONLY=1 unsloth studio -H 127.0.0.1 -p "$STUDIO_PORT" \ > logs/studio.log 2>&1 & @@ -207,7 +208,7 @@ jobs: echo "STUDIO_NEW_PW=$NEW" >> "$GITHUB_ENV" echo "STUDIO_NEW2_PW=$NEW2" >> "$GITHUB_ENV" - - name: Run Studio API & Auth tests + - name: Run Unsloth API & Auth tests # Do NOT pin STUDIO_AUTH_DIR here. The Mac/Linux mirrors # hardcode runner-specific paths (/Users/runner/..., # /home/runner/...), but on Windows the path is @@ -219,7 +220,7 @@ jobs: BASE_URL: http://127.0.0.1:18895 run: python tests/studio/studio_api_smoke.py - - name: Stop Studio + - name: Stop Unsloth if: always() run: | kill "${STUDIO_PID}" 2>/dev/null || true diff --git a/.github/workflows/studio-windows-inference-smoke.yml b/.github/workflows/studio-windows-inference-smoke.yml index 233292f7a3..d821664327 100644 --- a/.github/workflows/studio-windows-inference-smoke.yml +++ b/.github/workflows/studio-windows-inference-smoke.yml @@ -1,7 +1,7 @@ # SPDX-License-Identifier: AGPL-3.0-only # Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. -# Three end-to-end smoke jobs that boot a freshly-installed Studio and +# Three end-to-end smoke jobs that boot a freshly-installed Unsloth and # exercise the surfaces real users hit through the OpenAI / Anthropic # SDKs and curl, on the FREE windows-latest runner. Each job picks the # smallest model that exercises the behaviour under test, primes @@ -16,7 +16,7 @@ # Qwen3-VL-2B-Instruct UD-IQ2_XXS + mmproj-F16 (~1.4 GiB total). # Within the 14 GB windows-latest SSD budget. -name: Windows Studio GGUF CI +name: Windows Unsloth GGUF CI on: pull_request: @@ -57,7 +57,7 @@ jobs: STUDIO_PORT: '18888' HF_HOME: ${{ github.workspace }}/hf-cache # Force UTF-8 for stdio (Windows defaults to cp1252; hf - # download / Studio CLI print "✓" checkmarks and crash + # download / Unsloth CLI print "✓" checkmarks and crash # otherwise). PYTHONIOENCODING: utf-8 PYTHONUTF8: '1' @@ -160,7 +160,7 @@ jobs: # studio-windows-update-smoke.yml for the full rationale -- # creating an empty studio/frontend/dist trips setup.ps1's # mtime-based staleness check into "frontend up to date, skip - # rebuild" and Studio boots with an empty dist directory. + # rebuild" and Unsloth boots with an empty dist directory. # Add-MpPreference accepts paths that do not yet exist. foreach ($p in @( "$env:USERPROFILE\.unsloth", @@ -176,7 +176,7 @@ jobs: } } - - name: Install Studio (--local, --no-torch) + - name: Install Unsloth (--local, --no-torch) shell: pwsh env: GH_TOKEN: ${{ secrets.GITHUB_TOKEN }} @@ -214,7 +214,7 @@ jobs: echo "install.ps1 installed the Windows prebuilt llama.cpp:" cat "$INFO" - - name: Add Studio shim to GITHUB_PATH + - name: Add Unsloth shim to GITHUB_PATH run: | SHIM_DIR=~/.unsloth/studio/bin if [ ! -f "$SHIM_DIR/unsloth.exe" ]; then @@ -227,9 +227,10 @@ jobs: - name: Install OpenAI + Anthropic Python SDKs run: python -m pip install 'openai>=1.50' 'anthropic>=0.40' - - name: Reset auth + boot Studio (API-only) + - name: Reset auth + boot Unsloth (API-only) run: | - unsloth studio reset-password + # Wipe (not reset-password): the boot below must re-seed a fresh .bootstrap_password. + rm -rf ~/.unsloth/studio/auth mkdir -p logs UNSLOTH_API_ONLY=1 unsloth studio -H 127.0.0.1 -p "$STUDIO_PORT" \ > logs/studio.log 2>&1 & @@ -244,7 +245,7 @@ jobs: fi sleep 1 done - echo "Studio did not become healthy in 180s" + echo "Unsloth did not become healthy in 180s" tail -200 logs/studio.log exit 1 @@ -281,7 +282,7 @@ jobs: # Retry the load step a few times so a transient TCP RST during # llama-server warm-up (Windows runner image churn, # windows-latest -> windows-2025-vs2026 rollout) doesn't fail - # the whole job. The Studio backend's _wait_for_health now + # the whole job. The Unsloth backend's _wait_for_health now # catches httpx.ReadError too; this retry layer covers the # cases the backend can't recover from on its own. LOAD_OK=0 @@ -382,15 +383,15 @@ jobs: print(f"[{label}] OK -- 4 turns, run1 == run2, history grounded") PY - - name: Stop Studio + - name: Stop Unsloth if: always() # Run as cmd so we are not running through the Git Bash shell; # Git Bash on windows-latest has been observed to exit 143 # (SIGTERM) from any inline kill/sleep block, masking a green - # test run. The runner reclaims the Studio child process at + # test run. The runner reclaims the Unsloth child process at # job end either way, so just emit a marker and exit 0. shell: cmd - run: echo Stop Studio (no-op; runner reclaims STUDIO_PID=%STUDIO_PID% at job end) + run: echo Stop Unsloth (no-op; runner reclaims STUDIO_PID=%STUDIO_PID% at job end) - name: Collect llama-server logs if: always() @@ -398,10 +399,10 @@ jobs: # copy must not fail an otherwise-green job. continue-on-error: true shell: bash - # Copy llama-server's own stdout/stderr (teed by Studio under + # Copy llama-server's own stdout/stderr (teed by Unsloth under # ~/.unsloth/studio/logs/llama-server/) into the workspace so # upload-artifact can pick it up. Crucial for diagnosing a - # subprocess crash where Studio's traceback only shows the + # subprocess crash where Unsloth's traceback only shows the # symptom (httpx ReadError) but not the cause. run: | mkdir -p logs/llama-server @@ -439,14 +440,14 @@ jobs: # (211 s on first run; subsequent runs hit the cache, but the # one-time cost recurs every time the cache key bumps). Use # main's `--local-dir gguf-cache` pattern: cache the flat .gguf - # only, pass an absolute path to Studio's /api/inference/load. + # only, pass an absolute path to Unsloth's /api/inference/load. # The OpenAI/Anth and JSON+images jobs still cover the # gguf_variant resolution path. GGUF_REPO: unsloth/Qwen3.5-2B-GGUF GGUF_FILE: Qwen3.5-2B-UD-Q4_K_XL.gguf STUDIO_PORT: '18898' # Force UTF-8 for stdio (Windows defaults to cp1252; hf - # download / Studio CLI print "✓" checkmarks and crash + # download / Unsloth CLI print "✓" checkmarks and crash # otherwise). PYTHONIOENCODING: utf-8 PYTHONUTF8: '1' @@ -507,7 +508,7 @@ jobs: # studio-windows-update-smoke.yml for the full rationale -- # creating an empty studio/frontend/dist trips setup.ps1's # mtime-based staleness check into "frontend up to date, skip - # rebuild" and Studio boots with an empty dist directory. + # rebuild" and Unsloth boots with an empty dist directory. # Add-MpPreference accepts paths that do not yet exist. foreach ($p in @( "$env:USERPROFILE\.unsloth", @@ -523,7 +524,7 @@ jobs: } } - - name: Install Studio (--local, --no-torch) + - name: Install Unsloth (--local, --no-torch) shell: pwsh env: GH_TOKEN: ${{ secrets.GITHUB_TOKEN }} @@ -561,7 +562,7 @@ jobs: echo "install.ps1 installed the Windows prebuilt llama.cpp:" cat "$INFO" - - name: Add Studio shim to GITHUB_PATH + - name: Add Unsloth shim to GITHUB_PATH run: | SHIM_DIR=~/.unsloth/studio/bin if [ ! -f "$SHIM_DIR/unsloth.exe" ]; then @@ -571,9 +572,9 @@ jobs: fi cygpath -w "$SHIM_DIR" >> "$GITHUB_PATH" - - name: Reset auth + boot Studio (API-only, default tool policy) + - name: Reset auth + boot Unsloth (API-only, default tool policy) run: | - unsloth studio reset-password + rm -rf ~/.unsloth/studio/auth mkdir -p logs UNSLOTH_API_ONLY=1 unsloth studio -H 127.0.0.1 -p "$STUDIO_PORT" \ > logs/studio.log 2>&1 & @@ -607,7 +608,7 @@ jobs: # raw string, but we cannot embed `\a` etc. in JSON without # JSON-string-escaping every backslash. Replace `\` with `/` # via bash parameter expansion -- pathlib.Path on Windows - # accepts forward slashes natively, so Studio's loader sees + # accepts forward slashes natively, so Unsloth's loader sees # a normal path. GGUF_PATH="${GITHUB_WORKSPACE//\\//}/gguf-cache/${GGUF_FILE}" ls -lh "$GGUF_PATH" @@ -680,7 +681,7 @@ jobs: def post_sse(path, body, *, timeout = 600, retries = 1, soft = False): # The server-side agentic loop always answers over SSE. A # shared CI runner can stall the stream transport (the - # connection opening, or a mid-stream read) even when Studio + # connection opening, or a mid-stream read) even when Unsloth # is healthy, so harden the read three ways: # * retry a transport stall once with a fresh request, # capped at 300s (a healthy server answers a retry @@ -791,6 +792,7 @@ jobs: content = post_sse("/v1/chat/completions", { "messages": [{"role": "user", "content": "What is 123 * 456? Use the python tool to compute it and tell me the number."}], "enable_tools": True, + "permission_mode": "full", "enabled_tools": ["python"], "session_id": "ci-tool-calling-py", "temperature": TEMP, @@ -816,6 +818,7 @@ jobs: content = post_sse("/v1/chat/completions", { "messages": [{"role": "user", "content": "Use the terminal tool to run `echo hello-bash-tool` and tell me the exact output."}], "enable_tools": True, + "permission_mode": "full", "enabled_tools": ["terminal"], "session_id": "ci-tool-calling-bash", "temperature": TEMP, @@ -840,6 +843,7 @@ jobs: content = post_sse("/v1/chat/completions", { "messages": [{"role": "user", "content": "Search the web for 'unsloth ai github' and summarise."}], "enable_tools": True, + "permission_mode": "full", "enabled_tools": ["web_search"], "session_id": "ci-tool-calling-web", "temperature": TEMP, @@ -879,15 +883,15 @@ jobs: print(f"[tools] PASS thinking on/off (on={len(on_text)} chars, off={len(off_text)} chars)") PY - - name: Stop Studio + - name: Stop Unsloth if: always() # Run as cmd so we are not running through the Git Bash shell; # Git Bash on windows-latest has been observed to exit 143 # (SIGTERM) from any inline kill/sleep block, masking a green - # test run. The runner reclaims the Studio child process at + # test run. The runner reclaims the Unsloth child process at # job end either way, so just emit a marker and exit 0. shell: cmd - run: echo Stop Studio (no-op; runner reclaims STUDIO_PID=%STUDIO_PID% at job end) + run: echo Stop Unsloth (no-op; runner reclaims STUDIO_PID=%STUDIO_PID% at job end) - name: Collect llama-server logs if: always() @@ -895,10 +899,10 @@ jobs: # copy must not fail an otherwise-green job. continue-on-error: true shell: bash - # Copy llama-server's own stdout/stderr (teed by Studio under + # Copy llama-server's own stdout/stderr (teed by Unsloth under # ~/.unsloth/studio/logs/llama-server/) into the workspace so # upload-artifact can pick it up. Crucial for diagnosing a - # subprocess crash where Studio's traceback only shows the + # subprocess crash where Unsloth's traceback only shows the # symptom (httpx ReadError) but not the cause. run: | mkdir -p logs/llama-server @@ -936,7 +940,7 @@ jobs: STUDIO_PORT: '18899' HF_HOME: ${{ github.workspace }}/hf-cache # Force UTF-8 for stdio (Windows defaults to cp1252; hf - # download / Studio CLI print "✓" checkmarks and crash + # download / Unsloth CLI print "✓" checkmarks and crash # otherwise). PYTHONIOENCODING: utf-8 PYTHONUTF8: '1' @@ -1002,7 +1006,7 @@ jobs: # studio-windows-update-smoke.yml for the full rationale -- # creating an empty studio/frontend/dist trips setup.ps1's # mtime-based staleness check into "frontend up to date, skip - # rebuild" and Studio boots with an empty dist directory. + # rebuild" and Unsloth boots with an empty dist directory. # Add-MpPreference accepts paths that do not yet exist. foreach ($p in @( "$env:USERPROFILE\.unsloth", @@ -1018,7 +1022,7 @@ jobs: } } - - name: Install Studio (--local, --no-torch) + - name: Install Unsloth (--local, --no-torch) shell: pwsh env: GH_TOKEN: ${{ secrets.GITHUB_TOKEN }} @@ -1056,7 +1060,7 @@ jobs: echo "install.ps1 installed the Windows prebuilt llama.cpp:" cat "$INFO" - - name: Add Studio shim to GITHUB_PATH + - name: Add Unsloth shim to GITHUB_PATH run: | SHIM_DIR=~/.unsloth/studio/bin if [ ! -f "$SHIM_DIR/unsloth.exe" ]; then @@ -1069,9 +1073,9 @@ jobs: - name: Install OpenAI + Anthropic Python SDKs run: python -m pip install 'openai>=1.50' 'anthropic>=0.40' - - name: Reset auth + boot Studio (API-only) + - name: Reset auth + boot Unsloth (API-only) run: | - unsloth studio reset-password + rm -rf ~/.unsloth/studio/auth mkdir -p logs UNSLOTH_API_ONLY=1 unsloth studio -H 127.0.0.1 -p "$STUDIO_PORT" \ > logs/studio.log 2>&1 & @@ -1259,7 +1263,7 @@ jobs: except Exception as exc: print( f"[image/openai] WARN image_url SDK call raised: {type(exc).__name__}: " - f"{exc}. Studio successfully forwarded the request; failure here is " + f"{exc}. Unsloth successfully forwarded the request; failure here is " f"upstream llama.cpp vision behaviour." ) @@ -1300,19 +1304,19 @@ jobs: print( f"[image/anthropic] WARN anthropic image SDK call raised: " f"{type(exc).__name__}: {exc}. Likely upstream llama.cpp vision " - f"behaviour, NOT a Studio regression." + f"behaviour, NOT an Unsloth regression." ) PY - - name: Stop Studio + - name: Stop Unsloth if: always() # Run as cmd so we are not running through the Git Bash shell; # Git Bash on windows-latest has been observed to exit 143 # (SIGTERM) from any inline kill/sleep block, masking a green - # test run. The runner reclaims the Studio child process at + # test run. The runner reclaims the Unsloth child process at # job end either way, so just emit a marker and exit 0. shell: cmd - run: echo Stop Studio (no-op; runner reclaims STUDIO_PID=%STUDIO_PID% at job end) + run: echo Stop Unsloth (no-op; runner reclaims STUDIO_PID=%STUDIO_PID% at job end) - name: Collect llama-server logs if: always() @@ -1320,10 +1324,10 @@ jobs: # copy must not fail an otherwise-green job. continue-on-error: true shell: bash - # Copy llama-server's own stdout/stderr (teed by Studio under + # Copy llama-server's own stdout/stderr (teed by Unsloth under # ~/.unsloth/studio/logs/llama-server/) into the workspace so # upload-artifact can pick it up. Crucial for diagnosing a - # subprocess crash where Studio's traceback only shows the + # subprocess crash where Unsloth's traceback only shows the # symptom (httpx ReadError) but not the cause. run: | mkdir -p logs/llama-server @@ -1345,7 +1349,7 @@ jobs: # ── folded from studio-windows-no-vs-smoke.yml: install + run with no Visual Studio ── no-vs-cpu: - name: Studio install + inference without Visual Studio + name: Unsloth install + inference without Visual Studio runs-on: windows-latest timeout-minutes: 35 defaults: @@ -1499,7 +1503,7 @@ jobs: python -m pip install torch --index-url https://download.pytorch.org/whl/cpu --extra-index-url https://pypi.org/simple python -c "import torch; print('torch', torch.__version__, 'cuda?', torch.cuda.is_available())" - - name: Install Studio (--local, --no-torch) with no build tools present + - name: Install Unsloth (--local, --no-torch) with no build tools present shell: pwsh env: GH_TOKEN: ${{ secrets.GITHUB_TOKEN }} @@ -1535,15 +1539,15 @@ jobs: echo "Prebuilt installed with no build tools:" cat "$INFO" - - name: Add Studio shim to GITHUB_PATH + - name: Add Unsloth shim to GITHUB_PATH run: | SHIM_DIR=~/.unsloth/studio/bin [ -f "$SHIM_DIR/unsloth.exe" ] || { echo "::error::unsloth.exe shim not found"; ls -la ~/.unsloth/studio/ || true; exit 1; } cygpath -w "$SHIM_DIR" >> "$GITHUB_PATH" - - name: Reset auth + boot Studio (API-only) + - name: Reset auth + boot Unsloth (API-only) run: | - unsloth studio reset-password + rm -rf ~/.unsloth/studio/auth mkdir -p logs UNSLOTH_API_ONLY=1 unsloth studio -H 127.0.0.1 -p "$STUDIO_PORT" \ > logs/studio.log 2>&1 & @@ -1610,10 +1614,10 @@ jobs: } Remove-Item -LiteralPath $root -Recurse -Force -ErrorAction SilentlyContinue - - name: Stop Studio + - name: Stop Unsloth if: always() shell: cmd - run: echo Stop Studio (no-op; runner reclaims STUDIO_PID=%STUDIO_PID% at job end) + run: echo Stop Unsloth (no-op; runner reclaims STUDIO_PID=%STUDIO_PID% at job end) - name: Collect llama-server logs if: always() @@ -1885,8 +1889,11 @@ jobs: # (step/substep -> Write-StudioStdoutMirror / Get-StudioAnsi). $script:StudioVtOk = $false $script:UnslothVerbose = $false + # Get-HostMachineArch is reached only on the absent path, where + # Test-VCRedistInstalled consults it before trusting the System32 DLL, so + # part A passes without it and only the clean-box part fails. foreach ($fn in @('Get-StudioAnsi', 'Write-StudioStdoutMirror', 'step', 'substep', - 'Invoke-SetupCommand', 'Refresh-Environment', + 'Invoke-SetupCommand', 'Refresh-Environment', 'Get-HostMachineArch', 'Test-VCRedistInstalled', 'Ensure-VCRedist')) { $src = Get-FunctionSource -Path $setup -Name $fn if (-not $src) { throw "Function '$fn' not found in setup.ps1" } diff --git a/.github/workflows/studio-windows-ui-smoke.yml b/.github/workflows/studio-windows-ui-smoke.yml index 405309916a..d23cca323f 100644 --- a/.github/workflows/studio-windows-ui-smoke.yml +++ b/.github/workflows/studio-windows-ui-smoke.yml @@ -4,11 +4,11 @@ # Windows counterpart to studio-ui-smoke.yml / studio-mac-ui-smoke.yml. # Same Playwright + Chromium end-to-end chat UI flow + extra UI flow, # but on the FREE windows-latest runner so we catch Windows-specific -# regressions in the install path (install.ps1), the Studio CLI's +# regressions in the install path (install.ps1), the Unsloth CLI's # Windows process-management branches, and the llama.cpp prebuilt's # Windows HTTP layer. -name: Windows Studio UI CI +name: Windows Unsloth UI CI on: pull_request: @@ -19,6 +19,7 @@ on: - 'install.ps1' - 'pyproject.toml' - 'tests/studio/**' + - '.github/scripts/run-studio-permission-browser.sh' - '.github/workflows/studio-windows-ui-smoke.yml' push: branches: [main, pip] @@ -49,7 +50,7 @@ jobs: GGUF_FILE: gemma-3-270m-it-UD-Q4_K_XL.gguf STUDIO_PORT: '18896' HF_HOME: ${{ github.workspace }}/hf-cache - # Force UTF-8 for stdio so Python tools (hf download, Studio + # Force UTF-8 for stdio so Python tools (hf download, Unsloth # CLI, etc.) can print Unicode characters like the success # checkmark "✓". Windows defaults to cp1252 / charmap and # any tool that prints "OK ✓" hits a UnicodeEncodeError. @@ -121,7 +122,7 @@ jobs: # studio-windows-update-smoke.yml for the full rationale -- # creating an empty studio/frontend/dist trips setup.ps1's # mtime-based staleness check into "frontend up to date, skip - # rebuild" and Studio boots with an empty dist directory. + # rebuild" and Unsloth boots with an empty dist directory. # Add-MpPreference accepts paths that do not yet exist. foreach ($p in @( "$env:USERPROFILE\.unsloth", @@ -148,7 +149,7 @@ jobs: Set-Content -LiteralPath (Join-Path $appDir 'launch-studio.vbs') -Value 'WScript.Echo "legacy"' -Encoding Unicode Write-Host "seeded legacy launch-studio.vbs at $appDir" - - name: Install Studio (--local, --no-torch) + - name: Install Unsloth (--local, --no-torch) # install.ps1 is the supported Windows installer. install.sh # has no Windows branch (apt-get / brew calls). The PS1 # script's `Install-UnslothStudio @args` line at the bottom @@ -205,7 +206,7 @@ jobs: echo "install.ps1 installed the Windows prebuilt llama.cpp:" cat "$INFO" - - name: Assert Studio launcher chain (no VBS, hidden PowerShell shortcut) + - name: Assert Unsloth launcher chain (no VBS, hidden PowerShell shortcut) # The shortcut launch path is otherwise untested here (the steps below # boot `unsloth studio` directly). Guard against re-introducing the VBS # that tripped Kaspersky HEUR:Trojan.VBS.Agent.gen and against the .lnk @@ -234,7 +235,7 @@ jobs: } Write-Host "launcher chain OK (no VBS; hidden powershell over launch-studio.ps1)" - - name: Launch Studio via the shortcut and assert health + - name: Launch Unsloth via the shortcut and assert health # Run the exact command the .lnk stores (hidden PowerShell over # launch-studio.ps1) and confirm it brings the backend up. This is the # only step that proves the shortcut launch is not silently broken. @@ -265,10 +266,10 @@ jobs: $owner = (Get-NetTCPConnection -LocalPort $foundPort -State Listen -ErrorAction Stop | Select-Object -First 1).OwningProcess if ($owner) { taskkill /PID $owner /T /F 2>$null | Out-Null } } catch {} - if (-not $foundPort) { throw "Studio did not become healthy when launched via the shortcut" } - Write-Host "Studio healthy on port $foundPort (launched via the shortcut)" + if (-not $foundPort) { throw "Unsloth did not become healthy when launched via the shortcut" } + Write-Host "Unsloth healthy on port $foundPort (launched via the shortcut)" - - name: Add Studio shim to GITHUB_PATH + - name: Add Unsloth shim to GITHUB_PATH # install.ps1 puts unsloth.exe at $StudioHome\bin\unsloth.exe # and adds that dir to the User PATH via the Windows registry. # Registry-level PATH updates don't propagate to a running @@ -284,7 +285,7 @@ jobs: fi # GITHUB_PATH wants Windows-style paths; convert via cygpath. cygpath -w "$SHIM_DIR" >> "$GITHUB_PATH" - echo "Added Studio shim dir to PATH: $(cygpath -w "$SHIM_DIR")" + echo "Added Unsloth shim dir to PATH: $(cygpath -w "$SHIM_DIR")" - name: Install Playwright + Chromium # No --with-deps on Windows: that flag installs Linux apt @@ -294,9 +295,10 @@ jobs: python -m pip install 'playwright>=1.45' python -m playwright install chromium - - name: Reset auth + boot Studio + - name: Reset auth + boot Unsloth run: | - unsloth studio reset-password + # Wipe (not reset-password): the boot below must re-seed a fresh .bootstrap_password. + rm -rf ~/.unsloth/studio/auth mkdir -p logs UNSLOTH_API_ONLY=1 unsloth studio -H 127.0.0.1 -p "$STUDIO_PORT" \ > logs/studio.log 2>&1 & @@ -339,15 +341,19 @@ jobs: mkdir -p logs/playwright python tests/studio/playwright_chat_ui.py - - name: Stop Studio (chat-ui ends with Shutdown click; this is belt-and-suspenders) + - name: Stop Unsloth (chat-ui ends with Shutdown click; this is belt-and-suspenders) if: always() run: | kill "${STUDIO_PID}" 2>/dev/null || true sleep 2 - - name: Reset auth + boot Studio for extra UI tests (port 18897) + - name: Edge permission controls run: | - unsloth studio reset-password + bash .github/scripts/run-studio-permission-browser.sh 18895 chromium msedge + + - name: Reset auth + boot Unsloth for extra UI tests (port 18897) + run: | + rm -rf ~/.unsloth/studio/auth mkdir -p logs UNSLOTH_API_ONLY=1 unsloth studio -H 127.0.0.1 -p 18897 \ > logs/studio_extra.log 2>&1 & @@ -372,7 +378,7 @@ jobs: echo "STUDIO_EXTRA_OLD_PW=$OLD" >> "$GITHUB_ENV" echo "STUDIO_EXTRA_NEW_PW=$NEW" >> "$GITHUB_ENV" - - name: Drive Compare/Recipes/Export/Studio/Settings with Playwright + - name: Drive Compare/Recipes/Export/Unsloth/Settings with Playwright env: BASE_URL: http://127.0.0.1:18897 STUDIO_OLD_PW: ${{ env.STUDIO_EXTRA_OLD_PW }} @@ -386,7 +392,7 @@ jobs: mkdir -p logs/playwright_extra python tests/studio/playwright_extra_ui.py - - name: Stop second Studio + - name: Stop second Unsloth if: always() run: | kill "${STUDIO_EXTRA_PID}" 2>/dev/null || true @@ -402,5 +408,7 @@ jobs: logs/studio_extra.log logs/install.log logs/playwright + logs/playwright-permissions-* logs/playwright_extra + logs/studio-permissions-*.log retention-days: 7 diff --git a/.github/workflows/studio-windows-update-smoke.yml b/.github/workflows/studio-windows-update-smoke.yml index 5b92f1a3e0..0dcc828e6b 100644 --- a/.github/workflows/studio-windows-update-smoke.yml +++ b/.github/workflows/studio-windows-update-smoke.yml @@ -5,19 +5,19 @@ # studio-mac-update-smoke.yml. Verifies that on the FREE # windows-latest runner: # -# 1. install.ps1 --local --no-torch installs Studio AND auto-fetches +# 1. install.ps1 --local --no-torch installs Unsloth AND auto-fetches # the prebuilt llama.cpp Windows binary (app--windows-x64-cpu # from unslothai/llama.cpp). Hitting the source-build fallback is -# treated as an Unsloth bug -- Studio must always pick the +# treated as an Unsloth bug -- Unsloth must always pick the # prebuilt on Windows. # 2. unsloth studio update --local is idempotent. Two consecutive # runs both report "prebuilt up to date and validated", no # source-build fallback. The CLI's _find_setup_script picks # setup.ps1 on Windows automatically. -# 3. The installed Studio still boots and /api/health returns +# 3. The installed Unsloth still boots and /api/health returns # healthy after the update path. -name: Windows Studio Update CI +name: Windows Unsloth Update CI on: pull_request: @@ -45,7 +45,7 @@ permissions: jobs: update-idempotency: - name: Studio Updating Tests + name: Unsloth Updating Tests runs-on: windows-latest timeout-minutes: 30 defaults: @@ -53,7 +53,7 @@ jobs: shell: bash env: # Force UTF-8 for stdio (Windows defaults to cp1252; hf - # download / Studio CLI print "✓" checkmarks and crash + # download / Unsloth CLI print "✓" checkmarks and crash # otherwise). PYTHONIOENCODING: utf-8 PYTHONUTF8: '1' @@ -90,7 +90,7 @@ jobs: # reuses the existing Node with no download. # # (2) Defender. windows-latest's real-time scan opens / hashes - # every file Studio writes during install (Vite output = + # every file Unsloth writes during install (Vite output = # thousands of small chunks, uv pip = wheel-extraction = # thousands of small files). The latency dominates the # 200 s frontend build and the 90 s deps install. Adding @@ -109,7 +109,7 @@ jobs: # setup.ps1 line 1281-1296's mtime-based "is the frontend # stale?" check into "up to date, skip rebuild", because the # newly-created dist's mtime is younger than every source - # file. Studio then boots with an empty dist and 500s on + # file. Unsloth then boots with an empty dist and 500s on # GET / with FileNotFoundError: dist\index.html. See run # 25546676715 / job 74984469728. # Add-MpPreference accepts paths that do not yet exist; the @@ -129,7 +129,7 @@ jobs: } } - - name: Install Studio (--local, --no-torch) + - name: Install Unsloth (--local, --no-torch) shell: pwsh env: GH_TOKEN: ${{ secrets.GITHUB_TOKEN }} @@ -168,7 +168,7 @@ jobs: echo "install.ps1 installed the Windows prebuilt llama.cpp:" cat "$INFO" - - name: Add Studio shim to GITHUB_PATH + - name: Add Unsloth shim to GITHUB_PATH run: | SHIM_DIR=~/.unsloth/studio/bin if [ ! -f "$SHIM_DIR/unsloth.exe" ]; then @@ -198,6 +198,31 @@ jobs: fi echo "update path took the prebuilt fast path" + - name: Update must keep the --no-torch install GGUF-only + run: | + # `unsloth studio update` exports no UNSLOTH_NO_TORCH, so setup.ps1 has + # to recover the mode from the install manifest. Without that it reads + # the missing torch as a stale venv and tries to delete the venv it is + # running out of, and the shared dependency pass pulls torch back in. + # The skip line only prints when the dependency pass actually runs, so + # don't demand it if the fast path short-circuited that pass. + if grep -q "running ordered dependency installation" logs/update.log \ + && ! grep -q "skipping direct PyTorch and Triton installation (no-torch mode)" logs/update.log; then + echo "::error::studio update left no-torch mode; it would reinstall PyTorch." + grep -iE "no-torch|stale venv|PyTorch" logs/update.log | tail -40 + exit 1 + fi + PY="$HOME/.unsloth/studio/unsloth_studio/Scripts/python.exe" + if [ ! -f "$PY" ]; then + echo "::error::studio venv interpreter missing at $PY" + exit 1 + fi + if "$PY" -c "import torch" 2>/dev/null; then + echo "::error::torch was reinstalled into the --no-torch venv." + exit 1 + fi + echo "update preserved no-torch mode" + - name: Second update must also be a no-op env: GH_TOKEN: ${{ secrets.GITHUB_TOKEN }} @@ -212,7 +237,7 @@ jobs: grep -qE "prebuilt up to date and validated|prebuilt installed and validated" logs/update2.log echo "second update was clean" - - name: Boot Studio briefly to confirm the install is still usable + - name: Boot Unsloth briefly to confirm the install is still usable run: | mkdir -p logs UNSLOTH_API_ONLY=1 unsloth studio -H 127.0.0.1 -p 18891 \ @@ -239,13 +264,13 @@ jobs: sleep 1 done if [ -z "$HEALTHY" ]; then - echo "Studio failed to come up after \`update\`" + echo "Unsloth failed to come up after \`update\`" tail -200 logs/studio.log kill "$PID" 2>/dev/null || true exit 1 fi kill "$PID" 2>/dev/null || true - echo "post-update Studio /api/health OK" + echo "post-update Unsloth /api/health OK" - name: Uninstall and verify clean # Round-trip through scripts/uninstall.ps1 against the default diff --git a/.github/workflows/wheel-smoke.yml b/.github/workflows/wheel-smoke.yml index 3de3c33ca2..f7a7511616 100644 --- a/.github/workflows/wheel-smoke.yml +++ b/.github/workflows/wheel-smoke.yml @@ -3,7 +3,7 @@ # Builds the PyPI wheel from the PR branch, then verifies the built wheel # actually contains what we expect to ship and does NOT contain the broken -# Studio bundle that 2026.5.1 published. This is the single workflow that +# Unsloth bundle that 2026.5.1 published. This is the single workflow that # would have blocked the 2026.5.1 release before twine upload. # # Verified locally end-to-end against this branch: @@ -12,7 +12,7 @@ # lockfile shipped, frontend dist shipped, # no node_modules in wheel, no bun.lock in wheel, # main bundle has unstable_Provider hits=1 (assistant-ui internals only). -# - Studio backend imports cleanly from the installed wheel with the +# - Unsloth backend imports cleanly from the installed wheel with the # lightweight dep set below. name: Wheel CI @@ -101,7 +101,7 @@ jobs: hits = data.count("unstable_Provider:") print(f"main bundle: {js[0]}") print(f"unstable_Provider hits: {hits} (>=4 indicates 2026.5.1 regression)") - checks["bundle has no Studio unstable_Provider call site"] = (hits < 4) + checks["bundle has no Unsloth unstable_Provider call site"] = (hits < 4) print() for k, v in checks.items(): @@ -109,7 +109,7 @@ jobs: sys.exit(0 if all(checks.values()) else 1) PY - - name: Studio backend import smoke + - name: Unsloth backend import smoke # Imports `studio.backend.main:app` from the freshly-installed wheel in # a clean venv. This catches the class of bug that 2026.5.1 shipped with: # frontend dist missing, package-lock.json missing, or the wheel's Python @@ -125,7 +125,32 @@ jobs: /tmp/v/bin/pip install --no-deps dist/unsloth-*.whl # Run from /tmp so Python imports the installed package, not the source tree. cd /tmp - /tmp/v/bin/python -c "from studio.backend.main import app; print('Studio backend OK:', app.title)" + /tmp/v/bin/python -c "from studio.backend.main import app; print('Unsloth backend OK:', app.title)" + + - name: CLI without the Studio stack guides instead of tracebacking + # The smoke above installs studio.txt first, so it cannot catch a wheel + # that ships studio/ without declaring what it imports (#4701, #5260, + # #7147). Drop only structlog to reuse that venv without a re-download. + run: | + set -eu + /tmp/v/bin/pip uninstall -y structlog >/dev/null + cd /tmp + status=0 + for args in "export ./nope ./out" "list-checkpoints"; do + echo "--- unsloth $args" + out=$(/tmp/v/bin/unsloth $args 2>&1 || true) + printf '%s\n' "$out" + case "$out" in + *Traceback*) + echo "FAIL: raw traceback instead of guidance"; status=1 ;; + esac + case "$out" in + *'unsloth studio update'*) ;; + *) echo "FAIL: no remediation in the message"; status=1 ;; + esac + done + /tmp/v/bin/pip install -q structlog >/dev/null + exit "$status" - name: Upload wheel on failure if: failure() diff --git a/.gitignore b/.gitignore index 39ca2226ca..fa6997cb06 100644 --- a/.gitignore +++ b/.gitignore @@ -208,6 +208,9 @@ tmp/ **/node_modules/ auth.db +# Packaging snapshot of the root CHANGELOG.md (written by build.sh) +studio/CHANGELOG.md + # Tauri local build/generated output studio/src-tauri/target/ studio/src-tauri/gen/ @@ -238,4 +241,5 @@ package-lock.json !studio/package-lock.json llama.cpp/ # Stray "~" dir some tools create from a literal ~ TMPDIR; never part of the repo. -/~/ +~/ +/temp/ diff --git a/CHANGELOG.md b/CHANGELOG.md new file mode 100644 index 0000000000..241e013cea --- /dev/null +++ b/CHANGELOG.md @@ -0,0 +1,88 @@ +# Changelog + +Release notes for Unsloth and Unsloth Studio. + +Unsloth Studio reads this file to show release notes inside the "New Unsloth +version" update popup. Edit it here and the popup picks the change up on the +next update check, with no release or rebuild required. + +## Format + +Every release is a level-2 heading whose first token is the version, optionally +followed by a date: + +```md +## 2026.7.6 - 2026-07-22 +``` + +`## [2026.7.6] - 2026-07-22` and `## v2026.7.6` also work. Everything under a +heading, up to the next level-2 heading, is that release's notes and renders as +Markdown in the popup. + +Notes are matched to one exact version. When Studio offers an update to +`2026.7.6` it renders the `2026.7.6` section and nothing else. If that section +is missing, the popup links out to the online changelog rather than showing +notes from an unrelated release, so a new version needs its own section here +before its notes can appear. + +Keep the newest release at the top. Lead each bullet with the change itself: +the collapsed popup highlights the first sentence and dims the rest. +`## Unreleased` is ignored by the popup, so it is safe to stage notes there and +rename the heading at release time. + + + +## Unreleased + +## 2026.7.5 + +### What's Changed + +- AMD support is here. Train, run RL, chat with and deploy 500+ models on + Radeon, Instinct, Ryzen and data center GPUs across Windows, WSL and Linux, + up to 2x faster with 70% less VRAM and no accuracy loss. +- Intel XPU support lands in Studio, so Arc and Data Center GPUs run chat and + training alongside the NVIDIA, AMD and Apple paths. +- Local speech to text dictation runs fully offline, with slim Whisper bundles + and a picker for custom models. +- DoRA training is available in Studio, selectable next to LoRA and full + fine-tuning in the training tab. +- The update popup previews release notes inline, pulled from this file and + matched to the exact version being offered. + +### AMD, 23 July update + +Our AMD collaboration, custom Triton kernels and math algorithms bring local +training and inference to AMD hardware. The 23 July update builds on the +[AMD release](https://github.com/unslothai/unsloth/releases/tag/v0.1.501-beta): + +- RDNA2 and Gorgon Halo are supported, and the installer no longer fails to + detect GPUs on Strix Halo and other AMD cards. +- RDNA4 handling is better, and HIP and ROCm failures are caught and fixed + automatically instead of stopping the install. +- Unified memory safetensors loading is 2x faster, with much faster gradient + checkpointing on unified memory devices. +- Voice dictation through whisper.cpp has preliminary support. +- Rollback environments left by installs no longer eat 5GB of disk. They are + cleaned up automatically. + +Optimized ROCm builds cover GGUF and safetensors inference, and ROCm +compatibility is improved for MI300X and MI325X. Full guide: +[unsloth.ai/docs/basics/amd](https://unsloth.ai/docs/basics/amd). + +### Running larger models + +- Automatic GPU placement, or pick exactly which GPUs and layers to use. +- Move MoE expert layers into system memory so larger models fit. +- Split a model across several GPUs, or use tensor parallelism. +- Hardware settings are saved per model and quant. + +### Also in this release + +- Remote access with `unsloth studio --secure` over free HTTPS via Cloudflare. +- Web search reads PDF papers and manuals, and parallel tool calls, reasoning + output and tool retries are more reliable. +- The model download location is configurable, so weights can live on a second + drive instead of the default cache. +- Stalled Hugging Face XET downloads retry over standard HTTP, and existing + GGUF files are reused instead of downloaded again. diff --git a/MANIFEST.in b/MANIFEST.in new file mode 100644 index 0000000000..7bce036343 --- /dev/null +++ b/MANIFEST.in @@ -0,0 +1,2 @@ +include _changelog_build.py +include CHANGELOG.md diff --git a/README.md b/README.md index ef45b91430..e0fc8ee44c 100644 --- a/README.md +++ b/README.md @@ -11,6 +11,7 @@ Unsloth Studio lets you run and train models locally.

Features • + NewsQuickstartNotebooksDocumentation @@ -47,15 +48,51 @@ Unsloth Studio (Beta) lets you run and train text, [audio](https://unsloth.ai/do * [Auto set inference settings](https://unsloth.ai/docs/new/studio/chat#auto-parameter-tuning) and customize chat templates. * We work directly with teams behind [gpt-oss](https://docs.unsloth.ai/new/gpt-oss-how-to-run-and-fine-tune#unsloth-fixes-for-gpt-oss), [Qwen3](https://www.reddit.com/r/LocalLLaMA/comments/1kaodxu/qwen3_unsloth_dynamic_ggufs_128k_context_bug_fixes/), [Llama 4](https://github.com/ggml-org/llama.cpp/pull/12889), [Mistral](https://huggingface.co/mistralai/Mistral-Medium-3.5-128B/discussions/18), [Gemma 1-3](https://news.ycombinator.com/item?id=39671146), and [Phi-4](https://unsloth.ai/blog/phi4), where we’ve fixed bugs that improve model accuracy. * Chat with images, audio, PDFs, code, DOCX and more. [Connect API providers](https://unsloth.ai/docs/integrations/connections) (OpenAI, Anthropic) or servers (vLLM, Ollama). +* [**Compare any two models**](https://unsloth.ai/docs/new/studio/chat#model-arena) side by side with the same prompt. +* **OpenAI/Anthropic-compatible APIs**: Serve local models through `/v1/chat/completions`, `/v1/responses` and `/v1/messages`. +* **Connect local models to agents**: Use `unsloth start` with Claude Code, Codex, Hermes and more. +* **Web/PDF search** can read PDF papers, manuals and other PDF results. +* **GGUF hardware controls**: Choose GPUs/layers, offload MoE experts, use multi-GPU or Tensor Parallelism. +* The opt-in **MCP control endpoint** lets AI clients manage models, training, recipes and exports. ### Training -* Train and RL **500+ models** up to **2x faster** with up to **70% less VRAM**, with no accuracy loss. -* Custom Triton and mathematical **kernels**. See some collabs we did with [PyTorch](https://unsloth.ai/docs/get-started/reinforcement-learning-rl-guide/fp8-reinforcement-learning) and [Hugging Face](https://unsloth.ai/docs/new/faster-moe). +* Train and RL **500+ models** up to **2x faster** with **70% less VRAM**; MoE up to **12x faster**. +* Train and run RL on [AMD GPUs](https://unsloth.ai/docs/basics/amd) across Windows, WSL and Linux. * **Data Recipes**: [Auto-create datasets](https://unsloth.ai/docs/new/studio/data-recipe) from **PDF, CSV, DOCX** etc. Edit data in a visual-node workflow. -* **[Reinforcement Learning](https://unsloth.ai/docs/get-started/reinforcement-learning-rl-guide)** (RL): The most efficient [RL](https://unsloth.ai/docs/get-started/reinforcement-learning-rl-guide) library, using **80% less VRAM** for GRPO, [FP8](https://unsloth.ai/docs/get-started/reinforcement-learning-rl-guide/fp8-reinforcement-learning) etc. -* Supports full fine-tuning, RL, pretraining, 4-bit, 16-bit and, FP8 training. +* **[Reinforcement Learning](https://unsloth.ai/docs/get-started/reinforcement-learning-rl-guide)** uses **80% less VRAM** for GRPO, FP8 and vision RL, with 7x longer contexts. +* [**Long-context training**](https://unsloth.ai/docs/new/3x-faster-training-packing): **3x faster**, 30% less VRAM and 500K+ context. +* Supports LoRA/QLoRA, full fine-tuning, RL, pretraining, 4-bit, 16-bit and FP8. +* Custom Triton and mathematical **kernels** built with PyTorch and Hugging Face. * **Observability**: Monitor training live, track loss and GPU usage and customize graphs. * [Multi-GPU](https://unsloth.ai/docs/basics/multi-gpu-training-with-unsloth) training is supported, with major improvements coming soon. +## 🚀 Unsloth Start + +[Unsloth Start](https://unsloth.ai/docs/integrations/unsloth-start) connects [Claude Code](https://unsloth.ai/docs/basics/claude-code), [Codex](https://unsloth.ai/docs/basics/codex) and other agents to local models with one command. + +Start Unsloth, load a model, open your project folder, then run: + +```bash +unsloth start claude +``` + +Replace `claude` with any supported agent: + +| Agent | Command | +| --- | --- | +| Claude Code | `unsloth start claude` | +| OpenAI Codex | `unsloth start codex` | +| Hermes Agent | `unsloth start hermes` | +| OpenClaw | `unsloth start openclaw` | +| OpenCode | `unsloth start opencode` | +| Pi Coding Agent | `unsloth start pi` | + +Claude Code, Codex, OpenCode and Pi can keep their current model and use Unsloth as a local +subagent: + +```bash +unsloth start claude --as-subagent --model unsloth/model-GGUF:quant +``` + ## 📥 Install Unsloth can be used in two ways: through **[Unsloth Studio](https://unsloth.ai/docs/new/studio/)**, the web UI, or through **Unsloth Core**, the code-based version. Each has different requirements. @@ -65,7 +102,8 @@ Unsloth Studio (Beta) works on **Windows, Linux, WSL** and **macOS**. * **CPU:** Supported for Chat and Data Recipes currently * **NVIDIA:** Training works on RTX 30/40/50, Blackwell, DGX Spark, Station and more * **macOS:** Training, MLX and GGUF inference are ALL supported. -* **AMD:** Chat + Data works. Train with [Unsloth Core](#unsloth-core-code-based). Studio support is out soon. +* **AMD:** Training, RL, chat and deployment work on Windows, WSL and Linux. [Read the AMD guide](https://unsloth.ai/docs/basics/amd). +* **Vulkan:** GGUF inference is supported on [compatible GPUs, including Intel GPUs](https://github.com/unslothai/unsloth/pull/5819). Vulkan accelerates GGUF inference only; training still requires a supported PyTorch or MLX backend. * **Multi-GPU:** Available now, with a major upgrade on the way #### macOS, Linux, WSL: @@ -74,19 +112,35 @@ curl -fsSL https://unsloth.ai/install.sh | sh ``` Use the same command to update. +To force the Vulkan llama.cpp backend, set `UNSLOTH_FORCE_VULKAN=1` **before installing or updating**. The setting selects the llama.cpp binary bundle, so setting it only when launching Studio cannot replace an existing CPU bundle: + +```bash +export UNSLOTH_FORCE_VULKAN=1 +curl -fsSL https://unsloth.ai/install.sh | sh +``` + #### Windows: ```powershell irm https://unsloth.ai/install.ps1 | iex ``` Use the same command to update. +To force the Vulkan llama.cpp backend, set the environment variable before running the installer or updater: + +```powershell +$env:UNSLOTH_FORCE_VULKAN=1 +irm https://unsloth.ai/install.ps1 | iex +``` + +Re-running the current installer replaces a previously selected CPU bundle when the backend differs. A separate Vulkan SDK is not required; the GPU driver must provide a working Vulkan runtime. + #### Launch ```bash unsloth studio -p 8888 ``` For LAN or cloud access, add `-H 0.0.0.0` (raw port only; add `--cloudflare` for a public URL). By default, Unsloth is accessible only locally. -To reach Studio over HTTPS, use `unsloth studio --secure`. Studio stays bound to localhost and is reached only through a free Cloudflare tunnel, which publishes it at a public `https://*.trycloudflare.com` URL (it fails closed if the tunnel can't start, so the raw port is never exposed). This makes Studio reachable from the internet, so anyone with the link and API key can use it and run code: keep your API key private (see Remote access below). +To reach Unsloth over HTTPS, use `unsloth studio --secure`. Unsloth stays bound to localhost and is reached only through a free Cloudflare tunnel, which publishes it at a public `https://*.trycloudflare.com` URL (it fails closed if the tunnel can't start, so the raw port is never exposed). This makes Unsloth reachable from the internet, so anyone with the link and API key can use it and run code: keep your API key private (see Remote access below). #### Docker Use our [Docker image](https://hub.docker.com/r/unsloth/unsloth) ```unsloth/unsloth``` container. Run: @@ -122,7 +176,7 @@ You can use the same Docker image as Unsloth Studio. #### AMD, Intel: For RTX 50x, B200, 6000 GPUs: `uv pip install unsloth --torch-backend=auto`. Read our guides for: [Blackwell](https://unsloth.ai/docs/blog/fine-tuning-llms-with-blackwell-rtx-50-series-and-unsloth) and [DGX Spark](https://unsloth.ai/docs/blog/fine-tuning-llms-with-nvidia-dgx-spark-and-unsloth).
-To install Unsloth on **AMD** and **Intel** GPUs, follow our [AMD Guide](https://unsloth.ai/docs/get-started/install/amd) and [Intel Guide](https://unsloth.ai/docs/get-started/install/intel). +To install Unsloth on **AMD** and **Intel** GPUs, follow our [AMD Guide](https://unsloth.ai/docs/basics/amd) and [Intel Guide](https://unsloth.ai/docs/get-started/install/intel). ## 📒 Free Notebooks @@ -148,13 +202,20 @@ Read our [guide](https://unsloth.ai/docs/get-started/fine-tuning-llms-guide). Ad - See detailed documentation for Unsloth [here](https://unsloth.ai/docs) ## 🦥 Unsloth News -- **Connections**: Connect any API provider (OpenAI, Anthropic) or server (vLLM, Ollama). [Guide](https://unsloth.ai/docs/integrations/connections) -- **MTP**: Run Qwen3.6 MTP in Unsloth. MTP settings are autoset specific to your hardware. [Guide](https://unsloth.ai/docs/models/qwen3.6#mtp-guide) -- **API inference endpoint**: Deploy and run local LLMs in Claude Code, Codex tools. [Guide](https://unsloth.ai/docs/basics/api) -- **Qwen3.6**: Qwen3.6-35B-A3B can now be trained and run in Unsloth Studio. [Blog](https://unsloth.ai/docs/models/qwen3.6) -- **Gemma 4**: Run and train Google’s new models directly in Unsloth. [Blog](https://unsloth.ai/docs/models/gemma-4) +- **AMD training**: Train, run RL, chat and deploy on AMD GPUs across Windows, WSL and Linux. [Guide](https://unsloth.ai/docs/basics/amd) +- **GGUF hardware controls**: Choose GPU/layer placement, offload MoE experts and use multi-GPU or Tensor Parallelism. [#6414](https://github.com/unslothai/unsloth/pull/6414) +- **Local models for any agent**: Use `unsloth start` with Claude Code, Codex, Hermes, OpenCode, OpenClaw, Pi and more through Unsloth's OpenAI- and Anthropic-compatible APIs. [Guide](https://unsloth.ai/docs/basics/api) +- **MCP control endpoint**: Let compatible clients manage models, training, recipes, checkpoints and exports. [#7191](https://github.com/unslothai/unsloth/pull/7191) +- **Local inference reliability**: Resume long chats faster, recover stalled downloads and reuse existing GGUF files. [#7204](https://github.com/unslothai/unsloth/pull/7204) • [#6858](https://github.com/unslothai/unsloth/pull/6858) • [#7209](https://github.com/unslothai/unsloth/pull/7209) +- **New models**: [Qwen-AgentWorld](https://huggingface.co/unsloth/Qwen-AgentWorld-35B-A3B-GGUF), [Ornith](https://huggingface.co/unsloth/models?search=ornith), [Kimi K2.7 Code](https://unsloth.ai/docs/models/kimi-k2.7-code) and [MiniMax M3](https://unsloth.ai/docs/models/minimax-m3) +- **GLM-5.2**: Run Z.ai's 744B-parameter, 1M-context open model locally with Unsloth Dynamic GGUFs. [Guide](https://unsloth.ai/docs/models/glm-5.2) +- **DeepSeek-V4**: Run DeepSeek-V4-Flash locally with corrected multi-turn and tool-calling behavior. [Guide](https://unsloth.ai/docs/models/deepseek-v4) +- **DiffusionGemma**: Run and fine-tune Google's diffusion language model with 1.8x faster inference in Unsloth Studio. [Guide](https://unsloth.ai/docs/models/diffusiongemma) +- **Qwen3.6**: Run and train Qwen3.6 with MTP for 1.4-2.2x faster inference and NVFP4 quants for supported GPUs. [Guide](https://unsloth.ai/docs/models/qwen3.6) +- **Gemma 4**: Run and train Gemma 4 text, image and audio models with QAT, MTP, GGUF and MLX support. [Guide](https://unsloth.ai/docs/models/gemma-4) +- **MCP servers**: Connect local models to files, apps, databases and external tools through Model Context Protocol. [Guide](https://unsloth.ai/docs/basics/mcp) +- **Connections**: Mix local models with API providers (OpenAI, Anthropic) or servers (vLLM, Ollama) in the same interface. [Guide](https://unsloth.ai/docs/integrations/connections) - **Introducing Unsloth Studio**: our new web UI for running and training LLMs. [Blog](https://unsloth.ai/docs/new/studio) -- **Qwen3.5** - 0.8B, 2B, 4B, 9B, 27B, 35-A3B, 112B-A10B are now supported. [Guide + notebooks](https://unsloth.ai/docs/models/qwen3.5/fine-tune) - Train **MoE LLMs 12x faster** with 35% less VRAM - DeepSeek, GLM, Qwen and gpt-oss. [Blog](https://unsloth.ai/docs/new/faster-moe) - **Embedding models**: Unsloth now supports ~1.8-3.3x faster embedding fine-tuning. [Blog](https://unsloth.ai/docs/new/embedding-finetuning) • [Notebooks](https://unsloth.ai/docs/get-started/unsloth-notebooks#embedding-models) - New **7x longer context RL** vs. all other setups, via our new batching algorithms. [Blog](https://unsloth.ai/docs/new/grpo-long-context) @@ -208,7 +269,7 @@ unsloth studio -p 8888 #### Remote access: `--secure` (HTTPS tunnel) vs raw port By default `unsloth studio` binds to `127.0.0.1` (this machine only). To reach it from another device, pick one of: -- `--secure` (recommended): serve **only** through a free Cloudflare HTTPS link. Studio stays bound to localhost and the tunnel provides the public URL; it fails closed (does not start) if the tunnel can't come up, so the raw port is never exposed. +- `--secure` (recommended): serve **only** through a free Cloudflare HTTPS link. Unsloth stays bound to localhost and the tunnel provides the public URL; it fails closed (does not start) if the tunnel can't come up, so the raw port is never exposed. ```bash unsloth studio --secure -p 8888 ``` @@ -218,7 +279,9 @@ unsloth studio -H 0.0.0.0 -p 8888 ``` The Cloudflare tunnel is **off by default**: `-H 0.0.0.0` exposes the raw port only, not a public internet URL. Pair the wildcard bind with `--cloudflare` (`unsloth studio -H 0.0.0.0 --cloudflare`) to also publish a public `https://*.trycloudflare.com` link, or prefer `--secure` (above), which keeps the raw port private. `--cloudflare` has no effect on a loopback bind. -The first time Studio is published on a public URL (`--secure` or `--cloudflare`) with the auto-generated admin password still in place, it asks for a new admin password in the terminal (masked input with confirmation) before the public link goes up. Without an attached terminal it warns instead and keeps the bootstrap deadline: Studio shuts down after `UNSLOTH_STUDIO_BOOTSTRAP_TIMEOUT` (default 1 hour) unless the password is changed in the web UI. +On a wildcard bind Unsloth works out the address to share by asking `ifconfig.me` for the public IP, then asks `check-host.net` whether that port is reachable so it can tell you if a firewall is in the way. Both contact a third party. Set `UNSLOTH_STUDIO_DISABLE_PUBLIC_CHECK=1` to skip them; the banner then shows the LAN address and no reachability line. + +The first time Unsloth is published on a public URL (`--secure` or `--cloudflare`) with the auto-generated admin password still in place, it asks for a new admin password in the terminal (masked input with confirmation) before the public link goes up. Without an attached terminal it warns instead and keeps the bootstrap deadline: Unsloth shuts down after `UNSLOTH_STUDIO_BOOTSTRAP_TIMEOUT` (default 1 hour) unless the password is changed in the web UI. For headless setups that cannot answer that prompt, set the initial admin password non-interactively with `--password` (only takes effect when no password is set yet; if one already exists it is a hard error, so rotate later with `unsloth studio reset-password`): @@ -230,7 +293,7 @@ printf '%s\n' 'your-strong-password' | unsloth studio --secure --password - # A literal `--password VALUE` is visible in the process list and shell history, so prefer the `UNSLOTH_STUDIO_PASSWORD` env var or `--password -` (stdin) for automation. This applies to any launch (public or a headless `-H 0.0.0.0` bind), and the password is set in the parent before the server binds, so it never reaches a re-executed child process. -Server-side tools (web search, Python and terminal code execution) run as your user and are on by default. Anyone who can reach the server with the API key can run code on this machine, so keep your API key private and pass `--disable-tools` when exposing Studio. +Server-side tools (web search, Python and terminal code execution) run as your user and are on by default. Anyone who can reach the server with the API key can run code on this machine, so keep your API key private and pass `--disable-tools` when exposing Unsloth. #### Advanced launch options Installer options can be passed as environment variables. On macOS, Linux and WSL place the variable after the pipe so the shell passes it to `sh`; on Windows set it with `$env:` before piping to `iex`. @@ -243,7 +306,7 @@ curl -fsSL https://unsloth.ai/install.sh | UNSLOTH_NO_TORCH=1 sh $env:UNSLOTH_NO_TORCH=1; irm https://unsloth.ai/install.ps1 | iex ``` -Skip the post-install prompt that starts Studio (useful for automated installs): +Skip the post-install prompt that starts Unsloth (useful for automated installs): ```bash curl -fsSL https://unsloth.ai/install.sh | UNSLOTH_SKIP_AUTOSTART=1 sh ``` @@ -279,9 +342,9 @@ UNSLOTH_NPM_REGISTRY=https://artifactory.example.com/api/npm/npm/ ./install.sh - ```powershell $env:UNSLOTH_NPM_REGISTRY='https://artifactory.example.com/api/npm/npm/'; .\install.ps1 --local ``` -It is threaded as `--registry` into the Studio frontend `npm`/`bun` installs; the supply-chain locks (7-day `min-release-age`, exact version pins) stay in force. +It is threaded as `--registry` into the Unsloth frontend `npm`/`bun` installs; the supply-chain locks (7-day `min-release-age`, exact version pins) stay in force. -Cap Studio's native CPU thread pools on high-core hosts: `UNSLOTH_CPU_THREADS=8 unsloth studio -p 8888`. +Cap Unsloth's native CPU thread pools on high-core hosts: `UNSLOTH_CPU_THREADS=8 unsloth studio -p 8888`. #### Uninstall The recommended way to fully remove Unsloth Studio is the matching uninstall script for your OS. It stops any running servers, removes the install dir, the launcher data dir, the desktop shortcut, and any platform-specific entries (macOS `.app` bundle + Launch Services on Mac; Start Menu, `HKCU\Software\Unsloth` registry key and user `PATH` entries on Windows): diff --git a/_changelog_build.py b/_changelog_build.py new file mode 100644 index 0000000000..f5bcf2052c --- /dev/null +++ b/_changelog_build.py @@ -0,0 +1,36 @@ +# SPDX-License-Identifier: AGPL-3.0-only +# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. + +"""Snapshot CHANGELOG.md into the studio package at build time. + +CHANGELOG.md at the repo root stays the one file to edit. Copying it here, +rather than in build.sh, means every packaging path ships it, so release notes +still render when the popup cannot reach GitHub.""" + +from __future__ import annotations + +import shutil +from pathlib import Path + +from setuptools.command.build_py import build_py as _build_py + +ROOT = Path(__file__).resolve().parent +SOURCE = ROOT / "CHANGELOG.md" +SNAPSHOT = ROOT / "studio" / "CHANGELOG.md" + + +class build_py(_build_py): + def run(self) -> None: + # Beside the sources only if writable (PEP 517 may build an immutable + # checkout); into the staging directory always. + if SOURCE.is_file(): + try: + shutil.copyfile(SOURCE, SNAPSHOT) + except OSError: + pass + super().run() + if not SOURCE.is_file(): + return + staged = Path(self.build_lib) / "studio" / "CHANGELOG.md" + staged.parent.mkdir(parents = True, exist_ok = True) + shutil.copyfile(SOURCE, staged) diff --git a/build.sh b/build.sh index dc272f0de1..5b09a7791b 100644 --- a/build.sh +++ b/build.sh @@ -4,9 +4,9 @@ set -euo pipefail -# PyPI/Studio release publishing must use `./build.sh publish` (or an -# equivalent stamp -> build -> verify-dist -> upload flow) so packaged Studio -# artifacts include the display-only Studio release version. +# PyPI/Unsloth release publishing must use `./build.sh publish` (or an +# equivalent stamp -> build -> verify-dist -> upload flow) so packaged Unsloth +# artifacts include the display-only Unsloth release version. # 1. Build frontend (Vite outputs to dist/) cd studio/frontend @@ -87,7 +87,7 @@ cd ../.. # 2. Clean old artifacts rm -rf build dist *.egg-info -# 3. Stamp display-only Studio release metadata for packaged builds. +# 3. Stamp display-only Unsloth release metadata for packaged builds. _STUDIO_BUILD_INFO="studio/backend/utils/_studio_release_build.py" _STUDIO_BUILD_INFO_BACKUP="$(mktemp)" cp "$_STUDIO_BUILD_INFO" "$_STUDIO_BUILD_INFO_BACKUP" @@ -103,9 +103,13 @@ else STUDIO_STAMPED_VERSION="$(python scripts/stamp_studio_release.py)" fi -# 4. Build wheel/sdist +# 4. Build wheel/sdist. _changelog_build.py snapshots CHANGELOG.md into the studio +# package so release notes render offline. python -m build +# Drop the snapshot so a source checkout never serves a stale copy. +rm -f studio/CHANGELOG.md + if [ "${1:-}" = "publish" ]; then python scripts/stamp_studio_release.py --verify-dist dist --expected "$STUDIO_STAMPED_VERSION" fi diff --git a/install.ps1 b/install.ps1 index 4fa01bfa28..5b205df96d 100644 --- a/install.ps1 +++ b/install.ps1 @@ -28,6 +28,14 @@ function Install-UnslothStudio { } } + function Clear-TauriInstallError { + param([string]$Message) + if ($TauriMode) { + Write-TauriLog "ERROR_CLEAR" $Message + [Console]::Error.WriteLine("[TAURI:ERROR_CLEAR] $Message") + } + } + function Format-TauriDiagBool { param([bool]$Value) if ($Value) { return "true" } @@ -49,11 +57,32 @@ function Install-UnslothStudio { } } + # Machine arch; Get-TauriDiagArch above reports the process. An emulated x64 shell on + # ARM64 reports AMD64, but PROCESSOR_ARCHITEW6432 is ARM64 in exactly that case. + function Get-HostMachineArch { + $osArch = "" + try { $osArch = [System.Runtime.InteropServices.RuntimeInformation]::OSArchitecture.ToString() } catch { $osArch = "" } + $signals = @([string]$env:PROCESSOR_ARCHITEW6432, [string]$env:PROCESSOR_ARCHITECTURE, $osArch) + foreach ($s in $signals) { + if ($s.ToLowerInvariant() -eq "arm64") { return "arm64" } + } + foreach ($s in $signals) { + if ([string]::IsNullOrWhiteSpace($s)) { continue } + switch ($s.ToLowerInvariant()) { + "amd64" { return "x86_64" } + "x64" { return "x86_64" } + "x86" { return "x86" } + } + } + return "unknown" + } + function Get-TauriTorchIndexFamily { param([string]$TorchIndexUrl) if ($SkipTorch) { return "none" } if ([string]::IsNullOrWhiteSpace($TorchIndexUrl)) { return "none" } - $leaf = ($TorchIndexUrl.TrimEnd('/') -split '/')[-1].ToLowerInvariant() + # Drop query/fragment first so a token-authenticated pin classifies by family. + $leaf = (($TorchIndexUrl -split '[?#]', 2)[0].TrimEnd('/') -split '/')[-1].ToLowerInvariant() if (@("cpu", "cu118", "cu124", "cu126", "cu128", "cu130") -contains $leaf) { return $leaf } if ($leaf -match '^rocm[0-9]+\.[0-9]+$') { return $leaf } return "auto" @@ -62,7 +91,8 @@ function Install-UnslothStudio { function Get-TauriGpuBranch { param([string]$TorchIndexFamily) if ($SkipTorch) { return "no_torch" } - if ($TorchIndexFamily -like "cu*") { return "cuda" } + # Require a digit after "cu" so /current or /custom isn't branded CUDA (parity ^cu[0-9]). + if ($TorchIndexFamily -match '^cu[0-9]') { return "cuda" } if ($TorchIndexFamily -like "rocm*") { return "rocm" } if ($TorchIndexFamily -eq "cpu") { return "cpu" } return "unknown" @@ -84,7 +114,7 @@ function Install-UnslothStudio { [int]$Code = 1 ) if ($Code -eq 0) { $Code = 1 } - Write-TauriLog "ERROR" $Message + Write-TauriLog "ERROR_DEFAULT" $Message if (Get-Command Restore-StudioVenvRollback -CommandType Function -ErrorAction SilentlyContinue) { Restore-StudioVenvRollback } @@ -176,7 +206,7 @@ function Install-UnslothStudio { $envOverride = $env:STUDIO_HOME.Trim() } - # Custom Studio roots are not supported with --tauri (desktop app still + # Custom Unsloth roots are not supported with --tauri (desktop app still # resolves %USERPROFILE%\.unsloth\studio). Pass through if override == legacy. if ($TauriMode -and $envOverride) { $_tauriOverride = $envOverride @@ -467,43 +497,70 @@ function Install-UnslothStudio { } } + # Redact index-URL credentials (userinfo + ?query= + #fragment) from captured installer + # output before printing on failure; uv/pip errors echo the failing --index-url verbatim. + # Mirrors the other installers. Verbose mode streams uncaptured, so it isn't redacted. + function Redact-InstallOutput { + param([string]$Text) + if (-not $Text) { return $Text } + $Text = $Text -replace '(https?://)[^/@\s`]+@', '$1@' + $Text = $Text -replace '([?&][^=\s&`]+)=[^&#\s`]+', '$1=' + # A #token=... fragment is as sensitive as a query; URL-anchored. + return $Text -replace '(https?://[^\s`#]+)#[^\s`]+', '$1#' + } + # Run native commands quietly by default to match install.sh behavior. # Full command output is shown only when --verbose / UNSLOTH_VERBOSE=1. function Invoke-InstallCommand { param( - [Parameter(Mandatory = $true)][ScriptBlock]$Command + [Parameter(Mandatory = $true)][ScriptBlock]$Command, + [string]$Label = "install command" ) - # Installer-pinned index installs (torch) must beat an inherited uv mirror - # (#6898): when the command pins an index, clear every uv index env var so - # it wins, then restore in finally. Other installs keep the user's mirror. + # Installer-pinned index installs (torch) must beat an inherited uv mirror (#6898): + # for --default-index, clear the uv index env vars (restore in finally) and set + # UV_NO_CONFIG=1 so a uv.toml/pyproject index can't outrank the CLI pin (uv 0.10). $savedUvIndex = $null if ($Command.ToString() -match '--default-index') { $savedUvIndex = @{} - foreach ($n in 'UV_DEFAULT_INDEX', 'UV_INDEX_URL', 'UV_INDEX', 'UV_EXTRA_INDEX_URL') { + foreach ($n in 'UV_DEFAULT_INDEX', 'UV_INDEX_URL', 'UV_INDEX', 'UV_EXTRA_INDEX_URL', 'UV_TORCH_BACKEND', 'UV_FIND_LINKS', 'UV_CONFIG_FILE', 'UV_NO_CONFIG') { $savedUvIndex[$n] = [Environment]::GetEnvironmentVariable($n) Remove-Item "Env:$n" -ErrorAction SilentlyContinue } + $env:UV_NO_CONFIG = '1' } $prevEap = $ErrorActionPreference $ErrorActionPreference = "Continue" try { # Reset to avoid stale values from prior native commands. $global:LASTEXITCODE = 0 + Write-TauriLog "OUTPUT_CLEAR" $Label if ($script:UnslothVerbose) { # Merge stderr into stdout so progress/warning output stays visible # without flipping $? on successful native commands (PS 5.1 treats # stderr records as errors that set $? = $false even on exit code 0). - & $Command 2>&1 | Out-Host + # Redact per record: uv echoes index URLs (credentials and all) in + # its errors, and verbose mode must not bypass the quiet path's + # redaction. ForEach-Object/Out-Host leave $LASTEXITCODE untouched. + & $Command 2>&1 | ForEach-Object { Redact-InstallOutput "$_" } | Out-Host } else { $output = & $Command 2>&1 | Out-String if ($LASTEXITCODE -ne 0) { - Write-Host $output -ForegroundColor Red + Write-Host (Redact-InstallOutput $output) -ForegroundColor Red } } - return [int]$LASTEXITCODE + $exitCode = [int]$LASTEXITCODE + if ($exitCode -eq 0) { + Clear-TauriInstallError "$Label recovered" + } else { + Write-TauriLog "ERROR_OUTPUT" "$Label failed (exit code $exitCode)" + } + return $exitCode } finally { $ErrorActionPreference = $prevEap - if ($savedUvIndex) { foreach ($n in $savedUvIndex.Keys) { if ($null -ne $savedUvIndex[$n]) { Set-Item "Env:$n" $savedUvIndex[$n] } } } + if ($savedUvIndex) { + Remove-Item "Env:UV_NO_CONFIG" -ErrorAction SilentlyContinue + foreach ($n in $savedUvIndex.Keys) { if ($null -ne $savedUvIndex[$n]) { Set-Item "Env:$n" $savedUvIndex[$n] } } + } } } @@ -528,7 +585,7 @@ function Install-UnslothStudio { } $attempt = 1 while ($true) { - $code = Invoke-InstallCommand $Command + $code = Invoke-InstallCommand -Command $Command -Label $Label if ($code -eq 0) { return 0 } if ($attempt -ge $maxAttempts) { return $code } substep ("retrying ""$Label"" after transient failure (attempt $($attempt + 1)/$maxAttempts, waiting ${delay}s)...") "Yellow" @@ -756,7 +813,7 @@ function Find-FreeLaunchPort { return `$null } -# If Studio is already healthy on any expected port, just open it and exit. +# If Unsloth is already healthy on any expected port, just open it and exit. `$existingPort = Find-HealthyStudioPort if (`$existingPort) { Start-Process "http://localhost:`$existingPort" @@ -772,7 +829,7 @@ try { `$haveMutex = `$true } if (-not `$haveMutex) { - # Another launcher is already running; wait for it to bring Studio up + # Another launcher is already running; wait for it to bring Unsloth up `$deadline = (Get-Date).AddSeconds(`$timeoutSec) while ((Get-Date) -lt `$deadline) { `$port = Find-HealthyStudioPort @@ -1087,10 +1144,27 @@ exit 0 return $false } + # The interpreter's own arch, asked of it: win-amd64|win-arm64|win32|"". + function Get-PythonPlatformTag { + param([string]$Exe) + try { + return (& $Exe -c "import sysconfig; print(sysconfig.get_platform())" 2>$null | Out-String).Trim().ToLowerInvariant() + } catch { return "" } + } + # Returns @{ Version = "3.13"; Path = "C:\...\python.exe" } or $null. # The resolved Path is passed to `uv venv --python` to prevent uv from # re-resolving the version string back to a conda interpreter. function Find-CompatiblePython { + # -X64Only: best installed x64 interpreter or $null, never ARM64. Last resort for + # Install-X64Python, where x64 of a lower-priority minor beats ARM64. + param([switch]$X64Only) + # Windows on ARM: prefer x64. pyarrow (via datasets) and hf-transfer ship no + # win_arm64 wheel, so a native ARM64 Python source-builds both and dies on CMake / + # Rust minutes in; x64 runs fine emulated. ARM64 is still returned when it is all + # there is, and the caller then bootstraps x64 or warns. + $preferX64 = $X64Only -or ((Get-HostMachineArch) -eq "arm64") + $candidates = @() # Try the Python Launcher first (most reliable on Windows) # py.exe resolves to the standard CPython install, not conda. # Prefer the requested $PythonVersion, then newest-first fallback. @@ -1108,7 +1182,8 @@ exit 0 # Resolve the actual executable path and verify it is not conda-based $resolvedExe = (& $pyLauncher.Source "-$minor" -c "import sys; print(sys.executable)" 2>$null | Out-String).Trim() if ($resolvedExe -and (Test-Path $resolvedExe) -and -not (Test-IsCondaPython $resolvedExe)) { - return @{ Version = $ver; Path = $resolvedExe } + if (-not $preferX64) { return @{ Version = $ver; Path = $resolvedExe; Arch = "" } } + $candidates += @{ Version = $ver; Path = $resolvedExe } } } } catch {} @@ -1129,11 +1204,53 @@ exit 0 try { $out = & $cmd.Source --version 2>&1 | Out-String if ($out -match "Python (3\.1[1-3])\.\d+") { - return @{ Version = $Matches[1]; Path = $cmd.Source } + if (-not $preferX64) { return @{ Version = $Matches[1]; Path = $cmd.Source; Arch = "" } } + $candidates += @{ Version = $Matches[1]; Path = $cmd.Source } } } catch {} } } + # `py -3.12` runs the launcher's preferred build, normally the native ARM64 one, so + # a same-minor x64 install that is neither preferred nor on PATH never becomes a + # candidate. `-3.12-64` cannot disambiguate (deprecated, it only means "not + # 32-bit"), so enumerate every registration with -0p and probe each path. + if ($preferX64) { + foreach ($pyLauncher in @(Get-Command py -All -CommandType Application -ErrorAction SilentlyContinue)) { + if ($pyLauncher.Source -match $script:CondaSkipPattern) { continue } + $listed = @() + try { $listed = @(& $pyLauncher.Source "-0p" 2>$null) } catch {} + foreach ($line in $listed) { + # " -V:3.12 * C:\...\python.exe": tag, optional default marker, path. + $m = [regex]::Match([string]$line, '(?i)^\s*-\S+\s+\*?\s*"?(?

\S.*?\.exe)"?\s*$') + if (-not $m.Success) { continue } + $exe = $m.Groups['p'].Value.Trim() + if ($candidates | Where-Object { $_.Path -eq $exe }) { continue } + if (-not (Test-Path -LiteralPath $exe)) { continue } + if (Test-IsCondaPython $exe) { continue } + try { + $out = & $exe --version 2>&1 | Out-String + if ($out -match "Python (3\.1[1-3])\.\d+") { + $candidates += @{ Version = $Matches[1]; Path = $exe } + } + } catch {} + } + } + } + # Prefer x64, but only within one minor: $minors is the caller's version preference, + # so ranking on arch alone would answer UNSLOTH_PYTHON=3.12 with an x64 3.13 and + # never bootstrap x64 3.12. Probing costs a subprocess, so non-ARM returned above. + foreach ($c in $candidates) { + $tag = Get-PythonPlatformTag $c.Path + $c.Arch = if ($tag -eq "win-amd64") { "x86_64" } elseif ($tag -eq "win-arm64") { "arm64" } else { "unknown" } + } + foreach ($minor in $minors) { + $sameMinor = @($candidates | Where-Object { $_.Version -eq $minor }) + if ($sameMinor.Count -eq 0) { continue } + $x64 = $sameMinor | Where-Object { $_.Arch -eq "x86_64" } | Select-Object -First 1 + if ($x64) { return $x64 } + if (-not $X64Only) { return $sameMinor[0] } + } + if (-not $X64Only -and $candidates.Count -gt 0) { return $candidates[0] } return $null } @@ -1144,8 +1261,11 @@ exit 0 # (no UAC), putting python.exe + the py launcher on PATH. Mirrors the uv -> # astral.sh fallback below. Returns @{ Version; Path } or $null. function Install-PythonFromPythonOrg { + # $Arch overrides the host arch, to pull x64 onto an ARM64 box. + param([string]$Arch = "") # python.org ships one installer per architecture. - $archSuffix = switch (Get-TauriDiagArch) { + $targetArch = if ($Arch) { $Arch } else { Get-TauriDiagArch } + $archSuffix = switch ($targetArch) { "x86_64" { "-amd64" } "arm64" { "-arm64" } "x86" { "" } @@ -1210,6 +1330,28 @@ exit 0 return (Find-CompatiblePython) } + # ── Windows on ARM: get an x64 CPython ── + # --architecture x64 forces winget off the ARM64 build; python.org takes the same override. + function Install-X64Python { + if ($script:WingetAvailable) { + $prevEAP = $ErrorActionPreference + $ErrorActionPreference = "Continue" + try { + winget install -e --id "Python.Python.$PythonVersion" --source winget --architecture x64 --accept-package-agreements --accept-source-agreements + } catch { } + $ErrorActionPreference = $prevEAP + Refresh-SessionPath + $found = Find-CompatiblePython + if ($found -and $found.Arch -eq "x86_64") { return $found } + substep "winget could not provide an x64 Python -- trying python.org..." "Yellow" + } + $found = Install-PythonFromPythonOrg -Arch "x86_64" + if ($found -and $found.Arch -eq "x86_64") { return $found } + # Nothing installable (offline / no winget): an x64 build of another supported minor + # still runs the wheels ARM64 cannot, so take it over the native interpreter. + return (Find-CompatiblePython -X64Only) + } + # ── Install Python if no compatible version (3.11-3.13) found ── # Find-CompatiblePython returns @{ Version = "3.13"; Path = "C:\...\python.exe" } or $null. Write-TauriLog "STEP" "Installing Python" @@ -1281,6 +1423,26 @@ exit 0 return (Exit-InstallFailure "Python installation failed") } } + # ── Windows on ARM: swap a native ARM64 interpreter for x64 ── + # pyarrow and hf-transfer publish no win_arm64 wheel, so an ARM64 Python source-builds + # both and fails deep into the run. Warn up front if x64 is unobtainable. + if ($DetectedPython -and (Get-HostMachineArch) -eq "arm64" -and $DetectedPython.Arch -ne "x86_64") { + substep "windows on arm: only a native ARM64 Python $($DetectedPython.Version) was found." "Yellow" + substep "pyarrow and hf-transfer publish no win_arm64 wheels, so installing x64 Python..." "Yellow" + $X64Python = Install-X64Python + if ($X64Python) { + $DetectedPython = $X64Python + step "python" "using x64 Python $($DetectedPython.Version) under emulation" + } else { + Write-Host "[WARN] Could not install an x64 Python on this ARM64 machine." -ForegroundColor Yellow + Write-Host " Continuing with ARM64 Python $($DetectedPython.Version), but the install is likely to fail:" -ForegroundColor Yellow + Write-Host " pyarrow (via datasets) and hf-transfer ship no win_arm64 wheels and will be" -ForegroundColor Yellow + Write-Host " built from source, which needs CMake plus the MSVC and Rust toolchains." -ForegroundColor Yellow + Write-Host " Fix: install x64 Python from https://www.python.org/downloads/windows/" -ForegroundColor Yellow + Write-Host " (choose 'Windows installer (64-bit)', not ARM64), then re-run this installer." -ForegroundColor Yellow + } + } + $DiagPythonVersion = $PythonVersion if ($DetectedPython) { $DiagPythonVersion = $DetectedPython.Version } $InitialGpuBranch = "unknown" @@ -1395,13 +1557,82 @@ exit 0 $suffix++ $candidate = Join-Path $StudioHome "unsloth_studio.rollback.$stamp.$PID.$suffix" } - Move-Item -LiteralPath $ExistingDir -Destination $candidate -ErrorAction Stop $script:StudioVenvRollbackDir = $candidate $script:StudioVenvRollbackTarget = $ExistingDir $script:StudioVenvRollbackActive = $true + # Publish the rollback state before the atomic rename so interruption + # cannot land after Move-Item but before cleanup knows where the old venv went. + try { + Move-Item -LiteralPath $ExistingDir -Destination $candidate -ErrorAction Stop + } catch { + # A collision or ordinary rename failure leaves the original in place. + # Keep state active only when the rename happened before interruption. + if (Test-Path -LiteralPath $ExistingDir) { + $script:StudioVenvRollbackActive = $false + $script:StudioVenvRollbackDir = $null + } + throw + } substep "previous environment preserved for rollback" } + function Remove-StudioVenvTreeWithRetry { + param( + [Parameter(Mandatory = $true)][string]$Path, + [Parameter(Mandatory = $true)][string]$Label + ) + $lastError = $null + for ($attempt = 1; $attempt -le 3; $attempt++) { + try { + Remove-Item -LiteralPath $Path -Recurse -Force -ErrorAction Stop + } catch { + $lastError = $_.Exception.Message + } + if (-not (Test-Path -LiteralPath $Path)) { return $true } + if ($attempt -lt 3) { Start-Sleep -Milliseconds (250 * $attempt) } + } + Write-Host "[WARN] Could not remove $Label at $Path" -ForegroundColor Yellow + if ($lastError) { Write-Host " $lastError" -ForegroundColor Yellow } + return $false + } + + function Test-StudioVenvRollbackMustBePreserved { + param([Parameter(Mandatory = $true)][System.IO.FileSystemInfo]$Rollback) + # Preserve anything outside the installer's timestamp.PID[.suffix] format. + if ($Rollback.Name -notmatch '^unsloth_studio\.rollback\.[0-9]{14}\.([0-9]+)(?:\.[0-9]+)?$') { + return $true + } + $ownerPid = 0 + if (-not [int]::TryParse($Matches[1], [ref]$ownerPid)) { return $true } + if ($ownerPid -eq $PID) { return $true } + return $null -ne (Get-Process -Id $ownerPid -ErrorAction SilentlyContinue) + } + + function Remove-StaleStudioVenvRollbacks { + try { + $rollbacks = @( + Get-ChildItem -LiteralPath $StudioHome -Directory -Force -ErrorAction Stop | + Where-Object { $_.Name -like 'unsloth_studio.rollback.*' } + ) + } catch { + Write-Host "[WARN] Could not inspect stale environment rollbacks in $StudioHome" -ForegroundColor Yellow + Write-Host " $($_.Exception.Message)" -ForegroundColor Yellow + return + } + foreach ($rollback in $rollbacks) { + if (($rollback.Attributes -band [System.IO.FileAttributes]::ReparsePoint) -ne 0) { + Write-Host "[WARN] Refusing to remove rollback reparse point $($rollback.FullName)" -ForegroundColor Yellow + continue + } + # A concurrent installer may have moved its live venv aside. The PID + # in the generated name keeps this run from deleting its rescue copy. + if (Test-StudioVenvRollbackMustBePreserved -Rollback $rollback) { continue } + if (Remove-StudioVenvTreeWithRetry -Path $rollback.FullName -Label "stale environment rollback") { + substep "removed stale environment rollback $($rollback.Name)" + } + } + } + function Restore-StudioVenvRollback { if (-not $script:StudioVenvRollbackActive) { return } $backup = $script:StudioVenvRollbackDir @@ -1413,7 +1644,9 @@ exit 0 substep "restoring previous environment after failed install..." "Yellow" try { if (Test-Path -LiteralPath $target) { - Remove-Item -LiteralPath $target -Recurse -Force -ErrorAction SilentlyContinue + if (-not (Remove-StudioVenvTreeWithRetry -Path $target -Label "incomplete environment")) { + throw "Could not remove incomplete environment at $target" + } } Move-Item -LiteralPath $backup -Destination $target -Force -ErrorAction Stop substep "restored previous environment" @@ -1428,17 +1661,21 @@ exit 0 function Complete-StudioVenvRollback { if (-not $script:StudioVenvRollbackActive) { return } $backup = $script:StudioVenvRollbackDir - if ($backup -and (Test-Path -LiteralPath $backup)) { - Remove-Item -LiteralPath $backup -Recurse -Force -ErrorAction SilentlyContinue - } + # The replacement is committed. Disable restoration before deleting the + # backup so interruption cannot restore a partially deleted environment. $script:StudioVenvRollbackActive = $false $script:StudioVenvRollbackDir = $null + if ($backup -and (Test-Path -LiteralPath $backup)) { + Remove-StudioVenvTreeWithRetry -Path $backup -Label "environment rollback" | Out-Null + } } + $studioVenvReplacementCommitted = $false + try { if (Test-Path -LiteralPath $VenvPython) { # why: matching guard to the .venv branch below -- in env-mode # $StudioHome is a user-chosen workspace, so refuse to nuke an - # existing $StudioHome\unsloth_studio that lacks Studio sentinels. + # existing $StudioHome\unsloth_studio that lacks Unsloth sentinels. # -PathType Leaf rejects a directory at the sentinel path. Accept the # in-VENV ownership marker so partial-install retries are not blocked. if ( @@ -1449,7 +1686,7 @@ exit 0 ) { Write-Host "[ERROR] $VenvDir already exists but does not look like an Unsloth Studio install." -ForegroundColor Red Write-Host " Move it aside or choose an empty UNSLOTH_STUDIO_HOME." -ForegroundColor Yellow - throw "Refusing to delete non-Studio venv at $VenvDir" + throw "Refusing to delete non-Unsloth venv at $VenvDir" } # New layout already exists -- replace only after preserving rollback copy. substep "preserving existing environment for rollback..." @@ -1468,7 +1705,7 @@ exit 0 # workspace root (e.g. user's existing project Python venv). $OldVenv = Join-Path $StudioHome ".venv" $OldPy = Join-Path $OldVenv "Scripts\python.exe" - substep "found legacy Studio environment, validating..." + substep "found legacy Unsloth environment, validating..." $prevEAP2 = $ErrorActionPreference $ErrorActionPreference = "Continue" try { @@ -1498,7 +1735,7 @@ exit 0 # Skip in env-mode so we don't relocate the default-install venv into # the workspace root. $CwdVenv = Join-Path $env:USERPROFILE "unsloth_studio" - substep "found CWD-relative Studio environment, migrating to $VenvDir..." + substep "found CWD-relative Unsloth environment, migrating to $VenvDir..." Move-Item -LiteralPath $CwdVenv -Destination $VenvDir -Force substep "moved ~/unsloth_studio -> ~/.unsloth/studio/unsloth_studio" $_Migrated = $true @@ -1507,7 +1744,7 @@ exit 0 if (-not (Test-Path -LiteralPath $VenvPython)) { step "venv" "creating Python $($DetectedPython.Version) virtual environment" substep "$VenvDir" - $venvExit = Invoke-InstallCommand { uv venv $VenvDir --python "$($DetectedPython.Path)" } + $venvExit = Invoke-InstallCommand -Label "create virtual environment" { uv venv $VenvDir --python "$($DetectedPython.Path)" } if ($venvExit -ne 0) { Write-Host "[ERROR] Failed to create virtual environment (exit code $venvExit)" -ForegroundColor Red return (Exit-InstallFailure "Failed to create virtual environment (exit code $venvExit)" $venvExit) @@ -1517,7 +1754,7 @@ exit 0 substep "$VenvDir" } - # Mark the freshly-created venv as Studio-owned so a partial install can be + # Mark the freshly-created venv as Unsloth-owned so a partial install can be # repaired by re-running install.ps1; the env-mode deletion guard above # accepts this marker as the primary sentinel. if (Test-Path -LiteralPath $VenvDir -PathType Container) { @@ -1526,7 +1763,7 @@ exit 0 # ── Helper: run amd-smi without triggering a UAC elevation prompt ── # amd-smi on Windows auto-elevates to read GPU/APU memory, surfacing a confusing - # DiskPart UAC prompt mid-install (Studio backend amd.py hits the same). + # DiskPart UAC prompt mid-install (Unsloth backend amd.py hits the same). # __COMPAT_LAYER=RunAsInvoker forces it (and helpers it spawns) to run # un-elevated; on failure the WMI name -> gfx fallback still resolves the arch. function Invoke-AmdSmiNoElevate { @@ -1653,7 +1890,7 @@ exit 0 function Test-HipinfoIsVenvInternal { param([AllowNull()][string]$HipinfoPath) if ([string]::IsNullOrWhiteSpace($HipinfoPath)) { return $false } - # Also derive the venv from the setup python + default Studio home, so + # Also derive the venv from the setup python + default Unsloth home, so # the venv hipInfo is caught when VenvDir/VIRTUAL_ENV are unset. $venvRoots = @() if ($env:VIRTUAL_ENV) { $venvRoots += $env:VIRTUAL_ENV } @@ -1663,7 +1900,7 @@ exit 0 try { $venvRoots += (Split-Path -Parent (Split-Path -Parent $env:UNSLOTH_SETUP_PYTHON)) } catch {} } if ($env:USERPROFILE) { $venvRoots += (Join-Path $env:USERPROFILE ".unsloth\studio\unsloth_studio") } - # A custom Studio home (UNSLOTH_STUDIO_HOME / STUDIO_HOME alias) moves the + # A custom Unsloth home (UNSLOTH_STUDIO_HOME / STUDIO_HOME alias) moves the # venv off the default path; seed it too or its hipInfo escapes the filter. $studioHomeEnv = if (-not [string]::IsNullOrWhiteSpace($env:UNSLOTH_STUDIO_HOME)) { $env:UNSLOTH_STUDIO_HOME.Trim() } elseif (-not [string]::IsNullOrWhiteSpace($env:STUDIO_HOME)) { $env:STUDIO_HOME.Trim() } else { $null } if ($studioHomeEnv) { @@ -1821,12 +2058,14 @@ exit 0 # (gfx120X/110X/1151/1150/103X); unknown names fall back to CPU. elseif ($ROCmGpuLabel) { $nameArchTable = @( - @{ P = "9070 XT|9080"; A = "gfx1201" } # RDNA 4 (RX 9070 XT / 9080) - @{ P = "9070|9060"; A = "gfx1200" } # RDNA 4 (RX 9070 / 9060) - @{ P = "8060S|8050S|8040S|Strix Halo|Ryzen AI Max|AI Max"; A = "gfx1151" } # RDNA 3.5 (Strix Halo: Radeon 8060S/8050S/8040S iGPU, Ryzen AI Max+) - @{ P = "890M|880M|860M|840M|Strix Point|Krackan|HX 37[05]|AI 9 HX|AI 9 36[05]|AI 7 35[05]|AI 5 34[05]|AI 7 PRO 35|AI 5 33"; A = "gfx1150" } # RDNA 3.5 (Strix/Krackan Point: Radeon 890M/880M iGPU, Ryzen AI 9 HX 370/375) - @{ P = "RX 7900|RX 7800|RX 7700(?!S)|PRO W7900|PRO W7800|PRO W7700"; A = "gfx1100" } # RDNA 3 desktop/workstation (Navi 31) - @{ P = "RX 7600|RX 7700S|RX 7650|PRO W7600|PRO W7500|PRO V710"; A = "gfx1102" } # RDNA 3 (Navi 33) + @{ P = "9070|9080"; A = "gfx1201" } # RDNA 4 (Navi 48: RX 9070 XT / 9070 GRE / 9070 / 9080) + @{ P = "9060"; A = "gfx1200" } # RDNA 4 (Navi 44: RX 9060 XT / 9060) + @{ P = "8065S|8060S|8050S|8040S|Strix Halo|Ryzen AI Max|AI Max"; A = "gfx1151" } # RDNA 3.5 (Strix Halo + Gorgon Halo: Radeon 8065S/8060S/8050S/8040S iGPU, Ryzen AI Max / Max+) + @{ P = "890M|880M|Strix Point|HX 37[05]|AI 9 HX|AI 9 36[05]"; A = "gfx1150" } # RDNA 3.5 (Strix Point: Radeon 890M/880M, Ryzen AI 9 HX 370/375) + @{ P = "860M|840M|Krackan|AI 7 35[05]|AI 5 34[05]|AI 7 PRO 35|AI 5 33"; A = "gfx1152" } # RDNA 3.5 (Krackan Point: Radeon 860M/840M, Ryzen AI 7 350 / AI 5 340) + @{ P = "RX 7900|PRO W7900|PRO W7800"; A = "gfx1100" } # RDNA 3 desktop/workstation (Navi 31) + @{ P = "RX 7800|RX 7700(?!S)|PRO W7700|PRO V710"; A = "gfx1101" } # RDNA 3 (Navi 32) + @{ P = "RX 7600|RX 7700S|RX 7650|PRO W7600|PRO W7500"; A = "gfx1102" } # RDNA 3 (Navi 33) @{ P = "780M|760M|740M|Phoenix|Hawk Point|Z1 Extreme|Z2 Extreme"; A = "gfx1103" } # RDNA 3 iGPU (Phoenix / Hawk Point) @{ P = "RX 6900|RX 6800|RX 6750|RX 6700|PRO W6800|PRO W6900"; A = "gfx1030" } # RDNA 2 (Navi 21) -- gfx103X family @{ P = "RX 6650|RX 6600|PRO W6600|PRO W6650"; A = "gfx1032" } # RDNA 2 (Navi 23) -- gfx103X family @@ -1942,7 +2181,7 @@ exit 0 substep " Ensure the ROCm compute driver is installed alongside the display driver:" "Yellow" substep " https://rocm.docs.amd.com/en/latest/deploy/windows/index.html" "Yellow" } elseif ($ROCmGfxArch) { - # Known arch: Studio setup installs AMD's bundled-runtime ROCm PyTorch wheels + # Known arch: Unsloth setup installs AMD's bundled-runtime ROCm PyTorch wheels # (repo.amd.com), which ship their own runtime -- HIP SDK optional. step "gpu" "AMD ROCm ($ROCmGfxArch)" "Cyan" substep "Detected: $ROCmGpuLabel" "Cyan" @@ -1960,10 +2199,31 @@ exit 0 # On an AMD GPU (no NVIDIA), surface the optional WSL-ROCm driver hint. if (-not $HasNvidiaSmi -and ($ROCmGfxArch -or $ROCmGpuLabel)) { Show-AmdWslDriverHint } + # Trim trailing slashes from the URL PATH only, preserving ?query / #fragment: a whole-URL + # TrimEnd corrupts a token ending in "/", a single strip leaves .../cu128// empty. Shared. + function Trim-IndexPathSlashes { + param([string]$Url) + $value = $Url.Trim() + $idx = $value.IndexOfAny([char[]]@('?', '#')) + if ($idx -lt 0) { + return $value.TrimEnd('/') + } + return $value.Substring(0, $idx).TrimEnd('/') + $value.Substring($idx) + } + # ── Choose the correct PyTorch index URL based on driver CUDA version ── # Mirrors Get-PytorchCudaTag in setup.ps1. function Get-TorchIndexUrl { $baseUrl = if ($env:UNSLOTH_PYTORCH_MIRROR) { $env:UNSLOTH_PYTORCH_MIRROR.TrimEnd('/') } else { "https://download.pytorch.org/whl" } + # Explicit pin -- skip ALL GPU probing (headless / CI / cross-install). + # UNSLOTH_TORCH_INDEX_URL wins (full URL, verbatim); _FAMILY is the leaf appended + # to the mirror base. Matches install.sh / install_python_stack.py. + if (-not [string]::IsNullOrWhiteSpace($env:UNSLOTH_TORCH_INDEX_URL)) { + return (Trim-IndexPathSlashes $env:UNSLOTH_TORCH_INDEX_URL) + } + if (-not [string]::IsNullOrWhiteSpace($env:UNSLOTH_TORCH_INDEX_FAMILY)) { + return "$baseUrl/$($env:UNSLOTH_TORCH_INDEX_FAMILY.Trim().Trim('/'))" + } if (-not $NvidiaSmiExe) { return "$baseUrl/cpu" } try { $output = Invoke-NvidiaSmiBounded $NvidiaSmiExe @@ -1984,6 +2244,27 @@ exit 0 return "$baseUrl/cu126" } + # Strip userinfo AND query/fragment so an authenticated pin never leaks. Shared with + # _strip_index_url_credentials (install.sh / py / setup.ps1). + function Remove-IndexUrlCredentials { + param([string]$Url) + # Ordinal, not culture-aware: on non-English locales (e.g. th-TH) linguistic + # IndexOf treats "://" as ignorable, mis-locates it, and crashes Substring (issue #7279). + $sep = $Url.IndexOf('://', [System.StringComparison]::Ordinal) + if ($sep -lt 0) { return $Url } + $scheme = $Url.Substring(0, $sep) + $rest = $Url.Substring($sep + 3) + # Drop query / fragment (may hold auth tokens). + $q = $rest.IndexOfAny([char[]]('?', '#')) + if ($q -ge 0) { $rest = $rest.Substring(0, $q) } + $slash = $rest.IndexOf('/', [System.StringComparison]::Ordinal) + $authority = if ($slash -ge 0) { $rest.Substring(0, $slash) } else { $rest } + $at = $authority.LastIndexOf('@', [System.StringComparison]::Ordinal) + $host_ = if ($at -ge 0) { $authority.Substring($at + 1) } else { $authority } + if ($slash -ge 0) { return "${scheme}://${host_}$($rest.Substring($slash))" } + return "${scheme}://${host_}" + } + # ── Torch flavor helpers (to repair a stale CPU / wrong-CUDA wheel) ── # torch.__version__ -> flavor tag (cuXXX / rocm / cpu); untagged wheel = cpu, # matching setup.ps1's stale-venv parse. @@ -2002,11 +2283,13 @@ exit 0 param([string]$TorchIndexUrl, [string]$ROCmIndexUrl) if (-not [string]::IsNullOrWhiteSpace($ROCmIndexUrl)) { return 'rocm' } if ([string]::IsNullOrWhiteSpace($TorchIndexUrl)) { return $null } - $leaf = ($TorchIndexUrl.TrimEnd('/') -split '/')[-1].ToLowerInvariant() + # Drop query/fragment first so .../cu128?token=x classifies as cu128 (else it reinstalls every run). + $leaf = (($TorchIndexUrl -split '[?#]', 2)[0].TrimEnd('/') -split '/')[-1].ToLowerInvariant() if ($leaf -match '^cu\d+$') { return $leaf } if ($leaf -eq 'cpu') { return 'cpu' } if ($leaf -match '^rocm') { return 'rocm' } - if ($leaf -match '^gfx') { return 'rocm' } + # gfx must be followed by a digit (an architecture leaf); gfx-private is custom. + if ($leaf -match '^gfx[0-9]') { return 'rocm' } return $null } @@ -2041,6 +2324,10 @@ exit 0 } catch { return $null } } + # An explicit pin is authoritative: the AMD ROCm reroute below must not rewrite it + # (e.g. a deliberate cpu pin on an AMD host). + $TorchIndexPinned = (-not [string]::IsNullOrWhiteSpace($env:UNSLOTH_TORCH_INDEX_URL)) -or ` + (-not [string]::IsNullOrWhiteSpace($env:UNSLOTH_TORCH_INDEX_FAMILY)) $TorchIndexUrl = Get-TorchIndexUrl # ── GPU arch → newest compatible Windows ROCm wheel release ── @@ -2052,13 +2339,20 @@ exit 0 # Override with UNSLOTH_ROCM_WINDOWS_MIRROR for air-gapped / mirror installs. $ROCmIndexUrl = $null $ROCmTorchFloor = $null - if (($HasROCm -or $ROCmGfxArch) -and $TorchIndexUrl -like "*/cpu" -and -not $SkipTorch) { + $PinnedRocmVisionSpec = $null + $PinnedRocmAudioSpec = $null + if (-not $TorchIndexPinned -and ($HasROCm -or $ROCmGfxArch) -and $TorchIndexUrl -like "*/cpu" -and -not $SkipTorch) { $amdIndexBase = if ($env:UNSLOTH_ROCM_WINDOWS_MIRROR) { $env:UNSLOTH_ROCM_WINDOWS_MIRROR.TrimEnd('/') } else { "https://repo.amd.com/rocm/whl" } $archFamilyMap = @{ "gfx1201" = "gfx120X-all"; "gfx1200" = "gfx120X-all" # RDNA 4 "gfx1151" = "gfx1151"; "gfx1150" = "gfx1150" # RDNA 3.5 (Strix Halo/Point) + "gfx1152" = "gfx1152" # RDNA 3.5 (Krackan Point) "gfx1103" = "gfx110X-all"; "gfx1102" = "gfx110X-all" # RDNA 3 "gfx1101" = "gfx110X-all"; "gfx1100" = "gfx110X-all" + "gfx1036" = "gfx103X-all"; "gfx1035" = "gfx103X-all" # RDNA 2 (RX 6000) + "gfx1034" = "gfx103X-all"; "gfx1033" = "gfx103X-all" + "gfx1032" = "gfx103X-all"; "gfx1031" = "gfx103X-all" + "gfx1030" = "gfx103X-all" "gfx90a" = "gfx90a"; "gfx908" = "gfx908" # MI200/MI100 } # gfx120X (RDNA 4) and gfx1151/gfx1150 (Strix) have a null-pointer bug in @@ -2074,6 +2368,7 @@ exit 0 $torchFloorMap = @{ "gfx1201" = "torch>=2.11.0,<2.12.0"; "gfx1200" = "torch>=2.11.0,<2.12.0" "gfx1151" = "torch>=2.11.0,<2.12.0"; "gfx1150" = "torch>=2.11.0,<2.12.0" + "gfx1152" = "torch>=2.11.0,<2.12.0" } # Companion ranges track the torch ceiling so pip resolves a consistent # trio on AMD's per-arch index (each published independently). Mirrors @@ -2081,10 +2376,12 @@ exit 0 $torchvisionFloorMap = @{ "gfx1201" = "torchvision>=0.26.0,<0.27.0"; "gfx1200" = "torchvision>=0.26.0,<0.27.0" "gfx1151" = "torchvision>=0.26.0,<0.27.0"; "gfx1150" = "torchvision>=0.26.0,<0.27.0" + "gfx1152" = "torchvision>=0.26.0,<0.27.0" } $torchaudioFloorMap = @{ "gfx1201" = "torchaudio>=2.11.0,<2.12.0"; "gfx1200" = "torchaudio>=2.11.0,<2.12.0" "gfx1151" = "torchaudio>=2.11.0,<2.12.0"; "gfx1150" = "torchaudio>=2.11.0,<2.12.0" + "gfx1152" = "torchaudio>=2.11.0,<2.12.0" } $archFamily = if ($ROCmGfxArch -and $archFamilyMap.ContainsKey($ROCmGfxArch)) { $archFamilyMap[$ROCmGfxArch] } else { $null } if ($archFamily) { @@ -2102,6 +2399,32 @@ exit 0 } } + # A gfx*/rocm pin skips the auto-reroute above, but the generic CPU/CUDA install below + # would use torch>=2.4,<2.11 and pull a known-bad wheel on the gfx115x/gfx120x/rocm>=7.2 + # indexes (the _grouped_mm bug). Route a pinned ROCm index through the ROCm path. + if ($TorchIndexPinned -and -not $ROCmIndexUrl -and -not $SkipTorch) { + $_pinLeaf = (($TorchIndexUrl -split '[?#]', 2)[0].TrimEnd('/') -split '/')[-1].ToLower() + $_pinRocm211 = $false + # Anchor ($) so a suffixed custom leaf (rocm7.2-private) falls through to verbatim. + if ($_pinLeaf -match '^rocm(\d+)\.(\d+)$') { + # Only KNOWN-2.11 rocm (rocm7.2) gets the floor. Matches Test-RocmKnown211Version. + $_pinRocm211 = ([int]$Matches[1] -eq 7 -and [int]$Matches[2] -eq 2) + } + # Only the 2.11-allowlist gfx arches need the floor; others publish <2.11 and stay bare. + $_pinGfx211 = @('gfx120x-all', 'gfx1151', 'gfx1150', 'gfx1152') -contains $_pinLeaf + if ($_pinGfx211 -or $_pinRocm211) { + $ROCmIndexUrl = $TorchIndexUrl + $ROCmTorchFloor = "torch>=2.11.0,<2.12.0" + $PinnedRocmVisionSpec = "torchvision>=0.26.0,<0.27.0" + $PinnedRocmAudioSpec = "torchaudio>=2.11.0,<2.12.0" + substep "pinned ROCm index ($_pinLeaf) -- enforcing $ROCmTorchFloor" "Cyan" + } elseif ($_pinLeaf -match '^gfx[0-9]' -or $_pinLeaf -match '^rocm[0-9]+(\.[0-9]+)?$') { + # Other gfx / older rocm (<=7.1) ship torch <2.11; route via the ROCm path with + # bare specs. Only EXACT rocm/gfx* are families; a suffixed leaf is verbatim. + $ROCmIndexUrl = $TorchIndexUrl + } + } + if ($ROCmIndexUrl) { $TorchIndexFamily = "rocm" } else { @@ -2164,14 +2487,14 @@ exit 0 } if ($_Migrated) { - # Migrated env: force-reinstall unsloth+unsloth-zoo to ensure clean state - # in the new venv location, while preserving existing torch/CUDA + # Migrated env: force-reinstall unsloth+unsloth-zoo for a clean state, preserving + # existing torch/CUDA unless the flavor repair below re-lands it. Write-TauriLog "STEP" "Installing unsloth" substep "upgrading unsloth in migrated environment..." if ($SkipTorch) { # No-torch: install unsloth + unsloth-zoo with --no-deps, then # runtime deps (typer, safetensors, transformers, etc.) with --no-deps. - $baseInstallExit = Invoke-InstallCommandRetry -Label "install unsloth (migrated no-torch)" { uv pip install --python $VenvPython --no-deps --reinstall-package unsloth --reinstall-package unsloth-zoo "unsloth>=2026.7.2" "unsloth-zoo>=2026.7.2" } + $baseInstallExit = Invoke-InstallCommandRetry -Label "install unsloth (migrated no-torch)" { uv pip install --python $VenvPython --no-deps --reinstall-package unsloth --reinstall-package unsloth-zoo "unsloth>=2026.7.5" "unsloth-zoo>=2026.7.6" } if ($baseInstallExit -eq 0) { # Resolve pydantic WITH deps so pip pins pydantic-core # to the matching version (no-torch-runtime.txt below @@ -2185,7 +2508,7 @@ exit 0 } } } else { - $baseInstallExit = Invoke-InstallCommandRetry -Label "install unsloth (migrated)" { uv pip install --python $VenvPython --reinstall-package unsloth --reinstall-package unsloth-zoo "unsloth>=2026.7.2" "unsloth-zoo>=2026.7.2" } + $baseInstallExit = Invoke-InstallCommandRetry -Label "install unsloth (migrated)" { uv pip install --python $VenvPython --reinstall-package unsloth --reinstall-package unsloth-zoo "unsloth>=2026.7.5" "unsloth-zoo>=2026.7.6" } } if ($baseInstallExit -ne 0) { Write-Host "[ERROR] Failed to install unsloth (exit code $baseInstallExit)" -ForegroundColor Red @@ -2193,7 +2516,7 @@ exit 0 } if ($StudioLocalInstall) { substep "overlaying local repo (editable)..." - $overlayExit = Invoke-InstallCommand { uv pip install --python $VenvPython -e $RepoRoot --no-deps } + $overlayExit = Invoke-InstallCommand -Label "overlay local repo" { uv pip install --python $VenvPython -e $RepoRoot --no-deps } if ($overlayExit -ne 0) { Write-Host "[ERROR] Failed to overlay local repo (exit code $overlayExit)" -ForegroundColor Red return (Exit-InstallFailure "Failed to overlay local repo (exit code $overlayExit)" $overlayExit) @@ -2210,22 +2533,24 @@ exit 0 substep "skipping PyTorch (--no-torch flag set)." "Yellow" } elseif ($ROCmIndexUrl) { Write-TauriLog "STEP" "Installing PyTorch (AMD ROCm Windows)" - substep "installing PyTorch from $ROCmIndexUrl..." + substep "installing PyTorch from $(Remove-IndexUrlCredentials $ROCmIndexUrl)..." $torchSpec = if ($ROCmTorchFloor) { $ROCmTorchFloor } else { "torch" } # Pin the companions to match $torchSpec; bare names can resolve an # ABI-incompatible torchvision/torchaudio on AMD's per-arch index. - $visionSpec = if ($ROCmGfxArch -and $torchvisionFloorMap.ContainsKey($ROCmGfxArch)) { $torchvisionFloorMap[$ROCmGfxArch] } else { "torchvision" } - $audioSpec = if ($ROCmGfxArch -and $torchaudioFloorMap.ContainsKey($ROCmGfxArch)) { $torchaudioFloorMap[$ROCmGfxArch] } else { "torchaudio" } + $visionSpec = if ($PinnedRocmVisionSpec) { $PinnedRocmVisionSpec } elseif ($ROCmGfxArch -and $torchvisionFloorMap -and $torchvisionFloorMap.ContainsKey($ROCmGfxArch)) { $torchvisionFloorMap[$ROCmGfxArch] } else { "torchvision" } + $audioSpec = if ($PinnedRocmAudioSpec) { $PinnedRocmAudioSpec } elseif ($ROCmGfxArch -and $torchaudioFloorMap -and $torchaudioFloorMap.ContainsKey($ROCmGfxArch)) { $torchaudioFloorMap[$ROCmGfxArch] } else { "torchaudio" } $torchInstallExit = Invoke-InstallCommandRetry -Label "install PyTorch (AMD ROCm)" { uv pip install --python $VenvPython --force-reinstall --default-index $ROCmIndexUrl $torchSpec $visionSpec $audioSpec } if ($torchInstallExit -ne 0) { - # Transient AMD-index failure: fall back to a CPU base so the install - # still completes; Studio setup retries ROCm afterwards. - substep "ROCm PyTorch install failed (exit $torchInstallExit); using a CPU base, Studio setup retries ROCm." "Yellow" + # Transient AMD-index failure: fall back to a CPU base (Unsloth setup retries + # ROCm). Use an explicit CPU index -- for a pinned ROCm index $TorchIndexUrl IS + # the ROCm mirror, so reusing it would just retry it. + $CpuFallbackIndexUrl = if ($env:UNSLOTH_PYTORCH_MIRROR) { "$($env:UNSLOTH_PYTORCH_MIRROR.TrimEnd('/'))/cpu" } else { "https://download.pytorch.org/whl/cpu" } + substep "ROCm PyTorch install failed (exit $torchInstallExit); using a CPU base, Unsloth setup retries ROCm." "Yellow" # --force-reinstall: a failed ROCm install can leave an unpinned ROCm # torch (e.g. 2.10.0+rocm on gfx110X/gfx90a) that still satisfies the CPU # torch>= range, so without it uv would keep the ROCm build and only swap # the companions -- a mismatched venv the flavor-repair block won't fix. - $torchInstallExit = Invoke-InstallCommandRetry -Label "install PyTorch (CPU fallback)" { uv pip install --python $VenvPython --force-reinstall "torch>=2.4,<2.11.0" torchvision torchaudio --default-index $TorchIndexUrl } + $torchInstallExit = Invoke-InstallCommandRetry -Label "install PyTorch (CPU fallback)" { uv pip install --python $VenvPython --force-reinstall "torch>=2.4,<2.11.0" "torchvision>=0.19,<0.26.0" "torchaudio>=2.4,<2.11.0" --default-index $CpuFallbackIndexUrl } if ($torchInstallExit -ne 0) { Write-Host "[ERROR] Failed to install PyTorch (ROCm and CPU base both failed, exit code $torchInstallExit)" -ForegroundColor Red return (Exit-InstallFailure "Failed to install PyTorch (exit code $torchInstallExit)" $torchInstallExit) @@ -2238,8 +2563,27 @@ exit 0 } } else { Write-TauriLog "STEP" "Installing PyTorch" - substep "installing PyTorch ($TorchIndexUrl)..." - $torchInstallExit = Invoke-InstallCommandRetry -Label "install PyTorch" { uv pip install --python $VenvPython "torch>=2.4,<2.11.0" torchvision torchaudio --default-index $TorchIndexUrl } + # Windows on ARM lacks only torchaudio (whl/cpu win_arm64: torch 42, + # torchvision 60, torchaudio 0), so drop that pin instead of aborting. Ask the + # interpreter, not PROCESSOR_ARCHITECTURE; reached when no x64 Python exists. + $VenvPlatform = "" + try { + $VenvPlatform = (& $VenvPython -c "import sysconfig; print(sysconfig.get_platform())" 2>$null | Out-String).Trim().ToLowerInvariant() + } catch { $VenvPlatform = "" } + substep "installing PyTorch ($(Remove-IndexUrlCredentials $TorchIndexUrl))..." + # Bound the companions to the capped torch on EVERY index, cu + # families included: torchaudio 2.11 dropped its exact torch pin from + # the wheel metadata, so a bare companion next to torch<2.11 can + # resolve a mismatched 2.11.0 build. Mirrors install.sh. + $_pinVisionSpec = "torchvision>=0.19,<0.26.0" + $_pinAudioSpec = "torchaudio>=2.4,<2.11.0" + $_torchSpecs = @("torch>=2.4,<2.11.0", $_pinVisionSpec, $_pinAudioSpec) + if ($VenvPlatform -eq "win-arm64") { + substep "windows on arm: skipping torchaudio (upstream publishes no" + substep "win_arm64 wheel); torch and torchvision install normally." + $_torchSpecs = @("torch>=2.4,<2.11.0", $_pinVisionSpec) + } + $torchInstallExit = Invoke-InstallCommandRetry -Label "install PyTorch" { uv pip install --python $VenvPython @_torchSpecs --default-index $TorchIndexUrl } if ($torchInstallExit -ne 0) { Write-Host "[ERROR] Failed to install PyTorch (exit code $torchInstallExit)" -ForegroundColor Red return (Exit-InstallFailure "Failed to install PyTorch (exit code $torchInstallExit)" $torchInstallExit) @@ -2251,7 +2595,7 @@ exit 0 if ($SkipTorch) { # No-torch: install unsloth + unsloth-zoo with --no-deps, then # runtime deps (typer, safetensors, transformers, etc.) with --no-deps. - $baseInstallExit = Invoke-InstallCommandRetry -Label "install unsloth (no-torch)" { uv pip install --python $VenvPython --no-deps --upgrade-package unsloth --upgrade-package unsloth-zoo "unsloth>=2026.7.2" "unsloth-zoo>=2026.7.2" } + $baseInstallExit = Invoke-InstallCommandRetry -Label "install unsloth (no-torch)" { uv pip install --python $VenvPython --no-deps --upgrade-package unsloth --upgrade-package unsloth-zoo "unsloth>=2026.7.5" "unsloth-zoo>=2026.7.6" } if ($baseInstallExit -eq 0) { # Same pydantic-with-deps trick as the migrated branch. $baseInstallExit = Invoke-InstallCommandRetry -Label "install pydantic" { uv pip install --python $VenvPython pydantic } @@ -2263,7 +2607,7 @@ exit 0 } } } elseif ($StudioLocalInstall) { - $baseInstallExit = Invoke-InstallCommandRetry -Label "install unsloth (local)" { uv pip install --python $VenvPython --upgrade-package unsloth "unsloth>=2026.7.2" "unsloth-zoo>=2026.7.2" } + $baseInstallExit = Invoke-InstallCommandRetry -Label "install unsloth (local)" { uv pip install --python $VenvPython --upgrade-package unsloth "unsloth>=2026.7.5" "unsloth-zoo>=2026.7.6" } } else { $baseInstallExit = Invoke-InstallCommandRetry -Label "install unsloth" { uv pip install --python $VenvPython --upgrade-package unsloth -- "$PackageName" } } @@ -2274,7 +2618,7 @@ exit 0 if ($StudioLocalInstall) { substep "overlaying local repo (editable)..." - $overlayExit = Invoke-InstallCommand { uv pip install --python $VenvPython -e $RepoRoot --no-deps } + $overlayExit = Invoke-InstallCommand -Label "overlay local repo" { uv pip install --python $VenvPython -e $RepoRoot --no-deps } if ($overlayExit -ne 0) { Write-Host "[ERROR] Failed to overlay local repo (exit code $overlayExit)" -ForegroundColor Red return (Exit-InstallFailure "Failed to overlay local repo (exit code $overlayExit)" $overlayExit) @@ -2291,13 +2635,13 @@ exit 0 Write-TauriLog "STEP" "Installing unsloth" substep "installing unsloth (this may take a few minutes)..." if ($StudioLocalInstall) { - $baseInstallExit = Invoke-InstallCommandRetry -Label "install unsloth (auto torch backend)" { uv pip install --python $VenvPython "unsloth-zoo>=2026.7.2" "unsloth>=2026.7.2" --torch-backend=auto } + $baseInstallExit = Invoke-InstallCommandRetry -Label "install unsloth (auto torch backend)" { uv pip install --python $VenvPython "unsloth-zoo>=2026.7.6" "unsloth>=2026.7.5" --torch-backend=auto } if ($baseInstallExit -ne 0) { Write-Host "[ERROR] Failed to install unsloth (exit code $baseInstallExit)" -ForegroundColor Red return (Exit-InstallFailure "Failed to install unsloth (exit code $baseInstallExit)" $baseInstallExit) } substep "overlaying local repo (editable)..." - $overlayExit = Invoke-InstallCommand { uv pip install --python $VenvPython -e $RepoRoot --no-deps } + $overlayExit = Invoke-InstallCommand -Label "overlay local repo" { uv pip install --python $VenvPython -e $RepoRoot --no-deps } if ($overlayExit -ne 0) { Write-Host "[ERROR] Failed to overlay local repo (exit code $overlayExit)" -ForegroundColor Red return (Exit-InstallFailure "Failed to overlay local repo (exit code $overlayExit)" $overlayExit) @@ -2317,6 +2661,13 @@ exit 0 } } + $installedPackageVersion = (& $VenvPython -c "from importlib.metadata import version; import sys; print(version(sys.argv[1]))" $PackageName 2>$null | Out-String).Trim() + if ($LASTEXITCODE -eq 0 -and $installedPackageVersion) { + step $PackageName "$installedPackageVersion installed" + } else { + substep "[WARN] installed $PackageName version could not be determined" "Yellow" + } + # ── Enforce the installed torch flavor matches the detected GPU build ── # PEP 440 ignores the +cpu/+cuXXX/+rocm local label in a version range, so uv # keeps a stale torch==X+cpu against a CUDA index and setup.ps1 then loops on @@ -2335,10 +2686,10 @@ exit 0 $rocmSpec = if ($ROCmTorchFloor) { $ROCmTorchFloor } else { "torch" } # Pin companions like the fresh ROCm path (bare names can pull an # ABI-incompatible torchvision/torchaudio from the per-arch index). - $visionSpec = if ($ROCmGfxArch -and $torchvisionFloorMap.ContainsKey($ROCmGfxArch)) { $torchvisionFloorMap[$ROCmGfxArch] } else { "torchvision" } - $audioSpec = if ($ROCmGfxArch -and $torchaudioFloorMap.ContainsKey($ROCmGfxArch)) { $torchaudioFloorMap[$ROCmGfxArch] } else { "torchaudio" } + $visionSpec = if ($PinnedRocmVisionSpec) { $PinnedRocmVisionSpec } elseif ($ROCmGfxArch -and $torchvisionFloorMap -and $torchvisionFloorMap.ContainsKey($ROCmGfxArch)) { $torchvisionFloorMap[$ROCmGfxArch] } else { "torchvision" } + $audioSpec = if ($PinnedRocmAudioSpec) { $PinnedRocmAudioSpec } elseif ($ROCmGfxArch -and $torchaudioFloorMap -and $torchaudioFloorMap.ContainsKey($ROCmGfxArch)) { $torchaudioFloorMap[$ROCmGfxArch] } else { "torchaudio" } substep "PyTorch flavor mismatch (installed $installedTorchTag, need ROCm) -- reinstalling correct build..." "Yellow" - $torchFixExit = Invoke-InstallCommand { uv pip install --python $VenvPython --force-reinstall --default-index $ROCmIndexUrl $rocmSpec $visionSpec $audioSpec } + $torchFixExit = Invoke-InstallCommand -Label "reinstall PyTorch (ROCm)" { uv pip install --python $VenvPython --force-reinstall --default-index $ROCmIndexUrl $rocmSpec $visionSpec $audioSpec } if ($torchFixExit -ne 0) { Write-Host "[ERROR] Failed to reinstall PyTorch with the correct ROCm build (exit code $torchFixExit)" -ForegroundColor Red return (Exit-InstallFailure "Failed to reinstall PyTorch (ROCm) (exit code $torchFixExit)" $torchFixExit) @@ -2347,7 +2698,7 @@ exit 0 } elseif ($expectedTorchTag -ne 'rocm') { # CUDA: stale +cpu (or wrong cuXXX) against a CUDA index -> reinstall triplet. substep "PyTorch flavor mismatch (installed $installedTorchTag, need $expectedTorchTag) -- reinstalling correct build..." "Yellow" - $torchFixExit = Invoke-InstallCommand { uv pip install --python $VenvPython "torch>=2.4,<2.11.0" torchvision torchaudio --default-index $TorchIndexUrl --reinstall-package torch --reinstall-package torchvision --reinstall-package torchaudio } + $torchFixExit = Invoke-InstallCommand -Label "reinstall PyTorch ($expectedTorchTag)" { uv pip install --python $VenvPython "torch>=2.4,<2.11.0" "torchvision>=0.19,<0.26.0" "torchaudio>=2.4,<2.11.0" --default-index $TorchIndexUrl --reinstall-package torch --reinstall-package torchvision --reinstall-package torchaudio } if ($torchFixExit -ne 0) { Write-Host "[ERROR] Failed to reinstall PyTorch with the correct CUDA build (exit code $torchFixExit)" -ForegroundColor Red return (Exit-InstallFailure "Failed to reinstall PyTorch ($expectedTorchTag) (exit code $torchFixExit)" $torchFixExit) @@ -2422,7 +2773,7 @@ exit 0 Write-TauriLog "ERROR" "unsloth CLI was not installed correctly" Write-Host "[ERROR] unsloth CLI was not installed correctly." -ForegroundColor Red Write-Host " Expected: $UnslothExe" -ForegroundColor Yellow - Write-Host " This usually means an older unsloth version was installed that does not include the Studio CLI." -ForegroundColor Yellow + Write-Host " This usually means an older unsloth version was installed that does not include the Unsloth CLI." -ForegroundColor Yellow Write-Host " Try re-running the installer or see: https://github.com/unslothai/unsloth?tab=readme-ov-file#-quickstart" -ForegroundColor Yellow return (Exit-InstallFailure "unsloth CLI was not installed correctly") } @@ -2448,6 +2799,9 @@ exit 0 # an inherited value would put llama.cpp in the wrong place. $previousUnslothStudioHome = $env:UNSLOTH_STUDIO_HOME $hadPreviousUnslothStudioHome = ($null -ne $previousUnslothStudioHome) + $previousTauriMode = $env:UNSLOTH_TAURI_MODE + $hadPreviousTauriMode = ($null -ne $previousTauriMode) + $env:UNSLOTH_TAURI_MODE = if ($TauriMode) { "1" } else { "0" } if ($StudioRedirectMode -eq 'env') { $env:UNSLOTH_STUDIO_HOME = $StudioHome } else { @@ -2477,14 +2831,22 @@ exit 0 } else { Remove-Item Env:UNSLOTH_STUDIO_HOME -ErrorAction SilentlyContinue } + if ($hadPreviousTauriMode) { + $env:UNSLOTH_TAURI_MODE = $previousTauriMode + } else { + Remove-Item Env:UNSLOTH_TAURI_MODE -ErrorAction SilentlyContinue + } Remove-Item Env:UNSLOTH_LOCAL_LLAMA_CPP_DIR -ErrorAction SilentlyContinue Remove-Item Env:UNSLOTH_INSTALL_ROLLBACK_MANAGED -ErrorAction SilentlyContinue Remove-Item Env:UNSLOTH_SETUP_PYTHON -ErrorAction SilentlyContinue } if ($setupExit -ne 0) { - Write-Host "[ERROR] unsloth studio setup failed (exit code $setupExit)" -ForegroundColor Red + if (-not $TauriMode) { + Write-Host "[ERROR] unsloth studio setup failed (exit code $setupExit)" -ForegroundColor Red + } return (Exit-InstallFailure "unsloth studio setup failed (exit code $setupExit)" $setupExit) } + Clear-TauriInstallError "studio setup completed" # ── Expose `unsloth` via a shim dir containing only unsloth.exe ── # We do NOT add the venv Scripts dir to PATH (it also holds python.exe @@ -2533,7 +2895,7 @@ exit 0 Write-Host " Move or remove it manually, then re-run the installer." -ForegroundColor Yellow throw "Cannot create unsloth launcher: $ShimExe is a directory." } - # try/catch: if unsloth.exe is locked (Studio running), keep the old shim. + # try/catch: if unsloth.exe is locked (Unsloth running), keep the old shim. $shimUpdated = $false try { if (Test-Path -LiteralPath $ShimExe) { Remove-Item -LiteralPath $ShimExe -Force -ErrorAction Stop } @@ -2551,7 +2913,7 @@ exit 0 if (Test-Path -LiteralPath $ShimExe) { Write-Host "[WARN] Could not refresh unsloth launcher at $ShimExe." -ForegroundColor Yellow Write-Host " This usually means a running 'unsloth studio' process still holds the file open." -ForegroundColor Yellow - Write-Host " Close Studio and re-run the installer to pick up the latest launcher." -ForegroundColor Yellow + Write-Host " Close Unsloth and re-run the installer to pick up the latest launcher." -ForegroundColor Yellow Write-Host " Continuing with the existing launcher." -ForegroundColor Yellow } else { Write-Host "[WARN] Could not create unsloth launcher at $ShimExe" -ForegroundColor Yellow @@ -2572,6 +2934,13 @@ exit 0 } Refresh-SessionPath # sync current session with registry Complete-StudioVenvRollback + $studioVenvReplacementCommitted = $true + Remove-StaleStudioVenvRollbacks + } finally { + if (-not $studioVenvReplacementCommitted) { + Restore-StudioVenvRollback + } + } # Env-mode session export AFTER Refresh-SessionPath; otherwise a legacy # User PATH entry (Machine > User > current $env:Path) would win. @@ -2616,7 +2985,7 @@ exit 0 # Diagnostic only; never block install on a probe failure. } - # In interactive terminals, ask the user before starting Studio unless the + # In interactive terminals, ask the user before starting Unsloth unless the # caller explicitly disabled the post-install prompt. # In non-interactive environments (CI, Docker) just print instructions. $IsInteractive = (-not $SkipAutostart) -and [Environment]::UserInteractive -and (-not [Console]::IsInputRedirected) diff --git a/install.sh b/install.sh index f277d0cbfd..166beeb52c 100755 --- a/install.sh +++ b/install.sh @@ -19,6 +19,17 @@ # SPDX-License-Identifier: AGPL-3.0-only # Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0 set -e +# ── Why the installer lives in a function ── +# Under `curl ... | sh`, sh is the pipe READER. This file is ~150KB, so a top-level +# `exit` left most of it unread, the write end failed, and curl tacked +# "(56) Failure writing output to destination" onto our own error message. Wrapping +# the body forces sh to parse to the closing brace first, so the pipe always drains +# (install.ps1 has always had this shape). +# +# Body is deliberately NOT reindented: reflowing 4000+ lines would bury the change, +# and `exit` still exits the shell from inside a function. Do not add +# `exec < /dev/null`: for a piped shell that closes the script's own source. +_unsloth_main() { # ── Output style (aligned with studio/setup.sh) ── RULE="" @@ -97,7 +108,7 @@ if [ "$_VERBOSE" = true ]; then export UNSLOTH_VERBOSE=1 fi -# Custom Studio roots are not supported with --tauri (desktop app still +# Custom Unsloth roots are not supported with --tauri (desktop app still # resolves ~/.unsloth/studio). Pass through if the override == legacy default. if [ "$TAURI_MODE" = true ]; then _tauri_override_var="" @@ -159,26 +170,85 @@ run_maybe_quiet() { fi } +# Trim trailing slashes from the URL PATH only, preserving ?query / #fragment: a whole-URL +# strip corrupts a token ending in "/", a single strip leaves .../cu128// empty. Shared. +_trim_index_path_slashes() { + _tips_v="$1" + case "$_tips_v" in + *[?#]*) + _tips_head="${_tips_v%%[?#]*}" + _tips_tail="${_tips_v#"$_tips_head"}" + ;; + *) + _tips_head="$_tips_v" + _tips_tail="" + ;; + esac + while [ -n "$_tips_head" ] && [ "${_tips_head%/}" != "$_tips_head" ]; do + _tips_head="${_tips_head%/}" + done + printf '%s%s' "$_tips_head" "$_tips_tail" +} + +# Redact index-URL credentials (userinfo + ?query= + #fragment) from captured installer +# output before printing on failure; uv/pip errors echo the failing --index-url verbatim. +# Mirrors the other installers. Verbose mode streams uncaptured, so it isn't redacted. +_redact_install_output() { + sed -E \ + -e 's#(https?://)[^/@[:space:]`]+@#\1@#g' \ + -e 's#([?&][^=[:space:]&`]+)=[^&#[:space:]`]+#\1=#g' \ + -e 's|(https?://[^[:space:]`#]+)#[^[:space:]`]+|\1#|g' \ + "$@" +} + run_install_cmd() { _label="$1" shift - # Installer-pinned index installs (torch) must beat an inherited uv mirror - # (#6898): when we pass --default-index, neutralize every uv index env var so - # the pinned index wins. Other installs keep the user's mirror. + # Installer-pinned index installs (torch) must beat an inherited uv mirror (#6898): + # for --default-index, neutralize the uv index/backend/config vars (UV_TORCH_BACKEND + # redirects torch; UV_NO_CONFIG=1 + dropping UV_CONFIG_FILE stops a uv.toml/pyproject + # index outranking the CLI pin, uv 0.10). case " $* " in - *" --default-index "*) set -- env -u UV_DEFAULT_INDEX -u UV_INDEX_URL -u UV_INDEX -u UV_EXTRA_INDEX_URL "$@" ;; + *" --default-index "*) set -- env -u UV_DEFAULT_INDEX -u UV_INDEX_URL -u UV_INDEX -u UV_EXTRA_INDEX_URL -u UV_TORCH_BACKEND -u UV_FIND_LINKS -u UV_CONFIG_FILE UV_NO_CONFIG=1 "$@" ;; esac if _is_verbose; then - "$@" && return 0 - _rc=$? + # Stream through the redactor: uv echoes index URLs (credentials and + # all) in its errors, and verbose mode previously bypassed the + # redaction the quiet path applies. The rc file preserves the + # command's exit code across the pipe without relying on pipefail + # (this script runs under plain sh). + _rcf=$(mktemp) + tauri_stream_log stdout "OUTPUT_CLEAR" "$_label" + { + if "$@" 2>&1; then + _cmd_rc=0 + else + _cmd_rc=$? + fi + printf '%s' "$_cmd_rc" > "$_rcf" + } | _redact_install_output + _rc=$(cat "$_rcf" 2>/dev/null || echo 1) + rm -f "$_rcf" + _rc=${_rc:-1} + if [ "$_rc" -eq 0 ] 2>/dev/null; then + tauri_clear_install_error "$_label recovered" + return 0 + fi + tauri_stream_log stdout "ERROR_OUTPUT" "$_label failed (exit code $_rc)" step "error" "$_label failed (exit code $_rc)" "$C_ERR" >&2 return "$_rc" fi _log=$(mktemp) - "$@" >"$_log" 2>&1 && { rm -f "$_log"; return 0; } + tauri_stream_log stderr "OUTPUT_CLEAR" "$_label" + "$@" >"$_log" 2>&1 && { + rm -f "$_log" + tauri_clear_install_error "$_label recovered" + return 0 + } _rc=$? step "error" "$_label failed (exit code $_rc)" "$C_ERR" >&2 - cat "$_log" >&2 + _redact_install_output "$_log" >&2 + tauri_stream_log stderr "ERROR_OUTPUT" "$_label failed (exit code $_rc)" rm -f "$_log" return $_rc } @@ -217,10 +287,70 @@ run_install_cmd_retry() { done } -# 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. +# True when the runtime target is gfx906 (MI50/Radeon VII): the prebuilt AMD +# bitsandbytes wheel carries no gfx906 kernels, and force-reinstalling it would +# clobber a user's source-built bnb (the only 4-bit path on this arch) on every +# `studio update`. So skip the auto-install and leave whatever bnb is present. +# _gfx906_target is set during torch-index resolution; also honor an explicit +# UNSLOTH_ROCM_GFX_ARCH so a pinned-index install still skips. The override is +# normalized (gfx906:sramecc-:xnack- -> gfx906) so a copied HIP gcnArchName counts. +_is_gfx906_bnb_skip() { + [ "${_gfx906_target:-false}" = true ] && return 0 + _bnb_gfx_env=$(printf '%s' "${UNSLOTH_ROCM_GFX_ARCH:-}" | tr '[:upper:]' '[:lower:]' | tr -d '[:space:]') + _bnb_gfx_env=${_bnb_gfx_env%%:*} + [ "$_bnb_gfx_env" = "gfx906" ] && return 0 + # A pinned index (UNSLOTH_TORCH_INDEX_URL/_FAMILY) skips the reroute block that + # sets _gfx906_target, so a real gfx906 host with a pinned rocm6.3 index and no + # UNSLOTH_ROCM_GFX_ARCH would otherwise clobber a source-built bnb. Probe here + # in that gap; skip only when gfx906 is the SOLE distinct arch (mixed hosts + # opt in via the env var, mirroring the reroute block's de-dup rule). + if [ -z "$_bnb_gfx_env" ] && [ "${_torch_index_pinned:-false}" = true ]; then + _bnb_gfx_probe=$(_probe_amd_gfx_arch | awk 'NF && !seen[$0]++') + [ "$_bnb_gfx_probe" = "gfx906" ] && return 0 + fi + return 1 +} + +# `pip install unsloth` resolves its unconditional bitsandbytes dep to a generic +# CUDA wheel (no gfx906 kernels) once we skip the prebuilt one. Snapshot bnb before +# the unsloth install, then drop a freshly pulled wheel afterwards while leaving a +# pre-existing source build in place. +_gfx906_bnb_installed() { + "$_VENV_PY" -c "import importlib.util as u, sys; sys.exit(0 if u.find_spec('bitsandbytes') else 1)" >/dev/null 2>&1 +} +_gfx906_bnb_snapshot() { + _gfx906_bnb_absent_before=false + _is_gfx906_bnb_skip || return 0 + _gfx906_bnb_installed || _gfx906_bnb_absent_before=true +} +_gfx906_bnb_prune() { + _is_gfx906_bnb_skip || return 0 + [ "${_gfx906_bnb_absent_before:-false}" = true ] || return 0 + _gfx906_bnb_installed || return 0 + substep "gfx906: removing generic bitsandbytes pulled in as a dependency (no gfx906 kernels; build from source for 4-bit QLoRA)" "$C_WARN" + uv pip uninstall --python "$_VENV_PY" bitsandbytes >/dev/null 2>&1 \ + || "$_VENV_PY" -m pip uninstall -y bitsandbytes >/dev/null 2>&1 || true +} + +# Install bitsandbytes on AMD ROCm hosts. bnb <= 0.49.2 NaNs at 4-bit decode +# shape on every AMD GPU; the fix (bnb #1887) ships in continuous-release_main +# and, on PyPI, first in 0.50.0. Keep this floor in step with the amd extra in +# pyproject.toml and studio/install_python_stack.py. +_BNB_ROCM_PYPI_FALLBACK="bitsandbytes>=0.50.0" +# bitsandbytes ships no ROCm binary in its aarch64 wheel at any version: the PyPI +# 0.50.0 and continuous-release_main aarch64 wheels both carry only +# libbitsandbytes_cpu.so plus CUDA variants. So neither install path below gives +# aarch64 a 4-bit backend, and the messages must not claim one. Cf. gfx906. +_bnb_rocm_arch_has_binary() { + case "$_ARCH" in + aarch64|arm64) return 1 ;; + *) return 0 ;; + esac +} +_warn_bnb_no_rocm_binary() { + _bnb_rocm_arch_has_binary && return 0 + substep "[WARN] aarch64: bitsandbytes ships no ROCm kernels on this arch; 4-bit QLoRA needs a source build -- https://docs.unsloth.ai/get-started/install-and-update/amd" "$C_WARN" +} _install_bnb_rocm() { _label="$1" _venv_py="$2" @@ -235,9 +365,8 @@ _install_bnb_rocm() { _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. + # uv rejects the pre-release wheel: filename version (1.33.7rc0) does not + # match metadata (0.50.x.dev0). pip accepts it, 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 || \ @@ -253,18 +382,26 @@ _install_bnb_rocm() { --retries 8 --timeout 90 \ "$_bnb_whl_url" >"$_bnb_log" 2>&1; then rm -f "$_bnb_log" + _warn_bnb_no_rocm_binary return 0 fi _bnb_rc=$? if _is_verbose; then - cat "$_bnb_log" >&2 + _redact_install_output "$_bnb_log" >&2 fi rm -f "$_bnb_log" step "warning" "$_label (pre-release) failed (exit code $_bnb_rc)" "$C_WARN" >&2 - substep "[WARN] bnb pre-release install failed; falling back to PyPI (4-bit decode broken on ROCm)" "$C_WARN" + if _bnb_rocm_arch_has_binary; then + substep "[WARN] bnb pre-release install failed; falling back to PyPI $_BNB_ROCM_PYPI_FALLBACK, which carries the ROCm 4-bit fix" "$C_WARN" + else + substep "[WARN] bnb pre-release install failed; falling back to PyPI $_BNB_ROCM_PYPI_FALLBACK" "$C_WARN" + fi fi run_install_cmd "$_label (pypi fallback)" "$_venv_py" -m pip install \ - --force-reinstall --no-cache-dir --no-deps "bitsandbytes>=0.49.1" + --force-reinstall --no-cache-dir --no-deps "$_BNB_ROCM_PYPI_FALLBACK" + _bnb_pypi_rc=$? + _warn_bnb_no_rocm_binary + return $_bnb_pypi_rc } if [ "$_next_is_package" = true ]; then @@ -298,6 +435,34 @@ tauri_log() { fi } +tauri_stream_log() { + _tsl_stream="$1" + _tsl_tag="$2" + shift 2 + if [ "$TAURI_MODE" = true ]; then + if [ "$_tsl_stream" = stderr ]; then + printf '[TAURI:%s] %s\n' "$_tsl_tag" "$*" >&2 + else + printf '[TAURI:%s] %s\n' "$_tsl_tag" "$*" + fi + fi +} + +rollback_substep() { + if [ "$TAURI_MODE" = true ]; then + tauri_log "PROGRESS" "$1" + else + substep "$@" + fi +} + +tauri_clear_install_error() { + if [ "$TAURI_MODE" = true ]; then + tauri_log "ERROR_CLEAR" "$1" + printf '[TAURI:ERROR_CLEAR] %s\n' "$1" >&2 + fi +} + tauri_diag_marker() { _diag_gpu_branch="${1:-unknown}" _diag_torch_index_family="${2:-none}" @@ -310,6 +475,11 @@ _tauri_torch_index_family() { return fi _diag_url="${1:-}" + # Strip query/fragment AND a trailing slash before classifying (like _torch_index_url_leaf): + # a token isn't echoed into [TAURI:DIAG], and .../cu128/?token=x still classifies as cu128. + _diag_url="${_diag_url%%\?*}" + _diag_url="${_diag_url%%#*}" + _diag_url="${_diag_url%/}" case "$_diag_url" in */cu118) echo "cu118" ;; */cu124) echo "cu124" ;; @@ -343,7 +513,8 @@ _tauri_gpu_branch() { return fi case "$_diag_family" in - cu*) echo "cuda" ;; + # Require a digit after cu so /current or /custom isn't branded CUDA (parity ^cu[0-9]). + cu[0-9]*) echo "cuda" ;; rocm*) if [ "$_diag_radeon" = true ]; then echo "rocm_radeon" @@ -429,14 +600,20 @@ _start_studio_venv_replacement() { _stamp=$(date +%Y%m%d%H%M%S 2>/dev/null || echo "time") _candidate="$STUDIO_HOME/unsloth_studio.rollback.$_stamp.$$" _suffix=0 - while [ -e "$_candidate" ]; do + while [ -e "$_candidate" ] || [ -L "$_candidate" ]; do _suffix=$((_suffix + 1)) _candidate="$STUDIO_HOME/unsloth_studio.rollback.$_stamp.$$.$_suffix" done - mv "$_existing_dir" "$_candidate" _VENV_ROLLBACK_DIR="$_candidate" _VENV_ROLLBACK_TARGET="$_existing_dir" _VENV_ROLLBACK_ACTIVE=true + # Publish the rollback state before the atomic rename so a signal cannot + # land after mv but before the exit handlers know where the old venv went. + if ! mv "$_existing_dir" "$_candidate"; then + _VENV_ROLLBACK_ACTIVE=false + _VENV_ROLLBACK_DIR="" + return 1 + fi substep "previous environment preserved for rollback" } @@ -446,10 +623,10 @@ _restore_studio_venv_replacement() { _VENV_ROLLBACK_ACTIVE=false return 0 } - substep "restoring previous environment after failed install..." "$C_WARN" + rollback_substep "restoring previous environment after failed install..." "$C_WARN" rm -rf "$_VENV_ROLLBACK_TARGET" if mv "$_VENV_ROLLBACK_DIR" "$_VENV_ROLLBACK_TARGET"; then - substep "restored previous environment" + rollback_substep "restored previous environment" _VENV_ROLLBACK_ACTIVE=false _VENV_ROLLBACK_DIR="" else @@ -457,13 +634,68 @@ _restore_studio_venv_replacement() { fi } -_commit_studio_venv_replacement() { - [ "$_VENV_ROLLBACK_ACTIVE" = true ] || return 0 - if [ -n "$_VENV_ROLLBACK_DIR" ] && [ -d "$_VENV_ROLLBACK_DIR" ]; then - rm -rf "$_VENV_ROLLBACK_DIR" || true +_studio_venv_rollback_must_be_preserved() { + _rollback_name=${1##*/} + _rollback_metadata=${_rollback_name#unsloth_studio.rollback.} + _rollback_stamp=${_rollback_metadata%%.*} + _rollback_process=${_rollback_metadata#*.} + # Preserve anything outside the installer's timestamp.PID[.suffix] format. + [ "$_rollback_process" != "$_rollback_metadata" ] || return 0 + case "$_rollback_stamp" in + time) ;; + ''|*[!0-9]*) return 0 ;; + *) [ "${#_rollback_stamp}" -eq 14 ] || return 0 ;; + esac + _rollback_pid=${_rollback_process%%.*} + case "$_rollback_pid" in + ''|*[!0-9]*) return 0 ;; + esac + _rollback_suffix=${_rollback_process#*.} + if [ "$_rollback_suffix" != "$_rollback_process" ]; then + case "$_rollback_suffix" in ''|*[!0-9]*) return 0 ;; esac fi - _VENV_ROLLBACK_ACTIVE=false - _VENV_ROLLBACK_DIR="" + kill -0 "$_rollback_pid" 2>/dev/null +} + +_prune_stale_studio_venv_rollbacks() { + for _stale_rollback in "$STUDIO_HOME"/unsloth_studio.rollback.*; do + [ -d "$_stale_rollback" ] || continue + if [ -L "$_stale_rollback" ]; then + echo "⚠️ Refusing to remove rollback symlink $_stale_rollback" >&2 + continue + fi + # A concurrent installer may have moved its live venv aside. The PID in + # the generated name keeps this successful run from deleting its rescue copy. + _studio_venv_rollback_must_be_preserved "$_stale_rollback" && continue + if rm -rf "$_stale_rollback"; then + substep "removed stale environment rollback ${_stale_rollback##*/}" + else + echo "⚠️ Could not remove stale environment rollback $_stale_rollback" >&2 + fi + done +} + +_commit_studio_venv_replacement() { + if [ "$_VENV_ROLLBACK_ACTIVE" = true ]; then + _rollback_to_remove="$_VENV_ROLLBACK_DIR" + # The new environment is already committed. Clear the restore state + # before deletion so an interrupt cannot replace it with a half-deleted backup. + _VENV_ROLLBACK_ACTIVE=false + _VENV_ROLLBACK_DIR="" + if [ -n "$_rollback_to_remove" ] && [ -d "$_rollback_to_remove" ]; then + if ! rm -rf "$_rollback_to_remove"; then + echo "⚠️ Could not remove environment rollback $_rollback_to_remove" >&2 + fi + fi + fi + # Only prune older orphaned copies after the replacement has succeeded, so + # an interrupted install never discards the last known-good environment. + _prune_stale_studio_venv_rollbacks +} + +_cleanup_install_temporaries() { + [ -n "${_UV_OVERRIDE_TMPDIR:-}" ] && rm -rf "$_UV_OVERRIDE_TMPDIR" 2>/dev/null || true + [ -n "${_UNSLOTH_TORCH_OVERRIDES:-}" ] && rm -f "$_UNSLOTH_TORCH_OVERRIDES" 2>/dev/null || true } _on_install_exit() { @@ -471,13 +703,28 @@ _on_install_exit() { if [ "$_status" -ne 0 ]; then _restore_studio_venv_replacement fi - [ -n "${_UV_OVERRIDE_TMPDIR:-}" ] && rm -rf "$_UV_OVERRIDE_TMPDIR" 2>/dev/null || true + _cleanup_install_temporaries exit "$_status" } -# Empty so an inherited value can never reach the trap's rm; only a temp dir -# this script creates below (Apple Silicon, spaced path) is ever removed. + +_on_install_signal() { + _signal_status="$1" + # EXIT is disabled to avoid a second cleanup pass. Ignore further termination + # signals until the old environment is back in place. + trap - EXIT + trap '' HUP INT TERM + _restore_studio_venv_replacement + _cleanup_install_temporaries + exit "$_signal_status" +} +# Empty so an inherited value never reaches the trap's rm; only temp paths this +# script creates below (spaced-path dir, torch-trio overrides) are removed. _UV_OVERRIDE_TMPDIR="" +_UNSLOTH_TORCH_OVERRIDES="" trap _on_install_exit EXIT +trap '_on_install_signal 129' HUP +trap '_on_install_signal 130' INT +trap '_on_install_signal 143' TERM # ── Helper: download a URL to a file (supports curl and wget) ── download() { @@ -503,6 +750,45 @@ _is_pkg_installed() { esac } +# ── Helper: human-readable apt distro label for the sudo package prompt (#6207) ── +# Reads /etc/os-release so the Accept? prompt can say which distro we detected and +# that packages come from that distro's official apt repos (not a tarball). +_apt_distro_description() { + # Plain ( ... ) subshell — not $() — so case/;; stays bash-3.2-safe on macOS. + # Bash 3.2 misparses case arms inside command substitution and errors on `;;`. + ( + if [ ! -r /etc/os-release ]; then + printf 'a debian-like system' + exit 0 + fi + # shellcheck disable=SC1091 + . /etc/os-release 2>/dev/null || true + if [ -n "${NAME:-}" ] && [ -n "${VERSION_ID:-}" ]; then + _ad_label="$NAME $VERSION_ID" + elif [ -n "${PRETTY_NAME:-}" ]; then + _ad_label="$PRETTY_NAME" + elif [ -n "${NAME:-}" ]; then + _ad_label="$NAME" + else + printf 'a debian-like system' + exit 0 + fi + case " ${ID:-} ${ID_LIKE:-} " in + *" debian "*|*" ubuntu "*) _ad_label="${_ad_label} (debian-like)" ;; + esac + printf '%s' "$_ad_label" + ) +} + +# ── Helper: can the controlling terminal actually be opened for reading? ── +# `test -r` only checks permission bits, which look fine in containers and +# systemd units where open() then fails with ENXIO. Probe with a real open. +# The subshell is required: in dash a failed redirection on the special +# builtin `:` exits the whole script. +_can_read_tty() { + ( : /dev/null 2>&1 +} + # ── Helper: install packages via apt, escalating to sudo only if needed ── # Usage: _smart_apt_install pkg1 pkg2 pkg3 ... _smart_apt_install() { @@ -525,39 +811,90 @@ _smart_apt_install() { return 0 fi - # In Tauri mode, report needed packages and exit — Rust handles elevation + # Optional callers never elevate, in any mode: nothing on the consumer path + # builds anything, so neither the terminal sudo prompt below nor the Tauri + # NEED_SUDO dialog (whose Cancel leaves the user not installed) may gate the + # run over unused tools. The caller falls through to prebuilt llama.cpp. + # Required packages such as curl still escalate. + if [ "${_SMART_APT_OPTIONAL:-false}" = true ]; then + return 2 + fi + if [ "$TAURI_MODE" = true ]; then + # Report needed packages and exit — Rust handles elevation. tauri_log "NEED_SUDO" "$_STILL_MISSING" exit 2 fi # Step 3: Escalate -- need elevated permissions for remaining packages if command -v sudo >/dev/null 2>&1; then + _ad_desc="$(_apt_distro_description)" echo "" echo " !!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!" echo " WARNING: We require sudo elevated permissions to install:" echo " $_STILL_MISSING" - echo " If you accept, we'll run sudo now, and it'll prompt your password." + echo " Detected ${_ad_desc}." + echo " If you accept, we'll run sudo apt-get to install these packages" + echo " from your distro's official repositories (not a third-party tarball)." echo " !!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!" echo "" - printf " Accept? [Y/n] " - if [ -r /dev/tty ]; then - read -r REPLY /dev/null || true) case "$_p" in ''|*[!0-9]*) ;; @@ -901,7 +1238,7 @@ _acquire_lock() { # Lock dir exists -- check if owner is still alive _old_pid=$(cat "$LOCK_DIR/pid" 2>/dev/null || true) if [ -n "$_old_pid" ] && kill -0 "$_old_pid" 2>/dev/null; then - # Another launcher is running; wait for it to bring Studio up + # Another launcher is running; wait for it to bring Unsloth up _deadline=$(($(date +%s) + TIMEOUT_SEC)) while [ "$(date +%s)" -lt "$_deadline" ]; do _port=$(_find_healthy_port) && { @@ -1371,7 +1708,7 @@ WSLPS1_EOF # shortcut wasn't created; tell the user how to launch / re-enable it. if [ "$_css_created" -ne 1 ]; then substep "Couldn't create the Windows shortcut (WSL interop may be disabled)." "$C_WARN" - substep " Launch Studio from Windows: wsl -d \"$_css_distro\" -- bash -lc 'unsloth studio'" "$C_WARN" + substep " Launch Unsloth from Windows: wsl -d \"$_css_distro\" -- bash -lc 'unsloth studio'" "$C_WARN" substep " (re-enable shortcuts: turn WSL interop back on, e.g. run 'wsl --shutdown' then reopen WSL.)" "$C_WARN" fi fi @@ -1439,7 +1776,7 @@ if [ "$MAC_INTEL" = true ]; then echo "" echo " NOTE: Intel Mac (x86_64) detected." echo " PyTorch is unavailable for this platform (dropped Jan 2024)." - echo " Studio will install in GGUF-only mode." + echo " Unsloth will install in GGUF-only mode." echo " Chat, inference via GGUF, and data recipes will work." echo " Training requires Apple Silicon or Linux with GPU." echo "" @@ -1573,6 +1910,12 @@ _has_usable_nvidia_gpu() { # the STUDIO_HOME mkdir/venv so the origin distro is untouched. _maybe_reroute_strixhalo_to_2404() { [ "${OS:-}" = "wsl" ] || return 0 + # An explicit index pin skips every GPU-driven reroute (same contract as + # the later Radeon/Strix guard): the pin is honored in THIS distro rather + # than probing the GPU and switching distributions. Whitespace-only + # overrides do not gate (parity with get_torch_index_url). + _rr_pin=$(printf '%s' "${UNSLOTH_TORCH_INDEX_URL:-}${UNSLOTH_TORCH_INDEX_FAMILY:-}" | tr -d '[:space:]') + [ -n "$_rr_pin" ] && return 0 [ "${SKIP_TORCH:-false}" = "false" ] || return 0 [ "${UNSLOTH_SKIP_ROCM_WSL_SETUP:-0}" = "1" ] && return 0 [ "${UNSLOTH_WSL_REROUTED:-0}" = "1" ] && return 0 @@ -1582,7 +1925,7 @@ _maybe_reroute_strixhalo_to_2404() { # CUDA_VISIBLE_DEVICES=""/-1 and the /proc/driver/nvidia fallback for PATH/timeout gaps. if _has_usable_nvidia_gpu; then return 0; fi # Strix APUs show in /proc/cpuinfo; discrete cards don't, so also try WMI. Either reroutes. - if ! grep -qiE 'Ryzen AI Max|Radeon 80[0-9]0S|Strix Halo' /proc/cpuinfo 2>/dev/null \ + if ! grep -qiE 'Ryzen AI Max|Radeon 80[0-9][05]S|Strix Halo' /proc/cpuinfo 2>/dev/null \ && ! _wsl_amd_gpu_name >/dev/null 2>&1; then return 0 fi @@ -1634,6 +1977,10 @@ _maybe_reroute_strixhalo_to_2404() { # Forward explicit ROCm-bootstrap consent (e.g. Tauri) so the child auto-enables the # GPU instead of falling back to the desktop-app prompt path. [ "${UNSLOTH_ROCM_WSL_AUTO:-0}" = "1" ] && _rr_exports="$_rr_exports; export UNSLOTH_ROCM_WSL_AUTO=1" + # Forward a pinned torch index into the rerouted distro; dropping it would + # silently revert the child install to auto-detection. + [ -n "${UNSLOTH_TORCH_INDEX_URL:-}" ] && _rr_exports="$_rr_exports; export UNSLOTH_TORCH_INDEX_URL=$(_rr_q "$UNSLOTH_TORCH_INDEX_URL")" + [ -n "${UNSLOTH_TORCH_INDEX_FAMILY:-}" ] && _rr_exports="$_rr_exports; export UNSLOTH_TORCH_INDEX_FAMILY=$(_rr_q "$UNSLOTH_TORCH_INDEX_FAMILY")" [ "$_SKIP_AUTOSTART" = true ] && _rr_exports="$_rr_exports; export UNSLOTH_SKIP_AUTOSTART=1" _rr_args="" [ "$PACKAGE_NAME" != "unsloth" ] && _rr_args="$_rr_args --package $(_rr_q "$PACKAGE_NAME")" @@ -1671,67 +2018,142 @@ _maybe_reroute_strixhalo_to_2404() { _maybe_reroute_strixhalo_to_2404 || true # ── Check system dependencies ── -# cmake/git are only needed to *build* llama.cpp from source. Studio downloads a -# prebuilt by default, and setup.sh self-skips the source build when they're -# absent -- so macOS doesn't block on cmake (requiring it would force a manual -# Homebrew install). Linux keeps requiring them; its package manager has them. tauri_log "STEP" "Checking system dependencies" +# Without the Xcode CLT, macOS still ships /usr/bin/git as a stub that errors and pops +# a GUI dialog, so `command -v git` is not enough -- only running it tells the truth. +_has_working_git() { + command -v git >/dev/null 2>&1 || return 1 + git --version >/dev/null 2>&1 +} + +# macOS system-dependency check. A function so tests/sh can sed-extract it; the old +# inline form was untestable, which is why this gate shipped broken. +# +# The consumer install needs no developer toolchain: uv is a prebuilt binary, CPython +# is uv-managed, llama.cpp/whisper.cpp/Node are prebuilt downloads, and triton is +# skipped on macOS. Only `--local` needs git, for the unsloth-zoo git+https URL. +_check_macos_deps() { + _clt_missing=false + xcode-select -p >/dev/null 2>&1 || _clt_missing=true + + if [ "$STUDIO_LOCAL_INSTALL" = true ] && ! _has_working_git; then + echo "" + step "deps" "git is required for --local installs" "$C_ERR" + substep "--local installs unsloth-zoo from git+https://github.com/unslothai/unsloth-zoo," + substep "which needs a working git. Install the Xcode Command Line Tools:" + substep " xcode-select --install" + substep "Then re-run this script. A normal (non---local) install needs no compiler" + substep "and no git -- it uses prebuilt binaries and wheels only." + tauri_log "NEED_XCODE_CLT" "git" + return 1 + fi + + if [ "$_clt_missing" = true ]; then + # Not fatal, and no GUI dialog: firing xcode-select --install and exiting is + # what stranded clean Macs. + step "deps" "no Xcode Command Line Tools (not required)" "$C_WARN" + substep "Unsloth installs prebuilt binaries and wheels, so no compiler is needed." + substep "Install them only for a llama.cpp source build: xcode-select --install" + elif command -v cmake >/dev/null 2>&1; then + step "deps" "all system dependencies found" + else + # cmake is only for a source build, so its absence is not fatal. + step "deps" "using prebuilt llama.cpp (cmake not found)" "$C_WARN" + substep "Install cmake only if you want a source build: brew install cmake" + fi + return 0 +} + +# Linux/WSL system-dependency check. Same split as macOS, and a function for the same +# reason: tests/sh can extract it. +# +# Only a download transport is required. cmake, gcc and the libcurl headers exist +# solely for a llama.cpp source build the consumer path never does -- unslothai/ +# llama.cpp publishes linux-x64/arm64 prebuilts for cpu, cuda12, cuda13, rocm and +# vulkan. Requiring them turned every non-apt distro into a hard exit 1 over unused +# tooling. git follows macOS: --local only. +_check_linux_deps() { + _transport_missing=false + if ! command -v curl >/dev/null 2>&1 && ! command -v wget >/dev/null 2>&1; then + _transport_missing=true + fi + + # Wanted, never required: git fetches the triton_kernels git+https requirement (a + # training speedup), the rest serve the optional source build. Warn, never stop. + _optional_missing="" + command -v cmake >/dev/null 2>&1 || _optional_missing="$_optional_missing cmake" + _has_working_git || _optional_missing="$_optional_missing git" + command -v gcc >/dev/null 2>&1 || _optional_missing="$_optional_missing build-essential" + command -v curl-config >/dev/null 2>&1 || _optional_missing="$_optional_missing libcurl4-openssl-dev" + # Parameter expansion, not `sed`: sed may be absent on a minimal image, and a + # failed `$(... | sed ...)` yields "" -- "all found" on a machine that has none. + _optional_missing="${_optional_missing# }" + + if [ "$STUDIO_LOCAL_INSTALL" = true ] && ! _has_working_git; then + echo "" + step "deps" "git is required for --local installs" "$C_ERR" + substep "--local installs unsloth-zoo from git+https://github.com/unslothai/unsloth-zoo," + substep "which needs git. Install it with your package manager, then re-run." + substep "A normal (non---local) install needs no git and no compiler." + return 1 + fi + + # The one fatal case: nothing can be downloaded. apt is the only distro family we + # can drive unattended. + if [ "$_transport_missing" = true ]; then + if command -v apt-get >/dev/null 2>&1; then + echo "" + step "deps" "missing: curl" "$C_WARN" + substep "Needed to download uv, Python and the prebuilt inference engine." + _smart_apt_install curl + echo "" + else + echo "" + step "deps" "missing: curl (or wget)" "$C_ERR" + substep "Unsloth needs one of them to download uv, Python and the prebuilt" + substep "inference engine. Install one, then re-run setup:" + substep " Fedora/RHEL: sudo dnf install curl" + substep " Arch: sudo pacman -S --needed curl" + substep " openSUSE: sudo zypper install curl" + return 1 + fi + fi + + # Try apt for the optional set too; failing only costs the features warned about + # below. + if [ -n "$_optional_missing" ] && command -v apt-get >/dev/null 2>&1; then + step "deps" "installing optional build tools: $_optional_missing" "$C_DIM" + # Subshell because _smart_apt_install exits rather than returns, so `|| true` + # alone would not catch it. _SMART_APT_OPTIONAL suppresses every escalation + # path, so no install hinges on a prompt for tools nothing here needs. + ( _SMART_APT_OPTIONAL=true; _smart_apt_install $_optional_missing ) || true + _optional_missing="" + command -v cmake >/dev/null 2>&1 || _optional_missing="$_optional_missing cmake" + _has_working_git || _optional_missing="$_optional_missing git" + command -v gcc >/dev/null 2>&1 || _optional_missing="$_optional_missing build-essential" + command -v curl-config >/dev/null 2>&1 || _optional_missing="$_optional_missing libcurl4-openssl-dev" + _optional_missing="${_optional_missing# }" + fi + + if [ -n "$_optional_missing" ]; then + step "deps" "using prebuilt llama.cpp (missing: $_optional_missing)" "$C_WARN" + substep "Not required to run: Unsloth downloads a prebuilt inference engine." + case " $_optional_missing " in + *" git "*) substep "Without git the triton kernels training speedup is skipped." ;; + esac + else + step "deps" "all system dependencies found" + fi + return 0 +} + case "$OS" in macos) - # Xcode Command Line Tools provide the C/C++ compiler and git. - if ! xcode-select -p >/dev/null 2>&1; then - echo "" - echo "==> Xcode Command Line Tools are required." - echo " Installing (a system dialog will appear)..." - xcode-select --install /dev/null || true - echo " After the installation completes, please re-run this script." - exit 1 - fi - # cmake is only needed for a source build; the default prebuilt path - # doesn't use it, so its absence is not fatal -- no Homebrew prerequisite. - if command -v cmake >/dev/null 2>&1; then - step "deps" "all system dependencies found" - else - step "deps" "using prebuilt llama.cpp (cmake not found)" "$C_WARN" - substep "Install cmake only if you want a source build: brew install cmake" - fi + _check_macos_deps || exit 1 ;; linux|wsl) - MISSING="" - command -v cmake >/dev/null 2>&1 || MISSING="$MISSING cmake" - command -v git >/dev/null 2>&1 || MISSING="$MISSING git" - # curl or wget is needed for downloads; check both - if ! command -v curl >/dev/null 2>&1 && ! command -v wget >/dev/null 2>&1; then - MISSING="$MISSING curl" - fi - command -v gcc >/dev/null 2>&1 || MISSING="$MISSING build-essential" - # libcurl dev headers for llama.cpp HTTPS support - command -v curl-config >/dev/null 2>&1 || MISSING="$MISSING libcurl4-openssl-dev" - - MISSING=$(echo "$MISSING" | sed 's/^ *//') - if [ -n "$MISSING" ]; then - echo "" - step "deps" "missing: $MISSING" "$C_WARN" - substep "These are needed to build the GGUF inference engine." - if command -v apt-get >/dev/null 2>&1; then - _smart_apt_install $MISSING - else - echo " Automatic system package installation is supported on apt-based" - echo " Linux distributions (Ubuntu/Debian) only. Please install the" - echo " missing dependencies with your package manager, then re-run setup:" - echo " $MISSING" - echo "" - echo " Examples:" - echo " Fedora/RHEL: sudo dnf install cmake git gcc gcc-c++ make libcurl-devel" - echo " Arch: sudo pacman -S --needed cmake git base-devel curl" - echo " openSUSE: sudo zypper install cmake git gcc gcc-c++ make libcurl-devel" - exit 1 - fi - echo "" - else - step "deps" "all system dependencies found" - fi + _check_linux_deps || exit 1 ;; esac @@ -1821,11 +2243,13 @@ tauri_log "STEP" "Creating virtual environment" mkdir -p "$STUDIO_HOME" _MIGRATED=false +# Empty so an inherited value can never masquerade as a probed torch version. +_PREV_TORCH_VER="" if [ -x "$VENV_DIR/bin/python" ]; then # why: matching guard to the .venv branch below -- in env-mode # $STUDIO_HOME is a user-chosen workspace, so refuse to nuke an - # existing $STUDIO_HOME/unsloth_studio that lacks Studio sentinels. + # existing $STUDIO_HOME/unsloth_studio that lacks Unsloth sentinels. # Accept the in-VENV ownership marker so partial-install retries are # not blocked. Sentinels must be regular files: -f follows symlinks # to files (the legitimate ln -s shim shape) but rejects directories @@ -1838,6 +2262,12 @@ if [ -x "$VENV_DIR/bin/python" ]; then echo " Move it aside or choose an empty UNSLOTH_STUDIO_HOME." >&2 exit 1 fi + # Record the existing venv's torch BEFORE the replacement moves it aside: a re-run + # rebuilds the venv for clean state, but must keep the torch release the user + # already has (see _previous_torch_pin below). Last line only: sitecustomize or + # import-hook noise on stdout must not corrupt the version. + _PREV_TORCH_VER=$("$VENV_DIR/bin/python" -c \ + "import torch; print(torch.__version__)" 2>/dev/null | tail -n 1 || true) # New layout already exists — replace only after preserving rollback copy. substep "preserving existing environment for rollback..." _start_studio_venv_replacement "$VENV_DIR" @@ -1846,7 +2276,7 @@ elif [ "$_STUDIO_HOME_REDIRECT" != "env" ] && [ -x "$STUDIO_HOME/.venv/bin/pytho # Skip in env-mode so we don't rm -rf an unrelated .venv at the # workspace root (e.g. user's existing project Python venv). # In no-torch mode, a missing torch package is expected; validate Python only. - substep "found legacy Studio environment, validating..." + substep "found legacy Unsloth environment, validating..." _legacy_ok=false if [ "$SKIP_TORCH" = true ]; then if "$STUDIO_HOME/.venv/bin/python" -c "import sys; print(sys.executable)" >/dev/null 2>&1; then @@ -1903,7 +2333,7 @@ if [ ! -x "$VENV_DIR/bin/python" ]; then fi fi -# Mark the freshly-created venv as Studio-owned so a partial install can be +# Mark the freshly-created venv as Unsloth-owned so a partial install can be # repaired by re-running install.sh; the env-mode deletion guard above accepts # this marker as the primary sentinel. if [ -x "$VENV_DIR/bin/python" ]; then @@ -1991,6 +2421,15 @@ if [ "$SKIP_TORCH" = false ] && [ "$OS" = "macos" ] && [ "$_ARCH" = "arm64" ]; t TORCH_CONSTRAINT="torch>=2.6,<2.11.0" fi fi +# Companion (torchvision/torchaudio) constraints, bounded to torch's window. +# torchaudio 2.11 dropped its exact torch pin, so a bare companion next to a +# <2.11-capped torch resolves torchaudio 2.11 (verified: cpu leaf installed +# torch 2.10.0+cpu with torchaudio 2.11.0+cpu). torchvision still exact-pins +# torch and self-corrects, but is bounded for symmetry. Widened alongside the +# cu* torch window below; the torch-2.11 AMD paths (rocm7.2 / per-gfx / Strix) +# pin their own trio. +TORCHVISION_CONSTRAINT="torchvision>=0.19,<0.26.0" +TORCHAUDIO_CONSTRAINT="torchaudio>=2.4,<2.11.0" # ── Resolve repo root (for --local installs) ── _REPO_ROOT="$(cd "$(dirname "$0" 2>/dev/null || echo ".")" && pwd)" @@ -2040,18 +2479,155 @@ _has_amd_rocm_gpu() { amd-smi list 2>/dev/null | awk '/^GPU[[:space:]]*[:\[][[:space:]]*[0-9]/{ found=1 } END{ exit !found }'; then return 0 elif [ -e /dev/kfd ] && \ - awk 'FNR==1{ gpu=0; amd=0 } /gpu_id/{ gpu=($2+0>0) } /vendor_id/{ amd=($2==4098) } \ - gpu && amd { found=1 } END{ exit !found }' \ + awk '/vendor_id/ && $2 == 4098 { found = 1 } END { exit !found }' \ /sys/class/kfd/kfd/topology/nodes/*/properties 2>/dev/null; then - # vendor_id 4098 = 0x1002 (AMD). NVIDIA open kernel module (driver - # 560+) can register KFD topology nodes with non-zero gpu_id but - # vendor_id 4318 (0x10DE). Require AMD vendor to avoid misrouting - # NVIDIA-only hosts to the ROCm install path. + # vendor_id 4098 = 0x1002 (AMD) marks a GPU node: the KFD CPU node + # reports vendor_id 0, so any 4098 node is an AMD GPU. NVIDIA's open + # kernel module (driver 560+) registers KFD nodes as vendor_id 4318 + # (0x10DE), so this never false-positives on NVIDIA-only hosts. + # The prior check also required a gpu_id line, but gpu_id is a SIBLING + # sysfs file, not a line in properties -- it never matched, so the + # fallback silently missed every ROCm-less AMD host (issue: fresh + # Arch/CachyOS boxes reporting "no GPU detected"). return 0 fi return 1 } +# Returns 0 if an AMD display GPU is on the PCI bus even when ROCm can't use it +# (e.g. a Strix Halo iGPU with no /dev/kfd). Only sharpens the "no GPU detected" +# hint. vendor 0x1002 = AMD/ATI; class 0x03* = display controller. +_amd_gpu_present_via_pci() { + [ -d /sys/bus/pci/devices ] || return 1 + for _pci_vendor in /sys/bus/pci/devices/*/vendor; do + [ -r "$_pci_vendor" ] || continue + read -r _v < "$_pci_vendor" 2>/dev/null || continue + [ "$_v" = "0x1002" ] || continue + _cls="${_pci_vendor%vendor}class" + [ -r "$_cls" ] || continue + read -r _c < "$_cls" 2>/dev/null || continue + case "$_c" in 0x03*) return 0 ;; esac + done + return 1 +} + +# Map a gfx arch to the AMD pip index family (mirrors install.ps1 $archFamilyMap). +_amd_arch_index_family_for_gfx() { + case "$1" in + gfx1201|gfx1200) echo gfx120X-all ;; + gfx1151) echo gfx1151 ;; + gfx1150) echo gfx1150 ;; + gfx1152) echo gfx1152 ;; + gfx1103|gfx1102|gfx1101|gfx1100) echo gfx110X-all ;; + gfx1036|gfx1035|gfx1034|gfx1033|gfx1032|gfx1031|gfx1030) echo gfx103X-all ;; + gfx90a) echo gfx90a ;; + gfx908) echo gfx908 ;; + *) return 1 ;; + esac +} + +# Map a GPU marketing name to gfx arch (kept in sync with install.ps1 nameArchTable). +_infer_amd_gfx_arch_from_gpu_name() { + case "$1" in + *9070*|*9080*) echo gfx1201 ;; + *9060*) echo gfx1200 ;; + *"8065S"*|*"8060S"*|*"8050S"*|*"8040S"*|*"Strix Halo"*|*"Ryzen AI Max"*|*"AI Max"*) echo gfx1151 ;; + *"890M"*|*"880M"*|*"Strix Point"*|*"HX 37"*|*"AI 9 HX"*|*"AI 9 36"*) echo gfx1150 ;; + *"860M"*|*"840M"*|*"Krackan"*|*"AI 7 35"*|*"AI 5 34"*|*"AI 7 PRO 35"*|*"AI 5 33"*) echo gfx1152 ;; + *"RX 7600"*|*"RX 7700S"*|*"RX 7650"*|*"PRO W7600"*|*"PRO W7500"*) echo gfx1102 ;; + *"RX 7800"*|*"RX 7700"*|*"PRO W7700"*|*"PRO V710"*) echo gfx1101 ;; + *"RX 7900"*|*"PRO W7900"*|*"PRO W7800"*) echo gfx1100 ;; + *"780M"*|*"760M"*|*"740M"*|*"Phoenix"*|*"Hawk Point"*|*"Z1 Extreme"*|*"Z2 Extreme"*) echo gfx1103 ;; + *"RX 6900"*|*"RX 6800"*|*"RX 6750"*|*"RX 6700"*|*"PRO W6800"*|*"PRO W6900"*) echo gfx1030 ;; + *"RX 6650"*|*"RX 6600"*|*"PRO W6600"*|*"PRO W6650"*) echo gfx1032 ;; + *"RX 6500"*|*"RX 6400"*|*"RX 6300"*|*"PRO W6400"*|*"PRO W6500"*) echo gfx1034 ;; + *) return 1 ;; + esac +} + +# Best-effort gfx inference when ROCm tools can't see the GPU (unslothai#7301). +# Mirrors install.ps1 arch resolution on Windows ($HasROCm false, $ROCmGfxArch set). +_infer_linux_amd_gfx_arch() { + if [ -n "${UNSLOTH_ROCM_GFX_ARCH:-}" ]; then + printf '%s\n' "$(printf '%s' "$UNSLOTH_ROCM_GFX_ARCH" | tr '[:upper:]' '[:lower:]')" + return 0 + fi + # On WSL /proc/cpuinfo and lspci still report the host APU, but without the + # ROCDXG bridge (librocdxg over /dev/dxg) the AMD wheels can't reach the GPU; + # keep the CPU fallback there unless that runtime is present (the explicit + # override above still wins). Mirrors install_python_stack.py. + _gpu_evidence="" + if [ -e /dev/dxg ] || grep -qi microsoft /proc/version 2>/dev/null; then + for _d in /opt/rocm/lib /opt/rocm/lib64 /opt/rocm-*/lib /opt/rocm-*/lib64; do + { [ -e "$_d/librocdxg.so" ] || [ -e "$_d/librocdxg.so.1" ]; } && _rocdxg=1 && break + done + [ -n "${_rocdxg:-}" ] || return 1 + # WSL enumerates no PCI display device; /dev/dxg + librocdxg IS the + # GPU evidence there. + _gpu_evidence=1 + elif _amd_gpu_present_via_pci; then + _gpu_evidence=1 + fi + # /proc/cpuinfo leaks the HOST CPU model into VMs/containers that received + # no AMD GPU, so the CPU-model text alone is not GPU evidence: require an + # AMD display device (PCI vendor 0x1002, class 0x03*) before trusting it. + # The lspci fallback below needs no gate; an AMD display line IS evidence. + if [ -n "$_gpu_evidence" ] && grep -qiE 'Ryzen AI Max|Radeon 80[0-9][05]S|Strix Halo' /proc/cpuinfo 2>/dev/null; then + echo gfx1151 + return 0 + fi + if [ -n "$_gpu_evidence" ] && grep -qiE '890M|880M|Strix Point|HX 37[05]|AI 9 HX|AI 9 36[05]' /proc/cpuinfo 2>/dev/null; then + echo gfx1150 + return 0 + fi + if [ -n "$_gpu_evidence" ] && grep -qiE '860M|840M|Krackan|AI 7 35[05]|AI 5 34[05]|AI 7 PRO 35|AI 5 33' /proc/cpuinfo 2>/dev/null; then + echo gfx1152 + return 0 + fi + if command -v lspci >/dev/null 2>&1; then + # A non-AMD controller can enumerate first (Intel/ASPEED before an AMD + # dGPU), so scan every display-class line and take the first AMD one + # that maps. The vendor guard is case-SENSITIVE (a -i "ATI" would match + # "CorporATIon" on every Intel/NVIDIA line); whole-line matching also + # survives the 0000: PCI domain prefix. Mirrors install_python_stack.py. + _amd_disp=$(lspci -nn 2>/dev/null | grep -E 'VGA compatible controller|3D controller|Display controller' | grep -E 'AMD|ATI' || true) + while IFS= read -r _ln; do + [ -n "$_ln" ] || continue + if _gfx=$(_infer_amd_gfx_arch_from_gpu_name "$_ln"); then + echo "$_gfx" + return 0 + fi + done </dev/null 2>&1; then + _pg=$( (unset ROCR_VISIBLE_DEVICES HIP_VISIBLE_DEVICES; rocminfo 2>/dev/null) | grep -oE 'gfx[1-9][0-9a-z]{2,3}' || true) + fi + if [ -z "$_pg" ] && command -v amd-smi >/dev/null 2>&1; then + _pg=$( (unset ROCR_VISIBLE_DEVICES HIP_VISIBLE_DEVICES; amd-smi list 2>/dev/null) | grep -oE 'gfx[1-9][0-9a-z]{2,3}' || true) + if [ -z "$_pg" ]; then + _pg=$( (unset ROCR_VISIBLE_DEVICES HIP_VISIBLE_DEVICES; amd-smi static --asic 2>/dev/null) | grep -oE 'gfx[1-9][0-9a-z]{2,3}' || true) + fi + fi + printf '%s\n' "$_pg" +} + # ── 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 @@ -2059,6 +2635,24 @@ _has_amd_rocm_gpu() { get_torch_index_url() { _base="${UNSLOTH_PYTORCH_MIRROR:-https://download.pytorch.org/whl}" _base="${_base%/}" + # Explicit override -- skip ALL GPU probing (headless / container / CI / cross-install). + # UNSLOTH_TORCH_INDEX_URL wins (full URL, verbatim); _FAMILY is the leaf (cpu, cu128, ...) + # appended to the mirror base. Trim whitespace so a whitespace-only value is unset. + _url="${UNSLOTH_TORCH_INDEX_URL:-}" + _url="${_url#"${_url%%[![:space:]]*}"}"; _url="${_url%"${_url##*[![:space:]]}"}" + if [ -n "$_url" ]; then + # Trim trailing PATH slashes (a multi-slash path 404s on strict pip proxies) while + # preserving a ?query/#fragment token (a whole-URL strip would eat a "/"-ending token). + _url=$(_trim_index_path_slashes "$_url") + echo "$_url"; return + fi + _family="${UNSLOTH_TORCH_INDEX_FAMILY:-}" + _family="${_family#"${_family%%[![:space:]]*}"}"; _family="${_family%"${_family##*[![:space:]]}"}" + if [ -n "$_family" ]; then + while [ "${_family#/}" != "$_family" ]; do _family="${_family#/}"; done + while [ "${_family%/}" != "$_family" ]; do _family="${_family%/}"; done + echo "$_base/$_family"; return + fi # macOS: always CPU (no CUDA support) case "$(uname -s)" in Darwin) echo "$_base/cpu"; return ;; esac # Try nvidia-smi -- require the binary to actually list a usable GPU. @@ -2087,6 +2681,29 @@ get_torch_index_url() { if ! _has_amd_rocm_gpu; then echo "$_base/cpu"; return fi + # A generic rocm index is only safe when the gfx arch is readable: the + # Strix reroute (gfx1150/1151 -> arch-specific index) learns gfx from + # rocminfo/amd-smi, so if those are missing OR do not enumerate the GPU, an + # unknown-arch box might be Strix and would get the broken _grouped_mm + # wheels. Probe via the shared helper (override first, then rocminfo/amd-smi + # with visibility masks cleared); if the arch is unreadable, never guess a + # rocm index. A KFD-only host whose arch is still inferable from hardware + # IDs (PCI/cpuinfo/lspci) returns the cpu index and lets the runtime-less + # reroute below upgrade it to AMD per-arch wheels -- the reroute gate uses + # this same probe, so the handoff can't misfire. Only when inference fails + # too is CPU final, with the actionable warning. + _amd_gfx_probe=$(_probe_amd_gfx_arch) + if [ -z "$_amd_gfx_probe" ]; then + if _amd_inferred_gfx=$(_infer_linux_amd_gfx_arch 2>/dev/null) && \ + [ -n "$_amd_inferred_gfx" ] && \ + _amd_arch_index_family_for_gfx "$_amd_inferred_gfx" >/dev/null 2>&1; then + echo "[WARN] AMD GPU detected but rocminfo/amd-smi can't read its gfx arch -- inferring $_amd_inferred_gfx from hardware IDs." >&2 + echo "$_base/cpu"; return + fi + echo "[WARN] AMD GPU detected but its gfx arch can't be read (rocminfo/amd-smi missing or not enumerating the GPU) -- installing CPU-only PyTorch." >&2 + echo "[WARN] For GPU PyTorch, install or repair rocminfo/amd-smi (e.g. sudo pacman -S rocm-hip-sdk) and re-run this installer." >&2 + echo "$_base/cpu"; return + fi # AMD GPU confirmed -- detect ROCm version _rocm_tag="" _rocm_tag=$({ command -v amd-smi >/dev/null 2>&1 && \ @@ -2103,7 +2720,11 @@ get_torch_index_url() { { 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 + printf '%s\n' "$ver" | awk -F'[.-]' '{print "rocm"$1"."$2; exit}'; }) 2>/dev/null || _rocm_tag="" + # ^ || guard: when EVERY version source is missing (e.g. rocminfo present + # but rocm-core not installed, so dpkg-query/rpm exit 1), the whole || + # chain fails and set -e would kill the installer BEFORE the actionable + # no-version WARN below -- exactly the fresh-install case it exists for. # Validate _rocm_tag: must match "rocmX.Y" with major >= 1 case "$_rocm_tag" in rocm[1-9]*.[0-9]*) : ;; # valid (major >= 1) @@ -2139,12 +2760,27 @@ get_torch_index_url() { esac return fi - # AMD GPU confirmed by rocminfo/amd-smi but ROCm version could not be - # read from any source (amd-smi, /opt/rocm/.info/version, hipconfig, - # dpkg, rpm). Warn explicitly rather than silently installing CPU PyTorch. - echo "[WARN] AMD GPU detected but ROCm version could not be determined -- falling back to CPU-only PyTorch" >&2 - echo "[WARN] Ensure one of the following is accessible: amd-smi, hipconfig, /opt/rocm/.info/version, rocm-core package" >&2 - echo "[WARN] To install ROCm: https://rocm.docs.amd.com/en/latest/deploy/linux/index.html" >&2 + # AMD GPU confirmed (rocminfo/amd-smi or the KFD topology fallback) but + # no ROCm/HIP install was found to read the version from (amd-smi, + # /opt/rocm/.info/version, hipconfig, dpkg, rpm). This is the common + # fresh-install case: the GPU is real, but with no ROCm userspace the + # correct PyTorch build can't be selected. Warn with an actionable fix + # rather than silently installing CPU PyTorch. + # A user-set UNSLOTH_ROCM_GFX_ARCH seeded the probe above, so rocminfo/ + # amd-smi may still be unable to see the GPU; when the named arch maps to + # a wheel family, the runtime-less reroute (gated on the override) will + # install the AMD per-arch wheels -- a CPU-only warning here would be + # false for that path. Defer like the inferable-arch branch does. + if [ -n "${UNSLOTH_ROCM_GFX_ARCH:-}" ] && \ + _amd_arch_index_family_for_gfx "$_amd_gfx_probe" >/dev/null 2>&1; then + echo "[WARN] AMD GPU detected with no readable ROCm version, but UNSLOTH_ROCM_GFX_ARCH=$_amd_gfx_probe is set -- routing to AMD per-arch wheels." >&2 + echo "$_base/cpu"; return + fi + echo "[WARN] AMD GPU detected, but no ROCm/HIP install was found to select the matching GPU PyTorch build -- falling back to CPU-only PyTorch." >&2 + echo "[WARN] Install the ROCm/HIP SDK, then re-run this installer:" >&2 + echo "[WARN] Arch / CachyOS : sudo pacman -S rocm-hip-sdk" >&2 + echo "[WARN] other distros : https://rocm.docs.amd.com/en/latest/deploy/linux/index.html" >&2 + echo "[WARN] Minimum required for version detection: amd-smi, hipconfig, /opt/rocm/.info/version, or the rocm-core package." >&2 echo "$_base/cpu"; return fi # Parse CUDA version from nvidia-smi output (POSIX-safe, no grep -P). @@ -2187,16 +2823,155 @@ _torch_flavor_tag() { esac } +# Final path segment of a wheel index URL ($1), lowercased, query/fragment stripped first +# so a token-authenticated pin (.../cu128?token=x) classifies as cu128 (else it reinstalls +# every update). Classification only. Shared with the py / ps1 leaf extractors. +_torch_index_url_leaf() { + _tl_u="${1%%\?*}" + _tl_u="${_tl_u%%#*}" + # Strip ALL trailing slashes, not one: .../rocm7.2// must yield rocm7.2, not an empty leaf. + while [ -n "$_tl_u" ] && [ "${_tl_u%/}" != "$_tl_u" ]; do + _tl_u="${_tl_u%/}" + done + printf '%s' "${_tl_u##*/}" | tr '[:upper:]' '[:lower:]' +} + +# True (exit 0) when a lowercased leaf is an EXACT pip ROCm family: rocm[.] +# or a gfx ARCHITECTURE leaf (gfx followed by a digit: gfx90a, gfx1151, gfx120x-all). A leaf +# that merely starts with rocm/gfx (rocm7.2-private, gfx-private) is a custom verbatim pin. +# Matches the py / ps1 sides. +_is_pip_rocm_family_leaf() { + case "$1" in + gfx[0-9]*) return 0 ;; + rocm[0-9]*) + # Exact rocm[.]: both major and minor must be non-empty all-digits + # (rocm7., rocm7.2.1, rocm7.2-private are all custom pins, not a family). + _rocm_rest="${1#rocm}" + case "$_rocm_rest" in + *.*.*) return 1 ;; + *.*) + _rocm_minor="${_rocm_rest#*.}" + case "${_rocm_rest%%.*}" in "" | *[!0-9]*) return 1 ;; esac + case "$_rocm_minor" in "" | *[!0-9]*) return 1 ;; esac + ;; + *[!0-9]*) return 1 ;; + esac + return 0 + ;; + *) return 1 ;; + esac +} + +# Whether release base $1 (X.Y[.Z...]) falls inside constraint window $2 +# ("torch>=A.B[.C],="*",<"*) ;; + *) echo "no"; return ;; + esac + _trw_floor="${_trw_con#torch>=}"; _trw_floor="${_trw_floor%%,*}" + _trw_ceil="${_trw_con##*,<}" + _v_maj="${1%%.*}"; _v_rest="${1#*.}"; _v_min="${_v_rest%%.*}" + _f_maj="${_trw_floor%%.*}"; _f_rest="${_trw_floor#*.}"; _f_min="${_f_rest%%.*}" + _c_maj="${_trw_ceil%%.*}"; _c_rest="${_trw_ceil#*.}"; _c_min="${_c_rest%%.*}" + for _trw_n in "$_v_maj" "$_v_min" "$_f_maj" "$_f_min" "$_c_maj" "$_c_min"; do + case "$_trw_n" in ''|*[!0-9]*) echo "no"; return ;; esac + done + if [ "$_v_maj" -gt "$_f_maj" ] || { [ "$_v_maj" -eq "$_f_maj" ] && [ "$_v_min" -ge "$_f_min" ]; }; then + if [ "$_v_maj" -lt "$_c_maj" ] || { [ "$_v_maj" -eq "$_c_maj" ] && [ "$_v_min" -lt "$_c_min" ]; }; then + echo "yes" + return + fi + fi + echo "no" +} + +# Keep the previous venv's torch on a re-run: echo "torch==X.Y.Z" when the probed +# version ($1) is inside the active constraint window ($2), else "". The RELEASE is kept +# regardless of flavor tag; the pin installs from the freshly chosen index, so flavor +# follows the machine (cpu <-> cuda, cu126 -> cu130, PyPI bare -> +cu130) while the +# release follows the user. Gating on flavor was wrong: a PyPI torch reports a BARE +# version (on Linux the PyPI wheel IS CUDA), misclassified "cpu", so a healthy 2.10 on a +# cu130 host was moved to 2.11. Per-leaf floors still win (rocm7.2 / gfx >=2.11 for the +# Strix _grouped_mm fix, out-of-window manual installs) and are never pinned; the caller's +# _PREV_FALLBACK_CONSTRAINT installs the newest supported release when the index lacks the +# exact one. Opt out with UNSLOTH_TORCH_UPGRADE=1. +_previous_torch_pin() { + _ptp_ver="$1" + _ptp_con="$2" + [ -n "$_ptp_ver" ] || { echo ""; return; } + [ "${UNSLOTH_TORCH_UPGRADE:-0}" = "1" ] && { echo ""; return; } + _ptp_base="${_ptp_ver%%+*}" + # Base must be a plain numeric release (X.Y[.Z]); probe noise and + # nightly/dev/source builds (2.11.0.dev20250704, 2.9.0a0) must never + # become a pin -- no stable index carries them, so pinning would only + # print "keeping it" and then burn a doomed resolve before falling back. + case "$_ptp_base" in + *[!0-9.]* | *..* | .* | *.) echo ""; return ;; + [0-9]*.[0-9]*) ;; + *) echo ""; return ;; + esac + [ "$(_torch_release_in_window "$_ptp_base" "$_ptp_con")" = "yes" ] || { echo ""; return; } + echo "torch==$_ptp_base" +} + +# Install torch from TORCH_INDEX_URL honoring a kept-release pin: with _PREV_TORCH_PIN +# set, TORCH_CONSTRAINT is the exact previous release; fall back to the supported range +# if the index lacks it (pruned mirror) rather than failing. Used by every --default-index +# path (NVIDIA cu*, AMD rocm/gfx fallbacks, cpu/mac, ROCm repairs) so preservation is +# uniform. Extra args (e.g. --force-reinstall) are passed through to uv. +_install_torch_default_index() { + if [ -n "$_PREV_TORCH_PIN" ]; then + # Pair the companions with the kept torch minor: torchaudio no longer + # exact-pins torch in its metadata, so leaving it unconstrained resolves + # a newer mismatched build (a kept torch 2.9.0 pulled torchaudio 2.11.0). + _itdi_base="${_PREV_TORCH_PIN#torch==}" + _itdi_minor="${_itdi_base#*.}" + _itdi_minor="${_itdi_minor%%.*}" + _itdi_tv="torchvision" + _itdi_ta="torchaudio" + case "$_itdi_base" in + 2.*) + _itdi_tv="torchvision==0.$((_itdi_minor + 15)).*" + _itdi_ta="torchaudio==2.${_itdi_minor}.*" + ;; + esac + if ! run_install_cmd_retry "install PyTorch (kept release)" uv pip install --python "$_VENV_PY" "$TORCH_CONSTRAINT" "$_itdi_tv" "$_itdi_ta" \ + --default-index "$TORCH_INDEX_URL" "$@"; then + substep "[WARN] $_PREV_TORCH_PIN is not installable from $(_strip_index_url_credentials "$TORCH_INDEX_URL") -- installing the newest supported release instead" "$C_WARN" + TORCH_CONSTRAINT="$_PREV_FALLBACK_CONSTRAINT" + _PREV_TORCH_PIN="" + run_install_cmd_retry "install PyTorch" uv pip install --python "$_VENV_PY" "$TORCH_CONSTRAINT" "$TORCHVISION_CONSTRAINT" "$TORCHAUDIO_CONSTRAINT" \ + --default-index "$TORCH_INDEX_URL" "$@" + fi + else + run_install_cmd_retry "install PyTorch" uv pip install --python "$_VENV_PY" "$TORCH_CONSTRAINT" "$TORCHVISION_CONSTRAINT" "$TORCHAUDIO_CONSTRAINT" \ + --default-index "$TORCH_INDEX_URL" "$@" + fi +} + # Expected tag from the index leaf ($1): cuXXX / cpu / rocm (rocmX.Y and gfx* -> # rocm). Empty on an unknown leaf (odd mirror) so the repair safely no-ops. _expected_torch_flavor_tag() { - _u="${1%/}" - _leaf="${_u##*/}" + _leaf=$(_torch_index_url_leaf "$1") case "$_leaf" in - cu[0-9]*) echo "$_leaf" ;; - cpu) echo "cpu" ;; - rocm*|gfx*) echo "rocm" ;; - *) echo "" ;; + cu[0-9]*) + # Exact cu + digits only; a cu*-suffixed leaf (cu128-private) -> "" (custom), + # else a correct +cu128 wheel is force-reinstalled every run. + case "${_leaf#cu}" in + *[!0-9]*) echo "" ;; + *) echo "$_leaf" ;; + esac + ;; + cpu) echo "cpu" ;; + # Exact rocm/gfx families only; a custom rocm*-suffixed leaf -> "" (custom). + *) + if _is_pip_rocm_family_leaf "$_leaf"; then echo "rocm"; else echo ""; fi + ;; esac } @@ -2206,14 +2981,42 @@ _expected_torch_flavor_tag() { # fresh-install paths above already use -- so a stale wheel is auto-repairable. # Unknown/odd-mirror leaves -> no, so we warn rather than risk a wrong reinstall. _torch_index_repairable() { - _u="${1%/}" - _leaf="${_u##*/}" + _leaf=$(_torch_index_url_leaf "$1") case "$_leaf" in - cu[0-9]*|rocm[0-9]*|gfx*) echo "yes" ;; - *) echo "no" ;; + cu[0-9]*) echo "yes" ;; + # Only EXACT rocm/gfx families resolve via --default-index; a suffixed leaf is verbatim. + *) + if _is_pip_rocm_family_leaf "$_leaf"; then echo "yes"; else echo "no"; fi + ;; esac } +# Remove credentials from a wheel index URL ($1) so an authenticated pin never leaks: +# drops userinfo AND query/fragment; scheme/host/path stay exact. Shared with py / ps1. +_strip_index_url_credentials() { + _sic_url="$1" + case "$_sic_url" in + *://*) ;; + *) printf '%s' "$_sic_url"; return ;; + esac + _sic_scheme="${_sic_url%%://*}" + _sic_rest="${_sic_url#*://}" + # Drop query / fragment (may hold auth tokens). + _sic_rest="${_sic_rest%%\?*}" + _sic_rest="${_sic_rest%%#*}" + _sic_auth="${_sic_rest%%/*}" + # Drop user:pass@ userinfo if present. + case "$_sic_auth" in + *@*) _sic_host="${_sic_auth##*@}" ;; + *) _sic_host="$_sic_auth" ;; + esac + if [ "$_sic_auth" = "$_sic_rest" ]; then + printf '%s://%s' "$_sic_scheme" "$_sic_host" + else + printf '%s://%s/%s' "$_sic_scheme" "$_sic_host" "${_sic_rest#*/}" + 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/, @@ -2335,7 +3138,7 @@ _pick_radeon_wheel() { # the installer -- always returns 0. Runs the idempotent helper (ROCm 7.2 + # librocdxg), then sources the env it persisted so detection finds the GPU. # Export the ROCm-on-WSL env into this process and persist it to /etc/profile.d -# so non-login Studio/llama launches inherit it. Idempotent (writes only when +# so non-login Unsloth/llama launches inherit it. Idempotent (writes only when # the drop-in is missing); no-op without librocdxg, so never fires off WSL. # /etc/profile.d is root-owned -- sudo-tee when not root, else ROCm vanishes # after this shell on a non-root reinstall. Best-effort either way. @@ -2380,7 +3183,7 @@ _maybe_bootstrap_rocm_wsl() { rocminfo 2>/dev/null | awk '/Name:[[:space:]]*gfx[1-9]/ && !/generic/{found=1} END{exit !found}'; then # rocminfo may work only via the transient env _ensure_rocm_probe_env # just set, which dies with the installer. Persist the drop-in so login - # shells (Studio, llama.cpp) inherit it -- else a reinstall over an + # shells (Unsloth, llama.cpp) inherit it -- else a reinstall over an # existing /opt/rocm (uninstall keeps ROCm but drops it) loses the GPU. _persist_rocm_wsl_dropin return 0 @@ -2389,7 +3192,7 @@ _maybe_bootstrap_rocm_wsl() { [ -e /dev/dxg ] || return 0 # Strix APUs show in /proc/cpuinfo (the CPU model); discrete cards don't, so also # ask the Windows host. Either signal suffices; the bootstrap detects arch from rocminfo. - if ! grep -qiE 'Ryzen AI Max|Radeon 80[0-9]0S|Strix Halo' /proc/cpuinfo 2>/dev/null \ + if ! grep -qiE 'Ryzen AI Max|Radeon 80[0-9][05]S|Strix Halo' /proc/cpuinfo 2>/dev/null \ && ! _wsl_amd_gpu_name >/dev/null 2>&1; then return 0 fi @@ -2402,7 +3205,7 @@ _maybe_bootstrap_rocm_wsl() { # shellcheck disable=SC1091 . /etc/profile.d/unsloth-rocm-wsl.sh || true else - # librocdxg present but the env drop-in is gone (e.g. a Studio + # librocdxg present but the env drop-in is gone (e.g. an Unsloth # uninstall removed it while keeping shared ROCm). Restore the env. _persist_rocm_wsl_dropin fi @@ -2459,10 +3262,88 @@ _maybe_bootstrap_rocm_wsl() { [ -n "$_rw_tmp" ] && rm -f "$_rw_tmp" return 0 } -_maybe_bootstrap_rocm_wsl || true +# When the caller pins the wheel index (UNSLOTH_TORCH_INDEX_URL / _FAMILY), honour it +# everywhere: skip the WSL ROCm bootstrap and the Radeon/Strix reroute below (which would +# re-probe the GPU and overwrite the pin). Trim whitespace first (parity with +# get_torch_index_url): a whitespace-only override is unset there, so must not flip this true. +_torch_index_pinned=false +_ti_url_trim="${UNSLOTH_TORCH_INDEX_URL:-}" +_ti_url_trim="${_ti_url_trim#"${_ti_url_trim%%[![:space:]]*}"}"; _ti_url_trim="${_ti_url_trim%"${_ti_url_trim##*[![:space:]]}"}" +_ti_family_trim="${UNSLOTH_TORCH_INDEX_FAMILY:-}" +_ti_family_trim="${_ti_family_trim#"${_ti_family_trim%%[![:space:]]*}"}"; _ti_family_trim="${_ti_family_trim%"${_ti_family_trim##*[![:space:]]}"}" +if [ -n "$_ti_url_trim" ] || [ -n "$_ti_family_trim" ]; then + _torch_index_pinned=true +fi +[ "$_torch_index_pinned" = true ] || _maybe_bootstrap_rocm_wsl || true TORCH_INDEX_URL=$(get_torch_index_url) +# Linux: ROCm runtime missing but a supported AMD gfx arch is inferable (Strix Halo +# in /proc/cpuinfo, lspci marketing name, UNSLOTH_ROCM_GFX_ARCH). Route to AMD's +# per-arch wheels like install.ps1 does on Windows (unslothai#7301). +# Gated on the runtime probes NOT naming a gfx: either no AMD GPU is detected at +# all (_has_amd_rocm_gpu false), or the GPU is visible only through the +# env-independent KFD topology while rocminfo/amd-smi can't read its arch +# (KFD-only host, unslothai#7314 -- before the KFD detection fix these hosts +# reached this reroute via the false branch, so the empty-probe condition +# preserves that routing). A */cpu index chosen WITH a readable gfx +# (unsupported/unreadable ROCm version, after its own warning) is a deliberate +# fallback -- rerouting it would contradict that decision, and stays excluded +# because the shared probe returns its gfx. An explicit UNSLOTH_ROCM_GFX_ARCH +# override stays authoritative either way. +if [ "$_torch_index_pinned" = false ] && [ "$SKIP_TORCH" = false ] && \ + ! _has_usable_nvidia_gpu && \ + { [ -n "${UNSLOTH_ROCM_GFX_ARCH:-}" ] || ! _has_amd_rocm_gpu || \ + [ -z "$(_probe_amd_gfx_arch)" ]; } && \ + case "$(uname -s)" in Linux) true ;; *) false ;; esac && \ + case "$_ARCH" in x86_64|amd64) true ;; *) false ;; esac; then + # ROCm torch wheels are x86_64-only; get_torch_index_url returns CPU on other + # arches, so an inferred/overridden gfx must not reroute arm64 to AMD wheels. + case "$TORCH_INDEX_URL" in + */cpu) + _linux_inferred_gfx=$(_infer_linux_amd_gfx_arch 2>/dev/null || true) + if [ -n "$_linux_inferred_gfx" ]; then + _amd_family=$(_amd_arch_index_family_for_gfx "$_linux_inferred_gfx") || _amd_family="" + if [ -n "$_amd_family" ]; then + _amd_mirror="${UNSLOTH_AMD_ROCM_MIRROR:-https://repo.amd.com/rocm/whl}" + while [ "${_amd_mirror%/}" != "$_amd_mirror" ]; do + _amd_mirror="${_amd_mirror%/}" + done + TORCH_INDEX_URL="${_amd_mirror}/${_amd_family}/" + # Hand the inferred arch to setup.sh (llama.cpp): it re-probes + # ROCm on its own, and on these runtime-less hosts its probes + # find nothing, so without this it classifies the box as + # non-ROCm and installs the CPU prebuilt while torch just got + # AMD per-arch wheels. setup.sh and install_llama_prebuilt.py + # both honor UNSLOTH_ROCM_GFX_ARCH, so exporting it is the + # whole handoff (a user-set override re-exports unchanged). + export UNSLOTH_ROCM_GFX_ARCH="$_linux_inferred_gfx" + case "$_linux_inferred_gfx" in + gfx1201|gfx1200|gfx1151|gfx1150|gfx1152) + TORCH_CONSTRAINT="torch>=2.11.0,<2.12.0" + TORCHVISION_CONSTRAINT="torchvision>=0.26.0,<0.27.0" + TORCHAUDIO_CONSTRAINT="torchaudio>=2.11.0,<2.12.0" + ;; + esac + echo "" >&2 + # KFD-only hosts reach this reroute with /dev/kfd present + # (that's what detected them), so don't claim it's missing. + if _has_amd_rocm_gpu; then + echo " [WARN] AMD GPU visible via the kernel driver (KFD) but rocminfo/amd-smi can't read its gfx arch; using $_linux_inferred_gfx." >&2 + else + echo " [WARN] ROCm runtime not visible (/dev/kfd, rocminfo, amd-smi) but $_linux_inferred_gfx inferred." >&2 + fi + echo " [WARN] Routing to AMD arch-specific wheels ($(_strip_index_url_credentials "$TORCH_INDEX_URL"))." >&2 + echo " [WARN] These wheels bundle their own ROCm runtime; install the kernel stack for native compute:" >&2 + echo " [WARN] https://docs.unsloth.ai/get-started/install-and-update/amd" >&2 + echo " [WARN] Tip: set UNSLOTH_ROCM_GFX_ARCH=$_linux_inferred_gfx to skip inference next time." >&2 + echo "" >&2 + fi + fi + ;; + esac +fi + # Export the resolved torch backend ("cuda", "rocm", or "cpu") so that # downstream scripts (setup.sh -> install_python_stack.py) know what was # chosen here and can skip ROCm-specific repair steps on CUDA/CPU hosts. @@ -2470,24 +3351,74 @@ TORCH_INDEX_URL=$(get_torch_index_url) # whose base path happens to contain "rocm" or "gfx" must not mislabel a # cu*/cpu index as ROCm (radeon repo URLs end in rocm-rel-X.Y/, Strix # overrides in gfxNNNN/, so the trailing slash is stripped first). -_torch_index_leaf="${TORCH_INDEX_URL%/}" +# Lowercase the leaf so every gfx*/rocm*/cu* arm matches regardless of case (canonical AMD +# RDNA4 leaf is gfx120X-all). CUDA is branded only on a real cu[0-9]* leaf, so a mirror +# leaf (/current) does NOT commit a CUDA backend; an unknown leaf leaves the var unset so +# the stack probes the GPU. Query/fragment dropped first, then ALL trailing slashes (in +# lockstep with the shared _torch_index_url_leaf extractor). +_torch_index_leaf="${TORCH_INDEX_URL%%\?*}" +_torch_index_leaf="${_torch_index_leaf%%#*}" +# Strip ALL trailing slashes, not one: .../cu128// must yield cu128, not an empty leaf. +while [ -n "$_torch_index_leaf" ] && [ "${_torch_index_leaf%/}" != "$_torch_index_leaf" ]; do + _torch_index_leaf="${_torch_index_leaf%/}" +done _torch_index_leaf="${_torch_index_leaf##*/}" +_torch_index_leaf=$(printf '%s' "$_torch_index_leaf" | tr '[:upper:]' '[:lower:]') case "$_torch_index_leaf" in rocm*|gfx*) export UNSLOTH_TORCH_BACKEND="rocm" ;; cpu) export UNSLOTH_TORCH_BACKEND="cpu" ;; - *) export UNSLOTH_TORCH_BACKEND="cuda" ;; + cu[0-9]*) export UNSLOTH_TORCH_BACKEND="cuda" ;; + # Unknown leaf (odd mirror, /current): unset so a stale inherited value can't leak and + # the stack probes the GPU. + *) unset UNSLOTH_TORCH_BACKEND ;; esac -# rocm7.2 ships torch 2.11.0 -- adjust the constraint to allow it. -# All other ROCm tags and CUDA stay within <2.11.0. -case "$TORCH_INDEX_URL" in - */rocm7.2) TORCH_CONSTRAINT="torch>=2.11.0,<2.12.0" ;; +# Whether TORCH_INDEX_URL names an actual pip ROCm family (rocm* / gfx*), gating the +# ROCm-only side effects below (AMD bitsandbytes, ROCm-torch repair). Digit-gated so a leaf +# merely STARTING with "rocm" isn't force-repaired from the wrong path. +if _is_pip_rocm_family_leaf "$_torch_index_leaf"; then + _torch_index_is_rocm_family=true +else + _torch_index_is_rocm_family=false +fi + +# rocm7.2 and the per-gfx indexes with the _grouped_mm <2.11 bug (gfx120X-all, gfx1151, +# gfx1150) ship torch 2.11.0 -- raise the floor (also covers a pinned override that skipped +# the Strix reroute). Pin the companions too: the per-gfx index publishes them independently +# and a bare name can resolve a 2.12 ABI-mismatched wheel. Match on the FINAL leaf so a +# custom mirror with a gfx/rocm7.2 path segment but a cu*/cpu family isn't forced. +case "$_torch_index_leaf" in + rocm7.2|gfx120x-all|gfx1151|gfx1150|gfx1152) + TORCH_CONSTRAINT="torch>=2.11.0,<2.12.0" + TORCHVISION_CONSTRAINT="torchvision>=0.26.0,<0.27.0" + TORCHAUDIO_CONSTRAINT="torchaudio>=2.11.0,<2.12.0" + ;; + # CUDA cu12x/cu13x indexes ship torch 2.11.x: widen the ceiling to <2.12.0 (matches + # _CUDA_TORCH_PKG_SPEC) and widen the companions with it so the trio stays paired. + cu[0-9]*) + TORCH_CONSTRAINT="torch>=2.4,<2.12.0" + TORCHVISION_CONSTRAINT="torchvision>=0.19,<0.27.0" + TORCHAUDIO_CONSTRAINT="torchaudio>=2.4,<2.12.0" + ;; esac +# A pinned custom/unknown-leaf index (/simple, /current, /cu128-private) has no curated +# companion set, so bound torchvision/torchaudio to the same <2.11 range the Python path pins +# (else a mirror with newer companions resolves a 2.12 ABI-mismatched wheel). Known families +# keep their curated companions above (_expected_torch_flavor_tag returns "" only for custom). +if [ "$_torch_index_pinned" = true ] && \ + [ -z "$(_expected_torch_flavor_tag "$TORCH_INDEX_URL")" ]; then + TORCHVISION_CONSTRAINT="torchvision>=0.19,<0.26.0" + TORCHAUDIO_CONSTRAINT="torchaudio>=2.4,<2.11.0" +fi + # 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". +# Skipped when the index is pinned: an explicit override must not be rerouted to the +# Radeon/Strix repos by GPU probing. _amd_gpu_radeon=false +if [ "$_torch_index_pinned" = false ]; then case "$TORCH_INDEX_URL" in */rocm*) if _has_amd_rocm_gpu && command -v rocminfo >/dev/null 2>&1 && \ @@ -2496,29 +3427,64 @@ case "$TORCH_INDEX_URL" in fi ;; esac -# ── Strix Halo / Strix Point: force rocm7.2 wheels, bypass Radeon repo ─────── -# gfx1151 (Strix Halo) and gfx1150 (Strix Point) have a ROCm 7.1 driver bug -# that causes a segfault in torch._grouped_mm (moe_utils.py line 167). -# The Radeon repo now ships cp313 wheels for rocm-rel-7.1, so when -# _amd_gpu_radeon=true the installer silently lands on the broken combo. -# Detect these GPUs when TORCH_INDEX_URL is rocm7.1 and override to rocm7.2. -case "$TORCH_INDEX_URL" in - */rocm7.1|*/rocm7.1.*) +# 0 when a rocmX.Y index leaf ($1, the final path segment) is older than floor +# $2.$3 (int compare, so rocm7.2 < rocm7.13). Non-rocm leaves (gfx*, cu*, cpu) and +# non-numeric versions return 1. Leaf-based (like $_torch_index_leaf) so a mirror +# base holding its own rocm token compares the family leaf, not the base path. +_rocm_leaf_below() { + case "$1" in rocm[0-9]*.[0-9]*) : ;; *) return 1 ;; esac + _rb=${1#rocm}; _maj=${_rb%%.*}; _min=${_rb#*.}; _min=${_min%%.*} + case "$_maj$_min" in *[!0-9]*) return 1 ;; esac + if [ "$_maj" -lt "$2" ]; then return 0; fi + if [ "$_maj" -eq "$2" ] && [ "$_min" -lt "$3" ]; then return 0; fi + return 1 +} +# ── Strix Halo / Strix Point: route to the AMD arch-specific index ─────────── +# gfx1151/gfx1150 need torch 2.11+rocm7.13 from repo.amd.com/rocm/whl/gfx/, +# which carries AMD's real fixes (the rocm7.1 _grouped_mm segfault, moe_utils.py:167, +# and later Strix kernel bugs). Every generic pytorch.org index below rocm7.13 lacks +# them (and the Radeon repo can be offline, unslothai#7264), so reroute a detected +# Strix GPU whenever the picked index is older than the arch build -- covers today's +# rocm6.0-7.2 and any future 7.x < 7.13; rocm7.13+ already has the fixes, so leave it. +case "$_torch_index_leaf" in + rocm[0-9]*) # Collect every gfx token in rocminfo / amd-smi enumeration order # (skip duplicates), then index by HIP_VISIBLE_DEVICES / # ROCR_VISIBLE_DEVICES so a mixed Strix iGPU + non-Strix dGPU box # where the user selected the dGPU does NOT get rerouted to the # Strix per-gfx index. - _gfx_all="" - if command -v rocminfo >/dev/null 2>&1; then - _gfx_all=$(rocminfo 2>/dev/null | grep -oE 'gfx[1-9][0-9a-z]{2,3}') + # || true on each probe: no gfx match makes grep exit 1, which under + # set -euo pipefail would abort the installer before the next fallback + # runs (now that the case matches every rocm* index, not just rocm7.1). + # A user-supplied UNSLOTH_ROCM_GFX_ARCH overrides probing (mirrors setup.sh + # and the display block), so a Strix override still reaches the arch index. + _gfx_all=$(printf '%s' "${UNSLOTH_ROCM_GFX_ARCH:-}" | tr '[:upper:]' '[:lower:]') + if [ -z "$_gfx_all" ] && command -v rocminfo >/dev/null 2>&1; then + _gfx_all=$(rocminfo 2>/dev/null | grep -oE 'gfx[1-9][0-9a-z]{2,3}' || true) fi if [ -z "$_gfx_all" ] && command -v amd-smi >/dev/null 2>&1; then - _gfx_all=$(amd-smi list 2>/dev/null | grep -oE 'gfx[1-9][0-9a-z]{2,3}') + _gfx_all=$(amd-smi list 2>/dev/null | grep -oE 'gfx[1-9][0-9a-z]{2,3}' || true) # PowerShell paths also probe `amd-smi static --asic`; mirror it # so a host with hipinfo-less amd-smi reports the gfx target. if [ -z "$_gfx_all" ]; then - _gfx_all=$(amd-smi static --asic 2>/dev/null | grep -oE 'gfx[1-9][0-9a-z]{2,3}') + _gfx_all=$(amd-smi static --asic 2>/dev/null | grep -oE 'gfx[1-9][0-9a-z]{2,3}' || true) + fi + fi + # get_torch_index_url reads the arch with ROCR/HIP masks cleared, so a + # mask hiding every agent (e.g. ROCR_VISIBLE_DEVICES=-1) still lands + # here on a generic rocm index; re-probe unmasked or a masked-out Strix + # box keeps the broken generic wheels. Partial masks never get here + # (they enumerate at least one agent above) and keep their selection. + # ${VAR+x} (not :-): a SET-but-empty mask also hides every agent and + # must trigger the re-probe too. + if [ -z "$_gfx_all" ] && [ -n "${ROCR_VISIBLE_DEVICES+x}${HIP_VISIBLE_DEVICES+x}" ]; then + if command -v rocminfo >/dev/null 2>&1; then + _gfx_all=$( (unset ROCR_VISIBLE_DEVICES HIP_VISIBLE_DEVICES; rocminfo 2>/dev/null) | grep -oE 'gfx[1-9][0-9a-z]{2,3}' || true) + fi + if [ -z "$_gfx_all" ] && command -v amd-smi >/dev/null 2>&1; then + _gfx_all=$( (unset ROCR_VISIBLE_DEVICES HIP_VISIBLE_DEVICES; amd-smi list 2>/dev/null) | grep -oE 'gfx[1-9][0-9a-z]{2,3}' || true) + [ -z "$_gfx_all" ] && \ + _gfx_all=$( (unset ROCR_VISIBLE_DEVICES HIP_VISIBLE_DEVICES; amd-smi static --asic 2>/dev/null) | grep -oE 'gfx[1-9][0-9a-z]{2,3}' || true) fi fi _runtime_gfx="" @@ -2539,17 +3505,28 @@ case "$TORCH_INDEX_URL" in if (n > 0) print vals[idx] }') fi + # An explicit UNSLOTH_ROCM_GFX_ARCH=gfx906 pins the runtime target to the + # MI50 / Radeon VII path and must win over Strix probe-order detection on a + # mixed Strix + MI50 host, so the Strix reroute is suppressed when it is set. + # Normalize a copied HIP gcnArchName (gfx906:sramecc-:xnack- -> gfx906) and + # trim whitespace (mirrors the Python .strip()) so the feature-flag suffix or + # a stray newline does not defeat the exact gfx906 comparisons below. + _gfx906_env=$(printf '%s' "${UNSLOTH_ROCM_GFX_ARCH:-}" | tr '[:upper:]' '[:lower:]' | tr -d '[:space:]') + _gfx906_env=${_gfx906_env%%:*} _strix_gfx="" - case "$_runtime_gfx" in - gfx1151|gfx1150) _strix_gfx="$_runtime_gfx" ;; - esac - if [ -n "$_strix_gfx" ]; then + if [ "$_gfx906_env" != "gfx906" ]; then + case "$_runtime_gfx" in + gfx1151|gfx1150|gfx1152) _strix_gfx="$_runtime_gfx" ;; + esac + fi + # Skip rocm7.13+ generic indexes: they already ship the fixes, so the + # arch build (rocm7.13) would be a downgrade rather than a rescue. + if [ -n "$_strix_gfx" ] && _rocm_leaf_below "$_torch_index_leaf" 7 13; then echo "" >&2 - echo " [WARN] $_strix_gfx (Strix) + ROCm 7.1 detected -- known _grouped_mm segfault" >&2 - echo " [WARN] ROCm 7.1 wheels are broken for gfx1150/gfx1151 (moe_utils.py:167)" >&2 - echo " [WARN] Routing to AMD arch-specific index (torch 2.11+rocm7.13 has the real fix)" >&2 - echo " [WARN] Upgrade ROCm to 7.2+ to use the standard index:" >&2 - echo " [WARN] https://rocm.docs.amd.com/en/latest/deploy/linux/index.html" >&2 + echo " [WARN] $_strix_gfx (Strix) detected -- routing to the AMD arch-specific index" >&2 + echo " [WARN] torch 2.11+rocm7.13 has AMD's real gfx1150/gfx1151 fixes (the ROCm 7.1" >&2 + echo " [WARN] _grouped_mm segfault, moe_utils.py:167, and later Strix kernel bugs)," >&2 + echo " [WARN] and is more reliable than the rocm7.2 index or an offline Radeon repo." >&2 echo "" >&2 # AMD's arch-specific index serves torch 2.11.0+rocm7.13.0 which has AMD's # actual fix for the gfx1151/gfx1150 _grouped_mm kernel bug -- preferred @@ -2564,10 +3541,82 @@ case "$TORCH_INDEX_URL" in done TORCH_INDEX_URL="${_amd_strix_base}/${_strix_gfx}/" TORCH_CONSTRAINT="torch>=2.11.0,<2.12.0" + # Pin companions to 2.11 (per-gfx index publishes them independently). + TORCHVISION_CONSTRAINT="torchvision>=0.26.0,<0.27.0" + TORCHAUDIO_CONSTRAINT="torchaudio>=2.11.0,<2.12.0" _amd_gpu_radeon=false fi + # ── MI50 / Radeon VII (gfx906, Vega 20): legacy community-supported path ── + # Newer rocm wheel families bundle ROCm libraries whose Tensile kernels + # dropped gfx906 (rocBLAS "TensileLibrary.dat ... not read for gfx906", + # ROCm/TheRock#1844), so a rocm6.4+/7.x index installs a torch that fails + # at the first BLAS call. The rocm6.3 index is the last one whose wheels + # run on gfx906 (torch 2.7.0 verified on MI50 32GB; up to 2.9 in community + # use). Reroute any newer picked index; leave rocm6.0-6.3 alone. + # + # Target resolution: an explicit UNSLOTH_ROCM_GFX_ARCH wins (lets a host + # whose rocminfo/amd-smi emit no gfx token still opt in; _gfx906_env was + # lowercased above, before the Strix block it suppresses). Otherwise only + # treat gfx906 as the target when it is the SOLE distinct arch present: + # _gfx_all is de-duplicated by visible index, which loses per-device + # ordinals on a mixed host, so a non-gfx906 selection must never be + # downgraded to rocm6.3 -- such hosts set UNSLOTH_ROCM_GFX_ARCH to opt in. + _gfx906_target=false + if [ -n "$_gfx906_env" ]; then + [ "$_gfx906_env" = "gfx906" ] && _gfx906_target=true + elif [ -n "$_gfx_all" ]; then + _gfx906_uniq=$(printf '%s\n' "$_gfx_all" | awk 'NF && !seen[$0]++') + [ "$_gfx906_uniq" = "gfx906" ] && _gfx906_target=true + fi + # gfx906 always trains from the PyTorch rocm6.3 wheels, never the Radeon repo + # (repo.radeon.com wheels carry no gfx906 BLAS kernels). Clear the Radeon + # marketing-name flag as soon as gfx906 is the target -- even when the host + # already picks rocm6.0-6.3 and the reroute below is a no-op -- so a Radeon VII + # does not divert to the radeon branch on those versions. + if [ "$_gfx906_target" = true ]; then + _amd_gpu_radeon=false + fi + if [ "$_gfx906_target" = true ] && ! _rocm_leaf_below "$_torch_index_leaf" 6 4; then + echo "" >&2 + echo " [WARN] gfx906 (MI50 / Radeon VII / Vega 20) detected -- routing torch to the" >&2 + echo " [WARN] rocm6.3 index: it is the last wheel family that runs on gfx906 (newer" >&2 + echo " [WARN] rocm wheels ship without gfx906 BLAS kernels and fail at first use)." >&2 + echo " [WARN] gfx906 is a community-maintained legacy path: 16-bit LoRA and full" >&2 + echo " [WARN] finetuning work out of the box; bitsandbytes 4-bit QLoRA requires a" >&2 + echo " [WARN] source build of bitsandbytes for gfx906 (see docs.unsloth.ai/amd)." >&2 + echo "" >&2 + _amd_gfx906_base="${UNSLOTH_PYTORCH_MIRROR:-https://download.pytorch.org/whl}" + while [ "${_amd_gfx906_base%/}" != "$_amd_gfx906_base" ]; do + _amd_gfx906_base="${_amd_gfx906_base%/}" + done + TORCH_INDEX_URL="${_amd_gfx906_base}/rocm6.3" + # Reset to the default (<2.11) window: a rocm7.2 pick raised the floor + # to 2.11 above, which the rocm6.3 index (torch <= 2.9.x) cannot satisfy. + TORCH_CONSTRAINT="torch>=2.4,<2.11.0" + TORCHVISION_CONSTRAINT="torchvision>=0.19,<0.26.0" + TORCHAUDIO_CONSTRAINT="torchaudio>=2.4,<2.11.0" + # (_amd_gpu_radeon already cleared above for every gfx906 target.) + fi ;; esac +fi # _torch_index_pinned guard (Radeon + Strix reroute) +# Re-run over an existing install: keep the previous venv's torch RELEASE; the fresh +# index above supplies the right flavor for this machine. Evaluated HERE, after every +# index/constraint decision including the Strix reroute, so the window checked is the +# final one and a raised floor (rocm7.2 / Strix gfx) rejects an older release. +# _PREV_FALLBACK_CONSTRAINT keeps the range so the install can fall back when the exact +# release is not on the chosen index (mirrors may prune old wheels). Skipped for --no-torch. +_PREV_TORCH_PIN="" +_PREV_FALLBACK_CONSTRAINT="$TORCH_CONSTRAINT" +if [ "$SKIP_TORCH" = false ]; then + _prev_pin=$(_previous_torch_pin "$_PREV_TORCH_VER" "$TORCH_CONSTRAINT") + if [ -n "$_prev_pin" ]; then + _PREV_TORCH_PIN="$_prev_pin" + TORCH_CONSTRAINT="$_prev_pin" + substep "existing install has torch $_PREV_TORCH_VER -- keeping it (set UNSLOTH_TORCH_UPGRADE=1 to get the newest release)" + fi +fi + _TAURI_TORCH_INDEX_FAMILY=$(_tauri_torch_index_family "$TORCH_INDEX_URL") if [ "$_amd_gpu_radeon" = true ] && [ "$SKIP_TORCH" = false ]; then _TAURI_TORCH_INDEX_FAMILY="radeon" @@ -2615,12 +3664,14 @@ elif case "$TORCH_INDEX_URL" in */rocm*|*/gfx*) true ;; *) false ;; esac; then # gfx1102 matched BEFORE gfx1100 so the spaceless "RX 7700S" lands on # gfx1102 (bash case has no negative lookahead like the PS tables). case "$_gpu_disp_mkt" in - *"9070 XT"*|*9080*) _gpu_disp_gfx="gfx1201" ;; # RDNA 4 - *9070*|*9060*) _gpu_disp_gfx="gfx1200" ;; # RDNA 4 - *"8060S"*|*"8050S"*|*"8040S"*|*"Strix Halo"*|*"Ryzen AI Max"*|*"AI Max"*) _gpu_disp_gfx="gfx1151" ;; # RDNA 3.5 (Strix Halo: Radeon 8060S/8050S/8040S iGPU, Ryzen AI Max+) - *"890M"*|*"880M"*|*"860M"*|*"840M"*|*"Strix Point"*|*"Krackan"*|*"HX 37"*|*"AI 9 HX"*|*"AI 9 36"*|*"AI 7 35"*|*"AI 5 34"*|*"AI 7 PRO 35"*|*"AI 5 33"*) _gpu_disp_gfx="gfx1150" ;; # RDNA 3.5 (Strix/Krackan Point: Radeon 890M/880M iGPU, Ryzen AI 9 HX 370/375) - *"RX 7600"*|*"RX 7700S"*|*"RX 7650"*|*"PRO W7600"*|*"PRO W7500"*|*"PRO V710"*) _gpu_disp_gfx="gfx1102" ;; # RDNA 3 (Navi 33) - *"RX 7900"*|*"RX 7800"*|*"RX 7700"*|*"PRO W7900"*|*"PRO W7800"*|*"PRO W7700"*) _gpu_disp_gfx="gfx1100" ;; # RDNA 3 desktop / workstation (Navi 31) + *9070*|*9080*) _gpu_disp_gfx="gfx1201" ;; # RDNA 4 (Navi 48) + *9060*) _gpu_disp_gfx="gfx1200" ;; # RDNA 4 (Navi 44) + *"8065S"*|*"8060S"*|*"8050S"*|*"8040S"*|*"Strix Halo"*|*"Ryzen AI Max"*|*"AI Max"*) _gpu_disp_gfx="gfx1151" ;; # RDNA 3.5 (Strix Halo + Gorgon Halo: Radeon 8065S/8060S/8050S/8040S iGPU, Ryzen AI Max / Max+) + *"890M"*|*"880M"*|*"Strix Point"*|*"HX 37"*|*"AI 9 HX"*|*"AI 9 36"*) _gpu_disp_gfx="gfx1150" ;; # RDNA 3.5 (Strix Point: Radeon 890M/880M, Ryzen AI 9 HX 370/375) + *"860M"*|*"840M"*|*"Krackan"*|*"AI 7 35"*|*"AI 5 34"*|*"AI 7 PRO 35"*|*"AI 5 33"*) _gpu_disp_gfx="gfx1152" ;; # RDNA 3.5 (Krackan Point: Radeon 860M/840M, Ryzen AI 7 350 / AI 5 340) + *"RX 7600"*|*"RX 7700S"*|*"RX 7650"*|*"PRO W7600"*|*"PRO W7500"*) _gpu_disp_gfx="gfx1102" ;; # RDNA 3 (Navi 33) + *"RX 7800"*|*"RX 7700"*|*"PRO W7700"*|*"PRO V710"*) _gpu_disp_gfx="gfx1101" ;; # RDNA 3 (Navi 32) + *"RX 7900"*|*"PRO W7900"*|*"PRO W7800"*) _gpu_disp_gfx="gfx1100" ;; # RDNA 3 desktop / workstation (Navi 31) *"780M"*|*"760M"*|*"740M"*|*"Phoenix"*|*"Hawk Point"*|*"Z1 Extreme"*|*"Z2 Extreme"*) _gpu_disp_gfx="gfx1103" ;; # RDNA 3 iGPU (Phoenix / Hawk Point) *"RX 6900"*|*"RX 6800"*|*"RX 6750"*|*"RX 6700"*|*"PRO W6800"*|*"PRO W6900"*) _gpu_disp_gfx="gfx1030" ;; # RDNA 2 (Navi 21) *"RX 6650"*|*"RX 6600"*|*"PRO W6600"*|*"PRO W6650"*) _gpu_disp_gfx="gfx1032" ;; # RDNA 2 (Navi 23) @@ -2652,6 +3703,17 @@ elif case "$TORCH_INDEX_URL" in */rocm*|*/gfx*) true ;; *) false ;; esac; then elif [ "$OS" = "macos" ] && [ "$_ARCH" = "arm64" ]; then # Apple Silicon: PyTorch gets Metal (MPS) acceleration over unified memory, so not CPU-only. step "gpu" "Apple Silicon (Metal, unified memory)" +elif _has_amd_rocm_gpu; then + if [ "$_torch_index_pinned" = true ]; then + # An explicit UNSLOTH_TORCH_INDEX_URL/_FAMILY pin skipped all probing; + # do not claim ROCm is unusable when a CPU/other index was requested. + step "gpu" "AMD GPU (torch index pinned: $_torch_index_leaf)" "$C_WARN" + else + # AMD GPU visible to the kernel but the torch index stayed CPU: no usable + # ROCm userspace to pick a wheel. "none" would repeat the false diagnosis + # this installer used to give. + step "gpu" "AMD GPU (no usable ROCm -- CPU fallback)" "$C_WARN" + fi else step "gpu" "none (CPU-only)" "$C_WARN" fi @@ -2660,8 +3722,17 @@ fi case "$TORCH_INDEX_URL" in */cpu) if [ "$SKIP_TORCH" = false ] && [ "$OS" != "macos" ]; then - substep "No GPU detected -- installing CPU-only PyTorch." "$C_WARN" - if [ "$OS" = "wsl" ]; then + if [ "$_torch_index_pinned" = true ]; then + # An explicit CPU pin is a request, not a detection failure: + # skip the SDK guidance (ROCm may be perfectly healthy here). + substep "CPU-only PyTorch (index pinned via UNSLOTH_TORCH_INDEX_URL / _FAMILY)." + elif _has_amd_rocm_gpu; then + substep "AMD GPU detected, but no usable ROCm/HIP install -- installing CPU-only PyTorch." "$C_WARN" + substep "Install the ROCm/HIP SDK and re-run this installer for GPU PyTorch." "$C_WARN" + else + substep "No GPU detected -- installing CPU-only PyTorch." "$C_WARN" + fi + if [ "$OS" = "wsl" ] && [ "$_torch_index_pinned" = false ]; then # WSL + no GPU detected (detection above found nothing). Common # cause: an AMD GPU whose ROCm-on-WSL runtime isn't exposed yet -- # /dev/dxg present (graphics) but no ROCm runtime. @@ -2688,6 +3759,13 @@ case "$TORCH_INDEX_URL" in substep " driver is current; or run unsloth/scripts/install_rocm_wsl_strixhalo.sh yourself." else substep "AMD ROCm users: see https://docs.unsloth.ai/get-started/install-and-update/amd" + # Only when ROCm truly can't see the GPU: a detected-but-too-old + # ROCm (rocminfo works, wheels need 6.0+) has its own guidance. + if ! _has_amd_rocm_gpu && _amd_gpu_present_via_pci; then + substep "An AMD GPU is on the PCI bus but ROCm cannot see it (no /dev/kfd," "$C_WARN" + substep " rocminfo, or amd-smi). Install the ROCm kernel stack so /dev/kfd exists;" + substep " Strix Halo (gfx1151/gfx1150) needs a recent kernel (6.11+) and ROCm 7.x." + fi fi substep "Re-run with --no-torch for GGUF-only (faster, no PyTorch):" substep " curl -fsSL https://unsloth.ai/install.sh | sh -s -- --no-torch" @@ -2697,7 +3775,7 @@ case "$TORCH_INDEX_URL" in if [ "$_amd_gpu_radeon" = true ]; then substep "wheels: repo.radeon.com (Radeon)" else - substep "wheels: $TORCH_INDEX_URL" + substep "wheels: $(_strip_index_url_credentials "$TORCH_INDEX_URL")" fi ;; esac @@ -2705,9 +3783,47 @@ esac # ── Install unsloth directly into the venv (no activation needed) ── tauri_log "STEP" "Installing PyTorch" _VENV_PY="$VENV_DIR/bin/python" + +# A released unsloth wheel can pin an older torch (unsloth 2026.7.2 declares +# torch<2.11.0); a with-deps PyPI resolve then downgrades the whole trio, +# swapping the pinned +cuXXX/+rocm build for PyPI's default. The flavor guard +# below misses this (PyPI's torch 2.10 default is itself cu128-flavored), so +# freeze the trio via uv --overrides (overrides replace dependency requirements +# during resolution) while unsloth's other deps resolve normally. Sets +# _UNSLOTH_TORCH_OVERRIDES from the trio in the venv; every with-deps unsloth +# install (migrated and fresh) must call this before resolving and rm it after. +_build_unsloth_torch_overrides() { + _UNSLOTH_TORCH_OVERRIDES="" + [ "$SKIP_TORCH" = false ] || return 0 + _torch_trio_pins=$("$_VENV_PY" -c " +from importlib.metadata import version, PackageNotFoundError +for _p in ('torch', 'torchvision', 'torchaudio'): + try: + print(_p + '==' + version(_p)) + except PackageNotFoundError: + pass +" 2>/dev/null) || _torch_trio_pins="" + case "$_torch_trio_pins" in + torch==*) + _UNSLOTH_TORCH_OVERRIDES=$(mktemp) + printf '%s\n' "$_torch_trio_pins" > "$_UNSLOTH_TORCH_OVERRIDES" + # The CLI --overrides flag replaces any UV_OVERRIDE env file (same + # uv setting; macOS arm64 exports one here), so fold its pins in. + # awk, not cat: it drops inherited torch-trio lines (uv intersects + # duplicate overrides, so a conflicting pin would make resolution + # unsatisfiable) and newline-terminates the last line so an + # unterminated file cannot join two requirements into one. + for _ov_file in ${UV_OVERRIDE:-}; do + [ -f "$_ov_file" ] && awk '!/^[[:space:]]*torch(vision|audio)?([[:space:]<>=!~;@[]|$)/' "$_ov_file" >> "$_UNSLOTH_TORCH_OVERRIDES" + done + ;; + esac +} + if [ "$_MIGRATED" = true ]; then - # Migrated env: force-reinstall unsloth+unsloth-zoo to ensure clean state - # in the new venv location, while preserving existing torch/CUDA + # Migrated env: force-reinstall unsloth+unsloth-zoo for a clean state, preserving + # existing torch/CUDA unless the ROCm repair below fires. + _gfx906_bnb_snapshot substep "upgrading unsloth in migrated environment..." if [ "$SKIP_TORCH" = true ]; then # No-torch: install unsloth + unsloth-zoo with --no-deps (current @@ -2716,7 +3832,7 @@ if [ "$_MIGRATED" = true ]; then # to prevent transitive torch resolution. run_install_cmd_retry "install unsloth (migrated no-torch)" uv pip install --python "$_VENV_PY" --no-deps \ --reinstall-package unsloth --reinstall-package unsloth-zoo \ - "unsloth>=2026.7.2" "unsloth-zoo>=2026.7.2" + "unsloth>=2026.7.5" "unsloth-zoo>=2026.7.6" # Resolve pydantic WITH deps so pip pins pydantic-core to the # matching version (no-torch-runtime.txt below is --no-deps). # All transitive deps are torch-free. @@ -2729,9 +3845,13 @@ if [ "$_MIGRATED" = true ]; then else # Pin mlx-lm away from 0.31.3 here too: a curl-piped migration has no # overrides file, so UV_OVERRIDE is unset and this positional is the only cover. + _build_unsloth_torch_overrides run_install_cmd_retry "install unsloth (migrated)" uv pip install --python "$_VENV_PY" \ + ${_UNSLOTH_TORCH_OVERRIDES:+--overrides "$_UNSLOTH_TORCH_OVERRIDES"} \ --reinstall-package unsloth --reinstall-package unsloth-zoo \ - "unsloth>=2026.7.2" "unsloth-zoo>=2026.7.2" ${_MLX_LM_EXCLUDE_ARG:-} + "unsloth>=2026.7.5" "unsloth-zoo>=2026.7.6" ${_MLX_LM_EXCLUDE_ARG:-} + [ -n "$_UNSLOTH_TORCH_OVERRIDES" ] && rm -f "$_UNSLOTH_TORCH_OVERRIDES" + _UNSLOTH_TORCH_OVERRIDES="" fi if [ "$STUDIO_LOCAL_INSTALL" = true ]; then substep "overlaying local repo (editable)..." @@ -2744,21 +3864,19 @@ if [ "$_MIGRATED" = true ]; then # 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*|*/gfx*) - _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_retry "repair ROCm torch" uv pip install --python "$_VENV_PY" \ - "$TORCH_CONSTRAINT" torchvision torchaudio \ - --default-index "$TORCH_INDEX_URL" \ - --force-reinstall - fi - ;; - esac + if [ "$SKIP_TORCH" = false ] && [ "$_torch_index_is_rocm_family" = true ]; then + if _is_gfx906_bnb_skip; then + substep "gfx906: skipping prebuilt bitsandbytes (no gfx906 kernels); build from source for 4-bit QLoRA -- https://docs.unsloth.ai/get-started/install-and-update/amd" "$C_WARN" + else + _install_bnb_rocm "install bitsandbytes (AMD)" "$_VENV_PY" + fi + # 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)..." + _install_torch_default_index --force-reinstall + fi + _gfx906_bnb_prune fi elif [ -n "$TORCH_INDEX_URL" ]; then # Fresh: Step 1 - install torch from explicit index (skip when --no-torch or Intel Mac) @@ -2820,7 +3938,42 @@ elif [ -n "$TORCH_INDEX_URL" ]; then _ta_ver=$(_extract_version "$_ta_whl" "torchaudio") _radeon_versions_match=false - if [ -n "$_torch_ver" ] && [ -n "$_tv_ver" ] && [ -n "$_ta_ver" ]; then + # Kept release (_PREV_TORCH_PIN) wins here too: pick its exact + # patch (else the newest patch of its minor) plus the paired + # vision/audio wheels. Any gap falls back to the newest-trio + # search below, mirroring _install_torch_default_index, so a + # rerun never drifts to another release nor below the kept one. + if [ -n "$_PREV_TORCH_PIN" ]; then + _prev_kept_base="${_PREV_TORCH_PIN#torch==}" + _prev_kept_minor="${_prev_kept_base#*.}" + _prev_kept_minor="${_prev_kept_minor%%.*}" + case "$_prev_kept_minor" in + ''|*[!0-9]*) ;; + *) + _kept_torch=$(_pick_radeon_wheel "torch" "${_prev_kept_base}" 2>/dev/null) || _kept_torch="" + [ -z "$_kept_torch" ] && { _kept_torch=$(_pick_radeon_wheel "torch" "2.${_prev_kept_minor}." 2>/dev/null) || _kept_torch=""; } + _kept_tv=$(_pick_radeon_wheel "torchvision" "0.$((_prev_kept_minor + 15))." 2>/dev/null) || _kept_tv="" + _kept_ta=$(_pick_radeon_wheel "torchaudio" "2.${_prev_kept_minor}." 2>/dev/null) || _kept_ta="" + if [ -n "$_kept_torch" ] && [ -n "$_kept_tv" ] && [ -n "$_kept_ta" ]; then + _torch_whl=$_kept_torch + _tv_whl=$_kept_tv + _ta_whl=$_kept_ta + _tri_whl="" + _radeon_versions_match=true + # Say so when the listing pruned the exact patch + # and a same-series build is installed instead. + case "$(printf '%s' "${_kept_torch##*/}" | sed 's/%2[Bb]/+/g')" in + "torch-${_prev_kept_base}"[+-]*) ;; + *) substep "kept release ${_prev_kept_base} is not in the Radeon listing -- installing the closest 2.${_prev_kept_minor} series build instead" ;; + esac + else + substep "[WARN] Radeon repo lacks a complete wheel set for kept $_PREV_TORCH_PIN -- installing the newest compatible set instead" "$C_WARN" + fi + ;; + esac + fi + if [ "$_radeon_versions_match" != true ] && \ + [ -n "$_torch_ver" ] && [ -n "$_tv_ver" ] && [ -n "$_ta_ver" ]; then _torch_minor=${_torch_ver#*.} _ta_minor=${_ta_ver#*.} _tv_minor=${_tv_ver#*.} @@ -2877,10 +4030,8 @@ elif [ -n "$TORCH_INDEX_URL" ]; then 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_retry "install PyTorch" uv pip install --python "$_VENV_PY" \ - "$TORCH_CONSTRAINT" torchvision torchaudio \ - --default-index "$TORCH_INDEX_URL" + substep "[WARN] Radeon repo lacks a compatible wheel set for this Python; falling back to ROCm index ($(_strip_index_url_credentials "$TORCH_INDEX_URL"))" "$C_WARN" + _install_torch_default_index else substep "installing PyTorch from Radeon repo (${_RADEON_BASE_URL})..." # Pass explicit wheel URLs so the matched trio is @@ -2900,42 +4051,39 @@ elif [ -n "$TORCH_INDEX_URL" ]; then fi fi else - substep "[WARN] Radeon repo unavailable; falling back to ROCm index ($TORCH_INDEX_URL)" "$C_WARN" - run_install_cmd_retry "install PyTorch" uv pip install --python "$_VENV_PY" \ - "$TORCH_CONSTRAINT" torchvision torchaudio \ - --default-index "$TORCH_INDEX_URL" + substep "[WARN] Radeon repo unavailable; falling back to ROCm index ($(_strip_index_url_credentials "$TORCH_INDEX_URL"))" "$C_WARN" + _install_torch_default_index fi else substep "[WARN] Radeon GPU detected but could not detect full ROCm version; falling back to ROCm index" "$C_WARN" - run_install_cmd_retry "install PyTorch" uv pip install --python "$_VENV_PY" \ - "$TORCH_CONSTRAINT" torchvision torchaudio \ - --default-index "$TORCH_INDEX_URL" + _install_torch_default_index fi else - substep "installing PyTorch ($TORCH_INDEX_URL)..." - run_install_cmd_retry "install PyTorch" uv pip install --python "$_VENV_PY" "$TORCH_CONSTRAINT" torchvision torchaudio \ - --default-index "$TORCH_INDEX_URL" + substep "installing PyTorch ($(_strip_index_url_credentials "$TORCH_INDEX_URL"))..." + _install_torch_default_index 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*|*/gfx*) - _install_bnb_rocm "install bitsandbytes (AMD)" "$_VENV_PY" - ;; - esac + if [ "$SKIP_TORCH" = false ] && [ "$_torch_index_is_rocm_family" = true ]; then + if _is_gfx906_bnb_skip; then + substep "gfx906: skipping prebuilt bitsandbytes (no gfx906 kernels); build from source for 4-bit QLoRA -- https://docs.unsloth.ai/get-started/install-and-update/amd" "$C_WARN" + else + _install_bnb_rocm "install bitsandbytes (AMD)" "$_VENV_PY" + fi fi - # Fresh: Step 2 - install unsloth, preserving pre-installed torch + _gfx906_bnb_snapshot + # Fresh: Step 2 - install unsloth, preserving the torch Step 1 installed tauri_log "STEP" "Installing Unsloth" substep "installing unsloth (this may take a few minutes)..." + _build_unsloth_torch_overrides if [ "$SKIP_TORCH" = true ]; then # No-torch: install unsloth + unsloth-zoo with --no-deps, then # runtime deps (typer, safetensors, transformers, etc.) with --no-deps. run_install_cmd_retry "install unsloth (no-torch)" uv pip install --python "$_VENV_PY" --no-deps \ --upgrade-package unsloth --upgrade-package unsloth-zoo \ - "unsloth>=2026.7.2" "unsloth-zoo>=2026.7.2" + "unsloth>=2026.7.5" "unsloth-zoo>=2026.7.6" # Same pydantic-with-deps trick as the migrated branch. run_install_cmd_retry "install pydantic (with deps for compatible core)" \ uv pip install --python "$_VENV_PY" pydantic @@ -2953,7 +4101,8 @@ elif [ -n "$TORCH_INDEX_URL" ]; then fi elif [ "$STUDIO_LOCAL_INSTALL" = true ]; then run_install_cmd_retry "install unsloth (local)" uv pip install --python "$_VENV_PY" \ - --upgrade-package unsloth "unsloth>=2026.7.2" "unsloth-zoo>=2026.7.2" + ${_UNSLOTH_TORCH_OVERRIDES:+--overrides "$_UNSLOTH_TORCH_OVERRIDES"} \ + --upgrade-package unsloth "unsloth>=2026.7.5" "unsloth-zoo>=2026.7.6" substep "overlaying local repo (editable)..." run_install_cmd "overlay local repo" uv pip install --python "$_VENV_PY" -e "$_REPO_ROOT" --no-deps substep "overlaying unsloth-zoo from git main..." @@ -2962,30 +4111,27 @@ elif [ -n "$TORCH_INDEX_URL" ]; then "unsloth-zoo @ git+https://github.com/unslothai/unsloth-zoo" else run_install_cmd_retry "install unsloth" uv pip install --python "$_VENV_PY" \ + ${_UNSLOTH_TORCH_OVERRIDES:+--overrides "$_UNSLOTH_TORCH_OVERRIDES"} \ --upgrade-package unsloth -- "$PACKAGE_NAME" ${_MLX_LM_EXCLUDE_ARG:-} fi + [ -n "$_UNSLOTH_TORCH_OVERRIDES" ] && rm -f "$_UNSLOTH_TORCH_OVERRIDES" + _UNSLOTH_TORCH_OVERRIDES="" # 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*|*/gfx*) - _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_retry "repair ROCm torch" uv pip install --python "$_VENV_PY" \ - "$TORCH_CONSTRAINT" torchvision torchaudio \ - --default-index "$TORCH_INDEX_URL" \ - --force-reinstall - fi - ;; - esac + if [ "$SKIP_TORCH" = false ] && [ "$_torch_index_is_rocm_family" = true ]; then + _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)..." + _install_torch_default_index --force-reinstall + fi + _gfx906_bnb_prune fi else # Fallback: GPU detection failed to produce a URL -- let uv resolve torch tauri_log "STEP" "Installing Unsloth" substep "installing unsloth (this may take a few minutes)..." if [ "$STUDIO_LOCAL_INSTALL" = true ]; then - run_install_cmd_retry "install unsloth (auto torch backend)" uv pip install --python "$_VENV_PY" "unsloth-zoo>=2026.7.2" "unsloth>=2026.7.2" --torch-backend=auto + run_install_cmd_retry "install unsloth (auto torch backend)" uv pip install --python "$_VENV_PY" "unsloth-zoo>=2026.7.6" "unsloth>=2026.7.5" --torch-backend=auto substep "overlaying local repo (editable)..." run_install_cmd "overlay local repo" uv pip install --python "$_VENV_PY" -e "$_REPO_ROOT" --no-deps substep "overlaying unsloth-zoo from git main..." @@ -2997,6 +4143,15 @@ else fi fi +_installed_package_version=$("$_VENV_PY" -c \ + 'from importlib.metadata import version; import sys; print(version(sys.argv[1]))' \ + "$PACKAGE_NAME" 2>/dev/null || true) +if [ -n "$_installed_package_version" ]; then + step "$PACKAGE_NAME" "$_installed_package_version installed" +else + substep "[WARN] installed $PACKAGE_NAME version could not be determined" "$C_WARN" +fi + # ── Enforce the installed torch flavor matches the detected GPU build ── # PEP 440 ignores the +cpu/+cuXXX/+rocm local label in a version range, so uv # keeps a stale torch==X+cpu against a GPU index and the venv silently trains on @@ -3014,9 +4169,7 @@ if [ "$SKIP_TORCH" = false ] && [ -n "${TORCH_INDEX_URL:-}" ]; then if [ -n "$_installed_torch_tag" ] && [ "$_installed_torch_tag" != "$_expected_torch_tag" ] \ && [ "$(_torch_index_repairable "$TORCH_INDEX_URL")" = "yes" ]; then substep "PyTorch flavor mismatch (installed $_installed_torch_tag, need $_expected_torch_tag) -- reinstalling correct build..." - run_install_cmd "reinstall PyTorch ($_expected_torch_tag)" uv pip install --python "$_VENV_PY" \ - "$TORCH_CONSTRAINT" torchvision torchaudio \ - --default-index "$TORCH_INDEX_URL" \ + _install_torch_default_index \ --reinstall-package torch --reinstall-package torchvision --reinstall-package torchaudio _installed_torch_ver=$("$_VENV_PY" -c "import torch; print(torch.__version__)" 2>/dev/null || true) _installed_torch_tag="" @@ -3027,13 +4180,13 @@ if [ "$SKIP_TORCH" = false ] && [ -n "${TORCH_INDEX_URL:-}" ]; then substep "[WARN] PyTorch is CPU-only but a $_expected_torch_tag GPU build was expected for this machine." "$C_WARN" substep "[WARN] Training and GPU inference will run on CPU until this is fixed." "$C_WARN" substep "[WARN] Re-run this installer, or reinstall the GPU build manually:" "$C_WARN" - substep "[WARN] uv pip install --python \"$_VENV_PY\" \"$TORCH_CONSTRAINT\" torchvision torchaudio --default-index $TORCH_INDEX_URL --reinstall-package torch --reinstall-package torchvision --reinstall-package torchaudio" "$C_WARN" + substep "[WARN] uv pip install --python \"$_VENV_PY\" \"$TORCH_CONSTRAINT\" \"$TORCHVISION_CONSTRAINT\" \"$TORCHAUDIO_CONSTRAINT\" --default-index $(_strip_index_url_credentials "$TORCH_INDEX_URL") --reinstall-package torch --reinstall-package torchvision --reinstall-package torchaudio" "$C_WARN" fi fi fi # ── Run studio setup ── -tauri_log "STEP" "Running Studio setup" +tauri_log "STEP" "Running Unsloth setup" # When --local, use the repo's own setup.sh directly. # Otherwise, find it inside the installed package. SETUP_SH="" @@ -3066,6 +4219,7 @@ if [ -n "$VENV_ABS_BIN" ]; then fi if ! command -v bash >/dev/null 2>&1; then + tauri_log "ERROR" "bash is required to run studio setup" step "setup" "bash is required to run studio setup" "$C_ERR" substep "Please install bash and re-run install.sh" exit 1 @@ -3104,6 +4258,7 @@ if [ "$STUDIO_LOCAL_INSTALL" = true ]; then STUDIO_LOCAL_REPO="$_REPO_ROOT" \ UNSLOTH_NO_TORCH="$SKIP_TORCH" \ UNSLOTH_LOCAL_LLAMA_CPP_DIR="$_WITH_LLAMA_CPP_DIR" \ + UNSLOTH_TAURI_MODE="$TAURI_MODE" \ bash "$SETUP_SH" =0.12.0", "rich", "pydantic", "pyyaml", "nest-asyncio", + # Every CLI command imports studio.backend.*, which reaches structlog at + # module level. The rest of the server stack lives in the studio extra. + "structlog>=24.1.0", + # unsloth_cli/__init__.py reaches click via commands/start.py, so every + # command needs it. typer supplied it until 0.27 dropped the dependency. + "click>=8.0", ] [project.scripts] @@ -41,9 +47,14 @@ version = {attr = "unsloth.models._utils.__version__"} [tool.setuptools] include-package-data = true +[tool.setuptools.cmdclass] +# Snapshots CHANGELOG.md into studio/ so every build path ships it. +build_py = "_changelog_build.build_py" + [tool.setuptools.package-data] -unsloth_cli = ["codex_fallback_prompt.md"] +unsloth_cli = ["codex_fallback_prompt.md", "pi_subagent.ts"] studio = [ + "CHANGELOG.md", "*.sh", "*.ps1", "*.bat", @@ -68,13 +79,40 @@ include = ["unsloth*", "unsloth_cli*", "studio", "studio.backend*"] exclude = ["images*", "tests*", "*.node_modules", "*.node_modules.*"] [project.optional-dependencies] +# Studio's server stack, mirroring studio/backend/requirements/studio.txt. +# test_studio_extra_matches_requirements.py catches drift. +studio = [ + "typer", + "fastapi", + "uvicorn", + "pydantic", + "packaging", + "matplotlib==3.10.9", + "pandas", + "nest_asyncio", + "datasets==4.3.0", + "pyjwt", + "huggingface-hub==0.36.2", + "structlog>=24.1.0", + "diceware", + "ddgs", + "cryptography>=42.0.0", + "boto3>=1.34.0", + "httpx>=0.27.0", + "fastmcp>=3.0.2", + "sqlite-vec==0.1.9", + "pymupdf==1.27.2.3", + "pymupdf4llm==0.3.4", + "python-docx==1.2.0", +] + triton = [ "triton>=3.0.0 ; ('linux' in sys_platform)", "triton-windows ; (sys_platform == 'win32') and (platform_machine == 'AMD64' or platform_machine == 'x86_64')", ] huggingfacenotorch = [ - "unsloth_zoo>=2026.7.2", + "unsloth_zoo>=2026.7.6", "wheel>=0.42.0", "packaging", "numpy", @@ -93,9 +131,25 @@ huggingfacenotorch = [ "trl>=0.18.2,!=0.19.0,<=0.24.0", "sentence-transformers", ] +# torchcodec backend for Gemma audio / datasets>=4 (#7225). +# Pick the audio-torch* pin matching your torch minor (see TORCH_TORCHCODEC). +# torchcodec publishes no sdist and only manylinux_2_28_x86_64, macosx_*_arm64 +# and win_amd64 wheels, so Linux aarch64, Windows ARM64 and Intel Mac have +# nothing to resolve and pip fails the whole install rather than skipping audio. +# Gate on the platforms that have a wheel, matching +# PLATFORM_LACKS_TORCHCODEC_WHEEL in studio/install_python_stack.py. +audio-torch210 = [ + "torchcodec>=0.10.0,<0.11.0 ; python_version >= '3.10' and (((sys_platform == 'linux' or sys_platform == 'win32') and (platform_machine == 'x86_64' or platform_machine == 'AMD64')) or (sys_platform == 'darwin' and platform_machine == 'arm64'))", +] +audio-torch290 = [ + "torchcodec>=0.8.0,<0.10.0 ; python_version >= '3.10' and (((sys_platform == 'linux' or sys_platform == 'win32') and (platform_machine == 'x86_64' or platform_machine == 'AMD64')) or (sys_platform == 'darwin' and platform_machine == 'arm64'))", +] +audio-torch280 = [ + "torchcodec>=0.6.0,<0.8.0 ; python_version >= '3.9' and (((sys_platform == 'linux' or sys_platform == 'win32') and (platform_machine == 'x86_64' or platform_machine == 'AMD64')) or (sys_platform == 'darwin' and platform_machine == 'arm64'))", +] huggingface = [ "unsloth[huggingfacenotorch]", - "unsloth_zoo>=2026.7.2", + "unsloth_zoo>=2026.7.6", "torchvision", "unsloth[triton]", ] @@ -532,16 +586,19 @@ cu126-torch2100 = [ "unsloth[huggingface]", "bitsandbytes>=0.45.5,!=0.46.0,!=0.48.0", "unsloth[cu126onlytorch2100]", + "unsloth[audio-torch210]", ] cu128-torch2100 = [ "unsloth[huggingface]", "bitsandbytes>=0.45.5,!=0.46.0,!=0.48.0", "unsloth[cu128onlytorch2100]", + "unsloth[audio-torch210]", ] cu130-torch2100 = [ "unsloth[huggingface]", "bitsandbytes>=0.45.5,!=0.46.0,!=0.48.0", "unsloth[cu130onlytorch2100]", + "unsloth[audio-torch210]", ] kaggle = [ "unsloth[huggingface]", @@ -580,7 +637,7 @@ colab-ampere-torch220 = [ "flash-attn>=2.6.3 ; ('linux' in sys_platform)", ] colab-new = [ - "unsloth_zoo>=2026.7.2", + "unsloth_zoo>=2026.7.6", "packaging", "tyro", "transformers>=4.51.3,!=4.52.0,!=4.52.1,!=4.52.2,!=4.52.3,!=4.53.0,!=4.54.0,!=4.55.0,!=4.55.1,!=4.57.0,!=4.57.4,!=4.57.5,!=5.0.0,!=5.1.0,<=5.5.0", @@ -831,16 +888,19 @@ cu126-ampere-torch2100 = [ "unsloth[huggingface]", "bitsandbytes>=0.45.5,!=0.46.0,!=0.48.0", "unsloth[cu126onlytorch2100]", + "unsloth[audio-torch210]", ] cu128-ampere-torch2100 = [ "unsloth[huggingface]", "bitsandbytes>=0.45.5,!=0.46.0,!=0.48.0", "unsloth[cu128onlytorch2100]", + "unsloth[audio-torch210]", ] cu130-ampere-torch2100 = [ "unsloth[huggingface]", "bitsandbytes>=0.45.5,!=0.46.0,!=0.48.0", "unsloth[cu130onlytorch2100]", + "unsloth[audio-torch210]", ] flashattentiontorch260abiFALSEcu12x = [ "flash-attn @ https://github.com/Dao-AILab/flash-attention/releases/download/v2.7.4.post1/flash_attn-2.7.4.post1+cu12torch2.6cxx11abiFALSE-cp39-cp39-linux_x86_64.whl ; ('linux' in sys_platform) and python_version == '3.9'", @@ -1125,7 +1185,8 @@ intelgputorch210 = [ "torchvision @ https://download.pytorch.org/whl/xpu/torchvision-0.25.0%2Bxpu-cp313-cp313-win_amd64.whl#sha256=1c4b44b36a557f7381e3076fb8843366742238648441d607c8d049c6da0f8886 ; sys_platform == 'win32' and python_version == '3.13' and (platform_machine == 'AMD64' or platform_machine == 'x86_64')", ] intel-gpu-torch210 = [ - "unsloth[intelgputorch210]" + "unsloth[intelgputorch210]", + "unsloth[audio-torch210]", ] intelgputorch2110 = [ "unsloth_zoo[intelgpu]", @@ -1206,8 +1267,11 @@ intel = [ ] amd = [ "unsloth[huggingfacenotorch]", - "bitsandbytes>=0.49.1 ; ('linux' in sys_platform) and (platform_machine == 'AMD64' or platform_machine == 'x86_64' or platform_machine == 'aarch64')", - "bitsandbytes>=0.49.1 ; (sys_platform == 'win32') and (platform_machine == 'AMD64' or platform_machine == 'x86_64')", + # 4-bit decode is unreliable on ROCm before 0.50.0, the first PyPI release + # carrying the full path: blocksize/warp decoupling (bnb #1887), fused SIMT + # GEMM on RDNA (#1979), RDNA3/4 workgroup fix (#2012). + "bitsandbytes>=0.50.0 ; ('linux' in sys_platform) and (platform_machine == 'AMD64' or platform_machine == 'x86_64' or platform_machine == 'aarch64')", + "bitsandbytes>=0.50.0 ; (sys_platform == 'win32') and (platform_machine == 'AMD64' or platform_machine == 'x86_64')", ] rocm702-torch280 = [ "unsloth[amd]", @@ -1279,6 +1343,7 @@ rocm72-torch2100 = [ "torchvision @ https://repo.radeon.com/rocm/manylinux/rocm-rel-7.2/torchvision-0.25.0%2Brocm7.2.0.git82df5f59-cp311-cp311-linux_x86_64.whl ; platform_system == 'Linux' and python_version == '3.11' and platform_machine == 'x86_64'", "torchvision @ https://repo.radeon.com/rocm/manylinux/rocm-rel-7.2/torchvision-0.25.0%2Brocm7.2.0.git82df5f59-cp312-cp312-linux_x86_64.whl ; platform_system == 'Linux' and python_version == '3.12' and platform_machine == 'x86_64'", "torchvision @ https://repo.radeon.com/rocm/manylinux/rocm-rel-7.2/torchvision-0.25.0%2Brocm7.2.0.git82df5f59-cp313-cp313-linux_x86_64.whl ; platform_system == 'Linux' and python_version == '3.13' and platform_machine == 'x86_64'", + "unsloth[audio-torch210]", ] rocm711-torch2100 = [ "unsloth[amd]", @@ -1297,6 +1362,7 @@ rocm711-torch2100 = [ "torchvision @ https://repo.radeon.com/rocm/manylinux/rocm-rel-7.1.1/torchvision-0.25.0%2Brocm7.1.1.git82df5f59-cp311-cp311-linux_x86_64.whl ; platform_system == 'Linux' and python_version == '3.11' and platform_machine == 'x86_64'", "torchvision @ https://repo.radeon.com/rocm/manylinux/rocm-rel-7.1.1/torchvision-0.25.0%2Brocm7.1.1.git82df5f59-cp312-cp312-linux_x86_64.whl ; platform_system == 'Linux' and python_version == '3.12' and platform_machine == 'x86_64'", "torchvision @ https://repo.radeon.com/rocm/manylinux/rocm-rel-7.1.1/torchvision-0.25.0%2Brocm7.1.1.git82df5f59-cp313-cp313-linux_x86_64.whl ; platform_system == 'Linux' and python_version == '3.13' and platform_machine == 'x86_64'", + "unsloth[audio-torch210]", ] [project.urls] diff --git a/scripts/build_whisper_cpp.sh b/scripts/build_whisper_cpp.sh new file mode 100755 index 0000000000..9f7e4d4ef3 --- /dev/null +++ b/scripts/build_whisper_cpp.sh @@ -0,0 +1,71 @@ +#!/bin/sh +# Build whisper.cpp's whisper-server for Studio's GGUF dictation engine. +# +# Installs into the managed Studio home so the backend's binary discovery +# (core/inference/stt_ggml_sidecar.py::find_whisper_server_binary) picks it up: +# /whisper.cpp/build/bin/whisper-server (custom home) +# ~/.unsloth/whisper.cpp/build/bin/whisper-server (default) +# +# Usage: +# ./scripts/build_whisper_cpp.sh # build the pinned tag +# WHISPER_CPP_TAG=v1.9.0 ./scripts/build_whisper_cpp.sh +# +# Requires: git, cmake, a C/C++ toolchain (the same prerequisites as a +# llama.cpp source build). GPU backends are auto-detected by whisper.cpp's +# CMake (Metal on macOS; set GGML_CUDA=1 to force a CUDA build on Linux). + +set -eu + +WHISPER_CPP_SOURCE="${WHISPER_CPP_SOURCE:-https://github.com/ggml-org/whisper.cpp}" +WHISPER_CPP_TAG="${WHISPER_CPP_TAG:-v1.9.1}" + +STUDIO_HOME="${UNSLOTH_STUDIO_HOME:-${STUDIO_HOME:-}}" +CUSTOM_STUDIO_HOME=false +if [ -n "$STUDIO_HOME" ]; then + CUSTOM_STUDIO_HOME=true + INSTALL_DIR="$STUDIO_HOME/whisper.cpp" +else + INSTALL_DIR="$HOME/.unsloth/whisper.cpp" +fi + +command -v git >/dev/null 2>&1 || { echo "ERROR: git is required" >&2; exit 1; } +command -v cmake >/dev/null 2>&1 || { echo "ERROR: cmake is required" >&2; exit 1; } + +# Same policy as studio/setup.sh's _assert_studio_owned_or_absent: never delete +# a directory under a custom Studio home unless Studio itself created it (the +# marker file below). Protects a user-managed whisper.cpp/src from rm -rf. +STUDIO_OWNED_MARKER=".unsloth-studio-owned" +if [ "$CUSTOM_STUDIO_HOME" = true ] && [ -e "$INSTALL_DIR" ] && \ + [ ! -f "$INSTALL_DIR/$STUDIO_OWNED_MARKER" ]; then + echo "ERROR: $INSTALL_DIR already exists and is not marked as an Unsloth-owned whisper.cpp build tree." >&2 + echo " Move it aside or choose an empty UNSLOTH_STUDIO_HOME before re-running." >&2 + exit 1 +fi + +echo "==> Building whisper.cpp ($WHISPER_CPP_TAG) into $INSTALL_DIR" +mkdir -p "$INSTALL_DIR" +: > "$INSTALL_DIR/$STUDIO_OWNED_MARKER" + +if [ ! -d "$INSTALL_DIR/src/.git" ]; then + rm -rf "$INSTALL_DIR/src" + git clone --depth 1 --branch "$WHISPER_CPP_TAG" "$WHISPER_CPP_SOURCE" "$INSTALL_DIR/src" +else + git -C "$INSTALL_DIR/src" fetch --depth 1 origin "$WHISPER_CPP_TAG" + git -C "$INSTALL_DIR/src" checkout FETCH_HEAD +fi + +CMAKE_FLAGS="-DCMAKE_BUILD_TYPE=Release -DBUILD_SHARED_LIBS=OFF" +if [ "${GGML_CUDA:-0}" = "1" ]; then + CMAKE_FLAGS="$CMAKE_FLAGS -DGGML_CUDA=ON" +fi + +# shellcheck disable=SC2086 +cmake -S "$INSTALL_DIR/src" -B "$INSTALL_DIR/src/build" $CMAKE_FLAGS +NCPU="$(getconf _NPROCESSORS_ONLN 2>/dev/null || echo 4)" +cmake --build "$INSTALL_DIR/src/build" --config Release --target whisper-server -j"$NCPU" + +mkdir -p "$INSTALL_DIR/build/bin" +cp "$INSTALL_DIR/src/build/bin/whisper-server" "$INSTALL_DIR/build/bin/whisper-server" + +echo "==> Installed $INSTALL_DIR/build/bin/whisper-server" +"$INSTALL_DIR/build/bin/whisper-server" --help >/dev/null 2>&1 && echo "==> Binary runs OK" diff --git a/scripts/install_rocm_wsl_strixhalo.sh b/scripts/install_rocm_wsl_strixhalo.sh index aa560fc432..697aae933f 100644 --- a/scripts/install_rocm_wsl_strixhalo.sh +++ b/scripts/install_rocm_wsl_strixhalo.sh @@ -219,7 +219,7 @@ fi echo "${ROCM_DIR}/lib" | $SUDO tee /etc/ld.so.conf.d/rocm.conf >/dev/null $SUDO ldconfig -# ── Step 4: persist environment (system-wide so Studio's worker inherits it) ── +# ── Step 4: persist environment (system-wide so Unsloth's worker inherits it) ── say "Persisting ROCm-on-WSL environment" _envfile="/etc/profile.d/unsloth-rocm-wsl.sh" $SUDO tee "$_envfile" >/dev/null < list[str]: - text = path.read_text() + text = path.read_text(encoding = "utf-8") keys: list[str] = [] for m in re.finditer(r"(?:^|\n)\s*key:\s*([^\n]+)", text): keys.append(m.group(1).strip()) @@ -104,7 +104,7 @@ def main() -> int: for t in RESTRICTED_TRIGGERS: if t in triggers: - text = path.read_text() + text = path.read_text(encoding = "utf-8") if "lint:workflow_triggers-allow-workflow_run" not in text: findings.append( f"{path.name}: RESTRICTED trigger '{t}' requires an " diff --git a/scripts/lockfile_supply_chain_audit.py b/scripts/lockfile_supply_chain_audit.py index 66b48c094d..f9cf726dc1 100644 --- a/scripts/lockfile_supply_chain_audit.py +++ b/scripts/lockfile_supply_chain_audit.py @@ -2,7 +2,7 @@ # SPDX-License-Identifier: AGPL-3.0-only # Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. -"""Lockfile supply-chain audit for the Studio frontend and Tauri shell. +"""Lockfile supply-chain audit for the Unsloth frontend and Tauri shell. Runs BEFORE `npm ci` / `cargo fetch` in CI. Refuses to proceed when a lockfile contains patterns indicating supply-chain injection (npm @@ -294,7 +294,7 @@ CARGO_REGISTRY_SOURCE = "registry+https://github.com/rust-lang/crates.io-index" # Cargo non-registry source allowlist: `(crate_name, exact_source_string)`. # Both must match verbatim; bumping the pinned SHA forces a re-review. -# Studio's Tauri shell pulls `fix-path-env` from git because it is not +# Unsloth's Tauri shell pulls `fix-path-env` from git because it is not # published to crates.io; commit c4c45d5 was reviewed when it landed. CARGO_SOURCE_ALLOWLIST: tuple[tuple[str, str], ...] = ( ( diff --git a/scripts/notebook_validator.py b/scripts/notebook_validator.py index c1be7a63a4..7bcee47c66 100644 --- a/scripts/notebook_validator.py +++ b/scripts/notebook_validator.py @@ -95,8 +95,8 @@ COLAB_ORACLE_BASE_URL = "https://raw.githubusercontent.com/googlecolab/backend-i # Source: pytorch/torchcodec compatibility matrix on its README. TORCH_TORCHCODEC: dict[str, set[str]] = { "2.10": {"0.10"}, - "2.9": {"0.7", "0.8", "0.9"}, - "2.8": {"0.6"}, + "2.9": {"0.8", "0.9"}, + "2.8": {"0.6", "0.7"}, "2.7": {"0.3", "0.4", "0.5"}, "2.6": {"0.2", "0.3"}, "2.5": {"0.1", "0.2"}, diff --git a/scripts/profile_startup.py b/scripts/profile_startup.py new file mode 100644 index 0000000000..937d007ac1 --- /dev/null +++ b/scripts/profile_startup.py @@ -0,0 +1,377 @@ +#!/usr/bin/env python3 +# SPDX-License-Identifier: AGPL-3.0-only +# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0 + +"""Measure where Unsloth Studio's startup time goes, per platform. + +Nothing measured this before: the backend logs "lifespan startup completed in X ms" +but no test or CI job asserted a budget, and studio_test_kit discards the elapsed +time of its /healthz poll. A first local run (Linux, warm cache, fast server CPU) +found `import main` alone costs 6.6s before the server can bind, dominated by eager +module-level imports pulled in by the `routes` package: + + torch 1930 ms self + unsloth_zoo 914 ms self + routes 779 ms self + transformers 524 ms self + +Phases measured: + import `python -X importtime -c "import main"`, top cumulative + per-package self + spawn process start -> first byte on stdout + healthz process start -> /api/health (or /healthz) answers 200 + lifespan the backend's own "lifespan startup completed in X ms" log line + +Usage: + python scripts/profile_startup.py --repeats 3 --json out.json + python scripts/profile_startup.py --import-only # no server, no port needed + +Exit code is 0 unless --max-healthz-seconds is given and exceeded. +""" + +from __future__ import annotations + +import argparse +import json +import math +import os +import platform +import re +import shutil +import socket +import statistics +import subprocess +import sys +import threading +import time +import urllib.error +import urllib.request +from pathlib import Path + +REPO_ROOT = Path(__file__).resolve().parents[1] +BACKEND = REPO_ROOT / "studio" / "backend" + +_IMPORTTIME_RE = re.compile(r"import time:\s+(\d+)\s+\|\s+(\d+)\s+\|(\s*)(\S.*)") + + +def _free_port() -> int: + with socket.socket() as s: + s.bind(("127.0.0.1", 0)) + return int(s.getsockname()[1]) + + +def profile_imports(python: str, top: int = 15) -> dict: + """Cumulative and self import cost for the backend's module graph. + + Run in a subprocess with -X importtime: the numbers are only meaningful for a + cold interpreter, and importing in-process would measure a warm sys.modules. + """ + proc = subprocess.run( + [python, "-X", "importtime", "-c", "import sys; sys.path.insert(0, '.'); import main"], + cwd = BACKEND, + capture_output = True, + text = True, + timeout = 900, + ) + rows = [] + for line in proc.stderr.splitlines(): + m = _IMPORTTIME_RE.match(line) + if m: + rows.append((int(m.group(1)), int(m.group(2)), m.group(4).strip())) + if not rows: + return {"ok": False, "error": (proc.stderr or proc.stdout)[-2000:]} + if proc.returncode != 0: + # Rows survive up to the failure, so any total from a partial graph is wrong. + return { + "ok": False, + "error": (proc.stderr or proc.stdout)[-2000:], + "partial_rows": len(rows), + } + + by_cum = sorted(rows, key = lambda r: -r[1]) + # Total comes from the `main` row, not by_cum[0]: -X importtime also prints the + # interpreter's own startup graph (`site`), which can outrank a trivial main. + main_row = next((r for r in reversed(rows) if r[2] == "main"), None) + if main_row is None: + return { + "ok": False, + "error": "no `import main` row in -X importtime output\n" + + (proc.stderr or proc.stdout)[-2000:], + } + self_by_pkg: dict[str, int] = {} + for self_us, _cum, name in rows: + pkg = name.split(".")[0] + self_by_pkg[pkg] = self_by_pkg.get(pkg, 0) + self_us + + return { + "ok": True, + "total_seconds": round(main_row[1] / 1e6, 3), + "top_cumulative": [ + {"module": n, "seconds": round(c / 1e6, 3)} for _s, c, n in by_cum[:top] + ], + "self_by_package_ms": { + k: round(v / 1000) for k, v in sorted(self_by_pkg.items(), key = lambda x: -x[1])[:top] + }, + } + + +def _terminate_tree(proc: subprocess.Popen) -> None: + """Stop the server AND its children, which on Windows are a separate process. + + CI profiles `Scripts/unsloth.exe`, a distlib launcher stub that CreateProcess's + the venv python and waits, so terminate() reaps the stub only: the real backend + keeps the inherited stdout handle, the reader thread never sees EOF, and + --repeats strands one server per iteration on the shared UNSLOTH_STUDIO_HOME. + taskkill /T walks the tree, as unsloth_cli/commands/start.py already does. + """ + if proc.poll() is not None: + return + if os.name == "nt": + try: + killed = subprocess.run( + ["taskkill", "/PID", str(proc.pid), "/T", "/F"], + capture_output = True, + timeout = 30, + check = False, + ) + if killed.returncode == 0: + return + except Exception: + # taskkill missing or timed out; fall through so the stub still dies. + pass + # check=False: a nonzero taskkill does not raise, so fall through as well. + proc.terminate() + + +def profile_launch( + bin_path: str, + port: int, + timeout_s: int = 300, +) -> dict: + """Spawn the backend the way the desktop app does and time it to first 200.""" + log_lines: list[str] = [] + first_byte: list[float] = [] + t0 = time.perf_counter() + proc = subprocess.Popen( + [bin_path, "studio", "--api-only", "-H", "127.0.0.1", "-p", str(port)], + cwd = REPO_ROOT, + stdout = subprocess.PIPE, + stderr = subprocess.STDOUT, + text = True, + bufsize = 1, + ) + + def _drain() -> None: + # Runs alongside the health polling: the first read timestamps the spawn + # phase, and an undrained pipe blocks the backend before it binds. + for line in proc.stdout: + if not first_byte: + first_byte.append(time.perf_counter() - t0) + log_lines.append(line.rstrip("\n")) + + reader = threading.Thread(target = _drain, daemon = True) + reader.start() + + t_healthz = None + deadline = t0 + timeout_s + try: + while time.perf_counter() < deadline: + if proc.poll() is not None: + break + if t_healthz is None: + for url in ( + f"http://127.0.0.1:{port}/api/health", + f"http://127.0.0.1:{port}/healthz", + ): + try: + with urllib.request.urlopen(url, timeout = 2) as r: + if r.status == 200: + t_healthz = time.perf_counter() - t0 + break + except (urllib.error.URLError, OSError, TimeoutError): + pass + if t_healthz is not None: + break + time.sleep(0.25) + finally: + _terminate_tree(proc) + try: + # Safe: the reader drains the pipe, so the child cannot block on write(). + proc.wait(timeout = 30) + except subprocess.TimeoutExpired: + proc.kill() + proc.wait() + reader.join(timeout = 10) + + t_first_byte = first_byte[0] if first_byte else None + lifespan_ms = None + for line in log_lines: + m = re.search(r"lifespan startup completed in ([\d.]+)ms", line) + if m: + lifespan_ms = float(m.group(1)) + return { + "spawn_seconds": round(t_first_byte, 3) if t_first_byte is not None else None, + "healthz_seconds": round(t_healthz, 3) if t_healthz is not None else None, + "lifespan_ms": lifespan_ms, + "reached_healthz": t_healthz is not None, + "log_tail": log_lines[-25:], + } + + +def python_version_of(python: str) -> str: + """Version of the interpreter that runs the imports, not the one running us. + + --python points at the installed Studio venv while this script runs under the + runner's system python, so platform.python_version() would label it wrong. + """ + if python == sys.executable: + return platform.python_version() + try: + proc = subprocess.run( + [python, "-c", "import platform; print(platform.python_version())"], + capture_output = True, + text = True, + timeout = 60, + ) + if proc.returncode == 0 and proc.stdout.strip(): + return proc.stdout.strip() + except (OSError, subprocess.SubprocessError): + pass + return "unknown" + + +def find_bin() -> str | None: + home = os.environ.get("UNSLOTH_STUDIO_HOME") or str(Path.home() / ".unsloth" / "studio") + names = ["unsloth.exe", "unsloth"] if platform.system() == "Windows" else ["unsloth"] + subdirs = ["unsloth_studio/Scripts", "unsloth_studio/bin", "bin", "Scripts"] + for sd in subdirs: + for n in names: + p = Path(home) / sd / n + if p.exists(): + return str(p) + return shutil.which("unsloth") + + +def main(argv: list[str]) -> int: + ap = argparse.ArgumentParser( + description = __doc__, formatter_class = argparse.RawDescriptionHelpFormatter + ) + ap.add_argument( + "--repeats", + type = int, + default = 1, + help = "launch repeats; the median is reported (imports are measured once)", + ) + ap.add_argument( + "--python", + default = sys.executable, + help = "interpreter used for the import profile (default: this one)", + ) + ap.add_argument("--bin", help = "path to the unsloth CLI (default: autodetect)") + ap.add_argument( + "--import-only", + action = "store_true", + help = "skip the server phases (no install needed beyond the deps)", + ) + ap.add_argument( + "--max-healthz-seconds", + type = float, + help = "fail if the median time to a healthy port exceeds this", + ) + ap.add_argument("--json", help = "write the full report here") + a = ap.parse_args(argv) + # range(0) launches nothing, leaving the budget check with nothing to fail on. + if a.repeats < 1: + ap.error("--repeats must be at least 1") + # Same reason: --import-only never launches anything. + if a.import_only and a.max_healthz_seconds is not None: + ap.error("--max-healthz-seconds cannot be combined with --import-only") + # nan and inf parse fine as floats but `med > budget` is then always False, + # so the gate would report success without ever bounding anything. + if a.max_healthz_seconds is not None and not math.isfinite(a.max_healthz_seconds): + ap.error("--max-healthz-seconds must be a finite number") + + report: dict = { + "platform": platform.system().lower(), + "machine": platform.machine(), + "python": python_version_of(a.python), + "cpu_count": os.cpu_count(), + } + + print("== import graph ==") + report["imports"] = profile_imports(a.python) + imp = report["imports"] + if imp.get("ok"): + print(f" import main: {imp['total_seconds']}s") + for row in imp["top_cumulative"][:8]: + print(f" {row['seconds']:7.3f}s {row['module']}") + print(" self time by package (ms):") + for k, v in list(imp["self_by_package_ms"].items())[:8]: + print(f" {v:8} ms {k}") + else: + print(f" FAILED: {imp.get('error', '')[:400]}") + + if not a.import_only: + bin_path = a.bin or find_bin() + if not bin_path: + print( + "== launch == skipped: no unsloth CLI found " + "(set UNSLOTH_STUDIO_HOME or pass --bin)" + ) + report["launch"] = {"skipped": "no unsloth CLI found"} + else: + print(f"== launch == {bin_path}") + runs = [] + for i in range(a.repeats): + r = profile_launch(bin_path, _free_port()) + runs.append(r) + print( + f" run {i + 1}: healthz={r['healthz_seconds']}s " + f"lifespan={r['lifespan_ms']}ms reached={r['reached_healthz']}" + ) + got = [r["healthz_seconds"] for r in runs if r["healthz_seconds"] is not None] + report["launch"] = { + "runs": runs, + "failed_runs": sum(1 for r in runs if not r["reached_healthz"]), + "healthz_median_seconds": round(statistics.median(got), 3) if got else None, + "healthz_max_seconds": round(max(got), 3) if got else None, + } + if got: + print( + f" median time to healthy port: {report['launch']['healthz_median_seconds']}s" + ) + + if a.json: + Path(a.json).write_text(json.dumps(report, indent = 2), encoding = "utf-8") + print(f"\nwrote {a.json}") + + if a.max_healthz_seconds is not None: + launch = report.get("launch") or {} + med = launch.get("healthz_median_seconds") + failed = launch.get("failed_runs") or 0 + if failed: + # Failed launches fail the budget; dropping them would keep only the fast ones. + print( + f"::error::startup regression: {failed} of {len(launch.get('runs') or [])} " + f"launches never became healthy within the timeout" + ) + return 1 + if med is None: + # Nothing measured: exiting 0 would pass a requested budget without a + # single health request, so fail closed. + print( + "::error::startup regression: no healthz measurement, so the " + f"{a.max_healthz_seconds}s budget was never checked " + f"({launch.get('skipped') or 'launch phase produced no runs'})" + ) + return 1 + elif med > a.max_healthz_seconds: + print( + f"::error::startup regression: {med}s median to a healthy port " + f"exceeds the {a.max_healthz_seconds}s budget" + ) + return 1 + return 0 + + +if __name__ == "__main__": + raise SystemExit(main(sys.argv[1:])) diff --git a/scripts/scan_npm_packages.py b/scripts/scan_npm_packages.py index 47b85147ca..6c83552727 100644 --- a/scripts/scan_npm_packages.py +++ b/scripts/scan_npm_packages.py @@ -62,7 +62,7 @@ REPO_ROOT = Path(__file__).resolve().parents[1] # Hard caps (deliberately conservative; npm tarballs in this repo are # all well under these limits, so a packaging spike is noticeable). # ───────────────────────────────────────────────────────────────────── -# Caps calibrated against the real Studio frontend transitive closure: +# Caps calibrated against the real Unsloth frontend transitive closure: # - typescript.js is 9.1 MB (TS compiler bundled into one file) # - mermaid 11.x dist/mermaid.js.map is ~12 MB (sourcemap) # - lightningcss-linux-x64-{gnu,musl}.node is 10 MB diff --git a/scripts/scan_packages_baseline.json b/scripts/scan_packages_baseline.json index 929cc37bda..58b7f95ab1 100644 --- a/scripts/scan_packages_baseline.json +++ b/scripts/scan_packages_baseline.json @@ -1,5 +1,5 @@ { - "_comment": "scan_packages.py allowlist. Each entry is a CRITICAL/HIGH finding manually judged benign. Matched on (package, package-relative file, check, evidence_hash); evidence_hash is over the matched code with L: markers stripped, so version bumps and line shifts do not reopen an entry but changed code does. severity and evidence are for review only. Regenerate with --write-baseline AFTER reviewing every line.", + "_comment": "scan_packages.py allowlist (reviewed). Each entry is a CRITICAL/HIGH finding manually judged benign. Matched on (package, package-relative file, check, evidence_hash); evidence_hash is over the matched code with L: markers stripped, so version bumps and line shifts do not reopen an entry but changed code does. severity and evidence are for review only. Regenerate with --write-baseline AFTER reviewing every line.", "version": 1, "entries": [ { @@ -95,8 +95,16 @@ "file": "fastapi/routing.py", "check": "C2 polling/beaconing loop detected", "severity": "CRITICAL", - "evidence": "L586: while True: sha256:251135b5ebfdd1248916449f32262575e003ef64382501c65b7e4061d67bda45", - "evidence_hash": "365aef4449c8089753d9398417cd76ab762cef547d75db70d87bca9c0b550ab5" + "evidence": "L587: while True: sha256:06c2c7f15d73bf192e5e3272c5ff5fcaeff7f6774fef5f4eca6ef473ae50e2b3", + "evidence_hash": "57acd497f404c203e4450d0580ad85aa8a33406e8d64ad06fbac6cf47d97b24d" + }, + { + "package": "fastapi", + "file": "fastapi/routing.py", + "check": "C2 polling/beaconing loop detected", + "severity": "CRITICAL", + "evidence": "L592: while True: sha256:84283c09277ded3296998b2a6a838744457b606829cf5ab5d0da6f222ff020a0", + "evidence_hash": "a7295004315e26a8f3c64fb837521e9fdd7268219bb43e000fb0236ab0259223" }, { "package": "fastmcp-slim", @@ -303,8 +311,8 @@ "file": "openai/_base_client.py", "check": "C2 polling/beaconing loop detected", "severity": "CRITICAL", - "evidence": "L264: while True: sha256:95ca67e46d42354ae650abbdc5b0d97df8b0ed43187800bf40f5690c3901b94b", - "evidence_hash": "a57d8d15fed0bf04f9967dcc18a18b80bb19f4095675bccbb78ac0450d7fce14" + "evidence": "L274: while True: sha256:90a38e5c1e26893c7c273354143612640e9a9c0f079d3e2b60612d79f24e80a6", + "evidence_hash": "1022e8e8649436ec64a98a9d9141d085452c49549fd2157b0278fc369a83ac66" }, { "package": "openai", @@ -319,8 +327,8 @@ "file": "openai/auth/_workload.py", "check": "Accesses cloud metadata/IMDS AND makes network calls", "severity": "CRITICAL", - "evidence": "IMDS: L96: url = \"http://169.254.169.254/metadata/identity/oauth2/token\" | L149: url = \"http://metadata.google.internal/computeMetadata/v1/instance/service-accounts/default/identity\"\nNetwork: L77: http_client: httpx.Client | None = None, | L108: with httpx.Client() as client: | L133: http_client: httpx.Client | None = None, | L155: with httpx.Client() as client: | L248: with httpx.Client() as client:", - "evidence_hash": "1581d9f4a23393e9af23fbe5ef9f66807b22c5b5a3f1fe167254c9ebee108567" + "evidence": "IMDS: L97: url = \"http://169.254.169.254/metadata/identity/oauth2/token\" | L150: url = \"http://metadata.google.internal/computeMetadata/v1/instance/service-accounts/default/identity\"\nNetwork: L78: http_client: httpx.Client | None = None, | L109: with httpx.Client() as client: | L134: http_client: httpx.Client | None = None, | L156: with httpx.Client() as client: | L251: exchange_client = DefaultHttpx2Client(follow_redirects=False) if self._use_httpx2 else httpx.Client()", + "evidence_hash": "9717e51cb961dc14c458955d91a1e48e3753997346ecea0106bded3a8d64bfe0" }, { "package": "openai", @@ -343,8 +351,8 @@ "file": "openai/resources/beta/responses/responses.py", "check": "C2 polling/beaconing loop detected", "severity": "CRITICAL", - "evidence": "L3999: while True: sha256:df298b6eaf3416589b79f4ef283f8fb76e54d505bfda8840673f8e6419117e2e", - "evidence_hash": "10ce5cb5a7097fcff4042ddcfb4802edda60aa4b7b113c8b926a52ddb76f78c2" + "evidence": "L4000: while True: sha256:f8ab538118daba9ec06e27399dbdc90a4521c3390e6a47a6348a1f180a83effd", + "evidence_hash": "31481ea83c687acc27144d72d3832d4fb98dd1c79fb5e0ddd85080de95997b9f" }, { "package": "openai", @@ -359,16 +367,16 @@ "file": "openai/resources/realtime/realtime.py", "check": "C2 polling/beaconing loop detected", "severity": "CRITICAL", - "evidence": "L310: while True: sha256:458198ff3d3f05870bf98c9564cbfd68c739e57b9bbe4120ed81e3eb6af74a05", - "evidence_hash": "a3165d21e46b3ce553795daeae53e8f80e8e89c5cb228e68e6dcaff54bca5a89" + "evidence": "L311: while True: sha256:5b63313072aae9ca28677e03426513ccf12221e4f4e0ea6c31efbe09790633b5", + "evidence_hash": "05e1af469d651b51673763a7c4cdf759af9472fb627b7b470adc28cc237bd650" }, { "package": "openai", "file": "openai/resources/responses/responses.py", "check": "C2 polling/beaconing loop detected", "severity": "CRITICAL", - "evidence": "L3950: while True: sha256:1ce0b5a388c747945cdfda1a71b77afdfd03ae840d7aa9fa62f02eb00aa5e29f", - "evidence_hash": "6de300ebb5e6e17cb51c89cbcdf08515a44655182f0776f0908a9d1043ebbcd7" + "evidence": "L3951: while True: sha256:d68ef896bf0743ca430cfacb9a3353da1f3b9c51c3a21b6450a07a32b55aa2ac", + "evidence_hash": "160eecdd79b521bffbe8476f782b69a0724c35d1b19376a7600807165fd54f9f" }, { "package": "openai", @@ -1545,6 +1553,78 @@ "severity": "HIGH", "evidence": "Obfusc: L836: code = compile(module, \"\", \"exec\")\nExec: L736: exec(code, globs, locs)", "evidence_hash": "5c0992c90f05c772abd94d00784f157de337e1f8567f8b3aee1b15e46c96cd5d" + }, + { + "package": "unsloth-zoo", + "file": "tests/test_mlx_save_export_regressions.py", + "check": "Writes to /tmp and executes (staged dropper)", + "severity": "CRITICAL", + "evidence": "L165: temporary_location=\"/tmp/ignored\", sha256:ab5c587f9ec31a0cc10ee55698ab133a417148d9d3f371bbc81b1e13fa119c13", + "evidence_hash": "93a11159147aad94f353ec4d2e0b8486b256abef88cd96d741813222cd32b138" + }, + { + "package": "unsloth-zoo", + "file": "tests/test_vision_collator_audio.py", + "check": "Writes to /tmp and executes (staged dropper)", + "severity": "CRITICAL", + "evidence": "L111: out = extract_audio_info(msgs({\"type\": \"audio\", key: \"/tmp/a.wav\"})) sha256:2efe23ffbe2b91b8403aec9b700736919b59e5ca770f8e1f5501651b44b7d398", + "evidence_hash": "d416b79dd17b24214f3f7653ac01354507d7bf0fc464dee30a4a4b8998f063ba" + }, + { + "package": "openai", + "file": "openai/_base_client.py", + "check": "C2 polling/beaconing loop detected", + "severity": "CRITICAL", + "evidence": "L274: while True: sha256:90a38e5c1e26893c7c273354143612640e9a9c0f079d3e2b60612d79f24e80a6", + "evidence_hash": "1022e8e8649436ec64a98a9d9141d085452c49549fd2157b0278fc369a83ac66" + }, + { + "package": "openai", + "file": "openai/auth/_workload.py", + "check": "Accesses cloud metadata/IMDS AND makes network calls", + "severity": "CRITICAL", + "evidence": "IMDS: L97: url = \"http://169.254.169.254/metadata/identity/oauth2/token\" | L150: url = \"http://metadata.google.internal/computeMetadata/v1/instance/service-accounts/default/identity\"\nNetwork: L78: http_client: httpx.Client | None = None, | L109: with httpx.Client() as client: | L134: http_client: httpx.Client | None = None, | L156: with httpx.Client() as client: | L251: exchange_client = DefaultHttpx2Client(follow_redirects=False) if self._use_httpx2 else httpx.Client()", + "evidence_hash": "9717e51cb961dc14c458955d91a1e48e3753997346ecea0106bded3a8d64bfe0" + }, + { + "package": "openai", + "file": "openai/resources/beta/responses/responses.py", + "check": "C2 polling/beaconing loop detected", + "severity": "CRITICAL", + "evidence": "L4000: while True: sha256:f8ab538118daba9ec06e27399dbdc90a4521c3390e6a47a6348a1f180a83effd", + "evidence_hash": "31481ea83c687acc27144d72d3832d4fb98dd1c79fb5e0ddd85080de95997b9f" + }, + { + "package": "openai", + "file": "openai/resources/realtime/realtime.py", + "check": "C2 polling/beaconing loop detected", + "severity": "CRITICAL", + "evidence": "L311: while True: sha256:5b63313072aae9ca28677e03426513ccf12221e4f4e0ea6c31efbe09790633b5", + "evidence_hash": "05e1af469d651b51673763a7c4cdf759af9472fb627b7b470adc28cc237bd650" + }, + { + "package": "openai", + "file": "openai/resources/responses/responses.py", + "check": "C2 polling/beaconing loop detected", + "severity": "CRITICAL", + "evidence": "L3951: while True: sha256:d68ef896bf0743ca430cfacb9a3353da1f3b9c51c3a21b6450a07a32b55aa2ac", + "evidence_hash": "160eecdd79b521bffbe8476f782b69a0724c35d1b19376a7600807165fd54f9f" + }, + { + "package": "unsloth-zoo", + "file": "tests/test_gemma4_forced_float32_ple_dtype.py", + "check": "Advanced obfuscation (marshal/compile/zlib) + exec/eval", + "severity": "HIGH", + "evidence": "Obfusc: L277: compile(rewritten + _GEMMA4_PLE_CAST_HELPER, \"\", \"exec\") | L440: compile(on, \"\", \"exec\") | L468: compile(generated, \"\", \"exec\")\nExec: L19: exec(_GEMMA4_PLE_CAST_HELPER, namespace)", + "evidence_hash": "a85e24d8e7c431563cbd83b70f91a3b971abde0f37083d68e70984147960cc70" + }, + { + "package": "unsloth-zoo", + "file": "tests/test_vision_collator_audio.py", + "check": "Writes to /tmp and executes (staged dropper)", + "severity": "CRITICAL", + "evidence": "L111: out = extract_audio_info(msgs({\"type\": \"audio\", key: \"/tmp/a.wav\"})) sha256:022f81dd21acfc6a35a058de96132834c218404a9e37b3d09a7768a8c8f6c728", + "evidence_hash": "2d1e75446af120d9133a42aa8af426a839d3434d9dc109cc1d6c1b22ca1ddb75" } ] } diff --git a/scripts/stamp_studio_release.py b/scripts/stamp_studio_release.py index 7dab35ea8a..739f6d1063 100644 --- a/scripts/stamp_studio_release.py +++ b/scripts/stamp_studio_release.py @@ -2,7 +2,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 -"""Stamp and verify display-only Studio release metadata for builds.""" +"""Stamp and verify display-only Unsloth release metadata for builds.""" from __future__ import annotations @@ -50,7 +50,7 @@ MAX_VERSION_LENGTH = 64 PLACEHOLDER = """# SPDX-License-Identifier: AGPL-3.0-only # Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0 -\"\"\"Build-stamped Studio release metadata. +\"\"\"Build-stamped Unsloth release metadata. Release builds may rewrite this module in the build workspace before creating Python artifacts. Keep the committed value neutral so source checkouts do not @@ -145,7 +145,7 @@ def build_info_source(version: str | None) -> str: return f'''# SPDX-License-Identifier: AGPL-3.0-only # Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0 -"""Build-stamped Studio release metadata.""" +"""Build-stamped Unsloth release metadata.""" STUDIO_RELEASE_VERSION = {literal} ''' @@ -168,7 +168,7 @@ def stamp(require_release: bool) -> int: version, source = resolve_version() if version is not None and not is_valid_version(version): print( - f"Invalid Studio release version from {source}: {version!r}", + f"Invalid Unsloth release version from {source}: {version!r}", file = sys.stderr, ) return 2 @@ -196,9 +196,9 @@ def stamp(require_release: bool) -> int: if version is None: if require_release: print( - "No Studio release version available. Set " + "No Unsloth release version available. Set " "UNSLOTH_STUDIO_RELEASE_VERSION, build from a GitHub tag, " - "or run from an exact local Studio release tag.", + "or run from an exact local Unsloth release tag.", file = sys.stderr, ) return 2 @@ -207,7 +207,7 @@ def stamp(require_release: bool) -> int: return 0 _atomic_write_text(BUILD_INFO_PATH, build_info_source(version), encoding = "utf-8") - print(f"Stamping Studio release version {version} from {source}", file = sys.stderr) + print(f"Stamping Unsloth release version {version} from {source}", file = sys.stderr) print(version) return 0 @@ -233,7 +233,7 @@ def _read_sdist_member(path: Path) -> str | None: def verify_dist(expected: str, dist_dir: Path) -> int: if not is_valid_version(expected): - print(f"Invalid expected Studio release version: {expected!r}", file = sys.stderr) + print(f"Invalid expected Unsloth release version: {expected!r}", file = sys.stderr) return 2 artifacts = list(dist_dir.glob("*.whl")) + list(dist_dir.glob("*.tar.gz")) @@ -251,14 +251,14 @@ def verify_dist(expected: str, dist_dir: Path) -> int: if content is None: failures.append(f"{artifact.name}: missing {BUILD_INFO_SUFFIX}") elif expected_line not in content: - failures.append(f"{artifact.name}: Studio release version mismatch") + failures.append(f"{artifact.name}: Unsloth release version mismatch") if failures: for failure in failures: print(failure, file = sys.stderr) return 2 - print(f"Verified Studio release version {expected} in {len(artifacts)} artifact(s)") + print(f"Verified Unsloth release version {expected} in {len(artifacts)} artifact(s)") return 0 diff --git a/scripts/uninstall.ps1 b/scripts/uninstall.ps1 index 88defb9ea0..9b6e6ebb86 100644 --- a/scripts/uninstall.ps1 +++ b/scripts/uninstall.ps1 @@ -83,7 +83,7 @@ function Uninstall-UnslothStudio { } } - # A path is a Studio-owned root iff one of install.ps1's sentinels exists: + # A path is an Unsloth-owned root iff one of install.ps1's sentinels exists: # \share\studio.conf, \unsloth_studio\.unsloth-studio-owned, # or \bin\unsloth.exe. function _IsStudioRoot { @@ -164,7 +164,7 @@ function Uninstall-UnslothStudio { return $p } - # Discover non-default Studio roots from env vars + studio.conf files. + # Discover non-default Unsloth roots from env vars + studio.conf files. # Mirrors install.ps1's precedence: UNSLOTH_STUDIO_HOME wins, STUDIO_HOME # is ignored when both are set, so uninstalling install A doesn't also # delete install B if the user has a stale STUDIO_HOME pointing at B. @@ -207,7 +207,7 @@ function Uninstall-UnslothStudio { # Return $true iff the PID's image path lives under one of $KnownRoots. # Prevents killing an unrelated process that happens to listen on a stale - # Studio port. + # Unsloth port. function _PidUnderKnownRoot { param([int]$Pid_, [string[]]$KnownRoots) if (-not $KnownRoots -or $KnownRoots.Count -eq 0) { return $false } @@ -223,8 +223,8 @@ function Uninstall-UnslothStudio { return $false } - # Stop a Studio backend whose port is recorded in \studio.port. - # Only kills if the listening PID's exe path is under a known Studio root. + # Stop an Unsloth backend whose port is recorded in \studio.port. + # Only kills if the listening PID's exe path is under a known Unsloth root. function _StopByPortFile { param([string]$PortFile, [string[]]$KnownRoots) if (-not (Test-Path -LiteralPath $PortFile -PathType Leaf)) { return } @@ -372,7 +372,7 @@ function Uninstall-UnslothStudio { continue } if (-not (_IsStudioRoot $r)) { - _Substep "refusing to remove non-Studio path: $r" "Yellow" + _Substep "refusing to remove non-Unsloth path: $r" "Yellow" continue } _RemovePath $r @@ -436,7 +436,7 @@ function Uninstall-UnslothStudio { $entries = $rawPath -split ';' $kept = New-Object System.Collections.ArrayList $removedAny = $false - # Only remove PATH entries that live inside a Studio root we + # Only remove PATH entries that live inside an Unsloth root we # actually own (default or env-mode). A literal substring # match on `unsloth_studio` would clobber unrelated user # virtualenvs that happen to share the name. diff --git a/scripts/uninstall.sh b/scripts/uninstall.sh index 31e851fcbb..957d2b7af2 100755 --- a/scripts/uninstall.sh +++ b/scripts/uninstall.sh @@ -12,7 +12,7 @@ set -e -# Stop a Studio server via its PID file (written by install.sh's _spawn_terminal). +# Stop an Unsloth server via its PID file (written by install.sh's _spawn_terminal). _kill_pid_file() { _pid_file="$1" [ -f "$_pid_file" ] || return 0 @@ -47,7 +47,7 @@ _pkill_studio() { command -v pkill >/dev/null 2>&1 || return 0 # Scope fallback patterns to the install roots we are removing so a - # different Studio install (different UNSLOTH_STUDIO_HOME) is not touched. + # different Unsloth install (different UNSLOTH_STUDIO_HOME) is not touched. _kill_roots="$HOME/.unsloth/studio" _roots_from_conf=$(_custom_studio_roots 2>/dev/null || true) [ -n "$_roots_from_conf" ] && _kill_roots="$_kill_roots @@ -89,7 +89,7 @@ _remove_path() { fi } -# Accept as Studio root only if Studio sentinels exist (matches install.sh's +# Accept as Unsloth root only if Unsloth sentinels exist (matches install.sh's # env-mode ownership guard at install.sh:1358-1361). A bare unsloth_studio/ # directory is NOT enough -- require the install-time owner marker so a user # directory that happens to contain a folder named "unsloth_studio" is safe. @@ -175,8 +175,8 @@ _custom_studio_roots() { _from_conf "$HOME/.local/share/unsloth/studio.conf" } -# Remove $HOME/.local/bin/unsloth only if it's a Studio-managed symlink. -# Studio's install.sh writes this as a symlink into the studio venv +# Remove $HOME/.local/bin/unsloth only if it's an Unsloth-managed symlink. +# Unsloth's install.sh writes this as a symlink into the studio venv # (install.sh: `ln -sfn "$VENV_DIR/bin/unsloth" "$_shim_path"`). A # pip-installed `unsloth` CLI is a regular file — leave it alone to avoid # wiping an unrelated install. @@ -206,7 +206,7 @@ _custom_studio_roots | while IFS= read -r _custom_root; do continue fi if ! _is_studio_root "$_custom_root"; then - echo " refusing to remove non-Studio path: $_custom_root" >&2 + echo " refusing to remove non-Unsloth path: $_custom_root" >&2 continue fi _remove_path "$_custom_root" @@ -234,7 +234,7 @@ _remove_path "$HOME/.unsloth/rocm-smoketest" # Drop ~/.unsloth only if now empty (rmdir refuses non-empty, so user content is kept). rmdir "$HOME/.unsloth" 2>/dev/null || true _remove_path "$HOME/.local/share/unsloth" -# CLI shim: only the symlink Studio created, never a pip-installed file. +# CLI shim: only the symlink Unsloth created, never a pip-installed file. _remove_cli_shim echo "Removing desktop shortcut and launcher lock..." diff --git a/studio/MCP.md b/studio/MCP.md new file mode 100644 index 0000000000..127a85a116 --- /dev/null +++ b/studio/MCP.md @@ -0,0 +1,34 @@ +# Unsloth Studio MCP server + +Unsloth can expose a local MCP server so an MCP client can inspect models and +GPU state, validate recipes, start or stop training, inspect recipe output, and +export a loaded model. + +The server is disabled by default. Enable it for a local Unsloth process with: + +```bash +UNSLOTH_STUDIO_ENABLE_MCP=1 \ +UNSLOTH_STUDIO_MCP_TOKEN='use-a-local-secret' \ +unsloth studio +``` + +The endpoint is `http://127.0.0.1:8888/mcp/` when Unsloth uses its default port +(a request to `/mcp` redirects to the canonical `/mcp/`). Use the actual Unsloth +port when it is configured differently. + +The high-impact tools are: + +- `studio_status` and `list_local_models` for discovery +- `get_training_status`, `start_training`, `stop_training`, and `list_training_runs` +- `validate_recipe`, `get_recipe_job_status`, and `get_recipe_job_dataset` +- `load_checkpoint` and `export_gguf` + +`start_training` accepts the same fields as the Unsloth `TrainingStartRequest`. +The request is validated by the existing Pydantic model before a subprocess is +started. Export paths use the existing Unsloth validation as well. + +The endpoint always requires `UNSLOTH_STUDIO_MCP_TOKEN` and checks an exact +Bearer token for both HTTP and WebSocket connections. Keep it on localhost +unless the deployment has an authenticated reverse proxy. The MCP endpoint is +intentionally opt-in because tools can consume GPU memory, write model +artifacts, and stop active work. \ No newline at end of file diff --git a/studio/Unsloth_Studio_Colab.ipynb b/studio/Unsloth_Studio_Colab.ipynb index 619395bd6d..612d739806 100644 --- a/studio/Unsloth_Studio_Colab.ipynb +++ b/studio/Unsloth_Studio_Colab.ipynb @@ -1,134 +1,145 @@ { - "cells": [ - { - "cell_type": "markdown", - "metadata": { - "id": "view-in-github", - "colab_type": "text" - }, - "source": [ - "\"Open" - ] + "cells": [ + { + "cell_type": "markdown", + "metadata": { + "id": "view-in-github", + "colab_type": "text" + }, + "source": [ + "\"Open" + ] + }, + { + "cell_type": "markdown", + "metadata": { + "id": "6b87de59" + }, + "source": [ + "To run this, press \"*Runtime*\" and press \"*Run all*\" on a **free** Tesla T4 Google Colab instance!\n", + "

\n", + "\n", + "\n", + " Join Discord if you need help + ⭐ Star us on Github ⭐\n", + "
\n", + "\n", + "To install Unsloth Studio on your local device, follow [our guide](https://unsloth.ai/docs/new/unsloth-studio/install). Unsloth Studio is licensed [AGPL-3.0](https://github.com/unslothai/unsloth/blob/main/studio/LICENSE.AGPL-3.0).\n", + "\n", + "### Unsloth Studio\n", + "\n", + "Train and run open models with [**Unsloth Studio**](https://unsloth.ai/docs/new/unsloth-studio/start). NEW! Installation should now only take 2 mins!\n", + "\n", + "\n", + "We are actively working on making Unsloth Studio install on Colab T4 GPUs faster.\n", + "\n", + "[Features](https://unsloth.ai/docs/new/unsloth-studio#features) • [Quickstart](https://unsloth.ai/docs/new/unsloth-studio/start) • [Data Recipes](https://unsloth.ai/docs/new/unsloth-studio/data-recipe) • [Unsloth Chat](https://unsloth.ai/docs/new/unsloth-studio/chat) • [Export](https://unsloth.ai/docs/new/unsloth-studio/export)" + ], + "id": "6b87de59" + }, + { + "cell_type": "markdown", + "metadata": { + "id": "e4206349" + }, + "source": [ + "

" + ], + "id": "e4206349" + }, + { + "cell_type": "markdown", + "metadata": { + "id": "27da2957" + }, + "source": [ + "### Setup: Clone repo and run setup" + ], + "id": "27da2957" + }, + { + "cell_type": "code", + "metadata": { + "id": "27e68f91" + }, + "source": "!git clone --depth 1 --branch main https://github.com/unslothai/unsloth.git\n%cd /content/unsloth\n!chmod +x studio/setup.sh && ./studio/setup.sh --local", + "execution_count": null, + "outputs": [], + "id": "27e68f91" + }, + { + "cell_type": "markdown", + "metadata": { + "id": "3e1771a9" + }, + "source": [ + "### Start Unsloth Studio" + ], + "id": "3e1771a9" + }, + { + "cell_type": "code", + "metadata": { + "id": "277e431e" + }, + "source": [ + "import sys\n", + "sys.path.insert(0, \"/content/unsloth/studio/backend\")\n", + "from colab import start\n", + "\n", + "# On Colab, start() auto-opens a Cloudflare link and prints admin login credentials.\n", + "# Use the Cloudflare link above the ready card to open Studio (in-cell iframes often stay blank).\n", + "start()\n", + "\n", + "# To skip the Cloudflare tunnel and try the in-notebook proxy iframe only:\n", + "# start(cloudflare=False)" + ], + "execution_count": null, + "outputs": [], + "id": "277e431e" + }, + { + "cell_type": "markdown", + "metadata": { + "id": "f2b0c6a1" + }, + "source": [ + "And we're done! If you have any questions on Unsloth, we have a [Discord](https://discord.gg/unsloth) channel! If you find any bugs or want to keep updated with the latest LLM stuff, or need help, join projects etc, feel free to join our Discord!\n", + "\n", + "Some other resources:\n", + "1. Looking to use Unsloth locally? Read our [Installation Guide](https://unsloth.ai/docs/get-started/install) for details on installing Unsloth on Windows, Docker, AMD, Intel GPUs.\n", + "2. Learn how to do Reinforcement Learning with our [RL Guide and notebooks](https://unsloth.ai/docs/get-started/reinforcement-learning-rl-guide).\n", + "3. Read our guides and notebooks for [Text-to-speech (TTS)](https://unsloth.ai/docs/basics/text-to-speech-tts-fine-tuning) and [vision](https://unsloth.ai/docs/basics/vision-fine-tuning) model support.\n", + "4. Explore our [LLM Tutorials Directory](https://unsloth.ai/docs/models/tutorials-how-to-fine-tune-and-run-llms) to find dedicated guides for each model.\n", + "5. Need help with Inference? Read our [Inference & Deployment page](https://unsloth.ai/docs/basics/inference-and-deployment) for details on using vLLM, llama.cpp, Ollama etc.\n", + "\n", + "
\n", + " \n", + " \n", + " \n", + "\n", + " Join Discord if you need help + ⭐️ Star us on Github ⭐️\n", + "\n", + " This notebook is licensed AGPL-3.0\n", + "
" + ], + "id": "f2b0c6a1" + } + ], + "metadata": { + "accelerator": "GPU", + "colab": { + "gpuType": "T4", + "provenance": [], + "include_colab_link": true + }, + "kernelspec": { + "display_name": "Python 3", + "name": "python3" + }, + "language_info": { + "name": "python" + } }, - { - "cell_type": "markdown", - "id": "6b87de59", - "metadata": { - "id": "6b87de59" - }, - "source": [ - "To run this, press \"*Runtime*\" and press \"*Run all*\" on a **free** Tesla T4 Google Colab instance!\n", - "
\n", - "\n", - "\n", - " Join Discord if you need help + ⭐ Star us on Github ⭐\n", - "
\n", - "\n", - "To install Unsloth Studio on your local device, follow [our guide](https://unsloth.ai/docs/new/unsloth-studio/install). Unsloth Studio is licensed [AGPL-3.0](https://github.com/unslothai/unsloth/blob/main/studio/LICENSE.AGPL-3.0).\n", - "\n", - "### Unsloth Studio\n", - "\n", - "Train and run open models with [**Unsloth Studio**](https://unsloth.ai/docs/new/unsloth-studio/start). NEW! Installation should now only take 2 mins!\n", - "\n", - "\n", - "We are actively working on making Unsloth Studio install on Colab T4 GPUs faster.\n", - "\n", - "[Features](https://unsloth.ai/docs/new/unsloth-studio#features) • [Quickstart](https://unsloth.ai/docs/new/unsloth-studio/start) • [Data Recipes](https://unsloth.ai/docs/new/unsloth-studio/data-recipe) • [Studio Chat](https://unsloth.ai/docs/new/unsloth-studio/chat) • [Export](https://unsloth.ai/docs/new/unsloth-studio/export)" - ] - }, - { - "cell_type": "markdown", - "id": "e4206349", - "metadata": { - "id": "e4206349" - }, - "source": [ - "

" - ] - }, - { - "cell_type": "markdown", - "id": "27da2957", - "metadata": { - "id": "27da2957" - }, - "source": [ - "### Setup: Clone repo and run setup" - ] - }, - { - "cell_type": "code", - "execution_count": null, - "id": "27e68f91", - "metadata": { - "id": "27e68f91" - }, - "outputs": [], - "source": "!git clone --depth 1 --branch main https://github.com/unslothai/unsloth.git\n%cd /content/unsloth\n!chmod +x studio/setup.sh && ./studio/setup.sh --local" - }, - { - "cell_type": "markdown", - "id": "3e1771a9", - "metadata": { - "id": "3e1771a9" - }, - "source": [ - "### Start Unsloth Studio" - ] - }, - { - "cell_type": "code", - "execution_count": null, - "id": "277e431e", - "metadata": { - "id": "277e431e" - }, - "outputs": [], - "source": "import sys\nsys.path.insert(0, \"/content/unsloth/studio/backend\")\nfrom colab import start\n\n# Default: in-tab iframe only. start() blocks to keep the kernel alive.\nstart()\n\n# For a shareable Cloudflare link, replace start() above with:\n# start(cloudflare=True)" - }, - { - "cell_type": "markdown", - "id": "f2b0c6a1", - "metadata": { - "id": "f2b0c6a1" - }, - "source": [ - "And we're done! If you have any questions on Unsloth, we have a [Discord](https://discord.gg/unsloth) channel! If you find any bugs or want to keep updated with the latest LLM stuff, or need help, join projects etc, feel free to join our Discord!\n", - "\n", - "Some other resources:\n", - "1. Looking to use Unsloth locally? Read our [Installation Guide](https://unsloth.ai/docs/get-started/install) for details on installing Unsloth on Windows, Docker, AMD, Intel GPUs.\n", - "2. Learn how to do Reinforcement Learning with our [RL Guide and notebooks](https://unsloth.ai/docs/get-started/reinforcement-learning-rl-guide).\n", - "3. Read our guides and notebooks for [Text-to-speech (TTS)](https://unsloth.ai/docs/basics/text-to-speech-tts-fine-tuning) and [vision](https://unsloth.ai/docs/basics/vision-fine-tuning) model support.\n", - "4. Explore our [LLM Tutorials Directory](https://unsloth.ai/docs/models/tutorials-how-to-fine-tune-and-run-llms) to find dedicated guides for each model.\n", - "5. Need help with Inference? Read our [Inference & Deployment page](https://unsloth.ai/docs/basics/inference-and-deployment) for details on using vLLM, llama.cpp, Ollama etc.\n", - "\n", - "
\n", - " \n", - " \n", - " \n", - "\n", - " Join Discord if you need help + ⭐️ Star us on Github ⭐️\n", - "\n", - " This notebook is licensed AGPL-3.0\n", - "
" - ] - } - ], - "metadata": { - "accelerator": "GPU", - "colab": { - "gpuType": "T4", - "provenance": [], - "include_colab_link": true - }, - "kernelspec": { - "display_name": "Python 3", - "name": "python3" - }, - "language_info": { - "name": "python" - } - }, - "nbformat": 4, - "nbformat_minor": 5 + "nbformat": 4, + "nbformat_minor": 5 } \ No newline at end of file diff --git a/studio/backend/assets/chat_templates/gemma-4-edge.jinja b/studio/backend/assets/chat_templates/gemma-4-edge.jinja index 0266127233..74fa73ddd3 100644 --- a/studio/backend/assets/chat_templates/gemma-4-edge.jinja +++ b/studio/backend/assets/chat_templates/gemma-4-edge.jinja @@ -3,7 +3,7 @@ Source: google/gemma-4-31B-it HF discussion/PR #118 (adds the preserve_thinking flag plus null-rendering, string-arguments validation, balanced turn tags, empty messages handling, and OpenAI image_url/input_audio aliases). - Studio-local changes vs PR #118: + Unsloth-local changes vs PR #118: 1. preserve_thinking defaults to false (see SETUP block below). 2. The empty "<|channel>thought\n" block on enable_thinking=false is NOT emitted. Google ships a distinct template for E2B/E4B (google/gemma-4-E2B-it, diff --git a/studio/backend/assets/chat_templates/gemma-4.jinja b/studio/backend/assets/chat_templates/gemma-4.jinja index 65ab39df57..cc5f98065f 100644 --- a/studio/backend/assets/chat_templates/gemma-4.jinja +++ b/studio/backend/assets/chat_templates/gemma-4.jinja @@ -3,7 +3,7 @@ Source: google/gemma-4-31B-it HF discussion/PR #118 (adds the preserve_thinking flag plus null-rendering, string-arguments validation, balanced turn tags, empty messages handling, and OpenAI image_url/input_audio aliases). - Studio-local change: preserve_thinking defaults to false (see SETUP block below). + Unsloth-local change: preserve_thinking defaults to false (see SETUP block below). Applied to unsloth/gemma-4-*-GGUF models so the embedded GGUF template does not need re-downloading. Keep in sync with upstream if PR #118 changes. -#} diff --git a/studio/backend/assets/configs/full_finetune.yaml b/studio/backend/assets/configs/full_finetune.yaml index e398515f61..98c45dd851 100644 --- a/studio/backend/assets/configs/full_finetune.yaml +++ b/studio/backend/assets/configs/full_finetune.yaml @@ -30,6 +30,7 @@ lora: vision_all_linear: false use_rslora: false use_loftq: false + use_dora: false finetune_vision_layers: true finetune_language_layers: true finetune_attention_modules: true diff --git a/studio/backend/assets/configs/lora_text.yaml b/studio/backend/assets/configs/lora_text.yaml index 9cb6b8c700..6c6a4d8839 100644 --- a/studio/backend/assets/configs/lora_text.yaml +++ b/studio/backend/assets/configs/lora_text.yaml @@ -30,6 +30,7 @@ lora: vision_all_linear: false use_rslora: false use_loftq: false + use_dora: false finetune_vision_layers: true finetune_language_layers: true finetune_attention_modules: true diff --git a/studio/backend/assets/configs/model_defaults/default.yaml b/studio/backend/assets/configs/model_defaults/default.yaml index 841e8ba166..e569031a31 100644 --- a/studio/backend/assets/configs/model_defaults/default.yaml +++ b/studio/backend/assets/configs/model_defaults/default.yaml @@ -33,6 +33,7 @@ lora: - "down_proj" use_rslora: false use_loftq: false + use_dora: false finetune_vision_layers: true finetune_language_layers: true finetune_attention_modules: true diff --git a/studio/backend/assets/configs/model_defaults/embedding/unsloth_Qwen3-Embedding-0.6B.yaml b/studio/backend/assets/configs/model_defaults/embedding/unsloth_Qwen3-Embedding-0.6B.yaml index f7b49c75b7..7ac1c83e04 100644 --- a/studio/backend/assets/configs/model_defaults/embedding/unsloth_Qwen3-Embedding-0.6B.yaml +++ b/studio/backend/assets/configs/model_defaults/embedding/unsloth_Qwen3-Embedding-0.6B.yaml @@ -34,6 +34,7 @@ lora: - "down_proj" use_rslora: false use_loftq: false + use_dora: false logging: enable_wandb: false diff --git a/studio/backend/assets/configs/model_defaults/embedding/unsloth_all-MiniLM-L6-v2.yaml b/studio/backend/assets/configs/model_defaults/embedding/unsloth_all-MiniLM-L6-v2.yaml index be7da0f624..4cab9e9f96 100644 --- a/studio/backend/assets/configs/model_defaults/embedding/unsloth_all-MiniLM-L6-v2.yaml +++ b/studio/backend/assets/configs/model_defaults/embedding/unsloth_all-MiniLM-L6-v2.yaml @@ -30,6 +30,7 @@ lora: - "query" use_rslora: false use_loftq: false + use_dora: false logging: enable_wandb: false diff --git a/studio/backend/assets/configs/model_defaults/embedding/unsloth_bge-m3.yaml b/studio/backend/assets/configs/model_defaults/embedding/unsloth_bge-m3.yaml index d9e49bc0d5..c1f1c2a344 100644 --- a/studio/backend/assets/configs/model_defaults/embedding/unsloth_bge-m3.yaml +++ b/studio/backend/assets/configs/model_defaults/embedding/unsloth_bge-m3.yaml @@ -30,6 +30,7 @@ lora: - "value" use_rslora: false use_loftq: false + use_dora: false logging: enable_wandb: false diff --git a/studio/backend/assets/configs/model_defaults/embedding/unsloth_embeddinggemma-300m.yaml b/studio/backend/assets/configs/model_defaults/embedding/unsloth_embeddinggemma-300m.yaml index c3422d399f..7828feae81 100644 --- a/studio/backend/assets/configs/model_defaults/embedding/unsloth_embeddinggemma-300m.yaml +++ b/studio/backend/assets/configs/model_defaults/embedding/unsloth_embeddinggemma-300m.yaml @@ -33,6 +33,7 @@ lora: - "down_proj" use_rslora: false use_loftq: false + use_dora: false logging: enable_wandb: false diff --git a/studio/backend/assets/configs/model_defaults/embedding/unsloth_gte-modernbert-base.yaml b/studio/backend/assets/configs/model_defaults/embedding/unsloth_gte-modernbert-base.yaml index 529a56a527..5a4028f15b 100644 --- a/studio/backend/assets/configs/model_defaults/embedding/unsloth_gte-modernbert-base.yaml +++ b/studio/backend/assets/configs/model_defaults/embedding/unsloth_gte-modernbert-base.yaml @@ -29,6 +29,7 @@ lora: - "Wqkv" use_rslora: false use_loftq: false + use_dora: false logging: enable_wandb: false diff --git a/studio/backend/assets/configs/model_defaults/ernie/unsloth_ERNIE-4.5-21B-A3B-PT.yaml b/studio/backend/assets/configs/model_defaults/ernie/unsloth_ERNIE-4.5-21B-A3B-PT.yaml index 734115ec41..7645d11c98 100644 --- a/studio/backend/assets/configs/model_defaults/ernie/unsloth_ERNIE-4.5-21B-A3B-PT.yaml +++ b/studio/backend/assets/configs/model_defaults/ernie/unsloth_ERNIE-4.5-21B-A3B-PT.yaml @@ -34,6 +34,7 @@ lora: - "down_proj" use_rslora: false use_loftq: false + use_dora: false logging: enable_wandb: false diff --git a/studio/backend/assets/configs/model_defaults/ernie/unsloth_ERNIE-4.5-VL-28B-A3B-PT.yaml b/studio/backend/assets/configs/model_defaults/ernie/unsloth_ERNIE-4.5-VL-28B-A3B-PT.yaml index 1032449e8c..b746235f1f 100644 --- a/studio/backend/assets/configs/model_defaults/ernie/unsloth_ERNIE-4.5-VL-28B-A3B-PT.yaml +++ b/studio/backend/assets/configs/model_defaults/ernie/unsloth_ERNIE-4.5-VL-28B-A3B-PT.yaml @@ -35,6 +35,7 @@ lora: - "down_proj" use_rslora: false use_loftq: false + use_dora: false finetune_vision_layers: true finetune_language_layers: true finetune_attention_modules: true diff --git a/studio/backend/assets/configs/model_defaults/falcon/tiiuae_Falcon-H1-0.5B-Instruct.yaml b/studio/backend/assets/configs/model_defaults/falcon/tiiuae_Falcon-H1-0.5B-Instruct.yaml index c8e5f35841..4964fea276 100644 --- a/studio/backend/assets/configs/model_defaults/falcon/tiiuae_Falcon-H1-0.5B-Instruct.yaml +++ b/studio/backend/assets/configs/model_defaults/falcon/tiiuae_Falcon-H1-0.5B-Instruct.yaml @@ -34,6 +34,7 @@ lora: - "down_proj" use_rslora: false use_loftq: false + use_dora: false logging: enable_wandb: false diff --git a/studio/backend/assets/configs/model_defaults/gemma/unsloth_codegemma-7b-bnb-4bit.yaml b/studio/backend/assets/configs/model_defaults/gemma/unsloth_codegemma-7b-bnb-4bit.yaml index 251409c29d..e5f3344356 100644 --- a/studio/backend/assets/configs/model_defaults/gemma/unsloth_codegemma-7b-bnb-4bit.yaml +++ b/studio/backend/assets/configs/model_defaults/gemma/unsloth_codegemma-7b-bnb-4bit.yaml @@ -35,6 +35,7 @@ lora: - "down_proj" use_rslora: false use_loftq: false + use_dora: false logging: enable_wandb: false diff --git a/studio/backend/assets/configs/model_defaults/gemma/unsloth_functiongemma-270m-it.yaml b/studio/backend/assets/configs/model_defaults/gemma/unsloth_functiongemma-270m-it.yaml index 89b1d7f938..71c61f383a 100644 --- a/studio/backend/assets/configs/model_defaults/gemma/unsloth_functiongemma-270m-it.yaml +++ b/studio/backend/assets/configs/model_defaults/gemma/unsloth_functiongemma-270m-it.yaml @@ -35,6 +35,7 @@ lora: - "down_proj" use_rslora: false use_loftq: false + use_dora: false logging: enable_wandb: false diff --git a/studio/backend/assets/configs/model_defaults/gemma/unsloth_gemma-2-27b-bnb-4bit.yaml b/studio/backend/assets/configs/model_defaults/gemma/unsloth_gemma-2-27b-bnb-4bit.yaml index e3292b5972..3fe29cd800 100644 --- a/studio/backend/assets/configs/model_defaults/gemma/unsloth_gemma-2-27b-bnb-4bit.yaml +++ b/studio/backend/assets/configs/model_defaults/gemma/unsloth_gemma-2-27b-bnb-4bit.yaml @@ -33,6 +33,7 @@ lora: - "down_proj" use_rslora: false use_loftq: false + use_dora: false logging: enable_wandb: false diff --git a/studio/backend/assets/configs/model_defaults/gemma/unsloth_gemma-2-2b.yaml b/studio/backend/assets/configs/model_defaults/gemma/unsloth_gemma-2-2b.yaml index 98fe497912..cd4e3e0c4d 100644 --- a/studio/backend/assets/configs/model_defaults/gemma/unsloth_gemma-2-2b.yaml +++ b/studio/backend/assets/configs/model_defaults/gemma/unsloth_gemma-2-2b.yaml @@ -34,6 +34,7 @@ lora: - "down_proj" use_rslora: false use_loftq: false + use_dora: false logging: enable_wandb: false diff --git a/studio/backend/assets/configs/model_defaults/gemma/unsloth_gemma-3-270m-it.yaml b/studio/backend/assets/configs/model_defaults/gemma/unsloth_gemma-3-270m-it.yaml index bda5471643..97aa10e861 100644 --- a/studio/backend/assets/configs/model_defaults/gemma/unsloth_gemma-3-270m-it.yaml +++ b/studio/backend/assets/configs/model_defaults/gemma/unsloth_gemma-3-270m-it.yaml @@ -35,6 +35,7 @@ lora: - "down_proj" use_rslora: false use_loftq: false + use_dora: false logging: enable_wandb: false diff --git a/studio/backend/assets/configs/model_defaults/gemma/unsloth_gemma-3-27b-it.yaml b/studio/backend/assets/configs/model_defaults/gemma/unsloth_gemma-3-27b-it.yaml index 18392568bd..a1b1640fa2 100644 --- a/studio/backend/assets/configs/model_defaults/gemma/unsloth_gemma-3-27b-it.yaml +++ b/studio/backend/assets/configs/model_defaults/gemma/unsloth_gemma-3-27b-it.yaml @@ -29,6 +29,7 @@ lora: - "all-linear" use_rslora: false use_loftq: false + use_dora: false finetune_vision_layers: true finetune_language_layers: true finetune_attention_modules: true diff --git a/studio/backend/assets/configs/model_defaults/gemma/unsloth_gemma-3-4b-it.yaml b/studio/backend/assets/configs/model_defaults/gemma/unsloth_gemma-3-4b-it.yaml index 434ac41b46..dbf60f04d4 100644 --- a/studio/backend/assets/configs/model_defaults/gemma/unsloth_gemma-3-4b-it.yaml +++ b/studio/backend/assets/configs/model_defaults/gemma/unsloth_gemma-3-4b-it.yaml @@ -29,6 +29,7 @@ lora: - "all-linear" use_rslora: false use_loftq: false + use_dora: false finetune_vision_layers: true finetune_language_layers: true finetune_attention_modules: true diff --git a/studio/backend/assets/configs/model_defaults/gemma/unsloth_gemma-3-4b-pt.yaml b/studio/backend/assets/configs/model_defaults/gemma/unsloth_gemma-3-4b-pt.yaml index 5f0a7b26ce..54c7dd6cd4 100644 --- a/studio/backend/assets/configs/model_defaults/gemma/unsloth_gemma-3-4b-pt.yaml +++ b/studio/backend/assets/configs/model_defaults/gemma/unsloth_gemma-3-4b-pt.yaml @@ -29,6 +29,7 @@ lora: - "all-linear" use_rslora: false use_loftq: false + use_dora: false finetune_vision_layers: true finetune_language_layers: true finetune_attention_modules: true diff --git a/studio/backend/assets/configs/model_defaults/gemma/unsloth_gemma-3n-E4B-it.yaml b/studio/backend/assets/configs/model_defaults/gemma/unsloth_gemma-3n-E4B-it.yaml index dd5ae51ab0..119440a585 100644 --- a/studio/backend/assets/configs/model_defaults/gemma/unsloth_gemma-3n-E4B-it.yaml +++ b/studio/backend/assets/configs/model_defaults/gemma/unsloth_gemma-3n-E4B-it.yaml @@ -29,6 +29,7 @@ lora: - "all-linear" use_rslora: false use_loftq: false + use_dora: false finetune_vision_layers: true finetune_language_layers: true finetune_attention_modules: true diff --git a/studio/backend/assets/configs/model_defaults/gemma/unsloth_gemma-3n-E4B.yaml b/studio/backend/assets/configs/model_defaults/gemma/unsloth_gemma-3n-E4B.yaml index e53e163a04..d08e5e9547 100644 --- a/studio/backend/assets/configs/model_defaults/gemma/unsloth_gemma-3n-E4B.yaml +++ b/studio/backend/assets/configs/model_defaults/gemma/unsloth_gemma-3n-E4B.yaml @@ -29,6 +29,7 @@ lora: - "all-linear" use_rslora: false use_loftq: false + use_dora: false finetune_vision_layers: true finetune_language_layers: true finetune_attention_modules: true diff --git a/studio/backend/assets/configs/model_defaults/gemma/unsloth_gemma-4-26B-A4B-it.yaml b/studio/backend/assets/configs/model_defaults/gemma/unsloth_gemma-4-26B-A4B-it.yaml index ebe344e382..a266d7a39b 100644 --- a/studio/backend/assets/configs/model_defaults/gemma/unsloth_gemma-4-26B-A4B-it.yaml +++ b/studio/backend/assets/configs/model_defaults/gemma/unsloth_gemma-4-26B-A4B-it.yaml @@ -26,6 +26,7 @@ lora: - "all-linear" use_rslora: false use_loftq: false + use_dora: false finetune_vision_layers: true finetune_language_layers: true finetune_attention_modules: true diff --git a/studio/backend/assets/configs/model_defaults/gemma/unsloth_gemma-4-26B-A4B.yaml b/studio/backend/assets/configs/model_defaults/gemma/unsloth_gemma-4-26B-A4B.yaml index fb89a07133..970cac3259 100644 --- a/studio/backend/assets/configs/model_defaults/gemma/unsloth_gemma-4-26B-A4B.yaml +++ b/studio/backend/assets/configs/model_defaults/gemma/unsloth_gemma-4-26B-A4B.yaml @@ -26,6 +26,7 @@ lora: - "all-linear" use_rslora: false use_loftq: false + use_dora: false finetune_vision_layers: true finetune_language_layers: true finetune_attention_modules: true diff --git a/studio/backend/assets/configs/model_defaults/gemma/unsloth_gemma-4-31B-it.yaml b/studio/backend/assets/configs/model_defaults/gemma/unsloth_gemma-4-31B-it.yaml index 4a089992ac..5bba4ccdc0 100644 --- a/studio/backend/assets/configs/model_defaults/gemma/unsloth_gemma-4-31B-it.yaml +++ b/studio/backend/assets/configs/model_defaults/gemma/unsloth_gemma-4-31B-it.yaml @@ -26,6 +26,7 @@ lora: - "all-linear" use_rslora: false use_loftq: false + use_dora: false finetune_vision_layers: true finetune_language_layers: true finetune_attention_modules: true diff --git a/studio/backend/assets/configs/model_defaults/gemma/unsloth_gemma-4-31B.yaml b/studio/backend/assets/configs/model_defaults/gemma/unsloth_gemma-4-31B.yaml index ae7524b7c6..ac5c6eca22 100644 --- a/studio/backend/assets/configs/model_defaults/gemma/unsloth_gemma-4-31B.yaml +++ b/studio/backend/assets/configs/model_defaults/gemma/unsloth_gemma-4-31B.yaml @@ -26,6 +26,7 @@ lora: - "all-linear" use_rslora: false use_loftq: false + use_dora: false finetune_vision_layers: true finetune_language_layers: true finetune_attention_modules: true diff --git a/studio/backend/assets/configs/model_defaults/gemma/unsloth_gemma-4-E2B-it.yaml b/studio/backend/assets/configs/model_defaults/gemma/unsloth_gemma-4-E2B-it.yaml index 10c1abd8a5..68c2d35644 100644 --- a/studio/backend/assets/configs/model_defaults/gemma/unsloth_gemma-4-E2B-it.yaml +++ b/studio/backend/assets/configs/model_defaults/gemma/unsloth_gemma-4-E2B-it.yaml @@ -26,6 +26,7 @@ lora: - "all-linear" use_rslora: false use_loftq: false + use_dora: false finetune_vision_layers: true finetune_language_layers: true finetune_attention_modules: true diff --git a/studio/backend/assets/configs/model_defaults/gemma/unsloth_gemma-4-E2B.yaml b/studio/backend/assets/configs/model_defaults/gemma/unsloth_gemma-4-E2B.yaml index fb5c1d9dea..175f9c0f17 100644 --- a/studio/backend/assets/configs/model_defaults/gemma/unsloth_gemma-4-E2B.yaml +++ b/studio/backend/assets/configs/model_defaults/gemma/unsloth_gemma-4-E2B.yaml @@ -26,6 +26,7 @@ lora: - "all-linear" use_rslora: false use_loftq: false + use_dora: false finetune_vision_layers: true finetune_language_layers: true finetune_attention_modules: true diff --git a/studio/backend/assets/configs/model_defaults/gemma/unsloth_gemma-4-E4B-it.yaml b/studio/backend/assets/configs/model_defaults/gemma/unsloth_gemma-4-E4B-it.yaml index 189e5dc6b2..4f3834e7c0 100644 --- a/studio/backend/assets/configs/model_defaults/gemma/unsloth_gemma-4-E4B-it.yaml +++ b/studio/backend/assets/configs/model_defaults/gemma/unsloth_gemma-4-E4B-it.yaml @@ -26,6 +26,7 @@ lora: - "all-linear" use_rslora: false use_loftq: false + use_dora: false finetune_vision_layers: true finetune_language_layers: true finetune_attention_modules: true diff --git a/studio/backend/assets/configs/model_defaults/gemma/unsloth_gemma-4-E4B.yaml b/studio/backend/assets/configs/model_defaults/gemma/unsloth_gemma-4-E4B.yaml index aa51440b6a..d6d97f7e44 100644 --- a/studio/backend/assets/configs/model_defaults/gemma/unsloth_gemma-4-E4B.yaml +++ b/studio/backend/assets/configs/model_defaults/gemma/unsloth_gemma-4-E4B.yaml @@ -26,6 +26,7 @@ lora: - "all-linear" use_rslora: false use_loftq: false + use_dora: false finetune_vision_layers: true finetune_language_layers: true finetune_attention_modules: true diff --git a/studio/backend/assets/configs/model_defaults/gpt-oss/unsloth_gpt-oss-120b.yaml b/studio/backend/assets/configs/model_defaults/gpt-oss/unsloth_gpt-oss-120b.yaml index e2d67bcb0b..4f1f54a4e6 100644 --- a/studio/backend/assets/configs/model_defaults/gpt-oss/unsloth_gpt-oss-120b.yaml +++ b/studio/backend/assets/configs/model_defaults/gpt-oss/unsloth_gpt-oss-120b.yaml @@ -35,6 +35,7 @@ lora: - "down_proj" use_rslora: false use_loftq: false + use_dora: false logging: enable_wandb: false diff --git a/studio/backend/assets/configs/model_defaults/gpt-oss/unsloth_gpt-oss-20b.yaml b/studio/backend/assets/configs/model_defaults/gpt-oss/unsloth_gpt-oss-20b.yaml index aa436117a1..127700b53b 100644 --- a/studio/backend/assets/configs/model_defaults/gpt-oss/unsloth_gpt-oss-20b.yaml +++ b/studio/backend/assets/configs/model_defaults/gpt-oss/unsloth_gpt-oss-20b.yaml @@ -35,6 +35,7 @@ lora: - "down_proj" use_rslora: false use_loftq: false + use_dora: false logging: enable_wandb: false diff --git a/studio/backend/assets/configs/model_defaults/granite/unsloth_granite-4.0-350m-unsloth-bnb-4bit.yaml b/studio/backend/assets/configs/model_defaults/granite/unsloth_granite-4.0-350m-unsloth-bnb-4bit.yaml index 3f2cb84a94..2412b3accf 100644 --- a/studio/backend/assets/configs/model_defaults/granite/unsloth_granite-4.0-350m-unsloth-bnb-4bit.yaml +++ b/studio/backend/assets/configs/model_defaults/granite/unsloth_granite-4.0-350m-unsloth-bnb-4bit.yaml @@ -37,6 +37,7 @@ lora: - "shared_mlp.output_linear" use_rslora: false use_loftq: false + use_dora: false logging: enable_wandb: false diff --git a/studio/backend/assets/configs/model_defaults/granite/unsloth_granite-4.0-h-micro.yaml b/studio/backend/assets/configs/model_defaults/granite/unsloth_granite-4.0-h-micro.yaml index ab756fe764..81b59c4323 100644 --- a/studio/backend/assets/configs/model_defaults/granite/unsloth_granite-4.0-h-micro.yaml +++ b/studio/backend/assets/configs/model_defaults/granite/unsloth_granite-4.0-h-micro.yaml @@ -37,6 +37,7 @@ lora: - "shared_mlp.output_linear" use_rslora: false use_loftq: false + use_dora: false logging: enable_wandb: false diff --git a/studio/backend/assets/configs/model_defaults/llama/unsloth_Llama-3.2-11B-Vision-Instruct.yaml b/studio/backend/assets/configs/model_defaults/llama/unsloth_Llama-3.2-11B-Vision-Instruct.yaml index 1a7a91e56f..6110d84a6c 100644 --- a/studio/backend/assets/configs/model_defaults/llama/unsloth_Llama-3.2-11B-Vision-Instruct.yaml +++ b/studio/backend/assets/configs/model_defaults/llama/unsloth_Llama-3.2-11B-Vision-Instruct.yaml @@ -29,6 +29,7 @@ lora: - "all-linear" use_rslora: false use_loftq: false + use_dora: false finetune_vision_layers: true finetune_language_layers: true finetune_attention_modules: true diff --git a/studio/backend/assets/configs/model_defaults/llama/unsloth_Llama-3.2-1B-Instruct.yaml b/studio/backend/assets/configs/model_defaults/llama/unsloth_Llama-3.2-1B-Instruct.yaml index 7c7bb8dc3e..3c7fc7f238 100644 --- a/studio/backend/assets/configs/model_defaults/llama/unsloth_Llama-3.2-1B-Instruct.yaml +++ b/studio/backend/assets/configs/model_defaults/llama/unsloth_Llama-3.2-1B-Instruct.yaml @@ -34,6 +34,7 @@ lora: - "down_proj" use_rslora: false use_loftq: false + use_dora: false logging: enable_wandb: false diff --git a/studio/backend/assets/configs/model_defaults/llama/unsloth_Llama-3.2-3B-Instruct.yaml b/studio/backend/assets/configs/model_defaults/llama/unsloth_Llama-3.2-3B-Instruct.yaml index f73b0c09b6..2b0977e435 100644 --- a/studio/backend/assets/configs/model_defaults/llama/unsloth_Llama-3.2-3B-Instruct.yaml +++ b/studio/backend/assets/configs/model_defaults/llama/unsloth_Llama-3.2-3B-Instruct.yaml @@ -35,6 +35,7 @@ lora: - "down_proj" use_rslora: false use_loftq: false + use_dora: false logging: enable_wandb: false diff --git a/studio/backend/assets/configs/model_defaults/llama/unsloth_Llama-3.3-70B-Instruct.yaml b/studio/backend/assets/configs/model_defaults/llama/unsloth_Llama-3.3-70B-Instruct.yaml index ffefb29e24..1742c04a06 100644 --- a/studio/backend/assets/configs/model_defaults/llama/unsloth_Llama-3.3-70B-Instruct.yaml +++ b/studio/backend/assets/configs/model_defaults/llama/unsloth_Llama-3.3-70B-Instruct.yaml @@ -35,6 +35,7 @@ lora: - "down_proj" use_rslora: false use_loftq: false + use_dora: false logging: enable_wandb: false diff --git a/studio/backend/assets/configs/model_defaults/llama/unsloth_Meta-Llama-3.1-70B-bnb-4bit.yaml b/studio/backend/assets/configs/model_defaults/llama/unsloth_Meta-Llama-3.1-70B-bnb-4bit.yaml index cd986a6da1..f33726b0dd 100644 --- a/studio/backend/assets/configs/model_defaults/llama/unsloth_Meta-Llama-3.1-70B-bnb-4bit.yaml +++ b/studio/backend/assets/configs/model_defaults/llama/unsloth_Meta-Llama-3.1-70B-bnb-4bit.yaml @@ -34,6 +34,7 @@ lora: - "down_proj" use_rslora: false use_loftq: false + use_dora: false logging: enable_wandb: false diff --git a/studio/backend/assets/configs/model_defaults/llama/unsloth_Meta-Llama-3.1-8B-Instruct-bnb-4bit.yaml b/studio/backend/assets/configs/model_defaults/llama/unsloth_Meta-Llama-3.1-8B-Instruct-bnb-4bit.yaml index 55dd3144c6..79b30bd758 100644 --- a/studio/backend/assets/configs/model_defaults/llama/unsloth_Meta-Llama-3.1-8B-Instruct-bnb-4bit.yaml +++ b/studio/backend/assets/configs/model_defaults/llama/unsloth_Meta-Llama-3.1-8B-Instruct-bnb-4bit.yaml @@ -34,6 +34,7 @@ lora: - "down_proj" use_rslora: false use_loftq: false + use_dora: false logging: enable_wandb: false diff --git a/studio/backend/assets/configs/model_defaults/llama/unsloth_llama-3-8b-Instruct-bnb-4bit.yaml b/studio/backend/assets/configs/model_defaults/llama/unsloth_llama-3-8b-Instruct-bnb-4bit.yaml index 8c9cb07fb9..4ee9a5a8ed 100644 --- a/studio/backend/assets/configs/model_defaults/llama/unsloth_llama-3-8b-Instruct-bnb-4bit.yaml +++ b/studio/backend/assets/configs/model_defaults/llama/unsloth_llama-3-8b-Instruct-bnb-4bit.yaml @@ -34,6 +34,7 @@ lora: - "down_proj" use_rslora: false use_loftq: false + use_dora: false logging: enable_wandb: false diff --git a/studio/backend/assets/configs/model_defaults/llama/unsloth_llama-3-8b-bnb-4bit.yaml b/studio/backend/assets/configs/model_defaults/llama/unsloth_llama-3-8b-bnb-4bit.yaml index 32441c5674..da20663688 100644 --- a/studio/backend/assets/configs/model_defaults/llama/unsloth_llama-3-8b-bnb-4bit.yaml +++ b/studio/backend/assets/configs/model_defaults/llama/unsloth_llama-3-8b-bnb-4bit.yaml @@ -34,6 +34,7 @@ lora: - "down_proj" use_rslora: false use_loftq: false + use_dora: false logging: enable_wandb: false diff --git a/studio/backend/assets/configs/model_defaults/llasa/unsloth_Llasa-3B.yaml b/studio/backend/assets/configs/model_defaults/llasa/unsloth_Llasa-3B.yaml index 6bba9c9633..30e4440afb 100644 --- a/studio/backend/assets/configs/model_defaults/llasa/unsloth_Llasa-3B.yaml +++ b/studio/backend/assets/configs/model_defaults/llasa/unsloth_Llasa-3B.yaml @@ -30,6 +30,7 @@ lora: - "v_proj" use_rslora: false use_loftq: false + use_dora: false logging: enable_wandb: false diff --git a/studio/backend/assets/configs/model_defaults/mistral/unsloth_Magistral-Small-2509-unsloth-bnb-4bit.yaml b/studio/backend/assets/configs/model_defaults/mistral/unsloth_Magistral-Small-2509-unsloth-bnb-4bit.yaml index f9833ce705..9bb0a93e63 100644 --- a/studio/backend/assets/configs/model_defaults/mistral/unsloth_Magistral-Small-2509-unsloth-bnb-4bit.yaml +++ b/studio/backend/assets/configs/model_defaults/mistral/unsloth_Magistral-Small-2509-unsloth-bnb-4bit.yaml @@ -35,6 +35,7 @@ lora: - "down_proj" use_rslora: false use_loftq: false + use_dora: false finetune_vision_layers: true finetune_language_layers: true finetune_attention_modules: true diff --git a/studio/backend/assets/configs/model_defaults/mistral/unsloth_Ministral-3-3B-Instruct-2512.yaml b/studio/backend/assets/configs/model_defaults/mistral/unsloth_Ministral-3-3B-Instruct-2512.yaml index 0ba857cd40..ded3607a14 100644 --- a/studio/backend/assets/configs/model_defaults/mistral/unsloth_Ministral-3-3B-Instruct-2512.yaml +++ b/studio/backend/assets/configs/model_defaults/mistral/unsloth_Ministral-3-3B-Instruct-2512.yaml @@ -35,6 +35,7 @@ lora: - "down_proj" use_rslora: false use_loftq: false + use_dora: false finetune_vision_layers: true finetune_language_layers: true finetune_attention_modules: true diff --git a/studio/backend/assets/configs/model_defaults/mistral/unsloth_Mistral-Nemo-Base-2407-bnb-4bit.yaml b/studio/backend/assets/configs/model_defaults/mistral/unsloth_Mistral-Nemo-Base-2407-bnb-4bit.yaml index 3476f2dd6d..2ac72f1c88 100644 --- a/studio/backend/assets/configs/model_defaults/mistral/unsloth_Mistral-Nemo-Base-2407-bnb-4bit.yaml +++ b/studio/backend/assets/configs/model_defaults/mistral/unsloth_Mistral-Nemo-Base-2407-bnb-4bit.yaml @@ -34,6 +34,7 @@ lora: - "down_proj" use_rslora: false use_loftq: false + use_dora: false logging: enable_wandb: false diff --git a/studio/backend/assets/configs/model_defaults/mistral/unsloth_Mistral-Small-Instruct-2409.yaml b/studio/backend/assets/configs/model_defaults/mistral/unsloth_Mistral-Small-Instruct-2409.yaml index eda04d21f9..a087ced1f3 100644 --- a/studio/backend/assets/configs/model_defaults/mistral/unsloth_Mistral-Small-Instruct-2409.yaml +++ b/studio/backend/assets/configs/model_defaults/mistral/unsloth_Mistral-Small-Instruct-2409.yaml @@ -34,6 +34,7 @@ lora: - "down_proj" use_rslora: false use_loftq: false + use_dora: false logging: enable_wandb: false diff --git a/studio/backend/assets/configs/model_defaults/mistral/unsloth_Pixtral-12B-2409.yaml b/studio/backend/assets/configs/model_defaults/mistral/unsloth_Pixtral-12B-2409.yaml index bcd0d20c8c..c9811f4f06 100644 --- a/studio/backend/assets/configs/model_defaults/mistral/unsloth_Pixtral-12B-2409.yaml +++ b/studio/backend/assets/configs/model_defaults/mistral/unsloth_Pixtral-12B-2409.yaml @@ -29,6 +29,7 @@ lora: - "all-linear" use_rslora: false use_loftq: false + use_dora: false finetune_vision_layers: true finetune_language_layers: true finetune_attention_modules: false diff --git a/studio/backend/assets/configs/model_defaults/mistral/unsloth_mistral-7b-instruct-v0.3-bnb-4bit.yaml b/studio/backend/assets/configs/model_defaults/mistral/unsloth_mistral-7b-instruct-v0.3-bnb-4bit.yaml index 34a033e32f..e3659d9fb0 100644 --- a/studio/backend/assets/configs/model_defaults/mistral/unsloth_mistral-7b-instruct-v0.3-bnb-4bit.yaml +++ b/studio/backend/assets/configs/model_defaults/mistral/unsloth_mistral-7b-instruct-v0.3-bnb-4bit.yaml @@ -34,6 +34,7 @@ lora: - "down_proj" use_rslora: false use_loftq: false + use_dora: false logging: enable_wandb: false diff --git a/studio/backend/assets/configs/model_defaults/mistral/unsloth_mistral-7b-v0.3-bnb-4bit.yaml b/studio/backend/assets/configs/model_defaults/mistral/unsloth_mistral-7b-v0.3-bnb-4bit.yaml index 98105eaf38..ee17efc54d 100644 --- a/studio/backend/assets/configs/model_defaults/mistral/unsloth_mistral-7b-v0.3-bnb-4bit.yaml +++ b/studio/backend/assets/configs/model_defaults/mistral/unsloth_mistral-7b-v0.3-bnb-4bit.yaml @@ -33,6 +33,7 @@ lora: - "down_proj" use_rslora: false use_loftq: false + use_dora: false logging: enable_wandb: false diff --git a/studio/backend/assets/configs/model_defaults/other/OuteAI_Llama-OuteTTS-1.0-1B.yaml b/studio/backend/assets/configs/model_defaults/other/OuteAI_Llama-OuteTTS-1.0-1B.yaml index 72b5b018e1..ef836b9b55 100644 --- a/studio/backend/assets/configs/model_defaults/other/OuteAI_Llama-OuteTTS-1.0-1B.yaml +++ b/studio/backend/assets/configs/model_defaults/other/OuteAI_Llama-OuteTTS-1.0-1B.yaml @@ -33,6 +33,7 @@ lora: - "v_proj" use_rslora: false use_loftq: false + use_dora: false logging: enable_wandb: false diff --git a/studio/backend/assets/configs/model_defaults/other/Spark-TTS-0.5B_LLM.yaml b/studio/backend/assets/configs/model_defaults/other/Spark-TTS-0.5B_LLM.yaml index d20751b0c7..c80fad35a8 100644 --- a/studio/backend/assets/configs/model_defaults/other/Spark-TTS-0.5B_LLM.yaml +++ b/studio/backend/assets/configs/model_defaults/other/Spark-TTS-0.5B_LLM.yaml @@ -38,6 +38,7 @@ lora: - "down_proj" use_rslora: false use_loftq: false + use_dora: false logging: enable_wandb: false diff --git a/studio/backend/assets/configs/model_defaults/other/sesame_csm-1b.yaml b/studio/backend/assets/configs/model_defaults/other/sesame_csm-1b.yaml index 8a80282a2a..034b5bd131 100644 --- a/studio/backend/assets/configs/model_defaults/other/sesame_csm-1b.yaml +++ b/studio/backend/assets/configs/model_defaults/other/sesame_csm-1b.yaml @@ -37,6 +37,7 @@ lora: - "down_proj" use_rslora: false use_loftq: false + use_dora: false logging: enable_wandb: false diff --git a/studio/backend/assets/configs/model_defaults/other/unsloth_GLM-4.7-Flash.yaml b/studio/backend/assets/configs/model_defaults/other/unsloth_GLM-4.7-Flash.yaml index a973c2d4e4..d1a226be79 100644 --- a/studio/backend/assets/configs/model_defaults/other/unsloth_GLM-4.7-Flash.yaml +++ b/studio/backend/assets/configs/model_defaults/other/unsloth_GLM-4.7-Flash.yaml @@ -35,6 +35,7 @@ lora: - "out_proj" use_rslora: false use_loftq: false + use_dora: false logging: enable_wandb: false diff --git a/studio/backend/assets/configs/model_defaults/other/unsloth_LFM2-1.2B.yaml b/studio/backend/assets/configs/model_defaults/other/unsloth_LFM2-1.2B.yaml index b0feafbd6e..1b8df5ced9 100644 --- a/studio/backend/assets/configs/model_defaults/other/unsloth_LFM2-1.2B.yaml +++ b/studio/backend/assets/configs/model_defaults/other/unsloth_LFM2-1.2B.yaml @@ -29,6 +29,7 @@ lora: - "all-linear" use_rslora: false use_loftq: false + use_dora: false logging: enable_wandb: false diff --git a/studio/backend/assets/configs/model_defaults/other/unsloth_Nemotron-3-Nano-30B-A3B.yaml b/studio/backend/assets/configs/model_defaults/other/unsloth_Nemotron-3-Nano-30B-A3B.yaml index 2c44c91eab..cecab7f083 100644 --- a/studio/backend/assets/configs/model_defaults/other/unsloth_Nemotron-3-Nano-30B-A3B.yaml +++ b/studio/backend/assets/configs/model_defaults/other/unsloth_Nemotron-3-Nano-30B-A3B.yaml @@ -37,6 +37,7 @@ lora: - "out_proj" use_rslora: false use_loftq: false + use_dora: false logging: enable_wandb: false diff --git a/studio/backend/assets/configs/model_defaults/other/unsloth_PaddleOCR-VL.yaml b/studio/backend/assets/configs/model_defaults/other/unsloth_PaddleOCR-VL.yaml index e1fbc08e4d..730be338cf 100644 --- a/studio/backend/assets/configs/model_defaults/other/unsloth_PaddleOCR-VL.yaml +++ b/studio/backend/assets/configs/model_defaults/other/unsloth_PaddleOCR-VL.yaml @@ -35,6 +35,7 @@ lora: - "down_proj" use_rslora: false use_loftq: false + use_dora: false finetune_vision_layers: true finetune_language_layers: true finetune_attention_modules: true diff --git a/studio/backend/assets/configs/model_defaults/other/unsloth_answerdotai_ModernBERT-large.yaml b/studio/backend/assets/configs/model_defaults/other/unsloth_answerdotai_ModernBERT-large.yaml index 2abdfd8ac3..a70ac0bd49 100644 --- a/studio/backend/assets/configs/model_defaults/other/unsloth_answerdotai_ModernBERT-large.yaml +++ b/studio/backend/assets/configs/model_defaults/other/unsloth_answerdotai_ModernBERT-large.yaml @@ -33,6 +33,7 @@ lora: - "down_proj" use_rslora: false use_loftq: false + use_dora: false logging: enable_wandb: false diff --git a/studio/backend/assets/configs/model_defaults/other/unsloth_orpheus-3b-0.1-ft.yaml b/studio/backend/assets/configs/model_defaults/other/unsloth_orpheus-3b-0.1-ft.yaml index 5a3c4abb48..90ead037f6 100644 --- a/studio/backend/assets/configs/model_defaults/other/unsloth_orpheus-3b-0.1-ft.yaml +++ b/studio/backend/assets/configs/model_defaults/other/unsloth_orpheus-3b-0.1-ft.yaml @@ -38,6 +38,7 @@ lora: - "down_proj" use_rslora: false use_loftq: false + use_dora: false logging: enable_wandb: false diff --git a/studio/backend/assets/configs/model_defaults/other/unsloth_tinyllama-bnb-4bit.yaml b/studio/backend/assets/configs/model_defaults/other/unsloth_tinyllama-bnb-4bit.yaml index a6ce27620f..a97c557c31 100644 --- a/studio/backend/assets/configs/model_defaults/other/unsloth_tinyllama-bnb-4bit.yaml +++ b/studio/backend/assets/configs/model_defaults/other/unsloth_tinyllama-bnb-4bit.yaml @@ -34,6 +34,7 @@ lora: - "down_proj" use_rslora: false use_loftq: false + use_dora: false logging: enable_wandb: false diff --git a/studio/backend/assets/configs/model_defaults/other/unsloth_whisper-large-v3.yaml b/studio/backend/assets/configs/model_defaults/other/unsloth_whisper-large-v3.yaml index 050774a8cd..6855ed6a35 100644 --- a/studio/backend/assets/configs/model_defaults/other/unsloth_whisper-large-v3.yaml +++ b/studio/backend/assets/configs/model_defaults/other/unsloth_whisper-large-v3.yaml @@ -33,6 +33,7 @@ lora: - "v_proj" use_rslora: false use_loftq: false + use_dora: false logging: enable_wandb: false diff --git a/studio/backend/assets/configs/model_defaults/phi/unsloth_Phi-3-medium-4k-instruct.yaml b/studio/backend/assets/configs/model_defaults/phi/unsloth_Phi-3-medium-4k-instruct.yaml index c574714d78..1933fed2ba 100644 --- a/studio/backend/assets/configs/model_defaults/phi/unsloth_Phi-3-medium-4k-instruct.yaml +++ b/studio/backend/assets/configs/model_defaults/phi/unsloth_Phi-3-medium-4k-instruct.yaml @@ -34,6 +34,7 @@ lora: - "down_proj" use_rslora: false use_loftq: false + use_dora: false logging: enable_wandb: false diff --git a/studio/backend/assets/configs/model_defaults/phi/unsloth_Phi-3.5-mini-instruct.yaml b/studio/backend/assets/configs/model_defaults/phi/unsloth_Phi-3.5-mini-instruct.yaml index e803c842b3..fda4e64158 100644 --- a/studio/backend/assets/configs/model_defaults/phi/unsloth_Phi-3.5-mini-instruct.yaml +++ b/studio/backend/assets/configs/model_defaults/phi/unsloth_Phi-3.5-mini-instruct.yaml @@ -34,6 +34,7 @@ lora: - "down_proj" use_rslora: false use_loftq: false + use_dora: false logging: enable_wandb: false diff --git a/studio/backend/assets/configs/model_defaults/phi/unsloth_Phi-4.yaml b/studio/backend/assets/configs/model_defaults/phi/unsloth_Phi-4.yaml index 4de3d9437d..c3910e3e5b 100644 --- a/studio/backend/assets/configs/model_defaults/phi/unsloth_Phi-4.yaml +++ b/studio/backend/assets/configs/model_defaults/phi/unsloth_Phi-4.yaml @@ -35,6 +35,7 @@ lora: - "down_proj" use_rslora: false use_loftq: false + use_dora: false logging: enable_wandb: false diff --git a/studio/backend/assets/configs/model_defaults/qwen/imdatta0_tiny_qwen3_moe_2.8B_0.7B.yaml b/studio/backend/assets/configs/model_defaults/qwen/imdatta0_tiny_qwen3_moe_2.8B_0.7B.yaml index bb75b3ce52..765ffee938 100644 --- a/studio/backend/assets/configs/model_defaults/qwen/imdatta0_tiny_qwen3_moe_2.8B_0.7B.yaml +++ b/studio/backend/assets/configs/model_defaults/qwen/imdatta0_tiny_qwen3_moe_2.8B_0.7B.yaml @@ -36,6 +36,7 @@ lora: - "gate_up_proj" use_rslora: false use_loftq: false + use_dora: false logging: enable_wandb: false diff --git a/studio/backend/assets/configs/model_defaults/qwen/unsloth_Qwen2-7B.yaml b/studio/backend/assets/configs/model_defaults/qwen/unsloth_Qwen2-7B.yaml index c305d328c2..39b30e9cee 100644 --- a/studio/backend/assets/configs/model_defaults/qwen/unsloth_Qwen2-7B.yaml +++ b/studio/backend/assets/configs/model_defaults/qwen/unsloth_Qwen2-7B.yaml @@ -34,6 +34,7 @@ lora: - "down_proj" use_rslora: false use_loftq: false + use_dora: false logging: enable_wandb: false diff --git a/studio/backend/assets/configs/model_defaults/qwen/unsloth_Qwen2-VL-7B-Instruct.yaml b/studio/backend/assets/configs/model_defaults/qwen/unsloth_Qwen2-VL-7B-Instruct.yaml index 6cee3d0949..f97e525798 100644 --- a/studio/backend/assets/configs/model_defaults/qwen/unsloth_Qwen2-VL-7B-Instruct.yaml +++ b/studio/backend/assets/configs/model_defaults/qwen/unsloth_Qwen2-VL-7B-Instruct.yaml @@ -29,6 +29,7 @@ lora: - "all-linear" use_rslora: false use_loftq: false + use_dora: false finetune_vision_layers: true finetune_language_layers: true finetune_attention_modules: true diff --git a/studio/backend/assets/configs/model_defaults/qwen/unsloth_Qwen2.5-1.5B-Instruct.yaml b/studio/backend/assets/configs/model_defaults/qwen/unsloth_Qwen2.5-1.5B-Instruct.yaml index 20ba81df2c..e19b94ede2 100644 --- a/studio/backend/assets/configs/model_defaults/qwen/unsloth_Qwen2.5-1.5B-Instruct.yaml +++ b/studio/backend/assets/configs/model_defaults/qwen/unsloth_Qwen2.5-1.5B-Instruct.yaml @@ -34,6 +34,7 @@ lora: - "down_proj" use_rslora: false use_loftq: false + use_dora: false logging: enable_wandb: false diff --git a/studio/backend/assets/configs/model_defaults/qwen/unsloth_Qwen2.5-7B.yaml b/studio/backend/assets/configs/model_defaults/qwen/unsloth_Qwen2.5-7B.yaml index 9930786c24..982f54b32f 100644 --- a/studio/backend/assets/configs/model_defaults/qwen/unsloth_Qwen2.5-7B.yaml +++ b/studio/backend/assets/configs/model_defaults/qwen/unsloth_Qwen2.5-7B.yaml @@ -34,6 +34,7 @@ lora: - "down_proj" use_rslora: false use_loftq: false + use_dora: false logging: enable_wandb: false diff --git a/studio/backend/assets/configs/model_defaults/qwen/unsloth_Qwen2.5-Coder-1.5B-Instruct.yaml b/studio/backend/assets/configs/model_defaults/qwen/unsloth_Qwen2.5-Coder-1.5B-Instruct.yaml index 775c7ce08f..5242128004 100644 --- a/studio/backend/assets/configs/model_defaults/qwen/unsloth_Qwen2.5-Coder-1.5B-Instruct.yaml +++ b/studio/backend/assets/configs/model_defaults/qwen/unsloth_Qwen2.5-Coder-1.5B-Instruct.yaml @@ -34,6 +34,7 @@ lora: - "down_proj" use_rslora: false use_loftq: false + use_dora: false logging: enable_wandb: false diff --git a/studio/backend/assets/configs/model_defaults/qwen/unsloth_Qwen2.5-Coder-14B-Instruct.yaml b/studio/backend/assets/configs/model_defaults/qwen/unsloth_Qwen2.5-Coder-14B-Instruct.yaml index 856db0c1b3..3559b636c6 100644 --- a/studio/backend/assets/configs/model_defaults/qwen/unsloth_Qwen2.5-Coder-14B-Instruct.yaml +++ b/studio/backend/assets/configs/model_defaults/qwen/unsloth_Qwen2.5-Coder-14B-Instruct.yaml @@ -35,6 +35,7 @@ lora: - "down_proj" use_rslora: false use_loftq: false + use_dora: false logging: enable_wandb: false diff --git a/studio/backend/assets/configs/model_defaults/qwen/unsloth_Qwen2.5-Coder-7B-Instruct-bnb-4bit.yaml b/studio/backend/assets/configs/model_defaults/qwen/unsloth_Qwen2.5-Coder-7B-Instruct-bnb-4bit.yaml index 5900392547..3bc6d69afc 100644 --- a/studio/backend/assets/configs/model_defaults/qwen/unsloth_Qwen2.5-Coder-7B-Instruct-bnb-4bit.yaml +++ b/studio/backend/assets/configs/model_defaults/qwen/unsloth_Qwen2.5-Coder-7B-Instruct-bnb-4bit.yaml @@ -34,6 +34,7 @@ lora: - "down_proj" use_rslora: false use_loftq: false + use_dora: false logging: enable_wandb: false diff --git a/studio/backend/assets/configs/model_defaults/qwen/unsloth_Qwen2.5-VL-7B-Instruct-bnb-4bit.yaml b/studio/backend/assets/configs/model_defaults/qwen/unsloth_Qwen2.5-VL-7B-Instruct-bnb-4bit.yaml index bd54b1d015..604b86dacd 100644 --- a/studio/backend/assets/configs/model_defaults/qwen/unsloth_Qwen2.5-VL-7B-Instruct-bnb-4bit.yaml +++ b/studio/backend/assets/configs/model_defaults/qwen/unsloth_Qwen2.5-VL-7B-Instruct-bnb-4bit.yaml @@ -29,6 +29,7 @@ lora: - "all-linear" use_rslora: false use_loftq: false + use_dora: false finetune_vision_layers: true finetune_language_layers: true finetune_attention_modules: true diff --git a/studio/backend/assets/configs/model_defaults/qwen/unsloth_Qwen3-0.6B.yaml b/studio/backend/assets/configs/model_defaults/qwen/unsloth_Qwen3-0.6B.yaml index 9feb6dcaae..daed4ebccb 100644 --- a/studio/backend/assets/configs/model_defaults/qwen/unsloth_Qwen3-0.6B.yaml +++ b/studio/backend/assets/configs/model_defaults/qwen/unsloth_Qwen3-0.6B.yaml @@ -35,6 +35,7 @@ lora: - "down_proj" use_rslora: false use_loftq: false + use_dora: false logging: enable_wandb: false diff --git a/studio/backend/assets/configs/model_defaults/qwen/unsloth_Qwen3-14B-Base-unsloth-bnb-4bit.yaml b/studio/backend/assets/configs/model_defaults/qwen/unsloth_Qwen3-14B-Base-unsloth-bnb-4bit.yaml index a40eace253..05eef89b88 100644 --- a/studio/backend/assets/configs/model_defaults/qwen/unsloth_Qwen3-14B-Base-unsloth-bnb-4bit.yaml +++ b/studio/backend/assets/configs/model_defaults/qwen/unsloth_Qwen3-14B-Base-unsloth-bnb-4bit.yaml @@ -35,6 +35,7 @@ lora: - "down_proj" use_rslora: false use_loftq: false + use_dora: false logging: enable_wandb: false diff --git a/studio/backend/assets/configs/model_defaults/qwen/unsloth_Qwen3-14B.yaml b/studio/backend/assets/configs/model_defaults/qwen/unsloth_Qwen3-14B.yaml index c130771c32..b4580e6d71 100644 --- a/studio/backend/assets/configs/model_defaults/qwen/unsloth_Qwen3-14B.yaml +++ b/studio/backend/assets/configs/model_defaults/qwen/unsloth_Qwen3-14B.yaml @@ -35,6 +35,7 @@ lora: - "down_proj" use_rslora: false use_loftq: false + use_dora: false logging: enable_wandb: false diff --git a/studio/backend/assets/configs/model_defaults/qwen/unsloth_Qwen3-30B-A3B-Instruct-2507.yaml b/studio/backend/assets/configs/model_defaults/qwen/unsloth_Qwen3-30B-A3B-Instruct-2507.yaml index 2fb3a95c30..2eceb7d0de 100644 --- a/studio/backend/assets/configs/model_defaults/qwen/unsloth_Qwen3-30B-A3B-Instruct-2507.yaml +++ b/studio/backend/assets/configs/model_defaults/qwen/unsloth_Qwen3-30B-A3B-Instruct-2507.yaml @@ -36,6 +36,7 @@ lora: - "gate_up_proj" use_rslora: false use_loftq: false + use_dora: false logging: enable_wandb: false diff --git a/studio/backend/assets/configs/model_defaults/qwen/unsloth_Qwen3-32B.yaml b/studio/backend/assets/configs/model_defaults/qwen/unsloth_Qwen3-32B.yaml index 152f4ae06a..032091880c 100644 --- a/studio/backend/assets/configs/model_defaults/qwen/unsloth_Qwen3-32B.yaml +++ b/studio/backend/assets/configs/model_defaults/qwen/unsloth_Qwen3-32B.yaml @@ -35,6 +35,7 @@ lora: - "down_proj" use_rslora: false use_loftq: false + use_dora: false logging: enable_wandb: false diff --git a/studio/backend/assets/configs/model_defaults/qwen/unsloth_Qwen3-4B-Instruct-2507.yaml b/studio/backend/assets/configs/model_defaults/qwen/unsloth_Qwen3-4B-Instruct-2507.yaml index 94fe000708..e0e7f4ee3d 100644 --- a/studio/backend/assets/configs/model_defaults/qwen/unsloth_Qwen3-4B-Instruct-2507.yaml +++ b/studio/backend/assets/configs/model_defaults/qwen/unsloth_Qwen3-4B-Instruct-2507.yaml @@ -35,6 +35,7 @@ lora: - "down_proj" use_rslora: false use_loftq: false + use_dora: false logging: enable_wandb: false diff --git a/studio/backend/assets/configs/model_defaults/qwen/unsloth_Qwen3-4B-Thinking-2507.yaml b/studio/backend/assets/configs/model_defaults/qwen/unsloth_Qwen3-4B-Thinking-2507.yaml index 3c325485d2..bb463849ed 100644 --- a/studio/backend/assets/configs/model_defaults/qwen/unsloth_Qwen3-4B-Thinking-2507.yaml +++ b/studio/backend/assets/configs/model_defaults/qwen/unsloth_Qwen3-4B-Thinking-2507.yaml @@ -35,6 +35,7 @@ lora: - "down_proj" use_rslora: false use_loftq: false + use_dora: false logging: enable_wandb: false diff --git a/studio/backend/assets/configs/model_defaults/qwen/unsloth_Qwen3-VL-8B-Instruct-unsloth-bnb-4bit.yaml b/studio/backend/assets/configs/model_defaults/qwen/unsloth_Qwen3-VL-8B-Instruct-unsloth-bnb-4bit.yaml index 5b47c3bdd2..23e2b89dd0 100644 --- a/studio/backend/assets/configs/model_defaults/qwen/unsloth_Qwen3-VL-8B-Instruct-unsloth-bnb-4bit.yaml +++ b/studio/backend/assets/configs/model_defaults/qwen/unsloth_Qwen3-VL-8B-Instruct-unsloth-bnb-4bit.yaml @@ -29,6 +29,7 @@ lora: - "all-linear" use_rslora: false use_loftq: false + use_dora: false finetune_vision_layers: true finetune_language_layers: true finetune_attention_modules: true diff --git a/studio/backend/assets/configs/vision_lora.yaml b/studio/backend/assets/configs/vision_lora.yaml index 063a970316..a06f971523 100644 --- a/studio/backend/assets/configs/vision_lora.yaml +++ b/studio/backend/assets/configs/vision_lora.yaml @@ -30,6 +30,7 @@ lora: vision_all_linear: true use_rslora: false use_loftq: false + use_dora: false finetune_vision_layers: true finetune_language_layers: true finetune_attention_modules: true diff --git a/studio/backend/auth/authentication.py b/studio/backend/auth/authentication.py index b13cd1c851..2e9520827e 100644 --- a/studio/backend/auth/authentication.py +++ b/studio/backend/auth/authentication.py @@ -11,11 +11,12 @@ import jwt from .storage import ( API_KEY_PREFIX, + credential_generation, get_jwt_secret, get_user_and_secret, load_jwt_secret, save_refresh_token, - validate_api_key, + validate_api_key_with_credential, verify_refresh_token, ) @@ -54,11 +55,14 @@ def create_access_token( expires_delta: Optional[timedelta] = None, *, desktop: bool = False, + secret: Optional[str] = None, ) -> str: """ Create a signed JWT for the given subject (e.g. username). - Valid across restarts: the signing secret is stored in SQLite. + Valid across restarts: the signing secret is stored in SQLite. Callers that + already verified a credential pass ``secret`` so a rotation landing mid-request + cannot sign the token with the credential that just replaced it. """ to_encode = {"sub": subject} if desktop: @@ -69,7 +73,7 @@ def create_access_token( to_encode.update({"exp": expire}) return jwt.encode( to_encode, - _get_secret_for_subject(subject), + secret if secret is not None else _get_secret_for_subject(subject), algorithm = ALGORITHM, ) @@ -96,15 +100,28 @@ def is_desktop_access_token(token: str) -> bool: return payload.get("sub") == subject and payload.get("desktop") is True -def create_refresh_token(subject: str, *, desktop: bool = False) -> str: +def create_refresh_token( + subject: str, + *, + desktop: bool = False, + secret: Optional[str] = None, +) -> str: """ Create a random refresh token, store its hash in SQLite, and return it. Refresh tokens are opaque (not JWTs); expire after REFRESH_TOKEN_EXPIRE_DAYS. + ``secret`` stamps the token with the credential version the caller verified, + so a rotation cannot leave a token minted from the replaced credential valid. """ token = secrets.token_urlsafe(48) expires_at = datetime.now(timezone.utc) + timedelta(days = REFRESH_TOKEN_EXPIRE_DAYS) - save_refresh_token(token, subject, expires_at.isoformat(), is_desktop = desktop) + save_refresh_token( + token, + subject, + expires_at.isoformat(), + is_desktop = desktop, + secret_gen = credential_generation(secret) if secret is not None else None, + ) return token @@ -137,7 +154,22 @@ def reload_secret() -> None: async def get_current_subject(credentials: HTTPAuthorizationCredentials = Depends(security)) -> str: """Validate JWT and require the password-change flow to be completed.""" - return await _get_current_subject( + subject, _generation = await _get_current_credential( + credentials, + allow_password_change = False, + ) + return subject + + +async def get_current_credential( + credentials: HTTPAuthorizationCredentials = Depends(security), +) -> Tuple[str, Optional[str]]: + """As get_current_subject, but also returns the credential generation. + + For routes that persist a new credential and must not do so on behalf of one + a concurrent reset has revoked. + """ + return await _get_current_credential( credentials, allow_password_change = False, ) @@ -148,7 +180,7 @@ async def authenticated_via_api_key( ) -> bool: """True when the caller used an sk-unsloth API key, not a UI session JWT. - Lets routes treat programmatic API callers differently from the Studio UI + Lets routes treat programmatic API callers differently from the Unsloth UI (e.g. refuse a teardown the UI would allow). """ return bool(credentials and credentials.credentials.startswith(API_KEY_PREFIX)) @@ -158,27 +190,49 @@ async def get_current_subject_allow_password_change( credentials: HTTPAuthorizationCredentials = Depends(security), ) -> str: """Validate JWT but allow access to the password-change endpoint.""" - return await _get_current_subject( + subject, _generation = await _get_current_credential( credentials, allow_password_change = True, ) + return subject -async def _get_current_subject( +# The literal the examples ship with; pasted unedited more often than a revoked key. +API_KEY_PLACEHOLDER = f"{API_KEY_PREFIX}YOUR_KEY" + + +def _invalid_api_key_detail(token: str) -> str: + """Why the key failed. Only the example placeholder is called out; every real + key gets one indistinguishable message, so this leaks no key existence.""" + if token == API_KEY_PLACEHOLDER: + return ( + "This is the placeholder key from the example. Create an API key in " + f"Unsloth Studio under Settings > API and use it in place of {API_KEY_PLACEHOLDER}." + ) + return "Invalid or expired API key" + + +async def _get_current_credential( credentials: HTTPAuthorizationCredentials, *, allow_password_change: bool -) -> str: - """FastAPI dependency: validate the JWT and return the subject. Use on protected routes.""" +) -> Tuple[str, Optional[str]]: + """Validate the bearer and return ``(subject, credential generation)``. + + The generation is the credential version this request actually authenticated + against. Routes that persist new credentials must bind their write to it, or + a reset landing mid-request would bless what it just revoked. + """ token = credentials.credentials # --- API key path (sk-unsloth-...) --- if token.startswith(API_KEY_PREFIX): - username = validate_api_key(token) - if username is None: + verified = validate_api_key_with_credential(token) + if verified is None: raise HTTPException( status_code = status.HTTP_401_UNAUTHORIZED, - detail = "Invalid or expired API key", + detail = _invalid_api_key_detail(token), ) - return username + username, secret = verified + return username, credential_generation(secret) # --- JWT path --- subject = _decode_subject_without_verification(token) @@ -209,7 +263,7 @@ async def _get_current_subject( status_code = status.HTTP_403_FORBIDDEN, detail = "Password change required", ) - return subject + return subject, credential_generation(jwt_secret) except jwt.InvalidTokenError: raise HTTPException( status_code = status.HTTP_401_UNAUTHORIZED, diff --git a/studio/backend/auth/bootstrap_timeout.py b/studio/backend/auth/bootstrap_timeout.py index 728433dc54..97a8086f04 100644 --- a/studio/backend/auth/bootstrap_timeout.py +++ b/studio/backend/auth/bootstrap_timeout.py @@ -1,13 +1,13 @@ # SPDX-License-Identifier: AGPL-3.0-only # Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0 -"""Auto-shutdown for an exposed first-run Studio whose admin password is unchanged. +"""Auto-shutdown for an exposed first-run Unsloth whose admin password is unchanged. On a fresh install the seeded bootstrap admin password stays a valid login credential until first login changes it. When the web UI is put on the network (``--secure`` / ``0.0.0.0``) and nobody completes that first-login change within -a deadline, tear Studio down so a fresh, unconfigured instance does not stay -publicly reachable indefinitely. If the password was changed, Studio keeps +a deadline, tear Unsloth down so a fresh, unconfigured instance does not stay +publicly reachable indefinitely. If the password was changed, Unsloth keeps running. Scope: web UI launches only (never ``--api-only``, which authenticates by API @@ -98,7 +98,7 @@ def enforce_bootstrap_password_deadline( ) -> bool: """Deadline handler: shut down iff the seeded admin password is still unchanged. - Returns True if it shut Studio down, False if it left it running (the + Returns True if it shut Unsloth down, False if it left it running (the password was changed in time). """ try: @@ -106,7 +106,7 @@ def enforce_bootstrap_password_deadline( except Exception: return False if not still_default: - return False # password changed in time -> leave Studio running + return False # password changed in time -> leave Unsloth running message = ( "\nUnsloth Studio was exposed on the network but its default admin " diff --git a/studio/backend/auth/storage.py b/studio/backend/auth/storage.py index 9bb3ab5735..6cf4d44834 100644 --- a/studio/backend/auth/storage.py +++ b/studio/backend/auth/storage.py @@ -9,6 +9,7 @@ import ipaddress import os import secrets import sqlite3 +import tempfile import threading from datetime import datetime, timezone from typing import Optional, Tuple @@ -30,6 +31,97 @@ _BOOTSTRAP_PW_PATH = DB_PATH.parent / ".bootstrap_password" _bootstrap_password: Optional[str] = None +def _bootstrap_file_bytes(password: str) -> bytes: + """Exact on-disk form: the secret plus one LF. + + Bytes, not text: text mode writes CRLF on Windows, and `$(cat ...)` strips + the LF but leaves the CR attached to the credential. + """ + return (password + "\n").encode("utf-8") + + +def _persist_bootstrap_password(password: str) -> None: + """Atomically write the bootstrap password 0600, LF terminated on every OS. + + A partial write would destroy the only plaintext recovery credential. + """ + fd, tmp_name = tempfile.mkstemp( + prefix = f".{_BOOTSTRAP_PW_PATH.name}.", dir = _BOOTSTRAP_PW_PATH.parent + ) + try: + with os.fdopen(fd, "wb") as f: + f.write(_bootstrap_file_bytes(password)) + try: + os.chmod(tmp_name, 0o600) + except OSError: + pass + os.replace(tmp_name, _BOOTSTRAP_PW_PATH) + except BaseException: + try: + os.unlink(tmp_name) + except OSError: + pass + raise + + +def _normalise_bootstrap_file(raw: bytes, password: str) -> None: + """Append the LF a pre-newline release left off. + + Append-only, and only when the file is exactly the credential: + clear_bootstrap_password() may unlink or (when unlink fails, notably on + Windows while this descriptor is open) truncate through another descriptor + after we read, so a rewrite could restore revoked plaintext. An append + cannot: worst case is a lone "\\n" over a cleared file, which strips back to + no bootstrap password. Pre-newline releases wrote no terminator at all, so + that is the only shape in the wild; anything else reads fine, since every + reader strips, and is left alone. + """ + if raw != password.encode("utf-8"): + return + + # O_BINARY: without it Windows opens in text mode and turns the LF straight + # back into CRLF, the bug being fixed. + fd = os.open( + _BOOTSTRAP_PW_PATH, + os.O_WRONLY | os.O_APPEND | getattr(os, "O_BINARY", 0), + ) + try: + os.write(fd, b"\n") + try: + os.fchmod(fd, 0o600) + except (AttributeError, OSError): + # fchmod only reached Windows in 3.13. + pass + finally: + os.close(fd) + + +def _read_persisted_bootstrap_password() -> Optional[str]: + """Read the persisted password, normalising the file if it is malformed.""" + if not _BOOTSTRAP_PW_PATH.is_file(): + return None + + # No caller handles a raise, so an unreadable file has to mean "no bootstrap + # password", not a dead backend. We write UTF-8, so undecodable bytes are + # damage whose plaintext is worthless anyway. + try: + raw = _BOOTSTRAP_PW_PATH.read_bytes() + password = raw.decode("utf-8").strip() + except (OSError, UnicodeDecodeError): + return None + if not password: + return None + + # Older releases wrote no terminator; best-effort, a read-only auth dir must + # not fail startup. + if raw != _bootstrap_file_bytes(password): + try: + _normalise_bootstrap_file(raw, password) + except OSError: + pass + return password + + def generate_bootstrap_password() -> str: """Generate a 4-word diceware passphrase and persist it to disk. @@ -43,10 +135,10 @@ def generate_bootstrap_password() -> str: return _bootstrap_password # Persisted from a previous run? - if _BOOTSTRAP_PW_PATH.is_file(): - _bootstrap_password = _BOOTSTRAP_PW_PATH.read_text().strip() - if _bootstrap_password: - return _bootstrap_password + persisted = _read_persisted_bootstrap_password() + if persisted: + _bootstrap_password = persisted + return _bootstrap_password # First startup: generate a fresh passphrase. import diceware @@ -57,11 +149,7 @@ def generate_bootstrap_password() -> str: # Persist so the same passphrase survives restarts until password change. ensure_dir(_BOOTSTRAP_PW_PATH.parent) - _BOOTSTRAP_PW_PATH.write_text(_bootstrap_password) - try: - os.chmod(_BOOTSTRAP_PW_PATH, 0o600) - except OSError: - pass + _persist_bootstrap_password(_bootstrap_password) return _bootstrap_password @@ -72,13 +160,14 @@ def get_bootstrap_password() -> Optional[str]: def _load_bootstrap_password() -> Optional[str]: - """Load an existing bootstrap password without creating one.""" + """Load an existing bootstrap password without creating one. + + Upgrades take this path, not generate_bootstrap_password() + (ensure_default_admin short-circuits once the admin row exists), so it has + to normalise too. + """ global _bootstrap_password - _bootstrap_password = None - if _BOOTSTRAP_PW_PATH.is_file(): - bootstrap_password = _BOOTSTRAP_PW_PATH.read_text().strip() - if bootstrap_password: - _bootstrap_password = bootstrap_password + _bootstrap_password = _read_persisted_bootstrap_password() return _bootstrap_password @@ -97,9 +186,9 @@ def clear_bootstrap_password() -> None: # Removal failed (Windows AV, read-only auth dir). The hash is already # committed, so don't fail the change -- but truncate the file so its # stale plaintext can't be re-seeded by generate_bootstrap_password() - # if a later reset-password deletes auth.db and re-validates it. + # if auth.db is ever recreated. try: - _BOOTSTRAP_PW_PATH.write_text("") + _BOOTSTRAP_PW_PATH.write_text("", encoding = "utf-8") cleared = True except OSError: cleared = False @@ -132,6 +221,31 @@ def _hash_token(token: str) -> str: return hashlib.sha256(token.encode("utf-8")).hexdigest() +class CredentialRotated(Exception): + """A password reset revoked the credential this request authenticated with.""" + + +def credential_generation(jwt_secret: str) -> str: + """Marker for the credential version a refresh token was issued under. + + Every password change rotates ``jwt_secret``, so a token stamped with the + previous one is rejected even if it was inserted after the revoking DELETE. + """ + return hashlib.sha256(jwt_secret.encode("utf-8")).hexdigest() + + +def _current_secret(conn: sqlite3.Connection, username: str) -> Optional[str]: + row = conn.execute( + "SELECT jwt_secret FROM auth_user WHERE username = ?", (username,) + ).fetchone() + return row["jwt_secret"] if row else None + + +def _current_generation(conn: sqlite3.Connection, username: str) -> Optional[str]: + secret = _current_secret(conn, username) + return credential_generation(secret) if secret is not None else None + + def get_connection() -> sqlite3.Connection: """Get a connection to the auth database, creating tables if needed.""" ensure_dir(DB_PATH.parent) @@ -146,7 +260,7 @@ def get_connection() -> sqlite3.Connection: pass conn.row_factory = sqlite3.Row # WAL lets token reads run concurrently with refresh-token writes; - # busy_timeout bounds lock waits. Matches the other Studio SQLite stores. + # busy_timeout bounds lock waits. Matches the other Unsloth SQLite stores. # Set busy_timeout first: switching journal_mode needs a lock, so if a # refresh-token write already holds one, journal_mode=WAL raises SQLITE_BUSY; # with busy_timeout already in effect it waits instead of failing and leaving @@ -175,7 +289,8 @@ def get_connection() -> sqlite3.Connection: token_hash TEXT NOT NULL, username TEXT NOT NULL, expires_at TEXT NOT NULL, - is_desktop INTEGER NOT NULL DEFAULT 0 + is_desktop INTEGER NOT NULL DEFAULT 0, + secret_gen TEXT ); """ ) @@ -214,6 +329,8 @@ def get_connection() -> sqlite3.Connection: refresh_columns = {row["name"] for row in conn.execute("PRAGMA table_info(refresh_tokens)")} if "is_desktop" not in refresh_columns: conn.execute("ALTER TABLE refresh_tokens ADD COLUMN is_desktop INTEGER NOT NULL DEFAULT 0") + if "secret_gen" not in refresh_columns: + conn.execute("ALTER TABLE refresh_tokens ADD COLUMN secret_gen TEXT") conn.commit() return conn @@ -305,8 +422,8 @@ def get_or_create_identity_secret() -> bytes: def compute_identity_proof(nonce: bytes, host: str, port: int) -> str: """HMAC-SHA256 proof that the caller holds this install's identity secret, bound to the loopback address and port the connection landed on. A proof - relayed from a Studio on a different address/port (a squatter proxying to the - real one, e.g. localhost resolving to ::1 while Studio is on 127.0.0.1) was + relayed from an Unsloth on a different address/port (a squatter proxying to the + real one, e.g. localhost resolving to ::1 while Unsloth is on 127.0.0.1) was computed for that other endpoint and won't match the one the client dialed.""" try: host = ipaddress.ip_address(host).compressed # normalise 127.0.0.1 / ::1 forms @@ -587,12 +704,22 @@ def update_password( new_password: str, *, revoke_refresh_tokens: bool = False, -) -> bool: + expect_password_hash: Optional[str] = None, +) -> Optional[str]: """Update password, clear first-login requirement, rotate JWT secret. + Returns the new JWT secret, or None when nothing was updated. Callers that + mint tokens for the caller must sign with the returned secret: re-reading it + would pick up a reset that landed between this commit and the mint. + ``revoke_refresh_tokens`` deletes the user's refresh tokens in the SAME transaction: a separate delete could fail after the password commit and leave a pre-change token still able to mint access tokens. + + ``expect_password_hash`` makes the write conditional on the credential the + caller verified still being current, so a request that checked the old + password cannot overwrite a reset that landed while it was in flight. + Returns False when the credential moved underneath it. """ from .hashing import hash_password @@ -600,21 +727,32 @@ def update_password( jwt_secret = secrets.token_urlsafe(64) conn = get_connection() try: - cursor = conn.execute( - """ - UPDATE auth_user - SET password_salt = ?, password_hash = ?, jwt_secret = ?, must_change_password = 0 - WHERE username = ? - """, - (salt, pwd_hash, jwt_secret, username), - ) + if expect_password_hash is None: + cursor = conn.execute( + """ + UPDATE auth_user + SET password_salt = ?, password_hash = ?, jwt_secret = ?, must_change_password = 0 + WHERE username = ? + """, + (salt, pwd_hash, jwt_secret, username), + ) + else: + cursor = conn.execute( + """ + UPDATE auth_user + SET password_salt = ?, password_hash = ?, jwt_secret = ?, must_change_password = 0 + WHERE username = ? AND password_hash = ? + """, + (salt, pwd_hash, jwt_secret, username, expect_password_hash), + ) if revoke_refresh_tokens and cursor.rowcount > 0: conn.execute("DELETE FROM refresh_tokens WHERE username = ?", (username,)) conn.commit() if cursor.rowcount > 0: clear_bootstrap_password() clear_desktop_secret() - return cursor.rowcount > 0 + return jwt_secret + return None finally: conn.close() @@ -625,35 +763,49 @@ def save_refresh_token( expires_at: str, *, is_desktop: bool = False, + secret_gen: Optional[str] = None, ) -> None: """ Store a hashed refresh token with its associated username and expiry. + + ``secret_gen`` binds the token to a credential version; it defaults to the + current one, and callers that already verified a credential must pass the + version they verified rather than let this re-read a rotated one. """ token_hash = _hash_token(token) conn = get_connection() try: + if secret_gen is None: + secret_gen = _current_generation(conn, username) conn.execute( """ - INSERT INTO refresh_tokens (token_hash, username, expires_at, is_desktop) - VALUES (?, ?, ?, ?) + INSERT INTO refresh_tokens (token_hash, username, expires_at, is_desktop, secret_gen) + VALUES (?, ?, ?, ?, ?) """, - (token_hash, username, expires_at, int(is_desktop)), + (token_hash, username, expires_at, int(is_desktop), secret_gen), ) conn.commit() finally: conn.close() -def consume_refresh_token(token: str) -> Optional[Tuple[str, bool]]: +def consume_refresh_token(token: str) -> Optional[Tuple[str, bool, str]]: """Atomically validate-and-delete a refresh token for single-use rotation. DELETE RETURNING fuses validate and delete into one statement so two - concurrent refresh requests cannot both consume the same token. + concurrent refresh requests cannot both consume the same token. Returns + ``(username, is_desktop, jwt_secret)``; the caller must mint the replacement + tokens against that secret so a rotation landing mid-refresh cannot issue a + post-rotation session from a pre-rotation token. """ token_hash = _hash_token(token) now = datetime.now(timezone.utc).isoformat() conn = get_connection() try: + # One transaction with the delete: an unstamped legacy row has no + # generation to compare, so reading the credential after committing would + # hand a reset's new secret to a token issued before it. + conn.execute("BEGIN IMMEDIATE") conn.execute( "DELETE FROM refresh_tokens WHERE expires_at < ?", (now,), @@ -662,15 +814,21 @@ def consume_refresh_token(token: str) -> Optional[Tuple[str, bool]]: """ DELETE FROM refresh_tokens WHERE token_hash = ? AND expires_at >= ? - RETURNING username, is_desktop + RETURNING username, is_desktop, secret_gen """, (token_hash, now), ) row = cur.fetchone() - conn.commit() if row is None: + conn.commit() return None - return row["username"], bool(row["is_desktop"]) + secret = _current_secret(conn, row["username"]) + conn.commit() + if secret is None: + return None + if row["secret_gen"] is not None and row["secret_gen"] != credential_generation(secret): + return None + return row["username"], bool(row["is_desktop"]), secret finally: conn.close() @@ -694,7 +852,7 @@ def verify_refresh_token(token: str) -> Optional[Tuple[str, bool]]: cur = conn.execute( """ - SELECT id, username, expires_at, is_desktop FROM refresh_tokens + SELECT id, username, expires_at, is_desktop, secret_gen FROM refresh_tokens WHERE token_hash = ? """, (token_hash,), @@ -703,6 +861,13 @@ def verify_refresh_token(token: str) -> Optional[Tuple[str, bool]]: if row is None: return None + if row["secret_gen"] is not None and row["secret_gen"] != _current_generation( + conn, row["username"] + ): + conn.execute("DELETE FROM refresh_tokens WHERE id = ?", (row["id"],)) + conn.commit() + return None + # Check expiry expires_at = datetime.fromisoformat(row["expires_at"]) if datetime.now(timezone.utc) > expires_at: @@ -747,30 +912,41 @@ def create_desktop_secret() -> str: conn.close() -def validate_desktop_secret(raw_secret: str) -> Optional[str]: - """Return the real admin username when the desktop secret matches.""" +def validate_desktop_secret_with_credential(raw_secret: str) -> Optional[Tuple[str, str]]: + """Validate the desktop secret and return ``(username, jwt_secret)``. + + Both reads share one transaction so the returned secret is the credential + version the desktop secret was checked against; a reset landing mid-request + then invalidates the tokens minted from it rather than blessing them. + """ if not raw_secret.startswith(DESKTOP_SECRET_PREFIX): return None - if get_user_and_secret(DEFAULT_ADMIN_USERNAME) is None: - return None secret_hash = _pbkdf2_desktop_secret(raw_secret) conn = get_connection() try: - cur = conn.execute( + conn.execute("BEGIN") + row = conn.execute( "SELECT value FROM app_secrets WHERE key = ?", (_DESKTOP_SECRET_HASH_KEY,), - ) - row = cur.fetchone() - if row is None: + ).fetchone() + if row is None or not secrets.compare_digest(row["value"], secret_hash): return None - if not secrets.compare_digest(row["value"], secret_hash): + jwt_secret = _current_secret(conn, DEFAULT_ADMIN_USERNAME) + if jwt_secret is None: return None - return DEFAULT_ADMIN_USERNAME + return DEFAULT_ADMIN_USERNAME, jwt_secret finally: + conn.rollback() conn.close() +def validate_desktop_secret(raw_secret: str) -> Optional[str]: + """Return the real admin username when the desktop secret matches.""" + verified = validate_desktop_secret_with_credential(raw_secret) + return verified[0] if verified else None + + def clear_desktop_secret() -> None: """Remove backend-side desktop auth state.""" conn = get_connection() @@ -796,6 +972,7 @@ def create_api_key( name: str, expires_at: Optional[str] = None, internal: bool = False, + expect_gen: Optional[str] = None, ) -> Tuple[str, dict]: """Create a new API key for *username*. @@ -804,6 +981,10 @@ def create_api_key( Pass ``internal=True`` for keys minted by workflows (e.g. data-recipe runs) that should not appear in user-facing key listings. + + ``expect_gen`` ties the insert to the credential generation the request + authenticated under, so a session revoked by a concurrent password reset + cannot mint a key that outlives it. Raises ``CredentialRotated`` if it moved. """ raw_key = API_KEY_PREFIX + secrets.token_hex(16) key_hash = _pbkdf2_api_key(raw_key) @@ -812,6 +993,12 @@ def create_api_key( conn = get_connection() try: + if expect_gen is not None: + conn.execute("BEGIN IMMEDIATE") + if _current_generation(conn, username) != expect_gen: + raise CredentialRotated( + "The credential this request authenticated with was revoked." + ) conn.execute( """ INSERT INTO api_keys (username, key_prefix, key_hash, name, created_at, expires_at, is_internal) @@ -900,15 +1087,25 @@ def revoke_internal_api_key(key_id: int) -> bool: def validate_api_key(raw_key: str) -> Optional[str]: - """Validate *raw_key* and return the owning username, or ``None``. + """Validate *raw_key* and return the owning username, or ``None``.""" + verified = validate_api_key_with_credential(raw_key) + return verified[0] if verified else None - Also updates ``last_used_at`` on success. + +def validate_api_key_with_credential(raw_key: str) -> Optional[Tuple[str, str]]: + """Validate *raw_key* and return ``(username, jwt_secret)``, or ``None``. + + Also updates ``last_used_at`` on success. The key check and the credential + read share one write transaction, so the returned version is the one the key + was actually valid under: a reset committing right after cannot have its new + generation handed to a request the key it revoked authenticated. """ cache_id = _api_key_cache_id(raw_key) cached_hash = _api_key_hash_cache.get(cache_id) key_hash = cached_hash if cached_hash is not None else _pbkdf2_api_key(raw_key) conn = get_connection() try: + conn.execute("BEGIN IMMEDIATE") cur = conn.execute( "SELECT id, username, is_active, expires_at FROM api_keys WHERE key_hash = ?", (key_hash,), @@ -928,11 +1125,15 @@ def validate_api_key(raw_key: str) -> Optional[str]: expires = datetime.fromisoformat(row["expires_at"]) if datetime.now(timezone.utc) > expires: return None + secret = _current_secret(conn, row["username"]) + if secret is None: + 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"] + return row["username"], secret finally: + conn.rollback() conn.close() diff --git a/studio/backend/auth/terminal_prompt.py b/studio/backend/auth/terminal_prompt.py index 8491019ae9..925404f47d 100644 --- a/studio/backend/auth/terminal_prompt.py +++ b/studio/backend/auth/terminal_prompt.py @@ -2,14 +2,14 @@ # Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0 """Interactive terminal prompt that forces a bootstrap password change before -Studio is exposed on a public Cloudflare URL (``--secure`` / ``--cloudflare``). +Unsloth is exposed on a public Cloudflare URL (``--secure`` / ``--cloudflare``). Masked input echoes one ``*`` per keystroke (unlike ``getpass``). Works on Windows (``msvcrt``) and Linux/macOS (``termios``). All output goes to stderr so redirected stdout never swallows the prompt. Mirrored for the CLI at ``unsloth_cli/commands/_password_prompt.py`` (the CLI -cannot import the Studio backend package); keep the two in sync. +cannot import the Unsloth backend package); keep the two in sync. """ from __future__ import annotations @@ -236,6 +236,10 @@ def prompt_for_password_change( out.write(f"Password must be at least {min_length} characters; try again.\n") out.flush() continue + if any(ch.isspace() for ch in new_password): + out.write("Password cannot contain spaces; try again.\n") + out.flush() + continue if is_current_password(new_password): out.write( "New password must differ from the current bootstrap password; try again.\n" @@ -252,7 +256,7 @@ def prompt_for_password_change( out.flush() return True except (KeyboardInterrupt, EOFError): - out.write("Password change aborted; not exposing Studio.\n") + out.write("Password change aborted; not exposing Unsloth.\n") out.flush() return False diff --git a/studio/backend/cloudflare_tunnel.py b/studio/backend/cloudflare_tunnel.py index ef7bacba67..f7967e2faa 100644 --- a/studio/backend/cloudflare_tunnel.py +++ b/studio/backend/cloudflare_tunnel.py @@ -1,13 +1,13 @@ # SPDX-License-Identifier: AGPL-3.0-only # Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0 -"""Free Cloudflare quick tunnel for Studio's 0.0.0.0 launches. +"""Free Cloudflare quick tunnel for Unsloth's 0.0.0.0 launches. The raw http://: is often unreachable (https-vs-http, blocked ports, closed security groups); a cloudflared quick tunnel gives a free https://*.trycloudflare.com URL that works anywhere, with no account or domain. -Best-effort throughout: any failure collapses to "no URL" and Studio keeps +Best-effort throughout: any failure collapses to "no URL" and Unsloth keeps running. Stdlib only (back-end imports are lazy) so it is safe to import early. """ @@ -20,6 +20,7 @@ import shutil import subprocess import sys import threading +import time from pathlib import Path from typing import Optional, Tuple @@ -40,6 +41,22 @@ _RELEASE_BASE = "https://github.com/cloudflare/cloudflared/releases/latest/downl _READY_TIMEOUT = 15.0 # seconds to wait for the URL + a registered edge connection _DOWNLOAD_TIMEOUT = 60 # urlopen timeout for the one-time binary download +# A registered edge connection does not mean the hostname resolves yet, so the +# URL is fetched once before it is advertised. +_PUBLIC_PROBE_PATH = "/api/health" +_PUBLIC_PROBE_MARKER = "Unsloth UI Backend" +# One deadline for DNS propagation + the health probe, bounding the startup stall. +_PUBLIC_PROBE_TIMEOUT = 45.0 +_PUBLIC_PROBE_ATTEMPT_TIMEOUT = 5.0 +_PUBLIC_PROBE_RETRY_DELAY = 1.0 + +# Wait for the hostname via DoH first: an early OS lookup negative-caches the +# NXDOMAIN for up to 30 min. +_DNS_POLL_DELAY = 2.0 +# Retry transient DoH failures, but give up fast when DoH is blocked outright. +_DNS_MAX_DOH_ERRORS = 3 +_DOH_URL = "https://cloudflare-dns.com/dns-query?name={host}&type=A" + def _windows_hidden_kwargs() -> dict: """Suppress a child console window on Windows; no-op elsewhere.""" @@ -95,7 +112,7 @@ def _cache_path() -> Optional[Path]: def find_cloudflared() -> Optional[str]: - """Locate an existing cloudflared: PATH first, then the Studio bin cache.""" + """Locate an existing cloudflared: PATH first, then the Unsloth bin cache.""" on_path = shutil.which("cloudflared") if on_path: return on_path @@ -191,6 +208,59 @@ def ensure_cloudflared() -> Optional[str]: return None +def _wait_for_dns(host: str, deadline: float) -> None: + import json + import urllib.request + + errors = 0 + while True: + answered = False + try: + req = urllib.request.Request( + _DOH_URL.format(host = host), + headers = {"Accept": "application/dns-json", "User-Agent": "unsloth-studio"}, + ) + with urllib.request.urlopen(req, timeout = 5) as response: + answered = bool(json.loads(response.read(65536)).get("Answer")) + errors = 0 + except Exception: + errors += 1 + if errors >= _DNS_MAX_DOH_ERRORS: + return + if answered: + return + remaining = deadline - time.monotonic() + if remaining <= 0: + return + time.sleep(min(_DNS_POLL_DELAY, remaining)) + + +def verify_public_url(url: str, timeout: float = _PUBLIC_PROBE_TIMEOUT) -> bool: + import json + import urllib.request + from urllib.parse import urlsplit + + deadline = time.monotonic() + timeout + host = urlsplit(url).hostname + if host: + _wait_for_dns(host, deadline) + + probe_url = f"{url.rstrip('/')}{_PUBLIC_PROBE_PATH}" + while True: + try: + req = urllib.request.Request(probe_url, headers = {"User-Agent": "unsloth-studio"}) + with urllib.request.urlopen(req, timeout = _PUBLIC_PROBE_ATTEMPT_TIMEOUT) as response: + body = response.read(4096) + if json.loads(body).get("service") == _PUBLIC_PROBE_MARKER: + return True + except Exception: + pass + remaining = deadline - time.monotonic() + if remaining <= 0: + return False + time.sleep(min(_PUBLIC_PROBE_RETRY_DELAY, remaining)) + + class CloudflareTunnel: """A cloudflared quick tunnel to http://localhost:. Best-effort throughout. @@ -240,6 +310,7 @@ class CloudflareTunnel: stderr = subprocess.STDOUT, stdin = subprocess.DEVNULL, text = True, + encoding = "utf-8", errors = "replace", bufsize = 1, **_windows_hidden_kwargs(), @@ -309,7 +380,7 @@ class CloudflareTunnel: pass -# Single serving process per Studio launch, so one module-level tunnel handle is +# Single serving process per Unsloth launch, so one module-level tunnel handle is # enough; the lock guards the start/stop/shutdown races. _active_tunnel: Optional[CloudflareTunnel] = None _active_lock = threading.Lock() @@ -322,11 +393,12 @@ def start_studio_tunnel(port: int, timeout: float = _READY_TIMEOUT) -> Optional[ """Start a quick tunnel and return its public URL once it is actually serving, or None (best-effort). - Waits for cloudflared to both mint the URL and register an edge connection - before returning, so the caller never advertises a URL that yields Cloudflare - error 1033 (HTTP 530). If a URL is minted but no connection registers within - the window (e.g. quic is blocked on this network), retries once forcing the - http2 protocol. On any failure the tunnel is stopped and None is returned. + Waits for cloudflared to both mint the URL and register an edge connection, + then fetches /api/health over the public URL, so the caller never advertises + a link that yields Cloudflare error 1033 (HTTP 530) or an unresolvable host. + If a URL is minted but no connection registers within the window (e.g. quic + is blocked on this network), retries once forcing the http2 protocol. On any + failure the tunnel is stopped and None is returned. """ global _active_tunnel, _shutdown_requested binary = ensure_cloudflared() @@ -349,9 +421,13 @@ def start_studio_tunnel(port: int, timeout: float = _READY_TIMEOUT) -> Optional[ prior, _active_tunnel = _active_tunnel, tunnel if prior is not None: prior.stop() + registered = False try: tunnel.start() url = tunnel.wait_for_ready(timeout) + registered = url is not None + if url and not verify_public_url(url): + url = None except Exception: url = None if url: @@ -371,6 +447,9 @@ def start_studio_tunnel(port: int, timeout: float = _READY_TIMEOUT) -> Optional[ # http2 will not help, so do not burn another window on it. if not saw_url: return None + # probe failure after registering is DNS propagation; http2 would not help + if registered: + return None return None diff --git a/studio/backend/colab.py b/studio/backend/colab.py index e04543b3aa..bf4a6a44b5 100644 --- a/studio/backend/colab.py +++ b/studio/backend/colab.py @@ -1,9 +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 -""" -Colab helpers for Unsloth Studio. Uses Colab's built-in proxy. -""" +"""Colab helpers for Unsloth Studio. Uses Colab's built-in proxy.""" from pathlib import Path import sys @@ -22,11 +20,9 @@ logger = get_logger(__name__) def get_colab_url(port: int = 8888) -> str: - """ - Get the Colab proxy URL for a port. + """Get the Colab proxy URL for a port. - Retries up to 3 times, validating the result is a real HTTPS Colab URL. - Falls back to http://localhost:{port} only when all attempts fail. + Retries 3x validating a real HTTPS Colab URL; falls back to localhost on failure. """ import time as _time @@ -55,28 +51,243 @@ def get_colab_url(port: int = 8888) -> str: return fallback -def show_link(port: int = 8888, *, _url: "str | None" = None): - """Display a styled clickable link to the UI. - - *_url* is an optional pre-fetched proxy URL; pass it to avoid a second eval_js round-trip. - """ - from IPython.display import display, HTML - - url = _url if _url is not None else get_colab_url(port) - - # Truncated display URL; try/except so an odd URL shape still renders the link. +def _short_colab_url(url: str, port: int) -> str: + """Truncated display form of a Colab proxy URL; falls back to the full URL.""" try: port_prefix = f"{port}-" idx = url.index(port_prefix) next_dash = url.index("-", idx + len(port_prefix)) - short_url = url[: next_dash + 1] + "..." + return url[: next_dash + 1] + "..." except (ValueError, IndexError): - short_url = url + return url - # Plain-text line so the URL shows even if HTML display fails. - logger.info(f"🌐 Unsloth Studio URL: {url}") - html = f""" +def _is_colab_proxy_url(url: str, port: int) -> bool: + """True when *url* looks like a real Colab kernel proxy, not a localhost fallback.""" + return bool(url and isinstance(url, str) and url.startswith("https://") and str(port) in url) + + +def _is_colab_runtime() -> bool: + """True on a hosted Colab notebook kernel. + + Reuses the backend's main Colab detector (``/content`` + Colab env / ``google.colab``) + instead of a single env var, which is not always present on hosted runtimes. + """ + try: + from main import _IS_COLAB + return bool(_IS_COLAB) + except Exception: + return False + + +def _colab_login_credentials_path() -> Path: + from auth.storage import DB_PATH + return DB_PATH.parent / ".colab_notebook_login" + + +def _store_colab_login_credentials(username: str, password: str) -> None: + """Persist Colab admin credentials for notebook re-runs after interrupt.""" + path = _colab_login_credentials_path() + try: + path.parent.mkdir(parents = True, exist_ok = True) + path.write_text(f"{username}\n{password}\n", encoding = "utf-8") + try: + import os + os.chmod(path, 0o600) + except OSError: + pass + except OSError as e: + logger.info(f"Could not persist Colab login credentials ({e}).") + + +def _load_colab_login_credentials() -> "tuple[str, str] | None": + """Return stored Colab admin credentials from a previous ``start()`` run, if any.""" + path = _colab_login_credentials_path() + try: + if not path.is_file(): + return None + lines = path.read_text(encoding = "utf-8").splitlines() + if len(lines) >= 2 and lines[0] and lines[1]: + return lines[0], lines[1] + except (OSError, UnicodeDecodeError) as e: + logger.info(f"Could not load Colab login credentials ({e}).") + return None + + +def _clear_colab_login_credentials() -> None: + """Drop the cached Colab credentials once they no longer authenticate.""" + path = _colab_login_credentials_path() + try: + path.unlink(missing_ok = True) + except OSError as e: + logger.info(f"Could not clear Colab login credentials ({e}).") + + +def _colab_credentials_still_valid(username: str, password: str) -> bool: + """True when *password* still matches the stored admin hash. + + Guards against redisplaying a cached first-run password after the user has + changed the admin password through the app, which would print credentials + that no longer authenticate to the current Cloudflare tunnel. + """ + try: + from auth.storage import get_user_and_secret + from auth.hashing import verify_password + except Exception as e: + logger.info(f"Could not load auth to validate cached Colab credentials ({e}).") + return False + try: + row = get_user_and_secret(username) + if not row: + return False + salt, pwd_hash = row[0], row[1] + return bool(verify_password(password, salt, pwd_hash)) + except Exception as e: + logger.info(f"Could not validate cached Colab credentials ({e}).") + return False + + +def _colab_wants_cloudflare(cloudflare: "bool | None") -> bool: + """Resolve whether to open a Cloudflare tunnel. + + ``None`` auto-enables on real Colab (the in-cell proxy embed is often blank); + pass ``False`` to opt out. + """ + if cloudflare is not None: + return cloudflare + return _is_colab_runtime() + + +def _finalize_colab_admin_password() -> "tuple[str, str] | None": + """Clear the bootstrap-password gate on Colab so Cloudflare tunnels can start. + + Returns ``(username, password)`` for display in the notebook. On first run the + random admin password is finalized; on later runs (e.g. after interrupt) the + stored credentials are re-displayed so the Cloudflare link stays usable. + Anyone who can read this cell already controls the runtime. + """ + if not _is_colab_runtime(): + return None + try: + from auth.storage import ( + DEFAULT_ADMIN_USERNAME, + ensure_default_admin, + generate_bootstrap_password, + get_bootstrap_password, + requires_password_change, + update_password, + ) + except Exception as e: + logger.warning( + f"Could not load auth for Colab setup ({e}); Cloudflare link may be blocked." + ) + return None + + try: + ensure_default_admin() + username = DEFAULT_ADMIN_USERNAME + if not requires_password_change(username): + creds = _load_colab_login_credentials() + if creds is not None and _colab_credentials_still_valid(username, creds[1]): + return creds + # The admin password was changed through the app after the first run, + # so the cached copy is stale; drop it instead of printing dead credentials. + _clear_colab_login_credentials() + return None + password = get_bootstrap_password() or generate_bootstrap_password() + if not update_password(username, password): + logger.warning( + "Could not finalize Colab admin password; Cloudflare link may be blocked." + ) + return None + _store_colab_login_credentials(username, password) + return username, password + except Exception as e: + logger.warning( + f"Could not finalize Colab admin password ({e}); Cloudflare link may be blocked." + ) + return None + + +def _colab_login_html(username: str, password: str) -> str: + """Notebook card with Colab admin credentials (shown once after auto-finalize).""" + return f""" +
+

+ Unsloth Studio Login (Colab) +

+

+ Log in as {username} with this password. This cell is visible only in + your notebook session. +

+

+ Password: {password} +

+
+ """ + + +def _show_colab_login_credentials(username: str, password: str) -> None: + """Display Colab admin credentials in the notebook output.""" + from IPython.display import HTML, display + + logger.info(f"🔐 Unsloth Studio login — user: {username}") + display(HTML(_colab_login_html(username, password))) + + +def _ready_card_html( + url: str, + port: int, + *, + has_cloudflare_link: bool = False, + cloudflare_requested: bool = False, +) -> str: + """Branded ready card for the in-notebook Studio view. + + Colab ``*.prod.colab.dev`` proxy URLs are session-scoped and 404 when opened as a + top-level tab or on another device, so never ``window.open`` them. On real Colab the + Cloudflare link is the supported entry point because in-cell proxy embeds often stay blank. + """ + short_url = _short_colab_url(url, port) + if _is_colab_runtime() or _is_colab_proxy_url(url, port): + if has_cloudflare_link: + embed_note = ( + "Open Studio with the Cloudflare link above. In-cell proxy previews on " + "current Colab often stay blank, so the tunnel link is the supported path." + ) + elif cloudflare_requested: + embed_note = ( + "Could not open a Cloudflare tunnel, so Studio may be unreachable on Colab. " + "Check the logs above and re-run this cell. Pass " + '' + "cloudflare=True after fixing any tunnel errors." + ) + else: + embed_note = ( + "Colab proxy links cannot be opened in a new tab (they 404 outside this " + 'notebook). Re-run with start(cloudflare=True) for a working link.' + ) + return f""" +
+

+ + Unsloth Studio is Ready! +

+

+ {embed_note} +

+

+ {short_url} +

+
+ """ + + return f"""

""" - display(HTML(html)) + + +def show_link( + port: int = 8888, + *, + _url: "str | None" = None, + has_cloudflare_link: bool = False, + cloudflare_requested: bool = False, +): + """Display a styled ready card for the UI. + + Colab proxy URLs are informational only (no new-tab open; they 404 outside the cell); + non-proxy URLs keep a clickable open button. *_url* is an optional pre-fetched proxy + URL to avoid a second eval_js round-trip. + """ + from IPython.display import display, HTML + + url = _url if _url is not None else get_colab_url(port) + logger.info(f"🌐 Unsloth Studio URL: {url}") + display( + HTML( + _ready_card_html( + url, + port, + has_cloudflare_link = has_cloudflare_link, + cloudflare_requested = cloudflare_requested, + ) + ) + ) + + +def _warn_colab_cloudflare_missing(*, use_cloudflare: bool, cloudflare_url: "str | None") -> None: + """Log a prominent warning when Colab expected a tunnel but none was opened.""" + if not use_cloudflare or cloudflare_url or not _is_colab_runtime(): + return + logger.warning( + "Colab Cloudflare tunnel unavailable — Studio is unlikely to be reachable in this " + "notebook. Check the logs above for tunnel or auth errors, then re-run start()." + ) def _bootstrap_password_pending() -> bool: """True while the default admin still owes a bootstrap-password change. - While pending, main.py injects that password into same-origin GETs, and a public - tunnel GET (no Origin) reads as same-origin, so sharing the link would leak admin - access. Fails safe to pending if the state cannot be read. + While pending, a public tunnel GET (no Origin) reads as same-origin and gets the + injected password, so sharing the link would leak admin access. Fails safe to pending. """ try: from auth.storage import requires_password_change, DEFAULT_ADMIN_USERNAME @@ -121,15 +369,14 @@ def _bootstrap_password_pending() -> bool: def start_cloudflare_tunnel(port: int) -> "str | None": """Open a shareable Cloudflare quick tunnel to localhost:*port*, or None. - run_server suppresses the tunnel on Colab by design, so we start it directly. - Refused while the bootstrap password is pending; any failure collapses to None - and the Colab proxy still works. + run_server suppresses the tunnel on Colab, so we start it directly. Refused while the + bootstrap password is pending; any failure collapses to None (Colab proxy still works). """ if _bootstrap_password_pending(): logger.warning( "Cloudflare link not started: the admin account still has its temporary " "bootstrap password, which is exposed to anyone who can load the page. " - "Open Studio in this tab, log in and change the admin password, then re-run " + "Open Unsloth in this tab, log in and change the admin password, then re-run " "start(cloudflare=True) to get the shareable link." ) return None @@ -152,9 +399,9 @@ def start_cloudflare_tunnel(port: int) -> "str | None": def _publish_cloudflare_url(cloudflare_url: "str | None") -> None: """Publish a directly-started tunnel URL onto app.state so /api/health advertises it. - run_server only sets this when it opens the tunnel itself, which it skips on Colab, - so we set it here. Otherwise the frontend's API examples fall back to an - unreachable server_url. Best-effort. + run_server sets this only when it opens the tunnel itself (skipped on Colab), so we + set it here; otherwise the frontend's API examples fall back to an unreachable + server_url. Best-effort. """ if not cloudflare_url: return @@ -183,8 +430,7 @@ def _stop_cloudflare_tunnel() -> None: def _is_studio_healthy(port: int, timeout: float = 2.0) -> bool: """True only if Unsloth Studio (not some other app) answers /api/health on *port*. - The service-marker check stops the reuse path reusing or tunneling a foreign - process that merely serves /api/health. + The service-marker check stops the reuse path reusing or tunneling a foreign process. """ import json, urllib.request try: @@ -194,8 +440,29 @@ def _is_studio_healthy(port: int, timeout: float = 2.0) -> bool: return False -def _shareable_link_html(cloudflare_url: str) -> str: - """Branded card for the shareable Cloudflare link, styled like the show_link banner.""" +def _shareable_link_html( + cloudflare_url: str, + password: "str | None" = None, + username: "str | None" = None, +) -> str: + """Branded card for the shareable Cloudflare link, styled like the show_link banner. + + *password* renders under the link so the credential sits in the card with the button + it unlocks. The username is always the default admin, so it reads inline. + """ + login_block = "" + if password: + login_block = f""" +

+ Password +

+

{password}

+

+ Log in as {username} with this password. Shown only in your + notebook session, and never included in the shared link. +

""" return f"""
@@ -203,7 +470,7 @@ def _shareable_link_html(cloudflare_url: str) -> str: display: flex; align-items: center; gap: 12px;"> - Shareable Studio Link is Ready! + Shareable Unsloth Link is Ready! - This Cloudflare HTTPS link works from any device — share it with anyone. The Colab view below only works in this tab. + This Cloudflare HTTPS link works from any device, so you can share it with anyone.

- 🔗 {cloudflare_url} -

+ 🔗
{cloudflare_url} +

{login_block}
""" -def _show_and_embed(port: int, *, cloudflare_url: "str | None" = None): - """Render the Studio header + iframe for *port*, with a shareable-link card above - when *cloudflare_url* is set. Falls back to serve_kernel_port_as_iframe.""" - url = get_colab_url(port) - logger.info(f"🌐 Unsloth Studio URL: {url}") - if cloudflare_url: - logger.info(f"🔗 Shareable Cloudflare link: {cloudflare_url}") +# Height for serve_kernel_port_as_iframe (~82vh on a 1080p screen, clamped). +_COLAB_IFRAME_HEIGHT = 900 + +def _embed_kernel_port_iframe(port: int) -> bool: + """Embed Studio via Colab's native kernel-port iframe helper. + + Only trusted on a real Colab runtime: colabtools can import ``google.colab`` and + queue browser-side JS without appending an iframe, so callers outside Colab must use + the HTML iframe path instead. + """ + if not _is_colab_runtime(): + return False + try: + from google.colab import output as colab_output + except ImportError: + return False + try: + colab_output.serve_kernel_port_as_iframe( + port, + height = _COLAB_IFRAME_HEIGHT, + width = "100%", + ) + return True + except Exception as e: + logger.info(f"serve_kernel_port_as_iframe failed ({e}); trying HTML iframe.") + return False + + +def _embed_html_iframe(url: str, port: int) -> bool: + """Fallback embed: raw HTML iframe when the Colab helper is unavailable.""" try: from IPython.display import HTML, display + except ImportError: + return False - iframe_id = f"unsloth-studio-{port}" - - # Truncated header URL — best-effort, falls back to full URL. - try: - port_prefix = f"{port}-" - idx = url.index(port_prefix) - next_dash = url.index("-", idx + len(port_prefix)) - short_url = url[: next_dash + 1] + "..." - except (ValueError, IndexError): - short_url = url - - if cloudflare_url: - display(HTML(_shareable_link_html(cloudflare_url))) - + short_url = _short_colab_url(url, port) + iframe_id = f"unsloth-studio-{port}" + try: display( HTML(f"""
""") ) - except Exception: - # Fallback: Colab's built-in helper. + return True + except Exception as e: + logger.info(f"HTML iframe embed failed ({e}).") + return False + + +def _show_and_embed( + port: int, + *, + cloudflare_url: "str | None" = None, + colab_login: "tuple[str, str] | None" = None, + cloudflare_requested: bool = False, +): + """Render the Unsloth ready card + iframe for *port*. + + Prefer Colab's ``serve_kernel_port_as_iframe`` on real Colab; raw HTML iframe is the + fallback. Cloudflare cards stay clickable. + """ + url = get_colab_url(port) + logger.info(f"🌐 Unsloth Studio URL: {url}") + if cloudflare_url: + logger.info(f"🔗 Shareable Cloudflare link: {cloudflare_url}") + + _warn_colab_cloudflare_missing( + use_cloudflare = cloudflare_requested, + cloudflare_url = cloudflare_url, + ) + + # Fold the credentials into the link card rather than a second card below it. + credentials_shown = False + if cloudflare_url: try: - from google.colab import output as colab_output - colab_output.serve_kernel_port_as_iframe(port, height = 900, width = "100%") - except ImportError: - pass + from IPython.display import HTML, display + + username, password = colab_login if colab_login else (None, None) + display(HTML(_shareable_link_html(cloudflare_url, password, username))) + credentials_shown = bool(colab_login) + except Exception as e: + logger.info(f"Could not render Cloudflare link card ({e}).") + + if colab_login and not credentials_shown: + try: + _show_colab_login_credentials(*colab_login) + except Exception as e: + logger.info(f"Could not render Colab login card ({e}).") + + # With a tunnel up the embed below is skipped, so the ready card would only restate + # the link card and print a proxy URL that 404s outside this tab. + skip_ready_card = _is_colab_runtime() and bool(cloudflare_url) + if not skip_ready_card: + try: + show_link( + port, + _url = url, + has_cloudflare_link = bool(cloudflare_url), + cloudflare_requested = cloudflare_requested, + ) + except Exception as e: + logger.info(f"Could not render Unsloth link card ({e}).") + + # On Colab with a working tunnel, skip the in-cell proxy embed (often blank). + if _is_colab_runtime() and cloudflare_url: + return + + # Real Colab: kernel helper needs only the port (works when eval_js failed). + if _is_colab_runtime(): + if _embed_kernel_port_iframe(port): + return + _embed_html_iframe(url, port) -def start(port: int = 8888, *, cloudflare: bool = False): +def start(port: int = 8888, *, cloudflare: "bool | None" = None): """Start Unsloth Studio in Colab and display the URL. Args: port: Port to bind/serve on. - cloudflare: Opt in to a shareable Cloudflare HTTPS link reachable from any - device (default OFF). It exposes Studio's login page beyond Colab, so it - stays an explicit opt-in; the default shows only the in-tab proxy iframe. + cloudflare: Shareable Cloudflare HTTPS link. ``None`` (default) auto-enables on + real Colab because the in-cell proxy embed is often blank; pass ``False`` to + skip the tunnel or ``True`` to force it on other runtimes. Usage: - start() # Colab-proxy iframe only (default) - start(cloudflare=True) # also open a shareable Cloudflare link + start() # Cloudflare link on Colab (auto); proxy iframe elsewhere + start(cloudflare=False) # Colab proxy iframe only (often blank on current Colab) + start(cloudflare=True) # force Cloudflare link on any runtime """ import time logger.info("🦥 Starting Unsloth Studio...") + use_cloudflare = _colab_wants_cloudflare(cloudflare) - # Fast path: Studio already running (cell re-run). Re-launching would collide on - # the port, so just re-show the link and iframe. + # Fast path: already running (cell re-run); re-show link/iframe instead of rebinding the port. if _is_studio_healthy(port): - logger.info(f" Studio is already running on port {port} — reusing existing server.") + logger.info(f" Unsloth is already running on port {port} — reusing existing server.") # try/finally: tear the tunnel down even if interrupted mid-start/render. try: - cf_url = start_cloudflare_tunnel(port) if cloudflare else None + colab_login = _finalize_colab_admin_password() if use_cloudflare else None + cf_url = start_cloudflare_tunnel(port) if use_cloudflare else None _publish_cloudflare_url(cf_url) - _show_and_embed(port, cloudflare_url = cf_url) + _show_and_embed( + port, + cloudflare_url = cf_url, + colab_login = colab_login, + cloudflare_requested = use_cloudflare, + ) for _ in range(10000): time.sleep(300) print("=", end = "", flush = True) @@ -313,7 +664,6 @@ def start(port: int = 8888, *, cloudflare: bool = False): logger.info(" Loading backend...") from run import run_server - # Auto-detect frontend path repo_root = Path(__file__).parent.parent frontend_path = repo_root / "frontend" / "dist" @@ -323,8 +673,7 @@ def start(port: int = 8888, *, cloudflare: bool = False): logger.info(" Starting server...") try: - # cloudflare=False: this helper owns the tunnel (Colab's own - # start(cloudflare=...) drives it), so pin it off explicitly. + # cloudflare=False: this helper owns the tunnel (via start(cloudflare=...)), so pin it off. app = run_server( host = "0.0.0.0", port = port, @@ -339,14 +688,12 @@ def start(port: int = 8888, *, cloudflare: bool = False): logger.error(f"❌ Unsloth Studio failed to start: {exc}") return - # run_server auto-increments the port if in use; read back the bound port so the - # proxy URL and iframe point at the right place. + # run_server may auto-increment the port; read back the bound port for the proxy URL/iframe. actual_port: int = getattr(getattr(app, "state", None), "server_port", None) or port logger.info(f" Server started on port {actual_port}!") - # Poll health endpoint before showing the link — avoids the race where ready_event - # fires but the process hasn't finished binding. + # Poll health before showing the link: avoids the race where ready_event fires pre-bind. import urllib.request server_ready = False @@ -365,12 +712,17 @@ def start(port: int = 8888, *, cloudflare: bool = False): ) return - # Open the tunnel now the server is healthy, publish its URL for /api/health, and - # tear it down on interrupt (try/finally) rather than orphan the process. + # Server healthy: finalize Colab auth, open the tunnel, publish URL, tear down on interrupt. try: - cf_url = start_cloudflare_tunnel(actual_port) if cloudflare else None + colab_login = _finalize_colab_admin_password() if use_cloudflare else None + cf_url = start_cloudflare_tunnel(actual_port) if use_cloudflare else None _publish_cloudflare_url(cf_url) - _show_and_embed(actual_port, cloudflare_url = cf_url) + _show_and_embed( + actual_port, + cloudflare_url = cf_url, + colab_login = colab_login, + cloudflare_requested = use_cloudflare, + ) # Keep kernel alive so the daemon server thread runs. for _ in range(10000): diff --git a/studio/backend/core/data_recipe/jobs/manager.py b/studio/backend/core/data_recipe/jobs/manager.py index 0e0044702e..135c9fccf6 100644 --- a/studio/backend/core/data_recipe/jobs/manager.py +++ b/studio/backend/core/data_recipe/jobs/manager.py @@ -27,7 +27,6 @@ from .constants import ( ) from .parse import apply_update, coerce_event, parse_log_message from .types import Job -from .worker import run_job_process from loggers import get_logger logger = get_logger(__name__) @@ -169,12 +168,18 @@ class JobManager: native_path_secret_removed_for_child_start, run_without_native_path_secret, ) + from utils.hf_cache_settings import child_environment_for_spawn, get_hf_cache_paths - with native_path_secret_removed_for_child_start(): + cache_env = get_hf_cache_paths().child_env({}) + + with ( + child_environment_for_spawn(cache_env), + native_path_secret_removed_for_child_start(), + ): mp_q = _CTX.Queue() proc = _CTX.Process( target = run_without_native_path_secret, - args = (run_job_process,), + args = ("core.data_recipe.jobs.worker", "run_job_process", cache_env), kwargs = {"event_queue": mp_q, "recipe": recipe, "run": run_payload}, daemon = True, ) diff --git a/studio/backend/core/data_recipe/jobs/parse.py b/studio/backend/core/data_recipe/jobs/parse.py index 3be830d0e4..8c2e8a4d55 100644 --- a/studio/backend/core/data_recipe/jobs/parse.py +++ b/studio/backend/core/data_recipe/jobs/parse.py @@ -133,7 +133,7 @@ def parse_log_message(msg: str) -> ParsedUpdate | None: source = "github", status = "rate_limited", retry_after_sec = seconds, - message = ("Waiting for GitHub rate limit. Studio will resume automatically."), + message = ("Waiting for GitHub rate limit. Unsloth will resume automatically."), ), ) @@ -147,7 +147,7 @@ def parse_log_message(msg: str) -> ParsedUpdate | None: status = "rate_limited", retry_after_sec = seconds, message = ( - "Waiting for GitHub secondary rate limit. Studio will resume automatically." + "Waiting for GitHub secondary rate limit. Unsloth will resume automatically." ), ), ) @@ -161,7 +161,7 @@ def parse_log_message(msg: str) -> ParsedUpdate | None: source = "github", status = "rate_limited", retry_after_sec = seconds, - message = ("Waiting for GitHub rate limit. Studio will resume automatically."), + message = ("Waiting for GitHub rate limit. Unsloth will resume automatically."), ), ) diff --git a/studio/backend/core/data_recipe/local_callable_validators.py b/studio/backend/core/data_recipe/local_callable_validators.py index ebb1d39dfb..143895d781 100644 --- a/studio/backend/core/data_recipe/local_callable_validators.py +++ b/studio/backend/core/data_recipe/local_callable_validators.py @@ -238,7 +238,7 @@ def _run_oxc_batch( if not node_executable: return _fallback_results( len(code_values), - "Node.js not found (install Node >= 20.19, or re-run Studio setup to provision it).", + "Node.js not found (install Node >= 20.19, or re-run Unsloth setup to provision it).", ) try: tmp_dir = ensure_dir(oxc_validator_tmp_root()) @@ -257,6 +257,8 @@ def _run_oxc_batch( cwd = str(_OXC_TOOL_DIR), input = json.dumps(payload), text = True, + encoding = "utf-8", + errors = "replace", capture_output = True, check = False, env = env, diff --git a/studio/backend/core/data_recipe/service.py b/studio/backend/core/data_recipe/service.py index 4647dc098d..9770e88b7f 100644 --- a/studio/backend/core/data_recipe/service.py +++ b/studio/backend/core/data_recipe/service.py @@ -280,8 +280,8 @@ def create_data_designer(recipe: dict[str, Any], *, artifact_path: str | None = from data_designer.interface.data_designer import DataDesigner # pyright: ignore[reportMissingImports] if artifact_path is None: - # DataDesigner defaults to cwd/artifacts; packaged Studio can run with - # cwd=/, so keep default callers on Studio's writable recipe artifact root. + # DataDesigner defaults to cwd/artifacts; packaged Unsloth can run with + # cwd=/, so keep default callers on Unsloth's writable recipe artifact root. artifact_path = str(recipe_datasets_root()) recipe = _strip_frontend_model_config_metadata(recipe) diff --git a/studio/backend/core/export/export.py b/studio/backend/core/export/export.py index c8be50b08b..4979ebd48d 100644 --- a/studio/backend/core/export/export.py +++ b/studio/backend/core/export/export.py @@ -81,6 +81,82 @@ _PYTORCH_MISSING_MESSAGE = ( _LLAMA_CPP_SCRIPTS_WARNING_EMITTED = False +def _multi_gpu_device_map_kwargs() -> dict: + """``device_map`` kwargs for sharding a checkpoint across every visible GPU. + + unsloth's ``from_pretrained`` defaults to ``device_map="sequential"``, which stacks + the whole model on GPU0 and OOMs multi-GPU hosts whose other GPUs sit empty (#7053). + Returns ``{"device_map": "balanced"}`` only on a real multi-GPU CUDA/ROCm host + (mirroring the inference loader's ``get_device_map``), else empty so single-GPU, CPU + and MLX loads keep the loader default.""" + if _IS_MLX: + return {} + try: + from utils.hardware import get_device_map, get_parent_visible_gpu_ids + + visible = get_parent_visible_gpu_ids() + if len(visible) > 1: + device_map = get_device_map(visible) + elif not visible: + # UUID/MIG masks resolve to no numeric ids; get_device_map(None) falls back + # to the visible-GPU count, so a multi-GPU UUID/MIG host still shards. + device_map = get_device_map(None) + else: + return {} + if device_map == "balanced": + return {"device_map": device_map} + except Exception as exc: + logger.debug(f"multi-GPU device_map resolution failed; using loader default: {exc}") + return {} + + +def _is_oom_error(exc: BaseException) -> bool: + """True for an accelerator OOM, however it is spelled. + + accelerate and transformers re-raise it as a plain ``RuntimeError`` on several paths + and ROCm/XPU use their own classes, so match the message too. + """ + if torch is not None: + oom_types = tuple( + t + for t in ( + getattr(torch, "OutOfMemoryError", None), + getattr(getattr(torch, "cuda", None), "OutOfMemoryError", None), + getattr(getattr(torch, "xpu", None), "OutOfMemoryError", None), + ) + if isinstance(t, type) + ) + if oom_types and isinstance(exc, oom_types): + return True + return "out of memory" in f"{type(exc).__name__}: {exc}".lower() + + +def _is_cpu_spill_rejection(exc: BaseException) -> bool: + """bitsandbytes refuses a map that spills to CPU/disk with a plain ``ValueError``. + + Busy secondary GPUs can make ``balanced`` spill to CPU even where the old sequential + load fit on GPU0, and that message says nothing about memory, so the retry has to + match it explicitly. See transformers ``quantizers/quantizer_bnb_4bit.py``. + """ + return "dispatched on the cpu or the disk" in str(exc).lower() + + +class _CpuSpillRetry(Exception): + """A multi-GPU load that succeeded but left modules offloaded to CPU/disk.""" + + +def _cpu_offloaded_modules(model) -> int: + """Count the modules a load parked on CPU or disk. + + Only bitsandbytes refuses such a map; a full-precision load accepts it, leaves the + parameters on meta and dies much later in safetensors with "Cannot copy out of meta + tensor". Nothing raises at load time, so inspect the map directly. PEFT re-dispatches + when attaching an adapter, so in practice this catches merged checkpoints. + """ + device_map = getattr(model, "hf_device_map", None) or {} + return sum(1 for target in device_map.values() if str(target) in ("cpu", "disk")) + + def _supports_kwarg(fn, name): """True if `fn` accepts keyword `name` directly or via **kwargs.""" import inspect @@ -165,7 +241,7 @@ def _offline_window_if(local_files_only): def _is_wsl(): """Detect if running under Windows Subsystem for Linux.""" try: - return "microsoft" in open("/proc/version").read().lower() + return "microsoft" in open("/proc/version", encoding = "utf-8").read().lower() except Exception: return False @@ -271,6 +347,7 @@ class ExportBackend: load_in_4bit: bool = True, trust_remote_code: bool = False, hf_token: Optional[str] = None, + _device_map_override: Optional[dict] = None, ) -> Tuple[bool, str]: """ Load a checkpoint for export. @@ -303,6 +380,14 @@ class ExportBackend: # Skip the Hub when offline so a no-internet export uses the local cache. local_files_only = _hf_offline() + # Shard across every visible GPU instead of stacking on GPU0 (#7053); {} on + # single-GPU/CPU/MLX. _device_map_override is the single-device retry below. + _device_map_kw = ( + _multi_gpu_device_map_kwargs() + if _device_map_override is None + else _device_map_override + ) + # Run the type-detection probes in the forced-offline window (else a gated # base 404s); it covers is_vision_model's Hub reads + the transformers-5 # subprocess, and local_files_only makes detect_audio_type's requests.get skip. @@ -328,6 +413,7 @@ class ExportBackend: trust_remote_code = trust_remote_code, token = token, local_files_only = local_files_only, + **_device_map_kw, ) elif self._audio_type == "whisper": @@ -343,6 +429,7 @@ class ExportBackend: trust_remote_code = trust_remote_code, token = token, local_files_only = local_files_only, + **_device_map_kw, ) elif self._audio_type == "snac": @@ -355,6 +442,7 @@ class ExportBackend: trust_remote_code = trust_remote_code, token = token, local_files_only = local_files_only, + **_device_map_kw, ) elif self._audio_type == "bicodec": @@ -368,6 +456,7 @@ class ExportBackend: trust_remote_code = trust_remote_code, token = token, local_files_only = local_files_only, + **_device_map_kw, ) elif self._audio_type == "dac": @@ -380,6 +469,7 @@ class ExportBackend: trust_remote_code = trust_remote_code, token = token, local_files_only = local_files_only, + **_device_map_kw, ) elif self.is_vision: @@ -392,6 +482,7 @@ class ExportBackend: trust_remote_code = trust_remote_code, token = token, local_files_only = local_files_only, + **_device_map_kw, ) tokenizer = processor # vision: processor acts as tokenizer @@ -405,8 +496,16 @@ class ExportBackend: trust_remote_code = trust_remote_code, token = token, local_files_only = local_files_only, + **_device_map_kw, ) + # Only when we asked for the multi-GPU map: a single-GPU host has no second + # placement to retry on, so leave its behaviour untouched. + _offloaded = _cpu_offloaded_modules(model) if _device_map_kw else 0 + if _device_map_override is None and _offloaded: + del model + raise _CpuSpillRetry(f"{_offloaded} module(s) offloaded to CPU/disk") + if _IS_MLX: # MLX doesn't use PeftModel — detect LoRA via adapter_config.json self.is_peft = adapter_config.exists() @@ -429,11 +528,41 @@ class ExportBackend: return True, f"Loaded {model_type} model{peft_info} successfully" except Exception as e: - logger.error(f"Error loading checkpoint: {e}") - import traceback + # Sharding is an optimisation, never a requirement. "balanced" budgets from the + # free memory read BEFORE this process opens a CUDA context on each GPU, so when + # a training or chat job already owns the others the shard can OOM, or spill to + # CPU and be refused by bitsandbytes, where the old single-device load succeeded. + # Fall back once before giving up. + if ( + _device_map_override is None + and ( + isinstance(e, _CpuSpillRetry) or _is_oom_error(e) or _is_cpu_spill_rejection(e) + ) + and _multi_gpu_device_map_kwargs() + ): + # Retry outside this block: the live traceback pins the half-built model's + # frames, so an in-block retry inherits the exhausted device. + retry_reason = str(e) + else: + logger.error(f"Error loading checkpoint: {e}") + import traceback - logger.error(traceback.format_exc()) - return False, f"Failed to load checkpoint: {str(e)}" + logger.error(traceback.format_exc()) + return False, f"Failed to load checkpoint: {str(e)}" + + logger.warning( + f"Multi-GPU export load unusable ({retry_reason}); retrying on " + f"the single-device loader default." + ) + self.cleanup_memory() + return self.load_checkpoint( + checkpoint_path, + max_seq_length = max_seq_length, + load_in_4bit = load_in_4bit, + trust_remote_code = trust_remote_code, + hf_token = hf_token, + _device_map_override = {}, + ) def _write_export_metadata(self, save_directory: str): """Write export_metadata.json with base model info for Chat page discovery.""" @@ -445,7 +574,7 @@ class ExportBackend: ) metadata = {"base_model": base_model} metadata_path = os.path.join(save_directory, "export_metadata.json") - with open(metadata_path, "w") as f: + with open(metadata_path, "w", encoding = "utf-8") as f: json.dump(metadata, f, indent = 2) logger.info(f"Wrote export metadata to {metadata_path}") except Exception as e: @@ -1048,6 +1177,21 @@ class ExportBackend: "Use the safetensors adapter instead.", None, ) + # llama.cpp's convert_lora_to_gguf.py has no concept of DoRA's + # lora_magnitude_vector tensors: it only reads the standard + # lora_A/lora_B delta, so exporting a DoRA adapter would silently + # drop the magnitude rescaling and produce a GGUF LoRA file that + # loads fine but no longer matches the trained model. + _peft_config = getattr(self.current_model, "peft_config", {}).get("default") + if getattr(_peft_config, "use_dora", False): + return ( + False, + "GGUF LoRA export is not supported for DoRA adapters: the GGUF LoRA " + "format has no way to represent DoRA's magnitude vectors, so the " + "exported file would silently lose the DoRA behavior. Use the " + "safetensors adapter instead, or merge to a full GGUF model.", + None, + ) outtype = str(gguf_outtype).lower() if outtype not in _GGUF_LORA_OUTTYPES: return ( diff --git a/studio/backend/core/export/orchestrator.py b/studio/backend/core/export/orchestrator.py index 6d1a928f2e..aaf48615f0 100644 --- a/studio/backend/core/export/orchestrator.py +++ b/studio/backend/core/export/orchestrator.py @@ -230,16 +230,20 @@ class ExportOrchestrator: native_path_secret_removed_for_child_start, run_without_native_path_secret, ) + from utils.hf_cache_settings import child_environment_for_spawn, get_hf_cache_paths - from .worker import run_export_process + cache_env = get_hf_cache_paths().child_env({}) - with native_path_secret_removed_for_child_start(): + with ( + child_environment_for_spawn(cache_env), + native_path_secret_removed_for_child_start(), + ): self._cmd_queue = _CTX.Queue() self._resp_queue = _CTX.Queue() self._proc = _CTX.Process( target = run_without_native_path_secret, - args = (run_export_process,), + args = ("core.export.worker", "run_export_process", cache_env), kwargs = { "cmd_queue": self._cmd_queue, "resp_queue": self._resp_queue, diff --git a/studio/backend/core/inference/__init__.py b/studio/backend/core/inference/__init__.py index ad78157418..1491dfa749 100644 --- a/studio/backend/core/inference/__init__.py +++ b/studio/backend/core/inference/__init__.py @@ -11,7 +11,7 @@ subprocess and can be imported directly from .inference when needed. Public names are resolved lazily (PEP 562): importing this package -- or a dependency-light leaf like ``core.inference.chat_eos`` -- must NOT eagerly pull the orchestrator / llama_cpp import chain (httpx, subprocess plumbing, the ML -backend and its Studio dependencies). Those load only when a public name is +backend and its Unsloth dependencies). Those load only when a public name is actually accessed, so standalone helpers stay unit-testable without the full inference stack. """ diff --git a/studio/backend/core/inference/_vulkan_probe.py b/studio/backend/core/inference/_vulkan_probe.py index 706346daad..4bfefc21ce 100644 --- a/studio/backend/core/inference/_vulkan_probe.py +++ b/studio/backend/core/inference/_vulkan_probe.py @@ -6,12 +6,14 @@ Run in a short-lived subprocess (``python _vulkan_probe.py ``) so the Vulkan instance never lives in the long-running backend process. Loads the bundled ggml Vulkan backend from ```` and prints one -``\\t\\t\\t`` line per device to stdout. -Indices are ggml's own Vulkan device ordinals, which need not match nvidia-smi -order. ``is_igpu`` (from ggml's device type) is ``1`` for an integrated GPU -sharing system RAM. ``total_bytes`` is the device-local heap; the reader uses -it to reserve absolute headroom on a discrete card (parity with the CUDA/ROCm -fit) and ignores it for an iGPU, whose "VRAM" is shared system RAM. +``\\t\\t\\t\\t`` line per device to +stdout. Indices are ggml's own Vulkan device ordinals, which need not match +nvidia-smi order. ``is_igpu`` (from ggml's device type) is ``1`` for an +integrated GPU sharing system RAM. ``total_bytes`` is the device-local heap; +the reader uses it to reserve absolute headroom on a discrete card (parity +with the CUDA/ROCm fit) and ignores it for an iGPU, whose "VRAM" is shared +system RAM. ``name`` is ggml's device description (the marketing name, e.g. +"AMD Radeon RX 9070 XT"); empty when the registry lookup fails. Uses only the standard library so it stays runnable as a bare script. """ @@ -24,15 +26,30 @@ import sys _GGML_BACKEND_DEVICE_TYPE_IGPU = 2 -def _igpu_flags(base, lib, count: int) -> list[bool]: - """Per-device integrated-GPU flags via ggml's backend registry. +def _igpu_flags_and_names(base, lib, count: int) -> tuple[list[bool], list[str]]: + """Per-device integrated-GPU flags and descriptions via ggml's backend registry. The Vulkan reg enumerates devices in the same order as ``ggml_backend_vk_get_device_memory`` (each context uses ``ctx->device = - i``), so reg index == device ordinal. Returns all-False on any failure so - the reader never over-caps a discrete card. + i``), so reg index == device ordinal. Returns all-False / empty-name on any + failure so the reader never over-caps a discrete card and the memory + readings still get through. """ flags = [False] * count + names = [""] * count + + # The name lookup is bound OUTSIDE the type-detection try: a ggml-base + # without ggml_backend_dev_description (older/custom build) must degrade to + # unnamed devices, not abort before the iGPU flags are read (which would + # count an iGPU's shared RAM as VRAM). + describe = None + try: + base.ggml_backend_dev_description.restype = ctypes.c_char_p + base.ggml_backend_dev_description.argtypes = [ctypes.c_void_p] + describe = base.ggml_backend_dev_description + except Exception: + pass + try: lib.ggml_backend_vk_reg.restype = ctypes.c_void_p lib.ggml_backend_vk_reg.argtypes = [] @@ -45,17 +62,31 @@ def _igpu_flags(base, lib, count: int) -> list[bool]: reg = lib.ggml_backend_vk_reg() if not reg: - return flags + return flags, names dev_count = base.ggml_backend_reg_dev_count(reg) for i in range(min(count, dev_count)): dev = base.ggml_backend_reg_dev_get(reg, i) if dev: flags[i] = base.ggml_backend_dev_type(dev) == _GGML_BACKEND_DEVICE_TYPE_IGPU + if describe is not None: + try: + desc = describe(dev) + if desc: + # Tabs/newlines would corrupt the line protocol; + # spaces are safe. + names[i] = ( + desc.decode("utf-8", errors = "replace") + .replace("\t", " ") + .replace("\n", " ") + .strip() + ) + except Exception: + pass except Exception: - # Best-effort: any failure degrades to "discrete" so the memory - # readings still get through instead of crashing the probe. + # Best-effort: any failure degrades to "discrete"/"unnamed" so the + # memory readings still get through instead of crashing the probe. pass - return flags + return flags, names def main() -> int: @@ -63,6 +94,14 @@ def main() -> int: return 0 bindir = sys.argv[1] + # Device names can be non-ASCII (localized drivers); the platform-default + # stdout encoding (e.g. cp1252) would raise on them and lose the whole + # inventory. The reader decodes UTF-8 with the same error mode. + try: + sys.stdout.reconfigure(encoding = "utf-8", errors = "replace") + except Exception: + pass + # Hold add_dll_directory's handle for the rest of main() (the documented # idiom) so bindir stays on the search path while the sibling ggml DLLs # resolve below. @@ -96,12 +135,12 @@ def main() -> int: ] count = lib.ggml_backend_vk_get_device_count() - igpu = _igpu_flags(base, lib, count) + igpu, names = _igpu_flags_and_names(base, lib, count) rows = [] for i in range(count): free, total = ctypes.c_size_t(0), ctypes.c_size_t(0) lib.ggml_backend_vk_get_device_memory(i, ctypes.byref(free), ctypes.byref(total)) - rows.append("%d\t%d\t%d\t%d" % (i, free.value, int(igpu[i]), total.value)) + rows.append("%d\t%d\t%d\t%d\t%s" % (i, free.value, int(igpu[i]), total.value, names[i])) sys.stdout.write("\n".join(rows)) return 0 diff --git a/studio/backend/core/inference/anthropic_compat.py b/studio/backend/core/inference/anthropic_compat.py index 3c7a4cb182..a32e372d73 100644 --- a/studio/backend/core/inference/anthropic_compat.py +++ b/studio/backend/core/inference/anthropic_compat.py @@ -172,6 +172,136 @@ def anthropic_messages_to_openai( return result +_ANTHROPIC_SCHEMA_CLIENT_TOOL_PARAMETERS = { + "bash": { + "type": "object", + "properties": { + "command": {"type": "string"}, + "restart": {"type": "boolean"}, + }, + "anyOf": [ + {"required": ["command"]}, + {"properties": {"restart": {"const": True}}, "required": ["restart"]}, + ], + }, + "text_editor": { + "type": "object", + "properties": { + "command": { + "type": "string", + "enum": ["view", "str_replace", "create", "insert"], + }, + "path": {"type": "string"}, + "view_range": { + "type": "array", + "items": {"type": "integer"}, + "minItems": 2, + "maxItems": 2, + }, + "old_str": {"type": "string"}, + "new_str": {"type": "string"}, + "file_text": {"type": "string"}, + "insert_line": {"type": "integer"}, + "insert_text": {"type": "string"}, + }, + "required": ["command", "path"], + }, + "computer": { + "type": "object", + "properties": { + "action": {"type": "string"}, + "coordinate": { + "type": "array", + "items": {"type": "integer"}, + "minItems": 2, + "maxItems": 2, + }, + "text": {"type": "string"}, + "duration": {"type": "number"}, + "scroll_direction": {"type": "string"}, + "scroll_amount": {"type": "integer"}, + "start_coordinate": { + "type": "array", + "items": {"type": "integer"}, + "minItems": 2, + "maxItems": 2, + }, + "key": {"type": "string"}, + }, + "required": ["action"], + "additionalProperties": True, + }, + "memory": { + "type": "object", + "properties": { + "command": { + "type": "string", + "enum": ["view", "create", "str_replace", "insert", "delete", "rename"], + }, + "path": {"type": "string"}, + "view_range": { + "type": "array", + "items": {"type": "integer"}, + "minItems": 2, + "maxItems": 2, + }, + "file_text": {"type": "string"}, + "old_str": {"type": "string"}, + "new_str": {"type": "string"}, + "insert_line": {"type": "integer"}, + "insert_text": {"type": "string"}, + "old_path": {"type": "string"}, + "new_path": {"type": "string"}, + }, + "required": ["command"], + }, +} + +_ANTHROPIC_SCHEMA_CLIENT_TOOL_DESCRIPTIONS = { + "bash": "Run a command in the caller-owned persistent bash session, or restart it.", + "text_editor": "View, create, or edit files in the caller-owned filesystem.", + "computer": "Interact with the caller-owned computer using an action and its parameters.", + "memory": "Store and retrieve files in the caller-owned persistent memory directory.", +} + + +def anthropic_schema_client_tool_kind(tool) -> Optional[str]: + """Return the kind of a schema-less Anthropic client tool, if recognized.""" + td = tool if isinstance(tool, dict) else tool.model_dump() + if td.get("input_schema") is not None: + return None + type_ = td.get("type") + if not isinstance(type_, str): + return None + kind, separator, version = type_.rpartition("_") + if ( + separator + and kind in _ANTHROPIC_SCHEMA_CLIENT_TOOL_PARAMETERS + and len(version) == 8 + and version.isdigit() + ): + return kind + return None + + +def _anthropic_schema_client_tool_parameters(td: dict, kind: str) -> dict: + parameters = _ANTHROPIC_SCHEMA_CLIENT_TOOL_PARAMETERS[kind] + if kind != "text_editor": + return parameters + + version = td["type"].rpartition("_")[2] + commands = list(parameters["properties"]["command"]["enum"]) + if version < "20250429": + commands.append("undo_edit") + return { + **parameters, + "properties": { + **parameters["properties"], + "command": {**parameters["properties"]["command"], "enum": commands}, + }, + } + + def anthropic_tools_to_openai(tools: list) -> list[dict]: """Convert Anthropic client tools to OpenAI function-tool format.""" result = [] @@ -179,6 +309,9 @@ def anthropic_tools_to_openai(tools: list) -> list[dict]: td = t if isinstance(t, dict) else t.model_dump() name = td.get("name") input_schema = td.get("input_schema") + schema_client_kind = anthropic_schema_client_tool_kind(td) + if schema_client_kind is not None: + input_schema = _anthropic_schema_client_tool_parameters(td, schema_client_kind) if not name or input_schema is None: continue result.append( @@ -186,7 +319,8 @@ def anthropic_tools_to_openai(tools: list) -> list[dict]: "type": "function", "function": { "name": name, - "description": td.get("description", ""), + "description": td.get("description") + or _ANTHROPIC_SCHEMA_CLIENT_TOOL_DESCRIPTIONS.get(schema_client_kind, ""), "parameters": input_schema, }, } @@ -539,7 +673,7 @@ class AnthropicPassthroughEmitter: Only calls naming a tool in ``allowed_tools`` (the client's declared tools) are promoted; everything else streams as text exactly as before. - Never enabled for Studio's own tool loop. + Never enabled for Unsloth's own tool loop. """ from core.inference.passthrough_healing import StreamToolCallHealer diff --git a/studio/backend/core/inference/api_monitor.py b/studio/backend/core/inference/api_monitor.py index f76a38576f..b637ba56d1 100644 --- a/studio/backend/core/inference/api_monitor.py +++ b/studio/backend/core/inference/api_monitor.py @@ -5,6 +5,7 @@ from __future__ import annotations +import os import threading import time import uuid @@ -18,6 +19,14 @@ _MAX_PROMPT_CHARS = 12000 _MAX_REPLY_CHARS = 12000 _PREVIEW_CHARS = 360 +# Opt-in startup kill switch for Studio's in-memory API monitor. +_DISABLE_ENV = "UNSLOTH_STUDIO_DISABLE_API_MONITOR" +_TRUE_VALUES = frozenset({"1", "true", "yes", "on"}) + + +def _api_monitor_disabled() -> bool: + return os.environ.get(_DISABLE_ENV, "").strip().lower() in _TRUE_VALUES + def _trim(text: Optional[str], limit: int) -> str: if not text: @@ -52,6 +61,13 @@ class ApiMonitorEntry: total_tokens: Optional[int] = None total_tokens_authoritative: bool = False error: Optional[str] = None + # "request" (HTTP call) or "lifecycle" (model load/unload: event/reason, not a prompt; shared). + kind: str = "request" + event: Optional[str] = None + reason: Optional[str] = None + shared: bool = False + # 0-100 for a running download row; None when not applicable. + progress: Optional[float] = None def snapshot(self, *, include_details: bool = True) -> dict[str, Any]: duration_ms = None @@ -85,6 +101,10 @@ class ApiMonitorEntry: "completion_tokens": self.completion_tokens, "total_tokens": self.total_tokens, "error": self.error, + "kind": self.kind, + "event": self.event, + "reason": self.reason, + "progress": self.progress, } if include_details: payload["prompt"] = self.prompt @@ -93,10 +113,16 @@ class ApiMonitorEntry: class ApiMonitor: - def __init__(self, max_entries: int = _MAX_ENTRIES): + def __init__( + self, + max_entries: int = _MAX_ENTRIES, + *, + enabled: bool = True, + ): self._entries: deque[ApiMonitorEntry] = deque() self._max_entries = max(0, max_entries) self._lock = threading.Lock() + self._enabled = enabled def start( self, @@ -108,6 +134,8 @@ class ApiMonitor: context_length: Optional[int] = None, subject: Optional[str] = None, ) -> str: + if not self._enabled: + return "" now = time.time() entry = ApiMonitorEntry( id = f"apireq_{uuid.uuid4().hex[:12]}", @@ -127,6 +155,75 @@ class ApiMonitor: self._trim_terminal_locked() return entry.id + def record_lifecycle( + self, + *, + event: str, + model: str, + reason: Optional[str] = None, + running: bool = False, + ) -> str: + """Record a model load/unload alongside the request traffic that caused it. + + ``running=True`` opens the row for the caller to close with :meth:`finish` / + :meth:`fail`; an unload is terminal on arrival. Rows are shared (visible to + every subject) and share the request retention budget. + """ + if not self._enabled: + return "" + now = time.time() + entry = ApiMonitorEntry( + id = f"apievt_{uuid.uuid4().hex[:12]}", + endpoint = f"model.{event}", + method = "", + model = model or "default", + prompt = "", + status = "running" if running else "completed", + started_at = now, + updated_at = now, + started_monotonic = time.monotonic(), + finished_at = None if running else now, + finished_monotonic = None if running else time.monotonic(), + kind = "lifecycle", + event = event, + reason = reason, + shared = True, + ) + with self._lock: + self._entries.appendleft(entry) + self._trim_terminal_locked() + return entry.id + + def relabel(self, entry_id: Optional[str], model: str) -> None: + """Rename an open lifecycle row once the load resolves its real id: up front + the caller only has the load path, which may be an HF snapshot dir.""" + if not entry_id or not model: + return + with self._lock: + entry = self._find_locked(entry_id) + if entry is not None: + entry.model = model + entry.updated_at = time.time() + + def set_progress(self, entry_id: Optional[str], progress: Optional[float]) -> None: + """Update an open download row's percentage (clamped to 0-100).""" + if not entry_id or progress is None: + return + with self._lock: + entry = self._find_locked(entry_id) + if entry is not None and entry.status == "running": + entry.progress = min(100.0, max(0.0, float(progress))) + entry.updated_at = time.time() + + def discard(self, entry_id: Optional[str]) -> None: + """Drop a row that turned out not to be an event (an already-satisfied load).""" + if not entry_id: + return + with self._lock: + entry = self._find_locked(entry_id) + if entry is not None: + self._entries.remove(entry) + def append_reply(self, entry_id: Optional[str], text: str) -> None: if not entry_id or not text: return @@ -212,6 +309,18 @@ class ApiMonitor: self._entries.appendleft(entry) self._trim_terminal_locked() + def fail_open(self, entry_id: Optional[str], error: str) -> None: + """Fail only a still-open row: unlike :meth:`fail`, a catch-all in a + ``finally`` cannot stamp an error onto a request that already succeeded.""" + if not entry_id: + return + with self._lock: + entry = self._find_locked(entry_id) + if entry is None or entry.finished_at is not None: + return + # Same lock as the check, so a finish() cannot land in between. + self._fail_locked(entry, error) + def fail(self, entry_id: Optional[str], error: str) -> None: if not entry_id: return @@ -224,15 +333,18 @@ class ApiMonitor: if error: entry.error = _trim(error, 1000) return - now = time.time() - entry.status = "error" - entry.error = _trim(error, 1000) - entry.updated_at = now - entry.finished_at = now - entry.finished_monotonic = time.monotonic() - self._entries.remove(entry) - self._entries.appendleft(entry) - self._trim_terminal_locked() + self._fail_locked(entry, error) + + def _fail_locked(self, entry: ApiMonitorEntry, error: str) -> None: + now = time.time() + entry.status = "error" + entry.error = _trim(error, 1000) + entry.updated_at = now + entry.finished_at = now + entry.finished_monotonic = time.monotonic() + self._entries.remove(entry) + self._entries.appendleft(entry) + self._trim_terminal_locked() def snapshot( self, @@ -244,7 +356,7 @@ class ApiMonitor: return [ entry.snapshot(include_details = include_details) for entry in self._entries - if subject is None or entry.subject == subject + if self._visible(entry, subject) ] def get( @@ -257,22 +369,29 @@ class ApiMonitor: entry = self._find_locked(entry_id) if entry is None: return None - if subject is not None and entry.subject != subject: + if not self._visible(entry, subject): return None return entry.snapshot(include_details = True) def active_count(self, *, subject: Optional[str] = None) -> int: + # Lifecycle rows show as "running" while loading but are not in-flight API requests. with self._lock: return sum( 1 for entry in self._entries - if entry.status == "running" and (subject is None or entry.subject == subject) + if entry.status == "running" + and entry.kind != "lifecycle" + and (subject is None or entry.subject == subject) ) def clear(self) -> None: with self._lock: self._entries.clear() + @staticmethod + def _visible(entry: ApiMonitorEntry, subject: Optional[str]) -> bool: + return subject is None or entry.subject == subject or entry.shared + def _find_locked(self, entry_id: str) -> Optional[ApiMonitorEntry]: for entry in self._entries: if entry.id == entry_id: @@ -292,4 +411,4 @@ class ApiMonitor: self._entries = kept -api_monitor = ApiMonitor() +api_monitor = ApiMonitor(enabled = not _api_monitor_disabled()) diff --git a/studio/backend/core/inference/audio_codecs.py b/studio/backend/core/inference/audio_codecs.py index 93c7da72cb..b59f2bcce0 100644 --- a/studio/backend/core/inference/audio_codecs.py +++ b/studio/backend/core/inference/audio_codecs.py @@ -76,8 +76,14 @@ class AudioCodecManager: if self._snac_model is not None: return from snac import SNAC + from utils.hf_cache_settings import active_hf_hub_cache - self._snac_model = SNAC.from_pretrained("hubertsiuzdak/snac_24khz").to(device).eval() + # Route weights to the selected cache; this can run in the main process. + self._snac_model = ( + SNAC.from_pretrained("hubertsiuzdak/snac_24khz", cache_dir = active_hf_hub_cache()) + .to(device) + .eval() + ) logger.info("Loaded SNAC codec (24kHz)") def _load_bicodec( diff --git a/studio/backend/core/inference/chat_template_helpers.py b/studio/backend/core/inference/chat_template_helpers.py index 897db8262d..3a8463855b 100644 --- a/studio/backend/core/inference/chat_template_helpers.py +++ b/studio/backend/core/inference/chat_template_helpers.py @@ -10,10 +10,242 @@ native-chat-template fallback used by the transformers and MLX backends. import copy import json import logging +from dataclasses import dataclass from typing import Optional _THINK_OPEN = "" _THINK_CLOSE = "" +_GEMMA_CHANNEL_START = "<|channel>" +_GEMMA_THOUGHT_OPEN = "<|channel>thought" +_GEMMA_THOUGHT_CLOSE = "" +_GEMMA_TEMPLATE_OPENERS = ( + _GEMMA_THOUGHT_OPEN + "\n", + _GEMMA_THOUGHT_OPEN + "\\n", + _GEMMA_THOUGHT_OPEN + _GEMMA_THOUGHT_CLOSE, +) + + +def _tokenizer_objects(tokenizer) -> tuple: + """Return a processor/tokenizer and its distinct nested tokenizer.""" + if tokenizer is None: + return () + nested = getattr(tokenizer, "tokenizer", None) + return (tokenizer,) if nested is None or nested is tokenizer else (tokenizer, nested) + + +def _selected_template_strings_from_value( + template, + tools = None, + *, + prefer_tool_use: bool = True, +) -> tuple[str, ...]: + """Return the named chat template matching HF's default selection rules.""" + tools = tools or None + if isinstance(template, str): + return (template,) + if not isinstance(template, dict): + return () + if prefer_tool_use and tools and isinstance(template.get("tool_use"), str): + return (template["tool_use"],) + if isinstance(template.get("default"), str): + return (template["default"],) + values = tuple(value for value in template.values() if isinstance(value, str)) + return values if len(values) == 1 else () + + +def _selected_chat_template_strings(tokenizer, tools = None) -> tuple[str, ...]: + """Return the active chat template selected for this request.""" + tools = tools or None + getter = getattr(tokenizer, "get_chat_template", None) + if callable(getter): + for kwargs in ({"chat_template": None, "tools": tools}, {"tools": tools}, {}): + try: + selected = getter(**kwargs) + except Exception: + continue + if isinstance(selected, str): + return (selected,) + # ProcessorMixin.apply_chat_template does not switch to "tool_use" implicitly; + # it uses "default" unless chat_template= names another template. + is_processor = getattr(tokenizer, "tokenizer", None) is not None and callable( + getattr(tokenizer, "apply_chat_template", None) + ) + return _selected_template_strings_from_value( + getattr(tokenizer, "chat_template", None), + tools, + prefer_tool_use = not is_processor, + ) + + +def _detect_reasoning_channel_markers_from_templates( + templates: tuple[str, ...], +) -> Optional[tuple[str, str]]: + """Return Gemma native reasoning markers only when a template emits them.""" + if any(opener in template for template in templates for opener in _GEMMA_TEMPLATE_OPENERS): + return _GEMMA_THOUGHT_OPEN, _GEMMA_THOUGHT_CLOSE + return None + + +def detect_reasoning_channel_markers(tokenizer, tools = None) -> Optional[tuple[str, str]]: + """Return native Gemma thought-channel markers supported by a tokenizer. + + Detection uses the active chat template rather than model names or vocabulary + membership. Some models expose Gemma control tokens without using the native + thought-channel response protocol, and those must keep normal + ``skip_special_tokens`` streaming. + """ + for obj in _tokenizer_objects(tokenizer): + templates = _selected_chat_template_strings(obj, tools) + if templates: + return _detect_reasoning_channel_markers_from_templates(templates) + return None + + +def detect_reasoning_channel_markers_from_template( + template, tools = None +) -> Optional[tuple[str, str]]: + """Return native Gemma thought-channel markers from a raw template value.""" + return _detect_reasoning_channel_markers_from_templates( + _selected_template_strings_from_value(template, tools) + ) + + +def detect_reasoning_channel_markers_from_model_info( + tokenizer, + model_info: Optional[dict] = None, + tools = None, +) -> Optional[tuple[str, str]]: + """Return reasoning markers from the active or cached native template.""" + markers = detect_reasoning_channel_markers(tokenizer, tools = tools) + if markers is not None or not isinstance(model_info, dict): + return markers + + native_templates = ( + model_info.get("native_chat_template"), + (model_info.get("chat_template_info") or {}).get("template"), + ) + for template in native_templates: + markers = detect_reasoning_channel_markers_from_template(template, tools) + if markers is not None: + return markers + return None + + +@dataclass(frozen = True) +class ChatTemplateRenderResult: + """Prompt plus response-protocol metadata selected by the renderer.""" + + prompt: str + reasoning_channel_markers: Optional[tuple[str, str]] = None + + +def _split_partial_marker(text: str, marker: str) -> tuple[str, str]: + """Hold the longest suffix that may become ``marker`` in the next chunk.""" + for length in range(min(len(text), len(marker) - 1), 0, -1): + if text.endswith(marker[:length]): + return text[:-length], text[-length:] + return text, "" + + +class ReasoningChannelNormalizer: + """Incrementally convert one native reasoning channel to ````. + + The parser follows mlx-vlm's streaming boundary behavior but emits Unsloth's + established canonical text contract. Only the configured opening and + closing markers are consumed; tool-call and other control markers remain + available to downstream parsers. + """ + + def __init__(self, opening_marker: str, closing_marker: str): + self._opening_marker = opening_marker + self._closing_marker = closing_marker + self._buffer = "" + self._in_reasoning = False + self._reasoning_done = False + self._skip_opening_newline = False + + def feed(self, text: str) -> str: + """Consume a raw text delta and return the stable canonical delta.""" + self._buffer += text or "" + output: list[str] = [] + while self._buffer: + if self._reasoning_done: + output.append(self._buffer) + self._buffer = "" + break + + if self._in_reasoning and self._skip_opening_newline: + if self._buffer.startswith("\n"): + self._buffer = self._buffer[1:] + self._skip_opening_newline = False + if not self._buffer: + break + + marker = self._closing_marker if self._in_reasoning else self._opening_marker + index = self._buffer.find(marker) + if index < 0: + stable, self._buffer = _split_partial_marker(self._buffer, marker) + output.append(stable) + break + + output.append(self._buffer[:index]) + self._buffer = self._buffer[index + len(marker) :] + if self._in_reasoning: + output.append(_THINK_CLOSE) + self._in_reasoning = False + self._reasoning_done = True + else: + output.append(_THINK_OPEN) + self._in_reasoning = True + self._skip_opening_newline = True + return "".join(output) + + def finish(self) -> str: + """Flush a naturally completed stream and close an open think block.""" + output = self.drain() + if self._in_reasoning: + output += _THINK_CLOSE + self._in_reasoning = False + self._reasoning_done = True + return output + + def drain(self) -> str: + """Flush buffered literal text without synthesizing a closing tag.""" + output = self._buffer + self._buffer = "" + return output + + +def normalize_reasoning_snapshots( + stream, + tokenizer = None, + cancel_event = None, + markers: Optional[tuple[str, str]] = None, + tools = None, +): + """Normalize a prefix-monotonic cumulative text stream when supported.""" + markers = markers or detect_reasoning_channel_markers(tokenizer, tools = tools) + if markers is None: + yield from stream + return + + normalizer = ReasoningChannelNormalizer(*markers) + raw_output = "" + normalized_output = "" + for snapshot in stream: + if not snapshot.startswith(raw_output): + raise RuntimeError("Reasoning normalization requires cumulative text snapshots") + delta = normalizer.feed(snapshot[len(raw_output) :]) + raw_output = snapshot + if delta: + normalized_output += delta + yield normalized_output + + cancelled = cancel_event is not None and cancel_event.is_set() + tail = normalizer.drain() if cancelled else normalizer.finish() + if tail: + normalized_output += tail + yield normalized_output def detect_think_prefill(prompt: Optional[str], special_tokens = None) -> str: @@ -94,6 +326,58 @@ def _normalize_tool_call_arguments(messages: list) -> list: return out if mutated else messages +def _take_tool_result(pending: list, call_id) -> Optional[dict]: + if call_id: + for i, result in enumerate(pending): + if result.get("tool_call_id") == call_id: + return pending.pop(i) + for i, result in enumerate(pending): + if not result.get("tool_call_id"): + return pending.pop(i) + return None + + +def _split_parallel_tool_calls(messages: list) -> list: + """Llama 3.x templates render one call per message, so split parallel calls + into consecutive single-call messages, each followed by its own result.""" + if not any(isinstance(m, dict) and len(m.get("tool_calls") or ()) > 1 for m in messages): + return messages + + out: list = [] + i = 0 + total = len(messages) + while i < total: + msg = messages[i] + calls = msg.get("tool_calls") if isinstance(msg, dict) else None + if not calls or len(calls) <= 1: + out.append(msg) + i += 1 + continue + + # Tool results right after this message answer its calls. + j = i + 1 + pending: list = [] + while ( + j < total + and isinstance(messages[j], dict) + and messages[j].get("role") in ("tool", "ipython") + ): + pending.append(messages[j]) + j += 1 + + for idx, call in enumerate(calls): + piece = {**msg, "tool_calls": [call]} + if idx: + piece["content"] = "" + out.append(piece) + result = _take_tool_result(pending, call.get("id") if isinstance(call, dict) else None) + if result is not None: + out.append(result) + out.extend(pending) + i = j + return out + + def apply_chat_template_for_generation( tokenizer, messages: list, @@ -146,13 +430,21 @@ def apply_chat_template_for_generation( try: return _render(messages) except Exception: - # Strict tool templates reject the JSON-string ``arguments`` form via - # TypeError or a broad Jinja raise_exception, so retry with dicts coerced. - # Original messages render first, so working templates stay byte-identical. + # Retry with repairs applied cumulatively. Originals render first, so + # working templates stay byte-identical. + candidates: list = [] normalized = _normalize_tool_call_arguments(messages) - if normalized is messages: - raise - return _render(normalized) + if normalized is not messages: + candidates.append(normalized) + split = _split_parallel_tool_calls(normalized) + if split is not normalized: + candidates.append(split) + for candidate in candidates: + try: + return _render(candidate) + except Exception: + continue + raise def render_native_template( @@ -166,7 +458,8 @@ def render_native_template( preserve_thinking: Optional[bool] = None, apply_fn = None, hf_token: Optional[str] = None, -) -> Optional[str]: + return_metadata: bool = False, +): """Render ``messages`` + ``tools`` with the model's NATIVE chat template. Some Unsloth override templates (e.g. ``mistral``, ``gemma-4``) do not emit @@ -175,7 +468,9 @@ def render_native_template( tool-calling syntax. It is loaded straight from the repo (bypassing any override on the live tokenizer) and cached on ``model_info``. Returns the rendered prompt only if the native template actually emits the tools (render - differs with vs without tools); otherwise ``None``. + differs with vs without tools); otherwise ``None``. With ``return_metadata``, + returns ``ChatTemplateRenderResult`` so callers can stream with the response + protocol selected by this request's template. ``hf_token`` is the token the model was loaded with -- passed to the repo load so a gated/private model's native template can still be fetched (otherwise the @@ -261,7 +556,16 @@ def render_native_template( exc, ) return None - return with_tools if with_tools != no_tools else None + if with_tools == no_tools: + return None + if return_metadata: + return ChatTemplateRenderResult( + with_tools, + _detect_reasoning_channel_markers_from_templates( + _selected_template_strings_from_value(native_tpl, tools) + ), + ) + return with_tools def render_with_native_template_fallback( @@ -277,7 +581,8 @@ def render_with_native_template_fallback( preserve_thinking: Optional[bool] = None, apply_fn = None, hf_token: Optional[str] = None, -) -> str: + return_metadata: bool = False, +): """Return ``formatted_prompt``, swapping in a native-template render when an override template dropped the ``tools`` schema. @@ -285,9 +590,27 @@ def render_with_native_template_fallback( them (detected by comparison, robust against tool names in the system prompt), re-render with the model's native template. Shared by the transformers and MLX backends so both advertise tools consistently. ``hf_token`` is forwarded so a - gated/private model's native template can still be fetched.""" + gated/private model's native template can still be fetched. With + ``return_metadata``, returns the selected prompt plus reasoning-channel markers + for the exact template used by this request.""" + live_markers = detect_reasoning_channel_markers(tokenizer, tools = tools) + + def _result(prompt: str, markers = live_markers): + if return_metadata: + return ChatTemplateRenderResult(prompt, markers) + return prompt + if not tools: - return formatted_prompt + # Gemma 4 can emit its native reasoning protocol even when a generation-time + # Unsloth override rendered a marker-free prompt. Preserve the live-verified + # no-tools thinking behavior without letting cached native metadata describe + # unrelated tool prompts that kept the active override. + markers = live_markers + if markers is None: + markers = detect_reasoning_channel_markers_from_model_info( + tokenizer, model_info, tools = None + ) + return _result(formatted_prompt, markers) if apply_fn is None: apply_fn = apply_chat_template_for_generation # Probe whether the live template dropped the schema. A tools-requiring template @@ -307,9 +630,9 @@ def render_with_native_template_fallback( active_model_name, exc, ) - return formatted_prompt + return _result(formatted_prompt) if formatted_prompt != probe_no_tools: - return formatted_prompt # template already emits the tools schema + return _result(formatted_prompt) # template already emits the tools schema native_prompt = render_native_template( model_info = model_info, active_model_name = active_model_name, @@ -320,6 +643,7 @@ def render_with_native_template_fallback( preserve_thinking = preserve_thinking, apply_fn = apply_fn, hf_token = hf_token, + return_metadata = return_metadata, ) if native_prompt: logger.info( @@ -328,4 +652,4 @@ def render_with_native_template_fallback( active_model_name, ) return native_prompt - return formatted_prompt + return _result(formatted_prompt) diff --git a/studio/backend/core/inference/chat_templates.py b/studio/backend/core/inference/chat_templates.py index 58f63ff61b..04c0db6aae 100644 --- a/studio/backend/core/inference/chat_templates.py +++ b/studio/backend/core/inference/chat_templates.py @@ -4,13 +4,13 @@ """Bundled chat-template selection for GGUF inference. Some shipped GGUF quants embed an older chat template. Rather than re-cutting and -asking users to re-download every quant, Studio can override the embedded template +asking users to re-download every quant, Unsloth can override the embedded template at llama-server launch time with a bundled, up-to-date Jinja template for known model families. The override is wired through the existing ``chat_template_override`` -> ``--chat-template-file`` path in ``LlamaCppBackend.load_model``. Currently this covers ``unsloth/gemma-4-*-GGUF``, which gains the upstream PR #118 -``preserve_thinking`` flag (defaulted OFF here) so the Studio "Preserve thinking" +``preserve_thinking`` flag (defaulted OFF here) so the Unsloth "Preserve thinking" toggle appears while staying disabled by default. """ diff --git a/studio/backend/core/inference/external_provider.py b/studio/backend/core/inference/external_provider.py index 20312e067c..2debf946e9 100644 --- a/studio/backend/core/inference/external_provider.py +++ b/studio/backend/core/inference/external_provider.py @@ -473,7 +473,7 @@ def _apply_mistral_reasoning_controls( # handles every provider without storing credentials. def _create_shared_http_client() -> httpx.AsyncClient: # Unsupported env proxy schemes (socks:// etc) raise at construction and - # would crash Studio startup (#6090); retry ignoring env proxies instead. + # would crash Unsloth startup (#6090); retry ignoring env proxies instead. try: return httpx.AsyncClient() except (ImportError, ValueError) as exc: @@ -858,7 +858,7 @@ class ExternalProviderClient: if not self._is_openai_compatible(): # Gemini speaks its own native REST shape (contents/parts); # `_stream_gemini` translates request/response into the OpenAI - # Chat Completions chunk format the rest of Studio expects. + # Chat Completions chunk format the rest of Unsloth expects. # API ref: https://ai.google.dev/gemini-api/docs if self.provider_type == "gemini": async for line in self._stream_gemini( @@ -1706,7 +1706,7 @@ class ExternalProviderClient: # Translate OpenAI multimodal parts -> Anthropic native shapes. # - `image_url` -> `{type:"image", source:...}` # - `input_document` -> `{type:"document", source:...}` - # (Studio extension; mirrors Anthropic's document block, + # (Unsloth extension; mirrors Anthropic's document block, # which supports PDFs as base64 or URL per # https://platform.claude.com/docs/en/build-with-claude/vision) anthropic_parts: list[dict[str, Any]] = [] @@ -1749,7 +1749,7 @@ class ExternalProviderClient: } ) elif part.get("type") == "input_document": - # Studio's normalised PDF/doc type (file_data data-URI or + # Unsloth's normalised PDF/doc type (file_data data-URI or # file_url) -> Anthropic's native `document` block. url = part.get("file_url") or "" data_uri = part.get("file_data") or "" @@ -4704,7 +4704,7 @@ class ExternalProviderClient: {"type": "image_generation_call", "id": call_id} ) elif part_type == "input_document": - # Map Studio's `input_document` onto Responses' `input_file`. + # Map Unsloth's `input_document` onto Responses' `input_file`. # https://developers.openai.com/api/docs/guides/images-vision file_url = part.get("file_url") file_data = part.get("file_data") @@ -6010,7 +6010,7 @@ class ExternalProviderClient: if not models and self.provider_type == "ollama": models = await self._list_ollama_native_models() # Gemini's native /v1beta/models uses a different shape; repackage - # into the OpenAI-compatible one Studio expects. + # into the OpenAI-compatible one Unsloth expects. if not models and self.provider_type == "gemini": models = self._parse_gemini_models(data) return models @@ -6213,7 +6213,7 @@ def _friendly_provider_error_text( *, model: str | None = None, ) -> str: - """Rewrite common provider errors into actionable Studio copy.""" + """Rewrite common provider errors into actionable Unsloth copy.""" if status_code == 404 and model: lowered = raw_message.lower() if "not found" in lowered or "not_found" in lowered: diff --git a/studio/backend/core/inference/inference.py b/studio/backend/core/inference/inference.py index 172e9e5546..e78bf1be8d 100644 --- a/studio/backend/core/inference/inference.py +++ b/studio/backend/core/inference/inference.py @@ -5,9 +5,10 @@ from unsloth import FastLanguageModel, FastVisionModel from unsloth.chat_templates import get_chat_template -from transformers import TextStreamer +from transformers import TextIteratorStreamer, TextStreamer from peft import PeftModel, PeftModelForCausalLM +import contextlib import json import sys import torch @@ -32,6 +33,11 @@ from core.inference.chat_eos import ( chat_eos_repair, resolve_chat_turn_end_eos_ids_using, ) +from core.inference.chat_template_helpers import ( + ReasoningChannelNormalizer, + detect_reasoning_channel_markers, + detect_think_prefill, +) from core.inference.presence_penalty import _make_presence_penalty_processor from io import StringIO import structlog @@ -187,6 +193,53 @@ class HarmonyTextStreamer: self._queue.put(new_content) +class ReasoningTextIteratorStreamer(TextIteratorStreamer): + """TextIteratorStreamer that preserves native channel tokens until parsed.""" + + def __init__( + self, + tokenizer, + *, + markers: tuple[str, str], + skip_prompt: bool = True, + timeout: float = 0.2, + cancel_event = None, + **decode_kwargs, + ): + decode_kwargs["skip_special_tokens"] = False + super().__init__(tokenizer, skip_prompt = skip_prompt, timeout = timeout, **decode_kwargs) + self._normalizer = ReasoningChannelNormalizer(*markers) + self._cancel_event = cancel_event + self._aborted = False + + def abort(self): + """Mark generation as failed so ``end`` drains without closing.""" + self._aborted = True + + def on_finalized_text( + self, + text: str, + stream_end: bool = False, + ): + """Queue canonical deltas, closing only on natural stream completion.""" + delta = self._normalizer.feed(text) + if delta: + self.text_queue.put(delta, timeout = self.timeout) + + if stream_end: + cancelled = self._aborted or ( + self._cancel_event is not None and self._cancel_event.is_set() + ) + tail = self._normalizer.drain() if cancelled else self._normalizer.finish() + if tail: + self.text_queue.put(tail, timeout = self.timeout) + self.text_queue.put(self.stop_signal, timeout = self.timeout) + + +class _GenerationThreadError(RuntimeError): + """Generation worker failures that should propagate through stream routes.""" + + class InferenceBackend: """Unified inference backend supporting text, vision, and LoRA models""" @@ -514,7 +567,7 @@ class InferenceBackend: _meta_path = Path(config.path) / "export_metadata.json" try: if _meta_path.exists(): - _meta = json.loads(_meta_path.read_text()) + _meta = json.loads(_meta_path.read_text(encoding = "utf-8-sig")) if _meta.get("base_model"): processor_source = _meta["base_model"] except Exception: @@ -836,6 +889,7 @@ class InferenceBackend: thread_id: Optional[str] = None, rag_scope: Optional[dict] = None, presence_penalty: float = 0.0, + reasoning_prefilled: bool = False, ): """Run an agentic tool loop on top of ``generate_chat_response``. @@ -889,6 +943,7 @@ class InferenceBackend: session_id = session_id, thread_id = thread_id, rag_scope = rag_scope, + reasoning_prefilled = reasoning_prefilled, ) def generate_chat_response( @@ -960,8 +1015,7 @@ class InferenceBackend: thread can toggle adapters under the generation lock. """ if not self.active_model_name: - yield "Error: No active model" - return + raise RuntimeError("No active model") model_info = self.models[self.active_model_name] is_vision = model_info.get("is_vision", False) @@ -1049,6 +1103,7 @@ class InferenceBackend: template_messages = [{"role": "system", "content": system_prompt}] + messages else: template_messages = messages + reasoning_channel_markers_resolved = False try: if not (hasattr(tokenizer, "chat_template") and tokenizer.chat_template): raise ValueError( @@ -1058,6 +1113,7 @@ class InferenceBackend: f"Please use a model that includes a chat template, or manually set " f"one via tokenizer.chat_template before inference." ) + reasoning_channel_markers = None formatted_prompt = self._apply_chat_template_for_generation( tokenizer, template_messages, @@ -1073,7 +1129,7 @@ class InferenceBackend: render_with_native_template_fallback, ) - formatted_prompt = render_with_native_template_fallback( + render_result = render_with_native_template_fallback( formatted_prompt = formatted_prompt, tokenizer = tokenizer, model_info = model_info, @@ -1085,13 +1141,19 @@ class InferenceBackend: preserve_thinking = preserve_thinking, apply_fn = self._apply_chat_template_for_generation, hf_token = model_info.get("hf_token"), + return_metadata = True, ) + formatted_prompt = render_result.prompt + reasoning_channel_markers = render_result.reasoning_channel_markers + reasoning_channel_markers_resolved = True logger.debug(f"Formatted prompt: {formatted_prompt[:200]}...") except Exception as e: logger.error(f"Error applying chat template: {e}") # Fall back to manual formatting formatted_prompt = self.format_chat_prompt(messages, system_prompt) + reasoning_channel_markers = None + reasoning_channel_markers_resolved = True # Step 3: generate yield from self.generate_stream( @@ -1105,6 +1167,8 @@ class InferenceBackend: cancel_event = cancel_event, _adapter_state = _adapter_state, presence_penalty = presence_penalty, + reasoning_channel_markers = reasoning_channel_markers, + reasoning_channel_markers_resolved = reasoning_channel_markers_resolved, ) def _generate_vision_response( @@ -1190,21 +1254,27 @@ class InferenceBackend: # Stream with TextIteratorStreamer + background thread try: - from core.inference.chat_template_helpers import detect_think_prefill - # Re-emit an open prefill swallowed by skip_prompt (see # generate_stream). think_prefix = detect_think_prefill( prompt_text, getattr(raw_tokenizer, "all_special_tokens", None) ) - from transformers import TextIteratorStreamer import threading - streamer = TextIteratorStreamer( + streamer = self._make_text_streamer( raw_tokenizer, + protocol_source = processor, + # The text-only VLM fallback above did not render with the + # processor template, so its native markers do not describe + # this request's response protocol. + reasoning_channel_markers = detect_reasoning_channel_markers(processor) + if image + else None, + reasoning_channel_markers_resolved = True, skip_prompt = True, - skip_special_tokens = True, timeout = 0.2, + cancel_event = cancel_event, + use_harmony = self._is_gpt_oss_model(), ) generation_kwargs = dict( @@ -1226,6 +1296,10 @@ class InferenceBackend: ) if _pp is not None: generation_kwargs["logits_processor"] = _pp + stopping_criteria = self._cancel_stopping_criteria(cancel_event) + if stopping_criteria is not None: + generation_kwargs["stopping_criteria"] = stopping_criteria + active_stop_token_ids = self._generation_stop_token_ids(model, generation_kwargs) err: dict[str, str] = {} @@ -1235,6 +1309,8 @@ class InferenceBackend: model.generate(**generation_kwargs) except Exception as e: err["msg"] = str(e) + if hasattr(streamer, "abort"): + streamer.abort() logger.error(f"Vision generation error in thread: {e}") finally: try: @@ -1251,12 +1327,17 @@ class InferenceBackend: if think_prefix: yield think_prefix from queue import Empty + import time generation_complete = False + cancel_deadline = None try: while True: if cancel_event is not None and cancel_event.is_set(): - break + if cancel_deadline is None: + cancel_deadline = time.monotonic() + 10 + elif time.monotonic() >= cancel_deadline: + break try: new_token = next(streamer) except StopIteration: @@ -1265,27 +1346,48 @@ class InferenceBackend: except Empty: if not thread.is_alive(): generation_complete = True + output = yield from self._drain_streamer_tail( + streamer, output, active_stop_token_ids + ) + break + if cancel_deadline is not None: + remaining = cancel_deadline - time.monotonic() + if remaining <= 0: + break + thread.join(timeout = remaining) + if thread.is_alive(): + break + generation_complete = True + output = yield from self._drain_streamer_tail( + streamer, output, active_stop_token_ids + ) break continue if new_token: - output += new_token - cleaned = self._clean_generated_text(output) + output, cleaned = self._append_stream_delta( + output, new_token, active_stop_token_ids + ) yield cleaned finally: if cancel_event is not None and not generation_complete: cancel_event.set() - thread.join(timeout = 10) + join_timeout = 10 + if cancel_deadline is not None: + join_timeout = max(0, cancel_deadline - time.monotonic()) + thread.join(timeout = join_timeout) if thread.is_alive(): logger.warning( "Vision generation thread did not exit after cancel/join timeout" ) if err.get("msg"): - yield f"Error: {err['msg']}" + raise _GenerationThreadError(err["msg"]) + except _GenerationThreadError: + raise except Exception as e: logger.error(f"Vision generation error: {e}") - yield f"Error: {str(e)}" + raise def generate_audio_input_response( self, @@ -1410,11 +1512,13 @@ class InferenceBackend: ) if err.get("msg"): - yield f"Error: {err['msg']}" + raise _GenerationThreadError(err["msg"]) + except _GenerationThreadError: + raise except Exception as e: logger.error(f"Audio input generation error: {e}") - yield f"Error: {str(e)}" + raise def generate_whisper_response( self, @@ -1447,6 +1551,86 @@ class InferenceBackend: from utils.datasets import is_gpt_oss_model_name return is_gpt_oss_model_name(model_name or self.active_model_name or "") + def _make_text_streamer( + self, + tokenizer, + *, + protocol_source = None, + reasoning_channel_markers = None, + reasoning_channel_markers_resolved: bool = False, + skip_prompt: bool = True, + timeout: float = 0.2, + cancel_event = None, + use_harmony: bool = False, + ): + """Create the streamer matching this model's native response protocol.""" + if use_harmony: + try: + return HarmonyTextStreamer( + tokenizer, + skip_prompt = skip_prompt, + timeout = timeout, + ) + except Exception as e: + logger.warning(f"HarmonyTextStreamer init failed, falling back: {e}") + return TextIteratorStreamer( + tokenizer, + skip_prompt = skip_prompt, + skip_special_tokens = True, + timeout = timeout, + ) + + markers = ( + reasoning_channel_markers + if reasoning_channel_markers_resolved + else reasoning_channel_markers + or detect_reasoning_channel_markers(protocol_source or tokenizer) + ) + if markers is not None: + return ReasoningTextIteratorStreamer( + tokenizer, + markers = markers, + skip_prompt = skip_prompt, + timeout = timeout, + cancel_event = cancel_event, + ) + return TextIteratorStreamer( + tokenizer, + skip_prompt = skip_prompt, + skip_special_tokens = True, + timeout = timeout, + ) + + def _append_stream_delta( + self, + output: str, + new_token: str, + stop_token_ids = None, + ): + """Append a streamer delta and apply response-boundary cleanup.""" + output += new_token + return output, self._clean_generated_text(output, stop_token_ids = stop_token_ids) + + def _drain_streamer_tail( + self, + streamer, + output: str, + stop_token_ids = None, + ): + """Drain queued streamer text after the producer exits.""" + while True: + try: + new_token = next(streamer) + except StopIteration: + return output + except Exception: + return output + if new_token: + output, cleaned = self._append_stream_delta( + output, new_token, stop_token_ids = stop_token_ids + ) + yield cleaned + def generate_stream( self, prompt: str, @@ -1459,6 +1643,8 @@ class InferenceBackend: cancel_event = None, _adapter_state = None, presence_penalty: float = 0.0, + reasoning_channel_markers = None, + reasoning_channel_markers_resolved: bool = False, ) -> Generator[str, None, None]: """Generate a streaming text response (text models only). @@ -1467,8 +1653,7 @@ class InferenceBackend: ``presence_penalty`` matches the GGUF sampling path via a logits processor (0 disables it). """ if not self.active_model_name: - yield "Error: No active model" - return + raise RuntimeError("No active model") model_info = self.models[self.active_model_name] model = model_info["model"] @@ -1481,9 +1666,7 @@ class InferenceBackend: try: inputs = tokenizer(prompt, return_tensors = "pt").to(model.device) - from transformers import TextIteratorStreamer import threading - from core.inference.chat_template_helpers import detect_think_prefill # skip_prompt swallows an open prefilled by the template; # re-emit it so the frontend can render the thinking block. @@ -1494,30 +1677,16 @@ class InferenceBackend: else detect_think_prefill(prompt, getattr(tokenizer, "all_special_tokens", None)) ) - # gpt-oss models: HarmonyTextStreamer parses the multi-channel - # harmony protocol into tags - if self._is_gpt_oss_model(): - try: - streamer = HarmonyTextStreamer( - tokenizer, - skip_prompt = True, - timeout = 0.2, - ) - except Exception as e: - logger.warning(f"HarmonyTextStreamer init failed, falling back: {e}") - streamer = TextIteratorStreamer( - tokenizer, - skip_prompt = True, - skip_special_tokens = True, - timeout = 0.2, - ) - else: - streamer = TextIteratorStreamer( - tokenizer, - skip_prompt = True, - skip_special_tokens = True, - timeout = 0.2, - ) + streamer = self._make_text_streamer( + tokenizer, + protocol_source = model_info.get("tokenizer"), + reasoning_channel_markers = reasoning_channel_markers, + reasoning_channel_markers_resolved = reasoning_channel_markers_resolved, + skip_prompt = True, + timeout = 0.2, + cancel_event = cancel_event, + use_harmony = self._is_gpt_oss_model(), + ) generation_kwargs = dict( **inputs, @@ -1535,27 +1704,16 @@ class InferenceBackend: if tokenizer.pad_token_id is None else tokenizer.pad_token_id, ) + active_stop_token_ids = self._generation_stop_token_ids(model, generation_kwargs) # Presence penalty (GGUF parity); prompt_len excludes prompt tokens. _pp = _make_presence_penalty_processor( presence_penalty, int(inputs["input_ids"].shape[1]) ) if _pp is not None: generation_kwargs["logits_processor"] = _pp - if cancel_event is not None: - from transformers.generation.stopping_criteria import ( - StoppingCriteria, - StoppingCriteriaList, - ) - class _CancelCriteria(StoppingCriteria): - def __init__(self, ev): - self.ev = ev - - def __call__(self, input_ids, scores, **kwargs): - return self.ev.is_set() - - generation_kwargs["stopping_criteria"] = StoppingCriteriaList( - [_CancelCriteria(cancel_event)] - ) + stopping_criteria = self._cancel_stopping_criteria(cancel_event) + if stopping_criteria is not None: + generation_kwargs["stopping_criteria"] = stopping_criteria def generate_fn(): with self._generation_lock: @@ -1565,6 +1723,8 @@ class InferenceBackend: model.generate(**generation_kwargs) except Exception as e: err["msg"] = str(e) + if hasattr(streamer, "abort"): + streamer.abort() logger.error(f"Generation error: {e}") finally: try: @@ -1582,12 +1742,17 @@ class InferenceBackend: if think_prefix: yield think_prefix from queue import Empty + import time generation_complete = False + cancel_deadline = None try: while True: if cancel_event is not None and cancel_event.is_set(): - break + if cancel_deadline is None: + cancel_deadline = time.monotonic() + 10 + elif time.monotonic() >= cancel_deadline: + break try: new_token = next(streamer) except StopIteration: @@ -1596,11 +1761,27 @@ class InferenceBackend: except Empty: if not thread.is_alive(): generation_complete = True + output = yield from self._drain_streamer_tail( + streamer, output, active_stop_token_ids + ) + break + if cancel_deadline is not None: + remaining = cancel_deadline - time.monotonic() + if remaining <= 0: + break + thread.join(timeout = remaining) + if thread.is_alive(): + break + generation_complete = True + output = yield from self._drain_streamer_tail( + streamer, output, active_stop_token_ids + ) break continue if new_token: - output += new_token - cleaned = self._clean_generated_text(output) + output, cleaned = self._append_stream_delta( + output, new_token, active_stop_token_ids + ) yield cleaned finally: # Set cancel_event only on early exit (user cancel), NOT on @@ -1609,16 +1790,21 @@ class InferenceBackend: # disrupt the next serialized request (e.g. compare mode). if cancel_event is not None and not generation_complete: cancel_event.set() - thread.join(timeout = 10) + join_timeout = 10 + if cancel_deadline is not None: + join_timeout = max(0, cancel_deadline - time.monotonic()) + thread.join(timeout = join_timeout) if thread.is_alive(): logger.warning("Generation thread did not exit after cancel/join timeout") if err.get("msg"): - yield f"Error: {err['msg']}" + raise _GenerationThreadError(err["msg"]) + except _GenerationThreadError: + raise except Exception as e: logger.error(f"Error during generation: {e}") - yield f"Error: {str(e)}" + raise # ── Audio (TTS) Generation ──────────────────────────────────── @@ -1757,8 +1943,30 @@ class InferenceBackend: + text + "<|text_end|>\n<|audio_start|><|global_features_start|>\n" ) + with torch.inference_mode(): - with torch.amp.autocast("cuda", dtype = model.dtype): + # Derive the autocast device from the loaded model, not from the + # global backend: a CPU-fallback DAC on an XPU/CUDA host must not + # open a GPU autocast context around CPU tensors. + device_type = ( + model.device.type + if hasattr(model.device, "type") + else str(model.device).split(":", 1)[0] + ) + # Clamp to autocast-supported backends so exotic devices + # (e.g. "meta" during accelerate offloaded loading) do not raise. + # MPS is autocast-supported since torch 2.3, keep it in the set. + if device_type not in ("cuda", "xpu", "mps", "cpu"): + device_type = "cpu" + # CPU and XPU autocast only accept bfloat16/float16. For a + # float32 model, skip autocast entirely to avoid raising or + # producing a warning on every generate call. + autocast_dtype_supported = model.dtype in (torch.bfloat16, torch.float16) + if device_type in ("cpu", "xpu") and not autocast_dtype_supported: + autocast_ctx = contextlib.nullcontext() + else: + autocast_ctx = torch.amp.autocast(device_type, dtype = model.dtype) + with autocast_ctx: inputs = tokenizer([prompt], return_tensors = "pt").to(model.device) generated = model.generate( **inputs, @@ -2073,8 +2281,13 @@ class InferenceBackend: except Exception as e: logger.warning(f"Could not fully reset model state for {model_name}: {e}") - def reset_generation_state(self): - """Reset any cached generation state to prevent hanging after errors""" + def reset_generation_state(self, caller_cancel_event = None): + """Reset any cached generation state to prevent hanging after errors + + ``caller_cancel_event`` is accepted for signature parity with the + orchestrator, which uses it to drop a reset from a request that never + started. Nothing here cancels a live generation, so it is unused. + """ try: # Clear cached state for ALL loaded models for model_name in self.models.keys(): @@ -2107,8 +2320,42 @@ class InferenceBackend: return img.resize(new_size, Image.Resampling.LANCZOS) return img - def _clean_generated_text(self, text: str) -> str: - """Strip leaked special tokens using the tokenizer's own token list.""" + def _generation_stop_token_ids(self, model, generation_kwargs: dict): + """Return the stop-token ids active for a ``generate`` call.""" + if "eos_token_id" in generation_kwargs: + return generation_kwargs.get("eos_token_id") + generation_config = getattr(model, "generation_config", None) + eos_token_id = getattr(generation_config, "eos_token_id", None) + if eos_token_id is not None: + return eos_token_id + config = getattr(model, "config", None) + return getattr(config, "eos_token_id", None) + + def _cancel_stopping_criteria(self, cancel_event): + """Build a Transformers stopping criteria list for user cancellation.""" + if cancel_event is None: + return None + from transformers.generation.stopping_criteria import ( + StoppingCriteria, + StoppingCriteriaList, + ) + + class _CancelCriteria(StoppingCriteria): + def __init__(self, ev): + self.ev = ev + + def __call__(self, input_ids, scores, **kwargs): + return self.ev.is_set() + + return StoppingCriteriaList([_CancelCriteria(cancel_event)]) + + def _clean_generated_text( + self, + text: str, + *, + stop_token_ids = None, + ) -> str: + """Strip leaked response-boundary tokens after streaming.""" if self._is_gpt_oss_model(): # HarmonyTextStreamer emits clean .... Strip any # harmony protocol tokens and other gpt-oss tokens (e.g. @@ -2118,10 +2365,28 @@ class InferenceBackend: return text.strip() tokenizer = self.models.get(self.active_model_name, {}).get("tokenizer") + tokenizer = getattr(tokenizer, "tokenizer", tokenizer) if tokenizer: - for token in getattr(tokenizer, "all_special_tokens", []): - if token in text: - text = text.replace(token, "") + if stop_token_ids is None: + stop_token_ids = self.models.get(self.active_model_name, {}).get( + "chat_turn_end_eos_ids" + ) + if isinstance(stop_token_ids, int): + stop_token_ids = (stop_token_ids,) + for token_id in stop_token_ids or (): + try: + token = tokenizer.convert_ids_to_tokens(int(token_id)) + except Exception: + token = None + if isinstance(token, str) and token and text.endswith(token): + text = text[: -len(token)] + elif ( + isinstance(token, str) + and token + and text.endswith("") + and text[: -len("")].endswith(token) + ): + text = text[: -len("") - len(token)] + "" return text.strip() def _load_chat_template_info(self, model_name: str): diff --git a/studio/backend/core/inference/llama_admission.py b/studio/backend/core/inference/llama_admission.py index b6a939c87b..7bf0dd7429 100644 --- a/studio/backend/core/inference/llama_admission.py +++ b/studio/backend/core/inference/llama_admission.py @@ -13,37 +13,159 @@ from __future__ import annotations import asyncio import os +import sys import threading from collections import deque from dataclasses import dataclass from typing import Deque, Optional -ADMISSION_CONTROL_ENV = "UNSLOTH_OPENAI_COMPAT_ADMISSION_CONTROL" -ADMISSION_QUEUE_TIMEOUT_ENV = "UNSLOTH_OPENAI_COMPAT_ADMISSION_QUEUE_TIMEOUT" -ADMISSION_KEEPALIVE_INTERVAL_ENV = "UNSLOTH_OPENAI_COMPAT_ADMISSION_KEEPALIVE_INTERVAL" -ADMISSION_MAX_QUEUE_ENV = "UNSLOTH_OPENAI_COMPAT_ADMISSION_MAX_QUEUE" +# dataclass(slots = True) halves per-instance overhead. Measured as perf-neutral +# here, not a speed win: it costs a little on construction and gains it back on +# access. It is 3.10+ and this package declares >=3.9, so gate it rather than +# dropping it outright. Empty on 3.9 means a plain dataclass. +_SLOTS = {"slots": True} if sys.version_info >= (3, 10) else {} + + +ADMISSION_CONTROL_ENV = "UNSLOTH_LLAMA_ADMISSION_CONTROL" +ADMISSION_QUEUE_TIMEOUT_ENV = "UNSLOTH_LLAMA_ADMISSION_QUEUE_TIMEOUT" +ADMISSION_KEEPALIVE_INTERVAL_ENV = "UNSLOTH_LLAMA_ADMISSION_KEEPALIVE_INTERVAL" +ADMISSION_MAX_QUEUE_ENV = "UNSLOTH_LLAMA_ADMISSION_MAX_QUEUE" +ADMISSION_QUEUE_PER_SLOT_ENV = "UNSLOTH_LLAMA_ADMISSION_QUEUE_PER_SLOT" + +# The UNSLOTH_OPENAI_COMPAT_* spellings predate this queue being shared with the +# Anthropic /v1/messages route (same llama-server slots). Still honored; the +# neutral name above wins when both are set. +_LEGACY_ENV = { + ADMISSION_CONTROL_ENV: "UNSLOTH_OPENAI_COMPAT_ADMISSION_CONTROL", + ADMISSION_QUEUE_TIMEOUT_ENV: "UNSLOTH_OPENAI_COMPAT_ADMISSION_QUEUE_TIMEOUT", + ADMISSION_KEEPALIVE_INTERVAL_ENV: "UNSLOTH_OPENAI_COMPAT_ADMISSION_KEEPALIVE_INTERVAL", + ADMISSION_MAX_QUEUE_ENV: "UNSLOTH_OPENAI_COMPAT_ADMISSION_MAX_QUEUE", +} DEFAULT_ADMISSION_ENABLED = True +# None: a queued request waits for its slot indefinitely rather than timing out. DEFAULT_ADMISSION_QUEUE_TIMEOUT_S = None DEFAULT_ADMISSION_KEEPALIVE_INTERVAL_S = 5.0 -DEFAULT_ADMISSION_MAX_QUEUE = 64 +# None: no absolute cap, the wait line is sized from the pool instead. +DEFAULT_ADMISSION_MAX_QUEUE = None +# Wait line = 16 x the serving slots, so it tracks --parallel (4 slots -> 64 +# waiters, 8 -> 128). Purely a memory guard; waiting itself is never timed out. +DEFAULT_ADMISSION_QUEUE_PER_SLOT = 16 +# Floor for the scaled line, so a 1-slot backend (plain `unsloth studio`, or any +# load downshifted to fit VRAM) keeps the depth it had before scaling existed +# rather than dropping to 16 and rejecting callers that used to queue. +DEFAULT_ADMISSION_MIN_QUEUE = 64 -@dataclass(frozen = True) +def _executor_workers() -> int: + """Threads asyncio's default executor runs to_thread work on. + + Mirrors ThreadPoolExecutor's own default sizing, which is what + ``run_in_executor(None, ...)`` builds. 3.13 sizes it from + ``process_cpu_count()``, which honours CPU affinity and cgroup quotas; + ``cpu_count()`` would budget from the whole host inside a one-core container. + """ + cpus = getattr(os, "process_cpu_count", os.cpu_count)() or 1 + return min(32, cpus + 4) + + +def _executor_reserve(workers: int) -> int: + """Threads kept clear of parked approvals, for generation steps, stream + teardown and unrelated to_thread work. Scaled rather than flat: a flat count + would leave a 5-worker executor (one usable CPU) no budget at all. + """ + return max(2, workers // 8) + + +def _max_parked(capacity: int) -> int: + """How many holders may sit on an approval prompt with their slot given back. + + A pending prompt parks an executor thread (the loop blocks inside + to_thread(next, gen)) whether or not it parked its slot, the pool already + permits `capacity` of those, and every park admits one more, so budget only + what the executor has left over. Zero on a backend whose --parallel alone + fills it: the prompt then holds its slot, as it did before parking existed. + """ + workers = _executor_workers() + spare = workers - _executor_reserve(workers) - max(0, capacity) + # A quarter of the executor, floored at two while `spare` allows: a quarter of + # five is one, and one park cannot cover the two simultaneous prompts #7455 + # exists for. + return max(0, min(max(2, workers // 4), spare)) + + +# Process-wide, not per queue: there is one executor, and base_url takes a fresh +# port on every load, so a per-queue budget would hand the same allowance to each +# backend and to every reload, blind to the approvals parked on the old queue. +_PARK_LOCK = threading.Lock() +_parked_total = 0 + + +def _claim_park(limit: int) -> bool: + global _parked_total + with _PARK_LOCK: + if _parked_total >= limit: + return False + _parked_total += 1 + return True + + +def _drop_park() -> None: + global _parked_total + with _PARK_LOCK: + _parked_total = max(0, _parked_total - 1) + + +def _live_capacity(current: "LlamaAdmissionQueue") -> int: + """Slots across every backend still serving requests. + + One queue's capacity is the wrong denominator for a budget sized against the + one executor: a reload drains the old queue alongside the new one, and + prompts on both park threads. Idle queues hold nothing and are about to be + evicted. + """ + with _QUEUES_LOCK: + queues = list(_QUEUES.values()) + # is_idle takes each queue's own lock, so never while holding _QUEUES_LOCK. + total = sum(queue._capacity for queue in queues if queue is current or not queue.is_idle()) + return total if any(queue is current for queue in queues) else total + current._capacity + + +@dataclass(frozen = True, **_SLOTS) class LlamaAdmissionConfig: enabled: bool = DEFAULT_ADMISSION_ENABLED queue_timeout_s: Optional[float] = DEFAULT_ADMISSION_QUEUE_TIMEOUT_S keepalive_interval_s: float = DEFAULT_ADMISSION_KEEPALIVE_INTERVAL_S max_queue: Optional[int] = DEFAULT_ADMISSION_MAX_QUEUE + queue_per_slot: Optional[int] = DEFAULT_ADMISSION_QUEUE_PER_SLOT + # Unconditional floor on the scaled line. The env path clears it when the + # operator sets QUEUE_PER_SLOT, so only the default multiplier is floored. + min_queue: Optional[int] = DEFAULT_ADMISSION_MIN_QUEUE + + def queue_limit(self, capacity: int) -> Optional[int]: + """How many callers may line up for a pool of ``capacity`` slots. + + An explicit ``max_queue`` wins; otherwise the line scales with the slots + so it follows ``--parallel``. The default multiplier is floored, so a + 1-slot backend does not end up shallower than it was before scaling. None + (or any non-positive setting) means an unbounded line. + """ + if self.max_queue is not None: + return self.max_queue if self.max_queue > 0 else None + if not self.queue_per_slot or self.queue_per_slot <= 0: + return None + scaled = self.queue_per_slot * max(1, capacity) + return max(self.min_queue, scaled) if self.min_queue else scaled -@dataclass(frozen = True) +@dataclass(frozen = True, **_SLOTS) class LlamaAdmissionSnapshot: key: str capacity: int active: int queued: int + free: int = 0 class LlamaAdmissionError(Exception): @@ -69,8 +191,17 @@ class LlamaAdmissionCancelled(LlamaAdmissionError): pass -def _bool_env(name: str, default: bool) -> bool: +def _raw_env(name: str) -> Optional[str]: + """Value for a canonical name, falling back to its legacy spelling.""" value = os.environ.get(name) + if value is None or not value.strip(): + legacy = _LEGACY_ENV.get(name) + value = os.environ.get(legacy) if legacy else None + return value + + +def _bool_env(name: str, default: bool) -> bool: + value = _raw_env(name) if value is None or not value.strip(): return default value = value.strip().lower() @@ -82,7 +213,7 @@ def _bool_env(name: str, default: bool) -> bool: def _optional_positive_float_env(name: str, default: Optional[float]) -> Optional[float]: - value = os.environ.get(name) + value = _raw_env(name) if value is None or not value.strip(): return default try: @@ -93,7 +224,7 @@ def _optional_positive_float_env(name: str, default: Optional[float]) -> Optiona def _positive_float_env(name: str, default: float) -> float: - value = os.environ.get(name) + value = _raw_env(name) if value is None or not value.strip(): return default try: @@ -103,19 +234,38 @@ def _positive_float_env(name: str, default: float) -> float: return parsed if parsed > 0 else default -def _optional_positive_int_env(name: str, default: Optional[int]) -> Optional[int]: - value = os.environ.get(name) - if value is None or not value.strip(): - return default +def _queue_limits_from_env() -> tuple[Optional[int], Optional[int], Optional[int]]: + """(max_queue, queue_per_slot, min_queue) from the environment. + + An absolute MAX_QUEUE wins outright; MAX_QUEUE=0 asks for an unbounded line. + Unset leaves the per-slot multiplier in charge (itself 0 for unbounded). The + floor applies only to the default multiplier: setting QUEUE_PER_SLOT means + the operator wants that exact depth, however shallow. + """ + # Explicit means it parsed, not just that something was set: a typo falls back + # to the default multiplier, so it has to keep the default's floor too. + raw_per_slot = _raw_env(ADMISSION_QUEUE_PER_SLOT_ENV) try: - parsed = int(value.strip()) + per_slot = int((raw_per_slot or "").strip()) except ValueError: - return default - return parsed if parsed > 0 else None + per_slot, min_queue = DEFAULT_ADMISSION_QUEUE_PER_SLOT, DEFAULT_ADMISSION_MIN_QUEUE + else: + per_slot, min_queue = (per_slot if per_slot > 0 else None), None + raw = _raw_env(ADMISSION_MAX_QUEUE_ENV) + if raw is None or not raw.strip(): + return None, per_slot, min_queue + try: + parsed = int(raw.strip()) + except ValueError: + return None, per_slot, min_queue + return (parsed, None, None) if parsed > 0 else (None, None, None) def llama_admission_config_from_env() -> LlamaAdmissionConfig: + max_queue, queue_per_slot, min_queue = _queue_limits_from_env() return LlamaAdmissionConfig( + queue_per_slot = queue_per_slot, + min_queue = min_queue, enabled = _bool_env(ADMISSION_CONTROL_ENV, DEFAULT_ADMISSION_ENABLED), queue_timeout_s = _optional_positive_float_env( ADMISSION_QUEUE_TIMEOUT_ENV, @@ -125,14 +275,11 @@ def llama_admission_config_from_env() -> LlamaAdmissionConfig: ADMISSION_KEEPALIVE_INTERVAL_ENV, DEFAULT_ADMISSION_KEEPALIVE_INTERVAL_S, ), - max_queue = _optional_positive_int_env( - ADMISSION_MAX_QUEUE_ENV, - DEFAULT_ADMISSION_MAX_QUEUE, - ), + max_queue = max_queue, ) -@dataclass +@dataclass(**_SLOTS) class _Waiter: loop: asyncio.AbstractEventLoop future: asyncio.Future @@ -141,20 +288,130 @@ class _Waiter: class LlamaAdmissionLease: - def __init__(self, queue: Optional["LlamaAdmissionQueue"]): + __slots__ = ("_queue", "_slot", "_released", "_release_lock", "_parked", "_budgeted") + + def __init__( + self, + queue: Optional["LlamaAdmissionQueue"], + slot: Optional[int] = None, + ): self._queue = queue + self._slot = slot self._released = False self._release_lock = threading.Lock() + self._parked = False + self._budgeted = False + + @property + def slot(self) -> Optional[int]: + """Pool slot this lease holds, or None when admission is disabled.""" + return self._slot + + def park(self) -> bool: + """Hand the slot back while this holder waits on something off the GPU. + + A run stopped on a tool approval prompt is not decoding, so holding its + slot would let unanswered prompts fill the pool while llama-server idles. + The lease itself stays valid: releasing it after a park is still correct. + + False when the park budget is spent and nothing was given back: the + caller keeps its slot across the prompt, as it did before parking + existed. Slower for whoever is behind it, but each freed slot admits + another run that can park too, on the executor the generators run on. + """ + queue = self._queue + with self._release_lock: + if queue is None or self._released or self._parked: + return False + # Under the lease lock so the decision and the handover cannot split. + # Nothing takes the queue lock then a lease lock, so this order is + # the only one in play. + if not queue.try_park(self._slot): + return False + self._parked = True + self._budgeted = True + self._slot = None + return True + + def _drop_budget(self) -> None: + """Give the executor budget back now the prompt wait is over. + + Separate from the queue's parked count, which lasts until the slot is + back: the executor thread is free the moment the answer arrives. Holding + the budget until the resume lands would refuse someone else's park for a + finished wait, and that someone holds the slot the resumer wants. + """ + with self._release_lock: + if not self._budgeted: + return + self._budgeted = False + _drop_park() + + def unpark(self) -> None: + """Drop the parked state without reclaiming a slot. + + For a holder that is tearing down: it will not decode again. Resuming + holders must use ``unpark_async``, which waits for a slot instead of + going back to llama-server past the admission limit. + """ + with self._release_lock: + if not self._parked: + return + self._parked = False + self._drop_budget() + if self._queue is not None: + self._queue.unpark() + + async def unpark_async( + self, + *, + cancel_event = None, + poll_s: float = 0.02, + ) -> None: + """Take a slot back, waiting until the pool has room. + + ``park`` gave the slot to a waiter, so by the time the user answers the + prompt someone else may be decoding in it. Resuming regardless put two + holders on a one-slot server. Gives up if the caller is cancelled, since + the holder is then leaving anyway and must not be stuck here. + """ + queue = self._queue + if queue is None or not self._parked: + return + # Before the wait, not after: the prompt is answered, so this holder is + # already off the executor and must not keep anyone else off it. + self._drop_budget() + slot = await queue.acquire_parked_slot(cancel_event = cancel_event, poll_s = poll_s) + stranded = None + with self._release_lock: + # release() may have run during the wait; it clears the flag and does + # the unpark itself, so only the caller that clears it here repeats one. + parked, self._parked = self._parked, False + if self._released: + # Released while waiting: this lease will never hand the slot + # back, so return it here rather than strand it for good. + stranded = slot + else: + self._slot = slot + if parked: + queue.unpark() + if stranded is not None: + queue.release(stranded) def release(self) -> None: queue = None + parked = False with self._release_lock: if self._released: return self._released = True queue = self._queue + parked, self._parked = self._parked, False + self._drop_budget() if queue is not None: - queue.release() + if parked: + queue.unpark() + queue.release(self._slot) async def __aenter__(self) -> "LlamaAdmissionLease": return self @@ -164,6 +421,8 @@ class LlamaAdmissionLease: class LlamaAdmissionReservation: + __slots__ = ("_queue", "_lease", "_waiter", "snapshot") + def __init__( self, *, @@ -195,6 +454,13 @@ class LlamaAdmissionReservation: return self._lease async def wait(self, timeout_s: float) -> Optional[LlamaAdmissionLease]: + """Wait up to ``timeout_s`` for a slot. + + A timeout leaves this reservation queued so the caller can poll again. + Any exit that abandons the wait for good must call ``cancel()``, or the + slot granted later is delivered to a future nobody reads and is never + released. + """ lease = self.lease_nowait() if lease is not None: return lease @@ -229,12 +495,74 @@ class LlamaAdmissionReservation: class LlamaAdmissionQueue: + """A fixed pool of generation slots for one llama-server, plus a FIFO wait line. + + The pool mirrors llama-server's own ``--parallel`` slots: ``capacity`` slot ids + are each either free or held by exactly one caller. A caller that finds every + slot busy waits in arrival order and is handed the next slot to free, so no + caller is starved. This bounds only the callers that reserve: chat completions + and messages do, while /v1/completions, Studio's own chat endpoint and RAG + captioning all reach llama-server directly, so it is not a global cap. + Waiting is unbounded in time by default (``queue_timeout_s`` + None); the wait line itself is bounded, and only how many may line up before + new arrivals are rejected. By default that is ``16 x slots`` floored at 64, + not unlimited: an unbounded line takes ``max_queue`` or ``queue_per_slot`` + set to 0. See ``LlamaAdmissionConfig.queue_limit``. + """ + + __slots__ = ( + "key", + "_lock", + "_capacity", + "_free", + "_in_use", + "_held", + "_waiters", + "_parked", + "_unpark_tickets", + "_unpark_seq", + ) + def __init__(self, key: str): self.key = key self._lock = threading.Lock() - self._active = 0 self._capacity = 1 + self._free: list[int] = [0] + # Held slots as a bitmask: one int instead of a set, so the pool costs the + # same whether it is idle or saturated. _held is its popcount, kept as a + # counter because int.bit_count() is 3.10+ and this package targets 3.9. + self._in_use = 0 + self._held = 0 self._waiters: Deque[_Waiter] = deque() + # Holders parked on a tool approval prompt. They hold no slot, so this only + # keeps the queue off the idle-eviction list while they are away. + self._parked = 0 + # FIFO tickets for holders resuming from a park (see acquire_parked_slot). A + # bare count deadlocked: every approved holder blocked every other one. + self._unpark_tickets: Deque[int] = deque() + self._unpark_seq = 0 + + def _resize_pool_locked(self, capacity: int) -> None: + # Slots past a shrunk capacity retire when their holder releases them. + if capacity == self._capacity: + return + self._capacity = capacity + self._free = [slot for slot in range(capacity) if not self._in_use >> slot & 1] + + def _can_admit_locked(self, reserved: int) -> bool: + # Slots still held above a shrunk capacity keep occupying the backend, so + # count every held slot against the ceiling, not just the ids below it. + # ``reserved`` holds slots back for approved holders waiting to resume; + # without it a stream of new arrivals took the next slot, forever. + return bool(self._free) and (self._held + reserved) < self._capacity + + def _take_slot_locked(self, reserved: int) -> Optional[int]: + if not self._can_admit_locked(reserved): + return None + slot = self._free.pop() + self._in_use |= 1 << slot + self._held += 1 + return slot def reserve(self, *, capacity: int, config: LlamaAdmissionConfig) -> LlamaAdmissionReservation: capacity = max(1, int(capacity or 1)) @@ -242,22 +570,25 @@ class LlamaAdmissionQueue: return LlamaAdmissionReservation( queue = None, lease = LlamaAdmissionLease(None), - snapshot = LlamaAdmissionSnapshot(self.key, capacity, 0, 0), + snapshot = LlamaAdmissionSnapshot(self.key, capacity, 0, 0, capacity), ) loop = asyncio.get_running_loop() with self._lock: - self._capacity = capacity - self._prune_waiters_locked() + self._resize_pool_locked(capacity) self._grant_waiters_locked() - if self._active < self._capacity and not self._waiters: - self._active += 1 - return LlamaAdmissionReservation( - queue = self, - lease = LlamaAdmissionLease(self), - snapshot = self._snapshot_locked(), - ) - if config.max_queue is not None and len(self._waiters) >= config.max_queue: + if not self._waiters: + slot = self._take_slot_locked(len(self._unpark_tickets)) + if slot is not None: + # No snapshot here: callers read it through snapshot_now(), + # which re-reads the queue, so building one per admitted + # request would be pure allocation on the hot path. + return LlamaAdmissionReservation( + queue = self, + lease = LlamaAdmissionLease(self, slot), + ) + limit = config.queue_limit(self._capacity) + if limit is not None and self._live_waiters_locked() >= limit: raise LlamaAdmissionQueueFull( "llama-server generation queue is full", snapshot = self._snapshot_locked(), @@ -270,15 +601,82 @@ class LlamaAdmissionQueue: return LlamaAdmissionReservation( queue = self, waiter = waiter, - snapshot = self._snapshot_locked(), ) - def release(self) -> None: + def _release_slot_locked(self, slot: Optional[int]) -> None: + # A slot id at or past a shrunk capacity retires instead of returning. + if slot is None or not self._in_use >> slot & 1: + return + self._in_use &= ~(1 << slot) + self._held -= 1 + if slot < self._capacity: + self._free.append(slot) + + def release(self, slot: Optional[int]) -> None: with self._lock: - if self._active > 0: - self._active -= 1 + self._release_slot_locked(slot) self._grant_waiters_locked() + def try_park(self, slot: Optional[int]) -> bool: + """Return a parked holder's slot to the pool. See ``LlamaAdmissionLease.park``. + + False leaves the slot with its holder, so a refused park costs nothing to + undo. The per-queue count is only what ``is_idle`` reads; the budget and + the capacity it is sized from are both process-wide. + """ + if not _claim_park(_max_parked(_live_capacity(self))): + return False + with self._lock: + self._parked += 1 + self._release_slot_locked(slot) + self._grant_waiters_locked() + return True + + def unpark(self) -> None: + with self._lock: + if self._parked > 0: + self._parked -= 1 + + async def acquire_parked_slot( + self, + *, + cancel_event = None, + poll_s: float = 0.02, + ) -> Optional[int]: + """Wait for a slot for a holder resuming from a park, None if cancelled. + + Ordered by ticket rather than counted, so approvals resume in the order + they came back: counting them made every approved holder block every + other one, and with nothing decoding that never resolved. + """ + with self._lock: + self._unpark_seq += 1 + ticket = self._unpark_seq + self._unpark_tickets.append(ticket) + try: + while True: + with self._lock: + ahead = 0 + for queued in self._unpark_tickets: + if queued == ticket: + break + ahead += 1 + # Only the approvals ahead of this one hold slots back from it. + slot = self._take_slot_locked(ahead) + if slot is not None: + return slot + if cancel_event is not None and cancel_event.is_set(): + return None + await asyncio.sleep(poll_s) + finally: + with self._lock: + try: + self._unpark_tickets.remove(ticket) + except ValueError: + pass + # This ticket was holding a slot back from the wait line. + self._grant_waiters_locked() + def cancel(self, waiter: _Waiter) -> None: lease_to_release = None with self._lock: @@ -291,7 +689,13 @@ class LlamaAdmissionQueue: lease_to_release = waiter.granted_lease waiter.granted_lease = None if not waiter.future.done(): - waiter.loop.call_soon_threadsafe(waiter.future.cancel) + try: + waiter.loop.call_soon_threadsafe(waiter.future.cancel) + except RuntimeError: + # Loop gone. Routes call cancel() from finally blocks, so + # raising here would both mask their exception and skip the + # release below, stranding the slot for the process lifetime. + pass if lease_to_release is not None: lease_to_release.release() @@ -303,20 +707,32 @@ class LlamaAdmissionQueue: def is_idle(self) -> bool: with self._lock: self._prune_waiters_locked() - return self._active == 0 and not self._waiters + # A parked holder owns no slot but is coming back to this queue, so + # evicting it here would resume it against a fresh 1-slot pool. + return self._in_use == 0 and not self._waiters and not self._parked def _grant_waiters_locked(self) -> None: - self._prune_waiters_locked() - while self._waiters and self._active < self._capacity: + # Dead waiters are skipped as they are popped, so no prune is needed here. + while self._waiters and self._can_admit_locked(len(self._unpark_tickets)): waiter = self._waiters.popleft() if waiter.cancelled or waiter.future.done(): continue - self._active += 1 - lease = LlamaAdmissionLease(self) + slot = self._take_slot_locked(len(self._unpark_tickets)) + lease = LlamaAdmissionLease(self, slot) waiter.granted_lease = lease - waiter.loop.call_soon_threadsafe(self._deliver_lease, waiter, lease) + try: + waiter.loop.call_soon_threadsafe(self._deliver_lease, waiter, lease) + except RuntimeError: + # Waiter's loop is gone. Reclaim the slot; leaving the bit set + # would strand it, since _free is rebuilt from the bitmask. + waiter.granted_lease = None + self._release_slot_locked(slot) def _deliver_lease(self, waiter: _Waiter, lease: LlamaAdmissionLease) -> None: + # Runs on the waiter's own loop thread, which is also the only thread that + # cancels that reservation, so waiter state is safe to touch unlocked here. + # release() may be called from any thread, but only reaches this via + # call_soon_threadsafe. Cancelling off-loop would need this under _lock. if waiter.cancelled or waiter.future.done(): waiter.granted_lease = None if not waiter.future.done(): @@ -331,16 +747,32 @@ class LlamaAdmissionQueue: lease.release() def _prune_waiters_locked(self) -> None: + # Rebuilding the deque on every reserve/release dominated the hot path, so + # only pay it when a waiter actually died out of band (an externally + # cancelled future); cancel() already drops its own waiter eagerly. + for waiter in self._waiters: + if waiter.cancelled or waiter.future.done(): + break + else: + return self._waiters = deque( waiter for waiter in self._waiters if not waiter.cancelled and not waiter.future.done() ) + def _live_waiters_locked(self) -> int: + self._prune_waiters_locked() + return len(self._waiters) + def _snapshot_locked(self) -> LlamaAdmissionSnapshot: return LlamaAdmissionSnapshot( key = self.key, capacity = self._capacity, - active = self._active, + active = self._held, queued = len(self._waiters), + # What another caller could actually take, so the admission log never + # shows free slots next to queued requests: after a shrink, ids below + # the new capacity can be free while holdovers still fill the ceiling. + free = min(len(self._free), max(0, self._capacity - self._held)), ) @@ -364,5 +796,10 @@ def get_llama_admission_queue(key: str) -> LlamaAdmissionQueue: def reset_llama_admission_queues() -> None: + global _parked_total with _QUEUES_LOCK: _QUEUES.clear() + # The budget outlives the queues it was claimed against, so dropping them + # without it leaks the count and shrinks the budget for good. + with _PARK_LOCK: + _parked_total = 0 diff --git a/studio/backend/core/inference/llama_cpp.py b/studio/backend/core/inference/llama_cpp.py index 5622fed574..712caf43e5 100644 --- a/studio/backend/core/inference/llama_cpp.py +++ b/studio/backend/core/inference/llama_cpp.py @@ -9,7 +9,9 @@ OpenAI-compatible /v1/chat/completions endpoint. import atexit import contextlib +import functools import json +import math import os import re import struct @@ -21,13 +23,27 @@ import subprocess import sys import threading import time +import uuid from pathlib import Path -from typing import Callable, Collection, Generator, Iterable, List, Mapping, Optional, Union +from typing import ( + Callable, + Collection, + Generator, + Iterable, + List, + Literal, + Mapping, + MutableMapping, + Optional, + Union, +) import httpx from core.inference.llama_server_args import ( + _LAYER_OFFLOAD_FLAGS, _effective_tensor_parallel, + _flag_name, _tensor_parallel_matches_loaded, extra_args_disable_mmproj, parse_cache_override, @@ -69,6 +85,7 @@ from core.tool_healing import ( strip_outside_think, ) from utils.native_path_leases import child_env_without_native_path_secret +from utils.child_stdio import utf8_child_env from utils.hf_xet_fallback import hf_hub_download_with_xet_fallback from utils.subprocess_compat import ( windows_hidden_subprocess_kwargs as _windows_hidden_subprocess_kwargs, @@ -76,12 +93,17 @@ from utils.subprocess_compat import ( from utils.process_lifetime import child_popen_kwargs as _child_popen_kwargs from core.inference.tool_call_parser import ( MAX_ACT_REPROMPTS as _MAX_REPROMPTS, + NUDGE_TOOL_CALLS_STATUS as _NUDGE_TOOL_CALLS_STATUS, REPROMPT_MAX_CHARS as _REPROMPT_MAX_CHARS, + is_reprompt_repeat as _is_reprompt_repeat, + is_reprompt_restatement as _is_reprompt_restatement, is_short_intent_without_action as _is_short_intent_without_action, reprompt_to_act_message as _reprompt_to_act_message, ) from core.inference.tool_loop_controller import ( ToolLoopController, + append_deferred_nudges, + awaiting_approval_status, tool_event_provenance, ) from state.tool_approvals import ( @@ -111,10 +133,19 @@ LLAMA_SERVER_NOT_FOUND_DETAIL = ( "then try again. (Advanced: set LLAMA_SERVER_PATH to an existing binary.)" ) +# Shared by the route, pre-teardown and post-metadata rejections (#7205). +_VULKAN_DIFFUSION_GPU_IDS_ERROR = ( + "GPU selection (gpu_ids) is not supported for a DiffusionGemma " + "GGUF on a Vulkan llama.cpp build: the diffusion runner selects " + "its device by CUDA physical index, which has no defined mapping " + "to ggml Vulkan device ordinals. Omit gpu_ids to use the default " + "device." +) + # llama-server can serve HTTP 200 while running a model entirely on CPU when a # GPU backend fails to init (#5807 / #5106 / #5830). Classify the startup log so -# Studio can warn. Priority: explicit "offloaded N/M layers to GPU" counts +# Unsloth can warn. Priority: explicit "offloaded N/M layers to GPU" counts # (authoritative), then GPU "model buffer size" lines (host-pinned _Host # excluded), then the "device_info:" device table (disconfirm only). _GPU_OFFLOAD_MARKERS = ( @@ -221,7 +252,7 @@ def _wsl_system_rocm_lib_dirs() -> "list[str]": with open("/proc/version", encoding = "utf-8", errors = "replace") as fh: if "microsoft" not in fh.read().lower(): return [] - except OSError: + except (OSError, UnicodeDecodeError): return [] out: "list[str]" = [] for d in ("/opt/rocm/lib", "/opt/rocm/lib64"): @@ -232,13 +263,97 @@ def _wsl_system_rocm_lib_dirs() -> "list[str]": return out -# Plan-without-action re-prompt state (intent signal, caps, message) now lives -# in tool_call_parser, imported above under its old aliases. +def _bundled_hip_present(binary_dir: str) -> bool: + """True when a prebuilt bundle ships its own HIP backend library.""" + if not binary_dir: + return False + try: + # Glob the version suffix (libggml-hip.so, .so.0, .so.0.11.1) the same + # way the installer's runtime health check matches libggml-hip.so*. + return any(Path(str(binary_dir)).glob("libggml-hip.so*")) + except OSError: + return False + + +def _native_linux_system_rocm_lib_dirs(binary_dir: str = "") -> "list[str]": + """System ROCm lib dir(s) to prepend before a prebuilt's bundled HIP, on native Linux. + + The bundled bare-metal HIP runtime can mismatch the host amdkfd driver and crash + in hsa_init(); prepending the whole system ROCm lib dir loads a driver-matched, + version-consistent stack (libhsa-runtime64 / libamdhip64 / librocblas) ahead of it. + The whole dir is deliberate: mixing the bundle's rocBLAS with a different-version + system HIP/ROCR risks missing symbols. UNSLOTH_LLAMA_NO_SYSTEM_ROCM=1 keeps the pure + bundle (for a host whose system ROCm lacks this arch); no-op on WSL / non-Linux. + """ + if os.environ.get("UNSLOTH_LLAMA_NO_SYSTEM_ROCM") == "1": + return [] + if sys.platform != "linux" or os.path.exists("/dev/dxg"): + return [] + if not os.path.exists("/dev/kfd"): + return [] + if not _bundled_hip_present(binary_dir): + return [] + # Env-configured ROCm root first; /opt/rocm only as a fallback so a stale + # /opt/rocm doesn't shadow the driver-matching install these vars point at. + candidates = [] + for var in ("HIP_PATH", "HIP_PATH_57", "ROCM_PATH"): + val = os.environ.get(var) + if val: + candidates.append(val) + candidates.append("/opt/rocm") + out: "list[str]" = [] + seen: "set[str]" = set() + for base in candidates: + for lib_sub in ("lib", "lib64"): + d = os.path.join(base, lib_sub) + if d in seen: + continue + seen.add(d) + if os.path.exists(os.path.join(d, "libhsa-runtime64.so")) or os.path.exists( + os.path.join(d, "libhsa-runtime64.so.1") + ): + out.append(d) + # ROCm keeps LLVM's versioned runtime under /lib/llvm, so a + # lib64 host still finds it under lib. Probe both and keep them + # ahead of the bundle, else system libamd_comgr binds to the + # bundle's incompatible libLLVM.so.*. + for _sub in (lib_sub, "lib"): + llvm_lib = os.path.join(base, _sub, "llvm", "lib") + if llvm_lib not in seen and os.path.isdir(llvm_lib): + seen.add(llvm_lib) + out.append(llvm_lib) + return out + + +# Plan-without-action re-prompt state now lives in tool_call_parser (imported above). # Default max_tokens to the effective context when known. The floor is high # enough for reasoning-heavy GGUFs and max_tokens-omitting API clients. _DEFAULT_MAX_TOKENS_FLOOR = 32768 _DEFAULT_FIRST_TOKEN_TIMEOUT_S = 1200.0 # 20 min +# A transport error can arrive before the child is reapable; a request path cannot +# afford the 5s the background MTP reload spends on the same race. +_RESPAWN_REAP_GRACE_S = 1.0 + + +def _finalize_reasoning_only_cumulative( + cumulative: str, reasoning_text: str, finish_reason: Optional[str], promote_reasoning_only: bool +) -> str: + """Close a live thinking block and promote it only after a clean stop. + + Local inference streams cumulative snapshots. Replacing ``...`` with + bare reasoning at EOF makes the final snapshot shorter, so suffix-based + route consumers drop the intended fallback. Keep the snapshot append-only. + A length-truncated thought is not a final answer, so close it without + promotion and let the client surface the ``length`` terminal state. Raw + consumers that do not split reasoning from visible content can disable the + fallback to avoid returning the same reasoning twice. + """ + visible_fallback = ( + reasoning_text if promote_reasoning_only and finish_reason != "length" else "" + ) + return cumulative + "" + visible_fallback + # Only large streamed tool payloads get an early provisional card; render_html # is exempt because it needs immediate artifact feedback. @@ -248,12 +363,32 @@ _DEFAULT_STREAM_STALL_TIMEOUT_S = 120.0 # 2 min # loop). Structured delta.tool_calls are grammar-bounded by llama-server; text # parsed from content is not, so one runaway turn could fan out unbounded. _MAX_TOOL_CALLS_PER_TURN = 8 -_FORCED_REPEAT_PLAN_SIGNAL = re.compile( - r"\b(?:i\s+will|i'll|let\s+me|going\s+to|need\s+to|call|use|run|search|fetch|render)\b", +# Obligation phrasing INTENT_SIGNAL leaves alone ("I need to call ..."), paired with +# an action verb. Sentence-anchored: mid-sentence the same words are prose that names +# a tool ("The API I should invoke is foo() because ..."), and suppressing that loses +# a real answer. "should"/"must" sit outside the need|have|ought group because they +# take a bare infinitive. "invoke"/"query" stay out of the verb list: they read as +# technical prose far more often than as a stall. +_FORCED_PLAN_INTENT = re.compile( + r"(?:^|[.!?]\s+)\s*" + r"(?:i\s+(?:(?:need|have|ought)\s+to|should|must)|need\s+to|going\s+to|must|should)" + r"\s+(?:\w+\s+){0,2}?(?:call|use|run|search|fetch|render)\b", + re.I | re.M, +) +# "the answer is not in the context" announces a *missing* answer, so the negated +# forms are excluded or the plan behind them would ship as the final response. +_FINAL_ANSWER_SIGNAL = re.compile( + r"\b(?:final\s+answer|answer\s*:|here\s+is|here's|in\s+summary|result\s*:" + r"|(?:the\s+)?answer\s+is(?!\s+(?:not|unavailable|unknown|unclear|missing)\b))\b", re.I, ) -_FINAL_ANSWER_SIGNAL = re.compile( - r"\b(?:final\s+answer|answer\s*:|here\s+is|here's|in\s+summary|result\s*:)\b", +# A plan that pivots ("I should call web_search, but Tokyo is the capital") has an +# answer attached, so the turn must survive. Leaking a plan sentence is cosmetic; +# dropping an answer is not, so the doubtful case keeps the output. The pivot has to +# carry text of its own: "I should call web_search, though." answers nothing. +_ANSWER_PIVOT = re.compile( + r"\b(?:but|however|although|though|that\s+said|in\s+the\s+meantime|meanwhile)\b" + r"[\W_]*(?:\w+[\W_]+){1,}\w", re.I, ) @@ -345,14 +480,28 @@ def _held_rehearsal_tail_len(text: str, active_tools: list[dict]) -> int: return len(tail) if tail and _is_rehearsal_prefix(tail, active_tools) else 0 -def _should_suppress_forced_no_tool_output(text: str) -> bool: - """Suppress only repeated forced-turn planning text, not final answers.""" +def _should_suppress_forced_no_tool_output(text: str, previous: str = "") -> bool: + """Suppress only repeated forced-turn planning text, not final answers. + + ``previous`` is the stall text that triggered the nudge, so a retry that + moved on can be told from one that just said the same thing again. + """ stripped = text.strip() if not stripped or len(stripped) >= _REPROMPT_MAX_CHARS: return False if _FINAL_ANSWER_SIGNAL.search(stripped): return False - return _FORCED_REPEAT_PLAN_SIGNAL.search(stripped) is not None + plan = _FORCED_PLAN_INTENT.search(stripped) + if plan is not None: + # Only the plan itself is safe to drop; anything the turn pivots to after it + # is the answer the user is waiting for. + return _ANSWER_PIVOT.search(stripped[plan.end() :]) is None + if not _is_short_intent_without_action(stripped): + return False + # INTENT_SIGNAL also fires on lead-ins to a real answer ("Now I have the results. + # The capital is Tokyo."), so a bare intent match is a stall only when the retry + # adds nothing. No ``previous`` keeps the standalone "is this a stall?" contract. + return not previous or _is_reprompt_restatement(stripped, previous) # ── Pre-compiled patterns for GGUF shard detection ─────────── @@ -440,6 +589,23 @@ def _hf_offline_if_dns_dead(): os.environ.pop("TRANSFORMERS_OFFLINE", None) +try: + _SLOT_SAVE_MAX_BYTES = int(os.environ.get("UNSLOTH_SLOT_SAVE_MAX_BYTES") or (10 << 30)) +except ValueError: + _SLOT_SAVE_MAX_BYTES = 10 << 30 + +# The idle loop holds the lifecycle gate across a slot save, so a newly arriving +# request waits on the in-flight save's HTTP call. Bound it (was 120s) so a slow +# or stuck save can't stall the next request for minutes; best-effort save just +# falls back to a plain unload. Override with UNSLOTH_SLOT_SAVE_TIMEOUT (seconds). +try: + _SLOT_SAVE_HTTP_TIMEOUT = float(os.environ.get("UNSLOTH_SLOT_SAVE_TIMEOUT") or 30.0) +except ValueError: + _SLOT_SAVE_HTTP_TIMEOUT = 30.0 +if _SLOT_SAVE_HTTP_TIMEOUT <= 0: + _SLOT_SAVE_HTTP_TIMEOUT = 30.0 + + def _swa_cache_path() -> Path: home = os.environ.get("UNSLOTH_STUDIO_HOME") or os.environ.get("STUDIO_HOME") base = Path(home) if home else Path.home() / ".unsloth" / "studio" @@ -452,11 +618,11 @@ def _load_swa_cache() -> dict: if _SWA_CACHE is not None: return _SWA_CACHE try: - with open(_swa_cache_path()) as f: + with open(_swa_cache_path(), encoding = "utf-8-sig") as f: _SWA_CACHE = json.load(f) if not isinstance(_SWA_CACHE, dict): _SWA_CACHE = {} - except (FileNotFoundError, json.JSONDecodeError, OSError): + except (FileNotFoundError, json.JSONDecodeError, OSError, UnicodeDecodeError): _SWA_CACHE = {} return _SWA_CACHE @@ -466,10 +632,10 @@ def _save_swa_cache(cache: dict) -> None: path = _swa_cache_path() path.parent.mkdir(parents = True, exist_ok = True) tmp = path.with_suffix(".json.tmp") - with open(tmp, "w") as f: + with open(tmp, "w", encoding = "utf-8") as f: json.dump(cache, f, indent = 2, sort_keys = True) tmp.replace(path) - except OSError: + except (OSError, UnicodeDecodeError): pass @@ -495,8 +661,15 @@ def _swa_entry_from_layer_types(lt) -> Optional[object]: def _fetch_swa_entry_from_hf(repo_id: str) -> Optional[object]: try: from huggingface_hub import hf_hub_download - cfg_path = hf_hub_download(repo_id, "config.json", repo_type = "model") - with open(cfg_path) as f: + from utils.hf_cache_settings import active_hf_hub_cache + + cfg_path = hf_hub_download( + repo_id, + "config.json", + repo_type = "model", + cache_dir = active_hf_hub_cache(), + ) + with open(cfg_path, encoding = "utf-8-sig") as f: cfg = json.load(f) except Exception: return None @@ -897,6 +1070,7 @@ def _cached_hf_snapshot_file( filename: str, *, expected_size: Optional[int] = None, + cache_dir: Optional[str] = None, ) -> Optional[str]: """Return a cached snapshot file even when HF's current-ref probe misses it.""" if not filename: @@ -905,8 +1079,22 @@ def _cached_hf_snapshot_file( if not parts or any(part in (".", "..") for part in parts): return None try: - from utils.models.model_config import _iter_hf_cache_snapshots - for snap in _iter_hf_cache_snapshots(repo_id): + if cache_dir is None: + from utils.models.model_config import _iter_hf_cache_snapshots + snapshots = _iter_hf_cache_snapshots(repo_id) + else: + from hub.utils.hf_cache_state import iter_active_repo_cache_dirs + snapshots = ( + snapshot + for repo_dir in iter_active_repo_cache_dirs( + "model", + repo_id, + root = Path(cache_dir), + ) + for snapshot in (repo_dir / "snapshots").glob("*") + if snapshot.is_dir() + ) + for snap in snapshots: candidate = snap.joinpath(*parts) if not candidate.is_file(): continue @@ -998,6 +1186,293 @@ def _cached_colocated_split_main( return None +def _cached_variant_resolution(repo_id: str, hf_variant: str) -> tuple[Optional[str], list[str]]: + """Find a cached main GGUF and its shards for a variant.""" + candidate = next(_cached_variant_candidates(repo_id, hf_variant), None) + if candidate is None: + return None, [] + _, main, shards, _ = candidate + return main, shards + + +def _cached_variant_candidates( + repo_id: str, + hf_variant: str, + *, + require_mmproj: bool = False, +) -> Generator[tuple[str, str, list[str], Path], None, None]: + """Yield complete cached variant copies in snapshot preference order.""" + try: + from utils.models.model_config import _iter_hf_cache_snapshots + for snap in _iter_hf_cache_snapshots(repo_id): + cached_files = _gguf_snapshot_files(snap) + matches = _gguf_files_for_variant(cached_files, hf_variant) + if not matches: + continue + main = matches[0] + shards = _gguf_extra_shards(matches, main) + split = _SHARD_FULL_RE.match(main) + if split: + numbers = { + int(match.group(2)) + for path in [main, *shards] + if (match := _SHARD_FULL_RE.match(path)) + } + if numbers != set(range(1, int(split.group(3)) + 1)): + continue + main_path = snap.joinpath(*main.replace("\\", "/").split("/")) + if not main_path.is_file() or not _snapshot_has_all_shards( + str(main_path), main, shards, {} + ): + continue + if require_mmproj and not _pick_mmproj(cached_files): + continue + yield str(main_path), main, shards, snap + except Exception as e: + logger.debug(f"Cache lookup for variant failed: {e}") + + +def _cached_candidate_matches_revision_size( + repo_id: str, candidate: tuple[str, str, list[str], Path], hf_token: Optional[str] +) -> bool: + """Check cached byte sizes against the snapshot's own Hub revision. + + A snapshot pointer is normally published only after its blob is complete. + When the old revision is still queryable, also compare every weight file's + size so a manually truncated cache entry is not treated as reusable. If + metadata cannot be reached, retain the cache's normal offline semantics. + """ + main_path, main, shards, snap = candidate + paths = [main, *shards] + try: + from huggingface_hub import get_paths_info + infos = list( + get_paths_info( + repo_id, + paths, + revision = snap.name, + token = hf_token, + ) + ) + except Exception as e: + logger.debug( + "Could not size-check cached GGUF %s at revision %s: %s", + repo_id, + snap.name, + e, + ) + return True + + if not infos: + # The Hub answers an unknown (e.g. force-pushed away) revision with an + # empty result, not an error; treat it like unreachable metadata. + return True + expected_sizes = {info.path: info.size for info in infos if info.size is not None} + if any(path not in expected_sizes for path in paths): + return False + try: + if os.path.getsize(main_path) < expected_sizes[main]: + return False + except OSError: + return False + return _snapshot_has_all_shards(main_path, main, shards, expected_sizes) + + +def _cached_complete_candidate( + repo_id: str, gguf_filename: Optional[str], shards: list[str] +) -> Optional[tuple[str, str, list[str], Path]]: + """Return one complete exact-filename cache candidate with snapshot context.""" + if not gguf_filename: + return None + if shards: + main_path = _cached_colocated_split_main(repo_id, gguf_filename, shards, {}) + else: + m = _SHARD_FULL_RE.match(gguf_filename) + if m and int(m.group(3)) > 1: + return None + main_path = _cached_hf_snapshot_file(repo_id, gguf_filename) + if main_path is None: + return None + snap = _snapshot_dir_of(main_path) + if snap is None: + return None + return main_path, gguf_filename, shards, snap + + +def cached_gguf_for_load( + hf_repo: str, + hf_variant: Optional[str], + *, + require_mmproj: bool = False, + verify_sizes: bool = False, + hf_token: Optional[str] = None, +) -> Optional[str]: + """Return a cached GGUF that can be loaded without downloading.""" + if not hf_variant: + return None + hf_repo = _resolve_repo_id_casing(hf_repo) + for candidate in _cached_variant_candidates( + hf_repo, + hf_variant, + require_mmproj = require_mmproj, + ): + if verify_sizes and not _cached_candidate_matches_revision_size( + hf_repo, candidate, hf_token + ): + continue + return candidate[0] + return None + + +def _snapshot_dir_of(path: str) -> Optional[Path]: + """Return the HF cache snapshot containing path, if any.""" + try: + p = Path(os.path.abspath(path)) + except OSError: + return None + for ancestor in p.parents: + if ancestor.parent.name == "snapshots": + return ancestor + return None + + +def _hub_cache_dir_for_snapshot_path(path: Optional[str]) -> Optional[str]: + """Return the HF Hub cache root that owns a snapshot-contained path.""" + if not path: + return None + snapshot = _snapshot_dir_of(path) + if snapshot is None or snapshot.parent.name != "snapshots": + return None + return str(snapshot.parent.parent.parent) + + +def _companion_snapshot_sibling( + near_path: str, pick: Callable[[list[str]], Optional[str]] +) -> Optional[str]: + """Find a companion in the same snapshot as near_path.""" + snap = _snapshot_dir_of(near_path) + if snap is None: + return None + try: + sibling = pick(_gguf_snapshot_files(snap)) + except Exception: + return None + if not sibling: + return None + candidate = snap / sibling + return str(candidate) if candidate.is_file() else None + + +def _pick_mmproj(candidates: list[str]) -> Optional[str]: + mmproj_files = sorted( + f for f in candidates if f.lower().endswith(".gguf") and "mmproj" in Path(f).name.lower() + ) + if not mmproj_files: + return None + return next((f for f in mmproj_files if f.lower().endswith("-f16.gguf")), mmproj_files[0]) + + +def _hub_download_in_flight(hf_repo: str) -> bool: + try: + from hub.utils.download_registry import get_models_registry + return bool(get_models_registry().active_job_refs(hf_repo)) + except Exception: + return False + + +def _hub_download_blocks_gguf_load( + hf_repo: str, + hf_variant: Optional[str], + *, + require_mmproj: bool = False, + hf_token: Optional[str] = None, +) -> bool: + """Whether an active Hub job makes this GGUF load unsafe. + + Same-variant jobs can reclaim the stale snapshot a load would reuse, so + they always block. Other jobs block only when this load lacks a complete + cached copy and would write to the shared cache itself. + """ + try: + from hub.utils.download_registry import get_models_registry + + registry = get_models_registry() + if not registry.active_job_refs(hf_repo): + return False + if registry.has_active_variant(hf_repo, hf_variant): + return True + except Exception: + return False + return ( + cached_gguf_for_load( + hf_repo, + hf_variant, + require_mmproj = require_mmproj, + verify_sizes = True, + hf_token = hf_token, + ) + is None + ) + + +# Active GGUF loads by normalized repo ID. +_LOADS_IN_FLIGHT: dict[str, int] = {} +_LOADS_IN_FLIGHT_LOCK = threading.Lock() + + +@contextlib.contextmanager +def gguf_load_in_flight(hf_repo: Optional[str]): + """Track an HF GGUF load until the context exits.""" + key = (hf_repo or "").strip().lower() + if not key: + yield + return + with _LOADS_IN_FLIGHT_LOCK: + _LOADS_IN_FLIGHT[key] = _LOADS_IN_FLIGHT.get(key, 0) + 1 + try: + yield + finally: + with _LOADS_IN_FLIGHT_LOCK: + remaining = _LOADS_IN_FLIGHT.get(key, 1) - 1 + if remaining <= 0: + _LOADS_IN_FLIGHT.pop(key, None) + else: + _LOADS_IN_FLIGHT[key] = remaining + + +def hf_gguf_load_in_flight(hf_repo: str) -> bool: + """Return whether a GGUF load is active for hf_repo.""" + key = (hf_repo or "").strip().lower() + if not key: + return False + with _LOADS_IN_FLIGHT_LOCK: + return _LOADS_IN_FLIGHT.get(key, 0) > 0 + + +def _with_gguf_load_marker(load: Callable): + """Keep an HF repo marked for the full synchronous load call.""" + + @functools.wraps(load) + def wrapped(self, *args, **kwargs): + hf_repo = kwargs.get("hf_repo") + with gguf_load_in_flight(hf_repo): + if hf_repo and _hub_download_blocks_gguf_load( + hf_repo, + kwargs.get("hf_variant"), + require_mmproj = bool( + kwargs.get("is_vision") + and not extra_args_disable_mmproj(kwargs.get("extra_args")) + ), + hf_token = kwargs.get("hf_token"), + ): + raise RuntimeError( + f"'{hf_repo}' is currently being downloaded by the download manager" + ) + return load(self, *args, **kwargs) + + return wrapped + + def _gguf_extra_shards(files: Iterable[str], first_shard: str) -> list[str]: m = _SHARD_FULL_RE.match(first_shard) if not m: @@ -1082,9 +1557,24 @@ def _kv_bytes_per_elem(cache_type: Optional[str]) -> float: }.get((cache_type or "f16").strip().lower(), 2.0) +def _pad_kv_cells(cells: int) -> int: + return ((cells + 255) // 256) * 256 + + +def _kv_cache_cell_layout(n_ctx: int, n_parallel: int, kv_unified: bool) -> tuple[int, int, int]: + """Return llama.cpp's slot count, stream count, and cells per stream.""" + slots = max(1, n_parallel) + padded_ctx = _pad_kv_cells(n_ctx) + streams = 1 if kv_unified else slots + if padded_ctx <= 0: + return slots, streams, 0 + cells_per_stream = padded_ctx if kv_unified else _pad_kv_cells(padded_ctx // slots) + return slots, streams, cells_per_stream + + def _env_main_cache_type_for_budget(env: Optional[Mapping[str, str]] = None) -> Optional[str]: """Heavier of the inherited LLAMA_ARG_CACHE_TYPE_K/_V env types when it - exceeds the f16 default, else None. Studio emits --cache-type only for the + exceeds the f16 default, else None. Unsloth emits --cache-type only for the param/extras path, so a heavier env (f32) would otherwise reach the child unbudgeted; quantized env types stay over-reserved by f16 (-> None).""" e = os.environ if env is None else env @@ -1114,6 +1604,39 @@ def _extra_args_main_cache_type_for_budget(extra_args: Optional[Iterable[str]]) return max(candidates, key = _kv_bytes_per_elem) +def _effective_main_cache_types( + args: Optional[Iterable[str]], env: Optional[Mapping[str, str]] = None +) -> tuple[str, str]: + """Effective main K/V cache types after environment and CLI precedence.""" + source_env = os.environ if env is None else env + env_k = (source_env.get("LLAMA_ARG_CACHE_TYPE_K") or "f16").strip().lower() + env_v = (source_env.get("LLAMA_ARG_CACHE_TYPE_V") or "f16").strip().lower() + arg_k, arg_v = parse_cache_override_per_axis(args) + return ( + (arg_k or env_k).strip().lower(), + (arg_v or env_v).strip().lower(), + ) + + +def _planned_main_cache_types( + cache_type_kv: Optional[str], + extra_args: Optional[Iterable[str]], + env: Optional[Mapping[str, str]] = None, +) -> tuple[str, str]: + """Main K/V types the loader's managed flags and user extras will produce.""" + args = list(extra_args or ()) + emitted_type = _extra_args_main_cache_type_for_budget(args) or cache_type_kv + if emitted_type: + args = [ + "--cache-type-k", + emitted_type, + "--cache-type-v", + emitted_type, + *args, + ] + return _effective_main_cache_types(args, env) + + def _auto_mode_drops_mtp( req_mode: Optional[str], size_b: Optional[float], @@ -1152,28 +1675,95 @@ def _extra_args_set_spec_type(extra_args: Optional[Iterable[str]]) -> bool: return _extra_args_set_any_flag(extra_args, {"--spec-type", "--spec-default"}) -_GPU_OFFLOAD_OVERRIDE_FLAGS = frozenset({"-ngl", "--gpu-layers", "--n-gpu-layers", "-fit", "--fit"}) +# Layer-offload override detection. Single-sourced from llama_server_args, which +# also strips these (plus the MoE flags) from inherited extras; sharing the layer +# set keeps detection and stripping from drifting. +_GPU_OFFLOAD_OVERRIDE_FLAGS = _LAYER_OFFLOAD_FLAGS _THREAD_OVERRIDE_FLAGS = frozenset({"-t", "--threads"}) - - -def _extra_arg_flag_name(token: str) -> Optional[str]: - if not token.startswith("-") or token in {"-", "--"}: - return None - if len(token) >= 2 and (token[1].isdigit() or token[1] == "."): - return None - return token.split("=", 1)[0] +# common_params defaults in the bundled llama.cpp runtime. +_DEFAULT_LLAMA_N_BATCH = 2048 +_DEFAULT_LLAMA_N_UBATCH = 512 +_LLAMA_ARG_TRUE_VALUES = frozenset({"on", "enabled", "true", "1"}) +_LLAMA_ARG_FALSE_VALUES = frozenset({"off", "disabled", "false", "0"}) +_LLAMA_ARG_AUTO_VALUES = frozenset({"auto", "-1"}) +_LLAMA_ARG_TRUE_OR_AUTO_VALUES = _LLAMA_ARG_TRUE_VALUES | _LLAMA_ARG_AUTO_VALUES +_LLAMA_ARG_TRUE_FALSE_AUTO_VALUES = _LLAMA_ARG_TRUE_OR_AUTO_VALUES | _LLAMA_ARG_FALSE_VALUES def _extra_args_set_any_flag(extra_args: Optional[Iterable[str]], flags: Collection[str]) -> bool: if not extra_args: return False for raw in extra_args: - flag = _extra_arg_flag_name(str(raw)) + flag = _flag_name(str(raw)) if flag in flags: return True return False +def _swa_full_from_args_or_env( + extra_args: Optional[Iterable[str]], env: Optional[Mapping[str, str]] = None +) -> bool: + """Whether llama.cpp receives the enable-only full-size SWA option.""" + if _extra_args_set_any_flag(extra_args, {"--swa-full"}): + return True + value = (os.environ if env is None else env).get("LLAMA_ARG_SWA_FULL") + return value in _LLAMA_ARG_TRUE_VALUES + + +def _kv_unified_from_args( + extra_args: Optional[Iterable[str]], + default: bool = False, + env: Optional[Mapping[str, str]] = None, +) -> bool: + """Resolve llama.cpp's environment and last-wins unified KV flags.""" + enabled = False + value = (os.environ if env is None else env).get("LLAMA_ARG_KV_UNIFIED") + if value in _LLAMA_ARG_TRUE_VALUES: + enabled = True + elif value in _LLAMA_ARG_FALSE_VALUES: + enabled = False + if default: + # Studio's managed --kv-unified flag is appended after environment + # parsing and before user extras. + enabled = True + for raw in extra_args or (): + flag = _flag_name(str(raw)) + if flag in {"-kvu", "--kv-unified"}: + enabled = True + elif flag in {"-no-kvu", "--no-kv-unified"}: + enabled = False + return enabled + + +def _flash_attn_enabled_from_args( + args: Optional[Iterable[str]], + default: bool = True, + env: Optional[Mapping[str, str]] = None, +) -> bool: + """Resolve llama.cpp's environment and last-wins flash-attention settings.""" + enabled = default + # llama.cpp applies LLAMA_ARG_FLASH_ATTN before parsing argv (arg.cpp set_env), + # so the CLI still wins. --flash-attn has no args_neg, so no LLAMA_ARG_NO_ twin. + value = (os.environ if env is None else env).get("LLAMA_ARG_FLASH_ATTN") + if value in _LLAMA_ARG_FALSE_VALUES: + enabled = False + elif value in _LLAMA_ARG_TRUE_OR_AUTO_VALUES: + enabled = True + values = [str(arg) for arg in args] if args else [] + for i, raw in enumerate(values): + if _flag_name(raw) not in {"-fa", "--flash-attn"}: + continue + _, eq, inline = raw.partition("=") + value = inline if eq else "on" + if not eq and i + 1 < len(values) and values[i + 1] in _LLAMA_ARG_TRUE_FALSE_AUTO_VALUES: + value = values[i + 1] + if value in _LLAMA_ARG_FALSE_VALUES: + enabled = False + elif value in _LLAMA_ARG_TRUE_OR_AUTO_VALUES: + enabled = True + return enabled + + def _effective_spec_type( extra_args: Optional[Iterable[str]], env: Optional[Mapping[str, str]] = None ) -> Optional[str]: @@ -1185,7 +1775,8 @@ def _effective_spec_type( cli_present = False cli_value: Optional[str] = None for i, raw in enumerate(args): - flag, eq, inline = raw.partition("=") + flag = _flag_name(raw) + _, eq, inline = raw.partition("=") if flag == "--spec-default": cli_present = True cli_value = "default" @@ -1229,7 +1820,8 @@ def _extra_args_spec_draft_n_max(extra_args: Optional[Iterable[str]]) -> Optiona args = [str(a) for a in extra_args] found: Optional[int] = None for i, raw in enumerate(args): - flag, eq, inline = raw.partition("=") + flag = _flag_name(raw) + _, eq, inline = raw.partition("=") if flag not in ("--spec-draft-n-max", "--draft-max"): continue value = inline if eq else (args[i + 1] if i + 1 < len(args) else "") @@ -1259,7 +1851,8 @@ def _extra_args_mtp_draft_path( args = [str(a) for a in extra_args] if extra_args else [] found: Optional[str] = None for i, raw in enumerate(args): - flag, eq, inline = raw.partition("=") + flag = _flag_name(raw) + _, eq, inline = raw.partition("=") if flag not in flags: continue value = inline if eq else (args[i + 1] if i + 1 < len(args) else "") @@ -1283,7 +1876,8 @@ def _extra_args_draft_cache_types( k_type: Optional[str] = None v_type: Optional[str] = None for i, raw in enumerate(args): - flag, eq, inline = raw.partition("=") + flag = _flag_name(raw) + _, eq, inline = raw.partition("=") if flag not in k_flags and flag not in v_flags: continue value = inline if eq else (args[i + 1] if i + 1 < len(args) else "") @@ -1315,7 +1909,8 @@ def _extra_args_draft_offloaded_to_cpu( last_ngl: Optional[str] = None last_dev: Optional[str] = None for i, raw in enumerate(args): - flag, eq, inline = raw.partition("=") + flag = _flag_name(raw) + _, eq, inline = raw.partition("=") value = inline if eq else (args[i + 1] if i + 1 < len(args) else "") if flag in ngl_flags: last_ngl = value @@ -1337,31 +1932,61 @@ def _extra_args_draft_offloaded_to_cpu( def _extra_args_n_ubatch( - extra_args: Optional[Iterable[str]], env: Optional[Mapping[str, str]] = None + extra_args: Optional[Iterable[str]], + env: Optional[Mapping[str, str]] = None, + n_ctx: Optional[int] = None, ) -> Optional[int]: - """Physical micro-batch from extras (--ubatch-size/-ub) else the LLAMA_ARG_UBATCH - env, else None. It sizes the compute-graph buffer, so an override must reach - the VRAM reserve.""" + """Effective ubatch after llama.cpp normalizes it, or None at defaults.""" + values = { + "batch": _DEFAULT_LLAMA_N_BATCH, + "ubatch": _DEFAULT_LLAMA_N_UBATCH, + } + source_env = os.environ if env is None else env + overridden = False + for key, env_name in ( + ("batch", "LLAMA_ARG_BATCH"), + ("ubatch", "LLAMA_ARG_UBATCH"), + ): + raw = source_env.get(env_name) + if raw: + try: + values[key] = int(raw) + overridden = True + except (TypeError, ValueError): + pass + args = [str(a) for a in extra_args] if extra_args else [] - found: Optional[int] = None + flags = { + "-b": "batch", + "--batch-size": "batch", + "-ub": "ubatch", + "--ubatch-size": "ubatch", + } for i, raw in enumerate(args): - flag, eq, inline = raw.partition("=") - if flag not in ("--ubatch-size", "-ub"): + flag = _flag_name(raw) + _, eq, inline = raw.partition("=") + key = flags.get(flag) + if key is None: continue value = inline if eq else (args[i + 1] if i + 1 < len(args) else "") try: - found = int(value) + values[key] = int(value) + overridden = True except (TypeError, ValueError): continue - if found is not None: - return found - raw = (os.environ if env is None else env).get("LLAMA_ARG_UBATCH") - if raw: - try: - return int(raw) - except (TypeError, ValueError): - pass - return None + if not overridden: + return None + + # common_params stores signed values, then llama_context_params converts + # them to uint32_t. A zero ubatch means "use batch"; the context then caps + # ubatch at batch size. + batch = values["batch"] & 0xFFFFFFFF + raw_ubatch = values["ubatch"] + ubatch = batch if raw_ubatch == 0 else raw_ubatch & 0xFFFFFFFF + effective = min(batch, ubatch) + if n_ctx is not None and n_ctx > 0: + effective = min(effective, n_ctx) + return effective def _build_ngram_mod_flags( @@ -1403,7 +2028,7 @@ def _build_ngram_mod_flags( return [] -# Canonical Speculative Decoding modes exposed by the Studio chat UI. +# Canonical Speculative Decoding modes exposed by the Unsloth chat UI. # Dropdown renders five (auto, mtp, ngram, mtp+ngram, off); the load API # also accepts legacy values the original Switch and external callers emit # (default, draft-mtp, ngram-mod, ngram-simple). @@ -1452,7 +2077,7 @@ def _backfill_usage_from_timings(usage, timings): """Synthesize ``usage`` from llama-server's ``timings`` when the OpenAI-style usage block is missing or reports zero tokens. - The Studio chat UI computes generation t/s from + The Unsloth chat UI computes generation t/s from ``meta.usage.completion_tokens / totalStreamTime``. llama-server always populates ``timings.predicted_n`` (true decoded count) and ``timings.prompt_n``, but the final SSE chunk's ``usage`` can be absent @@ -1525,7 +2150,7 @@ def _llama_lib_dir(binary: str) -> Path: def _is_external_link(path: Path) -> bool: """True when ``path`` is a --with-llama-cpp-dir local link: a POSIX symlink or a Windows directory junction / reparse point. Such a link resolves into - the user's own llama.cpp checkout, which Studio does not own.""" + the user's own llama.cpp checkout, which Unsloth does not own.""" try: if os.path.islink(path): return True @@ -1546,10 +2171,16 @@ def _is_external_link(path: Path) -> bool: # Map OpenAI-style names to the values the model was trained on. Module-level # so duck-typed engine stand-ins in tests do not need the attribute. _INKLING_REASONING_EFFORT = { - "none": 0.0, "minimal": 0.2, "low": 0.2, "medium": 0.7, - "high": 0.9, "xhigh": 0.99, "max": 0.99, + "none": 0.0, + "minimal": 0.1, + "low": 0.2, + "medium": 0.7, + "high": 0.9, + "xhigh": 0.99, + "max": 0.99, } + def _coerce_reasoning_effort(architecture, kwargs: dict) -> dict: if architecture == "inkling": effort = kwargs.get("reasoning_effort") @@ -1599,6 +2230,8 @@ class LlamaCppBackend: self._effective_context_length: Optional[int] = None self._max_context_length: Optional[int] = None self._effective_parallel_slots: int = 1 + # --parallel the last load asked for, before any fit-time reduction. + self._requested_n_parallel: int = 1 self._chat_template: Optional[str] = None self._chat_template_override: Optional[str] = None self._supports_reasoning: bool = False @@ -1610,6 +2243,21 @@ class LlamaCppBackend: self._cache_type_kv: Optional[str] = None # Whether --split-mode tensor was applied on the active load. self._tensor_parallel: bool = False + # GPU memory strategy applied on the active load ("auto"/"manual"). + self._gpu_memory_mode: str = "auto" + # Manual-mode load options (echoed back so the UI round-trips them). + self._gpu_layers: int = -1 + # MoE expert layers to keep on CPU (--n-cpu-moe); 0 = none. + self._n_cpu_moe: int = 0 + # Relative model share per GPU (--tensor-split), in GPU order; None = + # default (llama.cpp splits by free VRAM). + self._tensor_split: Optional[List[float]] = None + # User-picked physical GPU indices (None = automatic selection). + self._gpu_ids: Optional[List[int]] = None + # RAW requested GPU pin, before the fit narrowed it. self._gpu_ids records the + # EFFECTIVE (fit-narrowed) pin for /status; dedupe compares this raw value so a + # [0, 1] narrowed to [0] and re-sent as [0, 1] still matches (#7239). + self._requested_gpu_ids: Optional[List[int]] = None # Layer load kept multi-GPU only to honor a downgraded tensor request, so a # later explicit tensor-off reloads instead of deduping to it (#6659). self._layer_preserves_tensor_intent: bool = False @@ -1624,6 +2272,11 @@ class LlamaCppBackend: self._spec_draft_n_max: Optional[int] = None # KV-cache estimation fields (populated by _read_gguf_metadata) self._n_layers: Optional[int] = None + # MoE metadata (populated by _read_gguf_metadata): expert count (>0 = + # MoE) and leading dense-layer count (offsets --n-cpu-moe, which counts + # from layer 0). See the n_moe_layers property. + self._n_experts: Optional[int] = None + self._leading_dense_block_count: Optional[int] = None self._n_kv_heads: Optional[int] = None self._n_kv_heads_by_layer: Optional[list[int]] = None self._n_heads: Optional[int] = None @@ -1656,6 +2309,9 @@ class LlamaCppBackend: # Serialises mid-session respawns so many generations hitting a killed # server trigger at most one reload (see _respawn_if_dead). self._respawn_lock = threading.Lock() + # Bumped by every unload. load_model clears _cancel_event, so a respawn that + # raced an unload needs a signal that survives the clear (see _respawn_if_dead). + self._unload_epoch = 0 # Set by the in-app updater while it swaps prebuilt binaries; load_model() # rejects fast so no server starts from a half-swapped binary. self._llama_update_in_progress = False @@ -1675,7 +2331,7 @@ class LlamaCppBackend: # observes it (direct proxy endpoints, or nothing in flight). self._mtp_watchdog_thread: Optional[threading.Thread] = None self._mtp_watchdog_stop = threading.Event() - # True when the launch actually runs MTP+tensor (Studio- or user/env-driven); + # True when the launch actually runs MTP+tensor (Unsloth- or user/env-driven); # gates the probe, watchdog, and recovery so pass-through MTP is covered. self._mtp_runtime_fallback_active = False self._stdout_lines: list[str] = [] @@ -1685,6 +2341,20 @@ class LlamaCppBackend: self._llama_log_path: Optional[Path] = None self._cancel_event = threading.Event() self._api_key: Optional[str] = None + self._slot_save_dir: Optional[str] = None + self._slot_save_binary: Optional[tuple[str, int]] = None + # (gguf_identity, launch_fingerprint) snapshotted at load, so a later slot + # save can tell whether the model files were swapped on disk since load. + self._slot_loaded_identity: Optional[tuple] = None + self._prompt_cache_disabled: bool = False + self._swa_full: bool = False + self._kv_cache_unified: bool = False + self._n_ubatch: int = self._DEFAULT_N_UBATCH + self._flash_attn_enabled: bool = True + self._effective_cache_types: tuple[str, str] = ("f16", "f16") + # Total KV allocation context across all slots. _effective_context_length + # becomes the per-slot request limit after /props reconciliation. + self._kv_cache_context_total: Optional[int] = None # True once a probe has completed; cleared on transient failure. self._is_audio: bool = False self._audio_type: Optional[str] = None @@ -1737,6 +2407,11 @@ class LlamaCppBackend: """True when the loaded GGUF is a block-diffusion model (DiffusionGemma).""" return self._is_diffusion + @property + def swa_full(self) -> bool: + """Whether the active llama-server received full-size SWA mode.""" + return self._swa_full + @property def hf_variant(self) -> Optional[str]: return self._hf_variant @@ -1793,6 +2468,17 @@ class LlamaCppBackend: slots = 1 return max(1, slots) + @property + def requested_parallel_slots(self) -> int: + """--parallel the last load asked for, before any fit-time reduction. + The reload dedupe compares requested-vs-requested (like requested_n_ctx); + the effective count would reload forever after a fitter reduction.""" + try: + slots = int(getattr(self, "_requested_n_parallel", 1)) + except (TypeError, ValueError): + slots = 1 + return max(1, slots) + @property def max_context_length(self) -> Optional[int]: """Return the largest context that fits on this hardware at load time. @@ -1818,6 +2504,8 @@ class LlamaCppBackend: def _reset_effective_parallel_slots(self) -> None: self._effective_parallel_slots = 1 + # Cleared with the effective count so a stale value can't skew the dedupe. + self._requested_n_parallel = 1 @staticmethod def _read_rss_bytes(pid: int) -> Optional[int]: @@ -1968,7 +2656,8 @@ class LlamaCppBackend: if self._reasoning_style == "reasoning_effort": return _coerce_reasoning_effort( getattr(self, "_architecture", None), - {"reasoning_effort": "high" if enable_thinking else "low"}) + {"reasoning_effort": "high" if enable_thinking else "low"}, + ) return {"enable_thinking": enable_thinking} def _request_reasoning_kwargs( @@ -2043,6 +2732,119 @@ class LlamaCppBackend: """Whether --split-mode tensor is active on the loaded server.""" return self._tensor_parallel + @property + def gpu_memory_mode(self) -> str: + """Active GPU memory strategy: 'auto' or 'manual' (gpu_layers < 0 = Auto/--fit, >= 0 = pinned).""" + return self._gpu_memory_mode + + @property + def gpu_layers(self) -> int: + """Requested --gpu-layers for manual mode (-1 when not manual).""" + return self._gpu_layers + + @property + def n_cpu_moe(self) -> int: + """MoE expert layers manual mode kept on CPU (--n-cpu-moe); 0 = none.""" + return self._n_cpu_moe + + @property + def tensor_split(self) -> Optional[List[float]]: + """Manual-mode relative model share per GPU (--tensor-split); None = + default (split by free VRAM).""" + return self._tensor_split + + @property + def gpu_ids(self) -> Optional[List[int]]: + """User-picked physical GPU indices, or None for automatic selection.""" + return self._gpu_ids + + @property + def requested_gpu_ids(self) -> Optional[List[int]]: + """RAW requested GPU pin (before the fit narrowed it), or None for auto. + gpu_ids echoes the EFFECTIVE pin for /status.""" + return self._requested_gpu_ids + + def matches_gpu_ids(self, gpu_ids: Optional[List[int]]) -> bool: + """Whether a requested pin is already satisfied by the active runner. + + A regular GGUF load may narrow the requested placement pool to the + smallest fitting subset. Accept both the original request and the + effective status-echoed subset so either can round-trip without a + needless reload. Diffusion drives one device and keeps its existing + lowest-device normalization. + """ + if self._is_diffusion: + requested = [sorted(int(x) for x in gpu_ids)[0]] if gpu_ids else None + return requested == (self._gpu_ids or None) + + requested = sorted(int(x) for x in gpu_ids) if gpu_ids else None + raw = self._requested_gpu_ids or None + effective = self._gpu_ids or None + return requested == raw or requested == effective + + def _record_matching_gpu_request(self, gpu_ids: Optional[List[int]]) -> None: + """Adopt the caller's explicit pool after a full already-loaded match. + + Matching an effective subset avoids a reload, but the incoming request + is still the user's latest placement intent. Record it so status and a + later reload do not restore GPUs the user just removed. + """ + if self._is_diffusion: + self._requested_gpu_ids = [sorted(int(x) for x in gpu_ids)[0]] if gpu_ids else None + else: + self._requested_gpu_ids = sorted(int(x) for x in gpu_ids) if gpu_ids else None + if self._last_load_kwargs is not None: + self._last_load_kwargs["gpu_ids"] = ( + list(self._requested_gpu_ids) if self._requested_gpu_ids else None + ) + + @property + def n_layers(self) -> Optional[int]: + """Model layer count (GGUF block_count), or None if unknown.""" + return self._n_layers + + @property + def n_moe_layers(self) -> int: + """Number of MoE expert layers (the --n-cpu-moe ceiling), 0 if not MoE. + + block_count minus the leading dense layers (which carry no experts): + --n-cpu-moe counts from layer 0, so those dense layers are no-ops. + """ + if not self._n_experts or not self._n_layers: + return 0 + return max(0, self._n_layers - (self._leading_dense_block_count or 0)) + + @staticmethod + def _resolve_cpu_moe_flag( + n_cpu_moe: int, n_moe_layers: int, leading_dense: int + ) -> Optional[int]: + """The --n-cpu-moe value (absolute first-N layers), or None to omit it. + + Clamps the requested count to the model's MoE layers, then offsets past + the leading dense layers (--n-cpu-moe counts from layer 0). Returns None + for nothing-to-offload (0 requested) or a non-MoE model. + """ + if n_cpu_moe <= 0 or n_moe_layers <= 0: + return None + return leading_dense + min(n_cpu_moe, n_moe_layers) + + @staticmethod + def _sanitize_tensor_split(tensor_split: Optional[List[float]]) -> List[float]: + """Per-GPU shares with negative and non-finite entries clamped to 0. + + A direct caller's negative entry would launch a placement different + from the ratio the UI showed, and inf would pass a plain ``> 0`` total + gate and emit ``--tensor-split inf,...``. Returns [] for input that + can't be read as floats (the length gate at the call site then drops + the split). + """ + try: + return [ + x if math.isfinite(x) and x > 0.0 else 0.0 for x in (float(v) for v in tensor_split) + ] + except (TypeError, ValueError, OverflowError): + return [] + @property def layer_preserves_tensor_intent(self) -> bool: """True when a downgraded tensor request kept this layer load multi-GPU.""" @@ -2067,7 +2869,7 @@ class LlamaCppBackend: @staticmethod def _resolved_studio_root_and_is_legacy() -> "tuple[Optional[Path], bool]": - """Resolve the Studio install root and classify it as the legacy + """Resolve the Unsloth install root and classify it as the legacy ~/.unsloth/studio root vs. a custom (env/venv-inferred) root. Returns (resolved_root, is_legacy). On any import/resolution failure the @@ -2239,15 +3041,18 @@ class LlamaCppBackend: "found": False, "mtp_token": None, "supports_mtp": False, + "mtp_probe_inconclusive": True, "ngram_mod_flavor": None, "supports_ngram_mod": False, "spec_draft_n_max_flag": None, "supports_kv_unified": False, "supports_fit_ctx": False, + "supports_fit_target": False, "supports_cache_ram": False, "supports_ctx_checkpoints": False, "supports_no_cache_prompt": False, "supports_metrics": False, + "supports_slot_save": False, } try: mtime = int(Path(bin_path).stat().st_mtime) @@ -2263,21 +3068,28 @@ class LlamaCppBackend: spec_draft_n_max_flag: Optional[str] = None supports_kv_unified = False supports_fit_ctx = False + supports_fit_target = False supports_cache_ram = False supports_ctx_checkpoints = False supports_no_cache_prompt = False supports_metrics = False + supports_slot_save = False + saw_spec_type = False + probe_ok = False + help_text = "" try: probe_env = cls._llama_server_env_for_binary(bin_path) result = subprocess.run( [bin_path, "--help"], capture_output = True, text = True, + encoding = "utf-8", errors = "replace", timeout = 10, check = False, env = probe_env, ) + probe_ok = result.returncode == 0 help_text = (result.stdout or "") + "\n" + (result.stderr or "") # Split into per-flag blocks (each --flag line + its indented # continuation), so the "argument has been removed" description @@ -2322,17 +3134,19 @@ class LlamaCppBackend: return False return "argument has been removed" not in desc - # MTP token from the --spec-type line. - spec_line = "" - for line in help_text.splitlines(): - if "--spec-type" in line: - spec_line = line - break - # PR #22673 used draft-mtp; later renamed to mtp. - if "draft-mtp" in spec_line: - mtp_token = "draft-mtp" - elif re.search(r"[|,\[]mtp[|,\]]", spec_line): - mtp_token = "mtp" + # MTP token from the full --spec-type help block (decl + indented + # continuation). First-line-only probing missed builds putting the + # enum on the next line (#7302). Prefer draft-mtp (PR #22673) over mtp. + spec_help = blocks.get("--spec-type") or "" + if not spec_help: + # Fallback: join --spec-type lines, avoiding incidental "mtp" in --help. + spec_help = "\n".join( + line for line in help_text.splitlines() if "--spec-type" in line + ) + mtp_token = cls._mtp_token_from_spec_help(spec_help) + # Only a resolved --spec-type block confirms missing MTP; empty/crash + # leaves saw_spec_type False so supports_mtp fails open. + saw_spec_type = bool(spec_help.strip()) and "--spec-type" in spec_help # ngram-mod flag flavor. Post-rename builds advertise both new # args (real) and legacy ones (stubs); pre-rename builds only @@ -2360,30 +3174,67 @@ class LlamaCppBackend: supports_kv_unified = _is_real("--kv-unified") supports_fit_ctx = _is_real("--fit-ctx") + supports_fit_target = _is_real("--fit-target") supports_cache_ram = _is_real("--cache-ram") supports_ctx_checkpoints = _is_real("--ctx-checkpoints") supports_no_cache_prompt = _is_real("--no-cache-prompt") supports_metrics = _is_real("--metrics") + supports_slot_save = _is_real("--slot-save-path") except (OSError, subprocess.SubprocessError) as exc: logger.debug(f"llama-server --help probe failed: {exc}") + saw_spec_type = False + probe_ok = False + help_text = "" + + help_nonempty = bool(help_text.strip()) + # Confirmed only when a successful --help lists a --spec-type block with + # mtp/draft-mtp; nonempty --help without it is a definitive pre-spec + # binary; failed/empty probes stay inconclusive (#7302). + if saw_spec_type and probe_ok: + supports_mtp = mtp_token is not None + mtp_probe_inconclusive = False + elif help_nonempty and probe_ok: + supports_mtp = False + mtp_probe_inconclusive = False + else: + supports_mtp = False + mtp_probe_inconclusive = True info = { "found": True, "mtp_token": mtp_token, - "supports_mtp": mtp_token is not None, + "supports_mtp": supports_mtp, + "mtp_probe_inconclusive": mtp_probe_inconclusive, "ngram_mod_flavor": ngram_mod_flavor, "supports_ngram_mod": ngram_mod_flavor is not None, "spec_draft_n_max_flag": spec_draft_n_max_flag, "supports_kv_unified": supports_kv_unified, "supports_fit_ctx": supports_fit_ctx, + "supports_fit_target": supports_fit_target, "supports_cache_ram": supports_cache_ram, "supports_ctx_checkpoints": supports_ctx_checkpoints, "supports_no_cache_prompt": supports_no_cache_prompt, "supports_metrics": supports_metrics, + "supports_slot_save": supports_slot_save, } cls._capability_cache[cache_key] = info return info + @staticmethod + def _mtp_token_from_spec_help(spec_help: str) -> Optional[str]: + """Extract ``draft-mtp`` / ``mtp`` from a ``--spec-type`` help snippet. + + Prefers ``draft-mtp`` (llama.cpp PR #22673) over the later bare ``mtp`` + rename. Returns ``None`` when neither token appears as an enum value. + """ + text = spec_help or "" + if "draft-mtp" in text: + return "draft-mtp" + # Bare `mtp` enum token (`|mtp|`, `,mtp,`, ...), not a substring. + if re.search(r"(?physical mapping.""" try: import torch - is_rocm = getattr(torch.version, "hip", None) is not None + + # Same ROCm detection as _emit_child_gpu_visibility: AMD SDK wheels + # leave version.hip unset but encode "rocm" in __version__. The two + # must agree, else an inherited ROCR mask reads back as "no mask", + # ordinal 0 is labelled physical 0, and the child's new ROCR pin + # re-exposes the GPU the inherited mask was hiding. + is_rocm = ( + getattr(torch.version, "hip", None) is not None + or "rocm" in getattr(torch, "__version__", "").lower() + ) except Exception: is_rocm = False if is_rocm: hip_v = os.environ.get("HIP_VISIBLE_DEVICES") - rocr_v = os.environ.get("ROCR_VISIBLE_DEVICES") + # ROCR_VISIBLE_DEVICES is a Linux ROCr variable; Windows HIP has no + # ROCr layer, so a stray ROCR var there does not mask the runtime and + # must not be read as the ordinal->physical mapping (mirrors the + # Windows gate in _emit_child_gpu_visibility). + rocr_v = None if sys.platform == "win32" else os.environ.get("ROCR_VISIBLE_DEVICES") cvd = ( hip_v if hip_v is not None @@ -2460,9 +3324,107 @@ class LlamaCppBackend: except ValueError: return None + @staticmethod + def _emit_child_gpu_visibility( + env: dict, + pinned: str, + *, + prefer_rocr: bool = False, + ) -> None: + """Write the child's GPU visibility mask: CUDA, plus a ROCm mirror on AMD + (masking only CUDA_VISIBLE_DEVICES leaves an AMD child seeing every GPU). + + Default: HIP_VISIBLE_DEVICES, clearing any inherited ROCR mask so the two + can't stack (ROCR re-indexes from 0, then a non-zero HIP pin points out of + range, HIP sees 0 devices, and llama.cpp falls back to CPU). + + prefer_rocr masks at the ROCr/HSA layer instead (clearing HIP). A HIP mask + filters only AFTER the HSA runtime enumerates every agent, and that + enumeration segfaults at startup on a GPU the build has no kernels for + (e.g. a gfx1036 iGPU under a gfx103X prebuilt: that bundle maps only + gfx1030/1031/1032/1034), before llama-server logs a line. ROCR drops the + device at the driver layer, consuming physical ids. + The CPU-only sentinel ("-1") has no portable ROCR spelling, so it keeps + the HIP mask. Windows keeps the HIP mask too: ROCR_VISIBLE_DEVICES is a + Linux ROCr variable (Windows HIP has no ROCr layer), so the ROCR pin + would be dead there while the cleared HIP mask stops selecting.""" + env["CUDA_VISIBLE_DEVICES"] = pinned + try: + import torch as _torch + + # torch.version.hip is set on ROCm, None on CUDA; AMD SDK wheels may + # leave it unset but encode "rocm" in __version__ (mirrors detect_hardware). + if ( + getattr(_torch.version, "hip", None) is not None + or "rocm" in getattr(_torch, "__version__", "").lower() + ): + if prefer_rocr and pinned != "-1" and sys.platform != "win32": + env["ROCR_VISIBLE_DEVICES"] = pinned + env.pop("HIP_VISIBLE_DEVICES", None) + # ROCR re-indexes the visible agents from 0, and with HIP + # cleared HIP honours CUDA_VISIBLE_DEVICES -- so it must carry + # the post-ROCR ordinals (0..N-1), not the physical ids, else a + # non-zero pick points out of range and HIP sees 0 devices (the + # same stacking the default path avoids by clearing ROCR). + env["CUDA_VISIBLE_DEVICES"] = ",".join( + str(i) for i in range(len(pinned.split(","))) + ) + else: + env["HIP_VISIBLE_DEVICES"] = pinned + env.pop("ROCR_VISIBLE_DEVICES", None) + except Exception as e: + logger.debug("Failed to set ROCm visibility env vars for child: %s", e) + + @staticmethod + def _pin_visible_gpu_order_for_split(env: dict) -> None: + """Pin the child's GPU enumeration to the picker's order for a manual + ``--tensor-split`` across the whole visible set. CUDA's default + FASTEST_FIRST enumeration applies the shares to the wrong cards on + heterogeneous hosts (#5025), and CUDA_DEVICE_ORDER only fixes the + numbering base: an inherited numeric visibility mask ALSO defines + enumeration order, so a reordered parent mask (CUDA_VISIBLE_DEVICES=3,1) + would still hand the shares to the wrong cards. The UI built the split + positionally over get_backend_visible_gpu_info's device list (ascending + physical via nvidia-smi, inherited mask order on the torch fallback), so + re-emit the same set in that report order -- not an assumed ascending + sort. The visible set itself never changes. No mask, an empty mask, or a + UUID/MIG mask (which resolves to None) is left alone -- the multi-GPU + controls are hidden for the latter.""" + env["CUDA_DEVICE_ORDER"] = "PCI_BUS_ID" + inherited = LlamaCppBackend._resolve_visible_physical_ids() + if not inherited: + return + order = None + try: + from utils.hardware import get_backend_visible_gpu_info + info = get_backend_visible_gpu_info() + if info.get("available") and info.get("index_kind") == "physical": + reported = [d["index"] for d in info.get("devices", [])] + if sorted(reported) == sorted(inherited): + order = reported + except Exception as e: + logger.debug("Could not read reported GPU order for split pin: %s", e) + if order is None: + order = sorted(inherited) + # Re-emit at the layer that produced the mapping. A parent masked only + # via ROCR_VISIBLE_DEVICES hides agents at the driver layer, and the + # default HIP re-emission clears that mask -- HSA then enumerates every + # agent again and can segfault at startup on an unsupported GPU the + # parent was hiding (the crash prefer_rocr exists to avoid). Linux-only, + # mirroring _resolve_visible_physical_ids: on Windows a stray ROCR var + # is dead and was not the mapping's source. + prefer_rocr = ( + sys.platform != "win32" + and env.get("HIP_VISIBLE_DEVICES") is None + and env.get("ROCR_VISIBLE_DEVICES") is not None + ) + LlamaCppBackend._emit_child_gpu_visibility( + env, ",".join(str(i) for i in order), prefer_rocr = prefer_rocr + ) + @staticmethod def _amd_apu_wants_unified_memory(gpu_indices = None) -> bool: - """True only for AMD unified-memory APUs (gfx1150/gfx1151), where + """True only for AMD unified-memory APUs (gfx1150/gfx1151/gfx1152), where GGML_CUDA_ENABLE_UNIFIED_MEMORY lets llama.cpp use shared system RAM (it hurts discrete GPUs). gpu_indices (PHYSICAL ids) scopes the check to the selected GPUs, so a dGPU on a mixed host is not treated as unified-memory; @@ -2492,7 +3454,9 @@ class LlamaCppBackend: ) arch_by_id[pid] = _arch.split(":")[0].strip().lower() for _i in list(gpu_indices) if gpu_indices is not None else list(arch_by_id): - if arch_by_id.get(_i) in {"gfx1150", "gfx1151"}: + # gfx1152 is Krackan Point (Radeon 860M/840M), the third RDNA 3.5 + # APU: same shared GPU/system-RAM pool as Strix Point/Halo. + if arch_by_id.get(_i) in {"gfx1150", "gfx1151", "gfx1152"}: return True except Exception: return False @@ -2692,6 +3656,8 @@ class LlamaCppBackend: ], capture_output = True, text = True, + encoding = "utf-8", + errors = "replace", timeout = 10, env = child_env_without_native_path_secret(), **_windows_hidden_subprocess_kwargs(), @@ -2763,18 +3729,17 @@ class LlamaCppBackend: return [] @staticmethod - def _get_gpu_free_memory_vulkan(binary: Optional[str] = None) -> list[tuple[int, int, int]]: - """Query free (and total) VRAM per device via the bundled ggml Vulkan backend. + def _run_vulkan_probe(binary: Optional[str] = None) -> list[dict]: + """Run ``_vulkan_probe.py`` and parse its per-device lines. - Loads ``libggml-vulkan`` in a short-lived subprocess (no Vulkan instance - in this process) and returns (device_index, free_mib, total_mib) sorted - by index. The index is ggml's compact Vulkan ordinal -- the one the - registry names ``Vulkan`` and load_model pins with ``--device``, - NOT the raw ``GGML_VK_VISIBLE_DEVICES`` space. A user-set - ``GGML_VK_VISIBLE_DEVICES`` is honored by ggml (passed through), so the - list already reflects it. iGPUs leave a host-RAM margin (see - ``_apply_igpu_host_reserve_mib``) and report total 0; discrete cards pass - their real total through. [] when no Vulkan build or device is reachable. + Returns raw (uncapped) rows sorted by index: + ``{"index", "free_mib", "total_mib", "is_igpu", "name"}``. The index is + ggml's compact Vulkan ordinal -- the one the registry names + ``Vulkan`` and load_model pins with ``--device``, NOT the raw + ``GGML_VK_VISIBLE_DEVICES`` space. A user-set ``GGML_VK_VISIBLE_DEVICES`` + is honored by ggml (passed through), so the list already reflects it. + ``name`` is ggml's device description; "" from an older 4-column probe. + [] when no Vulkan build or device is reachable. """ binary = binary or LlamaCppBackend._find_llama_server_binary() if not binary: @@ -2799,12 +3764,15 @@ class LlamaCppBackend: ) probe_script = Path(__file__).with_name("_vulkan_probe.py") try: + # UTF-8 to match the probe's stdout reconfigure: device names can be + # non-ASCII, and the platform-default decode (cp1252) could throw. result = subprocess.run( [sys.executable, str(probe_script), str(binary_dir)], capture_output = True, - text = True, + encoding = "utf-8", + errors = "replace", timeout = 15, - env = env, + env = utf8_child_env(env), **_windows_hidden_subprocess_kwargs(), ) if result.returncode != 0: @@ -2816,21 +3784,56 @@ class LlamaCppBackend: logger.debug(f"vulkan GPU probe failed: {e}") return [] - gpus: list[tuple[int, int, int]] = [] + rows: list[dict] = [] for line in result.stdout.strip().splitlines(): parts = line.split("\t") - if len(parts) != 4: + # 4 columns from an older probe (no name); 5 with the name column. + if len(parts) not in (4, 5): continue try: - idx = int(parts[0]) - free_mib = int(parts[1]) // (1024 * 1024) - is_igpu = parts[2] == "1" - # iGPU "total" is shared RAM, not a VRAM budget -> keep 0 so the - # fit stays on free*frac (the host reserve below is its - # headroom); a discrete card passes its real total through. - total_mib = 0 if is_igpu else int(parts[3]) // (1024 * 1024) + rows.append( + { + "index": int(parts[0]), + "free_mib": int(parts[1]) // (1024 * 1024), + "is_igpu": parts[2] == "1", + "total_mib": int(parts[3]) // (1024 * 1024), + "name": parts[4].strip() if len(parts) == 5 else "", + } + ) except ValueError: continue + rows.sort(key = lambda r: r["index"]) + return rows + + @staticmethod + def vulkan_device_inventory(binary: Optional[str] = None) -> list[dict]: + """UI-facing Vulkan device list: the devices llama-server will actually + use, with real totals (an iGPU keeps its shared-RAM total here -- the + caller labels it, unlike the fit which zeroes it). Same rows as + ``_run_vulkan_probe``; names fall back to ``Vulkan``. + """ + rows = LlamaCppBackend._run_vulkan_probe(binary) + for row in rows: + if not row["name"]: + row["name"] = f"Vulkan{row['index']}" + return rows + + @staticmethod + def _get_gpu_free_memory_vulkan(binary: Optional[str] = None) -> list[tuple[int, int, int]]: + """Query free (and total) VRAM per device via the bundled ggml Vulkan backend. + + Fit-oriented view of ``_run_vulkan_probe``: returns (device_index, + free_mib, total_mib) sorted by index. iGPUs leave a host-RAM margin (see + ``_apply_igpu_host_reserve_mib``) and report total 0; discrete cards pass + their real total through. [] when no Vulkan build or device is reachable. + """ + gpus: list[tuple[int, int, int]] = [] + for row in LlamaCppBackend._run_vulkan_probe(binary): + idx, free_mib, is_igpu = row["index"], row["free_mib"], row["is_igpu"] + # iGPU "total" is shared RAM, not a VRAM budget -> keep 0 so the + # fit stays on free*frac (the host reserve below is its + # headroom); a discrete card passes its real total through. + total_mib = 0 if is_igpu else row["total_mib"] capped = _apply_igpu_host_reserve_mib(free_mib, is_igpu) if capped < free_mib: logger.info( @@ -2839,7 +3842,6 @@ class LlamaCppBackend: f"({free_mib}->{capped}MiB usable)" ) gpus.append((idx, capped, total_mib)) - gpus.sort(key = lambda g: g[0]) if gpus: logger.info( "Vulkan GPU memory detected: " @@ -2858,7 +3860,7 @@ class LlamaCppBackend: except Exception: pass try: - with open("/proc/meminfo") as f: + with open("/proc/meminfo", encoding = "utf-8") as f: for line in f: if line.startswith("MemAvailable:"): return int(line.split()[1]) // 1024 # kB -> MiB @@ -2955,7 +3957,7 @@ class LlamaCppBackend: return prev = curr - # Free-VRAM fraction at which Studio pins the GPU directly instead of + # Free-VRAM fraction at which Unsloth pins the GPU directly instead of # deferring to ``--fit on``. 3% headroom: the compute buffer is now modelled in # the fit, so this only guards fragmentation + multi-GPU per-device CUDA context # (~2-3%); kept >= 3% as a floor (0.90 dropped 91-94% fits to CPU offload, #5106). @@ -2976,6 +3978,28 @@ class LlamaCppBackend: # aborts a --split-mode tensor load, so it's dropped for the tensor attempt. _TENSOR_PARALLEL_KV_TYPES = frozenset({"f16", "bf16", "f32"}) + # V cache types that llama.cpp can run WITHOUT flash attention. Only the V + # axis has the dependency: a quantized V cache (q8_0/q4_0/q4_1/q5_0/q5_1/ + # iq4_nl) aborts init with "V cache quantization requires flash_attn", while + # a quantized K cache runs fine without FA. So the flash-attn-off crash- + # recovery fallback must reset a quantized V cache to f16 before it can + # launch (and leaves K alone). These three are the only non-quantized types. + _NON_QUANTIZED_KV_TYPES = frozenset({"f16", "bf16", "f32"}) + + # Main-model placement settings that Manual mode owns. They must not leak + # from Studio's parent environment into llama-server and silently override + # the command assembled from the current request. Draft-model placement is + # intentionally separate and remains available to speculative decoding. + _MANUAL_PLACEMENT_ENV_VARS = ( + "LLAMA_ARG_CPU_MOE", + "LLAMA_ARG_N_CPU_MOE", + "LLAMA_ARG_N_GPU_LAYERS", + "LLAMA_ARG_TENSOR_SPLIT", + "LLAMA_ARG_FIT", + "LLAMA_ARG_FIT_TARGET", + "LLAMA_ARG_FIT_CTX", + ) + # (binary, mtime, model) that aborted on --split-mode tensor this process (#6415 # geometry limit, e.g. MQA n_head_kv=1). Model-keyed so one model's abort doesn't # skip tensor for others; tensor is tried by default, recorded only on a real abort. @@ -3106,6 +4130,9 @@ class LlamaCppBackend: lib_dirs.extend(_wsl_system_rocm_lib_dirs()) if lib_dirs: env.setdefault("HSA_ENABLE_DXG_DETECTION", "1") + # Native Linux AMD: system ROCm libs before the bundle's HIP runtime, + # which can be incompatible with the host amdkfd driver. + lib_dirs.extend(_native_linux_system_rocm_lib_dirs(binary_dir)) lib_dirs.append(binary_dir) _arch = platform.machine() # x86_64, aarch64, etc. @@ -3140,6 +4167,12 @@ class LlamaCppBackend: return env + @classmethod + def _clear_manual_placement_env(cls, env: dict[str, str]) -> None: + """Remove inherited main-model placement owned by Manual mode.""" + for name in cls._MANUAL_PLACEMENT_ENV_VARS: + env.pop(name, None) + @staticmethod def _select_gpus( model_size_bytes: int, @@ -3250,6 +4283,32 @@ class LlamaCppBackend: is non-None here.""" return self._embedding_length // self._n_heads if self._n_heads else 128 # type: ignore[operator] + def _max_kv_value_width( + self, + default_len: int, + swa_len: Optional[int] = None, + ) -> int: + """llama.cpp's hparams.n_embd_v_gqa_max() over every model layer.""" + n_layers = self._n_layers or 1 + n_kv = self._n_kv_heads or self._n_heads or 1 + if self._sliding_window_pattern is None: + max_len = max(default_len, swa_len or default_len) + return max( + self._kv_heads_for_layer(layer_idx, n_kv) * max_len for layer_idx in range(n_layers) + ) + return max( + self._kv_heads_for_layer(layer_idx, n_kv) + * ( + (swa_len or default_len) + if ( + layer_idx < len(self._sliding_window_pattern) + and self._sliding_window_pattern[layer_idx] + ) + else default_len + ) + for layer_idx in range(n_layers) + ) + def _estimate_kv_cache_bytes( self, n_ctx: int, @@ -3258,22 +4317,26 @@ class LlamaCppBackend: swa_full: bool = False, n_parallel: int = 1, kv_unified: bool = True, + n_ubatch: Optional[int] = None, ctx_checkpoints: int = 0, + flash_attn: bool = True, ) -> int: """Estimate KV cache VRAM for a given context length. 5-path architecture-aware estimation: 1. MLA -- compressed KV latent + RoPE, K-only (no separate V) 2. Hybrid -- only attention layers need KV (Mamba layers don't) - 3. SWA -- sliding-window layers cache min(ctx, window) tokens + 3. SWA -- sliding-window layers use compact or full cache cells 4. GQA -- standard full KV with explicit key/value dimensions 5. Legacy -- fallback using embed // n_heads Server-flag knobs (mirror llama-server's CLI): swa_full -- --swa-full: SWA layers cache full n_ctx (path 3->4). - n_parallel -- --parallel slots: non-SWA constant, SWA scale linearly. - kv_unified -- --kv-unified: memory no-op (API forward-compat). + n_parallel -- --parallel slots: controls per-slot stream padding. + kv_unified -- --kv-unified: one shared stream vs one per slot. + n_ubatch -- --ubatch-size: SWA cache's processing headroom. ctx_checkpoints -- --ctx-checkpoints: N SWA snapshots per slot. + flash_attn -- False pads variable-width V tensors to the model max. Returns 0 if metadata is insufficient. """ @@ -3288,9 +4351,17 @@ class LlamaCppBackend: n_kv = self._n_kv_heads or self._n_heads or 1 # type: ignore[assignment] # Bytes per element depends on KV cache quantization - bpe = _kv_bytes_per_elem(cache_type_kv) + bpe_k = _kv_bytes_per_elem(cache_type_kv) + # The automatic FA-off retry rewrites an invalid quantized V cache to + # f16. Pricing that viable retry here avoids under-reserving it. + bpe_v = bpe_k if flash_attn else max(bpe_k, _kv_bytes_per_elem("f16")) - slots = max(1, n_parallel) + slots, streams, cells_per_stream = _kv_cache_cell_layout(n_ctx, n_parallel, kv_unified) + total_cells = cells_per_stream * streams + ubatch = max( + 0, + int(self._DEFAULT_N_UBATCH if n_ubatch is None else n_ubatch), + ) # Path 1: MLA (DeepSeek-V2/V3, GLM-4.7, GLM-5, Kimi-K2.5) # One compressed KV latent per token/layer (shared across heads); V is @@ -3301,7 +4372,7 @@ class LlamaCppBackend: n_kv_mla = self._n_kv_heads or 1 rope_dim = self._key_length_mla or 64 key_len = self._kv_key_length or (self._kv_lora_rank + rope_dim) - return int(n_layers_kv * n_ctx * n_kv_mla * key_len * bpe) + return int(n_layers_kv * total_cells * n_kv_mla * key_len * bpe_k) key_len = self._kv_key_length val_len = self._kv_value_length @@ -3312,16 +4383,18 @@ class LlamaCppBackend: fai = self._full_attention_interval n_attn = -(-n_layers // fai) if fai > 0 else n_layers # ceiling division if key_len is not None and val_len is not None: - return int(n_attn * n_ctx * n_kv * (key_len + val_len) * bpe) + v_width = n_kv * val_len if flash_attn else self._max_kv_value_width(val_len) + return int(n_attn * total_cells * (n_kv * key_len * bpe_k + v_width * bpe_v)) head_dim = self._legacy_head_dim() - return int(n_attn * n_ctx * n_kv * 2 * head_dim * bpe) + return int(n_attn * total_cells * n_kv * 2 * head_dim * bpe_k) # Path 3: Sliding window (Gemma 2/3/3n/4, gpt-oss, Cohere2 ...). Pattern # from the resolver; if absent, falls through to the legacy 1/4-global # heuristic. --parallel N accounting (verified against llama-server): - # non-SWA cells = n_ctx split across slots (CONSTANT); SWA per-slot cells - # = 2*sliding_window (capped at n_ctx/per_slot_ctx) -> LINEAR in slots. - # --swa-full forces full n_ctx for SWA; --ctx-checkpoints N adds snapshots. + # non-SWA cells total n_ctx across streams. Compact SWA adds one processing + # micro-batch to the window allowance and pads to 256 cells; unified mode + # holds all slots in one stream, while non-unified mode has one stream per + # slot. --swa-full expands SWA to each stream's full context. if ( self._sliding_window is not None and self._sliding_window > 0 @@ -3329,15 +4402,19 @@ class LlamaCppBackend: and val_len is not None ): swa = self._sliding_window - per_slot_ctx = max(1, n_ctx // slots) - # --swa-full caches full per_slot_ctx (constant n_ctx total); else SWA - # caches 2*sliding_window per slot, clamped at per-slot ctx. - swa_cells_per_slot = per_slot_ctx if swa_full else min(n_ctx, 2 * swa, per_slot_ctx) + if swa_full: + swa_cells_total = total_cells + else: + swa_limit = swa * (slots if kv_unified else 1) + ubatch + swa_cells_per_stream = min(cells_per_stream, swa_limit) + swa_cells_per_stream = _pad_kv_cells(swa_cells_per_stream) + swa_cells_total = swa_cells_per_stream * streams key_len_swa = self._kv_key_length_swa or key_len val_len_swa = self._kv_value_length_swa or val_len + padded_v_width = None if flash_attn else self._max_kv_value_width(val_len, val_len_swa) if self._sliding_window_pattern is not None: - global_bytes = 0.0 # constant across slots - swa_bytes_per_slot = 0.0 # multiplied by slots + global_bytes = 0.0 + swa_bytes = 0.0 checkpoint_extra_per_slot = 0.0 # Only layers that allocate their own KV; trailing shared layers # reuse earlier caches. @@ -3347,41 +4424,48 @@ class LlamaCppBackend: layer_idx < len(self._sliding_window_pattern) and self._sliding_window_pattern[layer_idx] ) + layer_key_bytes = layer_n_kv * (key_len_swa if is_swa else key_len) * bpe_k + layer_value_bytes = ( + layer_n_kv * (val_len_swa if is_swa else val_len) + if padded_v_width is None + else padded_v_width + ) * bpe_v + layer_kv_bytes = layer_key_bytes + layer_value_bytes if is_swa: - swa_bytes_per_slot += ( - swa_cells_per_slot * layer_n_kv * (key_len_swa + val_len_swa) * bpe - ) + swa_bytes += swa_cells_total * layer_kv_bytes if ctx_checkpoints > 0 and not swa_full: - checkpoint_extra_per_slot += ( - ctx_checkpoints - * swa - * layer_n_kv - * (key_len_swa + val_len_swa) - * bpe - ) + checkpoint_extra_per_slot += ctx_checkpoints * swa * layer_kv_bytes else: - global_bytes += n_ctx * layer_n_kv * (key_len + val_len) * bpe - return int(global_bytes + slots * (swa_bytes_per_slot + checkpoint_extra_per_slot)) + global_bytes += total_cells * layer_kv_bytes + return int(global_bytes + swa_bytes + slots * checkpoint_extra_per_slot) n_global = max(1, n_layers_kv // 4) n_swa = n_layers_kv - n_global - kv_per_token = n_kv * (key_len + val_len) * bpe - kv_per_token_swa = n_kv * (key_len_swa + val_len_swa) * bpe - global_bytes = n_global * n_ctx * kv_per_token - swa_bytes_per_slot = n_swa * swa_cells_per_slot * kv_per_token_swa + global_v_width = n_kv * val_len if padded_v_width is None else padded_v_width + swa_v_width = n_kv * val_len_swa if padded_v_width is None else padded_v_width + kv_per_token = n_kv * key_len * bpe_k + global_v_width * bpe_v + kv_per_token_swa = n_kv * key_len_swa * bpe_k + swa_v_width * bpe_v + global_bytes = n_global * total_cells * kv_per_token + swa_bytes = n_swa * swa_cells_total * kv_per_token_swa checkpoint_extra_per_slot = ( ctx_checkpoints * n_swa * swa * kv_per_token_swa if ctx_checkpoints > 0 and not swa_full else 0.0 ) - return int(global_bytes + slots * (swa_bytes_per_slot + checkpoint_extra_per_slot)) + return int(global_bytes + swa_bytes + slots * checkpoint_extra_per_slot) # Path 4: Standard GQA with explicit key/value dimensions if key_len is not None and val_len is not None: - return int(n_layers_kv * n_ctx * n_kv * (key_len + val_len) * bpe) + padded_v_width = None if flash_attn else self._max_kv_value_width(val_len) + bytes_per_cell = 0.0 + for layer_idx in range(n_layers_kv): + layer_n_kv = self._kv_heads_for_layer(layer_idx, n_kv) + v_width = layer_n_kv * val_len if padded_v_width is None else padded_v_width + bytes_per_cell += layer_n_kv * key_len * bpe_k + v_width * bpe_v + return int(total_cells * bytes_per_cell) # Path 5: Legacy fallback (old GGUFs without explicit dimensions) head_dim = self._legacy_head_dim() - return int(2 * n_kv * head_dim * n_layers_kv * n_ctx * bpe) + return int(2 * n_kv * head_dim * n_layers_kv * total_cells * bpe_k) def _draft_backend_for(self, drafter_path: str) -> Optional["LlamaCppBackend"]: """Lightweight backend with a drafter GGUF's metadata, to size its own KV @@ -3429,6 +4513,10 @@ class LlamaCppBackend: draft_cache_type_k: Optional[str] = None, draft_cache_type_v: Optional[str] = None, n_parallel: int = 1, + swa_full: bool = False, + kv_unified: bool = True, + n_ubatch: Optional[int] = None, + flash_attn: bool = True, ) -> Optional[int]: """Draft KV cache bytes at n_ctx, sized from GGUF dims (K and V types are independent). Separate drafter (Gemma): its own KV via _estimate_kv_cache_bytes @@ -3442,12 +4530,23 @@ class LlamaCppBackend: db = self._draft_backend_for(drafter_path) if db is None or not db._can_estimate_kv(): return None + # Gemma 4 assistant layers share the target context's final global + # and SWA KV tensors, so only the drafter weights add memory. + if getattr(db, "_architecture", None) == "gemma4-assistant": + return 0 heavier = draft_cache_type_k if bpe_k >= bpe_v else draft_cache_type_v - # The drafter is served under the same --parallel slot count as the - # main model, so price its KV per slot too: a sliding-window drafter - # (Gemma) grows KV with slots and would otherwise be under-reserved. - kv = db._estimate_kv_cache_bytes(n_ctx, heavier, n_parallel = n_parallel) - return kv or None + # The drafter uses the main model's slot and stream layout, so its + # compact SWA and per-stream padding must follow the same settings. + kv = db._estimate_kv_cache_bytes( + n_ctx, + heavier, + n_parallel = n_parallel, + swa_full = swa_full, + kv_unified = kv_unified, + n_ubatch = n_ubatch, + flash_attn = flash_attn, + ) + return kv if kv > 0 else None nextn = self._nextn_predict_layers or 0 n_kv = self._n_kv_heads or self._n_heads k_len = self._kv_key_length @@ -3461,7 +4560,14 @@ class LlamaCppBackend: f16_bpe = _kv_bytes_per_elem("f16") bpe_k = max(bpe_k, f16_bpe) bpe_v = max(bpe_v, f16_bpe) - return int(nextn * n_kv * (k_len * bpe_k + v_len * bpe_v) * n_ctx) + _, streams, cells_per_stream = _kv_cache_cell_layout(n_ctx, n_parallel, kv_unified) + v_width = n_kv * v_len + if not flash_attn: + v_width = self._max_kv_value_width( + v_len, + self._kv_value_length_swa, + ) + return int(nextn * (n_kv * k_len * bpe_k + v_width * bpe_v) * cells_per_stream * streams) def _estimate_mtp_overhead_bytes( self, @@ -3474,6 +4580,10 @@ class LlamaCppBackend: draft_weights_bytes: int = 0, n_parallel: int = 1, mtp_keeps_target_ctx: bool = True, + swa_full: bool = False, + kv_unified: bool = True, + n_ubatch: Optional[int] = None, + flash_attn: bool = True, ) -> Optional[int]: """MTP draft reserve at ``n_ctx`` = draft KV (grows with ctx) + separate- drafter weights + (MTP + MLA only) a duplicated target KV context. The @@ -3489,6 +4599,10 @@ class LlamaCppBackend: draft_cache_type_k = draft_cache_type_k, draft_cache_type_v = draft_cache_type_v, n_parallel = n_parallel, + swa_full = swa_full, + kv_unified = kv_unified, + n_ubatch = n_ubatch, + flash_attn = flash_attn, ) weights = max(0, draft_weights_bytes) # MLA models (GLM-5.x, DeepSeek, Kimi-K2) under MTP keep a *second* full copy @@ -3504,7 +4618,15 @@ class LlamaCppBackend: # rather than duplicating the target, so they must not be charged for it. target_ctx_copy = 0 if mtp_keeps_target_ctx and self._kv_lora_rank is not None: - target_ctx_copy = self._estimate_kv_cache_bytes(n_ctx, "f16", n_parallel = n_parallel) + target_ctx_copy = self._estimate_kv_cache_bytes( + n_ctx, + "f16", + n_parallel = n_parallel, + swa_full = swa_full, + kv_unified = kv_unified, + n_ubatch = n_ubatch, + flash_attn = flash_attn, + ) if draft_kv is None: # KV unsized (exotic/remote drafter): still reserve known weights + any # MLA target copy so a large config can't launch over budget (the small @@ -3514,7 +4636,7 @@ class LlamaCppBackend: return total if total > 0 else None return draft_kv + weights + target_ctx_copy - _DEFAULT_N_UBATCH = 512 # llama.cpp --ubatch default; Studio does not override it + _DEFAULT_N_UBATCH = _DEFAULT_LLAMA_N_UBATCH _COMPUTE_BUFFER_SAFETY = 1.15 # upper-bound margin on the compute-buffer estimate # Soft VRAM the modeled terms omit; charged to the fit budget on tight tiers (#6682). _CUDA_CONTEXT_RESERVE_BYTES = 320 * 1024 * 1024 # CUDA ctx + cuBLAS workspace (~330 MiB) @@ -3572,7 +4694,10 @@ class LlamaCppBackend: n_embd = self._embedding_length or 0 if n_vocab <= 0 or n_embd <= 0: return 0 - ub = max(1, int(n_ubatch if n_ubatch else self._DEFAULT_N_UBATCH)) + ub = max( + 1, + int(self._DEFAULT_N_UBATCH if n_ubatch is None else n_ubatch), + ) par = max(1, int(n_parallel)) out_buffer = n_vocab * ub * 4 # f32 output/logits buffer act_scratch = 4 * n_embd * ub * 4 # a few resident hidden-width buffers @@ -3604,7 +4729,10 @@ class LlamaCppBackend: n_embd = self._embedding_length or 0 if n_embd <= 0 or n_ctx <= 0: return 0 - ub = max(1, int(n_ubatch if n_ubatch else self._DEFAULT_N_UBATCH)) + ub = max( + 1, + int(self._DEFAULT_N_UBATCH if n_ubatch is None else n_ubatch), + ) if getattr(self, "_architecture", None) == "deepseek4": # DSV4 indexer/CSA buffer (see constants): flat + linear, ub-scaled. Fires # for any KV type -- the indexer scratch is present even with an f16 cache. @@ -3652,9 +4780,12 @@ class LlamaCppBackend: per_device_overhead_bytes: int, min_gpus: int, n_ubatch: Optional[int] = None, + swa_full: bool = False, + kv_unified: bool = True, + flash_attn: bool = True, ) -> tuple[Optional[list[int]], bool, int]: """Largest serving-slot count in [1, n_parallel) whose fully-on-GPU footprint fits, - so Studio keeps the model on GPU (-ngl -1) instead of --fit on, which offloads layers + so Unsloth keeps the model on GPU (-ngl -1) instead of --fit on, which offloads layers to host and collapses decode ~3x (oobabooga #6718). ``base_footprint_bytes`` is the slot-independent footprint (weights + soft overhead + MTP + context-linear compute, minus the folded compute buffer); each candidate re-adds the slot-sized compute buffer @@ -3670,7 +4801,15 @@ class LlamaCppBackend: total = ( base_footprint_bytes + cb - + self._estimate_kv_cache_bytes(effective_ctx, cache_type_kv, n_parallel = slots) + + self._estimate_kv_cache_bytes( + effective_ctx, + cache_type_kv, + n_parallel = slots, + swa_full = swa_full, + kv_unified = kv_unified, + n_ubatch = n_ubatch, + flash_attn = flash_attn, + ) ) gpu_indices, use_fit = self._select_gpus( total, @@ -3695,7 +4834,9 @@ class LlamaCppBackend: swa_full: bool = False, n_parallel: int = 1, kv_unified: bool = True, + n_ubatch: Optional[int] = None, ctx_checkpoints: int = 0, + flash_attn: bool = True, kv_on_gpu: bool = True, mtp_engaged: bool = False, mtp_overhead_fn: Optional[Callable[[int], int]] = None, @@ -3732,7 +4873,9 @@ class LlamaCppBackend: swa_full = swa_full, n_parallel = n_parallel, kv_unified = kv_unified, + n_ubatch = n_ubatch, ctx_checkpoints = ctx_checkpoints, + flash_attn = flash_attn, ) # byte-accurate mtp_overhead_fn supersedes the flat fraction (the fallback @@ -3798,13 +4941,13 @@ class LlamaCppBackend: hf_repo: str, free_bytes: int, hf_token: Optional[str] = None, - ) -> Optional[tuple[str, int]]: + ) -> Optional[tuple[str, int, list[str]]]: """Find the smallest GGUF variant (including all shards) that fits. Groups split shards by variant prefix and sums their sizes (e.g. UD-Q4_K_XL with 9 shards of 50 GB each = 450 GB total). - Returns (first_shard_filename, total_size_bytes) or None. + Returns (first_shard_filename, total_size_bytes, extra_shards) or None. """ try: from huggingface_hub import get_paths_info, list_repo_files @@ -3840,9 +4983,13 @@ class LlamaCppBackend: # Smallest that fits variant_sizes.sort(key = lambda x: x[1]) - for first_file, total_size, _ in variant_sizes: + for first_file, total_size, shard_files in variant_sizes: if total_size > 0 and total_size <= free_bytes: - return first_file, total_size + return ( + first_file, + total_size, + [path for path in sorted(shard_files) if path != first_file], + ) return None except Exception: @@ -3938,6 +5085,21 @@ class LlamaCppBackend: LlamaCppBackend._gguf_skip_value(f, atype) return None + @classmethod + def _gguf_path_is_diffusion(cls, gguf_path: str, model_identifier: str) -> bool: + """Classify a downloaded GGUF without mutating the active backend.""" + probe = object.__new__(cls) + probe._model_identifier = model_identifier + probe._read_gguf_metadata(gguf_path) + return probe._is_diffusion + + def _reject_vulkan_diffusion_gpu_ids_before_teardown( + self, gguf_path: str, model_identifier: str + ) -> None: + """Reject Vulkan + gpu_ids for diffusion GGUFs before Phase 1 teardown.""" + if self._gguf_path_is_diffusion(gguf_path, model_identifier): + raise ValueError(_VULKAN_DIFFUSION_GPU_IDS_ERROR) + def _read_gguf_metadata(self, gguf_path: str) -> None: """Read context_length, architecture params, and chat_template from a GGUF header. @@ -3956,6 +5118,8 @@ class LlamaCppBackend: self._supports_preserve_thinking = False self._supports_tools = False self._n_layers = None + self._n_experts = None + self._leading_dense_block_count = None self._n_kv_heads = None self._n_kv_heads_by_layer = None self._n_heads = None @@ -4045,6 +5209,8 @@ class LlamaCppBackend: arch_keys = { f"{arch}.context_length": "context_length", f"{arch}.block_count": "n_layers", + f"{arch}.expert_count": "n_experts", + f"{arch}.leading_dense_block_count": "leading_dense_block_count", f"{arch}.attention.head_count_kv": "n_kv_heads", f"{arch}.attention.head_count": "n_heads", f"{arch}.embedding_length": "embedding_length", @@ -4126,7 +5292,7 @@ class LlamaCppBackend: ] # Otherwise hand off to the resolver (cache / bootstrap / transformers / HF). Diffusion models - # skip it: they do not use Studio's SWA pattern and the resolver can raise for them. + # skip it: they do not use Unsloth's SWA pattern and the resolver can raise for them. if ( self._sliding_window_pattern is None and self._sliding_window @@ -4233,6 +5399,28 @@ class LlamaCppBackend: return None + @staticmethod + def _diffusion_gpu_arg(gpu_ids: Optional[List[int]], *, cpu_only: bool = False) -> str: + """Device token passed to the diffusion visual-server child. + + The visual engine replaces its child's CUDA visibility mask with this + token, so an unpinned load must carry forward the first token from the + parent's mask rather than turning a parent-relative ordinal into a new + physical selection. + """ + if gpu_ids: + return str(sorted(gpu_ids)[0]) + if cpu_only: + return "" + if "DG_GPU" in os.environ: + return os.environ["DG_GPU"] + parent_mask = os.environ.get("CUDA_VISIBLE_DEVICES") + if parent_mask: + first = next((token.strip() for token in parent_mask.split(",") if token.strip()), "") + if first and first != "-1": + return first + return "0" + def _start_diffusion_server( self, *, @@ -4243,10 +5431,11 @@ class LlamaCppBackend: model_identifier: str, n_ctx: int, extra_args: Optional[List[str]], + gpu_ids: Optional[List[int]] = None, ) -> bool: """Launch the OpenAI-compat diffusion shim (which drives the on-device visual decoder) and wait for health. Presents the same /v1 + /health - interface as llama-server, so the rest of Studio is unchanged. + interface as llama-server, so the rest of Unsloth is unchanged. """ assets = self._find_diffusion_assets() if assets is None: @@ -4268,7 +5457,11 @@ class LlamaCppBackend: # CUDA_VISIBLE_DEVICES="" to force CPU serving. Keep the visual-server child # CPU-masked (empty --gpu) so the shim does not re-expose GPU 0 via its default. cpu_only = self._effective_gpu_count() == 0 - gpu = "" if cpu_only else os.environ.get("DG_GPU", "0") + # Honor the GPU picker first: the diffusion runner takes a single device, + # so use the lowest selected GPU (matches the sorted set recorded below, so + # the device used == the echoed gpu_ids[0]). With no pick, fall back to the + # CPU-only mask, else DG_GPU / 0. + gpu = self._diffusion_gpu_arg(gpu_ids, cpu_only = cpu_only) cmd = list(shim_cmd) + [ "--gguf", @@ -4296,6 +5489,11 @@ class LlamaCppBackend: env.setdefault("UNSLOTH_ALLOW_CPU", "1") env["DG_VISUAL_BIN"] = visual_bin env["DG_GPU"] = gpu + if gpu_ids: + # The visual server remasks via CUDA_VISIBLE_DEVICES=; pin PCI + # order (as the llama-server path does) so the picked physical id maps + # to the GPU the picker showed, not CUDA's default fastest-first order. + env["CUDA_DEVICE_ORDER"] = "PCI_BUS_ID" # The file-override shim imports its sibling visual_engine; put its dir on PYTHONPATH. # (The zoo-package shim is an installed module and needs no PYTHONPATH change.) if extra_pythonpath: @@ -4314,17 +5512,19 @@ class LlamaCppBackend: self._llama_log_path = log_dir / f"diffusion-{int(time.time())}-port-{self._port}.log" self._llama_log_fh = open(self._llama_log_path, "w", encoding = "utf-8", buffering = 1) logger.info(f"diffusion runner stdout/stderr -> {self._llama_log_path}") - except OSError as e: + except (OSError, UnicodeDecodeError) as e: logger.debug(f"Could not open diffusion runner log file: {e}") # The shim (and its visual server) die with this backend process, so a - # Studio crash/restart never orphans a GPU process. + # Unsloth crash/restart never orphans a GPU process. self._process = subprocess.Popen( cmd, stdout = subprocess.PIPE, stderr = subprocess.STDOUT, text = True, - env = env, + encoding = "utf-8", + errors = "replace", + env = utf8_child_env(env), **_windows_hidden_subprocess_kwargs(), **_child_popen_kwargs(), ) @@ -4340,7 +5540,33 @@ class LlamaCppBackend: self._is_audio = False # clear any prior TTS/audio model's routing flag self._model_identifier = model_identifier self._cache_type_kv = None + self._swa_full = False + self._kv_cache_unified = False + self._n_ubatch = self._DEFAULT_N_UBATCH + self._flash_attn_enabled = True + self._effective_cache_types = ("f16", "f16") + self._kv_cache_context_total = None self._gpu_offload_active = True + # Diffusion doesn't use the llama.cpp GPU-memory knobs; reset them to + # defaults (the picked device is still recorded below) so /load, /status + # and reload dedup don't report a previous GGUF's manual settings. + self._gpu_memory_mode = "auto" + self._gpu_layers = -1 + self._n_cpu_moe = 0 + self._tensor_split = None + # Diffusion is never tensor-parallel; clear any state left by a prior TP + # chat load (load_model phase 1 only kills the process, it doesn't run + # the unload reset) so /status doesn't misreport TP and an identical + # re-Apply doesn't reload against stale tensor-parallel state. + self._tensor_parallel = False + # The single-device runner records only the lowest selected GPU (chosen + # above), not the whole pick, and clears any explicit pin from a prior + # chat load; a multi-GPU list would misreport placement and mis-dedup. + self._gpu_ids = [sorted(gpu_ids)[0]] if gpu_ids else None + # The frontend prefers requested_gpu_ids when hydrating the picker. + # Diffusion uses only one device, so echo the collapsed effective pin, + # not unused members of the original request. + self._requested_gpu_ids = list(self._gpu_ids) if self._gpu_ids else None if hf_variant: self._hf_variant = hf_variant elif gguf_path: @@ -4406,6 +5632,9 @@ class LlamaCppBackend: touching the shared one; defaults to the shared event. """ cancel_event = cancel_event if cancel_event is not None else self._cancel_event + from utils.hf_cache_settings import get_hf_cache_paths + + download_cache_dir = str(get_hf_cache_paths().hub_cache) try: import huggingface_hub # noqa: F401 -- presence check only except ImportError: @@ -4438,42 +5667,52 @@ class LlamaCppBackend: except Exception as e: logger.warning(f"Could not list repo files: {e}") - # Offline: resolve variant -> filename from the local HF cache. - # The heuristic below assumes filenames echo the repo name, which - # breaks for e.g. Qwen3.6-27B-MTP-GGUF (no "MTP" in file). Match - # against the rel path (not just basename) so subdir layouts like - # ``BF16/foo.gguf`` are findable. + # Fall back to the local cache when the repo listing is unavailable. if not gguf_filename: - try: - from utils.models.model_config import _iter_hf_cache_snapshots - for snap in _iter_hf_cache_snapshots(hf_repo): - cached_files = _gguf_snapshot_files(snap) - matches = _gguf_files_for_variant(cached_files, hf_variant) - if not matches: - continue - gguf_filename = matches[0] - gguf_extra_shards = _gguf_extra_shards(matches, gguf_filename) - logger.info( - "Resolved variant %s -> %s from local HF cache", - hf_variant, - gguf_filename, - ) - break - except Exception as e: - logger.debug(f"Offline cache lookup for variant failed: {e}") + cached_name, cached_shards = _cached_variant_resolution(hf_repo, hf_variant) + if cached_name: + gguf_filename = cached_name + gguf_extra_shards = cached_shards + logger.info( + "Resolved variant %s -> %s from local HF cache", + hf_variant, + gguf_filename, + ) if not gguf_filename: repo_name = hf_repo.split("/")[-1].replace("-GGUF", "") gguf_filename = f"{repo_name}-{hf_variant}.gguf" + # Prefer the existing model. Updates use force=True to fetch a new revision. + if not force: + if hf_variant: + # Resolve by variant so a newer revision's filename does not hide + # the complete older copy. Size-check against that older snapshot's + # own revision when its metadata remains available. + cached_main = cached_gguf_for_load( + hf_repo, + hf_variant, + verify_sizes = True, + hf_token = hf_token, + ) + else: + candidate = _cached_complete_candidate(hf_repo, gguf_filename, gguf_extra_shards) + cached_main = ( + candidate[0] + if candidate is not None + and _cached_candidate_matches_revision_size(hf_repo, candidate, hf_token) + else None + ) + if cached_main is not None: + logger.info(f"Reusing cached GGUF: {cached_main}") + return cached_main + # Check disk space; fall back to a smaller variant if needed all_gguf_files = [gguf_filename] + gguf_extra_shards - expected_sizes: dict[str, int] = {} try: 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)) - expected_sizes = {p.path: p.size for p in path_infos if p.size} total_bytes = sum((p.size or 0) for p in path_infos) # Subtract bytes already in the HF cache so we only preflight @@ -4482,31 +5721,20 @@ class LlamaCppBackend: # cold whenever free disk is below the full weight footprint, # even though nothing needs downloading. already_cached_bytes = 0 - # Cross-snapshot / case-variant cache reuse is offline-only (see the download - # path below); online, hf_hub_download fetches the current revision and - # resumes partials, so an old snapshot must not be counted as cached here or - # the preflight would under-count the download and skip the disk fallback. + # Count only files that can resume this download. offline = _hf_env_offline() - # A split GGUF whose shards are not co-located in a single snapshot is - # refetched as a whole set later, so it must not be counted as cached here. - split_needs_refetch = False - if offline and not force and gguf_extra_shards: - # Scan all snapshots for one that holds the whole set co-located, so a - # newer snapshot with only the first shard does not mask an older - # complete one and needlessly trip the disk fallback. - if ( - _cached_colocated_split_main( - hf_repo, gguf_filename, gguf_extra_shards, expected_sizes - ) - is None - ): - split_needs_refetch = True + # Offline split sets are reusable only when every shard shares a snapshot. + split_needs_refetch = bool(offline and not force and gguf_extra_shards) if not force and not split_needs_refetch: for p in path_infos: if not p.size: continue try: - cached_path = try_to_load_from_cache(hf_repo, p.path) + cached_path = try_to_load_from_cache( + hf_repo, + p.path, + cache_dir = download_cache_dir, + ) except Exception: cached_path = None if ( @@ -4517,6 +5745,7 @@ class LlamaCppBackend: hf_repo, p.path, expected_size = p.size, + cache_dir = download_cache_dir, ) if isinstance(cached_path, str) and os.path.exists(cached_path): try: @@ -4530,12 +5759,8 @@ class LlamaCppBackend: total_download_bytes = max(0, total_bytes - already_cached_bytes) if total_download_bytes > 0: - cache_dir = os.environ.get( - "HF_HUB_CACHE", - str(Path.home() / ".cache" / "huggingface" / "hub"), - ) - Path(cache_dir).mkdir(parents = True, exist_ok = True) - free_bytes = shutil.disk_usage(cache_dir).free + Path(download_cache_dir).mkdir(parents = True, exist_ok = True) + free_bytes = shutil.disk_usage(download_cache_dir).free total_gb = total_download_bytes / (1024**3) free_gb = free_bytes / (1024**3) @@ -4553,7 +5778,7 @@ class LlamaCppBackend: # surface the disk shortfall for the requested variant. raise RuntimeError( f"Not enough disk space to download {gguf_filename}. " - f"Only {free_gb:.1f} GB free in {cache_dir}" + f"Only {free_gb:.1f} GB free in {download_cache_dir}" ) smaller = self._find_smallest_fitting_variant( hf_repo, @@ -4561,36 +5786,30 @@ class LlamaCppBackend: hf_token, ) if smaller: - fallback_file, fallback_size = smaller + fallback_file, fallback_size, fallback_shards = smaller logger.info( f"Selected variant too large ({total_gb:.1f} GB), " f"falling back to {fallback_file} ({fallback_size / (1024**3):.1f} GB)" ) gguf_filename = fallback_file - _m = _SHARD_RE.match(gguf_filename) - _prefix = _m.group(1) if _m else None - if _prefix: - prefix_lower = _prefix.lower() - gguf_extra_shards = sorted( - f - for f in all_gguf_files - if f.lower().startswith(prefix_lower) - and f != gguf_filename - and not _is_companion_gguf_path(f) + gguf_extra_shards = fallback_shards + + # The selected fallback is a new load target. Apply the + # same any-revision reuse policy before starting a fetch. + fallback_candidate = _cached_complete_candidate( + hf_repo, gguf_filename, gguf_extra_shards + ) + if fallback_candidate is not None and ( + _cached_candidate_matches_revision_size( + hf_repo, fallback_candidate, hf_token ) - else: - gguf_extra_shards = [] - # Record the fallback's size so the later cache-reuse probe can - # size-verify it; only for a single-file fallback, since - # _find_smallest_fitting_variant returns the whole-variant size - # and using that as the first shard's expected size would reject - # a valid cached first shard of a split fallback. - if not gguf_extra_shards: - expected_sizes[fallback_file] = fallback_size + ): + logger.info(f"Reusing cached fallback GGUF: {fallback_candidate[0]}") + return fallback_candidate[0] else: raise RuntimeError( f"Not enough disk space to download any variant. " - f"Only {free_gb:.1f} GB free in {cache_dir}" + f"Only {free_gb:.1f} GB free in {download_cache_dir}" ) except RuntimeError: raise @@ -4606,45 +5825,27 @@ class LlamaCppBackend: raise RuntimeError("Cancelled") dl_start = time.monotonic() # Xet primary, HTTP fallback on stall; per-file so finished shards stay cached. - local_path = None - # Reuse a cached copy from another snapshot / case-variant repo dir only when - # offline. Online, fall through to hf_hub_download so its revision/etag check - # fetches the current file (and resumes a partial) instead of serving a stale - # same-name blob from an older revision. - if not force and _hf_env_offline(): - if gguf_extra_shards: - # A split GGUF must load every shard from one snapshot; reuse only a - # snapshot that holds the whole set co-located, scanning past a newer - # snapshot that has just the first shard while an older one is complete. - local_path = _cached_colocated_split_main( - hf_repo, gguf_filename, gguf_extra_shards, expected_sizes - ) - else: - local_path = _cached_hf_snapshot_file( - hf_repo, - gguf_filename, - expected_size = expected_sizes.get(gguf_filename), - ) - if local_path is None: - local_path = hf_hub_download_with_xet_fallback( + local_path = hf_hub_download_with_xet_fallback( + hf_repo, + gguf_filename, + hf_token, + cancel_event = cancel_event, + on_status = lambda m: logger.info(m), + force_download = force, + cache_dir = download_cache_dir, + ) + for shard in gguf_extra_shards: + if cancel_event.is_set(): + raise RuntimeError("Cancelled") + logger.info(f"Resolving GGUF shard: {shard}") + hf_hub_download_with_xet_fallback( hf_repo, - gguf_filename, + shard, hf_token, cancel_event = cancel_event, - on_status = lambda m: logger.info(m), force_download = force, + cache_dir = download_cache_dir, ) - for shard in gguf_extra_shards: - if cancel_event.is_set(): - raise RuntimeError("Cancelled") - logger.info(f"Resolving GGUF shard: {shard}") - hf_hub_download_with_xet_fallback( - hf_repo, - shard, - hf_token, - cancel_event = cancel_event, - force_download = force, - ) except Exception as e: if isinstance(e, RuntimeError) and "Cancelled" in str(e): raise @@ -4667,10 +5868,12 @@ class LlamaCppBackend: pick: Callable[[list[str]], Optional[str]], label: str, cancel_event: Optional[threading.Event] = None, + near_path: Optional[str] = None, ) -> Optional[str]: """Resolve and fetch a companion GGUF (mmproj / MTP drafter) by name. - Tries the live repo file list, then the local HF cache snapshots + Prefers a companion co-located with ``near_path``'s cache snapshot, + then tries the live repo file list, then the local HF cache snapshots (offline, same fallback as _download_gguf), then hf_hub_download. Runs WITHOUT self._lock (like _download_gguf); honors _cancel_event so an /unload between the main download and here skips the fetch. @@ -4680,6 +5883,23 @@ class LlamaCppBackend: if cancel_event.is_set(): return None + # Keep companion files in the main GGUF's snapshot. + if near_path: + cached = _companion_snapshot_sibling(near_path, pick) + if cached: + logger.info("Reusing cached %s: %s", label, cached) + return cached + + from utils.hf_cache_settings import get_hf_cache_paths + + companion_cache_dir = _hub_cache_dir_for_snapshot_path(near_path) or str( + get_hf_cache_paths().hub_cache + ) + + if _hub_download_in_flight(hf_repo): + logger.info("Skipping %s download while a hub download is active", label) + return None + target: Optional[str] = None from huggingface_hub import list_repo_files @@ -4710,7 +5930,7 @@ class LlamaCppBackend: if target is None: try: from utils.models.model_config import _iter_hf_cache_snapshots - for snap in _iter_hf_cache_snapshots(hf_repo): + for snap in _iter_hf_cache_snapshots(hf_repo, companion_cache_dir): rel_files = _gguf_snapshot_files(snap) target = pick(rel_files) if target is not None: @@ -4728,7 +5948,11 @@ class LlamaCppBackend: # hf_hub_download with hf_repo would miss the canonical file and silently # drop the companion. _cached_hf_snapshot_file scans every case variant. if _hf_env_offline(): - cached = _cached_hf_snapshot_file(hf_repo, target) + cached = _cached_hf_snapshot_file( + hf_repo, + target, + cache_dir = companion_cache_dir, + ) if cached: logger.info("Resolved %s from local HF cache: %s", label, cached) return cached @@ -4741,6 +5965,7 @@ class LlamaCppBackend: target, hf_token, cancel_event = cancel_event, + cache_dir = companion_cache_dir, ) except Exception as e: logger.warning(f"Could not download {label}: {e}") @@ -4752,36 +5977,31 @@ class LlamaCppBackend: hf_repo: str, hf_token: Optional[str] = None, cancel_event: Optional[threading.Event] = None, + near_path: Optional[str] = None, ) -> Optional[str]: """Download the mmproj (vision projection) file from a GGUF repo. Prefers mmproj-F16.gguf, else any mmproj*.gguf. Returns the local path, or None if none exists. ``cancel_event`` overrides - ``self._cancel_event`` (defaults to it). + ``self._cancel_event`` (defaults to it). ``near_path`` prefers a + copy co-located with the main GGUF's cache snapshot. """ - def _pick_mmproj(candidates: list[str]) -> Optional[str]: - mmproj_files = sorted( - f - for f in candidates - if f.lower().endswith(".gguf") and "mmproj" in Path(f).name.lower() - ) - if not mmproj_files: - return None - for f in mmproj_files: - if f.lower().endswith("-f16.gguf"): - return f - return mmproj_files[0] - return self._download_companion_gguf( hf_repo = hf_repo, hf_token = hf_token, pick = _pick_mmproj, label = "mmproj", cancel_event = cancel_event, + near_path = near_path, ) - def _cached_repo_mtp_drafter(self, hf_repo: str) -> Optional[str]: + def _cached_repo_mtp_drafter( + self, + hf_repo: str, + *, + cache_dir: Optional[str] = None, + ) -> Optional[str]: """A drafter already in this repo's local HF cache, reused offline when a fresh copy can't be fetched. Prefers a repo-root ``mtp-*.gguf`` across all cached snapshots; else an existing ``MTP/`` copy (any precision -- the @@ -4791,7 +6011,12 @@ class LlamaCppBackend: roots: list[Path] = [] subdirs: list[Path] = [] - for snap in _iter_hf_cache_snapshots(hf_repo): # newest first + snapshots = ( + _iter_hf_cache_snapshots(hf_repo) + if cache_dir is None + else _iter_hf_cache_snapshots(hf_repo, cache_dir) + ) + for snap in snapshots: # newest first for f in sorted(_gguf_snapshot_files(snap)): if _is_companion_gguf_path(f) and "mmproj" not in f.lower(): (roots if "/" not in f else subdirs).append(snap / f) @@ -4809,6 +6034,7 @@ class LlamaCppBackend: *, hf_repo: str, hf_token: Optional[str] = None, + near_path: Optional[str] = None, ) -> Optional[str]: """Download the separate MTP drafter (speculative head) from a GGUF repo. @@ -4820,16 +6046,6 @@ class LlamaCppBackend: are intentionally skipped. Returns the local path, or None. """ - # Offline, reuse any drafter already on disk (a fresh copy can't be - # fetched). Online, _download_companion_gguf/hf_hub_download reuse the - # current cached file and refetch a changed one, so skip the probe here - # rather than pair new weights with a stale draft. - if _hf_env_offline(): - cached = self._cached_repo_mtp_drafter(hf_repo) - if cached: - logger.info(f"Reusing cached MTP drafter (offline): {cached}") - return cached - def _pick_mtp(candidates: list[str]) -> Optional[str]: # Root-level only: MTP/ subdir copies now share the mtp- prefix but # are explicit-selection, not auto-fetch (they'd sort ahead of root). @@ -4842,11 +6058,31 @@ class LlamaCppBackend: ) return mtp_files[0] if mtp_files else None + if near_path: + cached = _companion_snapshot_sibling(near_path, _pick_mtp) + if cached: + logger.info("Reusing cached MTP drafter: %s", cached) + return cached + + # Offline, reuse any drafter already on disk (a fresh copy can't be + # fetched). Online, _download_companion_gguf/hf_hub_download reuse the + # current cached file and refetch a changed one, so skip the probe here + # rather than pair new weights with a stale draft. + if _hf_env_offline(): + cached = self._cached_repo_mtp_drafter( + hf_repo, + cache_dir = _hub_cache_dir_for_snapshot_path(near_path), + ) + if cached: + logger.info(f"Reusing cached MTP drafter (offline): {cached}") + return cached + return self._download_companion_gguf( hf_repo = hf_repo, hf_token = hf_token, pick = _pick_mtp, label = "MTP drafter", + near_path = near_path, ) def _resolve_launch_mmproj_path( @@ -4972,7 +6208,7 @@ class LlamaCppBackend: return ( f"'{arch}' is a diffusion (image-generation) GGUF, which " "llama-server cannot run as a chat/completion model. Use " - "Studio's Images page to generate with local diffusion " + "Unsloth's Images page to generate with local diffusion " "GGUFs such as FLUX and Qwen-Image." ) if is_ollama: @@ -5051,6 +6287,9 @@ class LlamaCppBackend: total_by_idx: Optional[dict[int, int]] = None, n_ubatch: Optional[int] = None, soft_overhead_bytes: int = 0, + swa_full: bool = False, + kv_unified: bool = True, + flash_attn: bool = True, ) -> tuple[int, int, list[int], Optional[list[int]]]: """Plan a ``--split-mode tensor`` load. Pure: no model or GPU needed. @@ -5138,6 +6377,17 @@ class LlamaCppBackend: def _mtp_at(ctx: int) -> int: return mtp_overhead_fn(ctx) if mtp_overhead_fn is not None else 0 + def _kv_at(ctx: int) -> int: + return self._estimate_kv_cache_bytes( + ctx, + cache_type_kv, + n_parallel = n_parallel, + swa_full = swa_full, + kv_unified = kv_unified, + n_ubatch = n_ubatch, + flash_attn = flash_attn, + ) + # Context-linear compute buffer, summed over the split. Tensor mode # replicates the compute graph on EVERY device (measured: the per-device # buffer grows a flat n_ubatch*2 bytes/token, ~1024 B/tok on Qwen3.5-9B at @@ -5163,31 +6413,21 @@ class LlamaCppBackend: # Weights + buffers exceed the pool -> floor; the load then # falls back to layer split. return ctx_floor - if mtp_overhead_fn is not None: - # kv(ctx)+mtp(ctx)+compute(ctx) is not single-linear, so binary search. - def _consumer(c: int) -> int: - return ( - self._estimate_kv_cache_bytes(c, cache_type_kv, n_parallel = n_parallel) - + _mtp_at(c) - + _cc_ctx(c) - ) - if _consumer(ctx) <= kv_budget_b: - return ctx - lo, hi, best = ctx_floor, ctx, ctx_floor - while lo <= hi: - mid = (lo + hi) // 2 - if _consumer(mid) <= kv_budget_b: - best = mid - lo = mid + 1 - else: - hi = mid - 1 - return best - kv_at = self._estimate_kv_cache_bytes(ctx, cache_type_kv, n_parallel = n_parallel) - total_at = kv_at + _cc_ctx(ctx) # both ~linear through the origin - if total_at <= kv_budget_b: + def _consumer(c: int) -> int: + return _kv_at(c) + _mtp_at(c) + _cc_ctx(c) + + if _consumer(ctx) <= kv_budget_b: return ctx - return max(ctx_floor, int(ctx * kv_budget_b / total_at)) + lo, hi, best = ctx_floor, ctx, ctx_floor + while lo <= hi: + mid = (lo + hi) // 2 + if _consumer(mid) <= kv_budget_b: + best = mid + lo = mid + 1 + else: + hi = mid - 1 + return best # KV size unknown -> can't prove a safe cap; floor. return min(4096, ctx) if ctx > 0 else 4096 @@ -5199,11 +6439,7 @@ class LlamaCppBackend: effective_ctx = min(_fit_ctx(target_ctx), max_available_ctx) min_usable_mib = min(usable_by_idx.values()) - kv_bytes = ( - self._estimate_kv_cache_bytes(effective_ctx, cache_type_kv, n_parallel = n_parallel) - if (self._can_estimate_kv() and effective_ctx > 0) - else 0 - ) + kv_bytes = _kv_at(effective_ctx) if (self._can_estimate_kv() and effective_ctx > 0) else 0 # The MTP reserve also has to fit the even split (mirror the pooled budget): # byte-accurate per-ctx (0 when no fn) plus the same flat cushion as above. mtp_bytes = (_mtp_at(effective_ctx) if effective_ctx > 0 else 0) + flat_mtp_bytes @@ -5254,6 +6490,24 @@ class LlamaCppBackend: and ("unknown" in text or "unsupported" in text or "not supported" in text) ) + @staticmethod + def _mmproj_retry_failure_message(*, projector_confirmed: bool, detail: str) -> str: + """User-facing error when the text-only --mmproj strip retry also fails. + + Confirmed projector-format mismatches keep the historical wording. + Bare signal crashes (common on some ROCm/driver paths) must not be + reported as "Vision projector incompatible" — that misled #7302. + """ + if projector_confirmed: + return ( + "Vision projector incompatible with this llama.cpp " + "build, and the text-only retry also failed: " + detail + ) + return ( + "Vision model failed to start (llama-server crashed with " + "--mmproj), and the text-only retry also failed: " + detail + ) + @staticmethod def _output_has_nonprojector_diagnostic(output: str) -> bool: """True when the output already names a concrete non-projector cause (out @@ -5322,28 +6576,98 @@ class LlamaCppBackend: def explicit(i): nxt = out[i + 1] if i + 1 < len(out) else None - return nxt if nxt in ("on", "auto", "off") else None + return nxt if nxt in _LLAMA_ARG_TRUE_FALSE_AUTO_VALUES else None effective = None for i, tok in enumerate(out): - if tok.startswith(("--flash-attn=", "-fa=")): + name = _flag_name(tok) + if name in ("--flash-attn", "-fa") and "=" in tok: effective = tok.partition("=")[2] - elif tok in ("--flash-attn", "-fa"): + elif name in ("--flash-attn", "-fa"): effective = explicit(i) or "on" - if effective not in ("on", "auto"): + if effective not in _LLAMA_ARG_TRUE_OR_AUTO_VALUES: return None for i, tok in enumerate(out): - if tok.startswith(("--flash-attn=", "-fa=")): + name = _flag_name(tok) + if name in ("--flash-attn", "-fa") and "=" in tok: flag, _, value = tok.partition("=") - if value in ("on", "auto"): + if value in _LLAMA_ARG_TRUE_OR_AUTO_VALUES: out[i] = f"{flag}=off" - elif tok in ("--flash-attn", "-fa"): - if explicit(i) in ("on", "auto"): + elif name in ("--flash-attn", "-fa"): + if explicit(i) in _LLAMA_ARG_TRUE_OR_AUTO_VALUES: out[i + 1] = "off" elif explicit(i) is None: # bare flag (reads as on) -> explicit off out[i] = f"{tok}=off" + + # A quantized V cache requires flash attention in llama.cpp: the init + # aborts with "V cache quantization requires flash_attn". A quantized K + # cache has no such requirement and runs fine without FA, so it is left + # untouched -- resetting it would needlessly enlarge the K cache and can + # OOM a memory-constrained config. Studio launches with FA on, so a + # quantized --cache-type-v is legal at launch but would make THIS FA-off + # retry crash on init instead of recovering. Reset a quantized V cache -- + # main and draft (the draft context shares the global --flash-attn flag, + # so its V cache aborts too) -- to f16 (the llama.cpp default); + # non-quantized types -- f16/bf16/f32 -- run fine without FA and are left + # untouched. The value is rewritten in place so the list length is + # preserved for downstream slices, matching the flash-attn flip above. + _v_cache_flags = ( + "--cache-type-v", + "-ctv", + "--cache-type-v-draft", + "--spec-draft-type-v", + "-ctvd", + ) + _cache_reset = False + for i, tok in enumerate(out): + # llama.cpp rewrites '_' to '-' for any argv token starting with + # '--' before matching, so a legal pass-through spelling such as + # --cache_type_v parses as --cache-type-v and still enables a + # quantized V cache. Canonicalize the flag name the same way so the + # reset recognizes the underscore aliases too; short flags (-ctv) + # and the type value are left untouched. + name = _flag_name(tok) + if name not in _v_cache_flags: + continue + if "=" in tok: + flag, _, value = tok.partition("=") + if value.strip().lower() not in LlamaCppBackend._NON_QUANTIZED_KV_TYPES: + out[i] = f"{flag}=f16" + _cache_reset = True + elif i + 1 < len(out): + if out[i + 1].strip().lower() not in LlamaCppBackend._NON_QUANTIZED_KV_TYPES: + out[i + 1] = "f16" + _cache_reset = True + if _cache_reset: + logger.info( + "V cache dtype reset to f16 because flash attention was disabled " + "by the crash-recovery fallback (quantized V cache requires flash " + "attention in llama.cpp; the K cache is left untouched)." + ) return out + @staticmethod + def _drop_env_quantized_v_cache(env: MutableMapping[str, str]) -> bool: + """Drop an inherited quantized V-cache env var (main or draft) in place + before a flash-attn-off retry, returning True if anything was removed. + + The argv rewrite in ``_with_flash_attn_off`` only reaches flags on the + command line. Studio deliberately lets an env-only cache type reach the + child untouched (an asymmetric K/V env must survive), so a quantized V + cache set purely through ``LLAMA_ARG_CACHE_TYPE_V`` (or the draft + ``LLAMA_ARG_SPEC_DRAFT_CACHE_TYPE_V``) would still abort the FA-off retry + with "V cache quantization requires flash_attn". Dropping it lets + llama.cpp fall back to the f16 default. Only V is dropped: a quantized K + cache runs fine without flash attention, so its env var is preserved. + """ + dropped = False + for var in ("LLAMA_ARG_CACHE_TYPE_V", "LLAMA_ARG_SPEC_DRAFT_CACHE_TYPE_V"): + value = (env.get(var) or "").strip().lower() + if value and value not in LlamaCppBackend._NON_QUANTIZED_KV_TYPES: + env.pop(var, None) + dropped = True + return dropped + @staticmethod def _strip_mmproj_args(cmd: list[str]) -> list[str]: """Return cmd without the '--mmproj ' pair (text-only retry). @@ -5400,7 +6724,7 @@ class LlamaCppBackend: buffering = 1, ) logger.info(f"llama-server stdout/stderr -> {self._llama_log_path}") - except OSError as e: + except (OSError, UnicodeDecodeError) as e: # Best-effort; never block the load on logging. logger.debug(f"Could not open llama-server log file: {e}") self._llama_log_path = None @@ -5414,6 +6738,8 @@ class LlamaCppBackend: stdout = subprocess.PIPE, stderr = subprocess.STDOUT, text = True, + encoding = "utf-8", + errors = "replace", env = env, **_windows_hidden_subprocess_kwargs(), **_child_popen_kwargs(), @@ -5428,6 +6754,7 @@ class LlamaCppBackend: ) self._stdout_thread.start() + @_with_gguf_load_marker def load_model( self, *, @@ -5450,6 +6777,13 @@ class LlamaCppBackend: speculative_type: Optional[str] = None, spec_draft_n_max: Optional[int] = None, tensor_parallel: bool = False, + gpu_memory_mode: Literal["auto", "manual"] = "auto", + gpu_layers: int = -1, + n_cpu_moe: int = 0, + tensor_split: Optional[List[float]] = None, + # Explicit GPU placement pool (issue #7164). None/[] = auto-select; + # the fitter may pin the smallest subset of this pool that fits. + gpu_ids: Optional[List[int]] = None, n_threads: Optional[int] = None, n_gpu_layers: Optional[int] = None, # caller compat, unused n_parallel: int = 1, @@ -5482,6 +6816,14 @@ class LlamaCppBackend: "speculative_type": speculative_type, "spec_draft_n_max": spec_draft_n_max, "tensor_parallel": tensor_parallel, + # GPU-memory placement: replayed on respawn so a server SIGKILL'd by + # GPU/RAM pressure reloads onto the same devices with the same + # offload, not the auto defaults. + "gpu_memory_mode": gpu_memory_mode, + "gpu_layers": gpu_layers, + "n_cpu_moe": n_cpu_moe, + "tensor_split": list(tensor_split) if tensor_split is not None else None, + "gpu_ids": list(gpu_ids) if gpu_ids is not None else None, "n_threads": n_threads, "n_gpu_layers": n_gpu_layers, "n_parallel": n_parallel, @@ -5508,9 +6850,15 @@ class LlamaCppBackend: speculative_type = speculative_type, spec_draft_n_max = spec_draft_n_max, tensor_parallel = tensor_parallel, + gpu_memory_mode = gpu_memory_mode, + gpu_layers = gpu_layers, + n_cpu_moe = n_cpu_moe, + tensor_split = tensor_split, + gpu_ids = gpu_ids, chat_template_override = chat_template_override, extra_args = extra_args, is_vision = is_vision, + n_parallel = n_parallel, preserve_multi_gpu_on_layer = preserve_multi_gpu_on_layer, ): logger.info( @@ -5533,15 +6881,80 @@ class LlamaCppBackend: self._cancel_event.clear() - # ── Phase 1: kill old process (under lock, fast) ────────── - with self._lock: - self._kill_process() - # Resolve llama-server now but defer a not-found error: a block-diffusion # GGUF uses the diffusion runner, and its arch is only known after the header. binary = self._find_llama_server_binary() is_vulkan_backend = self._is_vulkan_backend(binary) + # Without --kv-unified an explicit --parallel N splits -c into windows of -c/N, so on a + # build lacking the flag the default of 4 would quarter every context window for a + # feature it cannot serve: fall back to one slot. Ahead of the KV estimates so the + # fit matches what launches. + if ( + n_parallel > 1 + and binary + and not self.probe_server_capabilities(binary).get("supports_kv_unified") + ): + logger.warning( + "llama-server at %s has no --kv-unified, so %d parallel slots would " + "split the context window %d ways. Using 1 slot instead; update " + "llama.cpp to run chats in parallel.", + binary, + n_parallel, + n_parallel, + ) + n_parallel = 1 + + # ── Vulkan-ordinal preflight (BEFORE the Phase 1 kill) ──────── + # An explicit Vulkan pin the ggml probe never enumerated cannot be honored. + # Validate it ABOVE the kill so an invalid selection leaves the live model + # untouched: CUDA ids are range-checked at the route, but Vulkan ordinals are + # not, so a stale gpu_ids=[99] used to kill the server then 400, leaving + # nothing running (#7239). _get_gpu_memory needs only the binary (safe pre- + # download) and reuses the later fit's issubset logic. Guarded on a found + # Vulkan build + a pin so a deferred not-found stays deferred for diffusion. + if is_vulkan_backend and gpu_ids and binary: + _pf_wanted = {int(x) for x in gpu_ids} + _pf_probed = {g[0] for g in self._get_gpu_memory(binary)} + if not _pf_wanted.issubset(_pf_probed): + raise ValueError( + f"Requested Vulkan GPU ordinal(s) {sorted(_pf_wanted)} not " + f"present. Available Vulkan devices: {sorted(_pf_probed)}." + ) + + # Classify before killing the healthy server (#7205); Phase 2 reuses this path. + _preflight_model_path = None + if is_vulkan_backend and gpu_ids and hf_repo: + _resolved_repo = _resolve_repo_id_casing(hf_repo) + if _resolved_repo != hf_repo: + logger.info( + "Using cached repo_id casing '%s' for requested '%s'", + _resolved_repo, + hf_repo, + ) + hf_repo = _resolved_repo + with _hf_offline_if_dns_dead(): + _preflight_model_path = self._download_gguf( + hf_repo = hf_repo, + hf_variant = hf_variant, + hf_token = hf_token, + ) + self._reject_vulkan_diffusion_gpu_ids_before_teardown( + _preflight_model_path, + model_identifier, + ) + elif is_vulkan_backend and gpu_ids and gguf_path and not hf_repo: + if not Path(gguf_path).is_file(): + raise FileNotFoundError(f"GGUF file not found: {gguf_path}") + self._reject_vulkan_diffusion_gpu_ids_before_teardown( + gguf_path, + model_identifier, + ) + + # ── Phase 1: kill old process (under lock, fast) ────────── + with self._lock: + self._kill_process() + # ── Phase 2: download (NO lock held, so cancel can proceed) ── # mtp_draft_path arrives set for local Gemma loads (detected # sibling); for -hf loads it's None here and resolved just below. @@ -5563,7 +6976,7 @@ class LlamaCppBackend: ) hf_repo = _resolved_repo with _hf_offline_if_dns_dead(): - model_path = self._download_gguf( + model_path = _preflight_model_path or self._download_gguf( hf_repo = hf_repo, hf_variant = hf_variant, hf_token = hf_token, @@ -5573,6 +6986,7 @@ class LlamaCppBackend: mmproj_path = self._download_mmproj( hf_repo = hf_repo, hf_token = hf_token, + near_path = model_path, ) # Auto-download the separate MTP drafter (e.g. Gemma) when # the requested spec mode can use it. Repos with the head @@ -5590,6 +7004,7 @@ class LlamaCppBackend: mtp_draft_path = self._download_mtp( hf_repo = hf_repo, hf_token = hf_token, + near_path = model_path, ) elif gguf_path: if not Path(gguf_path).is_file(): @@ -5614,6 +7029,20 @@ class LlamaCppBackend: # Not a tensor/layer GGUF: clear any preserved-fallback flag from a # prior load (this path skips the command builder that clears it). self._layer_preserves_tensor_intent = False + # On a Vulkan build gpu_ids are ggml Vulkan ordinals, but the diffusion + # runner selects its device by CUDA physical index (_diffusion_gpu_arg + # forwards gpu_ids[0] as a CUDA/DG_GPU token) with no mapping to them. + # The route rejects a CONFIRMED-diffusion pick up front; an uncached GGUF + # only classified as diffusion post-download still reaches here with a + # pin, so drop it and serve on the default device (like an unpinned load). + if gpu_ids and is_vulkan_backend: + logger.warning( + "Ignoring gpu_ids %s for diffusion GGUF on a Vulkan build: " + "the diffusion runner cannot map ggml Vulkan ordinals; " + "serving on the default device.", + gpu_ids, + ) + gpu_ids = None with self._lock: if self._cancel_event.is_set(): logger.info("Load cancelled before diffusion server start") @@ -5626,6 +7055,7 @@ class LlamaCppBackend: model_identifier = model_identifier, n_ctx = n_ctx, extra_args = extra_args, + gpu_ids = gpu_ids, ) if not binary: @@ -5644,6 +7074,8 @@ class LlamaCppBackend: # same message remote validation already shows. raise LlamaServerNotFoundError(LLAMA_SERVER_NOT_FOUND_DETAIL) + server_caps = self.probe_server_capabilities(binary) + # Outside ``self._lock`` so /unload, /cancel, /status aren't # blocked. ``unload_model`` also records the kill, so the # frontend /unload+/load Apply path engages the wait here even @@ -5664,6 +7096,18 @@ class LlamaCppBackend: # state to publish. ctx_override = parse_ctx_override(extra_args) requested_ctx = resolve_requested_ctx(extra_args, n_ctx) + swa_full = _swa_full_from_args_or_env(extra_args) + _effective_ubatch = _extra_args_n_ubatch( + extra_args, + n_ctx = (requested_ctx if requested_ctx > 0 else self._context_length), + ) + planned_kv_unified = _kv_unified_from_args( + extra_args, + default = n_parallel > 1 and server_caps.get("supports_kv_unified", False), + ) + # A hard-crash recovery may relaunch this same plan with FA off. + # Size that larger cache up front so the recovery cannot OOM. + planned_flash_attn = False cache_override = parse_cache_override(extra_args) # Budget the heavier of asymmetric --cache-type-k/-v extras (they # win per axis at launch, appended last); resolve_cache_type_kv only @@ -5687,6 +7131,59 @@ class LlamaCppBackend: # use the same helper so a healthy env-driven tensor server matches. split_mode_override = parse_split_mode_override(extra_args) tensor_parallel = _effective_tensor_parallel(extra_args, tensor_parallel) + # gpu_layers=0 leaves nothing to split, yet --split-mode tensor or + # a per-GPU ratio still launches tensor mode -- and under the + # CPU-only mask below (no visible devices) that aborts the server + # instead of loading on CPU. Drop both here (nothing to split). + if gpu_memory_mode == "manual" and gpu_layers == 0: + if tensor_parallel or tensor_split: + logger.info( + "Manual gpu_layers=0: dropping tensor split/parallel " + "flags (nothing to split on the GPU)" + ) + tensor_parallel = False + tensor_split = None + # Record the requested strategy for /status and the load + # response. 'manual' has no fallback, so the request value is the + # value actually applied. + self._gpu_memory_mode = gpu_memory_mode + # The layer/MoE/split knobs apply only with an explicit offload + # (manual + gpu_layers >= 0); else record defaults so /status and + # /load don't report knobs the server never applied. + if gpu_memory_mode == "manual" and gpu_layers >= 0: + self._gpu_layers = gpu_layers + self._n_cpu_moe = n_cpu_moe + self._tensor_split = tensor_split + else: + self._gpu_layers = -1 + self._n_cpu_moe = 0 + self._tensor_split = None + self._gpu_ids = sorted(gpu_ids) if gpu_ids else None + # Manual offload skips the TP planner but still emits --split-mode + # tensor at launch; drop it when fewer than 2 GPUs are in use -- + # tensor split is a no-op there and aborts on some architectures. + # Done before the cache-drop below so a quantized KV survives. + if ( + tensor_parallel + and gpu_memory_mode == "manual" + and gpu_layers >= 0 + and self._effective_gpu_count(sorted(gpu_ids) if gpu_ids else None) < 2 + ): + logger.info( + "Tensor parallelism requested in manual mode but fewer " + "than 2 GPUs are in use; ignoring (needs >= 2)." + ) + tensor_parallel = False + # Drop TP for manual + Auto layers before the cache-drop below (like + # the <2-GPU guard above), so a requested quantized KV survives into + # the --fit load rather than being stripped for a tensor attempt. + if tensor_parallel and gpu_memory_mode == "manual" and gpu_layers < 0: + logger.info( + "Manual mode with Auto layers hands memory management to " + "llama.cpp --fit, which is incompatible with tensor " + "parallelism; ignoring the tensor split." + ) + tensor_parallel = False # Tensor mode aborts on a quantized KV cache, so drop it for the # tensor attempt (and strip any inherited/explicit --cache-type # that would re-impose it when appended last). Layer split does @@ -5767,10 +7264,22 @@ class LlamaCppBackend: "Vision-capable GGUF loaded without a usable mmproj; " "image input will be disabled for this session" ) + # Seed before the try: the except (GPU-selection failure -> + # --fit on) falls through to the launch which reads this, and the + # probe that assigns it may throw first. Captured before manual + # empty `gpus` so the speculative defaults stay GPU-aware and the + # CPU-fallback check still knows GPUs were present. + _detected_gpus: list[tuple[int, int]] = [] model_size = None # set in the fit try; used by the APU RAM guard # Layer-fallback min GPUs; raised below on a tensor downgrade. Bound # before the try so the --fit-on except path still has it (no UnboundLocal). _layer_min_gpus = 1 + # An explicit Vulkan ordinal absent from the ggml probe cannot be + # honored; flag it in the fit and reject after the try (raising inside + # would be swallowed into the --fit-on fallback). Bound before the try. + _vulkan_explicit_unmatched = False + _vulkan_requested_ids: list[int] = [] + _vulkan_available_ordinals: list[int] = [] try: gguf_size = self._get_gguf_size_bytes(model_path) # Include GPU-loaded mmproj in the fit budget (#5825). @@ -5783,7 +7292,41 @@ class LlamaCppBackend: # Pass binary so a Vulkan build probes ggml's Vulkan ordinals. _gpu_mem = self._get_gpu_memory(binary) gpus = [(idx, free) for idx, free, _t in _gpu_mem] + # Restrict the fit (and thus the layer plan + pin env) to the + # selected GPUs; fail-open if none match so a stale UI choice + # can't strand the load on CPU (issue #7164). + if gpu_ids: + # A Vulkan build indexes by ggml ordinal. An explicit ordinal + # absent from the probe can't be pinned, so reject after the try + # rather than fail-open onto a device the user didn't pick. + _wanted_ids = {int(x) for x in gpu_ids} + # Reject if ANY requested ordinal is absent, not only when none + # match: [0, 99] against {0, 1} silently drops 99. Comparing the + # full requested set (before filter narrows) still lets the fitter + # pick a valid subset later -- that is narrowing, not absence. + _probed_ordinals = {g[0] for g in gpus} + if is_vulkan_backend and not _wanted_ids.issubset(_probed_ordinals): + _vulkan_explicit_unmatched = True + _vulkan_requested_ids = sorted(_wanted_ids) + _vulkan_available_ordinals = sorted(_probed_ordinals) + # Restrict the probed pool to the selection; fail-open (keep the + # full pool) if none match so a stale UI choice can't strand the + # load on CPU (issue #7164). + _sel_gpus = [g for g in gpus if g[0] in _wanted_ids] + gpus = _sel_gpus if _sel_gpus else gpus total_by_idx = {idx: total for idx, _f, total in _gpu_mem} + # GPU picker: restrict every mode to the chosen devices, so + # auto selection only considers them and manual mask to + # them (the env block below pins CUDA/HIP_VISIBLE_DEVICES). + if gpu_ids: + _picked = set(gpu_ids) + gpus = [g for g in gpus if g[0] in _picked] + + # GPUs the model will run on -- captured before manual + # empty `gpus` to bypass the planner. bool() drives the + # GPU-aware speculative defaults; the list feeds the + # CPU-fallback check. + _detected_gpus = list(gpus) def _gpu_usable(g, frac = _CTX_FIT_VRAM_FRACTION): # Per-GPU usable budget for ranking: free - (1-frac)*total. @@ -5815,6 +7358,44 @@ class LlamaCppBackend: # GPU/VRAM-fit logic below may shrink it on limited HW. max_available_ctx = self._context_length or effective_ctx + # Manual + Auto layers (the Manual default): hand memory + # management to llama.cpp's --fit. Emptying the probed GPU set + # no-ops the selection/TP planning below, leaving gpu_indices + # None (an explicit gpu_ids pick still pins below) and use_fit + # True. An explicit context is honored (--fit optimizes around + # it); 0 lets --fit size it. + if gpu_memory_mode == "manual" and gpu_layers < 0: + # Tensor parallelism was already dropped above (before the + # cache-drop), so a quantized KV survives into this --fit load. + gpus = [] + effective_ctx = requested_ctx if requested_ctx > 0 else 0 + original_ctx = effective_ctx + # --fit aborts under --split-mode tensor; a raw extras + # --split-mode/--tensor-split (appended last) would + # otherwise reach llama-server. Strip it like the TP + # downgrade does. + extra_args = strip_split_mode_only(extra_args) + elif gpu_memory_mode == "manual": + # Manual offload (--gpu-layers + --fit off): no automatic + # device masking (a gpu_ids pick still pins below) or + # context cap -- the user owns both. tensor_parallel is + # honored but skips the memory-based planner (gpus = []); + # the toggle just emits --split-mode tensor (split by free + # VRAM, or by the Split ratio if set). + gpus = [] + effective_ctx = ( + requested_ctx if requested_ctx > 0 else (self._context_length or 0) + ) + original_ctx = effective_ctx + # Strip the user --split-mode when the toggle owns the split + # (TP engaged -> Studio emits --split-mode tensor) or when the + # user asked for tensor (which aborts on a single GPU even if + # the manual <2-GPU guard downgraded TP). Otherwise keep their + # non-tensor mode (row/none/layer) -- the toggle can't express + # those. + if tensor_parallel or split_mode_override == "tensor": + extra_args = strip_split_mode_only(extra_args) + # Will MTP engage? If so, auto-fit reserves draft-model VRAM. # Mirrors _build_speculative_flags: forced mtp/mtp+ngram always # engage; auto only on an MTP model >= 3B; ngram/off never. A @@ -5830,7 +7411,7 @@ class LlamaCppBackend: and not bool(mtp_draft_path) ) # LLAMA_ARG_SPEC_TYPE only reaches the child when neither extras - # nor Studio emit a spec flag (mode "off", no user --spec-type), + # nor Unsloth emit a spec flag (mode "off", no user --spec-type), # since _build_speculative_flags emits one for every other mode. # Consult the env for the reserve only then, else a stale MTP env # would over-reserve. @@ -5839,7 +7420,7 @@ class LlamaCppBackend: if (not _extra_args_set_spec_type(extra_args) and _mtp_canonical == "off") else {} ) - # Extras can run MTP even when Studio suppresses its own emission. + # Extras can run MTP even when Unsloth suppresses its own emission. _user_mtp_via_extras = _extra_args_requests_mtp(extra_args, env = _spec_env) # A non-MTP model-based draft mode (draft-simple/draft-eagle3) in # extras also loads a separate draft model that needs reserving; @@ -5902,10 +7483,13 @@ class LlamaCppBackend: _extra_n_max = _extra_args_spec_draft_n_max(extra_args) _mtp_eff_n_max = _extra_n_max if _extra_n_max is not None else spec_draft_n_max if _mtp_eff_n_max is None: - _mtp_eff_n_max = 2 if gpus else 3 + # _detected_gpus (not gpus) so manual -- which empty + # gpus to bypass the planner -- keep the GPU draft depth the + # launch flags also use, instead of the CPU default. + _mtp_eff_n_max = 2 if _detected_gpus else 3 # Separate-drafter weights live on GPU (an embedded head is # already in model_size). Size the drafter the launch loads, by - # precedence: extras --model-draft (last-wins), else Studio's + # precedence: extras --model-draft (last-wins), else Unsloth's # emitted mtp_draft_path, else the env drafter. Sizing the wrong # one would under-reserve and OOM. _cli_draft_for_budget = _extra_args_mtp_draft_path(extra_args, env = {}) @@ -5954,6 +7538,10 @@ class LlamaCppBackend: draft_cache_type_k = _mtp_draft_ck, draft_cache_type_v = _mtp_draft_cv, n_parallel = n_parallel, + swa_full = swa_full, + kv_unified = planned_kv_unified, + n_ubatch = _effective_ubatch, + flash_attn = planned_flash_attn, ) if ( self._estimate_mtp_overhead_bytes( @@ -5965,6 +7553,10 @@ class LlamaCppBackend: draft_weights_bytes = _mtp_draft_weights, n_parallel = n_parallel, mtp_keeps_target_ctx = _engaged_is_mtp, + swa_full = swa_full, + kv_unified = planned_kv_unified, + n_ubatch = _effective_ubatch, + flash_attn = planned_flash_attn, ) is not None ): @@ -5981,6 +7573,10 @@ class LlamaCppBackend: _w: int = _mtp_draft_weights, _np: int = n_parallel, _mtp: bool = _engaged_is_mtp, + _swa_full: bool = swa_full, + _kv_unified: bool = planned_kv_unified, + _n_ubatch: Optional[int] = _effective_ubatch, + _flash_attn: bool = planned_flash_attn, ) -> int: v = self._estimate_mtp_overhead_bytes( ctx, @@ -5991,15 +7587,26 @@ class LlamaCppBackend: draft_weights_bytes = _w, n_parallel = _np, mtp_keeps_target_ctx = _mtp, + swa_full = _swa_full, + kv_unified = _kv_unified, + n_ubatch = _n_ubatch, + flash_attn = _flash_attn, ) return v if v is not None else 0 def _mtp_bytes(ctx: int) -> int: return mtp_overhead_fn(ctx) if mtp_overhead_fn is not None else 0 - # Effective micro-batch (a user --ubatch override scales the - # compute buffer); None -> the 512 default in the estimate. - _effective_ubatch = _extra_args_n_ubatch(extra_args) + def _kv_bytes(ctx: int) -> int: + return self._estimate_kv_cache_bytes( + ctx, + cache_type_kv, + n_parallel = n_parallel, + swa_full = swa_full, + kv_unified = planned_kv_unified, + n_ubatch = _effective_ubatch, + flash_attn = planned_flash_attn, + ) def _cc_bytes(ctx: int, n_gpus: int = 1) -> int: # Context-linear compute-buffer growth (flash-attn KQ mask + @@ -6040,7 +7647,8 @@ class LlamaCppBackend: # honor it, cap only if it fits no combination. Auto (native): # prefer fewer GPUs with reduced context (multi-GPU is slower). gpu_indices, use_fit = None, True - # Per-GPU weight proportions for tensor mode (None = even). + # Per-GPU weight proportions for tensor mode (None lets + # llama.cpp split by free VRAM). tp_tensor_split: Optional[list[int]] = None explicit_ctx = requested_ctx > 0 # Flat MTP reserve fraction: used only as the fallback when the @@ -6115,7 +7723,12 @@ class LlamaCppBackend: # GPUs below that reserve from the set up front (gpu_indices # becomes the CUDA_VISIBLE_DEVICES mask, fully excluding them). tp_gpus = gpus - if tensor_parallel: + # Manual mode owns the layer count and context, so it skips + # the memory-based planner; its toggle still emits + # --split-mode tensor below (split by free VRAM, or by the + # Split ratio if set). auto plans here. + plan_tp = tensor_parallel and gpu_memory_mode != "manual" + if plan_tp: # Deterministic per-device compute buffer (replicated on # every device in tensor mode); flat fallback when dims # are unavailable. _plan_tensor_parallel uses the same. @@ -6134,7 +7747,7 @@ class LlamaCppBackend: # free yet have no budget left. tp_gpus = [g for g in gpus if _gpu_usable(g) >= reserve_mib] - if tensor_parallel and len(tp_gpus) < 2: + if plan_tp and len(tp_gpus) < 2: # Tensor parallelism needs >= 2 usable GPUs. On a single # GPU --split-mode tensor is a no-op; with 0 GPUs (CPU-only # or probe failed) it must not reach llama-server; and a @@ -6233,6 +7846,9 @@ class LlamaCppBackend: total_by_idx = total_by_idx, n_ubatch = _effective_ubatch, soft_overhead_bytes = _soft_overhead, + swa_full = swa_full, + kv_unified = planned_kv_unified, + flash_attn = planned_flash_attn, ) use_fit = False elif gpus and self._can_estimate_kv() and effective_ctx > 0: @@ -6265,16 +7881,18 @@ class LlamaCppBackend: pool_budget, _ms, cache_type_kv, + swa_full = swa_full, n_parallel = n_parallel, + kv_unified = planned_kv_unified, + n_ubatch = _effective_ubatch, + flash_attn = planned_flash_attn, mtp_engaged = _mtp_reserves_gpu, mtp_overhead_fn = mtp_overhead_fn, compute_ctx_bytes_fn = _cc_sub, budget_frac = 1.0, total_mib = None, ) - kv = self._estimate_kv_cache_bytes( - capped, cache_type_kv, n_parallel = n_parallel - ) + kv = _kv_bytes(capped) footprint_mib = ( _ms + kv + _mtp_bytes(capped) + _cc_sub(capped) ) / (1024 * 1024) @@ -6294,9 +7912,7 @@ class LlamaCppBackend: # on and let llama-server flex -ngl (CPU offload). requested_total = ( model_size_fit - + self._estimate_kv_cache_bytes( - effective_ctx, cache_type_kv, n_parallel = n_parallel - ) + + _kv_bytes(effective_ctx) + _mtp_bytes(effective_ctx) + _cc_bytes(effective_ctx) ) @@ -6348,16 +7964,18 @@ class LlamaCppBackend: pool_budget, _ms, cache_type_kv, + swa_full = swa_full, n_parallel = n_parallel, + kv_unified = planned_kv_unified, + n_ubatch = _effective_ubatch, + flash_attn = planned_flash_attn, mtp_engaged = _mtp_reserves_gpu, mtp_overhead_fn = mtp_overhead_fn, compute_ctx_bytes_fn = _cc_sub, budget_frac = 1.0, total_mib = None, ) - kv = self._estimate_kv_cache_bytes( - capped, cache_type_kv, n_parallel = n_parallel - ) + kv = _kv_bytes(capped) footprint_mib = ( _ms + kv + _mtp_bytes(capped) + _cc_sub(capped) ) / (1024 * 1024) @@ -6374,11 +7992,7 @@ class LlamaCppBackend: if effective_ctx > 0: for n_gpus in range(_auto_min_gpus, len(ranked) + 1): subset = ranked[:n_gpus] - kv = self._estimate_kv_cache_bytes( - effective_ctx, - cache_type_kv, - n_parallel = n_parallel, - ) + kv = _kv_bytes(effective_ctx) footprint_mib = ( _subset_model_size(n_gpus) + kv @@ -6435,7 +8049,11 @@ class LlamaCppBackend: _apple_fit_budget_mib, model_size_fit, cache_type_kv, + swa_full = swa_full, n_parallel = n_parallel, + kv_unified = planned_kv_unified, + n_ubatch = _effective_ubatch, + flash_attn = planned_flash_attn, mtp_engaged = _mtp_reserves_gpu, mtp_overhead_fn = mtp_overhead_fn, compute_ctx_bytes_fn = _cc_bytes, @@ -6443,12 +8061,7 @@ class LlamaCppBackend: total_mib = None, ) _cap_footprint_mib = ( - model_size_fit - + self._estimate_kv_cache_bytes( - cap, cache_type_kv, n_parallel = n_parallel - ) - + _mtp_bytes(cap) - + _cc_bytes(cap) + model_size_fit + _kv_bytes(cap) + _mtp_bytes(cap) + _cc_bytes(cap) ) / (1024 * 1024) # Fit returns the request unchanged when it fits OR weights # exceed budget; only the latter over-commits, so floor to 4096. @@ -6495,6 +8108,9 @@ class LlamaCppBackend: _pipeline_overhead_bytes + _cc_bytes(effective_ctx), _layer_min_gpus, _effective_ubatch, + swa_full = swa_full, + kv_unified = planned_kv_unified, + flash_attn = planned_flash_attn, ) if not _uf_slots: logger.info( @@ -6519,9 +8135,7 @@ class LlamaCppBackend: _mtp_note = "" if effective_ctx < original_ctx: - kv_est = self._estimate_kv_cache_bytes( - effective_ctx, cache_type_kv, n_parallel = n_parallel - ) + kv_est = _kv_bytes(effective_ctx) logger.info( f"Context auto-reduced: {original_ctx} -> {effective_ctx} " f"(model: {model_size / (1024**3):.1f} GB, " @@ -6530,9 +8144,7 @@ class LlamaCppBackend: + ")" ) - kv_cache_bytes = self._estimate_kv_cache_bytes( - effective_ctx, cache_type_kv, n_parallel = n_parallel - ) + kv_cache_bytes = _kv_bytes(effective_ctx) mmproj_note = ( f"mmproj: {mmproj_size / (1024**3):.1f} GB, " if mmproj_size else "" ) @@ -6550,6 +8162,23 @@ class LlamaCppBackend: tp_tensor_split = None effective_ctx = requested_ctx # fall back to original + # An unenumerated explicit Vulkan ordinal can't be pinned; fail loudly + # instead of fitting onto an unselected device. Clear the raw selection + # the early state-publish recorded so it never leaks into gpu_ids (#7239). + if _vulkan_explicit_unmatched: + self._gpu_ids = None + self._requested_gpu_ids = None + raise ValueError( + f"Requested Vulkan GPU ordinal(s) {_vulkan_requested_ids} not " + f"present. Available Vulkan devices: {_vulkan_available_ordinals}." + ) + + # GPU picker: when no narrower subset was chosen (manual, or + # a failed/file-size selection), pin the whole picked set so the + # model can't spill onto an unpicked GPU. + if gpu_ids and gpu_indices is None: + gpu_indices = sorted(gpu_ids) + # Unified-memory APUs load weights into system RAM (under WSL the VM # cap, not the ROCm-reported VRAM, is the real ceiling); refuse an # oversize load the OS would otherwise kill mid-flight. Base model @@ -6586,8 +8215,6 @@ class LlamaCppBackend: model_path, "--port", str(self._port), - "-c", - str(effective_ctx) if effective_ctx > 0 else "0", "--parallel", str(n_parallel), "--flash-attn", @@ -6595,6 +8222,17 @@ class LlamaCppBackend: # 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", ] + # A positive context is always passed (in auto-fit, --fit then + # optimizes the gpu-layer offload around it). When auto-fit has + # no explicit context, omit -c so --fit sizes it to fit VRAM: + # "-c 0" would instead pin the FULL native context (llama.cpp's + # -c handler sets fit_params_min_ctx = UINT32_MAX on value 0, + # disabling --fit's reduction). See gpu_memory_mode. + auto_fit = gpu_memory_mode == "manual" and gpu_layers < 0 + if effective_ctx > 0: + cmd.extend(["-c", str(effective_ctx)]) + elif not auto_fit: + cmd.extend(["-c", "0"]) # Report a clean public model id (matching GET /v1/models) rather # than the raw -m path in llama-server's own /v1/models and the @@ -6606,7 +8244,63 @@ class LlamaCppBackend: cmd.extend(["--alias", _alias]) fully_gpu_offloaded = False - if use_fit: + # Set when a positional --tensor-split is emitted, so the env block + # can pin CUDA to PCI order even without a GPU subset (see below). + manual_tensor_split_emitted = False + if gpu_memory_mode == "manual" and gpu_layers >= 0: + # Pin the user's layer count and disable auto-fit. --fit off + # also means _ctx_integrity_flags must not add --fit-ctx. + use_fit = False + cmd.extend(["--gpu-layers", str(gpu_layers), "--fit", "off"]) + # Keep the first n_cpu_moe MoE layers' experts on CPU. + moe_flag = self._resolve_cpu_moe_flag( + n_cpu_moe, + self.n_moe_layers, + self._leading_dense_block_count or 0, + ) + if moe_flag is not None: + cmd.extend(["--n-cpu-moe", str(moe_flag)]) + elif n_cpu_moe: + # Requested on a dense model: nothing was emitted, so + # don't report a count llama-server never received. + self._n_cpu_moe = 0 + # Distribute the model across GPUs by the user's per-GPU shares + # (--tensor-split). Works with layer split and tensor + # parallelism; --fit off means no fit/tensor abort. Only emit + # when >1 GPU is in use AND the list length matches that count: + # the field is hidden (not cleared) when the picker narrows to + # one, and a direct caller can send a stale ratio for a different + # GPU set. Studio drops any mismatch to the free-VRAM default + # (llama.cpp would silently zero-pad a short list, or abort past + # its 16-device cap). + _split_gpus = self._effective_gpu_count(gpu_indices) + if tensor_split and _split_gpus > 1: + # An all-zero/non-positive sanitized split assigns nothing + # anywhere, so fall through to the free-VRAM default in + # that case. + _sanitized_split = self._sanitize_tensor_split(tensor_split) + _split_total = sum(_sanitized_split) + if len(_sanitized_split) == _split_gpus and _split_total > 0: + cmd.extend( + ["--tensor-split", ",".join(f"{x:g}" for x in _sanitized_split)] + ) + self._tensor_split = _sanitized_split + manual_tensor_split_emitted = True + else: + logger.warning( + "Dropping manual --tensor-split (%d entries for " + "%d GPUs, sanitized total %s); llama.cpp's " + "free-VRAM split applies instead", + len(tensor_split), + _split_gpus, + _split_total, + ) + self._tensor_split = None + elif tensor_split: + # Single effective GPU: the split is never emitted, so + # don't report it as active via /status and /load. + self._tensor_split = None + elif use_fit: cmd.extend(["--fit", "on"]) elif gpu_indices is not None: # Fits on selected GPU(s) -- force all layers on GPU. --fit off is @@ -6615,15 +8309,35 @@ class LlamaCppBackend: cmd.extend(["-ngl", "-1", "--fit", "off"]) fully_gpu_offloaded = True - server_caps = self.probe_server_capabilities(binary) # Expose Prometheus /metrics for the engine-stats logger, only # when the binary advertises it (older/custom binaries may not). if server_caps.get("supports_metrics"): cmd.append("--metrics") + self._slot_save_dir = None + self._slot_save_binary = None + self._prompt_cache_disabled = False + if server_caps.get("supports_slot_save"): + try: + from utils.paths.storage_roots import ( # noqa: WPS433 + llama_slot_cache_root, + ) + + slot_dir = llama_slot_cache_root() + slot_dir.mkdir(parents = True, exist_ok = True) + # Saved KV encodes chat content; keep it from other local users. + with contextlib.suppress(OSError): + os.chmod(slot_dir, 0o700) + cmd.extend(["--slot-save-path", str(slot_dir)]) + self._slot_save_dir = str(slot_dir) + self._slot_save_binary = (binary, Path(binary).stat().st_mtime_ns) + except OSError: + self._slot_save_dir = None + self._slot_save_binary = None cmd.extend( self._ctx_integrity_flags( n_parallel, use_fit, + auto_fit, requested_ctx, effective_ctx, server_caps, @@ -6666,6 +8380,11 @@ class LlamaCppBackend: "iq4_nl", "f32", } + # Normalize like the budget does (_planned_main_cache_types): a + # case-sensitive match drops "Q8_0", emitting no flag, so llama.cpp + # runs f16 while the estimate priced q8_0. Emit the normalized + # spelling; kv_cache_type_from_str is case-sensitive. + cache_type_kv = cache_type_kv.strip().lower() if cache_type_kv else cache_type_kv if ( cache_type_kv and cache_type_kv in _valid_cache_types @@ -6687,9 +8406,11 @@ class LlamaCppBackend: self._cache_type_kv = None # Tensor parallelism: split the model across GPUs by tensor - # rather than by layer. Multi-GPU only -- a no-op on a single - # GPU. Default (layer split) is left implicit by omitting the - # flag. See llama.cpp --split-mode. + # rather than by layer. The UI only offers it on multi-GPU; a + # direct single-GPU caller is redundant (supported archs no-op, + # unsupported ones abort and the /load path retries layer split). + # Default (layer split) is left implicit by omitting the flag. + # See llama.cpp --split-mode. if tensor_parallel: cmd.extend(["--split-mode", "tensor"]) if tp_tensor_split and len(tp_tensor_split) > 1: @@ -6721,7 +8442,7 @@ class LlamaCppBackend: extra_args = extra_args, model_identifier = model_identifier, model_path = model_path, - gpus = bool(gpus), + gpus = bool(_detected_gpus), binary = binary, mtp_draft_path = launch_mtp_draft_path, ) @@ -6800,8 +8521,9 @@ class LlamaCppBackend: else: self._api_key = None - # Windows + full offload: disable KV checkpoints (WDDM/PCI-E - # overhead). CPU/partial offload keeps prompt caching. #5692. + # Windows + full offload: drop the host-RAM KV checkpoints that cause + # WDDM/PCI-E overhead, but keep prompt caching (in-VRAM prefix reuse) so + # a repeated prompt is not re-prefilled on every request. #5692. if sys.platform == "win32" and full_offload_tuning_active: unsupported_cache_flags: list[str] = [] if server_caps.get("supports_cache_ram"): @@ -6812,42 +8534,76 @@ class LlamaCppBackend: cmd.extend(["--ctx-checkpoints", "0"]) else: unsupported_cache_flags.append("--ctx-checkpoints") - if server_caps.get("supports_no_cache_prompt"): - cmd.append("--no-cache-prompt") - else: - unsupported_cache_flags.append("--no-cache-prompt") if unsupported_cache_flags: logger.info( "Skipping unsupported Windows cache flags for llama-server: %s", ", ".join(unsupported_cache_flags), ) - # Vulkan pins via --device (a cmd arg, unlike the env-based - # CUDA/ROCm pin below), emitted BEFORE user extras so llama.cpp's - # last-wins parsing lets a user --device override Studio's pick. - if is_vulkan_backend and gpu_indices is not None: - cmd += LlamaCppBackend._vulkan_pin_args(gpu_indices) + # Vulkan pins via --device (a cmd arg), before user extras so a user + # --device wins. Fall back to raw ids when the fit did not narrow. + _vulkan_pin_ids = gpu_indices if gpu_indices is not None else (gpu_ids or None) + + # Record the pin actually applied (fit-narrowed gpu_indices, else the raw + # request) for the keep-warm loop, dedupe, and /status, so an explicit + # [0, 1] narrowed to [0] records [0] and /status never echoes an ordinal + # the child never saw. Auto selection (no gpu_ids) stays None (#7239). + if is_vulkan_backend: + # Only record an EXPLICIT Vulkan pin: an auto pick still narrows + + # pins below, but recording it would misreport an explicit pin and + # make dedupe miss the loaded server; mirrors the CUDA/ROCm branch. + self._gpu_ids = ( + sorted(int(x) for x in _vulkan_pin_ids) + if (gpu_ids and _vulkan_pin_ids) + else None + ) + elif gpu_ids: + # Physical pin: the fit-selected subset when the fit ran, else the raw + # user selection so an explicit choice is honoured even when the fit + # could not size the model. + _effective_pin_ids = ( + [int(x) for x in gpu_indices] + if gpu_indices is not None + else [int(x) for x in gpu_ids] + ) + self._gpu_ids = ( + sorted(int(x) for x in _effective_pin_ids) if _effective_pin_ids else None + ) + else: + self._gpu_ids = None + + # Also record the RAW requested pin (before the fit narrowed it). Load + # dedupe compares this so a [0, 1] narrowed to [0] and re-sent as [0, 1] + # still matches, while /status keeps echoing the effective pin (#7239). + self._requested_gpu_ids = sorted(int(x) for x in gpu_ids) if gpu_ids else None + + if is_vulkan_backend and _vulkan_pin_ids is not None: + cmd += LlamaCppBackend._vulkan_pin_args(_vulkan_pin_ids) # User pass-through args go last so llama.cpp's last-wins parsing - # lets the user override Studio's auto-set flags. Already + # lets the user override Unsloth's auto-set flags. Already # validated by the route via validate_extra_args(). if extra_args: cmd.extend(str(a) for a in extra_args) logger.info(f"Appending user extra args to llama-server: {list(extra_args)}") + kv_cache_unified = _kv_unified_from_args(cmd) + logger.info(f"Starting llama-server: {' '.join(self._redacted_cmd_for_log(cmd))}") # Library paths so llama-server finds its shared libs and CUDA DLLs. env = self._llama_server_env_for_binary(binary) + if gpu_memory_mode == "manual": + self._clear_manual_placement_env(env) # Omitting --threads relies on llama.cpp's physical-core default, so # drop an inherited LLAMA_ARG_THREADS that would otherwise feed the # arg handler and silently force hardware_concurrency(). #5692 if "--threads" not in cmd: env.pop("LLAMA_ARG_THREADS", None) - # Reconcile the inherited LLAMA_ARG_* env with Studio's final + # Reconcile the inherited LLAMA_ARG_* env with Unsloth's final # decision: stripping CLI extras on a tensor->layer downgrade - # can't remove env vars, so the child could run a mode/KV Studio + # can't remove env vars, so the child could run a mode/KV Unsloth # didn't budget. if not tensor_parallel: # Layer split: clear a non-layer inherited split mode (and any @@ -6857,7 +8613,7 @@ class LlamaCppBackend: env.pop("LLAMA_ARG_SPLIT_MODE", None) env.pop("LLAMA_ARG_TENSOR_SPLIT", None) else: - # Studio owns the tensor split: it emits --tensor-split when it + # Unsloth owns the tensor split: it emits --tensor-split when it # picks an uneven one (CLI wins) and nothing when an even split # is safe. Clear any inherited LLAMA_ARG_TENSOR_SPLIT so the even # case can't be overridden by a stale env (the layer branch above @@ -6893,32 +8649,48 @@ class LlamaCppBackend: f"Data-center GPU detected: applied DC llama.cpp env tuning (multi_gpu={multi_gpu})" ) - # Pin to selected GPU(s). On ROCm, narrowing only - # CUDA_VISIBLE_DEVICES leaves an AMD child seeing the full set, so - # set HIP_VISIBLE_DEVICES too. Vulkan is pinned via --device - # (above), not here. - if gpu_indices is not None and not is_vulkan_backend: - pinned = ",".join(str(i) for i in gpu_indices) - env["CUDA_VISIBLE_DEVICES"] = pinned - try: - import torch as _torch - if getattr(_torch.version, "hip", None) is not None: - env["HIP_VISIBLE_DEVICES"] = pinned - # Do NOT also set ROCR_VISIBLE_DEVICES to the same - # value. ROCR_VISIBLE_DEVICES filters at the HSA/ROCr - # layer and HIP_VISIBLE_DEVICES at the HIP layer, so - # setting both with the same physical indices applies - # the mask twice: ROCR reduces the visible set and - # re-indexes it from 0, then HIP indexes into the - # already-reduced set. A single non-zero pin (e.g. - # "1") then points out of range at the HIP layer, HIP - # enumerates 0 devices, and llama.cpp falls back to - # CPU ("ggml_cuda_init: no ROCm-capable device is - # detected"). The HIP mask alone narrows correctly; - # clear any inherited ROCR mask so it can't double up. - env.pop("ROCR_VISIBLE_DEVICES", None) - except Exception as e: - logger.debug("Failed to set ROCm visibility env vars for child: %s", e) + # Pin to selected GPU(s) (issue #7164; resolved above into gpu_indices). + # On ROCm, narrowing only CUDA_VISIBLE_DEVICES leaves the AMD child + # seeing the full set, so set HIP_VISIBLE_DEVICES too. Vulkan is pinned + # via --device (above), not here. + # A deliberate zero-offload load with no GPU companions runs + # entirely on CPU, yet a visible CUDA device still costs the child + # ~0.5 GB (context + compute scratch) that the CPU-only + # classification below reports as free. Hide the GPUs so the load + # is exactly what it claims: zero VRAM (verified: GPU stays at idle + # baseline and generation runs). Companion loads keep the normal + # masking, and a user device pin (in extras or an inherited + # LLAMA_ARG_DEVICE) keeps control of its own devices -- the child + # aborts on a pin it can't see. The draft-device forms count too: + # llama-server parses them even with no drafter loaded. + _cpu_only_zero_offload = ( + gpu_memory_mode == "manual" + and gpu_layers == 0 + and not is_vulkan_backend + and not self._zero_offload_keeps_gpu_visible(cmd, env) + ) + if _cpu_only_zero_offload: + self._emit_child_gpu_visibility(env, "-1") + elif gpu_indices is not None and not is_vulkan_backend: + # When the user picked GPUs by index, align CUDA's ordering + # with the PCI-bus order the picker enumerated (nvidia-smi), + # so "GPU 1" in the UI is GPU 1 to llama.cpp -- not CUDA's + # default FASTEST_FIRST order (#5025). + if gpu_ids: + env["CUDA_DEVICE_ORDER"] = "PCI_BUS_ID" + # Mask on AMD at the ROCr/HSA layer: HIP-only masking still + # enumerates every agent first, which segfaults on a deselected + # unsupported GPU (e.g. gfx1036 iGPU under a gfx103X prebuilt). + self._emit_child_gpu_visibility( + env, ",".join(str(i) for i in gpu_indices), prefer_rocr = True + ) + elif manual_tensor_split_emitted and not is_vulkan_backend: + # A manual per-GPU ratio across ALL GPUs (no explicit pick, so + # no CUDA_VISIBLE_DEVICES mask above): the UI built the + # --tensor-split list in ascending physical/PCI index order, + # so pin the child's enumeration to that order too. The whole + # visible set stays in use; only its ordering is fixed. + self._pin_visible_gpu_order_for_split(env) # Captured before any text-only fallback strips it from cmd. launched_with_mmproj = "--mmproj" in cmd @@ -6928,7 +8700,7 @@ class LlamaCppBackend: # 'on') even when -ngl is explicit. That step has aborted on # some ROCm hosts (ggml-cuda.cu ROCm error during worst-case # estimation, e.g. MTP + mmproj models on gfx1151). When - # Studio's own VRAM math already placed the model + # Unsloth's own VRAM math already placed the model # (use_fit=False), the step is redundant second-guessing -- # retry once with --fit off before declaring the load failed. # Never retry when fit was requested (use_fit) or the caller @@ -6974,7 +8746,7 @@ class LlamaCppBackend: buffering = 1, ) logger.info(f"llama-server stdout/stderr -> {self._llama_log_path}") - except OSError as e: + except (OSError, UnicodeDecodeError) as e: # Best-effort; never block the load on logging. logger.debug(f"Could not open llama-server log file: {e}") self._llama_log_path = None @@ -6984,6 +8756,8 @@ class LlamaCppBackend: stdout = subprocess.PIPE, stderr = subprocess.STDOUT, text = True, + encoding = "utf-8", + errors = "replace", env = env, **_windows_hidden_subprocess_kwargs(), **_child_popen_kwargs(), @@ -7011,7 +8785,7 @@ class LlamaCppBackend: and _startup_crashed and not _split_axis_crash ): - # We forced --fit off because Studio's (conservative) VRAM + # We forced --fit off because Unsloth's (conservative) VRAM # math placed the model fully on GPU. A startup crash here # means that estimate was optimistic, so fall back to --fit # on and let llama.cpp offload rather than fail the load. @@ -7023,7 +8797,7 @@ class LlamaCppBackend: self._process.returncode, self._llama_log_path, ) - # Flip Studio's own --fit off (added first, before any + # Flip Unsloth's own --fit off (added first, before any # user extra args) to on; a user's later --fit still wins # by last-arg. Defensive: if absent, the default is already # --fit on, so leave it. @@ -7040,7 +8814,7 @@ class LlamaCppBackend: ): logger.warning( "llama-server crashed during startup (exit code %s) " - "with the default memory-fit step enabled; Studio " + "with the default memory-fit step enabled; Unsloth " "already verified the model fits, retrying once " "with --fit off. Crash log: %s", self._process.returncode, @@ -7077,7 +8851,6 @@ class LlamaCppBackend: self._effective_context_length = ( effective_ctx if effective_ctx > 0 else self._context_length ) - self._reconcile_effective_ctx_with_server() self._max_context_length = ( max_available_ctx if max_available_ctx > 0 else self._effective_context_length ) @@ -7117,10 +8890,17 @@ class LlamaCppBackend: _fa_rc, ) self._kill_process() + # The argv rewrite can't reach an env-only quantized V + # cache; drop it so the FA-off child doesn't abort on it. + if self._drop_env_quantized_v_cache(env): + logger.info( + "Dropped inherited quantized V-cache env for the " + "--flash-attn off retry (requires flash attention)." + ) cmd = _fa_cmd healthy = _spawn_and_wait(_fa_cmd, label = "-noflash") - # MTP from Studio's spec flags or the user's (extra_args + # MTP from Unsloth's spec flags or the user's (extra_args # --spec-type / LLAMA_ARG_SPEC_TYPE). The env reaches the child # only when neither emits a spec flag, so consult it only then. _launch_spec_env: Mapping[str, str] = ( @@ -7162,6 +8942,13 @@ class LlamaCppBackend: _probe_rc, ) self._kill_process() + # The argv rewrite can't reach an env-only quantized V + # cache; drop it so the FA-off child doesn't abort on it. + if self._drop_env_quantized_v_cache(env): + logger.info( + "Dropped inherited quantized V-cache env for the " + "--flash-attn off retry (requires flash attention)." + ) cmd = _fa_cmd healthy = ( _spawn_and_wait(_fa_cmd, label = "-noflash-mtp") @@ -7244,24 +9031,34 @@ class LlamaCppBackend: self._kill_process() # The #6415 split-axis abort is latched earlier (first spawn). # Skip if a cancel/unload is pending (mirrors the MTP guard). + _projector_msg = self._is_projector_incompatibility(out) + _signal_mmproj_guess = self._is_signal_crash( + _crash_rc + ) and not self._output_has_nonprojector_diagnostic(out) if ( launched_with_mmproj and not self._cancel_event.is_set() - and ( - self._is_projector_incompatibility(out) - or ( - self._is_signal_crash(_crash_rc) - and not self._output_has_nonprojector_diagnostic(out) - ) - ) + and (_projector_msg or _signal_mmproj_guess) ): - logger.warning( - "llama-server could not load this model's vision " - "projector (--mmproj). The installed llama.cpp build is " - "likely too old for it. Loading text-only for this " - "session; run 'unsloth studio update' to enable vision." - ) + if _projector_msg: + logger.warning( + "llama-server could not load this model's vision " + "projector (--mmproj). The installed llama.cpp build is " + "likely too old for it. Loading text-only for this " + "session; run 'unsloth studio update' to enable vision." + ) + else: + logger.warning( + "llama-server crashed while loading this model's vision " + "projector (--mmproj). Retrying text-only for this " + "session; if this persists, run 'unsloth studio update' " + "or check GPU/driver logs." + ) cmd = self._strip_mmproj_args(_last_spawn_cmd) + # This retry bypasses _spawn_and_wait, so refresh the + # launched-argv snapshot itself -- the zero-offload + # classification below must not see the stripped --mmproj. + _last_spawn_cmd = list(cmd) self._is_vision = False self._mmproj_has_audio = False self._start_llama_process(cmd, env) @@ -7270,14 +9067,30 @@ class LlamaCppBackend: # an OS-killed text-only retry still gets the OOM message. _retry_rc = self._process.poll() if self._process is not None else None self._kill_process() + # If the text-only retry ALSO hard-crashed (a signal, not + # OOM/timeout), the vision projector was never the cause: + # llama-server is faulting during GPU/driver init. Say so + # -- with the ROCm fix -- instead of blaming the mmproj. + if self._is_signal_crash(_retry_rc): + raise RuntimeError( + "llama-server crashed at startup on both the vision " + "and text-only attempts -- a GPU driver/runtime " + "initialization crash, not a model or vision-projector " + "problem. This often means an unsupported secondary " + "GPU; on AMD/ROCm, hide it with ROCR_VISIBLE_DEVICES " + "(e.g. ROCR_VISIBLE_DEVICES=0 exposes only the first " + "GPU) before launching Unsloth Studio." + ) + _retry_detail = self._classify_llama_start_failure( + "\n".join(self._stdout_lines[-50:]), + gguf_path, + self._model_identifier, + _retry_rc, + ) raise RuntimeError( - "Vision projector incompatible with this llama.cpp " - "build, and the text-only retry also failed: " - + self._classify_llama_start_failure( - "\n".join(self._stdout_lines[-50:]), - gguf_path, - self._model_identifier, - _retry_rc, + self._mmproj_retry_failure_message( + projector_confirmed = _projector_msg, + detail = _retry_detail, ) ) else: @@ -7292,6 +9105,33 @@ class LlamaCppBackend: self._healthy = True self._commit_effective_parallel_slots(n_parallel) + self._swa_full = swa_full + self._kv_cache_unified = kv_cache_unified + self._n_ubatch = max( + 0, + int(self._DEFAULT_N_UBATCH if _effective_ubatch is None else _effective_ubatch), + ) + self._flash_attn_enabled = ( + _flash_attn_enabled_from_args(_last_spawn_cmd, env = env) + and self._architecture != "grok" + ) + self._effective_cache_types = _effective_main_cache_types( + _last_spawn_cmd, + env, + ) + self._kv_cache_context_total = effective_ctx if effective_ctx > 0 else None + + # Server is up: adopt the real per-request context it allocated + # -- the length --fit chose, or a --parallel slot split -- so the + # reported context_length matches reality. (Querying /props + # before the spawn above always failed; the seeded value was the + # requested/native length.) + self._reconcile_effective_ctx_with_server() + if self._kv_cache_context_total is not None: + self._n_ubatch = min( + self._n_ubatch, + self._kv_cache_context_total, + ) # Commit caller intent only after _healthy=True so a failed start # can't poison the next inheritance check. None keeps prior, [] @@ -7301,24 +9141,39 @@ class LlamaCppBackend: self._extra_args = list(extra_args) self._extra_args_source = (model_identifier, hf_variant) self._requested_n_ctx = int(n_ctx) + # Local n_parallel may have been reduced above; the snapshot has the ask. + self._requested_n_parallel = max(1, int(_pending_load_kwargs["n_parallel"])) # Commit the known-good snapshot + whether MTP+tensor is live, then # watch this load for a mid-generation crash. self._last_load_kwargs = _pending_load_kwargs self._mtp_runtime_fallback_active = _mtp_active_for_launched_server self._start_mtp_crash_watchdog() - # Catch silent CPU fallback when GPU was intended (#5106). - self._gpu_offload_active = self._classify_gpu_offload( - gpu_indices is not None or use_fit, gpus or [] - ) - if self._gpu_offload_active is False: + # Catch silent CPU fallback when GPU was intended (#5106). Manual + # offload (no picker) leaves gpu_indices None and use_fit False, so + # include its GPU-layer intent; use the preserved probe since + # auto-layers/manual empty `gpus`. A deliberate zero-offload load + # classifies by its launched argv instead: the main model is + # CPU-only by construction and must read False (not None), or + # training needlessly unloads a server holding no VRAM. + _deliberate_cpu_only = gpu_memory_mode == "manual" and gpu_layers == 0 + if _deliberate_cpu_only: + self._gpu_offload_active = self._zero_offload_gpu_flag( + _last_spawn_cmd, _detected_gpus, env + ) + else: + self._gpu_offload_active = self._classify_gpu_offload( + gpu_indices is not None or use_fit or gpu_memory_mode == "manual", + _detected_gpus, + ) + if self._gpu_offload_active is False and not _deliberate_cpu_only: logger.warning( "llama-server appears to have loaded the model entirely " - "on CPU even though Studio detected at least one GPU. " + "on CPU even though Unsloth detected at least one GPU. " "This usually means the prebuilt binary's GPU backend " "failed to load -- on Windows, cudart64_X.dll / " "cublas64_X.dll could not be resolved. Reinstall the " - "Studio llama.cpp prebuilt or install a matching CUDA " + "Unsloth llama.cpp prebuilt or install a matching CUDA " "toolkit (issue unslothai/unsloth#5106).", ) @@ -7354,6 +9209,15 @@ class LlamaCppBackend: if not self._healthy: return False + # Snapshot the files the server actually loaded. If a GGUF shard or a + # LoRA/control-vector sidecar is swapped on disk afterwards while the + # old weights stay mapped, save_slots_for_resume() compares against + # this and refuses to persist KV that a reload could misapply. + if self._slot_save_dir: + self._slot_loaded_identity = ( + self._gguf_file_identity(self._gguf_path), + self._slot_launch_fingerprint(), + ) return True def _build_speculative_flags( @@ -7478,18 +9342,29 @@ class LlamaCppBackend: caps = self.probe_server_capabilities(binary) mtp_token = caps.get("mtp_token") if caps else None if not mtp_token: - logger.warning( - "Requested MTP speculative decoding but " - "llama-server lacks --spec-type mtp/draft-mtp; " - "run `unsloth studio update`. Loading without " - "speculative decoding." - ) + inconclusive = bool(caps.get("mtp_probe_inconclusive")) if caps else True + if inconclusive: + logger.info( + "Requested MTP speculative decoding but llama-server MTP " + "capability probe was inconclusive; loading without " + "speculative decoding." + ) + else: + logger.warning( + "Requested MTP speculative decoding but " + "llama-server lacks --spec-type mtp/draft-mtp; " + "run `unsloth studio update`. Loading without " + "speculative decoding." + ) # Override an inherited LLAMA_ARG_SPEC_TYPE=draft-mtp (CLI wins # over env) so the child matches the binary-capability gate and # the no-MTP budget, like the sibling no-head/non-MTP fallbacks. flags.append("--spec-default") self._speculative_type = "default" - self._spec_fallback_reason = "binary_no_mtp" + if inconclusive: + self._spec_fallback_reason = None + else: + self._spec_fallback_reason = "binary_no_mtp" return False draft_n_max = _resolved_draft_n_max() n_max_flag = caps.get("spec_draft_n_max_flag") or "--spec-draft-n-max" @@ -7615,7 +9490,7 @@ class LlamaCppBackend: logger.info( "Auto: MLA embedded-MTP model detected; llama.cpp's MLA/DSA " "MTP path is slower than no speculation, so using ngram-mod " - "instead. Override via the Studio Speculative Decoding " + "instead. Override via the Unsloth Speculative Decoding " "dropdown or UNSLOTH_MLA_MTP_ENABLED=1." ) _emit_ngram_mod() @@ -7643,7 +9518,7 @@ class LlamaCppBackend: f"MTP GGUF detected but model size {_mtp_size_b:.1f}B " "is below the 3B speedup threshold; using ngram-mod " "only (zero-VRAM, no draft head). Override via " - "--spec-type or the Studio Speculative Decoding " + "--spec-type or the Unsloth Speculative Decoding " "dropdown." ) _emit_ngram_mod() @@ -7674,7 +9549,13 @@ class LlamaCppBackend: gguf_path: Optional[str] = None, spec_draft_n_max: Optional[int] = None, tensor_parallel: bool = False, + gpu_memory_mode: Literal["auto", "manual"] = "auto", + gpu_layers: int = -1, + n_cpu_moe: int = 0, + tensor_split: Optional[List[float]] = None, + gpu_ids: Optional[List[int]] = None, mtp_draft_path: Optional[str] = None, + n_parallel: int = 1, preserve_multi_gpu_on_layer: bool = False, ) -> bool: """True iff the live server already satisfies these load kwargs. @@ -7710,7 +9591,6 @@ class LlamaCppBackend: if _norm(self._cache_type_kv) != _norm(cache_type_kv): return False - # Reconcile a user --split-mode in extras AND an inherited tensor # LLAMA_ARG_SPLIT_MODE env, but only against a server that actually # launched tensor: if load_model downgraded to layer split it scrubbed @@ -7730,6 +9610,39 @@ class LlamaCppBackend: ): return False + # The diffusion runner is mode-agnostic (always "auto", ignores the + # layer/MoE/split knobs), so a standing manual preference in the + # request must not force a needless reload -- only the GPU pick matters. + if not self._is_diffusion: + requested_extra_args = extra_args if extra_args is not None else self._extra_args + if self._swa_full != _swa_full_from_args_or_env(requested_extra_args): + return False + # A GPU-memory-mode flip (Unsloth / manual) must always reload. + if self._gpu_memory_mode != gpu_memory_mode: + return False + # Requested-vs-requested (like n_ctx): comparing the effective count + # would reload forever whenever the fitter launched fewer slots. + if self._requested_n_parallel != max(1, int(n_parallel)): + return False + # Manual: a layer-count change always reloads (covers Auto(-1) <-> a + # pinned count); MoE/split only matter with an explicit offload. + if gpu_memory_mode == "manual" and ( + self._gpu_layers != gpu_layers + or ( + gpu_layers >= 0 + and ( + self._n_cpu_moe != n_cpu_moe + or (self._tensor_split or None) != (tensor_split or None) + ) + ) + ): + return False + # A changed GPU pick must reload. Regular GGUF accepts either the raw + # requested placement pool or the effective status-echoed subset; + # diffusion compares its normalized single-device pick. + if not self.matches_gpu_ids(gpu_ids): + return False + # Compare on the canonical requested mode. With --spec-type in # extra_args the backend stores None; mirror that here. if _extra_args_set_spec_type(extra_args): @@ -7785,6 +9698,7 @@ class LlamaCppBackend: current = list(self._extra_args) if self._extra_args is not None else [] if list(extra_args) != current: return False + self._record_matching_gpu_request(gpu_ids) return True def _classify_gpu_offload( @@ -7798,6 +9712,79 @@ class LlamaCppBackend: return None return classify_gpu_offload_lines(self._stdout_lines) + @staticmethod + def _cmd_has_gpu_companion(cmd: list, env: Optional[Mapping[str, str]] = None) -> bool: + """True when the argv/env carries a GPU companion: any --mmproj form, or + a drafter (Studio's --model-draft, the extras aliases, or the + LLAMA_ARG_SPEC_DRAFT_* env) -- these offload to the GPU regardless of + the main ``--gpu-layers``. A drafter explicitly forced to CPU + (--spec-draft-ngl 0 / --spec-draft-device cpu) doesn't count.""" + if any(str(a).startswith("--mmproj") for a in cmd): + return True + if _extra_args_mtp_draft_path(cmd, env) is None: + return False + return not _extra_args_draft_offloaded_to_cpu(cmd, env) + + @staticmethod + def _zero_offload_keeps_gpu_visible(cmd: list, env: Optional[Mapping[str, str]] = None) -> bool: + """Whether a zero-layer launch still has a reason to use visible GPUs. + + Keep this shared by child masking and post-launch residency bookkeeping: + a device pin, surviving tensor mode, mmproj, or GPU drafter prevents the + launch from being a confirmed zero-VRAM server. + """ + return ( + LlamaCppBackend._cmd_has_gpu_device_pin(cmd, env) + or _effective_tensor_parallel(cmd, False, env) + or LlamaCppBackend._cmd_has_gpu_companion(cmd, env) + ) + + @staticmethod + def _cmd_has_gpu_device_pin(cmd: list, env: Optional[Mapping[str, str]] = None) -> bool: + """True when the effective main or draft ``--device`` pin names a GPU.""" + main_flags = {"--device", "-dev"} + draft_flags = {"--spec-draft-device", "-devd", "--device-draft"} + last_main: Optional[str] = None + last_draft: Optional[str] = None + args = [str(arg) for arg in cmd] + for index, raw in enumerate(args): + flag = _flag_name(raw) + _, equals, inline = raw.partition("=") + if flag not in main_flags and flag not in draft_flags: + continue + value = inline if equals else (args[index + 1] if index + 1 < len(args) else "") + if flag in main_flags: + last_main = value + else: + last_draft = value + if last_main is None: + last_main = (env or {}).get("LLAMA_ARG_DEVICE") + + def _names_gpu(value: Optional[str]) -> bool: + if value is None: + return False + devices = [item.strip().lower() for item in value.split(",") if item.strip()] + return not devices or any(item not in ("cpu", "none") for item in devices) + + return _names_gpu(last_main) or _names_gpu(last_draft) + + @staticmethod + def _zero_offload_gpu_flag( + spawn_cmd: list, + detected_gpus: list, + env: Optional[Mapping[str, str]] = None, + ) -> Optional[bool]: + """GPU-residency flag for a deliberate manual zero-offload load. The + main model is CPU-only by construction, but device pins, tensor mode, + mmproj, and GPU drafters can still make the server hold VRAM. The counted + offload classifier cannot see those allocations. This uses the same + predicate as the launch-time zero-VRAM mask; None means no GPU signal.""" + if not detected_gpus: + return None + if LlamaCppBackend._is_vulkan_backend(): + return True + return LlamaCppBackend._zero_offload_keeps_gpu_visible(spawn_cmd, env) + def load_cancelled(self) -> bool: """True if a load was cancelled (e.g. via unload/_cancel_event) and not yet consumed by the next load_model. Lets the tensor->layer fallback @@ -7808,6 +9795,7 @@ class LlamaCppBackend: """Terminate the subprocess and cancel any in-flight download.""" self._cancel_event.set() with self._lock: + self._unload_epoch += 1 self._kill_process() logger.info(f"Unloaded GGUF model: {self._model_identifier}") self._model_identifier = None @@ -7830,6 +9818,16 @@ class LlamaCppBackend: self._effective_context_length = None self._max_context_length = None self._reset_effective_parallel_slots() + self._slot_save_dir = None + self._slot_save_binary = None + self._slot_loaded_identity = None + self._prompt_cache_disabled = False + self._swa_full = False + self._kv_cache_unified = False + self._n_ubatch = self._DEFAULT_N_UBATCH + self._flash_attn_enabled = True + self._effective_cache_types = ("f16", "f16") + self._kv_cache_context_total = None self._chat_template = None self._chat_template_override = None self._supports_reasoning = False @@ -7840,12 +9838,22 @@ class LlamaCppBackend: self._supports_preserve_thinking = False self._supports_tools = False self._cache_type_kv = None + # GPU-pin state describes the active runner only; clear it so an explicit + # pin never leaks into the next (or diffusion) runner. + self._gpu_ids = None + self._requested_gpu_ids = None self._tensor_parallel = False + self._gpu_memory_mode = "auto" + self._gpu_layers = -1 + self._n_cpu_moe = 0 + self._tensor_split = None self._layer_preserves_tensor_intent = False self._speculative_type = None self._requested_spec_mode = None self._spec_draft_n_max = None self._n_layers = None + self._n_experts = None + self._leading_dense_block_count = None self._n_kv_heads = None self._n_kv_heads_by_layer = None self._n_heads = None @@ -7908,6 +9916,10 @@ class LlamaCppBackend: # Clear healthy so a /load during the replacement's warm-up can't # short-circuit against the previous server's health (#5401). self._healthy = False + # Reset to unknown so the training guard treats the next (still + # loading) server as VRAM-resident rather than reading the killed + # server's stale zero-offload flag until the health probe reclassifies. + self._gpu_offload_active = None # Drives _wait_for_vram_settle in the next load_model; set in finally # so both in-process and frontend Apply paths record the kill. self._last_kill_monotonic = time.monotonic() @@ -7947,7 +9959,7 @@ class LlamaCppBackend: return try: path.parent.mkdir(parents = True, exist_ok = True) - path.write_text(f"{pid}:{cls._pid_start_identity(pid)}") + path.write_text(f"{pid}:{cls._pid_start_identity(pid)}", encoding = "utf-8") except Exception as e: logger.debug(f"Could not write llama-server pidfile: {e}") @@ -8021,7 +10033,7 @@ class LlamaCppBackend: def _pid_parent_is_alive(pid: int) -> bool: """True if the recorded server's parent is still running, i.e. the server is NOT orphaned. Lets the cross-session reap kill only a true orphan (parent - gone) and never a live server owned by a running Studio, regardless of which + gone) and never a live server owned by a running Unsloth, regardless of which process performs the sweep. Biased toward "alive" on uncertainty so a live server is never mistakenly reaped.""" try: @@ -8061,9 +10073,9 @@ class LlamaCppBackend: @classmethod def _reap_recorded_pid(cls) -> int: """Kill the exact llama-server PID recorded at spawn, but only when it is a - genuine orphan -- its parent (the Studio that spawned it) is gone. This is + genuine orphan -- its parent (the Unsloth that spawned it) is gone. This is the cross-session backstop the parent-death reaper (Job Object / - PR_SET_PDEATHSIG) cannot cover: an orphan left by an already-dead Studio + PR_SET_PDEATHSIG) cannot cover: an orphan left by an already-dead Unsloth (macOS, a best-effort failure, or a pre-existing orphan). Path-independent, so it also catches an orphan the install-root match would miss. @@ -8081,7 +10093,7 @@ class LlamaCppBackend: pid = -1 identity = "" try: - pid_str, _, identity = path.read_text().strip().partition(":") + pid_str, _, identity = path.read_text(encoding = "utf-8").strip().partition(":") pid = int(pid_str) except Exception: pid = -1 @@ -8120,7 +10132,7 @@ class LlamaCppBackend: """Kill orphaned llama-server processes started by studio. Only kills processes whose resolved binary lives under a known - Studio install dir (or matches an exact env-var override), to avoid + Unsloth install dir (or matches an exact env-var override), to avoid terminating unrelated llama-server instances. Mirrors every location _find_llama_server_binary() can return, so orphans from any supported install path are cleaned up. @@ -8140,7 +10152,7 @@ class LlamaCppBackend: try: # -- Build the ownership allowlist -------------------------------- # exact_binaries -- env var overrides (exact path match). - # install_roots -- Studio-owned dir trees (binary must be under one). + # install_roots -- Unsloth-owned dir trees (binary must be under one). install_roots: list[Path] = [] # Env-mode custom root (mirrors _find_llama_server_binary). @@ -8150,7 +10162,7 @@ class LlamaCppBackend: install_roots.append(_resolved_sr / "llama.cpp") # Primary install dir (default mode only). Env-mode skips this so a - # custom-root Studio can't kill a default-install Studio's server. + # custom-root Unsloth can't kill a default-install Unsloth's server. if not _is_custom_root: install_roots.append(Path.home() / ".unsloth" / "llama.cpp") @@ -8224,6 +10236,11 @@ class LlamaCppBackend: if not is_ours: continue + # A live parent means a running Unsloth (or the user's + # shell) still owns it -- not an orphan. + if LlamaCppBackend._pid_parent_is_alive(proc.info["pid"]): + continue + proc.kill() killed += 1 logger.info( @@ -8243,6 +10260,8 @@ class LlamaCppBackend: ["pgrep", "-a", "-f", "llama-server"], capture_output = True, text = True, + encoding = "utf-8", + errors = "replace", timeout = 5, env = child_env_without_native_path_secret(), ) @@ -8276,6 +10295,9 @@ class LlamaCppBackend: if not owned: continue + if LlamaCppBackend._pid_parent_is_alive(pid): + continue + try: os.kill(pid, signal.SIGKILL) killed += 1 @@ -8296,7 +10318,7 @@ class LlamaCppBackend: def _fit_off_retry_eligible(cmd: "list[str]", use_fit: bool) -> bool: """Whether a llama-server startup crash may be retried with --fit off. - Only when Studio's own VRAM math placed the model (use_fit=False) + Only when Unsloth's own VRAM math placed the model (use_fit=False) and nothing on the command line set the fit mode explicitly (-fit / --fit, space- or equals-form). --fit-ctx / --fit-target / -fitc / -fitt tune the fit step but do not select the mode, so @@ -8337,6 +10359,260 @@ class LlamaCppBackend: return False return True + def _slot_launch_fingerprint(self) -> tuple: + # KV validity keys on extra args, stat'd sidecar weights, effective ctx. + sidecars = [] + for path in self._sidecar_weight_files(): + try: + st = os.stat(path) + sidecars.append((path, st.st_size, st.st_mtime_ns)) + except OSError: + sidecars.append((path, None, None)) + return ( + tuple(self._extra_args or ()), + tuple(sidecars), + self._requested_n_ctx, + self._effective_context_length, + self._effective_cache_types, + self.effective_parallel_slots, + self._swa_full, + self._kv_cache_unified, + self._n_ubatch, + self._flash_attn_enabled, + ) + + def _gguf_file_identity(self, path) -> Optional[tuple]: + # (size, mtime_ns) per shard: a split GGUF keys KV validity on every sibling. + p = Path(path) + paths = [p] + m = _SHARD_FULL_RE.match(p.name) + if m: + prefix, _first, total = m.groups() + paths = [ + p.with_name(f"{prefix}-{i:05d}-of-{total}{p.suffix}") + for i in range(1, int(total) + 1) + ] + try: + return tuple((sp.stat().st_size, sp.stat().st_mtime_ns) for sp in paths) + except OSError: + return None + + _SIDECAR_WEIGHT_FLAGS = ( + "--lora", + "--lora-scaled", + "--control-vector", + "--control-vector-scaled", + ) + + def _sidecar_weight_files(self) -> list[str]: + # llama.cpp: comma-separated paths, FNAME:SCALE on -scaled (older builds: FNAME SCALE). + args = [str(a).strip() for a in (self._extra_args or ())] + files: list[str] = [] + for i, arg in enumerate(args): + flag = _flag_name(arg) + _, sep, inline = arg.partition("=") + if flag not in self._SIDECAR_WEIGHT_FLAGS: + continue + operand = inline if sep else (args[i + 1] if i + 1 < len(args) else "") + if not operand: + continue + candidates = [operand] + pieces = [p for p in operand.split(",") if p] + if len(pieces) > 1: + candidates.extend(pieces) + if flag.endswith("-scaled"): + for item in list(candidates): + # ":" tail is a scale; rpartition spares drive letters. + head, colon, tail = item.rpartition(":") + if not (colon and head): + continue + try: + float(tail) + except ValueError: + continue + candidates.append(head) + for cand in candidates: + if cand not in files: + files.append(cand) + return files + + def _prompt_cache_off(self) -> bool: + # Caching off makes restores useless; last prompt-cache flag wins, env only when unset. + last = None + for arg in self._extra_args or (): + flag = arg.strip().split("=", 1)[0] + if flag in ("--cache-prompt", "--no-cache-prompt"): + last = flag + if last is not None: + return last == "--no-cache-prompt" + if self._prompt_cache_disabled: + return True + if os.environ.get("LLAMA_ARG_NO_CACHE_PROMPT") is not None: + return True + env = (os.environ.get("LLAMA_ARG_CACHE_PROMPT") or "").strip().lower() + return env in _LLAMA_ARG_FALSE_VALUES + + def save_slots_for_resume( + self, should_abort: Optional[Callable[[], bool]] = None + ) -> Optional[dict]: + if ( + not self.is_loaded + or not self._slot_save_dir + or not self._gguf_path + or self._prompt_cache_off() + ): + return None + # Same predicate as the estimator's SWA path: a window alone is not enough. + # phi3 GGUFs carry attention.sliding_window but no key/value length, and + # llama.cpp forces them back to a non-SWA cache, so their slots do restore. + if ( + (self._sliding_window or 0) > 0 + and self._kv_key_length is not None + and self._kv_value_length is not None + and not self._swa_full + ): + logger.debug("Skipping slot save: compact SWA cache cannot be reused after restart") + return None + save_dir = Path(self._slot_save_dir) + gguf_stat = self._gguf_file_identity(self._gguf_path) + if gguf_stat is None: + return None + launch = self._slot_launch_fingerprint() + # If the GGUF or a sidecar was swapped on disk while the original weights + # stayed mapped, the live KV belongs to the old weights but a reload would + # load the new file. Persisting it would let restore misapply stale KV. + if self._slot_loaded_identity is not None and self._slot_loaded_identity != ( + gguf_stat, + launch, + ): + logger.debug("Skipping slot save: model files changed on disk since load") + return None + try: + estimate = self._estimate_kv_cache_bytes( + self._kv_cache_context_total + or self._effective_context_length + or self._context_length + or 0, + max(self._effective_cache_types, key = _kv_bytes_per_elem), + n_parallel = self.effective_parallel_slots, + swa_full = self._swa_full, + kv_unified = self._kv_cache_unified, + n_ubatch = self._n_ubatch, + flash_attn = self._flash_attn_enabled, + ) + # Skip before writing anything when the estimate alone blows the cap, + # rather than fully writing a slot and discarding it afterwards. + if estimate > _SLOT_SAVE_MAX_BYTES: + logger.debug( + "Skipping slot save: estimated %d bytes exceeds cap %d", + estimate, + _SLOT_SAVE_MAX_BYTES, + ) + return None + # A 0 estimate means metadata was insufficient, not a zero-byte cache: + # a slot can still be many GiB, so demand room for the whole cap before + # trusting the post-write check. + required = (estimate if estimate > 0 else _SLOT_SAVE_MAX_BYTES) + (1 << 30) + if shutil.disk_usage(save_dir).free < required: + logger.debug("Skipping slot save: insufficient free disk") + return None + except Exception: + pass + token = uuid.uuid4().hex[:8] + entries: list[dict] = [] + total_bytes = 0 + for slot in range(self.effective_parallel_slots): + # A request pending mid-save waits on the gate; stop wasting its time. + if should_abort is not None and should_abort(): + break + filename = f"resume-{token}-slot{slot}.bin" + path = save_dir / filename + try: + resp = httpx.post( + f"{self.base_url}/slots/{slot}", + params = {"action": "save"}, + json = {"filename": filename}, + headers = self._auth_headers, + timeout = _SLOT_SAVE_HTTP_TIMEOUT, + trust_env = False, + ) + except Exception as e: + logger.debug(f"slot {slot} save failed: {e}") + with contextlib.suppress(OSError): + path.unlink() + break + if resp.status_code != 200: + logger.debug(f"slot {slot} save returned HTTP {resp.status_code}") + with contextlib.suppress(OSError): + path.unlink() + continue + try: + body = resp.json() + if not isinstance(body, dict): + raise ValueError("slot save response was not a JSON object") + n_saved = int(body.get("n_saved") or 0) + except Exception as e: + # A 200 that still wrote a file but returns a malformed body must + # clean up like the transport/HTTP error paths above, or the file + # (which holds chat KV) is orphaned until the next startup sweep. + logger.debug(f"slot {slot} save returned an invalid response: {e}") + with contextlib.suppress(OSError): + path.unlink() + continue + if n_saved <= 0: + with contextlib.suppress(OSError): + path.unlink() + continue + # Account by the bytes actually on disk, not the server-reported + # count, so the cap holds even if a custom binary under-reports. + try: + n_written = path.stat().st_size + except OSError: + n_written = 0 + total_bytes += n_written + entries.append({"id": slot, "filename": filename, "n_saved": n_saved}) + if total_bytes > _SLOT_SAVE_MAX_BYTES: + break # already over the cap; the discard below cleans up + if not entries: + return None + if total_bytes > _SLOT_SAVE_MAX_BYTES: + logger.debug( + "Discarding slot save: %d bytes exceeds cap %d", + total_bytes, + _SLOT_SAVE_MAX_BYTES, + ) + for entry in entries: + with contextlib.suppress(OSError): + (save_dir / entry["filename"]).unlink() + return None + return { + "dir": self._slot_save_dir, + "binary": self._slot_save_binary, + "gguf": str(self._gguf_path), + "gguf_stat": gguf_stat, + "launch": launch, + "slots": entries, + } + + def restore_slots_for_resume(self, manifest: dict) -> None: + if not self.is_loaded or not self._slot_save_dir: + return + for entry in manifest.get("slots") or []: + try: + resp = httpx.post( + f"{self.base_url}/slots/{int(entry['id'])}", + params = {"action": "restore"}, + json = {"filename": str(entry["filename"])}, + headers = self._auth_headers, + timeout = _SLOT_SAVE_HTTP_TIMEOUT, + trust_env = False, + ) + except Exception as e: + logger.debug(f"slot restore failed: {e}") + break + if resp.status_code != 200: + logger.debug(f"slot {entry.get('id')} restore returned HTTP {resp.status_code}") + def _maybe_recover_from_mtp_crash(self, exc: Optional[BaseException] = None) -> bool: """Schedule one background reload without MTP after a mid-generation death. @@ -8350,15 +10626,18 @@ class LlamaCppBackend: return False if not self._mtp_runtime_fallback_active: return False - if not self._last_load_kwargs or self._process is None: + # Read before claiming: a raise after the claim strands the flag, and nothing + # else clears it, blocking every later respawn. + kwargs = self._last_load_kwargs + proc = self._process + if not kwargs or proc is None: return False # Single-flight: the first failure claims the reload. with self._mtp_runtime_fallback_lock: if self._mtp_runtime_fallback_in_progress: return False self._mtp_runtime_fallback_in_progress = True - snapshot = dict(self._last_load_kwargs) - proc = self._process + snapshot = dict(kwargs) def _recover(): try: @@ -8406,7 +10685,14 @@ class LlamaCppBackend: with self._mtp_runtime_fallback_lock: self._mtp_runtime_fallback_in_progress = False - threading.Thread(target = _recover, daemon = True, name = "mtp-crash-reload").start() + try: + threading.Thread(target = _recover, daemon = True, name = "mtp-crash-reload").start() + except RuntimeError as exc: + # Release the claim: a reload that never started would block respawn forever. + with self._mtp_runtime_fallback_lock: + self._mtp_runtime_fallback_in_progress = False + logger.error(f"Could not start the MTP-crash reload: {exc}") + return False return True def _start_mtp_crash_watchdog(self) -> None: @@ -8504,7 +10790,12 @@ class LlamaCppBackend: @staticmethod def _ctx_integrity_flags( - n_parallel: int, use_fit: bool, requested_ctx: int, effective_ctx: int, caps: dict + n_parallel: int, + use_fit: bool, + auto_fit: bool, + requested_ctx: int, + effective_ctx: int, + caps: dict, ) -> list[str]: """Flags that keep the per-request window equal to the advertised ctx. @@ -8512,14 +10803,28 @@ class LlamaCppBackend: ``--kv-unified`` default, silently splitting ``-c`` into per-slot windows of ``-c / N``; restore the shared pool so one request can use the full context. With ``--fit on``, ``--fit-ctx`` floors the fit step - at an explicitly requested ctx (default floor is 4096) so it offloads - or fails instead of silently shrinking the window. + at an explicitly requested ctx so it offloads or fails instead of + silently shrinking the window. The 8192 auto-floor and the tighter + ``--fit-target`` margin apply only under Manual + Auto (``auto_fit``), + which omits ``-c``: on the legacy auto path ``-c 0`` already pins the + native window and ``--fit-ctx 8192`` would override it down to 8192. """ flags: list[str] = [] if n_parallel > 1 and caps.get("supports_kv_unified"): flags.append("--kv-unified") - if use_fit and requested_ctx > 0 and effective_ctx > 0 and caps.get("supports_fit_ctx"): - flags.extend(["--fit-ctx", str(effective_ctx)]) + if use_fit and caps.get("supports_fit_ctx"): + if requested_ctx > 0 and effective_ctx > 0: + # Floor the fit step at the explicitly requested ctx. + flags.extend(["--fit-ctx", str(effective_ctx)]) + elif auto_fit: + # Manual + Auto omits -c, so floor at 8192 so --fit doesn't + # shrink the window below a usable size. + flags.extend(["--fit-ctx", "8192"]) + if use_fit and auto_fit and caps.get("supports_fit_target"): + # llama.cpp's --fit leaves 1 GiB free per device by default; + # tighten that to 512 MiB so it packs more of the model onto + # the GPU before spilling to system RAM. + flags.extend(["--fit-target", "512"]) return flags def _query_server_n_ctx(self) -> Optional[int]: @@ -8540,7 +10845,7 @@ class LlamaCppBackend: return None def _reconcile_effective_ctx_with_server(self) -> None: - """Adopt the server's real ``n_ctx`` when it is below Studio's value. + """Adopt the server's real ``n_ctx`` when it is below Unsloth's value. Keeps ``context_length`` (load response, status route, passthrough ``max_tokens`` ceiling) honest; clients sized to the requested value @@ -8549,6 +10854,8 @@ class LlamaCppBackend: actual_n_ctx = self._query_server_n_ctx() if not actual_n_ctx or actual_n_ctx <= 0: return + slots = 1 if self._kv_cache_unified else self.effective_parallel_slots + self._kv_cache_context_total = actual_n_ctx * slots if self._effective_context_length and actual_n_ctx < self._effective_context_length: logger.warning( "llama-server allocated a smaller per-request context than " @@ -8716,6 +11023,75 @@ class LlamaCppBackend: except Exception: logger.debug("Could not close httpx client", exc_info = True) + @staticmethod + def _install_cancel_aware_read( + client: "httpx.Client", + cancel_event: threading.Event, + response: Optional["httpx.Response"] = None, + poll_s: float = 0.2, + ) -> None: + """Wrap the httpcore stream so the reader interrupts its own blocked recv() on cancel. + + A cross-thread socket shutdown wakes a parked recv() on POSIX but not on + Windows (Winsock), so read in short slices and poll cancel_event between them + (plain or TLS); slice timeouts are swallowed so a slow-but-alive stream survives. + httpcore snapshots request.extensions["timeout"]["read"] once at body start, so + given ``response`` we re-read the live value per call to honor the post-first-token + stall timeout instead of the long prefill timeout.""" + import httpcore + + def _live_read_timeout() -> Optional[float]: + if response is None: + return None + try: + ext = response.request.extensions.get("timeout") + if isinstance(ext, dict): + value = ext.get("read") + if isinstance(value, (int, float)): + return float(value) + except Exception: + pass + return None + + try: + pool = getattr(getattr(client, "_transport", None), "_pool", None) + for connection in list(getattr(pool, "_connections", []) or []): + inner = getattr(connection, "_connection", None) + stream = getattr(inner, "_network_stream", None) + if stream is None or getattr(stream, "_unsloth_cancel_wrapped", False): + continue + orig_read = stream.read + + def read( + max_bytes, + timeout = None, + _orig = orig_read, + ): + live = _live_read_timeout() + effective = live if live is not None else timeout + deadline = None if effective is None else time.monotonic() + effective + while True: + if cancel_event.is_set(): + raise httpcore.ReadError("stream cancelled by user") + if deadline is None: + step = poll_s + else: + remaining = deadline - time.monotonic() + if remaining <= 0: + raise httpcore.ReadTimeout("read operation timed out") + step = min(poll_s, remaining) + try: + return _orig(max_bytes, timeout = step) + except httpcore.ReadTimeout: + if deadline is not None and time.monotonic() >= deadline: + raise + continue # slow but alive: keep reading + + stream.read = read + stream._unsloth_cancel_wrapped = True + except Exception: + logger.debug("Could not install cancel-aware read", exc_info = True) + @staticmethod @contextlib.contextmanager def _stream_with_retry( @@ -8773,6 +11149,11 @@ class LlamaCppBackend: headers = headers, ) as response: _response_ref[0] = response + if cancel_event is not None: + # Portable mid-stream cancel: the reader polls cancel itself, so + # Stop interrupts a stalled read where the watcher's Windows socket + # shutdown does not. Pass response to honor the live stall timeout. + LlamaCppBackend._install_cancel_aware_read(client, cancel_event, response) if cancel_event is not None and cancel_event.is_set(): raise _LlamaStreamCancelled yield response @@ -8785,6 +11166,21 @@ class LlamaCppBackend: finally: _cancel_closed.set() + def _server_socket_is_open(self, timeout_s: float = 0.15) -> bool: + """True if anything still accepts on the server port. + + The listening socket dies with the process, so this tells a live server + from a dead one without waiting for the child to become reapable. + """ + port = self._port + if not port: + return False + try: + with socket.create_connection(("127.0.0.1", port), timeout = timeout_s): + return True + except OSError: + return False + def _respawn_if_dead(self) -> bool: """Relaunch the llama-server if its process has exited. @@ -8794,28 +11190,114 @@ class LlamaCppBackend: recover, returning True once healthy. Serialised on ``_respawn_lock`` so many generations hitting the dead server trigger at most one reload. """ + # Read outside the lock so a queued caller can tell the replacement from the child + # its own error came from; otherwise each burns the grace wait below, and that + # sleep is held under the lock, so the waits serialise. + served_by = self._process with self._respawn_lock: proc = self._process if proc is None: return False - if proc.poll() is None: - # Process is alive: either a concurrent caller already respawned - # it (healthy), or this connection error wasn't a dead server. + if self._cancel_event.is_set(): + # unload_model sets this before it kills, so the child can still be + # accepting. Reporting it healthy would aim the retry at a server + # that is deliberately going away. + return False + if proc is not served_by: + # Replaced while we queued: this child never served our request. return self._healthy - kwargs = self._last_load_kwargs - if not kwargs: - return False - logger.warning( - f"llama-server for '{self._model_identifier}' exited " - f"(code {proc.returncode}); respawning to recover the session" - ) - with self._lock: - self._healthy = False + if proc.poll() is None: + # Still serving, so the error was transient. Charging it the grace below + # would cost a second per caller, serialised under this lock. + if self._server_socket_is_open(): + return self._healthy + # A closing server can beat its own exit status: calling it alive returns + # the stale _healthy and spends the retry on the corpse. + deadline = time.monotonic() + _RESPAWN_REAP_GRACE_S + while proc.poll() is None and time.monotonic() < deadline: + time.sleep(0.05) + if proc.poll() is None: + # Alive: either a concurrent caller already respawned it (healthy), or + # this connection error wasn't a dead server. + return self._healthy + with self._mtp_runtime_fallback_lock: + if self._mtp_runtime_fallback_in_progress: + # An MTP-free reload owns this corpse; replaying the old kwargs + # restarts the crashing config and aborts that reload. + logger.info("Respawn skipped: an MTP-free reload is already recovering.") + return False + # The RLock lets the load_model below re-enter it. + with self._serial_load_lock: + if self._process is not proc: + logger.info("Respawn skipped: a newer load is already active.") + return self._healthy + # Snapshot under _lock, the one unload_model holds, so a teardown is + # either wholly before us (flag set) or wholly after (epoch bumped). + # _serial_load_lock alone would not exclude it: unload never takes it. + with self._lock: + if self._cancel_event.is_set(): + logger.info("Respawn skipped: the model was unloaded.") + return False + kwargs = dict(self._last_load_kwargs or {}) + if not kwargs: + return False + epoch = self._unload_epoch + self._healthy = False + logger.warning( + f"llama-server for '{self._model_identifier}' exited " + f"(code {proc.returncode}); respawning to recover the session" + ) + try: + started = bool(self.load_model(**kwargs)) + except Exception as exc: + logger.error(f"Failed to respawn llama-server: {exc}") + return False + if started and self._unload_epoch != epoch: + # An unload landed mid-reload. load_model cleared _cancel_event on + # the way in, so the epoch is the only surviving evidence; undo the + # replacement rather than leave a model the user stopped running. + logger.info("Respawn undone: the model was unloaded during the reload.") + self.unload_model() + return False + return started + + @contextlib.contextmanager + def _open_chat_stream_with_respawn_retry(self, payload: dict, cancel_event): + """Open a chat stream, respawning a dead llama-server once before streaming. + + Retry only when opening the response fails: once it is open a consumer may + already have emitted content or tool events, so a replay could duplicate + output and side effects. ``base_url`` is resolved per attempt because a + respawn may pick a new port. The budget is one retry per model request, not + per chat turn, so a long tool loop never discards a completed tool. + + A child dying after the accept but before the headers surfaces as + ReadError/WriteError/RemoteProtocolError rather than ConnectError, and which + one differs per OS. llama-server flushes its 200 at slot start, so that window + is an upload still in flight or a request behind busy slots; a death during + decode arrives with the response open and is not replayed. Timeouts are + excluded: the server is slow, not dead, and a replay would spend the + first-token budget twice. + """ + for attempt in range(2): + response_opened = False try: - return bool(self.load_model(**kwargs)) - except Exception as exc: - logger.error(f"Failed to respawn llama-server: {exc}") - return False + url = f"{self.base_url}/v1/chat/completions" + with self._open_stream(url, payload, cancel_event) as opened: + response_opened = True + yield opened + return + except (httpx.NetworkError, httpx.RemoteProtocolError) as exc: + if response_opened: + raise + if self._maybe_recover_from_mtp_crash(exc): + raise RuntimeError("Lost connection to llama-server") from exc + if attempt == 0 and self._respawn_if_dead(): + logger.warning( + "llama-server was unreachable; respawned it and retrying the generation" + ) + continue + raise def generate_chat_completion( self, @@ -8834,6 +11316,7 @@ class LlamaCppBackend: reasoning_effort: Optional[str] = None, preserve_thinking: Optional[bool] = None, seed: Optional[int] = None, + promote_reasoning_only: bool = True, _allow_respawn_retry: bool = True, ) -> Generator[Union[str, dict], None, None]: """ @@ -8916,7 +11399,12 @@ class LlamaCppBackend: # model put its whole reply in reasoning # (e.g. Qwen3 always-think). Show it as # the main response, not a thinking block. - cumulative = reasoning_text + cumulative = _finalize_reasoning_only_cumulative( + cumulative, + reasoning_text, + _metadata_finish_reason, + promote_reasoning_only, + ) yield cumulative _stream_done = True break # exit inner while @@ -9013,6 +11501,7 @@ class LlamaCppBackend: reasoning_effort = reasoning_effort, preserve_thinking = preserve_thinking, seed = seed, + promote_reasoning_only = promote_reasoning_only, _allow_respawn_retry = False, ) return @@ -9054,6 +11543,7 @@ class LlamaCppBackend: confirm_tool_calls: bool = False, bypass_permissions: bool = False, permission_mode: Optional[str] = None, + promote_reasoning_only: bool = True, ) -> Generator[dict, None, None]: """ Agentic loop: let the model call tools, execute them, and continue. @@ -9072,17 +11562,22 @@ class LlamaCppBackend: from core.inference.tools import ( build_rag_autoinject, execute_tool, + has_text_only_provisional_card, is_always_safe_tool, - is_potentially_unsafe_tool_call, + is_high_risk_tool_call, ) - # Normalize the mode: "full" and bypass_permissions are the same - # switch, whichever arrives first wins toward the permissive side. - # "off" keeps the sandbox but never prompts. + # "full" and bypass_permissions are the same switch, whichever arrives + # first wins. "off" keeps the sandbox but never prompts. Unset defaults to + # "auto"; unknown falls back to the stricter "ask". An explicit + # confirm_tool_calls=True with no mode is already resolved to "ask" at the + # request layer, so it never arrives here as an ambiguous unset. if permission_mode == "full": bypass_permissions = True elif bypass_permissions: permission_mode = "full" + elif permission_mode is None: + permission_mode = "auto" elif permission_mode not in ("ask", "auto", "off"): permission_mode = "ask" @@ -9105,7 +11600,6 @@ class LlamaCppBackend: yield _ev conversation.extend(_auto["messages"]) - url = f"{self.base_url}/v1/chat/completions" _accumulated_completion_tokens = 0 _accumulated_predicted_ms = 0.0 _accumulated_predicted_n = 0 @@ -9266,6 +11760,10 @@ class LlamaCppBackend: # direct answer ("4", "Hello!") won't match. Pattern shared with the # safetensors loop (tool_call_parser.INTENT_SIGNAL). _reprompt_count = 0 + # Budgeted apart from _reprompt_count so a pre-tool nudge can't spend it. + _post_tool_reprompts = 0 + # Text that triggered the last nudge; if the retry restates it, stop. + _last_reprompt_text = "" # Gates ``max_tool_iterations`` on real tool turns (not the enlarged range) so reserved # re-prompt slots don't extend the budget. Mirrors the safetensors guard. _tool_iters_done = 0 @@ -9273,7 +11771,7 @@ class LlamaCppBackend: # Reserve extra iterations for re-prompts so they don't consume the # caller's tool-call budget; only when tool iterations are allowed. - _extra = _MAX_REPROMPTS if max_tool_iterations > 0 else 0 + _extra = _MAX_REPROMPTS + 1 if max_tool_iterations > 0 else 0 for iteration in range(max_tool_iterations + _extra): if cancel_event is not None and cancel_event.is_set(): return @@ -9339,6 +11837,7 @@ class LlamaCppBackend: # Time each reasoning pass so final answers can replace tool timing. _reasoning_started_at = None _reasoning_summary_emitted = False + _deferred_reasoning_summary = None cumulative_display = "" # Cumulative yielded text (with ) in_thinking = False has_content_tokens = False @@ -9365,7 +11864,7 @@ class LlamaCppBackend: _text_args_name = "" _confirm_gated_iteration = bool(confirm_tool_calls) and not bypass_permissions - with self._open_stream(url, payload, cancel_event) as ( + with self._open_chat_stream_with_respawn_retry(payload, cancel_event) as ( response, first_token_deadline, ): @@ -9396,7 +11895,12 @@ class LlamaCppBackend: ), } else: - cumulative_display = reasoning_accum + cumulative_display = _finalize_reasoning_only_cumulative( + cumulative_display, + reasoning_accum, + _iter_finish_reason, + promote_reasoning_only, + ) if not _suppress_visible_output: yield { "type": "content", @@ -9490,6 +11994,9 @@ class LlamaCppBackend: permission_mode == "auto" and is_always_safe_tool(current_name) ) + # A text-preview card still streams while gated; + # hiding it blanks the chat. + and not has_text_only_provisional_card(current_name) ) # Keep small-argument tools on the normal path. _args_len = len( @@ -9582,7 +12089,11 @@ class LlamaCppBackend: and not _reasoning_summary_emitted ): _reasoning_summary_emitted = True - yield _reasoning_summary_event(_reasoning_started_at) + _summary = _reasoning_summary_event(_reasoning_started_at) + if _suppress_visible_output: + _deferred_reasoning_summary = _summary + else: + yield _summary has_content_tokens = True content_accum += token @@ -9591,20 +12102,27 @@ class LlamaCppBackend: # TEXT call to a provisional card. Gated on an enabled-name # sniff + size floor so prose/small calls spawn no pane; id # matches the first call so the final tool_start reconciles. - if ( - not has_structured_tc - and not _confirm_gated_iteration - and _text_args_call_start >= 0 - ): + if not has_structured_tc and _text_args_call_start >= 0: if not _text_args_id: _call_text = content_accum[_text_args_call_start:] _sniffed = _sniff_text_tool_name( _call_text, _enabled_tool_names ) - if _sniffed and ( - _sniffed == "render_html" - or len(_call_text) - >= _PROVISIONAL_ARGS_MIN_CHARS + # Structured-path rule: gated calls + # stream only from a text-preview card. + if ( + _sniffed + and not ( + _confirm_gated_iteration + and not has_text_only_provisional_card( + _sniffed + ) + ) + and ( + _sniffed == "render_html" + or len(_call_text) + >= _PROVISIONAL_ARGS_MIN_CHARS + ) ): _text_args_id = "call_0" _text_args_name = _sniffed @@ -9859,8 +12377,17 @@ class LlamaCppBackend: # route's extractor closes the streamed ). if _reasoning_started_at is not None and not _reasoning_summary_emitted: _reasoning_summary_emitted = True - yield _reasoning_summary_event(_reasoning_started_at) - cumulative_display = reasoning_accum + _summary = _reasoning_summary_event(_reasoning_started_at) + if _suppress_visible_output: + _deferred_reasoning_summary = _summary + else: + yield _summary + cumulative_display = _finalize_reasoning_only_cumulative( + cumulative_display, + reasoning_accum, + _iter_finish_reason, + promote_reasoning_only, + ) if not _suppress_visible_output: yield { "type": "content", @@ -9889,12 +12416,10 @@ class LlamaCppBackend: ) if not _safety_tc: # ── Re-prompt on plan-without-action ── - # If the model described its intent (forward-looking - # language) without calling a tool, nudge it to act. - # Fires at most once per request, only on short - # responses with intent signals -- "4" or "Hello!" - # won't trigger it. Use content if available, else - # fall back to reasoning text (reasoning-only stalls). + # Intent described without a tool call: nudge it to act. Up + # to _MAX_REPROMPTS times, only on short responses with intent + # signals -- "4" or "Hello!" won't trigger it. Uses content, + # else reasoning text (reasoning-only stalls). _stripped = content_accum.strip() if not _stripped: _stripped = reasoning_accum.strip() @@ -9904,18 +12429,33 @@ class LlamaCppBackend: r"(?i)\brender[_\s-]?html\b", _stripped, ) + # A post-tool stall still deserves a nudge, but each retry + # re-runs tools, so allow only one. RAG autoinject never lands + # in history, so _auto keeps a doc-grounded turn from reading + # as pre-tool (mirrors safetensors rag_autoinjected). + _already_acted = bool(_auto) or any( + record.executed for record in tool_controller.history + ) + if _already_acted: + _reprompt_used, _reprompt_cap = _post_tool_reprompts, 1 + else: + _reprompt_used, _reprompt_cap = _reprompt_count, _MAX_REPROMPTS # None keeps the default-on re-prompt; False disables it. if ( auto_heal_tool_calls and (nudge_tool_calls is None or nudge_tool_calls) and active_tools and not _render_html_already_done_intent - and _reprompt_count < _MAX_REPROMPTS + and _reprompt_used < _reprompt_cap + and not _is_reprompt_repeat(_stripped, _last_reprompt_text) and _is_short_intent_without_action(_stripped) ): _reprompt_count += 1 + if _already_acted: + _post_tool_reprompts += 1 + _last_reprompt_text = _stripped logger.info( - f"Re-prompt {_reprompt_count}/{_MAX_REPROMPTS}: " + f"Re-prompt {_reprompt_used + 1}/{_reprompt_cap}: " f"model responded without calling tools " f"({len(_stripped)} chars)" ) @@ -9945,12 +12485,18 @@ class LlamaCppBackend: _it_r = _iter_timings or {} _accumulated_predicted_ms += _it_r.get("predicted_ms", 0) _accumulated_predicted_n += _it_r.get("predicted_n", 0) + # Blank first (the route resets its text cursor only on an + # empty status), then the badge so the retry is not a hang. yield {"type": "status", "text": ""} + yield {"type": "status", "text": _NUDGE_TOOL_CALLS_STATUS} continue if _forced_tool_call_pending: _forced_tool_call_pending = False - if not _should_suppress_forced_no_tool_output(_stripped): + if not _should_suppress_forced_no_tool_output( + _stripped, + _last_reprompt_text, + ): if cumulative_display: forced_visible_text = _strip_tool_markup( cumulative_display, @@ -9968,6 +12514,8 @@ class LlamaCppBackend: "type": "content", "text": forced_visible_text, } + if _deferred_reasoning_summary is not None: + yield _deferred_reasoning_summary elif not _suppress_visible_output: # Turn ended as a plain answer (no [ARGS] followed): the held # rehearsal tail is real prose, release it. @@ -10104,6 +12652,9 @@ class LlamaCppBackend: assistant_msg: dict = {"role": "assistant", "content": content_text} assistant_appended = False + # Collect no-op nudges and flush them after the batch, so a no-op + # doesn't abort it and drop the parallel calls that follow. + deferred_noop_msgs: list = [] # The text-path provisional card uses the parser's default id ("call_0"); # a Mistral-style call carries its own id and would open a duplicate. Reuse @@ -10146,14 +12697,14 @@ class LlamaCppBackend: "provenance": decision.provenance, } completion = tool_controller.record_noop(decision) - conversation.append(completion.model_message()) + deferred_noop_msgs.append(completion.model_message()) if _forced_tool_call_pending: _forced_tool_call_pending = False logger.info( "Suppressed local GGUF tool call as internal no-op: " f"action={decision.action} tool={decision.tool_name}" ) - break + continue if not assistant_appended: assistant_msg["tool_calls"] = [decision.as_assistant_tool_call()] @@ -10164,18 +12715,16 @@ class LlamaCppBackend: decision.as_assistant_tool_call() ) - # Bypass wins over the confirm gate at the loop level too, - # so a direct internal caller with both flags never prompts. - # In "auto" mode only calls detected as potentially unsafe - # pause; read-only calls run straight through. "off" never - # prompts (sandbox stays on). + # Bypass wins here too, so a direct internal caller with both + # flags never prompts. "auto" pauses only high-risk calls; + # "off" never prompts (sandbox stays on). needs_confirm = ( bool(confirm_tool_calls) and not bypass_permissions and permission_mode != "off" ) if needs_confirm and permission_mode == "auto": - needs_confirm = is_potentially_unsafe_tool_call( + needs_confirm = is_high_risk_tool_call( decision.tool_name, decision.arguments ) approval_id = new_approval_id() if needs_confirm else "" @@ -10187,18 +12736,31 @@ class LlamaCppBackend: start_event["awaiting_confirmation"] = needs_confirm try: - yield {"type": "status", "text": decision.status_text} + # Gated calls are not running yet; a "Running ..." badge + # counting up while it waits on a human reads as a hang. + yield { + "type": "status", + "text": ( + awaiting_approval_status(decision.tool_name) + if needs_confirm + else decision.status_text + ), + } yield start_event - if ( - decision_slot is not None - and wait_tool_decision( + _decision = ( + wait_tool_decision( decision_slot, approval_id, cancel_event = cancel_event, ) - == "deny" - ): + if decision_slot is not None + else None + ) + if _decision is not None and _decision != "deny": + # Approved: now it really is running. + yield {"type": "status", "text": decision.status_text} + if _decision == "deny": decision_slot = None resolved_provisional_tool_call_ids.add(decision.tool_call_id) yield { @@ -10264,6 +12826,10 @@ class LlamaCppBackend: _kb_search_count += 1 completion = tool_controller.record_result(decision, result) resolved_provisional_tool_call_ids.add(decision.tool_call_id) + # A real execution opens the post-tool phase; carrying the pre-tool + # stall text over would read the same sentence as a repeat and + # swallow the one post-tool nudge. + _last_reprompt_text = "" # A tool ran this turn, so it counts against the caller's budget. _turn_executed_real_tool = True yield completion.tool_end_event() @@ -10272,6 +12838,8 @@ class LlamaCppBackend: if _forced_tool_call_pending: _forced_tool_call_pending = False + append_deferred_nudges(conversation, deferred_noop_msgs) + # Close provisional cards not resolved by execution/no-op handling. for _pid, _pname in provisional_started_tool_calls.items(): if _pid not in resolved_provisional_tool_call_ids: @@ -10387,7 +12955,7 @@ class LlamaCppBackend: _stream_done = False try: - with self._open_stream(url, stream_payload, cancel_event) as ( + with self._open_chat_stream_with_respawn_retry(stream_payload, cancel_event) as ( response, first_token_deadline, ): @@ -10419,7 +12987,12 @@ class LlamaCppBackend: "text": _strip_tool_markup(cumulative, final = True), } else: - cumulative = reasoning_text + cumulative = _finalize_reasoning_only_cumulative( + cumulative, + reasoning_text, + _metadata_finish_reason, + promote_reasoning_only, + ) yield {"type": "content", "text": cumulative} _stream_done = True break # exit inner while @@ -10759,10 +13332,15 @@ class LlamaCppBackend: min_p: float = 0.0, max_new_tokens: int = 2048, repetition_penalty: float = 1.1, + cancel_event: Optional[threading.Event] = None, ) -> tuple: """ Generate TTS audio via llama-server /completion + codec decode. Returns (wav_bytes, sample_rate). + + ``cancel_event`` lets a Stop or a forced model swap end the request: the + decode is one blocking POST, so a watcher closes the client out from under + it rather than polling. Raises RuntimeError once cancelled. """ if audio_type not in self._TTS_PROMPTS: raise RuntimeError(f"GGUF TTS does not support '{audio_type}' codec.") @@ -10784,15 +13362,47 @@ class LlamaCppBackend: if need_ids: payload["n_probs"] = 1 + if cancel_event is not None and cancel_event.is_set(): + raise RuntimeError("Audio generation cancelled") + with httpx.Client( timeout = httpx.Timeout(300, connect = 10), headers = self._auth_headers, trust_env = False, ) as client: - resp = client.post(f"{self.base_url}/completion", json = payload) + finished = threading.Event() + watcher: Optional[threading.Thread] = None + if cancel_event is not None: + + def _close_when_cancelled() -> None: + while not finished.wait(0.05): + if cancel_event.is_set(): + # Closing mid-request makes the blocking post raise + # httpx.RequestError, the only way out of it. + with contextlib.suppress(Exception): + client.close() + return + + watcher = threading.Thread(target = _close_when_cancelled, daemon = True) + watcher.start() + try: + resp = client.post(f"{self.base_url}/completion", json = payload) + except httpx.RequestError: + if cancel_event is not None and cancel_event.is_set(): + raise RuntimeError("Audio generation cancelled") from None + raise + finally: + finished.set() + if watcher is not None: + watcher.join(timeout = 0.5) if resp.status_code != 200: raise RuntimeError(f"llama-server returned {resp.status_code}: {resp.text}") + # The codec decode below is GPU work with no interruption point, so check here: + # cancelling after this only wastes the decode it cannot stop. + if cancel_event is not None and cancel_event.is_set(): + raise RuntimeError("Audio generation cancelled") + data = resp.json() token_ids = ( [p["id"] for p in data.get("completion_probabilities", []) if "id" in p] diff --git a/studio/backend/core/inference/llama_keepwarm.py b/studio/backend/core/inference/llama_keepwarm.py index 4ce663c3ce..05b1271b27 100644 --- a/studio/backend/core/inference/llama_keepwarm.py +++ b/studio/backend/core/inference/llama_keepwarm.py @@ -15,6 +15,7 @@ import asyncio import contextlib import threading import time +from pathlib import Path from loggers import get_logger @@ -30,6 +31,8 @@ _last_active = time.monotonic() # otherwise 503 against an empty backend can reload it (set on unload, cleared on # reload). Storing the quant means the reload restores the exact freed variant. _last_unloaded_model = None +# Slot KV manifest saved by the idle unload; whoever pops it owns deleting its files. +_kv_resume = None # Guards inflight bumps against the idle-check-then-unload race, and blocks new # inference from starting mid-swap. Process-wide, not per-loop: the backend slot is # shared across every event loop in the process, so a per-loop gate would let a @@ -59,7 +62,7 @@ _INFERENCE_SUFFIXES = ( "/messages/count_tokens", # counts via the loaded tokenizer; protect like /messages "/embeddings", "/responses", - "/generate/stream", # Studio's own streaming route on the same llama-server + "/generate/stream", # Unsloth's own streaming route on the same llama-server "/audio/generate", # direct GGUF TTS; can outlive the idle TTL ) @@ -161,11 +164,17 @@ def inference_lifecycle_gate(): return _unload_gate() -def note_model_loaded() -> None: - """Record a successful GGUF load: stamp activity and drop any reload stash so - a manual load clears it synchronously, not only on the next idle poll.""" +def note_model_loaded(backend = None) -> None: + """Stamp activity and synchronously drop any reload stash.""" _note_activity() + resume = take_kv_resume() _set_last_unloaded(None) + if resume is None: + return + if backend is not None: + restore_kv_resume(backend, resume) + else: + _delete_resume_files(resume) def note_model_unloaded() -> None: @@ -182,9 +191,81 @@ def get_last_unloaded_model(): def _set_last_unloaded(value) -> None: - global _last_unloaded_model + global _last_unloaded_model, _kv_resume + stale = None with _lock: _last_unloaded_model = value + if value is None and _kv_resume is not None: + stale, _kv_resume = _kv_resume, None + if stale: + _delete_resume_files(stale) + + +def _delete_resume_files(manifest) -> None: + try: + base = Path(manifest.get("dir") or "") + for entry in manifest.get("slots") or []: + with contextlib.suppress(OSError): + (base / str(entry.get("filename"))).unlink() + except Exception: + pass + + +def _set_kv_resume(value) -> None: + global _kv_resume + stale = None + with _lock: + if _kv_resume is not None and _kv_resume is not value: + stale = _kv_resume + _kv_resume = value + if stale: + _delete_resume_files(stale) + + +def take_kv_resume(): + global _kv_resume + with _lock: + manifest, _kv_resume = _kv_resume, None + return manifest + + +def purge_kv_resume() -> None: + resume = take_kv_resume() + if resume: + _delete_resume_files(resume) + + +def restore_kv_resume(backend, manifest) -> None: + try: + gguf = manifest.get("gguf") + binary = manifest.get("binary") + current = getattr(backend, "_gguf_path", None) + same_gguf = bool(gguf and current) and Path(current).resolve() == Path(gguf).resolve() + if same_gguf: + # Same path is not enough: shards may have been rewritten meanwhile. + identity = getattr(backend, "_gguf_file_identity", None) + same_gguf = callable(identity) and identity(current) == manifest.get("gguf_stat") + if same_gguf: + # Nor the same file: launch overrides can invalidate KV numerics. + fingerprint = getattr(backend, "_slot_launch_fingerprint", None) + same_gguf = callable(fingerprint) and manifest.get("launch") == fingerprint() + if same_gguf and binary and binary == getattr(backend, "_slot_save_binary", None): + logger.info("Restoring saved slot KV onto the reloaded model") + backend.restore_slots_for_resume(manifest) + except Exception as exc: + logger.debug("slot restore after reload failed: %s", exc) + finally: + _delete_resume_files(manifest) + + +def sweep_slot_save_dir() -> None: + try: + from utils.paths.storage_roots import llama_slot_cache_root + for path in llama_slot_cache_root().glob("resume-*.bin"): + with contextlib.suppress(OSError): + path.unlink() + except Exception: + pass class LlamaKeepWarmMiddleware: @@ -264,9 +345,28 @@ def _loaded_identity(backend): return (backend.model_identifier, getattr(backend, "hf_variant", None), advertised) +def _note_idle_unload_event(freed) -> None: + """Monitor row for an idle auto-unload. Best-effort; uses the stash's + advertised repo id so the row never shows the on-disk load path.""" + try: + from core.inference.api_monitor import api_monitor + from core.inference.model_ids import public_model_id + + identifier, variant, advertised = (list(freed) + [None, None, None])[:3] + label = public_model_id(advertised or identifier) or "model" + if variant and ":" not in label: + label = f"{label}:{variant}" + api_monitor.record_lifecycle(event = "unload", model = label, reason = "idle") + except Exception as exc: + logger.debug("idle unload monitor event failed: %s", exc) + + async def idle_unload_loop(poll_seconds: float = 15.0) -> None: """Unload the loaded GGUF once idle past the configured TTL. Inert when off.""" - from utils.openai_auto_switch_settings import get_auto_unload_idle_seconds + from utils.openai_auto_switch_settings import ( + get_auto_unload_idle_seconds, + get_auto_unload_keep_kv, + ) seen_model = None while True: @@ -281,18 +381,50 @@ async def idle_unload_loop(poll_seconds: float = 15.0) -> None: # Track by (id, variant): a (re)loaded model -- including the same repo # at a different quant -- counts as activity so it survives one TTL # before its first request (loads bypass the activity middleware). - current = _loaded_identity(backend) - if current != seen_model: - seen_model = current - if current is not None: - _note_activity() - _set_last_unloaded(None) # a model is loaded; drop stale stash async with _unload_gate(): + # Purging the stash mid-reload would race the restore. + current = _loaded_identity(backend) + if current != seen_model: + seen_model = current + if current is not None: + _note_activity() + _set_last_unloaded(None) # a model is loaded; drop stale stash if backend.is_loaded and _is_idle(ttl): freed = _loaded_identity(backend) - await asyncio.to_thread(backend.unload_model) + manifest = None + if get_auto_unload_keep_kv(): + try: + manifest = await asyncio.to_thread( + backend.save_slots_for_resume, + lambda: not _is_idle(ttl), + ) + except Exception as exc: + logger.debug("slot save before idle unload failed: %s", exc) + # Re-read settings: the save can outlive a settings change. + ttl = get_auto_unload_idle_seconds() + if ttl <= 0 or not _is_idle(ttl): + if manifest: + _delete_resume_files(manifest) + continue + if manifest and not get_auto_unload_keep_kv(): + _delete_resume_files(manifest) + manifest = None + try: + await asyncio.to_thread(backend.unload_model) + except Exception: + # Failed unload means nothing will stash the manifest. + if manifest: + _delete_resume_files(manifest) + raise _set_last_unloaded(freed) # let an alias request reload it + if manifest and freed: + _set_kv_resume({"identity": freed, **manifest}) + logger.info("Idle auto-unload: saved slot KV for restore on reload") + elif manifest: + _delete_resume_files(manifest) logger.info("Idle auto-unload: freed GGUF after %ss idle", ttl) + # An idle unload stashes for reload and skips note_model_unloaded. + _note_idle_unload_event(freed) seen_model = None except Exception as exc: logger.debug("idle_unload_loop iteration failed: %s", exc) diff --git a/studio/backend/core/inference/llama_server_args.py b/studio/backend/core/inference/llama_server_args.py index f400d2ae40..7391e62516 100644 --- a/studio/backend/core/inference/llama_server_args.py +++ b/studio/backend/core/inference/llama_server_args.py @@ -3,10 +3,10 @@ """Boundary validator for user-supplied llama-server pass-through args. -Reject only flags Studio manages (model identity, auth, network, parallel +Reject only flags Unsloth manages (model identity, auth, network, parallel slots). Everything else (sampling, ``-c``, ``-ngl``, ``--flash-attn``, ``--cache-type-*``, ``--spec-*``, ``--jinja``, ...) is appended after -Studio's auto-set flags so llama.cpp's last-wins parser lets the user override. +Unsloth's auto-set flags so llama.cpp's last-wins parser lets the user override. Ref: https://github.com/ggml-org/llama.cpp/blob/master/tools/server/README.md """ @@ -16,18 +16,25 @@ from __future__ import annotations import os from typing import Iterable, Mapping, Optional +# Valid llama-server --parallel range, shared with LoadRequest.n_parallel. +# Mirrored by callers that cannot import this: run.py and unsloth_cli/commands/ +# studio.py (_PARALLEL_MIN/MAX), per-model-config.ts (N_PARALLEL_MIN/MAX); +# test_parallel_slots_per_load.py pins them together. +PARALLEL_MIN = 1 +PARALLEL_MAX = 64 + # Each group = every alias (short + long) of one hard-denied flag. # Extend the matching group when llama.cpp adds a new alias. _DENYLIST_GROUPS: tuple[frozenset[str], ...] = ( - # Parallel slots: owned by typer --parallel; a pass-through would desync - # app.state.llama_parallel_slots from llama-server. + # Parallel slots: owned by typer --parallel and LoadRequest.n_parallel; a + # pass-through would desync the slot bookkeeping from llama-server. frozenset({"-np", "--parallel", "--n-parallel"}), - # Model identity: Studio resolves it from LoadRequest; a second -m would - # load a different model than Studio thinks it loaded. + # Model identity: Unsloth resolves it from LoadRequest; a second -m would + # load a different model than Unsloth thinks it loaded. frozenset({"-m", "--model"}), - # Public model id: Studio sets a sanitized --alias so the OpenAI API never + # Public model id: Unsloth sets a sanitized --alias so the OpenAI API never # exposes the local .gguf path. A user-supplied alias is appended after - # Studio's and, with llama.cpp's last-wins parsing, would reintroduce the + # Unsloth's and, with llama.cpp's last-wins parsing, would reintroduce the # path leak this is meant to prevent. frozenset({"-a", "--alias"}), frozenset({"-mu", "--model-url"}), @@ -39,14 +46,14 @@ _DENYLIST_GROUPS: tuple[frozenset[str], ...] = ( frozenset({"-hft", "--hf-token"}), frozenset({"-mm", "--mmproj"}), frozenset({"-mmu", "--mmproj-url"}), - # Networking: Studio binds + proxies; retargeting orphans the proxy. + # Networking: Unsloth binds + proxies; retargeting orphans the proxy. frozenset({"--host"}), frozenset({"--port"}), frozenset({"--path"}), frozenset({"--api-prefix"}), frozenset({"--reuse-port"}), - # Auth / TLS: Studio terminates auth; upstream --api-key / TLS shadows - # Studio's key and breaks the proxy hop. + # Auth / TLS: Unsloth terminates auth; upstream --api-key / TLS shadows + # Unsloth's key and breaks the proxy hop. frozenset({"--api-key"}), frozenset({"--api-key-file"}), frozenset({"--ssl-key-file"}), @@ -64,12 +71,14 @@ _DENYLIST_GROUPS: tuple[frozenset[str], ...] = ( frozenset({"--models-max"}), frozenset({"--models-autoload", "--no-models-autoload"}), # Server-mode flips: --embedding / --rerank restrict llama-server to - # those endpoints, breaking Studio's /v1/chat/completions hop. + # those endpoints, breaking Unsloth's /v1/chat/completions hop. frozenset({"--embedding", "--embeddings"}), frozenset({"--rerank", "--reranking"}), # llama-server's own built-in tools flag would silently stack on top of - # Studio's --enable-tools / --disable-tools policy resolver. + # Unsloth's --enable-tools / --disable-tools policy resolver. frozenset({"--tools"}), + # Slot-state dir: Studio owns it for KV persistence across idle unload. + frozenset({"--slot-save-path"}), ) _DENYLIST: frozenset[str] = frozenset().union(*_DENYLIST_GROUPS) @@ -78,9 +87,10 @@ _DENYLIST: frozenset[str] = frozenset().union(*_DENYLIST_GROUPS) def _flag_name(token: str) -> Optional[str]: """Flag name for ``token``, or None if it isn't a flag. - Peels `--key=value` to `--key`, treats `-1`/`-0.5` as values (shorts - always start with a letter), and normalises attached `-np8` / `-np-1` / - `-np8x` to `-np`. Mirrors the CLI's `_expand_attached_np_short`. + Peels `--key=value` to `--key`, normalises long-option underscores like + llama.cpp, treats `-1`/`-0.5` as values (shorts always start with a letter), + and normalises attached `-np8` / `-np-1` / `-np8x` to `-np`. Mirrors the + CLI's `_expand_attached_np_short`. """ token = token.strip() if not token.startswith("-") or token in {"-", "--"}: @@ -88,6 +98,8 @@ def _flag_name(token: str) -> Optional[str]: if len(token) >= 2 and (token[1].isdigit() or token[1] == "."): return None name = token.split("=", 1)[0] + if name.startswith("--"): + name = name.replace("_", "-") if len(name) > 3 and name.startswith("-np"): suffix = name[3:] if suffix[0].isdigit() or ( @@ -116,11 +128,12 @@ def validate_extra_args(args: Optional[Iterable[str]]) -> list[str]: parse_ctx_override(out) parse_cache_override(out) parse_split_mode_override(out) + parse_gpu_layers_override(out) return out def is_managed_flag(flag: str) -> bool: - """True if ``flag`` is Studio-managed. Normalises via ``_flag_name`` so + """True if ``flag`` is Unsloth-managed. Normalises via ``_flag_name`` so `-np8` / `--parallel=8` classify like the canonical tokens.""" normalised = _flag_name(flag) return normalised is not None and normalised in _DENYLIST @@ -142,7 +155,7 @@ _SPEC_FLAGS: frozenset[str] = frozenset( "--draft-min", "--draft-max", # MTP path (llama.cpp #22673). The drafter selectors (local --model-draft - # and HF --spec-draft-hf aliases) are Studio-managed since the separate- + # and HF --spec-draft-hf aliases) are Unsloth-managed since the separate- # drafter support (Gemma 4): an inherited copy must not last-wins-override # the auto-detected drafter. Explicit extras for the current load are never # stripped. The per-drafter tuning knobs (--spec-draft-type-*, -ngld, @@ -179,25 +192,37 @@ _TEMPLATE_FLAGS: frozenset[str] = frozenset( # (--split-mode tensor). Pass-through stays allowed so users keep the # row/none/layer modes the toggle doesn't expose, but it's stripped on # inherit and reconciled into the round-tripped tensor_parallel state. -# --tensor-split is coupled to the split mode and is stripped with it: Studio +# --tensor-split is coupled to the split mode and is stripped with it: Unsloth # owns the tensor-mode split ratios, so an inherited/stale --tensor-split must -# not last-wins-override Studio's computed asymmetric split. +# not last-wins-override Unsloth's computed asymmetric split. _SPLIT_MODE_FLAGS: frozenset[str] = frozenset({"-sm", "--split-mode"}) _TENSOR_SPLIT_FLAGS: frozenset[str] = frozenset({"-ts", "--tensor-split"}) _SPLIT_SHADOWING_FLAGS: frozenset[str] = _SPLIT_MODE_FLAGS | _TENSOR_SPLIT_FLAGS +# GPU-offload flags. Stripped only when the GPU Memory mode owns offload +# (manual emits --fit / --gpu-layers / --n-cpu-moe); in auto, a user's +# inherited -ngl is respected (the offload_overridden path), so this group is +# opt-in, not default. Layer flags are shared with llama_cpp's override +# detection; the MoE flags are strip-only (manual's --n-cpu-moe slider owns them). +_GPU_LAYER_FLAGS: frozenset[str] = frozenset({"-ngl", "--gpu-layers", "--n-gpu-layers"}) +_LAYER_OFFLOAD_FLAGS: frozenset[str] = _GPU_LAYER_FLAGS | frozenset({"-fit", "--fit"}) +_MOE_OFFLOAD_FLAGS: frozenset[str] = frozenset({"-ncmoe", "--n-cpu-moe", "-cmoe", "--cpu-moe"}) +_OFFLOAD_SHADOWING_FLAGS: frozenset[str] = _LAYER_OFFLOAD_FLAGS | _MOE_OFFLOAD_FLAGS + _SHADOWING_FLAGS: frozenset[str] = ( _CONTEXT_FLAGS | _CACHE_FLAGS | _SPEC_FLAGS | _TEMPLATE_FLAGS | _SPLIT_SHADOWING_FLAGS ) # Shadowing flags that take no value -- strip the flag only, not the next token. -_BOOLEAN_SHADOWING_FLAGS: frozenset[str] = frozenset({"--spec-default", "--jinja", "--no-jinja"}) +_BOOLEAN_SHADOWING_FLAGS: frozenset[str] = frozenset( + {"--spec-default", "--jinja", "--no-jinja", "-cmoe", "--cpu-moe"} +) def parse_ctx_override(args: Optional[Iterable[str]]) -> Optional[int]: """Return the last user-supplied ``-c`` / ``--ctx-size`` value. - Mirrors llama.cpp's last-wins parsing for the one numeric knob Studio's + Mirrors llama.cpp's last-wins parsing for the one numeric knob Unsloth's load-time fit logic needs. """ if not args: @@ -286,11 +311,31 @@ def parse_cache_override(args: Optional[Iterable[str]]) -> Optional[str]: Mirrors parse_ctx_override but for cache type. Recognises both -ctk (key) and -ctv (value). When both flags appear, returns the last-wins value, treating key and value cache flags as the same setting because - Studio's KV estimate has a single cache_type_kv knob. + Unsloth's KV estimate has a single cache_type_kv knob. """ return _last_flag_value(args, _CACHE_FLAGS) +def parse_gpu_layers_override(args: Optional[Iterable[str]]) -> Optional[int]: + """Return the last user-supplied GPU layer count from extras. + + Manual GPU memory mode strips llama.cpp offload flags because the + first-class load fields own them. Callers use this parser first to preserve + an explicit ``-ngl`` / ``--gpu-layers`` / ``--n-gpu-layers`` value when + translating the extras into those fields. + """ + raw_value = _last_flag_value(args, _GPU_LAYER_FLAGS) + if raw_value is None: + return None + try: + value = int(raw_value) + except ValueError as exc: + raise ValueError("llama-server GPU layers flag requires an integer value") from exc + if value < -1: + raise ValueError("llama-server GPU layers flag requires an integer value of at least -1") + return value + + def parse_cache_override_per_axis( args: Optional[Iterable[str]], ) -> tuple[Optional[str], Optional[str]]: @@ -341,7 +386,7 @@ def resolve_tensor_parallel(args: Optional[Iterable[str]], fallback_tensor_paral def _env_split_mode_is_tensor(env: Optional[Mapping[str, str]] = None) -> bool: - """True when the inherited LLAMA_ARG_SPLIT_MODE env selects tensor. Studio + """True when the inherited LLAMA_ARG_SPLIT_MODE env selects tensor. Unsloth emits --split-mode only on its tensor branch, so a tensor env on the layer path would run the child tensor-parallel unbudgeted; this flips the budget to tensor. Only tensor is heavier, so other modes are ignored.""" @@ -424,14 +469,22 @@ def strip_shadowing_flags( strip_spec: bool = True, strip_template: bool = True, strip_split_mode: bool = True, + strip_tensor_split: bool = False, + strip_offload: bool = False, ) -> list[str]: - """Strip flags that shadow first-class Studio settings. + """Strip flags that shadow first-class Unsloth settings. Used when inheriting a previous load's ``llama_extra_args`` so an inherited `-c 4096` can't override the current `max_seq_length` (same for cache / spec / template / split-mode). Each ``strip_*`` toggle controls one group; the route only strips groups whose first-class field the caller actually supplied. + + ``strip_split_mode`` removes both ``--split-mode`` and the coupled + ``--tensor-split`` (the Tensor Parallelism toggle owns the whole split). + ``strip_tensor_split`` removes ``--tensor-split`` *alone*, so manual mode can + replace an inherited per-GPU ratio while leaving the user's ``--split-mode`` + row/none/layer choice intact. """ shadowing: set[str] = set() if strip_context: @@ -444,6 +497,10 @@ def strip_shadowing_flags( shadowing |= _TEMPLATE_FLAGS if strip_split_mode: shadowing |= _SPLIT_SHADOWING_FLAGS + if strip_tensor_split: + shadowing |= _TENSOR_SPLIT_FLAGS + if strip_offload: + shadowing |= _OFFLOAD_SHADOWING_FLAGS tokens = [str(a) for a in (args or [])] out: list[str] = [] diff --git a/studio/backend/core/inference/llama_stats.py b/studio/backend/core/inference/llama_stats.py index 6047aedbc0..ab0d287e8c 100644 --- a/studio/backend/core/inference/llama_stats.py +++ b/studio/backend/core/inference/llama_stats.py @@ -5,7 +5,7 @@ engine-stats log line (generation/prompt throughput, requests in flight). llama-server already computes these (it needs `--metrics`); this lifts them -into Studio's structured log so the terminal shows serving health, not just +into Unsloth's structured log so the terminal shows serving health, not just per-request access lines. Emitted only while there is activity. """ diff --git a/studio/backend/core/inference/local_model_resolver.py b/studio/backend/core/inference/local_model_resolver.py index 002cafe2c8..5d2a9e9c87 100644 --- a/studio/backend/core/inference/local_model_resolver.py +++ b/studio/backend/core/inference/local_model_resolver.py @@ -34,6 +34,15 @@ class _LocalGgufEntry: _CACHE_TTL_S = 5.0 _lock = threading.Lock() _scan: tuple[float, dict[str, _LocalGgufEntry]] = (0.0, {}) +# Not _lock: that is held for the whole scan, so the request path would wait on it. +_warm_lock = threading.Lock() +# Repos that finished downloading but are not in the published index yet: nothing +# else covers them until the next scan, and the request path must not call them absent. +_just_downloaded: set[str] = set() +_warming = False +_last_scan_s = 0.0 +# Rescan at most a tenth of the time: on the TTL alone a slow scan would run continuously. +_WARM_DUTY = 10.0 def _is_abs_path_id(value: str) -> bool: @@ -103,17 +112,26 @@ def _local_gguf_entry(loader_id: str, info) -> Optional[_LocalGgufEntry]: load_dir = _resolve_load_dir(p) variants, _ = list_local_gguf_variants(str(load_dir)) quants = tuple(v.quant for v in variants if getattr(v, "quant", None)) - return _LocalGgufEntry(loader_id, str(load_dir), quants) if quants else None + if not quants: + return None + # That call orders by descending size, so the head is the biggest quant (often + # F16). Downstream reads [0], and a bare id must mean whichever quant a plain + # load would take: answering with the largest can evict a model and then OOM. + from core.inference.openai_auto_download import preferred_quant + + best = preferred_quant(quants) + if best and quants[0] != best: + quants = (best, *(q for q in quants if q != best)) + return _LocalGgufEntry(loader_id, str(load_dir), quants) except Exception: return None -def info_has_local_gguf(info) -> bool: - """True when *info* (a LocalModelInfo) points to on-disk GGUF weights the - auto-switch path can load. Read from the files, not ``info.model_format``: the - HF-cache scanner leaves model_format unset for GGUF snapshots, so a - model_format filter would drop every cached GGUF. Lets /v1/models advertise - exactly what /v1 can serve.""" +def local_gguf_quants(info) -> Optional[tuple[str, ...]]: + """On-disk quant labels for *info*, or None when it is not a servable local + GGUF. Read from the files, not ``info.model_format``: the HF-cache scanner + leaves that unset for GGUF snapshots, so filtering on it drops every cached + GGUF. One scan tells /v1/models what it can serve and which quant to name.""" from pathlib import Path path = getattr(info, "path", None) @@ -123,14 +141,20 @@ def info_has_local_gguf(info) -> bool: if isinstance(path, str) and any( seg in (".studio_links", "ollama_links") for seg in Path(path).parts ): - return False - return _local_gguf_entry(getattr(info, "id", "") or "", info) is not None + return None + entry = _local_gguf_entry(getattr(info, "id", "") or "", info) + return entry.variants if entry is not None else None + + +def info_has_local_gguf(info) -> bool: + """True when *info* points to on-disk GGUF weights the auto-switch path can load.""" + return local_gguf_quants(info) is not None def _build_index() -> dict[str, _LocalGgufEntry]: """Map normalized id/model_id/display_name -> local GGUF entry. - Scans the same roots Studio's model picker lists (./models, the active plus + Scans the same roots Unsloth's model picker lists (./models, the active plus legacy/default HF caches, LM Studio dirs, and user scan folders) so a named local model is never missed and silently served as the loaded one. Ollama's scanner is skipped: it creates symlinks as a side effect and this runs on the @@ -146,10 +170,17 @@ def _build_index() -> dict[str, _LocalGgufEntry]: _is_hidden_model, ) from utils.paths import legacy_hf_cache_dir, hf_default_cache_dir, lmstudio_model_dirs + from utils.hf_cache_settings import known_hf_hub_caches + from core.inference.model_ids import public_model_id index: dict[str, _LocalGgufEntry] = {} seen_hf: set[str] = set() + try: + active_root = str(Path(_resolve_hf_cache_dir()).resolve()) + except Exception: + active_root = None + def _scan_hf_once(directory) -> list: if directory is None: return [] @@ -161,7 +192,13 @@ def _build_index() -> dict[str, _LocalGgufEntry]: if rp in seen_hf: return [] seen_hf.add(rp) - return _scan_hf_cache(directory) + # Only the active cache loads by repo id. Say so, or an inactive repo is + # indexed under an id it cannot load by, and its snapshot basename (what + # /v1/models advertises once loaded by path) is never a key at all. + # No format classification here: nothing on this path reads model_format, + # and its recursive walk would duplicate the one _local_gguf_entry already + # does per snapshot, on the request path. + return _scan_hf_cache(directory, active_cache = rp == active_root, classify_format = False) except Exception as exc: # a missing/malformed root must skip, never crash the index logger.debug("auto-switch: skipping HF cache dir %r: %s", directory, exc) return [] @@ -174,7 +211,12 @@ def _build_index() -> dict[str, _LocalGgufEntry]: except Exception as exc: logger.debug("auto-switch: ./models scan failed: %s", exc) try: - for hf_dir in (_resolve_hf_cache_dir(), legacy_hf_cache_dir(), hf_default_cache_dir()): + for hf_dir in ( + *known_hf_hub_caches(), + _resolve_hf_cache_dir(), + legacy_hf_cache_dir(), + hf_default_cache_dir(), + ): found += _scan_hf_once(hf_dir) except Exception as exc: logger.debug("auto-switch: HF cache scan failed: %s", exc) @@ -199,9 +241,13 @@ def _build_index() -> dict[str, _LocalGgufEntry]: raw_id = getattr(info, "id", None) if not raw_id: continue - # Skip what Studio hides from its pickers (validation probe, RAG embed + # Skip what Unsloth hides from its pickers (validation probe, RAG embed # weights): not chat models, so never an auto-switch target. - if _is_hidden_model(raw_id, getattr(info, "path", None)): + if _is_hidden_model( + raw_id, + getattr(info, "model_id", None), + getattr(info, "path", None), + ): continue # Advertise a client-facing alias, not an absolute filesystem path. loader_id = _advertised_loader_id(info) @@ -210,12 +256,91 @@ def _build_index() -> dict[str, _LocalGgufEntry]: continue # Index every alias (including the path) so a client can resolve by any of # them, even though only the non-path loader_id is advertised. - for key in (raw_id, getattr(info, "model_id", None), getattr(info, "display_name", None)): + for key in ( + raw_id, + getattr(info, "model_id", None), + getattr(info, "display_name", None), + public_model_id(raw_id), + ): if key: index.setdefault(key.strip().lower(), entry) + # Other revisions of the same repo resolve to their own weights, so a pin on + # one keeps working after Hugging Face writes a newer snapshot. + for name, sibling_entry in _sibling_revision_entries(raw_id, loader_id): + index.setdefault(name.strip().lower(), sibling_entry) return index +def _sibling_revision_entries(raw_id: str, loader_id: str): + """Yield ``(revision_name, entry)`` for the repo's OTHER cached revisions. + + An inactive-cache repo carries its snapshot path as the id, and /v1/models + advertises only that directory's basename once loaded, so anything durable + pinned to it (a subagent config) holds one revision hash. Hugging Face writes a + new snapshot dir on every update, and the scan emits a single entry per repo + pointed at the newest one, so that pin would otherwise stop resolving and drop + through to whatever model is loaded. + + Each revision gets an entry for its OWN directory rather than an alias onto the + scanned one: aliasing would redirect a pin that names an older complete revision + onto a newer half-downloaded snapshot and break a request that works today. + Incomplete revisions are skipped for the same reason. + + Sibling names are only revisions inside a real cache repo + (``/models--org--name/snapshots/``). A scan folder that merely happens + to be called ``snapshots`` holds unrelated models, and treating those as + revisions would silently serve one model in place of another. + """ + from pathlib import Path + from types import SimpleNamespace + + snapshots = Path(raw_id).parent + if snapshots.name != "snapshots" or not snapshots.parent.name.startswith("models--"): + return + from routes.models import snapshot_variants_all_complete + + try: + siblings = [p for p in snapshots.iterdir() if p.is_dir() and p.name != Path(raw_id).name] + except OSError: + return + for sibling in siblings: + if not snapshot_variants_all_complete(str(sibling)): + continue + entry = _local_gguf_entry(loader_id, SimpleNamespace(path = str(sibling))) + if entry is not None: + yield sibling.name, entry + + +def note_downloaded(repo_id: Optional[str]) -> None: + """Record a repo as present ahead of the scan that will index it.""" + if not repo_id: + return + with _lock: + _just_downloaded.add(repo_id.strip().lower()) + + +def recently_downloaded(repo_id: str) -> bool: + """Whether *repo_id* finished downloading since the last completed scan.""" + if not isinstance(repo_id, str) or not repo_id.strip(): + return False + return repo_id.strip().lower() in _just_downloaded + + +def invalidate_index() -> None: + """Mark the cached scan stale so the next resolve sees a just-finished download + instead of waiting out the TTL. + + Keeps the entries: the request path reads this cache without scanning, so + emptying it would leave it with no evidence about any local model until the + rebuild lands, and a bare request for one would be answered by whatever is + resident. Only a completed download invalidates, and that only adds, so the + retained entries stay true. + """ + global _scan + with _lock: + _scan = (0.0, _scan[1]) + + def _index() -> dict[str, _LocalGgufEntry]: global _scan # Build under the lock so concurrent callers with an expired cache don't all @@ -230,23 +355,74 @@ def _index() -> dict[str, _LocalGgufEntry]: # an install with many local models can itself exceed the TTL, which would # store the cache already expired and make every request rebuild the index. _scan = (time.monotonic(), fresh) + # The scan supersedes the notes: whatever landed is in the index now. + _just_downloaded.clear() return fresh -def resolve_local_gguf(requested: str) -> Optional[tuple[str, Optional[str], str]]: +def index_is_built() -> bool: + """Whether a scan has ever completed, freshness aside. + + Lock-free on purpose: ``_lock`` is held for the whole scan, so taking it would + park the request path on the scan it is trying to stay off. Safe because + ``_scan`` is only ever rebound, never mutated. + """ + return bool(_scan[0]) + + +def warm_index_soon() -> None: + """(Re)build the index off the request path when it is missing or past its TTL. + + The only refresh for callers using ``allow_scan=False``. Covers a stale index, + not just an absent one: a model downloaded through the Hub UI or dropped into a + scan folder has no invalidation hook and would otherwise stay invisible to them + for the life of the process. Never blocks, and never touches ``_lock``. + """ + global _warming + if time.monotonic() - _scan[0] < max(_CACHE_TTL_S, _last_scan_s * _WARM_DUTY): + return + with _warm_lock: + if _warming: + return + _warming = True + + def _run() -> None: + global _warming, _last_scan_s + started = time.monotonic() + try: + _index() + except Exception: + pass + finally: + _last_scan_s = time.monotonic() - started + with _warm_lock: + _warming = False + + threading.Thread(target = _run, name = "local-model-index-warm", daemon = True).start() + + +def resolve_local_gguf( + requested: str, *, allow_scan: bool = True +) -> Optional[tuple[str, Optional[str], str]]: """Return ``(load_path, gguf_variant, loader_id)`` for a local match, else None. ``load_path`` is the concrete on-disk path to hand /load (so it never fetches a remote), ``loader_id`` is the advertised id used as the launch-override key. ``requested`` is ``repo`` or ``repo:VARIANT``. An exact id match wins first (so ids containing a colon still resolve); else the last ``:VARIANT`` is split - off and resolves only when that quant is on disk. + off and resolves only when that quant is on disk, unless it names no quant at + all (an Ollama-style ":latest"), which means the repo. + + ``allow_scan=False`` answers from the last built index and never rebuilds, for + the request path: the scan walks several model dirs and HF caches, takes seconds + on a large install, and holds a lock everyone queues behind. Stale is fine there, + since disk barely moves and a finished download calls :func:`invalidate_index`. """ if not isinstance(requested, str) or not requested.strip(): return None requested = requested.strip() try: - index = _index() + index = _index() if allow_scan else _scan[1] entry = index.get(requested.lower()) if entry is not None: variant = entry.variants[0] if entry.variants else None @@ -262,8 +438,44 @@ def resolve_local_gguf(requested: str) -> Optional[tuple[str, Optional[str], str for v in entry.variants: if v.lower() == wanted: return entry.load_path, v, entry.loader_id - return None + from core.inference.openai_auto_download import looks_like_quant + + if looks_like_quant(variant): + return None + # ":latest" or ":8b" names no file, so it means the repo; a real quant that + # is not on disk still misses, or a swap would serve the wrong weights. + return entry.load_path, (entry.variants[0] if entry.variants else None), entry.loader_id except Exception: # Best-effort: any resolver failure falls through to the loaded model, # so a malformed name can never turn a servable request into a 500. return None + + +MISS_MODEL_NOT_FOUND = "model_not_found" +MISS_VARIANT_NOT_FOUND = "variant_not_found" + + +def describe_local_miss(requested: str) -> tuple[str, tuple[str, ...]]: + """Why :func:`resolve_local_gguf` missed, so an error can say "wrong quant" + instead of "no such model". + + ``(MISS_VARIANT_NOT_FOUND, )`` when the repo is downloaded but the + requested ``:VARIANT`` is not, else ``(MISS_MODEL_NOT_FOUND, ())``. Fail-safe: a + scan failure reports the generic miss rather than raising into the handler. + """ + if not isinstance(requested, str) or not requested.strip(): + return MISS_MODEL_NOT_FOUND, () + base, sep, variant = requested.strip().rpartition(":") + from core.inference.openai_auto_download import looks_like_quant + + # Split like the resolver or the two disagree: a tag naming no quant means the + # repo there, so reporting a missing quant for it would name one nobody asked for. + if not sep or not looks_like_quant(variant): + return MISS_MODEL_NOT_FOUND, () + try: + entry = _index().get(base.strip().lower()) + except Exception: + return MISS_MODEL_NOT_FOUND, () + if entry is None or not entry.variants: + return MISS_MODEL_NOT_FOUND, () + return MISS_VARIANT_NOT_FOUND, entry.variants diff --git a/studio/backend/core/inference/mcp_client.py b/studio/backend/core/inference/mcp_client.py index 6b5ce02216..98112c6d5b 100644 --- a/studio/backend/core/inference/mcp_client.py +++ b/studio/backend/core/inference/mcp_client.py @@ -906,7 +906,7 @@ def _call_stdio_tool( def _remaining() -> Optional[float]: return None if deadline is None else max(0.0, deadline - time.monotonic()) - # Callers without a Studio session id must retain the former one-shot + # Callers without an Unsloth session id must retain the former one-shot # behavior: no browser/cookie/tool state can leak into another request. # Use an ephemeral key (and close it below) rather than the shared empty # scope that the persistent-session cache used previously. @@ -971,7 +971,12 @@ def _call_stdio_tool( raise RuntimeError("MCP server connection is not available") else: rem = _remaining() - coro = _race_tool_call(session.client.call_tool(name, args), rem, cancel_event) + # raise_on_error=False for the same reason as the one-shot path. + coro = _race_tool_call( + session.client.call_tool(name, args, raise_on_error = False), + rem, + cancel_event, + ) return session.run(coro, rem) except (_MCPCancelled, asyncio.TimeoutError): # _race_tool_call cancels the pending call but cancellation is diff --git a/studio/backend/core/inference/mlx_inference.py b/studio/backend/core/inference/mlx_inference.py index 163ade10c4..2b300a32b1 100644 --- a/studio/backend/core/inference/mlx_inference.py +++ b/studio/backend/core/inference/mlx_inference.py @@ -8,14 +8,76 @@ instead of torch/transformers for model loading and generation. import json import os import threading +from contextlib import contextmanager from typing import Optional, Generator from core.inference.message_content import content_to_text from core.inference.runtime_context import runtime_context_length +from core.inference.chat_template_helpers import ( + ReasoningChannelNormalizer, + normalize_reasoning_snapshots, +) from loggers import get_logger logger = get_logger(__name__) +def _mlx_adapter_modules(model): + """Return bypassable adapter entries and unsupported wrapper paths.""" + adapters = [] + unsupported = [] + for path, module in model.named_modules(): + if not path or not (hasattr(module, "lora_a") and hasattr(module, "lora_b")): + continue + base = getattr(module, "linear", None) + if base is None: + base = getattr(module, "embedding", None) + if base is None: + unsupported.append(path) + else: + adapters.append((path, module, base)) + return adapters, unsupported + + +@contextmanager +def _temporary_mlx_adapter_state(model, use_adapter): + """Select base or adapter modules for one request, then restore the tree.""" + if use_adapter is None: + yield + return + if isinstance(use_adapter, str): + raise NotImplementedError( + "Unsloth MLX: named adapter selection is not supported; use True for " + "the loaded adapter or False for the base model." + ) + if use_adapter is not True and use_adapter is not False: + raise TypeError("Unsloth MLX: use_adapter must be None, True, False, or a string.") + + adapters, unsupported = _mlx_adapter_modules(model) + if use_adapter is True: + if not adapters and not unsupported: + logger.warning("MLX adapter requested, but the active model has no adapter layers") + yield + return + if unsupported: + raise RuntimeError( + "Unsloth MLX: cannot disable adapter layers without their base modules: " + + ", ".join(unsupported[:5]) + ) + if not adapters: + yield + return + + from mlx.utils import tree_unflatten + + base_modules = tree_unflatten([(path, base) for path, _, base in adapters]) + adapter_modules = tree_unflatten([(path, wrapper) for path, wrapper, _ in adapters]) + try: + model.update_modules(base_modules) + yield + finally: + model.update_modules(adapter_modules) + + def _mlx_vlm_model_config(model): """Return the loaded MLX model config and its type, preferring whichever of config / _config actually carries a model_type.""" @@ -119,19 +181,27 @@ def _vlm_messages_have_tool_history(messages): ) -def _build_generation_stats(prompt_n, prompt_tps, gen_n, gen_tps): +def _build_generation_stats( + prompt_n, + prompt_tps, + gen_n, + gen_tps, + cached_n = 0, +): """Map mlx stream stats onto the usage/timings shape llama-server emits.""" prompt_n = int(prompt_n or 0) gen_n = int(gen_n or 0) + cached_n = int(cached_n or 0) prompt_tps = float(prompt_tps or 0.0) gen_tps = float(gen_tps or 0.0) prompt_ms = (prompt_n / prompt_tps * 1000.0) if prompt_tps > 0 else 0.0 predicted_ms = (gen_n / gen_tps * 1000.0) if gen_tps > 0 else 0.0 + total_prompt_n = prompt_n + cached_n return { "usage": { - "prompt_tokens": prompt_n, + "prompt_tokens": total_prompt_n, "completion_tokens": gen_n, - "total_tokens": prompt_n + gen_n, + "total_tokens": total_prompt_n + gen_n, }, "timings": { "prompt_n": prompt_n, @@ -142,11 +212,123 @@ def _build_generation_stats(prompt_n, prompt_tps, gen_n, gen_tps): "predicted_ms": predicted_ms, "predicted_per_token_ms": (predicted_ms / gen_n) if gen_n > 0 else 0.0, "predicted_per_second": gen_tps, - "cache_n": 0, + "cache_n": cached_n, }, } +PROMPT_CACHE_ENTRIES = 6 +PROMPT_CACHE_MEMORY_FRACTION = 0.15 +PROMPT_CACHE_FALLBACK_BYTES = 2 * 1024**3 + + +def _mlx_prompt_cache_api(): + try: + from mlx_lm.models.cache import ( + LRUPromptCache, + can_trim_prompt_cache, + make_prompt_cache, + trim_prompt_cache, + ) + except ImportError: + return None + return LRUPromptCache, make_prompt_cache, can_trim_prompt_cache, trim_prompt_cache + + +def _prompt_cache_max_bytes(recommended_gb = None): + override = os.environ.get("UNSLOTH_MLX_PROMPT_CACHE_BYTES") + if override: + try: + return max(int(override), 0) + except ValueError: + logger.warning("Ignoring non-integer UNSLOTH_MLX_PROMPT_CACHE_BYTES=%r", override) + if recommended_gb: + return int(recommended_gb * 1e9 * PROMPT_CACHE_MEMORY_FRACTION) + return PROMPT_CACHE_FALLBACK_BYTES + + +def _flatten_kv_entries(cache): + for entry in cache: + nested = getattr(entry, "caches", None) + if nested is None: + yield entry + else: + yield from _flatten_kv_entries(nested) + + +def _kv_prefix_coverage(cache): + covered = None + for entry in _flatten_kv_entries(cache): + offset = getattr(entry, "offset", None) + if offset is None: + return None + if getattr(entry, "start_position", 0): + return None + window = getattr(entry, "max_size", None) + if window is not None and offset > window: + return None + if covered is None: + covered = offset + elif covered != offset: + return None + return covered + + +class _MLXPromptCacheHistory: + def __init__(self, max_entries, max_bytes): + api = _mlx_prompt_cache_api() + if api is None: + raise RuntimeError("mlx-lm is too old for LRUPromptCache") + lru_cls, make, can_trim, trim = api + self._make_prompt_cache = make + self._can_trim = can_trim + self._trim = trim + self._max_bytes = max_bytes + self._lru = lru_cls(max_size = max_entries, max_bytes = max_bytes) + + def fetch(self, model, key, tokens): + cache, rest = self._lru.fetch_nearest_cache(key, list(tokens)) + if cache is not None: + if rest: + return cache, list(rest) + if self._can_trim(cache) and self._trim(cache, 1) == 1: + return cache, list(tokens[-1:]) + if len(tokens) > 1: + head = list(tokens[:-1]) + cache, rest = self._lru.fetch_nearest_cache(key, head) + if cache is not None: + covered = len(head) - len(rest) + return cache, list(tokens[covered:]) + return self._make_prompt_cache(model), list(tokens) + + def insert(self, key, tokens, cache): + # An over-budget entry evicts itself and every other conversation. + nbytes = sum(getattr(entry, "nbytes", 0) for entry in cache) + if nbytes > self._max_bytes: + logger.debug( + "MLX prompt cache: skipping %.2f GB entry over the %.2f GB budget", + nbytes / 1e9, + self._max_bytes / 1e9, + ) + return + covered = _kv_prefix_coverage(cache) + if covered is None: + logger.debug("MLX prompt cache: skipping cache with unverifiable prefix coverage") + return + tokens = list(tokens) + if covered > len(tokens): + logger.debug( + "MLX prompt cache: cache covers %d tokens but only %d were tracked", + covered, + len(tokens), + ) + return + tokens = tokens[:covered] + if not tokens: + return + self._lru.insert_cache(key, tokens, cache) + + def _mlx_distributed_rank_size(group = None): """Return ``(rank, world_size)`` for an optional MLX distributed group.""" if group is None: @@ -251,6 +433,55 @@ class MLXInferenceBackend: # Recorded for unload to release pinned memory back to the OS. self._memory_limits_applied = {} + self._prompt_cache_history = None + self._prompt_cache_unavailable = False + + def _prompt_cache(self): + if self._prompt_cache_history is not None or self._prompt_cache_unavailable: + return self._prompt_cache_history + max_bytes = _prompt_cache_max_bytes(self._memory_limits_applied.get("recommended_gb")) + if max_bytes <= 0: + self._prompt_cache_unavailable = True + logger.info("MLX prompt cache disabled by budget") + return None + try: + self._prompt_cache_history = _MLXPromptCacheHistory( + PROMPT_CACHE_ENTRIES, + max_bytes, + ) + except Exception as exc: + self._prompt_cache_unavailable = True + logger.info("MLX prompt cache unavailable (%s); prefilling every request", exc) + return None + logger.info( + "MLX prompt cache: %d entries, %.2f GB budget", + PROMPT_CACHE_ENTRIES, + max_bytes / 1e9, + ) + return self._prompt_cache_history + + def _clear_prompt_cache(self): + self._prompt_cache_history = None + self._prompt_cache_unavailable = False + + def _prepare_prompt_cache(self, prompt, adapter_state): + history = self._prompt_cache() + if history is None: + return prompt, None, None, None, 0 + try: + tokenizer = self._tokenizer + bos = getattr(tokenizer, "bos_token", None) + add_special_tokens = bos is None or not prompt.startswith(bos) + tokens = list(tokenizer.encode(prompt, add_special_tokens = add_special_tokens)) + if not tokens: + return prompt, None, None, None, 0 + key = f"{self.active_model_name}|{adapter_state!r}" + cache, rest = history.fetch(self._model, key, tokens) + except Exception as exc: + logger.debug("MLX prompt cache lookup failed: %s", exc) + return prompt, None, None, None, 0 + return rest, cache, key, tokens, len(tokens) - len(rest) + def _configure_memory_limits(self): """Apply Metal memory caps before loading a model. @@ -473,6 +704,7 @@ class MLXInferenceBackend: self._distributed_world_size = 1 if self.active_model_name == model_name: self.active_model_name = None + self._clear_prompt_cache() gc.collect() mx.clear_cache() @@ -504,6 +736,7 @@ class MLXInferenceBackend: reasoning_effort = None, preserve_thinking = None, presence_penalty = 0.0, + _adapter_state = None, ) -> Generator[str, None, None]: if self._model is None: raise RuntimeError("No model loaded") @@ -533,7 +766,7 @@ class MLXInferenceBackend: break if self._is_vlm: - yield from self._generate_vlm( + stream = self._generate_vlm( full_messages, image, temperature, @@ -548,9 +781,10 @@ class MLXInferenceBackend: reasoning_effort = reasoning_effort, preserve_thinking = preserve_thinking, presence_penalty = presence_penalty, + _adapter_state = _adapter_state, ) else: - yield from self._generate_text( + stream = self._generate_text( full_messages, temperature, top_p, @@ -564,7 +798,9 @@ class MLXInferenceBackend: reasoning_effort = reasoning_effort, preserve_thinking = preserve_thinking, presence_penalty = presence_penalty, + _adapter_state = _adapter_state, ) + yield from stream def _generate_text( self, @@ -582,6 +818,7 @@ class MLXInferenceBackend: reasoning_effort = None, preserve_thinking = None, presence_penalty = 0.0, + _adapter_state = None, ): from mlx_lm import stream_generate from mlx_lm.sample_utils import make_sampler, make_logits_processors @@ -609,7 +846,7 @@ class MLXInferenceBackend: # probe and native render share a renderer. (VLM renders via the # processor for image tokens and is not wired here.) model_info = self.models.get(self.active_model_name, {}) - prompt = render_with_native_template_fallback( + render_result = render_with_native_template_fallback( formatted_prompt = prompt, tokenizer = self._tokenizer, model_info = model_info, @@ -620,17 +857,16 @@ class MLXInferenceBackend: reasoning_effort = reasoning_effort, preserve_thinking = preserve_thinking, hf_token = model_info.get("hf_token"), + return_metadata = True, ) + prompt = render_result.prompt + reasoning_channel_markers = render_result.reasoning_channel_markers # An open prefilled by the template lives in the prompt, not # the generated tokens; re-emit it so the frontend renders the block. think_prefix = detect_think_prefill( prompt, getattr(self._tokenizer, "all_special_tokens", None) ) - # Emit it before the first token so the block renders during prefill. - if think_prefix: - yield think_prefix - sampler = make_sampler( temp = temperature, top_p = top_p, @@ -654,22 +890,45 @@ class MLXInferenceBackend: if not logits_processors: logits_processors = None + preserve_native_channels = reasoning_channel_markers is not None token_ids = [] - logger.info( - "Generating: prompt_len=%d, max_tokens=%d, model=%s, tokenizer=%s", - len(prompt), - max_new_tokens, - type(self._model).__name__, - type(self._tokenizer).__name__, + normalizer = ( + ReasoningChannelNormalizer(*reasoning_channel_markers) + if reasoning_channel_markers is not None + else None ) - with self._generation_lock: + # MLX consumers diff cumulative snapshots. Keep a prompt-prefilled + # prefix on every native-protocol snapshot just as the normal + # decoding path does below. + normalized_output = think_prefix + with self._generation_lock, _temporary_mlx_adapter_state(self._model, _adapter_state): + ( + gen_prompt, + prompt_cache, + cache_key, + prompt_tokens, + cached_n, + ) = self._prepare_prompt_cache(prompt, _adapter_state) + logger.info( + "Generating: prompt_len=%d, cached=%d, max_tokens=%d, model=%s, tokenizer=%s", + len(prompt), + cached_n, + max_new_tokens, + type(self._model).__name__, + type(self._tokenizer).__name__, + ) final_response = None try: + # Enter request-scoped model state before yielding any response. + if think_prefix: + yield think_prefix gen_kwargs = dict( - prompt = prompt, + prompt = gen_prompt, max_tokens = max_new_tokens, sampler = sampler, ) + if prompt_cache is not None: + gen_kwargs["prompt_cache"] = prompt_cache if logits_processors is not None: gen_kwargs["logits_processors"] = logits_processors for response in stream_generate( @@ -679,14 +938,28 @@ class MLXInferenceBackend: ): final_response = response token_ids.append(response.token) - cumulative = self._tokenizer.decode( - token_ids, - skip_special_tokens = True, - ) - yield think_prefix + cumulative + if preserve_native_channels: + piece = getattr(response, "text", None) or "" + delta = normalizer.feed(piece) + if delta: + normalized_output += delta + yield normalized_output + else: + cumulative = self._tokenizer.decode( + token_ids, + skip_special_tokens = True, + ) + yield think_prefix + cumulative if cancel_event and cancel_event.is_set(): break + if prompt_cache is not None and prompt_tokens is not None: + history = self._prompt_cache_history + if history is not None: + try: + history.insert(cache_key, prompt_tokens + token_ids, prompt_cache) + except Exception as exc: + logger.debug("MLX prompt cache insert failed: %s", exc) except Exception as e: import traceback logger.error("stream_generate failed:\n%s", traceback.format_exc()) @@ -699,7 +972,14 @@ class MLXInferenceBackend: getattr(final_response, "prompt_tps", 0.0), getattr(final_response, "generation_tokens", 0), getattr(final_response, "generation_tps", 0.0), + cached_n, ) + if normalizer is not None: + cancelled = cancel_event is not None and cancel_event.is_set() + tail = normalizer.drain() if cancelled else normalizer.finish() + if tail: + normalized_output += tail + yield normalized_output def _generate_vlm( self, @@ -718,6 +998,7 @@ class MLXInferenceBackend: reasoning_effort = None, preserve_thinking = None, presence_penalty = 0.0, + _adapter_state = None, ): from mlx_vlm import stream_generate as vlm_stream @@ -821,9 +1102,6 @@ class MLXInferenceBackend: # Re-emit an open prefill from the prompt (see _generate_text). cumulative = detect_think_prefill(prompt, getattr(chat_target, "all_special_tokens", None)) - # Emit it before the first token so the block renders during prefill. - if cumulative: - yield cumulative logger.info( "VLM generating: prompt_len=%d, has_image=%s", len(prompt), @@ -858,31 +1136,46 @@ class MLXInferenceBackend: elif _rep_active: vlm_kwargs["repetition_penalty"] = float(repetition_penalty) - with self._generation_lock: - final_response = None - try: - for response in vlm_stream( - self._model, - self._processor, - prompt, - images, - **vlm_kwargs, - ): - final_response = response - token_text = response.text if hasattr(response, "text") else str(response) - cumulative += token_text - yield cumulative - if cancel_event and cancel_event.is_set(): - break - finally: - # mlx_vlm exposes the same stats fields as mlx_lm. - if final_response is not None: - self.last_generation_stats = _build_generation_stats( - getattr(final_response, "prompt_tokens", 0), - getattr(final_response, "prompt_tps", 0.0), - getattr(final_response, "generation_tokens", 0), - getattr(final_response, "generation_tps", 0.0), - ) + def _stream_vlm_snapshots(): + nonlocal cumulative + # Hold the generation lock AND the request-scoped adapter state for the + # whole stream so Base-vs-LoRA compare mode honors use_adapter and the + # wrapper tree is restored on completion, cancellation, or close. + with self._generation_lock, _temporary_mlx_adapter_state(self._model, _adapter_state): + final_response = None + try: + # Emit any prefilled block before the first token so the + # UI renders it during prefill, matching _generate_text. Done + # inside the adapter context so an unsupported request raises + # before any output escapes. + if cumulative: + yield cumulative + for response in vlm_stream( + self._model, + self._processor, + prompt, + images, + **vlm_kwargs, + ): + final_response = response + token_text = response.text if hasattr(response, "text") else str(response) + cumulative += token_text + yield cumulative + if cancel_event and cancel_event.is_set(): + break + finally: + # mlx_vlm exposes the same stats fields as mlx_lm. + if final_response is not None: + self.last_generation_stats = _build_generation_stats( + getattr(final_response, "prompt_tokens", 0), + getattr(final_response, "prompt_tps", 0.0), + getattr(final_response, "generation_tokens", 0), + getattr(final_response, "generation_tps", 0.0), + ) + + yield from normalize_reasoning_snapshots( + _stream_vlm_snapshots(), chat_target, cancel_event, tools = tools + ) def generate_with_adapter_control( self, @@ -890,10 +1183,14 @@ class MLXInferenceBackend: cancel_event = None, **gen_kwargs, ) -> Generator[str, None, None]: - # MLX LoRA adapter toggling not yet supported; generate normally - yield from self.generate_chat_response(cancel_event = cancel_event, **gen_kwargs) + yield from self.generate_chat_response( + cancel_event = cancel_event, + _adapter_state = use_adapter, + **gen_kwargs, + ) - def reset_generation_state(self): + def reset_generation_state(self, caller_cancel_event = None): + # caller_cancel_event: signature parity with the orchestrator; unused here. import mlx.core as mx import gc diff --git a/studio/backend/core/inference/model_ids.py b/studio/backend/core/inference/model_ids.py index 548cc60f94..3886307ae2 100644 --- a/studio/backend/core/inference/model_ids.py +++ b/studio/backend/core/inference/model_ids.py @@ -39,10 +39,29 @@ def _looks_like_path(identifier: str) -> bool: return False +def hf_cache_repo_id(path: Optional[str]) -> Optional[str]: + """``.../models--org--name/snapshots/`` -> ``org/name``, else None. + + A model loaded from the HF cache is identified by its snapshot dir, whose + basename is a commit hash; recover the repo id so callers don't show that. + """ + if not path: + return None + parts = str(path).replace("\\", "/").split("/") + for index, part in enumerate(parts): + # Only inside the real cache layout: a "models--" name alone is not a repo id. + if part.startswith("models--") and parts[index + 1 : index + 2] == ["snapshots"]: + return part[len("models--") :].replace("--", "/") + return None + + def public_model_id(identifier: Optional[str]) -> Optional[str]: """Return a clean, path-free public id for *identifier*. - - Local GGUF path -> the file stem with ``.gguf`` stripped, e.g. + - HF cache path -> the repo id it came from, e.g. + ``~/.cache/huggingface/hub/models--unsloth--X-GGUF/snapshots/`` -> + ``unsloth/X-GGUF``. + - Other local GGUF path -> the file stem with ``.gguf`` stripped, e.g. ``/srv/models/Qwen3-30B-A3B-Q4_K_M.gguf`` -> ``Qwen3-30B-A3B-Q4_K_M``. - HF repo id (``org/model``) and already-clean names -> returned unchanged. - ``None`` / empty -> returned unchanged. @@ -51,6 +70,9 @@ def public_model_id(identifier: Optional[str]) -> Optional[str]: return identifier if not _looks_like_path(identifier): return identifier + repo_id = hf_cache_repo_id(identifier) + if repo_id: + return repo_id name = os.path.basename(identifier.replace("\\", "/").rstrip("/")) if name.lower().endswith(_GGUF_SUFFIX): name = name[: -len(_GGUF_SUFFIX)] diff --git a/studio/backend/core/inference/openai_auto_download.py b/studio/backend/core/inference/openai_auto_download.py new file mode 100644 index 0000000000..cad5e40d14 --- /dev/null +++ b/studio/backend/core/inference/openai_auto_download.py @@ -0,0 +1,812 @@ +# SPDX-License-Identifier: AGPL-3.0-only +# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0 + +"""Opt-in: fetch a GGUF a /v1 request names but this server doesn't have. + +Auto-switch only loads models already on disk. With +``openai_api_auto_download_model`` on, a miss that looks like a real Hub repo is +fetched in the background and the request is told to retry rather than held +open: a quant is routinely tens of GB, far longer than any client (or the +Cloudflare edge on ``--secure``) will wait, and the inference lifecycle gate must +not be held meanwhile. The resident model keeps serving, and the retry that lands +after the download goes through the ordinary auto-switch path. + +Admission is deliberately narrow, since a request only needs an API key: +- ``namespace/name`` only, and only when the Hub confirms GGUF weights. A + namespace is not evidence of intent (LiteLLM and OpenRouter address every + provider that way), so ``gpt-4`` and ``anthropic/claude-3.5-sonnet`` alike + fall through to the resident model as before. +- GGUF only, decided from the remote file list, not the repo name: GGUF runs + under llama.cpp, which never imports repo Python. +- ``auto_map`` is refused, so ``trust_remote_code`` is only ever granted + deliberately in the UI, never by an API call. +- One download at a time, so a key holder cannot fan out fetches. +""" + +from __future__ import annotations + +import asyncio +import shutil +import threading +import time +from dataclasses import dataclass +from typing import Optional + +from loggers import get_logger + +logger = get_logger(__name__) + +# Keep the Hub probe short so a slow Hub can't stall the request path. +_MODEL_INFO_TIMEOUT_S = 8.0 +# auth_check and hf_hub_download take no timeout of their own and run while the +# provisional slot is held, so an unresponsive Hub would pin the single flight. The +# code probe fetches up to three configs, so it gets more room than the auth call. +_CODE_PROBE_TIMEOUT_S = 20.0 +# Headroom left free after the download, so filling the disk can't wedge the box. +_DISK_RESERVE_BYTES = 5 * 1024**3 +_WATCH_POLL_S = 2.0 +# A stalled watcher must not pin the single-flight slot forever. +_MAX_WATCH_S = 24 * 60 * 60 +# Past the watch window the row is resolved, so poll only to see whether the +# worker is still alive and still owns the slot. +_TIMED_OUT_POLL_S = 60.0 +_RETRY_AFTER_S = 30 +# Long enough for a client honouring Retry-After to come back and be told, short +# enough that one that never returns cannot hold the slot. +_FAILED_HOLD_S = 3 * _RETRY_AFTER_S +_MAX_LISTED_VARIANTS = 8 + + +@dataclass(frozen = True) +class AutoDownloadRefusal: + """Why this request cannot be served yet; the route raises it in the + surface's own error envelope.""" + + status: int + code: str + message: str + retry_after: Optional[int] = None + + +@dataclass +class _Active: + repo_id: str + # None while the Hub probe is still deciding which quant to fetch. + variant: Optional[str] = None + expected_bytes: int = 0 + monitor_id: Optional[str] = None + started_at: float = 0.0 + # Set when the worker failed. Held until a retry surfaces it: Retry-After is far + # longer than the watcher poll, so the client would restart the same failing download. + error: Optional[str] = None + failed_at: float = 0.0 + + +_lock = threading.Lock() +_active: Optional[_Active] = None + +# Repos the Hub says are not servable, so a "vendor/model" miss doesn't re-probe every request. +_NOT_SERVABLE_TTL_S = 10 * 60 +_NOT_SERVABLE_MAX = 256 +_cache_lock = threading.Lock() +_not_servable: dict[str, float] = {} + + +def _public_label(repo_id: str, variant: Optional[str]) -> str: + return f"{repo_id}:{variant}" if variant else repo_id + + +def split_model_ref(requested: str) -> tuple[str, Optional[str]]: + """``org/repo:QUANT`` -> ``("org/repo", "QUANT")``; no suffix -> variant None. + + Splits on the last colon. A slash-bearing suffix is only a variant when a real + Hub repo precedes it: "build/llama-13b" is a subdirectory GGUF key the catalog + advertises, while "C:/models/x.gguf" leaves a drive letter that is no repo id. + """ + text = (requested or "").strip() + base, sep, suffix = text.rpartition(":") + if not sep or not base or not suffix: + return text, None + stripped = base.strip() + if "/" in suffix: + from hub.utils.paths import is_valid_repo_id + if "/" not in stripped or not is_valid_repo_id(stripped): + return text, None + return stripped, suffix.strip() + + +def is_downloadable_ref(requested: str) -> bool: + """Whether *requested* is shaped like a Hub repo we may fetch. + + Requires an explicit namespace: keeps ``gpt-4`` and other foreign ids falling + through, and stops ModelConfig.from_identifier's bare-name ``unsloth/`` + prefixing from turning an unrelated label into a real repo. + """ + from hub.utils.paths import is_valid_repo_id + + repo_id, variant = split_model_ref(requested) + if "/" not in repo_id or not is_valid_repo_id(repo_id): + return False + if variant is not None: + from hub.utils.paths import is_valid_gguf_variant + return is_valid_gguf_variant(variant) + return True + + +def looks_like_quant(variant: Optional[str]) -> bool: + """Whether a ``:suffix`` names a GGUF quant rather than a foreign tag. + + Neither a namespace nor a colon proves a request was meant for this server + (``vendor/model`` is LiteLLM/OpenRouter, ``name:latest`` is Ollama). A real + quant label does. + """ + import re + + from utils.models.model_config import _GGUF_KNOWN_QUANT_RE + + if not variant: + return False + # _extract_quant_label can append a bpw modifier (IQ4_XS-3.53bpw); still a quant. + label = re.sub(r"-[0-9]+(?:\.[0-9]+)?bpw$", "", variant.strip(), flags = re.IGNORECASE) + return _GGUF_KNOWN_QUANT_RE.fullmatch(label) is not None + + +def _hub_token(hf_token: Optional[str]): + """The caller's token, or an explicit False. None makes huggingface_hub fall + back to a cached login (here the server owner's); only False is anonymous.""" + return hf_token or False + + +def _servable_key(repo_id: str, hf_token: Optional[str]) -> str: + """Cache key, per credential. + + The Hub 404s a private repo the caller cannot see, so a tokenless verdict says + nothing about a caller who has one. Digested, so no token is held here. + """ + import hashlib + + seen_as = hashlib.sha256(hf_token.encode()).hexdigest()[:16] if hf_token else "anon" + return f"{repo_id.lower()}\n{seen_as}" + + +def _mark_not_servable(repo_id: str, hf_token: Optional[str]) -> None: + with _cache_lock: + if len(_not_servable) >= _NOT_SERVABLE_MAX: + _not_servable.clear() + _not_servable[_servable_key(repo_id, hf_token)] = time.monotonic() + _NOT_SERVABLE_TTL_S + + +def _is_not_servable(repo_id: str, hf_token: Optional[str]) -> bool: + key = _servable_key(repo_id, hf_token) + with _cache_lock: + expires = _not_servable.get(key) + if expires is None: + return False + if expires <= time.monotonic(): + del _not_servable[key] + return False + return True + + +def _gated_refusal(repo_id: str) -> AutoDownloadRefusal: + return AutoDownloadRefusal( + status = 403, + code = "model_access_denied", + message = ( + f"'{repo_id}' is gated on Hugging Face. Accept its licence, then retry with " + "your own token in the X-Unsloth-HF-Token header: automatic download never " + "uses this server's Hugging Face identity." + ), + ) + + +async def _bounded_probe(fn, *args, timeout: float, default): + """Run a blocking Hub probe off the loop, bounding only the wait. + + The thread is left to finish (a blocking socket read cannot be cancelled); the + caller takes *default*, chosen per call site so a timeout errs the safe way. + """ + try: + return await asyncio.wait_for(asyncio.to_thread(fn, *args), timeout) + except (TimeoutError, asyncio.TimeoutError): + logger.debug("hub probe %s timed out after %ss", getattr(fn, "__name__", fn), timeout) + return default + + +def _auth_denied(repo_id: str, hf_token: Optional[str]) -> bool: + """Whether this token lacks file access to a gated repo. False when the + check is inconclusive: the download's own auth is the real gate.""" + from hub.utils.hf_errors import hf_error_status + + try: + from huggingface_hub import auth_check + auth_check(repo_id, token = _hub_token(hf_token)) + except Exception as exc: + return hf_error_status(exc) in (401, 403) + return False + + +def _gguf_variants(siblings) -> dict[str, int]: + """Quant label -> bytes the download will actually fetch. + + Mirrors list_gguf_variants for the selectable labels: companions (mmproj/MTP) + and big-endian builds are not quants, and sharded quants sum across shards. + Bytes come from the download plan, which folds companions back into every + quant, so the disk reserve is measured against what the worker fetches. + """ + from hub.utils.gguf import extract_quant_label as canonical_quant_label + from hub.utils.gguf_plan import build_gguf_variant_plans + from utils.models.model_config import ( + _extract_quant_label, + _is_big_endian_gguf_path, + _is_mmproj, + _is_mtp_drafter, + ) + + siblings = list(siblings or []) + plans = build_gguf_variant_plans(siblings) + sizes: dict[str, int] = {} + for sibling in siblings: + name = getattr(sibling, "rfilename", "") or "" + if not name.lower().endswith(".gguf"): + continue + quant = _extract_quant_label(name) + if not looks_like_quant(quant): + # With no recognized quant token the extractors part ways: this one takes + # the last hyphenated segment ("7b" of llama-7b) while the plan and worker + # key the whole stem, so advertising ours dispatches an unresolvable variant. + quant = canonical_quant_label(name) or quant + if _is_mmproj(name) or _is_mtp_drafter(name) or _is_big_endian_gguf_path(name, quant): + continue + plan = plans.get(quant.lower()) + if plan is not None: + sizes[quant] = plan.download_size_bytes + else: + sizes[quant] = sizes.get(quant, 0) + int(getattr(sibling, "size", 0) or 0) + return sizes + + +def _remaining_bytes(repo_id: str, plan, expected_bytes: int) -> int: + """Bytes still to fetch: a resumed quant or a companion shared with another + quant is already on disk, and charging for it can 507 a download that fits.""" + try: + from hub.utils.download_registry import existing_blob_bytes + + hashes = frozenset( + file.sha256 for file in getattr(plan, "expected_files", ()) or () if file.sha256 + ) + if not hashes: + return expected_bytes + return max(0, expected_bytes - existing_blob_bytes("model", repo_id, hashes)) + except Exception: + return expected_bytes + + +def _enough_disk(need_bytes: int) -> tuple[bool, int]: + """(fits, free_bytes). Fail-open on an unreadable cache root: the download + worker runs its own preflight, this only adds the reserve margin.""" + try: + from hub.utils.hf_cache_state import hf_cache_root + + root = hf_cache_root(create = True) + if root is None: + return True, 0 + free = shutil.disk_usage(root).free + except Exception: + return True, 0 + return free >= need_bytes + _DISK_RESERVE_BYTES, free + + +def _gb(num_bytes: int) -> str: + return f"{num_bytes / 1024**3:.1f} GB" + + +async def _job_state(repo_id: str, variant: Optional[str]) -> tuple[str, Optional[str]]: + from hub.services.models import downloads + try: + status = await downloads.get_download_status_response(repo_id, variant or "") + return status.state, status.error + except Exception as exc: + # "unknown", not "idle": idle ends the watch, and a failed probe proves nothing. + logger.debug("auto-download: status probe failed for %r: %s", repo_id, exc) + return "unknown", None + + +async def _progress_percent( + repo_id: str, variant: Optional[str], expected_bytes: int, hf_token: Optional[str] +) -> Optional[float]: + """0-100, or None. The hub service reports a 0-1 fraction, so scale it.""" + from hub.services.models import downloads + try: + payload = await downloads.get_gguf_download_progress_response( + repo_id, variant or "", expected_bytes, hf_token + ) + fraction = payload.get("progress") + if not isinstance(fraction, (int, float)): + return None + return min(100.0, max(0.0, float(fraction) * 100.0)) + except Exception: + return None + + +def _release(active: Optional[_Active]) -> None: + """Free the single-flight slot, but only while *active* still owns it. + + Keying on ``repo_id`` alone let a stale operation clear a newer one: variant A + errors, an adopting request frees the slot, a retry starts B, then A's watcher + matches the repo and clears B, admitting a second download alongside it. + """ + global _active + if active is None: + return + with _lock: + if _active is active: + _active = None + + +async def _watch(active: _Active, hf_token: Optional[str]) -> None: + """Poll a dispatched job so the monitor row resolves and the resolver cache + is dropped the moment the weights land.""" + from core.inference import api_monitor as monitor_module + + api_monitor = monitor_module.api_monitor + deadline = time.monotonic() + _MAX_WATCH_S + timed_out = False + try: + while True: + await asyncio.sleep(_TIMED_OUT_POLL_S if timed_out else _WATCH_POLL_S) + state, error = await _job_state(active.repo_id, active.variant) + if state in ("running", "cancelling", "unknown"): + if timed_out: + # A running worker still owns the slot: releasing on the clock alone + # would admit a second multi-GB download beside it. "unknown" cannot + # confirm it is alive, so release then, or a broken probe wedges us. + if state == "unknown": + return + continue + if time.monotonic() >= deadline: + api_monitor.fail_open(active.monitor_id, "Download timed out") + timed_out = True + continue + # Only "running" has progress; the others are still in flight, so keep the slot. + if state == "running": + api_monitor.set_progress( + active.monitor_id, + await _progress_percent( + active.repo_id, active.variant, active.expected_bytes, hf_token + ), + ) + continue + if state == "cancelled": + api_monitor.finish(active.monitor_id, status = "cancelled") + return + if state == "complete": + # No invalidate here: finalize_worker_exit already dropped the cache and + # warmed it; a second would mark that fresh scan stale and push a + # synchronous rescan onto the client's retry. + api_monitor.finish(active.monitor_id, status = "completed") + elif state == "idle": + # The job vanished without a terminal state (worker killed). + api_monitor.fail_open(active.monitor_id, "Download did not complete") + else: + api_monitor.fail_open(active.monitor_id, error or f"Download {state}") + # Keep the slot so the next retry is told it failed instead of + # silently restarting the same download. + active.error = error or f"Download {state}" + active.failed_at = time.monotonic() + return + return + except asyncio.CancelledError: + raise + except Exception as exc: + logger.warning("auto-download: watcher failed for %r: %s", active.repo_id, exc) + api_monitor.fail_open(active.monitor_id, "Download tracking failed") + finally: + if not active.failed_at: + _release(active) + + +def _downloading_refusal(label: str, percent: Optional[float]) -> AutoDownloadRefusal: + progress = f" ({percent:.0f}% done)" if percent is not None else "" + return AutoDownloadRefusal( + status = 503, + code = "model_downloading", + message = (f"Downloading '{label}'{progress}. Retry shortly. Track it in Unsloth Studio."), + retry_after = _RETRY_AFTER_S, + ) + + +async def _is_downloadable_model(repo_id: str, hf_token: Optional[str]) -> bool: + """Whether the Hub has this repo with GGUF weights we could fetch. + + Only asked while another download holds the slot, to tell a second download + apart from an ordinary foreign label. Any failure answers False: refusing + would strand normal traffic for the length of the download. + """ + if _is_not_servable(repo_id, hf_token): + return False + + def _probe(): + from huggingface_hub import HfApi + return HfApi(token = _hub_token(hf_token)).model_info(repo_id, timeout = _MODEL_INFO_TIMEOUT_S) + + try: + info = await asyncio.to_thread(_probe) + except Exception: + return False + # The same filter admission uses, not a bare .gguf test: mmproj, MTP drafters and + # big-endian builds are companions, not quants. Answering otherwise would hold an + # ordinary foreign label at model_download_busy for an unrelated download. + servable = bool(_gguf_variants(getattr(info, "siblings", None))) + if not servable: + _mark_not_servable(repo_id, hf_token) + return servable + + +async def maybe_auto_download( + requested_model: str, + *, + hf_token: Optional[str] = None, + require_vision: bool = False, +) -> Optional[AutoDownloadRefusal]: + """Start (or report on) a background fetch of *requested_model*. + + Returns None when the request should carry on unchanged, or a refusal the + caller must raise. Only called after the local resolver has already missed. + + ``require_vision`` refuses a target with no mmproj companion rather than spend + gigabytes on weights that cannot answer the request; the local capability guard + only ever sees an already-downloaded model. + """ + global _active + + repo_id, wanted_variant = split_model_ref(requested_model) + if not is_downloadable_ref(requested_model): + return None + if _is_not_servable(repo_id, hf_token) and not looks_like_quant(wanted_variant): + return None + + # Settle the single-flight slot before the network, so retries during a download stay cheap. + busy: Optional[_Active] = None + with _lock: + current = _active + if current is not None and current.failed_at: + # A held failure only owns the slot until someone is told about it. + if current.repo_id != repo_id and time.monotonic() - current.failed_at > _FAILED_HOLD_S: + _active = current = None + if current is not None and current.repo_id == repo_id: + adopted = current + elif current is not None: + adopted = None + busy = current + else: + adopted = None + provisional = _Active(repo_id = repo_id, started_at = time.time()) + _active = provisional + + if busy is not None: + # Refusing before the probe blocks ordinary drop-in traffic: a namespaced label + # that is no downloadable GGUF repo (LiteLLM/OpenRouter style) would be told to + # wait out a multi-hour download. Only a downloadable label is a 2nd download. + if not await _is_downloadable_model(repo_id, hf_token): + return None + return AutoDownloadRefusal( + status = 503, + code = "model_download_busy", + message = ( + f"Already downloading '{_public_label(busy.repo_id, busy.variant)}'. " + f"Retry '{requested_model}' once it finishes." + ), + retry_after = _RETRY_AFTER_S, + ) + + if adopted is not None: + if adopted.variant is None: + # Still probing: no job yet, and a stale whole-repo error would free the probe's slot. + return _downloading_refusal(adopted.repo_id, None) + state, error = await _job_state(adopted.repo_id, adopted.variant) + if state in ("running", "cancelling", "unknown"): + return _downloading_refusal( + _public_label(adopted.repo_id, adopted.variant), + await _progress_percent( + adopted.repo_id, adopted.variant, adopted.expected_bytes, hf_token + ), + ) + if state == "error" or adopted.error: + error = error or adopted.error + # Surface once, then free the slot so a retry can start over. + _release(adopted) + return AutoDownloadRefusal( + status = 502, + code = "model_download_failed", + message = f"Downloading '{requested_model}' failed: {error or 'unknown error'}", + ) + # complete/idle/cancelled: the watcher is about to free the slot, so retry once more. + return _downloading_refusal( + _public_label(adopted.repo_id, adopted.variant), + 100.0 if state == "complete" else None, + ) + + try: + return await _admit_and_start( + repo_id, wanted_variant, requested_model, hf_token, provisional, require_vision + ) + except BaseException: + # Not `except Exception`: a cancel mid-probe would otherwise wedge the provisional slot. + _release(provisional) + raise + + +async def _admit_and_start( + repo_id: str, + wanted_variant: Optional[str], + requested_model: str, + hf_token: Optional[str], + active: _Active, + require_vision: bool = False, +) -> Optional[AutoDownloadRefusal]: + from hub.utils.hf_errors import hf_error_status + + def _probe(): + from huggingface_hub import HfApi + return HfApi(token = _hub_token(hf_token)).model_info( + repo_id, files_metadata = True, timeout = _MODEL_INFO_TIMEOUT_S + ) + + try: + info = await asyncio.to_thread(_probe) + except Exception as exc: + _release(active) + status = hf_error_status(exc) + if status == 401: + return AutoDownloadRefusal( + status = 401, + code = "model_access_denied", + message = ( + f"Hugging Face rejected the token sent for '{repo_id}'. Replace the " + "X-Unsloth-HF-Token header with a valid token; retrying will not help." + ), + ) + if status == 403: + return _gated_refusal(repo_id) + if status == 404: + _mark_not_servable(repo_id, hf_token) + # Unknown to the Hub reads as a foreign label; only an explicit quant makes it ours. + if not looks_like_quant(wanted_variant): + return None + # A private repo reads as absent without a token; don't confirm either way. + return AutoDownloadRefusal( + status = 404, + code = "model_not_found", + message = ( + f"'{repo_id}' was not found on Hugging Face, or is not accessible. " + "If it is private, send a token in the X-Unsloth-HF-Token header." + ), + ) + logger.warning("auto-download: Hub lookup failed for %r: %s", repo_id, exc) + return AutoDownloadRefusal( + status = 503, + code = "model_lookup_failed", + message = f"Could not reach Hugging Face to look up '{repo_id}'. Retry shortly.", + retry_after = _RETRY_AFTER_S, + ) + + # Inconclusive on timeout: the download's own auth is the real gate. + if getattr(info, "gated", False) and await _bounded_probe( + _auth_denied, repo_id, hf_token, timeout = _MODEL_INFO_TIMEOUT_S, default = False + ): + # Metadata for a gated repo is not file access; unchecked, the config read below lies. + _release(active) + return _gated_refusal(repo_id) + + variants = _gguf_variants(getattr(info, "siblings", None)) + if not variants: + _release(active) + _mark_not_servable(repo_id, hf_token) + if not looks_like_quant(wanted_variant): + return None + return AutoDownloadRefusal( + status = 400, + code = "model_not_supported", + message = ( + f"'{repo_id}' has no GGUF weights. Automatic download serves GGUF only; " + "load other formats from Unsloth Studio." + ), + ) + + # trust_remote_code gate: _config_has_auto_map is tri-state, so refuse on True and on None. + from utils.security.consent import _config_has_auto_map + + # _hub_token, not the raw token: None lets huggingface_hub fall back to a cached + # server login, so a caller-named repo would be probed with this server's identity. + # Defaults to None on timeout, which refuses: unchecked is not cleared. + has_auto_map = await _bounded_probe( + _config_has_auto_map, + repo_id, + _hub_token(hf_token), + timeout = _CODE_PROBE_TIMEOUT_S, + default = None, + ) + if has_auto_map is not False: + _release(active) + unknown = has_auto_map is None + return AutoDownloadRefusal( + status = 403, + code = "remote_code_consent_required", + message = ( + f"'{repo_id}' " + + ( + "could not be checked for custom code" + if unknown + else "ships custom code that runs on load" + ) + + ". Load it once in Unsloth Studio to review and approve it, then retry." + ), + ) + + variant = _match_variant(wanted_variant, variants) + if variant is None: + _release(active) + listed = sorted(variants) + shown = ", ".join(listed[:_MAX_LISTED_VARIANTS]) + extra = len(listed) - _MAX_LISTED_VARIANTS + return AutoDownloadRefusal( + status = 404, + code = "model_not_found", + message = ( + f"'{repo_id}' has no quant '{wanted_variant}'. Available quants: " + f"{shown}{f' and {extra} more' if extra > 0 else ''}." + ), + ) + + expected_bytes = variants[variant] + from hub.utils.gguf_plan import build_gguf_variant_plans + + plan = build_gguf_variant_plans(list(getattr(info, "siblings", None) or [])).get( + variant.lower() + ) + if require_vision and not (plan and plan.mmproj_filenames): + _release(active) + return AutoDownloadRefusal( + status = 400, + code = "invalid_value", + message = ( + f"'{_public_label(repo_id, variant)}' ships no mmproj companion, so it " + "cannot answer the image or audio input in this request. It was not " + "downloaded." + ), + ) + + need_bytes = _remaining_bytes(repo_id, plan, expected_bytes) + fits, free = _enough_disk(need_bytes) + if not fits: + _release(active) + return AutoDownloadRefusal( + status = 507, + code = "insufficient_disk_space", + message = ( + f"'{_public_label(repo_id, variant)}' needs {_gb(need_bytes)} plus " + f"{_gb(_DISK_RESERVE_BYTES)} headroom, but only {_gb(free)} is free." + ), + ) + + return await _dispatch(repo_id, variant, expected_bytes, requested_model, hf_token, active) + + +def preferred_quant(labels) -> Optional[str]: + """The quant a plain load would pick from *labels*, or None. + + The one ranking for "which quant did they mean": local resolution, remote + admission and /v1/models must agree, or a bare id means a different quant + depending on which of them answered it. + """ + from utils.models.model_config import _pick_best_gguf + + # _pick_best_gguf ranks filenames and matches upper-case tokens, so feed "