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..2007789035 --- /dev/null +++ b/.github/scripts/run-studio-permission-browser.sh @@ -0,0 +1,69 @@ +#!/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" +unsloth studio reset-password +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/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..c48328e90f 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. @@ -256,7 +256,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 +359,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. @@ -448,7 +448,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 +543,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 }} @@ -620,7 +620,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 +706,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. @@ -764,7 +764,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 4daafae35d..081eda4e32 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: @@ -69,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*)' @@ -146,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 }} @@ -211,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 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/studio-api-smoke.yml b/.github/workflows/studio-api-smoke.yml index 15efee382e..cdf1f6bf12 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,7 +111,7 @@ 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 mkdir -p logs @@ -144,7 +144,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 +153,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..b8f587b63e 100644 --- a/.github/workflows/studio-backend-ci.yml +++ b/.github/workflows/studio-backend-ci.yml @@ -64,7 +64,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 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..3a9e373915 100644 --- a/.github/workflows/studio-frontend-ci.yml +++ b/.github/workflows/studio-frontend-ci.yml @@ -136,7 +136,7 @@ jobs: - 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 +144,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 58ef2558f3..c2d52eac22 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,7 +125,7 @@ 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 mkdir -p logs @@ -142,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 @@ -229,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. @@ -276,7 +276,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 +290,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 +323,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 +380,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 +390,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 @@ -503,7 +503,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 +575,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 +698,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. """ @@ -811,7 +811,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 @@ -834,7 +834,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): @@ -848,7 +848,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 "") @@ -868,7 +868,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 @@ -960,7 +960,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. @@ -973,7 +973,7 @@ 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. @@ -1076,13 +1076,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": [ @@ -1112,7 +1112,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 @@ -1148,9 +1148,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( @@ -1184,7 +1184,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..1968885a1d 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,7 +99,7 @@ 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 mkdir -p logs @@ -129,13 +129,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 946681706a..ce15eed5c8 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,7 +124,7 @@ 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 mkdir -p logs @@ -141,7 +141,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 +228,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 +283,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 +363,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 +376,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 @@ -478,7 +478,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 +574,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] @@ -660,7 +660,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): @@ -678,7 +678,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 "") @@ -704,7 +704,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 @@ -810,7 +810,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. @@ -826,7 +826,7 @@ 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. @@ -929,13 +929,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": [ @@ -1007,7 +1007,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 @@ -1023,11 +1023,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: @@ -1053,14 +1053,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( @@ -1099,11 +1099,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..7375e9bcbf 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,7 +144,7 @@ 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 mkdir -p logs @@ -188,7 +189,7 @@ 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 + # guard redirects mid-navigation. The retry FULLY resets Unsloth # (kill, reset-password, 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 @@ -209,7 +210,7 @@ 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 @@ -238,13 +239,17 @@ 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: | + bash .github/scripts/run-studio-permission-browser.sh 18895 webkit + + - name: Reset auth + boot Unsloth for extra UI tests (port 18897) run: | unsloth studio reset-password mkdir -p logs @@ -271,7 +276,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,7 +305,7 @@ 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 @@ -327,7 +332,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 +348,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..8e26b9fd0c 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: diff --git a/.github/workflows/studio-ui-smoke.yml b/.github/workflows/studio-ui-smoke.yml index 297a585430..0ad55ebd6d 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,15 +108,12 @@ 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 mkdir -p logs @@ -147,7 +145,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,29 +163,35 @@ 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 mkdir -p logs @@ -214,7 +218,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,16 +231,64 @@ 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 sleep 2 + # 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 + 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 Studio on its own port so a hang here cannot poison the + # 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 Studio for IME / i18n tests (port 18896) + - name: Reset auth + boot Unsloth for IME / i18n tests (port 18896) run: | unsloth studio reset-password mkdir -p logs @@ -256,7 +308,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 +325,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 +345,14 @@ 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_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..625c2c7811 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,13 @@ 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: 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..6dbcceebbd 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,7 +177,7 @@ 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 mkdir -p logs @@ -207,7 +207,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 +219,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 63a7e9dc8f..3ebe442f52 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,7 +227,7 @@ 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 mkdir -p logs @@ -244,7 +244,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 +281,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 +382,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 +398,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 +439,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 +507,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 +523,7 @@ jobs: } } - - name: Install Studio (--local, --no-torch) + - name: Install Unsloth (--local, --no-torch) shell: pwsh env: GH_TOKEN: ${{ secrets.GITHUB_TOKEN }} @@ -561,7 +561,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,7 +571,7 @@ 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 mkdir -p logs @@ -607,7 +607,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 +680,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 @@ -882,15 +882,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() @@ -898,10 +898,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 @@ -939,7 +939,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' @@ -1005,7 +1005,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", @@ -1021,7 +1021,7 @@ jobs: } } - - name: Install Studio (--local, --no-torch) + - name: Install Unsloth (--local, --no-torch) shell: pwsh env: GH_TOKEN: ${{ secrets.GITHUB_TOKEN }} @@ -1059,7 +1059,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 @@ -1072,7 +1072,7 @@ 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 mkdir -p logs @@ -1262,7 +1262,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." ) @@ -1303,19 +1303,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() @@ -1323,10 +1323,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 @@ -1348,7 +1348,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: @@ -1502,7 +1502,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 }} @@ -1538,13 +1538,13 @@ 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 mkdir -p logs @@ -1613,10 +1613,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() diff --git a/.github/workflows/studio-windows-ui-smoke.yml b/.github/workflows/studio-windows-ui-smoke.yml index 405309916a..f401f7be44 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,7 +295,7 @@ 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 mkdir -p logs @@ -339,13 +340,17 @@ 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: | + bash .github/scripts/run-studio-permission-browser.sh 18895 chromium msedge + + - name: Reset auth + boot Unsloth for extra UI tests (port 18897) run: | unsloth studio reset-password mkdir -p logs @@ -372,7 +377,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 +391,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 +407,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..42d74d47d2 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 @@ -212,7 +212,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 +239,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..cdad617027 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,7 @@ 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: Upload wheel on failure if: failure() diff --git a/README.md b/README.md index ef45b91430..6aa8f4f4c3 100644 --- a/README.md +++ b/README.md @@ -11,6 +11,7 @@ Unsloth Studio lets you run and train models locally.

Features • + News • Quickstart • Notebooks • Documentation @@ -47,15 +48,44 @@ 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` | + ## 📥 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 +95,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). * **Multi-GPU:** Available now, with a major upgrade on the way #### macOS, Linux, WSL: @@ -86,7 +117,7 @@ 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 +153,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 +179,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 +246,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 +256,7 @@ 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. +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 +268,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 +281,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 +317,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/build.sh b/build.sh index dc272f0de1..2a836e19d9 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" diff --git a/install.ps1 b/install.ps1 index c25d7e7b7f..cceefd2647 100644 --- a/install.ps1 +++ b/install.ps1 @@ -53,7 +53,8 @@ function Install-UnslothStudio { 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 +63,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" @@ -176,7 +178,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,22 +469,35 @@ 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 ) - # 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" @@ -493,17 +508,23 @@ function Install-UnslothStudio { # 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 } 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] } } + } } } @@ -756,7 +777,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 +793,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 @@ -1438,7 +1459,7 @@ exit 0 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 +1470,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 +1489,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 +1519,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 @@ -1517,7 +1538,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 +1547,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 +1674,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 +1684,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) { @@ -1942,7 +1963,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 +1981,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 +2026,25 @@ 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) + $sep = $Url.IndexOf('://') + 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('/') + $authority = if ($slash -ge 0) { $rest.Substring(0, $slash) } else { $rest } + $at = $authority.LastIndexOf('@') + $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 +2063,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 +2104,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 +2119,19 @@ 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) "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 @@ -2102,6 +2175,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') -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 +2263,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.3" "unsloth-zoo>=2026.7.3" } + $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.4" "unsloth-zoo>=2026.7.4" } if ($baseInstallExit -eq 0) { # Resolve pydantic WITH deps so pip pins pydantic-core # to the matching version (no-torch-runtime.txt below @@ -2185,7 +2284,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.3" "unsloth-zoo>=2026.7.3" } + $baseInstallExit = Invoke-InstallCommandRetry -Label "install unsloth (migrated)" { uv pip install --python $VenvPython --reinstall-package unsloth --reinstall-package unsloth-zoo "unsloth>=2026.7.4" "unsloth-zoo>=2026.7.4" } } if ($baseInstallExit -ne 0) { Write-Host "[ERROR] Failed to install unsloth (exit code $baseInstallExit)" -ForegroundColor Red @@ -2210,22 +2309,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 +2339,14 @@ 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 } + 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" + $torchInstallExit = Invoke-InstallCommandRetry -Label "install PyTorch" { uv pip install --python $VenvPython "torch>=2.4,<2.11.0" $_pinVisionSpec $_pinAudioSpec --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 +2358,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.3" "unsloth-zoo>=2026.7.3" } + $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.4" "unsloth-zoo>=2026.7.4" } 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 +2370,7 @@ exit 0 } } } elseif ($StudioLocalInstall) { - $baseInstallExit = Invoke-InstallCommandRetry -Label "install unsloth (local)" { uv pip install --python $VenvPython --upgrade-package unsloth "unsloth>=2026.7.3" "unsloth-zoo>=2026.7.3" } + $baseInstallExit = Invoke-InstallCommandRetry -Label "install unsloth (local)" { uv pip install --python $VenvPython --upgrade-package unsloth "unsloth>=2026.7.4" "unsloth-zoo>=2026.7.4" } } else { $baseInstallExit = Invoke-InstallCommandRetry -Label "install unsloth" { uv pip install --python $VenvPython --upgrade-package unsloth -- "$PackageName" } } @@ -2291,7 +2398,7 @@ 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.3" "unsloth>=2026.7.3" --torch-backend=auto } + $baseInstallExit = Invoke-InstallCommandRetry -Label "install unsloth (auto torch backend)" { uv pip install --python $VenvPython "unsloth-zoo>=2026.7.4" "unsloth>=2026.7.4" --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) @@ -2317,6 +2424,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,8 +2449,8 @@ 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 } if ($torchFixExit -ne 0) { @@ -2347,7 +2461,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 { 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 +2536,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") } @@ -2533,7 +2647,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 +2665,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 @@ -2616,7 +2730,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 4a3c4471fa..445bab616a 100755 --- a/install.sh +++ b/install.sh @@ -97,7 +97,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,18 +159,58 @@ 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) + { "$@" 2>&1; printf '%s' "$?" > "$_rcf"; } | _redact_install_output + _rc=$(cat "$_rcf" 2>/dev/null || echo 1) + rm -f "$_rcf" + [ "${_rc:-1}" -eq 0 ] 2>/dev/null && return 0 step "error" "$_label failed (exit code $_rc)" "$C_ERR" >&2 return "$_rc" fi @@ -178,7 +218,7 @@ run_install_cmd() { "$@" >"$_log" 2>&1 && { rm -f "$_log"; return 0; } _rc=$? step "error" "$_label failed (exit code $_rc)" "$C_ERR" >&2 - cat "$_log" >&2 + _redact_install_output "$_log" >&2 rm -f "$_log" return $_rc } @@ -257,7 +297,7 @@ _install_bnb_rocm() { 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 @@ -310,6 +350,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 +388,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" @@ -472,11 +518,13 @@ _on_install_exit() { _restore_studio_venv_replacement fi [ -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 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. +# 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 # ── Helper: download a URL to a file (supports curl and wget) ── @@ -663,7 +711,7 @@ POLL_INTERVAL_SEC=0.25 LOG_FILE="$DATA_DIR/studio.log" # why: in env-override mode multiple installs share an OS user; namespace the # lock and remember our own healthy port so we never attach to an unrelated -# Studio listening on the global 8888..8908 range. +# Unsloth listening on the global 8888..8908 range. LOCK_DIR="${XDG_RUNTIME_DIR:-/tmp}/unsloth-studio-launcher-$(id -u).lock" PORT_FILE="" # why: gate on the install-time mode (baked above) instead of the runtime env @@ -734,7 +782,7 @@ _candidate_ports() { _find_healthy_port() { if [ -n "$PORT_FILE" ] && [ -f "$PORT_FILE" ]; then # why: env-mode installs only attach to a port we previously launched - # ourselves; never to a sibling Studio that happens to be healthy. + # ourselves; never to a sibling Unsloth that happens to be healthy. _p=$(cat "$PORT_FILE" 2>/dev/null || true) case "$_p" in ''|*[!0-9]*) ;; @@ -901,7 +949,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 +1419,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 +1487,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 +1621,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 @@ -1634,6 +1688,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,7 +1729,7 @@ _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 +# cmake/git are only needed to *build* llama.cpp from source. Unsloth 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. @@ -1821,11 +1879,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 +1898,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 +1912,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 +1969,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 +2057,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)" @@ -2059,6 +2134,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. @@ -2187,16 +2280,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 +2438,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 +2595,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 +2640,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 @@ -2402,7 +2662,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,7 +2719,19 @@ _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) @@ -2470,24 +2742,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) + 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 && \ @@ -2564,10 +2886,31 @@ 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 ;; 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" @@ -2697,7 +3040,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 +3048,46 @@ 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. substep "upgrading unsloth in migrated environment..." if [ "$SKIP_TORCH" = true ]; then # No-torch: install unsloth + unsloth-zoo with --no-deps (current @@ -2716,7 +3096,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.3" "unsloth-zoo>=2026.7.3" + "unsloth>=2026.7.4" "unsloth-zoo>=2026.7.4" # 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 +3109,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.3" "unsloth-zoo>=2026.7.3" ${_MLX_LM_EXCLUDE_ARG:-} + "unsloth>=2026.7.4" "unsloth-zoo>=2026.7.4" ${_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 +3128,14 @@ 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 + _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)..." + _install_torch_default_index --force-reinstall + fi fi elif [ -n "$TORCH_INDEX_URL" ]; then # Fresh: Step 1 - install torch from explicit index (skip when --no-torch or Intel Mac) @@ -2820,7 +3197,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 +3289,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 +3310,34 @@ 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 + _install_bnb_rocm "install bitsandbytes (AMD)" "$_VENV_PY" fi - # Fresh: Step 2 - install unsloth, preserving pre-installed torch + # 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.3" "unsloth-zoo>=2026.7.3" + "unsloth>=2026.7.4" "unsloth-zoo>=2026.7.4" # 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 +3355,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.3" "unsloth-zoo>=2026.7.3" + ${_UNSLOTH_TORCH_OVERRIDES:+--overrides "$_UNSLOTH_TORCH_OVERRIDES"} \ + --upgrade-package unsloth "unsloth>=2026.7.4" "unsloth-zoo>=2026.7.4" 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 +3365,26 @@ 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 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.3" "unsloth>=2026.7.3" --torch-backend=auto + run_install_cmd_retry "install unsloth (auto torch backend)" uv pip install --python "$_VENV_PY" "unsloth-zoo>=2026.7.4" "unsloth>=2026.7.4" --torch-backend=auto substep "overlaying local repo (editable)..." run_install_cmd "overlay local repo" uv pip install --python "$_VENV_PY" -e "$_REPO_ROOT" --no-deps substep "overlaying unsloth-zoo from git main..." @@ -2997,6 +3396,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 +3422,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 +3433,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="" @@ -3227,7 +3633,7 @@ printf " ${C_TITLE}%s${C_RST}\n" "Unsloth Studio installed!" printf " ${C_DIM}%s${C_RST}\n" "$RULE" echo "" -# 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 (Docker, CI, cloud-init) just print instructions. if [ "$_SKIP_AUTOSTART" != true ] && [ -t 1 ]; then diff --git a/pyproject.toml b/pyproject.toml index 7b8fdd100d..071258eb8f 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -74,7 +74,7 @@ triton = [ ] huggingfacenotorch = [ - "unsloth_zoo>=2026.7.3", + "unsloth_zoo>=2026.7.4", "wheel>=0.42.0", "packaging", "numpy", @@ -95,7 +95,7 @@ huggingfacenotorch = [ ] huggingface = [ "unsloth[huggingfacenotorch]", - "unsloth_zoo>=2026.7.3", + "unsloth_zoo>=2026.7.4", "torchvision", "unsloth[triton]", ] @@ -580,7 +580,7 @@ colab-ampere-torch220 = [ "flash-attn>=2.6.3 ; ('linux' in sys_platform)", ] colab-new = [ - "unsloth_zoo>=2026.7.3", + "unsloth_zoo>=2026.7.4", "packaging", "tyro", "transformers>=4.51.3,!=4.52.0,!=4.52.1,!=4.52.2,!=4.52.3,!=4.53.0,!=4.54.0,!=4.55.0,!=4.55.1,!=4.57.0,!=4.57.4,!=4.57.5,!=5.0.0,!=5.1.0,<=5.5.0", diff --git a/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 < 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 index 91b39fcc69..127a85a116 100644 --- a/studio/MCP.md +++ b/studio/MCP.md @@ -1,10 +1,10 @@ # Unsloth Studio MCP server -Studio can expose a local MCP server so an MCP client can inspect models and +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 Studio process with: +The server is disabled by default. Enable it for a local Unsloth process with: ```bash UNSLOTH_STUDIO_ENABLE_MCP=1 \ @@ -12,8 +12,8 @@ UNSLOTH_STUDIO_MCP_TOKEN='use-a-local-secret' \ unsloth studio ``` -The endpoint is `http://127.0.0.1:8888/mcp/` when Studio uses its default port -(a request to `/mcp` redirects to the canonical `/mcp/`). Use the actual 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: @@ -23,9 +23,9 @@ The high-impact tools are: - `validate_recipe`, `get_recipe_job_status`, and `get_recipe_job_dataset` - `load_checkpoint` and `export_gguf` -`start_training` accepts the same fields as the Studio `TrainingStartRequest`. +`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 Studio validation as well. +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 diff --git a/studio/Unsloth_Studio_Colab.ipynb b/studio/Unsloth_Studio_Colab.ipynb index 619395bd6d..44282b2255 100644 --- a/studio/Unsloth_Studio_Colab.ipynb +++ b/studio/Unsloth_Studio_Colab.ipynb @@ -33,7 +33,7 @@ "\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)" + "[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)" ] }, { 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/auth/authentication.py b/studio/backend/auth/authentication.py index b13cd1c851..dfb8fc513e 100644 --- a/studio/backend/auth/authentication.py +++ b/studio/backend/auth/authentication.py @@ -148,7 +148,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)) 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..39fa691304 100644 --- a/studio/backend/auth/storage.py +++ b/studio/backend/auth/storage.py @@ -146,7 +146,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 @@ -305,8 +305,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 diff --git a/studio/backend/auth/terminal_prompt.py b/studio/backend/auth/terminal_prompt.py index 8491019ae9..e855f4078b 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 @@ -252,7 +252,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..b1ddc74c32 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. """ @@ -95,7 +95,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 @@ -309,7 +309,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() diff --git a/studio/backend/colab.py b/studio/backend/colab.py index e04543b3aa..1762469bcf 100644 --- a/studio/backend/colab.py +++ b/studio/backend/colab.py @@ -129,7 +129,7 @@ def start_cloudflare_tunnel(port: int) -> "str | None": 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 @@ -203,7 +203,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! 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..ffc81669ae 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()) 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/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/anthropic_compat.py b/studio/backend/core/inference/anthropic_compat.py index 3c7a4cb182..34445cc58e 100644 --- a/studio/backend/core/inference/anthropic_compat.py +++ b/studio/backend/core/inference/anthropic_compat.py @@ -539,7 +539,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/chat_template_helpers.py b/studio/backend/core/inference/chat_template_helpers.py index 5113eebb36..528c059fbc 100644 --- a/studio/backend/core/inference/chat_template_helpers.py +++ b/studio/backend/core/inference/chat_template_helpers.py @@ -150,7 +150,7 @@ def _split_partial_marker(text: str, marker: str) -> tuple[str, str]: class ReasoningChannelNormalizer: """Incrementally convert one native reasoning channel to ````. - The parser follows mlx-vlm's streaming boundary behavior but emits Studio's + 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. 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/llama_cpp.py b/studio/backend/core/inference/llama_cpp.py index fec5214cf9..2c7433f7a4 100644 --- a/studio/backend/core/inference/llama_cpp.py +++ b/studio/backend/core/inference/llama_cpp.py @@ -11,6 +11,7 @@ import atexit import contextlib import functools import json +import math import os import re import struct @@ -22,12 +23,24 @@ 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, + Optional, + Union, +) import httpx from core.inference.llama_server_args import ( + _LAYER_OFFLOAD_FLAGS, _effective_tensor_parallel, _tensor_parallel_matches_loaded, extra_args_disable_mmproj, @@ -116,7 +129,7 @@ LLAMA_SERVER_NOT_FOUND_DETAIL = ( # 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 = ( @@ -234,8 +247,7 @@ 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. +# 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. @@ -442,6 +454,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" @@ -1363,7 +1392,7 @@ def _kv_bytes_per_elem(cache_type: Optional[str]) -> float: 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 @@ -1431,7 +1460,10 @@ 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"}) @@ -1682,7 +1714,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). @@ -1731,7 +1763,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 @@ -1804,7 +1836,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 @@ -1895,6 +1927,17 @@ 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 # 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 @@ -1909,6 +1952,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 @@ -1960,7 +2008,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] = [] @@ -1970,6 +2018,12 @@ 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 # True once a probe has completed; cleared on transient failure. self._is_audio: bool = False self._audio_type: Optional[str] = None @@ -2329,6 +2383,79 @@ 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 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.""" @@ -2353,7 +2480,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 @@ -2530,10 +2657,12 @@ class LlamaCppBackend: "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) @@ -2549,10 +2678,12 @@ 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 try: probe_env = cls._llama_server_env_for_binary(bin_path) result = subprocess.run( @@ -2646,10 +2777,12 @@ 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}") @@ -2662,10 +2795,12 @@ class LlamaCppBackend: "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 @@ -2746,6 +2881,57 @@ class LlamaCppBackend: except ValueError: return None + @staticmethod + def _emit_child_gpu_visibility(env: dict, pinned: str) -> None: + """Write the child's GPU visibility mask (CUDA, plus the HIP mirror on + ROCm, where narrowing only CUDA_VISIBLE_DEVICES leaves an AMD child + seeing the full set). Do NOT also set ROCR_VISIBLE_DEVICES: ROCR and HIP + mask at different layers, so the same indices apply twice -- ROCR reduces + and re-indexes from 0, then a non-zero HIP pin points out of range, HIP + enumerates 0 devices, and llama.cpp falls back to CPU. The HIP mask alone + narrows correctly; clear any inherited ROCR mask so it can't double up.""" + env["CUDA_VISIBLE_DEVICES"] = pinned + try: + import torch as _torch + if getattr(_torch.version, "hip", None) is not None: + 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) + LlamaCppBackend._emit_child_gpu_visibility(env, ",".join(str(i) for i in order)) + @staticmethod def _amd_apu_wants_unified_memory(gpu_indices = None) -> bool: """True only for AMD unified-memory APUs (gfx1150/gfx1151), where @@ -3241,7 +3427,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). @@ -3262,6 +3448,20 @@ class LlamaCppBackend: # aborts a --split-mode tensor load, so it's dropped for the tensor attempt. _TENSOR_PARALLEL_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. @@ -3426,6 +3626,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, @@ -3800,7 +4006,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 = 512 # llama.cpp --ubatch default; Unsloth does not override it _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) @@ -3940,7 +4146,7 @@ class LlamaCppBackend: n_ubatch: Optional[int] = None, ) -> 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 @@ -4246,6 +4452,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 @@ -4335,6 +4543,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", @@ -4416,7 +4626,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 @@ -4523,6 +4733,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, *, @@ -4533,10 +4765,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: @@ -4558,7 +4791,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", @@ -4586,6 +4823,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: @@ -4608,7 +4850,7 @@ class LlamaCppBackend: 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, @@ -4631,6 +4873,23 @@ class LlamaCppBackend: self._model_identifier = model_identifier self._cache_type_kv = 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 + # Record only the single device the runner actually uses (the lowest + # selected GPU, chosen above) -- not the whole pick. The diffusion runner + # is single-device, so echoing a multi-GPU list would misreport placement + # in /status and let a re-Apply dedup against GPUs the runner never used. + self._gpu_ids = [sorted(gpu_ids)[0]] if gpu_ids else None if hf_variant: self._hf_variant = hf_variant elif gguf_path: @@ -5242,7 +5501,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: @@ -5721,6 +5980,11 @@ 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, + gpu_ids: Optional[List[int]] = None, n_threads: Optional[int] = None, n_gpu_layers: Optional[int] = None, # caller compat, unused n_parallel: int = 1, @@ -5753,6 +6017,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, @@ -5779,6 +6051,11 @@ 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, @@ -5899,6 +6176,7 @@ class LlamaCppBackend: model_identifier = model_identifier, n_ctx = n_ctx, extra_args = extra_args, + gpu_ids = gpu_ids, ) if not binary: @@ -5960,6 +6238,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 @@ -6040,6 +6371,12 @@ 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). @@ -6057,6 +6394,18 @@ class LlamaCppBackend: _gpu_mem = self._get_gpu_memory(binary) gpus = [(idx, free) for idx, free, _t in _gpu_mem] 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. @@ -6088,6 +6437,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 @@ -6103,7 +6490,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. @@ -6112,7 +6499,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; @@ -6175,10 +6562,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 = {}) @@ -6313,7 +6703,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 @@ -6388,7 +6779,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. @@ -6407,7 +6803,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 @@ -6823,6 +7219,12 @@ class LlamaCppBackend: tp_tensor_split = None effective_ctx = requested_ctx # fall back to original + # 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 @@ -6859,8 +7261,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", @@ -6868,6 +7268,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 @@ -6879,7 +7290,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 @@ -6893,10 +7360,31 @@ class LlamaCppBackend: # 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, @@ -6960,9 +7448,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: @@ -6994,7 +7484,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, ) @@ -7073,8 +7563,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"): @@ -7085,10 +7576,6 @@ 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", @@ -7097,12 +7584,12 @@ class LlamaCppBackend: # 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. + # last-wins parsing lets a user --device override Unsloth's pick. if is_vulkan_backend and gpu_indices is not None: cmd += LlamaCppBackend._vulkan_pin_args(gpu_indices) # 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) @@ -7112,15 +7599,17 @@ class LlamaCppBackend: # 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 @@ -7130,7 +7619,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 @@ -7170,28 +7659,39 @@ class LlamaCppBackend: # 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) + # 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" + self._emit_child_gpu_visibility(env, ",".join(str(i) for i in gpu_indices)) + 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 @@ -7201,7 +7701,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 @@ -7284,7 +7784,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. @@ -7296,7 +7796,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. @@ -7313,7 +7813,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, @@ -7350,7 +7850,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 ) @@ -7393,7 +7892,7 @@ class LlamaCppBackend: 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] = ( @@ -7535,6 +8034,10 @@ class LlamaCppBackend: "session; run 'unsloth studio update' to enable vision." ) 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) @@ -7566,6 +8069,13 @@ class LlamaCppBackend: self._healthy = True self._commit_effective_parallel_slots(n_parallel) + # 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() + # Commit caller intent only after _healthy=True so a failed start # can't poison the next inheritance check. None keeps prior, [] # clears, list sets. Source records hf_variant for the route's @@ -7580,18 +8090,31 @@ class LlamaCppBackend: 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).", ) @@ -7627,6 +8150,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( @@ -7888,7 +8420,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() @@ -7916,7 +8448,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() @@ -7947,6 +8479,11 @@ 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, preserve_multi_gpu_on_layer: bool = False, ) -> bool: @@ -8003,6 +8540,38 @@ 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: + # A GPU-memory-mode flip (Unsloth / manual) must always reload. + if self._gpu_memory_mode != gpu_memory_mode: + 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 (compare order-insensitively; None/[] + # both mean automatic). The diffusion runner collapses a multi-GPU pick + # to its single lowest device, so self._gpu_ids holds just that device; + # normalize the request the same way, or a multi-GPU pick that resolves + # to the same device needlessly reloads. + if self._is_diffusion: + requested_gpu_pick = [sorted(gpu_ids)[0]] if gpu_ids else None + else: + requested_gpu_pick = sorted(gpu_ids) if gpu_ids else None + if (self._gpu_ids or None) != requested_gpu_pick: + 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): @@ -8071,6 +8640,78 @@ 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, 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 @@ -8103,6 +8744,10 @@ 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._chat_template = None self._chat_template_override = None self._supports_reasoning = False @@ -8114,11 +8759,18 @@ class LlamaCppBackend: self._supports_tools = False self._cache_type_kv = None self._tensor_parallel = False + self._gpu_memory_mode = "auto" + self._gpu_layers = -1 + self._n_cpu_moe = 0 + self._tensor_split = None + self._gpu_ids = 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 @@ -8181,6 +8833,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() @@ -8294,7 +8950,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: @@ -8334,9 +8990,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. @@ -8393,7 +9049,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. @@ -8413,7 +9069,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). @@ -8423,7 +9079,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") @@ -8497,7 +9153,7 @@ class LlamaCppBackend: if not is_ours: continue - # A live parent means a running Studio (or the user's + # 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 @@ -8577,7 +9233,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 @@ -8618,6 +9274,237 @@ 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, + getattr(self, "_cache_type_kv", None), + self.effective_parallel_slots, + ) + + 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, 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 {"off", "disabled", "false", "0"} + + 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 + 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._effective_context_length or self._context_length or 0, + self._cache_type_kv, + n_parallel = self.effective_parallel_slots, + ) + # 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. @@ -8785,7 +9672,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. @@ -8793,14 +9685,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]: @@ -8821,7 +9727,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 @@ -8997,6 +9903,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( @@ -9054,6 +10029,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 diff --git a/studio/backend/core/inference/llama_keepwarm.py b/studio/backend/core/inference/llama_keepwarm.py index 4ce663c3ce..3380ebf5f5 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: @@ -266,7 +347,10 @@ def _loaded_identity(backend): 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,17 +365,47 @@ 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) seen_model = None except Exception as exc: diff --git a/studio/backend/core/inference/llama_server_args.py b/studio/backend/core/inference/llama_server_args.py index f400d2ae40..7b42d2f40d 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 """ @@ -22,12 +22,12 @@ _DENYLIST_GROUPS: tuple[frozenset[str], ...] = ( # Parallel slots: owned by typer --parallel; a pass-through would desync # app.state.llama_parallel_slots 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 +39,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 +64,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) @@ -120,7 +122,7 @@ def validate_extra_args(args: Optional[Iterable[str]]) -> list[str]: 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 +144,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 +181,38 @@ _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). +_LAYER_OFFLOAD_FLAGS: frozenset[str] = frozenset( + {"-ngl", "--gpu-layers", "--n-gpu-layers", "-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,7 +301,7 @@ 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) @@ -341,7 +356,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 +439,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 +467,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..64ab38ec75 100644 --- a/studio/backend/core/inference/local_model_resolver.py +++ b/studio/backend/core/inference/local_model_resolver.py @@ -130,7 +130,7 @@ def info_has_local_gguf(info) -> bool: 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 @@ -199,9 +199,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) diff --git a/studio/backend/core/inference/mcp_client.py b/studio/backend/core/inference/mcp_client.py index 6b5ce02216..0256df944e 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. diff --git a/studio/backend/core/inference/mlx_inference.py b/studio/backend/core/inference/mlx_inference.py index e7a90b4307..e78c93b6f3 100644 --- a/studio/backend/core/inference/mlx_inference.py +++ b/studio/backend/core/inference/mlx_inference.py @@ -8,6 +8,7 @@ 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 @@ -20,6 +21,63 @@ 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.""" @@ -508,6 +566,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") @@ -552,6 +611,7 @@ class MLXInferenceBackend: reasoning_effort = reasoning_effort, preserve_thinking = preserve_thinking, presence_penalty = presence_penalty, + _adapter_state = _adapter_state, ) else: stream = self._generate_text( @@ -568,6 +628,7 @@ class MLXInferenceBackend: reasoning_effort = reasoning_effort, preserve_thinking = preserve_thinking, presence_penalty = presence_penalty, + _adapter_state = _adapter_state, ) yield from stream @@ -587,6 +648,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 @@ -635,10 +697,6 @@ class MLXInferenceBackend: 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, @@ -680,9 +738,12 @@ class MLXInferenceBackend: type(self._model).__name__, type(self._tokenizer).__name__, ) - with self._generation_lock: + with self._generation_lock, _temporary_mlx_adapter_state(self._model, _adapter_state): final_response = None try: + # Enter request-scoped model state before yielding any response. + if think_prefix: + yield think_prefix gen_kwargs = dict( prompt = prompt, max_tokens = max_new_tokens, @@ -749,6 +810,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 @@ -852,9 +914,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), @@ -891,9 +950,18 @@ class MLXInferenceBackend: def _stream_vlm_snapshots(): nonlocal cumulative - with self._generation_lock: + # 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, @@ -927,8 +995,11 @@ 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): import mlx.core as mx diff --git a/studio/backend/core/inference/orchestrator.py b/studio/backend/core/inference/orchestrator.py index 3afda74411..75ef9c2399 100644 --- a/studio/backend/core/inference/orchestrator.py +++ b/studio/backend/core/inference/orchestrator.py @@ -54,9 +54,8 @@ class GenStreamError(str): """A stream chunk carrying a real backend/generation error, not model text. Subclasses str so existing display/logging consumers are unaffected, while - callers that must abort a distributed run on error (raise_on_streamed_error) - can distinguish a real error from model output whose visible text starts with - "Error:" by checking isinstance(chunk, GenStreamError). + callers can distinguish a real error from model output whose visible text + starts with "Error:" by checking isinstance(chunk, GenStreamError). """ __slots__ = ("public",) @@ -1502,14 +1501,27 @@ class InferenceOrchestrator: Uses the dispatcher path (no _gen_lock) so compare-mode requests don't block each other; the subprocess serializes them via its - sequential command loop. + sequential command loop. Backend failures raise instead of becoming + assistant text. """ - yield from self._generate_dispatched( + stream = self._generate_dispatched( use_adapter = use_adapter, cancel_event = cancel_event, stats_holder = stats_holder, **gen_kwargs, ) + try: + for chunk in stream: + if isinstance(chunk, GenStreamError): + # Preserve the public/operational flag so the route can surface + # the real message (e.g. "model is being unloaded") instead of a + # generic error. Mirrors the safetensors tool loop's _single_turn. + raise GenStreamErrorRaised(str(chunk), public = chunk.public) + yield chunk + finally: + close = getattr(stream, "close", None) + if callable(close): + close() def _generate_inner( self, diff --git a/studio/backend/core/inference/passthrough_healing.py b/studio/backend/core/inference/passthrough_healing.py index ed7c7ecfcf..e6da0a22b0 100644 --- a/studio/backend/core/inference/passthrough_healing.py +++ b/studio/backend/core/inference/passthrough_healing.py @@ -5,7 +5,7 @@ With server-side tools disabled (``unsloth run --disable-tools``, every ``unsloth start`` coding agent), requests carrying the client's own ``tools`` -bypass Studio's tool loop and are relayed to/from llama-server verbatim. Small +bypass Unsloth's tool loop and are relayed to/from llama-server verbatim. Small GGUF models often emit their tool calls as TEXT (``{...}``, Gemma ``<|tool_call>...``, ```` XML) instead of structured ``tool_calls`` -- on the passthrough that text reaches the agent as prose and @@ -18,7 +18,7 @@ promotes calls whose function name exactly matches a declared tool. Promotion removes EXACTLY the promoted calls' markup spans (the parser reports them): undeclared calls, unparseable blocks, and suppressed alternate formats keep every byte and relay as text, so healing can never silently delete model -output. Responses without a tool signal, requests without tools, and Studio's +output. Responses without a tool signal, requests without tools, and Unsloth's own enable-tools loop are untouched. Per-request opt-out: ``auto_heal_tool_calls: false``. Process kill-switch: ``UNSLOTH_DISABLE_TOOL_CALL_HEALING=1``. diff --git a/studio/backend/core/inference/pricing.py b/studio/backend/core/inference/pricing.py index 3b611d3596..30fec47723 100644 --- a/studio/backend/core/inference/pricing.py +++ b/studio/backend/core/inference/pricing.py @@ -122,12 +122,12 @@ def calculate_cost(provider: str, model: str, usage: dict[str, Any]) -> dict[str "priced": bool(prices), } - # Accept raw (input_tokens/output_tokens) and Studio chat-style + # Accept raw (input_tokens/output_tokens) and Unsloth chat-style # (prompt_tokens/completion_tokens) envelopes. Cache buckets differ: # raw Anthropic: input_tokens EXCLUDES cache buckets # raw OpenAI: input_tokens INCLUDES cache_read - # Studio Anthropic: prompt_tokens INCLUDES cache_creation + cache_read - # Studio OpenAI: prompt_tokens == raw input_tokens + # Unsloth Anthropic: prompt_tokens INCLUDES cache_creation + cache_read + # Unsloth OpenAI: prompt_tokens == raw input_tokens # Clamp >=0 so corrupted payloads can't produce a negative bill. cache_creation = max(0, int(usage.get("cache_creation_input_tokens") or 0)) cache_read_native_present = ( @@ -160,7 +160,7 @@ def calculate_cost(provider: str, model: str, usage: dict[str, Any]) -> dict[str output_tokens = max(0, int(usage.get("completion_tokens") or 0)) if provider == "openai": # Cached tokens land on input_tokens_details (raw Responses) or - # prompt_tokens_details (Studio chat-style). + # prompt_tokens_details (Unsloth chat-style). for key in ("input_tokens_details", "prompt_tokens_details"): details = usage.get(key) or {} if isinstance(details, dict): diff --git a/studio/backend/core/inference/safetensors_agentic.py b/studio/backend/core/inference/safetensors_agentic.py index 9110315815..40731de57b 100644 --- a/studio/backend/core/inference/safetensors_agentic.py +++ b/studio/backend/core/inference/safetensors_agentic.py @@ -995,7 +995,7 @@ def run_safetensors_tool_loop( if not safety_tc: # Re-prompt once on plan-without-action, before any tool runs # (GGUF loop parity). The retry is gated on nudge_tool_calls so - # Studio callers (which send True) always nudge, while API callers + # Unsloth callers (which send True) always nudge, while API callers # who omit the flag keep today's no-reprompt behavior (opt-in). intent_text = _reprompt_intent_text( content_accum, diff --git a/studio/backend/core/inference/sandbox_site/sitecustomize.py b/studio/backend/core/inference/sandbox_site/sitecustomize.py index d655e8e35a..244fa95145 100644 --- a/studio/backend/core/inference/sandbox_site/sitecustomize.py +++ b/studio/backend/core/inference/sandbox_site/sitecustomize.py @@ -4,7 +4,7 @@ """Sandbox-side compatibility shim for ChatGPT code-interpreter paths. Models habitually write to /mnt/data (or /mnt/outputs, /home/sandbox, -/workspace), none of which exist in the Studio sandbox. This module sits on the +/workspace), none of which exist in the Unsloth sandbox. This module sits on the sandbox subprocess PYTHONPATH (see ``tools._build_safe_env``), so it loads at interpreter startup in every sandboxed ``python`` run and any Python the ``terminal`` tool launches. diff --git a/studio/backend/core/inference/tool_loop_controller.py b/studio/backend/core/inference/tool_loop_controller.py index f7ed450d11..61643b5795 100644 --- a/studio/backend/core/inference/tool_loop_controller.py +++ b/studio/backend/core/inference/tool_loop_controller.py @@ -1,7 +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 -"""Shared controller state for Studio local agentic tool loops. +"""Shared controller state for Unsloth local agentic tool loops. This module is intentionally dependency-light: it owns only per-response ledger state and value objects used by the GGUF and safetensors loops. diff --git a/studio/backend/core/inference/tools.py b/studio/backend/core/inference/tools.py index 7f02392731..6ebda4fe98 100644 --- a/studio/backend/core/inference/tools.py +++ b/studio/backend/core/inference/tools.py @@ -2503,7 +2503,7 @@ def _build_safe_env(workdir: str) -> dict[str, str]: shim directory. """ # Start from the running interpreter's dir so 'python'/'pip' resolve to the - # same environment the Studio server runs in. + # same environment the Unsloth server runs in. exe_dir = os.path.dirname(sys.executable) path_entries = [exe_dir] if exe_dir else [] @@ -2793,7 +2793,7 @@ def _bypass_preexec(): """Minimal pre-exec for bypass exec: os.setsid() only. Required, not a restriction: _kill_process_tree does killpg(getpgid(child)), - so without a new session a timeout/cancel would kill the Studio server too. + so without a new session a timeout/cancel would kill the Unsloth server too. """ try: os.setsid() @@ -2801,13 +2801,13 @@ def _bypass_preexec(): pass -# Hardening the Studio parent is done once (PR_SET_DUMPABLE is process-global +# Hardening the Unsloth parent is done once (PR_SET_DUMPABLE is process-global # and sticky); guarded so repeated bypass calls do not re-issue the prctl. _parent_proc_hardened = False def _harden_parent_against_proc_env_leak() -> bool: - """Make the Studio process's /proc//environ unreadable to its children. + """Make the Unsloth process's /proc//environ unreadable to its children. Stripping the child env is not enough on Linux: a bypassed same-UID child can read /proc//environ to recover the parent's unfiltered @@ -5577,7 +5577,7 @@ def _truncate(text: str, limit: int = _MAX_OUTPUT_CHARS) -> str: # ChatGPT code-interpreter path conventions models write out of habit; none -# exist in the Studio sandbox, so a failure on one earns the retry hint. +# exist in the Unsloth sandbox, so a failure on one earns the retry hint. _MISSING_PATH_PREFIXES = ( "/mnt/data", "/mnt/outputs", @@ -5783,7 +5783,7 @@ def _python_exec( # Close the /proc//environ secret-recovery path first; if it # cannot be applied, fail closed rather than leak the parent environ. return ( - "Execution error: could not harden the Studio process against " + "Execution error: could not harden the Unsloth process against " "/proc environment reads; refusing bypass execution." ) @@ -5928,7 +5928,7 @@ def _bash_exec( # Close the /proc//environ secret-recovery path first; if it # cannot be applied, fail closed rather than leak the parent environ. return ( - "Execution error: could not harden the Studio process against " + "Execution error: could not harden the Unsloth process against " "/proc environment reads; refusing bypass execution." ) diff --git a/studio/backend/core/inference/worker.py b/studio/backend/core/inference/worker.py index e4628dcea8..9f301ba37e 100644 --- a/studio/backend/core/inference/worker.py +++ b/studio/backend/core/inference/worker.py @@ -513,20 +513,25 @@ def _handle_generate(backend, cmd: dict, resp_queue: Any, cancel_event) -> None: logger.info("Starting text generation for request_id=%s", request_id) - for cumulative_text in generator: - # cancel_event is an mp.Event — checked instantly, no queue polling. - if cancel_event.is_set(): - logger.info("Generation cancelled for request %s", request_id) - break + try: + for cumulative_text in generator: + # cancel_event is an mp.Event — checked instantly, no queue polling. + if cancel_event.is_set(): + logger.info("Generation cancelled for request %s", request_id) + break - _send_response( - resp_queue, - { - "type": "token", - "request_id": request_id, - "text": cumulative_text, - }, - ) + _send_response( + resp_queue, + { + "type": "token", + "request_id": request_id, + "text": cumulative_text, + }, + ) + finally: + close = getattr(generator, "close", None) + if callable(close): + close() _send_response( resp_queue, diff --git a/studio/backend/core/rag/captioner.py b/studio/backend/core/rag/captioner.py index 6d1512a770..8398506f21 100644 --- a/studio/backend/core/rag/captioner.py +++ b/studio/backend/core/rag/captioner.py @@ -6,7 +6,7 @@ Both turn pixels into indexable text and are a no-op (never raise) without a loaded vision model. They reuse the chat model's vision endpoint, so it must be served with ``--ubatch-size`` >= one image's tokens (some encoders, e.g. Gemma, attend -non-causally and abort otherwise); Studio's vision chat already requires this.""" +non-causally and abort otherwise); Unsloth's vision chat already requires this.""" from __future__ import annotations diff --git a/studio/backend/core/rag/config.py b/studio/backend/core/rag/config.py index 2de32a68e4..f54d795731 100644 --- a/studio/backend/core/rag/config.py +++ b/studio/backend/core/rag/config.py @@ -87,6 +87,22 @@ def _names_gguf(model: str) -> bool: return "gguf" in re.split(r"[^a-z0-9]+", model.lower()) +def gguf_repo_for_embedding_model(model: str) -> str: + """GGUF repo for ``model``, honoring an explicit companion override.""" + if "RAG_EMBED_GGUF_REPO" in os.environ: + return EMBED_GGUF_REPO + if model == DEFAULT_EMBEDDING_MODEL: + return EMBED_GGUF_REPO + if _names_gguf(model): + return model + return f"{model}-GGUF" + + +def default_gguf_repo() -> str: + """GGUF companion for the env/default embedding model.""" + return gguf_repo_for_embedding_model(EMBEDDING_MODEL) + + def effective_gguf_repo() -> str: """GGUF repo for the llama-server backend, tracking the effective model. @@ -95,14 +111,7 @@ def effective_gguf_repo() -> str: ``-GGUF`` companion repo (the unsloth convention the default pair follows), or is used as-is when it already names a GGUF repo. """ - if "RAG_EMBED_GGUF_REPO" in os.environ: - return EMBED_GGUF_REPO - model = effective_embedding_model() - if model == DEFAULT_EMBEDDING_MODEL: - return EMBED_GGUF_REPO - if _names_gguf(model): - return model - return f"{model}-GGUF" + return gguf_repo_for_embedding_model(effective_embedding_model()) # llama-server backend only. F16 over Q8_0: faster (no per-block dequant for this diff --git a/studio/backend/core/rag/embed_llama_server.py b/studio/backend/core/rag/embed_llama_server.py index 46a282c939..b141e59422 100644 --- a/studio/backend/core/rag/embed_llama_server.py +++ b/studio/backend/core/rag/embed_llama_server.py @@ -10,7 +10,7 @@ Opt-in (``RAG_EMBED_BACKEND=llama-server``). Runs a dedicated Device is ``auto`` (GPU when present, else CPU, falling back to CPU if a GPU start fails); ``RAG_EMBED_DEVICE`` forces it. We call only llama_cpp's *static* helpers (no torch), copying the instance-coupled bits locally, since constructing a -``LlamaCppBackend`` runs an ``__init__`` reaper that kills any Studio llama-server +``LlamaCppBackend`` runs an ``__init__`` reaper that kills any Unsloth llama-server -- so each request re-spawns ours if it died (self-heal). """ diff --git a/studio/backend/core/rag/embeddings.py b/studio/backend/core/rag/embeddings.py index b0ecedd593..15be7f1249 100644 --- a/studio/backend/core/rag/embeddings.py +++ b/studio/backend/core/rag/embeddings.py @@ -39,7 +39,7 @@ _model = None _name: str | None = None -# Studio device -> torch device string. Apple has no torch device -> CPU. +# Unsloth device -> torch device string. Apple has no torch device -> CPU. _TORCH_DEVICE = {DeviceType.CUDA: "cuda", DeviceType.XPU: "xpu"} diff --git a/studio/backend/core/rag/store.py b/studio/backend/core/rag/store.py index f9128d1715..1165b6bb0e 100644 --- a/studio/backend/core/rag/store.py +++ b/studio/backend/core/rag/store.py @@ -158,6 +158,16 @@ def list_documents(conn: sqlite3.Connection, scope: str) -> list[dict]: return [dict(r) for r in rows] +def list_all_documents(conn: sqlite3.Connection) -> list[dict]: + """Every uploaded document across all scopes (KBs, threads, projects).""" + rows = conn.execute( + "SELECT id, scope, kb_id, thread_id, project_id, filename, sha256, status, error, " + "num_chunks, stored_path, created_at " + "FROM documents ORDER BY created_at DESC" + ).fetchall() + return [dict(r) for r in rows] + + def get_document(conn: sqlite3.Connection, document_id: str) -> dict | None: row = conn.execute("SELECT * FROM documents WHERE id=?", (document_id,)).fetchone() return dict(row) if row else None diff --git a/studio/backend/core/training/resume.py b/studio/backend/core/training/resume.py index 2a4a198610..17183484c5 100644 --- a/studio/backend/core/training/resume.py +++ b/studio/backend/core/training/resume.py @@ -4,6 +4,8 @@ """Helpers for validating resumable training outputs.""" import json +import pickletools +import zipfile from pathlib import Path from typing import Optional @@ -33,27 +35,164 @@ def _checkpoint_step(path: Path) -> int: return -1 -def get_resume_checkpoint_path(path_value: str) -> Optional[str]: +_MODEL_FILES = ( + "adapter_model.safetensors", + "adapter_model.bin", + "model.safetensors", + "pytorch_model.bin", +) +_MODEL_INDEXES = ("model.safetensors.index.json", "pytorch_model.bin.index.json") + + +def _valid_state_file(path: Path, require_tensor: bool = True) -> bool: + try: + if not path.is_file() or path.stat().st_size == 0: + return False + if path.suffix == ".safetensors": + try: + from safetensors import SafetensorError, safe_open + except ImportError: + return False + try: + with safe_open(str(path), framework = "np") as state: + return bool(state.keys()) + except SafetensorError: + return False + if path.suffix in {".bin", ".pt"}: + with zipfile.ZipFile(path) as state: + infos = state.infolist() + names = [info.filename for info in infos] + data_name = next( + (name for name in names if name == "data.pkl" or name.endswith("/data.pkl")), + None, + ) + if data_name is None: + return False + data_prefix = data_name.removesuffix("data.pkl") + "data/" + operations = list(pickletools.genops(state.read(data_name))) + if not operations or operations[-1][0].name != "STOP": + return False + if not require_tensor: + return True + # Require a non-empty tensor record; a zero-byte one fails torch.load. + return any( + info.filename.startswith(data_prefix) + and not info.is_dir() + and info.file_size > 0 + for info in infos + ) + # Unrecognized state-file formats are not usable resume state. + return False + except (OSError, ValueError, zipfile.BadZipFile): + return False + + +def _checkpoint_state(path: Path) -> Optional[int]: + try: + state = json.loads((path / "trainer_state.json").read_text(encoding = "utf-8")) + step = state.get("global_step") if isinstance(state, dict) else None + except (OSError, UnicodeDecodeError, json.JSONDecodeError): + return None + if isinstance(step, bool) or not isinstance(step, int) or step < 0: + return None + directory_step = _checkpoint_step(path) + return step if directory_step < 0 or step == directory_step else None + + +_INDEX_SHARD_SUFFIX = { + "model.safetensors.index.json": ".safetensors", + "pytorch_model.bin.index.json": ".bin", +} + + +def _valid_indexed_shard(checkpoint: Path, shard: object, expected_suffix: str) -> bool: + # Shard must be a relative, in-format path contained in the checkpoint dir. + if not isinstance(shard, str) or not shard: + return False + if Path(shard).is_absolute() or Path(shard).suffix != expected_suffix: + return False + try: + root = checkpoint.resolve(strict = True) + candidate = (checkpoint / shard).resolve(strict = True) + candidate.relative_to(root) + except (OSError, ValueError): + return False + return _valid_state_file(candidate) + + +def _has_model_state(path: Path) -> bool: + if any(_valid_state_file(path / name) for name in _MODEL_FILES): + return True + for name in _MODEL_INDEXES: + try: + index = json.loads((path / name).read_text(encoding = "utf-8")) + shards = set(index["weight_map"].values()) + except ( + AttributeError, + OSError, + KeyError, + TypeError, + UnicodeDecodeError, + json.JSONDecodeError, + ): + continue + expected_suffix = _INDEX_SHARD_SUFFIX[name] + if shards and all(_valid_indexed_shard(path, shard, expected_suffix) for shard in shards): + return True + return False + + +def is_resume_checkpoint_valid( + path: Path, + expected_step: Optional[int] = None, + backend: Optional[str] = None, +) -> bool: + step = _checkpoint_state(path) if path.is_dir() else None + step_valid = step is not None and (expected_step is None or step == expected_step) + if backend == "mlx": + valid_bundle = _valid_state_file(path / "adapters.safetensors") and _valid_state_file( + path / "optimizer_state.safetensors" + ) + else: + valid_bundle = ( + _has_model_state(path) + # optimizer/scheduler state can be validly tensor-free (e.g. SGD without + # momentum); _has_model_state still requires real model tensors. + and _valid_state_file(path / "optimizer.pt", require_tensor = False) + and _valid_state_file(path / "scheduler.pt", require_tensor = False) + ) + if backend is None and not valid_bundle: + valid_bundle = _valid_state_file(path / "adapters.safetensors") and _valid_state_file( + path / "optimizer_state.safetensors" + ) + return step_valid and valid_bundle + + +def get_resume_checkpoint_path( + path_value: str, expected_step: Optional[int] = None +) -> Optional[str]: path = resolve_output_dir(path_value) if not _is_under_outputs(path) or not path.is_dir(): return None - if (path / "trainer_state.json").is_file(): + if is_resume_checkpoint_valid(path, expected_step): return str(path) - checkpoints = [ - child - for child in path.glob("checkpoint-*") - if child.is_dir() and (child / "trainer_state.json").is_file() - ] - if not checkpoints: - return None - return str(max(checkpoints, key = _checkpoint_step)) + checkpoints = sorted(path.glob("checkpoint-*"), key = _checkpoint_step, reverse = True) + return next( + ( + str(checkpoint) + for checkpoint in checkpoints + if _checkpoint_step(checkpoint) >= 0 + and is_resume_checkpoint_valid(checkpoint, expected_step) + ), + None, + ) def normalize_resume_output_dir(path_value: str) -> str: path = resolve_output_dir(path_value) if not _is_under_outputs(path): - raise ValueError("Resume checkpoint must be inside Studio outputs.") + raise ValueError("Resume checkpoint must be inside Unsloth outputs.") return str(path) @@ -78,9 +217,17 @@ def _uses_s3_dataset(run: dict) -> bool: def can_resume_run(run: dict) -> bool: if run.get("resumed_later"): return False + # Set when a stop-and-save failed to write a current-step checkpoint. + if run.get("resume_blocked"): + return False if _uses_s3_dataset(run): return False + status = run.get("status") + if status == "error": + # A save-time crash can report final_step == total_steps with no artifacts; checkpoint state alone decides resumability. + return has_resume_state(run.get("output_dir")) + final_step = run.get("final_step") total_steps = run.get("total_steps") has_remaining_steps = ( @@ -89,8 +236,4 @@ def can_resume_run(run: dict) -> bool: or total_steps <= 0 or final_step < total_steps ) - return ( - run.get("status") == "stopped" - and has_remaining_steps - and has_resume_state(run.get("output_dir")) - ) + return status == "stopped" and has_remaining_steps and has_resume_state(run.get("output_dir")) diff --git a/studio/backend/core/training/trainer.py b/studio/backend/core/training/trainer.py index 883a535a89..8e419849cb 100644 --- a/studio/backend/core/training/trainer.py +++ b/studio/backend/core/training/trainer.py @@ -797,7 +797,7 @@ class UnslothTrainer: ) logger.info("Loaded text model") - raise_if_offloaded(self.model, device_map, "Studio training") + raise_if_offloaded(self.model, device_map, "Unsloth training") if self.should_stop: return False @@ -3425,15 +3425,19 @@ class UnslothTrainer: logger.info( f"CPT: using UnslothTrainer with embedding_learning_rate={embedding_lr}\n" ) + cpt_args = _UnslothTrainingArguments( + embedding_learning_rate = embedding_lr, + **config_args, + ) + if config_args.get("packing", False): + cpt_args.packing_strategy = "wrapped" + logger.info("CPT packing strategy: wrapped\n") trainer_kwargs = { "model": self.model, "tokenizer": sft_tokenizer, "train_dataset": dataset["dataset"], "data_collator": data_collator, - "args": _UnslothTrainingArguments( - embedding_learning_rate = embedding_lr, - **config_args, - ), + "args": cpt_args, } if eval_dataset is not None: trainer_kwargs["eval_dataset"] = eval_dataset diff --git a/studio/backend/core/training/training.py b/studio/backend/core/training/training.py index 38f6b92f6d..26d8c23b66 100644 --- a/studio/backend/core/training/training.py +++ b/studio/backend/core/training/training.py @@ -140,7 +140,7 @@ def should_use_mlx_training_backend(*, device: Optional[Any] = None) -> bool: def _build_training_worker_config(values: dict[str, Any]) -> dict[str, Any]: - """Build the normalized worker config shared by Studio and the CLI adapter.""" + """Build the normalized worker config shared by Unsloth and the CLI adapter.""" config = { "model_name": values["model_name"], "project_name": values.get("project_name"), @@ -307,7 +307,7 @@ PLOT_HEIGHT = 3.5 @dataclass class TrainingProgress: - """Shared training progress payload for Studio and backend-aware trainers.""" + """Shared training progress payload for Unsloth and backend-aware trainers.""" epoch: float = 0 step: int = 0 @@ -328,7 +328,7 @@ class TrainingProgress: class _MLXTrainerAdapter: - """Adapts the legacy UnslothTrainer API to the shared Studio MLX worker path.""" + """Adapts the legacy UnslothTrainer API to the shared Unsloth MLX worker path.""" def __init__(self): self.model = None @@ -761,6 +761,7 @@ class TrainingBackend: # Left True after an abnormal death so _ensure_pump_alive spots a crash. self._pump_running: bool = False self._lock = threading.Lock() + self._run_intent_lock = threading.RLock() # Stop watchdog: after a stop is requested, escalates to force_terminate() # if the worker does not exit on its own within a bounded time. The watched @@ -773,6 +774,7 @@ class TrainingBackend: self._progress = TrainingProgress() self._should_stop = False self._cancel_requested = False # True only for stop(save=False) + self._cancel_cleanup_output_dir: Optional[str] = None # Throttled training-status logging to the server log (not one line/step). self._last_progress_log_ts: float = 0.0 @@ -792,6 +794,8 @@ class TrainingBackend: # Job metadata self.current_job_id: Optional[str] = None self._output_dir: Optional[str] = None + self._resume_source_run_id: Optional[str] = None + self._terminal_finalize_payload: Optional[dict] = None # DB persistence self._metric_buffer: list[dict] = [] @@ -819,6 +823,7 @@ class TrainingBackend: job_id: str, *, before_spawn = None, + resume_source_run_id: Optional[str] = None, **kwargs, ) -> bool: """Spawn a subprocess to run the full training pipeline. @@ -956,6 +961,7 @@ class TrainingBackend: self.current_job_id = job_id self._should_stop = False self._cancel_requested = False + self._cancel_cleanup_output_dir = None self._complete_seen.clear() self._progress = TrainingProgress( is_training = True, status_message = "Initializing training..." @@ -972,7 +978,10 @@ class TrainingBackend: self.eval_loss_history.clear() self.eval_step_history.clear() self.eval_enabled = False - self._output_dir = None + self._output_dir = config.get("output_dir") if resume_source_run_id else None + self._progress.output_dir = self._output_dir + self._resume_source_run_id = resume_source_run_id + self._terminal_finalize_payload = None self._metric_buffer.clear() self._run_finalized = False self._db_run_created = False @@ -990,6 +999,17 @@ class TrainingBackend: # in history during model loading and a fast terminal worker can't race the # pump into a duplicate create/finalize. From here the pump only finalizes. self._ensure_db_run_created() + if resume_source_run_id and not self._db_run_created: + if proc.is_alive(): + proc.terminate() + proc.join(timeout = 5.0) + if proc.is_alive(): + proc.kill() + proc.join(timeout = 2.0) + self._progress.is_training = False + self._progress.error = "Resume checkpoint is no longer available." + self._spawn_in_progress = False + return False # Assign handles and start the pump together under the lock so a concurrent # poll can't see a live _proc with no pump and spawn a duplicate. @@ -1011,28 +1031,75 @@ class TrainingBackend: def stop_training(self, save: bool = True) -> bool: """Send stop signal to the training subprocess.""" - self._should_stop = True - if not save: - self._cancel_requested = True - with self._lock: - if self._stop_queue is not None: - try: - self._stop_queue.put({"type": "stop", "save": save}) - except (OSError, ValueError): - pass - # Update progress immediately for responsive UI. - self._progress.status_message = ( - "Stopping training and saving checkpoint..." if save else "Cancelling training..." - ) - # Guarantee the run finalizes even if the worker wedges after saving. - self._start_stop_watchdog(cancel = not save) + with self._run_intent_lock: + with self._lock: + run_id = self.current_job_id + if not save and run_id: + persist_error: Optional[Exception] = None + for attempt in range(_DB_FINALIZE_RETRIES): + try: + from storage.studio_db import mark_run_cancel_requested + + self._ensure_db_run_created() + with self._lock: + terminal_payload = self._terminal_finalize_payload + if ( + terminal_payload + and terminal_payload.get("expected_job_id") == run_id + ): + return False + if not mark_run_cancel_requested(run_id): + if self._db_run_created: + return False + raise RuntimeError( + "Training run disappeared before cancellation persisted" + ) + if self.current_job_id != run_id: + return False + self._should_stop = self._cancel_requested = True + self._cancel_cleanup_output_dir = self._output_dir + self._output_dir = self._progress.output_dir = None + persist_error = None + break + except Exception as exc: + persist_error = exc + if attempt + 1 < _DB_FINALIZE_RETRIES: + time.sleep(_DB_FINALIZE_RETRY_S) + if persist_error is not None: + raise RuntimeError("Failed to persist Stop-without-Save") from persist_error + with self._lock: + if self.current_job_id != run_id: + return False + if save or not run_id: + self._should_stop = True + if not save and not run_id: + self._cancel_requested = True + self._cancel_cleanup_output_dir = self._output_dir + self._output_dir = self._progress.output_dir = None + if self._stop_queue is not None: + try: + self._stop_queue.put({"type": "stop", "save": save}) + except (OSError, ValueError): + pass + self._progress.status_message = ( + "Stopping training and saving checkpoint..." + if save + else "Cancelling training..." + ) + self._start_stop_watchdog(cancel = not save, expected_job_id = run_id) return True - def _start_stop_watchdog(self, cancel: bool) -> None: + def _start_stop_watchdog( + self, + cancel: bool, + expected_job_id: Optional[str] = None, + ) -> None: """Start a daemon that force-terminates the worker if a requested stop does not exit on its own. No-op if no worker is alive or a live watchdog already watches this proc (a stale watchdog on an old proc never blocks a new run's watcher).""" with self._lock: + if expected_job_id is not None and self.current_job_id != expected_job_id: + return proc = self._proc if proc is None or not proc.is_alive(): return @@ -1113,8 +1180,9 @@ class TrainingBackend: watched_job_id: Optional[str] = None, ) -> None: """Finalize parent state after a force-terminate so the UI leaves "Stopping..." - even if the worker is wedged in driver teardown; preserves output_dir so a saved - checkpoint is kept. No-ops if a new run already replaced the watched worker, so a + even if the worker is wedged in driver teardown; preserves output_dir on a save so + the checkpoint is kept, and clears it on a cancel (Stop without saving must not + offer resume/export). No-ops if a new run already replaced the watched worker, so a stale watchdog never marks a fresh run stopped or drops its handle. Supersession is checked on both the watched proc and job id: start_training sets @@ -1134,7 +1202,18 @@ class TrainingBackend: return # a new run is already starting up; leave its state alone run_id = self.current_job_id # == watched_job_id self._progress.is_training = False - self._progress.status_message = "Training stopped." + terminal_payload = self._terminal_finalize_kwargs() + status = terminal_payload["status"] + error_message = terminal_payload.get("error_message") + output_dir = terminal_payload["output_dir"] + clear_output_dir = terminal_payload["clear_output_dir"] + resume_blocked = bool(terminal_payload.get("resume_blocked")) + with self._lock: + if self.current_job_id != run_id: + return + self._progress.status_message = error_message or "Training stopped." + if error_message: + self._progress.error = error_message # Create the row if a start-time create failed (no-op otherwise; skips when the pump # is mid-create, in which case its create-then-finalize records the run instead). self._ensure_db_run_created() @@ -1148,7 +1227,8 @@ class TrainingBackend: batch: list = [] final_step = final_loss = duration = None loss_history: list = [] - output_dir = self._output_dir + if clear_output_dir: + self._output_dir = self._progress.output_dir = None if claim: self._run_finalized = True # claim this run's finalize batch = list(self._metric_buffer) @@ -1161,7 +1241,17 @@ class TrainingBackend: loss_history = list(self.loss_history) if claim: self._finish_stopped_run( - run_id, output_dir, batch, final_step, final_loss, duration, loss_history + run_id, + output_dir, + batch, + final_step, + final_loss, + duration, + loss_history, + status = status, + error_message = error_message, + clear_output_dir = clear_output_dir, + resume_blocked = resume_blocked, ) with self._lock: if target_proc is None or self._proc is target_proc: @@ -1176,6 +1266,10 @@ class TrainingBackend: final_loss: Optional[float], duration: Optional[float], loss_history: list, + status: str = "stopped", + error_message: Optional[str] = None, + clear_output_dir: bool = False, + resume_blocked: bool = False, ) -> None: """Record a force-stopped run finished by its captured id, from state snapshotted under the lock. insert_metrics_batch upserts and finish_run is an idempotent UPDATE, @@ -1194,14 +1288,16 @@ class TrainingBackend: sparkline = downsample(loss_history, 50) finish_run( id = run_id, - status = "stopped", + status = status, ended_at = datetime.now(timezone.utc).isoformat(), final_step = final_step, final_loss = final_loss, duration_seconds = duration, loss_sparkline = _json.dumps(sparkline), output_dir = output_dir, - error_message = None, + error_message = error_message, + clear_output_dir = clear_output_dir, + resume_blocked = resume_blocked, ) return except Exception: @@ -1231,7 +1327,7 @@ class TrainingBackend: logger.info("Force-terminating training subprocess (pid=%s)", proc.pid) proc.terminate() cancelled = self._cancel_requested - output_dir = self._output_dir + output_dir = self._cancel_cleanup_output_dir or self._output_dir if proc is not None: proc.join(timeout = 5.0) @@ -1595,17 +1691,60 @@ class TrainingBackend: ) self._ensure_db_run_created() - self._finalize_run_in_db( - status = "stopped" if self._should_stop else "error", - error_message = None - if self._should_stop - else "Training process terminated unexpectedly", - ) + terminal_payload = self._terminal_finalize_kwargs() + with self._lock: + if terminal_payload["clear_output_dir"]: + self._output_dir = self._progress.output_dir = None + if terminal_payload.get("error_message"): + self._progress.error = terminal_payload["error_message"] + self._progress.status_message = terminal_payload["error_message"] + self._finalize_run_in_db(**terminal_payload) except Exception: logger.exception("Training event pump: finalization after worker exit failed") self._pump_running = False return + def _has_current_resume_checkpoint(self, output_dir, step) -> bool: + # A valid checkpoint at the current step means the stop-and-save landed on + # disk even if the worker died before confirming it. + if not output_dir or not isinstance(step, int) or step <= 0: + return False + from core.training.resume import get_resume_checkpoint_path + return get_resume_checkpoint_path(output_dir, expected_step = step) is not None + + def _terminal_finalize_kwargs(self) -> dict: + with self._lock: + job_id = self.current_job_id + payload = self._terminal_finalize_payload + if payload and payload.get("expected_job_id") == job_id: + return dict(payload) + cancel, stopped = self._cancel_requested, self._should_stop + output_dir = None if cancel else self._output_dir + step = self._progress.step + existing_error = self._progress.error + status, error, blocked = ( + ("stopped", None, cancel) + if stopped + else ( + "error", + existing_error or "Training process terminated unexpectedly", + False, + ) + ) + # Block only when no valid current-step checkpoint actually landed. + if stopped and not cancel and not self._has_current_resume_checkpoint(output_dir, step): + status = "error" + error = "Stop and Save ended before a valid current-step checkpoint was written." + blocked = True + return { + "status": status, + "error_message": error, + "output_dir": output_dir, + "clear_output_dir": cancel, + "resume_blocked": blocked, + "expected_job_id": job_id, + } + def _handle_event(self, event: dict) -> None: """Apply a subprocess event to local state. @@ -1764,6 +1903,15 @@ class TrainingBackend: elif etype == "eval_configured": self.eval_enabled = True + elif etype == "output_dir": + event_output_dir = event.get("output_dir") + if self._cancel_requested: + self._cancel_cleanup_output_dir = event_output_dir + self._output_dir = self._progress.output_dir = None + else: + self._output_dir = event_output_dir + db_action = "persist_output_dir" + elif etype == "status": self._progress.status_message = event.get("message", "") self._progress.is_training = True @@ -1778,7 +1926,12 @@ class TrainingBackend: self._complete_seen.set() self._progress.is_training = False self._progress.is_completed = not stopped - self._output_dir = event.get("output_dir") + event_output_dir = event.get("output_dir") + if self._cancel_requested: + self._cancel_cleanup_output_dir = event_output_dir + self._output_dir = None + else: + self._output_dir = event_output_dir self._progress.output_dir = self._output_dir self._progress.status_message = msg if not self._db_run_created and self.current_job_id and self._db_config: @@ -1788,11 +1941,16 @@ class TrainingBackend: db_action_kwargs = { "status": "stopped" if stopped else "completed", "output_dir": self._output_dir, + "clear_output_dir": self._cancel_requested, + "expected_job_id": self.current_job_id, } + self._terminal_finalize_payload = dict(db_action_kwargs) elif etype == "error": self._progress.is_training = False self._progress.error = event.get("error", "Unknown error") + if self._cancel_requested: + self._output_dir = self._progress.output_dir = None logger.error("Training error: %s", event.get("error")) stack = event.get("stack", "") if stack: @@ -1801,29 +1959,36 @@ class TrainingBackend: db_action = "create_and_finalize" else: db_action = "finalize" + stop_save_failed = ( + self._should_stop + and not self._cancel_requested + and not self._has_current_resume_checkpoint( + self._output_dir, self._progress.step + ) + ) db_action_kwargs = { - "status": "stopped" if self._should_stop else "error", + "status": "stopped" + if self._should_stop + and not stop_save_failed + and not event.get("keep_error_status") + else "error", "error_message": event.get("error", "Unknown error"), + "output_dir": self._output_dir, + "clear_output_dir": self._cancel_requested, + "resume_blocked": stop_save_failed or bool(event.get("resume_blocked")), + "expected_job_id": self.current_job_id, } + self._terminal_finalize_payload = dict(db_action_kwargs) # --- DB I/O outside the lock --- if db_action == "create_run": - try: - from storage.studio_db import create_run - - create_run( - id = db_action_kwargs["job_id"], - model_name = db_action_kwargs["model_name"], - dataset_name = db_action_kwargs["dataset_name"], - config_json = db_action_kwargs["config_json"], - started_at = db_action_kwargs["started_at"], - total_steps = db_action_kwargs["total_steps"], - ) - self._db_run_created = True + self._ensure_db_run_created() + if self._db_run_created: if db_action_kwargs["total_steps"]: self._db_total_steps_set = True - except Exception: - logger.warning("Failed to create DB run record", exc_info = True) + self._persist_output_dir() + elif db_action == "persist_output_dir": + self._persist_output_dir() elif db_action == "create_and_finalize": self._ensure_db_run_created() self._finalize_run_in_db(**db_action_kwargs) @@ -1842,6 +2007,22 @@ class TrainingBackend: if etype == "progress": self._log_training_progress() + def _persist_output_dir(self) -> None: + with self._lock: + if ( + not self._output_dir + or not self.current_job_id + or not self._db_run_created + or self._cancel_requested + ): + return + run_id, output_dir = self.current_job_id, self._output_dir + try: + from storage.studio_db import update_run_output_dir + update_run_output_dir(run_id, output_dir) + except Exception: + logger.warning("Failed to persist output_dir", exc_info = True) + def _log_training_progress(self) -> None: """One throttled training-status line to the server log (the per-step stream still goes to the UI via SSE): first step, then at most every 30s, plus the @@ -1875,6 +2056,7 @@ class TrainingBackend: caller create at a time, and ``_db_run_created`` is published only after ``create_run`` commits, so a concurrent finalize never runs ``finish_run`` against a not-yet-inserted row (a zero-row UPDATE that would leave the run stuck as running).""" + self._run_intent_lock.acquire() with self._lock: if ( self._db_run_created @@ -1882,6 +2064,7 @@ class TrainingBackend: or not self.current_job_id or not self._db_config ): + self._run_intent_lock.release() return self._db_create_in_progress = True # only one caller creates job_id = self.current_job_id @@ -1898,6 +2081,12 @@ class TrainingBackend: or _s3_dataset_name(db_config.get("s3_dataset")) or "unknown" ) + with self._lock: + if self.current_job_id != job_id: + return + output_dir = self._output_dir + cancel_requested = self._cancel_requested + resumed_from_run_id = self._resume_source_run_id create_run( id = job_id, model_name = db_config["model_name"], @@ -1905,6 +2094,9 @@ class TrainingBackend: config_json = _json.dumps(db_config), started_at = started_at, total_steps = total_steps, + output_dir = output_dir, + cancel_requested = cancel_requested, + resumed_from_run_id = resumed_from_run_id, ) created = True except Exception: @@ -1919,12 +2111,15 @@ class TrainingBackend: if created: self._db_run_created = True # publish only after the insert commits self._db_create_in_progress = False + self._run_intent_lock.release() def _finalize_run_in_db( self, status: str, error_message: Optional[str] = None, output_dir: Optional[str] = None, + clear_output_dir: bool = False, + resume_blocked: bool = False, expected_job_id: Optional[str] = None, ) -> None: """Flush remaining metrics and mark a run finished in the DB. Claims the finalize @@ -1947,26 +2142,33 @@ class TrainingBackend: duration = self._progress.elapsed_seconds loss_history = list(self.loss_history) self._flush_metrics_to_db(run_id = run_id) - try: - from storage.studio_db import finish_run - from utils.downsample import downsample + for attempt in range(_DB_FINALIZE_RETRIES): + try: + from storage.studio_db import finish_run + from utils.downsample import downsample - sparkline = downsample(loss_history, 50) - finish_run( - id = run_id, - status = status, - ended_at = datetime.now(timezone.utc).isoformat(), - final_step = final_step, - final_loss = final_loss, - duration_seconds = duration, - loss_sparkline = _json.dumps(sparkline), - output_dir = output_dir, - error_message = error_message, - ) - except Exception: - with self._lock: - self._run_finalized = False # unclaim so a later flush can retry - logger.warning("Failed to finalize run in DB (status=%s)", status, exc_info = True) + finish_run( + id = run_id, + status = status, + ended_at = datetime.now(timezone.utc).isoformat(), + final_step = final_step, + final_loss = final_loss, + duration_seconds = duration, + loss_sparkline = _json.dumps(downsample(loss_history, 50)), + output_dir = output_dir, + error_message = error_message, + clear_output_dir = clear_output_dir, + resume_blocked = resume_blocked, + ) + return + except Exception: + if attempt + 1 < _DB_FINALIZE_RETRIES: + time.sleep(_DB_FINALIZE_RETRY_S) + continue + with self._lock: + if self.current_job_id == run_id: + self._run_finalized = False + logger.warning("Failed to finalize run in DB (status=%s)", status, exc_info = True) def _flush_metrics_to_db(self, run_id: Optional[str] = None) -> None: """Flush buffered metrics to the DB and update live progress. The target run id, diff --git a/studio/backend/core/training/worker.py b/studio/backend/core/training/worker.py index c52adbe8fa..2df4fa58c6 100644 --- a/studio/backend/core/training/worker.py +++ b/studio/backend/core/training/worker.py @@ -1100,7 +1100,7 @@ _MLX_VLM_RESIZED_IMAGE_LAYOUT_CACHE = {} def _mlx_vlm_resized_image_layout(processor = None) -> str | None: - """Return the numpy image layout expected after Studio-side VLM resizing.""" + """Return the numpy image layout expected after Unsloth-side VLM resizing.""" image_processor = getattr(processor, "image_processor", None) if image_processor is None: return None @@ -1257,7 +1257,7 @@ _MLX_STUDIO_LR_SCHEDULERS = {"linear", "cosine", "constant"} # Fallback alias map mirroring unsloth_zoo._normalize_mlx_optimizer_name, used -# only when mlx (Apple Silicon) is not importable so Studio config validation +# only when mlx (Apple Silicon) is not importable so Unsloth config validation # still works on non-MLX hosts. The zoo function stays the source of truth. _MLX_STUDIO_ADAMW_ALIASES = frozenset( ( @@ -1309,7 +1309,7 @@ def _normalize_mlx_studio_scheduler(value): def _resolve_mlx_local_dataset_files(file_paths: list) -> list[str]: - """Resolve CLI paths and Studio local dataset uploads without importing the GPU trainer.""" + """Resolve CLI paths and Unsloth local dataset uploads without importing the GPU trainer.""" from utils.paths import resolve_dataset_path all_files: list[str] = [] @@ -1840,8 +1840,15 @@ def _run_mlx_training(event_queue, stop_queue, config): # Resolve to ~/.unsloth/studio/outputs/ so the export page finds it from utils.paths import ensure_dir - output_dir = _resolve_mlx_output_dir(config, model_name) + # Resume must land in the original run dir even when config lacks output_dir. + resume_dir = config.get("output_dir", "") or _output_dir_from_resume_checkpoint( + resume_from_checkpoint + ) + output_dir = _resolve_mlx_output_dir( + {**config, "output_dir": resume_dir} if resume_dir else config, model_name + ) ensure_dir(Path(output_dir)) + _emit_output_dir(event_queue, output_dir) # ── 6. Create trainer ── eval_steps_val = config.get("eval_steps", 0) or 0 @@ -1912,7 +1919,7 @@ def _run_mlx_training(event_queue, stop_queue, config): if "max_grad_leaf_norm" in _supported_fields: mlx_config_kwargs["max_grad_leaf_norm"] = max_grad_leaf_norm if "append_eos" in _supported_fields: - # Studio SFT formatting owns rendered examples; raw/CPT text still + # Unsloth SFT formatting owns rendered examples; raw/CPT text still # needs MLX to append EOS like the CUDA raw-text path. mlx_config_kwargs["append_eos"] = bool(raw_text_mode) @@ -2067,6 +2074,17 @@ def _run_mlx_training(event_queue, stop_queue, config): trainer.add_eval_callback(_on_eval) + _opt_ref = [None] + _orig_build_optimizer = getattr(trainer, "_build_optimizer", None) + + if callable(_orig_build_optimizer): + + def _capture_optimizer(total_steps): + _opt_ref[0] = _orig_build_optimizer(total_steps) + return _opt_ref[0] + + trainer._build_optimizer = _capture_optimizer + # ── 11. Run training ── gc.collect() mx.synchronize() @@ -2082,31 +2100,58 @@ def _run_mlx_training(event_queue, stop_queue, config): trainer.save_model = _save_model # ── 12. Save and finalize ── - if trainer.stop_requested: - if not _stop_save[0]: - # Cancel (save=False): skip saving. - _send("complete", output_dir = None, status_message = "Training cancelled") + def _finish_tracking() -> None: + # Runs on every save/finalize exit so TB/W&B never leak on early return. + if tb_writer is not None: + try: + tb_writer.close() + except Exception: + pass + if wandb_run is not None: + try: + wandb_run.finish() + except Exception: + pass + + def _stop_checkpoint_ok() -> bool: + if _write_mlx_stop_checkpoint(trainer, _opt_ref[0], output_dir): + return True + _send( + "error", + error = ( + "Failed to save a resumable checkpoint after stop. " + "Model files were saved, but this run cannot be resumed." + ), + # A user stop finalizes as 'stopped'; keep this failure's error status so history explains it. + keep_error_status = True, + # Older checkpoints are stale; resuming would roll back past this stop. + resume_blocked = True, + ) + return False + + try: + if trainer.stop_requested: + if not _stop_save[0]: + # Cancel (save=False): skip saving. + _send("complete", output_dir = None, status_message = "Training cancelled") + else: + _send("status", status_message = "Saving stopped model...") + mx.synchronize() + trainer.save_model(output_dir) + # Stop-and-save promises a resumable checkpoint, not just model files. + if not _stop_checkpoint_ok(): + return + _send("complete", output_dir = output_dir, status_message = "Training stopped") else: - _send("status", status_message = "Saving stopped model...") + _send("status", status_message = "Saving model...") mx.synchronize() trainer.save_model(output_dir) - _send("complete", output_dir = output_dir, status_message = "Training stopped") - else: - _send("status", status_message = "Saving model...") - mx.synchronize() - trainer.save_model(output_dir) - _send("complete", output_dir = output_dir, status_message = "Training completed") - - if tb_writer is not None: - try: - tb_writer.close() - except Exception: - pass - if wandb_run is not None: - try: - wandb_run.finish() - except Exception: - pass + # A save-stop can race the natural final save; it made the same promise. + if trainer.stop_requested and _stop_save[0] and not _stop_checkpoint_ok(): + return + _send("complete", output_dir = output_dir, status_message = "Training completed") + finally: + _finish_tracking() def _is_current_process_apple_silicon() -> bool: @@ -2121,7 +2166,7 @@ def run_mlx_training_process( config: dict, transformers_activated: bool = False, ) -> None: - """MLX worker entrypoint shared by Studio subprocesses and the CLI adapter.""" + """MLX worker entrypoint shared by Unsloth subprocesses and the CLI adapter.""" model_name = config["model_name"] backend_path = str(Path(__file__).resolve().parent.parent.parent) @@ -2780,7 +2825,7 @@ def run_training_process(*, event_queue: Any, stop_queue: Any, config: dict) -> ) # Unified Windows APUs: the WDDM budget is user-raisable, but # nothing on the box says so -- users see "48 GB VRAM" on a - # 96 GB machine and assume a Studio bug. Say where the limit + # 96 GB machine and assume an Unsloth bug. Say where the limit # comes from and how to raise it. if _is_unified and sys.platform == "win32": try: @@ -3177,6 +3222,7 @@ def run_training_process(*, event_queue: Any, stop_queue: Any, config: dict) -> ) output_dir = str(resolve_output_dir(output_dir)) ensure_dir(Path(output_dir)) + _emit_output_dir(event_queue, output_dir) tensorboard_dir = config.get("tensorboard_dir") if config.get("enable_tensorboard", False): @@ -3296,6 +3342,61 @@ def _send_status(event_queue: Any, message: str) -> None: ) +def _emit_output_dir(event_queue: Any, output_dir: str) -> None: + try: + event_queue.put({"type": "output_dir", "output_dir": output_dir, "ts": time.time()}) + except Exception: + pass + + +def _mlx_has_checkpoint_at_step(output_dir, step: int) -> bool: + if step <= 0: + return False + from core.training.resume import is_resume_checkpoint_valid + return is_resume_checkpoint_valid( + Path(output_dir) / f"checkpoint-{step}", expected_step = step, backend = "mlx" + ) + + +def _write_mlx_stop_checkpoint(trainer, optimizer, output_dir) -> bool: + """Write a full resume checkpoint for a stopped MLX run. + + Returns True when a checkpoint for the current training step exists. + """ + step = int(getattr(trainer, "_global_step", 0) or 0) + # A periodic save or a resumed run may already cover the current step. + if _mlx_has_checkpoint_at_step(output_dir, step): + return True + if step <= 0 or optimizer is None: + return False + ckpt_dir = Path(output_dir) / f"checkpoint-{step}" + if ckpt_dir.is_symlink(): + # Refuse a symlinked dir: it could redirect writes outside output_dir. + logger.error("Refusing to write MLX stop checkpoint through symlink: %s", ckpt_dir) + return False + try: + ckpt_dir.mkdir(parents = True, exist_ok = True) + from unsloth_zoo.mlx.utils import ( + save_optimizer_state, + save_trainable_adapters, + save_trainer_state, + ) + + save_trainable_adapters(trainer.model, str(ckpt_dir)) + save_optimizer_state(optimizer, str(ckpt_dir)) + save_trainer_state( + { + "global_step": step, + "train_loss_history": list(getattr(trainer, "_train_loss_history", [])), + }, + str(ckpt_dir), + ) + logger.info("Saved stop checkpoint to %s", ckpt_dir) + except Exception: + logger.exception("Failed to write stop checkpoint under %s", output_dir) + return _mlx_has_checkpoint_at_step(output_dir, step) + + def _run_embedding_training(event_queue: Any, stop_queue: Any, config: dict) -> None: """Self-contained embedding model training pipeline. @@ -3660,6 +3761,7 @@ def _run_embedding_training(event_queue: Any, stop_queue: Any, config: dict) -> config.get("project_name"), ) output_dir = str(resolve_output_dir(output_dir)) + _emit_output_dir(event_queue, output_dir) num_epochs = config.get("num_epochs", 2) batch_size = config.get("batch_size", 256) diff --git a/studio/backend/hub/routes/__init__.py b/studio/backend/hub/routes/__init__.py index e9579635b0..7c5cfb9b3c 100644 --- a/studio/backend/hub/routes/__init__.py +++ b/studio/backend/hub/routes/__init__.py @@ -5,8 +5,10 @@ from hub.routes.inventory import router as inventory_router from hub.routes.datasets import router as datasets_router +from hub.routes.token import router as token_router __all__ = [ "inventory_router", "datasets_router", + "token_router", ] diff --git a/studio/backend/hub/routes/inventory.py b/studio/backend/hub/routes/inventory.py index 4b6c179a2b..1ffadf0544 100644 --- a/studio/backend/hub/routes/inventory.py +++ b/studio/backend/hub/routes/inventory.py @@ -28,6 +28,7 @@ from hub.schemas.inventory import ( CachedModelsResponse, DeleteCachedModelResponse, GgufVariantsResponse, + HiddenModelsResponse, LocalModelListResponse, ModelsFolderResponse, RecommendedFoldersResponse, @@ -214,6 +215,16 @@ async def list_cached_models( return await cache_inventory.list_cached_models_response(hf_token) +@router.get("/hidden-models", response_model = HiddenModelsResponse) +async def list_hidden_models(current_subject: str = Depends(get_current_subject)): + import asyncio + + from routes.models import hidden_model_matchers + + needles, exact_ids, exact_paths = await asyncio.to_thread(hidden_model_matchers) + return HiddenModelsResponse(needles = needles, exact_ids = exact_ids, exact_paths = exact_paths) + + @router.delete( "/delete-cached", response_model = DeleteCachedModelResponse, diff --git a/studio/backend/hub/routes/token.py b/studio/backend/hub/routes/token.py new file mode 100644 index 0000000000..1b7ad733a2 --- /dev/null +++ b/studio/backend/hub/routes/token.py @@ -0,0 +1,44 @@ +# SPDX-License-Identifier: AGPL-3.0-only +# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0 + +"""Hugging Face token validation endpoint.""" + +from __future__ import annotations + +import asyncio +from typing import Literal, Optional + +from fastapi import APIRouter, Depends, Request +from pydantic import BaseModel + +from auth.authentication import get_current_subject +from hub.dependencies import get_hf_token +from utils.client_ip import client_ip +from utils.hf_token_validation import validate_hf_token + + +router = APIRouter() + + +class HfTokenValidationResponse(BaseModel): + status: Literal["missing", "valid", "invalid", "rate_limited", "unavailable"] + retry_after_seconds: Optional[int] = None + + +@router.post("/token/validate", response_model = HfTokenValidationResponse) +async def validate_token( + request: Request, + hf_token: Optional[str] = Depends(get_hf_token), + current_subject: str = Depends(get_current_subject), +): + if not hf_token: + return HfTokenValidationResponse(status = "missing") + result = await asyncio.to_thread( + validate_hf_token, + hf_token, + rate_key = f"{current_subject}:{client_ip(request)}", + ) + return HfTokenValidationResponse( + status = result.status, + retry_after_seconds = result.retry_after_seconds, + ) diff --git a/studio/backend/hub/schemas/inventory.py b/studio/backend/hub/schemas/inventory.py index ef95efe2f2..19d6da3e11 100644 --- a/studio/backend/hub/schemas/inventory.py +++ b/studio/backend/hub/schemas/inventory.py @@ -160,6 +160,7 @@ class CachedRepoBase(BaseModel): repo_id: str size_bytes: int = 0 cache_path: Optional[str] = None + last_modified: Optional[float] = None partial: bool = False partial_transport: Optional[str] = None inventory_id: Optional[str] = None @@ -189,6 +190,12 @@ class CachedModelsResponse(BaseModel): cached: List[CachedModelRepo] = Field(default_factory = list) +class HiddenModelsResponse(BaseModel): + needles: List[str] = Field(default_factory = list) + exact_ids: List[str] = Field(default_factory = list) + exact_paths: List[str] = Field(default_factory = list) + + class AddScanFolderRequest(BaseModel): """Request body for adding a custom scan folder.""" diff --git a/studio/backend/hub/services/download_lifecycle.py b/studio/backend/hub/services/download_lifecycle.py index e5d48872c1..23f8c7c911 100644 --- a/studio/backend/hub/services/download_lifecycle.py +++ b/studio/backend/hub/services/download_lifecycle.py @@ -76,7 +76,7 @@ def spawn_worker( env["HF_HUB_DISABLE_PROGRESS_BARS"] = "1" env["HF_HUB_DISABLE_TELEMETRY"] = "1" env["HF_HUB_DISABLE_XET"] = "0" if use_xet else "1" - # No token in Studio settings: fall back to the backend's own HF_TOKEN so + # No token in Unsloth settings: fall back to the backend's own HF_TOKEN so # private repos stay downloadable (needed while inkling repos are private). if not hf_token: hf_token = os.environ.get("HF_TOKEN") or None diff --git a/studio/backend/hub/services/models/cache_inventory.py b/studio/backend/hub/services/models/cache_inventory.py index 1f38af9381..c1b864bb63 100644 --- a/studio/backend/hub/services/models/cache_inventory.py +++ b/studio/backend/hub/services/models/cache_inventory.py @@ -31,12 +31,20 @@ from hub.services.models.common import ( _is_checkpoint_weight_name, _is_gguf_filename, _is_main_gguf_filename, + _is_mmproj_filename, _is_transformers_safetensors_weight_name, _local_inventory_id, _prefer_complete_larger, _runtime_for_format, ) +# Imported at module scope (not inside the per-repo scan loop) so a broken +# import surfaces at startup instead of silently emptying the inventory: the +# scan loop swallows per-repo exceptions and would drop every repo. Lives under +# ``utils`` (not ``utils.models``) to avoid the eager model-config/checkpoint +# imports in ``utils/models/__init__.py``. +from utils.hidden_models import is_hidden_model + logger = get_logger(__name__) _repo_size_cache: "OrderedDict[tuple[str, str, str], tuple[int, frozenset[str], float]]" = ( @@ -125,6 +133,39 @@ def _repo_has_gguf_files(repo_info) -> bool: return _repo_gguf_size_bytes(repo_info) > 0 +def _blob_mtime(file_obj) -> float: + ts = getattr(file_obj, "blob_last_modified", None) + if isinstance(ts, (int, float)) and ts > 0: + return float(ts) + blob_path = getattr(file_obj, "blob_path", None) + if blob_path: + try: + return float(Path(blob_path).stat().st_mtime) + except OSError: + pass + return 0.0 + + +def _repo_gguf_last_modified(repo_info) -> float: + latest = 0.0 + for revision in repo_info.revisions: + for f in revision.files: + if _is_main_gguf_filename(f.file_name): + latest = max(latest, _blob_mtime(f)) + return latest + + +def _repo_has_mmproj(repo_info) -> bool: + # An mmproj file only makes a repo vision-capable when it is an actual GGUF + # projector; a non-GGUF sidecar (e.g. mmproj_config.json) does not, and the + # runtime's projector detection is GGUF-only. + return any( + _is_gguf_filename(f.file_name) and _is_mmproj_filename(f.file_name) + for revision in repo_info.revisions + for f in revision.files + ) + + def _cached_repo_file_name(file_obj) -> str: file_path = getattr(file_obj, "file_path", None) if file_path: @@ -243,6 +284,13 @@ def invalidate_hf_cache_scans() -> None: hf_cache_scan.invalidate_hf_cache_scans() +def _is_hidden_infra_repo(*values: str | None) -> bool: + """True for infra-only repos (the RAG embedder and the llama.cpp install + validation probe) that are cached as a side effect of Studio itself and are + not usable chat models.""" + return is_hidden_model(*values) + + def _scan_cached_gguf() -> list[dict]: """Synchronous HF-cache disk walk for GGUF repos; runs in a worker thread.""" cache_scans = all_hf_cache_scans() @@ -254,18 +302,30 @@ def _scan_cached_gguf() -> list[dict]: if str(repo_info.repo_type) != "model": continue repo_id = repo_info.repo_id + repo_path = Path(repo_info.repo_path) + snapshot_path = _cached_model_snapshot_path(repo_path) total_size = _repo_gguf_size_bytes(repo_info) has_variant_state, variant_state_size = _gguf_variant_state_summary(repo_id) + is_hidden_infra = _is_hidden_infra_repo( + repo_id, + str(repo_path), + str(snapshot_path) if snapshot_path is not None else None, + ) + # Hide infra repos unless the user downloaded a variant via + # the Hub; variant state only exists for user downloads. + if is_hidden_infra and not has_variant_state: + continue if total_size == 0 and not has_variant_state: continue partial = hf_cache_scan.is_gguf_repo_partial( repo_id, - Path(repo_info.repo_path), + repo_path, ) if total_size == 0 and not partial: continue key = repo_id.lower() existing = seen_lower.get(key) + last_modified = _repo_gguf_last_modified(repo_info) row = { "repo_id": repo_id, "size_bytes": max(total_size, variant_state_size), @@ -275,6 +335,9 @@ def _scan_cached_gguf() -> list[dict]: # per-variant detail lives on GgufVariantDetail. "partial_transport": None, } + last_modified = max(last_modified, (existing or {}).get("last_modified", 0.0)) + if last_modified > 0: + row["last_modified"] = last_modified row.update( _cache_inventory_fields( repo_id, @@ -283,8 +346,20 @@ def _scan_cached_gguf() -> list[dict]: requires_variant = True, ) ) + if _repo_has_mmproj(repo_info): + row["capabilities"]["supports_vision"] = True + # Visible infra variants remain management-only. + if is_hidden_infra: + row["capabilities"]["can_chat"] = False if _prefer_cache_row(row, existing): + if existing and existing["capabilities"].get("supports_vision"): + row["capabilities"]["supports_vision"] = True seen_lower[key] = row + else: + if last_modified > existing.get("last_modified", 0.0): + existing["last_modified"] = last_modified + if row["capabilities"].get("supports_vision"): + existing["capabilities"]["supports_vision"] = True except Exception as e: repo_label = getattr(repo_info, "repo_id", "") logger.warning(f"Skipping cached GGUF repo {repo_label}: {e}") @@ -312,13 +387,14 @@ class _CachedNonGgufPayload(NamedTuple): size_bytes: int has_runnable_weights: bool model_format: ModelFormat + last_modified: float def _repo_non_gguf_model_payload(repo_info) -> _CachedNonGgufPayload: - all_weight_blobs: dict[str, int] = {} - adapter_blobs: dict[str, int] = {} - safetensors_blobs: dict[str, int] = {} - checkpoint_blobs: dict[str, int] = {} + all_weight_blobs: dict[str, tuple[int, float]] = {} + adapter_blobs: dict[str, tuple[int, float]] = {} + safetensors_blobs: dict[str, tuple[int, float]] = {} + checkpoint_blobs: dict[str, tuple[int, float]] = {} has_config = False has_adapter_config = False has_adapter_weights = False @@ -326,12 +402,15 @@ def _repo_non_gguf_model_payload(repo_info) -> _CachedNonGgufPayload: has_transformers_safetensors = False has_checkpoint = False - def _record_blob(target: dict[str, int], file_obj, rev_id: str, file_name: str) -> None: + def _record_blob( + target: dict[str, tuple[int, float]], file_obj, rev_id: str, file_name: str + ) -> None: blob_path = getattr(file_obj, "blob_path", None) size = int(file_obj.size_on_disk or 0) key = str(blob_path) if blob_path else f"{rev_id}:{file_name}" - target[key] = size - all_weight_blobs[key] = size + value = (size, _blob_mtime(file_obj)) + target[key] = value + all_weight_blobs[key] = value for revision in repo_info.revisions: rev_id = getattr(revision, "commit_hash", None) or str(id(revision)) @@ -375,18 +454,19 @@ def _repo_non_gguf_model_payload(repo_info) -> _CachedNonGgufPayload: or "unknown" ) if model_format == "adapter": - size_bytes = sum(adapter_blobs.values()) + selected_blobs = adapter_blobs elif model_format == "safetensors": - size_bytes = sum(safetensors_blobs.values()) + selected_blobs = safetensors_blobs elif model_format == "checkpoint": - size_bytes = sum(checkpoint_blobs.values()) + selected_blobs = checkpoint_blobs else: - size_bytes = sum(all_weight_blobs.values()) + selected_blobs = all_weight_blobs return _CachedNonGgufPayload( - size_bytes = size_bytes, + size_bytes = sum(size for size, _mtime in selected_blobs.values()), has_runnable_weights = model_format != "unknown", model_format = model_format, + last_modified = max((mtime for _size, mtime in selected_blobs.values()), default = 0.0), ) @@ -475,6 +555,15 @@ def _scan_cached_models() -> list[dict]: if str(repo_info.repo_type) != "model": continue repo_id = repo_info.repo_id + repo_path = Path(repo_info.repo_path) + snapshot_path = _cached_model_snapshot_path(repo_path) + # The non-GGUF embedder has no variant downloads; always hide. + if _is_hidden_infra_repo( + repo_id, + str(repo_path), + str(snapshot_path) if snapshot_path is not None else None, + ): + continue has_main_gguf = _repo_has_gguf_files(repo_info) payload = _repo_non_gguf_model_payload(repo_info) if payload.size_bytes == 0: @@ -486,7 +575,6 @@ def _scan_cached_models() -> list[dict]: continue key = repo_id.lower() existing = seen_lower.get(key) - repo_path = Path(repo_info.repo_path) snapshot_partial = hf_cache_scan.is_snapshot_partial( "model", repo_id, @@ -508,6 +596,12 @@ def _scan_cached_models() -> list[dict]: ), **_cached_model_local_metadata(repo_path), } + last_modified = max( + payload.last_modified, + (existing or {}).get("last_modified", 0.0), + ) + if last_modified > 0: + row["last_modified"] = last_modified row.update( _cache_inventory_fields( repo_id, @@ -517,6 +611,8 @@ def _scan_cached_models() -> list[dict]: ) if _prefer_cache_row(row, existing): seen_lower[key] = row + elif last_modified > existing.get("last_modified", 0.0): + existing["last_modified"] = last_modified except Exception as e: repo_label = getattr(repo_info, "repo_id", "") logger.warning(f"Skipping cached model repo {repo_label}: {e}") diff --git a/studio/backend/hub/services/models/folder_browser.py b/studio/backend/hub/services/models/folder_browser.py index d56b62c318..7d9c3ac665 100644 --- a/studio/backend/hub/services/models/folder_browser.py +++ b/studio/backend/hub/services/models/folder_browser.py @@ -165,7 +165,7 @@ def _looks_like_model_dir(directory: Path) -> bool: def _build_browse_allowlist( media_roots: Optional[list[Path]] = None, drive_roots: Optional[list[Path]] = None ) -> list[Path]: - """Root directories the browser may walk (also seeds the suggestion chips): HOME, resolved HF cache dirs, Studio outputs/exports/root, registered scan folders, and well-known local-LLM dirs. Each is added only if it resolves to a real directory so the sandbox has no dead boundary. + """Root directories the browser may walk (also seeds the suggestion chips): HOME, resolved HF cache dirs, Unsloth outputs/exports/root, registered scan folders, and well-known local-LLM dirs. Each is added only if it resolves to a real directory so the sandbox has no dead boundary. *media_roots* / *drive_roots* let the caller pass already-probed removable-media and Windows drive roots so they aren't scanned again (a diff --git a/studio/backend/hub/services/models/local_inventory.py b/studio/backend/hub/services/models/local_inventory.py index a3782efead..b34532fa35 100644 --- a/studio/backend/hub/services/models/local_inventory.py +++ b/studio/backend/hub/services/models/local_inventory.py @@ -36,6 +36,7 @@ from hub.utils.paths import ( ) from hub.services.models import common as model_common from hub.services.models.ollama import scan_ollama_dir +from utils.hidden_models import is_hidden_model logger = get_logger(__name__) _MAX_MODELS_PER_CUSTOM_FOLDER = 200 @@ -623,6 +624,20 @@ def _dedupe_local_models(local_models: List[LocalModelInfo]) -> list[LocalModelI ) +def _filter_hidden_models(local_models: List[LocalModelInfo]) -> list[LocalModelInfo]: + """Remove infrastructure-only models from the shared local inventory.""" + visible: list[LocalModelInfo] = [] + for model in local_models: + resolved_cache_path = ( + hf_cache_scan.resolve_hf_cache_realpath(Path(model.path)) + if model.source == "hf_cache" + else None + ) + if not is_hidden_model(model.id, model.model_id, model.path, resolved_cache_path): + visible.append(model) + return visible + + async def list_local_models_response(models_dir: str = "./models") -> LocalModelListResponse: """List local model candidates from every supported on-device source.""" hf_cache_dir = _resolve_hf_cache_dir() @@ -653,7 +668,7 @@ async def list_local_models_response(models_dir: str = "./models") -> LocalModel ollama_dirs, ) local_models += await _collect_models_from_custom_folders() - models = _dedupe_local_models(local_models) + models = _dedupe_local_models(_filter_hidden_models(local_models)) return LocalModelListResponse( models_dir = str(models_root), diff --git a/studio/backend/hub/services/models/ollama.py b/studio/backend/hub/services/models/ollama.py index 96a4114620..2ccdbb44f1 100644 --- a/studio/backend/hub/services/models/ollama.py +++ b/studio/backend/hub/services/models/ollama.py @@ -85,7 +85,7 @@ def _contained_link_path(link_dir: Path, link_name: str) -> Optional[Path]: def _ollama_links_dir(ollama_dir: Path) -> Optional[Path]: - """Writable directory for Ollama ``.gguf`` symlinks. Prefers ``/.studio_links/`` next to the blobs; falls back to Studio's cache (read-only system installs), then the temp dir (sandboxed installs).""" + """Writable directory for Ollama ``.gguf`` symlinks. Prefers ``/.studio_links/`` next to the blobs; falls back to Unsloth's cache (read-only system installs), then the temp dir (sandboxed installs).""" def _ensure_writable_dir(path: Path) -> Optional[Path]: try: diff --git a/studio/backend/hub/tests/test_model_services.py b/studio/backend/hub/tests/test_model_services.py index 2c33e09b2b..693d945ee1 100644 --- a/studio/backend/hub/tests/test_model_services.py +++ b/studio/backend/hub/tests/test_model_services.py @@ -439,6 +439,287 @@ def test_cached_gguf_scan_includes_variant_state_without_completed_gguf(monkeypa assert row["capabilities"]["requires_variant"] is True +def test_cached_gguf_scan_hides_infra_repos_without_user_downloads(monkeypatch, tmp_path): + probe = _repo( + "ggml-org/models", + [_file("tinyllamas/stories260K.gguf", 1_200_000)], + tmp_path / "probe", + ) + embedder = _repo( + "unsloth/bge-small-en-v1.5-GGUF", + [_file("bge-small-en-v1.5-f16.gguf", 60_000_000)], + tmp_path / "embedder", + ) + chat = _repo("Org/Chat-GGUF", [_file("Q4_K_M.gguf", 100)], tmp_path / "chat") + monkeypatch.setattr( + cache_inventory, + "all_hf_cache_scans", + lambda: [SimpleNamespace(repos = [probe, embedder, chat])], + ) + monkeypatch.setattr( + cache_inventory.hf_cache_scan, + "is_gguf_repo_partial", + lambda _repo_id, _path: False, + ) + + result = {"cached": cache_inventory._scan_cached_gguf()} + + assert [row["repo_id"] for row in result["cached"]] == ["Org/Chat-GGUF"] + + +def test_cached_gguf_scan_keeps_infra_repo_with_user_downloaded_variant(monkeypatch, tmp_path): + monkeypatch.setattr(state_dir, "cache_root", lambda: tmp_path / "state") + embedder = _repo( + "unsloth/bge-small-en-v1.5-GGUF", + [ + _file("bge-small-en-v1.5-f16.gguf", 60_000_000), + _file("bge-small-en-v1.5-Q8_0.gguf", 35_000_000), + ], + tmp_path / "embedder", + ) + # Variant manifests only exist for user Hub downloads, not auto-downloads. + assert download_manifest.write_manifest( + "model", + "unsloth/bge-small-en-v1.5-GGUF", + "Q8_0", + [download_manifest.ExpectedFile(path = "bge-small-en-v1.5-Q8_0.gguf", size = 35_000_000)], + "http", + ) + monkeypatch.setattr( + cache_inventory, + "all_hf_cache_scans", + lambda: [SimpleNamespace(repos = [embedder])], + ) + monkeypatch.setattr( + cache_inventory.hf_cache_scan, + "is_gguf_repo_partial", + lambda _repo_id, _path: False, + ) + + result = {"cached": cache_inventory._scan_cached_gguf()} + + assert [row["repo_id"] for row in result["cached"]] == ["unsloth/bge-small-en-v1.5-GGUF"] + assert result["cached"][0]["capabilities"]["can_chat"] is False + + +def test_cached_models_scan_hides_non_gguf_embedder(monkeypatch, tmp_path): + embedder_path = tmp_path / "hub" / "models--unsloth--bge-small-en-v1.5" + embedder_path.mkdir(parents = True) + embedder = _repo( + "unsloth/bge-small-en-v1.5", + [_file("config.json", 12), _file("model.safetensors", 130_000_000)], + embedder_path, + ) + chat_path = tmp_path / "hub" / "models--Org--Chat" + chat_path.mkdir(parents = True) + chat = _repo( + "Org/Chat", + [_file("config.json", 12), _file("model.safetensors", 100)], + chat_path, + ) + monkeypatch.setattr( + cache_inventory, + "all_hf_cache_scans", + lambda: [SimpleNamespace(repos = [embedder, chat])], + ) + monkeypatch.setattr( + cache_inventory.hf_cache_scan, + "is_snapshot_partial", + lambda _kind, _repo_id, _path: False, + ) + + result = {"cached": cache_inventory._scan_cached_models()} + + assert [row["repo_id"] for row in result["cached"]] == ["Org/Chat"] + + +def test_cached_scans_hide_embedders_configured_by_cache_path(monkeypatch, tmp_path): + from core.rag import config as rag_config + + gguf_path = tmp_path / "hub" / "models--Org--PathEmbedder-GGUF" + gguf_path.mkdir(parents = True) + gguf = _repo( + "Org/PathEmbedder-GGUF", + [_file("model-F16.gguf", 60_000_000)], + gguf_path, + ) + model_path = tmp_path / "hub" / "models--Org--PathEmbedder" + model_path.mkdir(parents = True) + model = _repo( + "Org/PathEmbedder", + [_file("config.json", 12), _file("model.safetensors", 130_000_000)], + model_path, + ) + monkeypatch.setattr( + rag_config, + "effective_embedding_model", + lambda: str(model_path), + ) + monkeypatch.setattr( + rag_config, + "effective_gguf_repo", + lambda: str(gguf_path), + ) + monkeypatch.setattr( + cache_inventory, + "all_hf_cache_scans", + lambda: [SimpleNamespace(repos = [gguf, model])], + ) + monkeypatch.setattr( + cache_inventory.hf_cache_scan, + "is_gguf_repo_partial", + lambda _repo_id, _path: False, + ) + monkeypatch.setattr( + cache_inventory.hf_cache_scan, + "is_snapshot_partial", + lambda _kind, _repo_id, _path: False, + ) + + assert cache_inventory._scan_cached_gguf() == [] + assert cache_inventory._scan_cached_models() == [] + + +def test_cached_scans_hide_embedders_configured_by_snapshot_path(monkeypatch, tmp_path): + from core.rag import config as rag_config + + gguf_path = tmp_path / "hub" / "models--Org--SnapshotEmbedder-GGUF" + gguf_snapshot = gguf_path / "snapshots" / "gguf-revision" + gguf_snapshot.mkdir(parents = True) + gguf = _repo( + "Org/SnapshotEmbedder-GGUF", + [_file("model-F16.gguf", 60_000_000)], + gguf_path, + ) + model_path = tmp_path / "hub" / "models--Org--SnapshotEmbedder" + model_snapshot = model_path / "snapshots" / "model-revision" + model_snapshot.mkdir(parents = True) + model = _repo( + "Org/SnapshotEmbedder", + [_file("config.json", 12), _file("model.safetensors", 130_000_000)], + model_path, + ) + monkeypatch.setattr( + rag_config, + "effective_embedding_model", + lambda: str(model_snapshot), + ) + monkeypatch.setattr( + rag_config, + "effective_gguf_repo", + lambda: str(gguf_snapshot), + ) + monkeypatch.setattr( + cache_inventory, + "all_hf_cache_scans", + lambda: [SimpleNamespace(repos = [gguf, model])], + ) + + def _resolve_snapshot(repo_path): + return str( + { + gguf_path: gguf_snapshot, + model_path: model_snapshot, + }.get(Path(repo_path), Path(repo_path)) + ) + + monkeypatch.setattr( + cache_inventory.hf_cache_scan, + "resolve_hf_cache_realpath", + _resolve_snapshot, + ) + monkeypatch.setattr( + cache_inventory.hf_cache_scan, + "is_gguf_repo_partial", + lambda _repo_id, _path: False, + ) + monkeypatch.setattr( + cache_inventory.hf_cache_scan, + "is_snapshot_partial", + lambda _kind, _repo_id, _path: False, + ) + + assert cache_inventory._scan_cached_gguf() == [] + assert cache_inventory._scan_cached_models() == [] + + +def test_cached_models_scan_keeps_unrelated_repo_with_custom_generic_embedder( + monkeypatch, tmp_path +): + # A custom embedder with a generic basename ("org/model") must be hidden by + # EXACT repo-id match only. An unrelated cached chat model whose id merely + # contains "model" (e.g. "user/model-chat") must stay on device: substring + # basename matching used to drop real chat models from the inventory. + from core.rag import config as rag_config + + monkeypatch.setattr(rag_config, "effective_embedding_model", lambda: "org/model") + monkeypatch.setattr(rag_config, "effective_gguf_repo", lambda: "org/model-GGUF") + + def _model_repo(repo_id: str): + path = tmp_path / "hub" / f"models--{repo_id.replace('/', '--')}" + path.mkdir(parents = True) + return _repo( + repo_id, + [_file("config.json", 12), _file("model.safetensors", 100)], + path, + ) + + embedder = _model_repo("org/model") + chat = _model_repo("user/model-chat") + monkeypatch.setattr( + cache_inventory, + "all_hf_cache_scans", + lambda: [SimpleNamespace(repos = [embedder, chat])], + ) + monkeypatch.setattr( + cache_inventory.hf_cache_scan, + "is_snapshot_partial", + lambda _kind, _repo_id, _path: False, + ) + + result = {"cached": cache_inventory._scan_cached_models()} + + assert [row["repo_id"] for row in result["cached"]] == ["user/model-chat"] + + +def test_cached_scans_hide_stale_default_embedder_after_custom_setting(monkeypatch, tmp_path): + from core.rag import config as rag_config + + monkeypatch.setattr(rag_config, "effective_embedding_model", lambda: "org/custom") + monkeypatch.setattr(rag_config, "effective_gguf_repo", lambda: "org/custom-GGUF") + + gguf = _repo( + "unsloth/bge-small-en-v1.5-GGUF", + [_file("bge-small-en-v1.5-f16.gguf", 60_000_000)], + tmp_path / "default-gguf", + ) + weights_path = tmp_path / "hub" / "models--unsloth--bge-small-en-v1.5" + weights_path.mkdir(parents = True) + weights = _repo( + "unsloth/bge-small-en-v1.5", + [_file("config.json", 12), _file("model.safetensors", 130_000_000)], + weights_path, + ) + monkeypatch.setattr( + cache_inventory, + "all_hf_cache_scans", + lambda: [SimpleNamespace(repos = [gguf, weights])], + ) + monkeypatch.setattr( + cache_inventory.hf_cache_scan, + "is_gguf_repo_partial", + lambda _repo_id, _path: False, + ) + monkeypatch.setattr( + cache_inventory.hf_cache_scan, + "is_snapshot_partial", + lambda _kind, _repo_id, _path: False, + ) + + assert cache_inventory._scan_cached_gguf() == [] + assert cache_inventory._scan_cached_models() == [] + + def test_gguf_variant_requirements_include_split_files_and_preferred_mmproj(): requirements = gguf_variants._build_gguf_variant_requirements( [ @@ -1610,6 +1891,63 @@ def test_hf_cache_scan_uses_gguf_partial_row_for_variant_state(monkeypatch, tmp_ assert rows[0].capabilities.requires_variant is True +def test_local_inventory_filters_custom_embedder_hf_cache_row(monkeypatch, tmp_path): + from core.rag import config as rag_config + + monkeypatch.setattr(rag_config, "effective_embedding_model", lambda: "org/embedder") + monkeypatch.setattr(rag_config, "effective_gguf_repo", lambda: "org/embedder-GGUF") + + def _row(repo_id: str): + repo_path = tmp_path / f"models--{repo_id.replace('/', '--')}" + return model_common._local_model_info( + scan_path = repo_path, + load_path = repo_path, + source = "hf_cache", + model_format = "safetensors", + model_id = repo_id, + ) + + rows = local_inventory._filter_hidden_models([_row("org/embedder"), _row("org/chat-model")]) + + assert [row.model_id for row in rows] == ["org/chat-model"] + + +def test_local_inventory_filters_embedder_configured_by_snapshot_path(monkeypatch, tmp_path): + from core.rag import config as rag_config + + embedder_path = tmp_path / "hub" / "models--org--embedder" + embedder_snapshot = embedder_path / "snapshots" / "revision" + embedder_snapshot.mkdir(parents = True) + chat_path = tmp_path / "hub" / "models--org--chat-model" + chat_path.mkdir(parents = True) + monkeypatch.setattr( + rag_config, + "effective_embedding_model", + lambda: str(embedder_snapshot), + ) + monkeypatch.setattr(rag_config, "effective_gguf_repo", lambda: "org/embedder-GGUF") + monkeypatch.setattr( + local_inventory.hf_cache_scan, + "resolve_hf_cache_realpath", + lambda path: str(embedder_snapshot) if Path(path) == embedder_path else str(path), + ) + + def _row(repo_id: str, repo_path: Path): + return model_common._local_model_info( + scan_path = repo_path, + load_path = repo_path, + source = "hf_cache", + model_format = "safetensors", + model_id = repo_id, + ) + + rows = local_inventory._filter_hidden_models( + [_row("org/embedder", embedder_path), _row("org/chat-model", chat_path)] + ) + + assert [row.model_id for row in rows] == ["org/chat-model"] + + def test_model_download_job_helpers_preserve_idle_shape(): key = downloads._download_job_key("Org/Model", None) status = downloads._job_status(key) diff --git a/studio/backend/hub/utils/state_dir.py b/studio/backend/hub/utils/state_dir.py index 183e934724..898c03c87d 100644 --- a/studio/backend/hub/utils/state_dir.py +++ b/studio/backend/hub/utils/state_dir.py @@ -3,7 +3,7 @@ """Filesystem layout for Hub download state. -State directory sits beside HF's cache (under Studio's own cache root) +State directory sits beside HF's cache (under Unsloth's own cache root) so it survives ``huggingface-cli delete-cache`` and any other HF-side cache lifecycle. Two subdirectories: diff --git a/studio/backend/main.py b/studio/backend/main.py index 387f509ef9..74fe2c414f 100644 --- a/studio/backend/main.py +++ b/studio/backend/main.py @@ -19,7 +19,7 @@ os.environ["PYTHONWARNINGS"] = "ignore" # Pin GPU index ordering to PCI bus id before any torch import creates a CUDA # context. Without this, torch/CUDA default to FASTEST_FIRST while nvidia-smi -# (and Studio's VRAM probes) use PCI-bus order, so a GPU index chosen from +# (and Unsloth's VRAM probes) use PCI-bus order, so a GPU index chosen from # nvidia-smi data can resolve to a different physical card via # CUDA_VISIBLE_DEVICES. setdefault so an explicit user override wins. See # utils/hardware/hardware.py for the full rationale; set here too so the entry @@ -93,7 +93,7 @@ if sys.platform == "win32": # ── Windows AMD ROCm: make hipInfo.exe resolvable for subprocess probes ── # bitsandbytes' get_rocm_gpu_arch() runs `hipinfo.exe` via PATH at import # time; the AMD torch wheel ships it in the venv Scripts dir, which is on - # PATH only when the venv is activated -- Studio launches python directly. + # PATH only when the venv is activated -- Unsloth launches python directly. # Without this, every bitsandbytes import logs a scary (but harmless) # "Could not detect ROCm GPU architecture: [WinError 2]" ERROR + WARNING. # Gated on the file existing: only AMD ROCm wheels ship hipInfo.exe, so @@ -252,7 +252,7 @@ def _read_studio_install_id() -> str: Returns "" when absent or not a 64-char lowercase-hex token; then /api/health emits "" and the launcher accepts any healthy backend. - Carries no install-path info (matters when Studio runs -H 0.0.0.0).""" + Carries no install-path info (matters when Unsloth runs -H 0.0.0.0).""" try: token = (_STUDIO_ROOT_RESOLVED / "share" / "studio_install_id").read_text().strip() except (OSError, ValueError): @@ -289,6 +289,7 @@ from fastapi import Depends, FastAPI, HTTPException, Query, Request from fastapi.middleware.cors import CORSMiddleware from fastapi.staticfiles import StaticFiles from fastapi.responses import FileResponse, HTMLResponse, Response +from starlette.middleware.gzip import GZipMiddleware from pathlib import Path from datetime import datetime @@ -313,7 +314,9 @@ from routes.preview import router as preview_router from hub.routes import ( inventory_router as hub_inventory_router, datasets_router as hub_datasets_router, + token_router as hub_token_router, ) +from picker.routes import templates_router as picker_templates_router from hub.schemas.downloads import TransportCapabilities from hub.utils.download_registry import ( get_download_transport_capabilities, @@ -553,8 +556,9 @@ async def lifespan(app: FastAPI): app.state.research_supervisor.start() # Idle auto-unload loop (no-op unless the OpenAI auto-unload TTL is set). - from core.inference.llama_keepwarm import idle_unload_loop + from core.inference.llama_keepwarm import idle_unload_loop, sweep_slot_save_dir + sweep_slot_save_dir() app.state.idle_unload_task = asyncio.create_task(idle_unload_loop()) # Initialize RSA key pair for API key encryption (external providers). @@ -579,7 +583,7 @@ async def lifespan(app: FastAPI): print("DEFAULT ADMIN ACCOUNT CREATED") print(f" username: {storage.DEFAULT_ADMIN_USERNAME}") print(f" password saved to: {bootstrap_path}") - print(" Open the Studio UI to sign in and change it.") + print(" Open the Unsloth UI to sign in and change it.") print("=" * 60 + "\n") else: app.state.bootstrap_password = ( @@ -623,7 +627,7 @@ app = FastAPI( ) # The MCP surface is opt-in because it can start GPU jobs and write model -# artifacts. Mount it only when explicitly enabled by the Studio process. +# artifacts. Mount it only when explicitly enabled by the Unsloth process. if os.environ.get("UNSLOTH_STUDIO_ENABLE_MCP") == "1": from fastmcp.utilities.lifespan import combine_lifespans @@ -789,6 +793,7 @@ _BODY_PROTECTED_PREFIXES = ( "/v1/completions", "/p/", "/api/inference", + "/api/picker", "/api/data-recipe", "/api/datasets", "/api/hub", @@ -1002,7 +1007,7 @@ app.include_router(models_router, prefix = "/api/models", tags = ["models"]) app.include_router(chat_history_router, prefix = "/api/chat", tags = ["chat"]) app.include_router(research_runs_router, prefix = "/api/chat/research-runs", tags = ["research-runs"]) app.include_router(inference_router, prefix = "/api/inference", tags = ["inference"]) -# Studio-only inference endpoints (cancel, etc.) are NOT exposed on the /v1 +# Unsloth-only inference endpoints (cancel, etc.) are NOT exposed on the /v1 # OpenAI-compat prefix below. app.include_router(inference_studio_router, prefix = "/api/inference", tags = ["inference"]) @@ -1021,6 +1026,8 @@ app.include_router(rag_router, prefix = "/api/rag", tags = ["rag"]) app.include_router(training_history_router, prefix = "/api/train", tags = ["training-history"]) app.include_router(hub_inventory_router, prefix = "/api/hub", tags = ["hub"]) app.include_router(hub_datasets_router, prefix = "/api/hub/datasets", tags = ["hub"]) +app.include_router(picker_templates_router, prefix = "/api/picker", tags = ["picker"]) +app.include_router(hub_token_router, prefix = "/api/hub", tags = ["hub"]) # Re-wrap client-error responses on the /v1/* surface into OpenAI/Anthropic # error envelopes; non-/v1 paths keep FastAPI's default {"detail": ...} shape. @@ -1109,7 +1116,7 @@ def studio_install_source(_current_subject: str = Depends(get_current_subject)): @app.get("/api/studio/update-status") def studio_update_status(_current_subject: str = Depends(get_current_subject)): - """Return source-aware manual update status for browser-served Studio.""" + """Return source-aware manual update status for browser-served Unsloth.""" return get_studio_update_status(UNSLOTH_VERSION) @@ -1177,17 +1184,35 @@ def _get_cached_system_gpu_info(logger) -> dict[str, Any]: util = util_devices.get(idx, {}) total_vram = util.get("vram_total_gb") or dev.get("memory_total_gb") or 0 - used_vram = util.get("vram_used_gb") or 0 + # Keep None (usage unknown, e.g. Windows ROCm perf counter) so the UI + # shows unknown, not a fabricated 0 used / full free. + used_vram = util.get("vram_used_gb") enriched_dev = dict(dev) enriched_dev["vram_used_gb"] = used_vram - enriched_dev["vram_free_gb"] = round(total_vram - used_vram, 2) if total_vram else 0 + enriched_dev["vram_free_gb"] = ( + round(total_vram - used_vram, 2) if total_vram and used_vram is not None else None + ) enriched_dev["vram_utilization_pct"] = util.get("vram_utilization_pct") enriched_devices.append(enriched_dev) + # Whether GGUF loads accept an explicit gpu_ids pick: /load and + # /validate 400 picks on XPU hosts (no visibility mask speaks torch-xpu + # ordinals) and on Vulkan-only builds (--device pins ggml's own + # ordinals), so the picker must not offer them. + try: + from core.inference.llama_cpp import LlamaCppBackend + from utils.hardware import DeviceType, get_device + gpu_ids_supported = ( + get_device() != DeviceType.XPU and not LlamaCppBackend._is_vulkan_backend() + ) + except Exception as e: + logger.debug(f"Could not resolve gpu_ids support: {e}") + gpu_ids_supported = True gpu_info = { "available": visibility_info.get("available", False), "devices": enriched_devices, + "gguf_gpu_ids_supported": gpu_ids_supported, } _system_gpu_cache = (time.monotonic(), gpu_info) return gpu_info @@ -1517,6 +1542,34 @@ def _should_inject_bootstrap(request: Request) -> bool: return _is_local_bootstrap_request(request) +_IMMUTABLE_ASSET_CACHE_CONTROL = "public, max-age=31536000, immutable" + + +class ImmutableStaticFiles(StaticFiles): + """Serve Vite's content-hashed assets without browser revalidation.""" + + def file_response( + self, + full_path, + stat_result, + scope, + status_code = 200, + ): + response = super().file_response(full_path, stat_result, scope, status_code) + response.headers["Cache-Control"] = _IMMUTABLE_ASSET_CACHE_CONTROL + return response + + +class _AssetGZipMiddleware(GZipMiddleware): + """Serve range requests uncompressed; gzip + 206 mislabels Content-Range.""" + + async def __call__(self, scope, receive, send): + if scope["type"] == "http" and any(key == b"range" for key, _ in scope["headers"]): + await self.app(scope, receive, send) + return + await super().__call__(scope, receive, send) + + def setup_frontend(app: FastAPI, build_path: Path): """Mount frontend static files (optional)""" if not build_path.exists(): @@ -1524,7 +1577,12 @@ def setup_frontend(app: FastAPI, build_path: Path): assets_dir = build_path / "assets" if assets_dir.exists(): - app.mount("/assets", StaticFiles(directory = assets_dir), name = "assets") + assets_app = _AssetGZipMiddleware( + ImmutableStaticFiles(directory = assets_dir), + minimum_size = 1024, + compresslevel = 6, + ) + app.mount("/assets", assets_app, name = "assets") def _build_index_response(request: Request) -> Response: content = (build_path / "index.html").read_bytes() diff --git a/studio/backend/mcp_server.py b/studio/backend/mcp_server.py index f837f46425..e93490411d 100644 --- a/studio/backend/mcp_server.py +++ b/studio/backend/mcp_server.py @@ -3,7 +3,7 @@ """Curated MCP tools for driving an Unsloth Studio instance. -The MCP surface deliberately wraps the existing Studio services instead of +The MCP surface deliberately wraps the existing Unsloth services instead of duplicating training or export logic. It is opt-in because several tools can start GPU work or write model artifacts. """ @@ -17,14 +17,14 @@ from fastmcp import FastMCP class BearerTokenMiddleware: - """Require an exact bearer token when Studio MCP is exposed remotely.""" + """Require an exact bearer token when Unsloth MCP is exposed remotely.""" def __init__(self, app: Any, token: str) -> None: if not token or not token.strip(): - raise ValueError("Studio MCP bearer token must be a non-empty value") + raise ValueError("Unsloth MCP bearer token must be a non-empty value") if not token.isascii(): # A non-ASCII token cannot be sent in an HTTP header; reject it here. - raise ValueError("Studio MCP bearer token must contain ASCII characters only") + raise ValueError("Unsloth MCP bearer token must contain ASCII characters only") self.app = app # Compare on raw header bytes: str hmac.compare_digest raises on non-ASCII # input, which would surface as a 500 instead of a clean 401. @@ -76,18 +76,18 @@ def _dump(value: Any) -> Any: def _clamp(value: int, low: int, high: int) -> int: """Clamp an MCP-supplied integer into an inclusive range. - MCP tools call the Studio route functions directly, which skips FastAPI's + MCP tools call the Unsloth route functions directly, which skips FastAPI's Query(ge=, le=) validation, so we re-apply the same bounds here. """ return max(low, min(value, high)) def create_studio_mcp() -> FastMCP: - """Create the Studio MCP server and register the high-value tools.""" + """Create the Unsloth MCP server and register the high-value tools.""" mcp = FastMCP( "Unsloth Studio", instructions = ( - "Use read tools to inspect the local Studio state before starting GPU work. " + "Use read tools to inspect the local Unsloth state before starting GPU work. " "Training and export tools can consume substantial VRAM and write files. " "Never expose tokens or local paths from tool results unless the user asks." ), @@ -116,7 +116,7 @@ def create_studio_mcp() -> FastMCP: @mcp.tool async def list_local_models(models_dir: str = "./models") -> dict[str, Any]: - """List local and cached models available to Studio.""" + """List local and cached models available to Unsloth.""" from routes.models import list_local_models as list_models return _dump(await list_models(models_dir = models_dir, current_subject = "mcp")) @@ -128,9 +128,9 @@ def create_studio_mcp() -> FastMCP: @mcp.tool async def start_training(config: dict[str, Any]) -> dict[str, Any]: - """Start a validated Studio training job from a TrainingStartRequest-shaped object. + """Start a validated Unsloth training job from a TrainingStartRequest-shaped object. - The config is validated by the same Pydantic model used by the Studio UI. + The config is validated by the same Pydantic model used by the Unsloth UI. Call get_training_status first and do not start work while another job runs. """ from models import TrainingStartRequest @@ -138,7 +138,7 @@ def create_studio_mcp() -> FastMCP: request = TrainingStartRequest.model_validate(config) # Pass via_api_key explicitly (a direct call leaves it a Depends object). - # MCP drives Studio like the UI session, so it coexists and frees VRAM. + # MCP drives Unsloth like the UI session, so it coexists and frees VRAM. return _dump(await start(request, current_subject = "mcp", via_api_key = False)) @mcp.tool @@ -159,7 +159,7 @@ def create_studio_mcp() -> FastMCP: @mcp.tool def validate_recipe(recipe: dict[str, Any]) -> dict[str, Any]: - """Validate a Data Recipe with the same validator used by Studio.""" + """Validate a Data Recipe with the same validator used by Unsloth.""" from models.data_recipe import RecipePayload from routes.data_recipe.validate import validate @@ -225,7 +225,7 @@ def create_studio_mcp() -> FastMCP: imatrix: bool = False, imatrix_path: str | None = None, ) -> dict[str, Any]: - """Export the loaded model to GGUF using Studio's existing path validation. + """Export the loaded model to GGUF using Unsloth's existing path validation. quantization_method may be a single method or a list to produce several GGUFs from one load. Pass hf_token when push_to_hub is set (the backend diff --git a/studio/backend/models/inference.py b/studio/backend/models/inference.py index 3ae974448e..580a74dddf 100644 --- a/studio/backend/models/inference.py +++ b/studio/backend/models/inference.py @@ -18,6 +18,8 @@ from pydantic import ( model_validator, ) +from picker.schemas import MAX_CHAT_TEMPLATE_BYTES + class LoadRequest(BaseModel): """Request to load a model for inference""" @@ -54,8 +56,16 @@ class LoadRequest(BaseModel): @field_validator("chat_template_override") @classmethod def normalize_blank_chat_template_override(cls, value: Optional[str]) -> Optional[str]: - if value is not None and value.strip() == "": + if value is None: return None + # Char count is a lower bound on UTF-8 byte length: reject an oversized + # template before spending work encoding it. + if len(value) > MAX_CHAT_TEMPLATE_BYTES: + raise ValueError(f"Chat template exceeds the {MAX_CHAT_TEMPLATE_BYTES}-byte limit.") + if value.strip() == "": + return None + if len(value.encode("utf-8")) > MAX_CHAT_TEMPLATE_BYTES: + raise ValueError(f"Chat template exceeds the {MAX_CHAT_TEMPLATE_BYTES}-byte limit.") return value cache_type_kv: Optional[str] = Field( @@ -64,7 +74,7 @@ class LoadRequest(BaseModel): ) gpu_ids: Optional[List[int]] = Field( None, - description = "Physical GPU indices to use, for example [0, 1]. Omit or pass [] to use automatic selection. Explicit gpu_ids are unsupported when the parent CUDA_VISIBLE_DEVICES uses UUID/MIG entries. Not supported for GGUF models.", + description = "Physical GPU indices to use, for example [0, 1]. Omit or pass [] to use automatic selection. Explicit gpu_ids are unsupported when the parent CUDA_VISIBLE_DEVICES uses UUID/MIG entries. For GGUF models the picked devices are pinned via CUDA/HIP_VISIBLE_DEVICES.", ) speculative_type: Optional[str] = Field( None, @@ -100,12 +110,72 @@ class LoadRequest(BaseModel): "No effect on a single GPU. Ignored for non-GGUF models." ), ) + gpu_memory_mode: Literal["auto", "manual"] = Field( + "auto", + description = ( + "GPU memory strategy for GGUF models. 'auto' (default): Unsloth " + "selects GPUs and caps context to fit VRAM. 'manual': you own the " + "offload. Leave gpu_layers at -1 (Auto) to hand memory management to " + "llama.cpp's --fit (no device masking, no context auto-reduce, no " + "gpu-layer/tensor-split planning); set gpu_layers >= 0 to pin layers " + "and n_cpu_moe yourself (--fit off), with tensor_parallel still " + "applying (split by free VRAM unless tensor_split is set, no planner). " + "Ignored for non-GGUF." + ), + ) + gpu_layers: int = Field( + -1, + ge = -1, + description = ( + "Manual mode only: number of layers to offload to the GPU " + "(--gpu-layers, with --fit off). A value >= the model's layer count " + "offloads all of them. -1 = Auto: hand layer + context sizing to " + "llama.cpp's --fit. Ignored unless gpu_memory_mode is 'manual'." + ), + ) + n_cpu_moe: int = Field( + 0, + ge = 0, + description = ( + "Manual mode only: keep the first N MoE expert layers on the CPU " + "(--n-cpu-moe) to save VRAM on MoE models. 0 = none, N = number of " + "MoE layers offloaded (the backend offsets past any leading dense " + "layers). Ignored unless gpu_memory_mode is 'manual' with gpu_layers >= 0." + ), + ) + tensor_split: Optional[List[float]] = Field( + None, + description = ( + "Manual mode only: relative share of the model per GPU (--tensor-split), " + "in the order of the GPUs in use, e.g. [2, 1] for 2:1. Omit it to let " + "llama.cpp use its default, which splits by free VRAM. Any list given is " + "passed through as-is, so send [1, 1] to force an even split. Ignored " + "unless gpu_memory_mode is 'manual' with gpu_layers >= 0." + ), + ) + + @field_validator("tensor_split") + @classmethod + def _reject_degenerate_tensor_split(cls, value: Optional[List[float]]) -> Optional[List[float]]: + # A negative / non-finite / all-zero split is silently dropped at launch + # (stored as None) yet still compared raw in the reload dedupe, so an + # identical Apply reloads forever. Reject it up front; [] = no split. + if not value: + return value + import math + + if any((not math.isfinite(v)) or v < 0 for v in value): + raise ValueError("tensor_split entries must be finite and non-negative") + if sum(value) <= 0: + raise ValueError("tensor_split must have a positive total") + return value + llama_extra_args: Optional[List[str]] = Field( None, description = ( "Extra arguments forwarded verbatim to llama-server for GGUF models. " "One token per list entry, e.g. ['--top-k', '20', '--seed', '42']. " - "Studio-managed flags (model identity, port, context length, GPU placement, " + "Unsloth-managed flags (model identity, port, context length, GPU placement, " "auth, UI/server mode) are rejected. Ignored for non-GGUF models." ), ) @@ -133,11 +203,26 @@ class ValidateModelRequest(BaseModel): max_seq_length: int = Field(0, ge = 0, le = 1048576) load_in_4bit: bool = Field(True) gpu_ids: Optional[List[int]] = Field(None) + gpu_memory_mode: Literal["auto", "manual"] = Field( + "auto", + description = ( + "GGUF GPU-memory strategy intended for the follow-up load. Manual " + "placement bypasses the training coexistence estimate: Auto layers " + "delegate fitting to llama.cpp, while explicit layers are user-owned." + ), + ) include_context_length: bool = Field( False, description = "Also read the native context length from the local GGUF header. " "Opt-in so the normal load preflight doesn't pay for a cache scan it doesn't need.", ) + include_chat_template: bool = Field( + False, + description = "Also read the embedded chat template from the local GGUF header, so a " + "native (picked / drag-drop) file's default template can be shown before it is loaded. " + "Opt-in and, like include_context_length, a metadata-only probe that skips the training " + "guard. Only the leased file's own embedded template is read, never sibling sidecars.", + ) class TransformersUpgradeInfo(BaseModel): @@ -151,13 +236,13 @@ class TransformersUpgradeInfo(BaseModel): ) supported_in_pypi: bool = Field( False, - description = "True if the latest PyPI release ships this model_type; Studio can " + description = "True if the latest PyPI release ships this model_type; Unsloth can " "install it into a persistent sidecar after user consent.", ) supported_in_main: bool = Field( False, description = "True if transformers GitHub main ships this model_type (dev-only; " - "not installable through Studio yet).", + "not installable through Unsloth yet).", ) @@ -188,6 +273,21 @@ class ValidateModelResponse(BaseModel): description = "Native training context length, read from the GGUF header when the file " "is already downloaded locally; None for non-GGUF, gated, or not-yet-downloaded models.", ) + layer_count: Optional[int] = Field( + None, + description = "Total layer count (GGUF block_count), the manual gpu-layers ceiling, read " + "from the header alongside context_length; None when not read.", + ) + moe_layer_count: Optional[int] = Field( + None, + description = "MoE expert-layer count (the manual --n-cpu-moe ceiling), read from the GGUF " + "header alongside context_length; 0 for dense models, None when not read.", + ) + chat_template: Optional[str] = Field( + None, + description = "Embedded GGUF chat template, read from the header when include_chat_template " + "is set (native lease-backed picks); None for non-GGUF, over-cap, or not-read templates.", + ) # Additive fields; the consuming consent dialog ships in a follow-up frontend PR. requires_transformers_upgrade: bool = Field( False, @@ -333,6 +433,34 @@ class LoadResponse(BaseModel): False, description = "Whether tensor-parallel split (--split-mode tensor) is active.", ) + gpu_memory_mode: Literal["auto", "manual"] = Field( + "auto", + description = "Active GPU memory strategy ('auto' or 'manual').", + ) + gpu_layers: int = Field( + -1, + description = "Manual mode: requested --gpu-layers value (-1 = Auto/--fit, or when not manual).", + ) + n_cpu_moe: int = Field( + 0, + description = "Manual mode: MoE expert layers pinned to CPU (--n-cpu-moe); 0 = none.", + ) + tensor_split: Optional[List[float]] = Field( + None, + description = "Manual mode: relative model share per GPU (--tensor-split); None = default (split by free VRAM).", + ) + n_layers: Optional[int] = Field( + None, + description = "Model's layer count (GGUF block_count), for the manual gpu-layers ceiling.", + ) + n_moe_layers: int = Field( + 0, + description = "Model's MoE expert-layer count (the n_cpu_moe ceiling); 0 if not an MoE model.", + ) + gpu_ids: Optional[List[int]] = Field( + None, + description = "Physical GPU indices the model is pinned to, or None for automatic selection.", + ) class UnloadResponse(BaseModel): @@ -461,6 +589,42 @@ class InferenceStatusResponse(BaseModel): False, description = "Whether tensor-parallel split (--split-mode tensor) is active.", ) + gpu_memory_mode: Literal["auto", "manual"] = Field( + "auto", + description = "Active GPU memory strategy ('auto' or 'manual').", + ) + gpu_layers: int = Field( + -1, + description = "Manual mode: requested --gpu-layers value (-1 = Auto/--fit, or when not manual).", + ) + n_cpu_moe: int = Field( + 0, + description = "Manual mode: MoE expert layers pinned to CPU (--n-cpu-moe); 0 = none.", + ) + tensor_split: Optional[List[float]] = Field( + None, + description = "Manual mode: relative model share per GPU (--tensor-split); None = default (split by free VRAM).", + ) + requested_context_length: Optional[int] = Field( + None, + description = ( + "The n_ctx the active GGUF load was invoked with (0 = Auto). Lets the " + "UI re-seed a Manual + Auto-layers context pin on hydration, where " + "context_length only exposes the resolved value. None for non-GGUF." + ), + ) + n_layers: Optional[int] = Field( + None, + description = "Model's layer count (GGUF block_count), for the manual gpu-layers ceiling.", + ) + n_moe_layers: int = Field( + 0, + description = "Model's MoE expert-layer count (the n_cpu_moe ceiling); 0 if not an MoE model.", + ) + gpu_ids: Optional[List[int]] = Field( + None, + description = "Physical GPU indices the model is pinned to, or None for automatic selection.", + ) llama_cpp_supports_mtp: bool = Field( True, description = ( @@ -533,7 +697,7 @@ class ImageContentPart(BaseModel): class InputDocumentContentPart(BaseModel): """Document (PDF / file) content part in a multimodal message. - Studio-normalised shape (file_data or file_url, plus optional filename/media_type). + Unsloth-normalised shape (file_data or file_url, plus optional filename/media_type). Mapped onto Anthropic ``document`` / OpenAI ``input_file`` for vision providers; dropped for non-vision providers. """ @@ -689,7 +853,7 @@ class ThinkingConfig(BaseModel): """Anthropic-compatible thinking/reasoning configuration. Use type='disabled' to turn off thinking, or type='enabled' to turn it on. Only type is read; extra fields (e.g. budget_tokens) are ignored, since - Studio sets provider thinking budgets itself. + Unsloth sets provider thinking budgets itself. """ type: Literal["disabled", "enabled"] = "disabled" @@ -748,7 +912,7 @@ class ChatCompletionRequest(BaseModel): None, description = ( "OpenAI function-tool definitions. When provided without `enable_tools=true`, " - "Studio forwards the tools to the backend so the model returns structured " + "Unsloth forwards the tools to the backend so the model returns structured " "tool_calls for the client to execute (standard OpenAI function calling)." ), ) @@ -1160,7 +1324,7 @@ class ChatCompletionRequest(BaseModel): and (self.enable_tools is True or bool(self.mcp_enabled)) ): # "Ask" gates every call, so a direct API caller that omits the legacy - # confirm flag must still hit the confirmation gate for Studio's own + # confirm flag must still hit the confirmation gate for Unsloth's own # tool loop. An explicit confirm_tool_calls=False wins over the mode # (mirrors _permission_mode_confirm and the Anthropic pre-switch guard), # so only self-enable when the flag is unset. Only self-enable when that @@ -1168,7 +1332,7 @@ class ChatCompletionRequest(BaseModel): # (enable_tools / mcp_enabled) -- the router enters the loop on those # signals, not on enabled_tools alone (which merely filters which tools # run). A plain client-tool passthrough (client-supplied `tools` that - # Studio does not execute) must route verbatim, and external-provider + # Unsloth does not execute) must route verbatim, and external-provider # routing rejects confirm_tool_calls with tools, so skip the fold there. # # "auto" is deliberately NOT folded: it only prompts for a call the diff --git a/studio/backend/models/training.py b/studio/backend/models/training.py index ff815a2fa9..0f88b78f9f 100644 --- a/studio/backend/models/training.py +++ b/studio/backend/models/training.py @@ -446,7 +446,7 @@ class TrainingStartRequest(BaseModel): random_seed: int = Field( 3407, description = ( - "Random seed; matches the Studio backend / MLX worker default " + "Random seed; matches the Unsloth backend / MLX worker default " "and unsloth's historical recommended value." ), ) @@ -505,6 +505,13 @@ class TrainingStartRequest(BaseModel): description = "S3 bucket configuration for loading datasets from AWS S3. Requires boto3 to be installed.", ) + @field_validator("target_modules", mode = "before") + @classmethod + def _normalize_target_modules(cls, value: Any) -> Any: + # Sanitized non-LoRA history stores the unused value as null; treat it as a + # fresh request's omitted/default empty list on resume. + return [] if value is None else value + @model_validator(mode = "after") def _validate_streaming_splits(self) -> "TrainingStartRequest": # Streaming load_dataset does not accept HF slice syntax (e.g. "train[:50%]" diff --git a/studio/backend/picker/__init__.py b/studio/backend/picker/__init__.py new file mode 100644 index 0000000000..32014236c6 --- /dev/null +++ b/studio/backend/picker/__init__.py @@ -0,0 +1,2 @@ +# SPDX-License-Identifier: AGPL-3.0-only +# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0 diff --git a/studio/backend/picker/routes/__init__.py b/studio/backend/picker/routes/__init__.py new file mode 100644 index 0000000000..c0e988c8bb --- /dev/null +++ b/studio/backend/picker/routes/__init__.py @@ -0,0 +1,6 @@ +# SPDX-License-Identifier: AGPL-3.0-only +# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0 + +from .templates import router as templates_router + +__all__ = ["templates_router"] diff --git a/studio/backend/picker/routes/templates.py b/studio/backend/picker/routes/templates.py new file mode 100644 index 0000000000..02b8bf7184 --- /dev/null +++ b/studio/backend/picker/routes/templates.py @@ -0,0 +1,45 @@ +# SPDX-License-Identifier: AGPL-3.0-only +# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0 + +from __future__ import annotations + +import asyncio +from typing import Optional + +from fastapi import APIRouter, Body, Depends, Query + +from auth.authentication import get_current_subject +from hub.dependencies import get_hf_token + +from ..schemas import ( + MAX_CHAT_TEMPLATE_BYTES, + ModelTemplateResponse, + ValidateChatTemplateRequest, + ValidateChatTemplateResponse, +) +from ..service import read_default_chat_template, validate_chat_template + +router = APIRouter() + + +@router.post("/validate-chat-template", response_model = ValidateChatTemplateResponse) +async def validate_chat_template_route( + body: ValidateChatTemplateRequest = Body(...), + current_subject: str = Depends(get_current_subject), +) -> ValidateChatTemplateResponse: + return await asyncio.to_thread(validate_chat_template, body.template) + + +@router.get("/chat-template/{model_name:path}", response_model = ModelTemplateResponse) +async def get_default_chat_template_route( + model_name: str, + gguf_variant: Optional[str] = Query(None), + hf_token: Optional[str] = Depends(get_hf_token), + current_subject: str = Depends(get_current_subject), +) -> ModelTemplateResponse: + template = await asyncio.to_thread( + read_default_chat_template, model_name, hf_token, gguf_variant + ) + if template is not None and len(template.encode("utf-8")) > MAX_CHAT_TEMPLATE_BYTES: + template = None + return ModelTemplateResponse(model_name = model_name, chat_template = template) diff --git a/studio/backend/picker/schemas.py b/studio/backend/picker/schemas.py new file mode 100644 index 0000000000..b4f956188f --- /dev/null +++ b/studio/backend/picker/schemas.py @@ -0,0 +1,32 @@ +# SPDX-License-Identifier: AGPL-3.0-only +# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0 + +from typing import Optional + +from pydantic import BaseModel, Field, field_validator + +# Mirror the frontend's 64 KiB chat-template contract (per-model-config.ts) at +# the API boundary so a direct caller cannot make Jinja parse an oversized +# template. MaxBodyMiddleware only caps the whole request body, not this field. +MAX_CHAT_TEMPLATE_BYTES = 65_536 + + +class ValidateChatTemplateRequest(BaseModel): + template: str = Field(default = "") + + @field_validator("template") + @classmethod + def _enforce_template_size(cls, value: str) -> str: + if len(value.encode("utf-8")) > MAX_CHAT_TEMPLATE_BYTES: + raise ValueError(f"Chat template exceeds the {MAX_CHAT_TEMPLATE_BYTES}-byte limit.") + return value + + +class ValidateChatTemplateResponse(BaseModel): + valid: bool + error: Optional[str] = None + + +class ModelTemplateResponse(BaseModel): + model_name: str + chat_template: Optional[str] = None diff --git a/studio/backend/picker/service.py b/studio/backend/picker/service.py new file mode 100644 index 0000000000..13065b2920 --- /dev/null +++ b/studio/backend/picker/service.py @@ -0,0 +1,426 @@ +# SPDX-License-Identifier: AGPL-3.0-only +# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0 + +from __future__ import annotations + +import json +import logging +import os +import re +from pathlib import Path +from typing import Optional + +from hub.services.models.folder_browser import ( + _build_browse_allowlist, + _is_path_inside_allowlist, +) +from hub.utils.gguf import extract_quant_label, iter_hf_cache_snapshots +from utils.models.gguf_metadata import read_gguf_chat_template +from utils.models.model_config import ( + _extract_quant_label, + _is_big_endian_gguf_path, + _is_mmproj, + _is_mtp_drafter, +) +from utils.paths.path_utils import ( + is_local_path, + normalize_path, + resolve_cached_repo_id_case, +) + +from .schemas import MAX_CHAT_TEMPLATE_BYTES, ValidateChatTemplateResponse + +logger = logging.getLogger(__name__) + +_VALID_REPO_ID = re.compile(r"^[A-Za-z0-9._-]+/[A-Za-z0-9._-]+$") + + +def _is_valid_repo_id(repo_id: str) -> bool: + return bool(_VALID_REPO_ID.fullmatch(repo_id)) + + +_TOKENIZER_CONFIG_PATHS = ("tokenizer_config.json", "LLM/tokenizer_config.json") +_JINJA_TEMPLATE_PATHS = ("chat_template.jinja", "LLM/chat_template.jinja") +_PROCESSOR_TEMPLATE_PATHS = ("chat_template.json", "LLM/chat_template.json") + +# Cap sidecar reads so a malformed or hostile metadata file cannot exhaust memory +# before its template is size-checked. The JSON envelope may exceed a bare template +# (it carries other tokenizer metadata); the extracted template is still bounded by +# MAX_CHAT_TEMPLATE_BYTES downstream. +MAX_TEMPLATE_METADATA_BYTES = 4 * 1024 * 1024 + + +def _read_bounded_text(path: Path, limit: int) -> Optional[str]: + """Read at most `limit` bytes of UTF-8 text; None if larger or unreadable.""" + try: + with path.open("rb") as f: + data = f.read(limit + 1) + except OSError: + return None + if len(data) > limit: + return None + try: + return data.decode("utf-8") + except UnicodeError: + return None + + +def _leaf_inside_allowlist(path: Path, allow_roots: Optional[list[Path]]) -> bool: + # Block symlinked children from escaping the validated directory (realpath-checked). + # None = trusted caller (HF cache / remote download). + return allow_roots is None or _is_path_inside_allowlist(path, allow_roots) + + +def validate_chat_template(template: str) -> ValidateChatTemplateResponse: + text = (template or "").strip() + if not text: + return ValidateChatTemplateResponse(valid = True, error = None) + # Import Jinja lazily: optional at runtime (e.g. GGUF-only installs), so a + # missing dependency must not crash API startup. + try: + from jinja2 import TemplateError + from jinja2.ext import Extension + from jinja2.sandbox import ImmutableSandboxedEnvironment + except ImportError: + return ValidateChatTemplateResponse(valid = True, error = None) + + class _GenerationTag(Extension): + # Accept Transformers' {% generation %} assistant-mask tag so a pasted HF + # chat template validates (we only parse it). + tags = {"generation"} + + def parse(self, parser): + next(parser.stream) + return parser.parse_statements(["name:endgeneration"], drop_needle = True) + + try: + env = ImmutableSandboxedEnvironment( + trim_blocks = True, + lstrip_blocks = True, + extensions = ["jinja2.ext.loopcontrols", _GenerationTag], + ) + env.parse(text) + return ValidateChatTemplateResponse(valid = True, error = None) + except TemplateError as exc: + message = getattr(exc, "message", None) or str(exc) + lineno = getattr(exc, "lineno", None) + if lineno: + message = f"Line {lineno}: {message}" + return ValidateChatTemplateResponse(valid = False, error = message) + except Exception as exc: + return ValidateChatTemplateResponse(valid = False, error = str(exc)) + + +def _chat_template_from_tokenizer_config(config: dict) -> Optional[str]: + if not isinstance(config, dict): + return None + raw = config.get("chat_template") + if isinstance(raw, str) and raw.strip(): + return raw + if isinstance(raw, list): + fallback: Optional[str] = None + for entry in raw: + if not isinstance(entry, dict): + continue + template = entry.get("template") + if not isinstance(template, str): + continue + if entry.get("name") == "default": + return template + if fallback is None: + fallback = template + return fallback + return None + + +def _chat_template_from_jinja_file( + dir_path: Path, allow_roots: Optional[list[Path]] = None +) -> Optional[str]: + for rel in _JINJA_TEMPLATE_PATHS: + template_file = dir_path / rel + if not template_file.exists() or not _leaf_inside_allowlist(template_file, allow_roots): + continue + try: + if template_file.stat().st_size > MAX_CHAT_TEMPLATE_BYTES: + continue + template = template_file.read_text(encoding = "utf-8") + except Exception: + continue + if template.strip(): + return template + return None + + +def _chat_template_from_processor_payload(payload: object) -> Optional[str]: + # processor chat_template.json may be the template string itself or a + # {name: template} map, not only a tokenizer_config-shaped object. + if isinstance(payload, str): + return payload if payload.strip() else None + template = _chat_template_from_tokenizer_config(payload) # type: ignore[arg-type] + if template: + return template + if isinstance(payload, dict): + # Named-template map: prefer "default", else the first non-empty entry + # (mirrors the tokenizer-config list fallback). + default = payload.get("default") + if isinstance(default, str) and default.strip(): + return default + for value in payload.values(): + if isinstance(value, str) and value.strip(): + return value + return None + + +def _chat_template_from_processor_json( + dir_path: Path, allow_roots: Optional[list[Path]] = None +) -> Optional[str]: + for rel in _PROCESSOR_TEMPLATE_PATHS: + config_file = dir_path / rel + if not config_file.exists() or not _leaf_inside_allowlist(config_file, allow_roots): + continue + raw = _read_bounded_text(config_file, MAX_TEMPLATE_METADATA_BYTES) + if raw is None: + continue + try: + payload = json.loads(raw) + except Exception: + continue + template = _chat_template_from_processor_payload(payload) + if template: + return template + return None + + +def _chat_template_from_tokenizer_dir( + dir_path: Path, allow_roots: Optional[list[Path]] = None +) -> Optional[str]: + jinja = _chat_template_from_jinja_file(dir_path, allow_roots) + if jinja: + return jinja + for rel in _TOKENIZER_CONFIG_PATHS: + config_file = dir_path / rel + if not config_file.exists() or not _leaf_inside_allowlist(config_file, allow_roots): + continue + raw = _read_bounded_text(config_file, MAX_TEMPLATE_METADATA_BYTES) + if raw is None: + continue + try: + config = json.loads(raw) + except Exception: + continue + template = _chat_template_from_tokenizer_config(config) + if template: + return template + return _chat_template_from_processor_json(dir_path, allow_roots) + + +_GGUF_SCAN_MAX_DEPTH = 2 + + +def _iter_ggufs(dir_path: Path) -> list[Path]: + if dir_path == dir_path.parent: + return [] + root = str(dir_path) + found: list[Path] = [] + for current, dirs, files in os.walk(root, followlinks = False): + rel = os.path.relpath(current, root) + depth = 0 if rel == os.curdir else rel.count(os.sep) + 1 + if depth >= _GGUF_SCAN_MAX_DEPTH: + dirs[:] = [] + for name in files: + if not name.lower().endswith(".gguf") or _is_mmproj(name): + continue + path = Path(current) / name + try: + rel = path.relative_to(dir_path).as_posix() + except ValueError: + rel = name + quant = _extract_quant_label(rel) + if _is_mtp_drafter(rel) or _is_big_endian_gguf_path(rel, quant): + continue + found.append(path) + return found + + +def _variant_matches(relative_path: str, needle: str) -> bool: + quant = _extract_quant_label(relative_path).lower() + if quant == needle: + return True + if extract_quant_label(relative_path).lower() == needle: + return True + prefix = f"{needle}-" + if not quant.startswith(prefix): + return False + suffix = quant[len(prefix) :] + if not suffix.endswith("bpw"): + return False + value = suffix[:-3] + return bool(value) and value.replace(".", "", 1).isdigit() + + +_GGUF_SPLIT_INDEX_RE = re.compile(r"-(\d{3,})-of-\d{3,}$", re.IGNORECASE) + + +def _is_nonfirst_gguf_split(path: Path) -> bool: + match = _GGUF_SPLIT_INDEX_RE.search(path.stem) + return match is not None and int(match.group(1)) != 1 + + +def _find_gguf_in_dir(dir_path: Path, gguf_variant: Optional[str]) -> Optional[Path]: + try: + ggufs = sorted(_iter_ggufs(dir_path)) + except OSError: + return None + if not ggufs: + return None + needle = (gguf_variant or "").strip().lower() + if needle: + for path in ggufs: + try: + relative = path.relative_to(dir_path).as_posix() + except ValueError: + relative = path.name + if _variant_matches(relative, needle): + return path + return None + candidates = [path for path in ggufs if not _is_nonfirst_gguf_split(path)] or ggufs + try: + return max(candidates, key = lambda path: path.stat().st_size) + except OSError: + return candidates[0] + + +def _chat_template_from_dir( + dir_path: Path, + gguf_variant: Optional[str] = None, + allow_roots: Optional[list[Path]] = None, +) -> Optional[str]: + def from_gguf() -> Optional[str]: + gguf = _find_gguf_in_dir(dir_path, gguf_variant) + if gguf is None or not _leaf_inside_allowlist(gguf, allow_roots): + return None + return read_gguf_chat_template(str(gguf)) + + # Sidecar tokenizer files (chat_template.jinja / tokenizer_config.json) are the + # author's maintained template and supersede the GGUF's possibly-stale embedded + # copy. The variant only picks the GGUF fallback, so tokenizer-first precedence + # holds whether or not a variant is given. + return _chat_template_from_tokenizer_dir(dir_path, allow_roots) or from_gguf() + + +def read_default_chat_template( + model_name: str, + hf_token: Optional[str] = None, + gguf_variant: Optional[str] = None, +) -> Optional[str]: + if not isinstance(model_name, str) or not model_name.strip(): + return None + name = model_name.strip() + + if is_local_path(name): + try: + target = Path(normalize_path(name)).expanduser() + allow_roots = _build_browse_allowlist() + if not _is_path_inside_allowlist(target, allow_roots): + logger.debug("Refused chat template read outside allowed folders: %s", name) + return None + if name.lower().endswith(".gguf"): + # Prefer a maintained sidecar next to the file over the GGUF's + # embedded copy (tokenizer-first precedence, as elsewhere). + sidecar = _chat_template_from_tokenizer_dir(target.parent, allow_roots) + if sidecar: + return sidecar + return read_gguf_chat_template(str(target)) + return _chat_template_from_dir(target, gguf_variant, allow_roots) + except Exception as exc: + logger.debug("Could not read local chat template for %s: %s", name, exc) + return None + + if not _is_valid_repo_id(name): + return None + + resolved = resolve_cached_repo_id_case(name) + + try: + # Resolve within each cached revision, newest first. A revision's sidecar + # supersedes its own embedded GGUF copy, but must not override a newer + # revision, so precedence stays per-snapshot rather than global. + for snapshot in iter_hf_cache_snapshots(resolved): + template = _chat_template_from_dir(snapshot, gguf_variant) + if template: + return template + except Exception as exc: + logger.debug("Could not read cached chat template for %s: %s", resolved, exc) + + try: + from huggingface_hub import HfApi, hf_hub_download + + _api = HfApi() + + def _remote_exceeds_cap(rel: str) -> bool: + # Best-effort: skip the download when the remote's advertised size + # exceeds the cap, so a maliciously large sidecar is never fetched. + try: + infos = _api.get_paths_info(resolved, [rel], repo_type = "model", token = hf_token) + except Exception: + return False + for info in infos: + size = getattr(info, "size", None) + if ( + getattr(info, "path", None) == rel + and isinstance(size, int) + and size > MAX_TEMPLATE_METADATA_BYTES + ): + return True + return False + + def _download_text(rel: str) -> Optional[str]: + if _remote_exceeds_cap(rel): + return None + try: + path = hf_hub_download(resolved, rel, token = hf_token) + return _read_bounded_text(Path(path), MAX_TEMPLATE_METADATA_BYTES) + except Exception: + return None + + for rel in _JINJA_TEMPLATE_PATHS: + template = _download_text(rel) + if not template or not template.strip(): + continue + # A raw Jinja sidecar is the whole template, so it must fit the route's + # response cap (the local path skips oversized .jinja too). Download stays + # bounded at MAX_TEMPLATE_METADATA_BYTES so a large JSON embedding a small + # template still extracts below, but an over-cap Jinja is dropped so the + # search falls through to the tokenizer/processor template. + if len(template.encode("utf-8")) > MAX_CHAT_TEMPLATE_BYTES: + continue + return template + + for rel in _TOKENIZER_CONFIG_PATHS: + raw = _download_text(rel) + if not raw: + continue + try: + config = json.loads(raw) + except Exception: + continue + template = _chat_template_from_tokenizer_config(config) + if template: + return template + + for rel in _PROCESSOR_TEMPLATE_PATHS: + raw = _download_text(rel) + if not raw: + continue + try: + payload = json.loads(raw) + except Exception: + continue + template = _chat_template_from_processor_payload(payload) + if template: + return template + + return None + except Exception as exc: + logger.debug("Could not fetch chat template for %s: %s", resolved, exc) + return None diff --git a/studio/backend/plugins/data-designer-github-repo-seed/README.md b/studio/backend/plugins/data-designer-github-repo-seed/README.md index 346d94b305..44519496f5 100644 --- a/studio/backend/plugins/data-designer-github-repo-seed/README.md +++ b/studio/backend/plugins/data-designer-github-repo-seed/README.md @@ -4,7 +4,7 @@ A Data Designer seed-reader plugin for **Unsloth Studio** that scrapes real GitHub data (issues, pull requests, commits) from one or more repositories and hands it to the recipe pipeline as a seed dataset. -Designed to ship with Studio as a default seed source so any user with a +Designed to ship with Unsloth as a default seed source so any user with a GitHub token can build training datasets straight from live repos. ## What it does @@ -64,7 +64,7 @@ sleeps until reset when the budget drops below a safety threshold. ## Install -Shipped as a default Studio plugin. For development: +Shipped as a default Unsloth plugin. For development: ```bash pip install -e . diff --git a/studio/backend/plugins/data-designer-github-repo-seed/src/data_designer_github_repo_seed/__init__.py b/studio/backend/plugins/data-designer-github-repo-seed/src/data_designer_github_repo_seed/__init__.py index 62ecb2e280..d4d46da370 100644 --- a/studio/backend/plugins/data-designer-github-repo-seed/src/data_designer_github_repo_seed/__init__.py +++ b/studio/backend/plugins/data-designer-github-repo-seed/src/data_designer_github_repo_seed/__init__.py @@ -3,4 +3,4 @@ # Intentionally empty. Data-designer loads submodules lazily via qualified names # in plugin.py, so importing this package must not touch data_designer.engine.* -# during Studio bootstrap (circular import). +# during Unsloth bootstrap (circular import). diff --git a/studio/backend/plugins/data-designer-github-repo-seed/src/data_designer_github_repo_seed/scraper.py b/studio/backend/plugins/data-designer-github-repo-seed/src/data_designer_github_repo_seed/scraper.py index 637193e8b3..1af8133cc5 100644 --- a/studio/backend/plugins/data-designer-github-repo-seed/src/data_designer_github_repo_seed/scraper.py +++ b/studio/backend/plugins/data-designer-github-repo-seed/src/data_designer_github_repo_seed/scraper.py @@ -1,7 +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 -"""Multi-repo GitHub scraper for the Studio seed plugin. +"""Multi-repo GitHub scraper for the Unsloth seed plugin. Drives the GraphQL scraper in `scraper_impl/` per repo, capped via trial_limits to stop at `limit` items per resource. Then reads the per-resource JSONL shards diff --git a/studio/backend/requirements/extras-no-deps.txt b/studio/backend/requirements/extras-no-deps.txt index 5830a47789..3361af50dd 100644 --- a/studio/backend/requirements/extras-no-deps.txt +++ b/studio/backend/requirements/extras-no-deps.txt @@ -5,7 +5,7 @@ julius torchcodec==0.10.0 snac -# peft 0.19.0 causes export subprocess shutdown issues in Studio; +# peft 0.19.0 causes export subprocess shutdown issues in Unsloth; # installing with --no-deps to avoid pulling in torch>=0.11.0 peft==0.18.1 diff --git a/studio/backend/requirements/no-torch-runtime.txt b/studio/backend/requirements/no-torch-runtime.txt index de321f80ed..378fb33a60 100644 --- a/studio/backend/requirements/no-torch-runtime.txt +++ b/studio/backend/requirements/no-torch-runtime.txt @@ -70,7 +70,7 @@ cut_cross_entropy pillow # RAG store + document parsing, mirroring studio.txt. Pinned here because -# this file installs --no-deps; without them Studio runs with RAG disabled. +# this file installs --no-deps; without them Unsloth runs with RAG disabled. sqlite-vec==0.1.9 pymupdf==1.27.2.3 # 0.3.x keeps pymupdf-layout (which pulls onnxruntime) an optional extra; the diff --git a/studio/backend/requirements/single-env/constraints.txt b/studio/backend/requirements/single-env/constraints.txt index 0ed2bf8b26..0a5619924a 100644 --- a/studio/backend/requirements/single-env/constraints.txt +++ b/studio/backend/requirements/single-env/constraints.txt @@ -4,7 +4,7 @@ transformers==4.57.6 trl==0.23.1 huggingface-hub==0.36.2 -# Studio stack +# Unsloth stack datasets==4.3.0 pyarrow==23.0.1 diff --git a/studio/backend/requirements/studio.txt b/studio/backend/requirements/studio.txt index 6f4a5c3292..0c7503a5ca 100644 --- a/studio/backend/requirements/studio.txt +++ b/studio/backend/requirements/studio.txt @@ -1,4 +1,4 @@ -# Studio UI backend dependencies +# Unsloth UI backend dependencies typer fastapi uvicorn @@ -9,7 +9,7 @@ pandas nest_asyncio datasets==4.3.0 pyjwt -# gradio>=4.0.0 # 148 MB - Studio uses React + FastAPI, not Gradio +# gradio>=4.0.0 # 148 MB - Unsloth uses React + FastAPI, not Gradio huggingface-hub==0.36.2 structlog>=24.1.0 diceware diff --git a/studio/backend/routes/auth.py b/studio/backend/routes/auth.py index c61c1a16e4..d779c8784e 100644 --- a/studio/backend/routes/auth.py +++ b/studio/backend/routes/auth.py @@ -338,11 +338,11 @@ def _clear_login_bucket(key: tuple[str, str]) -> None: # so FastAPI runs it in the threadpool rather than blocking the event loop. @router.get("/identity") def identity(nonce: str, request: Request) -> dict: - """Challenge-response proof this is the real local Studio: caller sends a nonce, + """Challenge-response proof this is the real local Unsloth: caller sends a nonce, gets HMAC(install identity secret, nonce, connection address + port). Unauthenticated and side-effect free; a process that can't read the same-user secret can't forge a proof, and binding to the address/port the connection - landed on stops a squatter relaying a proof from the real Studio elsewhere.""" + landed on stops a squatter relaying a proof from the real Unsloth elsewhere.""" try: raw = base64.urlsafe_b64decode(nonce) except Exception: diff --git a/studio/backend/routes/chat_history.py b/studio/backend/routes/chat_history.py index b95ea9f4ce..1619961a2a 100644 --- a/studio/backend/routes/chat_history.py +++ b/studio/backend/routes/chat_history.py @@ -5,7 +5,7 @@ Chat history API routes backed by studio.db. """ -from typing import Any, Literal, Optional +from typing import Annotated, Any, Literal, Optional from fastapi import APIRouter, Depends, HTTPException, Query from pydantic import BaseModel, ConfigDict, Field, ValidationError @@ -20,13 +20,16 @@ from storage.studio_db import ( clear_chat_history, count_chat_threads, count_forks_for_message, + delete_chat_attachment, delete_chat_threads, delete_chat_project, ensure_chat_project_workspace, fork_chat_thread, + get_chat_attachment, get_chat_project, get_chat_thread, get_chat_message, + list_chat_attachments_page, list_chat_projects, list_chat_legacy_imports, list_chat_settings, @@ -280,6 +283,131 @@ async def delete_threads( return {"status": "deleted"} +@router.get("/attachments") +def list_attachments( + limit: Annotated[int, Query(ge = 1, le = 100)] = 50, + offset: Annotated[int, Query(ge = 0)] = 0, + current_subject: str = Depends(get_current_subject), +) -> dict: + """One bounded page of chat uploads for the settings Data tab.""" + attachments, next_offset = list_chat_attachments_page(limit = limit, offset = offset) + return {"attachments": attachments, "nextOffset": next_offset} + + +def _decode_attachment_base64(payload: str) -> bytes: + """Strict base64 decode of a stored payload. + + Normalizes first: strips whitespace, fixes padding, accepts the URL-safe + alphabet. validate=False would silently drop bad characters and serve + corrupted bytes instead of failing, so raise 422 on anything else. + """ + import base64 + + normalized = "".join(payload.split()) + altchars = b"-_" if ("-" in normalized or "_" in normalized) else None + normalized += "=" * (-len(normalized) % 4) + try: + return base64.b64decode(normalized, altchars = altchars, validate = True) + except Exception as exc: # noqa: BLE001 - corrupt stored payload + raise HTTPException(status_code = 422, detail = "Attachment data is corrupt") from exc + + +_AUDIO_FORMAT_MEDIA_TYPES = { + "mp3": "audio/mpeg", + "wav": "audio/wav", + "ogg": "audio/ogg", + "flac": "audio/flac", +} + + +def _safe_image_media_type(media_type: str) -> str: + """Clamp a data-URL media type to something inert to render. + + Imported chats store image parts verbatim, so the embedded type can be + text/html or image/svg+xml; echoing those would execute markup with the + app origin when opened. Anything not a plain raster type downloads as + bytes instead. + """ + lowered = media_type.strip().lower() + if lowered.startswith("image/") and lowered != "image/svg+xml": + return lowered + return "application/octet-stream" + + +@router.get("/attachments/{message_id}/{attachment_id}/file") +def get_attachment_file( + message_id: str, + attachment_id: str, + current_subject: str = Depends(get_current_subject), +): + """Serve one attachment's stored content: image or audio bytes, or + extracted text.""" + import urllib.parse + + from fastapi.responses import Response + + attachment = get_chat_attachment(message_id, attachment_id) + if attachment is None: + raise HTTPException(status_code = 404, detail = "Attachment not found") + + attachment_content_type = attachment.get("contentType") + texts: list[str] = [] + for part in attachment.get("content") or []: + if not isinstance(part, dict): + continue + image = part.get("image") + if isinstance(image, str) and image[:5].lower() == "data:": + header, _, payload = image.partition(",") + media_type = _safe_image_media_type( + header[5:].split(";", 1)[0] or "application/octet-stream" + ) + if "base64" not in header.lower(): + # RFC 2397 non-base64 form stores percent-encoded bytes. + data = urllib.parse.unquote_to_bytes(payload) + return Response(content = data, media_type = media_type) + data = _decode_attachment_base64(payload) + return Response(content = data, media_type = media_type) + # Audio parts: the attachment adapter stores {data, format} with raw + # base64; compare chats store a bare base64 string. + audio = part.get("audio") + if isinstance(audio, dict) or (isinstance(audio, str) and audio): + if isinstance(audio, dict): + payload = audio.get("data") + audio_format = audio.get("format") + else: + payload = audio.rsplit(",", 1)[-1] + audio_format = None + if isinstance(payload, str) and payload: + data = _decode_attachment_base64(payload) + media_type = ( + attachment_content_type + if isinstance(attachment_content_type, str) + and attachment_content_type.startswith("audio/") + else _AUDIO_FORMAT_MEDIA_TYPES.get( + str(audio_format or "").lower(), "application/octet-stream" + ) + ) + return Response(content = data, media_type = media_type) + text = part.get("text") + if isinstance(text, str) and text: + texts.append(text) + if texts: + return Response(content = "\n".join(texts), media_type = "text/plain; charset=utf-8") + raise HTTPException(status_code = 404, detail = "Attachment has no stored content") + + +@router.delete("/attachments/{message_id}/{attachment_id}") +def delete_attachment( + message_id: str, + attachment_id: str, + current_subject: str = Depends(get_current_subject), +) -> dict: + """Remove one attachment from its chat message.""" + if not delete_chat_attachment(message_id, attachment_id): + raise HTTPException(status_code = 404, detail = "Attachment not found") + return {"ok": True} + + @router.get("/projects", response_model = ChatProjectListResponse) async def list_projects( include_archived: bool = Query(False), current_subject: str = Depends(get_current_subject) @@ -410,7 +538,7 @@ async def get_thread_message( @router.put("/threads/{thread_id}/messages/{message_id}", response_model = ChatMessage) -async def save_thread_message( +def save_thread_message( thread_id: str, message_id: str, payload: ChatMessage, @@ -433,7 +561,7 @@ async def save_thread_message( @router.put("/threads/{thread_id}/messages", response_model = ChatMessageListResponse) -async def replace_thread_messages( +def replace_thread_messages( thread_id: str, payload: ChatMessageSyncRequest, current_subject: str = Depends(get_current_subject), diff --git a/studio/backend/routes/data_recipe/jobs.py b/studio/backend/routes/data_recipe/jobs.py index 59714380da..e870e8855e 100644 --- a/studio/backend/routes/data_recipe/jobs.py +++ b/studio/backend/routes/data_recipe/jobs.py @@ -37,7 +37,7 @@ def _resolve_local_v1_endpoint(request: Request) -> str: Resolution order: 1. ``app.state.server_port`` (run.py, post-bind) - survives proxies/tunnels. - 2. ``request.scope["server"]`` - when Studio starts outside ``run_server``. + 2. ``request.scope["server"]`` - when Unsloth starts outside ``run_server``. 3. parsed ``request.base_url`` - last resort for test fixtures. """ port: Any = getattr(request.app.state, "server_port", None) diff --git a/studio/backend/routes/datasets.py b/studio/backend/routes/datasets.py index 46319ca2ba..5456080f34 100644 --- a/studio/backend/routes/datasets.py +++ b/studio/backend/routes/datasets.py @@ -485,7 +485,7 @@ async def upload_dataset( # Stream to disk in chunks to avoid holding the whole file in memory. The # route-level cap gives a clear training-dataset error and avoids leaving - # oversized partial files in the Studio uploads directory. + # oversized partial files in the Unsloth uploads directory. upload_limit_bytes = get_upload_limit_bytes() total_bytes = 0 upload_complete = False diff --git a/studio/backend/routes/inference.py b/studio/backend/routes/inference.py index 52ea1f86a3..d3e588bb0b 100644 --- a/studio/backend/routes/inference.py +++ b/studio/backend/routes/inference.py @@ -13,7 +13,7 @@ from pathlib import Path from fastapi import APIRouter, Depends, HTTPException, Request, status from fastapi.responses import StreamingResponse, JSONResponse, Response from starlette.requests import ClientDisconnect -from typing import Any, Callable, List, Optional, Union +from typing import Any, Callable, List, Literal, Optional, Union import json import httpx from loggers import get_logger @@ -92,7 +92,7 @@ def _mlx_distributed_launch_detected() -> bool: def _install_httpcore_asyncgen_silencer() -> None: """Silence benign httpx/httpcore asyncgen GC noise on Python 3.13. - When Studio proxies a llama-server stream via httpx, the innermost + When Unsloth proxies a llama-server stream via httpx, the innermost ``HTTP11ConnectionByteStream.__aiter__`` async generator is finalised by the asyncgen GC hook on a task different from the one that opened it. Its ``aclose`` calls ``anyio.Lock.acquire`` → ``cancel_shielded_checkpoint``, @@ -229,14 +229,14 @@ def _friendly_upstream_error(text: str) -> str: parse grammar" / "failed to initialize samplers"). This surfaces to coding agents as a hard 400 on every tool-bearing turn. It is a llama-server limitation with some model/quant + tool-schema combinations, and recent llama.cpp builds handle the common - coding-agent tools, so point the user at updating Studio rather than the raw body. + coding-agent tools, so point the user at updating Unsloth rather than the raw body. """ lowered = text.lower() if "failed to parse grammar" in lowered or "failed to initialize samplers" in lowered: return ( "The model couldn't compile a tool-calling grammar for this request. This is a " "llama-server limitation with some model/quant and tool-schema combinations. " - "Update Studio (it installs the latest llama.cpp, which handles the common " + "Update Unsloth (it installs the latest llama.cpp, which handles the common " "coding-agent tools) or try a different GGUF model." ) return f"llama-server error: {text}" @@ -731,7 +731,7 @@ def _openai_passthrough_sse_line_terminal_state(raw_line: str) -> Optional[str]: Some llama-server builds can emit the logical final chunk (``finish_reason``) and optional usage chunk, then keep the HTTP stream open without sending the - OpenAI ``data: [DONE]`` sentinel. Classifying those chunks lets Studio close + OpenAI ``data: [DONE]`` sentinel. Classifying those chunks lets Unsloth close the client stream promptly while preserving an optional trailing usage chunk. """ if not raw_line.startswith("data:"): @@ -1778,7 +1778,7 @@ from core.inference.providers import get_base_url from core.inference.external_provider import ExternalProviderClient from core.inference.chat_templates import resolve_effective_chat_template_override from storage import providers_db -from utils.utils import safe_error_detail, log_and_http_error +from utils.utils import is_hf_authentication_error, safe_error_detail, log_and_http_error import io import base64 @@ -1786,7 +1786,7 @@ import numpy as np from datetime import date as _date router = APIRouter() -# Studio-only router (not mounted on /v1 OpenAI-compat). +# Unsloth-only router (not mounted on /v1 OpenAI-compat). studio_router = APIRouter() @@ -2108,9 +2108,9 @@ def _effective_enable_tools(payload) -> Optional[bool]: def _explicit_studio_tool_loop_requested(payload) -> bool: - """True when the request itself asks Studio to execute local tools. + """True when the request itself asks Unsloth to execute local tools. - Process-wide CLI policy can default Studio's tool loop on for ordinary chat, + Process-wide CLI policy can default Unsloth's tool loop on for ordinary chat, but it must not steal OpenAI-compatible client tools or response_format requests from the llama-server passthrough path. A policy of ``False`` (--disable-tools) vetoes even an explicit ``enable_tools: true`` ask. @@ -2122,7 +2122,7 @@ def _explicit_studio_tool_loop_requested(payload) -> bool: def _permission_mode_confirm(payload) -> bool: - """Effective confirm-gate intent for Studio's own local tool loop. + """Effective confirm-gate intent for Unsloth's own local tool loop. Honors the documented default that an unset permission_mode behaves as "ask". An explicit confirm_tool_calls (True or False) wins; explicit @@ -2144,7 +2144,7 @@ def _permission_mode_confirm(payload) -> bool: def _confirm_gate_needs_stream(payload) -> bool: - """Whether Studio's local tool-loop confirm gate still requires stream=true. + """Whether Unsloth's local tool-loop confirm gate still requires stream=true. The gate can only prompt while streaming, so a non-streaming request that will prompt must 400 up front. auto ("Approve for me") only prompts for a call the @@ -3115,13 +3115,16 @@ def _normalise_settings_str(value: Optional[str]) -> Optional[str]: def _should_strip_split_mode(request: LoadRequest, backend_extra: Optional[list[str]]) -> bool: - """Whether an inherited --split-mode should be stripped on reload. + """Whether an inherited --split-mode (and its coupled --tensor-split) should + be stripped on reload. The binary Tensor Parallelism toggle can't carry --split-mode's row/none/ layer modes, so only strip when the toggle overrides it: tensor being turned on, or the inherited mode is tensor (toggle turning it off). Non-tensor modes - survive. Shared by the inheritance strip and the already-loaded stale check - so they agree on what reload would do. + survive. A manual per-GPU ratio is handled by _should_strip_tensor_split, + which strips only --tensor-split so the inherited mode is kept. Shared by the + inheritance strip and the already-loaded stale check so they agree on what + reload would do. """ fields_set = getattr(request, "model_fields_set", set()) return "tensor_parallel" in fields_set and ( @@ -3129,6 +3132,25 @@ def _should_strip_split_mode(request: LoadRequest, backend_extra: Optional[list[ ) +def _should_strip_tensor_split(request: LoadRequest) -> bool: + """Whether an inherited --tensor-split alone should be stripped on reload. + + Manual explicit offload (gpu_layers >= 0) owns the per-GPU split: with a ratio + it emits its own --tensor-split (an inherited one, appended last, would + override it), and with the ratio cleared it wants llama.cpp's default + free-VRAM split. Either way an inherited --tensor-split must go, else the + cleared case silently keeps the stale ratio while status reports None. + Unlike _should_strip_split_mode this leaves --split-mode untouched, so a + user's row/none/layer mode survives a Studio split-ratio edit. When the + Tensor Parallelism toggle IS overriding the mode, _should_strip_split_mode + (called alongside this at every site) strips --split-mode anyway. + """ + return ( + getattr(request, "gpu_memory_mode", "auto") == "manual" + and getattr(request, "gpu_layers", -1) >= 0 + ) + + def _carry_preserved_tensor_intent( *, preserved: bool, same_model: bool, explicit_drop: bool ) -> bool: @@ -3143,7 +3165,7 @@ def _is_explicit_tensor_drop(request: LoadRequest) -> bool: """True only when the request explicitly selects a non-tensor --split-mode (e.g. layer/row/none), a deliberate departure from a preserved tensor->layer fallback. - A bare tensor_parallel field is NOT a drop: the Studio UI always sends it and echoes + A bare tensor_parallel field is NOT a drop: the Unsloth UI always sends it and echoes the /load response's resolved value back, so after a fallback every reload carries tensor_parallel=false even though the user never changed it -- treating that as a drop would collapse the preserved multi-GPU placement on the next ctx/settings reload. An @@ -3187,12 +3209,44 @@ def _request_matches_loaded_settings( else strip_shadowing_flags( backend_extra, strip_split_mode = _should_strip_split_mode(request, backend_extra), + strip_tensor_split = _should_strip_tensor_split(request), + strip_offload = request.gpu_memory_mode == "manual", ) ) if not _tensor_parallel_matches_loaded( effective_extra, request.tensor_parallel, llama_backend.tensor_parallel ): return False + # The diffusion runner is mode-agnostic (it always reports "auto" and 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 llama_backend.is_diffusion: + if request.gpu_memory_mode != llama_backend.gpu_memory_mode: + return False + # Manual: a layer-count change always reloads; MoE/split only matter with + # an explicit offload (gpu_layers >= 0), so a leftover value under Auto + # must not force one. Mirrors LlamaCppBackend._already_in_target_state. + if request.gpu_memory_mode == "manual" and ( + request.gpu_layers != llama_backend.gpu_layers + or ( + request.gpu_layers >= 0 + and ( + request.n_cpu_moe != llama_backend.n_cpu_moe + or (request.tensor_split or None) != (llama_backend.tensor_split or None) + ) + ) + ): + return False + # A changed GPU pick must reload. The diffusion runner collapses a multi-GPU + # request to its single lowest device (it drives one device only), so the + # backend records just that device; compare the request the same way, or a + # multi-GPU pick that resolves to the same device needlessly reloads. + if llama_backend.is_diffusion: + _req_gpu_ids = [sorted(request.gpu_ids)[0]] if request.gpu_ids else None + else: + _req_gpu_ids = sorted(request.gpu_ids) if request.gpu_ids else None + if _req_gpu_ids != llama_backend.gpu_ids: + return False # Preserved tensor->layer fallback (both report tensor=off, so the check above # matches): if the user now explicitly drops tensor intent, reload so placement # re-selects instead of keeping the all-GPU mask (#6659). The effective check @@ -3235,14 +3289,17 @@ def _request_matches_loaded_settings( # contain any shadow flag, so the reload path strips them rather than # leaving a stale override in effect. (backend_extra computed above.) if request.llama_extra_args is None: - # Mirror the reload's conditional split-mode strip, so a preserved - # non-tensor mode (row/none/layer) isn't seen as stale and doesn't - # trigger a needless reload of a healthy server. + # Mirror the reload's conditional strips, so a preserved non-tensor mode + # (row/none/layer) isn't seen as stale and doesn't trigger a needless + # reload of a healthy server, while an inherited offload/ratio flag that + # the reload *would* strip is correctly seen as stale. if ( backend_extra and strip_shadowing_flags( backend_extra, strip_split_mode = _should_strip_split_mode(request, backend_extra), + strip_tensor_split = _should_strip_tensor_split(request), + strip_offload = request.gpu_memory_mode == "manual", ) != backend_extra ): @@ -3861,6 +3918,46 @@ def _estimate_gguf_required_gb( return None +def _classify_diffusion_gguf(config: ModelConfig) -> Optional[bool]: + """Classify a GGUF as diffusion, normal, or unknown before it is loaded. + + ``None`` is important here: a remote GGUF whose header is not cached can + still be routed to the single-GPU diffusion runner after download. Treating + that case as normal would let Manual mode skip the training guard even + though the runner ignores Manual's llama-server placement controls. + """ + identity = " ".join( + str(getattr(config, attr, "") or "") for attr in ("identifier", "gguf_hf_repo", "gguf_file") + ).lower() + if "diffusion" in identity: + return True + + try: + main = getattr(config, "gguf_file", None) + if not (main and Path(main).is_file()): + repo = getattr(config, "gguf_hf_repo", None) + variant = getattr(config, "gguf_variant", None) + if repo and variant: + from hub.utils.gguf import resolve_local_gguf_path + main = resolve_local_gguf_path(repo, variant) + if not main or not Path(main).is_file(): + return None + + probe = LlamaCppBackend() + probe._read_gguf_metadata(str(main)) + if probe.is_diffusion: + return True + # A successfully decoded architecture proves that this is a normal + # llama-server GGUF. No architecture means the lightweight probe could + # not establish the routing decision, so preserve the unknown state. + if getattr(probe, "_architecture", None): + return False + return None + except Exception as e: + logger.debug("Could not identify diffusion GGUF for training guard: %s", e) + return None + + def _guard_chat_load_against_training( config: ModelConfig, *, @@ -3871,11 +3968,19 @@ def _guard_chat_load_against_training( requested_gpu_ids: Optional[List[int]], llama_extra_args: Optional[list[str]] = None, n_parallel: int = 1, + gpu_memory_mode: Literal["auto", "manual"] = "auto", ) -> None: - """Refuse loading a local chat model that would OOM an active training run. + """Protect active training from automatically placed chat-model loads. + No-op when training is inactive or unknown. `load_in_4bit` must be the - effective quantization (see _effective_load_in_4bit). Raises HTTP 409 when the - model would not fit alongside training.""" + effective quantization (see _effective_load_in_4bit). Manual chat-GGUF + placement is an explicit override: Auto layers delegate fitting to + llama.cpp's ``--fit`` and pinned layers are owned by the user, so neither is + estimated here. Diffusion is still guarded because its mode-agnostic runner + ignores those controls and uses one GPU. An unclassified GGUF is guarded as + potentially diffusion until its local header proves otherwise. Other loads + raise HTTP 409 when they would not fit beside training. + """ from core.training import get_training_backend from routes.training_vram import can_load_chat_during_training @@ -3887,6 +3992,19 @@ def _guard_chat_load_against_training( return is_gguf = bool(getattr(config, "is_gguf", False)) + diffusion_kind = _classify_diffusion_gguf(config) if is_gguf else False + if is_gguf and gpu_memory_mode == "manual" and diffusion_kind is False: + return + + diffusion_gpu = None + if is_gguf and diffusion_kind is not False: + # Use the same token selection as the runner: an explicit pick wins, + # followed by DG_GPU, the first parent-visible token, then GPU 0. + diffusion_gpu = LlamaCppBackend._diffusion_gpu_arg( + requested_gpu_ids, + cpu_only = LlamaCppBackend._effective_gpu_count() == 0, + ) + required_override_gb = ( _estimate_gguf_required_gb( config, @@ -3907,6 +4025,7 @@ def _guard_chat_load_against_training( requested_gpu_ids = requested_gpu_ids, is_gguf = is_gguf, required_override_gb = required_override_gb, + single_device_gpu = diffusion_gpu, ) if ok: return @@ -3934,6 +4053,98 @@ def _guard_chat_load_against_training( raise HTTPException(status_code = 409, detail = detail) +def _resolve_inherited_extra_args( + request, + config: ModelConfig, + model_identifier: str, + extra_llama_args: Optional[list[str]], + effective_chat_template_override: Optional[str] = None, +) -> Optional[list[str]]: + """Effective pass-through extras for a GGUF request that omitted the field: + the previous same-model load's extras, shadow-stripped, so a settings-Apply + reload (which does not round-trip the extras field) keeps them (#5401).""" + if getattr(request, "llama_extra_args", None) is not None: + return extra_llama_args + if not getattr(config, "is_gguf", False): + return extra_llama_args + llama_backend = get_llama_cpp_backend() + if not llama_backend.extra_args: + return extra_llama_args + # Inherit the previous load's extras (the chat-settings Apply path doesn't + # round-trip them; an explicit [] still clears). Gated on (model_identifier, + # hf_variant) to refuse cross-model pickup, and shadowing flags are + # stripped so an inherited override can't win the last-wins CLI + # parse against a freshly-supplied first-class field. + source = llama_backend.extra_args_source + # Compare against the resolved variant, not the request field: callers + # commonly omit gguf_variant for local ``.gguf`` paths and HF auto-pick + # flows. ``config.gguf_variant`` is the variant load_model was actually + # invoked with, so both sides of the comparison key off the same string. + resolved_variant = (config.gguf_variant or "").lower() + request_variant = (request.gguf_variant or "").lower() + stored_variant = (source[1] or "").lower() if source else "" + same_model = bool(source and source[0] and source[0].lower() == model_identifier.lower()) + if request.gguf_variant: + variant_mismatch = request_variant != stored_variant + else: + variant_mismatch = bool(stored_variant and resolved_variant != stored_variant) + same_source = same_model and not variant_mismatch + if not same_source: + logger.info( + "Not inheriting llama_extra_args: stored args came from %s, loading %s", + source, + (model_identifier, resolved_variant), + ) + # Cross-model: clear explicitly so the backend doesn't + # inherit via "no opinion" semantics. + extra_llama_args = [] + else: + # Strip only the groups whose first-class field was set by the caller, so + # an inherited --chat-template-file survives an Apply that omits + # chat_template_override. A bundled family template (e.g. gemma-4) counts as + # a first-class template even when the request omits chat_template_override, + # so strip the inherited --chat-template-file then too -- else the stale arg + # (appended last) shadows the bundled template while Studio reports its caps. + fields_set = getattr(request, "model_fields_set", set()) + stripped = strip_shadowing_flags( + llama_backend.extra_args, + strip_context = "max_seq_length" in fields_set, + strip_cache = "cache_type_kv" in fields_set, + strip_spec = ("speculative_type" in fields_set or "spec_draft_n_max" in fields_set), + strip_template = ( + "chat_template_override" in fields_set + or effective_chat_template_override is not None + ), + strip_split_mode = _should_strip_split_mode(request, llama_backend.extra_args), + # manual + per-GPU ratio emits its own --tensor-split; drop + # an inherited one (appended last would override it) while + # keeping the user's --split-mode row/none/layer choice. + strip_tensor_split = _should_strip_tensor_split(request), + # manual emits its own --fit/--gpu-layers, so an inherited offload flag + # must not last-wins-override it. auto leaves a user's inherited -ngl + # alone. getattr: a validate request reuses this resolver, no offload fields. + strip_offload = getattr(request, "gpu_memory_mode", "auto") == "manual", + ) + try: + extra_llama_args = validate_extra_args(stripped) + except ValueError: + # Shouldn't happen on already-validated args; degrade to + # no-extras rather than 400 if managed flags changed. + logger.warning( + "Stored llama_extra_args failed revalidation; loading without them: %s", + stripped, + ) + extra_llama_args = [] + else: + if extra_llama_args: + logger.info( + "Inheriting llama_extra_args from previous " + "load (same model, shadow-stripped): %s", + extra_llama_args, + ) + return extra_llama_args + + def _model_json_response(model, status_code: int = 200) -> Response: """Serialize a pydantic response once via pydantic-core. @@ -4040,6 +4251,35 @@ async def _load_model_impl(request: LoadRequest, fastapi_request: Request, curre None if request.llama_extra_args is None else extra_llama_args ) + # Manual mode owns the offload flags: strip them from EXPLICIT extras + # too (the inherited path already does), or a last-wins --gpu-layers / + # --fit in extras re-enables GPU offload on a load status reports as + # CPU-only. Manual + per-GPU ratio owns --tensor-split the same way. + if request.gpu_memory_mode == "manual" and extra_llama_args: + _stripped_explicit = strip_shadowing_flags( + extra_llama_args, + strip_context = False, + strip_cache = False, + strip_spec = False, + strip_template = False, + strip_split_mode = False, + strip_tensor_split = _should_strip_tensor_split(request), + strip_offload = True, + ) + if _stripped_explicit != extra_llama_args: + logger.info( + "Manual GPU memory owns the offload flags; stripping them " + "from explicit llama_extra_args: %s -> %s", + extra_llama_args, + _stripped_explicit, + ) + extra_llama_args = _stripped_explicit + + # Keep every downstream consumer on the normalized explicit list. In + # particular, the already-loaded comparator must not compare the raw + # request's managed offload flags against the stripped launch state. + request = request.model_copy(update = {"llama_extra_args": extra_llama_args}) + model_identifier, model_log_label, native_grant_backed = ( _resolve_model_identifier_for_request(request, operation = "load-model") ) @@ -4121,6 +4361,13 @@ async def _load_model_impl(request: LoadRequest, fastapi_request: Request, curre speculative_type = llama_backend.requested_spec_mode, spec_draft_n_max = llama_backend.spec_draft_n_max, tensor_parallel = llama_backend.tensor_parallel, + gpu_memory_mode = llama_backend.gpu_memory_mode, + gpu_layers = llama_backend.gpu_layers, + n_cpu_moe = llama_backend.n_cpu_moe, + tensor_split = llama_backend.tensor_split, + n_layers = llama_backend.n_layers, + n_moe_layers = llama_backend.n_moe_layers, + gpu_ids = llama_backend.gpu_ids, ) else: if ( @@ -4187,18 +4434,47 @@ async def _load_model_impl(request: LoadRequest, fastapi_request: Request, curre # Normalize gpu_ids: empty list means auto-selection, same as None effective_gpu_ids = request.gpu_ids if request.gpu_ids else None - # Reject GGUF + gpu_ids first so the guard can't mask it with a VRAM 409. + # GGUF supports gpu_ids: validate the pick up front (before the training + # guard) so a bad pick is a clean 400, not masked by a VRAM 409. Rejects + # negative / out-of-range / duplicate ids and UUID/MIG parents. XPU hosts + # are rejected outright: the picker's indices are torch-xpu ordinals neither + # applicator speaks (CUDA/HIP masks don't apply, the Vulkan --device pin + # uses ggml's own Vulkan ordinals), so a pick could land on the wrong device. if config.is_gguf and effective_gpu_ids is not None: - raise HTTPException( - status_code = 400, - detail = "gpu_ids is not supported for GGUF models yet.", - ) + from utils.hardware import DeviceType, get_device + from utils.hardware.hardware import resolve_requested_gpu_ids + + if get_device() == DeviceType.XPU: + raise HTTPException( + status_code = 400, + detail = ( + "GPU selection (gpu_ids) is not supported on Intel XPU. " + "Omit gpu_ids to use all devices." + ), + ) + # Same reasoning for a Vulkan-only build: --device pins ggml's own + # Vulkan ordinals, so a physical pick can land on the wrong card on + # masked or non-contiguous hosts. + if LlamaCppBackend._is_vulkan_backend(): + raise HTTPException( + status_code = 400, + detail = ( + "GPU selection (gpu_ids) is not supported with a Vulkan " + "llama.cpp build: physical GPU ids have no defined " + "mapping to Vulkan device ordinals. Omit gpu_ids to use " + "all devices." + ), + ) + try: + resolve_requested_gpu_ids(effective_gpu_ids) + except ValueError as exc: + raise HTTPException(status_code = 400, detail = str(exc)) from exc if not config.is_gguf and _mlx_distributed_launch_detected(): raise HTTPException( status_code = 400, detail = ( - "Studio does not support distributed MLX inference under " - "mlx.launch. Use `mlx.launch ... unsloth chat` or run Studio " + "Unsloth does not support distributed MLX inference under " + "mlx.launch. Use `mlx.launch ... unsloth chat` or run Unsloth " "without the distributed launcher." ), ) @@ -4222,8 +4498,20 @@ async def _load_model_impl(request: LoadRequest, fastapi_request: Request, curre "architectures)" ) - # Refuse a load that would OOM active training, before the unload step below - # frees the resident model. Off-loop: guard does sync nvidia-smi / HF work. + # Inherit the previous same-model load's pass-through extras when this + # request omits the field (a settings-Apply reload doesn't round-trip + # them); shadow-stripped so an inherited flag can't override a + # first-class field the caller did set (#5401). + extra_llama_args = _resolve_inherited_extra_args( + request, + config, + model_identifier, + extra_llama_args, + effective_chat_template_override, + ) + + # Apply the training coexistence policy before the unload step below + # frees the resident model. Off-loop: the default-mode guard does sync work. await asyncio.to_thread( _guard_chat_load_against_training, config, @@ -4234,6 +4522,7 @@ async def _load_model_impl(request: LoadRequest, fastapi_request: Request, curre requested_gpu_ids = effective_gpu_ids, llama_extra_args = extra_llama_args, n_parallel = getattr(fastapi_request.app.state, "llama_parallel_slots", 1), + gpu_memory_mode = request.gpu_memory_mode, ) # ── GGUF path: load via llama-server ────────────────────── @@ -4245,84 +4534,6 @@ async def _load_model_impl(request: LoadRequest, fastapi_request: Request, curre from core.inference.llama_cpp import gguf_load_in_flight gguf_load_stack.enter_context(gguf_load_in_flight(config.gguf_hf_repo)) - # Inherit llama_extra_args from the previous load when the request - # omits the field (the chat-settings Apply path doesn't round-trip - # them; explicit [] still clears). Gated on (model_identifier, - # hf_variant) to refuse cross-model pickup, and shadowing flags are - # stripped so an inherited override can't win the last-wins CLI - # parse against a freshly-supplied first-class field. - if request.llama_extra_args is None and llama_backend.extra_args: - source = llama_backend.extra_args_source - # Compare against the resolved variant, not the request - # field: callers commonly omit gguf_variant for local - # ``.gguf`` paths and HF auto-pick flows. ``config.gguf_ - # variant`` is the variant load_model was actually - # invoked with (see the HF / local branches below), so - # both sides of the comparison key off the same string. - resolved_variant = (config.gguf_variant or "").lower() - request_variant = (request.gguf_variant or "").lower() - stored_variant = (source[1] or "").lower() if source else "" - same_model = bool( - source and source[0] and source[0].lower() == model_identifier.lower() - ) - if request.gguf_variant: - variant_mismatch = request_variant != stored_variant - else: - variant_mismatch = bool(stored_variant and resolved_variant != stored_variant) - same_source = same_model and not variant_mismatch - if not same_source: - logger.info( - "Not inheriting llama_extra_args: stored args came from %s, loading %s", - source, - (model_identifier, resolved_variant), - ) - # Cross-model: clear explicitly so the backend doesn't - # inherit via "no opinion" semantics. - extra_llama_args = [] - else: - # Strip only the groups whose first-class field was set by - # the caller, so an inherited --chat-template-file survives - # an Apply that omits chat_template_override. A bundled family - # template (e.g. the gemma-4 override) is an effective - # first-class template setting even when the raw request - # omits chat_template_override, so strip the inherited - # --chat-template-file in that case too -- otherwise the stale - # extra arg (appended last) shadows the bundled template while - # Studio reports the bundled template's capabilities. - fields_set = getattr(request, "model_fields_set", set()) - stripped = strip_shadowing_flags( - llama_backend.extra_args, - strip_context = "max_seq_length" in fields_set, - strip_cache = "cache_type_kv" in fields_set, - strip_spec = ( - "speculative_type" in fields_set or "spec_draft_n_max" in fields_set - ), - strip_template = ( - "chat_template_override" in fields_set - or effective_chat_template_override is not None - ), - strip_split_mode = _should_strip_split_mode( - request, llama_backend.extra_args - ), - ) - try: - extra_llama_args = validate_extra_args(stripped) - except ValueError: - # Shouldn't happen on already-validated args; degrade to - # no-extras rather than 400 if managed flags changed. - logger.warning( - "Stored llama_extra_args failed revalidation; loading without them: %s", - stripped, - ) - extra_llama_args = [] - else: - if extra_llama_args: - logger.info( - "Inheriting llama_extra_args from previous " - "load (same model, shadow-stripped): %s", - extra_llama_args, - ) - # Block cache writes that would race the download manager. This runs # after pass-through argument inheritance so a carried --no-mmproj # changes the companion requirement exactly as it does for the load. @@ -4370,6 +4581,11 @@ async def _load_model_impl(request: LoadRequest, fastapi_request: Request, curre cache_type_kv = request.cache_type_kv, speculative_type = request.speculative_type, spec_draft_n_max = request.spec_draft_n_max, + gpu_memory_mode = request.gpu_memory_mode, + gpu_layers = request.gpu_layers, + n_cpu_moe = request.n_cpu_moe, + tensor_split = request.tensor_split, + gpu_ids = effective_gpu_ids, n_parallel = _n_parallel, ) if config.gguf_hf_repo: @@ -4494,7 +4710,7 @@ async def _load_model_impl(request: LoadRequest, fastapi_request: Request, curre # Clear any idle-unload reload stash now, not only on the next poll. from core.inference.llama_keepwarm import note_model_loaded - note_model_loaded() + await asyncio.to_thread(note_model_loaded, llama_backend) # A plain load advertises its own identifier; auto-switch overwrites # this with the repo id right after _load_model_impl returns. llama_backend._openai_advertised_id = None @@ -4537,6 +4753,13 @@ async def _load_model_impl(request: LoadRequest, fastapi_request: Request, curre speculative_type = llama_backend.requested_spec_mode, spec_draft_n_max = llama_backend.spec_draft_n_max, tensor_parallel = llama_backend.tensor_parallel, + gpu_memory_mode = llama_backend.gpu_memory_mode, + gpu_layers = llama_backend.gpu_layers, + n_cpu_moe = llama_backend.n_cpu_moe, + tensor_split = llama_backend.tensor_split, + n_layers = llama_backend.n_layers, + n_moe_layers = llama_backend.n_moe_layers, + gpu_ids = llama_backend.gpu_ids, ) # ── Standard path: load via Unsloth/transformers ────────── @@ -4727,7 +4950,7 @@ def _requires_trust_remote_code_for_model( model_identifier: str, hf_token: Optional[str] = None ) -> bool: """Whether loading this model would execute custom repo code, so the consent - dialog must run first. True if the Studio YAML default enables + dialog must run first. True if the Unsloth YAML default enables ``trust_remote_code`` OR the raw config declares an ``auto_map`` (Hub/local, config.json or tokenizer_config.json). Reads raw JSON only; never imports model code.""" @@ -4795,7 +5018,9 @@ def _requires_security_review_for_model( @router.post("/validate", response_model = ValidateModelResponse) async def validate_model( - request: ValidateModelRequest, current_subject: str = Depends(get_current_subject) + request: ValidateModelRequest, + fastapi_request: Request = None, + current_subject: str = Depends(get_current_subject), ): """ Lightweight validation endpoint for model identifiers. @@ -4823,15 +5048,39 @@ async def validate_model( detail = f"Invalid model identifier: {model_log_label}", ) - # Refuse early (before the frontend unloads to load this) if it can't fit - # alongside training, using the same settings /load uses so they agree. + # Apply the same training coexistence policy as /load before the frontend + # unloads the current model. effective_gpu_ids = request.gpu_ids if request.gpu_ids else None - # Mirror /load: reject GGUF + gpu_ids before the guard so both return 400. + # Mirror /load: GGUF supports gpu_ids, so validate the pick (a bad one is + # a clean 400) before the guard sizes the model against training VRAM. + # XPU-host picks are rejected like /load (no defined mapping from the + # picker's torch-xpu ordinals to the launcher's device spaces). if config.is_gguf and effective_gpu_ids is not None: - raise HTTPException( - status_code = 400, - detail = "gpu_ids is not supported for GGUF models yet.", - ) + from utils.hardware import DeviceType, get_device + from utils.hardware.hardware import resolve_requested_gpu_ids + + if get_device() == DeviceType.XPU: + raise HTTPException( + status_code = 400, + detail = ( + "GPU selection (gpu_ids) is not supported on Intel XPU. " + "Omit gpu_ids to use all devices." + ), + ) + if LlamaCppBackend._is_vulkan_backend(): + raise HTTPException( + status_code = 400, + detail = ( + "GPU selection (gpu_ids) is not supported with a Vulkan " + "llama.cpp build: physical GPU ids have no defined " + "mapping to Vulkan device ordinals. Omit gpu_ids to use " + "all devices." + ), + ) + try: + resolve_requested_gpu_ids(effective_gpu_ids) + except ValueError as exc: + raise HTTPException(status_code = 400, detail = str(exc)) from exc effective_load_in_4bit = _effective_load_in_4bit(config, request.load_in_4bit) # Both checks cover the [adapter, base] set (matching the scan route and workers): @@ -4895,16 +5144,32 @@ async def validate_model( latest_tier_active_for, config.identifier, request.hf_token ): effective_load_in_4bit = False - # Off-loop: guard does sync nvidia-smi / HF work. - await asyncio.to_thread( - _guard_chat_load_against_training, - config, - model_identifier = model_identifier, - hf_token = request.hf_token, - load_in_4bit = effective_load_in_4bit, - max_seq_length = request.max_seq_length, - requested_gpu_ids = effective_gpu_ids, - ) + # A metadata-only probe reads the GGUF header and allocates no VRAM, so the + # training guard must not refuse it. Real loads omit include_context_length / + # include_chat_template, and /load applies the guard again. + if not (request.include_context_length or request.include_chat_template): + # Match /load's inherited llama.cpp extras and parallel slot count so + # validation cannot pass a smaller estimate than the subsequent load. + effective_extra_args = _resolve_inherited_extra_args( + request, config, model_identifier, None + ) + # Off-loop: guard does sync nvidia-smi / HF work. + await asyncio.to_thread( + _guard_chat_load_against_training, + config, + model_identifier = model_identifier, + hf_token = request.hf_token, + load_in_4bit = effective_load_in_4bit, + max_seq_length = request.max_seq_length, + requested_gpu_ids = effective_gpu_ids, + llama_extra_args = effective_extra_args, + n_parallel = ( + getattr(fastapi_request.app.state, "llama_parallel_slots", 1) + if fastapi_request is not None + else 1 + ), + gpu_memory_mode = request.gpu_memory_mode, + ) # A selected GGUF loads via llama.cpp: auto_map Python and root pickle weights in a # mixed repo are inert for this load, so gating on them is a false positive. Only @@ -4918,10 +5183,21 @@ async def validate_model( # Native context length, read from the local GGUF header when present. # Lets the staged ("Load on selection" off) flow populate the context # slider before the GPU load; None until the file is downloaded. + # Staged header dims (one read): native context, total layer count, and + # MoE expert-layer count -- let the staged flow size the context, GPU- + # layers and manual --n-cpu-moe sliders before the load. context_length: Optional[int] = None - if request.include_context_length and is_gguf: + layer_count: Optional[int] = None + moe_layer_count: Optional[int] = None + chat_template: Optional[str] = None + # Both header probes read the same local GGUF, so resolve it once. + if (request.include_context_length or request.include_chat_template) and is_gguf: from hub.utils.gguf import resolve_local_gguf_path - from utils.models.gguf_metadata import read_gguf_context_length + from picker.schemas import MAX_CHAT_TEMPLATE_BYTES + from utils.models.gguf_metadata import ( + read_gguf_chat_template, + read_gguf_staged_dims, + ) # Best-effort: a header-read failure must never fail validation of an # otherwise-valid model (the outer except turns it into a 400). @@ -4937,9 +5213,26 @@ async def validate_model( model_identifier, request.gguf_variant ) if local_gguf: - context_length = read_gguf_context_length(local_gguf) + if request.include_context_length: + # Header walk reads tokenizer arrays (tens of ms); keep it + # off the event loop. + dims = await asyncio.to_thread(read_gguf_staged_dims, local_gguf) + if dims: + context_length = dims["context_length"] + layer_count = dims["layer_count"] + moe_layer_count = dims["moe_layer_count"] + if request.include_chat_template: + # Read only the leased GGUF's own embedded template (the copy + # llama.cpp loads), never a sibling sidecar: the native grant + # authorizes just this path, so neighbours would be scope escalation. + raw_template = await asyncio.to_thread(read_gguf_chat_template, local_gguf) + if ( + raw_template is not None + and len(raw_template.encode("utf-8")) <= MAX_CHAT_TEMPLATE_BYTES + ): + chat_template = raw_template except Exception as e: - logger.debug("Context-length probe failed for %s: %s", model_log_label, e) + logger.debug("Header probe failed for %s: %s", model_log_label, e) return ValidateModelResponse( valid = True, @@ -4954,6 +5247,9 @@ async def validate_model( requires_trust_remote_code = requires_trust_remote_code, requires_security_review = requires_security_review, context_length = context_length, + layer_count = layer_count, + moe_layer_count = moe_layer_count, + chat_template = chat_template, requires_transformers_upgrade = transformers_upgrade is not None, transformers_upgrade = transformers_upgrade, ) @@ -4966,6 +5262,14 @@ async def validate_model( raise HTTPException(status_code = 400, detail = str(e)) except Exception as e: redacted_msg = redact_native_paths(str(e)) + if is_hf_authentication_error(e): + raise HTTPException( + status_code = 400, + detail = ( + "Hugging Face authentication failed. Check or clear the token " + "in Settings, and confirm access to this gated repository." + ), + ) if _is_unsupported_nvfp4_inference_error(redacted_msg): logger.warning( "NVFP4 inference is not supported yet while validating '%s'", @@ -5360,7 +5664,7 @@ async def confirm_tool_call( @studio_router.get("/monitor") async def get_api_monitor(current_subject: str = Depends(get_current_subject)): - """Return recent OpenAI-compatible API activity for Studio.""" + """Return recent OpenAI-compatible API activity for Unsloth.""" active_model = _monitor_active_model() active_requests = api_monitor.active_count(subject = current_subject) if active_requests: @@ -5548,7 +5852,7 @@ async def get_status(current_subject: str = Depends(get_current_subject)): _display_model_id = os.path.basename(_model_id) _inference_cfg = load_inference_config(_model_id) if _model_id else None _audio_type = getattr(llama_backend, "_audio_type", None) - # Don't surface Studio's auto-applied bundled family template (e.g. the + # Don't surface Unsloth's auto-applied bundled family template (e.g. the # gemma-4 override) as a user-authored override: the frontend adopts # status.chat_template_override as editable state and would otherwise # re-send it as an explicit override for a later, unrelated model. Only @@ -5593,6 +5897,14 @@ async def get_status(current_subject: str = Depends(get_current_subject)): speculative_type = llama_backend.requested_spec_mode, spec_draft_n_max = llama_backend.spec_draft_n_max, tensor_parallel = llama_backend.tensor_parallel, + gpu_memory_mode = llama_backend.gpu_memory_mode, + gpu_layers = llama_backend.gpu_layers, + n_cpu_moe = llama_backend.n_cpu_moe, + tensor_split = llama_backend.tensor_split, + requested_context_length = llama_backend.requested_n_ctx, + n_layers = llama_backend.n_layers, + n_moe_layers = llama_backend.n_moe_layers, + gpu_ids = llama_backend.gpu_ids, llama_cpp_supports_mtp = _supports_mtp, spec_fallback_reason = llama_backend.spec_fallback_reason, llama_cpp_prebuilt_stale = _stale, @@ -6173,7 +6485,7 @@ def _build_external_messages( metadata; strip it for providers that can't parse the unknown key. 2. Marked server-side builtin cards (`_server_tool: true` on a canonical builtin name, or a Gemini `native_part` payload) are - Studio-internal tool cards from a prior native Gemini turn; + Unsloth-internal tool cards from a prior native Gemini turn; forwarding them to OpenAI / Anthropic / custom OAI-compat gateways sends an orphan `tool_calls` entry (no matching tool declaration, often no matching `role="tool"` reply) that can be rejected. We @@ -6859,7 +7171,7 @@ async def openai_chat_completions( # is invalid and must not evict the resident model first. # # Enter the local-loop arm exactly when the passthrough router below would - # run Studio's own tool loop. That gate is `_tools_on or _mcp_allowed` + # run Unsloth's own tool loop. That gate is `_tools_on or _mcp_allowed` # (see the use_tools block): _effective_enable_tools (which lets a # process-wide --enable-tools policy force the loop on) plus mcp_enabled # honoring --disable-tools, and tool_choice="none" disabling it unless the @@ -6884,7 +7196,7 @@ async def openai_chat_completions( or bool(payload.openai_code_exec_container_id) or bool(payload.anthropic_code_exec_container_id) # A JSON-schema response_format is guided-decoding structured output the - # router forwards to the llama-server passthrough, not Studio's tool + # router forwards to the llama-server passthrough, not Unsloth's tool # loop, so a --enable-tools policy must not 400 it as a local-confirm # request under ask/auto. or bool(_extract_response_format(payload)) @@ -6961,7 +7273,7 @@ async def openai_chat_completions( using_gguf = llama_backend.is_loaded # OpenAI-SDK clients send ``chat_template_kwargs`` via ``extra_body``, which - # the SDK spreads into the request body at the top level. Studio's + # the SDK spreads into the request body at the top level. Unsloth's # ChatCompletionRequest has ``extra="allow"`` so pydantic stashes them in # ``model_extra``, but downstream generators consume the typed # ``payload.enable_thinking``. Lift ``enable_thinking`` from the extra-body @@ -7215,7 +7527,7 @@ async def openai_chat_completions( # ── Standard OpenAI function-calling pass-through (GGUF only) ──── # When a client (opencode / Claude Code via OpenAI compat / Cursor / - # Continue / ...) sends standard OpenAI `tools` without Studio's + # Continue / ...) sends standard OpenAI `tools` without Unsloth's # `enable_tools` shorthand, forward the request to llama-server # verbatim so structured `tool_calls` flow back to the client. This # branch runs BEFORE `_extract_content_parts` because that helper is @@ -7238,7 +7550,7 @@ async def openai_chat_completions( _has_tool_catalog = bool(payload.tools and len(payload.tools) > 0) _has_active_tool_catalog = _has_tool_catalog and payload.tool_choice != "none" _has_client_tool_contract = _has_active_tool_catalog or _has_tool_messages - # The Studio tool loop needs a tool-capable backend, so a request that asks + # The Unsloth tool loop needs a tool-capable backend, so a request that asks # for it on a backend that can't run it (DiffusionGemma forces supports_tools # off) must not steal client tools from the passthrough (#6851). _studio_tool_loop_requested = ( @@ -7434,7 +7746,7 @@ async def openai_chat_completions( use_tools = False if use_tools: - # permission_mode ask/auto require the confirm gate for Studio's own + # permission_mode ask/auto require the confirm gate for Unsloth's own # tool loop. The request validator self-enables confirm only for # request-level tool signals (enable_tools/enabled_tools/mcp_enabled); # when a CLI policy (--enable-tools) forces the loop on without those, @@ -8711,7 +9023,7 @@ async def openai_chat_completions( _sf_model_info = backend.models.get(backend.active_model_name, {}) _sf_tpl = (_sf_model_info.get("chat_template_info") or {}).get("template") # Named templates may expose native reasoning only in their ``tool_use`` - # branch. Use a truthy placeholder for Studio-managed tools, whose concrete + # branch. Use a truthy placeholder for Unsloth-managed tools, whose concrete # schemas are selected below, and the request schemas for client passthrough. _sf_server_tool_intent = bool( _effective_enable_tools(payload) or _explicit_studio_tool_loop_requested(payload) @@ -8790,7 +9102,7 @@ async def openai_chat_completions( _sf_use_tools = False if _sf_use_tools: - # permission_mode ask/auto require the confirm gate for Studio's own tool + # permission_mode ask/auto require the confirm gate for Unsloth's own tool # loop; when a CLI policy (--enable-tools) forces the loop on without a # request-level tool signal, derive confirm here so the mode still gates # the call (matching the GGUF path). off/full never prompt. @@ -9443,6 +9755,13 @@ async def openai_chat_completions( backend.reset_generation_state() api_monitor.finish(monitor_id, "cancelled") raise + except GenStreamErrorRaised as exc: + # Adapter-controlled (compare-mode) backend failure. Honor the + # public flag so operational errors surface their real message. + backend.reset_generation_state() + _msg = _friendly_gen_stream_error(exc) + api_monitor.fail(monitor_id, _msg) + yield _openai_stream_error_sse({"error": {"message": _msg, "type": "server_error"}}) except Exception as e: backend.reset_generation_state() logger.error(f"Error during OpenAI streaming: {e}", exc_info = True) @@ -9591,6 +9910,13 @@ async def openai_chat_completions( except HTTPException: raise + except GenStreamErrorRaised as exc: + # Adapter-controlled (compare-mode) backend failure. Honor the public + # flag so operational errors surface their real message. + backend.reset_generation_state() + _msg = _friendly_gen_stream_error(exc) + api_monitor.fail(monitor_id, _msg) + raise HTTPException(status_code = 500, detail = _msg) except Exception as e: backend.reset_generation_state() logger.error(f"Error during OpenAI completion: {e}", exc_info = True) @@ -12036,7 +12362,7 @@ def _anthropic_requested_studio_tools(tools: Optional[list]) -> set[str]: def _select_anthropic_server_tools( all_tools: list[dict], requested_studio_tools: set[str], enabled_tools: Optional[list[str]] ) -> list[dict]: - """Select Studio tools requested through Anthropic tools and extensions.""" + """Select Unsloth tools requested through Anthropic tools and extensions.""" if not requested_studio_tools and enabled_tools is None: return all_tools @@ -12275,7 +12601,7 @@ async def anthropic_messages( ), ) - # Reject an unsupported confirm-gated permission mode for Studio's own + # Reject an unsupported confirm-gated permission mode for Unsloth's own # ("server") Anthropic tools before the switch, mirroring the malformed- and # mixed-tool checks above. ask always wants a per-call pause this passthrough # cannot offer, so it 400s whenever server tools are selected. auto only needs @@ -12764,11 +13090,11 @@ async def _anthropic_tool_stream( ends_on_tool_use = True elif etype == "tool_end": tool_blocks_emitted += 1 - # A tool_end means Studio executed the tool server-side, so + # A tool_end means Unsloth executed the tool server-side, so # the response no longer ends on a pending client action. # Without this, a server tool that produces no trailing text # would be mislabeled stop_reason "tool_use", telling the - # client to run a tool Studio already ran. + # client to run a tool Unsloth already ran. ends_on_tool_use = False elif etype == "content" and event.get("text"): ends_on_tool_use = False @@ -13684,7 +14010,7 @@ def _openai_messages_for_passthrough(payload) -> list[dict]: structured ``tool_calls``. Content-parts images already in the list are left untouched. - When a client uses Studio's legacy ``image_base64`` top-level field, the + When a client uses Unsloth's legacy ``image_base64`` top-level field, the image is re-encoded to PNG (llama-server's stb_image has limited format support) and spliced into the last user message as an OpenAI ``image_url`` content part so vision + function-calling requests work transparently. @@ -13818,7 +14144,7 @@ def _build_openai_passthrough_body( ) -> dict: """Assemble the llama-server request body from a ChatCompletionRequest. - Only known OpenAI / llama-server fields are forwarded, so Studio-specific + Only known OpenAI / llama-server fields are forwarded, so Unsloth-specific extensions (``enable_tools``, ``enabled_tools``, ``session_id``, ...) never leak to the backend. """ @@ -14068,7 +14394,7 @@ async def _openai_passthrough_stream_admitted( admission_lease: LlamaAdmissionLease, tracker, ): - """Streaming client-side pass-through after Studio granted an upstream slot. + """Streaming client-side pass-through after Unsloth granted an upstream slot. Forwards the client's OpenAI function-calling request to llama-server and relays the SSE stream back with minimal normalization (reasoning-only diff --git a/studio/backend/routes/mcp_servers.py b/studio/backend/routes/mcp_servers.py index 71f0fd2874..dc018d163a 100644 --- a/studio/backend/routes/mcp_servers.py +++ b/studio/backend/routes/mcp_servers.py @@ -82,7 +82,7 @@ def _validate_url(url: str) -> str: if _looks_like_command(trimmed): detail = ( "Local commands aren't enabled on this server. To allow them, " - "set UNSLOTH_STUDIO_ALLOW_STDIO_MCP=1 and restart Studio, or use " + "set UNSLOTH_STUDIO_ALLOW_STDIO_MCP=1 and restart Unsloth, or use " "an http:// or https:// URL instead." ) else: diff --git a/studio/backend/routes/models.py b/studio/backend/routes/models.py index b8526c75e7..a5ce1a72f0 100644 --- a/studio/backend/routes/models.py +++ b/studio/backend/routes/models.py @@ -59,59 +59,51 @@ def _safe_is_dir(path) -> bool: return False -# Hub repo id shape ("owner/name", no leading separator); anything else is -# treated as a local filesystem path. -_HF_REPO_ID_RE = re.compile(r"^[A-Za-z0-9][\w.\-]*/[\w.\-]+$") +# Shared with the hub inventory scans; keep the private aliases so existing +# importers stay valid. ``_HF_REPO_ID_RE`` is the Hub repo id shape ("owner/name"); +# anything else is treated as a local filesystem path. +from utils.hidden_models import ( + _HF_REPO_ID_RE, + _existing_resolved_path, + _safe_resolve, + is_hidden_model as _is_hidden_model, +) -def _is_hidden_model(*values: str | None) -> bool: - """True if any id/path is the RAG embedding model (EMBEDDING_MODEL or - EMBED_GGUF_REPO basename) or the llama.cpp install validation probe - (ggml-org/models / stories260K), so pickers hide them (GGUF and non-GGUF). - None are usable chat models; the probe can be cached as a side effect of - installing the prebuilt llama-server and otherwise sorts smallest, so it - would be auto-selected. A local-path embedder is matched by exact resolved - path only: a generic basename like "model" must not substring-hide - unrelated chat models.""" +def hidden_model_matchers() -> tuple[list[str], list[str], list[str]]: + """Substring needles, exact repo ids, and exact resolved paths identifying + infra models (the RAG embedder and the llama.cpp install validation probe) + that pickers hide. Served by the ``/api/hub/hidden-models`` endpoint. A + configured HF-repo embedder is published as its exact lowercased repo id + (mirroring ``utils.hidden_models.is_hidden_model``) and a local-path + embedder as its exact resolved path only: a generic basename like "model" + must not substring-hide unrelated chat models.""" from core.rag import config as rag_config needles = [ - # The validation probe's repo (matches the cached repo id) and its exact - # filename (matches the on-disk path). The filename carries the .gguf so - # it does not hide unrelated repos like ``user/stories260K-finetune-GGUF``. + # The validation probe's repo and its exact filename. The filename carries + # .gguf so it won't hide unrelated repos like ``user/stories260K-finetune-GGUF``. "ggml-org/models", "stories260k.gguf", ] + exact_ids: list[str] = [] exact_paths: list[str] = [] for model in ( rag_config.effective_embedding_model(), rag_config.effective_gguf_repo(), ): - if _HF_REPO_ID_RE.match(model): - needles.append(model.split("/")[-1].lower()) + # Resolve an existing local path before the repo-id regex: a local embedder + # shaped like "models/embedder" is an exact path, not a Hub repo id. + existing_path = _existing_resolved_path(model) + if existing_path: + exact_paths.append(existing_path.lower()) + elif _HF_REPO_ID_RE.match(model): + exact_ids.append(model.lower()) else: resolved = _safe_resolve(Path(model).expanduser()) if resolved: exact_paths.append(resolved.lower()) - for v in values: - if not v: - continue - low = v.lower() - if any(n in low for n in needles): - return True - if exact_paths: - resolved = _safe_resolve(Path(v).expanduser()) - if resolved and resolved.lower() in exact_paths: - return True - return False - - -def _safe_resolve(path: Path) -> Optional[str]: - """resolve() to a string, or None when the path is inaccessible.""" - try: - return str(path.resolve()) - except OSError: - return None + return needles, exact_ids, exact_paths backend_path = Path(__file__).parent.parent.parent @@ -138,6 +130,7 @@ try: _pick_best_gguf, _extract_quant_label, _is_big_endian_gguf_path, + _is_mtp_drafter, is_audio_input_type, ) from core.inference import get_inference_backend @@ -170,6 +163,7 @@ except ImportError: _pick_best_gguf, _extract_quant_label, _is_big_endian_gguf_path, + _is_mtp_drafter, is_audio_input_type, ) from core.inference import get_inference_backend @@ -544,7 +538,7 @@ def _ollama_links_dir(ollama_dir: Path) -> Optional[Path]: """Return a writable directory for Ollama ``.gguf`` symlinks. Prefers ``/.studio_links/`` so links sit next to their - blobs; falls back to a per-ollama-dir namespace under Studio's cache + blobs; falls back to a per-ollama-dir namespace under Unsloth's cache when the models dir is read-only (common for system installs). """ from utils.paths.storage_roots import cache_root @@ -555,7 +549,7 @@ def _ollama_links_dir(ollama_dir: Path) -> Optional[Path]: return primary except OSError as e: logger.debug( - "Ollama dir %s not writable for .studio_links (%s); falling back to Studio cache", + "Ollama dir %s not writable for .studio_links (%s); falling back to Unsloth cache", ollama_dir, e, ) @@ -594,7 +588,7 @@ def _scan_ollama_dir(ollama_dir: Path, limit: Optional[int] = None) -> List[Loca model, keyed by a short hash of the manifest path, so ``detect_mmproj_file`` only sees that model's projector). Links are symlinks when possible, else hardlinks; the link dir is - ``.studio_links/`` when writable, else Studio's cache. + ``.studio_links/`` when writable, else Unsloth's cache. """ manifests_root = ollama_dir / "manifests" if not manifests_root.is_dir(): @@ -850,10 +844,10 @@ def collect_local_models(models_root: Path) -> List[LocalModelInfo]: models = sorted( deduped.values(), - key = lambda item: (item.updated_at or 0), + key = lambda item: item.updated_at or 0, reverse = True, ) - return [m for m in models if not _is_hidden_model(m.id, m.path)] + return [m for m in models if not _is_hidden_model(m.id, m.model_id, m.path)] @router.get("/local", response_model = LocalModelListResponse) @@ -1194,7 +1188,7 @@ def _build_browse_allowlist( """Return the root directories the folder browser may walk. The same list seeds the sidebar suggestion chips, so chip targets are - always reachable. Roots: HOME, resolved HF cache dirs, Studio's + always reachable. Roots: HOME, resolved HF cache dirs, Unsloth's outputs/exports/studio root, registered scan folders, and well-known local-LLM dirs (LM Studio, Ollama, ``~/models``); each added only if it resolves to a real directory. @@ -1486,7 +1480,7 @@ def browse_folders( "Directory to list. If omitted, defaults to the current user's " "home directory. Tilde (`~`) and relative paths are expanded. " "Must resolve inside the allowlist of browseable roots (HOME, " - "HF cache, Studio dirs, registered scan folders, well-known " + "HF cache, Unsloth dirs, registered scan folders, well-known " "model dirs)." ), ), @@ -1797,9 +1791,11 @@ def _get_model_size_bytes(model_name: str, hf_token: Optional[str] = None) -> Op async def get_model_config( model_name: str, hf_token: Optional[str] = Query(None), + header_hf_token: Optional[str] = Depends(get_hf_token), current_subject: str = Depends(get_current_subject), ): """Get configuration for a specific model (wraps load_model_defaults).""" + hf_token = _normalize_hf_token(header_hf_token) or _normalize_hf_token(hf_token) try: if not is_local_path(model_name): resolved = resolve_cached_repo_id_case(model_name) @@ -2251,15 +2247,15 @@ async def delete_finetuned_model( gguf_variant: Optional[str] = Body(None), current_subject: str = Depends(get_current_subject), ): - """Delete a Studio-trained or exported model from disk. + """Delete an Unsloth-trained or exported model from disk. - Only paths under Studio's outputs/exports roots are accepted. + Only paths under Unsloth's outputs/exports roots are accepted. Exported GGUF entries can delete one quant variant at a time. """ if source not in {"training", "exported"}: raise HTTPException( status_code = 400, - detail = "Only trained or exported Studio models can be deleted", + detail = "Only trained or exported Unsloth models can be deleted", ) if not model_path or not model_path.strip(): @@ -2291,14 +2287,14 @@ async def delete_finetuned_model( if not _is_path_under_lexically(delete_path, allowed_root): raise HTTPException( status_code = 400, - detail = "Model path is outside Studio storage", + detail = "Model path is outside Unsloth storage", ) if export_type == "gguf" and gguf_variant: target_path = delete_path.resolve() if not _is_path_under(target_path, allowed_root): raise HTTPException( status_code = 400, - detail = "Model path is outside Studio storage", + detail = "Model path is outside Unsloth storage", ) else: target_path = delete_path @@ -2311,7 +2307,7 @@ async def delete_finetuned_model( if should_check_resolved_path and not _is_path_under(target_path, allowed_root): raise HTTPException( status_code = 400, - detail = "Model path is outside Studio storage", + detail = "Model path is outside Unsloth storage", ) if target_path == allowed_root: raise HTTPException( @@ -2518,6 +2514,7 @@ async def get_lora_base_model(lora_path: str, current_subject: str = Depends(get async def check_vision_model( model_name: str, hf_token: Optional[str] = Query(None), + header_hf_token: Optional[str] = Depends(get_hf_token), current_subject: str = Depends(get_current_subject), ): """ @@ -2525,6 +2522,7 @@ async def check_vision_model( This endpoint wraps the backend is_vision_model function. """ + hf_token = _normalize_hf_token(header_hf_token) or _normalize_hf_token(hf_token) try: logger.info(f"Checking if vision model: {model_name}") # Authenticate so a gated/private VLM classifies correctly (else 404 -> non-vision). @@ -2550,6 +2548,7 @@ async def check_vision_model( async def check_embedding_model( model_name: str, hf_token: Optional[str] = Query(None), + header_hf_token: Optional[str] = Depends(get_hf_token), current_subject: str = Depends(get_current_subject), ): """ @@ -2557,6 +2556,7 @@ async def check_embedding_model( This endpoint wraps the backend is_embedding_model function. """ + hf_token = _normalize_hf_token(header_hf_token) or _normalize_hf_token(hf_token) try: logger.info(f"Checking if embedding model: {model_name}") is_embedding = is_embedding_model(model_name, hf_token = hf_token) @@ -2620,12 +2620,6 @@ def _resolve_quant_gguf(repo_id: str, quant: str, is_local: bool) -> tuple[Optio Q8_0 weights). Never raises. """ try: - from utils.models.model_config import ( - _extract_quant_label, - _is_big_endian_gguf_path, - _is_mtp_drafter, - ) - if is_local: roots = [Path(repo_id)] else: @@ -2642,25 +2636,19 @@ def _resolve_quant_gguf(repo_id: str, quant: str, is_local: bool) -> tuple[Optio if snaps.is_dir(): roots.extend(s for s in snaps.iterdir() if s.is_dir()) - want = quant.lower().replace("-", "").replace("_", "") + want = _normalized_quant_label(quant) best_total = 0 best_first: Optional[str] = None for root in roots: matches: list[tuple[str, Path]] = [] total = 0 for f in _iter_gguf_paths(root): - if _is_mmproj_filename(f.name): - continue try: rel = f.relative_to(root).as_posix() except ValueError: rel = f.name - if _is_mtp_drafter(rel): - continue - q = _extract_quant_label(rel) - if _is_big_endian_gguf_path(rel, q): - continue - if q.lower().replace("-", "").replace("_", "") != want: + q = _main_variant_gguf_label(rel) + if q is None or _normalized_quant_label(q) != want: continue try: total += f.stat().st_size @@ -2778,7 +2766,11 @@ async def get_gguf_variants( ], has_vision = response.has_vision, default_variant = response.default_variant, - context_length = _read_native_context_length(repo_id, is_local = local), + # The header walk reads tokenizer arrays on dense models (tens of + # ms per uncached file); keep it off the event loop. + context_length = await asyncio.to_thread( + _read_native_context_length, repo_id, is_local = local + ), ) except HTTPException: raise @@ -3078,6 +3070,22 @@ def _is_main_gguf_filename(name: str) -> bool: return _is_gguf_filename(name) and not _is_mmproj_filename(name) +def _main_variant_gguf_label(rel_path: str) -> Optional[str]: + name = rel_path.rsplit("/", 1)[-1] + if not _is_main_gguf_filename(name): + return None + if _is_mtp_drafter(rel_path): + return None + label = _extract_quant_label(rel_path) + if _is_big_endian_gguf_path(rel_path, label): + return None + return label + + +def _normalized_quant_label(label: str) -> str: + return label.lower().replace("-", "").replace("_", "") + + def _repo_has_mmproj(repo_info) -> bool: """True if the repo ships a GGUF vision adapter (mmproj), so it can take image inputs. Cheap: scans already-listed file names only.""" @@ -3405,6 +3413,170 @@ async def delete_cached_model( ) +def _resolve_cached_model_path(repo_id: str, variant: Optional[str]) -> Path: + """Absolute path of a cached repo (newest snapshot dir) or, with *variant*, + that quant's main GGUF file (first split of a sharded quant). Paths come + from the HF cache scan only, so callers can't probe arbitrary paths.""" + cache_scans = _all_hf_cache_scans() + + matching_repos = [] + for hf_cache in cache_scans: + for repo_info in hf_cache.repos: + if repo_info.repo_type != "model": + continue + if repo_info.repo_id.lower() == repo_id.lower(): + matching_repos.append(repo_info) + if not matching_repos: + raise HTTPException(status_code = 404, detail = "Model not found in cache") + + if variant: + want = _normalized_quant_label(variant) + candidate_revisions = sorted( + (rev for repo_info in matching_repos for rev in repo_info.revisions), + key = lambda rev: getattr(rev, "last_modified", 0) or 0, + reverse = True, + ) + for rev in candidate_revisions: + snapshot = getattr(rev, "snapshot_path", None) + matches = [] + for f in rev.files: + p = Path(f.file_path) + rel = f.file_name + if snapshot: + try: + rel = p.relative_to(snapshot).as_posix() + except ValueError: + pass + label = _main_variant_gguf_label(rel) + if label is None or _normalized_quant_label(label) != want: + continue + if p.exists() or p.is_symlink(): + matches.append((rel, p)) + if matches: + # Path-sorted so a sharded quant deterministically yields its first split. + return sorted(matches, key = lambda m: m[0].lower())[0][1] + raise HTTPException( + status_code = 404, + detail = f"Variant {variant} not found in cache for {repo_id}", + ) + + def repo_size(repo_info) -> int: + gguf_size = _repo_gguf_size_bytes(repo_info) + if gguf_size > 0: + return gguf_size + return sum( + (getattr(f, "size_on_disk", None) or 0) + for rev in repo_info.revisions + for f in rev.files + ) + + def repo_last_modified(repo_info) -> float: + return max( + (getattr(rev, "last_modified", 0) or 0 for rev in repo_info.revisions), + default = 0, + ) + + target_repo = max( + matching_repos, + key = lambda repo_info: (repo_size(repo_info), repo_last_modified(repo_info)), + ) + + # Whole repo: the newest revision's snapshot dir holds the visible files. + revisions = sorted( + (rev for rev in target_repo.revisions if getattr(rev, "snapshot_path", None)), + key = lambda rev: getattr(rev, "last_modified", 0) or 0, + reverse = True, + ) + for rev in revisions: + p = Path(rev.snapshot_path) + if p.exists(): + return p + p = Path(target_repo.repo_path) + if p.exists(): + return p + raise HTTPException(status_code = 404, detail = "Cached model path not found") + + +def _wsl_reveal_in_explorer(path: Path) -> bool: + import subprocess + + from utils.paths.path_utils import _IS_WSL + + if not _IS_WSL: + return False + try: + windows_path = subprocess.run( + ["wslpath", "-w", str(path)], + capture_output = True, + text = True, + check = True, + timeout = 10, + ).stdout.strip() + if not windows_path: + return False + argument = f"/select,{windows_path}" if path.is_file() else windows_path + subprocess.Popen(["explorer.exe", argument]) + return True + except (OSError, subprocess.SubprocessError): + return False + + +def _reveal_in_file_manager(path: Path) -> None: + """Open the OS file manager with *path* selected (best effort per platform).""" + import subprocess + + target = str(path) + if sys.platform == "darwin": + cmd = ["open", "-R", target] if path.is_file() else ["open", target] + subprocess.Popen(cmd) + elif os.name == "nt": + if path.is_file(): + subprocess.Popen(["explorer", f"/select,{target}"]) + else: + os.startfile(target) # noqa: S606 - local user's own file manager + elif not _wsl_reveal_in_explorer(path): + # No cross-desktop "select file" standard on Linux; open the directory. + directory = target if path.is_dir() else str(path.parent) + subprocess.Popen(["xdg-open", directory]) + + +class CachedModelPathResponse(BaseModel): + path: str + is_dir: bool + + +@router.get("/cached-model-path", response_model = CachedModelPathResponse) +async def get_cached_model_path( + repo_id: str = Query(..., description = "HuggingFace repo ID"), + variant: str = Query("", description = "Quantization variant (empty for whole repo)"), + current_subject: str = Depends(get_current_subject), +): + """Absolute on-disk path of a cached repo or one of its GGUF variants.""" + if not _is_valid_repo_id(repo_id): + raise HTTPException(status_code = 400, detail = "Invalid repo_id format") + path = await asyncio.to_thread(_resolve_cached_model_path, repo_id, variant.strip() or None) + return {"path": str(path), "is_dir": path.is_dir()} + + +@router.post("/reveal-cached-model") +async def reveal_cached_model( + repo_id: str = Body(...), + variant: Optional[str] = Body(None), + current_subject: str = Depends(get_current_subject), +): + """Reveal a cached repo (or one GGUF variant's file) in the OS file manager.""" + if not _is_valid_repo_id(repo_id): + raise HTTPException(status_code = 400, detail = "Invalid repo_id format") + variant = (variant or "").strip() or None + path = await asyncio.to_thread(_resolve_cached_model_path, repo_id, variant) + try: + await asyncio.to_thread(_reveal_in_file_manager, path) + except Exception as e: + logger.error(f"Failed to reveal {path}: {e}") + raise HTTPException(status_code = 500, detail = "Failed to open file manager") + return {"status": "ok", "path": str(path)} + + @router.get("/checkpoints", response_model = CheckpointListResponse) async def list_checkpoints( outputs_dir: str = Query( @@ -3456,7 +3628,7 @@ _EXPORT_SIZE_CACHE: dict[str, tuple[int, int, str]] = {} def _is_sizable_local_path(model: str) -> bool: - """True only for local paths under a Studio data root. + """True only for local paths under an Unsloth data root. Containment is decided lexically (no filesystem access) before the path is touched, then the path is symlink-resolved and re-checked so a symlink diff --git a/studio/backend/routes/rag.py b/studio/backend/routes/rag.py index e20fea74a3..392a4e0d02 100644 --- a/studio/backend/routes/rag.py +++ b/studio/backend/routes/rag.py @@ -318,6 +318,39 @@ def list_project_documents(project_id: str, subject: str = Depends(get_current_s conn.close() +@router.get("/documents") +def list_all_uploaded_documents(subject: str = Depends(get_current_subject)) -> dict: + """Every uploaded file across chats, projects, and knowledge bases (settings + Data tab).""" + _require_rag() + conn = rag_db.get_connection() + try: + docs = store.list_all_documents(conn) + kb_names = {kb["id"]: kb["name"] for kb in store.list_kbs(conn)} + finally: + conn.close() + + from storage.studio_db import list_chat_projects + + project_names = {p["id"]: p["name"] for p in list_chat_projects(include_archived = True)} + + out = [] + for doc in docs: + view = _doc_view(doc) + stored_path = doc.get("stored_path") + size = None + if stored_path: + try: + size = os.path.getsize(stored_path) + except OSError: + size = None + view["sizeBytes"] = size + view["kbName"] = kb_names.get(doc.get("kb_id")) + view["projectName"] = project_names.get(doc.get("project_id")) + out.append(view) + return {"documents": out} + + @router.delete("/documents/{document_id}") def delete_document(document_id: str, subject: str = Depends(get_current_subject)) -> dict: _require_rag() @@ -424,8 +457,10 @@ _CONTENT_TYPES = { ".txt": "text/plain; charset=utf-8", ".md": "text/markdown; charset=utf-8", ".markdown": "text/markdown; charset=utf-8", - ".html": "text/html; charset=utf-8", - ".htm": "text/html; charset=utf-8", + # Served as plain text, never text/html: an uploaded HTML document rendered + # same-origin would execute its scripts with access to the app's storage. + ".html": "text/plain; charset=utf-8", + ".htm": "text/plain; charset=utf-8", ".docx": "application/vnd.openxmlformats-officedocument.wordprocessingml.document", } diff --git a/studio/backend/routes/settings.py b/studio/backend/routes/settings.py index 1ddfc0eacb..17e64df918 100644 --- a/studio/backend/routes/settings.py +++ b/studio/backend/routes/settings.py @@ -10,6 +10,7 @@ from pydantic import BaseModel, ConfigDict, Field, field_validator from auth.authentication import get_current_subject from auth.storage import rotate_preview_link_secret +from core.rag.config import default_gguf_repo, effective_gguf_repo from loggers import get_logger from utils.utils import safe_error_detail, log_and_http_error from utils.personalization_settings import ( @@ -35,9 +36,10 @@ from utils.helper_precache_settings import ( ) from utils.coding_agents import CODING_AGENTS, detect_installed_coding_agents from utils.openai_auto_switch_settings import ( - DEFAULT_AUTO_UNLOAD_IDLE_SECONDS, + DEFAULT_AUTO_UNLOAD_KEEP_KV, DEFAULT_OPENAI_AUTO_SWITCH_ENABLED, get_auto_unload_idle_seconds, + get_auto_unload_keep_kv, get_model_overrides, get_openai_auto_switch_enabled, get_stored_auto_unload_idle_seconds, @@ -89,7 +91,9 @@ class HelperPrecacheResponse(BaseModel): class OpenAIAutoSwitchPayload(BaseModel): enabled: bool - auto_unload_idle_seconds: int = Field(default = DEFAULT_AUTO_UNLOAD_IDLE_SECONDS, ge = 0) + # None leaves the stored value untouched (partial updates can't clobber it). + auto_unload_idle_seconds: Optional[int] = Field(default = None, ge = 0) + auto_unload_keep_kv: Optional[bool] = None class OpenAIAutoSwitchResponse(BaseModel): @@ -100,6 +104,7 @@ class OpenAIAutoSwitchResponse(BaseModel): # UNSLOTH_MODEL_IDLE_TTL set and nothing stored, this is true even while enabled # is false, so the UI can show idle-unload as active instead of "needs enable". idle_unload_active: bool = False + auto_unload_keep_kv: bool = DEFAULT_AUTO_UNLOAD_KEEP_KV class ModelOverridePayload(BaseModel): @@ -197,6 +202,7 @@ def get_openai_auto_switch( enabled = get_openai_auto_switch_enabled(), auto_unload_idle_seconds = get_stored_auto_unload_idle_seconds(), idle_unload_active = get_auto_unload_idle_seconds() > 0, + auto_unload_keep_kv = get_auto_unload_keep_kv(), ) @@ -205,8 +211,8 @@ def update_openai_auto_switch( payload: OpenAIAutoSwitchPayload, current_subject: str = Depends(get_current_subject) ) -> OpenAIAutoSwitchResponse: try: - enabled, idle_seconds = set_openai_auto_switch( - payload.enabled, payload.auto_unload_idle_seconds + enabled, idle_seconds, keep_kv = set_openai_auto_switch( + payload.enabled, payload.auto_unload_idle_seconds, payload.auto_unload_keep_kv ) except ValueError as exc: raise log_and_http_error( @@ -216,10 +222,16 @@ def update_openai_auto_switch( event = "settings.update_openai_auto_switch_failed", log = logger, ) from exc + idle_unload_active = get_auto_unload_idle_seconds() > 0 + if not keep_kv or not idle_unload_active: + # Keep-KV off or idle unload disabled: drop already-saved chat context too. + from core.inference.llama_keepwarm import purge_kv_resume + purge_kv_resume() return OpenAIAutoSwitchResponse( enabled = enabled, auto_unload_idle_seconds = idle_seconds, - idle_unload_active = get_auto_unload_idle_seconds() > 0, + idle_unload_active = idle_unload_active, + auto_unload_keep_kv = keep_kv, ) @@ -263,14 +275,18 @@ class EmbeddingModelPayload(BaseModel): class EmbeddingModelResponse(BaseModel): embedding_model: str + embedding_gguf_repo: str default_embedding_model: str + default_embedding_gguf_repo: str is_custom: bool def _embedding_model_response() -> EmbeddingModelResponse: return EmbeddingModelResponse( embedding_model = get_rag_embedding_model(), + embedding_gguf_repo = effective_gguf_repo(), default_embedding_model = default_embedding_model(), + default_embedding_gguf_repo = default_gguf_repo(), is_custom = get_stored_embedding_model() is not None, ) diff --git a/studio/backend/routes/training.py b/studio/backend/routes/training.py index d53e8f2bbc..a8a9874b1b 100644 --- a/studio/backend/routes/training.py +++ b/studio/backend/routes/training.py @@ -127,9 +127,9 @@ async def start_training( try: logger.info(f"Starting training job with model: {request.model_name}") - # When Studio is driven as an inference API (API-key auth), refuse to start + # When Unsloth is driven as an inference API (API-key auth), refuse to start # training while a request is in flight: training frees VRAM by unloading - # the chat model, which would kill the stream. The Studio UI (session auth) + # the chat model, which would kill the stream. The Unsloth UI (session auth) # still starts training and coexists/frees VRAM as before. (A mixed UI+API # session is not yet special-cased.) if via_api_key is True: @@ -139,7 +139,7 @@ async def start_training( status_code = 409, detail = ( "Cannot start training over the API while an inference request is in " - "progress. Wait for it to finish, or start training from the Studio UI." + "progress. Wait for it to finish, or start training from the Unsloth UI." ), ) @@ -196,6 +196,7 @@ async def start_training( request.local_eval_datasets, "Local eval dataset" ) resume_output_dir: Optional[str] = None + resume_run: Optional[dict] = None if request.resume_from_checkpoint: try: resume_output_dir = normalize_resume_output_dir(request.resume_from_checkpoint) @@ -208,7 +209,7 @@ async def start_training( if not resume_run or not can_resume_run(resume_run): raise HTTPException( status_code = 400, - detail = "Resume checkpoint must belong to a stopped run with saved trainer state.", + detail = "Resume checkpoint must belong to a stopped or errored run with complete saved trainer state.", ) resume_checkpoint = get_resume_checkpoint_path(resume_output_dir) if not resume_checkpoint: @@ -458,7 +459,10 @@ async def start_training( try: success = backend.start_training( - job_id = job_id, before_spawn = _free_vram_for_training, **training_kwargs + job_id = job_id, + before_spawn = _free_vram_for_training, + resume_source_run_id = resume_run["id"] if resume_run else None, + **training_kwargs, ) except SidecarSwapInProgress as exc: # Expected loss of the race against a sidecar install: a retryable @@ -521,7 +525,10 @@ async def stop_training( status = "idle", message = "No training job is currently running" ) - backend.stop_training(save = body.save) + if not backend.stop_training(save = body.save): + return TrainingStopResponse( + status = "idle", message = "No training job is currently running" + ) return TrainingStopResponse( status = "stopped", @@ -637,9 +644,9 @@ async def get_training_status(current_subject: str = Depends(get_current_subject "loss": getattr(progress, "loss", None), "learning_rate": getattr(progress, "learning_rate", None), } - output_dir = getattr(backend, "_output_dir", None) - if output_dir: - details["output_dir"] = output_dir + # Always present: an explicit null tells the client to drop a cached + # path (stop without save clears the run's output_dir). + details["output_dir"] = getattr(backend, "_output_dir", None) or None # Metric history for chart recovery after SSE reconnection. metric_history = None diff --git a/studio/backend/routes/training_vram.py b/studio/backend/routes/training_vram.py index fb361d3359..fd96fe2175 100644 --- a/studio/backend/routes/training_vram.py +++ b/studio/backend/routes/training_vram.py @@ -197,15 +197,18 @@ def can_load_chat_during_training( requested_gpu_ids: Optional[List[int]], is_gguf: bool = False, required_override_gb: Optional[float] = None, + single_device_gpu: Optional[str] = None, ) -> Tuple[bool, Dict[str, Any]]: """Decide if a NEW chat model can load without OOMing active training (inverse of can_keep_chat_during_training: training is already resident, so size the chat model against the free VRAM that remains). Sizes/places it the same way the loader will: HF auto reuses auto_select_gpu_ids; HF explicit requires an even-share per-GPU floor for device_map="balanced"; GGUF sizes from - required_override_gb over the visible pool. `load_in_4bit` must be effective - (LoRA can flip 4-bit -> 16-bit). Non-CUDA allows the load; default-deny on any - CUDA case it can't size, so a load never OOMs training.""" + required_override_gb over the visible pool. ``single_device_gpu`` is the + exact physical device token selected by a single-device runner. + `load_in_4bit` must be effective (LoRA can flip 4-bit -> 16-bit). Non-CUDA + allows the load; default-deny on any CUDA case it can't size, so a load never + OOMs training.""" try: from utils.hardware import ( DeviceType, @@ -251,26 +254,49 @@ def can_load_chat_during_training( } # Explicit GPUs, or GGUF: size directly and check live free VRAM. + if single_device_gpu is not None: + mode = "single_device" + elif is_gguf: + mode = "gguf" + else: + mode = "explicit" required_gb = required_override_gb if required_gb is None: required_gb, _meta = estimate_required_model_memory_gb(model_name, **est_kwargs) if required_gb is None: - mode = "explicit" if requested_gpu_ids else "gguf" return False, {"mode": mode, "reason": "estimate_unavailable"} free_by_index = _free_vram_by_index(get_visible_gpu_utilization().get("devices", [])) - if requested_gpu_ids: + if single_device_gpu is not None: + token = str(single_device_gpu).strip() + if not token: + # Empty token = a CPU-only single-device runner (e.g. a CPU + # diffusion GGUF): it uses no GPU VRAM, so it never threatens + # active training and can always load. + return True, {"mode": "single_device", "reason": "cpu_only"} + try: + selected_gpu = int(token) + if selected_gpu < 0: + raise ValueError + except (TypeError, ValueError): + # A non-numeric device token (e.g. a CUDA UUID / MIG handle) + # can't be mapped to a free-VRAM index, but the runner still + # drives ONE device. Size against the worst-case visible device + # (min free), never the aggregate pool, so a single-device load + # is never OK'd on capacity it can't use and OOMs training. + free_vals = [min(free_by_index.values())] if free_by_index else [] + else: + free_vals = [free_by_index.get(selected_gpu, 0.0)] + elif requested_gpu_ids: # Invalid ids -> load_model 400s first, so don't block; missing id = 0. try: resolved = resolve_requested_gpu_ids(requested_gpu_ids) except ValueError: - return True, {"mode": "explicit", "reason": "invalid_gpu_ids"} + return True, {"mode": mode, "reason": "invalid_gpu_ids"} free_vals = [free_by_index.get(i, 0.0) for i in resolved] - mode = "explicit" else: # GGUF: llama.cpp picks the GPU(s); any visible GPU is a candidate. free_vals = list(free_by_index.values()) - mode = "gguf" if not free_vals: return False, {"mode": mode, "reason": "no_visible_gpus"} diff --git a/studio/backend/run.py b/studio/backend/run.py index 4f105b53b1..398943cc2c 100644 --- a/studio/backend/run.py +++ b/studio/backend/run.py @@ -232,7 +232,7 @@ def _working_local_url(port: int) -> "str | None": def _localhost_ipv6_mismatch_url(bind_host: str, port: int) -> "str | None": """Return the IPv4 loopback URL when localhost won't reach 127.0.0.1. - Local Studio binds to 127.0.0.1. Where localhost resolves to IPv6 only (::1), + Local Unsloth binds to 127.0.0.1. Where localhost resolves to IPv6 only (::1), http://localhost: fails (or hits a different process on ::1) even though http://127.0.0.1: works. Return the IPv4 URL for the caller to surface. """ @@ -243,7 +243,7 @@ def _localhost_ipv6_mismatch_url(bind_host: str, port: int) -> "str | None": ipv4_url = f"http://127.0.0.1:{port}" - # Only warn once Studio is confirmed answering on IPv4 loopback. + # Only warn once Unsloth is confirmed answering on IPv4 loopback. if _working_local_url(port) != ipv4_url: return None @@ -265,7 +265,7 @@ def _localhost_ipv6_mismatch_url(bind_host: str, port: int) -> "str | None": if host == "::1": has_ipv6_loopback = True - # A connection to ::1 is NOT evidence Studio is reachable there: Studio binds + # A connection to ::1 is NOT evidence Unsloth is reachable there: Unsloth binds # 127.0.0.1 only, so anything on ::1 is a different process. Dual-stack # localhost is fine (browsers fall back to 127.0.0.1), so only the IPv6-only # case strands the user. @@ -287,7 +287,7 @@ def _stdout_color_ok() -> bool: def _print_localhost_ipv6_mismatch_warning(local_url: str, port: int) -> None: - """Warn that localhost points at ::1 while Studio is bound to 127.0.0.1.""" + """Warn that localhost points at ::1 while Unsloth is bound to 127.0.0.1.""" use_color = _stdout_color_ok() warn_c = "\033[38;5;215;1m" if use_color else "" reset = "\033[0m" if use_color else "" @@ -303,7 +303,7 @@ def _print_localhost_ipv6_mismatch_warning(local_url: str, port: int) -> None: def _verify_global_reachability(display_host: str, port: int) -> None: """Probe check-host.net to confirm display_host:port is reachable from the public internet. Synchronous so output lands between the banner URLs and the - stop hint. Bounded at ~15s; failures swallowed (verifier failing != Studio + stop hint. Bounded at ~15s; failures swallowed (verifier failing != Unsloth failing). Only meaningful for a wildcard bind.""" global _public_reachable # Reset to "unknown" each run; set True/False only when the probe decides. @@ -563,15 +563,15 @@ def _print_cloudflare_line(secure: bool = False, loopback_host: str = "127.0.0.1 " Cloudflare tunnel: ON. This Cloudflare URL is PUBLIC, and the " "raw port is also publicly reachable. --no-cloudflare disables " f"only the Cloudflare URL; bind {loopback_host} or close firewall " - "access to keep Studio private.", + "access to keep Unsloth private.", warn, ) else: _emit( " Cloudflare tunnel: ON. This is a PUBLIC internet URL: anyone " - "who has it can reach this Studio. Relaunch with --no-cloudflare " + "who has it can reach this Unsloth. Relaunch with --no-cloudflare " f"to disable the Cloudflare URL; bind {loopback_host} or close " - "firewall access to keep Studio private.", + "firewall access to keep Unsloth private.", warn, ) return @@ -580,12 +580,12 @@ def _print_cloudflare_line(secure: bool = False, loopback_host: str = "127.0.0.1 _emit( " Cloudflare tunnel: requested but failed to start. The raw port is " "still reachable from the public internet (see the reachability check " - "above): anyone who can reach it can access this Studio.", + "above): anyone who can reach it can access this Unsloth.", warn, ) elif _public_reachable is False: _emit( - " Cloudflare tunnel: requested but failed to start. Studio is reachable " + " Cloudflare tunnel: requested but failed to start. Unsloth is reachable " "on your local network only (no public link).", warn, ) @@ -593,7 +593,7 @@ def _print_cloudflare_line(secure: bool = False, loopback_host: str = "127.0.0.1 _emit( " Cloudflare tunnel: requested but failed to start. There is no " "Cloudflare public link. Raw port reachability was not verified; " - f"bind {loopback_host} or close firewall access to keep Studio private.", + f"bind {loopback_host} or close firewall access to keep Unsloth private.", warn, ) elif _cloudflare_flag: @@ -601,19 +601,19 @@ def _print_cloudflare_line(secure: bool = False, loopback_host: str = "127.0.0.1 _emit( " Cloudflare tunnel: OFF for this mode. The raw port is still " "reachable from the public internet (see the reachability check above): " - "anyone who can reach it can access this Studio.", + "anyone who can reach it can access this Unsloth.", warn, ) elif _public_reachable is False: _emit( - " Cloudflare tunnel: OFF for this mode. Studio is reachable on your " + " Cloudflare tunnel: OFF for this mode. Unsloth is reachable on your " "local network only (no public link)." ) else: _emit( " Cloudflare tunnel: OFF for this mode. There is no Cloudflare public " "link. Raw port reachability was not verified; " - f"bind {loopback_host} or close firewall access to keep Studio private.", + f"bind {loopback_host} or close firewall access to keep Unsloth private.", warn, ) elif _cloudflare_flag is False or _cloudflare_flag is None: @@ -624,12 +624,12 @@ def _print_cloudflare_line(secure: bool = False, loopback_host: str = "127.0.0.1 f" Cloudflare tunnel: OFF ({_reason}). The raw port is still " "reachable from the public internet (see the reachability check above): " "pass --cloudflare to also expose a public Cloudflare HTTPS link, or " - f"bind {loopback_host} to keep Studio private.", + f"bind {loopback_host} to keep Unsloth private.", warn, ) elif _public_reachable is False: _emit( - f" Cloudflare tunnel: OFF ({_reason}). Studio is reachable on your " + f" Cloudflare tunnel: OFF ({_reason}). Unsloth is reachable on your " "local network only. Pass --cloudflare to expose a public " "Cloudflare HTTPS link." ) @@ -638,7 +638,7 @@ def _print_cloudflare_line(secure: bool = False, loopback_host: str = "127.0.0.1 f" Cloudflare tunnel: OFF ({_reason}). There is no Cloudflare " "public link. Raw port reachability was not verified; pass --cloudflare " "to expose a public Cloudflare HTTPS link, or " - f"bind {loopback_host} or close firewall access to keep Studio private.", + f"bind {loopback_host} or close firewall access to keep Unsloth private.", warn, ) @@ -674,7 +674,7 @@ def _is_port_free(host: str, port: int) -> bool: For a ``0.0.0.0`` wildcard host, also check whether anything is listening on ``127.0.0.1`` (and ``::1`` when IPv6 exists): an SSH tunnel may hold loopback - while the wildcard bind succeeds, making Studio unreachable via ``localhost``. + while the wildcard bind succeeds, making Unsloth unreachable via ``localhost``. """ import socket @@ -1087,7 +1087,7 @@ def _terminal_password_gate( ) -> Tuple[bool, bool]: """Force a terminal password change before the public tunnel goes up. - When the tunnel is about to publish Studio and the seeded admin password was + When the tunnel is about to publish Unsloth and the seeded admin password was never changed, ask for a new one (masked, confirmed) before any public URL exists. The CLI normally does this before re-exec'ing the backend; this is the backstop for direct `python run.py` launches and older-CLI installs. @@ -1147,7 +1147,7 @@ def _terminal_password_gate( ) if not deadline_arms: print( - "Refusing to publish Studio on a public Cloudflare URL: the " + "Refusing to publish Unsloth on a public Cloudflare URL: the " "default admin password was never changed, no terminal is " "attached to change it here, and the bootstrap shutdown " "deadline does not apply to this launch (api-only, or " @@ -1163,11 +1163,11 @@ def _terminal_password_gate( # terminal-attached run / reset-password instead of reading it from disk. print( " WARNING: the default admin password is still active while " - "Studio is about to be published on a public Cloudflare URL, and " + "Unsloth is about to be published on a public Cloudflare URL, and " "no terminal is attached to change it here. The public page will " "NOT auto-fill the bootstrap credential. Set a new password by " "running `unsloth studio` locally with a terminal attached, or " - "`unsloth studio reset-password`. Studio shuts down after the " + "`unsloth studio reset-password`. Unsloth shuts down after the " "bootstrap deadline (UNSLOTH_STUDIO_BOOTSTRAP_TIMEOUT, default 1h) " "unless the password is changed.", file = sys.stderr, @@ -1222,7 +1222,7 @@ def _apply_supplied_password(password_value: "Optional[str]") -> None: _auth_storage.ensure_default_admin() if not _auth_storage.requires_password_change(_admin): print( - "Error: a Studio admin password is already set; --password only sets " + "Error: an Unsloth admin password is already set; --password only sets " "the initial password. Run `unsloth studio reset-password` first.", file = sys.stderr, flush = True, @@ -1337,7 +1337,7 @@ def run_server( pass # Persist a session log + native-crash stacks BEFORE importing main, so - # even import-time failures leave evidence on disk. Field report: Studio + # even import-time failures leave evidence on disk. Field report: Unsloth # "terminates without a warning" -- a native crash in the GPU runtime # kills the process with no Python traceback, and a desktop-shortcut # console closes before anything can be read. Console-only logging made @@ -1406,7 +1406,7 @@ def run_server( ensure_studio_directories() logger.info( - "Ensured Studio directories in %.1fms", + "Ensured Unsloth directories in %.1fms", (time.perf_counter() - boot_started) * 1000, ) @@ -1455,7 +1455,7 @@ def run_server( installer_bin = home / "unsloth_studio" / "bin" / "unsloth" tried_lines = "\n".join(f" - {p}" for p in attempted) or " (none)" raise SystemExit( - "[ERROR] Studio frontend build not found.\n" + "[ERROR] Unsloth frontend build not found.\n" f"Tried:\n{tried_lines}\n" "\n" "Likely cause: another 'unsloth' on PATH is shadowing the " @@ -1557,7 +1557,7 @@ def run_server( ) if not _pw_proceed: print( - "Not starting Studio; set a new admin password first, or launch " + "Not starting Unsloth; set a new admin password first, or launch " "without --secure/--cloudflare.", file = sys.stderr, flush = True, @@ -1695,7 +1695,7 @@ def run_server( logger = logger, ) logger.info( - "Studio will shut down in %ds unless the default admin password is changed.", + "Unsloth will shut down in %ds unless the default admin password is changed.", _bootstrap_timeout, ) except Exception as e: # best-effort: never block startup on the timeout @@ -1753,11 +1753,11 @@ def _build_arg_parser(): "--cloudflare", action = argparse.BooleanOptionalAction, default = None, - help = "Expose Studio on a PUBLIC internet URL via a free Cloudflare HTTPS " + help = "Expose Unsloth on a PUBLIC internet URL via a free Cloudflare HTTPS " "tunnel, for non-api-only wildcard binds (0.0.0.0 or ::). Off by default; " "pass --cloudflare to enable it (--secure implies it), --no-cloudflare to " "force it off. It does not change a raw wildcard bind. If the admin " - "password was never changed, Studio asks for a new one in the terminal " + "password was never changed, Unsloth asks for a new one in the terminal " "before publishing the URL.", ) parser.add_argument( @@ -1767,7 +1767,7 @@ def _build_arg_parser(): help = "Expose ONLY a Cloudflare HTTPS link: bind localhost and fail closed " "if the tunnel can't start. Without it, --no-secure also serves the raw " "0.0.0.0 port, which is reachable from anywhere on the network. If the " - "admin password was never changed, Studio asks for a new one in the " + "admin password was never changed, Unsloth asks for a new one in the " "terminal before publishing the URL.", ) # Back-compat: accept --not-secure as a hidden alias for --no-secure. diff --git a/studio/backend/startup_banner.py b/studio/backend/startup_banner.py index ea951a4325..9ec7a4f91c 100644 --- a/studio/backend/startup_banner.py +++ b/studio/backend/startup_banner.py @@ -1,7 +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 -"""Terminal banner for Studio startup. +"""Terminal banner for Unsloth startup. Stdlib only -- safe to import without the rest of the backend. """ @@ -172,7 +172,7 @@ def print_studio_access_banner( secondary, ), style( - " Only on trusted networks -- anyone who reaches this machine can use Studio.", + " Only on trusted networks -- anyone who reaches this machine can use Unsloth.", secondary, ), ] diff --git a/studio/backend/storage/studio_db.py b/studio/backend/storage/studio_db.py index 9de62cb9c4..43332b0986 100644 --- a/studio/backend/storage/studio_db.py +++ b/studio/backend/storage/studio_db.py @@ -7,6 +7,7 @@ Like auth/storage.py (module-level functions, raw sqlite3, per-function connections) plus WAL mode and PRAGMA foreign_keys = ON for CASCADE deletes. """ +import hashlib import json import logging import os @@ -100,6 +101,7 @@ _schema_lock = threading.Lock() _schema_ready = False _SQLITE_IN_CHUNK_SIZE = 900 _PROJECT_WORKSPACE_SUBDIRS = ("sandbox",) +_CHAT_ATTACHMENT_INVENTORY_VERSION = 1 def _project_slug(name: str) -> str: @@ -190,13 +192,18 @@ def _ensure_schema(conn: sqlite3.Connection) -> None: error_message TEXT, duration_seconds REAL, loss_sparkline TEXT, - display_name TEXT + display_name TEXT, + resume_blocked INTEGER NOT NULL DEFAULT 0 ) """ ) existing_cols = {row[1] for row in conn.execute("PRAGMA table_info(training_runs)").fetchall()} if "display_name" not in existing_cols: conn.execute("ALTER TABLE training_runs ADD COLUMN display_name TEXT") + if "resume_blocked" not in existing_cols: + conn.execute( + "ALTER TABLE training_runs ADD COLUMN resume_blocked INTEGER NOT NULL DEFAULT 0" + ) conn.execute( """ CREATE TABLE IF NOT EXISTS training_metrics ( @@ -313,6 +320,141 @@ def _ensure_schema(conn: sqlite3.Connection) -> None: ) """ ) + tombstone_schema = """ + CREATE TABLE chat_attachment_tombstones ( + thread_id TEXT NOT NULL REFERENCES chat_threads(id) ON DELETE CASCADE, + message_id TEXT NOT NULL, + attachment_id TEXT NOT NULL, + deleted_at INTEGER NOT NULL, + PRIMARY KEY(thread_id, message_id, attachment_id) + ) WITHOUT ROWID + """ + tombstone_table = conn.execute( + """ + SELECT 1 FROM sqlite_master + WHERE type = 'table' AND name = 'chat_attachment_tombstones' + """ + ).fetchone() + if tombstone_table is None: + conn.execute(tombstone_schema) + else: + tombstone_columns = { + row[1] for row in conn.execute("PRAGMA table_info(chat_attachment_tombstones)") + } + tombstone_fk_targets = { + row[2] for row in conn.execute("PRAGMA foreign_key_list(chat_attachment_tombstones)") + } + if "thread_id" not in tombstone_columns or "chat_threads" not in tombstone_fk_targets: + # The first implementation cascaded through chat_messages, which + # erased deletion knowledge during pruneMissing. Rebuild once, + # retaining every tombstone whose owning thread still exists. + conn.execute("SAVEPOINT migrate_chat_attachment_tombstones") + try: + conn.execute( + "ALTER TABLE chat_attachment_tombstones " + "RENAME TO chat_attachment_tombstones_legacy" + ) + conn.execute(tombstone_schema) + if "thread_id" in tombstone_columns: + conn.execute( + """ + INSERT OR IGNORE INTO chat_attachment_tombstones + (thread_id, message_id, attachment_id, deleted_at) + SELECT legacy.thread_id, legacy.message_id, + legacy.attachment_id, legacy.deleted_at + FROM chat_attachment_tombstones_legacy legacy + JOIN chat_threads thread ON thread.id = legacy.thread_id + """ + ) + else: + conn.execute( + """ + INSERT OR IGNORE INTO chat_attachment_tombstones + (thread_id, message_id, attachment_id, deleted_at) + SELECT message.thread_id, legacy.message_id, + legacy.attachment_id, legacy.deleted_at + FROM chat_attachment_tombstones_legacy legacy + JOIN chat_messages message ON message.id = legacy.message_id + """ + ) + conn.execute("DROP TABLE chat_attachment_tombstones_legacy") + conn.execute("RELEASE SAVEPOINT migrate_chat_attachment_tombstones") + except Exception: + conn.execute("ROLLBACK TO SAVEPOINT migrate_chat_attachment_tombstones") + conn.execute("RELEASE SAVEPOINT migrate_chat_attachment_tombstones") + raise + conn.execute( + """ + CREATE TABLE IF NOT EXISTS chat_attachment_inventory ( + message_id TEXT NOT NULL REFERENCES chat_messages(id) ON DELETE CASCADE, + attachment_id TEXT NOT NULL, + name TEXT NOT NULL, + type TEXT, + content_type TEXT, + size_bytes INTEGER, + PRIMARY KEY(message_id, attachment_id) + ) WITHOUT ROWID + """ + ) + conn.execute( + """ + CREATE TABLE IF NOT EXISTS chat_attachment_inventory_state ( + singleton INTEGER NOT NULL PRIMARY KEY CHECK(singleton = 1), + inventory_version INTEGER NOT NULL DEFAULT 0, + dirty INTEGER NOT NULL DEFAULT 1, + backfilled_at INTEGER NOT NULL + ) + """ + ) + inventory_state_columns = { + row[1] for row in conn.execute("PRAGMA table_info(chat_attachment_inventory_state)") + } + if "inventory_version" not in inventory_state_columns: + conn.execute( + "ALTER TABLE chat_attachment_inventory_state " + "ADD COLUMN inventory_version INTEGER NOT NULL DEFAULT 0" + ) + if "dirty" not in inventory_state_columns: + conn.execute( + "ALTER TABLE chat_attachment_inventory_state " + "ADD COLUMN dirty INTEGER NOT NULL DEFAULT 1" + ) + conn.execute( + """ + CREATE TRIGGER IF NOT EXISTS chat_attachment_inventory_dirty_insert + AFTER INSERT ON chat_messages + BEGIN + INSERT INTO chat_attachment_inventory_state + (singleton, inventory_version, dirty, backfilled_at) + VALUES (1, 0, 1, 0) + ON CONFLICT(singleton) DO UPDATE SET dirty = 1; + END + """ + ) + conn.execute( + """ + CREATE TRIGGER IF NOT EXISTS chat_attachment_inventory_dirty_update + AFTER UPDATE ON chat_messages + BEGIN + INSERT INTO chat_attachment_inventory_state + (singleton, inventory_version, dirty, backfilled_at) + VALUES (1, 0, 1, 0) + ON CONFLICT(singleton) DO UPDATE SET dirty = 1; + END + """ + ) + conn.execute( + """ + CREATE TRIGGER IF NOT EXISTS chat_attachment_inventory_dirty_delete + AFTER DELETE ON chat_messages + BEGIN + INSERT INTO chat_attachment_inventory_state + (singleton, inventory_version, dirty, backfilled_at) + VALUES (1, 0, 1, 0) + ON CONFLICT(singleton) DO UPDATE SET dirty = 1; + END + """ + ) conn.execute( "CREATE INDEX IF NOT EXISTS idx_chat_threads_model_type_created_at ON chat_threads(model_type, created_at)" ) @@ -567,6 +709,22 @@ def _ensure_schema(conn: sqlite3.Connection) -> None: "CREATE INDEX IF NOT EXISTS idx_research_document_sources_run " "ON research_document_sources(run_id, id)" ) + inventory_state = conn.execute( + """ + SELECT inventory_version, dirty + FROM chat_attachment_inventory_state + WHERE singleton = 1 + """ + ).fetchone() + # Positional read: works for raw tuple or sqlite3.Row (no row_factory precondition). + if ( + inventory_state is None + or inventory_state[0] != _CHAT_ATTACHMENT_INVENTORY_VERSION + or inventory_state[1] + ): + _rebuild_chat_attachment_inventory(conn) + _mark_chat_attachment_inventory_clean(conn) + conn.commit() def _prompt_entry_from_row(row: sqlite3.Row) -> dict: @@ -759,16 +917,43 @@ def create_run( config_json: str, started_at: str, total_steps: Optional[int], + *, + output_dir: Optional[str] = None, + cancel_requested: bool = False, + resumed_from_run_id: Optional[str] = None, ) -> None: conn = get_connection() try: conn.execute( """ - INSERT INTO training_runs (id, model_name, dataset_name, config_json, started_at, total_steps) - VALUES (?, ?, ?, ?, ?, ?) + INSERT INTO training_runs ( + id, model_name, dataset_name, config_json, started_at, total_steps, + output_dir, resume_blocked + ) + VALUES (?, ?, ?, ?, ?, ?, ?, ?) """, - (id, model_name, dataset_name, config_json, started_at, total_steps), + ( + id, + model_name, + dataset_name, + config_json, + started_at, + total_steps, + None if cancel_requested else output_dir, + int(cancel_requested), + ), ) + if resumed_from_run_id: + claimed = conn.execute( + """ + UPDATE training_runs SET resume_blocked = 1 + WHERE id = ? AND status IN ('stopped', 'error') + AND output_dir = ? AND resume_blocked = 0 + """, + (resumed_from_run_id, output_dir), + ) + if claimed.rowcount != 1: + raise RuntimeError("Resume source is no longer available") conn.commit() finally: conn.close() @@ -811,6 +996,8 @@ def finish_run( loss_sparkline: Optional[str] = None, output_dir: Optional[str] = None, error_message: Optional[str] = None, + clear_output_dir: bool = False, + resume_blocked: bool = False, ) -> None: conn = get_connection() try: @@ -818,9 +1005,16 @@ def finish_run( """ UPDATE training_runs SET status = ?, ended_at = ?, final_step = ?, final_loss = ?, - duration_seconds = ?, loss_sparkline = ?, output_dir = ?, - error_message = ? - WHERE id = ? + duration_seconds = ?, loss_sparkline = ?, + output_dir = CASE + WHEN resume_blocked = 1 OR ? = 1 THEN NULL + WHEN ? IS NOT NULL THEN ? + WHEN ? IN ('error', 'stopped') THEN output_dir + ELSE NULL + END, + error_message = ?, + resume_blocked = CASE WHEN resume_blocked = 1 OR ? = 1 THEN 1 ELSE ? END + WHERE id = ? AND status = 'running' """, ( status, @@ -829,8 +1023,13 @@ def finish_run( final_loss, duration_seconds, loss_sparkline, + int(clear_output_dir), output_dir, + output_dir, + status, error_message, + int(clear_output_dir), + int(resume_blocked), id, ), ) @@ -890,6 +1089,38 @@ def update_run_display_name(id: str, display_name: Optional[str]) -> None: conn.close() +def update_run_output_dir(id: str, output_dir: Optional[str]) -> None: + conn = get_connection() + try: + conn.execute( + """ + UPDATE training_runs SET output_dir = ? + WHERE id = ? AND status = 'running' AND resume_blocked = 0 + """, + (output_dir, id), + ) + conn.commit() + finally: + conn.close() + + +def mark_run_cancel_requested(id: str) -> bool: + """Clear resume/export state only while the exact run is still active.""" + conn = get_connection() + try: + cursor = conn.execute( + """ + UPDATE training_runs SET output_dir = NULL, resume_blocked = 1 + WHERE id = ? AND status = 'running' + """, + (id,), + ) + conn.commit() + return cursor.rowcount > 0 + finally: + conn.close() + + def list_runs(limit: int = 50, offset: int = 0) -> dict: conn = get_connection() try: @@ -899,15 +1130,15 @@ def list_runs(limit: int = 50, offset: int = 0) -> dict: SELECT r.id, r.status, r.model_name, r.dataset_name, r.started_at, r.ended_at, r.total_steps, r.final_step, r.final_loss, r.output_dir, r.duration_seconds, r.error_message, - r.loss_sparkline, r.display_name, r.config_json, + r.loss_sparkline, r.display_name, r.config_json, r.resume_blocked, CASE - WHEN r.status = 'stopped' + WHEN r.status IN ('stopped', 'error') AND r.output_dir IS NOT NULL AND EXISTS ( SELECT 1 FROM training_runs newer WHERE newer.output_dir = r.output_dir - AND newer.status IN ('stopped', 'completed') + AND newer.status IN ('stopped', 'completed', 'error', 'running') AND newer.started_at > r.started_at ) THEN 1 ELSE 0 @@ -942,13 +1173,13 @@ def get_run(id: str) -> Optional[dict]: """ SELECT r.*, CASE - WHEN r.status = 'stopped' + WHEN r.status IN ('stopped', 'error') AND r.output_dir IS NOT NULL AND EXISTS ( SELECT 1 FROM training_runs newer WHERE newer.output_dir = r.output_dir - AND newer.status IN ('stopped', 'completed') + AND newer.status IN ('stopped', 'completed', 'error', 'running') AND newer.started_at > r.started_at ) THEN 1 ELSE 0 @@ -983,12 +1214,12 @@ def get_resumable_run_by_output_dir(output_dir: str) -> Optional[dict]: 0 AS resumed_later FROM training_runs r WHERE r.output_dir = ? - AND r.status = 'stopped' + AND r.status IN ('stopped', 'error') AND NOT EXISTS ( SELECT 1 FROM training_runs newer WHERE newer.output_dir = r.output_dir - AND newer.status IN ('stopped', 'completed') + AND newer.status IN ('stopped', 'completed', 'error', 'running') AND newer.started_at > r.started_at ) ORDER BY r.started_at DESC @@ -1091,8 +1322,12 @@ def cleanup_orphaned_runs() -> None: conn.execute( """ UPDATE training_runs - SET status = 'error', - error_message = 'Server restarted during training', + SET status = CASE WHEN resume_blocked = 1 THEN 'stopped' ELSE 'error' END, + error_message = CASE + WHEN resume_blocked = 1 THEN NULL + ELSE 'Server restarted during training' + END, + output_dir = CASE WHEN resume_blocked = 1 THEN NULL ELSE output_dir END, ended_at = ? WHERE status = 'running' """, @@ -1396,7 +1631,14 @@ def delete_chat_threads(ids: list[str]) -> None: return conn = get_connection() try: + conn.execute("BEGIN IMMEDIATE") + _ensure_chat_attachment_inventory_current(conn) + conn.executemany( + "DELETE FROM chat_attachment_tombstones WHERE thread_id = ?", + [(id,) for id in ids], + ) conn.executemany("DELETE FROM chat_threads WHERE id = ?", [(id,) for id in ids]) + _mark_chat_attachment_inventory_clean(conn) conn.commit() finally: conn.close() @@ -1405,7 +1647,11 @@ def delete_chat_threads(ids: list[str]) -> None: def clear_chat_history() -> None: conn = get_connection() try: + conn.execute("BEGIN IMMEDIATE") + _ensure_chat_attachment_inventory_current(conn) + conn.execute("DELETE FROM chat_attachment_tombstones") conn.execute("DELETE FROM chat_threads") + _mark_chat_attachment_inventory_clean(conn) conn.commit() finally: conn.close() @@ -1531,6 +1777,7 @@ def delete_chat_project(id: str, delete_files: bool = False) -> Optional[dict]: conn = get_connection() try: conn.execute("BEGIN IMMEDIATE") + _ensure_chat_attachment_inventory_current(conn) row = conn.execute("SELECT * FROM chat_projects WHERE id = ?", (id,)).fetchone() if row is None: conn.rollback() @@ -1538,6 +1785,7 @@ def delete_chat_project(id: str, delete_files: bool = False) -> Optional[dict]: project = _chat_project_from_row(row) conn.execute("DELETE FROM chat_threads WHERE project_id = ?", (id,)) conn.execute("DELETE FROM chat_projects WHERE id = ?", (id,)) + _mark_chat_attachment_inventory_clean(conn) conn.commit() if delete_files: _delete_project_workspace(project) @@ -1712,10 +1960,265 @@ def _guard_research_messages( ) +_CONTENT_PART_ID_PREFIX = "content-part-sha256-" +_URI_SCHEME_RE = re.compile(r"^[A-Za-z][A-Za-z0-9+.-]*:") + + +def _is_locally_stored_blob(value: str) -> bool: + """True for data URIs or bare base64, never external/blob URI references.""" + candidate = value.lstrip() + if not candidate: + return False + if candidate[:5].lower() == "data:": + return True + if candidate.startswith(("//", "\\\\")): + return False + return _URI_SCHEME_RE.match(candidate) is None + + +def _managed_content_part_payload(part: dict) -> Optional[tuple[str, Any]]: + """Return the locally stored blob payload used to identify a content part.""" + image = part.get("image") + if isinstance(image, str) and image[:5].lower() == "data:": + return "image", image + + audio = part.get("audio") + if isinstance(audio, str) and _is_locally_stored_blob(audio): + return "audio", audio + if isinstance(audio, dict): + data = audio.get("data") + if isinstance(data, str) and _is_locally_stored_blob(data): + return "audio", audio + return None + + +def _content_part_id(part: dict) -> Optional[str]: + """Stable managed id derived from blob data, without mutating inference content.""" + payload = _managed_content_part_payload(part) + if payload is None: + return None + canonical = json.dumps( + payload, + ensure_ascii = False, + separators = (",", ":"), + sort_keys = True, + ).encode("utf-8") + return f"{_CONTENT_PART_ID_PREFIX}{hashlib.sha256(canonical).hexdigest()}" + + +def _chat_attachment_tombstones_for_messages( + conn: sqlite3.Connection, thread_id: str, message_ids: list[str] +) -> dict[str, set[str]]: + tombstones = {message_id: set() for message_id in message_ids} + unique_ids = list(dict.fromkeys(message_ids)) + for start in range(0, len(unique_ids), _SQLITE_IN_CHUNK_SIZE): + chunk = unique_ids[start : start + _SQLITE_IN_CHUNK_SIZE] + placeholders = ",".join("?" for _ in chunk) + rows = conn.execute( + f""" + SELECT message_id, attachment_id + FROM chat_attachment_tombstones + WHERE thread_id = ? AND message_id IN ({placeholders}) + """, + (thread_id, *chunk), + ).fetchall() + for row in rows: + tombstones[row["message_id"]].add(row["attachment_id"]) + return tombstones + + +def _reconcile_chat_message_uploads(message: dict, tombstones: set[str]) -> dict: + """Strip uploads previously deleted through the Data tab from a stale write.""" + if not tombstones: + return message + + reconciled = dict(message) + attachments = message.get("attachments") + if isinstance(attachments, list): + reconciled["attachments"] = [ + attachment + for attachment in attachments + if not (isinstance(attachment, dict) and str(attachment.get("id") or "") in tombstones) + ] + + content = message.get("content") + if isinstance(content, list): + reconciled["content"] = [ + part + for part in content + if not (isinstance(part, dict) and (_content_part_id(part) or "") in tombstones) + ] + return reconciled + + +def _chat_attachment_metadata_text(value, fallback: Optional[str] = None) -> Optional[str]: + """Keep untyped legacy/import metadata safe for SQLite binding.""" + if value is None: + return fallback + if isinstance(value, str): + return value or fallback + if isinstance(value, (bool, int, float)): + return str(value) + # Objects and arrays are not useful display metadata and sqlite3 rejects + # binding them directly. + return fallback + + +def _chat_attachment_inventory_entries( + attachments_json: Optional[str], + content_json: Optional[str], + tombstones: Optional[set[str]] = None, +) -> list[dict]: + tombstones = tombstones or set() + attachments = _json_loads(attachments_json, None) + if not isinstance(attachments, list): + attachments = [] + attachments = [ + attachment + for attachment in attachments + if isinstance(attachment, dict) and attachment.get("id") + ] + attachments.extend(_content_part_attachments(content_json)) + + entries: list[dict] = [] + seen: set[str] = set() + for attachment in attachments: + attachment_id = str(attachment["id"]) + if attachment_id in seen or attachment_id in tombstones: + continue + seen.add(attachment_id) + entries.append( + { + "id": attachment_id, + "name": _chat_attachment_metadata_text(attachment.get("name"), "attachment"), + "type": _chat_attachment_metadata_text(attachment.get("type")), + "contentType": _chat_attachment_metadata_text(attachment.get("contentType")), + "sizeBytes": _chat_attachment_size_bytes(attachment), + } + ) + return entries + + +def _replace_chat_attachment_inventory( + conn: sqlite3.Connection, + message_id: str, + attachments_json: Optional[str], + content_json: Optional[str], + tombstones: Optional[set[str]] = None, +) -> None: + conn.execute("DELETE FROM chat_attachment_inventory WHERE message_id = ?", (message_id,)) + entries = _chat_attachment_inventory_entries( + attachments_json, + content_json, + tombstones, + ) + conn.executemany( + """ + INSERT INTO chat_attachment_inventory + (message_id, attachment_id, name, type, content_type, size_bytes) + VALUES (?, ?, ?, ?, ?, ?) + """, + [ + ( + message_id, + entry["id"], + entry["name"], + entry["type"], + entry["contentType"], + entry["sizeBytes"], + ) + for entry in entries + ], + ) + + +def _mark_chat_attachment_inventory_clean(conn: sqlite3.Connection) -> None: + conn.execute( + """ + INSERT INTO chat_attachment_inventory_state + (singleton, inventory_version, dirty, backfilled_at) + VALUES (1, ?, 0, ?) + ON CONFLICT(singleton) DO UPDATE SET + inventory_version = excluded.inventory_version, + dirty = 0, + backfilled_at = excluded.backfilled_at + """, + ( + _CHAT_ATTACHMENT_INVENTORY_VERSION, + int(datetime.now(timezone.utc).timestamp() * 1000), + ), + ) + + +def _rebuild_chat_attachment_inventory(conn: sqlite3.Connection) -> None: + """Rebuild after schema upgrade or a write from an older Studio build.""" + conn.execute("DELETE FROM chat_attachment_inventory") + tombstones: dict[tuple[str, str], set[str]] = {} + for row in conn.execute( + "SELECT thread_id, message_id, attachment_id FROM chat_attachment_tombstones" + ).fetchall(): + tombstones.setdefault((row["thread_id"], row["message_id"]), set()).add( + row["attachment_id"] + ) + rows = conn.execute( + "SELECT id, thread_id, attachments_json, content_json FROM chat_messages" + ).fetchall() + for row in rows: + _replace_chat_attachment_inventory( + conn, + row["id"], + row["attachments_json"], + row["content_json"], + tombstones.get((row["thread_id"], row["id"]), set()), + ) + + +def _ensure_chat_attachment_inventory_current(conn: sqlite3.Connection) -> None: + state = conn.execute( + """ + SELECT inventory_version, dirty + FROM chat_attachment_inventory_state + WHERE singleton = 1 + """ + ).fetchone() + if ( + state is not None + and state["inventory_version"] == _CHAT_ATTACHMENT_INVENTORY_VERSION + and not state["dirty"] + ): + return + + owns_transaction = not conn.in_transaction + if owns_transaction: + conn.execute("BEGIN IMMEDIATE") + try: + state = conn.execute( + """ + SELECT inventory_version, dirty + FROM chat_attachment_inventory_state + WHERE singleton = 1 + """ + ).fetchone() + if ( + state is None + or state["inventory_version"] != _CHAT_ATTACHMENT_INVENTORY_VERSION + or state["dirty"] + ): + _rebuild_chat_attachment_inventory(conn) + _mark_chat_attachment_inventory_clean(conn) + if owns_transaction: + conn.commit() + except Exception: + if owns_transaction: + conn.rollback() + raise + + def upsert_chat_message(message: dict, *, allow_research_update: bool = False) -> dict: conn = get_connection() try: conn.execute("BEGIN IMMEDIATE") + _ensure_chat_attachment_inventory_current(conn) if not allow_research_update: _guard_research_messages(conn, message["threadId"], [message]) _raise_if_chat_message_thread_conflicts( @@ -1723,6 +2226,21 @@ def upsert_chat_message(message: dict, *, allow_research_update: bool = False) - message["threadId"], [message["id"]], ) + tombstones = _chat_attachment_tombstones_for_messages( + conn, + message["threadId"], + [message["id"]], + ) + reconciled = _reconcile_chat_message_uploads( + message, + tombstones.get(message["id"], set()), + ) + content_json = json.dumps(reconciled.get("content", [])) + attachments_json = ( + json.dumps(reconciled.get("attachments")) + if reconciled.get("attachments") is not None + else None + ) conn.execute( """ INSERT INTO chat_messages @@ -1738,23 +2256,32 @@ def upsert_chat_message(message: dict, *, allow_research_update: bool = False) - WHERE excluded.thread_id = chat_messages.thread_id """, ( - message["id"], - message["threadId"], - message.get("parentId"), - message["role"], - json.dumps(message.get("content", [])), - json.dumps(message.get("attachments")) - if message.get("attachments") is not None + reconciled["id"], + reconciled["threadId"], + reconciled.get("parentId"), + reconciled["role"], + content_json, + attachments_json, + json.dumps(reconciled.get("metadata")) + if reconciled.get("metadata") is not None else None, - json.dumps(message.get("metadata")) - if message.get("metadata") is not None - else None, - int(message["createdAt"]), + int(reconciled["createdAt"]), ), ) - _bump_chat_thread_updated_at(conn, message["threadId"], int(message["createdAt"])) + _replace_chat_attachment_inventory( + conn, + reconciled["id"], + attachments_json, + content_json, + ) + _bump_chat_thread_updated_at( + conn, + reconciled["threadId"], + int(reconciled["createdAt"]), + ) + _mark_chat_attachment_inventory_clean(conn) conn.commit() - return message + return reconciled except Exception: conn.rollback() raise @@ -1772,6 +2299,7 @@ def sync_chat_messages( conn = get_connection() try: conn.execute("BEGIN IMMEDIATE") + _ensure_chat_attachment_inventory_current(conn) if not allow_research_update: _guard_research_messages(conn, thread_id, messages) _raise_if_chat_message_thread_conflicts( @@ -1779,6 +2307,22 @@ def sync_chat_messages( thread_id, [m["id"] for m in messages], ) + tombstones = _chat_attachment_tombstones_for_messages( + conn, + thread_id, + [m["id"] for m in messages], + ) + reconciled_messages = [ + _reconcile_chat_message_uploads(m, tombstones.get(m["id"], set())) for m in messages + ] + serialized_messages = [ + ( + m, + json.dumps(m.get("content", [])), + json.dumps(m.get("attachments")) if m.get("attachments") is not None else None, + ) + for m in reconciled_messages + ] conn.executemany( """ INSERT INTO chat_messages @@ -1799,36 +2343,50 @@ def sync_chat_messages( thread_id, m.get("parentId"), m["role"], - json.dumps(m.get("content", [])), - json.dumps(m.get("attachments")) if m.get("attachments") is not None else None, + content_json, + attachments_json, json.dumps(m.get("metadata")) if m.get("metadata") is not None else None, int(m["createdAt"]), ) - for m in messages + for m, content_json, attachments_json in serialized_messages ], ) + for m, content_json, attachments_json in serialized_messages: + _replace_chat_attachment_inventory( + conn, + m["id"], + attachments_json, + content_json, + ) if prune_missing: - survivor_ids = {str(message["id"]) for message in messages} + retained_ids = {m["id"] for m in reconciled_messages} existing_ids = { - str(row["id"]) + row["id"] for row in conn.execute( - "SELECT id FROM chat_messages WHERE thread_id = ?", (thread_id,) + "SELECT id FROM chat_messages WHERE thread_id = ?", + (thread_id,), ).fetchall() } - removed_ids = existing_ids - survivor_ids - if removed_ids & _research_message_ids(conn, thread_id): + missing_ids = sorted(existing_ids - retained_ids) + if set(missing_ids) & _research_message_ids(conn, thread_id): raise ChatMessageProtectedError( "Research prompts and responses cannot be deleted from their original thread" ) - conn.executemany( - "DELETE FROM chat_messages WHERE thread_id = ? AND id = ?", - [(thread_id, message_id) for message_id in removed_ids], - ) + for start in range(0, len(missing_ids), _SQLITE_IN_CHUNK_SIZE): + chunk = missing_ids[start : start + _SQLITE_IN_CHUNK_SIZE] + placeholders = ",".join("?" for _ in chunk) + conn.execute( + f"DELETE FROM chat_messages WHERE thread_id = ? AND id IN ({placeholders})", + (thread_id, *chunk), + ) _recompute_chat_thread_updated_at(conn, thread_id) - elif messages: + elif reconciled_messages: _bump_chat_thread_updated_at( - conn, thread_id, max(int(m["createdAt"]) for m in messages) + conn, + thread_id, + max(int(m["createdAt"]) for m in reconciled_messages), ) + _mark_chat_attachment_inventory_clean(conn) conn.commit() return list_chat_messages(thread_id) except (ChatMessageConflictError, ChatMessageProtectedError): @@ -1911,6 +2469,7 @@ def fork_chat_thread( conn = get_connection() try: conn.execute("BEGIN IMMEDIATE") + _ensure_chat_attachment_inventory_current(conn) src = conn.execute( "SELECT * FROM chat_threads WHERE id = ?", (source_thread_id,) ).fetchone() @@ -1989,6 +2548,14 @@ def fork_chat_thread( """, fork_messages, ) + for row in ancestry: + _replace_chat_attachment_inventory( + conn, + id_map[row["id"]], + row["attachments_json"], + row["content_json"], + ) + _mark_chat_attachment_inventory_clean(conn) conn.commit() thread_row = conn.execute( "SELECT * FROM chat_threads WHERE id = ?", (new_thread_id,) @@ -2047,6 +2614,279 @@ def get_chat_message(thread_id: str, message_id: str) -> Optional[dict]: conn.close() +def _blob_part_base64_len(part: dict) -> int: + """Base64 payload length of an image or audio content part, or 0.""" + image = part.get("image") + if isinstance(image, str) and image[:5].lower() == "data:": + return len(image.rsplit(",", 1)[-1]) + audio = part.get("audio") + if isinstance(audio, str) and _is_locally_stored_blob(audio): + return len(audio.rsplit(",", 1)[-1]) + if isinstance(audio, dict): + data = audio.get("data") + if isinstance(data, str) and _is_locally_stored_blob(data): + return len(data) + return 0 + + +def _chat_attachment_size_bytes(attachment: dict) -> Optional[int]: + """Approximate stored size of one attachment's content parts. + + Image and audio parts hold base64 payloads (decoded bytes ~= 3/4 of the + encoded length); text parts count their character length. None when there + is no sizable content (e.g. a stripped/legacy attachment). + """ + total = 0 + found = False + for part in attachment.get("content") or []: + if not isinstance(part, dict): + continue + blob_len = _blob_part_base64_len(part) + if blob_len > 0: + total += (blob_len * 3) // 4 + found = True + continue + text = part.get("text") + if isinstance(text, str) and text: + total += len(text.encode("utf-8", errors = "ignore")) + found = True + return total if found else None + + +def _content_part_attachments(content_json: Optional[str]) -> list[dict]: + """Managed local blobs stored in content_json, with stable payload ids. + + Exact duplicate blobs intentionally share one inventory id. Deleting that + id removes every identical copy, avoiding ambiguous index-based addressing. + """ + content = _json_loads(content_json, None) + if not isinstance(content, list): + return [] + out: list[dict] = [] + seen: set[str] = set() + for part in content: + if not isinstance(part, dict): + continue + attachment_id = _content_part_id(part) + payload = _managed_content_part_payload(part) + if attachment_id is None or payload is None or attachment_id in seen: + continue + seen.add(attachment_id) + kind, value = payload + content_type = None + if kind == "image" and isinstance(value, str): + content_type = value[5:].split(";", 1)[0].split(",", 1)[0] or None + out.append( + { + "id": attachment_id, + "type": kind, + "name": "Chat image" if kind == "image" else "Chat audio", + "contentType": content_type, + "content": [part], + } + ) + return out + + +def list_chat_attachments_page( + limit: int = 50, offset: int = 0 +) -> tuple[list[dict], Optional[int]]: + """One bounded page from the normalized attachment inventory.""" + if not 1 <= limit <= 100: + raise ValueError("limit must be between 1 and 100") + if offset < 0: + raise ValueError("offset must be non-negative") + + conn = get_connection() + try: + _ensure_chat_attachment_inventory_current(conn) + rows = conn.execute( + """ + SELECT i.attachment_id, i.name, i.type, i.content_type, + i.size_bytes, m.id AS message_id, m.thread_id, + m.created_at, t.title AS thread_title, t.pair_id + FROM chat_attachment_inventory i + JOIN chat_messages m ON m.id = i.message_id + LEFT JOIN chat_threads t ON t.id = m.thread_id + ORDER BY m.created_at DESC, m.id ASC, i.attachment_id ASC + LIMIT ? OFFSET ? + """, + (limit + 1, offset), + ).fetchall() + finally: + conn.close() + + has_more = len(rows) > limit + page_rows = rows[:limit] + attachments = [ + { + "id": row["attachment_id"], + "messageId": row["message_id"], + "threadId": row["thread_id"], + "pairId": row["pair_id"], + "threadTitle": row["thread_title"], + "name": row["name"], + "type": row["type"], + "contentType": row["content_type"], + "sizeBytes": row["size_bytes"], + "createdAt": row["created_at"], + } + for row in page_rows + ] + return attachments, offset + limit if has_more else None + + +def list_chat_attachments() -> list[dict]: + """Compatibility helper returning the full normalized inventory.""" + attachments: list[dict] = [] + offset = 0 + while True: + page, next_offset = list_chat_attachments_page(limit = 100, offset = offset) + attachments.extend(page) + if next_offset is None: + return attachments + offset = next_offset + + +def get_chat_attachment(message_id: str, attachment_id: str) -> Optional[dict]: + """One attachment record (full content) from a message, or None.""" + conn = get_connection() + try: + row = conn.execute( + """ + SELECT message.attachments_json, message.content_json, + EXISTS( + SELECT 1 FROM chat_attachment_tombstones tombstone + WHERE tombstone.thread_id = message.thread_id + AND tombstone.message_id = message.id + AND tombstone.attachment_id = ? + ) AS tombstoned + FROM chat_messages message + WHERE message.id = ? + """, + (attachment_id, message_id), + ).fetchone() + finally: + conn.close() + if row is None or row["tombstoned"]: + return None + attachments = _json_loads(row["attachments_json"], None) + if isinstance(attachments, list): + for attachment in attachments: + if isinstance(attachment, dict) and str(attachment.get("id") or "") == attachment_id: + return attachment + if attachment_id.startswith(_CONTENT_PART_ID_PREFIX): + for attachment in _content_part_attachments(row["content_json"]): + if attachment["id"] == attachment_id: + return attachment + return None + + +def _record_chat_attachment_tombstone( + conn: sqlite3.Connection, thread_id: str, message_id: str, attachment_id: str +) -> None: + conn.execute( + """ + INSERT INTO chat_attachment_tombstones + (thread_id, message_id, attachment_id, deleted_at) + VALUES (?, ?, ?, ?) + ON CONFLICT(thread_id, message_id, attachment_id) DO UPDATE SET + deleted_at = excluded.deleted_at + """, + ( + thread_id, + message_id, + attachment_id, + int(datetime.now(timezone.utc).timestamp() * 1000), + ), + ) + + +def delete_chat_attachment(message_id: str, attachment_id: str) -> bool: + """Remove one stored upload from a message. + + The tombstone is retained while the thread exists, so pruning and later + recreating the same message id cannot restore the deleted upload. If an + ordinary attachment id collides with a content-blob id, both are deleted as + one managed item. + """ + conn = get_connection() + try: + conn.execute("BEGIN IMMEDIATE") + _ensure_chat_attachment_inventory_current(conn) + row = conn.execute( + """ + SELECT thread_id, attachments_json, content_json + FROM chat_messages WHERE id = ? + """, + (message_id,), + ).fetchone() + if row is None: + conn.rollback() + return False + + attachments = _json_loads(row["attachments_json"], None) + updated_attachments_json = row["attachments_json"] + deleted_attachment = False + if isinstance(attachments, list): + remaining_attachments = [ + attachment + for attachment in attachments + if not ( + isinstance(attachment, dict) + and str(attachment.get("id") or "") == attachment_id + ) + ] + deleted_attachment = len(remaining_attachments) != len(attachments) + if deleted_attachment: + updated_attachments_json = json.dumps(remaining_attachments) + + content = _json_loads(row["content_json"], None) + updated_content_json = row["content_json"] + deleted_content = False + if attachment_id.startswith(_CONTENT_PART_ID_PREFIX) and isinstance(content, list): + remaining_content = [ + part + for part in content + if not (isinstance(part, dict) and _content_part_id(part) == attachment_id) + ] + deleted_content = len(remaining_content) != len(content) + if deleted_content: + updated_content_json = json.dumps(remaining_content) + + if not deleted_attachment and not deleted_content: + conn.rollback() + return False + conn.execute( + """ + UPDATE chat_messages + SET attachments_json = ?, content_json = ? + WHERE id = ? + """, + (updated_attachments_json, updated_content_json, message_id), + ) + _record_chat_attachment_tombstone( + conn, + row["thread_id"], + message_id, + attachment_id, + ) + _replace_chat_attachment_inventory( + conn, + message_id, + updated_attachments_json, + updated_content_json, + ) + _mark_chat_attachment_inventory_clean(conn) + conn.commit() + return True + except Exception: + conn.rollback() + raise + finally: + conn.close() + + def list_chat_messages_for_threads(thread_ids: list[str]) -> list[dict]: if not thread_ids: return [] diff --git a/studio/backend/tests/conftest.py b/studio/backend/tests/conftest.py index b0b9ee309c..c2216104a3 100644 --- a/studio/backend/tests/conftest.py +++ b/studio/backend/tests/conftest.py @@ -101,13 +101,13 @@ def studio_server(request): @pytest.fixture def base_url(studio_server): - """Base URL for the e2e Studio server (from ``studio_server``).""" + """Base URL for the e2e Unsloth server (from ``studio_server``).""" return studio_server[0] @pytest.fixture def api_key(studio_server): - """API key for the e2e Studio server (from ``studio_server``).""" + """API key for the e2e Unsloth server (from ``studio_server``).""" return studio_server[1] diff --git a/studio/backend/tests/test_amd_apu_unified_memory.py b/studio/backend/tests/test_amd_apu_unified_memory.py index 4df9e85b30..9fd8260bf2 100644 --- a/studio/backend/tests/test_amd_apu_unified_memory.py +++ b/studio/backend/tests/test_amd_apu_unified_memory.py @@ -91,7 +91,7 @@ class TestApuRamShortfall: """On a unified-memory APU the weights load into system RAM, so a model larger than available RAM (the field case: a 64.6 GB GGUF on a WSL VM capped well below the ROCm-reported APU budget) must be refused before spawning, - not left to OOM-kill the Studio process.""" + not left to OOM-kill the Unsloth process.""" def test_field_case_wsl_cap_refuses(self): # 64.6 GB weights, ~46 GB available (WSL VM): refuse with guidance. diff --git a/studio/backend/tests/test_anthropic_compaction.py b/studio/backend/tests/test_anthropic_compaction.py index 1528eebe8b..acc0acc2e0 100644 --- a/studio/backend/tests/test_anthropic_compaction.py +++ b/studio/backend/tests/test_anthropic_compaction.py @@ -4,7 +4,7 @@ """Unit tests for Anthropic server-side context compaction wiring. Compaction is a beta (header ``compact-2026-01-12``) gated to Opus 4.6/4.7, -Sonnet 4.6, and Mythos preview. When enabled, Studio attaches +Sonnet 4.6, and Mythos preview. When enabled, Unsloth attaches ``context_management.edits[{type:"compact_20260112", trigger:{type:"input_tokens", value:N}}]``; the 50k-token minimum is clamped up so the request doesn't 400. diff --git a/studio/backend/tests/test_anthropic_fast_mode_edge.py b/studio/backend/tests/test_anthropic_fast_mode_edge.py index dd69d77590..03f5d1c0eb 100644 --- a/studio/backend/tests/test_anthropic_fast_mode_edge.py +++ b/studio/backend/tests/test_anthropic_fast_mode_edge.py @@ -330,7 +330,7 @@ def test_refusal_chunk_is_proper_openai_delta_shape(monkeypatch): def test_refusal_tool_event_chunk_shape(monkeypatch): - """Drop signal rides a Studio `_toolEvent` envelope (delta={}, + """Drop signal rides an Unsloth `_toolEvent` envelope (delta={}, finish_reason=null); the frontend latches on `_toolEvent.type == "anthropic_refusal"`.""" _, lines = _capture(monkeypatch, sse = _refusal_sse(), model = "claude-opus-4-7") @@ -409,7 +409,7 @@ def _fast_speed_sse(model: str = "claude-opus-4-7", speed: str = "fast") -> byte def test_usage_speed_propagates_to_final_usage_chunk_fast(monkeypatch): - """``usage.speed == "fast"`` from upstream must reach the Studio usage chunk.""" + """``usage.speed == "fast"`` from upstream must reach the Unsloth usage chunk.""" _, lines = _capture(monkeypatch, sse = _fast_speed_sse(speed = "fast")) usage_lines = [l for l in lines if l.startswith("data: ") and '"usage"' in l] assert usage_lines, lines @@ -428,7 +428,7 @@ def test_usage_speed_propagates_to_final_usage_chunk_standard(monkeypatch): def test_usage_speed_absent_when_anthropic_does_not_report(monkeypatch): - """Studio must not invent ``usage.speed`` when upstream omits it.""" + """Unsloth must not invent ``usage.speed`` when upstream omits it.""" _, lines = _capture(monkeypatch) parsed = [ json.loads(l[len("data: ") :]) for l in lines if l.startswith("data: ") and '"usage"' in l diff --git a/studio/backend/tests/test_anthropic_messages.py b/studio/backend/tests/test_anthropic_messages.py index 3b0ea37372..9ccc3f44dd 100644 --- a/studio/backend/tests/test_anthropic_messages.py +++ b/studio/backend/tests/test_anthropic_messages.py @@ -1418,7 +1418,7 @@ class TestNormalizeAnthropicOpenAIImages: # ===================================================================== -# Studio-tool alias detection (/v1/messages tool routing) +# Unsloth-tool alias detection (/v1/messages tool routing) # ===================================================================== @@ -1436,7 +1436,7 @@ class TestAnthropicRequestedStudioTools: def test_client_tool_named_python_is_not_misclassified(self): # input_schema is the client-tool discriminator; its presence must - # prevent the name from being treated as a Studio alias. + # prevent the name from being treated as an Unsloth alias. tools = [ { "name": "python", @@ -1747,9 +1747,9 @@ class TestAnthropicMessagesToolRouting: assert "name" in exc.value.detail def test_alias_named_client_tool_without_schema_rejected_with_400(self, monkeypatch): - # Regression: a typo'd client tool whose name collides with a Studio + # Regression: a typo'd client tool whose name collides with an Unsloth # alias (e.g. a custom "python" tool missing input_schema) must - # surface a 400, not silently switch into Studio's built-in python + # surface a 400, not silently switch into Unsloth's built-in python # execution. _mock_backend(monkeypatch) payload = _basic_payload(tools = [{"name": "python"}]) @@ -1770,7 +1770,7 @@ class TestAnthropicMessagesToolRouting: def test_disable_tools_policy_overrides_server_tool_alias(self, monkeypatch): # CLI `unsloth run --disable-tools` sets policy=False. A request with - # a Studio server-tool alias must NOT enter the agentic loop then. + # an Unsloth server-tool alias must NOT enter the agentic loop then. backend = _mock_backend(monkeypatch) set_tool_policy(False) payload = _basic_payload( diff --git a/studio/backend/tests/test_cached_gguf_routes.py b/studio/backend/tests/test_cached_gguf_routes.py index d4a7cae208..b3e6255d55 100644 --- a/studio/backend/tests/test_cached_gguf_routes.py +++ b/studio/backend/tests/test_cached_gguf_routes.py @@ -120,12 +120,151 @@ def test_is_hidden_model_hides_validation_probe_everywhere(): assert models_route._is_hidden_model( None, "/hf/models--ggml-org--models/snapshots/abc/tinyllamas/stories260K.gguf" ) + # A Windows-style snapshot path must match too, even on a POSIX interpreter + # (the filename check splits on both separators). + assert models_route._is_hidden_model( + r"C:\Users\u\.cache\huggingface\hub\models--ggml-org--models\snapshots\abc\tinyllamas\stories260K.gguf" + ) assert not models_route._is_hidden_model("unsloth/gemma-3-270m-it-GGUF") # The exact-filename needle must not hide a real repo that merely # references stories260K in its name. assert not models_route._is_hidden_model("user/stories260K-finetune-GGUF") +def test_is_hidden_model_matches_repo_ids_exactly(monkeypatch): + """A custom embedder with a generic basename is hidden by EXACT repo-id + match only, so unrelated cached repos that merely contain the basename stay + visible. Regression: substring basename matching hid real chat models like + ``user/model-chat`` from the On Device inventory.""" + from core.rag import config as rag_config + + monkeypatch.setattr(rag_config, "effective_embedding_model", lambda: "org/model") + monkeypatch.setattr(rag_config, "effective_gguf_repo", lambda: "org/model-GGUF") + + # The exact embedder repo and its GGUF companion are hidden. + assert models_route._is_hidden_model("org/model") + assert models_route._is_hidden_model("org/model-GGUF") + # Unrelated repos that merely contain "model" must NOT be hidden. + assert not models_route._is_hidden_model("user/model-chat") + assert not models_route._is_hidden_model("org/model-instruct") + assert not models_route._is_hidden_model("acme/remodelled-chat") + # The validation probe stays hidden regardless of embedder config. + assert models_route._is_hidden_model("ggml-org/models") + + +def test_is_hidden_model_matches_repo_derived_local_paths(monkeypatch): + """Match exact repo-derived cache and LM Studio paths.""" + from core.rag import config as rag_config + + monkeypatch.setattr(rag_config, "effective_embedding_model", lambda: "org/model") + monkeypatch.setattr(rag_config, "effective_gguf_repo", lambda: "org/model-GGUF") + + assert models_route._is_hidden_model( + "/cache/models--org--model/snapshots/abc/model.safetensors" + ) + assert models_route._is_hidden_model( + r"C:\Users\u\.cache\huggingface\hub\models--org--model-GGUF\snapshots\abc" + ) + assert models_route._is_hidden_model("/lm-studio/org/model-GGUF/model-Q8_0.gguf") + assert not models_route._is_hidden_model("/lm-studio/user/model-chat/model-Q8_0.gguf") + assert not models_route._is_hidden_model("/cache/models--org--model-instruct") + + +def test_is_hidden_model_prefers_existing_relative_path(monkeypatch, tmp_path): + """Prefer an existing relative path over repo-id syntax.""" + from core.rag import config as rag_config + + embedder = tmp_path / "models" / "embedder" + embedder.mkdir(parents = True) + monkeypatch.chdir(tmp_path) + monkeypatch.setattr(rag_config, "effective_embedding_model", lambda: "models/embedder") + monkeypatch.setattr(rag_config, "effective_gguf_repo", lambda: "org/embedder-GGUF") + + assert models_route._is_hidden_model(str(embedder)) + + +def test_is_hidden_model_keeps_stale_default_embedder_hidden(monkeypatch): + """Keep default embedders hidden after a settings change.""" + from core.rag import config as rag_config + + monkeypatch.setattr(rag_config, "effective_embedding_model", lambda: "org/custom") + monkeypatch.setattr(rag_config, "effective_gguf_repo", lambda: "org/custom-GGUF") + + assert models_route._is_hidden_model("unsloth/bge-small-en-v1.5") + assert models_route._is_hidden_model("unsloth/bge-small-en-v1.5-GGUF") + assert models_route._is_hidden_model("/models/bge-small-en-v1.5") + assert models_route._is_hidden_model("/models/bge-small-en-v1.5-F16.gguf") + assert models_route._is_hidden_model(r"C:\models\bge-small-en-v1.5-Q8_0.gguf") + # Repo IDs still use exact matching, and similar local basenames must have + # a real separator after the static default name. + assert not models_route._is_hidden_model("user/bge-small-en-v1.5-chat") + assert not models_route._is_hidden_model("/models/bge-small-en-v1.50") + + +def test_is_hidden_model_keeps_env_default_hidden_after_override(monkeypatch): + """A persisted override must not expose the deployment's env default.""" + from core.rag import config as rag_config + + monkeypatch.delenv("RAG_EMBED_GGUF_REPO", raising = False) + monkeypatch.setattr(rag_config, "EMBEDDING_MODEL", "org/env-default") + monkeypatch.setattr(rag_config, "effective_embedding_model", lambda: "org/custom") + monkeypatch.setattr(rag_config, "effective_gguf_repo", lambda: "org/custom-GGUF") + + assert models_route._is_hidden_model("org/env-default") + assert models_route._is_hidden_model("org/env-default-GGUF") + assert models_route._is_hidden_model("org/custom") + assert models_route._is_hidden_model("org/custom-GGUF") + assert not models_route._is_hidden_model("org/env-default-chat") + + +def test_hidden_models_importable_without_heavy_model_stack(): + """The hub cache scanner imports ``is_hidden_model`` at module scope, so it + must not drag in ``utils/models/__init__`` (the model-config + checkpoint + stack). Verify in a clean interpreter that importing the helper touches + neither ``utils.models`` nor those heavy submodules, and still classifies + the probe.""" + import os + import subprocess + import textwrap + + backend = Path(__file__).resolve().parents[1] + code = textwrap.dedent( + """ + import sys + + class _Blocker: + _blocked = ( + "utils.models", + "utils.models.model_config", + "utils.models.checkpoints", + ) + + def find_spec(self, name, path=None, target=None): + if name in self._blocked: + raise ImportError("blocked heavy import: " + name) + return None + + sys.meta_path.insert(0, _Blocker()) + from utils.hidden_models import is_hidden_model + + loaded = sorted(m for m in sys.modules if m.startswith("utils.models")) + assert not loaded, loaded + assert is_hidden_model("ggml-org/models") is True + assert is_hidden_model("unsloth/gemma-3-270m-it-GGUF") is False + print("HIDDEN_MODELS_IMPORT_OK") + """ + ) + env = dict(os.environ, PYTHONPATH = str(backend)) + proc = subprocess.run( + [sys.executable, "-c", code], + capture_output = True, + text = True, + env = env, + ) + assert proc.returncode == 0, proc.stderr + assert "HIDDEN_MODELS_IMPORT_OK" in proc.stdout + + def test_list_cached_gguf_hides_llama_validation_probe(monkeypatch, tmp_path): """The ggml-org/models / stories260K install validation probe can land in the HF cache as a side effect of installing the prebuilt llama-server. diff --git a/studio/backend/tests/test_chat_attachments.py b/studio/backend/tests/test_chat_attachments.py new file mode 100644 index 0000000000..459587ca9e --- /dev/null +++ b/studio/backend/tests/test_chat_attachments.py @@ -0,0 +1,634 @@ +# SPDX-License-Identifier: AGPL-3.0-only +# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0 + +import base64 +import json +import os +import sqlite3 +import sys + +import pytest +from fastapi import HTTPException + +_backend = os.path.join(os.path.dirname(__file__), "..") +sys.path.insert(0, _backend) + +from routes import chat_history +from storage import studio_db +from utils.paths import studio_db_path + +PNG_BYTES = base64.b64decode( + "iVBORw0KGgoAAAANSUhEUgAAAAEAAAABCAYAAAAfFcSJAAAADUlEQVR42mP8z8BQDwAEhQGAhKmMIQAAAABJRU5ErkJggg==" +) +PNG_DATA_URL = "data:image/png;base64," + base64.b64encode(PNG_BYTES).decode("ascii") + + +def _reset_studio_db(tmp_path, monkeypatch): + monkeypatch.setenv("UNSLOTH_STUDIO_HOME", str(tmp_path)) + monkeypatch.setenv("UNSLOTH_STUDIO_PROJECTS_HOME", str(tmp_path / "Projects")) + monkeypatch.setattr(studio_db, "_schema_ready", False) + + +def _thread( + thread_id: str = "thread-1", + title: str = "Test Chat", + pair_id: str | None = None, +) -> dict: + return { + "id": thread_id, + "title": title, + "modelType": "base", + "modelId": "test-model", + "pairId": pair_id, + "archived": False, + "createdAt": 1_700_000_000_000, + } + + +def _message( + message_id: str, + created_at: int = 1_700_000_000_000, + attachments = None, + thread_id: str = "thread-1", +) -> dict: + message = { + "id": message_id, + "threadId": thread_id, + "parentId": None, + "role": "user", + "content": [{"type": "text", "text": "hello"}], + "createdAt": created_at, + } + if attachments is not None: + message["attachments"] = attachments + return message + + +def _image_attachment(attachment_id: str = "att-1", name: str = "photo.png") -> dict: + return { + "id": attachment_id, + "type": "image", + "name": name, + "contentType": "image/png", + "content": [{"type": "image", "image": PNG_DATA_URL}], + "status": {"type": "complete"}, + } + + +def _seed( + tmp_path, + monkeypatch, + attachments, + message_id: str = "msg-1", +): + _reset_studio_db(tmp_path, monkeypatch) + studio_db.upsert_chat_thread(_thread()) + studio_db.upsert_chat_message(_message(message_id, attachments = attachments)) + + +def _set_raw_attachments_json(message_id: str, raw: str) -> None: + conn = sqlite3.connect(studio_db_path()) + try: + conn.execute( + "UPDATE chat_messages SET attachments_json = ? WHERE id = ?", + (raw, message_id), + ) + conn.commit() + finally: + conn.close() + + +def _raw_attachments_json(message_id: str): + conn = sqlite3.connect(studio_db_path()) + try: + row = conn.execute( + "SELECT attachments_json FROM chat_messages WHERE id = ?", + (message_id,), + ).fetchone() + return row[0] if row is not None else None + finally: + conn.close() + + +# --------------------------------------------------------------------------- +# Storage: list_chat_attachments +# --------------------------------------------------------------------------- + + +def test_list_chat_attachments_empty_db(tmp_path, monkeypatch): + _reset_studio_db(tmp_path, monkeypatch) + assert studio_db.list_chat_attachments() == [] + + +def test_list_chat_attachments_round_trip(tmp_path, monkeypatch): + _seed(tmp_path, monkeypatch, [_image_attachment()]) + records = studio_db.list_chat_attachments() + assert len(records) == 1 + record = records[0] + assert record["id"] == "att-1" + assert record["messageId"] == "msg-1" + assert record["threadId"] == "thread-1" + assert record["threadTitle"] == "Test Chat" + assert record["name"] == "photo.png" + assert record["type"] == "image" + assert record["contentType"] == "image/png" + assert record["createdAt"] == 1_700_000_000_000 + # Base64 length estimate is within padding error of the decoded size. + assert abs(record["sizeBytes"] - len(PNG_BYTES)) <= 2 + + +def test_list_chat_attachments_counts_text_utf8(tmp_path, monkeypatch): + text = "héllo wörld é世界" + attachment = { + "id": "att-txt", + "type": "document", + "name": "notes.txt", + "content": [{"type": "text", "text": text}], + } + _seed(tmp_path, monkeypatch, [attachment]) + records = studio_db.list_chat_attachments() + assert records[0]["sizeBytes"] == len(text.encode("utf-8")) + + +def test_list_chat_attachments_no_content_size_is_none(tmp_path, monkeypatch): + attachment = {"id": "att-empty", "name": "ghost.bin", "content": []} + _seed(tmp_path, monkeypatch, [attachment]) + records = studio_db.list_chat_attachments() + assert records[0]["sizeBytes"] is None + assert records[0]["name"] == "ghost.bin" + + +def test_list_chat_attachments_defaults_missing_name(tmp_path, monkeypatch): + attachment = {"id": "att-noname", "content": []} + _seed(tmp_path, monkeypatch, [attachment]) + assert studio_db.list_chat_attachments()[0]["name"] == "attachment" + + +def test_list_chat_attachments_sanitizes_structured_metadata(tmp_path, monkeypatch): + attachment = { + "id": "att-weird", + "name": {"nested": "name"}, + "type": ["image"], + "contentType": {"mime": "image/png"}, + "content": [], + } + _seed(tmp_path, monkeypatch, [attachment]) + record = studio_db.list_chat_attachments()[0] + assert record["name"] == "attachment" + assert record["type"] is None + assert record["contentType"] is None + + +def test_list_chat_attachments_skips_malformed_rows(tmp_path, monkeypatch): + _reset_studio_db(tmp_path, monkeypatch) + studio_db.upsert_chat_thread(_thread()) + for i, raw in enumerate( + [ + "not json at all", + '{"id": "att-obj"}', + "null", + "[]", + '[{"noid": true}, "just a string", 42]', + '[{"id": ""}]', + ] + ): + message_id = f"msg-bad-{i}" + studio_db.upsert_chat_message(_message(message_id)) + _set_raw_attachments_json(message_id, raw) + studio_db.upsert_chat_message(_message("msg-good", attachments = [_image_attachment("att-ok")])) + records = studio_db.list_chat_attachments() + assert [r["id"] for r in records] == ["att-ok"] + + +def test_list_chat_attachments_orders_newest_first(tmp_path, monkeypatch): + _reset_studio_db(tmp_path, monkeypatch) + studio_db.upsert_chat_thread(_thread()) + studio_db.upsert_chat_message( + _message("msg-old", 1_700_000_000_000, [_image_attachment("att-old")]) + ) + studio_db.upsert_chat_message( + _message("msg-new", 1_700_000_100_000, [_image_attachment("att-new")]) + ) + assert [r["id"] for r in studio_db.list_chat_attachments()] == ["att-new", "att-old"] + + +def test_list_chat_attachments_survives_missing_thread_row(tmp_path, monkeypatch): + _reset_studio_db(tmp_path, monkeypatch) + studio_db.upsert_chat_thread(_thread()) + studio_db.upsert_chat_message(_message("msg-1", attachments = [_image_attachment()])) + conn = sqlite3.connect(studio_db_path()) + try: + conn.execute("DELETE FROM chat_threads WHERE id = 'thread-1'") + conn.commit() + finally: + conn.close() + records = studio_db.list_chat_attachments() + assert len(records) == 1 + assert records[0]["threadTitle"] is None + + +def test_list_chat_attachments_includes_compare_pair_id(tmp_path, monkeypatch): + _reset_studio_db(tmp_path, monkeypatch) + studio_db.upsert_chat_thread(_thread(pair_id = "pair-1")) + studio_db.upsert_chat_message(_message("msg-compare", attachments = [_image_attachment()])) + record = studio_db.list_chat_attachments()[0] + assert record["threadId"] == "thread-1" + assert record["pairId"] == "pair-1" + + +def test_list_chat_attachments_gone_after_thread_delete(tmp_path, monkeypatch): + _seed(tmp_path, monkeypatch, [_image_attachment()]) + studio_db.delete_chat_threads(["thread-1"]) + assert studio_db.list_chat_attachments() == [] + + +# --------------------------------------------------------------------------- +# Storage: get_chat_attachment / delete_chat_attachment +# --------------------------------------------------------------------------- + + +def test_get_chat_attachment_found_and_missing(tmp_path, monkeypatch): + _seed(tmp_path, monkeypatch, [_image_attachment()]) + attachment = studio_db.get_chat_attachment("msg-1", "att-1") + assert attachment is not None + assert attachment["content"][0]["image"] == PNG_DATA_URL + assert studio_db.get_chat_attachment("msg-1", "att-missing") is None + assert studio_db.get_chat_attachment("msg-missing", "att-1") is None + + +def test_delete_chat_attachment_keeps_others(tmp_path, monkeypatch): + _seed( + tmp_path, + monkeypatch, + [_image_attachment("att-1"), _image_attachment("att-2", "other.png")], + ) + assert studio_db.delete_chat_attachment("msg-1", "att-1") is True + assert studio_db.get_chat_attachment("msg-1", "att-1") is None + assert studio_db.get_chat_attachment("msg-1", "att-2") is not None + assert [r["id"] for r in studio_db.list_chat_attachments()] == ["att-2"] + + +def test_delete_last_chat_attachment_stores_empty_list(tmp_path, monkeypatch): + _seed(tmp_path, monkeypatch, [_image_attachment()]) + assert studio_db.delete_chat_attachment("msg-1", "att-1") is True + # '[]' rather than NULL: a NULL attachments field reads back as missing + # and triggers the legacy IndexedDB backfill, resurrecting the deleted + # attachment on the next chat load. + assert _raw_attachments_json("msg-1") == "[]" + assert studio_db.list_chat_attachments() == [] + # The message itself must survive with its content intact. + message = studio_db.get_chat_message("thread-1", "msg-1") + assert message is not None + assert message["content"] == [{"type": "text", "text": "hello"}] + assert message["attachments"] == [] + + +def test_delete_chat_attachment_missing_targets(tmp_path, monkeypatch): + _seed(tmp_path, monkeypatch, [_image_attachment()]) + assert studio_db.delete_chat_attachment("msg-missing", "att-1") is False + assert studio_db.delete_chat_attachment("msg-1", "att-missing") is False + _set_raw_attachments_json("msg-1", "not json") + assert studio_db.delete_chat_attachment("msg-1", "att-1") is False + + +# --------------------------------------------------------------------------- +# Routes: /attachments endpoints (real storage, direct calls) +# --------------------------------------------------------------------------- + + +def test_list_attachments_route(tmp_path, monkeypatch): + _seed(tmp_path, monkeypatch, [_image_attachment()]) + result = chat_history.list_attachments(current_subject = "unsloth") + assert [a["id"] for a in result["attachments"]] == ["att-1"] + + +def test_attachment_file_serves_image_bytes(tmp_path, monkeypatch): + _seed(tmp_path, monkeypatch, [_image_attachment()]) + response = chat_history.get_attachment_file("msg-1", "att-1", current_subject = "unsloth") + assert response.body == PNG_BYTES + assert response.media_type == "image/png" + + +def test_attachment_file_tolerates_whitespace_in_base64(tmp_path, monkeypatch): + encoded = base64.b64encode(PNG_BYTES).decode("ascii") + wrapped = "\n".join(encoded[i : i + 8] for i in range(0, len(encoded), 8)) + attachment = _image_attachment() + attachment["content"] = [{"type": "image", "image": "data:image/png;base64," + wrapped}] + _seed(tmp_path, monkeypatch, [attachment]) + response = chat_history.get_attachment_file("msg-1", "att-1", current_subject = "unsloth") + assert response.body == PNG_BYTES + + +def test_attachment_file_corrupt_base64_is_422(tmp_path, monkeypatch): + attachment = _image_attachment() + attachment["content"] = [{"type": "image", "image": "data:image/png;base64,%%%"}] + _seed(tmp_path, monkeypatch, [attachment]) + with pytest.raises(HTTPException) as excinfo: + chat_history.get_attachment_file("msg-1", "att-1", current_subject = "unsloth") + assert excinfo.value.status_code == 422 + + +def test_attachment_file_accepts_urlsafe_base64(tmp_path, monkeypatch): + data = bytes(range(251, 256)) * 3 # encodes to characters remapped by urlsafe + payload = base64.urlsafe_b64encode(data).decode("ascii") + assert "-" in payload or "_" in payload + attachment = _image_attachment() + attachment["content"] = [{"type": "image", "image": "data:image/png;base64," + payload}] + _seed(tmp_path, monkeypatch, [attachment]) + response = chat_history.get_attachment_file("msg-1", "att-1", current_subject = "unsloth") + assert response.body == data + + +def test_attachment_file_accepts_missing_padding(tmp_path, monkeypatch): + payload = base64.b64encode(PNG_BYTES).decode("ascii").rstrip("=") + attachment = _image_attachment() + attachment["content"] = [{"type": "image", "image": "data:image/png;base64," + payload}] + _seed(tmp_path, monkeypatch, [attachment]) + response = chat_history.get_attachment_file("msg-1", "att-1", current_subject = "unsloth") + assert response.body == PNG_BYTES + + +def test_attachment_file_serves_percent_encoded_data_url(tmp_path, monkeypatch): + attachment = _image_attachment() + attachment["content"] = [{"type": "image", "image": "data:text/plain,hello%20world"}] + _seed(tmp_path, monkeypatch, [attachment]) + response = chat_history.get_attachment_file("msg-1", "att-1", current_subject = "unsloth") + assert response.body == b"hello world" + # Non-image data URL types are clamped so markup never renders same-origin. + assert response.media_type == "application/octet-stream" + + +def test_attachment_file_serves_text_parts(tmp_path, monkeypatch): + attachment = { + "id": "att-txt", + "type": "document", + "name": "notes.txt", + "content": [ + {"type": "text", "text": "first"}, + {"type": "text", "text": "second"}, + ], + } + _seed(tmp_path, monkeypatch, [attachment]) + response = chat_history.get_attachment_file("msg-1", "att-txt", current_subject = "unsloth") + assert response.body.decode("utf-8") == "first\nsecond" + assert response.media_type.startswith("text/plain") + + +def test_attachment_file_no_content_is_404(tmp_path, monkeypatch): + _seed(tmp_path, monkeypatch, [{"id": "att-empty", "name": "ghost", "content": []}]) + with pytest.raises(HTTPException) as excinfo: + chat_history.get_attachment_file("msg-1", "att-empty", current_subject = "unsloth") + assert excinfo.value.status_code == 404 + + +def test_attachment_file_missing_message_is_404(tmp_path, monkeypatch): + _reset_studio_db(tmp_path, monkeypatch) + with pytest.raises(HTTPException) as excinfo: + chat_history.get_attachment_file("nope", "att-1", current_subject = "unsloth") + assert excinfo.value.status_code == 404 + + +def test_attachment_file_non_data_url_image_is_404(tmp_path, monkeypatch): + attachment = _image_attachment() + attachment["content"] = [{"type": "image", "image": "https://example.com/a.png"}] + _seed(tmp_path, monkeypatch, [attachment]) + with pytest.raises(HTTPException) as excinfo: + chat_history.get_attachment_file("msg-1", "att-1", current_subject = "unsloth") + assert excinfo.value.status_code == 404 + + +def test_attachment_file_defaults_media_type(tmp_path, monkeypatch): + payload = base64.b64encode(b"raw-bytes").decode("ascii") + attachment = _image_attachment() + attachment["content"] = [{"type": "image", "image": "data:;base64," + payload}] + _seed(tmp_path, monkeypatch, [attachment]) + response = chat_history.get_attachment_file("msg-1", "att-1", current_subject = "unsloth") + assert response.body == b"raw-bytes" + assert response.media_type == "application/octet-stream" + + +def test_attachment_file_svg_media_type(tmp_path, monkeypatch): + svg = b"" + payload = base64.b64encode(svg).decode("ascii") + attachment = _image_attachment() + attachment["content"] = [{"type": "image", "image": "data:image/svg+xml;base64," + payload}] + _seed(tmp_path, monkeypatch, [attachment]) + response = chat_history.get_attachment_file("msg-1", "att-1", current_subject = "unsloth") + assert response.body == svg + # SVG can carry scripts, so it downloads as bytes instead of rendering. + assert response.media_type == "application/octet-stream" + + +def test_delete_attachment_route_then_404(tmp_path, monkeypatch): + _seed(tmp_path, monkeypatch, [_image_attachment()]) + result = chat_history.delete_attachment("msg-1", "att-1", current_subject = "unsloth") + assert result == {"ok": True} + with pytest.raises(HTTPException) as excinfo: + chat_history.delete_attachment("msg-1", "att-1", current_subject = "unsloth") + assert excinfo.value.status_code == 404 + + +# --------------------------------------------------------------------------- +# Audio attachments (adapter {data, format} and compare-chat bare base64) +# --------------------------------------------------------------------------- + +WAV_BYTES = b"RIFF$\x00\x00\x00WAVEfmt \x10\x00\x00\x00\x01\x00\x01\x00" +WAV_B64 = base64.b64encode(WAV_BYTES).decode("ascii") + + +def _audio_attachment(attachment_id: str = "att-audio") -> dict: + return { + "id": attachment_id, + "type": "file", + "name": "clip.wav", + "contentType": "audio/wav", + "content": [{"type": "audio", "audio": {"data": WAV_B64, "format": "wav"}}], + "status": {"type": "complete"}, + } + + +def test_audio_attachment_lists_with_size(tmp_path, monkeypatch): + _seed(tmp_path, monkeypatch, [_audio_attachment()]) + records = studio_db.list_chat_attachments() + assert len(records) == 1 + assert records[0]["id"] == "att-audio" + assert abs(records[0]["sizeBytes"] - len(WAV_BYTES)) <= 2 + + +def test_audio_attachment_file_serves_bytes(tmp_path, monkeypatch): + _seed(tmp_path, monkeypatch, [_audio_attachment()]) + response = chat_history.get_attachment_file("msg-1", "att-audio", current_subject = "unsloth") + assert response.body == WAV_BYTES + assert response.media_type == "audio/wav" + + +def test_audio_attachment_media_type_from_format(tmp_path, monkeypatch): + attachment = _audio_attachment() + attachment["contentType"] = None + attachment["content"] = [{"type": "audio", "audio": {"data": WAV_B64, "format": "mp3"}}] + _seed(tmp_path, monkeypatch, [attachment]) + response = chat_history.get_attachment_file("msg-1", "att-audio", current_subject = "unsloth") + assert response.media_type == "audio/mpeg" + + +def test_audio_attachment_corrupt_payload_is_422(tmp_path, monkeypatch): + attachment = _audio_attachment() + attachment["content"] = [{"type": "audio", "audio": {"data": "%%%", "format": "wav"}}] + _seed(tmp_path, monkeypatch, [attachment]) + with pytest.raises(HTTPException) as excinfo: + chat_history.get_attachment_file("msg-1", "att-audio", current_subject = "unsloth") + assert excinfo.value.status_code == 422 + + +# --------------------------------------------------------------------------- +# Compare-chat uploads stored as message content parts +# --------------------------------------------------------------------------- + + +def _compare_message(message_id: str = "msg-cmp") -> dict: + return { + "id": message_id, + "threadId": "thread-1", + "parentId": None, + "role": "user", + "content": [ + {"type": "image", "image": PNG_DATA_URL}, + {"type": "audio", "audio": WAV_B64}, + {"type": "text", "text": "compare these"}, + ], + "createdAt": 1_700_000_000_000, + } + + +def _seed_compare(tmp_path, monkeypatch): + _reset_studio_db(tmp_path, monkeypatch) + studio_db.upsert_chat_thread(_thread()) + studio_db.upsert_chat_message(_compare_message()) + + +_CONTENT_PART_PREFIX = "content-part-sha256-" + + +def _content_part_id_for(message_id: str, kind: str) -> str: + """Resolve the stable content-hash id for a message's stored blob. + + Content-part ids are SHA-256 hashes of the blob payload, not array + indices, so tests look them up from the listing instead of hardcoding an + index that would shift when an earlier part is deleted. + """ + for record in studio_db.list_chat_attachments(): + if record["messageId"] == message_id and record["type"] == kind: + return record["id"] + raise AssertionError(f"no {kind} content-part upload for {message_id}") + + +def test_content_part_uploads_are_listed(tmp_path, monkeypatch): + _seed_compare(tmp_path, monkeypatch) + records = studio_db.list_chat_attachments() + # Ids are stable content hashes, not array indices. + assert all(r["id"].startswith(_CONTENT_PART_PREFIX) for r in records) + assert {r["type"] for r in records} == {"image", "audio"} + image = next(r for r in records if r["type"] == "image") + assert image["contentType"] == "image/png" + assert abs(image["sizeBytes"] - len(PNG_BYTES)) <= 2 + audio = next(r for r in records if r["type"] == "audio") + assert audio["type"] == "audio" + + +def test_content_part_file_serves_image_bytes(tmp_path, monkeypatch): + _seed_compare(tmp_path, monkeypatch) + image_id = _content_part_id_for("msg-cmp", "image") + response = chat_history.get_attachment_file("msg-cmp", image_id, current_subject = "unsloth") + assert response.body == PNG_BYTES + assert response.media_type == "image/png" + + +def test_content_part_delete_keeps_text(tmp_path, monkeypatch): + _seed_compare(tmp_path, monkeypatch) + image_id = _content_part_id_for("msg-cmp", "image") + assert studio_db.delete_chat_attachment("msg-cmp", image_id) is True + message = studio_db.get_chat_message("thread-1", "msg-cmp") + types = [p["type"] for p in message["content"]] + assert types == ["audio", "text"] + # The surviving audio blob keeps its own stable hash id after the delete. + remaining = studio_db.list_chat_attachments() + assert [r["type"] for r in remaining] == ["audio"] + assert remaining[0]["id"].startswith(_CONTENT_PART_PREFIX) + assert remaining[0]["id"] != image_id + + +def test_content_part_delete_rejects_non_blob(tmp_path, monkeypatch): + _seed_compare(tmp_path, monkeypatch) + # The text part is not a stored upload, so it never gets an id: only the + # image and audio blobs are addressable. + assert len(studio_db.list_chat_attachments()) == 2 + # A well-formed but unknown content-hash id, and malformed ids, all no-op. + assert studio_db.delete_chat_attachment("msg-cmp", _CONTENT_PART_PREFIX + "0" * 64) is False + assert studio_db.delete_chat_attachment("msg-cmp", "content-part-99") is False + assert studio_db.delete_chat_attachment("msg-cmp", "content-part-x") is False + + +def test_text_only_messages_not_listed_as_uploads(tmp_path, monkeypatch): + _reset_studio_db(tmp_path, monkeypatch) + studio_db.upsert_chat_thread(_thread()) + # The word "image" inside text must not create phantom upload rows. + message = _message("msg-txt") + message["content"] = [{"type": "text", "text": 'discussing an "image" and "audio" here'}] + studio_db.upsert_chat_message(message) + assert studio_db.list_chat_attachments() == [] + + +def test_remote_image_urls_are_not_listed_as_uploads(tmp_path, monkeypatch): + _reset_studio_db(tmp_path, monkeypatch) + studio_db.upsert_chat_thread(_thread()) + message = _message("msg-remote") + message["content"] = [ + {"type": "image", "image": "https://example.com/cat.png"}, + {"type": "text", "text": "look at this"}, + ] + studio_db.upsert_chat_message(message) + # No stored bytes: nothing to list, open, or delete. + assert studio_db.list_chat_attachments() == [] + assert studio_db.get_chat_attachment("msg-remote", "content-part-0") is None + assert studio_db.delete_chat_attachment("msg-remote", "content-part-0") is False + stored = studio_db.get_chat_message("thread-1", "msg-remote") + assert [p["type"] for p in stored["content"]] == ["image", "text"] + + +def test_html_data_url_serves_as_octet_stream(tmp_path, monkeypatch): + _reset_studio_db(tmp_path, monkeypatch) + studio_db.upsert_chat_thread(_thread()) + html_b64 = base64.b64encode(b"").decode() + message = _message("msg-html") + message["content"] = [ + {"type": "image", "image": f"data:text/html;base64,{html_b64}"}, + ] + studio_db.upsert_chat_message(message) + attachment_id = _content_part_id_for("msg-html", "image") + response = chat_history.get_attachment_file( + "msg-html", attachment_id, current_subject = "unsloth" + ) + # Never echo a script-capable media type back under the app origin. + assert response.media_type == "application/octet-stream" + assert response.body == b"" + + +def test_svg_data_url_serves_as_octet_stream(tmp_path, monkeypatch): + _reset_studio_db(tmp_path, monkeypatch) + studio_db.upsert_chat_thread(_thread()) + svg_b64 = base64.b64encode(b"").decode() + message = _message("msg-svg") + message["content"] = [ + {"type": "image", "image": f"data:image/svg+xml;base64,{svg_b64}"}, + ] + studio_db.upsert_chat_message(message) + attachment_id = _content_part_id_for("msg-svg", "image") + response = chat_history.get_attachment_file("msg-svg", attachment_id, current_subject = "unsloth") + assert response.media_type == "application/octet-stream" + + +def test_png_data_url_keeps_its_media_type(tmp_path, monkeypatch): + _seed_compare(tmp_path, monkeypatch) + image_id = _content_part_id_for("msg-cmp", "image") + response = chat_history.get_attachment_file("msg-cmp", image_id, current_subject = "unsloth") + assert response.media_type == "image/png" diff --git a/studio/backend/tests/test_chat_load_during_training.py b/studio/backend/tests/test_chat_load_during_training.py index 63dba8579c..a5fd71b6a0 100644 --- a/studio/backend/tests/test_chat_load_during_training.py +++ b/studio/backend/tests/test_chat_load_during_training.py @@ -168,11 +168,14 @@ class TestCanLoadGGUF(_GpuCacheResetMixin, unittest.TestCase): devices, required_override = None, estimate = None, + single_device_gpu = None, + gpu_ids = None, ): with ( patch("utils.hardware.get_device", return_value = DeviceType.CUDA), patch("utils.hardware.estimate_required_model_memory_gb", return_value = (estimate, {})), patch("utils.hardware.get_visible_gpu_utilization", return_value = {"devices": devices}), + patch("utils.hardware.resolve_requested_gpu_ids", return_value = gpu_ids), patch("utils.hardware.auto_select_gpu_ids") as auto_mock, ): ok, info = tv.can_load_chat_during_training( @@ -180,9 +183,10 @@ class TestCanLoadGGUF(_GpuCacheResetMixin, unittest.TestCase): hf_token = None, load_in_4bit = True, max_seq_length = 0, - requested_gpu_ids = None, + requested_gpu_ids = gpu_ids, is_gguf = True, required_override_gb = required_override, + single_device_gpu = single_device_gpu, ) return ok, info, auto_mock @@ -198,6 +202,88 @@ class TestCanLoadGGUF(_GpuCacheResetMixin, unittest.TestCase): ok, _, _ = self._run(devices = _devices((0, 80, 35), (1, 80, 70)), required_override = 20.0) self.assertTrue(ok) + def test_no_per_gpu_floor_for_gguf_with_explicit_gpu_ids(self): + # gpu_ids narrows llama.cpp's candidate pool but does not turn its + # self-placement into HF device_map="balanced". The uneven selected + # pair therefore keeps the aggregate GGUF check without an even-share + # floor on the nearly-full card. + ok, info, _ = self._run( + devices = _devices((0, 80, 35), (1, 80, 70), (2, 80, 0)), + required_override = 20.0, + gpu_ids = [0, 1], + ) + self.assertTrue(ok) + self.assertEqual(info["mode"], "gguf") + + def test_single_device_uses_selected_gpu(self): + # The model needs 27 GB with headroom. GPU 0 has 45 GB free, while an + # unrelated training-heavy GPU 1 has only 10 GB free. + ok, info, _ = self._run( + devices = _devices((0, 80, 35), (1, 80, 70)), + required_override = 20.0, + single_device_gpu = "0", + ) + self.assertTrue(ok) + self.assertEqual(info["usable_gb"], 45.0) + + blocked, blocked_info, _ = self._run( + devices = _devices((0, 80, 35), (1, 80, 70)), + required_override = 20.0, + single_device_gpu = "1", + ) + self.assertFalse(blocked) + self.assertEqual(blocked_info["usable_gb"], 10.0) + + def test_single_device_unresolved_token_sizes_against_worst_device(self): + # A non-numeric device token (a CUDA UUID / MIG handle) can't map to a + # free-VRAM index. The runner still drives ONE device, so size against the + # worst-case visible device (min free), not the aggregate pool: one GPU + # with 80 GB free vs a 20 GB model -> allow. + ok, info, _ = self._run( + devices = _devices((0, 80, 0)), + required_override = 20.0, + single_device_gpu = "GPU-uuid", + ) + self.assertTrue(ok) + self.assertEqual(info["mode"], "single_device") + self.assertNotIn("reason", info) + + def test_single_device_unresolved_token_refuses_when_worst_device_full(self): + # Same UUID fallback, worst-case device nearly full (2 GB for a 20 GB + # model) -> refuse (default-deny), not on an unresolved-token technicality. + ok, info, _ = self._run( + devices = _devices((0, 80, 78)), + required_override = 20.0, + single_device_gpu = "GPU-uuid", + ) + self.assertFalse(ok) + self.assertNotEqual(info.get("reason"), "unresolved_gpu_id") + + def test_single_device_unresolved_token_uses_min_free_not_aggregate(self): + # The single-device runner uses ONE device but we can't tell which from a + # UUID token. Sizing against the aggregate pool would let a 20 GB model + # "fit" 160 GB of pooled free VRAM while landing on a 2 GB card and OOMing + # training. Min-free (2 GB) is the safe worst case -> refuse. + ok, info, _ = self._run( + devices = _devices((0, 80, 78), (1, 80, 0), (2, 80, 0)), + required_override = 20.0, + single_device_gpu = "GPU-uuid", + ) + self.assertFalse(ok) + self.assertEqual(info["mode"], "single_device") + + def test_single_device_cpu_token_allows(self): + # An empty device token = a CPU-only single-device runner (CPU diffusion + # GGUF): it uses no GPU VRAM, so it never threatens training -> allow + # regardless of how full the GPUs are. + ok, info, _ = self._run( + devices = _devices((0, 80, 78)), + required_override = 20.0, + single_device_gpu = "", + ) + self.assertTrue(ok) + self.assertEqual(info["reason"], "cpu_only") + def test_estimate_unavailable_refuses(self): # No override and the estimator can't size it -> default-deny. ok, info, _ = self._run(devices = _devices((0, 80, 0)), required_override = None, estimate = None) @@ -309,6 +395,8 @@ class TestChatLoadGuardRoute(unittest.TestCase): captured = None, training_active, decision, + gpu_memory_mode = "auto", + requested_gpu_ids = None, ): config = config or SimpleNamespace(is_gguf = False, is_lora = False, path = None) with _stub_guard_deps( @@ -320,7 +408,8 @@ class TestChatLoadGuardRoute(unittest.TestCase): hf_token = None, load_in_4bit = True, max_seq_length = 0, - requested_gpu_ids = None, + requested_gpu_ids = requested_gpu_ids, + gpu_memory_mode = gpu_memory_mode, ) def test_noop_when_training_inactive(self): @@ -332,6 +421,141 @@ class TestChatLoadGuardRoute(unittest.TestCase): def test_allows_when_fits(self): self._guard(training_active = True, decision = (True, {"mode": "auto"})) + def test_diffusion_detection_uses_name_before_download(self): + config = SimpleNamespace( + identifier = "unsloth/DiffusionGemma-GGUF", + gguf_hf_repo = "unsloth/DiffusionGemma-GGUF", + gguf_file = None, + ) + self.assertTrue(self.route._classify_diffusion_gguf(config)) + + def test_uncached_gguf_classification_remains_unknown(self): + config = SimpleNamespace( + identifier = "owner/renamed-model", + gguf_hf_repo = "owner/renamed-model", + gguf_variant = "Q4_K_M", + gguf_file = None, + ) + self.assertIsNone(self.route._classify_diffusion_gguf(config)) + + def test_diffusion_detection_reuses_loader_metadata_probe(self): + import tempfile + + seen = [] + + class _Probe: + is_diffusion = False + _architecture = None + + def _read_gguf_metadata(self, path): + seen.append(path) + self.is_diffusion = True + + with tempfile.TemporaryDirectory() as d: + model = Path(d) / "renamed.gguf" + model.write_bytes(b"GGUF") + config = SimpleNamespace(identifier = "local", gguf_file = str(model)) + with patch.object(self.route, "LlamaCppBackend", _Probe): + self.assertTrue(self.route._classify_diffusion_gguf(config)) + self.assertEqual(seen, [str(model)]) + + def test_local_chat_gguf_classification_is_definitive(self): + import tempfile + class _Probe: + is_diffusion = False + _architecture = "llama" + + def _read_gguf_metadata(self, _path): + pass + + with tempfile.TemporaryDirectory() as d: + model = Path(d) / "renamed.gguf" + model.write_bytes(b"GGUF") + config = SimpleNamespace(identifier = "local", gguf_file = str(model)) + with patch.object(self.route, "LlamaCppBackend", _Probe): + self.assertFalse(self.route._classify_diffusion_gguf(config)) + + def test_manual_known_normal_gguf_bypasses_training_estimate(self): + captured = [] + config = SimpleNamespace(is_gguf = True) + with patch.object(self.route, "_classify_diffusion_gguf", return_value = False): + self._guard( + config = config, + captured = captured, + training_active = True, + decision = (False, {"reason": "must not run"}), + gpu_memory_mode = "manual", + ) + self.assertEqual(captured, []) + + def test_manual_unknown_gguf_keeps_single_device_training_guard(self): + captured = [] + config = SimpleNamespace(is_gguf = True) + with ( + patch.object(self.route, "_classify_diffusion_gguf", return_value = None), + patch.object(self.route, "_estimate_gguf_required_gb", return_value = 12.5), + patch.object( + self.route.LlamaCppBackend, + "_diffusion_gpu_arg", + return_value = "2", + ), + ): + self._guard( + config = config, + captured = captured, + training_active = True, + decision = (True, {"mode": "single_device"}), + gpu_memory_mode = "manual", + ) + self.assertEqual(len(captured), 1) + self.assertEqual(captured[0]["single_device_gpu"], "2") + + def test_manual_diffusion_uses_single_device_guard(self): + captured = [] + config = SimpleNamespace(is_gguf = True) + with ( + patch.object(self.route, "_classify_diffusion_gguf", return_value = True), + patch.object(self.route, "_estimate_gguf_required_gb", return_value = 12.5), + ): + self._guard( + config = config, + captured = captured, + training_active = True, + decision = (True, {"mode": "gguf"}), + gpu_memory_mode = "manual", + requested_gpu_ids = [3, 1], + ) + self.assertEqual(len(captured), 1) + self.assertEqual(captured[0]["single_device_gpu"], "1") + self.assertEqual(captured[0]["requested_gpu_ids"], [3, 1]) + + def test_unpinned_diffusion_uses_runner_default_gpu(self): + captured = [] + config = SimpleNamespace(is_gguf = True) + with ( + patch.object(self.route, "_classify_diffusion_gguf", return_value = True), + patch.object(self.route, "_estimate_gguf_required_gb", return_value = 12.5), + patch.object( + self.route.LlamaCppBackend, + "_effective_gpu_count", + return_value = 2, + ), + patch.object( + self.route.LlamaCppBackend, + "_diffusion_gpu_arg", + return_value = "3", + ) as gpu_arg, + ): + self._guard( + config = config, + captured = captured, + training_active = True, + decision = (True, {"mode": "single_device"}), + gpu_memory_mode = "manual", + ) + gpu_arg.assert_called_once_with(None, cpu_only = False) + self.assertEqual(captured[0]["single_device_gpu"], "3") + def test_refuses_with_headroom_number(self): info = {"required_gb": 30.0, "usable_gb": 6.0, "needed_gb": 39.0, "mode": "auto"} with self.assertRaises(HTTPException) as exc: @@ -467,36 +691,189 @@ class TestValidateRefusesDuringTraining(unittest.TestCase): self.assertEqual(captured[0]["load_in_4bit"], False) self.assertEqual(captured[0]["max_seq_length"], 4096) - def test_rejects_gguf_with_gpu_ids_before_guard(self): - # /validate must mirror /load's GGUF + gpu_ids 400, before the VRAM guard. + def test_validate_forwards_manual_gpu_memory_mode_to_guard(self): from models.inference import ValidateModelRequest - request = ValidateModelRequest(model_path = "x.gguf", gpu_ids = [0]) + request = ValidateModelRequest( + model_path = "unsloth/model-GGUF", + gguf_variant = "Q4_K_M", + gpu_memory_mode = "manual", + ) cfg = SimpleNamespace( - identifier = "x.gguf", - display_name = "x", + identifier = "unsloth/model-GGUF", + display_name = "model-GGUF", is_gguf = True, is_lora = False, is_vision = False, path = None, base_model = None, ) - captured = [] + captured = {} with ( patch.object( self.route, "_resolve_model_identifier_for_request", - return_value = ("x.gguf", "x.gguf", False), + return_value = ("unsloth/model-GGUF", "unsloth/model-GGUF", False), ), patch.object(self.route.ModelConfig, "from_identifier", return_value = cfg), patch.object(self.route, "load_inference_config", return_value = {}), - _stub_guard_deps(training_active = True, decision = (True, {}), captured = captured), + patch.object( + self.route, + "_guard_chat_load_against_training", + lambda config, **kw: captured.update(kw), + ), ): - with self.assertRaises(HTTPException) as exc: - asyncio.run(self.route.validate_model(request, current_subject = "u")) - self.assertEqual(exc.exception.status_code, 400) - self.assertIn("gpu_ids is not supported for GGUF", exc.exception.detail) - self.assertEqual(captured, []) # guard never reached + asyncio.run(self.route.validate_model(request, current_subject = "u")) + self.assertEqual(captured.get("gpu_memory_mode"), "manual") + + def test_validate_forwards_inherited_extras_and_parallel_to_guard(self): + # Regression: /load resolves inherited same-model extras and passes the + # real slot count to the guard; validate must do the same, else it sizes + # a smaller estimate (no inherited -c/--model-draft, n_parallel=1) and + # /load then 409s after the frontend has already unloaded. + from models.inference import ValidateModelRequest + + request = ValidateModelRequest(model_path = "unsloth/Qwen3-1.7B", max_seq_length = 4096) + cfg = SimpleNamespace( + identifier = "unsloth/Qwen3-1.7B", + display_name = "Qwen3-1.7B", + is_gguf = False, + is_lora = False, + is_vision = False, + path = None, + base_model = None, + ) + captured = {} + with ( + patch.object( + self.route, + "_resolve_model_identifier_for_request", + return_value = ("unsloth/Qwen3-1.7B", "unsloth/Qwen3-1.7B", False), + ), + patch.object(self.route.ModelConfig, "from_identifier", return_value = cfg), + patch.object(self.route, "load_inference_config", return_value = {}), + patch.object(self.route, "_resolve_inherited_extra_args", return_value = ["-c", "32768"]), + patch.object( + self.route, + "_guard_chat_load_against_training", + lambda config, **kw: captured.update(kw), + ), + ): + asyncio.run(self.route.validate_model(request, current_subject = "u")) + self.assertEqual(captured.get("llama_extra_args"), ["-c", "32768"]) + self.assertIn("n_parallel", captured) + + def test_metadata_probe_skips_training_guard(self): + # A header-only probe (include_context_length) allocates no VRAM, so the + # training guard must not run -- else the staging GPU-layers / MoE sliders + # it feeds are hidden exactly when a during-training user needs them. + from models.inference import ValidateModelRequest + + request = ValidateModelRequest( + model_path = "unsloth/Qwen3-1.7B", + max_seq_length = 4096, + include_context_length = True, + ) + cfg = SimpleNamespace( + identifier = "unsloth/Qwen3-1.7B", + display_name = "Qwen3-1.7B", + is_gguf = False, + is_lora = False, + is_vision = False, + path = None, + base_model = None, + ) + guard_called = [] + with ( + patch.object( + self.route, + "_resolve_model_identifier_for_request", + return_value = ("unsloth/Qwen3-1.7B", "unsloth/Qwen3-1.7B", False), + ), + patch.object(self.route.ModelConfig, "from_identifier", return_value = cfg), + patch.object(self.route, "load_inference_config", return_value = {}), + patch.object( + self.route, + "_guard_chat_load_against_training", + lambda *a, **kw: guard_called.append(True), + ), + ): + asyncio.run(self.route.validate_model(request, current_subject = "u")) + self.assertEqual(guard_called, []) + + def _validate_gguf_template( + self, + *, + template, + canonical_path = "/picked/model.gguf", + ): + # Drive validate_model for a native lease-backed GGUF template probe and + # capture what the embedded-template reader was called with. + from models.inference import ValidateModelRequest + + request = ValidateModelRequest( + model_path = "model.gguf", + gguf_variant = "Q4_K_M", + native_path_lease = "signed-lease", + include_chat_template = True, + ) + cfg = SimpleNamespace( + identifier = canonical_path, + display_name = "model.gguf", + is_gguf = True, + is_lora = False, + is_vision = False, + gguf_file = canonical_path, + path = None, + base_model = None, + ) + import utils.models.gguf_metadata as gguf_meta + + seen = {} + + def _fake_read(path): + seen["path"] = path + return template + + guard_called = [] + with ( + patch.object( + self.route, + "_resolve_model_identifier_for_request", + return_value = (canonical_path, "model.gguf", True), + ), + patch.object(self.route.ModelConfig, "from_identifier", return_value = cfg), + patch.object(self.route, "load_inference_config", return_value = {}), + patch.object(gguf_meta, "read_gguf_chat_template", _fake_read), + patch.object( + self.route, + "_guard_chat_load_against_training", + lambda *a, **kw: guard_called.append(True), + ), + ): + resp = asyncio.run(self.route.validate_model(request, current_subject = "u")) + return resp, seen, guard_called + + def test_include_chat_template_reads_leased_gguf_embedded_template(self): + # The picker chat-template GET has no lease plumbing, so a native picked + # GGUF surfaces its default template through this lease-aware probe: the + # embedded template is read from the granted canonical path and returned. + resp, seen, _ = self._validate_gguf_template(template = "{{ messages }}") + self.assertEqual(resp.chat_template, "{{ messages }}") + # Read strictly the leased file's own embedded template, never a sibling + # sidecar: the grant authorizes just this one path. + self.assertEqual(seen["path"], "/picked/model.gguf") + + def test_include_chat_template_skips_training_guard(self): + # A template-only probe allocates no VRAM, so like include_context_length + # it must not be refused by the training guard. + _, _, guard_called = self._validate_gguf_template(template = "{{ messages }}") + self.assertEqual(guard_called, []) + + def test_include_chat_template_over_cap_is_dropped(self): + from picker.schemas import MAX_CHAT_TEMPLATE_BYTES + resp, _, _ = self._validate_gguf_template(template = "a" * (MAX_CHAT_TEMPLATE_BYTES + 1)) + self.assertIsNone(resp.chat_template) # ── _estimate_gguf_required_gb (sizes the same weights the loader loads) ────── diff --git a/studio/backend/tests/test_compute_buffer.py b/studio/backend/tests/test_compute_buffer.py index 8408f8203d..3e95acc98d 100644 --- a/studio/backend/tests/test_compute_buffer.py +++ b/studio/backend/tests/test_compute_buffer.py @@ -152,7 +152,7 @@ class TestFallback: class TestParallel1Default: - """At Studio's default --parallel 1 the buffer is negligible in pipeline.""" + """At Unsloth's default --parallel 1 the buffer is negligible in pipeline.""" def test_default_n_parallel(self): est = _backend()._estimate_compute_buffer_bytes() / MIB diff --git a/studio/backend/tests/test_cpu_threads.py b/studio/backend/tests/test_cpu_threads.py index 2930c9f081..9d8795b6c0 100644 --- a/studio/backend/tests/test_cpu_threads.py +++ b/studio/backend/tests/test_cpu_threads.py @@ -1,7 +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 -"""Tests for Studio's early CPU thread-pool configuration.""" +"""Tests for Unsloth's early CPU thread-pool configuration.""" import ast import os @@ -30,7 +30,7 @@ def test_cpu_thread_cap_seeds_native_pool_limits(): } -# Explicit per-library values win over the Studio knob via setdefault. +# Explicit per-library values win over the Unsloth knob via setdefault. def test_cpu_thread_cap_preserves_runtime_specific_override(): env = {"UNSLOTH_CPU_THREADS": "4", "OMP_NUM_THREADS": "2"} diff --git a/studio/backend/tests/test_cuda_torch_spec.py b/studio/backend/tests/test_cuda_torch_spec.py new file mode 100644 index 0000000000..928cef787e --- /dev/null +++ b/studio/backend/tests/test_cuda_torch_spec.py @@ -0,0 +1,73 @@ +# SPDX-License-Identifier: AGPL-3.0-only +# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0 + +"""Tests for _CUDA_TORCH_PKG_SPEC in install_python_stack.py. + +The CUDA repair path installs the torch trio from an exclusive --index-url (no +PyPI fallback), so these pinned ranges decide which torch the venv gets. The +upper bound is locked to the 2.11.x family to match the base image and rocm7.2 +spec and to keep the companions off a torch-2.12 wheel that would ABI-mismatch. +""" + +from __future__ import annotations + +import sys +from pathlib import Path + +import pytest +from packaging.requirements import Requirement + +# install_python_stack.py lives at repo_root/studio/install_python_stack.py +_INSTALL_SCRIPT = Path(__file__).resolve().parents[2] / "install_python_stack.py" + + +def _load_module(monkeypatch): + """(Re-)import and return install_python_stack (mirrors test_torchao_select).""" + sys.modules.pop("install_python_stack", None) + monkeypatch.syspath_prepend(str(_INSTALL_SCRIPT.parent)) + import install_python_stack + + return install_python_stack + + +def _spec_of(pkg_spec: str): + """Parse 'torch>=2.4,<2.12.0' into a packaging SpecifierSet.""" + return Requirement(pkg_spec).specifier + + +@pytest.mark.parametrize( + "index, allowed, rejected", + [ + # torch: 2.11.x allowed (matches base image); 2.12.x excluded. + (0, ["2.11.0", "2.11.2", "2.10.0", "2.4.0"], ["2.12.0", "2.3.0", "1.13.1"]), + # torchvision: 0.26.x (torch 2.11 companion) allowed; 0.27.x (torch 2.12) out. + (1, ["0.26.0", "0.26.1", "0.19.0"], ["0.27.0", "0.18.0"]), + # torchaudio: same 2.11.x window as torch. + (2, ["2.11.0", "2.10.0", "2.4.0"], ["2.12.0", "2.3.0"]), + ], +) +def test_cuda_spec_bounds(monkeypatch, index, allowed, rejected): + mod = _load_module(monkeypatch) + spec = _spec_of(mod._CUDA_TORCH_PKG_SPEC[index]) + for v in allowed: + assert spec.contains(v, prereleases = True), f"{v} should satisfy {spec}" + for v in rejected: + assert not spec.contains(v, prereleases = True), f"{v} should not satisfy {spec}" + + +def test_cuda_spec_matches_rocm72_upper_bound(monkeypatch): + """CUDA and rocm7.2 target the same torch 2.11.x family, so their upper + bounds must stay in lockstep (bump both together at 2.12.x).""" + mod = _load_module(monkeypatch) + rocm72 = mod._ROCM_TORCH_PKG_SPECS["rocm7.2"] + + def _upper(pkg_spec: str) -> str: + for clause in _spec_of(pkg_spec): + if clause.operator == "<": + return clause.version + raise AssertionError(f"no upper bound in {pkg_spec!r}") + + for cuda_pkg, rocm_pkg in zip(mod._CUDA_TORCH_PKG_SPEC, rocm72, strict = True): + assert _upper(cuda_pkg) == _upper( + rocm_pkg + ), f"CUDA {cuda_pkg!r} upper bound must match rocm7.2 {rocm_pkg!r}" diff --git a/studio/backend/tests/test_embedding_model_security_gate.py b/studio/backend/tests/test_embedding_model_security_gate.py index 940b35d7ba..b3fa98b604 100644 --- a/studio/backend/tests/test_embedding_model_security_gate.py +++ b/studio/backend/tests/test_embedding_model_security_gate.py @@ -52,6 +52,16 @@ def client(monkeypatch): monkeypatch.setattr(settings, "_resolves_as_local_gguf", lambda m: False) monkeypatch.setattr(settings, "get_rag_embedding_model", lambda: saved.get("model", "")) monkeypatch.setattr(settings, "get_stored_embedding_model", lambda: saved.get("model")) + monkeypatch.setattr( + settings, + "effective_gguf_repo", + lambda: f"{saved.get('model', 'unsloth/default-embed')}-GGUF", + ) + monkeypatch.setattr( + settings, + "default_gguf_repo", + lambda: "unsloth/default-embed-GGUF", + ) app = FastAPI() app.include_router(settings.router) @@ -257,6 +267,13 @@ def test_clean_repo_saves_under_force(client, monkeypatch): r = c.put("/embedding-model", json = {"embedding_model": "acme/clean-embed", "force": True}) assert r.status_code == 200 assert saved.get("model") == "acme/clean-embed" + assert r.json() == { + "embedding_model": "acme/clean-embed", + "embedding_gguf_repo": "acme/clean-embed-GGUF", + "default_embedding_model": "unsloth/default-embed", + "default_embedding_gguf_repo": "unsloth/default-embed-GGUF", + "is_custom": True, + } def test_load_sink_refuses_flagged_model(monkeypatch): diff --git a/studio/backend/tests/test_embedding_model_settings.py b/studio/backend/tests/test_embedding_model_settings.py index 3be4af0e32..bcf3ded71c 100644 --- a/studio/backend/tests/test_embedding_model_settings.py +++ b/studio/backend/tests/test_embedding_model_settings.py @@ -53,3 +53,10 @@ def test_custom_model_overrides_default_and_derives_gguf(settings_store, monkeyp assert ems.reset_rag_embedding_model() == rag_config.EMBEDDING_MODEL assert ems.get_stored_embedding_model() is None + + +def test_env_default_derives_its_gguf_companion(monkeypatch): + monkeypatch.delenv("RAG_EMBED_GGUF_REPO", raising = False) + monkeypatch.setattr(rag_config, "EMBEDDING_MODEL", "org/env-default-embedder") + + assert rag_config.default_gguf_repo() == "org/env-default-embedder-GGUF" diff --git a/studio/backend/tests/test_export_absolute_paths.py b/studio/backend/tests/test_export_absolute_paths.py index 761ea08e3f..5097f9f53a 100644 --- a/studio/backend/tests/test_export_absolute_paths.py +++ b/studio/backend/tests/test_export_absolute_paths.py @@ -158,6 +158,7 @@ def _install_lightweight_backend_stubs(monkeypatch): utils_model_config._pick_best_gguf = lambda variants: variants[0] if variants else None utils_model_config._extract_quant_label = lambda value: value utils_model_config._is_big_endian_gguf_path = lambda *args, **kwargs: False + utils_model_config._is_mtp_drafter = lambda *args, **kwargs: False utils_model_config.is_audio_input_type = lambda *args, **kwargs: None monkeypatch.setitem( sys.modules, diff --git a/studio/backend/tests/test_frontend_resolution.py b/studio/backend/tests/test_frontend_resolution.py index c3e0524a30..7ac2717aae 100644 --- a/studio/backend/tests/test_frontend_resolution.py +++ b/studio/backend/tests/test_frontend_resolution.py @@ -218,7 +218,7 @@ def test_systemexit_message_contains_actionable_fixes(tmp_path, monkeypatch): installer_bin = home / "unsloth_studio" / "bin" / "unsloth" tried_lines = "\n".join(f" - {p}" for p in attempted) message = ( - "[ERROR] Studio frontend build not found.\n" + "[ERROR] Unsloth frontend build not found.\n" f"Tried:\n{tried_lines}\n" "\n" "Likely cause: another 'unsloth' on PATH is shadowing the " diff --git a/studio/backend/tests/test_gemini_provider.py b/studio/backend/tests/test_gemini_provider.py index 85ceb04d27..c6ffa798d0 100644 --- a/studio/backend/tests/test_gemini_provider.py +++ b/studio/backend/tests/test_gemini_provider.py @@ -768,7 +768,7 @@ def test_cached_content_pass_through(monkeypatch): def test_boolean_caching_does_not_set_cached_content(monkeypatch): - """Studio's existing True/False signals shouldn't fabricate a cache id.""" + """Unsloth's existing True/False signals shouldn't fabricate a cache id.""" captured = _capture_body(monkeypatch, enable_prompt_caching = True) assert "cachedContent" not in captured["body"] @@ -2613,7 +2613,7 @@ def test_gemini_native_skips_orphan_function_response_for_native_part_replay(mon def test_gemini_native_part_falls_back_to_args_google(monkeypatch): """Round 27: a direct OpenAI-compat API caller (or imported third-party - thread) cannot use Studio's non-standard `tool_calls[].extra_content` + thread) cannot use Unsloth's non-standard `tool_calls[].extra_content` field, so the native_part payload round-trips through `function.arguments` as `{"google": {"native_part": {...}}}`. The synthetic-builtin detector recognizes that location, but the replay branch was only reading from diff --git a/studio/backend/tests/test_gemma4_chat_template_override.py b/studio/backend/tests/test_gemma4_chat_template_override.py index f726741aa5..9fb24a4cf6 100644 --- a/studio/backend/tests/test_gemma4_chat_template_override.py +++ b/studio/backend/tests/test_gemma4_chat_template_override.py @@ -3,7 +3,7 @@ """Auto-override of the chat template for ``unsloth/gemma-4-*-GGUF``. -Studio ships a bundled ``gemma-4.jinja`` (PR #118 based, ``preserve_thinking`` +Unsloth ships a bundled ``gemma-4.jinja`` (PR #118 based, ``preserve_thinking`` defaulted off) and applies it to gemma-4 GGUF loads via the existing ``chat_template_override`` -> ``--chat-template-file`` path, so users do not need to re-download quants. Pins the family matcher, the resolver precedence, the diff --git a/studio/backend/tests/test_gguf_load_cache_reuse.py b/studio/backend/tests/test_gguf_load_cache_reuse.py index 15d91cd324..62596fcc8a 100644 --- a/studio/backend/tests/test_gguf_load_cache_reuse.py +++ b/studio/backend/tests/test_gguf_load_cache_reuse.py @@ -728,9 +728,13 @@ class TestLoadHubDownloadExclusion: source = (Path(__file__).resolve().parent.parent / "routes" / "inference.py").read_text() gguf_branch = source[source.index("if config.is_gguf:") :] + # The gguf_load_in_flight marker must be entered before the hub-download + # guard and the unload so a concurrent load can't race the download + # manager. The llama_extra_args inheritance that used to sit between the + # marker and the guard now runs in _guard_chat_load_against_training, ahead + # of the GGUF branch, so it is no longer a landmark inside this slice. assert ( gguf_branch.index("enter_context(gguf_load_in_flight") - < gguf_branch.index("if request.llama_extra_args is None") < gguf_branch.index("_hub_download_blocks_gguf_load") < gguf_branch.index("unsloth_backend.unload_model") ) diff --git a/studio/backend/tests/test_gguf_metadata.py b/studio/backend/tests/test_gguf_metadata.py index a5be07f8e3..ec0330ce05 100644 --- a/studio/backend/tests/test_gguf_metadata.py +++ b/studio/backend/tests/test_gguf_metadata.py @@ -15,6 +15,7 @@ from utils.models.gguf_metadata import ( pairing_score, read_gguf_context_length, read_gguf_general_metadata, + read_gguf_staged_dims, read_mmproj_audio_capability, ) @@ -153,6 +154,78 @@ def test_context_length_ignores_foreign_arch_key(tmp_path: Path): assert read_gguf_context_length(str(p)) is None +# --- read_gguf_staged_dims (one pass: context + layer + moe counts) ---- + + +def test_staged_dims_none_for_missing_or_non_gguf(tmp_path: Path): + assert read_gguf_staged_dims(str(tmp_path / "nope.gguf")) is None + p = tmp_path / "garbage.gguf" + p.write_bytes(b"not a gguf at all") + assert read_gguf_staged_dims(str(p)) is None + + +def test_staged_dims_moe_with_leading_dense(tmp_path: Path): + # GLM-4.7-Flash shape: context + total layers + MoE layers in one read. + p = _write_synthetic_gguf( + tmp_path / "glm.gguf", + {"general.architecture": "deepseek2"}, + extra_uint32 = { + "deepseek2.context_length": 202752, + "deepseek2.block_count": 47, + "deepseek2.expert_count": 64, + "deepseek2.leading_dense_block_count": 1, + }, + ) + assert read_gguf_staged_dims(str(p)) == { + "context_length": 202752, + "layer_count": 47, + "moe_layer_count": 46, + } + + +def test_staged_dims_dense_model(tmp_path: Path): + # Dense: layer_count present, moe_layer_count 0 (slider hidden). + p = _write_synthetic_gguf( + tmp_path / "dense.gguf", + {"general.architecture": "qwen3"}, + extra_uint32 = {"qwen3.context_length": 40960, "qwen3.block_count": 36}, + ) + assert read_gguf_staged_dims(str(p)) == { + "context_length": 40960, + "layer_count": 36, + "moe_layer_count": 0, + } + + +def test_staged_dims_all_moe_no_leading_dense(tmp_path: Path): + # Experts present, no leading_dense key -> every block is a MoE layer. + p = _write_synthetic_gguf( + tmp_path / "moe.gguf", + {"general.architecture": "qwen35moe"}, + extra_uint32 = {"qwen35moe.block_count": 40, "qwen35moe.expert_count": 256}, + ) + assert read_gguf_staged_dims(str(p)) == { + "context_length": None, + "layer_count": 40, + "moe_layer_count": 40, + } + + +def test_staged_dims_uint64_block_count(tmp_path: Path): + # block_count stored as uint64 (vtype 10) still parses; moe == block_count. + p = _write_synthetic_gguf( + tmp_path / "moe64.gguf", + {"general.architecture": "gpt-oss"}, + extra_uint32 = {"gpt-oss.expert_count": 32}, + extra_uint64 = {"gpt-oss.block_count": 24}, + ) + assert read_gguf_staged_dims(str(p)) == { + "context_length": None, + "layer_count": 24, + "moe_layer_count": 24, + } + + def test_context_length_read_from_uint64(tmp_path: Path): # Some models store context_length as a uint64 (vtype 10). p = _write_synthetic_gguf( diff --git a/studio/backend/tests/test_gpu_memory_mode.py b/studio/backend/tests/test_gpu_memory_mode.py new file mode 100644 index 0000000000..b17274197f --- /dev/null +++ b/studio/backend/tests/test_gpu_memory_mode.py @@ -0,0 +1,879 @@ +# SPDX-License-Identifier: AGPL-3.0-only +# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0 + +"""Backend contract for the GPU Memory mode dropdown. + +The dropdown threads a single ``gpu_memory_mode`` ("auto" | "manual") from the +chat UI through the load request. "manual" lets the user own the offload: with +``gpu_layers < 0`` (Auto, the default) it hands all memory management to +llama.cpp's ``--fit on`` (no CUDA/HIP device masking, no context auto-reduce, no +gpu-layer or tensor-split planning); with ``gpu_layers >= 0`` it pins the layers +and MoE offload itself (``--fit off``). These tests pin: + + * the pydantic request/response/status contract (snake_case key, default + "auto", unknown values rejected), + * the backend ``gpu_memory_mode`` property and its reset on unload, + * the ``_already_in_target_state`` reload-detection branch, and + * that the manual + Auto-layers branch in ``load_model`` empties the probed + GPU set and drops tensor parallelism so the selection below no-ops, while + the explicit-offload branch emits ``--gpu-layers`` / ``--fit off``. +""" + +from __future__ import annotations + +import inspect +import sys +import types as _types +from pathlib import Path + +import pytest + +_BACKEND_DIR = str(Path(__file__).resolve().parent.parent) +if _BACKEND_DIR not in sys.path: + sys.path.insert(0, _BACKEND_DIR) + +# Same external-dep stubs as the other llama_cpp unit tests so importing +# the backend doesn't drag in structlog / httpx / loggers. +_loggers_stub = _types.ModuleType("loggers") +_loggers_stub.get_logger = lambda name: __import__("logging").getLogger(name) +sys.modules.setdefault("loggers", _loggers_stub) + +_structlog_stub = _types.ModuleType("structlog") +_structlog_stub.get_logger = lambda *a, **k: __import__("logging").getLogger("stub") +sys.modules.setdefault("structlog", _structlog_stub) + +# httpx is a real, installed backend dependency: import it so the genuine module +# is in sys.modules. A hand-rolled stub here is inevitably incomplete and, since +# setdefault installs it before real httpx loads, would poison a combined pytest +# run -- routes/inference references httpx.Response (and other attrs) at def time. +import httpx # noqa: F401 + +from core.inference import llama_cpp as llama_cpp_module +from core.inference.llama_cpp import LlamaCppBackend +from models.inference import ( + InferenceStatusResponse, + LoadRequest, + LoadResponse, +) + + +# ── Pydantic contract (snake_case key, default "auto") ─────────────── + + +def test_load_request_defaults_gpu_memory_mode_auto(): + assert LoadRequest(model_path = "owner/repo").gpu_memory_mode == "auto" + + +def test_load_request_round_trips_json_key(): + req = LoadRequest.model_validate({"model_path": "owner/repo", "gpu_memory_mode": "manual"}) + assert req.gpu_memory_mode == "manual" + assert req.model_dump()["gpu_memory_mode"] == "manual" + + +def test_load_request_rejects_unknown_mode(): + with pytest.raises(ValueError): + LoadRequest(model_path = "owner/repo", gpu_memory_mode = "bogus") + + +@pytest.mark.parametrize("model_cls", [LoadResponse, InferenceStatusResponse]) +def test_response_models_emit_gpu_memory_mode(model_cls): + if model_cls is LoadResponse: + default = model_cls( + status = "loaded", + model = "owner/repo", + display_name = "repo", + inference = {}, + ) + manual = model_cls( + status = "loaded", + model = "owner/repo", + display_name = "repo", + inference = {}, + gpu_memory_mode = "manual", + ) + else: + default = model_cls() + manual = model_cls(gpu_memory_mode = "manual") + assert default.model_dump()["gpu_memory_mode"] == "auto" + assert manual.model_dump()["gpu_memory_mode"] == "manual" + + +# ── Backend property + reset ───────────────────────────────────────── + + +class _FakeProcess: + """Stand-in for subprocess.Popen so _kill_process is a no-op.""" + + def terminate(self): + pass + + def wait(self, timeout = None): + return 0 + + def kill(self): + pass + + def poll(self): + return 0 + + +def test_gpu_memory_mode_property_defaults_auto(): + assert LlamaCppBackend().gpu_memory_mode == "auto" + + +def test_gpu_memory_mode_property_reflects_field(): + backend = LlamaCppBackend() + backend._gpu_memory_mode = "manual" + assert backend.gpu_memory_mode == "manual" + + +def test_unload_resets_gpu_memory_mode(): + backend = LlamaCppBackend() + backend._process = _FakeProcess() + backend._gpu_memory_mode = "manual" + backend.unload_model() + assert backend.gpu_memory_mode == "auto" + + +# ── _already_in_target_state reload-detection branch ───────────────── + + +def _loaded_backend(gpu_memory_mode: str) -> LlamaCppBackend: + backend = LlamaCppBackend() + backend._process = _FakeProcess() # is_loaded only checks "is not None" + backend._healthy = True + backend._model_identifier = "owner/repo" + backend._hf_variant = "Q4_K_M" + backend._requested_n_ctx = 8192 + backend._cache_type_kv = None + backend._requested_spec_mode = "auto" + backend._chat_template_override = None + backend._is_vision = False + backend._extra_args = None + backend._gguf_path = None + backend._gpu_memory_mode = gpu_memory_mode + return backend + + +def _target_state(backend: LlamaCppBackend, gpu_memory_mode: str) -> bool: + return backend._already_in_target_state( + gguf_path = None, + model_identifier = "owner/repo", + hf_variant = "Q4_K_M", + n_ctx = 8192, + cache_type_kv = None, + speculative_type = "auto", + chat_template_override = None, + extra_args = None, + is_vision = False, + gpu_memory_mode = gpu_memory_mode, + ) + + +@pytest.mark.parametrize("mode", ["auto", "manual"]) +def test_already_in_target_state_matches_same_mode(mode): + assert _target_state(_loaded_backend(mode), mode) is True + + +@pytest.mark.parametrize("loaded,requested", [("auto", "manual"), ("manual", "auto")]) +def test_already_in_target_state_reloads_on_mode_change(loaded, requested): + # Flipping the dropdown either direction must force a reload so the command + # is rebuilt with/without the Unsloth GPU masking. + assert _target_state(_loaded_backend(loaded), requested) is False + + +def test_already_in_target_state_ignores_mode_for_diffusion(): + # The diffusion runner is mode-agnostic (always "auto"), so a standing manual + # preference must not force a needless reload. + backend = _loaded_backend("auto") + backend._is_diffusion = True + assert _target_state(backend, "manual") is True + + +# ── load_model: manual + Auto layers bypasses Unsloth GPU management ── + + +def _load_model_source() -> str: + return inspect.getsource(llama_cpp_module.LlamaCppBackend.load_model) + + +def test_auto_layers_branch_empties_gpus_and_drops_tensor_parallel(): + # Emptying the probed set makes the selection / TP planning below no-op, so + # gpu_indices stays None and use_fit True (--fit on). + src = _load_model_source() + gate = src.find('if gpu_memory_mode == "manual" and gpu_layers < 0:') + assert gate != -1, "load_model must branch on manual + Auto layers (gpu_layers < 0)" + block = src[gate : gate + 1400] + assert "gpus = []" in block, "Auto-layers branch must empty the probed GPU set" + # --fit aborts under --split-mode tensor, so a raw-extras split-mode is stripped. + assert "strip_split_mode_only(extra_args)" in block + assert "requested_ctx if requested_ctx > 0 else 0" in block + # The branch sits before GPU selection assigns gpu_indices; --fit on is its emission. + assert gate < src.find("gpu_indices, use_fit = None, True") + assert 'cmd.extend(["--fit", "on"])' in src + # TP drops for this path, but at a guard BEFORE the quantized-KV cache-drop, so + # a requested quantized cache survives into the --fit load. + tp_drop = src.find('if tensor_parallel and gpu_memory_mode == "manual" and gpu_layers < 0:') + assert tp_drop != -1, "manual + Auto layers must drop tensor_parallel" + assert "tensor_parallel = False" in src[tp_drop : tp_drop + 400] + cache_drop = src.find("Tensor parallelism requires a non-quantized KV cache") + assert cache_drop != -1 + assert ( + tp_drop < cache_drop + ), "TP must drop before the cache-drop so a quantized KV survives --fit" + + +def test_auto_layers_never_sends_ctx_size_zero(): + # Sending "-c 0" sets fit_params_min_ctx = UINT32_MAX in llama.cpp, pinning + # the full native context and disabling --fit's reduction. So the base cmd + # must never carry -c, "-c 0" is emitted only outside the Auto-layers (--fit) + # case, and a positive context is passed through (which --fit optimizes + # layers around). + src = _load_model_source() + base_start = src.find("cmd = [") + base_end = src.find("\n ]", base_start) + base_block = src[base_start:base_end] + assert '"-c"' not in base_block, "-c must be conditional, not in the base cmd list" + assert 'cmd.extend(["-c", str(effective_ctx)])' in src, "positive ctx must pass -c" + assert 'auto_fit = gpu_memory_mode == "manual" and gpu_layers < 0' in src + zero = src.find('cmd.extend(["-c", "0"])') + assert zero != -1, '"-c 0" emission must exist outside the Auto-layers case' + guard = src.rfind("elif not auto_fit:", 0, zero) + assert guard != -1 and zero - guard < 120, '"-c 0" must sit under the not-auto_fit guard' + + +def test_manual_mode_clears_inherited_main_model_placement_env(): + env = {name: "inherited" for name in LlamaCppBackend._MANUAL_PLACEMENT_ENV_VARS} + env["LLAMA_ARG_N_GPU_LAYERS_DRAFT"] = "7" + env["UNRELATED"] = "kept" + + LlamaCppBackend._clear_manual_placement_env(env) + + assert not (set(env) & set(LlamaCppBackend._MANUAL_PLACEMENT_ENV_VARS)) + assert env["LLAMA_ARG_N_GPU_LAYERS_DRAFT"] == "7" + assert env["UNRELATED"] == "kept" + + +def test_load_model_sanitizes_manual_env_after_building_child_env(): + src = _load_model_source() + env_build = src.find("env = self._llama_server_env_for_binary(binary)") + env_clear = src.find("self._clear_manual_placement_env(env)", env_build) + launch = src.find("subprocess.Popen", env_build) + assert env_build != -1 + assert env_build < env_clear < launch + + +# ── Manual offload (--gpu-layers + --fit off + --n-cpu-moe) ─────────── + + +def test_load_request_accepts_manual(): + req = LoadRequest( + model_path = "owner/repo", + gpu_memory_mode = "manual", + gpu_layers = 20, + n_cpu_moe = 8, + tensor_split = [2, 1], + ) + assert req.gpu_memory_mode == "manual" + assert req.gpu_layers == 20 + assert req.n_cpu_moe == 8 + assert req.tensor_split == [2, 1] + + +def test_load_request_manual_defaults(): + req = LoadRequest(model_path = "owner/repo") + assert req.gpu_layers == -1 + assert req.n_cpu_moe == 0 + assert req.tensor_split is None + + +@pytest.mark.parametrize("bad", [[0, 0], [-1, 2], [float("inf"), 1], [float("nan"), 1]]) +def test_load_request_rejects_degenerate_tensor_split(bad): + # A negative/non-finite/all-zero split is dropped at launch but compared raw + # in the reload dedupe, so it would reload forever -- reject it up front. + with pytest.raises(ValueError): + LoadRequest(model_path = "owner/repo", tensor_split = bad) + + +@pytest.mark.parametrize("good", [[2, 1], [1, 1], [], None]) +def test_load_request_accepts_valid_tensor_split(good): + assert LoadRequest(model_path = "owner/repo", tensor_split = good).tensor_split == good + + +def test_route_normalizes_explicit_extras_before_reload_dedupe(): + route_src = (Path(_BACKEND_DIR) / "routes" / "inference.py").read_text(encoding = "utf-8") + load_impl = route_src[route_src.index("async def _load_model_impl") :] + strip = load_impl.index("_stripped_explicit = strip_shadowing_flags") + normalize = load_impl.index( + 'request = request.model_copy(update = {"llama_extra_args": extra_llama_args})' + ) + dedupe = load_impl.index("and _request_matches_loaded_settings(") + assert strip < normalize < dedupe + + +@pytest.mark.parametrize("model_cls", [LoadResponse, InferenceStatusResponse]) +def test_response_models_emit_manual_fields(model_cls): + if model_cls is LoadResponse: + obj = model_cls( + status = "loaded", + model = "owner/repo", + display_name = "repo", + inference = {}, + gpu_memory_mode = "manual", + gpu_layers = 20, + n_cpu_moe = 8, + tensor_split = [2, 1], + n_layers = 32, + n_moe_layers = 32, + ) + else: + obj = model_cls( + gpu_memory_mode = "manual", + gpu_layers = 20, + n_cpu_moe = 8, + tensor_split = [2, 1], + n_layers = 32, + n_moe_layers = 32, + ) + dumped = obj.model_dump() + assert dumped["gpu_memory_mode"] == "manual" + assert dumped["gpu_layers"] == 20 + assert dumped["n_cpu_moe"] == 8 + assert dumped["tensor_split"] == [2, 1] + assert dumped["n_layers"] == 32 + assert dumped["n_moe_layers"] == 32 + + +def test_manual_properties_default_and_reflect_and_reset(): + backend = LlamaCppBackend() + assert backend.gpu_layers == -1 and backend.n_cpu_moe == 0 + assert backend.tensor_split is None + backend._gpu_layers = 20 + backend._n_cpu_moe = 8 + backend._tensor_split = [2, 1] + assert backend.gpu_layers == 20 and backend.n_cpu_moe == 8 + assert backend.tensor_split == [2, 1] + backend._process = _FakeProcess() + backend.unload_model() + assert backend.gpu_layers == -1 and backend.n_cpu_moe == 0 + assert backend.tensor_split is None + + +def test_n_moe_layers_property(): + # 0 for a dense model (hides the slider); block_count for all-MoE; + # block_count - leading_dense otherwise (GLM-4.7-Flash: 47 - 1 -> 46). + b = LlamaCppBackend() + b._n_layers = 36 + b._n_experts = None + assert b.n_moe_layers == 0 + b._n_experts = 128 + b._leading_dense_block_count = None + assert b.n_moe_layers == 36 + b._n_layers = 47 + b._leading_dense_block_count = 1 + assert b.n_moe_layers == 46 + + +def _target_state_manual( + backend, + *, + gpu_layers, + n_cpu_moe, + tensor_split = None, +): + return backend._already_in_target_state( + gguf_path = None, + model_identifier = "owner/repo", + hf_variant = "Q4_K_M", + n_ctx = 8192, + cache_type_kv = None, + speculative_type = "auto", + chat_template_override = None, + extra_args = None, + is_vision = False, + gpu_memory_mode = "manual", + gpu_layers = gpu_layers, + n_cpu_moe = n_cpu_moe, + tensor_split = tensor_split, + ) + + +def test_manual_reloads_on_gpu_layers_or_n_cpu_moe_or_split_change(): + backend = _loaded_backend("manual") + backend._gpu_layers = 20 + backend._n_cpu_moe = 0 + backend._tensor_split = None + # Same knobs -> no reload. + assert _target_state_manual(backend, gpu_layers = 20, n_cpu_moe = 0) is True + # Changed layer count -> reload. + assert _target_state_manual(backend, gpu_layers = 16, n_cpu_moe = 0) is False + # Changed MoE offload -> reload. + assert _target_state_manual(backend, gpu_layers = 20, n_cpu_moe = 8) is False + # Added a GPU split -> reload. + assert _target_state_manual(backend, gpu_layers = 20, n_cpu_moe = 0, tensor_split = [2, 1]) is False + # Same GPU split -> no reload. + backend._tensor_split = [2, 1] + assert _target_state_manual(backend, gpu_layers = 20, n_cpu_moe = 0, tensor_split = [2, 1]) is True + + +def test_auto_layers_reload_tracks_only_gpu_layers(): + # Under Auto (gpu_layers < 0) the MoE/split knobs don't apply, so a leftover + # request value must not reload -- only a gpu_layers change (Auto -> pinned) does. + backend = _loaded_backend("manual") + backend._gpu_layers = -1 + backend._n_cpu_moe = 0 + backend._tensor_split = None + # Same Auto, leftover MoE/split in the request -> still no reload. + assert _target_state_manual(backend, gpu_layers = -1, n_cpu_moe = 8, tensor_split = [2, 1]) is True + # Auto -> explicit offload reloads. + assert _target_state_manual(backend, gpu_layers = 20, n_cpu_moe = 0) is False + + +def test_manual_offload_emits_gpu_layers_fit_off_and_n_cpu_moe(): + src = _load_model_source() + gate = src.find('elif gpu_memory_mode == "manual":') + assert gate != -1, "load_model must have an explicit-offload manual branch" + block = src[gate : gate + 700] + # Empties the probed set (skips the planner) but keeps the user's TP choice + # (only the Auto-layers branch above drops TP). + assert "gpus = []" in block + assert "tensor_parallel = False" not in block + # The cmd emits the layer count with fit disabled, gated on gpu_layers >= 0. + assert 'if gpu_memory_mode == "manual" and gpu_layers >= 0:' in src + assert 'cmd.extend(["--gpu-layers", str(gpu_layers), "--fit", "off"])' in src + # MoE offload uses --n-cpu-moe via _resolve_cpu_moe_flag (tested behaviorally below). + assert "_resolve_cpu_moe_flag(" in src + assert 'cmd.extend(["--n-cpu-moe", str(moe_flag)])' in src + # A count requested on a dense model is never emitted, so it must also be + # dropped from the recorded state -- else /status and /load report a count + # llama-server never received (same rule as the tensor-split drop below). + moe_emit = src.find('cmd.extend(["--n-cpu-moe", str(moe_flag)])') + assert "elif n_cpu_moe:" in src[moe_emit : moe_emit + 300] + assert "self._n_cpu_moe = 0" in src[moe_emit : moe_emit + 300] + # The offload path forces use_fit False so --fit-ctx is never added under --fit off. + emit = src.find('cmd.extend(["--gpu-layers", str(gpu_layers), "--fit", "off"])') + assert "use_fit = False" in src[src.rfind("\n", 0, emit) - 200 : emit + 80] + + +def test_status_reports_requested_context_length(): + # The hydration path re-seeds a Manual+Auto context pin from the REQUESTED + # n_ctx (0 = Auto); context_length only exposes the resolved value. + assert "requested_context_length" in InferenceStatusResponse.model_fields + s = InferenceStatusResponse(requested_context_length = 8192) + assert s.model_dump()["requested_context_length"] == 8192 + assert InferenceStatusResponse().model_dump()["requested_context_length"] is None + # The /status route must actually wire it from the backend (a declared-but- + # never-populated field would leave hydration silently reverting the pin). + from pathlib import Path as _P + + route_src = (_P(_BACKEND_DIR) / "routes" / "inference.py").read_text(encoding = "utf-8") + assert "requested_context_length = llama_backend.requested_n_ctx" in route_src + + +def test_manual_offload_emits_tensor_split(): + # The offload path emits --tensor-split from the per-GPU shares, only when + # provided, with >1 GPU in use, AND matching that count (a stale ratio on a + # narrowed picker or a mismatched direct-API list must not emit -- llama- + # server aborts on a split/GPU-count mismatch). + src = _load_model_source() + assert "if tensor_split and _split_gpus > 1:" in src + # Emit only on a length match AND a positive sanitized total: a mismatched + # or all-zero split aborts llama-server / assigns nothing, so it's dropped. + # The emitted list is the sanitized one (clamping tested behaviorally below). + assert "_sanitized_split = self._sanitize_tensor_split(tensor_split)" in src + assert "if len(_sanitized_split) == _split_gpus and _split_total > 0:" in src + assert '"--tensor-split"' in src + # Joined as a comma list (e.g. "2,1") within the explicit-offload cmd branch. + gate = src.find('if gpu_memory_mode == "manual" and gpu_layers >= 0:') + nxt = src.find("elif use_fit:", gate) + assert '","' in src[gate:nxt] and "tensor_split" in src[gate:nxt] + # A split with a single effective GPU is never emitted, so it must also be + # dropped from the recorded state -- else /status and /load report a ratio + # llama-server never received and the dedupe baseline preserves it. + assert "elif tensor_split:" in src[gate:nxt] + drop = src.find("elif tensor_split:", gate, nxt) + assert "self._tensor_split = None" in src[drop : drop + 250] + + +def test_sanitize_tensor_split_clamps_negative_and_non_finite(): + # Negative entries would launch a placement different from the ratio the + # UI showed; inf passes a plain > 0 total gate and would emit + # "--tensor-split inf,..." (llama.cpp normalizes shares by the running + # total, so an inf poisons the shares from that entry on). Both clamp to 0. + sanitize = LlamaCppBackend._sanitize_tensor_split + assert sanitize([2, 1]) == [2.0, 1.0] + assert sanitize([-1, 2]) == [0.0, 2.0] + assert sanitize([float("inf"), 1]) == [0.0, 1.0] + assert sanitize([float("nan"), 1]) == [0.0, 1.0] + # All-zero survives sanitization; the call site's total gate drops it. + assert sanitize([0, 0]) == [0.0, 0.0] + # Unreadable input -> []; the call site's length gate drops it. + assert sanitize(["x", 1]) == [] + assert sanitize([10**400, 1]) == [] + + +def test_zero_offload_mask_honors_device_pin_spellings(): + # A user device pin must keep the GPUs visible: llama-server aborts on a + # pin it can't see ('error: invalid device'). The pin can arrive as + # --device or its -dev alias, as the draft forms (parsed even with no + # drafter loaded), or as an inherited LLAMA_ARG_DEVICE env var. + load_src = _load_model_source() + assert "self._zero_offload_keeps_gpu_visible(cmd, env)" in load_src + block = inspect.getsource(LlamaCppBackend._cmd_has_gpu_device_pin) + for flag in ( + '"--device"', + '"-dev"', + '"--spec-draft-device"', + '"-devd"', + '"--device-draft"', + ): + assert flag in block + assert '"LLAMA_ARG_DEVICE"' in block + + +def test_resolve_cpu_moe_flag(): + # Clamp the requested MoE-layer count to the model's MoE layers, then offset + # past leading dense layers (--n-cpu-moe counts from layer 0). + R = LlamaCppBackend._resolve_cpu_moe_flag + assert R(0, 40, 0) is None # nothing requested + assert R(8, 0, 0) is None # dense model (no MoE layers) + assert R(8, 40, 0) == 8 # all-MoE: direct + assert R(100, 40, 0) == 40 # clamp to the MoE layer count + # GLM-4.7-Flash (deepseek2): block_count 47, leading_dense 1, n_moe 46. + assert R(5, 46, 1) == 6 # offset past the 1 dense layer + assert R(46, 46, 1) == 47 # all MoE on CPU == block_count + + +def test_manual_allows_tensor_parallel_via_split_mode(): + # Manual offload keeps the user's TP choice but skips the memory-based planner + # (plan_tp excludes manual, so its empty gpu set can't downgrade TP). The + # --split-mode tensor emission gates on tensor_parallel alone, so manual + # reaches it -- with tp_tensor_split None it's an even split (no + # --tensor-split). --fit off means no fit/tensor abort. + src = _load_model_source() + assert 'plan_tp = tensor_parallel and gpu_memory_mode != "manual"' in src + assert "if plan_tp:" in src + assert "if plan_tp and len(tp_gpus) < 2:" in src + sm = src.find('cmd.extend(["--split-mode", "tensor"])') + assert sm != -1, "TP must emit --split-mode tensor" + guard = src.rfind("if tensor_parallel:", 0, sm) + assert guard != -1 and sm - guard < 200, "split-mode gates on tensor_parallel" + # The tensor-split is only emitted for a planned (non-even) split, which + # manual never produces, so manual stays an even split. + assert "if tp_tensor_split and len(tp_tensor_split) > 1:" in src + + +def test_fit_sets_target_margin(): + # Manual + Auto (auto_fit) tightens the per-device VRAM margin to 512 MiB. + caps = {"supports_fit_target": True} + flags = LlamaCppBackend._ctx_integrity_flags(1, True, True, 0, 0, caps) + assert flags[flags.index("--fit-target") + 1] == "512" + # Not emitted on the legacy auto path (fit on but not auto_fit): -c 0 pins + # native there, so the tighter margin must not ride along. + assert "--fit-target" not in LlamaCppBackend._ctx_integrity_flags(1, True, False, 0, 0, caps) + # Not emitted when fit is off. + assert "--fit-target" not in LlamaCppBackend._ctx_integrity_flags(1, False, False, 0, 0, caps) + # Not emitted when the binary lacks support. + assert "--fit-target" not in LlamaCppBackend._ctx_integrity_flags( + 1, True, True, 0, 0, {"supports_fit_target": False} + ) + + +# ── GPU picker (gpu_ids -> CUDA_VISIBLE_DEVICES) ───────────────────── + + +def test_load_request_accepts_gpu_ids(): + req = LoadRequest(model_path = "owner/repo", gpu_ids = [1, 0]) + assert req.gpu_ids == [1, 0] + assert LoadRequest(model_path = "owner/repo").gpu_ids is None + + +@pytest.mark.parametrize("model_cls", [LoadResponse, InferenceStatusResponse]) +def test_response_models_emit_gpu_ids(model_cls): + if model_cls is LoadResponse: + obj = model_cls(status = "loaded", model = "m", display_name = "m", inference = {}, gpu_ids = [1]) + else: + obj = model_cls(gpu_ids = [1]) + assert obj.model_dump()["gpu_ids"] == [1] + + +def test_gpu_ids_property_default_and_reset(): + backend = LlamaCppBackend() + assert backend.gpu_ids is None + backend._gpu_ids = [0, 1] + assert backend.gpu_ids == [0, 1] + backend._process = _FakeProcess() + backend.unload_model() + assert backend.gpu_ids is None + + +def _target_state_gpu_ids(backend, gpu_ids): + return backend._already_in_target_state( + gguf_path = None, + model_identifier = "owner/repo", + hf_variant = "Q4_K_M", + n_ctx = 8192, + cache_type_kv = None, + speculative_type = "auto", + chat_template_override = None, + extra_args = None, + is_vision = False, + gpu_ids = gpu_ids, + ) + + +def test_gpu_ids_reload_detection_is_order_insensitive(): + backend = _loaded_backend("auto") + backend._gpu_ids = [0, 1] + # Same set, different order -> no reload. + assert _target_state_gpu_ids(backend, [1, 0]) is True + # Different set -> reload. + assert _target_state_gpu_ids(backend, [0]) is False + # Dropping the pick (auto) -> reload. + assert _target_state_gpu_ids(backend, None) is False + + +def test_gpu_ids_reload_detection_collapses_diffusion_to_single_device(): + # The diffusion runner drives only its single lowest device, so the backend + # records [lowest]. A later multi-GPU request that still resolves to that + # same lowest device must dedupe (no needless reload); a request whose lowest + # device moves, or that drops the pick, must reload. + backend = _loaded_backend("auto") + backend._is_diffusion = True + backend._gpu_ids = [1] # loaded on the lowest of an earlier [3, 1] pick + assert _target_state_gpu_ids(backend, [3, 1]) is True + assert _target_state_gpu_ids(backend, [1]) is True + # Lowest device changes (2, not 1) -> reload. + assert _target_state_gpu_ids(backend, [3, 2]) is False + # Dropping the pick (auto) -> reload. + assert _target_state_gpu_ids(backend, None) is False + + +def test_start_diffusion_server_resets_tensor_parallel(): + # A prior tensor-parallel chat load leaves self._tensor_parallel True (load_model + # phase 1 only kills the process, it skips the unload reset). Diffusion is never + # TP, so startup must clear it -- else /status misreports TP and an identical + # diffusion re-Apply reloads against stale tensor-parallel state. + src = inspect.getsource(llama_cpp_module.LlamaCppBackend._start_diffusion_server) + assert "self._tensor_parallel = False" in src + + +def test_route_matches_loaded_settings_collapses_diffusion_gpu_ids(): + # The route-level reload dedupe mirrors the backend: for a loaded diffusion + # model it compares the request against the single recorded device, not the + # full requested list, or a same-device multi-GPU pick reloads needlessly. + route_src = (Path(_BACKEND_DIR) / "routes" / "inference.py").read_text(encoding = "utf-8") + match_impl = route_src[route_src.index("def _request_matches_loaded_settings") :] + guard = match_impl.index("if llama_backend.is_diffusion:") + collapse = match_impl.index("[sorted(request.gpu_ids)[0]] if request.gpu_ids else None") + compare = match_impl.index("if _req_gpu_ids != llama_backend.gpu_ids:") + assert guard < collapse < compare + + +# ── Manual tensor split: child enumeration pinned to the picker's order ────── + + +def _patch_split_pin_env(monkeypatch, *, inherited, reported): + """Point the pin helper at a fake inherited mask and picker report. + ``reported`` None = enumeration unavailable (falls back to ascending).""" + import utils.hardware as hw + + monkeypatch.setattr( + LlamaCppBackend, "_resolve_visible_physical_ids", staticmethod(lambda: inherited) + ) + info = ( + {"available": False} + if reported is None + else { + "available": True, + "index_kind": "physical", + "devices": [{"index": i} for i in reported], + } + ) + monkeypatch.setattr(hw, "get_backend_visible_gpu_info", lambda: info) + + +def test_split_pin_reorders_inherited_numeric_mask(monkeypatch): + # Parent CUDA_VISIBLE_DEVICES=3,1 makes the child enumerate dev0=phys3, but + # nvidia-smi reported the picker's list ascending -- the mask must be + # re-emitted in that order or the per-GPU shares land on the wrong cards. + _patch_split_pin_env(monkeypatch, inherited = [3, 1], reported = [1, 3]) + env = {"CUDA_VISIBLE_DEVICES": "3,1"} + LlamaCppBackend._pin_visible_gpu_order_for_split(env) + assert env["CUDA_DEVICE_ORDER"] == "PCI_BUS_ID" + assert env["CUDA_VISIBLE_DEVICES"] == "1,3" + + +def test_split_pin_keeps_mask_order_when_picker_reported_it(monkeypatch): + # Torch-fallback enumeration (no nvidia-smi) reports devices in inherited + # mask order, so the picker's split list follows the mask -- the pin must + # keep that order, not re-sort it into a mismatch. + _patch_split_pin_env(monkeypatch, inherited = [3, 1], reported = [3, 1]) + env = {"CUDA_VISIBLE_DEVICES": "3,1"} + LlamaCppBackend._pin_visible_gpu_order_for_split(env) + assert env["CUDA_VISIBLE_DEVICES"] == "3,1" + + +def test_split_pin_falls_back_to_ascending_without_report(monkeypatch): + # Enumeration unavailable: ascending physical is the best guess (it matches + # the dominant nvidia-smi report order). + _patch_split_pin_env(monkeypatch, inherited = [3, 1], reported = None) + env = {"CUDA_VISIBLE_DEVICES": "3,1"} + LlamaCppBackend._pin_visible_gpu_order_for_split(env) + assert env["CUDA_VISIBLE_DEVICES"] == "1,3" + + +def test_split_pin_without_mask_only_sets_pci_order(monkeypatch): + # No inherited mask (or a UUID/MIG one resolving to None): enumeration order + # is fully fixed by CUDA_DEVICE_ORDER, so no mask is written. + _patch_split_pin_env(monkeypatch, inherited = None, reported = None) + env = {} + LlamaCppBackend._pin_visible_gpu_order_for_split(env) + assert env == {"CUDA_DEVICE_ORDER": "PCI_BUS_ID"} + + +def test_split_pin_mirrors_hip_mask_on_rocm(monkeypatch): + # ROCm: the pin must land in HIP_VISIBLE_DEVICES too, and an inherited ROCR + # mask is cleared so the mask can't apply twice (ROCR re-indexes, then HIP + # would index into the already-reduced set). + _patch_split_pin_env(monkeypatch, inherited = [3, 1], reported = [1, 3]) + torch_stub = _types.ModuleType("torch") + torch_stub.version = _types.SimpleNamespace(hip = "6.0") + monkeypatch.setitem(sys.modules, "torch", torch_stub) + env = {"CUDA_VISIBLE_DEVICES": "3,1", "ROCR_VISIBLE_DEVICES": "3,1"} + LlamaCppBackend._pin_visible_gpu_order_for_split(env) + assert env["CUDA_VISIBLE_DEVICES"] == "1,3" + assert env["HIP_VISIBLE_DEVICES"] == "1,3" + assert "ROCR_VISIBLE_DEVICES" not in env + + +# ── Diffusion single-device selection ─────────────────────────────────────── + + +def test_diffusion_gpu_arg_uses_lowest_explicit_physical_id(monkeypatch): + monkeypatch.setenv("CUDA_VISIBLE_DEVICES", "3,1") + monkeypatch.setenv("DG_GPU", "7") + assert LlamaCppBackend._diffusion_gpu_arg([3, 1]) == "1" + + +def test_diffusion_gpu_arg_preserves_parent_mask_order(monkeypatch): + monkeypatch.delenv("DG_GPU", raising = False) + monkeypatch.setenv("CUDA_VISIBLE_DEVICES", "3,1") + assert LlamaCppBackend._diffusion_gpu_arg(None) == "3" + + +def test_diffusion_gpu_arg_honors_override_and_cpu_mask(monkeypatch): + monkeypatch.setenv("DG_GPU", "GPU-abc") + assert LlamaCppBackend._diffusion_gpu_arg(None) == "GPU-abc" + assert LlamaCppBackend._diffusion_gpu_arg(None, cpu_only = True) == "" + + +# ── Deliberate zero-offload (manual gpu_layers=0): training-skip flag ───────── + + +def test_zero_offload_flag_false_without_companions(): + # CPU-only by construction: False lets training skip unloading a server that + # holds no VRAM. + cmd = ["llama-server", "-m", "model.gguf", "--gpu-layers", "0", "--fit", "off"] + assert LlamaCppBackend._zero_offload_gpu_flag(cmd, [(0, 8000, 24000)], {}) is False + + +@pytest.mark.parametrize( + "companion", + ["--mmproj", "--model-draft", "-md", "--spec-draft-model", "-hfd"], +) +def test_zero_offload_flag_true_with_companion(companion): + # mmproj / a drafter offload to GPU regardless of --gpu-layers, so the + # server still holds VRAM and training must unload it. Drafter detection + # reuses the extras parser, so pass-through aliases count too. + cmd = ["llama-server", "-m", "model.gguf", "--gpu-layers", "0", companion, "x.gguf"] + assert LlamaCppBackend._zero_offload_gpu_flag(cmd, [(0, 8000, 24000)], {}) is True + + +def test_zero_offload_flag_true_with_inline_companion_forms(): + cmd = ["llama-server", "-m", "model.gguf", "--spec-draft-model=x.gguf"] + assert LlamaCppBackend._zero_offload_gpu_flag(cmd, [(0, 8000, 24000)], {}) is True + cmd = ["llama-server", "-m", "model.gguf", "--mmproj=proj.gguf"] + assert LlamaCppBackend._zero_offload_gpu_flag(cmd, [(0, 8000, 24000)], {}) is True + + +def test_zero_offload_flag_true_with_env_drafter(): + cmd = ["llama-server", "-m", "model.gguf", "--gpu-layers", "0"] + env = {"LLAMA_ARG_SPEC_DRAFT_MODEL": "x.gguf"} + assert LlamaCppBackend._zero_offload_gpu_flag(cmd, [(0, 8000, 24000)], env) is True + + +@pytest.mark.parametrize( + "device_args", + [ + ["--device", "CUDA0"], + ["--device=CUDA0"], + ["-dev", "CUDA0"], + ["--spec-draft-device", "CUDA0"], + ["--device-draft=CUDA0"], + ], +) +def test_zero_offload_flag_true_with_device_pin(device_args): + cmd = ["llama-server", "-m", "model.gguf", "--gpu-layers", "0", *device_args] + assert LlamaCppBackend._zero_offload_gpu_flag(cmd, [(0, 8000, 24000)], {}) is True + + +def test_zero_offload_flag_true_with_env_device_pin(): + cmd = ["llama-server", "-m", "model.gguf", "--gpu-layers", "0"] + env = {"LLAMA_ARG_DEVICE": "CUDA0"} + assert LlamaCppBackend._zero_offload_gpu_flag(cmd, [(0, 8000, 24000)], env) is True + + +@pytest.mark.parametrize( + ("device_args", "env"), + [ + (["--device", "cpu"], {}), + (["--device=none"], {}), + (["--spec-draft-device", "cpu"], {}), + ([], {"LLAMA_ARG_DEVICE": "none"}), + (["--device", "CUDA0", "--device", "cpu"], {}), + ], +) +def test_zero_offload_flag_false_with_cpu_device_pin(device_args, env): + cmd = ["llama-server", "-m", "model.gguf", "--gpu-layers", "0", *device_args] + assert LlamaCppBackend._zero_offload_gpu_flag(cmd, [(0, 8000, 24000)], env) is False + + +def test_zero_offload_flag_true_with_surviving_tensor_mode(): + cmd = ["llama-server", "-m", "model.gguf", "--gpu-layers", "0", "--split-mode", "tensor"] + assert LlamaCppBackend._zero_offload_gpu_flag(cmd, [(0, 8000, 24000)], {}) is True + + +def test_zero_offload_flag_true_for_unmasked_vulkan(monkeypatch): + monkeypatch.setattr(LlamaCppBackend, "_is_vulkan_backend", staticmethod(lambda: True)) + cmd = ["llama-server", "-m", "model.gguf", "--gpu-layers", "0"] + assert LlamaCppBackend._zero_offload_gpu_flag(cmd, [(0, 8000, 24000)], {}) is True + + +def test_zero_offload_flag_none_without_gpus(): + cmd = ["llama-server", "-m", "model.gguf", "--gpu-layers", "0"] + assert LlamaCppBackend._zero_offload_gpu_flag(cmd, [], {}) is None + + +def test_cmd_has_gpu_companion_detection(): + # The env mask for CPU-only zero-offload loads keys off this scan: any + # --mmproj form or a drafter (flag aliases / env) keeps the GPUs visible. + has = LlamaCppBackend._cmd_has_gpu_companion + assert has(["llama-server", "-m", "m.gguf"], {}) is False + assert has(["llama-server", "--mmproj", "p.gguf"], {}) is True + assert has(["llama-server", "--mmproj=p.gguf"], {}) is True + assert has(["llama-server", "-md", "d.gguf"], {}) is True + assert has(["llama-server"], {"LLAMA_ARG_SPEC_DRAFT_MODEL": "d.gguf"}) is True + + +def test_cmd_companion_ignores_cpu_forced_drafter(): + # A CPU-pinned drafter holds no VRAM: the zero-offload mask may hide the GPUs + # and training may leave the server alone. + has = LlamaCppBackend._cmd_has_gpu_companion + cmd = ["llama-server", "-md", "d.gguf", "--spec-draft-ngl", "0"] + assert has(cmd, {}) is False + cmd = ["llama-server", "-md", "d.gguf", "--spec-draft-device", "cpu"] + assert has(cmd, {}) is False + # mmproj still counts even alongside a CPU drafter. + cmd = ["llama-server", "-md", "d.gguf", "--spec-draft-ngl", "0", "--mmproj", "p.gguf"] + assert has(cmd, {}) is True diff --git a/studio/backend/tests/test_gpu_selection.py b/studio/backend/tests/test_gpu_selection.py index 69ad560788..d4f2fbe993 100644 --- a/studio/backend/tests/test_gpu_selection.py +++ b/studio/backend/tests/test_gpu_selection.py @@ -853,7 +853,13 @@ class TestRouteErrors(unittest.TestCase): self.assertIn("only supported on CUDA devices", str(exc_info.exception)) - def test_inference_route_rejects_gpu_ids_for_gguf(self): + def test_inference_route_validates_gpu_ids_for_gguf(self): + # gpu_ids is now SUPPORTED for GGUF (the GPU picker), but still + # validated: a rejected pick surfaces as a clean 400, not the old + # "not supported for GGUF" rejection. Patch the validator so the test + # is deterministic regardless of the host's (or a prior test's) GPU env. + import utils.hardware.hardware as hardware_mod + inference_route = _load_route_module( "inference_route_module_for_gguf_gpu_ids_test", "routes/inference.py", @@ -887,6 +893,11 @@ class TestRouteErrors(unittest.TestCase): ), patch.object(inference_route.asyncio, "to_thread", new = _inline_to_thread), patch.object(inference_route, "_hf_offline_if_dns_dead", nullcontext), + patch.object( + hardware_mod, + "resolve_requested_gpu_ids", + side_effect = ValueError("Invalid gpu_ids [0, 1]: rejected by test"), + ), ): with self.assertRaises(HTTPException) as exc_info: asyncio.run( @@ -901,8 +912,11 @@ class TestRouteErrors(unittest.TestCase): ) ) + # The validator's ValueError becomes a clean 400 (not the removed + # "not supported for GGUF" rejection). self.assertEqual(exc_info.exception.status_code, 400) - self.assertIn("GGUF", exc_info.exception.detail) + self.assertIn("gpu_ids", exc_info.exception.detail.lower()) + self.assertNotIn("not supported", exc_info.exception.detail.lower()) def test_training_route_returns_400_for_invalid_gpu_ids(self): training_route = _load_route_module( diff --git a/studio/backend/tests/test_hf_token_validation.py b/studio/backend/tests/test_hf_token_validation.py new file mode 100644 index 0000000000..31b30fc37d --- /dev/null +++ b/studio/backend/tests/test_hf_token_validation.py @@ -0,0 +1,165 @@ +# SPDX-License-Identifier: AGPL-3.0-only +# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0 + +"""Focused coverage for cached, rate-limited HF token validation.""" + +from __future__ import annotations + +from pathlib import Path +import sys + +import httpx +import pytest + + +_BACKEND_DIR = str(Path(__file__).resolve().parent.parent) +if _BACKEND_DIR not in sys.path: + sys.path.insert(0, _BACKEND_DIR) + +import utils.hf_token_validation as validation + + +@pytest.fixture(autouse = True) +def _reset_validation_state(): + validation.reset_hf_token_validation_state() + yield + validation.reset_hf_token_validation_state() + + +def test_cached_token_does_not_spend_another_attempt(monkeypatch): + calls = [] + + def _check(token): + calls.append(token) + return validation.TokenValidationResult(status = "valid") + + monkeypatch.setattr(validation, "_check_remote", _check) + first = validation.validate_hf_token("hf_valid", rate_key = "user:ip") + second = validation.validate_hf_token("hf_valid", rate_key = "user:ip") + + assert first.status == second.status == "valid" + assert calls == ["hf_valid"] + + +def test_three_uncached_attempts_per_hour(monkeypatch): + monkeypatch.setattr( + validation, + "_check_remote", + lambda _token: validation.TokenValidationResult(status = "invalid"), + ) + + for index in range(3): + result = validation.validate_hf_token(f"hf_bad_{index}", rate_key = "user:ip") + assert result.status == "invalid" + + limited = validation.validate_hf_token("hf_bad_4", rate_key = "user:ip") + assert limited.status == "rate_limited" + assert limited.retry_after_seconds is not None + assert limited.retry_after_seconds > 0 + + other_user = validation.validate_hf_token("hf_other", rate_key = "other:ip") + assert other_user.status == "invalid" + + +def test_window_rolls_forward(monkeypatch): + clock = {"now": 100.0} + monkeypatch.setattr(validation.time, "monotonic", lambda: clock["now"]) + monkeypatch.setattr(validation, "_MAX_ATTEMPTS", 1) + monkeypatch.setattr(validation, "_WINDOW_SECONDS", 10.0) + monkeypatch.setattr( + validation, + "_check_remote", + lambda _token: validation.TokenValidationResult(status = "invalid"), + ) + + assert validation.validate_hf_token("hf_a", rate_key = "user:ip").status == "invalid" + assert validation.validate_hf_token("hf_b", rate_key = "user:ip").status == "rate_limited" + clock["now"] += 11.0 + assert validation.validate_hf_token("hf_b", rate_key = "user:ip").status == "invalid" + + +@pytest.mark.parametrize( + ("status_code", "expected"), + [(200, "valid"), (401, "invalid"), (429, "rate_limited"), (500, "unavailable")], +) +def test_remote_status_classification(monkeypatch, status_code, expected): + response = httpx.Response( + status_code, + request = httpx.Request("GET", "https://huggingface.co/api/whoami-v2"), + headers = {"Retry-After": "42"} if status_code == 429 else None, + ) + + class _Session: + def get(self, url, *, headers, timeout): + assert url == "https://huggingface.co/api/whoami-v2" + assert headers["authorization"] == "Bearer hf_test" + assert timeout == validation._REMOTE_TIMEOUT_SECONDS + return response + + monkeypatch.setattr(validation, "get_session", lambda: _Session()) + result = validation._check_remote("hf_test") + assert result.status == expected + if status_code == 429: + assert result.retry_after_seconds == 42 + + +def test_wrapped_http_401_is_invalid(monkeypatch): + response = httpx.Response( + 401, + request = httpx.Request("GET", "https://huggingface.co/api/whoami-v2"), + ) + + class _Session: + def get(self, _url, **_kwargs): + error = RuntimeError("Invalid user token.") + error.response = response + raise error + + monkeypatch.setattr(validation, "get_session", lambda: _Session()) + assert validation._check_remote("hf_test").status == "invalid" + + +def test_remote_timeout_is_bounded_and_unavailable(monkeypatch): + class _Session: + def get(self, _url, *, headers, timeout): + assert headers["authorization"] == "Bearer hf_test" + assert timeout == validation._REMOTE_TIMEOUT_SECONDS + raise TimeoutError("timed out") + + monkeypatch.setattr(validation, "get_session", lambda: _Session()) + assert validation._check_remote("hf_test").status == "unavailable" + + +def test_raw_token_is_not_retained(monkeypatch): + monkeypatch.setattr( + validation, + "_check_remote", + lambda _token: validation.TokenValidationResult(status = "valid"), + ) + token = "hf_do_not_store_this_value" + validation.validate_hf_token(token, rate_key = "user:ip") + + assert token not in repr(validation._cache) + assert token not in repr(validation._attempts) + + +def test_unexpected_remote_exception_releases_singleflight(monkeypatch): + calls = 0 + monkeypatch.setattr(validation, "_INFLIGHT_WAIT_SECONDS", 0.0) + + def _check(_token): + nonlocal calls + calls += 1 + if calls == 1: + raise RuntimeError("unexpected failure") + return validation.TokenValidationResult(status = "valid") + + monkeypatch.setattr(validation, "_check_remote", _check) + + with pytest.raises(RuntimeError, match = "unexpected failure"): + validation.validate_hf_token("hf_test", rate_key = "user:ip") + + result = validation.validate_hf_token("hf_test", rate_key = "user:ip") + assert result.status == "valid" + assert calls == 2 + assert validation._inflight == {} diff --git a/studio/backend/tests/test_hf_xet_fallback.py b/studio/backend/tests/test_hf_xet_fallback.py index 2fff744b64..48aff29659 100644 --- a/studio/backend/tests/test_hf_xet_fallback.py +++ b/studio/backend/tests/test_hf_xet_fallback.py @@ -1,10 +1,10 @@ # SPDX-License-Identifier: AGPL-3.0-only # Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0 -"""Tests for the Studio shim over the shared unsloth_zoo Xet -> HTTP fallback. +"""Tests for the Unsloth shim over the shared unsloth_zoo Xet -> HTTP fallback. The transport-policy matrix is tested once in unsloth_zoo; here we assert only the -Studio seam: re-exporting the shared API and injecting the marker-aware +Unsloth seam: re-exporting the shared API and injecting the marker-aware prepare_cache_for_transport on the HTTP retry. CPU-only, no network, no real subprocess. """ @@ -69,7 +69,7 @@ def test_child_should_disable_xet_truth_table(): def test_shim_injects_studio_prepare_on_http_retry(monkeypatch): - """A Xet stall retries over HTTP and the shim runs Studio's marker-aware + """A Xet stall retries over HTTP and the shim runs Unsloth's marker-aware ``prepare_cache_for_transport(..., 'http')`` before the retry.""" _requires_shared() for var in ("UNSLOTH_DISABLE_XET", "UNSLOTH_STABLE_DOWNLOADS", "HF_HUB_DISABLE_XET"): @@ -107,11 +107,11 @@ def test_shim_injects_studio_prepare_on_http_retry(monkeypatch): out = xf.hf_hub_download_with_xet_fallback(DL_REPO, FILE, None) assert out == "/cache/model.gguf" assert seen_disable_xet == [False, True] # Xet first, then HTTP - assert prepared == [("model", DL_REPO, "http")], "shim must run Studio's marker-aware prep" + assert prepared == [("model", DL_REPO, "http")], "shim must run Unsloth's marker-aware prep" def test_shim_snapshot_injects_studio_prepare(monkeypatch): - """The snapshot wrapper forwards Studio's marker-aware prep, like the file wrapper.""" + """The snapshot wrapper forwards Unsloth's marker-aware prep, like the file wrapper.""" captured = {} def fake_snapshot(repo_id, **kwargs): @@ -127,7 +127,7 @@ def test_shim_snapshot_injects_studio_prepare(monkeypatch): def test_degrades_gracefully_without_shared_helper(monkeypatch): - """On an older unsloth_zoo lacking the shared helper, the shim still imports (Studio + """On an older unsloth_zoo lacking the shared helper, the shim still imports (Unsloth boots) and exposes stub API doing plain HF downloads with the watchdog disabled.""" import importlib @@ -206,7 +206,7 @@ def test_degrades_gracefully_without_shared_helper(monkeypatch): def test_degrades_when_unsloth_zoo_entirely_absent(): """When unsloth_zoo is absent entirely, the import raises ModuleNotFoundError(name='unsloth_zoo') (top-level package). Guard that the shim still - degrades and does not re-raise, breaking every Studio import that pulls it in.""" + degrades and does not re-raise, breaking every Unsloth import that pulls it in.""" import importlib class _BlockZoo: @@ -248,7 +248,7 @@ def test_degrades_when_unsloth_zoo_entirely_absent(): def test_degrades_when_shared_helper_import_raises_importerror(): """unsloth_zoo can be installed yet fail to import when torch is missing (llama.cpp/GGUF-only - Studio), raising ImportError not ModuleNotFoundError. The shim must degrade for that too.""" + Unsloth), raising ImportError not ModuleNotFoundError. The shim must degrade for that too.""" import importlib class _BlockWithImportError: @@ -329,7 +329,7 @@ def test_retries_under_light_gpu_init_when_import_fails(monkeypatch): # with it set); accessing DownloadStallError drives it via __getattr__. stall_error = degraded.DownloadStallError assert seen_env == [None, "1"], seen_env - # Both attempts raised -> Studio still boots in degraded mode. + # Both attempts raised -> Unsloth still boots in degraded mode. assert issubclass(stall_error, RuntimeError) # The env override must not leak past the load. assert os.environ.get("UNSLOTH_ZOO_DISABLE_GPU_INIT") is None diff --git a/studio/backend/tests/test_identity.py b/studio/backend/tests/test_identity.py index 1e84ddef35..712348f7ca 100644 --- a/studio/backend/tests/test_identity.py +++ b/studio/backend/tests/test_identity.py @@ -3,7 +3,7 @@ """Tests for the server identity handshake (`GET /api/auth/identity`). -The endpoint lets a client confirm an endpoint is really this Studio install +The endpoint lets a client confirm an endpoint is really this Unsloth install before sending it a credential: the client sends a random nonce and checks the returned HMAC against one computed from the install identity secret. A process that cannot read this same-user secret cannot forge a matching proof. diff --git a/studio/backend/tests/test_index_bootstrap_origin_extra.py b/studio/backend/tests/test_index_bootstrap_origin_extra.py index feda88c14c..e1c52a653e 100644 --- a/studio/backend/tests/test_index_bootstrap_origin_extra.py +++ b/studio/backend/tests/test_index_bootstrap_origin_extra.py @@ -26,7 +26,7 @@ def _build_request( def test_is_same_origin_request_ipv6_loopback_same_origin(): - """Studio supports ``-H ::1`` binds; netloc is ``[::1]:8902``. Bare + """Unsloth supports ``-H ::1`` binds; netloc is ``[::1]:8902``. Bare ``partition(":")`` mis-parses the bracketed form and would refuse the bootstrap on legitimate same-origin navigation. """ diff --git a/studio/backend/tests/test_install_resolve_prebuilt.py b/studio/backend/tests/test_install_resolve_prebuilt.py index e97ca47717..3ebad861ad 100644 --- a/studio/backend/tests/test_install_resolve_prebuilt.py +++ b/studio/backend/tests/test_install_resolve_prebuilt.py @@ -445,6 +445,101 @@ def test_route_to_vulkan_prebuilt_cpu_fallback_wins(): assert routed is host +@pytest.mark.parametrize("cpu_flag", ["--cpu-fallback", "--force-cpu"]) +def test_resolve_prebuilt_cpu_fallback_overrides_intel_vulkan(monkeypatch, capsys, cpu_flag): + """Either CPU flag via CLI must suppress Vulkan even on an Intel GPU host: both + drop GPU detection (--force-cpu additionally persists, on the install path).""" + monkeypatch.setattr( + ilp, + "detect_host", + lambda: _host(is_linux = True, is_x86_64 = True, has_intel_gpu = True), + ) + seen = {} + + def _resolver(tag, host, repo, published_release_tag): + seen["host"] = host + seen["repo"] = repo + raise ilp.PrebuiltFallback("no asset") + + monkeypatch.setattr(ilp, "resolve_simple_install_release_plans", _resolver) + monkeypatch.setattr( + sys, + "argv", + [ + "install_llama_prebuilt.py", + "--resolve-prebuilt", + "latest", + cpu_flag, + "--output-format", + "json", + ], + ) + assert ilp.main() == ilp.EXIT_SUCCESS + # The CPU flag must suppress Intel GPU, route to fork (not upstream Vulkan) + assert seen["host"].has_intel_gpu is False + assert seen["repo"] == FORK + + +@pytest.mark.parametrize( + "flags, expect_force, expect_persist", + [ + ([], False, False), + # Automatic/transient last resort (arm64 GPU-build recovery): drops GPU but + # does NOT persist, so a later update heals to a GPU bundle (#6097). + (["--cpu-fallback"], True, False), + # Deliberate CPU-only (UNSLOTH_LLAMA_CPP_BACKEND=cpu): drops GPU AND persists so + # the updater re-asserts it and never revives the Intel iGPU crash (#7213). + (["--force-cpu"], True, True), + (["--cpu-fallback", "--force-cpu"], True, True), + ], +) +def test_cli_cpu_flags_thread_force_and_persist( + monkeypatch, tmp_path, flags, expect_force, expect_persist +): + captured = {} + monkeypatch.setattr(ilp, "install_prebuilt", lambda **kw: captured.update(kw)) + monkeypatch.setattr( + sys, + "argv", + ["install_llama_prebuilt.py", "--install-dir", str(tmp_path / "llama.cpp"), *flags], + ) + assert ilp.main() == ilp.EXIT_SUCCESS + assert captured["force_cpu"] is expect_force + assert captured["persist_force_cpu"] is expect_persist + + +@pytest.mark.parametrize( + "existing, requested, expected", + [ + # A deliberate --force-cpu on top of a naturally-installed CPU bundle (same + # asset, install skipped) must still flip the marker to true (#7213). + (False, True, True), + (None, True, True), + # No spurious writes when already in sync, and a released force syncs down. + (True, True, True), + (False, False, False), + (True, False, False), + ], +) +def test_sync_marker_force_cpu(tmp_path, existing, requested, expected): + marker = {"tag": "b9585", "asset": "llama-b9585-bin-ubuntu-x64.tar.gz"} + if existing is not None: + marker["force_cpu"] = existing + marker_path = tmp_path / "UNSLOTH_PREBUILT_INFO.json" + marker_path.write_text(json.dumps(marker)) + ilp.sync_marker_force_cpu(tmp_path, requested) + written = json.loads(marker_path.read_text()) + assert written["force_cpu"] is expected + # Unrelated fields are preserved. + assert written["asset"] == "llama-b9585-bin-ubuntu-x64.tar.gz" + + +def test_sync_marker_force_cpu_missing_marker_is_noop(tmp_path): + # No marker (or unreadable) must not crash the reuse path. + ilp.sync_marker_force_cpu(tmp_path, True) + assert not (tmp_path / "UNSLOTH_PREBUILT_INFO.json").exists() + + def test_route_to_vulkan_prebuilt_hidden_nvidia_not_rerouted(): # A mixed NVIDIA+Intel host that hid NVIDIA (CUDA_VISIBLE_DEVICES=""/-1): # physical NVIDIA present but not usable. Must NOT auto-route to Vulkan, or diff --git a/studio/backend/tests/test_llama_cpp_context_fit.py b/studio/backend/tests/test_llama_cpp_context_fit.py index d3a10df8ca..2a4f6d19d2 100644 --- a/studio/backend/tests/test_llama_cpp_context_fit.py +++ b/studio/backend/tests/test_llama_cpp_context_fit.py @@ -567,7 +567,7 @@ class TestClassifyGpuOffload: assert inst._classify_gpu_offload(False, []) is None def test_user_did_not_intend_gpu_returns_none(self): - # Studio called start_llama_server without expecting GPU; don't warn. + # Unsloth called start_llama_server without expecting GPU; don't warn. inst = self._backend( [ "load_tensors: CPU_Mapped model buffer size = 21000.0 MiB", diff --git a/studio/backend/tests/test_llama_cpp_mmproj_fallback.py b/studio/backend/tests/test_llama_cpp_mmproj_fallback.py index 04d4aac9e1..049058e511 100644 --- a/studio/backend/tests/test_llama_cpp_mmproj_fallback.py +++ b/studio/backend/tests/test_llama_cpp_mmproj_fallback.py @@ -222,7 +222,7 @@ class TestFlashAttnOff: assert _flash_off(["llama-server", "-fa=on"]) == ["llama-server", "-fa=off"] def test_flips_every_occurrence_last_wins(self): - # extra_args can re-enable FA after Studio's flag; llama.cpp is last-wins, + # extra_args can re-enable FA after Unsloth's flag; llama.cpp is last-wins, # so one leftover 'on' would re-crash the retry. Every enable must flip. cmd = ["llama-server", "--flash-attn", "on", "--mmproj", "/p", "--flash-attn", "on"] out = _flash_off(cmd) @@ -234,7 +234,7 @@ class TestFlashAttnOff: assert _flash_off(["llama-server", "--flash-attn=off"]) is None def test_none_when_user_off_wins_last(self): - # User appended 'off' after Studio's 'on'; effective (last-wins) is off, + # User appended 'off' after Unsloth's 'on'; effective (last-wins) is off, # so there is nothing to retry. assert _flash_off(["llama-server", "--flash-attn", "on", "--flash-attn", "off"]) is None diff --git a/studio/backend/tests/test_llama_cpp_mtp_detection.py b/studio/backend/tests/test_llama_cpp_mtp_detection.py index 3f9d2a8f50..1d15647967 100644 --- a/studio/backend/tests/test_llama_cpp_mtp_detection.py +++ b/studio/backend/tests/test_llama_cpp_mtp_detection.py @@ -9,6 +9,7 @@ the _already_in_target_state mirror that prevents needless reloads. from __future__ import annotations +import ast import inspect import os import struct @@ -345,10 +346,62 @@ def test_windows_full_offload_flags_use_current_llama_server_args(): stale_checkpoint_flag = "--checkpoint-" + "every-n-tokens" assert '"--cache-ram"' in src assert '"--ctx-checkpoints"' in src - assert '"--no-cache-prompt"' in src + # Prompt caching stays on (in-VRAM prefix reuse); #5692 only needed the host-RAM + # checkpoints (--cache-ram / --ctx-checkpoints) disabled, not prompt reuse. + assert '"--no-cache-prompt"' not in src assert stale_checkpoint_flag not in src +# Backend-wide guard: Unsloth must never inject --no-cache-prompt into a llama-server +# command. It disables in-VRAM prompt-prefix reuse, re-prefilling every repeated prompt +# (#5692 only needed --cache-ram / --ctx-checkpoints off; #7260 dropped the stray flag). +# Detecting it (_is_real) or honouring a user-supplied one (_prompt_cache_off) is fine. +_NO_CACHE_PROMPT_FLAG = "--no-cache-prompt" +_LIST_MUTATORS = frozenset({"append", "extend", "insert"}) + + +def _has_flag_literal(node: ast.AST) -> bool: + return any( + isinstance(n, ast.Constant) and n.value == _NO_CACHE_PROMPT_FLAG for n in ast.walk(node) + ) + + +def _no_cache_prompt_injections(source: str, filename: str) -> list[tuple[str, int]]: + """(file, lineno) for each spot adding --no-cache-prompt to a list.""" + hits: list[tuple[str, int]] = [] + for node in ast.walk(ast.parse(source, filename = filename)): + # cmd.append/extend/insert(... flag ...) or cmd += [... flag ...] + if ( + isinstance(node, ast.Call) + and isinstance(node.func, ast.Attribute) + and node.func.attr in _LIST_MUTATORS + and any(_has_flag_literal(a) for a in node.args) + ) or ( + isinstance(node, ast.AugAssign) + and isinstance(node.op, ast.Add) + and _has_flag_literal(node.value) + ): + hits.append((filename, node.lineno)) + return hits + + +def test_unsloth_never_injects_no_cache_prompt_into_any_command(): + root = Path(_BACKEND_DIR) + files = [p for p in root.rglob("*.py") if "tests" not in p.relative_to(root).parts] + violations: list[tuple[str, int]] = [] + for path in files: + try: + violations += _no_cache_prompt_injections(path.read_text(encoding = "utf-8"), str(path)) + except (OSError, UnicodeDecodeError, SyntaxError): + continue + assert files, "no backend source files were scanned" + assert violations == [], ( + "Unsloth must never add --no-cache-prompt to a llama-server command " + "(it disables prompt-prefix reuse); detecting or honouring a user-supplied " + f"one is fine. Offending sites: {violations}" + ) + + def test_load_model_sets_threads_once(): src = inspect.getsource(LlamaCppBackend.load_model) assert src.count('cmd.extend(["--threads", str(') == 1 @@ -741,6 +794,25 @@ def test_probe_reports_windows_cache_flags_absent_for_older_binary(tmp_path): assert caps["supports_no_cache_prompt"] is False +@_NEEDS_BASH +def test_probe_detects_slot_save_path(tmp_path): + fake = _make_fake_llama_server( + tmp_path / "llama-server", + "--slot-save-path PATH path to save slot kv cache\n--threads N\n", + ) + _clear_caps_cache() + caps = LlamaCppBackend.probe_server_capabilities(str(fake)) + assert caps["supports_slot_save"] is True + + +@_NEEDS_BASH +def test_probe_reports_slot_save_absent_for_older_binary(tmp_path): + fake = _make_fake_llama_server(tmp_path / "llama-server", "--threads N\n") + _clear_caps_cache() + caps = LlamaCppBackend.probe_server_capabilities(str(fake)) + assert caps["supports_slot_save"] is False + + def test_build_ngram_mod_flags_new(): flags = _build_ngram_mod_flags({"ngram_mod_flavor": "new"}) assert flags == [ @@ -1014,7 +1086,7 @@ def test_already_in_target_state_2b_falls_back_to_ngram_below_threshold(monkeypa ) -# usage backfill from timings (Studio UI t/s widget fix). +# usage backfill from timings (Unsloth UI t/s widget fix). def test_backfill_usage_from_timings_fills_when_completion_tokens_zero(): @@ -1606,7 +1678,7 @@ def test_reload_forced_mtp_bounces_auto_mla(): ) -# ── Full named-repo resolver matrix (the shipping Studio families) ───── +# ── Full named-repo resolver matrix (the shipping Unsloth families) ───── # # Locks auto / off / forced-mtp routing for every Qwen3.5 (MTP + plain) and # gemma-4 (regular + QAT) GGUF repo, including the giant MoEs that stay diff --git a/studio/backend/tests/test_llama_cpp_no_context_shift.py b/studio/backend/tests/test_llama_cpp_no_context_shift.py index 10b1dc7ff6..662c918305 100644 --- a/studio/backend/tests/test_llama_cpp_no_context_shift.py +++ b/studio/backend/tests/test_llama_cpp_no_context_shift.py @@ -5,7 +5,7 @@ With llama-server's default context-shift behavior, the UI cannot tell the user the KV cache was rotated -- earlier turns silently vanish from the conversation. -The Studio backend always passes ``--no-context-shift`` so the server returns a +The Unsloth backend always passes ``--no-context-shift`` so the server returns a clean error instead, and the chat adapter can point the user at the ``Context Length`` input in the settings panel. @@ -118,9 +118,17 @@ def test_flag_sits_inside_the_base_cmd_list(): "conditional branch -- otherwise some code paths would still " "run with silent context shift enabled." ) - # Pin that it sits next to -c / --ctx so the grouping makes sense. - assert '"-c"' in block assert '"--flash-attn"' in block + # -c is emitted in the conditional right after the base list, not inside + # it: auto-fit (--fit on with no pinned context) must omit -c entirely, + # because "-c 0" pins the full native context and disables --fit's + # VRAM-based sizing. Pin that it still sits next to the base block so the + # context grouping stays intact. + after = rest[end_rel : end_rel + 1000] + assert '"-c"' in after, ( + "-c must still be emitted in the conditional immediately after the " + "base cmd list (omitted only in auto-fit, where --fit sizes context)." + ) def _iter_lines_with_offset(text: str): diff --git a/studio/backend/tests/test_llama_cpp_props_readback.py b/studio/backend/tests/test_llama_cpp_props_readback.py index 316956325f..fe1e67edad 100644 --- a/studio/backend/tests/test_llama_cpp_props_readback.py +++ b/studio/backend/tests/test_llama_cpp_props_readback.py @@ -4,7 +4,7 @@ """Tests for the post-launch /props context readback. llama-server's memory-fit step or --parallel slot split can allocate less -context than the requested -c while Studio keeps advertising the requested +context than the requested -c while Unsloth keeps advertising the requested value; clients sized to it then die on exceed_context_size_error 400s. ``_reconcile_effective_ctx_with_server`` must adopt the server's real ``default_generation_settings.n_ctx`` whenever it is smaller. @@ -223,33 +223,48 @@ _CAPS_NONE = {"supports_kv_unified": False, "supports_fit_ctx": False} def test_kv_unified_added_for_multi_slot(): """Explicit --parallel N disables llama-server's auto-slots kv-unified - default, splitting -c into per-slot windows of -c/N; Studio must restore + default, splitting -c into per-slot windows of -c/N; Unsloth must restore the shared pool so one request can use the full advertised context.""" - flags = LlamaCppBackend._ctx_integrity_flags(4, False, 98304, 98304, _CAPS_ALL) + flags = LlamaCppBackend._ctx_integrity_flags(4, False, False, 98304, 98304, _CAPS_ALL) assert "--kv-unified" in flags def test_kv_unified_skipped_for_single_slot_or_old_build(): assert "--kv-unified" not in LlamaCppBackend._ctx_integrity_flags( - 1, False, 98304, 98304, _CAPS_ALL + 1, False, False, 98304, 98304, _CAPS_ALL ) assert "--kv-unified" not in LlamaCppBackend._ctx_integrity_flags( - 4, False, 98304, 98304, _CAPS_NONE + 4, False, False, 98304, 98304, _CAPS_NONE ) def test_fit_ctx_floors_explicit_request_under_fit(): - flags = LlamaCppBackend._ctx_integrity_flags(1, True, 98304, 98304, _CAPS_ALL) + # An explicit requested ctx floors --fit-ctx at that value on any --fit + # path, including legacy auto (auto_fit False). + flags = LlamaCppBackend._ctx_integrity_flags(1, True, False, 98304, 98304, _CAPS_ALL) assert flags[flags.index("--fit-ctx") + 1] == "98304" -def test_fit_ctx_skipped_without_fit_or_explicit_ctx_or_support(): +def test_fit_ctx_skipped_without_fit_or_support(): + # No --fit on -> no --fit-ctx. assert "--fit-ctx" not in LlamaCppBackend._ctx_integrity_flags( - 1, False, 98304, 98304, _CAPS_ALL + 1, False, False, 98304, 98304, _CAPS_ALL ) - assert "--fit-ctx" not in LlamaCppBackend._ctx_integrity_flags(1, True, 0, 262144, _CAPS_ALL) + # --fit on but the binary doesn't support --fit-ctx. assert "--fit-ctx" not in LlamaCppBackend._ctx_integrity_flags( - 1, True, 98304, 98304, _CAPS_NONE + 1, True, True, 98304, 98304, _CAPS_NONE + ) + + +def test_fit_ctx_floors_auto_request_at_8192_only_under_auto_fit(): + # Manual + Auto (auto_fit) floors the auto window at 8192 so --fit can't + # shrink it to a tiny size. + flags = LlamaCppBackend._ctx_integrity_flags(1, True, True, 0, 262144, _CAPS_ALL) + assert flags[flags.index("--fit-ctx") + 1] == "8192" + # Legacy auto (fit on but not auto_fit) emits -c 0 to pin native, so the + # 8192 floor must NOT ride along and override that pin. + assert "--fit-ctx" not in LlamaCppBackend._ctx_integrity_flags( + 1, True, False, 0, 262144, _CAPS_ALL ) diff --git a/studio/backend/tests/test_llama_cpp_slot_resume.py b/studio/backend/tests/test_llama_cpp_slot_resume.py new file mode 100644 index 0000000000..8b20c952c4 --- /dev/null +++ b/studio/backend/tests/test_llama_cpp_slot_resume.py @@ -0,0 +1,494 @@ +# SPDX-License-Identifier: AGPL-3.0-only +# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0 + +import os +from types import SimpleNamespace + +import core.inference.llama_cpp as llama_cpp +from core.inference.llama_cpp import LlamaCppBackend + + +def _resume_backend(tmp_path, n_slots = 1): + backend = LlamaCppBackend() + backend._healthy = True + # No-op lifecycle methods so the atexit cleanup can kill the fake quietly. + backend._process = SimpleNamespace( + poll = lambda: None, + terminate = lambda: None, + wait = lambda *a, **k: 0, + kill = lambda: None, + pid = 0, + ) + backend._port = 8081 + backend._slot_save_dir = str(tmp_path) + backend._slot_save_binary = ("/bin/llama-server", 1) + (tmp_path / "model.gguf").write_bytes(b"gguf") + backend._gguf_path = str(tmp_path / "model.gguf") + backend._effective_parallel_slots = n_slots + backend._estimate_kv_cache_bytes = lambda *a, **k: 0 + return backend + + +def _fake_disk(monkeypatch, free = 1 << 40): + monkeypatch.setattr(llama_cpp.shutil, "disk_usage", lambda _p: SimpleNamespace(free = free)) + + +class _Resp: + def __init__( + self, + status_code = 200, + body = None, + ): + self.status_code = status_code + self._body = body or {} + + def json(self): + return self._body + + +def test_save_returns_none_when_slot_save_disabled(monkeypatch, tmp_path): + backend = _resume_backend(tmp_path) + backend._slot_save_dir = None + monkeypatch.setattr( + llama_cpp.httpx, + "post", + lambda *a, **k: (_ for _ in ()).throw(AssertionError), + raising = False, + ) + assert backend.save_slots_for_resume() is None + + +def test_save_skipped_when_prompt_cache_disabled(monkeypatch, tmp_path): + backend = _resume_backend(tmp_path) + backend._prompt_cache_disabled = True + monkeypatch.setattr( + llama_cpp.httpx, + "post", + lambda *a, **k: (_ for _ in ()).throw(AssertionError), + raising = False, + ) + assert backend.save_slots_for_resume() is None + + +def test_save_skipped_when_insufficient_free_disk(monkeypatch, tmp_path): + backend = _resume_backend(tmp_path) + backend._estimate_kv_cache_bytes = lambda *a, **k: 1 << 40 + _fake_disk(monkeypatch, free = 1 << 20) + monkeypatch.setattr( + llama_cpp.httpx, + "post", + lambda *a, **k: (_ for _ in ()).throw(AssertionError), + raising = False, + ) + assert backend.save_slots_for_resume() is None + + +def test_save_collects_manifest_across_slots(monkeypatch, tmp_path): + backend = _resume_backend(tmp_path, n_slots = 2) + _fake_disk(monkeypatch) + calls = [] + + def fake_post(url, **kwargs): + calls.append((url, kwargs["params"], kwargs["json"])) + return _Resp(200, {"n_saved": 40, "n_written": 100}) + + monkeypatch.setattr(llama_cpp.httpx, "post", fake_post, raising = False) + manifest = backend.save_slots_for_resume() + assert manifest is not None + assert manifest["dir"] == str(tmp_path) + assert manifest["binary"] == ("/bin/llama-server", 1) + assert manifest["gguf"] == str(tmp_path / "model.gguf") + st = os.stat(manifest["gguf"]) + assert manifest["gguf_stat"] == ((st.st_size, st.st_mtime_ns),) + assert manifest["launch"] == backend._slot_launch_fingerprint() + assert [e["id"] for e in manifest["slots"]] == [0, 1] + assert all(e["n_saved"] == 40 for e in manifest["slots"]) + assert [c[1] for c in calls] == [{"action": "save"}] * 2 + assert "/slots/0" in calls[0][0] and "/slots/1" in calls[1][0] + + +def test_save_unlinks_empty_slot_and_returns_none(monkeypatch, tmp_path): + backend = _resume_backend(tmp_path) + _fake_disk(monkeypatch) + + def fake_post(url, **kwargs): + (tmp_path / kwargs["json"]["filename"]).write_bytes(b"") + return _Resp(200, {"n_saved": 0, "n_written": 0}) + + monkeypatch.setattr(llama_cpp.httpx, "post", fake_post, raising = False) + assert backend.save_slots_for_resume() is None + assert list(tmp_path.glob("resume-*.bin")) == [] # empty-slot file removed + + +def test_save_cap_breach_discards_all_files(monkeypatch, tmp_path): + backend = _resume_backend(tmp_path, n_slots = 2) + _fake_disk(monkeypatch) + monkeypatch.setattr(llama_cpp, "_SLOT_SAVE_MAX_BYTES", 150) + + def fake_post(url, **kwargs): + (tmp_path / kwargs["json"]["filename"]).write_bytes(b"x" * 100) + return _Resp(200, {"n_saved": 40, "n_written": 100}) + + monkeypatch.setattr(llama_cpp.httpx, "post", fake_post, raising = False) + assert backend.save_slots_for_resume() is None # 200 bytes > 150 cap + assert list(tmp_path.glob("resume-*.bin")) == [] + + +def test_save_transport_error_aborts_remaining_slots(monkeypatch, tmp_path): + backend = _resume_backend(tmp_path, n_slots = 3) + _fake_disk(monkeypatch) + calls = [] + + def fake_post(url, **kwargs): + calls.append(url) + raise OSError("connection refused") + + monkeypatch.setattr(llama_cpp.httpx, "post", fake_post, raising = False) + assert backend.save_slots_for_resume() is None + assert len(calls) == 1 # no retries against a dead server + + +def test_save_transport_error_unlinks_partial_file(monkeypatch, tmp_path): + backend = _resume_backend(tmp_path) + _fake_disk(monkeypatch) + + def fake_post(url, **kwargs): + (tmp_path / kwargs["json"]["filename"]).write_bytes(b"partial") + raise OSError("timed out") + + monkeypatch.setattr(llama_cpp.httpx, "post", fake_post, raising = False) + assert backend.save_slots_for_resume() is None + assert list(tmp_path.glob("resume-*.bin")) == [] + + +def test_fingerprint_tracks_lora_sidecar_rewrite(tmp_path): + backend = _resume_backend(tmp_path) + adapter = tmp_path / "adapter.gguf" + adapter.write_bytes(b"v1") + backend._extra_args = ["--lora", str(adapter)] + + before = backend._slot_launch_fingerprint() + adapter.write_bytes(b"v2-different") # re-exported adapter, same path + assert backend._slot_launch_fingerprint() != before + + backend._extra_args = [f"--lora={adapter}"] + assert backend._sidecar_weight_files() == [str(adapter)] + backend._extra_args = ["--lora-scaled", str(adapter), "0.5"] + assert backend._sidecar_weight_files() == [str(adapter)] + backend._extra_args = ["--control-vector", str(adapter), "--threads", "4"] + assert backend._sidecar_weight_files() == [str(adapter)] + + +def test_sidecar_files_parse_csv_and_colon_scale(tmp_path): + backend = _resume_backend(tmp_path) + a, b = tmp_path / "a.gguf", tmp_path / "b.gguf" + + backend._extra_args = ["--lora", f"{a},{b}"] + files = backend._sidecar_weight_files() + assert str(a) in files and str(b) in files + + backend._extra_args = ["--lora-scaled", f"{a}:0.5"] + assert str(a) in backend._sidecar_weight_files() + + backend._extra_args = ["--control-vector-scaled", f"{a}:1.0,{b}:2.0"] + files = backend._sidecar_weight_files() + assert str(a) in files and str(b) in files + + # Windows drive letter must not be mistaken for a scale separator. + backend._extra_args = ["--lora-scaled", "C:\\adapters\\a.gguf:0.75"] + assert "C:\\adapters\\a.gguf" in backend._sidecar_weight_files() + backend._extra_args = ["--lora", "C:\\adapters\\a.gguf"] + assert backend._sidecar_weight_files() == ["C:\\adapters\\a.gguf"] + + +def test_fingerprint_tracks_colon_scaled_adapter_rewrite(tmp_path): + backend = _resume_backend(tmp_path) + adapter = tmp_path / "adapter.gguf" + adapter.write_bytes(b"v1") + backend._extra_args = ["--lora-scaled", f"{adapter}:0.5"] + + before = backend._slot_launch_fingerprint() + adapter.write_bytes(b"v2-different") # re-exported adapter, same path + assert backend._slot_launch_fingerprint() != before + + +def test_fingerprint_tracks_effective_context_length(tmp_path): + backend = _resume_backend(tmp_path) + backend._effective_context_length = 8192 + + before = backend._slot_launch_fingerprint() + backend._effective_context_length = 4096 # auto-fit landed smaller on reload + assert backend._slot_launch_fingerprint() != before + + +def test_gguf_file_identity_covers_split_shards(tmp_path): + backend = _resume_backend(tmp_path) + first = tmp_path / "m-00001-of-00002.gguf" + second = tmp_path / "m-00002-of-00002.gguf" + first.write_bytes(b"a") + second.write_bytes(b"bb") + + before = backend._gguf_file_identity(str(first)) + st1, st2 = os.stat(first), os.stat(second) + assert before == ((st1.st_size, st1.st_mtime_ns), (st2.st_size, st2.st_mtime_ns)) + + second.write_bytes(b"rewritten") # sibling changes, primary untouched + after = backend._gguf_file_identity(str(first)) + assert after is not None and after != before + assert after[0] == before[0] # primary shard unchanged + + second.unlink() + assert backend._gguf_file_identity(str(first)) is None # missing shard + + +def test_save_skipped_when_user_disabled_prompt_cache(monkeypatch, tmp_path): + backend = _resume_backend(tmp_path) + backend._extra_args = ["--no-cache-prompt"] + monkeypatch.setattr( + llama_cpp.httpx, + "post", + lambda *a, **k: (_ for _ in ()).throw(AssertionError), + raising = False, + ) + assert backend.save_slots_for_resume() is None + + +def test_save_skipped_when_env_disables_prompt_cache(monkeypatch, tmp_path): + backend = _resume_backend(tmp_path) + monkeypatch.setenv("LLAMA_ARG_CACHE_PROMPT", "0") + monkeypatch.setattr( + llama_cpp.httpx, + "post", + lambda *a, **k: (_ for _ in ()).throw(AssertionError), + raising = False, + ) + assert backend.save_slots_for_resume() is None + monkeypatch.delenv("LLAMA_ARG_CACHE_PROMPT") + monkeypatch.setenv("LLAMA_ARG_NO_CACHE_PROMPT", "1") # legacy negative form + assert backend.save_slots_for_resume() is None + + +def test_explicit_cache_prompt_flag_overrides_env(monkeypatch, tmp_path): + backend = _resume_backend(tmp_path) + monkeypatch.setenv("LLAMA_ARG_CACHE_PROMPT", "0") + backend._extra_args = ["--cache-prompt"] # CLI wins over env in llama.cpp + _fake_disk(monkeypatch) + monkeypatch.setattr( + llama_cpp.httpx, + "post", + lambda *a, **k: _Resp(200, {"n_saved": 1, "n_written": 1}), + raising = False, + ) + assert backend.save_slots_for_resume() is not None + + +def test_user_cache_prompt_overrides_studio_no_cache_flag(monkeypatch, tmp_path): + # User extras follow Studio's flags, so an explicit --cache-prompt wins. + backend = _resume_backend(tmp_path) + backend._prompt_cache_disabled = True + backend._extra_args = ["--cache-prompt"] + _fake_disk(monkeypatch) + monkeypatch.setattr( + llama_cpp.httpx, + "post", + lambda *a, **k: _Resp(200, {"n_saved": 1, "n_written": 1}), + raising = False, + ) + assert backend.save_slots_for_resume() is not None + # Last flag wins when both appear in extras. + backend._extra_args = ["--cache-prompt", "--no-cache-prompt"] + assert backend.save_slots_for_resume() is None + + +def test_save_stops_writing_once_cap_exceeded(monkeypatch, tmp_path): + backend = _resume_backend(tmp_path, n_slots = 3) + _fake_disk(monkeypatch) + monkeypatch.setattr(llama_cpp, "_SLOT_SAVE_MAX_BYTES", 150) + calls = [] + + def fake_post(url, **kwargs): + calls.append(url) + (tmp_path / kwargs["json"]["filename"]).write_bytes(b"x" * 100) + return _Resp(200, {"n_saved": 1, "n_written": 100}) + + monkeypatch.setattr(llama_cpp.httpx, "post", fake_post, raising = False) + assert backend.save_slots_for_resume() is None + assert len(calls) == 2 # cap blown after slot 1; slot 2 never attempted + assert list(tmp_path.glob("resume-*.bin")) == [] + + +def test_save_aborts_between_slots_when_no_longer_idle(monkeypatch, tmp_path): + backend = _resume_backend(tmp_path, n_slots = 3) + _fake_disk(monkeypatch) + calls = [] + + def fake_post(url, **kwargs): + calls.append(url) + return _Resp(200, {"n_saved": 5, "n_written": 10}) + + monkeypatch.setattr(llama_cpp.httpx, "post", fake_post, raising = False) + aborts = iter([False, True, True]) + manifest = backend.save_slots_for_resume(should_abort = lambda: next(aborts)) + assert len(calls) == 1 # slots 1 and 2 skipped + assert manifest is not None + assert [e["id"] for e in manifest["slots"]] == [0] + + +def test_save_non_200_slot_is_skipped_but_others_kept(monkeypatch, tmp_path): + backend = _resume_backend(tmp_path, n_slots = 2) + _fake_disk(monkeypatch) + + def fake_post(url, **kwargs): + if "/slots/0" in url: + return _Resp(500) + return _Resp(200, {"n_saved": 5, "n_written": 10}) + + monkeypatch.setattr(llama_cpp.httpx, "post", fake_post, raising = False) + manifest = backend.save_slots_for_resume() + assert manifest is not None + assert [e["id"] for e in manifest["slots"]] == [1] + + +def test_restore_posts_each_slot_and_tolerates_failures(monkeypatch, tmp_path): + backend = _resume_backend(tmp_path) + calls = [] + + def fake_post(url, **kwargs): + calls.append((url, kwargs["params"], kwargs["json"])) + return _Resp(500 if "/slots/0" in url else 200, {"n_restored": 5}) + + monkeypatch.setattr(llama_cpp.httpx, "post", fake_post, raising = False) + backend.restore_slots_for_resume( + { + "slots": [ + {"id": 0, "filename": "resume-a-slot0.bin", "n_saved": 5}, + {"id": 1, "filename": "resume-a-slot1.bin", "n_saved": 5}, + ] + } + ) + assert [c[1] for c in calls] == [{"action": "restore"}] * 2 + assert calls[0][2] == {"filename": "resume-a-slot0.bin"} + + +def test_restore_transport_error_stops_early(monkeypatch, tmp_path): + backend = _resume_backend(tmp_path) + calls = [] + + def fake_post(url, **kwargs): + calls.append(url) + raise OSError("connection refused") + + monkeypatch.setattr(llama_cpp.httpx, "post", fake_post, raising = False) + backend.restore_slots_for_resume( + {"slots": [{"id": 0, "filename": "a.bin"}, {"id": 1, "filename": "b.bin"}]} + ) + assert len(calls) == 1 + + +def test_save_deletes_orphan_on_malformed_response(monkeypatch, tmp_path): + # A 200 that writes a file but returns a non-numeric counter must be cleaned + # up like any other save failure, not left orphaned holding chat KV. + backend = _resume_backend(tmp_path) + _fake_disk(monkeypatch) + + def fake_post(url, **kwargs): + (tmp_path / kwargs["json"]["filename"]).write_bytes(b"chat-kv") + return _Resp(200, {"n_saved": "not-an-int"}) + + monkeypatch.setattr(llama_cpp.httpx, "post", fake_post, raising = False) + assert backend.save_slots_for_resume() is None + assert list(tmp_path.glob("resume-*.bin")) == [] + + +def test_save_deletes_orphan_on_non_dict_response(monkeypatch, tmp_path): + backend = _resume_backend(tmp_path) + _fake_disk(monkeypatch) + + def fake_post(url, **kwargs): + (tmp_path / kwargs["json"]["filename"]).write_bytes(b"chat-kv") + return _Resp(200, ["unexpected", "list"]) + + monkeypatch.setattr(llama_cpp.httpx, "post", fake_post, raising = False) + assert backend.save_slots_for_resume() is None + assert list(tmp_path.glob("resume-*.bin")) == [] + + +def test_save_cap_uses_actual_file_size_not_reported_bytes(monkeypatch, tmp_path): + # A binary under-reporting n_written must not slip past the disk cap: the + # cap is enforced against the bytes actually on disk. + backend = _resume_backend(tmp_path) + _fake_disk(monkeypatch) + monkeypatch.setattr(llama_cpp, "_SLOT_SAVE_MAX_BYTES", 150) + + def fake_post(url, **kwargs): + (tmp_path / kwargs["json"]["filename"]).write_bytes(b"x" * 200) + return _Resp(200, {"n_saved": 5, "n_written": 1}) # under-reported + + monkeypatch.setattr(llama_cpp.httpx, "post", fake_post, raising = False) + assert backend.save_slots_for_resume() is None # 200 real bytes > 150 cap + assert list(tmp_path.glob("resume-*.bin")) == [] + + +def test_save_skipped_when_estimate_exceeds_cap(monkeypatch, tmp_path): + # An estimate over the cap skips before writing any slot at all. + backend = _resume_backend(tmp_path) + backend._estimate_kv_cache_bytes = lambda *a, **k: 1 << 40 + monkeypatch.setattr(llama_cpp, "_SLOT_SAVE_MAX_BYTES", 1 << 20) + _fake_disk(monkeypatch) + monkeypatch.setattr( + llama_cpp.httpx, + "post", + lambda *a, **k: (_ for _ in ()).throw(AssertionError), + raising = False, + ) + assert backend.save_slots_for_resume() is None + + +def test_save_skipped_when_model_file_changed_since_load(monkeypatch, tmp_path): + # The GGUF/sidecars were swapped on disk after the server loaded them, so the + # live KV belongs to the old weights: refuse to persist it (no POST at all). + backend = _resume_backend(tmp_path) + backend._slot_loaded_identity = ((("stale", 0),), ()) # != current identity + _fake_disk(monkeypatch) + monkeypatch.setattr( + llama_cpp.httpx, + "post", + lambda *a, **k: (_ for _ in ()).throw(AssertionError), + raising = False, + ) + assert backend.save_slots_for_resume() is None + + +def test_save_proceeds_when_load_identity_matches(monkeypatch, tmp_path): + # Matching load-time snapshot: the save runs normally. + backend = _resume_backend(tmp_path) + backend._slot_loaded_identity = ( + backend._gguf_file_identity(backend._gguf_path), + backend._slot_launch_fingerprint(), + ) + _fake_disk(monkeypatch) + + def fake_post(url, **kwargs): + (tmp_path / kwargs["json"]["filename"]).write_bytes(b"kv") + return _Resp(200, {"n_saved": 5, "n_written": 2}) + + monkeypatch.setattr(llama_cpp.httpx, "post", fake_post, raising = False) + manifest = backend.save_slots_for_resume() + assert manifest is not None + assert [e["id"] for e in manifest["slots"]] == [0] + + +def test_save_skipped_when_estimate_unavailable_and_low_disk(monkeypatch, tmp_path): + # A 0 estimate means metadata was insufficient, not a zero-byte cache: the save + # must demand room for the whole cap, not just 1 GiB, on a low-disk host. + backend = _resume_backend(tmp_path) + backend._estimate_kv_cache_bytes = lambda *a, **k: 0 # metadata unavailable + monkeypatch.setattr(llama_cpp, "_SLOT_SAVE_MAX_BYTES", 8 << 30) # 8 GiB cap + _fake_disk(monkeypatch, free = 2 << 30) # 2 GiB free < 8 + 1 GiB required + monkeypatch.setattr( + llama_cpp.httpx, + "post", + lambda *a, **k: (_ for _ in ()).throw(AssertionError), + raising = False, + ) + assert backend.save_slots_for_resume() is None diff --git a/studio/backend/tests/test_llama_cpp_stall_timeout.py b/studio/backend/tests/test_llama_cpp_stall_timeout.py new file mode 100644 index 0000000000..da36f75e8e --- /dev/null +++ b/studio/backend/tests/test_llama_cpp_stall_timeout.py @@ -0,0 +1,125 @@ +# SPDX-License-Identifier: AGPL-3.0-only +# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0 + +"""Regression test for the post-first-token stall timeout in the cancel-aware read. + +httpcore snapshots ``request.extensions["timeout"]["read"]`` once at body start, so +when ``_iter_text_cancellable`` lowers it after the first token, a one-token-then-silent +server hangs for the full prefill window. The fix re-reads the live extensions timeout +per call; a fake clock and always-silent stream check the read gives up after the live +stall timeout, not the stale prefill one. +""" + +from __future__ import annotations + +import inspect +import sys +import threading +import types as _types +from pathlib import Path + +import pytest + +_BACKEND_DIR = str(Path(__file__).resolve().parent.parent) +if _BACKEND_DIR not in sys.path: + sys.path.insert(0, _BACKEND_DIR) + +# Mirror sibling tests' stubbing so the module imports without fastapi. +_loggers_stub = _types.ModuleType("loggers") +_loggers_stub.get_logger = lambda name: __import__("logging").getLogger(name) +sys.modules.setdefault("loggers", _loggers_stub) +sys.modules.setdefault("structlog", _types.ModuleType("structlog")) + +import httpcore # noqa: E402 + +from core.inference import llama_cpp as llama_cpp_mod # noqa: E402 +from core.inference.llama_cpp import LlamaCppBackend # noqa: E402 + +_PREFILL_TIMEOUT = 1200.0 # what httpcore snapshots from the prefill timeout +_STALL_TIMEOUT = 120.0 # the post-first-token stall timeout the wrapper must honor + + +class _Obj: + pass + + +def _install(response, clock, silent_stream): + """Wire fake client/pool so _install_cancel_aware_read finds the stream; return the wrapped stream.read.""" + inner = _Obj() + inner._network_stream = silent_stream + connection = _Obj() + connection._connection = inner + pool = _Obj() + pool._connections = [connection] + transport = _Obj() + transport._pool = pool + client = _Obj() + client._transport = transport + + cancel_event = threading.Event() # never set: we test the stall path, not cancel + sig = inspect.signature(LlamaCppBackend._install_cancel_aware_read) + if "response" in sig.parameters: + # Fixed signature: wrapper reads the live extensions timeout. + LlamaCppBackend._install_cancel_aware_read(client, cancel_event, response) + else: + # Pre-fix signature: no response, so the stall assertion fails (proves the bug). + LlamaCppBackend._install_cancel_aware_read(client, cancel_event) + return silent_stream.read + + +def test_stall_timeout_honored_after_first_token(monkeypatch): + clock = {"t": 0.0} + monkeypatch.setattr(llama_cpp_mod.time, "monotonic", lambda: clock["t"]) + + # One token then silence: every read times out, advancing fake time by its timeout. + def silent_read(max_bytes, timeout = None): + clock["t"] += timeout if timeout is not None else 0.0 + raise httpcore.ReadTimeout("slice timed out on silence") + + stream = _Obj() + stream.read = silent_read + + # First token seen: the live read timeout is lowered to the stall timeout. + request = _Obj() + request.extensions = {"timeout": {"read": _STALL_TIMEOUT}} + response = _Obj() + response.request = request + + wrapped_read = _install(response, clock, stream) + + # httpcore still passes the stale prefill timeout it snapshotted at body start. + with pytest.raises(httpcore.ReadTimeout): + wrapped_read(65536, timeout = _PREFILL_TIMEOUT) + + # Must give up ~stall timeout after the last token, not the prefill window. + assert clock["t"] <= _STALL_TIMEOUT * 1.5, ( + f"stall timeout not honored: waited {clock['t']}s " + f"(expected ~{_STALL_TIMEOUT}s, not {_PREFILL_TIMEOUT}s)" + ) + assert clock["t"] >= _STALL_TIMEOUT * 0.5 + + +def test_prefill_timeout_used_when_no_live_override(monkeypatch): + """Without a lowered live timeout, the wrapper honors the passed prefill timeout, so the normal first-token wait is unchanged.""" + clock = {"t": 0.0} + monkeypatch.setattr(llama_cpp_mod.time, "monotonic", lambda: clock["t"]) + + def silent_read(max_bytes, timeout = None): + clock["t"] += timeout if timeout is not None else 0.0 + raise httpcore.ReadTimeout("slice timed out on silence") + + stream = _Obj() + stream.read = silent_read + + # No timeout extension: wrapper falls back to httpcore's passed timeout. + request = _Obj() + request.extensions = {} + response = _Obj() + response.request = request + + wrapped_read = _install(response, clock, stream) + + with pytest.raises(httpcore.ReadTimeout): + wrapped_read(65536, timeout = _PREFILL_TIMEOUT) + + assert clock["t"] >= _PREFILL_TIMEOUT * 0.9 diff --git a/studio/backend/tests/test_llama_cpp_tool_loop.py b/studio/backend/tests/test_llama_cpp_tool_loop.py index bd2c008589..e99e227d40 100644 --- a/studio/backend/tests/test_llama_cpp_tool_loop.py +++ b/studio/backend/tests/test_llama_cpp_tool_loop.py @@ -122,7 +122,7 @@ def _structured_tool_call(tool_name: str, arguments: dict, call_id: str) -> list def test_structured_tool_call_after_visible_preface_is_executed(monkeypatch): """llama-server may emit content first and then native delta.tool_calls. - Studio must not drop that tool call after it has streamed the preface. + Unsloth must not drop that tool call after it has streamed the preface. """ tool_call_id = "call_render_late" diff --git a/studio/backend/tests/test_llama_cpp_update.py b/studio/backend/tests/test_llama_cpp_update.py index 83ea07a066..f12384231f 100644 --- a/studio/backend/tests/test_llama_cpp_update.py +++ b/studio/backend/tests/test_llama_cpp_update.py @@ -83,6 +83,7 @@ def _write_install( repo: str = "unslothai/llama.cpp", asset: str | None = None, release_tag: str | None = None, + force_cpu: bool | None = None, ) -> str: """Create a fake prebuilt install and return the llama-server path.""" bin_dir = dir_ / "build" / "bin" @@ -99,6 +100,8 @@ def _write_install( } if asset is not None: marker["asset"] = asset + if force_cpu is not None: + marker["force_cpu"] = force_cpu (dir_ / MARKER).write_text(json.dumps(marker)) return str(binary) @@ -493,6 +496,47 @@ def test_start_update_preserves_vulkan_via_env(monkeypatch, tmp_path): assert popen_kwargs["env"]["UNSLOTH_FORCE_VULKAN"] == "1" +@pytest.mark.parametrize( + "force_cpu, expect_flag", + [ + # A deliberate CPU install (marker force_cpu=True) re-asserts --force-cpu on + # update so detect_host on a GPU host cannot re-route and revive the crash + # (#7213); --force-cpu also re-persists the flag for the next update. + (True, True), + # A transient fallback (or a legacy marker without the flag) stays free to + # heal to a GPU bundle (#6097). + (False, False), + (None, False), + ], +) +def test_start_update_cpu_fallback_preserved_by_flag(monkeypatch, tmp_path, force_cpu, expect_flag): + asset = "llama-b9493-bin-ubuntu-x64.tar.gz" + install_dir = tmp_path / "llama.cpp" + binary = _write_install(install_dir, "b9493", asset = asset, force_cpu = force_cpu) + monkeypatch.setattr(upd, "_find_binary", lambda: binary) + monkeypatch.setattr(upd, "_installer_script", lambda: tmp_path / "install_llama_prebuilt.py") + monkeypatch.setattr(freshness, "_fetch_latest_release_tag", lambda repo, timeout = 5.0: "b9518") + + captured: dict = {} + + def _on_start(cmd): + captured["cmd"] = cmd + _write_install(install_dir, "b9518", asset = asset, force_cpu = force_cpu) + + _patch_installer_popen(monkeypatch, lines = ["installed\n"], on_start = _on_start) + + assert upd.start_update()["started"] is True + deadline = time.time() + 10 + while time.time() < deadline: + job = upd.get_update_status()["job"] + if job["state"] in ("success", "error"): + break + time.sleep(0.05) + assert job["state"] == "success", job + assert ("--force-cpu" in captured["cmd"]) is expect_flag + assert "--cpu-fallback" not in captured["cmd"] + + def test_start_update_reports_full_release_tag(monkeypatch, tmp_path): install_dir = tmp_path / "llama.cpp" binary = _write_install(install_dir, "b9595") @@ -676,7 +720,7 @@ def test_install_cmd_rocm_marker_forwards_gfx(monkeypatch, tmp_path): assert "--rocm-gfx" in cmd assert cmd[cmd.index("--rocm-gfx") + 1] == "gfx110x" assert "--has-rocm" not in cmd - assert "--cpu-fallback" not in cmd + assert "--force-cpu" not in cmd assert "--simple-policy" not in cmd assert "--published-repo" in cmd and "unslothai/llama.cpp" in cmd @@ -690,17 +734,17 @@ def test_install_cmd_fork_rocm_marker_forwards_has_rocm(monkeypatch, tmp_path): def test_install_cmd_ggml_cpu_marker_has_no_cpu_fallback(monkeypatch, tmp_path): - # Legacy CPU installs recorded a ggml-org marker (new installs use the fork). - # Re-running into the same install-dir/repo reproduces the same CPU bundle; - # --cpu-fallback (which force-drops GPU detection) is reserved for setup.sh's - # arm64 rescue and must not appear here. + # Legacy CPU installs recorded a ggml-org marker (new installs use the fork) with + # no force_cpu field. Re-running into the same install-dir/repo reproduces the same + # CPU bundle; --force-cpu (the persisted-CPU re-assert) must not appear for a marker + # that never recorded a deliberate CPU choice, so it can still heal to GPU (#6097). cmd = _capture_install_cmd( monkeypatch, tmp_path, repo = "ggml-org/llama.cpp", asset = "llama-b9334-bin-ubuntu-x64.tar.gz", ) - assert "--cpu-fallback" not in cmd + assert "--force-cpu" not in cmd assert "--rocm-gfx" not in cmd assert "--has-rocm" not in cmd assert "--simple-policy" not in cmd @@ -714,7 +758,7 @@ def test_install_cmd_cuda_marker_minimal_and_backward_compatible(monkeypatch, tm assert "--simple-policy" not in cmd assert "--rocm-gfx" not in cmd assert "--has-rocm" not in cmd - assert "--cpu-fallback" not in cmd + assert "--force-cpu" not in cmd def test_install_cmd_pins_offered_release_tag(monkeypatch, tmp_path): diff --git a/studio/backend/tests/test_llama_cpp_wait_for_health.py b/studio/backend/tests/test_llama_cpp_wait_for_health.py index 82c5b4931a..423c3dd009 100644 --- a/studio/backend/tests/test_llama_cpp_wait_for_health.py +++ b/studio/backend/tests/test_llama_cpp_wait_for_health.py @@ -224,7 +224,7 @@ class TestRetryLogFilenameUnique: class TestFitOffRetryEligible: """Gate for the one-shot --fit off startup-crash retry. - Retry only when Studio's own VRAM math placed the model and nothing + Retry only when Unsloth's own VRAM math placed the model and nothing on the command line chose the fit mode explicitly.""" def test_eligible_for_plain_ngl_launch(self): diff --git a/studio/backend/tests/test_llama_cpp_wait_for_vram_settle.py b/studio/backend/tests/test_llama_cpp_wait_for_vram_settle.py index d0213f6079..b28df7ec3f 100644 --- a/studio/backend/tests/test_llama_cpp_wait_for_vram_settle.py +++ b/studio/backend/tests/test_llama_cpp_wait_for_vram_settle.py @@ -346,7 +346,7 @@ def test_helper_is_static_method_callable_off_class(): def test_kill_orphaned_servers_returns_count(): """The reaper reports how many owned orphans it killed, so __init__ can - arm the settle wait. Only Studio-owned llama-server procs count.""" + arm the settle wait. Only Unsloth-owned llama-server procs count.""" import os mypid = os.getpid() @@ -376,7 +376,7 @@ def test_kill_orphaned_servers_returns_count(): patch.object(LlamaCppBackend, "_pid_parent_is_alive", staticmethod(lambda pid: False)), ): n = LlamaCppBackend._kill_orphaned_servers() - assert n == 1, "only the Studio-owned orphan should be counted" + assert n == 1, "only the Unsloth-owned orphan should be counted" assert killed == [mypid + 1] # No owned orphans -> zero, so __init__ leaves the cold-start sentinel. @@ -392,8 +392,8 @@ def test_kill_orphaned_servers_returns_count(): def test_kill_orphaned_servers_spares_live_parent(): - """A Studio-owned llama-server whose parent is still running is not an - orphan (a live Studio or the user's shell owns it) and must never be + """An Unsloth-owned llama-server whose parent is still running is not an + orphan (a live Unsloth or the user's shell owns it) and must never be killed; only the true orphan (parent gone) is reaped.""" import os @@ -548,7 +548,7 @@ def test_record_then_reap_round_trip_identity_matches(tmp_path): def test_reap_recorded_pid_spares_live_server(tmp_path): - """A recorded server whose parent is still alive (the running Studio) is NEVER + """A recorded server whose parent is still alive (the running Unsloth) is NEVER reaped, and its pidfile is kept. This is the finding-3 guard: a helper backend constructed in-process must not kill the active chat server. Uses the REAL _pid_parent_is_alive (the child's parent is this live test process).""" diff --git a/studio/backend/tests/test_llama_cpp_windows_nvidia_path.py b/studio/backend/tests/test_llama_cpp_windows_nvidia_path.py index 957de4bad6..489d9eb8d1 100644 --- a/studio/backend/tests/test_llama_cpp_windows_nvidia_path.py +++ b/studio/backend/tests/test_llama_cpp_windows_nvidia_path.py @@ -3,7 +3,7 @@ """Tests for the Windows pip-nvidia DLL dir resolver. -Studio installs torch with bundled CUDA wheels (nvidia-cuda-runtime-cu13, +Unsloth installs torch with bundled CUDA wheels (nvidia-cuda-runtime-cu13, nvidia-cublas-cu13, etc.) and the prebuilt llama-server.exe must find those DLLs at runtime to load CUDA. Mirrors the Linux LD_LIBRARY_PATH block. See unslothai/unsloth#5106. diff --git a/studio/backend/tests/test_llama_server_args.py b/studio/backend/tests/test_llama_server_args.py index deeb228026..fa4ba71791 100644 --- a/studio/backend/tests/test_llama_server_args.py +++ b/studio/backend/tests/test_llama_server_args.py @@ -75,7 +75,7 @@ validate_extra_args = _lsa.validate_extra_args # Reasoning controls ["--reasoning-format", "deepseek"], ["-rea", "auto"], - # Soft-managed: user flags last-wins over Studio's auto-set version. + # Soft-managed: user flags last-wins over Unsloth's auto-set version. # --parallel / -np / --n-parallel are hard-denied (KV-cache + slot # count would desync); use `unsloth studio run --parallel N` instead. ["-c", "131072"], @@ -150,7 +150,7 @@ def test_non_flag_token_passes_through(): "--mmproj", "-mmu", "--mmproj-url", - # Networking (Studio binds + proxies) + # Networking (Unsloth binds + proxies) "--host", "--port", "--path", @@ -176,13 +176,15 @@ def test_non_flag_token_passes_through(): "--models-autoload", "--no-models-autoload", # Server-mode flips: --embedding / --rerank restrict llama-server to - # those endpoints and break Studio's chat hop. + # those endpoints and break Unsloth's chat hop. "--embedding", "--embeddings", "--rerank", "--reranking", - # llama-server's own --tools clashes with Studio's tool policy. + # llama-server's own --tools clashes with Unsloth's tool policy. "--tools", + # Slot-state dir: Studio owns it for KV persistence across idle unload. + "--slot-save-path", ], ) def test_denylist_rejects_all_aliases(denied): @@ -194,7 +196,7 @@ def test_denylist_rejects_all_aliases(denied): "args,offending", [ # Pass-through --parallel would last-wins-override the real slot - # count while Studio's KV-cache fit + llama_parallel_slots stay at + # count while Unsloth's KV-cache fit + llama_parallel_slots stay at # the typer value -- plan vs. process disagree. (["--parallel", "8"], "--parallel"), (["--parallel=8"], "--parallel"), @@ -224,6 +226,16 @@ def test_denylist_rejects_equals_form(): validate_extra_args(["--port=9000"]) +def test_slot_save_path_is_managed_in_all_forms(): + for args in (["--slot-save-path", "/tmp/x"], ["--slot-save-path=/tmp/x"], ["--slot-save-path"]): + with pytest.raises(ValueError, match = "--slot-save-path"): + validate_extra_args(args) + assert is_managed_flag("--slot-save-path") is True + assert is_managed_flag("--slot-save-path=/tmp/x") is True + # --slots (read-only diagnostics endpoint) stays a user choice. + assert is_managed_flag("--slots") is False + + @pytest.mark.parametrize( "padded", [" --parallel", "--parallel ", "\t--parallel", " -np", "-np \n", "-np\t"], @@ -656,7 +668,7 @@ def test_extra_args_disable_mmproj_last_wins(): def test_strip_shadowing_flags_drops_model_draft_with_spec(): - # --model-draft (and aliases) are Studio-managed since the separate + # --model-draft (and aliases) are Unsloth-managed since the separate # MTP drafter support: an inherited copy must not last-wins-override # the auto-detected drafter. out = strip_shadowing_flags( @@ -681,7 +693,7 @@ def test_strip_shadowing_flags_drops_model_draft_with_spec(): ) def test_strip_shadowing_flags_drops_hf_drafter_selectors_with_spec(selector): # HF drafter selectors must reset on inherit like local --model-draft, or a - # stale inherited HF drafter last-wins over Studio's re-derived spec choice. + # stale inherited HF drafter last-wins over Unsloth's re-derived spec choice. out = strip_shadowing_flags( selector + ["--top-k", "20"], strip_context = False, @@ -747,6 +759,34 @@ def test_strip_shadowing_flags_defaults_strip_split_mode_too(): assert strip_shadowing_flags(["--split-mode", "tensor"]) == [] +def test_strip_offload_is_opt_in_and_covers_moe(): + base = dict( + strip_context = False, + strip_cache = False, + strip_spec = False, + strip_template = False, + strip_split_mode = False, + ) + # Default: offload (incl. MoE) flags are NOT stripped. + assert strip_shadowing_flags(["--n-cpu-moe", "8", "--top-k", "20"], **base) == [ + "--n-cpu-moe", + "8", + "--top-k", + "20", + ] + # Opt-in strips layer AND MoE offload flags (value-aware), keeps the rest. + assert strip_shadowing_flags( + ["--n-cpu-moe", "8", "--gpu-layers", "33", "--fit", "off", "--top-k", "20"], + **base, + strip_offload = True, + ) == ["--top-k", "20"] + # Boolean --cpu-moe drops the flag only, not the following value. + assert strip_shadowing_flags(["--cpu-moe", "--seed", "-1"], **base, strip_offload = True) == [ + "--seed", + "-1", + ] + + @pytest.mark.parametrize( "args", [ @@ -769,7 +809,7 @@ def test_strip_split_mode_only_preserves_none_and_empty(): def test_strip_shadowing_flags_drops_tensor_split_with_split_mode(): # --tensor-split is coupled to the split mode: stripped together so a stale - # ratio can't override Studio's computed tensor split. Other flags survive. + # ratio can't override Unsloth's computed tensor split. Other flags survive. out = strip_shadowing_flags( ["--split-mode", "row", "--tensor-split", "1,1", "--top-k", "20"], strip_context = False, @@ -796,6 +836,23 @@ def test_strip_split_mode_only_drops_tensor_split_too(): assert strip_split_mode_only(["-sm=tensor", "-ts=3,1"]) == [] +def test_strip_tensor_split_alone_preserves_split_mode(): + # Manual mode emits its own --tensor-split, so an inherited ratio is dropped + # -- but the user's --split-mode row/none/layer choice (which the manual + # ratio toggle can't express) must survive. strip_tensor_split removes only + # the ratio, unlike strip_split_mode which removes the whole group. + out = strip_shadowing_flags( + ["--split-mode", "row", "--tensor-split", "1,1", "--top-k", "20"], + strip_context = False, + strip_cache = False, + strip_spec = False, + strip_template = False, + strip_split_mode = False, + strip_tensor_split = True, + ) + assert out == ["--split-mode", "row", "--top-k", "20"] + + def test_strip_shadowing_flags_keeps_model_draft_without_spec(): out = strip_shadowing_flags( ["--model-draft", "/custom/mtp.gguf"], diff --git a/studio/backend/tests/test_local_llama_cpp_link.py b/studio/backend/tests/test_local_llama_cpp_link.py index c78c029d91..6b44f61972 100644 --- a/studio/backend/tests/test_local_llama_cpp_link.py +++ b/studio/backend/tests/test_local_llama_cpp_link.py @@ -4,7 +4,7 @@ """Behavioral tests for the --with-llama-cpp-dir 'unmanaged local link' contract. When the canonical llama.cpp dir is a symlink (POSIX) / junction (Windows) to a -user's own checkout, Studio must treat it as externally managed: +user's own checkout, Unsloth must treat it as externally managed: - the in-app updater must not offer or apply a prebuilt over the link - orphan cleanup must not kill a llama-server the user launched from that tree @@ -67,7 +67,7 @@ def test_active_install_is_local_link(tmp_path: Path) -> None: binary = str(link / _server_subpath()) assert u._active_install_is_local_link(binary) is True - # A plain (non-link) llama.cpp dir is Studio-managed, not a local link. + # A plain (non-link) llama.cpp dir is Unsloth-managed, not a local link. plain = tmp_path / "plain" / "llama.cpp" plain.mkdir(parents = True) assert u._active_install_is_local_link(str(plain / _server_subpath())) is False diff --git a/studio/backend/tests/test_mcp_servers.py b/studio/backend/tests/test_mcp_servers.py index 6432ffb8e1..c5c37f098f 100644 --- a/studio/backend/tests/test_mcp_servers.py +++ b/studio/backend/tests/test_mcp_servers.py @@ -577,7 +577,7 @@ def test_clear_oauth_tokens_swallows_constructor_errors(tmp_path, monkeypatch): def test_tool_xml_parser_handles_hyphenated_function_names(): """Hyphenated tool names like `mcp__srv__list-issues` must parse, else the - model can call the tool but Studio can't dispatch.""" + model can call the tool but Unsloth can't dispatch.""" from core.inference.tool_call_parser import parse_tool_calls_from_text calls = parse_tool_calls_from_text( diff --git a/studio/backend/tests/test_mcp_stdio_improvements.py b/studio/backend/tests/test_mcp_stdio_improvements.py index b0bfd45135..745c2cc447 100644 --- a/studio/backend/tests/test_mcp_stdio_improvements.py +++ b/studio/backend/tests/test_mcp_stdio_improvements.py @@ -188,7 +188,7 @@ def test_validate_url_allows_url_in_argument(monkeypatch): # ── P6: Data Recipe stdio path obeys the same host gate ───────────── -# build_mcp_providers needs the Studio-only data_designer plugin; skip if absent. +# build_mcp_providers needs the Unsloth-only data_designer plugin; skip if absent. _STDIO_RECIPE = { "mcp_providers": [ diff --git a/studio/backend/tests/test_middleware.py b/studio/backend/tests/test_middleware.py index 024e5d3677..9b585e2cb7 100644 --- a/studio/backend/tests/test_middleware.py +++ b/studio/backend/tests/test_middleware.py @@ -14,6 +14,7 @@ import pytest from fastapi import FastAPI, HTTPException, Request from fastapi.responses import Response from fastapi.testclient import TestClient +from starlette.middleware.gzip import GZipMiddleware _BACKEND_ROOT = Path(__file__).resolve().parents[1] @@ -514,6 +515,71 @@ class TestResearchPortMiddleware: assert seen["server"] == ("127.0.0.1", 4321) +class TestFrontendAssets: + def test_hashed_assets_are_compressed_and_cached(self, tmp_path, main_module): + content = b"export const value = 'responsive';\n" * 200 + (tmp_path / "page-abc123.js").write_bytes(content) + app = FastAPI() + assets_app = GZipMiddleware( + main_module.ImmutableStaticFiles(directory = tmp_path), + minimum_size = 1024, + compresslevel = 6, + ) + app.mount("/assets", assets_app, name = "assets") + + response = TestClient(app).get( + "/assets/page-abc123.js", + headers = {"Accept-Encoding": "gzip"}, + ) + + assert response.status_code == 200 + assert response.content == content + assert response.headers["content-encoding"] == "gzip" + assert response.headers["cache-control"] == (main_module._IMMUTABLE_ASSET_CACHE_CONTROL) + assert "accept-encoding" in response.headers["vary"].lower() + + def test_asset_revalidation_keeps_immutable_cache_header(self, tmp_path, main_module): + (tmp_path / "page-abc123.js").write_text("export {};", encoding = "utf-8") + app = FastAPI() + app.mount( + "/assets", + main_module.ImmutableStaticFiles(directory = tmp_path), + name = "assets", + ) + client = TestClient(app) + first = client.get("/assets/page-abc123.js") + + response = client.get( + "/assets/page-abc123.js", + headers = {"If-None-Match": first.headers["etag"]}, + ) + + assert response.status_code == 304 + assert response.headers["cache-control"] == (main_module._IMMUTABLE_ASSET_CACHE_CONTROL) + + def test_range_request_is_not_compressed(self, tmp_path, main_module): + content = b"export const value = 'responsive';\n" * 200 + (tmp_path / "page-abc123.js").write_bytes(content) + app = FastAPI() + assets_app = main_module._AssetGZipMiddleware( + main_module.ImmutableStaticFiles(directory = tmp_path), + minimum_size = 1024, + compresslevel = 6, + ) + app.mount("/assets", assets_app, name = "assets") + + response = TestClient(app).get( + "/assets/page-abc123.js", + headers = {"Accept-Encoding": "gzip", "Range": "bytes=0-99"}, + ) + + assert response.status_code == 206 + assert response.headers.get("content-encoding") != "gzip" + assert response.headers["content-range"] == f"bytes 0-99/{len(content)}" + assert response.content == content[:100] + assert response.headers["cache-control"] == (main_module._IMMUTABLE_ASSET_CACHE_CONTROL) + + # /api/health auth gate diff --git a/studio/backend/tests/test_mlx_inference_backend.py b/studio/backend/tests/test_mlx_inference_backend.py index fa50cd84d6..fafaea0043 100644 --- a/studio/backend/tests/test_mlx_inference_backend.py +++ b/studio/backend/tests/test_mlx_inference_backend.py @@ -2,6 +2,7 @@ import sys import types +from contextlib import contextmanager from types import SimpleNamespace import pytest @@ -40,12 +41,16 @@ class _DummyModel: def _install_fake_mlx(monkeypatch): mlx_pkg = types.ModuleType("mlx") mlx_core = types.ModuleType("mlx.core") + mlx_utils = types.ModuleType("mlx.utils") mlx_core.metal = _DummyMetal() mlx_core.set_wired_limit = _DummyMX.set_wired_limit mlx_core.device_info = _DummyMX.device_info + mlx_utils.tree_unflatten = dict mlx_pkg.core = mlx_core + mlx_pkg.utils = mlx_utils monkeypatch.setitem(sys.modules, "mlx", mlx_pkg) monkeypatch.setitem(sys.modules, "mlx.core", mlx_core) + monkeypatch.setitem(sys.modules, "mlx.utils", mlx_utils) def _install_fake_fast_mlx(monkeypatch, calls): @@ -68,6 +73,99 @@ def _install_fake_fast_mlx(monkeypatch, calls): monkeypatch.setitem(sys.modules, "unsloth_zoo.mlx.loader", mlx_loader) +class _AdapterTree: + def __init__(self, modules): + self.modules = dict(modules) + + def named_modules(self): + return list(self.modules.items()) + + def update_modules(self, modules): + self.modules.update(modules) + + +def test_temporary_mlx_adapter_state_bypasses_and_restores_wrappers(monkeypatch): + _install_fake_mlx(monkeypatch) + from core.inference.mlx_inference import _temporary_mlx_adapter_state + + base = object() + wrapper = SimpleNamespace(lora_a = object(), lora_b = object(), linear = base, m = object()) + model = _AdapterTree({"model.layers.0.proj": wrapper}) + + with pytest.raises(RuntimeError, match = "generation failed"): + with _temporary_mlx_adapter_state(model, False): + assert model.modules["model.layers.0.proj"] is base + raise RuntimeError("generation failed") + assert model.modules["model.layers.0.proj"] is wrapper + + +def test_temporary_mlx_adapter_state_validates_requests(): + from core.inference.mlx_inference import _temporary_mlx_adapter_state + + wrapper = SimpleNamespace(lora_a = object(), lora_b = object(), embedding = object()) + model = _AdapterTree({"embed_tokens": wrapper}) + with _temporary_mlx_adapter_state(model, True): + assert model.modules["embed_tokens"] is wrapper + with pytest.raises(NotImplementedError, match = "named adapter"): + with _temporary_mlx_adapter_state(model, "other"): + pass + + base_model = _AdapterTree({"proj": object()}) + with _temporary_mlx_adapter_state(base_model, None): + pass + with _temporary_mlx_adapter_state(base_model, True): + pass + + unsupported = _AdapterTree({"proj": SimpleNamespace(lora_a = object(), lora_b = object())}) + with _temporary_mlx_adapter_state(unsupported, True): + pass + with pytest.raises(RuntimeError, match = "without their base modules"): + with _temporary_mlx_adapter_state(unsupported, False): + pass + + +def test_temporary_mlx_adapter_state_uses_real_mlx_module_tree(): + nn = pytest.importorskip("mlx.nn") + pytest.importorskip("mlx_lm") + from mlx_lm.models.switch_layers import SwitchLinear + from mlx_lm.tuner.dora import DoRALinear + from mlx_lm.tuner.lora import LoRAEmbedding, LoRALinear, LoRASwitchLinear + + from core.inference.mlx_inference import _temporary_mlx_adapter_state + + class _Layer(nn.Module): + def __init__(self): + super().__init__() + quantized = nn.QuantizedLinear.from_linear(nn.Linear(32, 32), group_size = 32, bits = 4) + self.quantized_proj = LoRALinear.from_base(quantized) + self.dora_proj = DoRALinear.from_base(nn.Linear(4, 4)) + + class _Model(nn.Module): + def __init__(self): + super().__init__() + self.layers = [_Layer()] + self.embed_tokens = LoRAEmbedding.from_base(nn.Embedding(16, 4)) + self.experts = LoRASwitchLinear.from_base(SwitchLinear(4, 4, 2)) + + model = _Model() + wrappers = { + path: module + for path, module in model.named_modules() + if hasattr(module, "lora_a") and hasattr(module, "lora_b") + } + bases = { + path: getattr(module, "linear", getattr(module, "embedding", None)) + for path, module in wrappers.items() + } + + with _temporary_mlx_adapter_state(model, False): + live = dict(model.named_modules()) + assert all(live[path] is base for path, base in bases.items()) + + restored = dict(model.named_modules()) + assert all(restored[path] is wrapper for path, wrapper in wrappers.items()) + + def test_mlx_inference_text_load_forwards_studio_settings(monkeypatch): _install_fake_mlx(monkeypatch) calls = [] @@ -138,7 +236,7 @@ def test_mlx_inference_vlm_lora_uses_unsloth_loader_without_native_adapter_rewri _install_fake_fast_mlx(monkeypatch, calls) def _native_vlm_load(*_args, **_kwargs): - raise AssertionError("Studio MLX VLM inference must use FastMLXModel") + raise AssertionError("Unsloth MLX VLM inference must use FastMLXModel") mlx_vlm = types.ModuleType("mlx_vlm") mlx_vlm.load = _native_vlm_load @@ -333,10 +431,87 @@ def test_mlx_generate_chat_response_accepts_template_kwargs(): ), f"{name!r} must default to None so existing callers stay valid" +def test_mlx_vlm_reemits_think_prefill_inside_adapter_context(monkeypatch): + """A prefilled block must be re-emitted as the first VLM snapshot, + inside the adapter context (so unsupported requests still raise first), so + the UI renders the thinking block during prefill and a pre-first-token + cancel does not drop it. Mirrors _generate_text.""" + from core.inference import mlx_inference + + MLXInferenceBackend = mlx_inference.MLXInferenceBackend + + order = [] + + @contextmanager + def _adapter_state(_model, state): + assert backend._generation_lock.locked() + order.append("adapter_enter") + try: + yield + finally: + order.append("adapter_exit") + + monkeypatch.setattr(mlx_inference, "_temporary_mlx_adapter_state", _adapter_state) + monkeypatch.setattr( + "core.inference.chat_template_helpers.detect_think_prefill", + lambda *_a, **_k: "\n", + ) + + prompt_utils = SimpleNamespace( + MODEL_CONFIG = {"deepseek_vl_v2": object()}, + apply_chat_template = lambda *_a, **_k: " model-aware", + ) + mlx_vlm = types.ModuleType("mlx_vlm") + mlx_vlm.prompt_utils = prompt_utils + + def _vlm_stream(*_a, **_k): + # The prefill must have been emitted before any generated token. + assert order[-1] == "adapter_enter" + yield SimpleNamespace(text = "ok", prompt_tokens = 3, generation_tokens = 1) + + mlx_vlm.stream_generate = _vlm_stream + monkeypatch.setitem(sys.modules, "mlx_vlm", mlx_vlm) + monkeypatch.setattr( + "core.inference.chat_template_helpers.apply_chat_template_for_generation", + lambda _t, _m, **_k: " model-aware", + ) + + backend = MLXInferenceBackend() + backend._model = SimpleNamespace(config = {"model_type": "deepseek_vl_v2"}) + backend._processor = SimpleNamespace(tokenizer = SimpleNamespace()) + args = ([{"role": "user", "content": [{"type": "image"}]}], object(), 0, 1, 0, 0, 1, 1, None) + + gen = backend._generate_vlm(*args, _adapter_state = False) + # First snapshot is the prefill alone, emitted after entering the adapter context. + assert next(gen) == "\n" + assert order == ["adapter_enter"] + # Subsequent snapshots are cumulative (prefill + generated text). + assert next(gen) == "\nok" + gen.close() + assert order == ["adapter_enter", "adapter_exit"] + + def test_mlx_vlm_generation_selects_renderer_by_capability(monkeypatch): - from core.inference.mlx_inference import MLXInferenceBackend + from core.inference import mlx_inference + + MLXInferenceBackend = mlx_inference.MLXInferenceBackend calls = {"generic": [], "model": [], "stream": []} + adapter_events = [] + adapter_active = {"value": False} + + @contextmanager + def _adapter_state(_model, state): + assert backend._generation_lock.locked() + adapter_events.append(("enter", state)) + adapter_active["value"] = True + try: + yield + finally: + adapter_active["value"] = False + adapter_events.append(("exit", state)) + + monkeypatch.setattr(mlx_inference, "_temporary_mlx_adapter_state", _adapter_state) state = {"generic": "serialized", "model": " model-aware"} prompt_utils = SimpleNamespace( MODEL_CONFIG = {"deepseek_vl_v2": object()}, @@ -346,10 +521,13 @@ def test_mlx_vlm_generation_selects_renderer_by_capability(monkeypatch): ) mlx_vlm = types.ModuleType("mlx_vlm") mlx_vlm.prompt_utils = prompt_utils - mlx_vlm.stream_generate = lambda *_args, **kwargs: ( - calls["stream"].append((_args, kwargs)) - or iter([SimpleNamespace(text = "ok", prompt_tokens = 3, generation_tokens = 1)]) - ) + + def _vlm_stream(*args, **kwargs): + assert adapter_active["value"] + calls["stream"].append((args, kwargs)) + yield SimpleNamespace(text = "ok", prompt_tokens = 3, generation_tokens = 1) + + mlx_vlm.stream_generate = _vlm_stream monkeypatch.setitem(sys.modules, "mlx_vlm", mlx_vlm) def generic(_target, _messages, **kwargs): @@ -369,7 +547,11 @@ def test_mlx_vlm_generation_selects_renderer_by_capability(monkeypatch): backend._processor = SimpleNamespace(tokenizer = SimpleNamespace()) args = ([{"role": "user", "content": [{"type": "image"}]}], object(), 0, 1, 0, 0, 1, 1, None) tools = [{"function": {"name": "search"}}] - assert list(backend._generate_vlm(*args)) == ["ok"] + generator = backend._generate_vlm(*args, _adapter_state = False) + assert next(generator) == "ok" + assert adapter_active["value"] and backend._generation_lock.locked() + generator.close() + assert adapter_events == [("enter", False), ("exit", False)] assert calls["model"][0]["num_images"] == 1 assert calls["stream"][0][0][2] == " model-aware" with pytest.raises(RuntimeError, match = "dropping requested tools"): @@ -449,7 +631,10 @@ def test_mlx_generate_text_forwards_kwargs_into_template_helper(monkeypatch): """Mac text path must route through apply_chat_template_for_generation so reasoning / tool kwargs reach the tokenizer.""" _install_fake_mlx(monkeypatch) - from core.inference.mlx_inference import MLXInferenceBackend + from core.inference import mlx_inference + + MLXInferenceBackend = mlx_inference.MLXInferenceBackend + real_adapter_state = mlx_inference._temporary_mlx_adapter_state # The text path renders once with tools, then the native-template fallback makes a second no- # tools probe call (tools=None) to detect whether the template dropped the schema. @@ -474,11 +659,31 @@ def test_mlx_generate_text_forwards_kwargs_into_template_helper(monkeypatch): mlx_lm_sample.make_sampler = lambda **_kw: object() mlx_lm_sample.make_logits_processors = lambda **_kw: None + adapter_events = [] + adapter_active = {"value": False} + stream_state = {"fail": False} + + @contextmanager + def _adapter_state(_model, state): + assert backend._generation_lock.locked() + adapter_events.append(("enter", state)) + adapter_active["value"] = True + try: + yield + finally: + adapter_active["value"] = False + adapter_events.append(("exit", state)) + + monkeypatch.setattr(mlx_inference, "_temporary_mlx_adapter_state", _adapter_state) + class _Resp: def __init__(self, tok): self.token = tok def _stream_generate(_model, _tokenizer, **_kw): + assert adapter_active["value"] + if stream_state["fail"]: + raise RuntimeError("generation failed") yield _Resp(1) mlx_lm_pkg.stream_generate = _stream_generate @@ -500,17 +705,45 @@ def test_mlx_generate_text_forwards_kwargs_into_template_helper(monkeypatch): backend._tokenizer = _Tok() backend._is_vlm = False - out = list( - backend.generate_chat_response( - messages = [{"role": "user", "content": "ping"}], - tools = [{"function": {"name": "web_search"}}], - enable_thinking = True, - reasoning_effort = "medium", - preserve_thinking = True, - max_new_tokens = 1, - ) + generator = backend.generate_with_adapter_control( + use_adapter = False, + messages = [{"role": "user", "content": "ping"}], + tools = [{"function": {"name": "web_search"}}], + enable_thinking = True, + reasoning_effort = "medium", + preserve_thinking = True, + max_new_tokens = 1, ) - assert out == ["hi"] + assert next(generator) == "hi" + assert adapter_active["value"] and backend._generation_lock.locked() + generator.close() + assert adapter_events == [("enter", False), ("exit", False)] + stream_state["fail"] = True + with pytest.raises(RuntimeError, match = "generation failed"): + list( + backend.generate_with_adapter_control( + use_adapter = False, + messages = [{"role": "user", "content": "ping"}], + max_new_tokens = 1, + ) + ) + assert adapter_events[-2:] == [("enter", False), ("exit", False)] + assert not backend._generation_lock.locked() + + monkeypatch.setattr(mlx_inference, "_temporary_mlx_adapter_state", real_adapter_state) + monkeypatch.setattr( + "core.inference.chat_template_helpers.detect_think_prefill", + lambda *_args, **_kwargs: "", + ) + stream_state["fail"] = False + named = backend.generate_with_adapter_control( + use_adapter = "named", + messages = [{"role": "user", "content": "ping"}], + max_new_tokens = 1, + ) + with pytest.raises(NotImplementedError, match = "named adapter"): + next(named) + assert not adapter_active["value"] and not backend._generation_lock.locked() # The toggled kwargs must reach the chat-template helper on the real render # (one of the calls carries the tools; the fallback probe passes tools=None). tool_renders = [ diff --git a/studio/backend/tests/test_mlx_repair.py b/studio/backend/tests/test_mlx_repair.py index 365cc46410..47a695ccbd 100644 --- a/studio/backend/tests/test_mlx_repair.py +++ b/studio/backend/tests/test_mlx_repair.py @@ -103,7 +103,7 @@ def test_repair_install_pins_transformers_and_cleans_up(monkeypatch): assert mr.attempt_mlx_repair() is True cmd = captured["cmd"] # transformers is pinned via a constraint file so the mlx install cannot - # upgrade it underneath Studio, and the temp constraint file is cleaned up. + # upgrade it underneath Unsloth, and the temp constraint file is cleaned up. assert "--constraint" in cmd assert "--upgrade" in cmd reinstall_pairs = set(zip(cmd, cmd[1:])) @@ -123,7 +123,7 @@ def test_install_requires_prebuilt_wheels(monkeypatch): # A source distribution's PEP 517 build backend runs arbitrary code at install # time, before the post-install stack check. The unattended self-heal must # require pre-built wheels so a malicious resolver-selected sdist cannot execute - # during ordinary Studio startup. mlx/mlx-metal ship wheels only and + # during ordinary Unsloth startup. mlx/mlx-metal ship wheels only and # mlx-lm/mlx-vlm publish py3-none-any wheels, so a healthy self-heal still works. pytest.importorskip("transformers") captured = {} @@ -143,7 +143,7 @@ def test_install_requires_prebuilt_wheels(monkeypatch): def test_install_env_drops_secrets_and_source_redirects(monkeypatch): - # The unattended self-heal must not hand resolver/build code the full Studio + # The unattended self-heal must not hand resolver/build code the full Unsloth # environment: secrets and package-source redirects are dropped, while the # variables uv genuinely needs are forwarded. monkeypatch.setenv("HF_TOKEN", "secret-hf") diff --git a/studio/backend/tests/test_mlx_stop_checkpoint.py b/studio/backend/tests/test_mlx_stop_checkpoint.py new file mode 100644 index 0000000000..d4a00cc6c8 --- /dev/null +++ b/studio/backend/tests/test_mlx_stop_checkpoint.py @@ -0,0 +1,137 @@ +# SPDX-License-Identifier: AGPL-3.0-only +# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0 + +"""Regression tests for MLX stop-and-save checkpoint handling.""" + +import importlib.util +import json +import sys +import types +from pathlib import Path + +import numpy as np +from safetensors.numpy import save_file + + +_BACKEND = Path(__file__).resolve().parents[1] + + +def _load_worker_module(): + spec = importlib.util.spec_from_file_location( + "training_worker_under_test", + _BACKEND / "core" / "training" / "worker.py", + ) + module = importlib.util.module_from_spec(spec) + assert spec.loader is not None + spec.loader.exec_module(module) + return module + + +worker = _load_worker_module() + + +class _FakeTrainer: + def __init__(self, step: int): + self._global_step = step + self._train_loss_history = [] + self.model = object() + + +def _write_checkpoint(out: Path, step: int) -> Path: + checkpoint = out / f"checkpoint-{step}" + checkpoint.mkdir(parents = True, exist_ok = True) + (checkpoint / "trainer_state.json").write_text( + json.dumps({"global_step": step}), encoding = "utf-8" + ) + save_file({"weight": np.ones(1, dtype = np.float32)}, checkpoint / "adapters.safetensors") + save_file( + {"state": np.ones(1, dtype = np.float32)}, + checkpoint / "optimizer_state.safetensors", + ) + return checkpoint + + +def test_mlx_has_checkpoint_at_step_requires_complete_state(tmp_path): + out = tmp_path / "outputs" / "run_x" + _write_checkpoint(out, 5) + + assert worker._mlx_has_checkpoint_at_step(out, 5) is True + + +def test_write_mlx_stop_checkpoint_returns_true_when_current_step_checkpoint_exists(tmp_path): + out = tmp_path / "outputs" / "run_x" + _write_checkpoint(out, 5) + + assert worker._write_mlx_stop_checkpoint(_FakeTrainer(step = 5), object(), out) is True + + +def test_write_mlx_stop_checkpoint_writes_current_step_when_only_older_checkpoint_exists( + tmp_path, monkeypatch +): + out = tmp_path / "outputs" / "run_x" + _write_checkpoint(out, 5) + + saved_steps: list[int] = [] + + def _save_state(_value, path, name): + save_file({"state": np.ones(1, dtype = np.float32)}, Path(path, name)) + + def _save_trainer_state(state, ckpt_dir, **_kwargs): + Path(ckpt_dir, "trainer_state.json").write_text(json.dumps(state), encoding = "utf-8") + saved_steps.append(int(state["global_step"])) + + fake_utils = types.SimpleNamespace( + save_trainable_adapters = lambda model, path: _save_state( + model, path, "adapters.safetensors" + ), + save_optimizer_state = lambda optimizer, path: _save_state( + optimizer, path, "optimizer_state.safetensors" + ), + save_trainer_state = _save_trainer_state, + ) + monkeypatch.setitem(sys.modules, "unsloth_zoo.mlx.utils", fake_utils) + + assert worker._write_mlx_stop_checkpoint(_FakeTrainer(step = 10), object(), out) is True + assert saved_steps == [10] + assert (out / "checkpoint-10" / "trainer_state.json").is_file() + + +def test_write_mlx_stop_checkpoint_returns_false_without_optimizer(tmp_path): + out = tmp_path / "outputs" / "run_x" + out.mkdir(parents = True) + + assert worker._write_mlx_stop_checkpoint(_FakeTrainer(step = 5), None, out) is False + + +def test_write_mlx_stop_checkpoint_rejects_incomplete_current_checkpoint(tmp_path): + out = tmp_path / "outputs" / "run_x" + ckpt = out / "checkpoint-5" + ckpt.mkdir(parents = True) + (ckpt / "trainer_state.json").write_text('{"global_step": 5}', encoding = "utf-8") + + assert worker._write_mlx_stop_checkpoint(_FakeTrainer(step = 5), None, out) is False + + +def test_write_mlx_stop_checkpoint_ignores_stale_checkpoint_without_optimizer(tmp_path): + # An older checkpoint does not cover the current step, so this still fails. + out = tmp_path / "outputs" / "run_x" + _write_checkpoint(out, 5) + + assert worker._write_mlx_stop_checkpoint(_FakeTrainer(step = 10), None, out) is False + + +def test_write_mlx_stop_checkpoint_returns_false_when_save_fails(tmp_path, monkeypatch): + out = tmp_path / "outputs" / "run_x" + out.mkdir(parents = True) + + def _boom(*_args, **_kwargs): + raise RuntimeError("save failed") + + fake_utils = types.SimpleNamespace( + save_trainable_adapters = _boom, + save_optimizer_state = lambda *_a, **_k: None, + save_trainer_state = lambda *_a, **_k: None, + ) + monkeypatch.setitem(sys.modules, "unsloth_zoo.mlx.utils", fake_utils) + + assert worker._write_mlx_stop_checkpoint(_FakeTrainer(step = 5), object(), out) is False diff --git a/studio/backend/tests/test_model_picker_regression.py b/studio/backend/tests/test_model_picker_regression.py new file mode 100644 index 0000000000..f38a4d0b8d --- /dev/null +++ b/studio/backend/tests/test_model_picker_regression.py @@ -0,0 +1,232 @@ +# SPDX-License-Identifier: AGPL-3.0-only +# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0 + +"""Regression guards for the model-picker per-model-config feature (the set of +bugs that got the predecessor PR reverted). Pure-function / validation checks +only, so they run on CPU in the backend pytest job with no model download. + +Covers, at the backend layer: + - infra-model hiding: the RAG embedder (bge-small-en-v1.5) and the llama.cpp + install-validation probe (ggml-org/models / stories260K) stay hidden, while + normal chat repos are not hidden; + - the HF token is honored from the dedicated header with the query string as a + fallback, never the other way around; + - the chat-template byte caps reject oversized overrides (both the char-count + fast path and the UTF-8 byte path) and the sidecar reader is size-bounded. +""" + +from __future__ import annotations + +import sys +import types + +import pytest + +# Keep this test runnable without the optional structlog dependency (mirrors +# tests/test_cached_gguf_routes.py), since importing routes.models pulls it in. +if "structlog" not in sys.modules: + + class _DummyLogger: + def __getattr__(self, _name): + return lambda *args, **kwargs: None + + sys.modules["structlog"] = types.SimpleNamespace( + BoundLogger = _DummyLogger, + get_logger = lambda *args, **kwargs: _DummyLogger(), + ) + +import routes.models as models_route +from core.rag import config as rag_config +from hub.dependencies import get_hf_token +from models.inference import LoadRequest +from picker.schemas import MAX_CHAT_TEMPLATE_BYTES +from picker.service import _read_bounded_text +from utils.hidden_models import is_hidden_model + + +@pytest.fixture(autouse = True) +def _pin_default_embedder(monkeypatch): + """Pin the effective embedder to Studio's static default so hiding is + deterministic and cannot depend on ambient RAG config / env.""" + default = "unsloth/bge-small-en-v1.5" + monkeypatch.setattr(rag_config, "EMBEDDING_MODEL", default, raising = False) + monkeypatch.setattr(rag_config, "effective_embedding_model", lambda: default) + monkeypatch.setattr(rag_config, "effective_gguf_repo", lambda: default) + monkeypatch.setattr(rag_config, "default_gguf_repo", lambda: default) + + +# --------------------------------------------------------------------------- # +# Infra-model hiding (the "infra models resurfaced in the picker" regression) # +# --------------------------------------------------------------------------- # + + +@pytest.mark.parametrize( + "value", + [ + "ggml-org/models", # the probe repo id + "unsloth/bge-small-en-v1.5", # the RAG embedder repo + "unsloth/bge-small-en-v1.5-GGUF", # its GGUF companion + "/root/.cache/huggingface/hub/x/stories260K.gguf", # probe on disk + "/root/.cache/x/Stories260K.GGUF", # case-insensitive + r"C:\\models\\stories260K.gguf", # windows-style path + "/opt/models/bge-small-en-v1.5", # embedder basename folder + "/opt/models/bge-small-en-v1.5-Q8_0.gguf", # suffixed local weight + ], +) +def test_infra_models_are_hidden(value): + assert is_hidden_model(value) is True + + +@pytest.mark.parametrize( + "value", + [ + "unsloth/gemma-3-270m-it-GGUF", # a normal small chat GGUF + "unsloth/Qwen3-0.6B", # a normal non-GGUF chat model + "user/stories260K-finetune-GGUF", # repo id merely contains "stories260k" + "user/model-chat", # generic repo must not be hidden + "meta-llama/Llama-3.1-8B-Instruct", + ], +) +def test_normal_models_are_not_hidden(value): + assert is_hidden_model(value) is False + + +def test_is_hidden_model_ignores_empty_values(): + assert is_hidden_model(None) is False + assert is_hidden_model("") is False + assert is_hidden_model(None, "", "unsloth/gemma-3-270m-it-GGUF") is False + + +def test_hidden_model_matchers_expose_probe_needles(): + needles, exact_ids, _exact_paths = models_route.hidden_model_matchers() + lowered = [n.lower() for n in needles] + assert "ggml-org/models" in lowered + assert "stories260k.gguf" in lowered + # The configured embedder is exposed as an exact repo id, never as a + # basename needle that would substring-hide unrelated chat models. + assert "bge-small-en-v1.5" not in lowered + assert "unsloth/bge-small-en-v1.5" in exact_ids + + +def test_hidden_model_matchers_custom_repo_publishes_exact_ids(monkeypatch): + monkeypatch.setattr(rag_config, "effective_embedding_model", lambda: "org/model") + monkeypatch.setattr(rag_config, "effective_gguf_repo", lambda: "org/model-GGUF") + needles, exact_ids, exact_paths = models_route.hidden_model_matchers() + assert needles == ["ggml-org/models", "stories260k.gguf"] + assert "org/model" in exact_ids + assert "org/model-gguf" in exact_ids + assert exact_paths == [] + + +def test_hidden_model_matchers_local_owner_name_path_is_exact_path(monkeypatch, tmp_path): + # A local embedder shaped like owner/name that exists on disk must be an + # exact resolved path, not a Hub repo id (mirroring is_hidden_model), so the + # local row stays hidden instead of showing as a chat model. + (tmp_path / "models" / "embedder").mkdir(parents = True) + monkeypatch.chdir(tmp_path) + monkeypatch.setattr(rag_config, "effective_embedding_model", lambda: "models/embedder") + monkeypatch.setattr(rag_config, "effective_gguf_repo", lambda: "ggml-org/models") + _needles, exact_ids, exact_paths = models_route.hidden_model_matchers() + resolved = str((tmp_path / "models" / "embedder").resolve()).lower() + assert resolved in exact_paths + assert "models/embedder" not in exact_ids + + +# --------------------------------------------------------------------------- # +# HF token via header, query string only as a fallback (the token-leak fix) # +# --------------------------------------------------------------------------- # + + +def test_get_hf_token_strips_and_returns(): + assert get_hf_token(" hf_abc ") == "hf_abc" + + +@pytest.mark.parametrize("value", [None, "", " ", "\n\t"]) +def test_get_hf_token_blank_is_none(value): + assert get_hf_token(value) is None + + +@pytest.mark.parametrize( + "value,expected", + [(" hf_x ", "hf_x"), ("", None), (" ", None), (None, None), (1234, None)], +) +def test_normalize_hf_token(value, expected): + assert models_route._normalize_hf_token(value) == expected + + +def test_header_token_wins_over_query(): + header, query = "hf_header", "hf_query" + resolved = models_route._normalize_hf_token(header) or models_route._normalize_hf_token(query) + assert resolved == "hf_header" + + +def test_query_token_is_fallback_when_header_absent(): + resolved = models_route._normalize_hf_token(None) or models_route._normalize_hf_token( + "hf_query" + ) + assert resolved == "hf_query" + + +# --------------------------------------------------------------------------- # +# Chat-template byte caps (the unbounded-template hardening) # +# --------------------------------------------------------------------------- # + + +def _load_request(**overrides): + data = {"model_path": "unsloth/test-model-GGUF", "gguf_variant": "Q4_K_M"} + data.update(overrides) + return LoadRequest.model_validate(data) + + +def test_blank_chat_template_override_normalizes_to_none(): + assert _load_request(chat_template_override = " \n\t").chat_template_override is None + + +def test_nonblank_chat_template_override_preserved_verbatim(): + template = " {{ messages }} " + assert _load_request(chat_template_override = template).chat_template_override == template + + +def test_chat_template_at_byte_limit_is_accepted(): + template = "a" * MAX_CHAT_TEMPLATE_BYTES # exactly the limit, 1 byte/char + assert ( + len(_load_request(chat_template_override = template).chat_template_override) + == MAX_CHAT_TEMPLATE_BYTES + ) + + +def test_chat_template_over_char_limit_is_rejected(): + with pytest.raises(Exception): # pydantic ValidationError wrapping ValueError + _load_request(chat_template_override = "a" * (MAX_CHAT_TEMPLATE_BYTES + 1)) + + +def test_chat_template_over_byte_limit_is_rejected(): + # Char count stays under the limit but UTF-8 bytes exceed it (3 bytes/char), + # so only the byte-count branch can catch this. + multibyte = "€" * (MAX_CHAT_TEMPLATE_BYTES // 2) # euro sign, 3 bytes each + assert len(multibyte) <= MAX_CHAT_TEMPLATE_BYTES + assert len(multibyte.encode("utf-8")) > MAX_CHAT_TEMPLATE_BYTES + with pytest.raises(Exception): + _load_request(chat_template_override = multibyte) + + +def test_read_bounded_text_reads_within_limit(tmp_path): + p = tmp_path / "t.json" + p.write_text("hello", encoding = "utf-8") + assert _read_bounded_text(p, 16) == "hello" + + +def test_read_bounded_text_rejects_over_limit(tmp_path): + p = tmp_path / "big.json" + p.write_bytes(b"x" * 100) + assert _read_bounded_text(p, 50) is None + + +def test_read_bounded_text_at_limit_is_read(tmp_path): + p = tmp_path / "exact.json" + p.write_bytes(b"x" * 50) + assert _read_bounded_text(p, 50) == "x" * 50 + + +def test_read_bounded_text_missing_file_is_none(tmp_path): + assert _read_bounded_text(tmp_path / "nope.json", 50) is None diff --git a/studio/backend/tests/test_model_update_robustness.py b/studio/backend/tests/test_model_update_robustness.py index edf55812e2..7d98766616 100644 --- a/studio/backend/tests/test_model_update_robustness.py +++ b/studio/backend/tests/test_model_update_robustness.py @@ -314,6 +314,7 @@ def test_cached_model_scan_keeps_local_safetensors_repo(monkeypatch, tmp_path): file_name = "model.safetensors", size_on_disk = 100, blob_path = str(repo_path / "blobs" / "modelsha"), + blob_last_modified = 3_000.0, ), ] ) @@ -336,6 +337,51 @@ def test_cached_model_scan_keeps_local_safetensors_repo(monkeypatch, tmp_path): assert rows[0]["repo_id"] == "Org/SafeTensorRepo" assert rows[0]["model_format"] == "safetensors" assert rows[0]["size_bytes"] == 100 + assert rows[0]["last_modified"] == 3_000.0 + + +def test_cached_gguf_scan_keeps_download_timestamp(monkeypatch, tmp_path): + repo_path = tmp_path / "models--Org--GgufRepo" + repo = SimpleNamespace( + repo_id = "Org/GgufRepo", + repo_type = "model", + repo_path = repo_path, + revisions = [ + SimpleNamespace( + files = [ + SimpleNamespace( + file_name = "model-Q4_K_M.gguf", + size_on_disk = 100, + blob_path = None, + blob_last_modified = 5_000.0, + ), + ] + ) + ], + ) + monkeypatch.setattr( + CI, + "all_hf_cache_scans", + lambda: [SimpleNamespace(repos = [repo])], + ) + monkeypatch.setattr( + CI.hf_cache_scan, + "is_gguf_repo_partial", + lambda *args, **kwargs: False, + ) + monkeypatch.setattr( + CI, + "_gguf_variant_state_summary", + lambda _repo_id: (False, 0), + ) + + rows = CI._scan_cached_gguf() + + assert len(rows) == 1 + assert rows[0]["repo_id"] == "Org/GgufRepo" + assert rows[0]["model_format"] == "gguf" + assert rows[0]["size_bytes"] == 100 + assert rows[0]["last_modified"] == 5_000.0 # ── hf_hub_download_with_xet_fallback force_download bypass (X2/F2) ─── @@ -636,3 +682,18 @@ def test_reclaim_replaced_gguf_variant_keeps_no_symlink_current_file(monkeypatch assert snap.exists() is True # the current file must survive assert result["removed_snapshots"] == 0 assert result["deleted_blobs"] == 0 + + +def _mmproj_repo(*file_names: str): + return SimpleNamespace( + revisions = [SimpleNamespace(files = [SimpleNamespace(file_name = n) for n in file_names])] + ) + + +def test_repo_has_mmproj_requires_gguf_projector(): + # A non-GGUF sidecar whose name merely contains "mmproj" must NOT mark the + # repo vision-capable; the runtime's projector detection is GGUF-only. + assert CI._repo_has_mmproj(_mmproj_repo("model-Q4_K_M.gguf", "mmproj_config.json")) is False + assert CI._repo_has_mmproj(_mmproj_repo("model-Q4_K_M.gguf", "README-mmproj.md")) is False + # A real GGUF projector still marks the repo vision-capable. + assert CI._repo_has_mmproj(_mmproj_repo("model-Q4_K_M.gguf", "mmproj-F16.gguf")) is True diff --git a/studio/backend/tests/test_mtp_vram_budget.py b/studio/backend/tests/test_mtp_vram_budget.py index 0efbbf596d..694d60cfc6 100644 --- a/studio/backend/tests/test_mtp_vram_budget.py +++ b/studio/backend/tests/test_mtp_vram_budget.py @@ -502,7 +502,7 @@ class TestExtraArgsMtpDetection: assert _extra_args_mtp_draft_path([], env = dict(os.environ)) == "/large.gguf" def test_load_model_gates_env_spec_type_on_off_mode(self): - # LLAMA_ARG_SPEC_TYPE only reaches the child when Studio emits no spec + # LLAMA_ARG_SPEC_TYPE only reaches the child when Unsloth emits no spec # flag (UI mode "off", no user --spec-type); otherwise the emitted # --spec-type/--spec-default overrides the env, so the reserve must not # consult it or a stale MTP env over-reserves (Finding F3). Whitespace- @@ -530,8 +530,8 @@ class TestExtraArgsMtpDetection: def test_load_model_drafter_budget_precedence(self): # The budget sizes the drafter the launch actually loads: CLI extras win, - # then Studio's emitted mtp_draft_path (overrides LLAMA_ARG_SPEC_DRAFT_MODEL), - # then the env drafter -- not the env before Studio's (reviewer.py R3). + # then Unsloth's emitted mtp_draft_path (overrides LLAMA_ARG_SPEC_DRAFT_MODEL), + # then the env drafter -- not the env before Unsloth's (reviewer.py R3). compact = "".join(inspect.getsource(LlamaCppBackend.load_model).split()) assert "_cli_draft_for_budget=_extra_args_mtp_draft_path(extra_args,env={})" in compact assert "_env_draft_for_budget=_extra_args_mtp_draft_path([],env=os.environ)" in compact @@ -732,7 +732,7 @@ class TestExtraArgsMtpDetection: assert _extra_args_n_ubatch([], env = {"LLAMA_ARG_UBATCH": "notint"}) is None def test_env_main_cache_type_for_budget(self): - # The child inherits LLAMA_ARG_CACHE_TYPE_K/_V, but Studio emits no + # The child inherits LLAMA_ARG_CACHE_TYPE_K/_V, but Unsloth emits no # --cache-type when neither param nor extras set it -> a heavier env # main KV (f32) must be adopted so the reserve matches the child. assert _env_main_cache_type_for_budget(env = {}) is None @@ -765,7 +765,7 @@ class TestExtraArgsMtpDetection: assert "cache_type_kv=_env_main_cache_type_for_budget()" in compact def test_env_split_mode_is_tensor(self): - # The child inherits LLAMA_ARG_SPLIT_MODE, but Studio emits --split-mode + # The child inherits LLAMA_ARG_SPLIT_MODE, but Unsloth emits --split-mode # only on its tensor branch -> a tensor env must flip the budget so the # heavier per-device compute buffer is reserved (not layer overhead). assert _env_split_mode_is_tensor(env = {}) is False @@ -918,7 +918,7 @@ class TestExtraArgsMtpDetection: # Cluster A: when the final decision is layer split, an inherited # non-layer LLAMA_ARG_SPLIT_MODE (and paired LLAMA_ARG_TENSOR_SPLIT) must # be popped from the child env so the child cannot run tensor/row/none - # against Studio's layer budget. Whitespace-stripped for formatter. + # against Unsloth's layer budget. Whitespace-stripped for formatter. compact = "".join(inspect.getsource(LlamaCppBackend.load_model).split()) assert 'env.get("LLAMA_ARG_SPLIT_MODE")' in compact assert '_inherited_sm!="layer"' in compact @@ -936,10 +936,10 @@ class TestExtraArgsMtpDetection: assert "env.pop(_ct_var,None)" in compact def test_load_model_clears_tensor_split_env_in_tensor_mode(self): - # review run3 #2: Studio owns the tensor split. When it emits no + # review run3 #2: Unsloth owns the tensor split. When it emits no # --tensor-split (even split), a stale inherited LLAMA_ARG_TENSOR_SPLIT must # be cleared in the TENSOR branch too (not just the layer downgrade), or the - # child runs a split Studio didn't budget. The else (tensor) branch pops it. + # child runs a split Unsloth didn't budget. The else (tensor) branch pops it. src = inspect.getsource(LlamaCppBackend.load_model) compact = "".join(src.split()) # appears in both the layer branch and the tensor branch. @@ -1005,14 +1005,14 @@ def test_qwen36_class_regression_picks_lower_ctx_with_mtp(): def test_mtp_draft_budget_prefers_user_extras_drafter(): # A user --model-draft in extras is appended last and wins at launch, so the - # VRAM budget must size it first; then Studio's emitted mtp_draft_path (which + # VRAM budget must size it first; then Unsloth's emitted mtp_draft_path (which # overrides LLAMA_ARG_SPEC_DRAFT_MODEL), then the env drafter (load_model is too # entangled to drive end-to-end; assert the precedence at the source level). # Whitespace-stripped so the check survives any formatter line-wrapping. compact = "".join(inspect.getsource(LlamaCppBackend.load_model).split()) - # CLI extras sized first (env={} so the env doesn't pre-empt Studio's drafter). + # CLI extras sized first (env={} so the env doesn't pre-empt Unsloth's drafter). assert "_cli_draft_for_budget=_extra_args_mtp_draft_path(extra_args,env={})" in compact - # Order: CLI extras, then Studio's mtp_draft_path, then the env drafter. + # Order: CLI extras, then Unsloth's mtp_draft_path, then the env drafter. assert "_cli_draft_for_budgetor_studio_draft_for_budgetor_env_draft_for_budget" in compact - # The env must not be consulted before Studio's resolved drafter. + # The env must not be consulted before Unsloth's resolved drafter. assert "_extra_args_mtp_draft_path(extra_args)ormtp_draft_path" not in compact diff --git a/studio/backend/tests/test_multimodal_document.py b/studio/backend/tests/test_multimodal_document.py index 5cd7c876cc..b347c4aef8 100644 --- a/studio/backend/tests/test_multimodal_document.py +++ b/studio/backend/tests/test_multimodal_document.py @@ -3,7 +3,7 @@ """Tests for PDF / document attachment translation on external providers. -Studio adds a normalised `input_document` content part on +Unsloth adds a normalised `input_document` content part on ChatCompletionRequest so the frontend needn't know the per-provider attachment shape: diff --git a/studio/backend/tests/test_nudge_tool_calls_wiring.py b/studio/backend/tests/test_nudge_tool_calls_wiring.py index e03fd0c7d7..82a6543aeb 100644 --- a/studio/backend/tests/test_nudge_tool_calls_wiring.py +++ b/studio/backend/tests/test_nudge_tool_calls_wiring.py @@ -3,7 +3,7 @@ """Wiring guard for the plan-without-action ``nudge_tool_calls`` policy. -Decided policy: the re-prompt is ALWAYS ON for the Studio inference paths +Decided policy: the re-prompt is ALWAYS ON for the Unsloth inference paths (safetensors, GGUF/llama_cpp, MLX) and OPT-IN for the API (/v1 OpenAI-compat + Anthropic-compat, controlled by the request's ``nudge_tool_calls``, default off). @@ -16,7 +16,7 @@ Mechanism (verified here without loading a model): opt-in), while the GGUF loop keeps its pre-existing default-on behaviour (``None`` keeps nudging) so an omitted flag never disables GGUF; * the API request models default the flag to ``None`` (opt-in / off); - * the Studio-facing routes forward the request's flag, and the Studio frontend + * the Unsloth-facing routes forward the request's flag, and the Unsloth frontend sends ``nudge_tool_calls: true`` -- exercised behaviourally in ``test_safetensors_tool_loop.py`` and ``test_llama_cpp_tool_loop.py``. """ @@ -87,7 +87,7 @@ def test_api_request_models_default_the_flag_off(): def test_studio_routes_forward_the_request_flag(): - # The Studio chat frontend posts to /v1/chat/completions and /v1/messages + # The Unsloth chat frontend posts to /v1/chat/completions and /v1/messages # with nudge_tool_calls=true; the route handlers forward the request value # (external API clients that omit it fall back to the opt-in default). from routes import inference as routes_inference diff --git a/studio/backend/tests/test_offline_gguf_cache_fallback.py b/studio/backend/tests/test_offline_gguf_cache_fallback.py index 0b9a1e704f..35dd3979b1 100644 --- a/studio/backend/tests/test_offline_gguf_cache_fallback.py +++ b/studio/backend/tests/test_offline_gguf_cache_fallback.py @@ -119,6 +119,13 @@ def _build_cache( return snap +def _symlink_or_skip(link: Path, target: Path) -> None: + try: + link.symlink_to(target) + except OSError as exc: + pytest.skip(f"symlinks unavailable: {exc}") + + @pytest.fixture def hf_cache(tmp_path, monkeypatch): """Point ``huggingface_hub.constants.HF_HUB_CACHE`` at a temp dir.""" @@ -897,7 +904,7 @@ class TestHfOfflineIfDnsDead: assert "HF_HUB_OFFLINE" not in os.environ def test_user_set_hf_hub_offline_is_preserved(self, dns, clean_offline_env, monkeypatch): - # User explicitly set offline before launching Studio. + # User explicitly set offline before launching Unsloth. monkeypatch.setenv("HF_HUB_OFFLINE", "1") dns.fail() with _hf_offline_if_dns_dead() as did_set: @@ -1084,7 +1091,7 @@ class TestListLocalGgufVariantsSubdir: target.write_bytes(b"\0" * 20) out = _find_local_gguf_by_variant(str(tmp_path), "Q4_K_M") - assert out == str(target.resolve()) + assert out == str(target.absolute()) def test_find_local_gguf_by_variant_skips_big_endian_only_match(self, tmp_path): from utils.models.model_config import _find_local_gguf_by_variant @@ -1094,6 +1101,57 @@ class TestListLocalGgufVariantsSubdir: assert _find_local_gguf_by_variant(str(tmp_path), "Q4_K_M") is None + def test_find_local_gguf_by_variant_keeps_split_symlink_name(self, tmp_path): + from utils.models.model_config import _find_local_gguf_by_variant + + blobs = tmp_path / "blobs" + blobs.mkdir() + snap = tmp_path / "snapshots" / "rev" / "BF16" + snap.mkdir(parents = True) + (tmp_path / "snapshots" / "rev" / "config.json").write_text("{}") + for i, sha in enumerate(("aa" * 32, "bb" * 32), start = 1): + (blobs / sha).write_bytes(b"\0" * 10) + _symlink_or_skip(snap / f"model-BF16-0000{i}-of-00002.gguf", blobs / sha) + + out = _find_local_gguf_by_variant(str(tmp_path / "snapshots" / "rev"), "BF16") + assert out is not None + assert Path(out).name == "model-BF16-00001-of-00002.gguf" + + def test_detect_gguf_model_keeps_split_symlink_name(self, tmp_path): + from utils.models.model_config import detect_gguf_model + + blobs = tmp_path / "blobs" + blobs.mkdir() + snap = tmp_path / "snapshots" / "rev" + snap.mkdir(parents = True) + for i, (sha, size) in enumerate((("cc" * 32, 10), ("dd" * 32, 20)), start = 1): + (blobs / sha).write_bytes(b"\0" * size) + _symlink_or_skip(snap / f"model-BF16-0000{i}-of-00002.gguf", blobs / sha) + + out = detect_gguf_model(str(snap)) + assert out is not None + assert Path(out).name == "model-BF16-00001-of-00002.gguf" + + def test_lone_split_symlink_uses_colocated_target_shards(self, tmp_path): + from utils.models.model_config import _find_local_gguf_by_variant, detect_gguf_model + + target_dir = tmp_path / "external" / "BF16" + target_dir.mkdir(parents = True) + target = target_dir / "model-BF16-00001-of-00002.gguf" + target.write_bytes(b"\0" * 10) + (target_dir / "model-BF16-00002-of-00002.gguf").write_bytes(b"\0" * 10) + + local = tmp_path / "local" + local.mkdir() + (local / "config.json").write_text("{}") + link = local / target.name + _symlink_or_skip(link, target) + + expected = str(target.absolute()) + assert _find_local_gguf_by_variant(str(local), "BF16") == expected + assert detect_gguf_model(str(local)) == expected + assert detect_gguf_model(str(link)) == expected + def test_model_config_variant_ignores_big_endian_sibling(self, tmp_path): from utils.models.model_config import ModelConfig diff --git a/studio/backend/tests/test_offline_inference_parent.py b/studio/backend/tests/test_offline_inference_parent.py index 71331220d6..bd0014ea64 100644 --- a/studio/backend/tests/test_offline_inference_parent.py +++ b/studio/backend/tests/test_offline_inference_parent.py @@ -139,7 +139,7 @@ class TestLoraDetectOffline: monkeypatch.setenv("HF_HUB_OFFLINE", "1") - # Studio catches Exception broadly; pin that the call still happens + # Unsloth catches Exception broadly; pin that the call still happens # (so cached LoRAs aren't missed) and returns fast via the mock. class _OfflineModeIsEnabled(Exception): pass diff --git a/studio/backend/tests/test_openai_auto_switch.py b/studio/backend/tests/test_openai_auto_switch.py index d02a2a4f7e..1ee9ef36d3 100644 --- a/studio/backend/tests/test_openai_auto_switch.py +++ b/studio/backend/tests/test_openai_auto_switch.py @@ -8,6 +8,7 @@ tests/test_gguf_completion_usage.py. """ import asyncio +import os import pytest @@ -18,6 +19,10 @@ from utils import openai_auto_switch_settings as settings class _FakeBackend: + effective_parallel_slots = 1 + _slot_save_binary = None + _gguf_path = None + def __init__( self, loaded_id = None, @@ -29,6 +34,22 @@ class _FakeBackend: self.hf_variant = hf_variant self._openai_advertised_id = advertised_id + def save_slots_for_resume(self, should_abort = None): + return None + + def restore_slots_for_resume(self, manifest): + return None + + def _slot_launch_fingerprint(self): + return ((), None, None, 1) + + def _gguf_file_identity(self, path): + try: + st = os.stat(path) + except OSError: + return None + return ((st.st_size, st.st_mtime_ns),) + class _LoadRecorder: """Stand-in for the load route: records calls and simulates a load.""" @@ -53,10 +74,15 @@ class _LoadRecorder: from fastapi import HTTPException raise HTTPException(status_code = 503, detail = "load failed") self.backend.model_identifier = request.model_path + self.backend.hf_variant = getattr(request, "gguf_variant", None) + self.backend._gguf_path = request.model_path self.backend.is_loaded = True # Mirror _load_model_impl: a load advertises its own id until the # auto-switch caller overwrites it with the repo id. self.backend._openai_advertised_id = None + from core.inference import llama_keepwarm as kw + + kw.note_model_loaded(self.backend) return None @@ -446,6 +472,75 @@ def test_idle_loop_unloads_after_ttl_and_stashes_for_reload(monkeypatch): assert stash is not None and stash[0] == "unsloth/Idle-GGUF" and stash[1] == "Q4_K_M" +def test_idle_loop_deletes_saved_kv_when_unload_fails(monkeypatch, tmp_path): + import time + from core.inference import llama_keepwarm as kw + + monkeypatch.setattr(settings, "get_auto_unload_idle_seconds", lambda: 0.005) + monkeypatch.setattr(settings, "get_auto_unload_keep_kv", lambda: True) + kw._inflight = 0 + kw._pending = 0 + kw._last_active = time.monotonic() - 3600 + kw._last_unloaded_model = None + kw._kv_resume = None + + saved = tmp_path / "resume-abc-slot0.bin" + backend = _FakeBackend("unsloth/Idle-GGUF") + manifests = [] + + def _save(should_abort = None): + if manifests: + return None + saved.write_bytes(b"kv") + manifest = {"dir": str(tmp_path), "slots": [{"id": 0, "filename": saved.name}]} + manifests.append(manifest) + return manifest + + def _unload(): + raise RuntimeError("cuda teardown failed") + + backend.save_slots_for_resume = _save + backend.unload_model = _unload + monkeypatch.setattr(inference_route, "get_llama_cpp_backend", lambda: backend) + + async def _drive(): + task = asyncio.create_task(kw.idle_unload_loop(poll_seconds = 0.01)) + for _ in range(200): + await asyncio.sleep(0.01) + if manifests and not saved.exists(): + break + task.cancel() + try: + await task + except asyncio.CancelledError: + pass + + asyncio.run(_drive()) + assert manifests and not saved.exists() + assert kw._kv_resume is None + + +def test_disabling_idle_unload_purges_saved_kv(monkeypatch, tmp_path): + # PUT leaves keep-KV on but makes idle unload inactive: saved KV must go too. + import routes.settings as settings_route + from core.inference import llama_keepwarm as kw + + saved = tmp_path / "resume-abc-slot0.bin" + saved.write_bytes(b"kv") + kw._kv_resume = { + "identity": ("m", None, "m"), + "dir": str(tmp_path), + "slots": [{"id": 0, "filename": saved.name}], + } + monkeypatch.setattr(settings_route, "set_openai_auto_switch", lambda *a: (False, 300, True)) + monkeypatch.setattr(settings_route, "get_auto_unload_idle_seconds", lambda: 0) + + payload = settings_route.OpenAIAutoSwitchPayload(enabled = False) + resp = settings_route.update_openai_auto_switch(payload, "tester") + assert resp.idle_unload_active is False and resp.auto_unload_keep_kv is True + assert kw._kv_resume is None and not saved.exists() + + def test_audio_generate_is_tracked_as_inference_path(): # Direct GGUF TTS uses the llama backend and can outlive the idle TTL, so # the keep-warm middleware must count it as in-flight inference. @@ -689,7 +784,7 @@ def test_v1_models_retrieve_is_case_insensitive(monkeypatch): def test_index_excludes_hidden_models(tmp_path, monkeypatch): # The llama.cpp validation probe and RAG embedding weights are hidden from - # Studio's pickers; they must never become auto-switch targets. + # Unsloth's pickers; they must never become auto-switch targets. from types import SimpleNamespace import routes.models as models_route @@ -697,6 +792,10 @@ def test_index_excludes_hidden_models(tmp_path, monkeypatch): normal.write_bytes(b"x" * 32) probe = tmp_path / "stories260K.gguf" # llama.cpp install-validation probe probe.write_bytes(b"x" * 32) + embedder = tmp_path / "embedding-Q8_0.gguf" + embedder.write_bytes(b"x" * 32) + local_default_embedder = tmp_path / "bge-small-en-v1.5-F16.gguf" + local_default_embedder.write_bytes(b"x" * 32) def _info(mid, path): return SimpleNamespace(id = mid, path = str(path), model_id = mid, display_name = mid) @@ -704,7 +803,22 @@ def test_index_excludes_hidden_models(tmp_path, monkeypatch): monkeypatch.setattr( models_route, "_scan_models_dir", - lambda *a, **k: [_info("org/Normal-GGUF", normal), _info("ggml-org/models", probe)], + lambda *a, **k: [ + _info("org/Normal-GGUF", normal), + _info("ggml-org/models", probe), + SimpleNamespace( + id = str(embedder), + path = str(embedder), + model_id = "unsloth/bge-small-en-v1.5-GGUF", + display_name = "embedding-Q8_0", + ), + SimpleNamespace( + id = str(local_default_embedder), + path = str(local_default_embedder), + model_id = None, + display_name = local_default_embedder.name, + ), + ], ) monkeypatch.setattr(models_route, "_scan_hf_cache", lambda *a, **k: []) monkeypatch.setattr(models_route, "_resolve_hf_cache_dir", lambda: tmp_path) @@ -713,6 +827,8 @@ def test_index_excludes_hidden_models(tmp_path, monkeypatch): index = resolver._index() assert "org/normal-gguf" in index # keys are normalized to lowercase assert "ggml-org/models" not in index + assert "unsloth/bge-small-en-v1.5-gguf" not in index + assert str(local_default_embedder).lower() not in index # And the hidden probe cannot be auto-switched to by name. resolver._scan = (0.0, {}) assert resolver.resolve_local_gguf("ggml-org/models") is None @@ -1609,11 +1725,11 @@ def test_env_idle_ttl_standalone_when_no_stored_value(monkeypatch): def test_stored_idle_value_overrides_env_and_stays_gated(monkeypatch): # An explicit stored value wins over the env default and remains gated on the # auto-switch toggle. - store = {settings.AUTO_UNLOAD_IDLE_SETTING_KEY: 30} + store = {settings.AUTO_UNLOAD_IDLE_SETTING_KEY: 90} monkeypatch.setattr(settings, "_cached_setting", lambda k, d = None: store.get(k, d)) monkeypatch.setenv("UNSLOTH_MODEL_IDLE_TTL", "600") monkeypatch.setattr(settings, "get_openai_auto_switch_enabled", lambda: True) - assert settings.get_auto_unload_idle_seconds() == 30 # stored wins, not env + assert settings.get_auto_unload_idle_seconds() == 90 # stored wins, not env monkeypatch.setattr(settings, "get_openai_auto_switch_enabled", lambda: False) assert settings.get_auto_unload_idle_seconds() == 0 # explicit value still gated off @@ -1729,6 +1845,8 @@ def test_index_advertises_alias_not_filesystem_path(tmp_path, monkeypatch): # host path in /v1/models, yet the model stays resolvable by that path too. from types import SimpleNamespace import routes.models as models_route + from storage import studio_db + import utils.paths as paths gguf = tmp_path / "model-Q4_K_M.gguf" gguf.write_bytes(b"x" * 32) @@ -1742,6 +1860,8 @@ def test_index_advertises_alias_not_filesystem_path(tmp_path, monkeypatch): monkeypatch.setattr(models_route, "_scan_hf_cache", lambda *a, **k: []) monkeypatch.setattr(models_route, "_resolve_hf_cache_dir", lambda: tmp_path) monkeypatch.setattr(models_route, "_is_hidden_model", lambda *a, **k: False) + monkeypatch.setattr(paths, "lmstudio_model_dirs", lambda: []) + monkeypatch.setattr(studio_db, "list_scan_folders", lambda: []) resolver._scan = (0.0, {}) # The advertised id is the alias, never the absolute path. @@ -2887,8 +3007,10 @@ def test_non_gguf_load_clears_reload_stash(): # A non-GGUF (Transformers/Unsloth) load must clear the stash like the GGUF # branch, so it never lingers until the idle poll (or forever, idle-unload off). import inspect + src = inspect.getsource(inference_route._load_model_impl) - assert src.count("note_model_loaded()") >= 2 + assert src.count("note_model_loaded()") >= 1 # non-GGUF branch + assert "to_thread(note_model_loaded, llama_backend)" in src # GGUF branch def test_chat_rejects_malformed_tool_choice_before_switch(monkeypatch): @@ -3094,3 +3216,542 @@ def test_responses_stream_hint_matches_toggle_regardless_of_active_model(monkeyp monkeypatch, enabled = False, active_model_name = "unsloth/Llama-3.2-1B-Instruct" ) assert "Model auto-switch" in non_gguf_loaded + + +# ── idle-unload KV persistence (slot save/restore) ────────────────── + + +def _seed_kv_manifest( + tmp_path, + identity = ("unsloth/A-GGUF", "Q4_K_M", "unsloth/A-GGUF"), + gguf = None, +): + if gguf is None: + gguf_file = tmp_path / "model.gguf" + gguf_file.write_bytes(b"gguf") + gguf = str(gguf_file) + st = os.stat(gguf) + state_file = tmp_path / "resume-abc-slot0.bin" + state_file.write_bytes(b"kv") + return state_file, { + "identity": identity, + "dir": str(tmp_path), + "binary": ("/bin/llama-server", 111), + "gguf": gguf, + "gguf_stat": ((st.st_size, st.st_mtime_ns),), + "launch": ((), None, None, 1), + "slots": [{"id": 0, "filename": state_file.name, "n_saved": 42}], + } + + +def _drive_idle_loop( + kw, + poll_seconds = 0.02, + run_for = 0.2, +): + async def _drive(): + task = asyncio.create_task(kw.idle_unload_loop(poll_seconds = poll_seconds)) + await asyncio.sleep(run_for) + task.cancel() + try: + await task + except asyncio.CancelledError: + pass + + asyncio.run(_drive()) + + +def test_idle_unload_saves_slots_before_unload_and_stashes_manifest(monkeypatch, tmp_path): + import time + from core.inference import llama_keepwarm as kw + + monkeypatch.setattr(settings, "get_auto_unload_idle_seconds", lambda: 0.005) + monkeypatch.setattr(settings, "get_auto_unload_keep_kv", lambda: True) + kw._inflight = 0 + kw._pending = 0 + kw._last_active = time.monotonic() - 3600 + kw._last_unloaded_model = None + kw._kv_resume = None + + events = [] + backend = _FakeBackend("unsloth/Idle-GGUF", hf_variant = "Q4_K_M") + manifest = { + "dir": str(tmp_path), + "binary": ("bin", 1), + "slots": [{"id": 0, "filename": "f.bin", "n_saved": 42}], + } + + def _save(should_abort = None): + events.append("save") + return manifest + + def _unload(): + events.append("unload") + backend.is_loaded = False + + backend.save_slots_for_resume = _save + backend.unload_model = _unload + monkeypatch.setattr(inference_route, "get_llama_cpp_backend", lambda: backend) + + _drive_idle_loop(kw) + # KV must be saved while the server is still alive, then exactly one unload. + assert events == ["save", "unload"] + assert kw.get_last_unloaded_model()[:2] == ("unsloth/Idle-GGUF", "Q4_K_M") + resume = kw.take_kv_resume() + assert resume is not None + assert resume["identity"][:2] == ("unsloth/Idle-GGUF", "Q4_K_M") + assert resume["slots"][0]["filename"] == "f.bin" + + +def test_idle_save_failure_still_unloads_plain(monkeypatch): + import time + from core.inference import llama_keepwarm as kw + + monkeypatch.setattr(settings, "get_auto_unload_idle_seconds", lambda: 0.005) + monkeypatch.setattr(settings, "get_auto_unload_keep_kv", lambda: True) + kw._inflight = 0 + kw._pending = 0 + kw._last_active = time.monotonic() - 3600 + kw._last_unloaded_model = None + kw._kv_resume = None + + unloads = [] + backend = _FakeBackend("unsloth/Idle-GGUF", hf_variant = "Q4_K_M") + + def _save(should_abort = None): + raise RuntimeError("slot save exploded") + + def _unload(): + unloads.append(1) + backend.is_loaded = False + + backend.save_slots_for_resume = _save + backend.unload_model = _unload + monkeypatch.setattr(inference_route, "get_llama_cpp_backend", lambda: backend) + + _drive_idle_loop(kw) + assert unloads == [1] # the save failure must not skip the unload + assert kw.get_last_unloaded_model() is not None + assert kw.take_kv_resume() is None + + +def test_keep_kv_setting_off_skips_save(monkeypatch): + import time + from core.inference import llama_keepwarm as kw + + monkeypatch.setattr(settings, "get_auto_unload_idle_seconds", lambda: 0.005) + monkeypatch.setattr(settings, "get_auto_unload_keep_kv", lambda: False) + kw._inflight = 0 + kw._pending = 0 + kw._last_active = time.monotonic() - 3600 + kw._last_unloaded_model = None + kw._kv_resume = None + + saves, unloads = [], [] + backend = _FakeBackend("unsloth/Idle-GGUF") + + def _unload(): + unloads.append(1) + backend.is_loaded = False + + backend.save_slots_for_resume = lambda *a, **k: saves.append(1) + backend.unload_model = _unload + monkeypatch.setattr(inference_route, "get_llama_cpp_backend", lambda: backend) + + _drive_idle_loop(kw) + assert saves == [] + assert unloads == [1] + assert kw.take_kv_resume() is None + + +def test_keep_kv_disabled_mid_save_discards_manifest(monkeypatch, tmp_path): + import time + from core.inference import llama_keepwarm as kw + + keep = {"on": True} + monkeypatch.setattr(settings, "get_auto_unload_idle_seconds", lambda: 0.005) + monkeypatch.setattr(settings, "get_auto_unload_keep_kv", lambda: keep["on"]) + kw._inflight = 0 + kw._pending = 0 + kw._last_active = time.monotonic() - 3600 + kw._last_unloaded_model = None + kw._kv_resume = None + + unloads = [] + backend = _FakeBackend("unsloth/Idle-GGUF", hf_variant = "Q4_K_M") + state_file = tmp_path / "resume-mid-slot0.bin" + state_file.write_bytes(b"kv") + manifest = { + "dir": str(tmp_path), + "binary": ("bin", 1), + "slots": [{"id": 0, "filename": state_file.name, "n_saved": 1}], + } + + def _save(should_abort = None): + keep["on"] = False # user flips the toggle while the save runs + return manifest + + def _unload(): + unloads.append(1) + backend.is_loaded = False + + backend.save_slots_for_resume = _save + backend.unload_model = _unload + monkeypatch.setattr(inference_route, "get_llama_cpp_backend", lambda: backend) + + _drive_idle_loop(kw) + assert unloads == [1] # still unloads; only the stash is dropped + assert kw.take_kv_resume() is None + assert not state_file.exists() + + +def test_idle_ttl_disabled_mid_save_skips_unload(monkeypatch, tmp_path): + import time + from core.inference import llama_keepwarm as kw + + ttl = {"v": 0.005} + monkeypatch.setattr(settings, "get_auto_unload_idle_seconds", lambda: ttl["v"]) + monkeypatch.setattr(settings, "get_auto_unload_keep_kv", lambda: True) + kw._inflight = 0 + kw._pending = 0 + kw._last_active = time.monotonic() - 3600 + kw._last_unloaded_model = None + kw._kv_resume = None + + unloads = [] + backend = _FakeBackend("unsloth/Idle-GGUF", hf_variant = "Q4_K_M") + state_file = tmp_path / "resume-mid-slot0.bin" + state_file.write_bytes(b"kv") + manifest = { + "dir": str(tmp_path), + "binary": ("bin", 1), + "slots": [{"id": 0, "filename": state_file.name, "n_saved": 1}], + } + + def _save(should_abort = None): + ttl["v"] = 0 # user turns idle unload off while the save runs + return manifest + + backend.save_slots_for_resume = _save + backend.unload_model = lambda: unloads.append(1) + monkeypatch.setattr(inference_route, "get_llama_cpp_backend", lambda: backend) + + _drive_idle_loop(kw) + assert unloads == [] # the unload was cancelled by the setting change + assert kw.take_kv_resume() is None + assert not state_file.exists() + + +def test_alias_reload_restores_slots_and_deletes_files(monkeypatch, tmp_path): + from core.inference import llama_keepwarm as kw + + backend = _FakeBackend(None) # idle-unload emptied the backend + backend._slot_save_binary = ("/bin/llama-server", 111) + restored = [] + backend.restore_slots_for_resume = lambda manifest: restored.append(manifest) + + rec = _LoadRecorder(backend) + _wire(monkeypatch, enabled = True, resolves_to = None, backend = backend, recorder = rec) + monkeypatch.setattr(kw, "_inflight", 0) + state_file, manifest = _seed_kv_manifest(tmp_path) + monkeypatch.setattr(kw, "_last_unloaded_model", (manifest["gguf"], "Q4_K_M")) + monkeypatch.setattr(kw, "_kv_resume", manifest) + + _run_hook("gpt-4o-mini") + assert len(rec.calls) == 1 + assert len(restored) == 1 # same model + binary: restore ran + assert not state_file.exists() # state file deleted after the restore + assert kw._kv_resume is None + + +def test_no_restore_when_different_model_loads(monkeypatch, tmp_path): + from core.inference import llama_keepwarm as kw + + backend = _FakeBackend(None) + backend._slot_save_binary = ("/bin/llama-server", 111) + restored = [] + backend.restore_slots_for_resume = lambda manifest: restored.append(manifest) + rec = _LoadRecorder(backend) + _wire( + monkeypatch, + enabled = True, + resolves_to = ("unsloth/B-GGUF", None, "unsloth/B-GGUF"), + backend = backend, + recorder = rec, + ) + monkeypatch.setattr(kw, "_inflight", 0) + state_file, manifest = _seed_kv_manifest(tmp_path) # manifest is for model A + monkeypatch.setattr(kw, "_kv_resume", manifest) + + _run_hook("unsloth/B-GGUF") + assert len(rec.calls) == 1 + assert restored == [] # different model: never restored + assert not state_file.exists() # but the stale files are gone + assert kw._kv_resume is None + + +def test_restore_skipped_when_binary_changed(monkeypatch, tmp_path): + from core.inference import llama_keepwarm as kw + + state_file, manifest = _seed_kv_manifest(tmp_path) + backend = _FakeBackend("unsloth/A-GGUF", hf_variant = "Q4_K_M") + backend._gguf_path = manifest["gguf"] + backend._slot_save_binary = ("/bin/llama-server", 222) # newer mtime + restored = [] + backend.restore_slots_for_resume = lambda manifest: restored.append(manifest) + + kw.restore_kv_resume(backend, manifest) + assert restored == [] + assert not state_file.exists() + + +def test_restore_skipped_when_launch_config_changed(tmp_path): + from core.inference import llama_keepwarm as kw + + state_file, manifest = _seed_kv_manifest(tmp_path) + backend = _FakeBackend("unsloth/A-GGUF", hf_variant = "Q4_K_M") + backend._gguf_path = manifest["gguf"] + backend._slot_save_binary = ("/bin/llama-server", 111) + backend._slot_launch_fingerprint = lambda: (("--rope-freq-scale", "0.5"), None, None, 1) + restored = [] + backend.restore_slots_for_resume = lambda manifest: restored.append(manifest) + + kw.restore_kv_resume(backend, manifest) + assert restored == [] + assert not state_file.exists() + + +def test_restore_skipped_when_gguf_rewritten_in_place(tmp_path): + from core.inference import llama_keepwarm as kw + + state_file, manifest = _seed_kv_manifest(tmp_path) + with open(manifest["gguf"], "wb") as fh: + fh.write(b"different weights") # same path, new content + backend = _FakeBackend("unsloth/A-GGUF", hf_variant = "Q4_K_M") + backend._gguf_path = manifest["gguf"] + backend._slot_save_binary = ("/bin/llama-server", 111) + restored = [] + backend.restore_slots_for_resume = lambda manifest: restored.append(manifest) + + kw.restore_kv_resume(backend, manifest) + assert restored == [] + assert not state_file.exists() + + +def test_note_model_unloaded_purges_manifest_and_files(tmp_path): + from core.inference import llama_keepwarm as kw + + state_file, manifest = _seed_kv_manifest(tmp_path) + kw._set_last_unloaded(("org/A-GGUF", "Q4_K_M")) + kw._set_kv_resume(manifest) + kw.note_model_unloaded() + assert kw.get_last_unloaded_model() is None + assert kw.take_kv_resume() is None + assert not state_file.exists() + + +def test_note_model_loaded_purges_manifest_and_files(tmp_path): + from core.inference import llama_keepwarm as kw + + state_file, manifest = _seed_kv_manifest(tmp_path) + kw._set_last_unloaded(("org/A-GGUF", "Q4_K_M")) + kw._set_kv_resume(manifest) + kw.note_model_loaded() + assert kw.get_last_unloaded_model() is None + assert kw.take_kv_resume() is None + assert not state_file.exists() + + +def test_new_idle_save_purges_previous_manifest_files(tmp_path): + from core.inference import llama_keepwarm as kw + + old_file, old_manifest = _seed_kv_manifest(tmp_path) + kw._set_kv_resume(old_manifest) + new_file = tmp_path / "resume-def-slot0.bin" + new_file.write_bytes(b"kv2") + kw._set_kv_resume( + { + "identity": ("unsloth/B-GGUF", None, "unsloth/B-GGUF"), + "dir": str(tmp_path), + "binary": ("/bin/llama-server", 111), + "slots": [{"id": 0, "filename": new_file.name, "n_saved": 7}], + } + ) + assert not old_file.exists() # replaced manifest's files purged + assert new_file.exists() + assert kw.take_kv_resume()["slots"][0]["filename"] == new_file.name + + +def test_sweep_slot_save_dir_removes_only_resume_files(monkeypatch, tmp_path): + from core.inference import llama_keepwarm as kw + from utils.paths import storage_roots + + monkeypatch.setattr(storage_roots, "llama_slot_cache_root", lambda: tmp_path) + stale = tmp_path / "resume-old-slot0.bin" + stale.write_bytes(b"kv") + other = tmp_path / "unrelated.txt" + other.write_text("keep") + kw.sweep_slot_save_dir() + assert not stale.exists() + assert other.exists() + + +def test_keep_kv_setting_roundtrip_and_default(monkeypatch): + import storage.studio_db as db + + store = {} + monkeypatch.setattr(db, "upsert_app_settings", lambda m: store.update(m)) + monkeypatch.setattr(settings, "_cached_setting", lambda k, d = None: store.get(k, d)) + + assert settings.get_auto_unload_keep_kv() is True # default when never stored + assert settings.set_openai_auto_switch(True, 60, False)[2] is False + assert store[settings.AUTO_UNLOAD_KEEP_KV_SETTING_KEY] is False + assert settings.get_auto_unload_keep_kv() is False + # None leaves the stored value untouched (older clients can't reset it). + assert settings.set_openai_auto_switch(True, 60, None)[2] is False + assert store[settings.AUTO_UNLOAD_KEEP_KV_SETTING_KEY] is False + with pytest.raises(ValueError, match = "true or false"): + settings.set_openai_auto_switch(True, 60, "garbage") + + +def test_stale_stash_cleanup_waits_for_lifecycle_gate(monkeypatch, tmp_path): + # The loop's stale-stash purge must wait on the gate a mid-reload holds. + import time + from core.inference import llama_keepwarm as kw + + monkeypatch.setattr(settings, "get_auto_unload_idle_seconds", lambda: 3600) + kw._inflight = 0 + kw._pending = 0 + kw._last_active = time.monotonic() + backend = _FakeBackend("unsloth/New-GGUF") + monkeypatch.setattr(inference_route, "get_llama_cpp_backend", lambda: backend) + state_file, manifest = _seed_kv_manifest(tmp_path) + kw._kv_resume = manifest + kw._last_unloaded_model = ("unsloth/A-GGUF", "Q4_K_M") + + assert kw._lifecycle_lock.acquire(blocking = False) # simulate in-flight reload + try: + _drive_idle_loop(kw) + assert kw._kv_resume is manifest # purge deferred while the gate is held + assert state_file.exists() + finally: + kw._lifecycle_lock.release() + _drive_idle_loop(kw) + assert kw._kv_resume is None # gate freed: genuinely stale stash purged + assert not state_file.exists() + + +def test_put_route_disabling_keep_kv_purges_saved_state(monkeypatch, tmp_path): + import routes.settings as settings_route + import storage.studio_db as db + from core.inference import llama_keepwarm as kw + + store = {} + monkeypatch.setattr(db, "upsert_app_settings", lambda m: store.update(m)) + monkeypatch.setattr(settings, "_cached_setting", lambda k, d = None: store.get(k, d)) + state_file, manifest = _seed_kv_manifest(tmp_path) + monkeypatch.setattr(kw, "_kv_resume", manifest) + + payload = settings_route.OpenAIAutoSwitchPayload(enabled = True, auto_unload_keep_kv = False) + resp = settings_route.update_openai_auto_switch(payload, "tester") + assert resp.auto_unload_keep_kv is False + assert kw._kv_resume is None + assert not state_file.exists() + + +def test_keep_kv_only_update_leaves_env_idle_ttl_active(monkeypatch): + # A keep-KV-only update must not materialize the env TTL as a stored value. + import routes.settings as settings_route + import storage.studio_db as db + + store = {} + monkeypatch.setattr(db, "upsert_app_settings", lambda m: store.update(m)) + monkeypatch.setattr(settings, "_cached_setting", lambda k, d = None: store.get(k, d)) + monkeypatch.setenv(settings.MODEL_IDLE_TTL_ENV_VAR, "600") + + assert settings_route.OpenAIAutoSwitchPayload(enabled = False).auto_unload_idle_seconds is None + enabled, idle, keep_kv = settings.set_openai_auto_switch(False, None, False) + assert settings.AUTO_UNLOAD_IDLE_SETTING_KEY not in store # idle untouched + assert settings.get_auto_unload_idle_seconds() == 600 # env TTL still active + assert (enabled, idle, keep_kv) == (False, 600, False) + + +def test_load_impl_notes_loaded_with_backend_off_loop(): + import inspect + src = inspect.getsource(inference_route._load_model_impl) + assert "to_thread(note_model_loaded, llama_backend)" in src + + +def test_restore_matches_gguf_realpath_across_naming(tmp_path): + from core.inference import llama_keepwarm as kw + + blob = tmp_path / "blob.gguf" + blob.write_bytes(b"gguf") + link = tmp_path / "snapshot.gguf" + try: + link.symlink_to(blob) + except OSError: + pytest.skip("symlinks unsupported on this host") + + backend = _FakeBackend("/hf/snapshots/d7f5", hf_variant = None) + backend._gguf_path = str(link) # reload resolved the symlink spelling + backend._slot_save_binary = ("/bin/llama-server", 111) + restored = [] + backend.restore_slots_for_resume = lambda manifest: restored.append(manifest) + state_file, manifest = _seed_kv_manifest( + tmp_path, identity = ("unsloth/A-GGUF", None, "unsloth/A-GGUF"), gguf = str(blob) + ) + + kw.restore_kv_resume(backend, manifest) + assert len(restored) == 1 # names differ, file identical: restore ran + assert not state_file.exists() + + +def test_setter_rejects_idle_below_floor(monkeypatch): + import storage.studio_db as db + + writes = [] + monkeypatch.setattr(db, "upsert_app_settings", lambda m: writes.append(dict(m))) + settings._cache.clear() + + with pytest.raises(ValueError, match = "at least 60"): + settings.set_openai_auto_switch(True, 30) + assert writes == [] # rejected before any persist + # 0 (off) and >= 60 pass through unchanged. + assert settings.set_openai_auto_switch(True, 0)[1] == 0 + assert settings.set_openai_auto_switch(True, 60)[1] == 60 + assert settings.set_openai_auto_switch(True, 3600)[1] == 3600 + + +def test_put_route_rejects_idle_below_floor(): + import routes.settings as settings_route + from fastapi import HTTPException + + payload = settings_route.OpenAIAutoSwitchPayload(enabled = True, auto_unload_idle_seconds = 30) + with pytest.raises(HTTPException) as excinfo: + settings_route.update_openai_auto_switch(payload, "tester") + assert excinfo.value.status_code == 400 + + +def test_stored_legacy_idle_below_floor_is_clamped(monkeypatch): + # Values persisted before the floor existed are raised to it on read, for + # both the effective TTL and the value the settings UI displays. + store = {settings.AUTO_UNLOAD_IDLE_SETTING_KEY: 5} + monkeypatch.setattr(settings, "_cached_setting", lambda k, d = None: store.get(k, d)) + monkeypatch.setattr(settings, "get_openai_auto_switch_enabled", lambda: True) + assert settings.get_auto_unload_idle_seconds() == 60 + assert settings.get_stored_auto_unload_idle_seconds() == 60 + store[settings.AUTO_UNLOAD_IDLE_SETTING_KEY] = 90 + assert settings.get_auto_unload_idle_seconds() == 90 + + +def test_env_idle_below_floor_is_clamped(monkeypatch): + monkeypatch.setattr(settings, "_cached_setting", lambda k, d = None: d) + monkeypatch.setenv(settings.MODEL_IDLE_TTL_ENV_VAR, "5") + assert settings.get_auto_unload_idle_seconds() == 60 + monkeypatch.setenv(settings.MODEL_IDLE_TTL_ENV_VAR, "0") + assert settings.get_auto_unload_idle_seconds() == 0 + monkeypatch.setenv(settings.MODEL_IDLE_TTL_ENV_VAR, "600") + assert settings.get_auto_unload_idle_seconds() == 600 + monkeypatch.delenv(settings.MODEL_IDLE_TTL_ENV_VAR) + assert settings.get_auto_unload_idle_seconds() == 0 diff --git a/studio/backend/tests/test_openai_compaction.py b/studio/backend/tests/test_openai_compaction.py index c7de0a9aed..6fad2c5eaf 100644 --- a/studio/backend/tests/test_openai_compaction.py +++ b/studio/backend/tests/test_openai_compaction.py @@ -86,7 +86,7 @@ def test_cloud_openai_sets_compaction_block(monkeypatch): def test_cloud_openai_below_default_threshold_passes_through(monkeypatch): - # Studio doesn't clamp the OpenAI side -- the API accepts whatever the + # Unsloth doesn't clamp the OpenAI side -- the API accepts whatever the # caller sends, so a small probe like 60k still goes through. captured = _capture( monkeypatch, diff --git a/studio/backend/tests/test_openai_image_generation.py b/studio/backend/tests/test_openai_image_generation.py index ace57588d3..c2eef0381f 100644 --- a/studio/backend/tests/test_openai_image_generation.py +++ b/studio/backend/tests/test_openai_image_generation.py @@ -4,7 +4,7 @@ """Unit tests for OpenAI Responses API image_generation tool wiring. The tool is a server-side Responses-API tool (``{type: "image_generation"}``); -the result comes back as an ``image_generation_call`` output item, which Studio +the result comes back as an ``image_generation_call`` output item, which Unsloth translates into ``_toolEvent`` chunks so the chat adapter renders it inline. Tests pin: the tool is added to the body only on a cloud OpenAI base when asked for, the done event produces the expected chunks, and non-cloud bases drop it. diff --git a/studio/backend/tests/test_openai_tool_passthrough.py b/studio/backend/tests/test_openai_tool_passthrough.py index 8725b28ac8..161c8743c4 100644 --- a/studio/backend/tests/test_openai_tool_passthrough.py +++ b/studio/backend/tests/test_openai_tool_passthrough.py @@ -119,7 +119,7 @@ class TestFriendlyUpstreamError: raw = '{"error":{"code":400,"message":"Failed to initialize samplers: failed to parse grammar","type":"invalid_request_error"}}' msg = _friendly_upstream_error(raw) assert "failed to parse grammar" not in msg # raw body is not surfaced verbatim - assert "tool-calling grammar" in msg and "Update Studio" in msg + assert "tool-calling grammar" in msg and "Update Unsloth" in msg def test_failed_to_initialize_samplers_alone_matches(self): assert "tool-calling grammar" in _friendly_upstream_error("Failed to initialize samplers") @@ -262,7 +262,7 @@ class TestChatMessageToolRoles: def test_tool_empty_content_accepted(self): # Empty tool output (mkdir, git add, ...) is routine in agentic loops; - # OpenAI and llama-server both accept it, so Studio must not 400. + # OpenAI and llama-server both accept it, so Unsloth must not 400. msg = ChatMessage(role = "tool", tool_call_id = "call_1", content = "") assert msg.content == "" @@ -400,7 +400,7 @@ class TestChatCompletionRequestToolFields: assert req.session_id == "abc" def test_stream_defaults_false_matching_openai_spec(self): - # OpenAI defaults `stream` to false. Studio used to default true, + # OpenAI defaults `stream` to false. Unsloth used to default true, # breaking naive curl/.NET clients (#5047) that omit it. Pin the fix. req = self._make() assert req.stream is False @@ -664,7 +664,7 @@ class TestChatCompletionRequestToolFields: raise AssertionError("client tools must use passthrough") def generate_chat_completion_with_tools(self, **_kwargs): - raise AssertionError("Studio tool loop must stay disabled") + raise AssertionError("Unsloth tool loop must stay disabled") async def fake_passthrough(llama_backend, payload, model_name, **kwargs): captured["body"] = inference_route._build_openai_passthrough_body( @@ -707,11 +707,11 @@ class TestChatCompletionRequestToolFields: assert monitor.active_count() == 0 def test_permission_mode_does_not_reject_client_tool_passthrough(self, monkeypatch): - # A non-streaming client-tool passthrough (client tools, no Studio tool + # A non-streaming client-tool passthrough (client tools, no Unsloth tool # loop) that also carries permission_mode "ask"/"auto" must reach the # provider passthrough, not the confirm-without-stream guard: the # validator leaves confirm_tool_calls unset for passthrough, and a bare - # permission_mode only gates Studio's own local tool loop. An explicit + # permission_mode only gates Unsloth's own local tool loop. An explicit # confirm_tool_calls=True still forces the local-confirm rejection. # The pre-switch guard only runs when an automatic load may run, so force # that predicate on to exercise it against a resident passthrough backend. @@ -732,7 +732,7 @@ class TestChatCompletionRequestToolFields: raise AssertionError("client tools must use passthrough") def generate_chat_completion_with_tools(self, **_kwargs): - raise AssertionError("Studio tool loop must stay disabled") + raise AssertionError("Unsloth tool loop must stay disabled") async def fake_passthrough(llama_backend, payload, model_name, **kwargs): inference_route.api_monitor.finish(kwargs.get("monitor_id")) @@ -757,7 +757,7 @@ class TestChatCompletionRequestToolFields: return self._v1_client(monkeypatch, _GGUFBackend()) # A process --enable-tools policy must not turn a client-tool passthrough - # into a Studio local loop, so a policy of None or True both keep the + # into an Unsloth local loop, so a policy of None or True both keep the # passthrough (the guard mirrors _explicit_studio_tool_loop_requested). for policy in (None, True): for mode in ("ask", "auto"): @@ -810,7 +810,7 @@ class TestChatCompletionRequestToolFields: assert "requires stream=true" in resp.json()["error"]["message"] def test_permission_mode_policy_forced_local_loop_rejected_before_switch(self, monkeypatch): - # A process --enable-tools policy forces Studio's own tool loop on even + # A process --enable-tools policy forces Unsloth's own tool loop on even # when the request omits enable_tools and carries no client tools. A # non-streaming ask/auto request is then confirm-gated with no stream to # prompt on, so it must 400 at the pre-switch guard -- before @@ -863,7 +863,7 @@ class TestChatCompletionRequestToolFields: def test_enable_tools_on_non_tool_backend_keeps_client_tools_on_passthrough(self, monkeypatch): # DiffusionGemma forces supports_tools off while passthrough stays # available (#6851): enable_tools=True must not steal client tools - # from the passthrough into a Studio tool loop that cannot run. + # from the passthrough into an Unsloth tool loop that cannot run. import routes.inference as inference_route captured = {} @@ -883,7 +883,7 @@ class TestChatCompletionRequestToolFields: raise AssertionError("client tools must use passthrough") def generate_chat_completion_with_tools(self, **_kwargs): - raise AssertionError("Studio tool loop cannot run on a non-tool backend") + raise AssertionError("Unsloth tool loop cannot run on a non-tool backend") async def fake_passthrough(llama_backend, payload, model_name, **kwargs): captured["body"] = inference_route._build_openai_passthrough_body( @@ -2581,7 +2581,7 @@ class TestGgufVisionToolRouting: raise AssertionError("plain GGUF path should not be used") def _tools(**_kwargs): - raise AssertionError("Studio tool loop should not steal response_format") + raise AssertionError("Unsloth tool loop should not steal response_format") backend = SimpleNamespace( is_loaded = True, @@ -2654,7 +2654,7 @@ class TestGgufVisionToolRouting: raise AssertionError("plain GGUF path should not be used") def _tools(**_kwargs): - raise AssertionError("Studio tool loop should not replace client tools") + raise AssertionError("Unsloth tool loop should not replace client tools") backend = SimpleNamespace( is_loaded = True, @@ -2726,7 +2726,7 @@ class TestGgufVisionToolRouting: yield "plain response" def _tools(**_kwargs): - raise AssertionError("tool_choice='none' must not start Studio's tool loop") + raise AssertionError("tool_choice='none' must not start Unsloth's tool loop") backend = SimpleNamespace( is_loaded = True, @@ -2780,7 +2780,7 @@ class TestGgufVisionToolRouting: raise AssertionError("plain GGUF path should not be used") def _tools(**_kwargs): - raise AssertionError("enabled_tools alone must not start Studio's tool loop") + raise AssertionError("enabled_tools alone must not start Unsloth's tool loop") backend = SimpleNamespace( is_loaded = True, @@ -2844,7 +2844,7 @@ class TestGgufVisionToolRouting: raise AssertionError("plain GGUF path should not be used") def _tools(**_kwargs): - raise AssertionError("enabled_tools alone must not start Studio's tool loop") + raise AssertionError("enabled_tools alone must not start Unsloth's tool loop") backend = SimpleNamespace( is_loaded = True, diff --git a/studio/backend/tests/test_orchestrator_unload_cancel.py b/studio/backend/tests/test_orchestrator_unload_cancel.py index fb80b6d061..3a36500aee 100644 --- a/studio/backend/tests/test_orchestrator_unload_cancel.py +++ b/studio/backend/tests/test_orchestrator_unload_cancel.py @@ -34,6 +34,70 @@ def _bare_orchestrator(): return o +def test_adapter_control_raises_stream_errors(monkeypatch): + o = _bare_orchestrator() + monkeypatch.setattr( + o, + "_generate_dispatched", + lambda **_kwargs: iter([orch_mod.GenStreamError("Error: adapter failed")]), + ) + + with pytest.raises(RuntimeError, match = "adapter failed"): + list(o.generate_with_adapter_control(use_adapter = False)) + + closed = [] + + def _stream(**_kwargs): + try: + yield "token" + yield "late token" + finally: + closed.append(True) + + monkeypatch.setattr(o, "_generate_dispatched", _stream) + generator = o.generate_with_adapter_control(use_adapter = False) + assert next(generator) == "token" + generator.close() + assert closed == [True] + + +def test_worker_closes_cancelled_generator_before_gen_done(): + from core.inference.worker import _handle_generate + + events = [] + + class _Backend: + last_generation_stats = None + + def generate_with_adapter_control(self, **_kwargs): + try: + yield "token" + yield "late token" + finally: + events.append("closed") + + class _Responses: + def __init__(self): + self.items = [] + + def put(self, item): + if item["type"] == "gen_done": + assert events == ["closed"] + self.items.append(item) + + responses = _Responses() + cancel = threading.Event() + cancel.set() + _handle_generate( + _Backend(), + {"request_id": "r1", "messages": [], "use_adapter": False}, + responses, + cancel, + ) + + assert [item["type"] for item in responses.items] == ["gen_done"] + + def test_unload_cancels_inflight_generation_then_unloads(monkeypatch): o = _bare_orchestrator() monkeypatch.setattr(o, "_ensure_subprocess_alive", lambda: True) diff --git a/studio/backend/tests/test_password_prompt_backstop.py b/studio/backend/tests/test_password_prompt_backstop.py index 597eac1625..3c2c1956f9 100644 --- a/studio/backend/tests/test_password_prompt_backstop.py +++ b/studio/backend/tests/test_password_prompt_backstop.py @@ -3,7 +3,7 @@ """Pre-tunnel terminal password gate: never publish a public Cloudflare URL while the seeded default admin password is active. Imports run.py directly, -so run under the Studio venv.""" +so run under the Unsloth venv.""" from __future__ import annotations diff --git a/studio/backend/tests/test_permission_mode.py b/studio/backend/tests/test_permission_mode.py index 3b7197fc49..4fc64a6291 100644 --- a/studio/backend/tests/test_permission_mode.py +++ b/studio/backend/tests/test_permission_mode.py @@ -1438,7 +1438,7 @@ def test_unknown_permission_mode_normalizes_to_ask_on_request_models(): def test_ask_auto_self_enable_confirm_on_chat_request(): # "Ask" gates every call, so a direct /chat/completions caller that requests - # ask but omits the legacy confirm flag self-enables it when Studio's own tool + # ask but omits the legacy confirm flag self-enables it when Unsloth's own tool # loop is requested. Only the router's loop-entry signals count (enable_tools / # mcp_enabled); enabled_tools alone never starts the loop. for loop in ({"enable_tools": True}, {"mcp_enabled": True}): @@ -1481,7 +1481,7 @@ def test_ask_auto_self_enable_confirm_on_chat_request(): confirm_tool_calls = False, ) assert req.confirm_tool_calls is False - # A plain client-tool passthrough (client-supplied tools that Studio does not + # A plain client-tool passthrough (client-supplied tools that Unsloth does not # execute) must NOT self-enable confirm, or the route rejects the passthrough. req = ChatCompletionRequest( messages = [{"role": "user", "content": "hi"}], diff --git a/studio/backend/tests/test_picker_service.py b/studio/backend/tests/test_picker_service.py new file mode 100644 index 0000000000..be7ea18f03 --- /dev/null +++ b/studio/backend/tests/test_picker_service.py @@ -0,0 +1,266 @@ +# SPDX-License-Identifier: AGPL-3.0-only +# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0 + +import json +from types import SimpleNamespace + +from picker.service import ( + MAX_TEMPLATE_METADATA_BYTES, + _chat_template_from_dir, + _chat_template_from_processor_json, + _chat_template_from_tokenizer_config, + _chat_template_from_tokenizer_dir, + _find_gguf_in_dir, + _iter_ggufs, + read_default_chat_template, + validate_chat_template, +) + + +def test_iter_ggufs_skips_gguf_companions(tmp_path): + mtp_dir = tmp_path / "MTP" + mtp_dir.mkdir() + main = tmp_path / "model-Q8_0.gguf" + main.write_bytes(b"") + (tmp_path / "mmproj-F16.gguf").write_bytes(b"") + (tmp_path / "mtp-model-Q8_0.gguf").write_bytes(b"") + (mtp_dir / "model-Q8_0-MTP.gguf").write_bytes(b"") + (tmp_path / "model-Q8_0-be.gguf").write_bytes(b"") + + assert _iter_ggufs(tmp_path) == [main] + + +def test_find_gguf_in_dir_matches_quant_label(tmp_path): + mtp_dir = tmp_path / "MTP" + mtp_dir.mkdir() + main = tmp_path / "model-Q8_0.gguf" + main.write_bytes(b"") + (mtp_dir / "model-Q8_0-MTP.gguf").write_bytes(b"") + (tmp_path / "model-Q4_K_M.gguf").write_bytes(b"") + + assert _find_gguf_in_dir(tmp_path, "Q8_0") == main + assert _find_gguf_in_dir(tmp_path, "Q4_K") is None + + +def test_find_gguf_in_dir_without_variant_prefers_largest_model(tmp_path): + smaller = tmp_path / "a-model-Q4_K_M.gguf" + larger = tmp_path / "z-model-Q8_0.gguf" + smaller.write_bytes(b"0") + larger.write_bytes(b"00") + + assert _find_gguf_in_dir(tmp_path, None) == larger + + +def test_find_gguf_in_dir_without_variant_prefers_first_split(tmp_path): + first = tmp_path / "model-Q4_K_M-00001-of-00003.gguf" + second = tmp_path / "model-Q4_K_M-00002-of-00003.gguf" + third = tmp_path / "model-Q4_K_M-00003-of-00003.gguf" + first.write_bytes(b"0") + second.write_bytes(b"000") + third.write_bytes(b"00") + + assert _find_gguf_in_dir(tmp_path, None) == first + + first.unlink() + assert _find_gguf_in_dir(tmp_path, None) == second + + +def test_find_gguf_in_dir_matches_bpw_variant_base_label(tmp_path): + target = tmp_path / "model-IQ4_XS-3.53bpw.gguf" + target.write_bytes(b"") + (tmp_path / "model-Q4_K_M.gguf").write_bytes(b"") + + assert _find_gguf_in_dir(tmp_path, "IQ4_XS") == target + assert _find_gguf_in_dir(tmp_path, "IQ4_XS-3.53bpw") == target + assert _find_gguf_in_dir(tmp_path, "Q4_K") is None + + +def test_validate_chat_template_accepts_valid_and_empty(): + assert validate_chat_template("{{ messages[0].content }}").valid is True + assert validate_chat_template("").valid is True + assert validate_chat_template(" ").valid is True + + +def test_validate_chat_template_reports_syntax_error_with_line(): + result = validate_chat_template("{% if %}{% endif %}") + assert result.valid is False + assert result.error is not None + assert result.error.startswith("Line ") + + +def test_chat_template_from_tokenizer_config_reads_string(): + assert _chat_template_from_tokenizer_config({"chat_template": "HELLO"}) == "HELLO" + assert _chat_template_from_tokenizer_config({"chat_template": " "}) is None + assert _chat_template_from_tokenizer_config({}) is None + + +def test_chat_template_from_tokenizer_config_prefers_named_default(): + config = { + "chat_template": [ + {"name": "tool_use", "template": "TOOL"}, + {"name": "default", "template": "DEFAULT"}, + ] + } + assert _chat_template_from_tokenizer_config(config) == "DEFAULT" + + +def test_chat_template_from_tokenizer_config_falls_back_to_first_entry(): + config = { + "chat_template": [ + {"name": "tool_use", "template": "TOOL"}, + {"name": "other", "template": "OTHER"}, + ] + } + assert _chat_template_from_tokenizer_config(config) == "TOOL" + + +def test_chat_template_from_tokenizer_dir_prefers_jinja_file(tmp_path): + (tmp_path / "chat_template.jinja").write_text("FROM_JINJA", encoding = "utf-8") + (tmp_path / "tokenizer_config.json").write_text( + json.dumps({"chat_template": "FROM_CONFIG"}), encoding = "utf-8" + ) + assert _chat_template_from_tokenizer_dir(tmp_path) == "FROM_JINJA" + + +def test_chat_template_from_tokenizer_dir_reads_tokenizer_config(tmp_path): + (tmp_path / "tokenizer_config.json").write_text( + json.dumps({"chat_template": "FROM_CONFIG"}), encoding = "utf-8" + ) + assert _chat_template_from_tokenizer_dir(tmp_path) == "FROM_CONFIG" + + +def test_chat_template_from_dir_without_variant_prefers_tokenizer(tmp_path): + (tmp_path / "tokenizer_config.json").write_text( + json.dumps({"chat_template": "FROM_CONFIG"}), encoding = "utf-8" + ) + assert _chat_template_from_dir(tmp_path) == "FROM_CONFIG" + + +def test_chat_template_from_dir_with_variant_still_prefers_tokenizer(tmp_path, monkeypatch): + (tmp_path / "tokenizer_config.json").write_text( + json.dumps({"chat_template": "FROM_CONFIG"}), encoding = "utf-8" + ) + (tmp_path / "model-Q4_K_M.gguf").write_bytes(b"") + monkeypatch.setattr("picker.service.read_gguf_chat_template", lambda _path: "FROM_GGUF") + # Selecting a variant must not flip precedence to the embedded GGUF template. + assert _chat_template_from_dir(tmp_path, "Q4_K_M") == "FROM_CONFIG" + + +def test_chat_template_from_dir_with_variant_falls_back_to_gguf(tmp_path, monkeypatch): + (tmp_path / "model-Q4_K_M.gguf").write_bytes(b"") + monkeypatch.setattr("picker.service.read_gguf_chat_template", lambda _path: "FROM_GGUF") + # With no tokenizer sidecar, the embedded GGUF template is still the fallback. + assert _chat_template_from_dir(tmp_path, "Q4_K_M") == "FROM_GGUF" + + +def test_chat_template_from_dir_returns_none_when_absent(tmp_path): + assert _chat_template_from_dir(tmp_path) is None + + +def test_read_default_chat_template_direct_gguf_prefers_sidecar(tmp_path, monkeypatch): + gguf = tmp_path / "model-Q4_K_M.gguf" + gguf.write_bytes(b"") + (tmp_path / "tokenizer_config.json").write_text( + json.dumps({"chat_template": "FROM_CONFIG"}), encoding = "utf-8" + ) + monkeypatch.setattr("picker.service._build_browse_allowlist", lambda: [tmp_path]) + monkeypatch.setattr("picker.service.read_gguf_chat_template", lambda _path: "FROM_GGUF") + # A directly selected .gguf must prefer a maintained sidecar over its embedded copy. + assert read_default_chat_template(str(gguf)) == "FROM_CONFIG" + + +def test_read_default_chat_template_direct_gguf_falls_back_to_embedded(tmp_path, monkeypatch): + gguf = tmp_path / "model-Q4_K_M.gguf" + gguf.write_bytes(b"") + monkeypatch.setattr("picker.service._build_browse_allowlist", lambda: [tmp_path]) + monkeypatch.setattr("picker.service.read_gguf_chat_template", lambda _path: "FROM_GGUF") + # With no sidecar next to the file, the embedded GGUF template is the fallback. + assert read_default_chat_template(str(gguf)) == "FROM_GGUF" + + +def test_tokenizer_config_over_size_limit_is_skipped_not_parsed(tmp_path): + # An oversized tokenizer_config.json must be skipped before json.loads so a + # hostile sidecar cannot exhaust memory. + padding = "x" * (MAX_TEMPLATE_METADATA_BYTES + 1024) + (tmp_path / "tokenizer_config.json").write_text( + json.dumps({"chat_template": "HELLO", "_pad": padding}), encoding = "utf-8" + ) + assert _chat_template_from_tokenizer_dir(tmp_path) is None + + +def test_processor_json_over_size_limit_is_skipped_not_parsed(tmp_path): + padding = "x" * (MAX_TEMPLATE_METADATA_BYTES + 1024) + (tmp_path / "chat_template.json").write_text( + json.dumps({"default": "HELLO", "_pad": padding}), encoding = "utf-8" + ) + assert _chat_template_from_processor_json(tmp_path) is None + + +def test_tokenizer_config_at_size_limit_is_still_read(tmp_path): + # A normal-sized config is unaffected by the bound (regression guard). + (tmp_path / "tokenizer_config.json").write_text( + json.dumps({"chat_template": "FROM_CONFIG"}), encoding = "utf-8" + ) + assert _chat_template_from_tokenizer_dir(tmp_path) == "FROM_CONFIG" + + +def test_remote_template_over_size_limit_is_skipped_before_download(monkeypatch): + # An uncached Hub repo whose template exceeds the cap must be skipped via the + # remote size pre-check, never downloaded. + import huggingface_hub + + monkeypatch.setattr("picker.service.resolve_cached_repo_id_case", lambda name: name) + monkeypatch.setattr("picker.service.iter_hf_cache_snapshots", lambda resolved: []) + + def _fail_download(*args, **kwargs): + raise AssertionError("oversized remote template must not be downloaded") + + def _fake_get_paths_info(self, repo_id, paths, **kwargs): + return [SimpleNamespace(path = p, size = MAX_TEMPLATE_METADATA_BYTES + 1) for p in paths] + + monkeypatch.setattr(huggingface_hub, "hf_hub_download", _fail_download) + monkeypatch.setattr(huggingface_hub.HfApi, "get_paths_info", _fake_get_paths_info) + + assert read_default_chat_template("org/oversized-model") is None + + +def test_remote_oversized_jinja_falls_through_to_tokenizer_template(tmp_path, monkeypatch): + # A raw chat_template.jinja between the response cap (MAX_CHAT_TEMPLATE_BYTES) + # and the download bound (MAX_TEMPLATE_METADATA_BYTES) must not be returned: the + # route drops it, so the remote path must skip the oversized Jinja and fall + # through to the smaller tokenizer_config.json. + import huggingface_hub + from picker.schemas import MAX_CHAT_TEMPLATE_BYTES + + big_jinja = tmp_path / "chat_template.jinja" + big_jinja.write_text("{{ x }}" * (MAX_CHAT_TEMPLATE_BYTES // 4), encoding = "utf-8") + assert MAX_CHAT_TEMPLATE_BYTES < big_jinja.stat().st_size < MAX_TEMPLATE_METADATA_BYTES + tokenizer_config = tmp_path / "tokenizer_config.json" + tokenizer_config.write_text(json.dumps({"chat_template": "SMALL_TEMPLATE"}), encoding = "utf-8") + files = { + "chat_template.jinja": big_jinja, + "tokenizer_config.json": tokenizer_config, + } + + monkeypatch.setattr("picker.service.resolve_cached_repo_id_case", lambda name: name) + monkeypatch.setattr("picker.service.iter_hf_cache_snapshots", lambda resolved: []) + + def _fake_download(repo_id, rel, **kwargs): + target = files.get(rel) + if target is None: + raise FileNotFoundError(rel) + return str(target) + + def _fake_get_paths_info(self, repo_id, paths, **kwargs): + return [ + SimpleNamespace( + path = p, + size = files[p].stat().st_size if p in files else 0, + ) + for p in paths + ] + + monkeypatch.setattr(huggingface_hub, "hf_hub_download", _fake_download) + monkeypatch.setattr(huggingface_hub.HfApi, "get_paths_info", _fake_get_paths_info) + + assert read_default_chat_template("org/big-jinja-model") == "SMALL_TEMPLATE" diff --git a/studio/backend/tests/test_providers_api.py b/studio/backend/tests/test_providers_api.py index 5e24ed752d..7cac3a9e99 100644 --- a/studio/backend/tests/test_providers_api.py +++ b/studio/backend/tests/test_providers_api.py @@ -38,11 +38,11 @@ BASE_URL = os.getenv("STUDIO_TEST_URL", "http://localhost:8000") USERNAME = os.getenv("STUDIO_TEST_USER", "unsloth") PASSWORD = os.getenv("STUDIO_TEST_PASSWORD", "") -# Skip the whole module when no live Studio server / bootstrap password is +# Skip the whole module when no live Unsloth server / bootstrap password is # available (e.g. on CI) so pytest discovery does not error out. pytestmark = pytest.mark.skipif( not PASSWORD, - reason = "Integration test requires a running Studio server; set STUDIO_TEST_PASSWORD to enable.", + reason = "Integration test requires a running Unsloth server; set STUDIO_TEST_PASSWORD to enable.", ) # provider_type → (env var name, model for inference test) diff --git a/studio/backend/tests/test_rag_embed_llama_server.py b/studio/backend/tests/test_rag_embed_llama_server.py index 0e1f74cefe..3a332ee19b 100644 --- a/studio/backend/tests/test_rag_embed_llama_server.py +++ b/studio/backend/tests/test_rag_embed_llama_server.py @@ -149,7 +149,7 @@ def test_build_env_gpu_inherits_devices(monkeypatch): monkeypatch.setenv("CUDA_VISIBLE_DEVICES", "0,1") b = LlamaServerBackend() env = b._build_env("/bin/llama-server", use_gpu = True) - assert env.get("CUDA_VISIBLE_DEVICES") == "0,1" # inherit Studio's selection + assert env.get("CUDA_VISIBLE_DEVICES") == "0,1" # inherit Unsloth's selection def test_use_gpu_explicit_modes(monkeypatch): diff --git a/studio/backend/tests/test_recommended_folders_permission.py b/studio/backend/tests/test_recommended_folders_permission.py index 33a457755e..b65695ad93 100644 --- a/studio/backend/tests/test_recommended_folders_permission.py +++ b/studio/backend/tests/test_recommended_folders_permission.py @@ -112,7 +112,7 @@ def test_path_under_unreadable_parent_returns_false_not_raises(tmp_path): ) def test_demonstrates_the_underlying_stdlib_regression(tmp_path): """Documents *why* _safe_is_dir exists: the old bare pattern raises on - the interpreters Studio ships on (3.12+).""" + the interpreters Unsloth ships on (3.12+).""" parent = tmp_path / "ollama" parent.mkdir() os.chmod(parent, 0o000) diff --git a/studio/backend/tests/test_responses_tool_passthrough.py b/studio/backend/tests/test_responses_tool_passthrough.py index 46dd0d42e4..69715649b7 100644 --- a/studio/backend/tests/test_responses_tool_passthrough.py +++ b/studio/backend/tests/test_responses_tool_passthrough.py @@ -120,7 +120,7 @@ class TestResponsesRequestTools: def test_builtin_tool_type_passes_validation(self): """Non-function built-in tools (web_search, file_search, mcp, ...) must not raise at validation so SDKs that default to them don't - fail on Studio; they're filtered out during translation.""" + fail on Unsloth; they're filtered out during translation.""" req = ResponsesRequest( input = "hi", tools = [{"type": "web_search_preview"}], diff --git a/studio/backend/tests/test_rocm_oom_guard.py b/studio/backend/tests/test_rocm_oom_guard.py index 6e70c7cde4..699d0b74f5 100644 --- a/studio/backend/tests/test_rocm_oom_guard.py +++ b/studio/backend/tests/test_rocm_oom_guard.py @@ -36,7 +36,7 @@ class TestIsIntegratedSignal: """hipDeviceProp_t.integrated wins when truthy; 0/absent never downgrades. Same universal gate PR #5988's UMA safetensors fast-load uses -- keeps - Studio's two unified-memory consumers on one signal.""" + Unsloth's two unified-memory consumers on one signal.""" def test_integrated_upgrades_unknown_apu(self) -> None: # gfx1103 Phoenix iGPU: outside the hardcoded arch set, but the diff --git a/studio/backend/tests/test_rocm_windows_vram_7072.py b/studio/backend/tests/test_rocm_windows_vram_7072.py new file mode 100644 index 0000000000..b4079831b7 --- /dev/null +++ b/studio/backend/tests/test_rocm_windows_vram_7072.py @@ -0,0 +1,361 @@ +# SPDX-License-Identifier: AGPL-3.0-only +# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0 + +"""Regression tests for issue #7072 -- "VRAM Usage in System Tab is wrong". + +Reporter: dual AMD (Radeon PRO W7900 ~48GB + W7500 8GB), Windows 10, ROCm 7.13, +torch 2.11.0+rocm7.13. On Windows without a HIP SDK, amd-smi is permanently +disabled (avoids a UAC/DiskPart prompt) and hipMemGetInfo returns free==total +(used 0). Two symptoms followed: + + * System tab (/api/system -> get_visible_gpu_utilization) showed ~0 VRAM used + on every GPU (torch mem_get_info free==total quirk; ROCm/ROCm#1909). + * get_gpu_utilization()'s Windows fallback SUMMED "GPU Adapter Memory\\Dedicated + Usage" across all adapters into ONE fake device with only GPU 0's total, so + the second GPU never appeared. + +The fix reads the per-adapter (LUID-instanced) Dedicated Usage performance +counter -- Task Manager's source -- for per-GPU used, takes per-GPU total from +torch device properties, and guards the free==total mem_get_info quirk. CI has no +AMD GPU/Windows, so torch, the performance counter, and platform are all mocked. +""" + +from __future__ import annotations + +import subprocess +import sys +import types + +import pytest + +from utils.hardware import hardware as hw + +GB = 1024**3 +MiB = 1024**2 + + +# ----------------------------------------------------------------------------- # +# Fakes +# ----------------------------------------------------------------------------- # +def _fake_torch( + devices, + *, + free_equals_total = False, + used_per_device = None, +): + """Build a fake `torch` module. devices: list of (name, total_bytes).""" + dev = list(devices) + + class _Props: + def __init__(self, name, total): + self.name = name + self.total_memory = total + + def get_device_properties(i): + name, total = dev[i] + return _Props(name, total) + + def mem_get_info(i): + _, total = dev[i] + if free_equals_total: + return (total, total) + used = used_per_device[i] if used_per_device is not None else 0 + return (total - used, total) + + t = types.ModuleType("torch") + t.__version__ = "2.11.0+rocm7.13" + t.version = types.SimpleNamespace(hip = "7.13", cuda = None) + t.cuda = types.SimpleNamespace( + is_available = lambda: len(dev) > 0, + device_count = lambda: len(dev), + current_device = lambda: 0, + get_device_properties = get_device_properties, + mem_get_info = mem_get_info, + memory_allocated = lambda i: 0, + memory_reserved = lambda i: 0, + ) + return t + + +def _adapter_output(adapters): + if not adapters: + return "__NONE__\n" + return "".join(f"{name}|{int(used)}\n" for name, used in adapters) + + +def _subprocess_run(*, adapter_output = "__NONE__\n", util_output = "12.0\n"): + def fake_run(cmd, *a, **k): + joined = " ".join(cmd) if isinstance(cmd, list) else str(cmd) + if "GPU Adapter Memory" in joined and "InstanceName" in joined: + out = adapter_output + elif "engtype_3D" in joined or "GPU Engine" in joined: + out = util_output + else: + out = "-1\n" + return subprocess.CompletedProcess(args = cmd, returncode = 0, stdout = out, stderr = "") + + return fake_run + + +@pytest.fixture +def win_rocm(monkeypatch): + """Configure the hardware module as a Windows ROCm host with 2 visible GPUs.""" + monkeypatch.setattr(hw, "get_device", lambda: hw.DeviceType.CUDA) + monkeypatch.setattr(hw, "IS_ROCM", True) + monkeypatch.setattr(hw.platform, "system", lambda: "Windows") + monkeypatch.setattr(hw.sys, "platform", "win32") + monkeypatch.setattr(hw, "_smi_query", lambda *a, **k: None) # amd-smi disabled + # Visible set via HIP mask so we don't shell out to amd-smi for the count. + monkeypatch.setenv("HIP_VISIBLE_DEVICES", "0,1") + monkeypatch.delenv("CUDA_VISIBLE_DEVICES", raising = False) + monkeypatch.delenv("ROCR_VISIBLE_DEVICES", raising = False) + return monkeypatch + + +REPORTER_ADAPTERS = [ + ("luid_0x00000000_0x0000d1e2_phys_0", 40.0 * GB), # W7900, model loaded + ("luid_0x00000000_0x0000e34a_phys_0", 0.5 * GB), # W7500, idle + ("luid_0x00000000_0x0000f001_phys_0", 3 * MiB), # Basic Render Driver +] +DEVICES = [("AMD Radeon PRO W7900", 48 * GB), ("AMD Radeon PRO W7500", 8 * GB)] + + +# ----------------------------------------------------------------------------- # +# System tab (get_visible_gpu_utilization) -- the reporter's screenshot +# ----------------------------------------------------------------------------- # +def test_system_tab_shows_per_gpu_used(win_rocm, monkeypatch): + monkeypatch.setitem(sys.modules, "torch", _fake_torch(DEVICES, free_equals_total = True)) + monkeypatch.setattr( + hw.subprocess, "run", _subprocess_run(adapter_output = _adapter_output(REPORTER_ADAPTERS)) + ) + + devices = hw.get_visible_gpu_utilization()["devices"] + by_idx = {d["index"]: d for d in devices} + assert len(devices) == 2 + assert by_idx[0]["vram_total_gb"] == 48.0 + assert by_idx[0]["vram_used_gb"] == pytest.approx(40.0, abs = 0.01) # not 0 + assert by_idx[1]["vram_total_gb"] == 8.0 # own total + # The 3 MiB Basic Render Driver counter makes this a hidden-adapter case: only + # the 40 GiB is forced onto the 48 GiB card; the idle card reads Unknown. + assert by_idx[1]["vram_used_gb"] is None + assert by_idx[1]["vram_utilization_pct"] is None + assert all( + d["vram_used_gb"] <= d["vram_total_gb"] for d in devices if d["vram_used_gb"] is not None + ) + + +def test_gpu_utilization_does_not_collapse(win_rocm, monkeypatch): + monkeypatch.setitem(sys.modules, "torch", _fake_torch(DEVICES, free_equals_total = True)) + monkeypatch.setattr( + hw.subprocess, "run", _subprocess_run(adapter_output = _adapter_output(REPORTER_ADAPTERS)) + ) + + result = hw.get_gpu_utilization() + devices = result["devices"] + assert sorted(d["index"] for d in devices) == [0, 1] # both GPUs, no collapse + assert {d["vram_total_gb"] for d in devices} == {48.0, 8.0} + assert result["vram_total_gb"] == 48.0 # legacy primary mirror preserved + + +def test_localized_counter_reports_unknown_not_zero(win_rocm, monkeypatch): + monkeypatch.setitem(sys.modules, "torch", _fake_torch(DEVICES, free_equals_total = True)) + monkeypatch.setattr(hw.subprocess, "run", _subprocess_run(adapter_output = "__NONE__\n")) + + devices = hw.get_visible_gpu_utilization()["devices"] + assert len(devices) == 2 # both still shown with correct totals + assert {d["vram_total_gb"] for d in devices} == {48.0, 8.0} + assert all(d["vram_used_gb"] is None for d in devices) # unknown, not fake 0 + assert all(d["vram_utilization_pct"] is None for d in devices) + + +# ----------------------------------------------------------------------------- # +# mem_get_info free==total guard scoping +# ----------------------------------------------------------------------------- # +def test_mem_get_info_guard_scopes_to_windows_rocm(monkeypatch): + torch_mod = _fake_torch(DEVICES, free_equals_total = True) + monkeypatch.setattr(hw, "get_device", lambda: hw.DeviceType.CUDA) + monkeypatch.setitem(sys.modules, "torch", torch_mod) + + # Windows ROCm -> used unknown (None), total kept. + monkeypatch.setattr(hw, "IS_ROCM", True) + monkeypatch.setattr(hw.sys, "platform", "win32") + win = hw._torch_get_per_device_info([0, 1]) + assert [d["used_gb"] for d in win] == [None, None] + assert [d["total_gb"] for d in win] == [48.0, 8.0] + + # Linux ROCm -> unchanged numeric used. + monkeypatch.setattr(hw.sys, "platform", "linux") + assert [d["used_gb"] for d in hw._torch_get_per_device_info([0, 1])] == [0.0, 0.0] + + # Windows NVIDIA -> guard must not fire. + monkeypatch.setattr(hw, "IS_ROCM", False) + monkeypatch.setattr(hw.sys, "platform", "win32") + assert [d["used_gb"] for d in hw._torch_get_per_device_info([0, 1])] == [0.0, 0.0] + + +# ----------------------------------------------------------------------------- # +# Per-adapter attribution helpers (pure unit) +# ----------------------------------------------------------------------------- # +def test_match_adapter_pairs_and_clamps(): + assert hw._match_adapter_used_to_devices([40 * GB, 0.5 * GB], [48 * GB, 8 * GB]) == [ + 40 * GB, + 0.5 * GB, + ] + assert hw._match_adapter_used_to_devices([100 * GB], [48 * GB]) == [48 * GB] # clamp + assert hw._match_adapter_used_to_devices([40 * GB], [48 * GB, 8 * GB]) == [40 * GB, None] + + +def test_match_adapter_reports_unknown_when_more_active_than_visible(): + # More adapters actively using VRAM than are visible (a GPU outside the mask): + # attribution would fabricate a value, so report unknown for every device. + assert hw._match_adapter_used_to_devices([40 * GB, 0.5 * GB], [8 * GB]) == [None] + + +def test_match_adapter_reports_unknown_when_hidden_high_use_adapter_survives_filter(): + # Idle 8 GiB card (10 MiB noise) beside a hidden 48 GiB card at 40 GiB: the + # 40 GiB can't fit the 8 GiB device, so clamping there would fabricate. Unknown. + assert hw._match_adapter_used_to_devices([40 * GB, 10 * MiB], [8 * GB]) == [None] + # Order of the counters must not matter. + assert hw._match_adapter_used_to_devices([10 * MiB, 40 * GB], [8 * GB]) == [None] + + +def test_match_adapter_reports_unknown_for_placeholder_fallback(): + # Every counter below the 64 MiB floor plus a placeholder: no LUID-to-ordinal + # mapping tells placeholder from idle GPU, so report unknown, not fabricate. + # Single visible 8 GiB card idle (10 MiB) beside a 50 MiB placeholder counter. + assert hw._match_adapter_used_to_devices([50 * MiB, 10 * MiB], [8 * GB]) == [None] + # Order of the counters must not matter. + assert hw._match_adapter_used_to_devices([10 * MiB, 50 * MiB], [8 * GB]) == [None] + # Two idle visible GPUs plus a placeholder: all three counters below the floor. + assert hw._match_adapter_used_to_devices([50 * MiB, 10 * MiB, 5 * MiB], [48 * GB, 8 * GB]) == [ + None, + None, + ] + + +def test_match_adapter_reports_unknown_when_usage_not_capacity_ordered(): + # 8 GiB card at 7 GiB beside a 48 GiB card at 5 GiB: the bigger usage still fits + # the smaller card, so both pairings are feasible -> unknown. + assert hw._match_adapter_used_to_devices([7 * GB, 5 * GB], [8 * GB, 48 * GB]) == [None, None] + # Device order must not matter (same physical situation, ordinals flipped). + assert hw._match_adapter_used_to_devices([7 * GB, 5 * GB], [48 * GB, 8 * GB]) == [None, None] + # Same-capacity cards with unequal usage are equally unattributable. + assert hw._match_adapter_used_to_devices([12 * GB, 8 * GB], [24 * GB, 24 * GB]) == [None, None] + # A single usage that fits both cards can sit on either -> unknown. + assert hw._match_adapter_used_to_devices([5 * GB], [48 * GB, 8 * GB]) == [None, None] + # But a capacity-forced assignment (usage exceeds the smaller card) is kept: + # 40 GiB can only be the 48 GiB card, so it is not fabrication. + assert hw._match_adapter_used_to_devices([40 * GB], [48 * GB, 8 * GB]) == [40 * GB, None] + + +def test_match_adapter_reports_unknown_when_hidden_usage_fits_visible_card(): + # A survivor that merely *fits* a visible card must not be pinned onto it. Two + # cards (48/8 GiB) at 40 GiB / 10 MiB beside a hidden 6 GiB adapter: the 6 GiB + # fits the idle 8 GiB card but isn't forced -> Unknown; only 40 GiB is forced. + assert hw._match_adapter_used_to_devices([40 * GB, 10 * MiB, 6 * GB], [48 * GB, 8 * GB]) == [ + 40 * GB, + None, + ] + # Counter order must not matter. + assert hw._match_adapter_used_to_devices([6 * GB, 40 * GB, 10 * MiB], [48 * GB, 8 * GB]) == [ + 40 * GB, + None, + ] + # A single visible card with a hidden adapter is never attributable: a fitting + # survivor could be the hidden GPU's while the visible card is idle. + assert hw._match_adapter_used_to_devices([6 * GB, 10 * MiB], [8 * GB]) == [None] + + +def test_match_adapter_capacity_forced_matrix(): + """Exhaustive hidden-adapter matrix for the capacity-forced rule. + + A value is emitted only when the supra-threshold counters number exactly the + visible devices AND a device's ranked usage strictly exceeds every smaller + card's capacity. Otherwise (a visible card idle, a merely-fitting usage, or the + smallest card) every device reports unknown. + """ + m = hw._match_adapter_used_to_devices + # -- exactly-n supra-threshold counters, capacity-forced survivors are kept - # + # Both visible cards have a real reading (the 3 MiB is a placeholder): 40 GiB + # forced onto the 48 GiB card, 0.5 GiB not forced -> None. + assert m([40 * GB, 0.5 * GB, 3 * MiB], [48 * GB, 8 * GB]) == [40 * GB, None] + # Three visible cards all active (supra-threshold) + placeholder: 40 > 24 and + # 20 > 8, both forced; the 8 GiB card is not forced -> None. + assert m([40 * GB, 20 * GB, 5 * GB, 3 * MiB], [48 * GB, 24 * GB, 8 * GB]) == [ + 40 * GB, + 20 * GB, + None, + ] + # -- fewer supra-threshold counters than visible cards -> all unknown ------ # + # A visible card is idle, so even a "forced" 40 could be the hidden GPU's. + assert m([40 * GB, 3 * MiB, 3 * MiB], [48 * GB, 8 * GB]) == [None, None] + assert m([40 * GB, 10 * MiB, 10 * MiB], [48 * GB, 8 * GB]) == [None, None] + assert m([40 * GB, 20 * GB, 3 * MiB, 3 * MiB], [48 * GB, 24 * GB, 8 * GB]) == [ + None, + None, + None, + ] + # Middle usage (6 GiB) fits both the 24 and 8 GiB cards, and only two cards are + # active for three visible -> not a bijection -> all unknown. + assert m([40 * GB, 6 * GB, 3 * MiB, 3 * MiB], [48 * GB, 24 * GB, 8 * GB]) == [ + None, + None, + None, + ] + # -- hidden larger than every visible card -> all unknown ----------------- # + assert m([40 * GB, 10 * MiB], [8 * GB]) == [None] + assert m([48 * GB, 3 * MiB, 3 * MiB], [24 * GB, 8 * GB]) == [None, None] + # -- more active adapters than visible cards -> all unknown --------------- # + assert m([40 * GB, 7 * GB, 6 * GB, 3 * MiB], [48 * GB, 8 * GB]) == [None, None] + assert m([40 * GB, 7 * GB, 6 * GB, 3 * MiB, 3 * MiB], [48 * GB, 8 * GB]) == [None, None] + # -- every counter below the noise floor (placeholder fallback) -> unknown - # + assert m([50 * MiB, 10 * MiB], [8 * GB]) == [None] + assert m([50 * MiB, 10 * MiB, 5 * MiB], [48 * GB, 8 * GB]) == [None, None] + # -- equal-capacity cards with a hidden adapter: nothing is forced -------- # + assert m([40 * GB, 40 * GB, 3 * MiB], [48 * GB, 48 * GB]) == [None, None] + assert m([40 * GB, 30 * GB, 3 * MiB], [48 * GB, 48 * GB]) == [None, None] + + +def test_perf_counter_parser_and_sentinel(monkeypatch): + monkeypatch.setattr(hw.platform, "system", lambda: "Windows") + monkeypatch.setattr( + hw.subprocess, "run", _subprocess_run(adapter_output = _adapter_output(REPORTER_ADAPTERS)) + ) + parsed = hw._rocm_windows_perf_counter_vram_by_adapter() + assert parsed is not None and len(parsed) == 3 + assert parsed[0][0].startswith("luid_") + monkeypatch.setattr(hw.subprocess, "run", _subprocess_run(adapter_output = "__NONE__\n")) + assert hw._rocm_windows_perf_counter_vram_by_adapter() is None + + +# ----------------------------------------------------------------------------- # +# Unified-memory (Strix Halo APU) total reconciliation (Codex #7238) +# ----------------------------------------------------------------------------- # +def test_unified_memory_adopts_torch_total_even_when_used_unknown(): + """Windows ROCm unified-memory APU: torch's used is None but its total (the full + GTT pool) is authoritative. The correction must still adopt the larger total; + used stays at amd-smi's figure when torch's is unknown.""" + metrics = {"vram_total_gb": 8.0, "vram_used_gb": 2.0, "vram_utilization_pct": 25.0} + hw._apply_unified_memory_correction(metrics, {"total_gb": 124.0, "used_gb": None, "index": 0}) + assert metrics["vram_total_gb"] == 124.0 # full unified pool, not the 8 GB carve-out + assert metrics["vram_used_gb"] == 2.0 # amd-smi used preserved (torch's was None) + assert metrics["vram_utilization_pct"] == pytest.approx(round(2.0 / 124.0 * 100, 1)) + + +def test_unified_memory_overwrites_used_when_torch_used_known(): + """When torch reports both a larger total and a known used, both are adopted + and utilization is recomputed against the corrected total (unchanged path).""" + metrics = {"vram_total_gb": 8.0, "vram_used_gb": 2.0, "vram_utilization_pct": 25.0} + hw._apply_unified_memory_correction(metrics, {"total_gb": 124.0, "used_gb": 40.0, "index": 0}) + assert metrics["vram_total_gb"] == 124.0 + assert metrics["vram_used_gb"] == 40.0 + assert metrics["vram_utilization_pct"] == pytest.approx(round(40.0 / 124.0 * 100, 1)) + + +def test_unified_memory_no_op_when_torch_total_not_larger(): + """A discrete GPU where torch total does not exceed amd-smi's is left untouched.""" + metrics = {"vram_total_gb": 48.0, "vram_used_gb": 10.0, "vram_utilization_pct": 20.8} + hw._apply_unified_memory_correction(metrics, {"total_gb": 48.0, "used_gb": None, "index": 0}) + assert metrics["vram_total_gb"] == 48.0 + assert metrics["vram_used_gb"] == 10.0 + assert metrics["vram_utilization_pct"] == 20.8 diff --git a/studio/backend/tests/test_safetensors_tool_loop.py b/studio/backend/tests/test_safetensors_tool_loop.py index 915f82ac8e..31c728afca 100644 --- a/studio/backend/tests/test_safetensors_tool_loop.py +++ b/studio/backend/tests/test_safetensors_tool_loop.py @@ -2230,8 +2230,8 @@ def _reprompt_loop(*, auto_heal_tool_calls): tools = [{"type": "function", "function": {"name": "search_knowledge_base"}}], execute_tool = exec_fn, auto_heal_tool_calls = auto_heal_tool_calls, - # Studio always nudges (always-on for the Studio inference paths); the - # API opts in per request. Model the Studio caller here. + # Unsloth always nudges (always-on for the Unsloth inference paths); the + # API opts in per request. Model the Unsloth caller here. nudge_tool_calls = True, max_tool_iterations = 3, ) @@ -3203,7 +3203,7 @@ class TestLoopBehaviour: class TestLoopRePrompt: - """Plan-without-action re-prompt parity with GGUF: nudge instead of terminating, up to ``MAX_ACT_REPROMPTS`` extra slots. Studio always nudges, so these drive the loop with ``nudge_tool_calls=True``.""" + """Plan-without-action re-prompt parity with GGUF: nudge instead of terminating, up to ``MAX_ACT_REPROMPTS`` extra slots. Unsloth always nudges, so these drive the loop with ``nudge_tool_calls=True``.""" def test_reasoning_intent_does_not_reprompt_a_visible_answer(self): generations = 0 @@ -4258,7 +4258,7 @@ class TestPlanWithoutActionReprompt: def test_omitted_nudge_flag_is_not_reprompted(self): # The retry is new on this loop: API callers who do not send the flag - # must keep today's behavior. Studio opts in explicitly. + # must keep today's behavior. Unsloth opts in explicitly. loop, exec_fn = _make_loop( turns = [ ["I'll search the web for that."], diff --git a/studio/backend/tests/test_secure_tunnel_gate.py b/studio/backend/tests/test_secure_tunnel_gate.py index 2c13e13bbb..a8c0c2305f 100644 --- a/studio/backend/tests/test_secure_tunnel_gate.py +++ b/studio/backend/tests/test_secure_tunnel_gate.py @@ -2,7 +2,7 @@ # Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0 """Cloudflare tunnel start gate, incl. --secure on loopback. Imports run.py -directly, so run under the Studio venv.""" +directly, so run under the Unsloth venv.""" from __future__ import annotations diff --git a/studio/backend/tests/test_server_disk_logging.py b/studio/backend/tests/test_server_disk_logging.py index 05d03d869c..ce733c2aaa 100644 --- a/studio/backend/tests/test_server_disk_logging.py +++ b/studio/backend/tests/test_server_disk_logging.py @@ -3,7 +3,7 @@ """Tests for the server session log + native-crash capture in run.py. -Field regression: Studio "terminates without a warning" -- a native crash in +Field regression: Unsloth "terminates without a warning" -- a native crash in the GPU runtime kills the process with no Python traceback, and a desktop- shortcut console closes before anything can be read. The server must tee its console output to disk and aim faulthandler at the same file so even hard diff --git a/studio/backend/tests/test_setup_llama_cpp_backend.py b/studio/backend/tests/test_setup_llama_cpp_backend.py new file mode 100644 index 0000000000..36928c680c --- /dev/null +++ b/studio/backend/tests/test_setup_llama_cpp_backend.py @@ -0,0 +1,154 @@ +# SPDX-License-Identifier: AGPL-3.0-only +# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0 + +"""setup.sh and setup.ps1 must map UNSLOTH_LLAMA_CPP_BACKEND=cpu to +install_llama_prebuilt.py's --force-cpu so users can force the CPU-only prebuilt +on GPU hosts (#7213). The match is case-insensitive and whitespace-trimmed, an +unrecognized value warns instead of silently falling back, and macOS warns (no +CPU-only bundle). Runs the real block extracted from each script so the tests +track the shipped logic. +""" + +import os +import re +import shutil +import subprocess +from pathlib import Path + +import pytest + +_STUDIO = Path(__file__).resolve().parents[2] +_SETUP_SH = _STUDIO / "setup.sh" +_SETUP_PS1 = _STUDIO / "setup.ps1" +_SKIP_NO_BASH = pytest.mark.skipif(shutil.which("bash") is None, reason = "bash unavailable") +_SKIP_NO_PWSH = pytest.mark.skipif(shutil.which("pwsh") is None, reason = "pwsh unavailable") + + +def _backend_block() -> str: + text = _SETUP_SH.read_text(encoding = "utf-8") + m = re.search(r"_llama_backend=.*?esac", text, re.DOTALL) + assert m, "UNSLOTH_LLAMA_CPP_BACKEND block not found in setup.sh" + return m.group(0) + + +def _run(value: str | None, system: str = "Linux") -> tuple[list[str], str]: + # Pass the value through env (not the script text) so whitespace survives, and + # stub the setup.sh logging helpers the unknown-value branch calls. system sets + # _HOST_SYSTEM so the macOS (Darwin) no-op branch can be exercised. + env = {k: v for k, v in os.environ.items() if k != "UNSLOTH_LLAMA_CPP_BACKEND"} + if value is not None: + env["UNSLOTH_LLAMA_CPP_BACKEND"] = value + harness = ( + f'_PREBUILT_CMD=()\nC_WARN=""\n_HOST_SYSTEM="{system}"\n' + 'step() { printf "STEP: %s\\n" "$*" >&2; }\n' + f"{_backend_block()}\n" + 'printf "%s\\n" "${_PREBUILT_CMD[@]}"' + ) + out = subprocess.run( + ["bash", "-c", harness], capture_output = True, text = True, env = env, check = True + ) + return out.stdout.split(), out.stderr + + +@_SKIP_NO_BASH +@pytest.mark.parametrize("value", ["cpu", "CPU", "Cpu", " cpu ", "CPU\t"]) +def test_backend_cpu_appends_flag(value): + # A deliberate CPU choice persists, so it uses --force-cpu (not the transient + # --cpu-fallback the arm64 GPU-build recovery uses). + args, stderr = _run(value) + assert "--force-cpu" in args + assert "--cpu-fallback" not in args + assert "Ignoring" not in stderr + + +@_SKIP_NO_BASH +@pytest.mark.parametrize("value", ["cpu", "CPU", " cpu "]) +def test_backend_cpu_macos_warns_no_flag(value): + # macOS has no CPU-only bundle (the universal build already runs on CPU), so the + # override warns instead of writing a misleading forced-CPU marker. + args, stderr = _run(value, system = "Darwin") + assert "--force-cpu" not in args + assert "--cpu-fallback" not in args + assert "macOS" in stderr + + +@_SKIP_NO_BASH +@pytest.mark.parametrize("value", [None, "", "auto", "AUTO", " "]) +def test_backend_auto_no_flag_no_warn(value): + args, stderr = _run(value) + assert "--force-cpu" not in args + assert "Ignoring" not in stderr + + +@_SKIP_NO_BASH +@pytest.mark.parametrize("value", ["vulkan", "gpu", "cuda"]) +def test_backend_unknown_warns_and_no_flag(value): + args, stderr = _run(value) + assert "--force-cpu" not in args + assert "Ignoring" in stderr + + +@_SKIP_NO_BASH +def test_arm64_recovery_uses_transient_cpu_fallback(): + # The arm64 Linux GPU-build recovery must stay transient (--cpu-fallback), never + # the persisted --force-cpu, so a later update can still heal to a GPU bundle (#6097). + text = _SETUP_SH.read_text(encoding = "utf-8") + m = re.search(r"_ARM64_CPU_CMD=\((.*?)\)", text, re.DOTALL) + assert m, "arm64 CPU recovery command not found in setup.sh" + block = m.group(1) + assert "--cpu-fallback" in block + assert "--force-cpu" not in block + + +def _ps1_search(pattern: str, flags = 0) -> str: + m = re.search(pattern, _SETUP_PS1.read_text(encoding = "utf-8"), flags) + assert m, f"setup.ps1 block not found: {pattern}" + return m.group(0) + + +def _run_ps1(value: str | None) -> str: + # The override is normalized (assign + warn) at the top of the prebuilt block and + # applied to $prebuiltArgs lower down; compose both real snippets. + normalize = _ps1_search( + r'\$llamaBackend = "\$\(\$env:UNSLOTH_LLAMA_CPP_BACKEND\)".*?Write-Host.*?\n\s*\}', + re.DOTALL, + ) + apply_flag = _ps1_search( + r'if \(\$llamaBackend -eq "cpu"\) \{\s*\$prebuiltArgs \+= "--force-cpu"\s*\}' + ) + env = {k: v for k, v in os.environ.items() if k != "UNSLOTH_LLAMA_CPP_BACKEND"} + if value is not None: + env["UNSLOTH_LLAMA_CPP_BACKEND"] = value + harness = f'$prebuiltArgs = @()\n{normalize}\n{apply_flag}\n"ARGS:" + ($prebuiltArgs -join ",")' + out = subprocess.run( + ["pwsh", "-NoProfile", "-Command", harness], + capture_output = True, + text = True, + env = env, + check = True, + ) + return out.stdout + + +@_SKIP_NO_PWSH +@pytest.mark.parametrize("value", ["cpu", "CPU", "Cpu", " cpu ", "CPU\t"]) +def test_ps1_backend_cpu_appends_flag(value): + out = _run_ps1(value) + assert "--force-cpu" in out + assert "Ignoring" not in out + + +@_SKIP_NO_PWSH +@pytest.mark.parametrize("value", [None, "", "auto", "AUTO", " "]) +def test_ps1_backend_auto_no_flag_no_warn(value): + out = _run_ps1(value) + assert "--force-cpu" not in out + assert "Ignoring" not in out + + +@_SKIP_NO_PWSH +@pytest.mark.parametrize("value", ["vulkan", "gpu", "cuda"]) +def test_ps1_backend_unknown_warns_and_no_flag(value): + out = _run_ps1(value) + assert "--force-cpu" not in out + assert "Ignoring" in out diff --git a/studio/backend/tests/test_slot_offload_fit.py b/studio/backend/tests/test_slot_offload_fit.py index ac606e4627..d354c7e113 100644 --- a/studio/backend/tests/test_slot_offload_fit.py +++ b/studio/backend/tests/test_slot_offload_fit.py @@ -3,7 +3,7 @@ """Tests for the offload-avoidance serving-slot reduction (`_slots_that_fit_on_gpu`). -When a pinned context does not fit at the requested `--parallel` slot count, Studio would +When a pinned context does not fit at the requested `--parallel` slot count, Unsloth would flip to `--fit on` and llama-server offloads layers to host RAM, collapsing decode ~3x (oobabooga #6718). Instead the loader retries the on-GPU fit at fewer slots and keeps the largest count that stays fully on GPU (`-ngl -1`). These tests drive the real helper with diff --git a/studio/backend/tests/test_studio_api.py b/studio/backend/tests/test_studio_api.py index 928b636e3e..087c00b648 100644 --- a/studio/backend/tests/test_studio_api.py +++ b/studio/backend/tests/test_studio_api.py @@ -11,7 +11,7 @@ the CLI's ``--help`` output: 1. curl -- basic chat completions (non-streaming) 2. curl -- streaming chat completions 3. Python OpenAI SDK -- streaming completions - 4. curl -- Studio server-side tools (enable_tools=true) + 4. curl -- Unsloth server-side tools (enable_tools=true) 5. curl -- Standard OpenAI function calling (non-streaming) 6. curl -- Standard OpenAI function calling (streaming) 7. curl -- Standard OpenAI function calling (multi-turn tool loop) @@ -31,7 +31,7 @@ Usage: python tests/test_studio_api.py python tests/test_studio_api.py --model unsloth/... --gguf-variant ... - # Pytest mode, external server — start a Studio server yourself, + # Pytest mode, external server — start an Unsloth server yourself, # then point pytest at it. Fastest iteration loop. unsloth studio run --model unsloth/Qwen3-1.7B-GGUF --gguf-variant UD-Q4_K_XL & export UNSLOTH_E2E_BASE_URL=http://127.0.0.1:8080 @@ -341,7 +341,7 @@ def _final_finish_reason(chunks: list[dict]) -> str | None: def test_openai_tools_nonstream(base_url: str, api_key: str): """Standard OpenAI function calling, non-streaming, tool_choice='required'. - Regression: before the fix, Studio stripped `tools` and the model + Regression: before the fix, Unsloth stripped `tools` and the model returned plain text with finish_reason='stop'. After the fix, llama-server's response is forwarded verbatim so the client sees finish_reason='tool_calls' with a structured tool_calls array and diff --git a/studio/backend/tests/test_tensor_parallel.py b/studio/backend/tests/test_tensor_parallel.py index 0d71b89d87..00c7aeac69 100644 --- a/studio/backend/tests/test_tensor_parallel.py +++ b/studio/backend/tests/test_tensor_parallel.py @@ -262,9 +262,12 @@ def test_proportional_tensor_split_is_emitted_in_tensor_mode(): src = _load_model_source() assert '"--tensor-split"' in src gate = src.find("if tensor_parallel:") - ts = src.find('"--tensor-split"') + # Find the TP block's emission (after the gate); manual mode emits its own + # --tensor-split earlier in the source from the user's per-GPU shares. + ts = src.find('"--tensor-split"', gate) nxt_else = src.find("self._tensor_parallel = False") assert 0 <= gate < ts < nxt_else, "--tensor-split must be emitted under `if tensor_parallel:`" + assert "tp_tensor_split" in src[gate:nxt_else] def test_mtp_decode_probe_wired_under_tensor_parallel(): @@ -420,7 +423,7 @@ def test_runtime_recovery_fires_for_user_env_mtp(monkeypatch): # MTP driven by user extra_args / LLAMA_ARG_SPEC_TYPE leaves _speculative_type # unset, but the launch flag still gates recovery on (pass-through MTP). b = _recovery_backend() - b._speculative_type = None # Studio stepped back; user/env owns the spec + b._speculative_type = None # Unsloth stepped back; user/env owns the spec done = threading.Event() captured = {} diff --git a/studio/backend/tests/test_tool_confirm_stream.py b/studio/backend/tests/test_tool_confirm_stream.py index b8e0472e12..0813f6b68d 100644 --- a/studio/backend/tests/test_tool_confirm_stream.py +++ b/studio/backend/tests/test_tool_confirm_stream.py @@ -3,12 +3,12 @@ """End-to-end handshake test for the tool-confirmation gate, no model. -The real Studio stream wrappers in ``routes/inference.py`` drive the +The real Unsloth stream wrappers in ``routes/inference.py`` drive the synchronous agentic generator with ``await asyncio.to_thread(next, gen, ...)`` so the blocking ``threading.Event`` wait runs off the event loop. This test rebuilds that exact pattern around the real ``state.tool_approvals`` functions, served by a real uvicorn process on -loopback (the same server Studio uses), and proves the load-bearing +loopback (the same server Unsloth uses), and proves the load-bearing property: * ``tool_start`` reaches the client before the gate blocks, and diff --git a/studio/backend/tests/test_tool_message_empty_content.py b/studio/backend/tests/test_tool_message_empty_content.py index d63b16ce80..636a35f5a9 100644 --- a/studio/backend/tests/test_tool_message_empty_content.py +++ b/studio/backend/tests/test_tool_message_empty_content.py @@ -4,7 +4,7 @@ """Empty ``role="tool"`` content must be accepted on the OpenAI-compat surface. Agentic clients send ``content: ""`` when a command produced no output; -OpenAI and llama-server both accept it. Studio used to 400, which standard +OpenAI and llama-server both accept it. Unsloth used to 400, which standard clients treat as non-retryable and kill the session. The validator must normalize empty/missing tool content to ``""`` instead of raising. """ diff --git a/studio/backend/tests/test_tp_vision_regression.py b/studio/backend/tests/test_tp_vision_regression.py index 09af876da6..d1372ca415 100644 --- a/studio/backend/tests/test_tp_vision_regression.py +++ b/studio/backend/tests/test_tp_vision_regression.py @@ -126,10 +126,21 @@ _ALLOWED_TP_DROP_GUARDS = { # Capability: --split-mode tensor aborted for this (binary, model) (#6415). # Self-healing -- tried by default, skipped only after a real abort (vs #6416). "tensor_parallel and self._tensor_split_aborts(binary, model_identifier)", - # Capacity: tensor needs >= 2 GPUs clearing the compute-buffer reserve. - "tensor_parallel and len(tp_gpus) < 2", + # Capacity: tensor needs >= 2 GPUs clearing the compute-buffer reserve. Gated + # on plan_tp (not raw tensor_parallel) so manual mode skips this planner (#6414). + "plan_tp and len(tp_gpus) < 2", # Capacity: pooled usable VRAM can't hold weights + MTP reserve -> layer split. "_tp_weight_budget_mib <= _tp_required_mib", + # Manual mode, Auto layers: --fit owns memory and is incompatible with a + # tensor split, so TP is dropped (surfaced via logger.info) before the + # cache-drop, so a quantized KV survives into the --fit load (#6414). + "tensor_parallel and gpu_memory_mode == 'manual' and (gpu_layers < 0)", + # Manual mode, explicit layers: a tensor split still needs >= 2 GPUs in use. + "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)", + # Manual mode, zero layers: nothing to split on the GPU, and a tensor-mode + # launch under the CPU-only GPU mask (no visible devices) aborts the server + # instead of the intended CPU-only load (#6414). + "gpu_memory_mode == 'manual' and gpu_layers == 0", } @@ -364,7 +375,7 @@ def test_compute_buffer_downgrade_preserves_multi_gpu_intent(): full GPU set too, so it is symmetric with the budget/geometry downgrades and doesn't collapse a multi-GPU layer load to one card (reviewer.py P1 on #6659).""" src = inspect.getsource(LlamaCppBackend.load_model) - gate = src.find("tensor_parallel and len(tp_gpus) < 2") + gate = src.find("plan_tp and len(tp_gpus) < 2") assert gate != -1 # Bound to exactly this block: from its gate to the next (budget) downgrade. nxt = src.find("_tp_weight_budget_mib <= _tp_required_mib", gate) @@ -625,7 +636,7 @@ def _fallback_loaded_backend(layer_preserves_tensor_intent: bool) -> LlamaCppBac def test_tensor_off_echo_preserves_multi_gpu_fallback(): - """The Studio UI always sends tensor_parallel and echoes the /load response's + """The Unsloth UI always sends tensor_parallel and echoes the /load response's resolved value, so after a fallback a ctx/settings reload carries tensor_parallel= false even though the user never changed it. That echo must NOT collapse the preserved multi-GPU placement -- it dedupes (Codex #6659).""" diff --git a/studio/backend/tests/test_trained_model_scan.py b/studio/backend/tests/test_trained_model_scan.py index 7bf572e214..5d74bb7d28 100644 --- a/studio/backend/tests/test_trained_model_scan.py +++ b/studio/backend/tests/test_trained_model_scan.py @@ -1,7 +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 -"""Tests for Studio trained-model discovery used by Chat.""" +"""Tests for Unsloth trained-model discovery used by Chat.""" import json from pathlib import Path diff --git a/studio/backend/tests/test_training_config_popover_source.py b/studio/backend/tests/test_training_config_popover_source.py new file mode 100644 index 0000000000..4263b012eb --- /dev/null +++ b/studio/backend/tests/test_training_config_popover_source.py @@ -0,0 +1,109 @@ +# SPDX-License-Identifier: AGPL-3.0-only +# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0 + +"""Source-level regression guards for the Training Config popover data source +(#6853). + +The live Training Progress popover used to read the editable form store +(useTrainingConfigStore) while a run was active, so it showed stale/static +values whenever the user touched the form after starting the run; only the +History view read the run's saved config snapshot. These guards pin the fixed +wiring: both views feed ProgressSection a config override mapped from +GET /api/train/runs/{id}, and ProgressSection prefers that override whenever +one is present -- not only for historical views. +""" + +from __future__ import annotations + +from pathlib import Path + +_STUDIO_FRONTEND = Path(__file__).resolve().parents[2] / "frontend" / "src" / "features" / "studio" + + +def _read(rel: str) -> str: + return (_STUDIO_FRONTEND / rel).read_text(encoding = "utf-8") + + +def test_progress_section_prefers_override_over_form_store(): + src = _read("sections/progress-section.tsx") + # Fields key on the override's presence, not isHistorical: a live view passing + # an override wins over the store; without one, live keeps the store while + # History shows blanks rather than unrelated live form values. + assert "const cfg = configOverride ?? (isHistorical ? undefined : config)" in src + assert "const cfgEpochs = cfg?.epochs" in src + assert "isHistorical ? configOverride?.epochs" not in src + + +def test_live_view_fetches_the_active_run_config(): + src = _read("live-training-view.tsx") + # Live view resolves the run's saved config snapshot by job id... + assert "getTrainingRun(" in src + assert "mapRunConfigToOverride(" in src + # ...and hands it to the popover. + assert "configOverride={runConfigOverride}" in src + + +def test_live_view_fetches_as_soon_as_the_job_id_exists(): + # start_training() inserts the run row BEFORE the pump consumes any event, so + # the saved config is available during configuring/loading/downloading. The + # job id is therefore the whole readiness condition: gating on a first step + # or a terminal phase would show the wrong config for the entire pre-step + # window of a long load, or for a run adopted from another client. + src = _read("live-training-view.tsx") + assert "if (!runtime.jobId) {" in src + assert "[runtime.jobId, fetchedRunConfig, fetchAttempt]" in src + # No step/phase readiness gate may creep back in. + assert "runRowReady" not in src + + +def test_live_view_retries_the_transient_row_miss(): + # start_training() creates the row before the pump, but a lookup racing that + # commit can still 404. Nothing else in the effect deps changes on failure, so + # the retry must be explicit and bounded, else a genuinely absent row would + # poll forever instead of falling back to the form store. + src = _read("live-training-view.tsx") + assert "RUN_CONFIG_FETCH_RETRIES" in src + assert "RUN_CONFIG_FETCH_RETRY_MS" in src + assert "setFetchAttempt(" in src + assert "attempts >= RUN_CONFIG_FETCH_RETRIES" in src + # The budget is keyed by job so a new run always starts fresh. + assert "fetchAttempt?.jobId === jobId ? fetchAttempt.count : 0" in src + # The pending retry must be cancelled with the effect. + assert "clearTimeout(retryTimer)" in src + + +def test_live_view_prefers_saved_training_method(): + # The method label / LoRA-row visibility must come from the run snapshot, + # not the editable form (which may have changed since the run started). + src = _read("live-training-view.tsx") + assert "runConfigOverride?.trainingMethod ?? config.trainingMethod" in src + + +def test_history_view_uses_the_shared_mapper(): + src = _read("historical-training-view.tsx") + # Shared mapper, not a re-inlined field-by-field copy that could drift. + assert "mapRunConfigToOverride(detail.config)" in src + assert "num_epochs" not in src + + +def test_shared_mapper_matches_backend_config_keys(): + src = _read("sections/run-config-override.ts") + # The mapper reads the run config JSON the backend snapshots at job start; + # keep the key set pinned so a silent rename breaks loudly here. + for key in ( + "training_type", + "load_in_4bit", + "num_epochs", + "batch_size", + "learning_rate", + "max_steps", + "max_seq_length", + "warmup_steps", + "optim", + "lora_r", + "lora_alpha", + "lora_dropout", + "use_rslora", + "use_loftq", + ): + assert key in src, f"run-config mapper lost backend key {key}" diff --git a/studio/backend/tests/test_training_nan_loss_handling.py b/studio/backend/tests/test_training_nan_loss_handling.py index a2dc78bee2..5a477a084d 100644 --- a/studio/backend/tests/test_training_nan_loss_handling.py +++ b/studio/backend/tests/test_training_nan_loss_handling.py @@ -1,7 +1,7 @@ # SPDX-License-Identifier: AGPL-3.0-only # Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. -"""Pin Studio's behavior when a training event reports non-finite (NaN/Inf) loss. +"""Pin Unsloth's behavior when a training event reports non-finite (NaN/Inf) loss. The training event handler used to filter NaN/Inf to None silently while leaving the previous finite loss in progress.loss — so the API kept reporting diff --git a/studio/backend/tests/test_training_pump_resilience.py b/studio/backend/tests/test_training_pump_resilience.py index d75b205f35..5d2a218482 100644 --- a/studio/backend/tests/test_training_pump_resilience.py +++ b/studio/backend/tests/test_training_pump_resilience.py @@ -310,6 +310,80 @@ def test_pump_finalizes_when_read_keeps_raising_on_dead_worker(monkeypatch): assert b._pump_running is False +def test_interrupted_cancel_clears_in_memory_output_dir(monkeypatch): + # Stop-without-save interrupted before its complete event: /status must not + # keep serving the cleared run's output_dir. + b = TrainingBackend() + finalized: dict = {} + monkeypatch.setattr(b, "_ensure_db_run_created", lambda: None) + monkeypatch.setattr(b, "_finalize_run_in_db", lambda **kw: finalized.update(kw)) + + b._proc = _FakeProc(alive = False) + b._event_queue = _IdleQueue() + b._progress.is_training = True + b._should_stop = True + b._cancel_requested = True + b._output_dir = "/out/x" + + b._pump_loop() + + assert b._output_dir is None + assert finalized.get("status") == "stopped" + assert finalized.get("output_dir") is None + assert finalized.get("clear_output_dir") is True + + +def test_worker_exit_reuses_terminal_stop_save_error(monkeypatch): + b = TrainingBackend() + finalized: dict = {} + monkeypatch.setattr(b, "_ensure_db_run_created", lambda: None) + monkeypatch.setattr(b, "_finalize_run_in_db", lambda **kw: finalized.update(kw)) + + b._proc = _FakeProc(alive = False) + b._event_queue = _IdleQueue() + b._progress.is_training = True + b._should_stop = True + b._cancel_requested = False + b._output_dir = "/out/x" + b.current_job_id = "job-x" + b._terminal_finalize_payload = { + "status": "error", + "error_message": "checkpoint failed", + "output_dir": "/out/x", + "clear_output_dir": False, + "resume_blocked": True, + "expected_job_id": "job-x", + } + + b._pump_loop() + + assert b._output_dir == "/out/x" + assert finalized.get("status") == "error" + assert finalized.get("output_dir") == "/out/x" + assert finalized.get("clear_output_dir") is False + assert finalized.get("resume_blocked") is True + + +def test_dead_worker_crash_preserves_output_dir(monkeypatch): + # A crash (no stop requested) after output_dir was emitted must keep the dir + # in the error finalize: checkpoints under it may still exist. + b = TrainingBackend() + finalized: dict = {} + monkeypatch.setattr(b, "_ensure_db_run_created", lambda: None) + monkeypatch.setattr(b, "_finalize_run_in_db", lambda **kw: finalized.update(kw)) + + b._proc = _FakeProc(alive = False) + b._event_queue = _IdleQueue() + b._progress.is_training = True + b._output_dir = "/out/x" + + b._pump_loop() + + assert finalized.get("status") == "error" + assert finalized.get("output_dir") == "/out/x" + assert finalized.get("clear_output_dir") is False + + def test_start_training_clears_stale_pump_running_flag(): # A prior pump that died abnormally leaves _pump_running True. The next # start_training must clear it during reset so the start-time watchdog can't diff --git a/studio/backend/tests/test_training_resume.py b/studio/backend/tests/test_training_resume.py index 91fdac9961..51425b0428 100644 --- a/studio/backend/tests/test_training_resume.py +++ b/studio/backend/tests/test_training_resume.py @@ -7,6 +7,9 @@ import importlib.util import json from pathlib import Path +import pytest +import torch + _BACKEND = Path(__file__).resolve().parents[1] @@ -25,6 +28,30 @@ def _load_resume_module(): resume = _load_resume_module() +def test_resume_request_accepts_sanitized_null_target_modules(): + from models.training import TrainingStartRequest + request = TrainingStartRequest( + model_name = "unsloth/Qwen3-0.6B", + training_type = "Full Finetuning", + format_type = "alpaca", + target_modules = None, + ) + + assert request.target_modules == [] + + +def _write_checkpoint(out: Path, step: int) -> Path: + checkpoint = out / f"checkpoint-{step}" + checkpoint.mkdir(parents = True, exist_ok = True) + (checkpoint / "trainer_state.json").write_text( + json.dumps({"global_step": step}), encoding = "utf-8" + ) + torch.save({"weight": torch.ones(1)}, checkpoint / "adapter_model.bin") + torch.save({"state": {0: torch.ones(1)}}, checkpoint / "optimizer.pt") + torch.save({"last_epoch": step}, checkpoint / "scheduler.pt") + return checkpoint + + def _stopped_run(**overrides): run = { "status": "stopped", @@ -44,6 +71,36 @@ def test_can_resume_run_allows_checkpointed_non_s3_run(monkeypatch): assert resume.can_resume_run(_stopped_run()) is True +def test_can_resume_run_allows_errored_run_with_checkpoint(monkeypatch): + monkeypatch.setattr(resume, "has_resume_state", lambda _path: True) + + assert resume.can_resume_run(_stopped_run(status = "error")) is True + + +def test_can_resume_run_rejects_errored_run_without_checkpoint(monkeypatch): + monkeypatch.setattr(resume, "has_resume_state", lambda _path: False) + + assert resume.can_resume_run(_stopped_run(status = "error")) is False + + +def test_can_resume_run_allows_errored_run_at_final_step(monkeypatch): + # A save-time crash records final_step == total_steps; resuming re-runs the + # final-save path from the checkpoint. + monkeypatch.setattr(resume, "has_resume_state", lambda _path: True) + + run = _stopped_run(status = "error", final_step = 10, total_steps = 10) + + assert resume.can_resume_run(run) is True + + +def test_can_resume_run_rejects_stopped_run_at_final_step(monkeypatch): + monkeypatch.setattr(resume, "has_resume_state", lambda _path: True) + + run = _stopped_run(final_step = 10, total_steps = 10) + + assert resume.can_resume_run(run) is False + + def test_can_resume_run_rejects_s3_dataset_source(monkeypatch): monkeypatch.setattr(resume, "has_resume_state", lambda _path: True) @@ -91,3 +148,444 @@ def test_list_runs_includes_config_json_for_resume_policy(monkeypatch, tmp_path) result = studio_db.list_runs() assert result["runs"][0]["config_json"] == config_json + + +def test_crashed_run_with_persisted_output_dir_is_resumable(monkeypatch, tmp_path): + from storage import studio_db + + monkeypatch.setenv("UNSLOTH_STUDIO_HOME", str(tmp_path)) + monkeypatch.setattr(studio_db, "_schema_ready", False) + + out = tmp_path / "outputs" / "run_x" + _write_checkpoint(out, 10) + + studio_db.create_run( + id = "run-crash", + model_name = "m", + dataset_name = "d", + config_json = "{}", + started_at = "2026-01-01T00:00:00Z", + total_steps = 20, + ) + studio_db.update_run_output_dir("run-crash", str(out)) + conn = studio_db.get_connection() + conn.execute("UPDATE training_runs SET status = 'error' WHERE id = 'run-crash'") + conn.commit() + conn.close() + + run = studio_db.get_run("run-crash") + assert run["output_dir"] == str(out) + assert resume.can_resume_run(run) is True + + +def test_checkpoint_discovery_skips_malformed_newest(monkeypatch, tmp_path): + monkeypatch.setenv("UNSLOTH_STUDIO_HOME", str(tmp_path)) + out = tmp_path / "outputs" / "run_x" + valid = _write_checkpoint(out, 5) + (_write_checkpoint(out, 8) / "scheduler.pt").unlink() + malformed = out / "checkpoint-10" + malformed.mkdir() + (malformed / "trainer_state.json").write_text(json.dumps({"global_step": 10}), encoding = "utf-8") + (malformed / "adapter_model.bin").write_bytes(b"not a torch archive") + (malformed / "optimizer.pt").write_bytes(b"not a torch archive") + + assert resume.get_resume_checkpoint_path(str(out)) == str(valid) + + +def test_completed_run_keeps_output_dir_and_rejects_stale_cancel(monkeypatch, tmp_path): + from storage import studio_db + + monkeypatch.setenv("UNSLOTH_STUDIO_HOME", str(tmp_path)) + monkeypatch.setattr(studio_db, "_schema_ready", False) + + studio_db.create_run( + id = "r", + model_name = "m", + dataset_name = "d", + config_json = "{}", + started_at = "2026-01-01T00:00:00Z", + total_steps = 10, + ) + studio_db.update_run_output_dir("r", "/out/x") + studio_db.finish_run( + id = "r", + status = "completed", + ended_at = "t", + final_step = 2, + final_loss = None, + duration_seconds = 1, + loss_sparkline = "[]", + output_dir = "/out/x", + error_message = None, + ) + + assert studio_db.get_run("r")["output_dir"] == "/out/x" + assert studio_db.mark_run_cancel_requested("r") is False + assert studio_db.get_run("r")["output_dir"] == "/out/x" + assert studio_db.get_run("r")["resume_blocked"] == 0 + + +def test_finish_run_clears_output_dir_for_stop_without_save(monkeypatch, tmp_path): + from storage import studio_db + + monkeypatch.setenv("UNSLOTH_STUDIO_HOME", str(tmp_path)) + monkeypatch.setattr(studio_db, "_schema_ready", False) + + studio_db.create_run( + id = "r", + model_name = "m", + dataset_name = "d", + config_json = "{}", + started_at = "2026-01-01T00:00:00Z", + total_steps = 10, + ) + studio_db.update_run_output_dir("r", "/out/x") + studio_db.finish_run( + id = "r", + status = "stopped", + ended_at = "t", + final_step = 2, + final_loss = None, + duration_seconds = 1, + loss_sparkline = "[]", + output_dir = None, + error_message = None, + clear_output_dir = True, + ) + + assert studio_db.get_run("r")["output_dir"] is None + conn = studio_db.get_connection() + conn.execute( + "UPDATE training_runs SET status = 'running', output_dir = '/out/x', resume_blocked = 0 WHERE id = 'r'" + ) + conn.commit() + conn.close() + studio_db.mark_run_cancel_requested("r") + studio_db.cleanup_orphaned_runs() + assert studio_db.get_run("r")["status"] == "stopped" + assert studio_db.get_run("r")["output_dir"] is None + + +def test_finish_run_clears_output_dir_on_cancel_error_finalize(monkeypatch, tmp_path): + from storage import studio_db + + monkeypatch.setenv("UNSLOTH_STUDIO_HOME", str(tmp_path)) + monkeypatch.setattr(studio_db, "_schema_ready", False) + + studio_db.create_run( + id = "r", + model_name = "m", + dataset_name = "d", + config_json = "{}", + started_at = "2026-01-01T00:00:00Z", + total_steps = 10, + ) + studio_db.update_run_output_dir("r", "/out/x") + studio_db.finish_run( + id = "r", + status = "stopped", + ended_at = "t", + final_step = 2, + final_loss = None, + duration_seconds = 1, + loss_sparkline = "[]", + output_dir = "/out/x", + error_message = "worker failed during cancel", + clear_output_dir = True, + ) + + assert studio_db.get_run("r")["output_dir"] is None + + +def test_finish_run_preserves_output_dir_for_interrupted_stop_and_save(monkeypatch, tmp_path): + from storage import studio_db + + monkeypatch.setenv("UNSLOTH_STUDIO_HOME", str(tmp_path)) + monkeypatch.setattr(studio_db, "_schema_ready", False) + + studio_db.create_run( + id = "r", + model_name = "m", + dataset_name = "d", + config_json = "{}", + started_at = "2026-01-01T00:00:00Z", + total_steps = 10, + ) + studio_db.update_run_output_dir("r", "/out/x") + studio_db.finish_run( + id = "r", + status = "stopped", + ended_at = "t", + final_step = 2, + final_loss = None, + duration_seconds = 1, + loss_sparkline = "[]", + output_dir = None, + error_message = None, + ) + + assert studio_db.get_run("r")["output_dir"] == "/out/x" + + +def test_resumed_errored_run_is_not_offered_again(monkeypatch, tmp_path): + from storage import studio_db + + monkeypatch.setenv("UNSLOTH_STUDIO_HOME", str(tmp_path)) + monkeypatch.setattr(studio_db, "_schema_ready", False) + + out = tmp_path / "outputs" / "run_x" + _write_checkpoint(out, 10) + + studio_db.create_run( + id = "run-old", + model_name = "m", + dataset_name = "d", + config_json = "{}", + started_at = "2026-01-01T00:00:00Z", + total_steps = 20, + ) + studio_db.update_run_output_dir("run-old", str(out)) + studio_db.finish_run( + id = "run-old", + status = "error", + ended_at = "2026-01-01T00:05:00Z", + final_step = 10, + final_loss = None, + duration_seconds = 1, + loss_sparkline = "[]", + output_dir = None, + error_message = "killed", + ) + studio_db.create_run( + id = "run-new", + model_name = "m", + dataset_name = "d", + config_json = "{}", + started_at = "2026-01-02T00:00:00Z", + total_steps = 20, + output_dir = str(out), + resumed_from_run_id = "run-old", + ) + with pytest.raises(RuntimeError, match = "no longer available"): + studio_db.create_run( + id = "run-duplicate", + model_name = "m", + dataset_name = "d", + config_json = "{}", + started_at = "2026-01-02T00:00:01Z", + total_steps = 20, + output_dir = str(out), + resumed_from_run_id = "run-old", + ) + assert studio_db.get_run("run-duplicate") is None + studio_db.finish_run( + id = "run-new", + status = "error", + ended_at = "2026-01-02T00:05:00Z", + final_step = 15, + final_loss = None, + duration_seconds = 1, + loss_sparkline = "[]", + output_dir = None, + error_message = "killed again", + ) + + old_run = studio_db.get_run("run-old") + new_run = studio_db.get_run("run-new") + assert old_run["resumed_later"] == 1 + assert resume.can_resume_run(old_run) is False + assert new_run["resumed_later"] == 0 + assert resume.can_resume_run(new_run) is True + assert studio_db.get_resumable_run_by_output_dir(str(out))["id"] == "run-new" + + +def test_running_continuation_blocks_older_resume(monkeypatch, tmp_path): + from storage import studio_db + + monkeypatch.setenv("UNSLOTH_STUDIO_HOME", str(tmp_path)) + monkeypatch.setattr(studio_db, "_schema_ready", False) + + out = tmp_path / "outputs" / "run_x" + _write_checkpoint(out, 10) + + studio_db.create_run( + id = "run-old", + model_name = "m", + dataset_name = "d", + config_json = "{}", + started_at = "2026-01-01T00:00:00Z", + total_steps = 20, + ) + studio_db.update_run_output_dir("run-old", str(out)) + studio_db.finish_run( + id = "run-old", + status = "error", + ended_at = "2026-01-01T00:05:00Z", + final_step = 10, + final_loss = None, + duration_seconds = 1, + loss_sparkline = "[]", + output_dir = None, + error_message = "killed", + ) + studio_db.create_run( + id = "run-new", + model_name = "m", + dataset_name = "d", + config_json = "{}", + started_at = "2026-01-02T00:00:00Z", + total_steps = 20, + output_dir = str(out), + resumed_from_run_id = "run-old", + ) + + old_run = studio_db.get_run("run-old") + assert old_run["resumed_later"] == 1 + assert resume.can_resume_run(old_run) is False + assert studio_db.get_resumable_run_by_output_dir(str(out)) is None + + +def test_stop_save_checkpoint_failure_keeps_error_status(monkeypatch, tmp_path): + # A stop-and-save whose checkpoint write failed must finalize as an error so + # history explains the missing resume state (keep_error_status flag). + from core.training.training import TrainingBackend + from storage import studio_db + + monkeypatch.setenv("UNSLOTH_STUDIO_HOME", str(tmp_path)) + monkeypatch.setattr(studio_db, "_schema_ready", False) + + studio_db.create_run( + id = "run-failed-save", + model_name = "m", + dataset_name = "d", + config_json = "{}", + started_at = "2026-01-01T00:00:00Z", + total_steps = 10, + ) + backend = TrainingBackend() + backend.current_job_id = "run-failed-save" + backend._db_run_created = True + backend._should_stop = True + backend._handle_event( + { + "type": "error", + "error": "Failed to save a resumable checkpoint after stop.", + "keep_error_status": True, + } + ) + + run = studio_db.get_run("run-failed-save") + assert run["status"] == "error" + assert "resumable checkpoint" in run["error_message"] + + +def test_can_resume_run_rejects_resume_blocked_run(monkeypatch): + monkeypatch.setattr(resume, "has_resume_state", lambda _path: True) + + assert resume.can_resume_run(_stopped_run(status = "error", resume_blocked = 1)) is False + + +def test_stop_save_checkpoint_failure_with_stale_checkpoint_is_not_resumable(monkeypatch, tmp_path): + # A failed stop-and-save must not offer Resume from an older periodic checkpoint; + # that would roll back past the recorded final step. + from core.training.training import TrainingBackend + from storage import studio_db + + monkeypatch.setenv("UNSLOTH_STUDIO_HOME", str(tmp_path)) + monkeypatch.setattr(studio_db, "_schema_ready", False) + + out = tmp_path / "outputs" / "run_x" + _write_checkpoint(out, 10) + + studio_db.create_run( + id = "run-stale-ckpt", + model_name = "m", + dataset_name = "d", + config_json = "{}", + started_at = "2026-01-01T00:00:00Z", + total_steps = 20, + ) + studio_db.update_run_output_dir("run-stale-ckpt", str(out)) + backend = TrainingBackend() + backend.current_job_id = "run-stale-ckpt" + backend._db_run_created = True + backend._should_stop = True + backend._output_dir = str(out) + backend._handle_event( + { + "type": "error", + "error": "Failed to save a resumable checkpoint after stop.", + "keep_error_status": True, + "resume_blocked": True, + } + ) + + run = studio_db.get_run("run-stale-ckpt") + assert run["status"] == "error" + assert run["resume_blocked"] == 1 + assert run["output_dir"] == str(out) + assert resume.can_resume_run(run) is False + + +def test_user_stop_error_without_checkpoint_ack_is_blocked(monkeypatch, tmp_path): + from core.training.training import TrainingBackend + from storage import studio_db + + monkeypatch.setenv("UNSLOTH_STUDIO_HOME", str(tmp_path)) + monkeypatch.setattr(studio_db, "_schema_ready", False) + + studio_db.create_run( + id = "run-user-stop", + model_name = "m", + dataset_name = "d", + config_json = "{}", + started_at = "2026-01-01T00:00:00Z", + total_steps = 10, + ) + backend = TrainingBackend() + backend.current_job_id = "run-user-stop" + backend._db_run_created = True + backend._should_stop = True + backend._handle_event({"type": "error", "error": "interrupted"}) + + run = studio_db.get_run("run-user-stop") + assert run["status"] == "error" and run["resume_blocked"] == 1 + + +def test_terminal_fallback_keeps_resumable_when_current_checkpoint_landed(monkeypatch, tmp_path): + # Worker died before its terminal event, but a valid current-step checkpoint + # is on disk: the fallback must keep the run resumable, not block it. + from core.training.training import TrainingBackend + + monkeypatch.setenv("UNSLOTH_STUDIO_HOME", str(tmp_path)) + out = tmp_path / "outputs" / "run_ok" + _write_checkpoint(out, 7) + + backend = TrainingBackend() + backend.current_job_id = "run-ok" + backend._should_stop = True + backend._output_dir = str(out) + backend._progress.step = 7 + + kwargs = backend._terminal_finalize_kwargs() + assert kwargs["status"] == "stopped" + assert kwargs["resume_blocked"] is False + + +def test_terminal_fallback_blocks_when_no_current_checkpoint(monkeypatch, tmp_path): + # Same path, but only a stale (older-step) checkpoint exists: must block. + from core.training.training import TrainingBackend + + monkeypatch.setenv("UNSLOTH_STUDIO_HOME", str(tmp_path)) + out = tmp_path / "outputs" / "run_stale" + _write_checkpoint(out, 5) + + backend = TrainingBackend() + backend.current_job_id = "run-stale" + backend._should_stop = True + backend._output_dir = str(out) + backend._progress.step = 7 + + kwargs = backend._terminal_finalize_kwargs() + assert kwargs["status"] == "error" + assert kwargs["resume_blocked"] is True diff --git a/studio/backend/tests/test_training_stop_watchdog.py b/studio/backend/tests/test_training_stop_watchdog.py index 0cd702bce2..cbe2082e82 100644 --- a/studio/backend/tests/test_training_stop_watchdog.py +++ b/studio/backend/tests/test_training_stop_watchdog.py @@ -353,7 +353,7 @@ def test_finalize_after_escalation_clears_state(monkeypatch): # stopped so the UI leaves "Stopping..." and a new run can start. b = TrainingBackend() finstop: list = [] - monkeypatch.setattr(b, "_finish_stopped_run", lambda *a: finstop.append(a)) + monkeypatch.setattr(b, "_finish_stopped_run", lambda *a, **k: finstop.append(a)) b._proc = _FakeProc(alive = True) # wedged: still reports alive b._should_stop = True @@ -365,7 +365,7 @@ def test_finalize_after_escalation_clears_state(monkeypatch): assert b._proc is None, "the wedged handle must be dropped so is_training_active clears" assert b._progress.is_training is False - assert b._progress.status_message == "Training stopped." + assert "valid current-step checkpoint" in b._progress.status_message assert finstop and finstop[0][0] == "job_c", "the captured run must be finalized by id" assert b.is_training_active() is False @@ -375,7 +375,7 @@ def test_finalize_after_escalation_preserves_output_dir(monkeypatch): # must record it even if the watchdog wins the finalize race against the pump. b = TrainingBackend() finstop: list = [] - monkeypatch.setattr(b, "_finish_stopped_run", lambda *a: finstop.append(a)) + monkeypatch.setattr(b, "_finish_stopped_run", lambda *a, **k: finstop.append(a)) b._proc = _FakeProc(alive = True) b._should_stop = True @@ -390,6 +390,28 @@ def test_finalize_after_escalation_preserves_output_dir(monkeypatch): assert finstop[0][1] == "/tmp/outputs/run-123" +def test_finalize_after_escalation_clears_output_dir_on_cancel(monkeypatch): + # Stop-without-saving promises no resume: a cancel that escalates through the + # watchdog clears the persisted output_dir, not a checkpoint path. + b = TrainingBackend() + finstop: list = [] + monkeypatch.setattr(b, "_finish_stopped_run", lambda *a, **k: finstop.append((a, k))) + + b._proc = _FakeProc(alive = True) + b._should_stop = True + b._cancel_requested = True + b.current_job_id = "job_c" + b._db_run_created = True + b._output_dir = "/tmp/outputs/run-123" + + b._finalize_stopped_after_escalation(watched_job_id = "job_c") + + assert finstop and finstop[0][0][0] == "job_c" + assert finstop[0][0][1] is None, "a cancelled run must not record a checkpoint path" + assert finstop[0][1].get("clear_output_dir") is True + assert b._output_dir is None, "/status must stop exposing the cancelled run's dir" + + def test_stop_training_starts_watchdog_only_when_worker_alive(monkeypatch): # No worker -> nothing to escalate; the watchdog must not spawn. b = TrainingBackend() @@ -409,7 +431,7 @@ def test_finalize_after_escalation_no_ops_when_superseded(monkeypatch): # The escalation finalize must then leave the NEW run untouched, not drop its handle. b = TrainingBackend() finstop: list = [] - monkeypatch.setattr(b, "_finish_stopped_run", lambda *a: finstop.append(a)) + monkeypatch.setattr(b, "_finish_stopped_run", lambda *a, **k: finstop.append(a)) old_proc = _FakeProc(alive = False) # force-terminated worker we were watching new_proc = _FakeProc(alive = True) # a new run already took over @@ -430,7 +452,7 @@ def test_finalize_after_escalation_runs_for_its_own_worker(monkeypatch): # finalizes the captured run by id. b = TrainingBackend() finstop: list = [] - monkeypatch.setattr(b, "_finish_stopped_run", lambda *a: finstop.append(a)) + monkeypatch.setattr(b, "_finish_stopped_run", lambda *a, **k: finstop.append(a)) proc = _FakeProc(alive = False) b._proc = proc @@ -451,7 +473,7 @@ def test_finalize_after_escalation_no_ops_on_job_change_during_startup(monkeypat # catch this even though the proc-only guard would not. b = TrainingBackend() finstop: list = [] - monkeypatch.setattr(b, "_finish_stopped_run", lambda *a: finstop.append(a)) + monkeypatch.setattr(b, "_finish_stopped_run", lambda *a, **k: finstop.append(a)) old_proc = _FakeProc(alive = False) # old worker, dead; new _proc not installed yet b._proc = old_proc # still the old handle (== target), so proc guard would pass @@ -509,6 +531,7 @@ def _install_fake_db(monkeypatch): recs["insert_ids"].append(job_id), ) fake_db.update_run_progress = lambda **kw: recs["progress_ids"].append(kw.get("id")) + fake_db.mark_run_cancel_requested = lambda _run_id: True fake_storage.studio_db = fake_db monkeypatch.setitem(sys.modules, "storage", fake_storage) monkeypatch.setitem(sys.modules, "storage.studio_db", fake_db) @@ -518,9 +541,48 @@ def _install_fake_db(monkeypatch): return recs +def test_stop_without_save_creates_missing_row_before_signal(monkeypatch): + recs = _install_fake_db(monkeypatch) + b = TrainingBackend() + b.current_job_id, b._db_config = "job_missing", {"model_name": "m"} + b._stop_queue = queue.Queue() + assert b.stop_training(save = False) is True + assert [run["id"] for run in recs["created"]] == ["job_missing"] + assert b._stop_queue.get_nowait() == {"type": "stop", "save": False} + + b._cancel_requested = b._should_stop = False + sys.modules["storage.studio_db"].mark_run_cancel_requested = lambda _run_id: False + assert b.stop_training(save = False) is False + assert not b._cancel_requested and b._stop_queue.empty() + + new_queue = queue.Queue() + b.current_job_id, b._db_run_created = "job_old", True + b._cancel_requested = b._should_stop = False + + def _supersede(_run_id): + b.current_job_id = "job_new" + b._stop_queue = new_queue + return True + + sys.modules["storage.studio_db"].mark_run_cancel_requested = _supersede + assert b.stop_training(save = False) is False + assert not b._cancel_requested and new_queue.empty() + + def test_finalize_run_in_db_single_winner_under_concurrency(monkeypatch): # The watchdog and pump can both finalize; only one call may reach finish_run. recs = _install_fake_db(monkeypatch) + monkeypatch.setitem(_G, "_DB_FINALIZE_RETRY_S", 0.0) + attempts = 0 + + def flaky_finish(**kw): + nonlocal attempts + attempts += 1 + if attempts < 3: + raise RuntimeError("database is locked") + recs["finished"].append(kw) + + sys.modules["storage.studio_db"].finish_run = flaky_finish b = TrainingBackend() b.current_job_id = "job_x" b._db_run_created = True @@ -539,6 +601,7 @@ def test_finalize_run_in_db_single_winner_under_concurrency(monkeypatch): t.join(timeout = 5) assert len(recs["finished"]) == 1, f"finalize must run once, got {len(recs['finished'])}" + assert attempts == 3 assert b._run_finalized is True @@ -646,7 +709,13 @@ def test_ensure_db_run_created_publishes_only_after_insert(monkeypatch): monkeypatch.setitem(sys.modules, "storage", fake_storage) monkeypatch.setitem(sys.modules, "storage.studio_db", fake_db) - b._ensure_db_run_created() + b._run_intent_lock.acquire() + creator = threading.Thread(target = b._ensure_db_run_created) + creator.start() + time.sleep(0.02) + assert b._db_create_in_progress is False + b._run_intent_lock.release() + creator.join(timeout = 5) assert observed["flag_during_create"] is False, "flag must not be published before insert" assert observed["in_progress_during_create"] is True @@ -718,6 +787,7 @@ def test_escalation_finalizes_watched_run_by_id_end_to_end(monkeypatch): b = TrainingBackend() b.current_job_id = "job_old" b._db_run_created = True + b._should_stop = True b._proc = _FakeProc(alive = False) b._progress.is_training = True b._progress.step = 42 @@ -726,7 +796,8 @@ def test_escalation_finalizes_watched_run_by_id_end_to_end(monkeypatch): b._finalize_stopped_after_escalation(target_proc = b._proc, watched_job_id = "job_old") assert [f["id"] for f in recs["finished"]] == ["job_old"], "must finish the captured run by id" - assert recs["finished"][0]["status"] == "stopped" + assert recs["finished"][0]["status"] == "error" + assert recs["finished"][0]["resume_blocked"] is True assert recs["insert_ids"] == ["job_old"], "buffered metrics must land on the captured run" assert b._metric_buffer == [], "the captured batch must be drained" @@ -737,7 +808,7 @@ def test_escalation_defers_when_row_cannot_be_created_here(monkeypatch): # so the pump's create-then-finalize records the run. Parent state still clears. b = TrainingBackend() called: list = [] - monkeypatch.setattr(b, "_finish_stopped_run", lambda *a: called.append(a)) + monkeypatch.setattr(b, "_finish_stopped_run", lambda *a, **k: called.append(a)) b._proc = _FakeProc(alive = False) b.current_job_id = "job_q" @@ -785,7 +856,7 @@ def test_escalation_does_not_drop_a_new_runs_handle(monkeypatch): new_proc = _FakeProc(alive = True) b._proc = old_proc - def hijack(*a): + def hijack(*a, **k): b._proc = new_proc # a new run takes over during the finalize monkeypatch.setattr(b, "_finish_stopped_run", hijack) diff --git a/studio/backend/tests/test_transformers_latest.py b/studio/backend/tests/test_transformers_latest.py index 20616dccba..af48d674cc 100644 --- a/studio/backend/tests/test_transformers_latest.py +++ b/studio/backend/tests/test_transformers_latest.py @@ -1036,7 +1036,7 @@ def test_upgrade_check_mixed_pypi_main_reports_dev_only(monkeypatch): def test_install_endpoint_not_mounted_on_v1(): - """The consented pip-install endpoint is a Studio admin action; it must live + """The consented pip-install endpoint is an Unsloth admin action; it must live on studio_router (kept off the OpenAI-compatible /v1 mount), not router.""" from routes import inference as ri diff --git a/studio/backend/tests/test_utils.py b/studio/backend/tests/test_utils.py index 64a3c62156..741f19c67a 100644 --- a/studio/backend/tests/test_utils.py +++ b/studio/backend/tests/test_utils.py @@ -38,7 +38,7 @@ from utils.hardware import ( DeviceType, ) import utils.hardware.hardware as _hw_module -from utils.utils import format_error_message +from utils.utils import format_error_message, is_hf_authentication_error # ========== Helpers ========== @@ -439,6 +439,20 @@ class TestFormatErrorMessage: msg = format_error_message(err, "any/model") assert "invalid" in msg.lower() + def test_hf_authentication_error_follows_wrapped_401(self): + response = type("Response", (), {"status_code": 401})() + auth_error = Exception("request failed") + auth_error.response = response + wrapper = RuntimeError("model validation failed") + wrapper.__cause__ = auth_error + assert is_hf_authentication_error(wrapper) is True + + def test_hf_authentication_error_does_not_treat_429_as_invalid(self): + response = type("Response", (), {"status_code": 429})() + rate_error = Exception("too many requests") + rate_error.response = response + assert is_hf_authentication_error(rate_error) is False + # --- OOM on CUDA --- @needs_torch diff --git a/studio/backend/utils/_studio_release_build.py b/studio/backend/utils/_studio_release_build.py index 267197a202..07ede36912 100644 --- a/studio/backend/utils/_studio_release_build.py +++ b/studio/backend/utils/_studio_release_build.py @@ -1,7 +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 -"""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 diff --git a/studio/backend/utils/api_errors.py b/studio/backend/utils/api_errors.py index cae8daf287..a3686c3a26 100644 --- a/studio/backend/utils/api_errors.py +++ b/studio/backend/utils/api_errors.py @@ -20,7 +20,7 @@ client-error responses on the ``/v1/*`` surface: CRITICAL: the exception handlers installed by :func:`install_api_error_handlers` are global, but they ONLY transform responses for paths that start with ``/v1/``. For every other path (``/api/...``, frontend routes) they reproduce FastAPI's -default behavior byte-for-byte, because the Studio frontend depends on the +default behavior byte-for-byte, because the Unsloth frontend depends on the ``{"detail": ...}`` shape for ``/api/*``. Public contract (other modules depend on these): @@ -107,7 +107,7 @@ def anthropic_error_body( Returns ``{"type": "error", "request_id": None, "error": {"type", "message"}}``. ``request_id`` is a required (nullable) field on the spec's ErrorResponse; - Studio has no request-id system, so it is null. ``err_type`` defaults to + Unsloth has no request-id system, so it is null. ``err_type`` defaults to :data:`ANTHROPIC_TYPE_BY_STATUS` for ``status`` (``"api_error"`` fallback). """ return { @@ -192,7 +192,7 @@ def install_api_error_handlers(app) -> None: Both handlers are global but only transform responses for OpenAI/Anthropic- compatible surfaces (see :func:`wants_api_error_envelope`: the ``/v1/*`` mount and the preview ``/p/.../v1/*`` mount). Every other path reproduces FastAPI's - default ``{"detail": ...}`` behavior exactly so the Studio frontend keeps working. + default ``{"detail": ...}`` behavior exactly so the Unsloth frontend keeps working. """ @app.exception_handler(RequestValidationError) diff --git a/studio/backend/utils/client_ip.py b/studio/backend/utils/client_ip.py index 94acbf1809..cc48a096d2 100644 --- a/studio/backend/utils/client_ip.py +++ b/studio/backend/utils/client_ip.py @@ -4,12 +4,12 @@ """Resolve the caller's IP for rate limiting. Trust model, in order: - 1. If the operator opts in via ``UNSLOTH_STUDIO_TRUST_FORWARDED`` (Studio behind + 1. If the operator opts in via ``UNSLOTH_STUDIO_TRUST_FORWARDED`` (Unsloth behind their own reverse proxy), honor the *rightmost* ``X-Forwarded-For`` hop -- the one the trusted proxy appended. The leftmost entry is client-controlled and spoofable, so this assumes a proxy that appends (or overwrites) the header; only enable the env var behind such a proxy. - 2. If the socket peer is loopback, honor ``CF-Connecting-IP``. Studio's managed + 2. If the socket peer is loopback, honor ``CF-Connecting-IP``. Unsloth's managed Cloudflare tunnel terminates at 127.0.0.1, so every tunneled visitor would otherwise collapse onto the same socket peer (the local cloudflared process) and share one rate-limit bucket. ``CF-Connecting-IP`` is set by Cloudflare's diff --git a/studio/backend/utils/cpu_threads.py b/studio/backend/utils/cpu_threads.py index 4ed0021054..91d577408d 100644 --- a/studio/backend/utils/cpu_threads.py +++ b/studio/backend/utils/cpu_threads.py @@ -1,7 +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 -"""Early CPU thread-pool configuration for Studio processes.""" +"""Early CPU thread-pool configuration for Unsloth processes.""" import os from typing import MutableMapping, Optional diff --git a/studio/backend/utils/datasets/cache_safe.py b/studio/backend/utils/datasets/cache_safe.py index e629210f33..d2dc7b737a 100644 --- a/studio/backend/utils/datasets/cache_safe.py +++ b/studio/backend/utils/datasets/cache_safe.py @@ -7,7 +7,7 @@ A shared HF datasets cache can contain subtrees owned by another user (for example populated by an earlier root-run job). datasets then raises "[Errno 13] Permission denied: ..._builder.lock" while locking the cached builder, killing the training run even though the dataset itself is fine. -Retry such loads in a Studio-owned cache so the run proceeds; the worst case +Retry such loads in an Unsloth-owned cache so the run proceeds; the worst case is one rebuild of the dataset in the fallback location. """ @@ -26,7 +26,7 @@ def studio_datasets_cache() -> str: def load_dataset_cache_safe(*args, **kwargs): - """datasets.load_dataset, retried in a Studio-owned cache on EACCES.""" + """datasets.load_dataset, retried in an Unsloth-owned cache on EACCES.""" from datasets import load_dataset try: return load_dataset(*args, **kwargs) diff --git a/studio/backend/utils/hardware/VRAM_ESTIMATION.md b/studio/backend/utils/hardware/VRAM_ESTIMATION.md index a6b4de29d2..68ca1d5ffd 100644 --- a/studio/backend/utils/hardware/VRAM_ESTIMATION.md +++ b/studio/backend/utils/hardware/VRAM_ESTIMATION.md @@ -106,7 +106,7 @@ Non_flash_attention = B * num_attention_heads * S^2 * 2 * 12.0 * effective_layer Activations = max(Per_layer_with_gc, Non_flash_attention) ``` -Studio resolves the attention implementation with Unsloth's +Unsloth resolves the attention implementation with Unsloth's `resolve_attention_implementation` helper and uses that result directly. The estimator does not duplicate model-family attention policy. diff --git a/studio/backend/utils/hardware/amd.py b/studio/backend/utils/hardware/amd.py index f5b64c45d0..91a06c9a2a 100644 --- a/studio/backend/utils/hardware/amd.py +++ b/studio/backend/utils/hardware/amd.py @@ -125,7 +125,7 @@ def _run_amd_smi(*args: str, timeout: int = _AMD_SMI_DEFAULT_TIMEOUT) -> Optiona # amd-smi does not exist on Windows (neither Adrenalin nor the HIP SDK # ship a CLI) and can be absent on minimal Linux installs. Disable the # poller in one step instead of burning the 3-strike circuit breaker - # on guaranteed FileNotFoundError spawns. Studio's VRAM display falls + # on guaranteed FileNotFoundError spawns. Unsloth's VRAM display falls # back to torch mem_get_info. if not _amd_smi_disabled: logger.info( diff --git a/studio/backend/utils/hardware/hardware.py b/studio/backend/utils/hardware/hardware.py index 117ad7b780..9fef53e65e 100644 --- a/studio/backend/utils/hardware/hardware.py +++ b/studio/backend/utils/hardware/hardware.py @@ -37,7 +37,7 @@ logger = get_logger(__name__) # ── GPU index ordering ────────────────────────────────────────────────────── # CUDA defaults to CUDA_DEVICE_ORDER=FASTEST_FIRST, numbering GPUs by compute -# performance. nvidia-smi -- and every free-VRAM probe in Studio -- numbers GPUs +# performance. nvidia-smi -- and every free-VRAM probe in Unsloth -- numbers GPUs # by PCI bus id instead. On a mixed-GPU host (e.g. an RTX 5090 alongside an RTX # PRO 6000) the two orderings disagree, so an index picked from nvidia-smi data # ("the emptiest card is GPU 1") gets written into CUDA_VISIBLE_DEVICES and then @@ -49,7 +49,7 @@ logger = get_logger(__name__) # and spawn workers copy os.environ. setdefault so an explicit user override wins. os.environ.setdefault("CUDA_DEVICE_ORDER", "PCI_BUS_ID") -# Studio workers can import MLX without importing unsloth first, so mirror the +# Unsloth workers can import MLX without importing unsloth first, so mirror the # package bootstrap here. Keep an explicit user value authoritative. if platform.system() == "Darwin" and platform.machine() == "arm64": os.environ.setdefault("AGX_RELAX_CDM_CTXSTORE_TIMEOUT", "1") @@ -117,7 +117,7 @@ def _has_mlx() -> bool: def _has_usable_mlx_stack() -> bool: - """True only when the FULL Studio MLX training/export stack is usable + """True only when the FULL Unsloth MLX training/export stack is usable (mlx + mlx-lm + mlx-vlm at the minimum versions unsloth-zoo requires), not just a bare ``import mlx.core``. A backtracked/old mlx-vlm still imports but breaks VLM Train/Export, so the training gate must match the self-heal's own @@ -538,21 +538,31 @@ def _torch_get_physical_gpu_count() -> Optional[int]: def _torch_get_per_device_info(device_indices: list[int]) -> list[Dict[str, Any]]: - """Query torch for per-GPU name, total VRAM, and used VRAM.""" + """Query torch for per-GPU name, total VRAM, and used VRAM. + + ``used_gb`` is ``None`` on Windows ROCm when ``hipMemGetInfo`` reports + ``free == total`` (ROCm/ROCm#1909): that 0 means unknown, not empty. + """ mod, _ = _torch_get_device_module() if mod is None: return [] + # free==total is a Windows-ROCm-only quirk. + _win_rocm = sys.platform == "win32" and IS_ROCM devices = [] for ordinal, phys_idx in enumerate(device_indices): try: # torch ordinals are 0-based relative to CUDA_VISIBLE_DEVICES. props = mod.get_device_properties(ordinal) total_bytes = props.total_memory + used_bytes: Optional[int] # Prefer mem_get_info (system-wide) so auto-select sees other consumers. if hasattr(mod, "mem_get_info"): free_bytes, total_bytes = mod.mem_get_info(ordinal) used_bytes = total_bytes - free_bytes + # free==total is the broken-API sentinel, not an idle GPU. + if _win_rocm and free_bytes == total_bytes: + used_bytes = None else: used_bytes = mod.memory_allocated(ordinal) devices.append( @@ -561,7 +571,7 @@ def _torch_get_per_device_info(device_indices: list[int]) -> list[Dict[str, Any] "visible_ordinal": ordinal, "name": props.name, "total_gb": round(total_bytes / (1024**3), 2), - "used_gb": round(used_bytes / (1024**3), 2), + "used_gb": round(used_bytes / (1024**3), 2) if used_bytes is not None else None, } ) except Exception as e: @@ -724,20 +734,30 @@ def _rocm_linux_sysfs_vram_gb() -> tuple[Optional[float], Optional[float]]: return None, None -def _rocm_windows_perf_counter_vram_gb() -> tuple[Optional[float], Optional[float]]: - """Query system-wide dedicated GPU VRAM via Windows Performance Counters. +# ── Windows AMD/ROCm per-adapter VRAM (issue #7072) ────────────────────────── +# amd-smi is disabled and hipMemGetInfo reports free==total, so read used from the +# per-LUID "GPU Adapter Memory" perf counters and take each total from torch, so +# every GPU shows instead of one fake device with GPU 0's total. +# Placeholder adapters (Basic Render Driver / idle iGPU) drop only when they would +# outnumber the real torch devices. +_ROCM_WIN_ADAPTER_MIN_BYTES = 64 * 1024 * 1024 # 64 MiB - Same data source as Task Manager, so cross-process usage is accurate. - Works for any GPU vendor without amd-smi or nvidia-smi. - Returns (used_gb, total_gb) or (None, None) on failure. + +def _rocm_windows_perf_counter_vram_by_adapter() -> Optional[list[tuple[str, float]]]: + """Per-adapter dedicated VRAM usage on Windows via Performance Counters. + + Returns ``[(instance_name, used_bytes)]`` (one per LUID-named adapter), or + ``None`` when the counter is unavailable/localized/empty so callers fall back. """ if platform.system() != "Windows": - return None, None + return None try: + # Emit "|" per sample, or a __NONE__ sentinel. ps = ( "$s=(Get-Counter '\\GPU Adapter Memory(*)\\Dedicated Usage'" " -ErrorAction SilentlyContinue).CounterSamples;" - "if($s){($s|Measure-Object CookedValue -Sum).Sum}else{-1}" + "if($s){$s|ForEach-Object{'{0}|{1}' -f $_.InstanceName,[int64]$_.CookedValue}}" + "else{'__NONE__'}" ) r = subprocess.run( ["powershell", "-NoProfile", "-NonInteractive", "-Command", ps], @@ -746,16 +766,167 @@ def _rocm_windows_perf_counter_vram_gb() -> tuple[Optional[float], Optional[floa timeout = 5, ) if r.returncode != 0 or not r.stdout.strip(): - return None, None - used_bytes = float(r.stdout.strip()) - if used_bytes < 0: - return None, None - import torch as _torch - - total_bytes = _torch.cuda.get_device_properties(0).total_memory - return round(used_bytes / (1024**3), 2), round(total_bytes / (1024**3), 2) + return None + adapters: list[tuple[str, float]] = [] + for line in r.stdout.splitlines(): + line = line.strip() + if not line or line == "__NONE__" or "|" not in line: + continue + instance, _, raw = line.rpartition("|") + try: + used = float(raw.strip()) + except (ValueError, TypeError): + continue + if used < 0: + continue + adapters.append((instance.strip(), used)) + return adapters or None except Exception: - return None, None + return None + + +def _match_adapter_used_to_devices( + adapter_useds: list[float], device_totals: list[float] +) -> list[Optional[float]]: + """Attribute per-adapter used bytes to torch devices by capacity ranking. + + Windows shares no key between LUID counters and torch ordinals, so usages are + ranked against device totals and each is trusted only when capacity *forces* it + (it exceeds every smaller device); an ambiguous ranking reports unknown + (``None``) rather than fabricate a per-index free. + + Extra counters mean a hidden/display adapter, and the noise filter may have + dropped a real reading, so values are emitted only when the supra-threshold + counters number EXACTLY the visible devices AND capacity forces the mapping; + otherwise every device is unknown. Best-effort but correct for the common + loaded-card case (#7072). Returns a list aligned to ``device_totals``. + """ + n = len(device_totals) + if n == 0: + return [] + useds = sorted(adapter_useds, reverse = True) + ranked_positions = sorted(range(n), key = lambda i: -device_totals[i]) + ranked_totals = [device_totals[pos] for pos in ranked_positions] + assigned: list[Optional[float]] + # More counters than devices -> a hidden/display adapter (check before noise filter). + if len(useds) > n: + non_trivial = [u for u in useds if u >= _ROCM_WIN_ADAPTER_MIN_BYTES] + if len(non_trivial) != n: + # Not a clean bijection (a masked GPU is busy or a visible card idle): + # no counter maps to a specific card, so report unknown. + return [None] * n + # Exactly n supra-threshold counters: extras were placeholders, so a + # capacity-ranked bijection is plausible. + useds = non_trivial + ranked_useds = [useds[rank] for rank in range(n)] + # A usage above its ranked capacity is a hidden larger GPU; clamping onto the + # smaller card would fabricate a fully-used reading. + for rank in range(n): + if ranked_useds[rank] > ranked_totals[rank]: + return [None] * n + # Capacity forces the mapping only when the usage exceeds the next-smaller + # capacity; the smallest card and merely-fitting usages stay unknown. + # Keeps 40 GiB over 48/8 GiB -> [40, None]. + assigned = [None] * n + for rank, pos in enumerate(ranked_positions): + if rank + 1 < n and ranked_useds[rank] > ranked_totals[rank + 1]: + assigned[pos] = min(ranked_useds[rank], device_totals[pos]) + return assigned + # No hidden adapters: every counter is a visible card, so ranking is a permutation. + ranked_useds = [useds[rank] if rank < len(useds) else 0.0 for rank in range(n)] + # Ambiguous if a strictly larger usage also fits the next smaller card: the two + # could be swapped without breaking capacity, so ranking can't tell them apart. + for rank in range(n - 1): + upper, lower = ranked_useds[rank], ranked_useds[rank + 1] + if upper > lower and upper <= ranked_totals[rank + 1]: + return [None] * n + assigned = [None] * n + for rank, pos in enumerate(ranked_positions): + if rank < len(useds): + assigned[pos] = min(useds[rank], device_totals[pos]) + return assigned + + +def _rocm_windows_per_device_vram(device_indices: list[int]) -> list[Dict[str, Any]]: + """Per-GPU VRAM on Windows AMD/ROCm: total from torch properties (reliable), + used from the per-adapter Dedicated Usage counter. + + Returns ``{index, visible_ordinal, name, used_gb, total_gb}`` per visible GPU + (``used_gb`` may be ``None`` when the counter is unavailable), or ``[]`` when + torch can't enumerate devices so callers fall through to the torch last resort. + """ + if platform.system() != "Windows": + return [] + mod, _ = _torch_get_device_module() + if mod is None: + return [] + # Totals/names from torch properties (mem_get_info's free==total quirk zeroes used). + dev_meta: list[Dict[str, Any]] = [] + for ordinal, phys_idx in enumerate(device_indices): + try: + props = mod.get_device_properties(ordinal) + dev_meta.append( + { + "index": phys_idx, + "visible_ordinal": ordinal, + "name": props.name, + "total_bytes": int(props.total_memory), + } + ) + except Exception as e: + logger.debug("torch property probe failed for ordinal %d: %s", ordinal, e) + if not dev_meta: + return [] + + adapters = _rocm_windows_perf_counter_vram_by_adapter() + if adapters: + assigned = _match_adapter_used_to_devices( + [used for _, used in adapters], + [d["total_bytes"] for d in dev_meta], + ) + else: + # Counter unavailable: show every GPU with a correct total, used unknown. + assigned = [None] * len(dev_meta) + + devices: list[Dict[str, Any]] = [] + for meta, used_bytes in zip(dev_meta, assigned): + total_gb = round(meta["total_bytes"] / (1024**3), 2) + used_gb = round(used_bytes / (1024**3), 2) if used_bytes is not None else None + devices.append( + { + "index": meta["index"], + "visible_ordinal": meta["visible_ordinal"], + "name": meta["name"], + "used_gb": used_gb, + "total_gb": total_gb, + } + ) + return devices + + +def _rocm_windows_device_payload_entry( + device: DeviceType, dev: Dict[str, Any], gpu_util_pct: Optional[float] +) -> Dict[str, Any]: + """Build a ``get_gpu_utilization`` device entry from a per-device VRAM dict.""" + total_gb = dev["total_gb"] + used_gb = dev["used_gb"] + return { + "available": True, + "backend": _backend_label(device), + "index": dev["index"], + "visible_ordinal": dev["visible_ordinal"], + "name": dev.get("name", "Unknown"), + "gpu_utilization_pct": gpu_util_pct, + "temperature_c": None, + "vram_used_gb": used_gb, + "vram_total_gb": total_gb, + "vram_utilization_pct": round((used_gb / total_gb) * 100, 1) + if total_gb and total_gb > 0 and used_gb is not None + else None, + "power_draw_w": None, + "power_limit_w": None, + "power_utilization_pct": None, + } def _gpu_utilization_payload( @@ -821,30 +992,24 @@ def get_gpu_utilization() -> Dict[str, Any]: index_kind = result.get("index_kind"), ) - # Fallback Windows ROCm + # Fallback Windows ROCm: per-adapter VRAM attribution (issue #7072), so + # every visible GPU is shown instead of a sum collapsed onto one device. if IS_ROCM and platform.system() == "Windows": - _win_used, _win_total = _rocm_windows_perf_counter_vram_gb() - if _win_used is not None and _win_total is not None: - _win_util = _rocm_windows_perf_counter_gpu_util_pct() + _win_ids = _get_parent_visible_gpu_spec().get("numeric_ids") + if not _win_ids: + _win_ids = list(range(_torch_get_physical_gpu_count() or 0)) + _win_devices = _rocm_windows_per_device_vram(_win_ids) + if _win_devices: + # A single visible GPU can own the aggregate 3D-engine utilization; + # across several GPUs the sum isn't per-device, so leave it unset. + _win_util = ( + _rocm_windows_perf_counter_gpu_util_pct() if len(_win_devices) == 1 else None + ) return _gpu_utilization_payload( device, [ - { - "available": True, - "backend": _backend_label(device), - "index": 0, - "visible_ordinal": 0, - "gpu_utilization_pct": _win_util, - "temperature_c": None, - "vram_used_gb": _win_used, - "vram_total_gb": _win_total, - "vram_utilization_pct": round((_win_used / _win_total) * 100, 1) - if _win_total > 0 - else None, - "power_draw_w": None, - "power_limit_w": None, - "power_utilization_pct": None, - } + _rocm_windows_device_payload_entry(device, _wd, _win_util) + for _wd in _win_devices ], ) @@ -901,7 +1066,7 @@ def get_gpu_utilization() -> Dict[str, Any]: "vram_used_gb": _used, "vram_total_gb": _total, "vram_utilization_pct": round((_used / _total) * 100, 1) - if _total > 0 + if _total > 0 and _used is not None else None, "power_draw_w": None, "power_limit_w": None, @@ -995,19 +1160,27 @@ def _apply_unified_memory_correction( endpoints stay in sync on AMD iGPUs with unified memory. """ torch_total_gb = torch_info["total_gb"] + torch_used_gb = torch_info.get("used_gb") smi_total_gb = device_metrics.get("vram_total_gb") or 0.0 + # torch sees the full unified (GTT) pool; amd-smi only the dedicated carve-out. + # Adopt torch's larger total regardless of used: on Windows ROCm torch_used is + # None (free==total sentinel) but its total stays authoritative. Overwrite used + # only when torch's is known, then recompute utilization against whatever remains. if torch_total_gb > smi_total_gb: - torch_used_gb = torch_info["used_gb"] device_metrics["vram_total_gb"] = torch_total_gb - device_metrics["vram_used_gb"] = torch_used_gb + if torch_used_gb is not None: + device_metrics["vram_used_gb"] = torch_used_gb + _used_for_pct = device_metrics.get("vram_used_gb") device_metrics["vram_utilization_pct"] = ( - round((torch_used_gb / torch_total_gb) * 100, 1) if torch_total_gb > 0 else None + round((_used_for_pct / torch_total_gb) * 100, 1) + if torch_total_gb > 0 and _used_for_pct is not None + else None ) logger.debug( - "ROCm unified memory: replaced amd-smi VRAM (%.2f GB) with " - "torch mem_get_info total (%.2f GB) for device %s", - smi_total_gb, + "ROCm unified memory: adopted torch mem_get_info total (%.2f GB) over " + "amd-smi (%.2f GB) for device %s", torch_total_gb, + smi_total_gb, torch_info.get("index"), ) @@ -1067,6 +1240,49 @@ def get_visible_gpu_utilization() -> Dict[str, Any]: _reconcile_rocm_unified_memory(result, numeric_ids) return result + # Windows AMD/ROCm (issue #7072): the System tab's VRAM source. The torch + # fallback below would report used==0 (free==total), so read per-adapter + # Dedicated Usage instead; total from torch properties. + if IS_ROCM and platform.system() == "Windows": + win_numeric_ids = parent_visible_spec.get("numeric_ids") + if win_numeric_ids: + win_ids = win_numeric_ids + win_index_kind = "physical" + else: + win_ids = list(range(_torch_get_physical_gpu_count() or 0)) + win_index_kind = "relative" + win_devices = _rocm_windows_per_device_vram(win_ids) + if win_devices: + devices = [] + for wd in win_devices: + total = wd["total_gb"] + used = wd["used_gb"] + devices.append( + { + "index": wd["index"], + "index_kind": win_index_kind, + "visible_ordinal": wd["visible_ordinal"], + "name": wd.get("name"), + "gpu_utilization_pct": None, + "temperature_c": None, + "vram_used_gb": used, + "vram_total_gb": total, + "vram_utilization_pct": round((used / total) * 100, 1) + if total and total > 0 and used is not None + else None, + "power_draw_w": None, + "power_limit_w": None, + "power_utilization_pct": None, + } + ) + return { + "available": True, + "backend": _backend_label(device), + "parent_visible_gpu_ids": win_numeric_ids or [], + "devices": devices, + "index_kind": win_index_kind, + } + # Torch-based fallback for CUDA (nvidia-smi unavailable, AMD ROCm) and XPU (Intel) if device in (DeviceType.CUDA, DeviceType.XPU): parent_ids = get_parent_visible_gpu_ids() @@ -1094,7 +1310,7 @@ def get_visible_gpu_utilization() -> Dict[str, Any]: "vram_used_gb": used, "vram_total_gb": total, "vram_utilization_pct": round((used / total) * 100, 1) - if total > 0 + if total > 0 and used is not None else None, "power_draw_w": None, "power_limit_w": None, diff --git a/studio/backend/utils/helper_precache_settings.py b/studio/backend/utils/helper_precache_settings.py index db19a2d028..e7d3c0e6dd 100644 --- a/studio/backend/utils/helper_precache_settings.py +++ b/studio/backend/utils/helper_precache_settings.py @@ -32,7 +32,7 @@ def helper_model_disabled_by_env() -> bool: def get_helper_precache_enabled() -> bool: """Read the persisted startup pre-cache preference. - Missing or unreadable settings default to False so Studio startup never + Missing or unreadable settings default to False so Unsloth startup never performs optional network work unless the user explicitly opted in. """ try: @@ -45,7 +45,7 @@ def get_helper_precache_enabled() -> bool: def set_helper_precache_enabled(value: Any) -> bool: - """Persist whether Studio should pre-cache the Helper LLM at startup.""" + """Persist whether Unsloth should pre-cache the Helper LLM at startup.""" parsed = _coerce_bool(value) if parsed is None: raise ValueError("Helper LLM startup pre-cache must be true or false.") diff --git a/studio/backend/utils/hf_token_validation.py b/studio/backend/utils/hf_token_validation.py new file mode 100644 index 0000000000..7247c6e756 --- /dev/null +++ b/studio/backend/utils/hf_token_validation.py @@ -0,0 +1,208 @@ +# SPDX-License-Identifier: AGPL-3.0-only +# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0 + +"""Cached, rate-limited Hugging Face token validation.""" + +from __future__ import annotations + +import hashlib +import threading +import time +from collections import deque +from dataclasses import dataclass +from typing import Literal + +from huggingface_hub import HfApi +from huggingface_hub.utils import build_hf_headers, get_session + + +TokenValidationStatus = Literal["valid", "invalid", "rate_limited", "unavailable"] + + +@dataclass(frozen = True) +class TokenValidationResult: + status: TokenValidationStatus + retry_after_seconds: int | None = None + + +_WINDOW_SECONDS = 3600.0 +_MAX_ATTEMPTS = 3 +_CACHE_TTL_SECONDS = 3600.0 +_TEMPORARY_CACHE_TTL_SECONDS = 15.0 +_MAX_BUCKETS = 4096 +_MAX_CACHE_ENTRIES = 4096 +_INFLIGHT_WAIT_SECONDS = 30.0 +_REMOTE_TIMEOUT_SECONDS = 10.0 + +_attempts: dict[str, deque[float]] = {} +_cache: dict[str, tuple[float, TokenValidationResult]] = {} +_inflight: dict[str, threading.Event] = {} +_lock = threading.Lock() + + +def _fingerprint(token: str) -> str: + return hashlib.sha256(token.encode("utf-8")).hexdigest() + + +def _prune_attempts(bucket: deque[float], now: float) -> None: + while bucket and now - bucket[0] >= _WINDOW_SECONDS: + bucket.popleft() + + +def _prune_locked(now: float) -> None: + for key in list(_attempts): + bucket = _attempts[key] + _prune_attempts(bucket, now) + if not bucket: + del _attempts[key] + for key, (expires_at, _result) in list(_cache.items()): + if expires_at <= now: + del _cache[key] + + +def _cached_locked(fingerprint: str, now: float) -> TokenValidationResult | None: + cached = _cache.get(fingerprint) + if cached is None: + return None + expires_at, result = cached + if expires_at <= now: + del _cache[fingerprint] + return None + return result + + +def _retry_after(bucket: deque[float], now: float) -> int: + return max(1, int(_WINDOW_SECONDS - (now - bucket[0])) + 1) + + +def _reserve_attempt_locked(rate_key: str, now: float) -> TokenValidationResult | None: + bucket = _attempts.get(rate_key) + if bucket is None: + if len(_attempts) >= _MAX_BUCKETS: + _prune_locked(now) + if len(_attempts) >= _MAX_BUCKETS: + return TokenValidationResult( + status = "rate_limited", + retry_after_seconds = max(1, int(_WINDOW_SECONDS)), + ) + bucket = _attempts[rate_key] = deque() + _prune_attempts(bucket, now) + if len(bucket) >= _MAX_ATTEMPTS: + return TokenValidationResult( + status = "rate_limited", + retry_after_seconds = _retry_after(bucket, now), + ) + bucket.append(now) + return None + + +def _http_status(response: object | None) -> int | None: + status = getattr(response, "status_code", None) + try: + return int(status) if status is not None else None + except (TypeError, ValueError): + return None + + +def _remote_retry_after(response: object | None) -> int | None: + headers = getattr(response, "headers", None) + if not headers: + return None + raw = headers.get("Retry-After") + try: + return max(1, int(float(raw))) if raw is not None else None + except (TypeError, ValueError): + return None + + +def _classify_response(response: object | None) -> TokenValidationResult: + status = _http_status(response) + if status is not None and 200 <= status < 300: + return TokenValidationResult(status = "valid") + if status == 401: + return TokenValidationResult(status = "invalid") + if status == 429: + return TokenValidationResult( + status = "rate_limited", + retry_after_seconds = _remote_retry_after(response), + ) + return TokenValidationResult(status = "unavailable") + + +def _check_remote(token: str) -> TokenValidationResult: + api = HfApi() + try: + # HfApi.whoami has no timeout parameter in the pinned Hub client. + # Use its session and headers against the same whoami endpoint. + response = get_session().get( + f"{api.endpoint}/api/whoami-v2", + headers = build_hf_headers(token = token), + timeout = _REMOTE_TIMEOUT_SECONDS, + ) + except Exception as exc: + # huggingface-hub 0.36.x can wrap a 401 as requests.HTTPError. + return _classify_response(getattr(exc, "response", None)) + return _classify_response(response) + + +def validate_hf_token(token: str, *, rate_key: str) -> TokenValidationResult: + """Validate ``token`` without retaining it, sharing results across callers. + + Cached checks do not consume the caller's three-per-hour network budget. A + single-flight event also prevents simultaneously mounted UI surfaces from + sending duplicate ``whoami`` requests for the same token. + """ + normalized = token.strip() + if not normalized: + return TokenValidationResult(status = "invalid") + token_fingerprint = _fingerprint(normalized) + owner_event: threading.Event | None = None + + try: + while True: + now = time.monotonic() + with _lock: + cached = _cached_locked(token_fingerprint, now) + if cached is not None: + return cached + waiting = _inflight.get(token_fingerprint) + if waiting is None: + limited = _reserve_attempt_locked(rate_key, now) + if limited is not None: + return limited + owner_event = threading.Event() + _inflight[token_fingerprint] = owner_event + break + if not waiting.wait(_INFLIGHT_WAIT_SECONDS): + return TokenValidationResult(status = "unavailable") + + result = _check_remote(normalized) + now = time.monotonic() + ttl = ( + _CACHE_TTL_SECONDS + if result.status in ("valid", "invalid") + else max(_TEMPORARY_CACHE_TTL_SECONDS, float(result.retry_after_seconds or 0)) + ) + with _lock: + if len(_cache) >= _MAX_CACHE_ENTRIES: + _prune_locked(now) + if len(_cache) < _MAX_CACHE_ENTRIES: + _cache[token_fingerprint] = (now + ttl, result) + return result + finally: + if owner_event is not None: + with _lock: + event = _inflight.get(token_fingerprint) + if event is owner_event: + _inflight.pop(token_fingerprint, None) + event.set() + + +def reset_hf_token_validation_state() -> None: + """Clear process state for test isolation.""" + with _lock: + for event in _inflight.values(): + event.set() + _inflight.clear() + _attempts.clear() + _cache.clear() diff --git a/studio/backend/utils/hf_xet_fallback.py b/studio/backend/utils/hf_xet_fallback.py index 9bc4a60fad..2628b99a2d 100644 --- a/studio/backend/utils/hf_xet_fallback.py +++ b/studio/backend/utils/hf_xet_fallback.py @@ -1,9 +1,9 @@ # SPDX-License-Identifier: AGPL-3.0-only # Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0 -"""Studio shim over the shared ``unsloth_zoo.hf_xet_fallback`` Xet -> HTTP stall fallback. +"""Unsloth shim over the shared ``unsloth_zoo.hf_xet_fallback`` Xet -> HTTP stall fallback. -Re-exports the shared API and injects Studio's marker-aware cache purge +Re-exports the shared API and injects Unsloth's marker-aware cache purge (``prepare_cache_for_transport``) so the download manager keeps its ``.transport`` marker semantics on the HTTP retry. @@ -68,7 +68,7 @@ def _load_shared() -> bool: _shared_available = True _shared_import_error = None return True - except Exception as exc2: # noqa: BLE001 - degrade so Studio still boots with plain HF + except Exception as exc2: # noqa: BLE001 - degrade so Unsloth still boots with plain HF _shared_import_error = exc2 _shared_available = False import logging as _logging @@ -263,7 +263,7 @@ __all__ = [ def _studio_prepare_for_http(repo_type: str, repo_id: str) -> None: - """Studio's marker-aware purge before an HTTP resume, keeping the download manager's ``.transport`` + """Unsloth's marker-aware purge before an HTTP resume, keeping the download manager's ``.transport`` accounting consistent (vs unsloth_zoo's generic default). Guarded: a purge failure is logged, not fatal to the retry.""" try: @@ -273,7 +273,7 @@ def _studio_prepare_for_http(repo_type: str, repo_id: str) -> None: try: from loggers import get_logger get_logger(__name__).debug( - "Studio prepare_cache_for_transport failed for %s: %s", repo_id, exc + "Unsloth prepare_cache_for_transport failed for %s: %s", repo_id, exc ) except ModuleNotFoundError as logger_exc: if logger_exc.name != "loggers": @@ -294,8 +294,8 @@ def hf_hub_download_with_xet_fallback( on_status: Optional[Callable[[str], None]] = None, force_download: bool = False, ) -> str: - """Single-file download via the shared fallback with Studio's marker-aware HTTP-retry prep. - ``force_download`` re-fetches a newer blob over a cached one (Studio's model-update path).""" + """Single-file download via the shared fallback with Unsloth's marker-aware HTTP-retry prep. + ``force_download`` re-fetches a newer blob over a cached one (Unsloth's model-update path).""" return _shared_hf_hub_download_with_xet_fallback( repo_id, filename, @@ -313,6 +313,6 @@ def hf_hub_download_with_xet_fallback( def snapshot_download_with_xet_fallback(repo_id: str, **kwargs: Any) -> str: - """Whole-repo download via the shared fallback with Studio's marker-aware HTTP-retry prep.""" + """Whole-repo download via the shared fallback with Unsloth's marker-aware HTTP-retry prep.""" kwargs.setdefault("prepare_for_http_fn", _studio_prepare_for_http) return _shared_snapshot_download_with_xet_fallback(repo_id, **kwargs) diff --git a/studio/backend/utils/hidden_models.py b/studio/backend/utils/hidden_models.py new file mode 100644 index 0000000000..20d0bb966e --- /dev/null +++ b/studio/backend/utils/hidden_models.py @@ -0,0 +1,142 @@ +# SPDX-License-Identifier: AGPL-3.0-only +# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0 + +"""Infra-only model detection shared by the model routes and the hub +inventory. Lives directly under ``utils`` (not ``utils.models``) so the hub +cache scanner can import it without pulling in ``utils/models/__init__.py``, +which eagerly loads the model-config/checkpoint stack, and without importing +``routes.models`` (import-time side effects, would cycle).""" + +from __future__ import annotations + +import re +from pathlib import Path +from typing import Optional + +# Hub repo id shape ("owner/name", no leading separator); anything else is +# treated as a local filesystem path. +_HF_REPO_ID_RE = re.compile(r"^[A-Za-z0-9][\w.\-]*/[\w.\-]+$") + +# The llama.cpp install-validation probe repo. Always hidden. +_PROBE_REPO_ID = "ggml-org/models" +# The probe's on-disk filename. Carries the ".gguf" so it stays specific and +# does not hide unrelated repos like ``user/stories260K-finetune-GGUF``. +_PROBE_FILENAME = "stories260k.gguf" +# Keep previously cached defaults hidden after settings changes. +_DEFAULT_EMBEDDING_REPO_IDS = { + "unsloth/bge-small-en-v1.5", + "unsloth/bge-small-en-v1.5-GGUF", +} +# Local copies do not always retain the repo id. Keep a narrow basename +# fallback for Studio's static default embedder only; configured custom repos +# remain exact-match-only. +_DEFAULT_EMBEDDING_PATH_BASENAMES = {"bge-small-en-v1.5"} + + +def _safe_resolve(path: Path) -> Optional[str]: + """resolve() to a string, or None when the path is inaccessible.""" + try: + return str(path.resolve()) + except OSError: + return None + + +def _existing_resolved_path(value: str) -> Optional[str]: + """Resolve an existing local path.""" + path = Path(value).expanduser() + try: + if not path.exists(): + return None + except OSError: + return None + return _safe_resolve(path) + + +def _path_contains_repo_id(value: str, repo_ids: set[str]) -> bool: + """Match exact repo-derived path segments.""" + parts = [part for part in value.lower().replace("\\", "/").split("/") if part] + for repo_id in repo_ids: + owner, name = repo_id.split("/", 1) + if f"models--{owner}--{name}" in parts: + return True + if any( + parts[index] == owner and parts[index + 1] == name for index in range(len(parts) - 1) + ): + return True + return False + + +def _path_basename_is_default_embedder(value: str) -> bool: + """Match a default embedder folder or a suffixed local weight filename.""" + normalized = value.lower().replace("\\", "/").rstrip("/") + basename = normalized.rsplit("/", 1)[-1] + return any( + basename == needle + or any(basename.startswith(f"{needle}{separator}") for separator in ("-", "_", ".")) + for needle in _DEFAULT_EMBEDDING_PATH_BASENAMES + ) + + +def is_hidden_model(*values: str | None) -> bool: + """True if any id/path is the RAG embedding model (the effective embedder + or its GGUF companion repo) or the llama.cpp install validation probe + (ggml-org/models / stories260K), so pickers hide them (GGUF and non-GGUF). + None are usable chat models; the probe can be cached as a side effect of + installing the prebuilt llama-server and otherwise sorts smallest, so it + would be auto-selected. + + Hub repo ids are matched EXACTLY (case-insensitive full "owner/name"), so a + custom embedder with a generic basename like "org/model" cannot substring + hide unrelated cached repos such as "user/model-chat" or "org/model-GGUF". + Existing paths take precedence over the identical ``owner/name`` repo + shape. Cache and LM Studio paths use exact repo-derived segments. Local + copies of the static default embedder also use a boundary-aware basename + fallback; configured custom repos never do.""" + from core.rag import config as rag_config + + hidden_repo_ids = { + _PROBE_REPO_ID.lower(), + *(repo_id.lower() for repo_id in _DEFAULT_EMBEDDING_REPO_IDS), + } + exact_paths: list[str] = [] + for model in { + rag_config.EMBEDDING_MODEL, + rag_config.default_gguf_repo(), + rag_config.effective_embedding_model(), + rag_config.effective_gguf_repo(), + }: + existing_path = _existing_resolved_path(model) + if existing_path: + exact_paths.append(existing_path.lower()) + elif _HF_REPO_ID_RE.match(model): + hidden_repo_ids.add(model.lower()) + else: + resolved = _safe_resolve(Path(model).expanduser()) + if resolved: + exact_paths.append(resolved.lower()) + for v in values: + if not v: + continue + low = v.lower() + if _HF_REPO_ID_RE.match(v): + # A repo id ("owner/name"): match the hidden set exactly. It is + # never a filesystem path, so skip the path/filename checks. + if low in hidden_repo_ids: + return True + continue + # Anything else is treated as a filesystem path (the cached snapshot + # path, or a local model id). Match the probe by its exact filename and + # any configured local-path embedder by exact resolved path. Split on + # both separators so a Windows-style path ("...\\stories260K.gguf") is + # matched even when this runs on a POSIX interpreter (and vice versa). + if low.replace("\\", "/").rsplit("/", 1)[-1] == _PROBE_FILENAME: + return True + if _path_basename_is_default_embedder(v): + return True + if _path_contains_repo_id(v, hidden_repo_ids): + return True + if exact_paths: + resolved = _safe_resolve(Path(v).expanduser()) + if resolved and resolved.lower() in exact_paths: + return True + return False diff --git a/studio/backend/utils/host_policy.py b/studio/backend/utils/host_policy.py index f506eadc03..55565bb338 100644 --- a/studio/backend/utils/host_policy.py +++ b/studio/backend/utils/host_policy.py @@ -1,7 +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 -"""Bind-host trust policy for the Studio backend. +"""Bind-host trust policy for the Unsloth backend. Stdlib only -- safe to import without the rest of the backend. diff --git a/studio/backend/utils/llama_cpp_update.py b/studio/backend/utils/llama_cpp_update.py index f6d3635301..67733bde35 100644 --- a/studio/backend/utils/llama_cpp_update.py +++ b/studio/backend/utils/llama_cpp_update.py @@ -82,7 +82,7 @@ def _utcnow() -> str: def _find_binary() -> Optional[str]: """Locate the active llama-server binary via the inference backend's own - resolver, so update targets exactly what Studio runs. Lazy import keeps the + resolver, so update targets exactly what Unsloth runs. Lazy import keeps the heavy inference module off this module's import path.""" try: from core.inference.llama_cpp import LlamaCppBackend @@ -109,7 +109,7 @@ def _installer_script() -> Optional[Path]: """Locate install_llama_prebuilt.py. Honours UNSLOTH_LLAMA_INSTALLER, then searches up from this file for both ``/install_llama_prebuilt.py`` and ``/studio/install_llama_prebuilt.py`` so it works in the dev tree and - in an installed Studio layout.""" + in an installed Unsloth layout.""" env = os.environ.get("UNSLOTH_LLAMA_INSTALLER") if env and Path(env).is_file(): return Path(env) @@ -227,7 +227,7 @@ def _is_under(path: Path, root: Path) -> bool: def _llama_install_root(binary: Optional[str]) -> Optional[Path]: - """The Studio-managed llama.cpp root the active binary lives under, or None + """The Unsloth-managed llama.cpp root the active binary lives under, or None when the binary is unmanaged. Installing anywhere the active binary is not would not replace what _find_llama_server_binary runs (which prefers a pinned LLAMA_SERVER_PATH, then UNSLOTH_LLAMA_CPP_PATH, then a llama.cpp tree), so we @@ -327,7 +327,7 @@ def _source_build_status(binary: str, *, force_refresh: bool) -> Optional[dict]: def _is_external_link(path: Optional[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, so Studio must never auto-update it.""" + the user's own llama.cpp checkout, so Unsloth must never auto-update it.""" if path is None: return False try: @@ -479,6 +479,7 @@ def _run_update( asset: Optional[str], script: Path, pin_release_tag: Optional[str] = None, + force_cpu: bool = False, ) -> None: """Worker: put the backend into a maintenance state, run the installer for the latest prebuilt, then refresh caches so the next load uses the new build. @@ -522,6 +523,12 @@ def _run_update( if pin_release_tag: cmd.extend(["--published-release-tag", pin_release_tag]) cmd.extend(_rocm_install_args(asset)) + # Re-assert a deliberate CPU install (--force-cpu) so detect_host on a GPU host + # does not re-route to a GPU/Vulkan bundle and revive the crash (#7213). --force-cpu + # (not --cpu-fallback) also re-persists force_cpu, keeping the choice across future + # updates. A natural fallback (or a legacy marker without the flag) heals to GPU (#6097). + if force_cpu: + cmd.append("--force-cpu") logger.info("llama update: installing", cmd = " ".join(cmd)) # Stream progress lines into job["progress"]. env = dict(os.environ, UNSLOTH_PROGRESS_PERCENT_STEP = "5") @@ -635,7 +642,7 @@ def start_update() -> dict: "reason": "local_link", "message": ( "llama.cpp is a local directory linked with --with-llama-cpp-dir; " - "Studio won't replace it. Update your own llama.cpp checkout instead." + "Unsloth won't replace it. Update your own llama.cpp checkout instead." ), "job": get_update_status()["job"], } @@ -671,6 +678,7 @@ def start_update() -> dict: repo = marker.get("published_repo") or DEFAULT_PUBLISHED_REPO from_tag = marker.get("tag") or marker.get("release_tag") asset = marker.get("asset") + force_cpu = bool(marker.get("force_cpu")) # Install exactly the release the banner offered: the installer's own # "latest" is commit-date ordered and can lag the published_at pick # above, reinstalling the current build in a loop (the #6219 class). @@ -705,6 +713,8 @@ def start_update() -> dict: repo = (res or {}).get("repo") or DEFAULT_PUBLISHED_REPO from_tag = None asset = (res or {}).get("asset") + # Source builds carry no forced-CPU marker, so nothing to preserve here. + force_cpu = False # No pin: source-build detection resolves via --resolve-prebuilt latest, # the same resolver the unpinned apply uses, so the two already agree. pin_release_tag = None @@ -735,7 +745,7 @@ def start_update() -> dict: thread = threading.Thread( target = _run_update, - args = (install_dir, repo, asset, script, pin_release_tag), + args = (install_dir, repo, asset, script, pin_release_tag, force_cpu), name = "llama-cpp-update", daemon = True, ) diff --git a/studio/backend/utils/mlx_repair.py b/studio/backend/utils/mlx_repair.py index 7e1c9864c9..4ea1ec62f5 100644 --- a/studio/backend/utils/mlx_repair.py +++ b/studio/backend/utils/mlx_repair.py @@ -3,7 +3,7 @@ """Best-effort MLX self-heal for Apple Silicon. -On macOS, Studio enables Train/Export only when the MLX training/export stack is +On macOS, Unsloth enables Train/Export only when the MLX training/export stack is usable (see utils.hardware.hardware.detect_hardware -> CHAT_ONLY). MLX is pulled only transitively via unsloth-zoo, and a resolver backtrack (mlx-vlm -> transformers>=5 vs the single-env transformers pin) can silently drop it, leaving @@ -13,7 +13,7 @@ a background thread, then re-detects so the gate re-opens without a manual The install mirrors the main Apple Silicon installer (install_python_stack.py): it points UV_OVERRIDE at overrides-darwin-arm64.txt so the resolver keeps the -Studio transformers pin AND installs a current mlx-vlm, and it requires the same +Unsloth transformers pin AND installs a current mlx-vlm, and it requires the same minimum versions unsloth-zoo declares so a backtracked old mlx-vlm (which still imports but breaks VLM Train/Export) is never accepted as healthy. @@ -69,11 +69,11 @@ _MLX_REINSTALL_ARGS = tuple( # reject anything. mlx/mlx-metal ship wheels only (no sdist on PyPI) and # mlx-lm/mlx-vlm publish py3-none-any wheels, so requiring wheels does not break a # healthy self-heal; if a wheel is genuinely unavailable the install fails and -# Studio stays chat-only (the existing safe fallback) until `unsloth studio update`. +# Unsloth stays chat-only (the existing safe fallback) until `unsloth studio update`. _ONLY_BINARY_ARG = "--only-binary=:all:" # Allowlist of environment variables forwarded to the install subprocess. The # self-heal runs without confirmation on the default startup path, so it must not -# hand resolver/build code the full Studio environment. Everything outside this +# hand resolver/build code the full Unsloth environment. Everything outside this # set is dropped, which excludes three dangerous classes by construction: # * secrets (HF_TOKEN, AWS_*, WANDB_API_KEY, ...) that a malicious wheel/sdist # build hook would otherwise read straight out of os.environ; @@ -207,13 +207,13 @@ def _mlx_install_env() -> dict[str, str]: The self-heal runs without confirmation on the default startup path, so it forwards only the variables uv genuinely needs (see _MLX_ENV_ALLOWLIST) instead - of the full Studio environment: secrets and package-source redirects in + of the full Unsloth environment: secrets and package-source redirects in os.environ are dropped so a malicious resolver-selected artifact cannot read - Studio secrets or be steered to a hostile index. + Unsloth secrets or be steered to a hostile index. Mirror the main installer (install_python_stack.py) by pointing UV_OVERRIDE at overrides-darwin-arm64.txt, which relaxes mlx-vlm/mlx-lm's transformers>=5 - requirement to >=4.57.6. Without it, uv keeps the Studio transformers pin only + requirement to >=4.57.6. Without it, uv keeps the Unsloth transformers pin only by silently backtracking mlx-vlm to an old, unsupported version (uv honours UV_OVERRIDE; plain pip ignores it, so the transformers constraint below is the pip-path safety net). We set UV_OVERRIDE ourselves, so a poisoned one in the @@ -234,17 +234,17 @@ def _mlx_install_env() -> dict[str, str]: def _transformers_constraint_args() -> tuple[list[str], str | None]: """Pin transformers to the running version for the mlx install. - The install must never upgrade transformers underneath a running Studio + The install must never upgrade transformers underneath a running Unsloth (the single-env install pins transformers==4.57.6). With UV_OVERRIDE set this is belt-and-suspenders; on the plain-pip path (no UV_OVERRIDE support) it is the actual guard -- the resolver either finds an mlx build compatible with the - pin or fails, leaving us chat-only rather than breaking Studio. Returns + pin or fails, leaving us chat-only rather than breaking Unsloth. Returns (pip args, temp file path to clean up). Read the version from installed metadata rather than `import transformers`: transformers can have valid metadata yet fail to import (e.g. an incompatible huggingface_hub), and in that case we still want to pin it so the mlx install - cannot quietly upgrade it out from under Studio.""" + cannot quietly upgrade it out from under Unsloth.""" from importlib.metadata import PackageNotFoundError, version as _dist_version try: @@ -263,10 +263,10 @@ def attempt_mlx_repair(*, timeout: int = _REPAIR_TIMEOUT_S) -> bool: """Install a usable mlx/mlx-lm/mlx-vlm stack by name into the running venv. Best-effort; returns True iff the resulting stack meets unsloth-zoo's minimums (so a backtracked old mlx-vlm is rejected, not accepted). transformers is held - at its pinned version so the install can never upgrade it underneath Studio.""" + at its pinned version so the install can never upgrade it underneath Unsloth.""" # Prepare the constraint inside the try: this runs on a daemon thread, so an # exception here (e.g. tempfile.mkstemp failing on a full disk or bad TMPDIR) - # must leave Studio chat-only, not crash the background self-heal thread. + # must leave Unsloth chat-only, not crash the background self-heal thread. constraint_path = None try: constraint_args, constraint_path = _transformers_constraint_args() @@ -279,7 +279,7 @@ def attempt_mlx_repair(*, timeout: int = _REPAIR_TIMEOUT_S) -> bool: ) if cmd is None: logger.warning( - "MLX self-heal requires uv so Studio can apply dependency overrides; " + "MLX self-heal requires uv so Unsloth can apply dependency overrides; " "staying chat-only. Run `unsloth studio update` to restore uv." ) return False diff --git a/studio/backend/utils/models/checkpoints.py b/studio/backend/utils/models/checkpoints.py index b6b080b1c4..f2125ad034 100644 --- a/studio/backend/utils/models/checkpoints.py +++ b/studio/backend/utils/models/checkpoints.py @@ -37,7 +37,7 @@ def _checkpoint_sort_key(checkpoint_path: Path) -> tuple[int, int, str]: def _infer_base_model_from_history(checkpoint_dir: Path) -> Optional[str]: - """Best-effort base-model lookup using persisted Studio run metadata.""" + """Best-effort base-model lookup using persisted Unsloth run metadata.""" checkpoint_name = checkpoint_dir.name resolved_checkpoint_dir = str(checkpoint_dir.resolve()) diff --git a/studio/backend/utils/models/gguf_metadata.py b/studio/backend/utils/models/gguf_metadata.py index c24ec28e1d..749f2c9234 100644 --- a/studio/backend/utils/models/gguf_metadata.py +++ b/studio/backend/utils/models/gguf_metadata.py @@ -50,9 +50,15 @@ _CACHE_MAX_ENTRIES = 4096 # keyed by (file cache key, wanted key). None = key absent / file unreadable. _BOOL_CACHE: Dict[Tuple[_CacheKey, str], Optional[bool]] = {} -# Native training context length (``{arch}.context_length``). None = absent / -# unreadable. Lets the UI show the real context ceiling before a model loads. -_CONTEXT_CACHE: Dict[_CacheKey, Optional[int]] = {} +_STRING_CACHE: Dict[Tuple[_CacheKey, str], Optional[str]] = {} + +# GGUF header dims for the staged/deferred-load UI: context_length, layer_count +# (block_count), and moe_layer_count (block_count minus leading dense layers; 0 +# if not MoE). One cached pass fills all three so the staged sheet can size every +# slider before the model loads. None = unreadable / not a GGUF. The native +# training context length (``{arch}.context_length``) the UI shows before a model +# loads is read from here via read_gguf_context_length. +_DIMS_CACHE: Dict[_CacheKey, Optional[Dict[str, Optional[int]]]] = {} def _cache_key(path: str) -> Optional[_CacheKey]: @@ -142,32 +148,45 @@ def _parse_gguf_header(path: str) -> Optional[Dict[str, str]]: return out -def read_gguf_context_length(path: str) -> Optional[int]: - """Return the GGUF's native training context length (``{arch}.context_length``), - or ``None`` if missing/unreadable/not a GGUF. Cached by (path, mtime, size). - Lets the UI populate the context slider before the model is loaded.""" +def read_gguf_staged_dims(path: str) -> Optional[Dict[str, Optional[int]]]: + """GGUF header dims for the staged-load UI in one cached pass: + ``{"context_length", "layer_count", "moe_layer_count"}``. Each may be None + when absent (moe_layer_count is 0 for a dense model). Returns ``None`` if not + a GGUF / unreadable. Cached by (path, mtime, size). Lets the staged sheet size + the context, GPU-layers and MoE sliders before the model loads.""" key = _cache_key(path) if key is None: return None with _CACHE_LOCK: - if key in _CONTEXT_CACHE: - return _CONTEXT_CACHE[key] - result = _parse_gguf_context_length(path) + if key in _DIMS_CACHE: + return _DIMS_CACHE[key] + result = _parse_gguf_staged_dims(path) with _CACHE_LOCK: - while len(_CONTEXT_CACHE) >= _CACHE_MAX_ENTRIES: + while len(_DIMS_CACHE) >= _CACHE_MAX_ENTRIES: try: - _CONTEXT_CACHE.pop(next(iter(_CONTEXT_CACHE))) + _DIMS_CACHE.pop(next(iter(_DIMS_CACHE))) except StopIteration: break - _CONTEXT_CACHE[key] = result + _DIMS_CACHE[key] = result return result -def _parse_gguf_context_length(path: str) -> Optional[int]: - # The context key is architecture-namespaced (``llama.context_length`` etc.), - # so we learn the key only after reading ``general.architecture``. GGUF writes - # general.* before arch.* keys, matching the loader's own parser. - ctx_key: Optional[str] = None +def read_gguf_context_length(path: str) -> Optional[int]: + """Native training context length (``{arch}.context_length``), or ``None``. + Thin accessor over read_gguf_staged_dims.""" + dims = read_gguf_staged_dims(path) + return dims["context_length"] if dims else None + + +def _parse_gguf_arch_uints(path: str, wanted_suffixes: frozenset[str]) -> Optional[Dict[str, int]]: + """Walk a GGUF header once and return the requested architecture-namespaced + uint (vtype 4/10) keys, e.g. ``{"block_count": 32}``. Keys are + ``{arch}.``; the arch is learned from ``general.architecture`` (GGUF + writes general.* before arch.* keys, matching the loader's own parser). + Returns ``None`` if not a GGUF / unreadable, else a dict (possibly empty or + partial when some keys are absent).""" + arch: Optional[str] = None + found: Dict[str, int] = {} try: with open(path, "rb") as f: head = f.read(24) @@ -204,28 +223,68 @@ def _parse_gguf_context_length(path: str) -> Optional[int]: sbytes = f.read(slen) if len(sbytes) < slen: break - ctx_key = f"{sbytes.decode('utf-8', 'replace')}.context_length" - elif ctx_key is not None and key == ctx_key and vtype in (4, 10): + arch = sbytes.decode("utf-8", "replace") + elif ( + arch is not None + and vtype in (4, 10) + and key.startswith(f"{arch}.") + and key[len(arch) + 1 :] in wanted_suffixes + ): width = 4 if vtype == 4 else 8 n_bytes = f.read(width) if len(n_bytes) < width: break - value = struct.unpack(" 0 else None + found[key[len(arch) + 1 :]] = struct.unpack( + " Optional[Dict[str, Optional[int]]]: + vals = _parse_gguf_arch_uints( + path, + frozenset( + { + "context_length", + "block_count", + "expert_count", + "leading_dense_block_count", + } + ), + ) + if vals is None: + return None + ctx = vals.get("context_length") + block = vals.get("block_count") + # A real context/layer count is positive; treat 0/garbage as absent so the + # UI never builds a slider with max < min. + context_length = ctx if ctx and ctx > 0 else None + layer_count = block if block and block > 0 else None + # MoE layer count = block_count - leading dense layers, only when experts + # exist; else 0 (dense -> slider hidden). Mirrors n_moe_layers in + # core/inference/llama_cpp.py. + if not vals.get("expert_count") or not block: + moe_layer_count: Optional[int] = 0 + else: + moe_layer_count = max(0, block - (vals.get("leading_dense_block_count") or 0)) + return { + "context_length": context_length, + "layer_count": layer_count, + "moe_layer_count": moe_layer_count, + } # Strings (8) and arrays (9) are handled inline. @@ -353,6 +412,83 @@ def _read_gguf_bool(path: str, wanted_key: str) -> Optional[bool]: return result +def _parse_gguf_string(path: str, wanted_key: str) -> Optional[str]: + try: + with open(path, "rb") as f: + head = f.read(24) + if len(head) < 24: + return None + magic, _version, _tcount, kv_count = struct.unpack(" 1 << 20: + break + kbytes = f.read(klen) + if len(kbytes) < klen: + break + key = kbytes.decode("utf-8", "replace") + vt_bytes = f.read(4) + if len(vt_bytes) < 4: + break + vtype = struct.unpack(" 1 << 22: + break + sbytes = f.read(slen) + if len(sbytes) < slen: + break + return sbytes.decode("utf-8", "replace") + if not _skip_gguf_value(f, vtype): + break + except (struct.error, UnicodeDecodeError): + break + except OSError as e: + logger.debug(f"_parse_gguf_string: cannot open {path}: {e}") + return None + except Exception as e: + logger.debug(f"_parse_gguf_string: parse failure on {path}: {e}") + return None + return None + + +def _read_gguf_string(path: str, wanted_key: str) -> Optional[str]: + fkey = _cache_key(path) + if fkey is None: + return None + ckey = (fkey, wanted_key) + with _CACHE_LOCK: + if ckey in _STRING_CACHE: + return _STRING_CACHE[ckey] + result = _parse_gguf_string(path, wanted_key) + with _CACHE_LOCK: + while len(_STRING_CACHE) >= _CACHE_MAX_ENTRIES: + try: + _STRING_CACHE.pop(next(iter(_STRING_CACHE))) + except StopIteration: + break + _STRING_CACHE[ckey] = result + return result + + +def read_gguf_chat_template(path: str) -> Optional[str]: + template = _read_gguf_string(path, "tokenizer.chat_template") + if isinstance(template, str) and template.strip(): + return template + return None + + def read_mmproj_audio_capability(path: str) -> Optional[bool]: """``clip.has_audio_encoder`` from an mmproj GGUF (e.g. Gemma 4's gemma4ua): ``True``/``False`` if present, ``None`` if absent/unreadable. diff --git a/studio/backend/utils/models/model_config.py b/studio/backend/utils/models/model_config.py index 284bbb5745..821529083d 100644 --- a/studio/backend/utils/models/model_config.py +++ b/studio/backend/utils/models/model_config.py @@ -1249,6 +1249,77 @@ def _iter_gguf_files(directory: Path, recursive: bool = False): yield f +_GGUF_SPLIT_FILE_RE = re.compile( + r"^(?P.+)-(?P\d{5})-of-(?P\d{5})\.gguf$", + re.IGNORECASE, +) + + +def _colocated_first_split_shard(path: Path) -> tuple[Optional[Path], bool]: + """Return shard 1 and whether every shard is beside *path*.""" + match = _GGUF_SPLIT_FILE_RE.match(path.name) + if match is None: + return None, False + + prefix = match.group("prefix").casefold() + total_text = match.group("total") + total = int(total_text) + if total < 1: + return None, False + + first: Optional[Path] = None + indices: set[int] = set() + try: + siblings = path.parent.iterdir() + for sibling in siblings: + sibling_match = _GGUF_SPLIT_FILE_RE.match(sibling.name) + if ( + sibling_match is None + or sibling_match.group("prefix").casefold() != prefix + or sibling_match.group("total") != total_text + ): + continue + try: + if not sibling.is_file(): + continue + except OSError: + continue + index = int(sibling_match.group("index")) + if not 1 <= index <= total: + continue + indices.add(index) + if index == 1: + first = sibling + except OSError: + return None, False + + return first, first is not None and len(indices) == total + + +def _local_gguf_load_path(path: Path) -> Path: + """Choose a loadable local path while preserving complete symlink sets.""" + if _GGUF_SPLIT_FILE_RE.match(path.name) is None: + return path.absolute() + + first, complete = _colocated_first_split_shard(path) + if complete and first is not None: + return first.absolute() + + try: + is_symlink = path.is_symlink() + except OSError: + is_symlink = False + if is_symlink: + try: + target = path.resolve() + except OSError: + return (first or path).absolute() + target_first, _ = _colocated_first_split_shard(target) + return (target_first or target).absolute() + + return (first or path).absolute() + + def detect_mmproj_file(path: str, search_root: Optional[str] = None) -> Optional[str]: """Find the mmproj GGUF for a model. @@ -1434,7 +1505,7 @@ def detect_gguf_model(path: str) -> Optional[str]: except OSError: is_dir = False # stat() unavailable in the lock window if not is_dir: - return str(p.absolute()) # absolute() keeps symlink names readable + return str(_local_gguf_load_path(p)) # Directory named "*.gguf": fall through to the dir scan below. # Case 2: directory containing .gguf files (skip mmproj / MTP drafter) @@ -1452,7 +1523,7 @@ def detect_gguf_model(path: str) -> Optional[str]: gguf_files.append(f) gguf_files.sort(key = lambda f: f.stat().st_size, reverse = True) if gguf_files: - return str(gguf_files[0].resolve()) + return str(_local_gguf_load_path(gguf_files[0])) return None @@ -1879,7 +1950,7 @@ def _find_local_gguf_by_variant(directory: str, variant: str) -> Optional[str]: For sharded GGUFs (multiple files sharing a quant label), returns the first shard (sorted by name), which is what ``llama-server -m`` expects. - Returns the resolved absolute path, or ``None`` if no match. + Returns the absolute path, or ``None`` if no match. """ p = _resolve_gguf_dir(Path(directory)) if p is None: @@ -1900,7 +1971,7 @@ def _find_local_gguf_by_variant(directory: str, variant: str) -> Optional[str]: matches.append(f) matches.sort() if matches: - return str(matches[0].resolve()) + return str(_local_gguf_load_path(matches[0])) return None @@ -2083,7 +2154,7 @@ def _has_model_weight_files(model_dir: Path) -> bool: def _detect_training_output_type(model_dir: Path) -> Optional[str]: - """Classify a Studio training output as LoRA or full finetune.""" + """Classify an Unsloth training output as LoRA or full finetune.""" adapter_config = model_dir / "adapter_config.json" adapter_model = model_dir / "adapter_model.safetensors" if adapter_config.exists() or adapter_model.exists(): @@ -2105,7 +2176,7 @@ def _looks_like_lora_adapter(model_dir: Path) -> bool: def scan_trained_models(outputs_dir: str = str(outputs_root())) -> List[Tuple[str, str, str]]: - """Scan outputs folder for trained Studio models. + """Scan outputs folder for trained Unsloth models. Returns: List of (display_name, model_path, model_type), where model_type is diff --git a/studio/backend/utils/openai_auto_switch_settings.py b/studio/backend/utils/openai_auto_switch_settings.py index 1689395f40..7007440f4c 100644 --- a/studio/backend/utils/openai_auto_switch_settings.py +++ b/studio/backend/utils/openai_auto_switch_settings.py @@ -8,7 +8,9 @@ Two settings, both off by default so existing API behavior is unchanged: names a downloaded local GGUF different from the loaded one transparently loads it before serving (llama-swap-style). Unknown names pass through. - ``openai_api_auto_unload_idle_seconds``: when > 0, the loaded GGUF is - unloaded after this many idle seconds to free VRAM. + unloaded after this many idle seconds to free VRAM. Enabled values have a + 60s floor (0 stays "off"): a tiny TTL tears the model down between turns of + an active chat, forcing a full weight reload + prompt re-prefill per turn. The idle TTL can also be set at startup via the ``UNSLOTH_MODEL_IDLE_TTL`` env var. Unlike the stored setting (which stays gated on auto-switch), the env value @@ -28,11 +30,14 @@ from typing import Any, Optional OPENAI_AUTO_SWITCH_SETTING_KEY = "openai_api_auto_switch_model" AUTO_UNLOAD_IDLE_SETTING_KEY = "openai_api_auto_unload_idle_seconds" +AUTO_UNLOAD_KEEP_KV_SETTING_KEY = "openai_api_auto_unload_keep_kv" MODEL_OVERRIDES_SETTING_KEY = "openai_api_auto_switch_overrides" MODEL_IDLE_TTL_ENV_VAR = "UNSLOTH_MODEL_IDLE_TTL" DEFAULT_OPENAI_AUTO_SWITCH_ENABLED = False DEFAULT_AUTO_UNLOAD_IDLE_SECONDS = 0 +DEFAULT_AUTO_UNLOAD_KEEP_KV = True +MIN_AUTO_UNLOAD_IDLE_SECONDS = 60 _CACHE_TTL_S = 2.0 _cache_lock = threading.Lock() @@ -58,6 +63,10 @@ def _coerce_int(value: Any) -> int | None: return None +def _apply_idle_floor(seconds: int) -> int: + return 0 if seconds <= 0 else max(MIN_AUTO_UNLOAD_IDLE_SECONDS, seconds) + + def _cached_setting(key: str, default: Any) -> Any: """Read an app setting, memoized for _CACHE_TTL_S to spare the hot path.""" now = time.monotonic() @@ -91,12 +100,34 @@ def _stored_idle_seconds() -> Optional[int]: return _coerce_int(_cached_setting(AUTO_UNLOAD_IDLE_SETTING_KEY, None)) +_env_floor_warned = False + + def _env_idle_seconds() -> Optional[int]: - """UNSLOTH_MODEL_IDLE_TTL as a non-negative seconds value, or None if unset/invalid.""" + """UNSLOTH_MODEL_IDLE_TTL as a non-negative seconds value, or None if unset/invalid. + + Floored to MIN_AUTO_UNLOAD_IDLE_SECONDS here (with a one-time warning) since + headless/container deploys have no UI to surface a validation error.""" raw = os.environ.get(MODEL_IDLE_TTL_ENV_VAR) if raw is None or not raw.strip(): return None - return _coerce_int(raw) + parsed = _coerce_int(raw) + if parsed is None: + return None + floored = _apply_idle_floor(parsed) + if floored != parsed: + global _env_floor_warned + if not _env_floor_warned: + _env_floor_warned = True + from loggers import get_logger + get_logger(__name__).warning( + "%s=%s is below the %ss minimum; using %ss", + MODEL_IDLE_TTL_ENV_VAR, + parsed, + MIN_AUTO_UNLOAD_IDLE_SECONDS, + floored, + ) + return floored def get_stored_auto_unload_idle_seconds() -> int: @@ -108,7 +139,9 @@ def get_stored_auto_unload_idle_seconds() -> int: """ stored = _stored_idle_seconds() if stored is not None: - return stored + # Floor legacy values persisted before the minimum existed, so the UI + # displays the effective TTL and round-trips it cleanly. + return _apply_idle_floor(stored) env = _env_idle_seconds() return env if env is not None else DEFAULT_AUTO_UNLOAD_IDLE_SECONDS @@ -118,32 +151,63 @@ def get_auto_unload_idle_seconds() -> int: stored = _stored_idle_seconds() if stored is not None: # An explicit UI/API value stays gated on auto-switch: off reports 0 so the - # off state is identical to pre-feature. - return stored if get_openai_auto_switch_enabled() else 0 + # off state is identical to pre-feature. Floored to cover values persisted + # before the minimum existed. + return _apply_idle_floor(stored) if get_openai_auto_switch_enabled() else 0 # No stored value: UNSLOTH_MODEL_IDLE_TTL is a standalone startup default that # enables idle-unload even with auto-switch off (headless/container deploys). env = _env_idle_seconds() return env if env is not None else 0 -def set_openai_auto_switch(enabled: Any, idle_seconds: Any) -> tuple[bool, int]: - """Set both auto-switch flags in one transaction so a settings PUT can't leave - one key updated and the other stale. Both values are coerced before any write, - so an invalid value raises without persisting either.""" +def get_auto_unload_keep_kv() -> bool: + """Whether the idle unload persists slot KV to disk for restore on reload.""" + parsed = _coerce_bool(_cached_setting(AUTO_UNLOAD_KEEP_KV_SETTING_KEY, None)) + return parsed if parsed is not None else DEFAULT_AUTO_UNLOAD_KEEP_KV + + +def set_openai_auto_switch( + enabled: Any, + idle_seconds: Any, + keep_kv: Any = None, +) -> tuple[bool, int, bool]: + """One-transaction write; ``None`` leaves a stored value untouched.""" parsed_enabled = _coerce_bool(enabled) if parsed_enabled is None: raise ValueError("OpenAI auto-switch must be true or false.") - parsed_idle = _coerce_int(idle_seconds) - if parsed_idle is None: - raise ValueError("Auto-unload idle seconds must be a non-negative integer.") + parsed_idle = None + if idle_seconds is not None: + parsed_idle = _coerce_int(idle_seconds) + if parsed_idle is None: + raise ValueError("Auto-unload idle seconds must be a non-negative integer.") + if 0 < parsed_idle < MIN_AUTO_UNLOAD_IDLE_SECONDS: + raise ValueError( + f"Auto-unload idle seconds must be 0 (off) or at least " + f"{MIN_AUTO_UNLOAD_IDLE_SECONDS}." + ) + parsed_keep_kv = None + if keep_kv is not None: + parsed_keep_kv = _coerce_bool(keep_kv) + if parsed_keep_kv is None: + raise ValueError("Keep KV on idle unload must be true or false.") from storage.studio_db import upsert_app_settings - upsert_app_settings( - {OPENAI_AUTO_SWITCH_SETTING_KEY: parsed_enabled, AUTO_UNLOAD_IDLE_SETTING_KEY: parsed_idle} - ) + updates: dict[str, Any] = {OPENAI_AUTO_SWITCH_SETTING_KEY: parsed_enabled} + if parsed_idle is not None: + updates[AUTO_UNLOAD_IDLE_SETTING_KEY] = parsed_idle + if parsed_keep_kv is not None: + updates[AUTO_UNLOAD_KEEP_KV_SETTING_KEY] = parsed_keep_kv + upsert_app_settings(updates) _invalidate(OPENAI_AUTO_SWITCH_SETTING_KEY) - _invalidate(AUTO_UNLOAD_IDLE_SETTING_KEY) - return parsed_enabled, parsed_idle + if parsed_idle is not None: + _invalidate(AUTO_UNLOAD_IDLE_SETTING_KEY) + if parsed_keep_kv is not None: + _invalidate(AUTO_UNLOAD_KEEP_KV_SETTING_KEY) + return ( + parsed_enabled, + parsed_idle if parsed_idle is not None else get_stored_auto_unload_idle_seconds(), + parsed_keep_kv if parsed_keep_kv is not None else get_auto_unload_keep_kv(), + ) def get_model_overrides() -> dict[str, dict]: diff --git a/studio/backend/utils/paths/storage_roots.py b/studio/backend/utils/paths/storage_roots.py index 759681da3f..35b8c57e9b 100644 --- a/studio/backend/utils/paths/storage_roots.py +++ b/studio/backend/utils/paths/storage_roots.py @@ -36,7 +36,7 @@ def _infer_studio_home_from_venv() -> Path | None: def studio_root() -> Path: - """Studio install root. + """Unsloth install root. Priority: UNSLOTH_STUDIO_HOME, then STUDIO_HOME alias, then sys.prefix inference, then legacy ~/.unsloth/studio. UNSLOTH_STUDIO_HOME wins if @@ -61,8 +61,13 @@ def cache_root() -> Path: return studio_root() / "cache" +def llama_slot_cache_root() -> Path: + """Dir llama-server saves/restores slot KV state in across idle unloads.""" + return cache_root() / "llama-slots" + + def studio_bin_root() -> Path: - """Dir for Studio-managed executables (the `unsloth` shim, downloaded tools like cloudflared).""" + """Dir for Unsloth-managed executables (the `unsloth` shim, downloaded tools like cloudflared).""" return studio_root() / "bin" @@ -443,7 +448,7 @@ def resolve_export_write_dir(path_value: str | None = None) -> Path: Unlike :func:`resolve_export_dir`, this function passes absolute paths through as-is so users can target a different drive when - their Studio install lives on a constrained system volume + their Unsloth install lives on a constrained system volume (see :gh-issue:`6082`). Used only by the export write path. """ if not path_value or not str(path_value).strip(): diff --git a/studio/backend/utils/preview_rate_limit.py b/studio/backend/utils/preview_rate_limit.py index dd38cfd5e7..c59a1bf5b3 100644 --- a/studio/backend/utils/preview_rate_limit.py +++ b/studio/backend/utils/preview_rate_limit.py @@ -5,7 +5,7 @@ A signed link stops ref guessing, but anyone with a link can still drive GPU generation. This bounds sustained abuse from a single source. In-process and -single-worker only (like the login limiter in ``routes/auth.py``); Studio runs as +single-worker only (like the login limiter in ``routes/auth.py``); Unsloth runs as one uvicorn process, so a shared store isn't needed. """ diff --git a/studio/backend/utils/process_lifetime.py b/studio/backend/utils/process_lifetime.py index 3ffd54cc26..c63227ae86 100644 --- a/studio/backend/utils/process_lifetime.py +++ b/studio/backend/utils/process_lifetime.py @@ -1,7 +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 -"""Bind Studio child processes to the parent's lifetime so none survive an +"""Bind Unsloth child processes to the parent's lifetime so none survive an abnormal parent exit (terminal-window close, Task Manager "End Task", SIGKILL, crash) -- the cooperative shutdown path only runs on graceful exits. @@ -139,7 +139,7 @@ def _install_windows_job() -> None: kernel32.CloseHandle(job) return # AssignProcessToJobObject(parent) makes children inherit the job. May - # fail if Studio already runs inside an incompatible host job (pre-Win8); + # fail if Unsloth already runs inside an incompatible host job (pre-Win8); # degrade to the cooperative path rather than blocking startup. if not kernel32.AssignProcessToJobObject(job, kernel32.GetCurrentProcess()): kernel32.CloseHandle(job) diff --git a/studio/backend/utils/studio_version.py b/studio/backend/utils/studio_version.py index 9c18070fbb..82ade74bba 100644 --- a/studio/backend/utils/studio_version.py +++ b/studio/backend/utils/studio_version.py @@ -1,7 +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 -"""Network-free Studio release version resolution for display-only UI.""" +"""Network-free Unsloth release version resolution for display-only UI.""" from __future__ import annotations @@ -20,7 +20,7 @@ _MAX_VERSION_LENGTH = 64 def is_valid_studio_release_version(value: object) -> bool: - """Return True for Studio release tags such as ``v0.1.39-beta``.""" + """Return True for Unsloth release tags such as ``v0.1.39-beta``.""" if not isinstance(value, str): return False version = value.strip() @@ -102,7 +102,7 @@ def _git_branch(repo_root: Path) -> str | None: def get_studio_version(repo_root: Path | None = None) -> str: - """Return the installed Studio release tag for display, or ``dev``. + """Return the installed Unsloth release tag for display, or ``dev``. Intentionally separate from the PyPI ``unsloth`` package version used by update checks. Never performs network requests. diff --git a/studio/backend/utils/training_runs.py b/studio/backend/utils/training_runs.py index dc2535e570..dcdfa1395d 100644 --- a/studio/backend/utils/training_runs.py +++ b/studio/backend/utils/training_runs.py @@ -1,7 +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 -"""Helpers for naming and describing Studio training runs.""" +"""Helpers for naming and describing Unsloth training runs.""" from __future__ import annotations diff --git a/studio/backend/utils/transformers_latest.py b/studio/backend/utils/transformers_latest.py index 40c8f729a5..9f1d11be5b 100644 --- a/studio/backend/utils/transformers_latest.py +++ b/studio/backend/utils/transformers_latest.py @@ -5,7 +5,7 @@ When a model's ``model_type`` is absent from every installed transformers overlay (base 4.57.x plus the .venv_t5_530/550/510 sidecars and, if provisioned, .venv_t5_latest), -Studio cannot load it today. This module answers, without authentication, code execution, +Unsloth cannot load it today. This module answers, without authentication, code execution, or trust_remote_code: 1. Does the LATEST transformers release on PyPI ship this ``model_type``? @@ -387,7 +387,7 @@ def check_upgrade_for_model(model_name: str, hf_token: str | None = None) -> dic _SHADOWABLE_DEPS = frozenset({"tokenizers", "safetensors"}) # Provided by the sidecar recipe; checked against its pin, not the base env. _SIDECAR_PROVIDED = {"huggingface-hub": "1.8.0", "hf-xet": "1.4.2"} -# CLI-only; never imported at runtime in Studio's workers. +# CLI-only; never imported at runtime in Unsloth's workers. _IGNORED_DEPS = frozenset({"typer"}) @@ -538,7 +538,7 @@ def _install_latest_transformers_locked(version: str, before_swap = None) -> dic return { "success": False, "version": version, - "message": "Cannot install: Studio is in offline mode.", + "message": "Cannot install: Unsloth is in offline mode.", } # Re-verify against a LIVE snapshot (a release may land inside the cache TTL); # fall back to the cached one on fetch failure. @@ -573,13 +573,13 @@ def _install_latest_transformers_locked(version: str, before_swap = None) -> dic "version": version, "message": "Cannot install transformers " f"{version}: this environment does not satisfy {', '.join(blockers)}. " - "A Studio update is required first.", + "An Unsloth update is required first.", } if not ensure_latest_transformers_venv(version, extra_packages, before_swap = before_swap): return { "success": False, "version": version, - "message": f"Installing transformers {version} failed; see the Studio logs.", + "message": f"Installing transformers {version} failed; see the Unsloth logs.", } _invalidate_capability_caches() return { diff --git a/studio/backend/utils/transformers_version.py b/studio/backend/utils/transformers_version.py index 9f9f8aa3de..1fbcc9f46f 100644 --- a/studio/backend/utils/transformers_version.py +++ b/studio/backend/utils/transformers_version.py @@ -2151,7 +2151,7 @@ def end_sidecar_swap() -> None: def sidecar_swap_in_progress() -> bool: """True while a .venv_t5_latest install or repair holds the reservation, - in this process or any other Studio process (lock file).""" + in this process or any other Unsloth process (lock file).""" return sidecar_swap_kind() is not None diff --git a/studio/backend/utils/upload_limits.py b/studio/backend/utils/upload_limits.py index c21ea69af7..fff0ac423f 100644 --- a/studio/backend/utils/upload_limits.py +++ b/studio/backend/utils/upload_limits.py @@ -1,7 +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 -"""Shared Studio upload/request size limits.""" +"""Shared Unsloth upload/request size limits.""" from __future__ import annotations diff --git a/studio/backend/utils/utils.py b/studio/backend/utils/utils.py index 3818253ac9..31f5f31bee 100644 --- a/studio/backend/utils/utils.py +++ b/studio/backend/utils/utils.py @@ -123,6 +123,26 @@ def without_hf_auth(): os.environ.pop("HF_HUB_DISABLE_IMPLICIT_TOKEN", None) +def is_hf_authentication_error(error: Exception) -> bool: + """Return whether an exception chain contains a definitive HF auth failure.""" + seen: set[int] = set() + current: BaseException | None = error + while current is not None and id(current) not in seen: + seen.add(id(current)) + response = getattr(current, "response", None) + status = getattr(response, "status_code", None) + try: + if status is not None and int(status) == 401: + return True + except (TypeError, ValueError): + pass + message = str(current).lower() + if "invalid user token" in message or "invalid hf token" in message: + return True + current = current.__cause__ or current.__context__ + return False + + def format_error_message(error: Exception, model_name: str) -> str: """ Format a user-friendly error message for common load issues. diff --git a/studio/frontend/.npmrc b/studio/frontend/.npmrc index 19783b5ff4..414379da6e 100644 --- a/studio/frontend/.npmrc +++ b/studio/frontend/.npmrc @@ -1,4 +1,4 @@ -# Studio frontend npm configuration. +# Unsloth frontend npm configuration. # # Mini Shai-Hulud / Axios-style supply chain defense. # Requires npm >=11.10.0. Refuses tarballs published less than 7 days ago, diff --git a/studio/frontend/src/app/auth-guards.ts b/studio/frontend/src/app/auth-guards.ts index 6849f380b8..a3523ac580 100644 --- a/studio/frontend/src/app/auth-guards.ts +++ b/studio/frontend/src/app/auth-guards.ts @@ -23,19 +23,47 @@ interface AuthStatus { requires_password_change: boolean; } +const AUTH_STATUS_TTL_MS = 30_000; +let authStatusCheckedAt = 0; +let authStatusRequest: Promise | null = null; + +function hasFreshAuthStatus(): boolean { + return ( + authStatusCheckedAt !== 0 && + Date.now() - authStatusCheckedAt < AUTH_STATUS_TTL_MS + ); +} + async function fetchAuthStatus(): Promise { - try { - const res = await fetch(apiUrl("/api/auth/status")); - if (!res.ok) return { initialized: true, requires_password_change: mustChangePassword() }; - const status = (await res.json()) as AuthStatus; - // Server truth wins; keep localStorage in sync both ways. - if (status.requires_password_change !== mustChangePassword()) { - setMustChangePassword(status.requires_password_change); + if (authStatusRequest) return authStatusRequest; + + const request = (async () => { + try { + const res = await fetch(apiUrl("/api/auth/status")); + if (!res.ok) { + return { + initialized: true, + requires_password_change: mustChangePassword(), + }; + } + const status = (await res.json()) as AuthStatus; + authStatusCheckedAt = Date.now(); + // Server truth wins; keep localStorage in sync both ways. + if (status.requires_password_change !== mustChangePassword()) { + setMustChangePassword(status.requires_password_change); + } + return status; + } catch { + return { + initialized: true, + requires_password_change: mustChangePassword(), + }; } - return status; - } catch { - return { initialized: true, requires_password_change: mustChangePassword() }; - } + })().finally(() => { + authStatusRequest = null; + }); + authStatusRequest = request; + return request; } function authRedirect(to: "/login" | "/change-password"): never { @@ -49,12 +77,17 @@ export async function requireAuth(): Promise { } if (await hasActiveSession()) { - const { requires_password_change } = await fetchAuthStatus(); - if (requires_password_change || mustChangePassword()) { - authRedirect("/change-password"); + // Reconcile periodically so local-only routes cannot outlive a server-side + // password-change requirement, while nearby route switches stay local. + if (mustChangePassword() || !hasFreshAuthStatus()) { + const { requires_password_change } = await fetchAuthStatus(); + if (requires_password_change || mustChangePassword()) { + authRedirect("/change-password"); + } } return; } + const status = await fetchAuthStatus(); if (status.requires_password_change || mustChangePassword()) { authRedirect("/change-password"); diff --git a/studio/frontend/src/app/provider.tsx b/studio/frontend/src/app/provider.tsx index c35706e50a..e6c89b9cd7 100644 --- a/studio/frontend/src/app/provider.tsx +++ b/studio/frontend/src/app/provider.tsx @@ -51,9 +51,12 @@ async function showSetupWindow(isCurrent: WindowLayoutGuard): Promise { const { getCurrentWindow, LogicalSize } = await import( "@tauri-apps/api/window" ); + const { invoke } = await import("@tauri-apps/api/core"); if (!isCurrent()) return; const win = getCurrentWindow(); + await invoke("reset_app_window_layout_initialized"); + if (!isCurrent()) return; await win.setResizable(false); if (!isCurrent()) return; await win.setSize(new LogicalSize(SETUP_WINDOW_WIDTH, SETUP_WINDOW_HEIGHT)); @@ -98,18 +101,20 @@ async function applyAppWindowLayout( if (!isCurrent()) return; const win = getCurrentWindow(); - // Decide first-launch vs restore from the on-disk state file BEFORE touching the - // window. Probing the window after restoreStateCurrent is unreliable: on GTK, - // set_size on a hidden window is deferred until show(), so innerSize() reads a - // stale value and a baseline fallback would overwrite the queued restore. On - // macOS the same probe works, hence the inconsistency between prior iterations. - const hasSavedState = await invoke("has_saved_window_state"); + // Setup-window activity may create plugin state before the full app is ever + // shown, so use a dedicated full-app marker to decide whether restoration is + // appropriate. Keep checking plugin state so a missing/corrupt state file + // falls back to a monitor-safe centered layout. + const [hasInitializedAppLayout, hasSavedState] = await Promise.all([ + invoke("has_initialized_app_window_layout"), + invoke("has_saved_window_state"), + ]); if (!isCurrent()) return; await win.setResizable(true); if (!isCurrent()) return; - if (hasSavedState) { + if (hasInitializedAppLayout && hasSavedState) { // Subsequent launch: plugin restores size/position/maximized, with built-in // off-screen protection for positions saved on a now-disconnected display. await restoreStateCurrent( @@ -144,6 +149,9 @@ async function applyAppWindowLayout( }); if (!isCurrent()) return; await enforceMinimumWindowSize(win, LogicalSize, isCurrent); + + if (!isCurrent()) return; + await invoke("mark_app_window_layout_initialized"); } async function showWindowFallback(): Promise { @@ -412,7 +420,7 @@ function TauriWrapper({ children }: { children: ReactNode }) { {desktopBooting ? (

-
Preparing Studio
+
Preparing Unsloth
The local backend is ready. Signing in to your desktop session before loading chats. diff --git a/studio/frontend/src/app/routes/__root.tsx b/studio/frontend/src/app/routes/__root.tsx index ba56ce7525..7137fd6f96 100644 --- a/studio/frontend/src/app/routes/__root.tsx +++ b/studio/frontend/src/app/routes/__root.tsx @@ -16,6 +16,7 @@ import { type ChatSearch, } from "@/features/chat"; import { RemoteCodeConsentDialog } from "@/features/security"; +import { HfTokenWarningDialog } from "@/features/hf-auth"; import { TransformersUpgradeDialog } from "@/features/transformers-upgrade"; import { useTrainingUnloadGuard } from "@/features/training"; import { useExportRuntimeLifecycle } from "@/features/export"; @@ -195,9 +196,6 @@ function RootLayout() { chatRuntime.setActiveThreadId(null); chatRuntime.setActiveProjectId(null); chatRuntime.setIncognito(false); - // Detach the staging UI but keep any in-flight download running, like Hub. - if (chatRuntime.pendingSelection) - chatRuntime.abandonStagedModel({ keepDownload: true }); void navigate({ to: "/chat", search: { new: crypto.randomUUID() }, @@ -220,16 +218,13 @@ function RootLayout() { chatRuntime.setActiveProjectId(null); chatRuntime.setActiveThreadId(null); chatRuntime.setIncognito(false); - // Leaving chat must not kill an in-flight download: detach the staging UI - // but keep the transfer running in the manager, like a Hub download. - if (chatRuntime.pendingSelection) - chatRuntime.abandonStagedModel({ keepDownload: true }); }, [isChatRoute]); return ( {!isAuthFlowRoute && } + {hideNavbar ? ( @@ -279,7 +274,7 @@ function RootLayout() { initial={{ opacity: 0 }} animate={{ opacity: 1 }} exit={{ opacity: 0 }} - transition={{ duration: 0.15 }} + transition={{ duration: 0.06 }} className="flex min-h-0 min-w-0 flex-1 basis-0 flex-col overflow-visible" > }> diff --git a/studio/frontend/src/app/routes/data-recipes.tsx b/studio/frontend/src/app/routes/data-recipes.tsx index c35e63da5f..22f87821af 100644 --- a/studio/frontend/src/app/routes/data-recipes.tsx +++ b/studio/frontend/src/app/routes/data-recipes.tsx @@ -1,15 +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 -import { createRoute } from "@tanstack/react-router"; -import { lazy } from "react"; +import { createRoute, lazyRouteComponent } from "@tanstack/react-router"; import { requireAuth } from "../auth-guards"; import { Route as rootRoute } from "./__root"; -const DataRecipesPage = lazy(() => - import("@/features/data-recipes").then((m) => ({ - default: m.DataRecipesPage, - })), +const DataRecipesPage = lazyRouteComponent( + () => import("@/features/data-recipes"), + "DataRecipesPage", ); export const Route = createRoute({ diff --git a/studio/frontend/src/app/routes/export.tsx b/studio/frontend/src/app/routes/export.tsx index 40118c6a92..5a7b586f19 100644 --- a/studio/frontend/src/app/routes/export.tsx +++ b/studio/frontend/src/app/routes/export.tsx @@ -1,15 +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 -import { createRoute } from "@tanstack/react-router"; -import { lazy } from "react"; +import { createRoute, lazyRouteComponent } from "@tanstack/react-router"; import { requireAuth } from "../auth-guards"; import { Route as rootRoute } from "./__root"; -const ExportPage = lazy(() => - import("@/features/export/export-page").then((m) => ({ - default: m.ExportPage, - })), +const ExportPage = lazyRouteComponent( + () => import("@/features/export/export-page"), + "ExportPage", ); export type ExportSearch = { diff --git a/studio/frontend/src/app/routes/hub.tsx b/studio/frontend/src/app/routes/hub.tsx index c623ef9848..2207490e44 100644 --- a/studio/frontend/src/app/routes/hub.tsx +++ b/studio/frontend/src/app/routes/hub.tsx @@ -1,15 +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 -import { createRoute } from "@tanstack/react-router"; -import { lazy } from "react"; +import { createRoute, lazyRouteComponent } from "@tanstack/react-router"; import { requireAuth } from "../auth-guards"; import { Route as rootRoute } from "./__root"; -const ModelsPage = lazy(() => - import("@/features/hub/hub-page").then((m) => ({ - default: m.ModelsPage, - })), +const ModelsPage = lazyRouteComponent( + () => import("@/features/hub/hub-page"), + "ModelsPage", ); export interface ModelsSearch { @@ -31,7 +29,11 @@ export const Route = createRoute({ const model = search.model; if (typeof model === "string" && model.length > 0) next.model = model; const section = search.section; - if (section === "trending" || section === "latest" || section === "finetune") { + if ( + section === "trending" || + section === "latest" || + section === "finetune" + ) { next.section = section; } const kind = search.kind; diff --git a/studio/frontend/src/app/routes/projects.tsx b/studio/frontend/src/app/routes/projects.tsx index c63b1d5838..17f58ef631 100644 --- a/studio/frontend/src/app/routes/projects.tsx +++ b/studio/frontend/src/app/routes/projects.tsx @@ -1,15 +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 -import { createRoute } from "@tanstack/react-router"; -import { lazy } from "react"; +import { createRoute, lazyRouteComponent } from "@tanstack/react-router"; import { requireAuth } from "../auth-guards"; import { Route as rootRoute } from "./__root"; -const ProjectsPage = lazy(() => - import("@/features/chat/projects-page").then((m) => ({ - default: m.ProjectsPage, - })), +const ProjectsPage = lazyRouteComponent( + () => import("@/features/chat/projects-page"), + "ProjectsPage", ); export const Route = createRoute({ diff --git a/studio/frontend/src/app/routes/studio.tsx b/studio/frontend/src/app/routes/studio.tsx index ae7f445e94..798044bf64 100644 --- a/studio/frontend/src/app/routes/studio.tsx +++ b/studio/frontend/src/app/routes/studio.tsx @@ -1,15 +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 -import { createRoute } from "@tanstack/react-router"; -import { lazy } from "react"; +import { createRoute, lazyRouteComponent } from "@tanstack/react-router"; import { requireAuth } from "../auth-guards"; import { Route as rootRoute } from "./__root"; -const StudioPage = lazy(() => - import("@/features/studio/studio-page").then((m) => ({ - default: m.StudioPage, - })), +const StudioPage = lazyRouteComponent( + () => import("@/features/studio/studio-page"), + "StudioPage", ); export const Route = createRoute({ diff --git a/studio/frontend/src/components/app-sidebar.tsx b/studio/frontend/src/components/app-sidebar.tsx index b6e218793a..8dd7bcdd9a 100644 --- a/studio/frontend/src/components/app-sidebar.tsx +++ b/studio/frontend/src/components/app-sidebar.tsx @@ -93,7 +93,12 @@ import { import { Tooltip as TooltipPrimitive } from "radix-ui"; import { HugeiconsIcon } from "@hugeicons/react"; import { ChevronDown, Moon } from "lucide-react"; -import { Link, useNavigate, useRouterState } from "@tanstack/react-router"; +import { + Link, + useNavigate, + useRouter, + useRouterState, +} from "@tanstack/react-router"; import { archiveChatItem, ChatSearchDialog, @@ -136,6 +141,7 @@ import { import type { TrainingRunSummary } from "@/features/training"; import { useExportRuntimeStore } from "@/features/export"; import { useEffect, useMemo, useRef, useState, type ReactNode } from "react"; +import { isDownloadCancelled } from "@/lib/native-files"; import { toast } from "@/lib/toast"; import { ShutdownDialog } from "@/components/shutdown-dialog"; import { translate, useT, type TranslationKey } from "@/i18n"; @@ -256,6 +262,10 @@ function createNavigationNonce(): string { return `${Date.now()}-${Math.random().toString(36).slice(2, 10)}`; } +function preloadSilently(request: Promise): void { + void request.catch(() => undefined); +} + function NavItem({ icon, label, @@ -267,6 +277,7 @@ function NavItem({ className, spinner, tooltip, + onIntent, }: { icon: typeof ZapIcon; label: string; @@ -277,6 +288,7 @@ function NavItem({ dataTour?: string; className?: string; spinner?: boolean; + onIntent?: () => void; // Overrides the hover tooltip (defaults to `label`). Used to explain why a // disabled item (e.g. Train/Export on a chat-only host) is greyed out. tooltip?: string; @@ -288,6 +300,8 @@ function NavItem({ tooltip={tooltip ?? label} disabled={disabled} onClick={onClick} + onPointerEnter={disabled ? undefined : onIntent} + onFocus={disabled ? undefined : onIntent} isActive={active} data-tour={dataTour} className="sidebar-nav-btn h-[33px] rounded-full gap-[8.5px] pl-3 pr-2.5 font-medium group-data-[collapsible=icon]:px-2.5 group-data-[collapsible=icon]:!w-[32px] group-data-[collapsible=icon]:mx-auto" @@ -324,6 +338,7 @@ export function AppSidebar() { }); const { togglePinned, isMobile, setOpenMobile } = useSidebar(); const navigate = useNavigate(); + const router = useRouter(); // Web update detection: `webUpdate` is non-null only when the installed // (PyPI) version is behind the latest release, so the card is hidden by @@ -957,11 +972,13 @@ export function AppSidebar() { const ids = item.type === "single" ? [item.id] : (await listStoredChatThreads({ pairId: item.id })).map((t) => t.id); - await Promise.all( - ids.map((id) => exportConversationByFormat(id, format)), - ); - } catch { - toast.error("Export failed."); + for (const id of ids) { + await exportConversationByFormat(id, format); + } + } catch (error) { + if (!isDownloadCancelled(error)) { + toast.error("Export failed."); + } } }} > @@ -969,10 +986,10 @@ export function AppSidebar() { ))} - {/* Bulk export and import live in Settings -> Chat -> Data. */} + {/* Bulk export and import live in Settings -> Data. */} - useSettingsDialogStore.getState().openDialog("chat") + useSettingsDialogStore.getState().openDialog("data") } > Export all chats… @@ -997,28 +1014,6 @@ export function AppSidebar() { - {isPinned ? ( - - - - - - Unpin - - - ) : null} ); } @@ -1218,6 +1213,9 @@ export function AppSidebar() { navigate({ to: "/projects" }); closeMobileIfOpen(); }} + onIntent={() => { + preloadSilently(router.preloadRoute({ to: "/projects" })); + }} className="group/projects-item relative" > - )} -
- )} - {(speculativeType === "mtp" || - speculativeType === "mtp+ngram") && ( -
-
- - Draft Tokens - - - Max MTP draft tokens per step - (--spec-draft-n-max). Lower = less wasted - draft decode; higher = bigger speedup when - acceptance stays high. Default: 2 on GPU, - 3 on CPU/Mac. - -
- { - const raw = e.target.value; - if (raw === "") { - setSpecDraftNMax(null); - return; - } - const parsed = Number.parseInt(raw, 10); - if (Number.isFinite(parsed)) { - const clamped = Math.max(1, Math.min(16, parsed)); - setSpecDraftNMax(clamped); - } - }} - data-test-id="spec-draft-n-max-input" - aria-label="Speculative decoding draft tokens" - className="h-7 w-[88px] rounded-full border-border bg-background hover:bg-accent/50 dark:border-transparent dark:bg-white/[0.05] dark:hover:bg-white/[0.1] pl-3 py-0 text-[13px] font-medium text-nav-fg outline-none focus-visible:ring-0" - /> -
- )} - - )} -
-
- - Tensor Parallelism - - - No effect on a single GPU. On multi-GPU setups, improves - tokens/sec during generation when using dense models. MoE - models don't benefit and can be much slower. - -
- -
- - )} - {/* No persistent "enable custom code" toggle: it is consented per model - via the load-time review dialog. */} - {/* Apply/Reset belongs to the model-reload settings above (context - length, KV cache, speculative decoding). Render it here, before - the Chat Template row, so it never reads as attached to Chat - Template (which is edited via its own dialog). When a model is - staged (deferred load), Load/Cancel takes its place: there's - nothing loaded to "apply" against yet. */} - {pendingSelection ? ( -
- {stagedDownloading && ( -

- Downloading…{" "} - {Math.round((stagedDownloadFraction ?? 0) * 100)}% + : "" + }`}

- )} - - {stagedLoading ? ( - // Mid-load: nothing to load or abandon until it settles, so disable. - - ) : ( -
+ {mtpUpdatable && llamaUpdateStatus?.update_available && ( - -
- )} -
- ) : modelSettingsDirty ? ( -
- - -
- ) : null} - -
- + )} +
+ )} + {showContextVramWarning && ( +

+ Context length exceeds the estimated VRAM capacity ( + {ggufMaxContextLength?.toLocaleString()} tokens). The + model may use system RAM. +

+ )} + + )}
- +
savePresetWithName(presetNameInput)} disabled={!(settingsHydrated && presetSaveState.canSubmit)} - variant={presetSaveState.isSaveReady ? "default" : "outline"} + variant={ + presetSaveState.isSaveReady ? "default" : "outline" + } size="sm" className={cn( "h-9 w-full rounded-full text-[13px] font-medium tracking-nav", @@ -1456,7 +912,8 @@ export function ChatSettingsPanel({ Prompt caching - Reuse compatible prompt prefixes for lower latency and cost. + Reuse compatible prompt prefixes for lower latency and + cost.
Anthropic exposes a 5 minute and a 1 hour ephemeral - cache pool. The 1 hour pool costs 2x base input on - write vs 1.25x for 5 minute, but reads stay 0.1x for - both, so a single read landing more than 5 minutes - after the write pays off the premium. + cache pool. The 1 hour pool costs 2x base input on write + vs 1.25x for 5 minute, but reads stay 0.1x for both, so + a single read landing more than 5 minutes after the + write pays off the premium.