diff --git a/.github/scripts/run-studio-permission-browser.sh b/.github/scripts/run-studio-permission-browser.sh
index 2007789035..e5a9a4c135 100755
--- a/.github/scripts/run-studio-permission-browser.sh
+++ b/.github/scripts/run-studio-permission-browser.sh
@@ -17,7 +17,8 @@ if [ -n "${STUDIO_PERMISSION_FRONTEND:-}" ]; then
fi
mkdir -p "$artifact_dir"
-unsloth studio reset-password
+# Wipe (not reset-password): the boot below must re-seed a fresh .bootstrap_password.
+rm -rf "$studio_home/auth"
UNSLOTH_API_ONLY=1 unsloth studio -H 127.0.0.1 -p "$port" "$@" \
>"$server_log" 2>&1 &
studio_pid=$!
diff --git a/.github/workflows/consolidated-tests-ci.yml b/.github/workflows/consolidated-tests-ci.yml
index d1bea819eb..afad1b6c46 100644
--- a/.github/workflows/consolidated-tests-ci.yml
+++ b/.github/workflows/consolidated-tests-ci.yml
@@ -7,7 +7,7 @@
#
# Why a separate workflow:
# - studio-backend-ci.yml's "Repo tests (CPU)" job already auto-discovers
-# tests/ minus tests/qlora, tests/saving, tests/utils, tests/sh. The 16
+# tests/ minus tests/qlora, tests/saving, tests/utils, tests/sh. The 17
# Bucket-A tests below live inside those --ignore dirs (CPU-runnable but
# historically excluded with their GPU siblings); pulling them out into
# a sibling job keeps the existing 760-passed baseline stable while we
@@ -274,6 +274,7 @@ jobs:
tests/saving/test_export_dispatch.py \
tests/saving/test_imatrix_export.py \
tests/saving/test_gguf_single_pass_export.py \
+ tests/saving/test_offline_gguf_vlm_tokenizer_7481.py \
tests/utils/test_attention_masks.py \
tests/utils/test_trunc_normal_patch.py \
tests/python/test_fast_language_model_text_only.py
@@ -365,17 +366,17 @@ jobs:
tests/saving/test_export_dispatch.py \
tests/saving/test_imatrix_export.py \
tests/saving/test_gguf_single_pass_export.py \
+ tests/saving/test_offline_gguf_vlm_tokenizer_7481.py \
tests/utils/test_attention_masks.py \
tests/utils/test_trunc_normal_patch.py \
tests/python/test_fast_language_model_text_only.py \
tests/test_bad_mappings_redirect.py \
tests/test_prefetch_snapshot_scope.py \
tests/test_gemma_2b_mapper_key.py \
- --deselect 'tests/utils/test_attention_masks.py::test_run_attention_flash_varlen_receives_window_and_softcap'
- # The deselected test monkeypatches flash_attn_varlen_func, which is
- # only bound on the module when `flash_attn` is importable. flash_attn
- # requires CUDA + dev toolchain, which the CPU-only ubuntu-latest
- # runner does not have. The other Bucket-A tests pass cleanly.
+ tests/test_raw_text_json_loading.py
+ # test_run_attention_flash_varlen_receives_window_and_softcap was deselected
+ # until attention_dispatch.py predefined flash_attn_varlen_func as None; it
+ # monkeypatches that name, so it no longer needs flash_attn on this runner.
- name: unsloth_zoo @ ${{ env.UNSLOTH_ZOO_REF }} — full pytest (CPU)
# 106 of 111 test_* in unsloth_zoo are CPU-only. The two CUDA-skip
@@ -2129,7 +2130,7 @@ jobs:
pip show unsloth_zoo
echo "::endgroup::"
echo "Consolidated job done. Coverage:"
- echo " - 16 unsloth Bucket-A tests under tests/saving/ + tests/utils/"
+ echo " - 17 unsloth Bucket-A tests under tests/saving/ + tests/utils/"
echo " - unsloth_zoo @ ${UNSLOTH_ZOO_REF} pytest tests/ (5 GPU cases deselected)"
echo " - unsloth_zoo.compiler.test_apply_fused_lm_head"
diff --git a/.github/workflows/cross-platform-parity-ci.yml b/.github/workflows/cross-platform-parity-ci.yml
index bb7dcbf8e4..45ce231743 100644
--- a/.github/workflows/cross-platform-parity-ci.yml
+++ b/.github/workflows/cross-platform-parity-ci.yml
@@ -1,18 +1,16 @@
# SPDX-License-Identifier: AGPL-3.0-only
# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved.
-# Runs installer parity and autostart opt-out tests on Windows and macOS.
+# Runs installer parity and autostart opt-out tests across all three platforms.
#
-# Why: that test is the guard that install.sh and install.ps1 stay in
-# sync, but today it only runs on ubuntu-latest (auto-discovered by
-# studio-backend-ci.yml's "Repo tests (CPU)" job). The test reads both
-# installer scripts, and on Windows Path.read_text() defaults to the
-# cp1252 locale encoding, so a non-cp1252 byte in install.sh (it already
-# contains a U+274C) raises UnicodeDecodeError there even though Linux and
-# macOS default to UTF-8. The reads were pinned to encoding="utf-8" in
-# #6166; this job keeps that from silently regressing by exercising the
-# test on the platforms it claims parity for. Pure pytest, no GPU,
-# sub-second, so the matrix is cheap.
+# Why: the parity test guards that install.sh and install.ps1 stay in sync.
+# It originally ran only on ubuntu-latest through studio-backend-ci.yml.
+# On Windows, Path.read_text() defaults to the cp1252 locale encoding, so a
+# non-cp1252 byte in install.sh raises UnicodeDecodeError even though Linux
+# and macOS default to UTF-8. The reads were pinned to encoding="utf-8" in
+# #6166; this matrix keeps that from silently regressing. Pure pytest, no GPU,
+# sub-second, so the matrix is cheap. Linux also runs the POSIX rollback test
+# under dash, matching the supported curl-to-sh installer path.
name: Cross-platform parity
@@ -23,6 +21,8 @@ on:
- 'install.ps1'
- 'tests/test_installer_skip_autostart.py'
- 'tests/python/test_cross_platform_parity.py'
+ - 'tests/sh/test_install_rollback_lifecycle.sh'
+ - 'tests/studio/test_install_rollback_lifecycle.ps1'
- '.github/workflows/cross-platform-parity-ci.yml'
push:
branches: [main]
@@ -31,6 +31,8 @@ on:
- 'install.ps1'
- 'tests/test_installer_skip_autostart.py'
- 'tests/python/test_cross_platform_parity.py'
+ - 'tests/sh/test_install_rollback_lifecycle.sh'
+ - 'tests/studio/test_install_rollback_lifecycle.ps1'
- '.github/workflows/cross-platform-parity-ci.yml'
workflow_dispatch:
@@ -47,7 +49,7 @@ jobs:
strategy:
fail-fast: false
matrix:
- os: [windows-latest, macos-latest]
+ os: [ubuntu-latest, windows-latest, macos-latest]
runs-on: ${{ matrix.os }}
timeout-minutes: 10
steps:
@@ -67,3 +69,10 @@ jobs:
tests/python/test_cross_platform_parity.py
tests/test_installer_skip_autostart.py
-q
+ - name: PowerShell rollback lifecycle tests
+ if: runner.os == 'Windows'
+ shell: pwsh
+ run: pwsh -NoProfile -File tests/studio/test_install_rollback_lifecycle.ps1
+ - name: POSIX rollback lifecycle tests
+ if: runner.os == 'Linux'
+ run: sh tests/sh/test_install_rollback_lifecycle.sh
diff --git a/.github/workflows/local-agent-guides-ci.yml b/.github/workflows/local-agent-guides-ci.yml
index c48328e90f..0dc0cc66d7 100644
--- a/.github/workflows/local-agent-guides-ci.yml
+++ b/.github/workflows/local-agent-guides-ci.yml
@@ -167,7 +167,9 @@ jobs:
# ── boot the server under test (factored helper) ──────────────────
- name: Serve unsloth run --disable-tools (gemma-4-E4B)
run: |
- unsloth studio reset-password
+ # Wipe, not reset-password: since #7573 the reset rotates in place and
+ # prints the new passphrase, which would land unmasked in the job log.
+ rm -rf ~/.unsloth/studio/auth
bash .github/scripts/serve-unsloth-run.sh \
--gguf-file "$GITHUB_WORKSPACE/gguf-cache/${GGUF_FILE}" \
--port "$STUDIO_PORT" --log-dir logs \
@@ -371,7 +373,7 @@ jobs:
- name: Serve unsloth run --disable-tools (gemma-4-E4B)
run: |
- unsloth studio reset-password
+ rm -rf ~/.unsloth/studio/auth
bash .github/scripts/serve-unsloth-run.sh \
--gguf-file "$GITHUB_WORKSPACE/gguf-cache/${GGUF_FILE}" \
--port "$STUDIO_PORT" --log-dir logs \
@@ -554,7 +556,7 @@ jobs:
- name: Serve unsloth run --disable-tools (gemma-4-E4B)
run: |
- unsloth studio reset-password
+ rm -rf ~/.unsloth/studio/auth
bash .github/scripts/serve-unsloth-run.sh \
--gguf-file "$GITHUB_WORKSPACE/gguf-cache/${GGUF_FILE}" \
--port "$STUDIO_PORT" --log-dir logs \
@@ -718,7 +720,7 @@ jobs:
- name: Serve unsloth run --disable-tools (gemma-3-270m)
run: |
- unsloth studio reset-password
+ rm -rf ~/.unsloth/studio/auth
bash .github/scripts/serve-unsloth-run.sh \
--model "$GGUF_REPO" --gguf-variant "$GGUF_VARIANT" \
--port "$STUDIO_PORT" --log-dir logs \
diff --git a/.github/workflows/release-desktop.yml b/.github/workflows/release-desktop.yml
index 081eda4e32..0a8d71610d 100644
--- a/.github/workflows/release-desktop.yml
+++ b/.github/workflows/release-desktop.yml
@@ -766,6 +766,7 @@ jobs:
env:
GH_REPO: ${{ github.repository }}
APP_VERSION: ${{ needs.prepare-version.outputs.app_version }}
+ PYPI_VERSION: ${{ needs.prepare-version.outputs.pypi_version }}
STUDIO_VERSION: ${{ needs.prepare-version.outputs.studio_version }}
DESKTOP_RELEASE_TAG: ${{ needs.prepare-version.outputs.desktop_release_tag }}
DESKTOP_PRERELEASE: ${{ needs.prepare-version.outputs.prerelease }}
@@ -911,6 +912,8 @@ jobs:
notes = pathlib.Path(os.environ['RUNNER_TEMP'], 'desktop-release-notes.md').read_text()
metadata = {
'version': os.environ['APP_VERSION'],
+ # App version is SemVer; CHANGELOG.md is keyed by the backend release.
+ 'pypi_version': os.environ['PYPI_VERSION'],
'notes': notes,
'pub_date': datetime.datetime.now(datetime.timezone.utc).isoformat(timespec='milliseconds').replace('+00:00', 'Z'),
'platforms': {
diff --git a/.github/workflows/startup-profile-ci.yml b/.github/workflows/startup-profile-ci.yml
new file mode 100644
index 0000000000..fbde99836d
--- /dev/null
+++ b/.github/workflows/startup-profile-ci.yml
@@ -0,0 +1,156 @@
+# SPDX-License-Identifier: AGPL-3.0-only
+# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved.
+
+# Measures where Studio's startup time goes, on each platform.
+#
+# Nothing recorded a number before: main.py logs "lifespan startup completed in X ms"
+# and studio_test_kit polls /healthz, but both throw the elapsed time away. A first
+# local run (Linux, warm cache, 18-core server) put `import main` at 5.7-6.6s BEFORE
+# the server can bind, dominated by eager module-level imports pulled in by routes:
+# torch ~1.9s self, unsloth_zoo ~0.8s, routes ~0.6s, transformers ~0.5s.
+#
+# Not a gate yet: --max-healthz-seconds exists, but a budget should come from
+# observed numbers rather than a guess.
+
+name: Startup profile
+
+on:
+ pull_request:
+ paths:
+ # The measured import graph is the whole backend tree: main.py imports auth,
+ # core, hub, loggers, models, picker, routes and utils at module scope.
+ - 'studio/backend/**'
+ - '!studio/backend/tests/**'
+ # The launch phase spawns `unsloth studio --api-only`, so the CLI counts too.
+ - 'unsloth_cli/**'
+ - 'studio/src-tauri/src/preflight**'
+ # The profiler hardcodes the desktop argv that process.rs::backend_args builds,
+ # so a change there must schedule a run or the two silently diverge.
+ - 'studio/src-tauri/src/process.rs'
+ - 'scripts/profile_startup.py'
+ - '.github/workflows/startup-profile-ci.yml'
+ # The job profiles whatever `install.sh --local` built: the installers pick the
+ # venv's Python and the dependency specs, and pyproject's include list is what
+ # makes --local overlay studio.backend*.
+ - 'install.sh'
+ - 'install.ps1'
+ - 'pyproject.toml'
+ # --local also runs the checkout's setup scripts (install.sh picks
+ # $_REPO_ROOT/studio/setup.sh, the editable install resolves setup.ps1 to the
+ # repo), and both call install_python_stack.py, which picks the dependencies.
+ - 'studio/setup.sh'
+ - 'studio/setup.ps1'
+ - 'studio/install_python_stack.py'
+ workflow_dispatch:
+ inputs:
+ repeats:
+ description: 'launch repeats per OS (median reported)'
+ type: string
+ default: '3'
+
+concurrency:
+ group: ${{ github.workflow }}-${{ github.ref }}
+ cancel-in-progress: true
+
+permissions:
+ contents: read
+
+jobs:
+ profile:
+ name: startup ${{ matrix.os }}
+ runs-on: ${{ matrix.os }}
+ timeout-minutes: 60
+ continue-on-error: true
+ strategy:
+ fail-fast: false
+ matrix:
+ os: [ubuntu-latest, macos-14, windows-latest]
+
+ env:
+ UNSLOTH_STUDIO_HOME: ${{ github.workspace }}/.studio-home
+ # A wildcard bind calls ifconfig.me on the startup path; loopback times our code.
+ UNSLOTH_STUDIO_DISABLE_PUBLIC_CHECK: '1'
+
+ steps:
+ - uses: actions/checkout@de0fac2e4500dabe0009e67214ff5f5447ce83dd # v6.0.2
+ with:
+ persist-credentials: false
+
+ - name: Install Studio
+ shell: bash
+ env:
+ GH_TOKEN: ${{ secrets.GITHUB_TOKEN }}
+ run: |
+ set -o pipefail
+ mkdir -p logs
+ # --local is load-bearing: it overlays the checkout, so the profiled server
+ # is this diff. Without it install.sh resolves unsloth from PyPI.
+ if [ "${{ runner.os }}" = "Windows" ]; then
+ pwsh -NoProfile -File ./install.ps1 --local 2>&1 | tee logs/install.log
+ else
+ bash install.sh --local 2>&1 | tee logs/install.log
+ fi
+
+ - name: Profile startup
+ shell: bash
+ run: |
+ BIN="$UNSLOTH_STUDIO_HOME/unsloth_studio/bin/unsloth"
+ [ -x "$BIN" ] || BIN="$UNSLOTH_STUDIO_HOME/unsloth_studio/Scripts/unsloth.exe"
+ [ -x "$BIN" ] || BIN=""
+ # Profile imports with the INSTALLED interpreter: that venv is what launches.
+ PY="$UNSLOTH_STUDIO_HOME/unsloth_studio/bin/python"
+ [ -x "$PY" ] || PY="$UNSLOTH_STUDIO_HOME/unsloth_studio/Scripts/python.exe"
+ [ -x "$PY" ] || PY="$(command -v python3 || command -v python)"
+ python3 scripts/profile_startup.py \
+ --python "$PY" \
+ ${BIN:+--bin "$BIN"} \
+ --repeats "${{ inputs.repeats || '3' }}" \
+ --json "startup-${{ matrix.os }}.json" 2>&1 | tee logs/profile.log
+
+ - name: Summary
+ if: always()
+ shell: bash
+ run: |
+ f="startup-${{ matrix.os }}.json"
+ [ -f "$f" ] || { echo "no profile produced"; exit 0; }
+ python3 - "$f" >> "$GITHUB_STEP_SUMMARY" <<'PY'
+ import json, sys
+ d = json.load(open(sys.argv[1]))
+ print(f"### {d['platform']} / {d['machine']} (py {d['python']}, {d['cpu_count']} cpu)\n")
+ imp = d.get("imports", {})
+ # Gate on ok: a failed `import main` still leaves rows, so a total can lie.
+ if imp.get("ok"):
+ print(f"**`import main`: {imp['total_seconds']}s**\n")
+ print("| package | self ms |")
+ print("|---|---:|")
+ for k, v in list(imp.get("self_by_package_ms", {}).items())[:8]:
+ print(f"| {k} | {v} |")
+ print()
+ else:
+ print("**`import main` failed - no valid import profile**\n")
+ print("```\n" + (imp.get("error") or "")[-1500:] + "\n```\n")
+ lau = d.get("launch") or {}
+ runs = len(lau.get("runs") or [])
+ failed = lau.get("failed_runs") or 0
+ if lau.get("healthz_median_seconds") is not None:
+ # The aggregates cover only the runs that reached healthz, so flag the
+ # failures: bare numbers would read as a normal fast startup.
+ note = f" _({runs - failed} of {runs} launches; {failed} never became healthy)_" if failed else ""
+ print(f"**time to a healthy port: {lau['healthz_median_seconds']}s median, "
+ f"{lau['healthz_max_seconds']}s max**{note}\n")
+ elif lau.get("skipped"):
+ print(f"_launch phase skipped: {lau['skipped']}_\n")
+ elif runs:
+ print(f"**no launch measurement: all {runs} launches failed to become healthy**\n")
+ PY
+
+ - name: Upload profile
+ if: always()
+ uses: actions/upload-artifact@043fb46d1a93c77aae656e7c1c64a875d1fc6a0a # v7.0.1
+ with:
+ name: startup-profile-${{ matrix.os }}
+ path: |
+ startup-*.json
+ logs/
+ retention-days: 14
+ if-no-files-found: warn
diff --git a/.github/workflows/studio-api-smoke.yml b/.github/workflows/studio-api-smoke.yml
index cdf1f6bf12..1cfa66fea4 100644
--- a/.github/workflows/studio-api-smoke.yml
+++ b/.github/workflows/studio-api-smoke.yml
@@ -113,7 +113,8 @@ jobs:
- name: Reset auth + boot Unsloth (API-only)
run: |
- unsloth studio reset-password
+ # Wipe (not reset-password): the boot below must re-seed a fresh .bootstrap_password.
+ rm -rf ~/.unsloth/studio/auth
mkdir -p logs
UNSLOTH_API_ONLY=1 unsloth studio -H 127.0.0.1 -p "$STUDIO_PORT" \
> logs/studio.log 2>&1 &
diff --git a/.github/workflows/studio-backend-ci.yml b/.github/workflows/studio-backend-ci.yml
index b8f587b63e..dd5efbb299 100644
--- a/.github/workflows/studio-backend-ci.yml
+++ b/.github/workflows/studio-backend-ci.yml
@@ -30,6 +30,13 @@ on:
- 'unsloth/**'
- 'unsloth_cli/**'
- 'tests/**'
+ # The root installers: tests/sh/*.sh and tests/studio/install/* assert
+ # against these two files, so a change here must run the suite that
+ # covers it. Without them an install-only edit (the shape most AMD/ROCm
+ # routing fixes take) skipped Backend CI entirely.
+ - 'install.sh'
+ - 'install.ps1'
+ - 'scripts/**'
- 'pyproject.toml'
- '.github/workflows/studio-backend-ci.yml'
push:
@@ -193,6 +200,7 @@ jobs:
--ignore=tests/sh \
--ignore=tests/studio/test_hardware_dispatch_matrix.py \
--ignore=tests/studio/test_is_mlx_dispatch_gate.py \
+ --ignore=tests/studio/test_xpu_spoof_pipeline.py \
--ignore=tests/vllm_compat \
--ignore=tests/version_compat \
-m 'not server and not e2e' \
@@ -205,36 +213,53 @@ jobs:
env:
PYTHONPATH: ${{ github.workspace }}/studio
UNSLOTH_COMPILE_DISABLE: '1'
- # These two files mutate hardware.py module globals at runtime
- # via the spoof fixtures, which leaks state into any other test
- # that imports hardware. Run them in their own pytest invocation
- # so the leak does not cross file boundaries.
+ # These files mutate hardware.py module globals at runtime via the
+ # spoof fixtures (CUDA/ROCm/XPU/MLX/CPU), which leaks state into any
+ # other test that imports hardware. Run them in their own pytest
+ # invocation so the leak does not cross file boundaries.
run: |
python -m pytest -q --tb=short \
tests/studio/test_hardware_dispatch_matrix.py \
- tests/studio/test_is_mlx_dispatch_gate.py
+ tests/studio/test_is_mlx_dispatch_gate.py \
+ tests/studio/test_xpu_spoof_pipeline.py
+
+ - name: CLI tests (unsloth_cli)
+ # unsloth_cli/tests had no CI at all: `unsloth_cli/**` was only a paths
+ # trigger and a ruff target, so 673 tests covering the studio launcher,
+ # the pre-exposure gate and the auth secret writers ran nowhere, and
+ # four of them had been failing on main unnoticed.
+ # Own step, not folded into the tests/ discovery above: pyproject's
+ # testpaths is tests/, and this suite needs no PYTHONPATH or CUDA spoof
+ # (it self-bootstraps sys.path and imports neither unsloth nor torch).
+ run: python -m pytest unsloth_cli/tests -q --tb=short
- name: Shell installer tests
- # Subset that does not depend on a writable / pristine install.sh
- # tree; test_install_host_defaults.sh checks install.ps1 layout
- # which has drifted (separate followup).
+ # Auto-discovered rather than allowlisted. The old hardcoded list had
+ # silently fallen seven files behind tests/run_all.sh, including
+ # test_strixhalo_wsl_reroute.sh -- the only shell coverage of the ROCm
+ # WSL reroute -- so that suite never ran on a PR. Skips are explicit,
+ # each with a reason, and tests/studio/test_ci_shell_suite_coverage.py
+ # fails if this step stops discovering the directory or the skip list
+ # grows without one.
+ #
+ # Skipped:
+ # test_install_host_defaults.sh: asserts an install.ps1 layout that
+ # has drifted (separate followup).
+ # test_install_rollback_lifecycle.sh: already runs on both platforms
+ # in cross-platform-parity-ci.yml.
run: |
set -e
- for s in \
- tests/sh/test_get_torch_index_url.sh \
- tests/sh/test_mac_intel_compat.sh \
- tests/sh/test_node_decision.sh \
- tests/sh/test_studio_home_node_dir.sh \
- tests/sh/test_system_node_readonly.sh \
- tests/sh/test_nvcc_meets_llama_minimum.sh \
- tests/sh/test_resolve_cuda_archs.sh \
- tests/sh/test_tauri_install_exit_order.sh \
- tests/sh/test_torch_constraint.sh \
- tests/sh/test_torch_flavor.sh \
- tests/sh/test_with_llama_cpp_dir_flag.sh \
- tests/sh/test_with_llama_cpp_dir_link_behavior.sh; do
+ skip="test_install_host_defaults.sh test_install_rollback_lifecycle.sh"
+ found=0
+ for s in tests/sh/test_*.sh; do
+ case " $skip " in
+ *" $(basename "$s") "*) echo "skipping $s (see workflow comment)"; continue ;;
+ esac
+ found=$((found + 1))
echo "::group::$s"
bash "$s"
echo "::endgroup::"
done
+ [ "$found" -gt 0 ] || { echo "::error::no shell tests discovered under tests/sh"; exit 1; }
+ echo "ran $found shell installer test files"
diff --git a/.github/workflows/studio-frontend-ci.yml b/.github/workflows/studio-frontend-ci.yml
index 3a9e373915..773e555c8b 100644
--- a/.github/workflows/studio-frontend-ci.yml
+++ b/.github/workflows/studio-frontend-ci.yml
@@ -133,6 +133,9 @@ jobs:
- name: Typecheck
run: npm run typecheck
+ - name: Unit tests
+ run: npm test
+
- name: Build
run: npm run build
diff --git a/.github/workflows/studio-inference-smoke.yml b/.github/workflows/studio-inference-smoke.yml
index c2d52eac22..c37c9555bf 100644
--- a/.github/workflows/studio-inference-smoke.yml
+++ b/.github/workflows/studio-inference-smoke.yml
@@ -127,7 +127,8 @@ jobs:
- name: Reset auth + boot Unsloth (API-only)
run: |
- unsloth studio reset-password
+ # Wipe (not reset-password): the boot below must re-seed a fresh .bootstrap_password.
+ rm -rf ~/.unsloth/studio/auth
mkdir -p logs
UNSLOTH_API_ONLY=1 unsloth studio -H 127.0.0.1 -p "$STUDIO_PORT" \
> logs/studio.log 2>&1 &
@@ -400,7 +401,7 @@ jobs:
# tool_policy=None so each request's `enable_tools` field is
# honoured.
run: |
- unsloth studio reset-password
+ rm -rf ~/.unsloth/studio/auth
mkdir -p logs
UNSLOTH_API_ONLY=1 unsloth studio -H 127.0.0.1 -p "$STUDIO_PORT" \
> logs/studio.log 2>&1 &
@@ -978,7 +979,7 @@ jobs:
# response_format requests aren't routed through the agentic
# tool loop.
run: |
- unsloth studio reset-password
+ rm -rf ~/.unsloth/studio/auth
mkdir -p logs
UNSLOTH_API_ONLY=1 unsloth studio -H 127.0.0.1 -p "$STUDIO_PORT" \
> logs/studio.log 2>&1 &
diff --git a/.github/workflows/studio-mac-api-smoke.yml b/.github/workflows/studio-mac-api-smoke.yml
index 1968885a1d..c2307f17a1 100644
--- a/.github/workflows/studio-mac-api-smoke.yml
+++ b/.github/workflows/studio-mac-api-smoke.yml
@@ -101,7 +101,8 @@ jobs:
- name: Reset auth + boot Unsloth (API-only)
run: |
- unsloth studio reset-password
+ # Wipe (not reset-password): the boot below must re-seed a fresh .bootstrap_password.
+ rm -rf ~/.unsloth/studio/auth
mkdir -p logs
UNSLOTH_API_ONLY=1 unsloth studio -H 127.0.0.1 -p "$STUDIO_PORT" \
> logs/studio.log 2>&1 &
diff --git a/.github/workflows/studio-mac-inference-smoke.yml b/.github/workflows/studio-mac-inference-smoke.yml
index ce15eed5c8..1dbf86ae98 100644
--- a/.github/workflows/studio-mac-inference-smoke.yml
+++ b/.github/workflows/studio-mac-inference-smoke.yml
@@ -126,7 +126,8 @@ jobs:
- name: Reset auth + boot Unsloth (API-only)
run: |
- unsloth studio reset-password
+ # Wipe (not reset-password): the boot below must re-seed a fresh .bootstrap_password.
+ rm -rf ~/.unsloth/studio/auth
mkdir -p logs
UNSLOTH_API_ONLY=1 unsloth studio -H 127.0.0.1 -p "$STUDIO_PORT" \
> logs/studio.log 2>&1 &
@@ -386,7 +387,7 @@ jobs:
# tool_policy=None so each request's `enable_tools` field is
# honoured.
run: |
- unsloth studio reset-password
+ rm -rf ~/.unsloth/studio/auth
mkdir -p logs
UNSLOTH_API_ONLY=1 unsloth studio -H 127.0.0.1 -p "$STUDIO_PORT" \
> logs/studio.log 2>&1 &
@@ -831,7 +832,7 @@ jobs:
# response_format requests aren't routed through the agentic
# tool loop.
run: |
- unsloth studio reset-password
+ rm -rf ~/.unsloth/studio/auth
mkdir -p logs
UNSLOTH_API_ONLY=1 unsloth studio -H 127.0.0.1 -p "$STUDIO_PORT" \
> logs/studio.log 2>&1 &
diff --git a/.github/workflows/studio-mac-ui-smoke.yml b/.github/workflows/studio-mac-ui-smoke.yml
index 7375e9bcbf..3bed2fcdff 100644
--- a/.github/workflows/studio-mac-ui-smoke.yml
+++ b/.github/workflows/studio-mac-ui-smoke.yml
@@ -146,7 +146,8 @@ jobs:
- name: Reset auth + boot Unsloth
run: |
- unsloth studio reset-password
+ # Wipe (not reset-password): the boot below must re-seed a fresh .bootstrap_password.
+ rm -rf ~/.unsloth/studio/auth
mkdir -p logs
UNSLOTH_API_ONLY=1 unsloth studio -H 127.0.0.1 -p "$STUDIO_PORT" \
> logs/studio.log 2>&1 &
@@ -190,7 +191,7 @@ jobs:
# 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 Unsloth
- # (kill, reset-password, reboot, wait /api/health, re-export
+ # (kill, wipe auth, reboot, wait /api/health, re-export
# bootstrap pw) before re-running the script. A real test failure
# (assertion / timeout) does NOT match any pattern so it bypasses
# retry and surfaces immediately.
@@ -213,7 +214,7 @@ jobs:
echo "::warning::Playwright flake on attempt ${attempt}; resetting Unsloth and retrying..."
kill "${STUDIO_PID}" 2>/dev/null || true
sleep 2
- unsloth studio reset-password
+ rm -rf ~/.unsloth/studio/auth
UNSLOTH_API_ONLY=1 unsloth studio -H 127.0.0.1 -p "$STUDIO_PORT" \
> "logs/studio_retry_${attempt}.log" 2>&1 &
STUDIO_PID=$!
@@ -251,7 +252,7 @@ jobs:
- name: Reset auth + boot Unsloth for extra UI tests (port 18897)
run: |
- unsloth studio reset-password
+ rm -rf ~/.unsloth/studio/auth
mkdir -p logs
UNSLOTH_API_ONLY=1 unsloth studio -H 127.0.0.1 -p 18897 \
> logs/studio_extra.log 2>&1 &
@@ -308,7 +309,7 @@ jobs:
echo "::warning::Playwright flake on attempt ${attempt}; resetting Unsloth and retrying..."
kill "${STUDIO_EXTRA_PID}" 2>/dev/null || true
sleep 2
- unsloth studio reset-password
+ rm -rf ~/.unsloth/studio/auth
UNSLOTH_API_ONLY=1 unsloth studio -H 127.0.0.1 -p 18897 \
> "logs/studio_extra_retry_${attempt}.log" 2>&1 &
STUDIO_EXTRA_PID=$!
diff --git a/.github/workflows/studio-tauri-smoke.yml b/.github/workflows/studio-tauri-smoke.yml
index 8e26b9fd0c..c6dad07f37 100644
--- a/.github/workflows/studio-tauri-smoke.yml
+++ b/.github/workflows/studio-tauri-smoke.yml
@@ -91,6 +91,16 @@ jobs:
npm run build
test -f dist/index.html
+ # The crate carries ~100 unit tests (native_file_dialogs, preflight,
+ # install, desktop_auth, ...) that nothing ran until now: this workflow
+ # only ever built. Run them here, where the toolchain and the WebKit dev
+ # packages are already installed, so a broken assertion fails the PR
+ # instead of sitting unnoticed. `--no-fail-fast` reports every failing
+ # test in one run rather than stopping at the first.
+ - name: Rust unit tests (studio/src-tauri)
+ working-directory: studio/src-tauri
+ run: cargo test --no-fail-fast
+
- name: Tauri debug build (Linux, no bundle, no codesign)
# `--debug` + `--no-bundle` keeps this lean: compiles the Rust crate,
# confirms the frontend dist is wired into Tauri, but skips the AppImage
diff --git a/.github/workflows/studio-ui-smoke.yml b/.github/workflows/studio-ui-smoke.yml
index 0ad55ebd6d..3a0713f301 100644
--- a/.github/workflows/studio-ui-smoke.yml
+++ b/.github/workflows/studio-ui-smoke.yml
@@ -115,7 +115,8 @@ jobs:
- name: Reset auth + boot Unsloth
run: |
- unsloth studio reset-password
+ # Wipe (not reset-password): the boot below must re-seed a fresh .bootstrap_password.
+ rm -rf ~/.unsloth/studio/auth
mkdir -p logs
UNSLOTH_API_ONLY=1 unsloth studio -H 127.0.0.1 -p "$STUDIO_PORT" \
> logs/studio.log 2>&1 &
@@ -193,7 +194,7 @@ jobs:
# warm install we already did) so this adds little wall time.
- name: Reset auth + boot Unsloth for extra UI tests (port 18894)
run: |
- unsloth studio reset-password
+ rm -rf ~/.unsloth/studio/auth
mkdir -p logs
UNSLOTH_API_ONLY=1 unsloth studio -H 127.0.0.1 -p 18894 \
> logs/studio_extra.log 2>&1 &
@@ -231,6 +232,15 @@ jobs:
mkdir -p logs/playwright_extra
python tests/studio/playwright_extra_ui.py
+ - name: UI font size scaling regression (Playwright)
+ env:
+ BASE_URL: http://127.0.0.1:18894
+ STUDIO_PW: ${{ env.STUDIO_EXTRA_NEW_PW }}
+ PW_ART_DIR: logs/playwright_fontscale
+ run: |
+ mkdir -p logs/playwright_fontscale
+ python tests/studio/playwright_ui_font_scale.py
+
- name: Stop second Unsloth
if: always()
run: |
@@ -244,7 +254,7 @@ jobs:
# (RAG embedder + llama.cpp probe) stay hidden from the picker.
- name: Reset auth + boot Unsloth for model-config tests (port 18898)
run: |
- unsloth studio reset-password
+ rm -rf ~/.unsloth/studio/auth
mkdir -p logs
UNSLOTH_API_ONLY=1 unsloth studio -H 127.0.0.1 -p 18898 \
> logs/studio_modelcfg.log 2>&1 &
@@ -290,7 +300,7 @@ jobs:
# earlier UI tests. No GGUF -- the bug surface is the composer.
- name: Reset auth + boot Unsloth for IME / i18n tests (port 18896)
run: |
- unsloth studio reset-password
+ rm -rf ~/.unsloth/studio/auth
mkdir -p logs
UNSLOTH_API_ONLY=1 unsloth studio -H 127.0.0.1 -p 18896 \
> logs/studio_ime.log 2>&1 &
@@ -352,6 +362,7 @@ jobs:
logs/playwright
logs/playwright-permissions-*
logs/playwright_extra
+ logs/playwright_fontscale
logs/playwright_modelcfg
logs/playwright_ime
logs/studio-permissions-*.log
diff --git a/.github/workflows/studio-update-smoke.yml b/.github/workflows/studio-update-smoke.yml
index 625c2c7811..047840e41c 100644
--- a/.github/workflows/studio-update-smoke.yml
+++ b/.github/workflows/studio-update-smoke.yml
@@ -146,6 +146,46 @@ jobs:
kill "$PID" 2>/dev/null || true
echo "post-update Unsloth /api/health OK"
+ - name: A complete install reports itself complete
+ run: |
+ set -o pipefail
+ unsloth studio verify-install
+ unsloth studio desktop-capabilities --json | tee /tmp/caps.json
+ jq -e '.studio_install_ok == true' /tmp/caps.json
+ jq -e '.desktop_manageability_version >= 2' /tmp/caps.json
+
+ - name: An incomplete install must not report itself ready
+ # An installer killed part-way leaves a working CLI but no studio.txt
+ # deps, which the old preflight called ManagedReady. The manifest is
+ # written last, so removing it reproduces that state.
+ run: |
+ set -o pipefail
+ # install.sh's default root, resolved explicitly: `python` on PATH
+ # here is setup-python's, not the managed venv.
+ MANIFEST="$HOME/.unsloth/studio/unsloth_studio/unsloth_install_manifest.json"
+ test -f "$MANIFEST" || { echo "::error::installer never wrote $MANIFEST"; exit 1; }
+ rm -f "$MANIFEST"
+ unsloth studio desktop-capabilities --json | tee /tmp/caps_bad.json
+ jq -e '.studio_install_ok == false' /tmp/caps_bad.json
+ if unsloth studio verify-install; then
+ echo "::error::verify-install passed on an install with no manifest"
+ exit 1
+ fi
+ echo "incomplete install correctly reported not-ready"
+
+ - name: Update repairs an incomplete install
+ # `--local` bypasses setup.sh's PyPI version compare, so this asserts
+ # the repair OUTCOME. The non-local fast path the desktop Repair button
+ # uses is covered by tests/studio/install/test_setup_fast_path_guard.py.
+ env:
+ GH_TOKEN: ${{ secrets.GITHUB_TOKEN }}
+ run: |
+ set -o pipefail
+ unsloth studio update --local 2>&1 | tee logs/update_repair.log
+ unsloth studio verify-install
+ unsloth studio desktop-capabilities --json | jq -e '.studio_install_ok == true'
+ echo "update repaired the incomplete install"
+
- name: Uninstall and verify clean
# Round-trip the installer through scripts/uninstall.sh: confirms the
# uninstaller actually finds and removes everything install.sh +
diff --git a/.github/workflows/studio-windows-api-smoke.yml b/.github/workflows/studio-windows-api-smoke.yml
index 6dbcceebbd..b328939846 100644
--- a/.github/workflows/studio-windows-api-smoke.yml
+++ b/.github/workflows/studio-windows-api-smoke.yml
@@ -179,7 +179,8 @@ jobs:
- name: Reset auth + boot Unsloth (API-only)
run: |
- unsloth studio reset-password
+ # Wipe (not reset-password): the boot below must re-seed a fresh .bootstrap_password.
+ rm -rf ~/.unsloth/studio/auth
mkdir -p logs
UNSLOTH_API_ONLY=1 unsloth studio -H 127.0.0.1 -p "$STUDIO_PORT" \
> logs/studio.log 2>&1 &
diff --git a/.github/workflows/studio-windows-inference-smoke.yml b/.github/workflows/studio-windows-inference-smoke.yml
index 3ebe442f52..d821664327 100644
--- a/.github/workflows/studio-windows-inference-smoke.yml
+++ b/.github/workflows/studio-windows-inference-smoke.yml
@@ -229,7 +229,8 @@ jobs:
- name: Reset auth + boot Unsloth (API-only)
run: |
- unsloth studio reset-password
+ # Wipe (not reset-password): the boot below must re-seed a fresh .bootstrap_password.
+ rm -rf ~/.unsloth/studio/auth
mkdir -p logs
UNSLOTH_API_ONLY=1 unsloth studio -H 127.0.0.1 -p "$STUDIO_PORT" \
> logs/studio.log 2>&1 &
@@ -573,7 +574,7 @@ jobs:
- name: Reset auth + boot Unsloth (API-only, default tool policy)
run: |
- unsloth studio reset-password
+ rm -rf ~/.unsloth/studio/auth
mkdir -p logs
UNSLOTH_API_ONLY=1 unsloth studio -H 127.0.0.1 -p "$STUDIO_PORT" \
> logs/studio.log 2>&1 &
@@ -1074,7 +1075,7 @@ jobs:
- name: Reset auth + boot Unsloth (API-only)
run: |
- unsloth studio reset-password
+ rm -rf ~/.unsloth/studio/auth
mkdir -p logs
UNSLOTH_API_ONLY=1 unsloth studio -H 127.0.0.1 -p "$STUDIO_PORT" \
> logs/studio.log 2>&1 &
@@ -1546,7 +1547,7 @@ jobs:
- name: Reset auth + boot Unsloth (API-only)
run: |
- unsloth studio reset-password
+ rm -rf ~/.unsloth/studio/auth
mkdir -p logs
UNSLOTH_API_ONLY=1 unsloth studio -H 127.0.0.1 -p "$STUDIO_PORT" \
> logs/studio.log 2>&1 &
@@ -1888,8 +1889,11 @@ jobs:
# (step/substep -> Write-StudioStdoutMirror / Get-StudioAnsi).
$script:StudioVtOk = $false
$script:UnslothVerbose = $false
+ # Get-HostMachineArch is reached only on the absent path, where
+ # Test-VCRedistInstalled consults it before trusting the System32 DLL, so
+ # part A passes without it and only the clean-box part fails.
foreach ($fn in @('Get-StudioAnsi', 'Write-StudioStdoutMirror', 'step', 'substep',
- 'Invoke-SetupCommand', 'Refresh-Environment',
+ 'Invoke-SetupCommand', 'Refresh-Environment', 'Get-HostMachineArch',
'Test-VCRedistInstalled', 'Ensure-VCRedist')) {
$src = Get-FunctionSource -Path $setup -Name $fn
if (-not $src) { throw "Function '$fn' not found in setup.ps1" }
diff --git a/.github/workflows/studio-windows-ui-smoke.yml b/.github/workflows/studio-windows-ui-smoke.yml
index f401f7be44..d23cca323f 100644
--- a/.github/workflows/studio-windows-ui-smoke.yml
+++ b/.github/workflows/studio-windows-ui-smoke.yml
@@ -297,7 +297,8 @@ jobs:
- name: Reset auth + boot Unsloth
run: |
- unsloth studio reset-password
+ # Wipe (not reset-password): the boot below must re-seed a fresh .bootstrap_password.
+ rm -rf ~/.unsloth/studio/auth
mkdir -p logs
UNSLOTH_API_ONLY=1 unsloth studio -H 127.0.0.1 -p "$STUDIO_PORT" \
> logs/studio.log 2>&1 &
@@ -352,7 +353,7 @@ jobs:
- name: Reset auth + boot Unsloth for extra UI tests (port 18897)
run: |
- unsloth studio reset-password
+ rm -rf ~/.unsloth/studio/auth
mkdir -p logs
UNSLOTH_API_ONLY=1 unsloth studio -H 127.0.0.1 -p 18897 \
> logs/studio_extra.log 2>&1 &
diff --git a/.github/workflows/studio-windows-update-smoke.yml b/.github/workflows/studio-windows-update-smoke.yml
index 42d74d47d2..0dcc828e6b 100644
--- a/.github/workflows/studio-windows-update-smoke.yml
+++ b/.github/workflows/studio-windows-update-smoke.yml
@@ -198,6 +198,31 @@ jobs:
fi
echo "update path took the prebuilt fast path"
+ - name: Update must keep the --no-torch install GGUF-only
+ run: |
+ # `unsloth studio update` exports no UNSLOTH_NO_TORCH, so setup.ps1 has
+ # to recover the mode from the install manifest. Without that it reads
+ # the missing torch as a stale venv and tries to delete the venv it is
+ # running out of, and the shared dependency pass pulls torch back in.
+ # The skip line only prints when the dependency pass actually runs, so
+ # don't demand it if the fast path short-circuited that pass.
+ if grep -q "running ordered dependency installation" logs/update.log \
+ && ! grep -q "skipping direct PyTorch and Triton installation (no-torch mode)" logs/update.log; then
+ echo "::error::studio update left no-torch mode; it would reinstall PyTorch."
+ grep -iE "no-torch|stale venv|PyTorch" logs/update.log | tail -40
+ exit 1
+ fi
+ PY="$HOME/.unsloth/studio/unsloth_studio/Scripts/python.exe"
+ if [ ! -f "$PY" ]; then
+ echo "::error::studio venv interpreter missing at $PY"
+ exit 1
+ fi
+ if "$PY" -c "import torch" 2>/dev/null; then
+ echo "::error::torch was reinstalled into the --no-torch venv."
+ exit 1
+ fi
+ echo "update preserved no-torch mode"
+
- name: Second update must also be a no-op
env:
GH_TOKEN: ${{ secrets.GITHUB_TOKEN }}
diff --git a/.github/workflows/wheel-smoke.yml b/.github/workflows/wheel-smoke.yml
index cdad617027..f7a7511616 100644
--- a/.github/workflows/wheel-smoke.yml
+++ b/.github/workflows/wheel-smoke.yml
@@ -127,6 +127,31 @@ jobs:
cd /tmp
/tmp/v/bin/python -c "from studio.backend.main import app; print('Unsloth backend OK:', app.title)"
+ - name: CLI without the Studio stack guides instead of tracebacking
+ # The smoke above installs studio.txt first, so it cannot catch a wheel
+ # that ships studio/ without declaring what it imports (#4701, #5260,
+ # #7147). Drop only structlog to reuse that venv without a re-download.
+ run: |
+ set -eu
+ /tmp/v/bin/pip uninstall -y structlog >/dev/null
+ cd /tmp
+ status=0
+ for args in "export ./nope ./out" "list-checkpoints"; do
+ echo "--- unsloth $args"
+ out=$(/tmp/v/bin/unsloth $args 2>&1 || true)
+ printf '%s\n' "$out"
+ case "$out" in
+ *Traceback*)
+ echo "FAIL: raw traceback instead of guidance"; status=1 ;;
+ esac
+ case "$out" in
+ *'unsloth studio update'*) ;;
+ *) echo "FAIL: no remediation in the message"; status=1 ;;
+ esac
+ done
+ /tmp/v/bin/pip install -q structlog >/dev/null
+ exit "$status"
+
- name: Upload wheel on failure
if: failure()
uses: actions/upload-artifact@043fb46d1a93c77aae656e7c1c64a875d1fc6a0a # v7.0.1
diff --git a/.gitignore b/.gitignore
index 39ca2226ca..fa6997cb06 100644
--- a/.gitignore
+++ b/.gitignore
@@ -208,6 +208,9 @@ tmp/
**/node_modules/
auth.db
+# Packaging snapshot of the root CHANGELOG.md (written by build.sh)
+studio/CHANGELOG.md
+
# Tauri local build/generated output
studio/src-tauri/target/
studio/src-tauri/gen/
@@ -238,4 +241,5 @@ package-lock.json
!studio/package-lock.json
llama.cpp/
# Stray "~" dir some tools create from a literal ~ TMPDIR; never part of the repo.
-/~/
+~/
+/temp/
diff --git a/CHANGELOG.md b/CHANGELOG.md
new file mode 100644
index 0000000000..241e013cea
--- /dev/null
+++ b/CHANGELOG.md
@@ -0,0 +1,88 @@
+# Changelog
+
+Release notes for Unsloth and Unsloth Studio.
+
+Unsloth Studio reads this file to show release notes inside the "New Unsloth
+version" update popup. Edit it here and the popup picks the change up on the
+next update check, with no release or rebuild required.
+
+## Format
+
+Every release is a level-2 heading whose first token is the version, optionally
+followed by a date:
+
+```md
+## 2026.7.6 - 2026-07-22
+```
+
+`## [2026.7.6] - 2026-07-22` and `## v2026.7.6` also work. Everything under a
+heading, up to the next level-2 heading, is that release's notes and renders as
+Markdown in the popup.
+
+Notes are matched to one exact version. When Studio offers an update to
+`2026.7.6` it renders the `2026.7.6` section and nothing else. If that section
+is missing, the popup links out to the online changelog rather than showing
+notes from an unrelated release, so a new version needs its own section here
+before its notes can appear.
+
+Keep the newest release at the top. Lead each bullet with the change itself:
+the collapsed popup highlights the first sentence and dims the rest.
+`## Unreleased` is ignored by the popup, so it is safe to stage notes there and
+rename the heading at release time.
+
+
+
+## Unreleased
+
+## 2026.7.5
+
+### What's Changed
+
+- AMD support is here. Train, run RL, chat with and deploy 500+ models on
+ Radeon, Instinct, Ryzen and data center GPUs across Windows, WSL and Linux,
+ up to 2x faster with 70% less VRAM and no accuracy loss.
+- Intel XPU support lands in Studio, so Arc and Data Center GPUs run chat and
+ training alongside the NVIDIA, AMD and Apple paths.
+- Local speech to text dictation runs fully offline, with slim Whisper bundles
+ and a picker for custom models.
+- DoRA training is available in Studio, selectable next to LoRA and full
+ fine-tuning in the training tab.
+- The update popup previews release notes inline, pulled from this file and
+ matched to the exact version being offered.
+
+### AMD, 23 July update
+
+Our AMD collaboration, custom Triton kernels and math algorithms bring local
+training and inference to AMD hardware. The 23 July update builds on the
+[AMD release](https://github.com/unslothai/unsloth/releases/tag/v0.1.501-beta):
+
+- RDNA2 and Gorgon Halo are supported, and the installer no longer fails to
+ detect GPUs on Strix Halo and other AMD cards.
+- RDNA4 handling is better, and HIP and ROCm failures are caught and fixed
+ automatically instead of stopping the install.
+- Unified memory safetensors loading is 2x faster, with much faster gradient
+ checkpointing on unified memory devices.
+- Voice dictation through whisper.cpp has preliminary support.
+- Rollback environments left by installs no longer eat 5GB of disk. They are
+ cleaned up automatically.
+
+Optimized ROCm builds cover GGUF and safetensors inference, and ROCm
+compatibility is improved for MI300X and MI325X. Full guide:
+[unsloth.ai/docs/basics/amd](https://unsloth.ai/docs/basics/amd).
+
+### Running larger models
+
+- Automatic GPU placement, or pick exactly which GPUs and layers to use.
+- Move MoE expert layers into system memory so larger models fit.
+- Split a model across several GPUs, or use tensor parallelism.
+- Hardware settings are saved per model and quant.
+
+### Also in this release
+
+- Remote access with `unsloth studio --secure` over free HTTPS via Cloudflare.
+- Web search reads PDF papers and manuals, and parallel tool calls, reasoning
+ output and tool retries are more reliable.
+- The model download location is configurable, so weights can live on a second
+ drive instead of the default cache.
+- Stalled Hugging Face XET downloads retry over standard HTTP, and existing
+ GGUF files are reused instead of downloaded again.
diff --git a/MANIFEST.in b/MANIFEST.in
new file mode 100644
index 0000000000..7bce036343
--- /dev/null
+++ b/MANIFEST.in
@@ -0,0 +1,2 @@
+include _changelog_build.py
+include CHANGELOG.md
diff --git a/README.md b/README.md
index 6aa8f4f4c3..e0fc8ee44c 100644
--- a/README.md
+++ b/README.md
@@ -86,6 +86,13 @@ Replace `claude` with any supported agent:
| OpenCode | `unsloth start opencode` |
| Pi Coding Agent | `unsloth start pi` |
+Claude Code, Codex, OpenCode and Pi can keep their current model and use Unsloth as a local
+subagent:
+
+```bash
+unsloth start claude --as-subagent --model unsloth/model-GGUF:quant
+```
+
## 📥 Install
Unsloth can be used in two ways: through **[Unsloth Studio](https://unsloth.ai/docs/new/studio/)**, the web UI, or through **Unsloth Core**, the code-based version. Each has different requirements.
@@ -96,7 +103,7 @@ Unsloth Studio (Beta) works on **Windows, Linux, WSL** and **macOS**.
* **NVIDIA:** Training works on RTX 30/40/50, Blackwell, DGX Spark, Station and more
* **macOS:** Training, MLX and GGUF inference are ALL supported.
* **AMD:** Training, RL, chat and deployment work on Windows, WSL and Linux. [Read the AMD guide](https://unsloth.ai/docs/basics/amd).
-* **Vulkan:** GGUF inference is supported on [compatible GPUs, including Intel GPUs](https://github.com/unslothai/unsloth/pull/5819).
+* **Vulkan:** GGUF inference is supported on [compatible GPUs, including Intel GPUs](https://github.com/unslothai/unsloth/pull/5819). Vulkan accelerates GGUF inference only; training still requires a supported PyTorch or MLX backend.
* **Multi-GPU:** Available now, with a major upgrade on the way
#### macOS, Linux, WSL:
@@ -105,12 +112,28 @@ curl -fsSL https://unsloth.ai/install.sh | sh
```
Use the same command to update.
+To force the Vulkan llama.cpp backend, set `UNSLOTH_FORCE_VULKAN=1` **before installing or updating**. The setting selects the llama.cpp binary bundle, so setting it only when launching Studio cannot replace an existing CPU bundle:
+
+```bash
+export UNSLOTH_FORCE_VULKAN=1
+curl -fsSL https://unsloth.ai/install.sh | sh
+```
+
#### Windows:
```powershell
irm https://unsloth.ai/install.ps1 | iex
```
Use the same command to update.
+To force the Vulkan llama.cpp backend, set the environment variable before running the installer or updater:
+
+```powershell
+$env:UNSLOTH_FORCE_VULKAN=1
+irm https://unsloth.ai/install.ps1 | iex
+```
+
+Re-running the current installer replaces a previously selected CPU bundle when the backend differs. A separate Vulkan SDK is not required; the GPU driver must provide a working Vulkan runtime.
+
#### Launch
```bash
unsloth studio -p 8888
@@ -256,6 +279,8 @@ 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.
+On a wildcard bind Unsloth works out the address to share by asking `ifconfig.me` for the public IP, then asks `check-host.net` whether that port is reachable so it can tell you if a firewall is in the way. Both contact a third party. Set `UNSLOTH_STUDIO_DISABLE_PUBLIC_CHECK=1` to skip them; the banner then shows the LAN address and no reachability line.
+
The first time Unsloth is published on a public URL (`--secure` or `--cloudflare`) with the auto-generated admin password still in place, it asks for a new admin password in the terminal (masked input with confirmation) before the public link goes up. Without an attached terminal it warns instead and keeps the bootstrap deadline: Unsloth shuts down after `UNSLOTH_STUDIO_BOOTSTRAP_TIMEOUT` (default 1 hour) unless the password is changed in the web UI.
For headless setups that cannot answer that prompt, set the initial admin password non-interactively with `--password` (only takes effect when no password is set yet; if one already exists it is a hard error, so rotate later with `unsloth studio reset-password`):
diff --git a/_changelog_build.py b/_changelog_build.py
new file mode 100644
index 0000000000..f5bcf2052c
--- /dev/null
+++ b/_changelog_build.py
@@ -0,0 +1,36 @@
+# SPDX-License-Identifier: AGPL-3.0-only
+# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved.
+
+"""Snapshot CHANGELOG.md into the studio package at build time.
+
+CHANGELOG.md at the repo root stays the one file to edit. Copying it here,
+rather than in build.sh, means every packaging path ships it, so release notes
+still render when the popup cannot reach GitHub."""
+
+from __future__ import annotations
+
+import shutil
+from pathlib import Path
+
+from setuptools.command.build_py import build_py as _build_py
+
+ROOT = Path(__file__).resolve().parent
+SOURCE = ROOT / "CHANGELOG.md"
+SNAPSHOT = ROOT / "studio" / "CHANGELOG.md"
+
+
+class build_py(_build_py):
+ def run(self) -> None:
+ # Beside the sources only if writable (PEP 517 may build an immutable
+ # checkout); into the staging directory always.
+ if SOURCE.is_file():
+ try:
+ shutil.copyfile(SOURCE, SNAPSHOT)
+ except OSError:
+ pass
+ super().run()
+ if not SOURCE.is_file():
+ return
+ staged = Path(self.build_lib) / "studio" / "CHANGELOG.md"
+ staged.parent.mkdir(parents = True, exist_ok = True)
+ shutil.copyfile(SOURCE, staged)
diff --git a/build.sh b/build.sh
index 2a836e19d9..5b09a7791b 100644
--- a/build.sh
+++ b/build.sh
@@ -103,9 +103,13 @@ else
STUDIO_STAMPED_VERSION="$(python scripts/stamp_studio_release.py)"
fi
-# 4. Build wheel/sdist
+# 4. Build wheel/sdist. _changelog_build.py snapshots CHANGELOG.md into the studio
+# package so release notes render offline.
python -m build
+# Drop the snapshot so a source checkout never serves a stale copy.
+rm -f studio/CHANGELOG.md
+
if [ "${1:-}" = "publish" ]; then
python scripts/stamp_studio_release.py --verify-dist dist --expected "$STUDIO_STAMPED_VERSION"
fi
diff --git a/install.ps1 b/install.ps1
index cceefd2647..5b205df96d 100644
--- a/install.ps1
+++ b/install.ps1
@@ -28,6 +28,14 @@ function Install-UnslothStudio {
}
}
+ function Clear-TauriInstallError {
+ param([string]$Message)
+ if ($TauriMode) {
+ Write-TauriLog "ERROR_CLEAR" $Message
+ [Console]::Error.WriteLine("[TAURI:ERROR_CLEAR] $Message")
+ }
+ }
+
function Format-TauriDiagBool {
param([bool]$Value)
if ($Value) { return "true" }
@@ -49,6 +57,26 @@ function Install-UnslothStudio {
}
}
+ # Machine arch; Get-TauriDiagArch above reports the process. An emulated x64 shell on
+ # ARM64 reports AMD64, but PROCESSOR_ARCHITEW6432 is ARM64 in exactly that case.
+ function Get-HostMachineArch {
+ $osArch = ""
+ try { $osArch = [System.Runtime.InteropServices.RuntimeInformation]::OSArchitecture.ToString() } catch { $osArch = "" }
+ $signals = @([string]$env:PROCESSOR_ARCHITEW6432, [string]$env:PROCESSOR_ARCHITECTURE, $osArch)
+ foreach ($s in $signals) {
+ if ($s.ToLowerInvariant() -eq "arm64") { return "arm64" }
+ }
+ foreach ($s in $signals) {
+ if ([string]::IsNullOrWhiteSpace($s)) { continue }
+ switch ($s.ToLowerInvariant()) {
+ "amd64" { return "x86_64" }
+ "x64" { return "x86_64" }
+ "x86" { return "x86" }
+ }
+ }
+ return "unknown"
+ }
+
function Get-TauriTorchIndexFamily {
param([string]$TorchIndexUrl)
if ($SkipTorch) { return "none" }
@@ -86,7 +114,7 @@ function Install-UnslothStudio {
[int]$Code = 1
)
if ($Code -eq 0) { $Code = 1 }
- Write-TauriLog "ERROR" $Message
+ Write-TauriLog "ERROR_DEFAULT" $Message
if (Get-Command Restore-StudioVenvRollback -CommandType Function -ErrorAction SilentlyContinue) {
Restore-StudioVenvRollback
}
@@ -485,7 +513,8 @@ function Install-UnslothStudio {
# Full command output is shown only when --verbose / UNSLOTH_VERBOSE=1.
function Invoke-InstallCommand {
param(
- [Parameter(Mandatory = $true)][ScriptBlock]$Command
+ [Parameter(Mandatory = $true)][ScriptBlock]$Command,
+ [string]$Label = "install command"
)
# Installer-pinned index installs (torch) must beat an inherited uv mirror (#6898):
# for --default-index, clear the uv index env vars (restore in finally) and set
@@ -504,6 +533,7 @@ function Install-UnslothStudio {
try {
# Reset to avoid stale values from prior native commands.
$global:LASTEXITCODE = 0
+ Write-TauriLog "OUTPUT_CLEAR" $Label
if ($script:UnslothVerbose) {
# Merge stderr into stdout so progress/warning output stays visible
# without flipping $? on successful native commands (PS 5.1 treats
@@ -518,7 +548,13 @@ function Install-UnslothStudio {
Write-Host (Redact-InstallOutput $output) -ForegroundColor Red
}
}
- return [int]$LASTEXITCODE
+ $exitCode = [int]$LASTEXITCODE
+ if ($exitCode -eq 0) {
+ Clear-TauriInstallError "$Label recovered"
+ } else {
+ Write-TauriLog "ERROR_OUTPUT" "$Label failed (exit code $exitCode)"
+ }
+ return $exitCode
} finally {
$ErrorActionPreference = $prevEap
if ($savedUvIndex) {
@@ -549,7 +585,7 @@ function Install-UnslothStudio {
}
$attempt = 1
while ($true) {
- $code = Invoke-InstallCommand $Command
+ $code = Invoke-InstallCommand -Command $Command -Label $Label
if ($code -eq 0) { return 0 }
if ($attempt -ge $maxAttempts) { return $code }
substep ("retrying ""$Label"" after transient failure (attempt $($attempt + 1)/$maxAttempts, waiting ${delay}s)...") "Yellow"
@@ -1108,10 +1144,27 @@ exit 0
return $false
}
+ # The interpreter's own arch, asked of it: win-amd64|win-arm64|win32|"".
+ function Get-PythonPlatformTag {
+ param([string]$Exe)
+ try {
+ return (& $Exe -c "import sysconfig; print(sysconfig.get_platform())" 2>$null | Out-String).Trim().ToLowerInvariant()
+ } catch { return "" }
+ }
+
# Returns @{ Version = "3.13"; Path = "C:\...\python.exe" } or $null.
# The resolved Path is passed to `uv venv --python` to prevent uv from
# re-resolving the version string back to a conda interpreter.
function Find-CompatiblePython {
+ # -X64Only: best installed x64 interpreter or $null, never ARM64. Last resort for
+ # Install-X64Python, where x64 of a lower-priority minor beats ARM64.
+ param([switch]$X64Only)
+ # Windows on ARM: prefer x64. pyarrow (via datasets) and hf-transfer ship no
+ # win_arm64 wheel, so a native ARM64 Python source-builds both and dies on CMake /
+ # Rust minutes in; x64 runs fine emulated. ARM64 is still returned when it is all
+ # there is, and the caller then bootstraps x64 or warns.
+ $preferX64 = $X64Only -or ((Get-HostMachineArch) -eq "arm64")
+ $candidates = @()
# Try the Python Launcher first (most reliable on Windows)
# py.exe resolves to the standard CPython install, not conda.
# Prefer the requested $PythonVersion, then newest-first fallback.
@@ -1129,7 +1182,8 @@ exit 0
# Resolve the actual executable path and verify it is not conda-based
$resolvedExe = (& $pyLauncher.Source "-$minor" -c "import sys; print(sys.executable)" 2>$null | Out-String).Trim()
if ($resolvedExe -and (Test-Path $resolvedExe) -and -not (Test-IsCondaPython $resolvedExe)) {
- return @{ Version = $ver; Path = $resolvedExe }
+ if (-not $preferX64) { return @{ Version = $ver; Path = $resolvedExe; Arch = "" } }
+ $candidates += @{ Version = $ver; Path = $resolvedExe }
}
}
} catch {}
@@ -1150,11 +1204,53 @@ exit 0
try {
$out = & $cmd.Source --version 2>&1 | Out-String
if ($out -match "Python (3\.1[1-3])\.\d+") {
- return @{ Version = $Matches[1]; Path = $cmd.Source }
+ if (-not $preferX64) { return @{ Version = $Matches[1]; Path = $cmd.Source; Arch = "" } }
+ $candidates += @{ Version = $Matches[1]; Path = $cmd.Source }
}
} catch {}
}
}
+ # `py -3.12` runs the launcher's preferred build, normally the native ARM64 one, so
+ # a same-minor x64 install that is neither preferred nor on PATH never becomes a
+ # candidate. `-3.12-64` cannot disambiguate (deprecated, it only means "not
+ # 32-bit"), so enumerate every registration with -0p and probe each path.
+ if ($preferX64) {
+ foreach ($pyLauncher in @(Get-Command py -All -CommandType Application -ErrorAction SilentlyContinue)) {
+ if ($pyLauncher.Source -match $script:CondaSkipPattern) { continue }
+ $listed = @()
+ try { $listed = @(& $pyLauncher.Source "-0p" 2>$null) } catch {}
+ foreach ($line in $listed) {
+ # " -V:3.12 * C:\...\python.exe": tag, optional default marker, path.
+ $m = [regex]::Match([string]$line, '(?i)^\s*-\S+\s+\*?\s*"?(?
\S.*?\.exe)"?\s*$')
+ if (-not $m.Success) { continue }
+ $exe = $m.Groups['p'].Value.Trim()
+ if ($candidates | Where-Object { $_.Path -eq $exe }) { continue }
+ if (-not (Test-Path -LiteralPath $exe)) { continue }
+ if (Test-IsCondaPython $exe) { continue }
+ try {
+ $out = & $exe --version 2>&1 | Out-String
+ if ($out -match "Python (3\.1[1-3])\.\d+") {
+ $candidates += @{ Version = $Matches[1]; Path = $exe }
+ }
+ } catch {}
+ }
+ }
+ }
+ # Prefer x64, but only within one minor: $minors is the caller's version preference,
+ # so ranking on arch alone would answer UNSLOTH_PYTHON=3.12 with an x64 3.13 and
+ # never bootstrap x64 3.12. Probing costs a subprocess, so non-ARM returned above.
+ foreach ($c in $candidates) {
+ $tag = Get-PythonPlatformTag $c.Path
+ $c.Arch = if ($tag -eq "win-amd64") { "x86_64" } elseif ($tag -eq "win-arm64") { "arm64" } else { "unknown" }
+ }
+ foreach ($minor in $minors) {
+ $sameMinor = @($candidates | Where-Object { $_.Version -eq $minor })
+ if ($sameMinor.Count -eq 0) { continue }
+ $x64 = $sameMinor | Where-Object { $_.Arch -eq "x86_64" } | Select-Object -First 1
+ if ($x64) { return $x64 }
+ if (-not $X64Only) { return $sameMinor[0] }
+ }
+ if (-not $X64Only -and $candidates.Count -gt 0) { return $candidates[0] }
return $null
}
@@ -1165,8 +1261,11 @@ exit 0
# (no UAC), putting python.exe + the py launcher on PATH. Mirrors the uv ->
# astral.sh fallback below. Returns @{ Version; Path } or $null.
function Install-PythonFromPythonOrg {
+ # $Arch overrides the host arch, to pull x64 onto an ARM64 box.
+ param([string]$Arch = "")
# python.org ships one installer per architecture.
- $archSuffix = switch (Get-TauriDiagArch) {
+ $targetArch = if ($Arch) { $Arch } else { Get-TauriDiagArch }
+ $archSuffix = switch ($targetArch) {
"x86_64" { "-amd64" }
"arm64" { "-arm64" }
"x86" { "" }
@@ -1231,6 +1330,28 @@ exit 0
return (Find-CompatiblePython)
}
+ # ── Windows on ARM: get an x64 CPython ──
+ # --architecture x64 forces winget off the ARM64 build; python.org takes the same override.
+ function Install-X64Python {
+ if ($script:WingetAvailable) {
+ $prevEAP = $ErrorActionPreference
+ $ErrorActionPreference = "Continue"
+ try {
+ winget install -e --id "Python.Python.$PythonVersion" --source winget --architecture x64 --accept-package-agreements --accept-source-agreements
+ } catch { }
+ $ErrorActionPreference = $prevEAP
+ Refresh-SessionPath
+ $found = Find-CompatiblePython
+ if ($found -and $found.Arch -eq "x86_64") { return $found }
+ substep "winget could not provide an x64 Python -- trying python.org..." "Yellow"
+ }
+ $found = Install-PythonFromPythonOrg -Arch "x86_64"
+ if ($found -and $found.Arch -eq "x86_64") { return $found }
+ # Nothing installable (offline / no winget): an x64 build of another supported minor
+ # still runs the wheels ARM64 cannot, so take it over the native interpreter.
+ return (Find-CompatiblePython -X64Only)
+ }
+
# ── Install Python if no compatible version (3.11-3.13) found ──
# Find-CompatiblePython returns @{ Version = "3.13"; Path = "C:\...\python.exe" } or $null.
Write-TauriLog "STEP" "Installing Python"
@@ -1302,6 +1423,26 @@ exit 0
return (Exit-InstallFailure "Python installation failed")
}
}
+ # ── Windows on ARM: swap a native ARM64 interpreter for x64 ──
+ # pyarrow and hf-transfer publish no win_arm64 wheel, so an ARM64 Python source-builds
+ # both and fails deep into the run. Warn up front if x64 is unobtainable.
+ if ($DetectedPython -and (Get-HostMachineArch) -eq "arm64" -and $DetectedPython.Arch -ne "x86_64") {
+ substep "windows on arm: only a native ARM64 Python $($DetectedPython.Version) was found." "Yellow"
+ substep "pyarrow and hf-transfer publish no win_arm64 wheels, so installing x64 Python..." "Yellow"
+ $X64Python = Install-X64Python
+ if ($X64Python) {
+ $DetectedPython = $X64Python
+ step "python" "using x64 Python $($DetectedPython.Version) under emulation"
+ } else {
+ Write-Host "[WARN] Could not install an x64 Python on this ARM64 machine." -ForegroundColor Yellow
+ Write-Host " Continuing with ARM64 Python $($DetectedPython.Version), but the install is likely to fail:" -ForegroundColor Yellow
+ Write-Host " pyarrow (via datasets) and hf-transfer ship no win_arm64 wheels and will be" -ForegroundColor Yellow
+ Write-Host " built from source, which needs CMake plus the MSVC and Rust toolchains." -ForegroundColor Yellow
+ Write-Host " Fix: install x64 Python from https://www.python.org/downloads/windows/" -ForegroundColor Yellow
+ Write-Host " (choose 'Windows installer (64-bit)', not ARM64), then re-run this installer." -ForegroundColor Yellow
+ }
+ }
+
$DiagPythonVersion = $PythonVersion
if ($DetectedPython) { $DiagPythonVersion = $DetectedPython.Version }
$InitialGpuBranch = "unknown"
@@ -1416,13 +1557,82 @@ exit 0
$suffix++
$candidate = Join-Path $StudioHome "unsloth_studio.rollback.$stamp.$PID.$suffix"
}
- Move-Item -LiteralPath $ExistingDir -Destination $candidate -ErrorAction Stop
$script:StudioVenvRollbackDir = $candidate
$script:StudioVenvRollbackTarget = $ExistingDir
$script:StudioVenvRollbackActive = $true
+ # Publish the rollback state before the atomic rename so interruption
+ # cannot land after Move-Item but before cleanup knows where the old venv went.
+ try {
+ Move-Item -LiteralPath $ExistingDir -Destination $candidate -ErrorAction Stop
+ } catch {
+ # A collision or ordinary rename failure leaves the original in place.
+ # Keep state active only when the rename happened before interruption.
+ if (Test-Path -LiteralPath $ExistingDir) {
+ $script:StudioVenvRollbackActive = $false
+ $script:StudioVenvRollbackDir = $null
+ }
+ throw
+ }
substep "previous environment preserved for rollback"
}
+ function Remove-StudioVenvTreeWithRetry {
+ param(
+ [Parameter(Mandatory = $true)][string]$Path,
+ [Parameter(Mandatory = $true)][string]$Label
+ )
+ $lastError = $null
+ for ($attempt = 1; $attempt -le 3; $attempt++) {
+ try {
+ Remove-Item -LiteralPath $Path -Recurse -Force -ErrorAction Stop
+ } catch {
+ $lastError = $_.Exception.Message
+ }
+ if (-not (Test-Path -LiteralPath $Path)) { return $true }
+ if ($attempt -lt 3) { Start-Sleep -Milliseconds (250 * $attempt) }
+ }
+ Write-Host "[WARN] Could not remove $Label at $Path" -ForegroundColor Yellow
+ if ($lastError) { Write-Host " $lastError" -ForegroundColor Yellow }
+ return $false
+ }
+
+ function Test-StudioVenvRollbackMustBePreserved {
+ param([Parameter(Mandatory = $true)][System.IO.FileSystemInfo]$Rollback)
+ # Preserve anything outside the installer's timestamp.PID[.suffix] format.
+ if ($Rollback.Name -notmatch '^unsloth_studio\.rollback\.[0-9]{14}\.([0-9]+)(?:\.[0-9]+)?$') {
+ return $true
+ }
+ $ownerPid = 0
+ if (-not [int]::TryParse($Matches[1], [ref]$ownerPid)) { return $true }
+ if ($ownerPid -eq $PID) { return $true }
+ return $null -ne (Get-Process -Id $ownerPid -ErrorAction SilentlyContinue)
+ }
+
+ function Remove-StaleStudioVenvRollbacks {
+ try {
+ $rollbacks = @(
+ Get-ChildItem -LiteralPath $StudioHome -Directory -Force -ErrorAction Stop |
+ Where-Object { $_.Name -like 'unsloth_studio.rollback.*' }
+ )
+ } catch {
+ Write-Host "[WARN] Could not inspect stale environment rollbacks in $StudioHome" -ForegroundColor Yellow
+ Write-Host " $($_.Exception.Message)" -ForegroundColor Yellow
+ return
+ }
+ foreach ($rollback in $rollbacks) {
+ if (($rollback.Attributes -band [System.IO.FileAttributes]::ReparsePoint) -ne 0) {
+ Write-Host "[WARN] Refusing to remove rollback reparse point $($rollback.FullName)" -ForegroundColor Yellow
+ continue
+ }
+ # A concurrent installer may have moved its live venv aside. The PID
+ # in the generated name keeps this run from deleting its rescue copy.
+ if (Test-StudioVenvRollbackMustBePreserved -Rollback $rollback) { continue }
+ if (Remove-StudioVenvTreeWithRetry -Path $rollback.FullName -Label "stale environment rollback") {
+ substep "removed stale environment rollback $($rollback.Name)"
+ }
+ }
+ }
+
function Restore-StudioVenvRollback {
if (-not $script:StudioVenvRollbackActive) { return }
$backup = $script:StudioVenvRollbackDir
@@ -1434,7 +1644,9 @@ exit 0
substep "restoring previous environment after failed install..." "Yellow"
try {
if (Test-Path -LiteralPath $target) {
- Remove-Item -LiteralPath $target -Recurse -Force -ErrorAction SilentlyContinue
+ if (-not (Remove-StudioVenvTreeWithRetry -Path $target -Label "incomplete environment")) {
+ throw "Could not remove incomplete environment at $target"
+ }
}
Move-Item -LiteralPath $backup -Destination $target -Force -ErrorAction Stop
substep "restored previous environment"
@@ -1449,13 +1661,17 @@ exit 0
function Complete-StudioVenvRollback {
if (-not $script:StudioVenvRollbackActive) { return }
$backup = $script:StudioVenvRollbackDir
- if ($backup -and (Test-Path -LiteralPath $backup)) {
- Remove-Item -LiteralPath $backup -Recurse -Force -ErrorAction SilentlyContinue
- }
+ # The replacement is committed. Disable restoration before deleting the
+ # backup so interruption cannot restore a partially deleted environment.
$script:StudioVenvRollbackActive = $false
$script:StudioVenvRollbackDir = $null
+ if ($backup -and (Test-Path -LiteralPath $backup)) {
+ Remove-StudioVenvTreeWithRetry -Path $backup -Label "environment rollback" | Out-Null
+ }
}
+ $studioVenvReplacementCommitted = $false
+ try {
if (Test-Path -LiteralPath $VenvPython) {
# why: matching guard to the .venv branch below -- in env-mode
# $StudioHome is a user-chosen workspace, so refuse to nuke an
@@ -1528,7 +1744,7 @@ exit 0
if (-not (Test-Path -LiteralPath $VenvPython)) {
step "venv" "creating Python $($DetectedPython.Version) virtual environment"
substep "$VenvDir"
- $venvExit = Invoke-InstallCommand { uv venv $VenvDir --python "$($DetectedPython.Path)" }
+ $venvExit = Invoke-InstallCommand -Label "create virtual environment" { uv venv $VenvDir --python "$($DetectedPython.Path)" }
if ($venvExit -ne 0) {
Write-Host "[ERROR] Failed to create virtual environment (exit code $venvExit)" -ForegroundColor Red
return (Exit-InstallFailure "Failed to create virtual environment (exit code $venvExit)" $venvExit)
@@ -1842,12 +2058,14 @@ exit 0
# (gfx120X/110X/1151/1150/103X); unknown names fall back to CPU.
elseif ($ROCmGpuLabel) {
$nameArchTable = @(
- @{ P = "9070 XT|9080"; A = "gfx1201" } # RDNA 4 (RX 9070 XT / 9080)
- @{ P = "9070|9060"; A = "gfx1200" } # RDNA 4 (RX 9070 / 9060)
- @{ P = "8060S|8050S|8040S|Strix Halo|Ryzen AI Max|AI Max"; A = "gfx1151" } # RDNA 3.5 (Strix Halo: Radeon 8060S/8050S/8040S iGPU, Ryzen AI Max+)
- @{ P = "890M|880M|860M|840M|Strix Point|Krackan|HX 37[05]|AI 9 HX|AI 9 36[05]|AI 7 35[05]|AI 5 34[05]|AI 7 PRO 35|AI 5 33"; A = "gfx1150" } # RDNA 3.5 (Strix/Krackan Point: Radeon 890M/880M iGPU, Ryzen AI 9 HX 370/375)
- @{ P = "RX 7900|RX 7800|RX 7700(?!S)|PRO W7900|PRO W7800|PRO W7700"; A = "gfx1100" } # RDNA 3 desktop/workstation (Navi 31)
- @{ P = "RX 7600|RX 7700S|RX 7650|PRO W7600|PRO W7500|PRO V710"; A = "gfx1102" } # RDNA 3 (Navi 33)
+ @{ P = "9070|9080"; A = "gfx1201" } # RDNA 4 (Navi 48: RX 9070 XT / 9070 GRE / 9070 / 9080)
+ @{ P = "9060"; A = "gfx1200" } # RDNA 4 (Navi 44: RX 9060 XT / 9060)
+ @{ P = "8065S|8060S|8050S|8040S|Strix Halo|Ryzen AI Max|AI Max"; A = "gfx1151" } # RDNA 3.5 (Strix Halo + Gorgon Halo: Radeon 8065S/8060S/8050S/8040S iGPU, Ryzen AI Max / Max+)
+ @{ P = "890M|880M|Strix Point|HX 37[05]|AI 9 HX|AI 9 36[05]"; A = "gfx1150" } # RDNA 3.5 (Strix Point: Radeon 890M/880M, Ryzen AI 9 HX 370/375)
+ @{ P = "860M|840M|Krackan|AI 7 35[05]|AI 5 34[05]|AI 7 PRO 35|AI 5 33"; A = "gfx1152" } # RDNA 3.5 (Krackan Point: Radeon 860M/840M, Ryzen AI 7 350 / AI 5 340)
+ @{ P = "RX 7900|PRO W7900|PRO W7800"; A = "gfx1100" } # RDNA 3 desktop/workstation (Navi 31)
+ @{ P = "RX 7800|RX 7700(?!S)|PRO W7700|PRO V710"; A = "gfx1101" } # RDNA 3 (Navi 32)
+ @{ P = "RX 7600|RX 7700S|RX 7650|PRO W7600|PRO W7500"; A = "gfx1102" } # RDNA 3 (Navi 33)
@{ P = "780M|760M|740M|Phoenix|Hawk Point|Z1 Extreme|Z2 Extreme"; A = "gfx1103" } # RDNA 3 iGPU (Phoenix / Hawk Point)
@{ P = "RX 6900|RX 6800|RX 6750|RX 6700|PRO W6800|PRO W6900"; A = "gfx1030" } # RDNA 2 (Navi 21) -- gfx103X family
@{ P = "RX 6650|RX 6600|PRO W6600|PRO W6650"; A = "gfx1032" } # RDNA 2 (Navi 23) -- gfx103X family
@@ -2030,16 +2248,18 @@ exit 0
# _strip_index_url_credentials (install.sh / py / setup.ps1).
function Remove-IndexUrlCredentials {
param([string]$Url)
- $sep = $Url.IndexOf('://')
+ # Ordinal, not culture-aware: on non-English locales (e.g. th-TH) linguistic
+ # IndexOf treats "://" as ignorable, mis-locates it, and crashes Substring (issue #7279).
+ $sep = $Url.IndexOf('://', [System.StringComparison]::Ordinal)
if ($sep -lt 0) { return $Url }
$scheme = $Url.Substring(0, $sep)
$rest = $Url.Substring($sep + 3)
# Drop query / fragment (may hold auth tokens).
$q = $rest.IndexOfAny([char[]]('?', '#'))
if ($q -ge 0) { $rest = $rest.Substring(0, $q) }
- $slash = $rest.IndexOf('/')
+ $slash = $rest.IndexOf('/', [System.StringComparison]::Ordinal)
$authority = if ($slash -ge 0) { $rest.Substring(0, $slash) } else { $rest }
- $at = $authority.LastIndexOf('@')
+ $at = $authority.LastIndexOf('@', [System.StringComparison]::Ordinal)
$host_ = if ($at -ge 0) { $authority.Substring($at + 1) } else { $authority }
if ($slash -ge 0) { return "${scheme}://${host_}$($rest.Substring($slash))" }
return "${scheme}://${host_}"
@@ -2126,6 +2346,7 @@ exit 0
$archFamilyMap = @{
"gfx1201" = "gfx120X-all"; "gfx1200" = "gfx120X-all" # RDNA 4
"gfx1151" = "gfx1151"; "gfx1150" = "gfx1150" # RDNA 3.5 (Strix Halo/Point)
+ "gfx1152" = "gfx1152" # RDNA 3.5 (Krackan Point)
"gfx1103" = "gfx110X-all"; "gfx1102" = "gfx110X-all" # RDNA 3
"gfx1101" = "gfx110X-all"; "gfx1100" = "gfx110X-all"
"gfx1036" = "gfx103X-all"; "gfx1035" = "gfx103X-all" # RDNA 2 (RX 6000)
@@ -2147,6 +2368,7 @@ exit 0
$torchFloorMap = @{
"gfx1201" = "torch>=2.11.0,<2.12.0"; "gfx1200" = "torch>=2.11.0,<2.12.0"
"gfx1151" = "torch>=2.11.0,<2.12.0"; "gfx1150" = "torch>=2.11.0,<2.12.0"
+ "gfx1152" = "torch>=2.11.0,<2.12.0"
}
# Companion ranges track the torch ceiling so pip resolves a consistent
# trio on AMD's per-arch index (each published independently). Mirrors
@@ -2154,10 +2376,12 @@ exit 0
$torchvisionFloorMap = @{
"gfx1201" = "torchvision>=0.26.0,<0.27.0"; "gfx1200" = "torchvision>=0.26.0,<0.27.0"
"gfx1151" = "torchvision>=0.26.0,<0.27.0"; "gfx1150" = "torchvision>=0.26.0,<0.27.0"
+ "gfx1152" = "torchvision>=0.26.0,<0.27.0"
}
$torchaudioFloorMap = @{
"gfx1201" = "torchaudio>=2.11.0,<2.12.0"; "gfx1200" = "torchaudio>=2.11.0,<2.12.0"
"gfx1151" = "torchaudio>=2.11.0,<2.12.0"; "gfx1150" = "torchaudio>=2.11.0,<2.12.0"
+ "gfx1152" = "torchaudio>=2.11.0,<2.12.0"
}
$archFamily = if ($ROCmGfxArch -and $archFamilyMap.ContainsKey($ROCmGfxArch)) { $archFamilyMap[$ROCmGfxArch] } else { $null }
if ($archFamily) {
@@ -2187,7 +2411,7 @@ exit 0
$_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
+ $_pinGfx211 = @('gfx120x-all', 'gfx1151', 'gfx1150', 'gfx1152') -contains $_pinLeaf
if ($_pinGfx211 -or $_pinRocm211) {
$ROCmIndexUrl = $TorchIndexUrl
$ROCmTorchFloor = "torch>=2.11.0,<2.12.0"
@@ -2270,7 +2494,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 (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" }
+ $baseInstallExit = Invoke-InstallCommandRetry -Label "install unsloth (migrated no-torch)" { uv pip install --python $VenvPython --no-deps --reinstall-package unsloth --reinstall-package unsloth-zoo "unsloth>=2026.7.5" "unsloth-zoo>=2026.7.6" }
if ($baseInstallExit -eq 0) {
# Resolve pydantic WITH deps so pip pins pydantic-core
# to the matching version (no-torch-runtime.txt below
@@ -2284,7 +2508,7 @@ exit 0
}
}
} else {
- $baseInstallExit = Invoke-InstallCommandRetry -Label "install unsloth (migrated)" { uv pip install --python $VenvPython --reinstall-package unsloth --reinstall-package unsloth-zoo "unsloth>=2026.7.4" "unsloth-zoo>=2026.7.4" }
+ $baseInstallExit = Invoke-InstallCommandRetry -Label "install unsloth (migrated)" { uv pip install --python $VenvPython --reinstall-package unsloth --reinstall-package unsloth-zoo "unsloth>=2026.7.5" "unsloth-zoo>=2026.7.6" }
}
if ($baseInstallExit -ne 0) {
Write-Host "[ERROR] Failed to install unsloth (exit code $baseInstallExit)" -ForegroundColor Red
@@ -2292,7 +2516,7 @@ exit 0
}
if ($StudioLocalInstall) {
substep "overlaying local repo (editable)..."
- $overlayExit = Invoke-InstallCommand { uv pip install --python $VenvPython -e $RepoRoot --no-deps }
+ $overlayExit = Invoke-InstallCommand -Label "overlay local repo" { uv pip install --python $VenvPython -e $RepoRoot --no-deps }
if ($overlayExit -ne 0) {
Write-Host "[ERROR] Failed to overlay local repo (exit code $overlayExit)" -ForegroundColor Red
return (Exit-InstallFailure "Failed to overlay local repo (exit code $overlayExit)" $overlayExit)
@@ -2339,6 +2563,13 @@ exit 0
}
} else {
Write-TauriLog "STEP" "Installing PyTorch"
+ # Windows on ARM lacks only torchaudio (whl/cpu win_arm64: torch 42,
+ # torchvision 60, torchaudio 0), so drop that pin instead of aborting. Ask the
+ # interpreter, not PROCESSOR_ARCHITECTURE; reached when no x64 Python exists.
+ $VenvPlatform = ""
+ try {
+ $VenvPlatform = (& $VenvPython -c "import sysconfig; print(sysconfig.get_platform())" 2>$null | Out-String).Trim().ToLowerInvariant()
+ } catch { $VenvPlatform = "" }
substep "installing PyTorch ($(Remove-IndexUrlCredentials $TorchIndexUrl))..."
# Bound the companions to the capped torch on EVERY index, cu
# families included: torchaudio 2.11 dropped its exact torch pin from
@@ -2346,7 +2577,13 @@ exit 0
# 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 }
+ $_torchSpecs = @("torch>=2.4,<2.11.0", $_pinVisionSpec, $_pinAudioSpec)
+ if ($VenvPlatform -eq "win-arm64") {
+ substep "windows on arm: skipping torchaudio (upstream publishes no"
+ substep "win_arm64 wheel); torch and torchvision install normally."
+ $_torchSpecs = @("torch>=2.4,<2.11.0", $_pinVisionSpec)
+ }
+ $torchInstallExit = Invoke-InstallCommandRetry -Label "install PyTorch" { uv pip install --python $VenvPython @_torchSpecs --default-index $TorchIndexUrl }
if ($torchInstallExit -ne 0) {
Write-Host "[ERROR] Failed to install PyTorch (exit code $torchInstallExit)" -ForegroundColor Red
return (Exit-InstallFailure "Failed to install PyTorch (exit code $torchInstallExit)" $torchInstallExit)
@@ -2358,7 +2595,7 @@ exit 0
if ($SkipTorch) {
# No-torch: install unsloth + unsloth-zoo with --no-deps, then
# runtime deps (typer, safetensors, transformers, etc.) with --no-deps.
- $baseInstallExit = Invoke-InstallCommandRetry -Label "install unsloth (no-torch)" { uv pip install --python $VenvPython --no-deps --upgrade-package unsloth --upgrade-package unsloth-zoo "unsloth>=2026.7.4" "unsloth-zoo>=2026.7.4" }
+ $baseInstallExit = Invoke-InstallCommandRetry -Label "install unsloth (no-torch)" { uv pip install --python $VenvPython --no-deps --upgrade-package unsloth --upgrade-package unsloth-zoo "unsloth>=2026.7.5" "unsloth-zoo>=2026.7.6" }
if ($baseInstallExit -eq 0) {
# Same pydantic-with-deps trick as the migrated branch.
$baseInstallExit = Invoke-InstallCommandRetry -Label "install pydantic" { uv pip install --python $VenvPython pydantic }
@@ -2370,7 +2607,7 @@ exit 0
}
}
} elseif ($StudioLocalInstall) {
- $baseInstallExit = Invoke-InstallCommandRetry -Label "install unsloth (local)" { uv pip install --python $VenvPython --upgrade-package unsloth "unsloth>=2026.7.4" "unsloth-zoo>=2026.7.4" }
+ $baseInstallExit = Invoke-InstallCommandRetry -Label "install unsloth (local)" { uv pip install --python $VenvPython --upgrade-package unsloth "unsloth>=2026.7.5" "unsloth-zoo>=2026.7.6" }
} else {
$baseInstallExit = Invoke-InstallCommandRetry -Label "install unsloth" { uv pip install --python $VenvPython --upgrade-package unsloth -- "$PackageName" }
}
@@ -2381,7 +2618,7 @@ exit 0
if ($StudioLocalInstall) {
substep "overlaying local repo (editable)..."
- $overlayExit = Invoke-InstallCommand { uv pip install --python $VenvPython -e $RepoRoot --no-deps }
+ $overlayExit = Invoke-InstallCommand -Label "overlay local repo" { uv pip install --python $VenvPython -e $RepoRoot --no-deps }
if ($overlayExit -ne 0) {
Write-Host "[ERROR] Failed to overlay local repo (exit code $overlayExit)" -ForegroundColor Red
return (Exit-InstallFailure "Failed to overlay local repo (exit code $overlayExit)" $overlayExit)
@@ -2398,13 +2635,13 @@ exit 0
Write-TauriLog "STEP" "Installing unsloth"
substep "installing unsloth (this may take a few minutes)..."
if ($StudioLocalInstall) {
- $baseInstallExit = Invoke-InstallCommandRetry -Label "install unsloth (auto torch backend)" { uv pip install --python $VenvPython "unsloth-zoo>=2026.7.4" "unsloth>=2026.7.4" --torch-backend=auto }
+ $baseInstallExit = Invoke-InstallCommandRetry -Label "install unsloth (auto torch backend)" { uv pip install --python $VenvPython "unsloth-zoo>=2026.7.6" "unsloth>=2026.7.5" --torch-backend=auto }
if ($baseInstallExit -ne 0) {
Write-Host "[ERROR] Failed to install unsloth (exit code $baseInstallExit)" -ForegroundColor Red
return (Exit-InstallFailure "Failed to install unsloth (exit code $baseInstallExit)" $baseInstallExit)
}
substep "overlaying local repo (editable)..."
- $overlayExit = Invoke-InstallCommand { uv pip install --python $VenvPython -e $RepoRoot --no-deps }
+ $overlayExit = Invoke-InstallCommand -Label "overlay local repo" { uv pip install --python $VenvPython -e $RepoRoot --no-deps }
if ($overlayExit -ne 0) {
Write-Host "[ERROR] Failed to overlay local repo (exit code $overlayExit)" -ForegroundColor Red
return (Exit-InstallFailure "Failed to overlay local repo (exit code $overlayExit)" $overlayExit)
@@ -2452,7 +2689,7 @@ exit 0
$visionSpec = if ($PinnedRocmVisionSpec) { $PinnedRocmVisionSpec } elseif ($ROCmGfxArch -and $torchvisionFloorMap -and $torchvisionFloorMap.ContainsKey($ROCmGfxArch)) { $torchvisionFloorMap[$ROCmGfxArch] } else { "torchvision" }
$audioSpec = if ($PinnedRocmAudioSpec) { $PinnedRocmAudioSpec } elseif ($ROCmGfxArch -and $torchaudioFloorMap -and $torchaudioFloorMap.ContainsKey($ROCmGfxArch)) { $torchaudioFloorMap[$ROCmGfxArch] } else { "torchaudio" }
substep "PyTorch flavor mismatch (installed $installedTorchTag, need ROCm) -- reinstalling correct build..." "Yellow"
- $torchFixExit = Invoke-InstallCommand { uv pip install --python $VenvPython --force-reinstall --default-index $ROCmIndexUrl $rocmSpec $visionSpec $audioSpec }
+ $torchFixExit = Invoke-InstallCommand -Label "reinstall PyTorch (ROCm)" { uv pip install --python $VenvPython --force-reinstall --default-index $ROCmIndexUrl $rocmSpec $visionSpec $audioSpec }
if ($torchFixExit -ne 0) {
Write-Host "[ERROR] Failed to reinstall PyTorch with the correct ROCm build (exit code $torchFixExit)" -ForegroundColor Red
return (Exit-InstallFailure "Failed to reinstall PyTorch (ROCm) (exit code $torchFixExit)" $torchFixExit)
@@ -2461,7 +2698,7 @@ exit 0
} elseif ($expectedTorchTag -ne 'rocm') {
# CUDA: stale +cpu (or wrong cuXXX) against a CUDA index -> reinstall triplet.
substep "PyTorch flavor mismatch (installed $installedTorchTag, need $expectedTorchTag) -- reinstalling correct build..." "Yellow"
- $torchFixExit = Invoke-InstallCommand { uv pip install --python $VenvPython "torch>=2.4,<2.11.0" "torchvision>=0.19,<0.26.0" "torchaudio>=2.4,<2.11.0" --default-index $TorchIndexUrl --reinstall-package torch --reinstall-package torchvision --reinstall-package torchaudio }
+ $torchFixExit = Invoke-InstallCommand -Label "reinstall PyTorch ($expectedTorchTag)" { uv pip install --python $VenvPython "torch>=2.4,<2.11.0" "torchvision>=0.19,<0.26.0" "torchaudio>=2.4,<2.11.0" --default-index $TorchIndexUrl --reinstall-package torch --reinstall-package torchvision --reinstall-package torchaudio }
if ($torchFixExit -ne 0) {
Write-Host "[ERROR] Failed to reinstall PyTorch with the correct CUDA build (exit code $torchFixExit)" -ForegroundColor Red
return (Exit-InstallFailure "Failed to reinstall PyTorch ($expectedTorchTag) (exit code $torchFixExit)" $torchFixExit)
@@ -2562,6 +2799,9 @@ exit 0
# an inherited value would put llama.cpp in the wrong place.
$previousUnslothStudioHome = $env:UNSLOTH_STUDIO_HOME
$hadPreviousUnslothStudioHome = ($null -ne $previousUnslothStudioHome)
+ $previousTauriMode = $env:UNSLOTH_TAURI_MODE
+ $hadPreviousTauriMode = ($null -ne $previousTauriMode)
+ $env:UNSLOTH_TAURI_MODE = if ($TauriMode) { "1" } else { "0" }
if ($StudioRedirectMode -eq 'env') {
$env:UNSLOTH_STUDIO_HOME = $StudioHome
} else {
@@ -2591,14 +2831,22 @@ exit 0
} else {
Remove-Item Env:UNSLOTH_STUDIO_HOME -ErrorAction SilentlyContinue
}
+ if ($hadPreviousTauriMode) {
+ $env:UNSLOTH_TAURI_MODE = $previousTauriMode
+ } else {
+ Remove-Item Env:UNSLOTH_TAURI_MODE -ErrorAction SilentlyContinue
+ }
Remove-Item Env:UNSLOTH_LOCAL_LLAMA_CPP_DIR -ErrorAction SilentlyContinue
Remove-Item Env:UNSLOTH_INSTALL_ROLLBACK_MANAGED -ErrorAction SilentlyContinue
Remove-Item Env:UNSLOTH_SETUP_PYTHON -ErrorAction SilentlyContinue
}
if ($setupExit -ne 0) {
- Write-Host "[ERROR] unsloth studio setup failed (exit code $setupExit)" -ForegroundColor Red
+ if (-not $TauriMode) {
+ Write-Host "[ERROR] unsloth studio setup failed (exit code $setupExit)" -ForegroundColor Red
+ }
return (Exit-InstallFailure "unsloth studio setup failed (exit code $setupExit)" $setupExit)
}
+ Clear-TauriInstallError "studio setup completed"
# ── Expose `unsloth` via a shim dir containing only unsloth.exe ──
# We do NOT add the venv Scripts dir to PATH (it also holds python.exe
@@ -2686,6 +2934,13 @@ exit 0
}
Refresh-SessionPath # sync current session with registry
Complete-StudioVenvRollback
+ $studioVenvReplacementCommitted = $true
+ Remove-StaleStudioVenvRollbacks
+ } finally {
+ if (-not $studioVenvReplacementCommitted) {
+ Restore-StudioVenvRollback
+ }
+ }
# Env-mode session export AFTER Refresh-SessionPath; otherwise a legacy
# User PATH entry (Machine > User > current $env:Path) would win.
diff --git a/install.sh b/install.sh
index 445bab616a..166beeb52c 100755
--- a/install.sh
+++ b/install.sh
@@ -19,6 +19,17 @@
# SPDX-License-Identifier: AGPL-3.0-only
# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0
set -e
+# ── Why the installer lives in a function ──
+# Under `curl ... | sh`, sh is the pipe READER. This file is ~150KB, so a top-level
+# `exit` left most of it unread, the write end failed, and curl tacked
+# "(56) Failure writing output to destination" onto our own error message. Wrapping
+# the body forces sh to parse to the closing brace first, so the pipe always drains
+# (install.ps1 has always had this shape).
+#
+# Body is deliberately NOT reindented: reflowing 4000+ lines would bury the change,
+# and `exit` still exits the shell from inside a function. Do not add
+# `exec < /dev/null`: for a piped shell that closes the script's own source.
+_unsloth_main() {
# ── Output style (aligned with studio/setup.sh) ──
RULE=""
@@ -207,18 +218,37 @@ run_install_cmd() {
# 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
+ tauri_stream_log stdout "OUTPUT_CLEAR" "$_label"
+ {
+ if "$@" 2>&1; then
+ _cmd_rc=0
+ else
+ _cmd_rc=$?
+ fi
+ printf '%s' "$_cmd_rc" > "$_rcf"
+ } | _redact_install_output
_rc=$(cat "$_rcf" 2>/dev/null || echo 1)
rm -f "$_rcf"
- [ "${_rc:-1}" -eq 0 ] 2>/dev/null && return 0
+ _rc=${_rc:-1}
+ if [ "$_rc" -eq 0 ] 2>/dev/null; then
+ tauri_clear_install_error "$_label recovered"
+ return 0
+ fi
+ tauri_stream_log stdout "ERROR_OUTPUT" "$_label failed (exit code $_rc)"
step "error" "$_label failed (exit code $_rc)" "$C_ERR" >&2
return "$_rc"
fi
_log=$(mktemp)
- "$@" >"$_log" 2>&1 && { rm -f "$_log"; return 0; }
+ tauri_stream_log stderr "OUTPUT_CLEAR" "$_label"
+ "$@" >"$_log" 2>&1 && {
+ rm -f "$_log"
+ tauri_clear_install_error "$_label recovered"
+ return 0
+ }
_rc=$?
step "error" "$_label failed (exit code $_rc)" "$C_ERR" >&2
_redact_install_output "$_log" >&2
+ tauri_stream_log stderr "ERROR_OUTPUT" "$_label failed (exit code $_rc)"
rm -f "$_log"
return $_rc
}
@@ -257,10 +287,70 @@ run_install_cmd_retry() {
done
}
-# Install bitsandbytes on AMD ROCm hosts. Uses the continuous-release_main
-# wheel for the ROCm 4-bit GEMV fix (bnb PR #1887, post-0.49.2); bnb <= 0.49.2
-# NaNs at decode shape on every AMD GPU. Falls back to PyPI >=0.49.1 if the
-# pre-release URL is unreachable. Drop the pin once bnb 0.50+ ships on PyPI.
+# True when the runtime target is gfx906 (MI50/Radeon VII): the prebuilt AMD
+# bitsandbytes wheel carries no gfx906 kernels, and force-reinstalling it would
+# clobber a user's source-built bnb (the only 4-bit path on this arch) on every
+# `studio update`. So skip the auto-install and leave whatever bnb is present.
+# _gfx906_target is set during torch-index resolution; also honor an explicit
+# UNSLOTH_ROCM_GFX_ARCH so a pinned-index install still skips. The override is
+# normalized (gfx906:sramecc-:xnack- -> gfx906) so a copied HIP gcnArchName counts.
+_is_gfx906_bnb_skip() {
+ [ "${_gfx906_target:-false}" = true ] && return 0
+ _bnb_gfx_env=$(printf '%s' "${UNSLOTH_ROCM_GFX_ARCH:-}" | tr '[:upper:]' '[:lower:]' | tr -d '[:space:]')
+ _bnb_gfx_env=${_bnb_gfx_env%%:*}
+ [ "$_bnb_gfx_env" = "gfx906" ] && return 0
+ # A pinned index (UNSLOTH_TORCH_INDEX_URL/_FAMILY) skips the reroute block that
+ # sets _gfx906_target, so a real gfx906 host with a pinned rocm6.3 index and no
+ # UNSLOTH_ROCM_GFX_ARCH would otherwise clobber a source-built bnb. Probe here
+ # in that gap; skip only when gfx906 is the SOLE distinct arch (mixed hosts
+ # opt in via the env var, mirroring the reroute block's de-dup rule).
+ if [ -z "$_bnb_gfx_env" ] && [ "${_torch_index_pinned:-false}" = true ]; then
+ _bnb_gfx_probe=$(_probe_amd_gfx_arch | awk 'NF && !seen[$0]++')
+ [ "$_bnb_gfx_probe" = "gfx906" ] && return 0
+ fi
+ return 1
+}
+
+# `pip install unsloth` resolves its unconditional bitsandbytes dep to a generic
+# CUDA wheel (no gfx906 kernels) once we skip the prebuilt one. Snapshot bnb before
+# the unsloth install, then drop a freshly pulled wheel afterwards while leaving a
+# pre-existing source build in place.
+_gfx906_bnb_installed() {
+ "$_VENV_PY" -c "import importlib.util as u, sys; sys.exit(0 if u.find_spec('bitsandbytes') else 1)" >/dev/null 2>&1
+}
+_gfx906_bnb_snapshot() {
+ _gfx906_bnb_absent_before=false
+ _is_gfx906_bnb_skip || return 0
+ _gfx906_bnb_installed || _gfx906_bnb_absent_before=true
+}
+_gfx906_bnb_prune() {
+ _is_gfx906_bnb_skip || return 0
+ [ "${_gfx906_bnb_absent_before:-false}" = true ] || return 0
+ _gfx906_bnb_installed || return 0
+ substep "gfx906: removing generic bitsandbytes pulled in as a dependency (no gfx906 kernels; build from source for 4-bit QLoRA)" "$C_WARN"
+ uv pip uninstall --python "$_VENV_PY" bitsandbytes >/dev/null 2>&1 \
+ || "$_VENV_PY" -m pip uninstall -y bitsandbytes >/dev/null 2>&1 || true
+}
+
+# Install bitsandbytes on AMD ROCm hosts. bnb <= 0.49.2 NaNs at 4-bit decode
+# shape on every AMD GPU; the fix (bnb #1887) ships in continuous-release_main
+# and, on PyPI, first in 0.50.0. Keep this floor in step with the amd extra in
+# pyproject.toml and studio/install_python_stack.py.
+_BNB_ROCM_PYPI_FALLBACK="bitsandbytes>=0.50.0"
+# bitsandbytes ships no ROCm binary in its aarch64 wheel at any version: the PyPI
+# 0.50.0 and continuous-release_main aarch64 wheels both carry only
+# libbitsandbytes_cpu.so plus CUDA variants. So neither install path below gives
+# aarch64 a 4-bit backend, and the messages must not claim one. Cf. gfx906.
+_bnb_rocm_arch_has_binary() {
+ case "$_ARCH" in
+ aarch64|arm64) return 1 ;;
+ *) return 0 ;;
+ esac
+}
+_warn_bnb_no_rocm_binary() {
+ _bnb_rocm_arch_has_binary && return 0
+ substep "[WARN] aarch64: bitsandbytes ships no ROCm kernels on this arch; 4-bit QLoRA needs a source build -- https://docs.unsloth.ai/get-started/install-and-update/amd" "$C_WARN"
+}
_install_bnb_rocm() {
_label="$1"
_venv_py="$2"
@@ -275,9 +365,8 @@ _install_bnb_rocm() {
_bnb_whl_url=""
;;
esac
- # uv rejects the continuous-release_main bitsandbytes wheel because the
- # filename version (1.33.7rc0) does not match the embedded metadata version
- # (0.50.0.dev0). pip accepts the mismatch, so bootstrap pip and use it.
+ # uv rejects the pre-release wheel: filename version (1.33.7rc0) does not
+ # match metadata (0.50.x.dev0). pip accepts it, so bootstrap pip and use it.
if ! "$_venv_py" -m pip --version >/dev/null 2>&1; then
if ! run_maybe_quiet "$_venv_py" -m ensurepip --upgrade; then
run_maybe_quiet uv pip install --python "$_venv_py" pip || \
@@ -293,6 +382,7 @@ _install_bnb_rocm() {
--retries 8 --timeout 90 \
"$_bnb_whl_url" >"$_bnb_log" 2>&1; then
rm -f "$_bnb_log"
+ _warn_bnb_no_rocm_binary
return 0
fi
_bnb_rc=$?
@@ -301,10 +391,17 @@ _install_bnb_rocm() {
fi
rm -f "$_bnb_log"
step "warning" "$_label (pre-release) failed (exit code $_bnb_rc)" "$C_WARN" >&2
- substep "[WARN] bnb pre-release install failed; falling back to PyPI (4-bit decode broken on ROCm)" "$C_WARN"
+ if _bnb_rocm_arch_has_binary; then
+ substep "[WARN] bnb pre-release install failed; falling back to PyPI $_BNB_ROCM_PYPI_FALLBACK, which carries the ROCm 4-bit fix" "$C_WARN"
+ else
+ substep "[WARN] bnb pre-release install failed; falling back to PyPI $_BNB_ROCM_PYPI_FALLBACK" "$C_WARN"
+ fi
fi
run_install_cmd "$_label (pypi fallback)" "$_venv_py" -m pip install \
- --force-reinstall --no-cache-dir --no-deps "bitsandbytes>=0.49.1"
+ --force-reinstall --no-cache-dir --no-deps "$_BNB_ROCM_PYPI_FALLBACK"
+ _bnb_pypi_rc=$?
+ _warn_bnb_no_rocm_binary
+ return $_bnb_pypi_rc
}
if [ "$_next_is_package" = true ]; then
@@ -338,6 +435,34 @@ tauri_log() {
fi
}
+tauri_stream_log() {
+ _tsl_stream="$1"
+ _tsl_tag="$2"
+ shift 2
+ if [ "$TAURI_MODE" = true ]; then
+ if [ "$_tsl_stream" = stderr ]; then
+ printf '[TAURI:%s] %s\n' "$_tsl_tag" "$*" >&2
+ else
+ printf '[TAURI:%s] %s\n' "$_tsl_tag" "$*"
+ fi
+ fi
+}
+
+rollback_substep() {
+ if [ "$TAURI_MODE" = true ]; then
+ tauri_log "PROGRESS" "$1"
+ else
+ substep "$@"
+ fi
+}
+
+tauri_clear_install_error() {
+ if [ "$TAURI_MODE" = true ]; then
+ tauri_log "ERROR_CLEAR" "$1"
+ printf '[TAURI:ERROR_CLEAR] %s\n' "$1" >&2
+ fi
+}
+
tauri_diag_marker() {
_diag_gpu_branch="${1:-unknown}"
_diag_torch_index_family="${2:-none}"
@@ -475,14 +600,20 @@ _start_studio_venv_replacement() {
_stamp=$(date +%Y%m%d%H%M%S 2>/dev/null || echo "time")
_candidate="$STUDIO_HOME/unsloth_studio.rollback.$_stamp.$$"
_suffix=0
- while [ -e "$_candidate" ]; do
+ while [ -e "$_candidate" ] || [ -L "$_candidate" ]; do
_suffix=$((_suffix + 1))
_candidate="$STUDIO_HOME/unsloth_studio.rollback.$_stamp.$$.$_suffix"
done
- mv "$_existing_dir" "$_candidate"
_VENV_ROLLBACK_DIR="$_candidate"
_VENV_ROLLBACK_TARGET="$_existing_dir"
_VENV_ROLLBACK_ACTIVE=true
+ # Publish the rollback state before the atomic rename so a signal cannot
+ # land after mv but before the exit handlers know where the old venv went.
+ if ! mv "$_existing_dir" "$_candidate"; then
+ _VENV_ROLLBACK_ACTIVE=false
+ _VENV_ROLLBACK_DIR=""
+ return 1
+ fi
substep "previous environment preserved for rollback"
}
@@ -492,10 +623,10 @@ _restore_studio_venv_replacement() {
_VENV_ROLLBACK_ACTIVE=false
return 0
}
- substep "restoring previous environment after failed install..." "$C_WARN"
+ rollback_substep "restoring previous environment after failed install..." "$C_WARN"
rm -rf "$_VENV_ROLLBACK_TARGET"
if mv "$_VENV_ROLLBACK_DIR" "$_VENV_ROLLBACK_TARGET"; then
- substep "restored previous environment"
+ rollback_substep "restored previous environment"
_VENV_ROLLBACK_ACTIVE=false
_VENV_ROLLBACK_DIR=""
else
@@ -503,13 +634,68 @@ _restore_studio_venv_replacement() {
fi
}
-_commit_studio_venv_replacement() {
- [ "$_VENV_ROLLBACK_ACTIVE" = true ] || return 0
- if [ -n "$_VENV_ROLLBACK_DIR" ] && [ -d "$_VENV_ROLLBACK_DIR" ]; then
- rm -rf "$_VENV_ROLLBACK_DIR" || true
+_studio_venv_rollback_must_be_preserved() {
+ _rollback_name=${1##*/}
+ _rollback_metadata=${_rollback_name#unsloth_studio.rollback.}
+ _rollback_stamp=${_rollback_metadata%%.*}
+ _rollback_process=${_rollback_metadata#*.}
+ # Preserve anything outside the installer's timestamp.PID[.suffix] format.
+ [ "$_rollback_process" != "$_rollback_metadata" ] || return 0
+ case "$_rollback_stamp" in
+ time) ;;
+ ''|*[!0-9]*) return 0 ;;
+ *) [ "${#_rollback_stamp}" -eq 14 ] || return 0 ;;
+ esac
+ _rollback_pid=${_rollback_process%%.*}
+ case "$_rollback_pid" in
+ ''|*[!0-9]*) return 0 ;;
+ esac
+ _rollback_suffix=${_rollback_process#*.}
+ if [ "$_rollback_suffix" != "$_rollback_process" ]; then
+ case "$_rollback_suffix" in ''|*[!0-9]*) return 0 ;; esac
fi
- _VENV_ROLLBACK_ACTIVE=false
- _VENV_ROLLBACK_DIR=""
+ kill -0 "$_rollback_pid" 2>/dev/null
+}
+
+_prune_stale_studio_venv_rollbacks() {
+ for _stale_rollback in "$STUDIO_HOME"/unsloth_studio.rollback.*; do
+ [ -d "$_stale_rollback" ] || continue
+ if [ -L "$_stale_rollback" ]; then
+ echo "⚠️ Refusing to remove rollback symlink $_stale_rollback" >&2
+ continue
+ fi
+ # A concurrent installer may have moved its live venv aside. The PID in
+ # the generated name keeps this successful run from deleting its rescue copy.
+ _studio_venv_rollback_must_be_preserved "$_stale_rollback" && continue
+ if rm -rf "$_stale_rollback"; then
+ substep "removed stale environment rollback ${_stale_rollback##*/}"
+ else
+ echo "⚠️ Could not remove stale environment rollback $_stale_rollback" >&2
+ fi
+ done
+}
+
+_commit_studio_venv_replacement() {
+ if [ "$_VENV_ROLLBACK_ACTIVE" = true ]; then
+ _rollback_to_remove="$_VENV_ROLLBACK_DIR"
+ # The new environment is already committed. Clear the restore state
+ # before deletion so an interrupt cannot replace it with a half-deleted backup.
+ _VENV_ROLLBACK_ACTIVE=false
+ _VENV_ROLLBACK_DIR=""
+ if [ -n "$_rollback_to_remove" ] && [ -d "$_rollback_to_remove" ]; then
+ if ! rm -rf "$_rollback_to_remove"; then
+ echo "⚠️ Could not remove environment rollback $_rollback_to_remove" >&2
+ fi
+ fi
+ fi
+ # Only prune older orphaned copies after the replacement has succeeded, so
+ # an interrupted install never discards the last known-good environment.
+ _prune_stale_studio_venv_rollbacks
+}
+
+_cleanup_install_temporaries() {
+ [ -n "${_UV_OVERRIDE_TMPDIR:-}" ] && rm -rf "$_UV_OVERRIDE_TMPDIR" 2>/dev/null || true
+ [ -n "${_UNSLOTH_TORCH_OVERRIDES:-}" ] && rm -f "$_UNSLOTH_TORCH_OVERRIDES" 2>/dev/null || true
}
_on_install_exit() {
@@ -517,15 +703,28 @@ _on_install_exit() {
if [ "$_status" -ne 0 ]; then
_restore_studio_venv_replacement
fi
- [ -n "${_UV_OVERRIDE_TMPDIR:-}" ] && rm -rf "$_UV_OVERRIDE_TMPDIR" 2>/dev/null || true
- [ -n "${_UNSLOTH_TORCH_OVERRIDES:-}" ] && rm -f "$_UNSLOTH_TORCH_OVERRIDES" 2>/dev/null || true
+ _cleanup_install_temporaries
exit "$_status"
}
+
+_on_install_signal() {
+ _signal_status="$1"
+ # EXIT is disabled to avoid a second cleanup pass. Ignore further termination
+ # signals until the old environment is back in place.
+ trap - EXIT
+ trap '' HUP INT TERM
+ _restore_studio_venv_replacement
+ _cleanup_install_temporaries
+ exit "$_signal_status"
+}
# Empty so an inherited value never reaches the trap's rm; only temp paths this
# script creates below (spaced-path dir, torch-trio overrides) are removed.
_UV_OVERRIDE_TMPDIR=""
_UNSLOTH_TORCH_OVERRIDES=""
trap _on_install_exit EXIT
+trap '_on_install_signal 129' HUP
+trap '_on_install_signal 130' INT
+trap '_on_install_signal 143' TERM
# ── Helper: download a URL to a file (supports curl and wget) ──
download() {
@@ -551,6 +750,45 @@ _is_pkg_installed() {
esac
}
+# ── Helper: human-readable apt distro label for the sudo package prompt (#6207) ──
+# Reads /etc/os-release so the Accept? prompt can say which distro we detected and
+# that packages come from that distro's official apt repos (not a tarball).
+_apt_distro_description() {
+ # Plain ( ... ) subshell — not $() — so case/;; stays bash-3.2-safe on macOS.
+ # Bash 3.2 misparses case arms inside command substitution and errors on `;;`.
+ (
+ if [ ! -r /etc/os-release ]; then
+ printf 'a debian-like system'
+ exit 0
+ fi
+ # shellcheck disable=SC1091
+ . /etc/os-release 2>/dev/null || true
+ if [ -n "${NAME:-}" ] && [ -n "${VERSION_ID:-}" ]; then
+ _ad_label="$NAME $VERSION_ID"
+ elif [ -n "${PRETTY_NAME:-}" ]; then
+ _ad_label="$PRETTY_NAME"
+ elif [ -n "${NAME:-}" ]; then
+ _ad_label="$NAME"
+ else
+ printf 'a debian-like system'
+ exit 0
+ fi
+ case " ${ID:-} ${ID_LIKE:-} " in
+ *" debian "*|*" ubuntu "*) _ad_label="${_ad_label} (debian-like)" ;;
+ esac
+ printf '%s' "$_ad_label"
+ )
+}
+
+# ── Helper: can the controlling terminal actually be opened for reading? ──
+# `test -r` only checks permission bits, which look fine in containers and
+# systemd units where open() then fails with ENXIO. Probe with a real open.
+# The subshell is required: in dash a failed redirection on the special
+# builtin `:` exits the whole script.
+_can_read_tty() {
+ ( : /dev/null 2>&1
+}
+
# ── Helper: install packages via apt, escalating to sudo only if needed ──
# Usage: _smart_apt_install pkg1 pkg2 pkg3 ...
_smart_apt_install() {
@@ -573,39 +811,90 @@ _smart_apt_install() {
return 0
fi
- # In Tauri mode, report needed packages and exit — Rust handles elevation
+ # Optional callers never elevate, in any mode: nothing on the consumer path
+ # builds anything, so neither the terminal sudo prompt below nor the Tauri
+ # NEED_SUDO dialog (whose Cancel leaves the user not installed) may gate the
+ # run over unused tools. The caller falls through to prebuilt llama.cpp.
+ # Required packages such as curl still escalate.
+ if [ "${_SMART_APT_OPTIONAL:-false}" = true ]; then
+ return 2
+ fi
+
if [ "$TAURI_MODE" = true ]; then
+ # Report needed packages and exit — Rust handles elevation.
tauri_log "NEED_SUDO" "$_STILL_MISSING"
exit 2
fi
# Step 3: Escalate -- need elevated permissions for remaining packages
if command -v sudo >/dev/null 2>&1; then
+ _ad_desc="$(_apt_distro_description)"
echo ""
echo " !!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!"
echo " WARNING: We require sudo elevated permissions to install:"
echo " $_STILL_MISSING"
- echo " If you accept, we'll run sudo now, and it'll prompt your password."
+ echo " Detected ${_ad_desc}."
+ echo " If you accept, we'll run sudo apt-get to install these packages"
+ echo " from your distro's official repositories (not a third-party tarball)."
echo " !!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!"
echo ""
- printf " Accept? [Y/n] "
- if [ -r /dev/tty ]; then
- read -r REPLY /dev/null \
+ if ! grep -qiE 'Ryzen AI Max|Radeon 80[0-9][05]S|Strix Halo' /proc/cpuinfo 2>/dev/null \
&& ! _wsl_amd_gpu_name >/dev/null 2>&1; then
return 0
fi
@@ -1729,67 +2018,142 @@ _maybe_reroute_strixhalo_to_2404() {
_maybe_reroute_strixhalo_to_2404 || true
# ── Check system dependencies ──
-# cmake/git are only needed to *build* llama.cpp from source. 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.
tauri_log "STEP" "Checking system dependencies"
+# Without the Xcode CLT, macOS still ships /usr/bin/git as a stub that errors and pops
+# a GUI dialog, so `command -v git` is not enough -- only running it tells the truth.
+_has_working_git() {
+ command -v git >/dev/null 2>&1 || return 1
+ git --version >/dev/null 2>&1
+}
+
+# macOS system-dependency check. A function so tests/sh can sed-extract it; the old
+# inline form was untestable, which is why this gate shipped broken.
+#
+# The consumer install needs no developer toolchain: uv is a prebuilt binary, CPython
+# is uv-managed, llama.cpp/whisper.cpp/Node are prebuilt downloads, and triton is
+# skipped on macOS. Only `--local` needs git, for the unsloth-zoo git+https URL.
+_check_macos_deps() {
+ _clt_missing=false
+ xcode-select -p >/dev/null 2>&1 || _clt_missing=true
+
+ if [ "$STUDIO_LOCAL_INSTALL" = true ] && ! _has_working_git; then
+ echo ""
+ step "deps" "git is required for --local installs" "$C_ERR"
+ substep "--local installs unsloth-zoo from git+https://github.com/unslothai/unsloth-zoo,"
+ substep "which needs a working git. Install the Xcode Command Line Tools:"
+ substep " xcode-select --install"
+ substep "Then re-run this script. A normal (non---local) install needs no compiler"
+ substep "and no git -- it uses prebuilt binaries and wheels only."
+ tauri_log "NEED_XCODE_CLT" "git"
+ return 1
+ fi
+
+ if [ "$_clt_missing" = true ]; then
+ # Not fatal, and no GUI dialog: firing xcode-select --install and exiting is
+ # what stranded clean Macs.
+ step "deps" "no Xcode Command Line Tools (not required)" "$C_WARN"
+ substep "Unsloth installs prebuilt binaries and wheels, so no compiler is needed."
+ substep "Install them only for a llama.cpp source build: xcode-select --install"
+ elif command -v cmake >/dev/null 2>&1; then
+ step "deps" "all system dependencies found"
+ else
+ # cmake is only for a source build, so its absence is not fatal.
+ step "deps" "using prebuilt llama.cpp (cmake not found)" "$C_WARN"
+ substep "Install cmake only if you want a source build: brew install cmake"
+ fi
+ return 0
+}
+
+# Linux/WSL system-dependency check. Same split as macOS, and a function for the same
+# reason: tests/sh can extract it.
+#
+# Only a download transport is required. cmake, gcc and the libcurl headers exist
+# solely for a llama.cpp source build the consumer path never does -- unslothai/
+# llama.cpp publishes linux-x64/arm64 prebuilts for cpu, cuda12, cuda13, rocm and
+# vulkan. Requiring them turned every non-apt distro into a hard exit 1 over unused
+# tooling. git follows macOS: --local only.
+_check_linux_deps() {
+ _transport_missing=false
+ if ! command -v curl >/dev/null 2>&1 && ! command -v wget >/dev/null 2>&1; then
+ _transport_missing=true
+ fi
+
+ # Wanted, never required: git fetches the triton_kernels git+https requirement (a
+ # training speedup), the rest serve the optional source build. Warn, never stop.
+ _optional_missing=""
+ command -v cmake >/dev/null 2>&1 || _optional_missing="$_optional_missing cmake"
+ _has_working_git || _optional_missing="$_optional_missing git"
+ command -v gcc >/dev/null 2>&1 || _optional_missing="$_optional_missing build-essential"
+ command -v curl-config >/dev/null 2>&1 || _optional_missing="$_optional_missing libcurl4-openssl-dev"
+ # Parameter expansion, not `sed`: sed may be absent on a minimal image, and a
+ # failed `$(... | sed ...)` yields "" -- "all found" on a machine that has none.
+ _optional_missing="${_optional_missing# }"
+
+ if [ "$STUDIO_LOCAL_INSTALL" = true ] && ! _has_working_git; then
+ echo ""
+ step "deps" "git is required for --local installs" "$C_ERR"
+ substep "--local installs unsloth-zoo from git+https://github.com/unslothai/unsloth-zoo,"
+ substep "which needs git. Install it with your package manager, then re-run."
+ substep "A normal (non---local) install needs no git and no compiler."
+ return 1
+ fi
+
+ # The one fatal case: nothing can be downloaded. apt is the only distro family we
+ # can drive unattended.
+ if [ "$_transport_missing" = true ]; then
+ if command -v apt-get >/dev/null 2>&1; then
+ echo ""
+ step "deps" "missing: curl" "$C_WARN"
+ substep "Needed to download uv, Python and the prebuilt inference engine."
+ _smart_apt_install curl
+ echo ""
+ else
+ echo ""
+ step "deps" "missing: curl (or wget)" "$C_ERR"
+ substep "Unsloth needs one of them to download uv, Python and the prebuilt"
+ substep "inference engine. Install one, then re-run setup:"
+ substep " Fedora/RHEL: sudo dnf install curl"
+ substep " Arch: sudo pacman -S --needed curl"
+ substep " openSUSE: sudo zypper install curl"
+ return 1
+ fi
+ fi
+
+ # Try apt for the optional set too; failing only costs the features warned about
+ # below.
+ if [ -n "$_optional_missing" ] && command -v apt-get >/dev/null 2>&1; then
+ step "deps" "installing optional build tools: $_optional_missing" "$C_DIM"
+ # Subshell because _smart_apt_install exits rather than returns, so `|| true`
+ # alone would not catch it. _SMART_APT_OPTIONAL suppresses every escalation
+ # path, so no install hinges on a prompt for tools nothing here needs.
+ ( _SMART_APT_OPTIONAL=true; _smart_apt_install $_optional_missing ) || true
+ _optional_missing=""
+ command -v cmake >/dev/null 2>&1 || _optional_missing="$_optional_missing cmake"
+ _has_working_git || _optional_missing="$_optional_missing git"
+ command -v gcc >/dev/null 2>&1 || _optional_missing="$_optional_missing build-essential"
+ command -v curl-config >/dev/null 2>&1 || _optional_missing="$_optional_missing libcurl4-openssl-dev"
+ _optional_missing="${_optional_missing# }"
+ fi
+
+ if [ -n "$_optional_missing" ]; then
+ step "deps" "using prebuilt llama.cpp (missing: $_optional_missing)" "$C_WARN"
+ substep "Not required to run: Unsloth downloads a prebuilt inference engine."
+ case " $_optional_missing " in
+ *" git "*) substep "Without git the triton kernels training speedup is skipped." ;;
+ esac
+ else
+ step "deps" "all system dependencies found"
+ fi
+ return 0
+}
+
case "$OS" in
macos)
- # Xcode Command Line Tools provide the C/C++ compiler and git.
- if ! xcode-select -p >/dev/null 2>&1; then
- echo ""
- echo "==> Xcode Command Line Tools are required."
- echo " Installing (a system dialog will appear)..."
- xcode-select --install /dev/null || true
- echo " After the installation completes, please re-run this script."
- exit 1
- fi
- # cmake is only needed for a source build; the default prebuilt path
- # doesn't use it, so its absence is not fatal -- no Homebrew prerequisite.
- if command -v cmake >/dev/null 2>&1; then
- step "deps" "all system dependencies found"
- else
- step "deps" "using prebuilt llama.cpp (cmake not found)" "$C_WARN"
- substep "Install cmake only if you want a source build: brew install cmake"
- fi
+ _check_macos_deps || exit 1
;;
linux|wsl)
- MISSING=""
- command -v cmake >/dev/null 2>&1 || MISSING="$MISSING cmake"
- command -v git >/dev/null 2>&1 || MISSING="$MISSING git"
- # curl or wget is needed for downloads; check both
- if ! command -v curl >/dev/null 2>&1 && ! command -v wget >/dev/null 2>&1; then
- MISSING="$MISSING curl"
- fi
- command -v gcc >/dev/null 2>&1 || MISSING="$MISSING build-essential"
- # libcurl dev headers for llama.cpp HTTPS support
- command -v curl-config >/dev/null 2>&1 || MISSING="$MISSING libcurl4-openssl-dev"
-
- MISSING=$(echo "$MISSING" | sed 's/^ *//')
- if [ -n "$MISSING" ]; then
- echo ""
- step "deps" "missing: $MISSING" "$C_WARN"
- substep "These are needed to build the GGUF inference engine."
- if command -v apt-get >/dev/null 2>&1; then
- _smart_apt_install $MISSING
- else
- echo " Automatic system package installation is supported on apt-based"
- echo " Linux distributions (Ubuntu/Debian) only. Please install the"
- echo " missing dependencies with your package manager, then re-run setup:"
- echo " $MISSING"
- echo ""
- echo " Examples:"
- echo " Fedora/RHEL: sudo dnf install cmake git gcc gcc-c++ make libcurl-devel"
- echo " Arch: sudo pacman -S --needed cmake git base-devel curl"
- echo " openSUSE: sudo zypper install cmake git gcc gcc-c++ make libcurl-devel"
- exit 1
- fi
- echo ""
- else
- step "deps" "all system dependencies found"
- fi
+ _check_linux_deps || exit 1
;;
esac
@@ -2115,18 +2479,155 @@ _has_amd_rocm_gpu() {
amd-smi list 2>/dev/null | awk '/^GPU[[:space:]]*[:\[][[:space:]]*[0-9]/{ found=1 } END{ exit !found }'; then
return 0
elif [ -e /dev/kfd ] && \
- awk 'FNR==1{ gpu=0; amd=0 } /gpu_id/{ gpu=($2+0>0) } /vendor_id/{ amd=($2==4098) } \
- gpu && amd { found=1 } END{ exit !found }' \
+ awk '/vendor_id/ && $2 == 4098 { found = 1 } END { exit !found }' \
/sys/class/kfd/kfd/topology/nodes/*/properties 2>/dev/null; then
- # vendor_id 4098 = 0x1002 (AMD). NVIDIA open kernel module (driver
- # 560+) can register KFD topology nodes with non-zero gpu_id but
- # vendor_id 4318 (0x10DE). Require AMD vendor to avoid misrouting
- # NVIDIA-only hosts to the ROCm install path.
+ # vendor_id 4098 = 0x1002 (AMD) marks a GPU node: the KFD CPU node
+ # reports vendor_id 0, so any 4098 node is an AMD GPU. NVIDIA's open
+ # kernel module (driver 560+) registers KFD nodes as vendor_id 4318
+ # (0x10DE), so this never false-positives on NVIDIA-only hosts.
+ # The prior check also required a gpu_id line, but gpu_id is a SIBLING
+ # sysfs file, not a line in properties -- it never matched, so the
+ # fallback silently missed every ROCm-less AMD host (issue: fresh
+ # Arch/CachyOS boxes reporting "no GPU detected").
return 0
fi
return 1
}
+# Returns 0 if an AMD display GPU is on the PCI bus even when ROCm can't use it
+# (e.g. a Strix Halo iGPU with no /dev/kfd). Only sharpens the "no GPU detected"
+# hint. vendor 0x1002 = AMD/ATI; class 0x03* = display controller.
+_amd_gpu_present_via_pci() {
+ [ -d /sys/bus/pci/devices ] || return 1
+ for _pci_vendor in /sys/bus/pci/devices/*/vendor; do
+ [ -r "$_pci_vendor" ] || continue
+ read -r _v < "$_pci_vendor" 2>/dev/null || continue
+ [ "$_v" = "0x1002" ] || continue
+ _cls="${_pci_vendor%vendor}class"
+ [ -r "$_cls" ] || continue
+ read -r _c < "$_cls" 2>/dev/null || continue
+ case "$_c" in 0x03*) return 0 ;; esac
+ done
+ return 1
+}
+
+# Map a gfx arch to the AMD pip index family (mirrors install.ps1 $archFamilyMap).
+_amd_arch_index_family_for_gfx() {
+ case "$1" in
+ gfx1201|gfx1200) echo gfx120X-all ;;
+ gfx1151) echo gfx1151 ;;
+ gfx1150) echo gfx1150 ;;
+ gfx1152) echo gfx1152 ;;
+ gfx1103|gfx1102|gfx1101|gfx1100) echo gfx110X-all ;;
+ gfx1036|gfx1035|gfx1034|gfx1033|gfx1032|gfx1031|gfx1030) echo gfx103X-all ;;
+ gfx90a) echo gfx90a ;;
+ gfx908) echo gfx908 ;;
+ *) return 1 ;;
+ esac
+}
+
+# Map a GPU marketing name to gfx arch (kept in sync with install.ps1 nameArchTable).
+_infer_amd_gfx_arch_from_gpu_name() {
+ case "$1" in
+ *9070*|*9080*) echo gfx1201 ;;
+ *9060*) echo gfx1200 ;;
+ *"8065S"*|*"8060S"*|*"8050S"*|*"8040S"*|*"Strix Halo"*|*"Ryzen AI Max"*|*"AI Max"*) echo gfx1151 ;;
+ *"890M"*|*"880M"*|*"Strix Point"*|*"HX 37"*|*"AI 9 HX"*|*"AI 9 36"*) echo gfx1150 ;;
+ *"860M"*|*"840M"*|*"Krackan"*|*"AI 7 35"*|*"AI 5 34"*|*"AI 7 PRO 35"*|*"AI 5 33"*) echo gfx1152 ;;
+ *"RX 7600"*|*"RX 7700S"*|*"RX 7650"*|*"PRO W7600"*|*"PRO W7500"*) echo gfx1102 ;;
+ *"RX 7800"*|*"RX 7700"*|*"PRO W7700"*|*"PRO V710"*) echo gfx1101 ;;
+ *"RX 7900"*|*"PRO W7900"*|*"PRO W7800"*) echo gfx1100 ;;
+ *"780M"*|*"760M"*|*"740M"*|*"Phoenix"*|*"Hawk Point"*|*"Z1 Extreme"*|*"Z2 Extreme"*) echo gfx1103 ;;
+ *"RX 6900"*|*"RX 6800"*|*"RX 6750"*|*"RX 6700"*|*"PRO W6800"*|*"PRO W6900"*) echo gfx1030 ;;
+ *"RX 6650"*|*"RX 6600"*|*"PRO W6600"*|*"PRO W6650"*) echo gfx1032 ;;
+ *"RX 6500"*|*"RX 6400"*|*"RX 6300"*|*"PRO W6400"*|*"PRO W6500"*) echo gfx1034 ;;
+ *) return 1 ;;
+ esac
+}
+
+# Best-effort gfx inference when ROCm tools can't see the GPU (unslothai#7301).
+# Mirrors install.ps1 arch resolution on Windows ($HasROCm false, $ROCmGfxArch set).
+_infer_linux_amd_gfx_arch() {
+ if [ -n "${UNSLOTH_ROCM_GFX_ARCH:-}" ]; then
+ printf '%s\n' "$(printf '%s' "$UNSLOTH_ROCM_GFX_ARCH" | tr '[:upper:]' '[:lower:]')"
+ return 0
+ fi
+ # On WSL /proc/cpuinfo and lspci still report the host APU, but without the
+ # ROCDXG bridge (librocdxg over /dev/dxg) the AMD wheels can't reach the GPU;
+ # keep the CPU fallback there unless that runtime is present (the explicit
+ # override above still wins). Mirrors install_python_stack.py.
+ _gpu_evidence=""
+ if [ -e /dev/dxg ] || grep -qi microsoft /proc/version 2>/dev/null; then
+ for _d in /opt/rocm/lib /opt/rocm/lib64 /opt/rocm-*/lib /opt/rocm-*/lib64; do
+ { [ -e "$_d/librocdxg.so" ] || [ -e "$_d/librocdxg.so.1" ]; } && _rocdxg=1 && break
+ done
+ [ -n "${_rocdxg:-}" ] || return 1
+ # WSL enumerates no PCI display device; /dev/dxg + librocdxg IS the
+ # GPU evidence there.
+ _gpu_evidence=1
+ elif _amd_gpu_present_via_pci; then
+ _gpu_evidence=1
+ fi
+ # /proc/cpuinfo leaks the HOST CPU model into VMs/containers that received
+ # no AMD GPU, so the CPU-model text alone is not GPU evidence: require an
+ # AMD display device (PCI vendor 0x1002, class 0x03*) before trusting it.
+ # The lspci fallback below needs no gate; an AMD display line IS evidence.
+ if [ -n "$_gpu_evidence" ] && grep -qiE 'Ryzen AI Max|Radeon 80[0-9][05]S|Strix Halo' /proc/cpuinfo 2>/dev/null; then
+ echo gfx1151
+ return 0
+ fi
+ if [ -n "$_gpu_evidence" ] && grep -qiE '890M|880M|Strix Point|HX 37[05]|AI 9 HX|AI 9 36[05]' /proc/cpuinfo 2>/dev/null; then
+ echo gfx1150
+ return 0
+ fi
+ if [ -n "$_gpu_evidence" ] && grep -qiE '860M|840M|Krackan|AI 7 35[05]|AI 5 34[05]|AI 7 PRO 35|AI 5 33' /proc/cpuinfo 2>/dev/null; then
+ echo gfx1152
+ return 0
+ fi
+ if command -v lspci >/dev/null 2>&1; then
+ # A non-AMD controller can enumerate first (Intel/ASPEED before an AMD
+ # dGPU), so scan every display-class line and take the first AMD one
+ # that maps. The vendor guard is case-SENSITIVE (a -i "ATI" would match
+ # "CorporATIon" on every Intel/NVIDIA line); whole-line matching also
+ # survives the 0000: PCI domain prefix. Mirrors install_python_stack.py.
+ _amd_disp=$(lspci -nn 2>/dev/null | grep -E 'VGA compatible controller|3D controller|Display controller' | grep -E 'AMD|ATI' || true)
+ while IFS= read -r _ln; do
+ [ -n "$_ln" ] || continue
+ if _gfx=$(_infer_amd_gfx_arch_from_gpu_name "$_ln"); then
+ echo "$_gfx"
+ return 0
+ fi
+ done </dev/null 2>&1; then
+ _pg=$( (unset ROCR_VISIBLE_DEVICES HIP_VISIBLE_DEVICES; rocminfo 2>/dev/null) | grep -oE 'gfx[1-9][0-9a-z]{2,3}' || true)
+ fi
+ if [ -z "$_pg" ] && command -v amd-smi >/dev/null 2>&1; then
+ _pg=$( (unset ROCR_VISIBLE_DEVICES HIP_VISIBLE_DEVICES; amd-smi list 2>/dev/null) | grep -oE 'gfx[1-9][0-9a-z]{2,3}' || true)
+ if [ -z "$_pg" ]; then
+ _pg=$( (unset ROCR_VISIBLE_DEVICES HIP_VISIBLE_DEVICES; amd-smi static --asic 2>/dev/null) | grep -oE 'gfx[1-9][0-9a-z]{2,3}' || true)
+ fi
+ fi
+ printf '%s\n' "$_pg"
+}
+
# ── Detect GPU and choose PyTorch index URL ──
# Mirrors Get-TorchIndexUrl in install.ps1.
# On CPU-only machines this returns the cpu index, avoiding the solver
@@ -2180,6 +2681,29 @@ get_torch_index_url() {
if ! _has_amd_rocm_gpu; then
echo "$_base/cpu"; return
fi
+ # A generic rocm index is only safe when the gfx arch is readable: the
+ # Strix reroute (gfx1150/1151 -> arch-specific index) learns gfx from
+ # rocminfo/amd-smi, so if those are missing OR do not enumerate the GPU, an
+ # unknown-arch box might be Strix and would get the broken _grouped_mm
+ # wheels. Probe via the shared helper (override first, then rocminfo/amd-smi
+ # with visibility masks cleared); if the arch is unreadable, never guess a
+ # rocm index. A KFD-only host whose arch is still inferable from hardware
+ # IDs (PCI/cpuinfo/lspci) returns the cpu index and lets the runtime-less
+ # reroute below upgrade it to AMD per-arch wheels -- the reroute gate uses
+ # this same probe, so the handoff can't misfire. Only when inference fails
+ # too is CPU final, with the actionable warning.
+ _amd_gfx_probe=$(_probe_amd_gfx_arch)
+ if [ -z "$_amd_gfx_probe" ]; then
+ if _amd_inferred_gfx=$(_infer_linux_amd_gfx_arch 2>/dev/null) && \
+ [ -n "$_amd_inferred_gfx" ] && \
+ _amd_arch_index_family_for_gfx "$_amd_inferred_gfx" >/dev/null 2>&1; then
+ echo "[WARN] AMD GPU detected but rocminfo/amd-smi can't read its gfx arch -- inferring $_amd_inferred_gfx from hardware IDs." >&2
+ echo "$_base/cpu"; return
+ fi
+ echo "[WARN] AMD GPU detected but its gfx arch can't be read (rocminfo/amd-smi missing or not enumerating the GPU) -- installing CPU-only PyTorch." >&2
+ echo "[WARN] For GPU PyTorch, install or repair rocminfo/amd-smi (e.g. sudo pacman -S rocm-hip-sdk) and re-run this installer." >&2
+ echo "$_base/cpu"; return
+ fi
# AMD GPU confirmed -- detect ROCm version
_rocm_tag=""
_rocm_tag=$({ command -v amd-smi >/dev/null 2>&1 && \
@@ -2196,7 +2720,11 @@ get_torch_index_url() {
{ command -v rpm >/dev/null 2>&1 && \
ver="$(rpm -q --qf '%{VERSION}\n' rocm-core 2>/dev/null)" && \
[ -n "$ver" ] && \
- printf '%s\n' "$ver" | awk -F'[.-]' '{print "rocm"$1"."$2; exit}'; }) 2>/dev/null
+ printf '%s\n' "$ver" | awk -F'[.-]' '{print "rocm"$1"."$2; exit}'; }) 2>/dev/null || _rocm_tag=""
+ # ^ || guard: when EVERY version source is missing (e.g. rocminfo present
+ # but rocm-core not installed, so dpkg-query/rpm exit 1), the whole ||
+ # chain fails and set -e would kill the installer BEFORE the actionable
+ # no-version WARN below -- exactly the fresh-install case it exists for.
# Validate _rocm_tag: must match "rocmX.Y" with major >= 1
case "$_rocm_tag" in
rocm[1-9]*.[0-9]*) : ;; # valid (major >= 1)
@@ -2232,12 +2760,27 @@ get_torch_index_url() {
esac
return
fi
- # AMD GPU confirmed by rocminfo/amd-smi but ROCm version could not be
- # read from any source (amd-smi, /opt/rocm/.info/version, hipconfig,
- # dpkg, rpm). Warn explicitly rather than silently installing CPU PyTorch.
- echo "[WARN] AMD GPU detected but ROCm version could not be determined -- falling back to CPU-only PyTorch" >&2
- echo "[WARN] Ensure one of the following is accessible: amd-smi, hipconfig, /opt/rocm/.info/version, rocm-core package" >&2
- echo "[WARN] To install ROCm: https://rocm.docs.amd.com/en/latest/deploy/linux/index.html" >&2
+ # AMD GPU confirmed (rocminfo/amd-smi or the KFD topology fallback) but
+ # no ROCm/HIP install was found to read the version from (amd-smi,
+ # /opt/rocm/.info/version, hipconfig, dpkg, rpm). This is the common
+ # fresh-install case: the GPU is real, but with no ROCm userspace the
+ # correct PyTorch build can't be selected. Warn with an actionable fix
+ # rather than silently installing CPU PyTorch.
+ # A user-set UNSLOTH_ROCM_GFX_ARCH seeded the probe above, so rocminfo/
+ # amd-smi may still be unable to see the GPU; when the named arch maps to
+ # a wheel family, the runtime-less reroute (gated on the override) will
+ # install the AMD per-arch wheels -- a CPU-only warning here would be
+ # false for that path. Defer like the inferable-arch branch does.
+ if [ -n "${UNSLOTH_ROCM_GFX_ARCH:-}" ] && \
+ _amd_arch_index_family_for_gfx "$_amd_gfx_probe" >/dev/null 2>&1; then
+ echo "[WARN] AMD GPU detected with no readable ROCm version, but UNSLOTH_ROCM_GFX_ARCH=$_amd_gfx_probe is set -- routing to AMD per-arch wheels." >&2
+ echo "$_base/cpu"; return
+ fi
+ echo "[WARN] AMD GPU detected, but no ROCm/HIP install was found to select the matching GPU PyTorch build -- falling back to CPU-only PyTorch." >&2
+ echo "[WARN] Install the ROCm/HIP SDK, then re-run this installer:" >&2
+ echo "[WARN] Arch / CachyOS : sudo pacman -S rocm-hip-sdk" >&2
+ echo "[WARN] other distros : https://rocm.docs.amd.com/en/latest/deploy/linux/index.html" >&2
+ echo "[WARN] Minimum required for version detection: amd-smi, hipconfig, /opt/rocm/.info/version, or the rocm-core package." >&2
echo "$_base/cpu"; return
fi
# Parse CUDA version from nvidia-smi output (POSIX-safe, no grep -P).
@@ -2649,7 +3192,7 @@ _maybe_bootstrap_rocm_wsl() {
[ -e /dev/dxg ] || return 0
# Strix APUs show in /proc/cpuinfo (the CPU model); discrete cards don't, so also
# ask the Windows host. Either signal suffices; the bootstrap detects arch from rocminfo.
- if ! grep -qiE 'Ryzen AI Max|Radeon 80[0-9]0S|Strix Halo' /proc/cpuinfo 2>/dev/null \
+ if ! grep -qiE 'Ryzen AI Max|Radeon 80[0-9][05]S|Strix Halo' /proc/cpuinfo 2>/dev/null \
&& ! _wsl_amd_gpu_name >/dev/null 2>&1; then
return 0
fi
@@ -2735,6 +3278,72 @@ fi
TORCH_INDEX_URL=$(get_torch_index_url)
+# Linux: ROCm runtime missing but a supported AMD gfx arch is inferable (Strix Halo
+# in /proc/cpuinfo, lspci marketing name, UNSLOTH_ROCM_GFX_ARCH). Route to AMD's
+# per-arch wheels like install.ps1 does on Windows (unslothai#7301).
+# Gated on the runtime probes NOT naming a gfx: either no AMD GPU is detected at
+# all (_has_amd_rocm_gpu false), or the GPU is visible only through the
+# env-independent KFD topology while rocminfo/amd-smi can't read its arch
+# (KFD-only host, unslothai#7314 -- before the KFD detection fix these hosts
+# reached this reroute via the false branch, so the empty-probe condition
+# preserves that routing). A */cpu index chosen WITH a readable gfx
+# (unsupported/unreadable ROCm version, after its own warning) is a deliberate
+# fallback -- rerouting it would contradict that decision, and stays excluded
+# because the shared probe returns its gfx. An explicit UNSLOTH_ROCM_GFX_ARCH
+# override stays authoritative either way.
+if [ "$_torch_index_pinned" = false ] && [ "$SKIP_TORCH" = false ] && \
+ ! _has_usable_nvidia_gpu && \
+ { [ -n "${UNSLOTH_ROCM_GFX_ARCH:-}" ] || ! _has_amd_rocm_gpu || \
+ [ -z "$(_probe_amd_gfx_arch)" ]; } && \
+ case "$(uname -s)" in Linux) true ;; *) false ;; esac && \
+ case "$_ARCH" in x86_64|amd64) true ;; *) false ;; esac; then
+ # ROCm torch wheels are x86_64-only; get_torch_index_url returns CPU on other
+ # arches, so an inferred/overridden gfx must not reroute arm64 to AMD wheels.
+ case "$TORCH_INDEX_URL" in
+ */cpu)
+ _linux_inferred_gfx=$(_infer_linux_amd_gfx_arch 2>/dev/null || true)
+ if [ -n "$_linux_inferred_gfx" ]; then
+ _amd_family=$(_amd_arch_index_family_for_gfx "$_linux_inferred_gfx") || _amd_family=""
+ if [ -n "$_amd_family" ]; then
+ _amd_mirror="${UNSLOTH_AMD_ROCM_MIRROR:-https://repo.amd.com/rocm/whl}"
+ while [ "${_amd_mirror%/}" != "$_amd_mirror" ]; do
+ _amd_mirror="${_amd_mirror%/}"
+ done
+ TORCH_INDEX_URL="${_amd_mirror}/${_amd_family}/"
+ # Hand the inferred arch to setup.sh (llama.cpp): it re-probes
+ # ROCm on its own, and on these runtime-less hosts its probes
+ # find nothing, so without this it classifies the box as
+ # non-ROCm and installs the CPU prebuilt while torch just got
+ # AMD per-arch wheels. setup.sh and install_llama_prebuilt.py
+ # both honor UNSLOTH_ROCM_GFX_ARCH, so exporting it is the
+ # whole handoff (a user-set override re-exports unchanged).
+ export UNSLOTH_ROCM_GFX_ARCH="$_linux_inferred_gfx"
+ case "$_linux_inferred_gfx" in
+ gfx1201|gfx1200|gfx1151|gfx1150|gfx1152)
+ TORCH_CONSTRAINT="torch>=2.11.0,<2.12.0"
+ TORCHVISION_CONSTRAINT="torchvision>=0.26.0,<0.27.0"
+ TORCHAUDIO_CONSTRAINT="torchaudio>=2.11.0,<2.12.0"
+ ;;
+ esac
+ echo "" >&2
+ # KFD-only hosts reach this reroute with /dev/kfd present
+ # (that's what detected them), so don't claim it's missing.
+ if _has_amd_rocm_gpu; then
+ echo " [WARN] AMD GPU visible via the kernel driver (KFD) but rocminfo/amd-smi can't read its gfx arch; using $_linux_inferred_gfx." >&2
+ else
+ echo " [WARN] ROCm runtime not visible (/dev/kfd, rocminfo, amd-smi) but $_linux_inferred_gfx inferred." >&2
+ fi
+ echo " [WARN] Routing to AMD arch-specific wheels ($(_strip_index_url_credentials "$TORCH_INDEX_URL"))." >&2
+ echo " [WARN] These wheels bundle their own ROCm runtime; install the kernel stack for native compute:" >&2
+ echo " [WARN] https://docs.unsloth.ai/get-started/install-and-update/amd" >&2
+ echo " [WARN] Tip: set UNSLOTH_ROCM_GFX_ARCH=$_linux_inferred_gfx to skip inference next time." >&2
+ echo "" >&2
+ fi
+ fi
+ ;;
+ esac
+fi
+
# Export the resolved torch backend ("cuda", "rocm", or "cpu") so that
# downstream scripts (setup.sh -> install_python_stack.py) know what was
# chosen here and can skip ROCm-specific repair steps on CUDA/CPU hosts.
@@ -2779,7 +3388,7 @@ fi
# 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)
+ rocm7.2|gfx120x-all|gfx1151|gfx1150|gfx1152)
TORCH_CONSTRAINT="torch>=2.11.0,<2.12.0"
TORCHVISION_CONSTRAINT="torchvision>=0.26.0,<0.27.0"
TORCHAUDIO_CONSTRAINT="torchaudio>=2.11.0,<2.12.0"
@@ -2818,29 +3427,64 @@ case "$TORCH_INDEX_URL" in
fi
;;
esac
-# ── Strix Halo / Strix Point: force rocm7.2 wheels, bypass Radeon repo ───────
-# gfx1151 (Strix Halo) and gfx1150 (Strix Point) have a ROCm 7.1 driver bug
-# that causes a segfault in torch._grouped_mm (moe_utils.py line 167).
-# The Radeon repo now ships cp313 wheels for rocm-rel-7.1, so when
-# _amd_gpu_radeon=true the installer silently lands on the broken combo.
-# Detect these GPUs when TORCH_INDEX_URL is rocm7.1 and override to rocm7.2.
-case "$TORCH_INDEX_URL" in
- */rocm7.1|*/rocm7.1.*)
+# 0 when a rocmX.Y index leaf ($1, the final path segment) is older than floor
+# $2.$3 (int compare, so rocm7.2 < rocm7.13). Non-rocm leaves (gfx*, cu*, cpu) and
+# non-numeric versions return 1. Leaf-based (like $_torch_index_leaf) so a mirror
+# base holding its own rocm token compares the family leaf, not the base path.
+_rocm_leaf_below() {
+ case "$1" in rocm[0-9]*.[0-9]*) : ;; *) return 1 ;; esac
+ _rb=${1#rocm}; _maj=${_rb%%.*}; _min=${_rb#*.}; _min=${_min%%.*}
+ case "$_maj$_min" in *[!0-9]*) return 1 ;; esac
+ if [ "$_maj" -lt "$2" ]; then return 0; fi
+ if [ "$_maj" -eq "$2" ] && [ "$_min" -lt "$3" ]; then return 0; fi
+ return 1
+}
+# ── Strix Halo / Strix Point: route to the AMD arch-specific index ───────────
+# gfx1151/gfx1150 need torch 2.11+rocm7.13 from repo.amd.com/rocm/whl/gfx/,
+# which carries AMD's real fixes (the rocm7.1 _grouped_mm segfault, moe_utils.py:167,
+# and later Strix kernel bugs). Every generic pytorch.org index below rocm7.13 lacks
+# them (and the Radeon repo can be offline, unslothai#7264), so reroute a detected
+# Strix GPU whenever the picked index is older than the arch build -- covers today's
+# rocm6.0-7.2 and any future 7.x < 7.13; rocm7.13+ already has the fixes, so leave it.
+case "$_torch_index_leaf" in
+ rocm[0-9]*)
# Collect every gfx token in rocminfo / amd-smi enumeration order
# (skip duplicates), then index by HIP_VISIBLE_DEVICES /
# ROCR_VISIBLE_DEVICES so a mixed Strix iGPU + non-Strix dGPU box
# where the user selected the dGPU does NOT get rerouted to the
# Strix per-gfx index.
- _gfx_all=""
- if command -v rocminfo >/dev/null 2>&1; then
- _gfx_all=$(rocminfo 2>/dev/null | grep -oE 'gfx[1-9][0-9a-z]{2,3}')
+ # || true on each probe: no gfx match makes grep exit 1, which under
+ # set -euo pipefail would abort the installer before the next fallback
+ # runs (now that the case matches every rocm* index, not just rocm7.1).
+ # A user-supplied UNSLOTH_ROCM_GFX_ARCH overrides probing (mirrors setup.sh
+ # and the display block), so a Strix override still reaches the arch index.
+ _gfx_all=$(printf '%s' "${UNSLOTH_ROCM_GFX_ARCH:-}" | tr '[:upper:]' '[:lower:]')
+ if [ -z "$_gfx_all" ] && command -v rocminfo >/dev/null 2>&1; then
+ _gfx_all=$(rocminfo 2>/dev/null | grep -oE 'gfx[1-9][0-9a-z]{2,3}' || true)
fi
if [ -z "$_gfx_all" ] && command -v amd-smi >/dev/null 2>&1; then
- _gfx_all=$(amd-smi list 2>/dev/null | grep -oE 'gfx[1-9][0-9a-z]{2,3}')
+ _gfx_all=$(amd-smi list 2>/dev/null | grep -oE 'gfx[1-9][0-9a-z]{2,3}' || true)
# PowerShell paths also probe `amd-smi static --asic`; mirror it
# so a host with hipinfo-less amd-smi reports the gfx target.
if [ -z "$_gfx_all" ]; then
- _gfx_all=$(amd-smi static --asic 2>/dev/null | grep -oE 'gfx[1-9][0-9a-z]{2,3}')
+ _gfx_all=$(amd-smi static --asic 2>/dev/null | grep -oE 'gfx[1-9][0-9a-z]{2,3}' || true)
+ fi
+ fi
+ # get_torch_index_url reads the arch with ROCR/HIP masks cleared, so a
+ # mask hiding every agent (e.g. ROCR_VISIBLE_DEVICES=-1) still lands
+ # here on a generic rocm index; re-probe unmasked or a masked-out Strix
+ # box keeps the broken generic wheels. Partial masks never get here
+ # (they enumerate at least one agent above) and keep their selection.
+ # ${VAR+x} (not :-): a SET-but-empty mask also hides every agent and
+ # must trigger the re-probe too.
+ if [ -z "$_gfx_all" ] && [ -n "${ROCR_VISIBLE_DEVICES+x}${HIP_VISIBLE_DEVICES+x}" ]; then
+ if command -v rocminfo >/dev/null 2>&1; then
+ _gfx_all=$( (unset ROCR_VISIBLE_DEVICES HIP_VISIBLE_DEVICES; rocminfo 2>/dev/null) | grep -oE 'gfx[1-9][0-9a-z]{2,3}' || true)
+ fi
+ if [ -z "$_gfx_all" ] && command -v amd-smi >/dev/null 2>&1; then
+ _gfx_all=$( (unset ROCR_VISIBLE_DEVICES HIP_VISIBLE_DEVICES; amd-smi list 2>/dev/null) | grep -oE 'gfx[1-9][0-9a-z]{2,3}' || true)
+ [ -z "$_gfx_all" ] && \
+ _gfx_all=$( (unset ROCR_VISIBLE_DEVICES HIP_VISIBLE_DEVICES; amd-smi static --asic 2>/dev/null) | grep -oE 'gfx[1-9][0-9a-z]{2,3}' || true)
fi
fi
_runtime_gfx=""
@@ -2861,17 +3505,28 @@ case "$TORCH_INDEX_URL" in
if (n > 0) print vals[idx]
}')
fi
+ # An explicit UNSLOTH_ROCM_GFX_ARCH=gfx906 pins the runtime target to the
+ # MI50 / Radeon VII path and must win over Strix probe-order detection on a
+ # mixed Strix + MI50 host, so the Strix reroute is suppressed when it is set.
+ # Normalize a copied HIP gcnArchName (gfx906:sramecc-:xnack- -> gfx906) and
+ # trim whitespace (mirrors the Python .strip()) so the feature-flag suffix or
+ # a stray newline does not defeat the exact gfx906 comparisons below.
+ _gfx906_env=$(printf '%s' "${UNSLOTH_ROCM_GFX_ARCH:-}" | tr '[:upper:]' '[:lower:]' | tr -d '[:space:]')
+ _gfx906_env=${_gfx906_env%%:*}
_strix_gfx=""
- case "$_runtime_gfx" in
- gfx1151|gfx1150) _strix_gfx="$_runtime_gfx" ;;
- esac
- if [ -n "$_strix_gfx" ]; then
+ if [ "$_gfx906_env" != "gfx906" ]; then
+ case "$_runtime_gfx" in
+ gfx1151|gfx1150|gfx1152) _strix_gfx="$_runtime_gfx" ;;
+ esac
+ fi
+ # Skip rocm7.13+ generic indexes: they already ship the fixes, so the
+ # arch build (rocm7.13) would be a downgrade rather than a rescue.
+ if [ -n "$_strix_gfx" ] && _rocm_leaf_below "$_torch_index_leaf" 7 13; then
echo "" >&2
- echo " [WARN] $_strix_gfx (Strix) + ROCm 7.1 detected -- known _grouped_mm segfault" >&2
- echo " [WARN] ROCm 7.1 wheels are broken for gfx1150/gfx1151 (moe_utils.py:167)" >&2
- echo " [WARN] Routing to AMD arch-specific index (torch 2.11+rocm7.13 has the real fix)" >&2
- echo " [WARN] Upgrade ROCm to 7.2+ to use the standard index:" >&2
- echo " [WARN] https://rocm.docs.amd.com/en/latest/deploy/linux/index.html" >&2
+ echo " [WARN] $_strix_gfx (Strix) detected -- routing to the AMD arch-specific index" >&2
+ echo " [WARN] torch 2.11+rocm7.13 has AMD's real gfx1150/gfx1151 fixes (the ROCm 7.1" >&2
+ echo " [WARN] _grouped_mm segfault, moe_utils.py:167, and later Strix kernel bugs)," >&2
+ echo " [WARN] and is more reliable than the rocm7.2 index or an offline Radeon repo." >&2
echo "" >&2
# AMD's arch-specific index serves torch 2.11.0+rocm7.13.0 which has AMD's
# actual fix for the gfx1151/gfx1150 _grouped_mm kernel bug -- preferred
@@ -2891,6 +3546,57 @@ case "$TORCH_INDEX_URL" in
TORCHAUDIO_CONSTRAINT="torchaudio>=2.11.0,<2.12.0"
_amd_gpu_radeon=false
fi
+ # ── MI50 / Radeon VII (gfx906, Vega 20): legacy community-supported path ──
+ # Newer rocm wheel families bundle ROCm libraries whose Tensile kernels
+ # dropped gfx906 (rocBLAS "TensileLibrary.dat ... not read for gfx906",
+ # ROCm/TheRock#1844), so a rocm6.4+/7.x index installs a torch that fails
+ # at the first BLAS call. The rocm6.3 index is the last one whose wheels
+ # run on gfx906 (torch 2.7.0 verified on MI50 32GB; up to 2.9 in community
+ # use). Reroute any newer picked index; leave rocm6.0-6.3 alone.
+ #
+ # Target resolution: an explicit UNSLOTH_ROCM_GFX_ARCH wins (lets a host
+ # whose rocminfo/amd-smi emit no gfx token still opt in; _gfx906_env was
+ # lowercased above, before the Strix block it suppresses). Otherwise only
+ # treat gfx906 as the target when it is the SOLE distinct arch present:
+ # _gfx_all is de-duplicated by visible index, which loses per-device
+ # ordinals on a mixed host, so a non-gfx906 selection must never be
+ # downgraded to rocm6.3 -- such hosts set UNSLOTH_ROCM_GFX_ARCH to opt in.
+ _gfx906_target=false
+ if [ -n "$_gfx906_env" ]; then
+ [ "$_gfx906_env" = "gfx906" ] && _gfx906_target=true
+ elif [ -n "$_gfx_all" ]; then
+ _gfx906_uniq=$(printf '%s\n' "$_gfx_all" | awk 'NF && !seen[$0]++')
+ [ "$_gfx906_uniq" = "gfx906" ] && _gfx906_target=true
+ fi
+ # gfx906 always trains from the PyTorch rocm6.3 wheels, never the Radeon repo
+ # (repo.radeon.com wheels carry no gfx906 BLAS kernels). Clear the Radeon
+ # marketing-name flag as soon as gfx906 is the target -- even when the host
+ # already picks rocm6.0-6.3 and the reroute below is a no-op -- so a Radeon VII
+ # does not divert to the radeon branch on those versions.
+ if [ "$_gfx906_target" = true ]; then
+ _amd_gpu_radeon=false
+ fi
+ if [ "$_gfx906_target" = true ] && ! _rocm_leaf_below "$_torch_index_leaf" 6 4; then
+ echo "" >&2
+ echo " [WARN] gfx906 (MI50 / Radeon VII / Vega 20) detected -- routing torch to the" >&2
+ echo " [WARN] rocm6.3 index: it is the last wheel family that runs on gfx906 (newer" >&2
+ echo " [WARN] rocm wheels ship without gfx906 BLAS kernels and fail at first use)." >&2
+ echo " [WARN] gfx906 is a community-maintained legacy path: 16-bit LoRA and full" >&2
+ echo " [WARN] finetuning work out of the box; bitsandbytes 4-bit QLoRA requires a" >&2
+ echo " [WARN] source build of bitsandbytes for gfx906 (see docs.unsloth.ai/amd)." >&2
+ echo "" >&2
+ _amd_gfx906_base="${UNSLOTH_PYTORCH_MIRROR:-https://download.pytorch.org/whl}"
+ while [ "${_amd_gfx906_base%/}" != "$_amd_gfx906_base" ]; do
+ _amd_gfx906_base="${_amd_gfx906_base%/}"
+ done
+ TORCH_INDEX_URL="${_amd_gfx906_base}/rocm6.3"
+ # Reset to the default (<2.11) window: a rocm7.2 pick raised the floor
+ # to 2.11 above, which the rocm6.3 index (torch <= 2.9.x) cannot satisfy.
+ TORCH_CONSTRAINT="torch>=2.4,<2.11.0"
+ TORCHVISION_CONSTRAINT="torchvision>=0.19,<0.26.0"
+ TORCHAUDIO_CONSTRAINT="torchaudio>=2.4,<2.11.0"
+ # (_amd_gpu_radeon already cleared above for every gfx906 target.)
+ fi
;;
esac
fi # _torch_index_pinned guard (Radeon + Strix reroute)
@@ -2958,12 +3664,14 @@ elif case "$TORCH_INDEX_URL" in */rocm*|*/gfx*) true ;; *) false ;; esac; then
# gfx1102 matched BEFORE gfx1100 so the spaceless "RX 7700S" lands on
# gfx1102 (bash case has no negative lookahead like the PS tables).
case "$_gpu_disp_mkt" in
- *"9070 XT"*|*9080*) _gpu_disp_gfx="gfx1201" ;; # RDNA 4
- *9070*|*9060*) _gpu_disp_gfx="gfx1200" ;; # RDNA 4
- *"8060S"*|*"8050S"*|*"8040S"*|*"Strix Halo"*|*"Ryzen AI Max"*|*"AI Max"*) _gpu_disp_gfx="gfx1151" ;; # RDNA 3.5 (Strix Halo: Radeon 8060S/8050S/8040S iGPU, Ryzen AI Max+)
- *"890M"*|*"880M"*|*"860M"*|*"840M"*|*"Strix Point"*|*"Krackan"*|*"HX 37"*|*"AI 9 HX"*|*"AI 9 36"*|*"AI 7 35"*|*"AI 5 34"*|*"AI 7 PRO 35"*|*"AI 5 33"*) _gpu_disp_gfx="gfx1150" ;; # RDNA 3.5 (Strix/Krackan Point: Radeon 890M/880M iGPU, Ryzen AI 9 HX 370/375)
- *"RX 7600"*|*"RX 7700S"*|*"RX 7650"*|*"PRO W7600"*|*"PRO W7500"*|*"PRO V710"*) _gpu_disp_gfx="gfx1102" ;; # RDNA 3 (Navi 33)
- *"RX 7900"*|*"RX 7800"*|*"RX 7700"*|*"PRO W7900"*|*"PRO W7800"*|*"PRO W7700"*) _gpu_disp_gfx="gfx1100" ;; # RDNA 3 desktop / workstation (Navi 31)
+ *9070*|*9080*) _gpu_disp_gfx="gfx1201" ;; # RDNA 4 (Navi 48)
+ *9060*) _gpu_disp_gfx="gfx1200" ;; # RDNA 4 (Navi 44)
+ *"8065S"*|*"8060S"*|*"8050S"*|*"8040S"*|*"Strix Halo"*|*"Ryzen AI Max"*|*"AI Max"*) _gpu_disp_gfx="gfx1151" ;; # RDNA 3.5 (Strix Halo + Gorgon Halo: Radeon 8065S/8060S/8050S/8040S iGPU, Ryzen AI Max / Max+)
+ *"890M"*|*"880M"*|*"Strix Point"*|*"HX 37"*|*"AI 9 HX"*|*"AI 9 36"*) _gpu_disp_gfx="gfx1150" ;; # RDNA 3.5 (Strix Point: Radeon 890M/880M, Ryzen AI 9 HX 370/375)
+ *"860M"*|*"840M"*|*"Krackan"*|*"AI 7 35"*|*"AI 5 34"*|*"AI 7 PRO 35"*|*"AI 5 33"*) _gpu_disp_gfx="gfx1152" ;; # RDNA 3.5 (Krackan Point: Radeon 860M/840M, Ryzen AI 7 350 / AI 5 340)
+ *"RX 7600"*|*"RX 7700S"*|*"RX 7650"*|*"PRO W7600"*|*"PRO W7500"*) _gpu_disp_gfx="gfx1102" ;; # RDNA 3 (Navi 33)
+ *"RX 7800"*|*"RX 7700"*|*"PRO W7700"*|*"PRO V710"*) _gpu_disp_gfx="gfx1101" ;; # RDNA 3 (Navi 32)
+ *"RX 7900"*|*"PRO W7900"*|*"PRO W7800"*) _gpu_disp_gfx="gfx1100" ;; # RDNA 3 desktop / workstation (Navi 31)
*"780M"*|*"760M"*|*"740M"*|*"Phoenix"*|*"Hawk Point"*|*"Z1 Extreme"*|*"Z2 Extreme"*) _gpu_disp_gfx="gfx1103" ;; # RDNA 3 iGPU (Phoenix / Hawk Point)
*"RX 6900"*|*"RX 6800"*|*"RX 6750"*|*"RX 6700"*|*"PRO W6800"*|*"PRO W6900"*) _gpu_disp_gfx="gfx1030" ;; # RDNA 2 (Navi 21)
*"RX 6650"*|*"RX 6600"*|*"PRO W6600"*|*"PRO W6650"*) _gpu_disp_gfx="gfx1032" ;; # RDNA 2 (Navi 23)
@@ -2995,6 +3703,17 @@ elif case "$TORCH_INDEX_URL" in */rocm*|*/gfx*) true ;; *) false ;; esac; then
elif [ "$OS" = "macos" ] && [ "$_ARCH" = "arm64" ]; then
# Apple Silicon: PyTorch gets Metal (MPS) acceleration over unified memory, so not CPU-only.
step "gpu" "Apple Silicon (Metal, unified memory)"
+elif _has_amd_rocm_gpu; then
+ if [ "$_torch_index_pinned" = true ]; then
+ # An explicit UNSLOTH_TORCH_INDEX_URL/_FAMILY pin skipped all probing;
+ # do not claim ROCm is unusable when a CPU/other index was requested.
+ step "gpu" "AMD GPU (torch index pinned: $_torch_index_leaf)" "$C_WARN"
+ else
+ # AMD GPU visible to the kernel but the torch index stayed CPU: no usable
+ # ROCm userspace to pick a wheel. "none" would repeat the false diagnosis
+ # this installer used to give.
+ step "gpu" "AMD GPU (no usable ROCm -- CPU fallback)" "$C_WARN"
+ fi
else
step "gpu" "none (CPU-only)" "$C_WARN"
fi
@@ -3003,8 +3722,17 @@ fi
case "$TORCH_INDEX_URL" in
*/cpu)
if [ "$SKIP_TORCH" = false ] && [ "$OS" != "macos" ]; then
- substep "No GPU detected -- installing CPU-only PyTorch." "$C_WARN"
- if [ "$OS" = "wsl" ]; then
+ if [ "$_torch_index_pinned" = true ]; then
+ # An explicit CPU pin is a request, not a detection failure:
+ # skip the SDK guidance (ROCm may be perfectly healthy here).
+ substep "CPU-only PyTorch (index pinned via UNSLOTH_TORCH_INDEX_URL / _FAMILY)."
+ elif _has_amd_rocm_gpu; then
+ substep "AMD GPU detected, but no usable ROCm/HIP install -- installing CPU-only PyTorch." "$C_WARN"
+ substep "Install the ROCm/HIP SDK and re-run this installer for GPU PyTorch." "$C_WARN"
+ else
+ substep "No GPU detected -- installing CPU-only PyTorch." "$C_WARN"
+ fi
+ if [ "$OS" = "wsl" ] && [ "$_torch_index_pinned" = false ]; then
# WSL + no GPU detected (detection above found nothing). Common
# cause: an AMD GPU whose ROCm-on-WSL runtime isn't exposed yet --
# /dev/dxg present (graphics) but no ROCm runtime.
@@ -3031,6 +3759,13 @@ case "$TORCH_INDEX_URL" in
substep " driver is current; or run unsloth/scripts/install_rocm_wsl_strixhalo.sh yourself."
else
substep "AMD ROCm users: see https://docs.unsloth.ai/get-started/install-and-update/amd"
+ # Only when ROCm truly can't see the GPU: a detected-but-too-old
+ # ROCm (rocminfo works, wheels need 6.0+) has its own guidance.
+ if ! _has_amd_rocm_gpu && _amd_gpu_present_via_pci; then
+ substep "An AMD GPU is on the PCI bus but ROCm cannot see it (no /dev/kfd," "$C_WARN"
+ substep " rocminfo, or amd-smi). Install the ROCm kernel stack so /dev/kfd exists;"
+ substep " Strix Halo (gfx1151/gfx1150) needs a recent kernel (6.11+) and ROCm 7.x."
+ fi
fi
substep "Re-run with --no-torch for GGUF-only (faster, no PyTorch):"
substep " curl -fsSL https://unsloth.ai/install.sh | sh -s -- --no-torch"
@@ -3088,6 +3823,7 @@ for _p in ('torch', 'torchvision', 'torchaudio'):
if [ "$_MIGRATED" = true ]; then
# Migrated env: force-reinstall unsloth+unsloth-zoo for a clean state, preserving
# existing torch/CUDA unless the ROCm repair below fires.
+ _gfx906_bnb_snapshot
substep "upgrading unsloth in migrated environment..."
if [ "$SKIP_TORCH" = true ]; then
# No-torch: install unsloth + unsloth-zoo with --no-deps (current
@@ -3096,7 +3832,7 @@ if [ "$_MIGRATED" = true ]; then
# to prevent transitive torch resolution.
run_install_cmd_retry "install unsloth (migrated no-torch)" uv pip install --python "$_VENV_PY" --no-deps \
--reinstall-package unsloth --reinstall-package unsloth-zoo \
- "unsloth>=2026.7.4" "unsloth-zoo>=2026.7.4"
+ "unsloth>=2026.7.5" "unsloth-zoo>=2026.7.6"
# Resolve pydantic WITH deps so pip pins pydantic-core to the
# matching version (no-torch-runtime.txt below is --no-deps).
# All transitive deps are torch-free.
@@ -3113,7 +3849,7 @@ if [ "$_MIGRATED" = true ]; then
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.4" "unsloth-zoo>=2026.7.4" ${_MLX_LM_EXCLUDE_ARG:-}
+ "unsloth>=2026.7.5" "unsloth-zoo>=2026.7.6" ${_MLX_LM_EXCLUDE_ARG:-}
[ -n "$_UNSLOTH_TORCH_OVERRIDES" ] && rm -f "$_UNSLOTH_TORCH_OVERRIDES"
_UNSLOTH_TORCH_OVERRIDES=""
fi
@@ -3129,13 +3865,18 @@ if [ "$_MIGRATED" = true ]; then
# existing ROCm installs gain the AMD bitsandbytes build without a
# fresh reinstall.
if [ "$SKIP_TORCH" = false ] && [ "$_torch_index_is_rocm_family" = true ]; then
- _install_bnb_rocm "install bitsandbytes (AMD)" "$_VENV_PY"
+ if _is_gfx906_bnb_skip; then
+ substep "gfx906: skipping prebuilt bitsandbytes (no gfx906 kernels); build from source for 4-bit QLoRA -- https://docs.unsloth.ai/get-started/install-and-update/amd" "$C_WARN"
+ else
+ _install_bnb_rocm "install bitsandbytes (AMD)" "$_VENV_PY"
+ fi
# Repair ROCm torch if overwritten during migrated install
_has_hip=$("$_VENV_PY" -c "import torch; print(getattr(torch.version,'hip','') or '')" 2>/dev/null || true)
if [ -z "$_has_hip" ]; then
substep "repairing ROCm torch (overwritten by dependency resolution)..."
_install_torch_default_index --force-reinstall
fi
+ _gfx906_bnb_prune
fi
elif [ -n "$TORCH_INDEX_URL" ]; then
# Fresh: Step 1 - install torch from explicit index (skip when --no-torch or Intel Mac)
@@ -3326,8 +4067,13 @@ elif [ -n "$TORCH_INDEX_URL" ]; then
# 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 ] && [ "$_torch_index_is_rocm_family" = true ]; then
- _install_bnb_rocm "install bitsandbytes (AMD)" "$_VENV_PY"
+ if _is_gfx906_bnb_skip; then
+ substep "gfx906: skipping prebuilt bitsandbytes (no gfx906 kernels); build from source for 4-bit QLoRA -- https://docs.unsloth.ai/get-started/install-and-update/amd" "$C_WARN"
+ else
+ _install_bnb_rocm "install bitsandbytes (AMD)" "$_VENV_PY"
+ fi
fi
+ _gfx906_bnb_snapshot
# Fresh: Step 2 - install unsloth, preserving the torch Step 1 installed
tauri_log "STEP" "Installing Unsloth"
substep "installing unsloth (this may take a few minutes)..."
@@ -3337,7 +4083,7 @@ elif [ -n "$TORCH_INDEX_URL" ]; 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.4" "unsloth-zoo>=2026.7.4"
+ "unsloth>=2026.7.5" "unsloth-zoo>=2026.7.6"
# Same pydantic-with-deps trick as the migrated branch.
run_install_cmd_retry "install pydantic (with deps for compatible core)" \
uv pip install --python "$_VENV_PY" pydantic
@@ -3356,7 +4102,7 @@ elif [ -n "$TORCH_INDEX_URL" ]; then
elif [ "$STUDIO_LOCAL_INSTALL" = true ]; then
run_install_cmd_retry "install unsloth (local)" uv pip install --python "$_VENV_PY" \
${_UNSLOTH_TORCH_OVERRIDES:+--overrides "$_UNSLOTH_TORCH_OVERRIDES"} \
- --upgrade-package unsloth "unsloth>=2026.7.4" "unsloth-zoo>=2026.7.4"
+ --upgrade-package unsloth "unsloth>=2026.7.5" "unsloth-zoo>=2026.7.6"
substep "overlaying local repo (editable)..."
run_install_cmd "overlay local repo" uv pip install --python "$_VENV_PY" -e "$_REPO_ROOT" --no-deps
substep "overlaying unsloth-zoo from git main..."
@@ -3378,13 +4124,14 @@ elif [ -n "$TORCH_INDEX_URL" ]; then
substep "repairing ROCm torch (overwritten by dependency resolution)..."
_install_torch_default_index --force-reinstall
fi
+ _gfx906_bnb_prune
fi
else
# Fallback: GPU detection failed to produce a URL -- let uv resolve torch
tauri_log "STEP" "Installing Unsloth"
substep "installing unsloth (this may take a few minutes)..."
if [ "$STUDIO_LOCAL_INSTALL" = true ]; then
- run_install_cmd_retry "install unsloth (auto torch backend)" uv pip install --python "$_VENV_PY" "unsloth-zoo>=2026.7.4" "unsloth>=2026.7.4" --torch-backend=auto
+ run_install_cmd_retry "install unsloth (auto torch backend)" uv pip install --python "$_VENV_PY" "unsloth-zoo>=2026.7.6" "unsloth>=2026.7.5" --torch-backend=auto
substep "overlaying local repo (editable)..."
run_install_cmd "overlay local repo" uv pip install --python "$_VENV_PY" -e "$_REPO_ROOT" --no-deps
substep "overlaying unsloth-zoo from git main..."
@@ -3472,6 +4219,7 @@ if [ -n "$VENV_ABS_BIN" ]; then
fi
if ! command -v bash >/dev/null 2>&1; then
+ tauri_log "ERROR" "bash is required to run studio setup"
step "setup" "bash is required to run studio setup" "$C_ERR"
substep "Please install bash and re-run install.sh"
exit 1
@@ -3510,6 +4258,7 @@ if [ "$STUDIO_LOCAL_INSTALL" = true ]; then
STUDIO_LOCAL_REPO="$_REPO_ROOT" \
UNSLOTH_NO_TORCH="$SKIP_TORCH" \
UNSLOTH_LOCAL_LLAMA_CPP_DIR="$_WITH_LLAMA_CPP_DIR" \
+ UNSLOTH_TAURI_MODE="$TAURI_MODE" \
bash "$SETUP_SH" =0.12.0",
"rich",
"pydantic",
"pyyaml",
"nest-asyncio",
+ # Every CLI command imports studio.backend.*, which reaches structlog at
+ # module level. The rest of the server stack lives in the studio extra.
+ "structlog>=24.1.0",
+ # unsloth_cli/__init__.py reaches click via commands/start.py, so every
+ # command needs it. typer supplied it until 0.27 dropped the dependency.
+ "click>=8.0",
]
[project.scripts]
@@ -41,9 +47,14 @@ version = {attr = "unsloth.models._utils.__version__"}
[tool.setuptools]
include-package-data = true
+[tool.setuptools.cmdclass]
+# Snapshots CHANGELOG.md into studio/ so every build path ships it.
+build_py = "_changelog_build.build_py"
+
[tool.setuptools.package-data]
-unsloth_cli = ["codex_fallback_prompt.md"]
+unsloth_cli = ["codex_fallback_prompt.md", "pi_subagent.ts"]
studio = [
+ "CHANGELOG.md",
"*.sh",
"*.ps1",
"*.bat",
@@ -68,13 +79,40 @@ include = ["unsloth*", "unsloth_cli*", "studio", "studio.backend*"]
exclude = ["images*", "tests*", "*.node_modules", "*.node_modules.*"]
[project.optional-dependencies]
+# Studio's server stack, mirroring studio/backend/requirements/studio.txt.
+# test_studio_extra_matches_requirements.py catches drift.
+studio = [
+ "typer",
+ "fastapi",
+ "uvicorn",
+ "pydantic",
+ "packaging",
+ "matplotlib==3.10.9",
+ "pandas",
+ "nest_asyncio",
+ "datasets==4.3.0",
+ "pyjwt",
+ "huggingface-hub==0.36.2",
+ "structlog>=24.1.0",
+ "diceware",
+ "ddgs",
+ "cryptography>=42.0.0",
+ "boto3>=1.34.0",
+ "httpx>=0.27.0",
+ "fastmcp>=3.0.2",
+ "sqlite-vec==0.1.9",
+ "pymupdf==1.27.2.3",
+ "pymupdf4llm==0.3.4",
+ "python-docx==1.2.0",
+]
+
triton = [
"triton>=3.0.0 ; ('linux' in sys_platform)",
"triton-windows ; (sys_platform == 'win32') and (platform_machine == 'AMD64' or platform_machine == 'x86_64')",
]
huggingfacenotorch = [
- "unsloth_zoo>=2026.7.4",
+ "unsloth_zoo>=2026.7.6",
"wheel>=0.42.0",
"packaging",
"numpy",
@@ -93,9 +131,25 @@ huggingfacenotorch = [
"trl>=0.18.2,!=0.19.0,<=0.24.0",
"sentence-transformers",
]
+# torchcodec backend for Gemma audio / datasets>=4 (#7225).
+# Pick the audio-torch* pin matching your torch minor (see TORCH_TORCHCODEC).
+# torchcodec publishes no sdist and only manylinux_2_28_x86_64, macosx_*_arm64
+# and win_amd64 wheels, so Linux aarch64, Windows ARM64 and Intel Mac have
+# nothing to resolve and pip fails the whole install rather than skipping audio.
+# Gate on the platforms that have a wheel, matching
+# PLATFORM_LACKS_TORCHCODEC_WHEEL in studio/install_python_stack.py.
+audio-torch210 = [
+ "torchcodec>=0.10.0,<0.11.0 ; python_version >= '3.10' and (((sys_platform == 'linux' or sys_platform == 'win32') and (platform_machine == 'x86_64' or platform_machine == 'AMD64')) or (sys_platform == 'darwin' and platform_machine == 'arm64'))",
+]
+audio-torch290 = [
+ "torchcodec>=0.8.0,<0.10.0 ; python_version >= '3.10' and (((sys_platform == 'linux' or sys_platform == 'win32') and (platform_machine == 'x86_64' or platform_machine == 'AMD64')) or (sys_platform == 'darwin' and platform_machine == 'arm64'))",
+]
+audio-torch280 = [
+ "torchcodec>=0.6.0,<0.8.0 ; python_version >= '3.9' and (((sys_platform == 'linux' or sys_platform == 'win32') and (platform_machine == 'x86_64' or platform_machine == 'AMD64')) or (sys_platform == 'darwin' and platform_machine == 'arm64'))",
+]
huggingface = [
"unsloth[huggingfacenotorch]",
- "unsloth_zoo>=2026.7.4",
+ "unsloth_zoo>=2026.7.6",
"torchvision",
"unsloth[triton]",
]
@@ -532,16 +586,19 @@ cu126-torch2100 = [
"unsloth[huggingface]",
"bitsandbytes>=0.45.5,!=0.46.0,!=0.48.0",
"unsloth[cu126onlytorch2100]",
+ "unsloth[audio-torch210]",
]
cu128-torch2100 = [
"unsloth[huggingface]",
"bitsandbytes>=0.45.5,!=0.46.0,!=0.48.0",
"unsloth[cu128onlytorch2100]",
+ "unsloth[audio-torch210]",
]
cu130-torch2100 = [
"unsloth[huggingface]",
"bitsandbytes>=0.45.5,!=0.46.0,!=0.48.0",
"unsloth[cu130onlytorch2100]",
+ "unsloth[audio-torch210]",
]
kaggle = [
"unsloth[huggingface]",
@@ -580,7 +637,7 @@ colab-ampere-torch220 = [
"flash-attn>=2.6.3 ; ('linux' in sys_platform)",
]
colab-new = [
- "unsloth_zoo>=2026.7.4",
+ "unsloth_zoo>=2026.7.6",
"packaging",
"tyro",
"transformers>=4.51.3,!=4.52.0,!=4.52.1,!=4.52.2,!=4.52.3,!=4.53.0,!=4.54.0,!=4.55.0,!=4.55.1,!=4.57.0,!=4.57.4,!=4.57.5,!=5.0.0,!=5.1.0,<=5.5.0",
@@ -831,16 +888,19 @@ cu126-ampere-torch2100 = [
"unsloth[huggingface]",
"bitsandbytes>=0.45.5,!=0.46.0,!=0.48.0",
"unsloth[cu126onlytorch2100]",
+ "unsloth[audio-torch210]",
]
cu128-ampere-torch2100 = [
"unsloth[huggingface]",
"bitsandbytes>=0.45.5,!=0.46.0,!=0.48.0",
"unsloth[cu128onlytorch2100]",
+ "unsloth[audio-torch210]",
]
cu130-ampere-torch2100 = [
"unsloth[huggingface]",
"bitsandbytes>=0.45.5,!=0.46.0,!=0.48.0",
"unsloth[cu130onlytorch2100]",
+ "unsloth[audio-torch210]",
]
flashattentiontorch260abiFALSEcu12x = [
"flash-attn @ https://github.com/Dao-AILab/flash-attention/releases/download/v2.7.4.post1/flash_attn-2.7.4.post1+cu12torch2.6cxx11abiFALSE-cp39-cp39-linux_x86_64.whl ; ('linux' in sys_platform) and python_version == '3.9'",
@@ -1125,7 +1185,8 @@ intelgputorch210 = [
"torchvision @ https://download.pytorch.org/whl/xpu/torchvision-0.25.0%2Bxpu-cp313-cp313-win_amd64.whl#sha256=1c4b44b36a557f7381e3076fb8843366742238648441d607c8d049c6da0f8886 ; sys_platform == 'win32' and python_version == '3.13' and (platform_machine == 'AMD64' or platform_machine == 'x86_64')",
]
intel-gpu-torch210 = [
- "unsloth[intelgputorch210]"
+ "unsloth[intelgputorch210]",
+ "unsloth[audio-torch210]",
]
intelgputorch2110 = [
"unsloth_zoo[intelgpu]",
@@ -1206,8 +1267,11 @@ intel = [
]
amd = [
"unsloth[huggingfacenotorch]",
- "bitsandbytes>=0.49.1 ; ('linux' in sys_platform) and (platform_machine == 'AMD64' or platform_machine == 'x86_64' or platform_machine == 'aarch64')",
- "bitsandbytes>=0.49.1 ; (sys_platform == 'win32') and (platform_machine == 'AMD64' or platform_machine == 'x86_64')",
+ # 4-bit decode is unreliable on ROCm before 0.50.0, the first PyPI release
+ # carrying the full path: blocksize/warp decoupling (bnb #1887), fused SIMT
+ # GEMM on RDNA (#1979), RDNA3/4 workgroup fix (#2012).
+ "bitsandbytes>=0.50.0 ; ('linux' in sys_platform) and (platform_machine == 'AMD64' or platform_machine == 'x86_64' or platform_machine == 'aarch64')",
+ "bitsandbytes>=0.50.0 ; (sys_platform == 'win32') and (platform_machine == 'AMD64' or platform_machine == 'x86_64')",
]
rocm702-torch280 = [
"unsloth[amd]",
@@ -1279,6 +1343,7 @@ rocm72-torch2100 = [
"torchvision @ https://repo.radeon.com/rocm/manylinux/rocm-rel-7.2/torchvision-0.25.0%2Brocm7.2.0.git82df5f59-cp311-cp311-linux_x86_64.whl ; platform_system == 'Linux' and python_version == '3.11' and platform_machine == 'x86_64'",
"torchvision @ https://repo.radeon.com/rocm/manylinux/rocm-rel-7.2/torchvision-0.25.0%2Brocm7.2.0.git82df5f59-cp312-cp312-linux_x86_64.whl ; platform_system == 'Linux' and python_version == '3.12' and platform_machine == 'x86_64'",
"torchvision @ https://repo.radeon.com/rocm/manylinux/rocm-rel-7.2/torchvision-0.25.0%2Brocm7.2.0.git82df5f59-cp313-cp313-linux_x86_64.whl ; platform_system == 'Linux' and python_version == '3.13' and platform_machine == 'x86_64'",
+ "unsloth[audio-torch210]",
]
rocm711-torch2100 = [
"unsloth[amd]",
@@ -1297,6 +1362,7 @@ rocm711-torch2100 = [
"torchvision @ https://repo.radeon.com/rocm/manylinux/rocm-rel-7.1.1/torchvision-0.25.0%2Brocm7.1.1.git82df5f59-cp311-cp311-linux_x86_64.whl ; platform_system == 'Linux' and python_version == '3.11' and platform_machine == 'x86_64'",
"torchvision @ https://repo.radeon.com/rocm/manylinux/rocm-rel-7.1.1/torchvision-0.25.0%2Brocm7.1.1.git82df5f59-cp312-cp312-linux_x86_64.whl ; platform_system == 'Linux' and python_version == '3.12' and platform_machine == 'x86_64'",
"torchvision @ https://repo.radeon.com/rocm/manylinux/rocm-rel-7.1.1/torchvision-0.25.0%2Brocm7.1.1.git82df5f59-cp313-cp313-linux_x86_64.whl ; platform_system == 'Linux' and python_version == '3.13' and platform_machine == 'x86_64'",
+ "unsloth[audio-torch210]",
]
[project.urls]
diff --git a/scripts/build_whisper_cpp.sh b/scripts/build_whisper_cpp.sh
new file mode 100755
index 0000000000..9f7e4d4ef3
--- /dev/null
+++ b/scripts/build_whisper_cpp.sh
@@ -0,0 +1,71 @@
+#!/bin/sh
+# Build whisper.cpp's whisper-server for Studio's GGUF dictation engine.
+#
+# Installs into the managed Studio home so the backend's binary discovery
+# (core/inference/stt_ggml_sidecar.py::find_whisper_server_binary) picks it up:
+# /whisper.cpp/build/bin/whisper-server (custom home)
+# ~/.unsloth/whisper.cpp/build/bin/whisper-server (default)
+#
+# Usage:
+# ./scripts/build_whisper_cpp.sh # build the pinned tag
+# WHISPER_CPP_TAG=v1.9.0 ./scripts/build_whisper_cpp.sh
+#
+# Requires: git, cmake, a C/C++ toolchain (the same prerequisites as a
+# llama.cpp source build). GPU backends are auto-detected by whisper.cpp's
+# CMake (Metal on macOS; set GGML_CUDA=1 to force a CUDA build on Linux).
+
+set -eu
+
+WHISPER_CPP_SOURCE="${WHISPER_CPP_SOURCE:-https://github.com/ggml-org/whisper.cpp}"
+WHISPER_CPP_TAG="${WHISPER_CPP_TAG:-v1.9.1}"
+
+STUDIO_HOME="${UNSLOTH_STUDIO_HOME:-${STUDIO_HOME:-}}"
+CUSTOM_STUDIO_HOME=false
+if [ -n "$STUDIO_HOME" ]; then
+ CUSTOM_STUDIO_HOME=true
+ INSTALL_DIR="$STUDIO_HOME/whisper.cpp"
+else
+ INSTALL_DIR="$HOME/.unsloth/whisper.cpp"
+fi
+
+command -v git >/dev/null 2>&1 || { echo "ERROR: git is required" >&2; exit 1; }
+command -v cmake >/dev/null 2>&1 || { echo "ERROR: cmake is required" >&2; exit 1; }
+
+# Same policy as studio/setup.sh's _assert_studio_owned_or_absent: never delete
+# a directory under a custom Studio home unless Studio itself created it (the
+# marker file below). Protects a user-managed whisper.cpp/src from rm -rf.
+STUDIO_OWNED_MARKER=".unsloth-studio-owned"
+if [ "$CUSTOM_STUDIO_HOME" = true ] && [ -e "$INSTALL_DIR" ] && \
+ [ ! -f "$INSTALL_DIR/$STUDIO_OWNED_MARKER" ]; then
+ echo "ERROR: $INSTALL_DIR already exists and is not marked as an Unsloth-owned whisper.cpp build tree." >&2
+ echo " Move it aside or choose an empty UNSLOTH_STUDIO_HOME before re-running." >&2
+ exit 1
+fi
+
+echo "==> Building whisper.cpp ($WHISPER_CPP_TAG) into $INSTALL_DIR"
+mkdir -p "$INSTALL_DIR"
+: > "$INSTALL_DIR/$STUDIO_OWNED_MARKER"
+
+if [ ! -d "$INSTALL_DIR/src/.git" ]; then
+ rm -rf "$INSTALL_DIR/src"
+ git clone --depth 1 --branch "$WHISPER_CPP_TAG" "$WHISPER_CPP_SOURCE" "$INSTALL_DIR/src"
+else
+ git -C "$INSTALL_DIR/src" fetch --depth 1 origin "$WHISPER_CPP_TAG"
+ git -C "$INSTALL_DIR/src" checkout FETCH_HEAD
+fi
+
+CMAKE_FLAGS="-DCMAKE_BUILD_TYPE=Release -DBUILD_SHARED_LIBS=OFF"
+if [ "${GGML_CUDA:-0}" = "1" ]; then
+ CMAKE_FLAGS="$CMAKE_FLAGS -DGGML_CUDA=ON"
+fi
+
+# shellcheck disable=SC2086
+cmake -S "$INSTALL_DIR/src" -B "$INSTALL_DIR/src/build" $CMAKE_FLAGS
+NCPU="$(getconf _NPROCESSORS_ONLN 2>/dev/null || echo 4)"
+cmake --build "$INSTALL_DIR/src/build" --config Release --target whisper-server -j"$NCPU"
+
+mkdir -p "$INSTALL_DIR/build/bin"
+cp "$INSTALL_DIR/src/build/bin/whisper-server" "$INSTALL_DIR/build/bin/whisper-server"
+
+echo "==> Installed $INSTALL_DIR/build/bin/whisper-server"
+"$INSTALL_DIR/build/bin/whisper-server" --help >/dev/null 2>&1 && echo "==> Binary runs OK"
diff --git a/scripts/lint_workflow_triggers.py b/scripts/lint_workflow_triggers.py
index 0688f6c65c..8f22fcaf45 100644
--- a/scripts/lint_workflow_triggers.py
+++ b/scripts/lint_workflow_triggers.py
@@ -52,14 +52,14 @@ def _normalise_on(on_field):
def _load_workflow(path: Path):
try:
- return yaml.safe_load(path.read_text())
+ return yaml.safe_load(path.read_text(encoding = "utf-8"))
except Exception as exc:
print(f"ERROR: failed to parse {path}: {exc}", file = sys.stderr)
sys.exit(2)
def _extract_cache_keys(path: Path) -> list[str]:
- text = path.read_text()
+ text = path.read_text(encoding = "utf-8")
keys: list[str] = []
for m in re.finditer(r"(?:^|\n)\s*key:\s*([^\n]+)", text):
keys.append(m.group(1).strip())
@@ -104,7 +104,7 @@ def main() -> int:
for t in RESTRICTED_TRIGGERS:
if t in triggers:
- text = path.read_text()
+ text = path.read_text(encoding = "utf-8")
if "lint:workflow_triggers-allow-workflow_run" not in text:
findings.append(
f"{path.name}: RESTRICTED trigger '{t}' requires an "
diff --git a/scripts/notebook_validator.py b/scripts/notebook_validator.py
index c1be7a63a4..7bcee47c66 100644
--- a/scripts/notebook_validator.py
+++ b/scripts/notebook_validator.py
@@ -95,8 +95,8 @@ COLAB_ORACLE_BASE_URL = "https://raw.githubusercontent.com/googlecolab/backend-i
# Source: pytorch/torchcodec compatibility matrix on its README.
TORCH_TORCHCODEC: dict[str, set[str]] = {
"2.10": {"0.10"},
- "2.9": {"0.7", "0.8", "0.9"},
- "2.8": {"0.6"},
+ "2.9": {"0.8", "0.9"},
+ "2.8": {"0.6", "0.7"},
"2.7": {"0.3", "0.4", "0.5"},
"2.6": {"0.2", "0.3"},
"2.5": {"0.1", "0.2"},
diff --git a/scripts/profile_startup.py b/scripts/profile_startup.py
new file mode 100644
index 0000000000..937d007ac1
--- /dev/null
+++ b/scripts/profile_startup.py
@@ -0,0 +1,377 @@
+#!/usr/bin/env python3
+# SPDX-License-Identifier: AGPL-3.0-only
+# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0
+
+"""Measure where Unsloth Studio's startup time goes, per platform.
+
+Nothing measured this before: the backend logs "lifespan startup completed in X ms"
+but no test or CI job asserted a budget, and studio_test_kit discards the elapsed
+time of its /healthz poll. A first local run (Linux, warm cache, fast server CPU)
+found `import main` alone costs 6.6s before the server can bind, dominated by eager
+module-level imports pulled in by the `routes` package:
+
+ torch 1930 ms self
+ unsloth_zoo 914 ms self
+ routes 779 ms self
+ transformers 524 ms self
+
+Phases measured:
+ import `python -X importtime -c "import main"`, top cumulative + per-package self
+ spawn process start -> first byte on stdout
+ healthz process start -> /api/health (or /healthz) answers 200
+ lifespan the backend's own "lifespan startup completed in X ms" log line
+
+Usage:
+ python scripts/profile_startup.py --repeats 3 --json out.json
+ python scripts/profile_startup.py --import-only # no server, no port needed
+
+Exit code is 0 unless --max-healthz-seconds is given and exceeded.
+"""
+
+from __future__ import annotations
+
+import argparse
+import json
+import math
+import os
+import platform
+import re
+import shutil
+import socket
+import statistics
+import subprocess
+import sys
+import threading
+import time
+import urllib.error
+import urllib.request
+from pathlib import Path
+
+REPO_ROOT = Path(__file__).resolve().parents[1]
+BACKEND = REPO_ROOT / "studio" / "backend"
+
+_IMPORTTIME_RE = re.compile(r"import time:\s+(\d+)\s+\|\s+(\d+)\s+\|(\s*)(\S.*)")
+
+
+def _free_port() -> int:
+ with socket.socket() as s:
+ s.bind(("127.0.0.1", 0))
+ return int(s.getsockname()[1])
+
+
+def profile_imports(python: str, top: int = 15) -> dict:
+ """Cumulative and self import cost for the backend's module graph.
+
+ Run in a subprocess with -X importtime: the numbers are only meaningful for a
+ cold interpreter, and importing in-process would measure a warm sys.modules.
+ """
+ proc = subprocess.run(
+ [python, "-X", "importtime", "-c", "import sys; sys.path.insert(0, '.'); import main"],
+ cwd = BACKEND,
+ capture_output = True,
+ text = True,
+ timeout = 900,
+ )
+ rows = []
+ for line in proc.stderr.splitlines():
+ m = _IMPORTTIME_RE.match(line)
+ if m:
+ rows.append((int(m.group(1)), int(m.group(2)), m.group(4).strip()))
+ if not rows:
+ return {"ok": False, "error": (proc.stderr or proc.stdout)[-2000:]}
+ if proc.returncode != 0:
+ # Rows survive up to the failure, so any total from a partial graph is wrong.
+ return {
+ "ok": False,
+ "error": (proc.stderr or proc.stdout)[-2000:],
+ "partial_rows": len(rows),
+ }
+
+ by_cum = sorted(rows, key = lambda r: -r[1])
+ # Total comes from the `main` row, not by_cum[0]: -X importtime also prints the
+ # interpreter's own startup graph (`site`), which can outrank a trivial main.
+ main_row = next((r for r in reversed(rows) if r[2] == "main"), None)
+ if main_row is None:
+ return {
+ "ok": False,
+ "error": "no `import main` row in -X importtime output\n"
+ + (proc.stderr or proc.stdout)[-2000:],
+ }
+ self_by_pkg: dict[str, int] = {}
+ for self_us, _cum, name in rows:
+ pkg = name.split(".")[0]
+ self_by_pkg[pkg] = self_by_pkg.get(pkg, 0) + self_us
+
+ return {
+ "ok": True,
+ "total_seconds": round(main_row[1] / 1e6, 3),
+ "top_cumulative": [
+ {"module": n, "seconds": round(c / 1e6, 3)} for _s, c, n in by_cum[:top]
+ ],
+ "self_by_package_ms": {
+ k: round(v / 1000) for k, v in sorted(self_by_pkg.items(), key = lambda x: -x[1])[:top]
+ },
+ }
+
+
+def _terminate_tree(proc: subprocess.Popen) -> None:
+ """Stop the server AND its children, which on Windows are a separate process.
+
+ CI profiles `Scripts/unsloth.exe`, a distlib launcher stub that CreateProcess's
+ the venv python and waits, so terminate() reaps the stub only: the real backend
+ keeps the inherited stdout handle, the reader thread never sees EOF, and
+ --repeats strands one server per iteration on the shared UNSLOTH_STUDIO_HOME.
+ taskkill /T walks the tree, as unsloth_cli/commands/start.py already does.
+ """
+ if proc.poll() is not None:
+ return
+ if os.name == "nt":
+ try:
+ killed = subprocess.run(
+ ["taskkill", "/PID", str(proc.pid), "/T", "/F"],
+ capture_output = True,
+ timeout = 30,
+ check = False,
+ )
+ if killed.returncode == 0:
+ return
+ except Exception:
+ # taskkill missing or timed out; fall through so the stub still dies.
+ pass
+ # check=False: a nonzero taskkill does not raise, so fall through as well.
+ proc.terminate()
+
+
+def profile_launch(
+ bin_path: str,
+ port: int,
+ timeout_s: int = 300,
+) -> dict:
+ """Spawn the backend the way the desktop app does and time it to first 200."""
+ log_lines: list[str] = []
+ first_byte: list[float] = []
+ t0 = time.perf_counter()
+ proc = subprocess.Popen(
+ [bin_path, "studio", "--api-only", "-H", "127.0.0.1", "-p", str(port)],
+ cwd = REPO_ROOT,
+ stdout = subprocess.PIPE,
+ stderr = subprocess.STDOUT,
+ text = True,
+ bufsize = 1,
+ )
+
+ def _drain() -> None:
+ # Runs alongside the health polling: the first read timestamps the spawn
+ # phase, and an undrained pipe blocks the backend before it binds.
+ for line in proc.stdout:
+ if not first_byte:
+ first_byte.append(time.perf_counter() - t0)
+ log_lines.append(line.rstrip("\n"))
+
+ reader = threading.Thread(target = _drain, daemon = True)
+ reader.start()
+
+ t_healthz = None
+ deadline = t0 + timeout_s
+ try:
+ while time.perf_counter() < deadline:
+ if proc.poll() is not None:
+ break
+ if t_healthz is None:
+ for url in (
+ f"http://127.0.0.1:{port}/api/health",
+ f"http://127.0.0.1:{port}/healthz",
+ ):
+ try:
+ with urllib.request.urlopen(url, timeout = 2) as r:
+ if r.status == 200:
+ t_healthz = time.perf_counter() - t0
+ break
+ except (urllib.error.URLError, OSError, TimeoutError):
+ pass
+ if t_healthz is not None:
+ break
+ time.sleep(0.25)
+ finally:
+ _terminate_tree(proc)
+ try:
+ # Safe: the reader drains the pipe, so the child cannot block on write().
+ proc.wait(timeout = 30)
+ except subprocess.TimeoutExpired:
+ proc.kill()
+ proc.wait()
+ reader.join(timeout = 10)
+
+ t_first_byte = first_byte[0] if first_byte else None
+ lifespan_ms = None
+ for line in log_lines:
+ m = re.search(r"lifespan startup completed in ([\d.]+)ms", line)
+ if m:
+ lifespan_ms = float(m.group(1))
+ return {
+ "spawn_seconds": round(t_first_byte, 3) if t_first_byte is not None else None,
+ "healthz_seconds": round(t_healthz, 3) if t_healthz is not None else None,
+ "lifespan_ms": lifespan_ms,
+ "reached_healthz": t_healthz is not None,
+ "log_tail": log_lines[-25:],
+ }
+
+
+def python_version_of(python: str) -> str:
+ """Version of the interpreter that runs the imports, not the one running us.
+
+ --python points at the installed Studio venv while this script runs under the
+ runner's system python, so platform.python_version() would label it wrong.
+ """
+ if python == sys.executable:
+ return platform.python_version()
+ try:
+ proc = subprocess.run(
+ [python, "-c", "import platform; print(platform.python_version())"],
+ capture_output = True,
+ text = True,
+ timeout = 60,
+ )
+ if proc.returncode == 0 and proc.stdout.strip():
+ return proc.stdout.strip()
+ except (OSError, subprocess.SubprocessError):
+ pass
+ return "unknown"
+
+
+def find_bin() -> str | None:
+ home = os.environ.get("UNSLOTH_STUDIO_HOME") or str(Path.home() / ".unsloth" / "studio")
+ names = ["unsloth.exe", "unsloth"] if platform.system() == "Windows" else ["unsloth"]
+ subdirs = ["unsloth_studio/Scripts", "unsloth_studio/bin", "bin", "Scripts"]
+ for sd in subdirs:
+ for n in names:
+ p = Path(home) / sd / n
+ if p.exists():
+ return str(p)
+ return shutil.which("unsloth")
+
+
+def main(argv: list[str]) -> int:
+ ap = argparse.ArgumentParser(
+ description = __doc__, formatter_class = argparse.RawDescriptionHelpFormatter
+ )
+ ap.add_argument(
+ "--repeats",
+ type = int,
+ default = 1,
+ help = "launch repeats; the median is reported (imports are measured once)",
+ )
+ ap.add_argument(
+ "--python",
+ default = sys.executable,
+ help = "interpreter used for the import profile (default: this one)",
+ )
+ ap.add_argument("--bin", help = "path to the unsloth CLI (default: autodetect)")
+ ap.add_argument(
+ "--import-only",
+ action = "store_true",
+ help = "skip the server phases (no install needed beyond the deps)",
+ )
+ ap.add_argument(
+ "--max-healthz-seconds",
+ type = float,
+ help = "fail if the median time to a healthy port exceeds this",
+ )
+ ap.add_argument("--json", help = "write the full report here")
+ a = ap.parse_args(argv)
+ # range(0) launches nothing, leaving the budget check with nothing to fail on.
+ if a.repeats < 1:
+ ap.error("--repeats must be at least 1")
+ # Same reason: --import-only never launches anything.
+ if a.import_only and a.max_healthz_seconds is not None:
+ ap.error("--max-healthz-seconds cannot be combined with --import-only")
+ # nan and inf parse fine as floats but `med > budget` is then always False,
+ # so the gate would report success without ever bounding anything.
+ if a.max_healthz_seconds is not None and not math.isfinite(a.max_healthz_seconds):
+ ap.error("--max-healthz-seconds must be a finite number")
+
+ report: dict = {
+ "platform": platform.system().lower(),
+ "machine": platform.machine(),
+ "python": python_version_of(a.python),
+ "cpu_count": os.cpu_count(),
+ }
+
+ print("== import graph ==")
+ report["imports"] = profile_imports(a.python)
+ imp = report["imports"]
+ if imp.get("ok"):
+ print(f" import main: {imp['total_seconds']}s")
+ for row in imp["top_cumulative"][:8]:
+ print(f" {row['seconds']:7.3f}s {row['module']}")
+ print(" self time by package (ms):")
+ for k, v in list(imp["self_by_package_ms"].items())[:8]:
+ print(f" {v:8} ms {k}")
+ else:
+ print(f" FAILED: {imp.get('error', '')[:400]}")
+
+ if not a.import_only:
+ bin_path = a.bin or find_bin()
+ if not bin_path:
+ print(
+ "== launch == skipped: no unsloth CLI found "
+ "(set UNSLOTH_STUDIO_HOME or pass --bin)"
+ )
+ report["launch"] = {"skipped": "no unsloth CLI found"}
+ else:
+ print(f"== launch == {bin_path}")
+ runs = []
+ for i in range(a.repeats):
+ r = profile_launch(bin_path, _free_port())
+ runs.append(r)
+ print(
+ f" run {i + 1}: healthz={r['healthz_seconds']}s "
+ f"lifespan={r['lifespan_ms']}ms reached={r['reached_healthz']}"
+ )
+ got = [r["healthz_seconds"] for r in runs if r["healthz_seconds"] is not None]
+ report["launch"] = {
+ "runs": runs,
+ "failed_runs": sum(1 for r in runs if not r["reached_healthz"]),
+ "healthz_median_seconds": round(statistics.median(got), 3) if got else None,
+ "healthz_max_seconds": round(max(got), 3) if got else None,
+ }
+ if got:
+ print(
+ f" median time to healthy port: {report['launch']['healthz_median_seconds']}s"
+ )
+
+ if a.json:
+ Path(a.json).write_text(json.dumps(report, indent = 2), encoding = "utf-8")
+ print(f"\nwrote {a.json}")
+
+ if a.max_healthz_seconds is not None:
+ launch = report.get("launch") or {}
+ med = launch.get("healthz_median_seconds")
+ failed = launch.get("failed_runs") or 0
+ if failed:
+ # Failed launches fail the budget; dropping them would keep only the fast ones.
+ print(
+ f"::error::startup regression: {failed} of {len(launch.get('runs') or [])} "
+ f"launches never became healthy within the timeout"
+ )
+ return 1
+ if med is None:
+ # Nothing measured: exiting 0 would pass a requested budget without a
+ # single health request, so fail closed.
+ print(
+ "::error::startup regression: no healthz measurement, so the "
+ f"{a.max_healthz_seconds}s budget was never checked "
+ f"({launch.get('skipped') or 'launch phase produced no runs'})"
+ )
+ return 1
+ elif med > a.max_healthz_seconds:
+ print(
+ f"::error::startup regression: {med}s median to a healthy port "
+ f"exceeds the {a.max_healthz_seconds}s budget"
+ )
+ return 1
+ return 0
+
+
+if __name__ == "__main__":
+ raise SystemExit(main(sys.argv[1:]))
diff --git a/scripts/scan_packages_baseline.json b/scripts/scan_packages_baseline.json
index 1f7bc8dcc0..58b7f95ab1 100644
--- a/scripts/scan_packages_baseline.json
+++ b/scripts/scan_packages_baseline.json
@@ -1,5 +1,5 @@
{
- "_comment": "scan_packages.py allowlist. Each entry is a CRITICAL/HIGH finding manually judged benign. Matched on (package, package-relative file, check, evidence_hash); evidence_hash is over the matched code with L: markers stripped, so version bumps and line shifts do not reopen an entry but changed code does. severity and evidence are for review only. Regenerate with --write-baseline AFTER reviewing every line.",
+ "_comment": "scan_packages.py allowlist (reviewed). Each entry is a CRITICAL/HIGH finding manually judged benign. Matched on (package, package-relative file, check, evidence_hash); evidence_hash is over the matched code with L: markers stripped, so version bumps and line shifts do not reopen an entry but changed code does. severity and evidence are for review only. Regenerate with --write-baseline AFTER reviewing every line.",
"version": 1,
"entries": [
{
@@ -98,6 +98,14 @@
"evidence": "L587: while True: sha256:06c2c7f15d73bf192e5e3272c5ff5fcaeff7f6774fef5f4eca6ef473ae50e2b3",
"evidence_hash": "57acd497f404c203e4450d0580ad85aa8a33406e8d64ad06fbac6cf47d97b24d"
},
+ {
+ "package": "fastapi",
+ "file": "fastapi/routing.py",
+ "check": "C2 polling/beaconing loop detected",
+ "severity": "CRITICAL",
+ "evidence": "L592: while True: sha256:84283c09277ded3296998b2a6a838744457b606829cf5ab5d0da6f222ff020a0",
+ "evidence_hash": "a7295004315e26a8f3c64fb837521e9fdd7268219bb43e000fb0236ab0259223"
+ },
{
"package": "fastmcp-slim",
"file": "fastmcp/cli/apps_dev.py",
@@ -303,8 +311,8 @@
"file": "openai/_base_client.py",
"check": "C2 polling/beaconing loop detected",
"severity": "CRITICAL",
- "evidence": "L264: while True: sha256:95ca67e46d42354ae650abbdc5b0d97df8b0ed43187800bf40f5690c3901b94b",
- "evidence_hash": "a57d8d15fed0bf04f9967dcc18a18b80bb19f4095675bccbb78ac0450d7fce14"
+ "evidence": "L274: while True: sha256:90a38e5c1e26893c7c273354143612640e9a9c0f079d3e2b60612d79f24e80a6",
+ "evidence_hash": "1022e8e8649436ec64a98a9d9141d085452c49549fd2157b0278fc369a83ac66"
},
{
"package": "openai",
@@ -319,8 +327,8 @@
"file": "openai/auth/_workload.py",
"check": "Accesses cloud metadata/IMDS AND makes network calls",
"severity": "CRITICAL",
- "evidence": "IMDS: L96: url = \"http://169.254.169.254/metadata/identity/oauth2/token\" | L149: url = \"http://metadata.google.internal/computeMetadata/v1/instance/service-accounts/default/identity\"\nNetwork: L77: http_client: httpx.Client | None = None, | L108: with httpx.Client() as client: | L133: http_client: httpx.Client | None = None, | L155: with httpx.Client() as client: | L248: with httpx.Client() as client:",
- "evidence_hash": "1581d9f4a23393e9af23fbe5ef9f66807b22c5b5a3f1fe167254c9ebee108567"
+ "evidence": "IMDS: L97: url = \"http://169.254.169.254/metadata/identity/oauth2/token\" | L150: url = \"http://metadata.google.internal/computeMetadata/v1/instance/service-accounts/default/identity\"\nNetwork: L78: http_client: httpx.Client | None = None, | L109: with httpx.Client() as client: | L134: http_client: httpx.Client | None = None, | L156: with httpx.Client() as client: | L251: exchange_client = DefaultHttpx2Client(follow_redirects=False) if self._use_httpx2 else httpx.Client()",
+ "evidence_hash": "9717e51cb961dc14c458955d91a1e48e3753997346ecea0106bded3a8d64bfe0"
},
{
"package": "openai",
@@ -343,8 +351,8 @@
"file": "openai/resources/beta/responses/responses.py",
"check": "C2 polling/beaconing loop detected",
"severity": "CRITICAL",
- "evidence": "L3999: while True: sha256:df298b6eaf3416589b79f4ef283f8fb76e54d505bfda8840673f8e6419117e2e",
- "evidence_hash": "10ce5cb5a7097fcff4042ddcfb4802edda60aa4b7b113c8b926a52ddb76f78c2"
+ "evidence": "L4000: while True: sha256:f8ab538118daba9ec06e27399dbdc90a4521c3390e6a47a6348a1f180a83effd",
+ "evidence_hash": "31481ea83c687acc27144d72d3832d4fb98dd1c79fb5e0ddd85080de95997b9f"
},
{
"package": "openai",
@@ -359,16 +367,16 @@
"file": "openai/resources/realtime/realtime.py",
"check": "C2 polling/beaconing loop detected",
"severity": "CRITICAL",
- "evidence": "L310: while True: sha256:458198ff3d3f05870bf98c9564cbfd68c739e57b9bbe4120ed81e3eb6af74a05",
- "evidence_hash": "a3165d21e46b3ce553795daeae53e8f80e8e89c5cb228e68e6dcaff54bca5a89"
+ "evidence": "L311: while True: sha256:5b63313072aae9ca28677e03426513ccf12221e4f4e0ea6c31efbe09790633b5",
+ "evidence_hash": "05e1af469d651b51673763a7c4cdf759af9472fb627b7b470adc28cc237bd650"
},
{
"package": "openai",
"file": "openai/resources/responses/responses.py",
"check": "C2 polling/beaconing loop detected",
"severity": "CRITICAL",
- "evidence": "L3950: while True: sha256:1ce0b5a388c747945cdfda1a71b77afdfd03ae840d7aa9fa62f02eb00aa5e29f",
- "evidence_hash": "6de300ebb5e6e17cb51c89cbcdf08515a44655182f0776f0908a9d1043ebbcd7"
+ "evidence": "L3951: while True: sha256:d68ef896bf0743ca430cfacb9a3353da1f3b9c51c3a21b6450a07a32b55aa2ac",
+ "evidence_hash": "160eecdd79b521bffbe8476f782b69a0724c35d1b19376a7600807165fd54f9f"
},
{
"package": "openai",
@@ -1545,6 +1553,78 @@
"severity": "HIGH",
"evidence": "Obfusc: L836: code = compile(module, \"\", \"exec\")\nExec: L736: exec(code, globs, locs)",
"evidence_hash": "5c0992c90f05c772abd94d00784f157de337e1f8567f8b3aee1b15e46c96cd5d"
+ },
+ {
+ "package": "unsloth-zoo",
+ "file": "tests/test_mlx_save_export_regressions.py",
+ "check": "Writes to /tmp and executes (staged dropper)",
+ "severity": "CRITICAL",
+ "evidence": "L165: temporary_location=\"/tmp/ignored\", sha256:ab5c587f9ec31a0cc10ee55698ab133a417148d9d3f371bbc81b1e13fa119c13",
+ "evidence_hash": "93a11159147aad94f353ec4d2e0b8486b256abef88cd96d741813222cd32b138"
+ },
+ {
+ "package": "unsloth-zoo",
+ "file": "tests/test_vision_collator_audio.py",
+ "check": "Writes to /tmp and executes (staged dropper)",
+ "severity": "CRITICAL",
+ "evidence": "L111: out = extract_audio_info(msgs({\"type\": \"audio\", key: \"/tmp/a.wav\"})) sha256:2efe23ffbe2b91b8403aec9b700736919b59e5ca770f8e1f5501651b44b7d398",
+ "evidence_hash": "d416b79dd17b24214f3f7653ac01354507d7bf0fc464dee30a4a4b8998f063ba"
+ },
+ {
+ "package": "openai",
+ "file": "openai/_base_client.py",
+ "check": "C2 polling/beaconing loop detected",
+ "severity": "CRITICAL",
+ "evidence": "L274: while True: sha256:90a38e5c1e26893c7c273354143612640e9a9c0f079d3e2b60612d79f24e80a6",
+ "evidence_hash": "1022e8e8649436ec64a98a9d9141d085452c49549fd2157b0278fc369a83ac66"
+ },
+ {
+ "package": "openai",
+ "file": "openai/auth/_workload.py",
+ "check": "Accesses cloud metadata/IMDS AND makes network calls",
+ "severity": "CRITICAL",
+ "evidence": "IMDS: L97: url = \"http://169.254.169.254/metadata/identity/oauth2/token\" | L150: url = \"http://metadata.google.internal/computeMetadata/v1/instance/service-accounts/default/identity\"\nNetwork: L78: http_client: httpx.Client | None = None, | L109: with httpx.Client() as client: | L134: http_client: httpx.Client | None = None, | L156: with httpx.Client() as client: | L251: exchange_client = DefaultHttpx2Client(follow_redirects=False) if self._use_httpx2 else httpx.Client()",
+ "evidence_hash": "9717e51cb961dc14c458955d91a1e48e3753997346ecea0106bded3a8d64bfe0"
+ },
+ {
+ "package": "openai",
+ "file": "openai/resources/beta/responses/responses.py",
+ "check": "C2 polling/beaconing loop detected",
+ "severity": "CRITICAL",
+ "evidence": "L4000: while True: sha256:f8ab538118daba9ec06e27399dbdc90a4521c3390e6a47a6348a1f180a83effd",
+ "evidence_hash": "31481ea83c687acc27144d72d3832d4fb98dd1c79fb5e0ddd85080de95997b9f"
+ },
+ {
+ "package": "openai",
+ "file": "openai/resources/realtime/realtime.py",
+ "check": "C2 polling/beaconing loop detected",
+ "severity": "CRITICAL",
+ "evidence": "L311: while True: sha256:5b63313072aae9ca28677e03426513ccf12221e4f4e0ea6c31efbe09790633b5",
+ "evidence_hash": "05e1af469d651b51673763a7c4cdf759af9472fb627b7b470adc28cc237bd650"
+ },
+ {
+ "package": "openai",
+ "file": "openai/resources/responses/responses.py",
+ "check": "C2 polling/beaconing loop detected",
+ "severity": "CRITICAL",
+ "evidence": "L3951: while True: sha256:d68ef896bf0743ca430cfacb9a3353da1f3b9c51c3a21b6450a07a32b55aa2ac",
+ "evidence_hash": "160eecdd79b521bffbe8476f782b69a0724c35d1b19376a7600807165fd54f9f"
+ },
+ {
+ "package": "unsloth-zoo",
+ "file": "tests/test_gemma4_forced_float32_ple_dtype.py",
+ "check": "Advanced obfuscation (marshal/compile/zlib) + exec/eval",
+ "severity": "HIGH",
+ "evidence": "Obfusc: L277: compile(rewritten + _GEMMA4_PLE_CAST_HELPER, \"\", \"exec\") | L440: compile(on, \"\", \"exec\") | L468: compile(generated, \"\", \"exec\")\nExec: L19: exec(_GEMMA4_PLE_CAST_HELPER, namespace)",
+ "evidence_hash": "a85e24d8e7c431563cbd83b70f91a3b971abde0f37083d68e70984147960cc70"
+ },
+ {
+ "package": "unsloth-zoo",
+ "file": "tests/test_vision_collator_audio.py",
+ "check": "Writes to /tmp and executes (staged dropper)",
+ "severity": "CRITICAL",
+ "evidence": "L111: out = extract_audio_info(msgs({\"type\": \"audio\", key: \"/tmp/a.wav\"})) sha256:022f81dd21acfc6a35a058de96132834c218404a9e37b3d09a7768a8c8f6c728",
+ "evidence_hash": "2d1e75446af120d9133a42aa8af426a839d3434d9dc109cc1d6c1b22ca1ddb75"
}
]
}
diff --git a/studio/Unsloth_Studio_Colab.ipynb b/studio/Unsloth_Studio_Colab.ipynb
index 44282b2255..612d739806 100644
--- a/studio/Unsloth_Studio_Colab.ipynb
+++ b/studio/Unsloth_Studio_Colab.ipynb
@@ -1,134 +1,145 @@
{
- "cells": [
- {
- "cell_type": "markdown",
- "metadata": {
- "id": "view-in-github",
- "colab_type": "text"
- },
- "source": [
- ""
- ]
+ "cells": [
+ {
+ "cell_type": "markdown",
+ "metadata": {
+ "id": "view-in-github",
+ "colab_type": "text"
+ },
+ "source": [
+ ""
+ ]
+ },
+ {
+ "cell_type": "markdown",
+ "metadata": {
+ "id": "6b87de59"
+ },
+ "source": [
+ "To run this, press \"*Runtime*\" and press \"*Run all*\" on a **free** Tesla T4 Google Colab instance!\n",
+ "
\n",
+ "\n",
+ "\n",
+ " Join Discord if you need help + ⭐ Star us on Github ⭐\n",
+ "
\n",
+ "\n",
+ "To install Unsloth Studio on your local device, follow [our guide](https://unsloth.ai/docs/new/unsloth-studio/install). Unsloth Studio is licensed [AGPL-3.0](https://github.com/unslothai/unsloth/blob/main/studio/LICENSE.AGPL-3.0).\n",
+ "\n",
+ "### Unsloth Studio\n",
+ "\n",
+ "Train and run open models with [**Unsloth Studio**](https://unsloth.ai/docs/new/unsloth-studio/start). NEW! Installation should now only take 2 mins!\n",
+ "\n",
+ "\n",
+ "We are actively working on making Unsloth Studio install on Colab T4 GPUs faster.\n",
+ "\n",
+ "[Features](https://unsloth.ai/docs/new/unsloth-studio#features) • [Quickstart](https://unsloth.ai/docs/new/unsloth-studio/start) • [Data Recipes](https://unsloth.ai/docs/new/unsloth-studio/data-recipe) • [Unsloth Chat](https://unsloth.ai/docs/new/unsloth-studio/chat) • [Export](https://unsloth.ai/docs/new/unsloth-studio/export)"
+ ],
+ "id": "6b87de59"
+ },
+ {
+ "cell_type": "markdown",
+ "metadata": {
+ "id": "e4206349"
+ },
+ "source": [
+ "
"
+ ],
+ "id": "e4206349"
+ },
+ {
+ "cell_type": "markdown",
+ "metadata": {
+ "id": "27da2957"
+ },
+ "source": [
+ "### Setup: Clone repo and run setup"
+ ],
+ "id": "27da2957"
+ },
+ {
+ "cell_type": "code",
+ "metadata": {
+ "id": "27e68f91"
+ },
+ "source": "!git clone --depth 1 --branch main https://github.com/unslothai/unsloth.git\n%cd /content/unsloth\n!chmod +x studio/setup.sh && ./studio/setup.sh --local",
+ "execution_count": null,
+ "outputs": [],
+ "id": "27e68f91"
+ },
+ {
+ "cell_type": "markdown",
+ "metadata": {
+ "id": "3e1771a9"
+ },
+ "source": [
+ "### Start Unsloth Studio"
+ ],
+ "id": "3e1771a9"
+ },
+ {
+ "cell_type": "code",
+ "metadata": {
+ "id": "277e431e"
+ },
+ "source": [
+ "import sys\n",
+ "sys.path.insert(0, \"/content/unsloth/studio/backend\")\n",
+ "from colab import start\n",
+ "\n",
+ "# On Colab, start() auto-opens a Cloudflare link and prints admin login credentials.\n",
+ "# Use the Cloudflare link above the ready card to open Studio (in-cell iframes often stay blank).\n",
+ "start()\n",
+ "\n",
+ "# To skip the Cloudflare tunnel and try the in-notebook proxy iframe only:\n",
+ "# start(cloudflare=False)"
+ ],
+ "execution_count": null,
+ "outputs": [],
+ "id": "277e431e"
+ },
+ {
+ "cell_type": "markdown",
+ "metadata": {
+ "id": "f2b0c6a1"
+ },
+ "source": [
+ "And we're done! If you have any questions on Unsloth, we have a [Discord](https://discord.gg/unsloth) channel! If you find any bugs or want to keep updated with the latest LLM stuff, or need help, join projects etc, feel free to join our Discord!\n",
+ "\n",
+ "Some other resources:\n",
+ "1. Looking to use Unsloth locally? Read our [Installation Guide](https://unsloth.ai/docs/get-started/install) for details on installing Unsloth on Windows, Docker, AMD, Intel GPUs.\n",
+ "2. Learn how to do Reinforcement Learning with our [RL Guide and notebooks](https://unsloth.ai/docs/get-started/reinforcement-learning-rl-guide).\n",
+ "3. Read our guides and notebooks for [Text-to-speech (TTS)](https://unsloth.ai/docs/basics/text-to-speech-tts-fine-tuning) and [vision](https://unsloth.ai/docs/basics/vision-fine-tuning) model support.\n",
+ "4. Explore our [LLM Tutorials Directory](https://unsloth.ai/docs/models/tutorials-how-to-fine-tune-and-run-llms) to find dedicated guides for each model.\n",
+ "5. Need help with Inference? Read our [Inference & Deployment page](https://unsloth.ai/docs/basics/inference-and-deployment) for details on using vLLM, llama.cpp, Ollama etc.\n",
+ "\n",
+ "
\n",
+ " \n",
+ " \n",
+ " \n",
+ "\n",
+ " Join Discord if you need help + ⭐️ Star us on Github ⭐️\n",
+ "\n",
+ " This notebook is licensed AGPL-3.0\n",
+ "
\n",
- "\n",
- "\n",
- " Join Discord if you need help + ⭐ Star us on Github ⭐\n",
- "
\n",
- "\n",
- "To install Unsloth Studio on your local device, follow [our guide](https://unsloth.ai/docs/new/unsloth-studio/install). Unsloth Studio is licensed [AGPL-3.0](https://github.com/unslothai/unsloth/blob/main/studio/LICENSE.AGPL-3.0).\n",
- "\n",
- "### Unsloth Studio\n",
- "\n",
- "Train and run open models with [**Unsloth Studio**](https://unsloth.ai/docs/new/unsloth-studio/start). NEW! Installation should now only take 2 mins!\n",
- "\n",
- "\n",
- "We are actively working on making Unsloth Studio install on Colab T4 GPUs faster.\n",
- "\n",
- "[Features](https://unsloth.ai/docs/new/unsloth-studio#features) • [Quickstart](https://unsloth.ai/docs/new/unsloth-studio/start) • [Data Recipes](https://unsloth.ai/docs/new/unsloth-studio/data-recipe) • [Unsloth Chat](https://unsloth.ai/docs/new/unsloth-studio/chat) • [Export](https://unsloth.ai/docs/new/unsloth-studio/export)"
- ]
- },
- {
- "cell_type": "markdown",
- "id": "e4206349",
- "metadata": {
- "id": "e4206349"
- },
- "source": [
- "
"
- ]
- },
- {
- "cell_type": "markdown",
- "id": "27da2957",
- "metadata": {
- "id": "27da2957"
- },
- "source": [
- "### Setup: Clone repo and run setup"
- ]
- },
- {
- "cell_type": "code",
- "execution_count": null,
- "id": "27e68f91",
- "metadata": {
- "id": "27e68f91"
- },
- "outputs": [],
- "source": "!git clone --depth 1 --branch main https://github.com/unslothai/unsloth.git\n%cd /content/unsloth\n!chmod +x studio/setup.sh && ./studio/setup.sh --local"
- },
- {
- "cell_type": "markdown",
- "id": "3e1771a9",
- "metadata": {
- "id": "3e1771a9"
- },
- "source": [
- "### Start Unsloth Studio"
- ]
- },
- {
- "cell_type": "code",
- "execution_count": null,
- "id": "277e431e",
- "metadata": {
- "id": "277e431e"
- },
- "outputs": [],
- "source": "import sys\nsys.path.insert(0, \"/content/unsloth/studio/backend\")\nfrom colab import start\n\n# Default: in-tab iframe only. start() blocks to keep the kernel alive.\nstart()\n\n# For a shareable Cloudflare link, replace start() above with:\n# start(cloudflare=True)"
- },
- {
- "cell_type": "markdown",
- "id": "f2b0c6a1",
- "metadata": {
- "id": "f2b0c6a1"
- },
- "source": [
- "And we're done! If you have any questions on Unsloth, we have a [Discord](https://discord.gg/unsloth) channel! If you find any bugs or want to keep updated with the latest LLM stuff, or need help, join projects etc, feel free to join our Discord!\n",
- "\n",
- "Some other resources:\n",
- "1. Looking to use Unsloth locally? Read our [Installation Guide](https://unsloth.ai/docs/get-started/install) for details on installing Unsloth on Windows, Docker, AMD, Intel GPUs.\n",
- "2. Learn how to do Reinforcement Learning with our [RL Guide and notebooks](https://unsloth.ai/docs/get-started/reinforcement-learning-rl-guide).\n",
- "3. Read our guides and notebooks for [Text-to-speech (TTS)](https://unsloth.ai/docs/basics/text-to-speech-tts-fine-tuning) and [vision](https://unsloth.ai/docs/basics/vision-fine-tuning) model support.\n",
- "4. Explore our [LLM Tutorials Directory](https://unsloth.ai/docs/models/tutorials-how-to-fine-tune-and-run-llms) to find dedicated guides for each model.\n",
- "5. Need help with Inference? Read our [Inference & Deployment page](https://unsloth.ai/docs/basics/inference-and-deployment) for details on using vLLM, llama.cpp, Ollama etc.\n",
- "\n",
- "
\n",
- " \n",
- " \n",
- " \n",
- "\n",
- " Join Discord if you need help + ⭐️ Star us on Github ⭐️\n",
- "\n",
- " This notebook is licensed AGPL-3.0\n",
- "
"
- ]
- }
- ],
- "metadata": {
- "accelerator": "GPU",
- "colab": {
- "gpuType": "T4",
- "provenance": [],
- "include_colab_link": true
- },
- "kernelspec": {
- "display_name": "Python 3",
- "name": "python3"
- },
- "language_info": {
- "name": "python"
- }
- },
- "nbformat": 4,
- "nbformat_minor": 5
+ "nbformat": 4,
+ "nbformat_minor": 5
}
\ No newline at end of file
diff --git a/studio/backend/assets/configs/full_finetune.yaml b/studio/backend/assets/configs/full_finetune.yaml
index e398515f61..98c45dd851 100644
--- a/studio/backend/assets/configs/full_finetune.yaml
+++ b/studio/backend/assets/configs/full_finetune.yaml
@@ -30,6 +30,7 @@ lora:
vision_all_linear: false
use_rslora: false
use_loftq: false
+ use_dora: false
finetune_vision_layers: true
finetune_language_layers: true
finetune_attention_modules: true
diff --git a/studio/backend/assets/configs/lora_text.yaml b/studio/backend/assets/configs/lora_text.yaml
index 9cb6b8c700..6c6a4d8839 100644
--- a/studio/backend/assets/configs/lora_text.yaml
+++ b/studio/backend/assets/configs/lora_text.yaml
@@ -30,6 +30,7 @@ lora:
vision_all_linear: false
use_rslora: false
use_loftq: false
+ use_dora: false
finetune_vision_layers: true
finetune_language_layers: true
finetune_attention_modules: true
diff --git a/studio/backend/assets/configs/model_defaults/default.yaml b/studio/backend/assets/configs/model_defaults/default.yaml
index 841e8ba166..e569031a31 100644
--- a/studio/backend/assets/configs/model_defaults/default.yaml
+++ b/studio/backend/assets/configs/model_defaults/default.yaml
@@ -33,6 +33,7 @@ lora:
- "down_proj"
use_rslora: false
use_loftq: false
+ use_dora: false
finetune_vision_layers: true
finetune_language_layers: true
finetune_attention_modules: true
diff --git a/studio/backend/assets/configs/model_defaults/embedding/unsloth_Qwen3-Embedding-0.6B.yaml b/studio/backend/assets/configs/model_defaults/embedding/unsloth_Qwen3-Embedding-0.6B.yaml
index f7b49c75b7..7ac1c83e04 100644
--- a/studio/backend/assets/configs/model_defaults/embedding/unsloth_Qwen3-Embedding-0.6B.yaml
+++ b/studio/backend/assets/configs/model_defaults/embedding/unsloth_Qwen3-Embedding-0.6B.yaml
@@ -34,6 +34,7 @@ lora:
- "down_proj"
use_rslora: false
use_loftq: false
+ use_dora: false
logging:
enable_wandb: false
diff --git a/studio/backend/assets/configs/model_defaults/embedding/unsloth_all-MiniLM-L6-v2.yaml b/studio/backend/assets/configs/model_defaults/embedding/unsloth_all-MiniLM-L6-v2.yaml
index be7da0f624..4cab9e9f96 100644
--- a/studio/backend/assets/configs/model_defaults/embedding/unsloth_all-MiniLM-L6-v2.yaml
+++ b/studio/backend/assets/configs/model_defaults/embedding/unsloth_all-MiniLM-L6-v2.yaml
@@ -30,6 +30,7 @@ lora:
- "query"
use_rslora: false
use_loftq: false
+ use_dora: false
logging:
enable_wandb: false
diff --git a/studio/backend/assets/configs/model_defaults/embedding/unsloth_bge-m3.yaml b/studio/backend/assets/configs/model_defaults/embedding/unsloth_bge-m3.yaml
index d9e49bc0d5..c1f1c2a344 100644
--- a/studio/backend/assets/configs/model_defaults/embedding/unsloth_bge-m3.yaml
+++ b/studio/backend/assets/configs/model_defaults/embedding/unsloth_bge-m3.yaml
@@ -30,6 +30,7 @@ lora:
- "value"
use_rslora: false
use_loftq: false
+ use_dora: false
logging:
enable_wandb: false
diff --git a/studio/backend/assets/configs/model_defaults/embedding/unsloth_embeddinggemma-300m.yaml b/studio/backend/assets/configs/model_defaults/embedding/unsloth_embeddinggemma-300m.yaml
index c3422d399f..7828feae81 100644
--- a/studio/backend/assets/configs/model_defaults/embedding/unsloth_embeddinggemma-300m.yaml
+++ b/studio/backend/assets/configs/model_defaults/embedding/unsloth_embeddinggemma-300m.yaml
@@ -33,6 +33,7 @@ lora:
- "down_proj"
use_rslora: false
use_loftq: false
+ use_dora: false
logging:
enable_wandb: false
diff --git a/studio/backend/assets/configs/model_defaults/embedding/unsloth_gte-modernbert-base.yaml b/studio/backend/assets/configs/model_defaults/embedding/unsloth_gte-modernbert-base.yaml
index 529a56a527..5a4028f15b 100644
--- a/studio/backend/assets/configs/model_defaults/embedding/unsloth_gte-modernbert-base.yaml
+++ b/studio/backend/assets/configs/model_defaults/embedding/unsloth_gte-modernbert-base.yaml
@@ -29,6 +29,7 @@ lora:
- "Wqkv"
use_rslora: false
use_loftq: false
+ use_dora: false
logging:
enable_wandb: false
diff --git a/studio/backend/assets/configs/model_defaults/ernie/unsloth_ERNIE-4.5-21B-A3B-PT.yaml b/studio/backend/assets/configs/model_defaults/ernie/unsloth_ERNIE-4.5-21B-A3B-PT.yaml
index 734115ec41..7645d11c98 100644
--- a/studio/backend/assets/configs/model_defaults/ernie/unsloth_ERNIE-4.5-21B-A3B-PT.yaml
+++ b/studio/backend/assets/configs/model_defaults/ernie/unsloth_ERNIE-4.5-21B-A3B-PT.yaml
@@ -34,6 +34,7 @@ lora:
- "down_proj"
use_rslora: false
use_loftq: false
+ use_dora: false
logging:
enable_wandb: false
diff --git a/studio/backend/assets/configs/model_defaults/ernie/unsloth_ERNIE-4.5-VL-28B-A3B-PT.yaml b/studio/backend/assets/configs/model_defaults/ernie/unsloth_ERNIE-4.5-VL-28B-A3B-PT.yaml
index 1032449e8c..b746235f1f 100644
--- a/studio/backend/assets/configs/model_defaults/ernie/unsloth_ERNIE-4.5-VL-28B-A3B-PT.yaml
+++ b/studio/backend/assets/configs/model_defaults/ernie/unsloth_ERNIE-4.5-VL-28B-A3B-PT.yaml
@@ -35,6 +35,7 @@ lora:
- "down_proj"
use_rslora: false
use_loftq: false
+ use_dora: false
finetune_vision_layers: true
finetune_language_layers: true
finetune_attention_modules: true
diff --git a/studio/backend/assets/configs/model_defaults/falcon/tiiuae_Falcon-H1-0.5B-Instruct.yaml b/studio/backend/assets/configs/model_defaults/falcon/tiiuae_Falcon-H1-0.5B-Instruct.yaml
index c8e5f35841..4964fea276 100644
--- a/studio/backend/assets/configs/model_defaults/falcon/tiiuae_Falcon-H1-0.5B-Instruct.yaml
+++ b/studio/backend/assets/configs/model_defaults/falcon/tiiuae_Falcon-H1-0.5B-Instruct.yaml
@@ -34,6 +34,7 @@ lora:
- "down_proj"
use_rslora: false
use_loftq: false
+ use_dora: false
logging:
enable_wandb: false
diff --git a/studio/backend/assets/configs/model_defaults/gemma/unsloth_codegemma-7b-bnb-4bit.yaml b/studio/backend/assets/configs/model_defaults/gemma/unsloth_codegemma-7b-bnb-4bit.yaml
index 251409c29d..e5f3344356 100644
--- a/studio/backend/assets/configs/model_defaults/gemma/unsloth_codegemma-7b-bnb-4bit.yaml
+++ b/studio/backend/assets/configs/model_defaults/gemma/unsloth_codegemma-7b-bnb-4bit.yaml
@@ -35,6 +35,7 @@ lora:
- "down_proj"
use_rslora: false
use_loftq: false
+ use_dora: false
logging:
enable_wandb: false
diff --git a/studio/backend/assets/configs/model_defaults/gemma/unsloth_functiongemma-270m-it.yaml b/studio/backend/assets/configs/model_defaults/gemma/unsloth_functiongemma-270m-it.yaml
index 89b1d7f938..71c61f383a 100644
--- a/studio/backend/assets/configs/model_defaults/gemma/unsloth_functiongemma-270m-it.yaml
+++ b/studio/backend/assets/configs/model_defaults/gemma/unsloth_functiongemma-270m-it.yaml
@@ -35,6 +35,7 @@ lora:
- "down_proj"
use_rslora: false
use_loftq: false
+ use_dora: false
logging:
enable_wandb: false
diff --git a/studio/backend/assets/configs/model_defaults/gemma/unsloth_gemma-2-27b-bnb-4bit.yaml b/studio/backend/assets/configs/model_defaults/gemma/unsloth_gemma-2-27b-bnb-4bit.yaml
index e3292b5972..3fe29cd800 100644
--- a/studio/backend/assets/configs/model_defaults/gemma/unsloth_gemma-2-27b-bnb-4bit.yaml
+++ b/studio/backend/assets/configs/model_defaults/gemma/unsloth_gemma-2-27b-bnb-4bit.yaml
@@ -33,6 +33,7 @@ lora:
- "down_proj"
use_rslora: false
use_loftq: false
+ use_dora: false
logging:
enable_wandb: false
diff --git a/studio/backend/assets/configs/model_defaults/gemma/unsloth_gemma-2-2b.yaml b/studio/backend/assets/configs/model_defaults/gemma/unsloth_gemma-2-2b.yaml
index 98fe497912..cd4e3e0c4d 100644
--- a/studio/backend/assets/configs/model_defaults/gemma/unsloth_gemma-2-2b.yaml
+++ b/studio/backend/assets/configs/model_defaults/gemma/unsloth_gemma-2-2b.yaml
@@ -34,6 +34,7 @@ lora:
- "down_proj"
use_rslora: false
use_loftq: false
+ use_dora: false
logging:
enable_wandb: false
diff --git a/studio/backend/assets/configs/model_defaults/gemma/unsloth_gemma-3-270m-it.yaml b/studio/backend/assets/configs/model_defaults/gemma/unsloth_gemma-3-270m-it.yaml
index bda5471643..97aa10e861 100644
--- a/studio/backend/assets/configs/model_defaults/gemma/unsloth_gemma-3-270m-it.yaml
+++ b/studio/backend/assets/configs/model_defaults/gemma/unsloth_gemma-3-270m-it.yaml
@@ -35,6 +35,7 @@ lora:
- "down_proj"
use_rslora: false
use_loftq: false
+ use_dora: false
logging:
enable_wandb: false
diff --git a/studio/backend/assets/configs/model_defaults/gemma/unsloth_gemma-3-27b-it.yaml b/studio/backend/assets/configs/model_defaults/gemma/unsloth_gemma-3-27b-it.yaml
index 18392568bd..a1b1640fa2 100644
--- a/studio/backend/assets/configs/model_defaults/gemma/unsloth_gemma-3-27b-it.yaml
+++ b/studio/backend/assets/configs/model_defaults/gemma/unsloth_gemma-3-27b-it.yaml
@@ -29,6 +29,7 @@ lora:
- "all-linear"
use_rslora: false
use_loftq: false
+ use_dora: false
finetune_vision_layers: true
finetune_language_layers: true
finetune_attention_modules: true
diff --git a/studio/backend/assets/configs/model_defaults/gemma/unsloth_gemma-3-4b-it.yaml b/studio/backend/assets/configs/model_defaults/gemma/unsloth_gemma-3-4b-it.yaml
index 434ac41b46..dbf60f04d4 100644
--- a/studio/backend/assets/configs/model_defaults/gemma/unsloth_gemma-3-4b-it.yaml
+++ b/studio/backend/assets/configs/model_defaults/gemma/unsloth_gemma-3-4b-it.yaml
@@ -29,6 +29,7 @@ lora:
- "all-linear"
use_rslora: false
use_loftq: false
+ use_dora: false
finetune_vision_layers: true
finetune_language_layers: true
finetune_attention_modules: true
diff --git a/studio/backend/assets/configs/model_defaults/gemma/unsloth_gemma-3-4b-pt.yaml b/studio/backend/assets/configs/model_defaults/gemma/unsloth_gemma-3-4b-pt.yaml
index 5f0a7b26ce..54c7dd6cd4 100644
--- a/studio/backend/assets/configs/model_defaults/gemma/unsloth_gemma-3-4b-pt.yaml
+++ b/studio/backend/assets/configs/model_defaults/gemma/unsloth_gemma-3-4b-pt.yaml
@@ -29,6 +29,7 @@ lora:
- "all-linear"
use_rslora: false
use_loftq: false
+ use_dora: false
finetune_vision_layers: true
finetune_language_layers: true
finetune_attention_modules: true
diff --git a/studio/backend/assets/configs/model_defaults/gemma/unsloth_gemma-3n-E4B-it.yaml b/studio/backend/assets/configs/model_defaults/gemma/unsloth_gemma-3n-E4B-it.yaml
index dd5ae51ab0..119440a585 100644
--- a/studio/backend/assets/configs/model_defaults/gemma/unsloth_gemma-3n-E4B-it.yaml
+++ b/studio/backend/assets/configs/model_defaults/gemma/unsloth_gemma-3n-E4B-it.yaml
@@ -29,6 +29,7 @@ lora:
- "all-linear"
use_rslora: false
use_loftq: false
+ use_dora: false
finetune_vision_layers: true
finetune_language_layers: true
finetune_attention_modules: true
diff --git a/studio/backend/assets/configs/model_defaults/gemma/unsloth_gemma-3n-E4B.yaml b/studio/backend/assets/configs/model_defaults/gemma/unsloth_gemma-3n-E4B.yaml
index e53e163a04..d08e5e9547 100644
--- a/studio/backend/assets/configs/model_defaults/gemma/unsloth_gemma-3n-E4B.yaml
+++ b/studio/backend/assets/configs/model_defaults/gemma/unsloth_gemma-3n-E4B.yaml
@@ -29,6 +29,7 @@ lora:
- "all-linear"
use_rslora: false
use_loftq: false
+ use_dora: false
finetune_vision_layers: true
finetune_language_layers: true
finetune_attention_modules: true
diff --git a/studio/backend/assets/configs/model_defaults/gemma/unsloth_gemma-4-26B-A4B-it.yaml b/studio/backend/assets/configs/model_defaults/gemma/unsloth_gemma-4-26B-A4B-it.yaml
index ebe344e382..a266d7a39b 100644
--- a/studio/backend/assets/configs/model_defaults/gemma/unsloth_gemma-4-26B-A4B-it.yaml
+++ b/studio/backend/assets/configs/model_defaults/gemma/unsloth_gemma-4-26B-A4B-it.yaml
@@ -26,6 +26,7 @@ lora:
- "all-linear"
use_rslora: false
use_loftq: false
+ use_dora: false
finetune_vision_layers: true
finetune_language_layers: true
finetune_attention_modules: true
diff --git a/studio/backend/assets/configs/model_defaults/gemma/unsloth_gemma-4-26B-A4B.yaml b/studio/backend/assets/configs/model_defaults/gemma/unsloth_gemma-4-26B-A4B.yaml
index fb89a07133..970cac3259 100644
--- a/studio/backend/assets/configs/model_defaults/gemma/unsloth_gemma-4-26B-A4B.yaml
+++ b/studio/backend/assets/configs/model_defaults/gemma/unsloth_gemma-4-26B-A4B.yaml
@@ -26,6 +26,7 @@ lora:
- "all-linear"
use_rslora: false
use_loftq: false
+ use_dora: false
finetune_vision_layers: true
finetune_language_layers: true
finetune_attention_modules: true
diff --git a/studio/backend/assets/configs/model_defaults/gemma/unsloth_gemma-4-31B-it.yaml b/studio/backend/assets/configs/model_defaults/gemma/unsloth_gemma-4-31B-it.yaml
index 4a089992ac..5bba4ccdc0 100644
--- a/studio/backend/assets/configs/model_defaults/gemma/unsloth_gemma-4-31B-it.yaml
+++ b/studio/backend/assets/configs/model_defaults/gemma/unsloth_gemma-4-31B-it.yaml
@@ -26,6 +26,7 @@ lora:
- "all-linear"
use_rslora: false
use_loftq: false
+ use_dora: false
finetune_vision_layers: true
finetune_language_layers: true
finetune_attention_modules: true
diff --git a/studio/backend/assets/configs/model_defaults/gemma/unsloth_gemma-4-31B.yaml b/studio/backend/assets/configs/model_defaults/gemma/unsloth_gemma-4-31B.yaml
index ae7524b7c6..ac5c6eca22 100644
--- a/studio/backend/assets/configs/model_defaults/gemma/unsloth_gemma-4-31B.yaml
+++ b/studio/backend/assets/configs/model_defaults/gemma/unsloth_gemma-4-31B.yaml
@@ -26,6 +26,7 @@ lora:
- "all-linear"
use_rslora: false
use_loftq: false
+ use_dora: false
finetune_vision_layers: true
finetune_language_layers: true
finetune_attention_modules: true
diff --git a/studio/backend/assets/configs/model_defaults/gemma/unsloth_gemma-4-E2B-it.yaml b/studio/backend/assets/configs/model_defaults/gemma/unsloth_gemma-4-E2B-it.yaml
index 10c1abd8a5..68c2d35644 100644
--- a/studio/backend/assets/configs/model_defaults/gemma/unsloth_gemma-4-E2B-it.yaml
+++ b/studio/backend/assets/configs/model_defaults/gemma/unsloth_gemma-4-E2B-it.yaml
@@ -26,6 +26,7 @@ lora:
- "all-linear"
use_rslora: false
use_loftq: false
+ use_dora: false
finetune_vision_layers: true
finetune_language_layers: true
finetune_attention_modules: true
diff --git a/studio/backend/assets/configs/model_defaults/gemma/unsloth_gemma-4-E2B.yaml b/studio/backend/assets/configs/model_defaults/gemma/unsloth_gemma-4-E2B.yaml
index fb5c1d9dea..175f9c0f17 100644
--- a/studio/backend/assets/configs/model_defaults/gemma/unsloth_gemma-4-E2B.yaml
+++ b/studio/backend/assets/configs/model_defaults/gemma/unsloth_gemma-4-E2B.yaml
@@ -26,6 +26,7 @@ lora:
- "all-linear"
use_rslora: false
use_loftq: false
+ use_dora: false
finetune_vision_layers: true
finetune_language_layers: true
finetune_attention_modules: true
diff --git a/studio/backend/assets/configs/model_defaults/gemma/unsloth_gemma-4-E4B-it.yaml b/studio/backend/assets/configs/model_defaults/gemma/unsloth_gemma-4-E4B-it.yaml
index 189e5dc6b2..4f3834e7c0 100644
--- a/studio/backend/assets/configs/model_defaults/gemma/unsloth_gemma-4-E4B-it.yaml
+++ b/studio/backend/assets/configs/model_defaults/gemma/unsloth_gemma-4-E4B-it.yaml
@@ -26,6 +26,7 @@ lora:
- "all-linear"
use_rslora: false
use_loftq: false
+ use_dora: false
finetune_vision_layers: true
finetune_language_layers: true
finetune_attention_modules: true
diff --git a/studio/backend/assets/configs/model_defaults/gemma/unsloth_gemma-4-E4B.yaml b/studio/backend/assets/configs/model_defaults/gemma/unsloth_gemma-4-E4B.yaml
index aa51440b6a..d6d97f7e44 100644
--- a/studio/backend/assets/configs/model_defaults/gemma/unsloth_gemma-4-E4B.yaml
+++ b/studio/backend/assets/configs/model_defaults/gemma/unsloth_gemma-4-E4B.yaml
@@ -26,6 +26,7 @@ lora:
- "all-linear"
use_rslora: false
use_loftq: false
+ use_dora: false
finetune_vision_layers: true
finetune_language_layers: true
finetune_attention_modules: true
diff --git a/studio/backend/assets/configs/model_defaults/gpt-oss/unsloth_gpt-oss-120b.yaml b/studio/backend/assets/configs/model_defaults/gpt-oss/unsloth_gpt-oss-120b.yaml
index e2d67bcb0b..4f1f54a4e6 100644
--- a/studio/backend/assets/configs/model_defaults/gpt-oss/unsloth_gpt-oss-120b.yaml
+++ b/studio/backend/assets/configs/model_defaults/gpt-oss/unsloth_gpt-oss-120b.yaml
@@ -35,6 +35,7 @@ lora:
- "down_proj"
use_rslora: false
use_loftq: false
+ use_dora: false
logging:
enable_wandb: false
diff --git a/studio/backend/assets/configs/model_defaults/gpt-oss/unsloth_gpt-oss-20b.yaml b/studio/backend/assets/configs/model_defaults/gpt-oss/unsloth_gpt-oss-20b.yaml
index aa436117a1..127700b53b 100644
--- a/studio/backend/assets/configs/model_defaults/gpt-oss/unsloth_gpt-oss-20b.yaml
+++ b/studio/backend/assets/configs/model_defaults/gpt-oss/unsloth_gpt-oss-20b.yaml
@@ -35,6 +35,7 @@ lora:
- "down_proj"
use_rslora: false
use_loftq: false
+ use_dora: false
logging:
enable_wandb: false
diff --git a/studio/backend/assets/configs/model_defaults/granite/unsloth_granite-4.0-350m-unsloth-bnb-4bit.yaml b/studio/backend/assets/configs/model_defaults/granite/unsloth_granite-4.0-350m-unsloth-bnb-4bit.yaml
index 3f2cb84a94..2412b3accf 100644
--- a/studio/backend/assets/configs/model_defaults/granite/unsloth_granite-4.0-350m-unsloth-bnb-4bit.yaml
+++ b/studio/backend/assets/configs/model_defaults/granite/unsloth_granite-4.0-350m-unsloth-bnb-4bit.yaml
@@ -37,6 +37,7 @@ lora:
- "shared_mlp.output_linear"
use_rslora: false
use_loftq: false
+ use_dora: false
logging:
enable_wandb: false
diff --git a/studio/backend/assets/configs/model_defaults/granite/unsloth_granite-4.0-h-micro.yaml b/studio/backend/assets/configs/model_defaults/granite/unsloth_granite-4.0-h-micro.yaml
index ab756fe764..81b59c4323 100644
--- a/studio/backend/assets/configs/model_defaults/granite/unsloth_granite-4.0-h-micro.yaml
+++ b/studio/backend/assets/configs/model_defaults/granite/unsloth_granite-4.0-h-micro.yaml
@@ -37,6 +37,7 @@ lora:
- "shared_mlp.output_linear"
use_rslora: false
use_loftq: false
+ use_dora: false
logging:
enable_wandb: false
diff --git a/studio/backend/assets/configs/model_defaults/llama/unsloth_Llama-3.2-11B-Vision-Instruct.yaml b/studio/backend/assets/configs/model_defaults/llama/unsloth_Llama-3.2-11B-Vision-Instruct.yaml
index 1a7a91e56f..6110d84a6c 100644
--- a/studio/backend/assets/configs/model_defaults/llama/unsloth_Llama-3.2-11B-Vision-Instruct.yaml
+++ b/studio/backend/assets/configs/model_defaults/llama/unsloth_Llama-3.2-11B-Vision-Instruct.yaml
@@ -29,6 +29,7 @@ lora:
- "all-linear"
use_rslora: false
use_loftq: false
+ use_dora: false
finetune_vision_layers: true
finetune_language_layers: true
finetune_attention_modules: true
diff --git a/studio/backend/assets/configs/model_defaults/llama/unsloth_Llama-3.2-1B-Instruct.yaml b/studio/backend/assets/configs/model_defaults/llama/unsloth_Llama-3.2-1B-Instruct.yaml
index 7c7bb8dc3e..3c7fc7f238 100644
--- a/studio/backend/assets/configs/model_defaults/llama/unsloth_Llama-3.2-1B-Instruct.yaml
+++ b/studio/backend/assets/configs/model_defaults/llama/unsloth_Llama-3.2-1B-Instruct.yaml
@@ -34,6 +34,7 @@ lora:
- "down_proj"
use_rslora: false
use_loftq: false
+ use_dora: false
logging:
enable_wandb: false
diff --git a/studio/backend/assets/configs/model_defaults/llama/unsloth_Llama-3.2-3B-Instruct.yaml b/studio/backend/assets/configs/model_defaults/llama/unsloth_Llama-3.2-3B-Instruct.yaml
index f73b0c09b6..2b0977e435 100644
--- a/studio/backend/assets/configs/model_defaults/llama/unsloth_Llama-3.2-3B-Instruct.yaml
+++ b/studio/backend/assets/configs/model_defaults/llama/unsloth_Llama-3.2-3B-Instruct.yaml
@@ -35,6 +35,7 @@ lora:
- "down_proj"
use_rslora: false
use_loftq: false
+ use_dora: false
logging:
enable_wandb: false
diff --git a/studio/backend/assets/configs/model_defaults/llama/unsloth_Llama-3.3-70B-Instruct.yaml b/studio/backend/assets/configs/model_defaults/llama/unsloth_Llama-3.3-70B-Instruct.yaml
index ffefb29e24..1742c04a06 100644
--- a/studio/backend/assets/configs/model_defaults/llama/unsloth_Llama-3.3-70B-Instruct.yaml
+++ b/studio/backend/assets/configs/model_defaults/llama/unsloth_Llama-3.3-70B-Instruct.yaml
@@ -35,6 +35,7 @@ lora:
- "down_proj"
use_rslora: false
use_loftq: false
+ use_dora: false
logging:
enable_wandb: false
diff --git a/studio/backend/assets/configs/model_defaults/llama/unsloth_Meta-Llama-3.1-70B-bnb-4bit.yaml b/studio/backend/assets/configs/model_defaults/llama/unsloth_Meta-Llama-3.1-70B-bnb-4bit.yaml
index cd986a6da1..f33726b0dd 100644
--- a/studio/backend/assets/configs/model_defaults/llama/unsloth_Meta-Llama-3.1-70B-bnb-4bit.yaml
+++ b/studio/backend/assets/configs/model_defaults/llama/unsloth_Meta-Llama-3.1-70B-bnb-4bit.yaml
@@ -34,6 +34,7 @@ lora:
- "down_proj"
use_rslora: false
use_loftq: false
+ use_dora: false
logging:
enable_wandb: false
diff --git a/studio/backend/assets/configs/model_defaults/llama/unsloth_Meta-Llama-3.1-8B-Instruct-bnb-4bit.yaml b/studio/backend/assets/configs/model_defaults/llama/unsloth_Meta-Llama-3.1-8B-Instruct-bnb-4bit.yaml
index 55dd3144c6..79b30bd758 100644
--- a/studio/backend/assets/configs/model_defaults/llama/unsloth_Meta-Llama-3.1-8B-Instruct-bnb-4bit.yaml
+++ b/studio/backend/assets/configs/model_defaults/llama/unsloth_Meta-Llama-3.1-8B-Instruct-bnb-4bit.yaml
@@ -34,6 +34,7 @@ lora:
- "down_proj"
use_rslora: false
use_loftq: false
+ use_dora: false
logging:
enable_wandb: false
diff --git a/studio/backend/assets/configs/model_defaults/llama/unsloth_llama-3-8b-Instruct-bnb-4bit.yaml b/studio/backend/assets/configs/model_defaults/llama/unsloth_llama-3-8b-Instruct-bnb-4bit.yaml
index 8c9cb07fb9..4ee9a5a8ed 100644
--- a/studio/backend/assets/configs/model_defaults/llama/unsloth_llama-3-8b-Instruct-bnb-4bit.yaml
+++ b/studio/backend/assets/configs/model_defaults/llama/unsloth_llama-3-8b-Instruct-bnb-4bit.yaml
@@ -34,6 +34,7 @@ lora:
- "down_proj"
use_rslora: false
use_loftq: false
+ use_dora: false
logging:
enable_wandb: false
diff --git a/studio/backend/assets/configs/model_defaults/llama/unsloth_llama-3-8b-bnb-4bit.yaml b/studio/backend/assets/configs/model_defaults/llama/unsloth_llama-3-8b-bnb-4bit.yaml
index 32441c5674..da20663688 100644
--- a/studio/backend/assets/configs/model_defaults/llama/unsloth_llama-3-8b-bnb-4bit.yaml
+++ b/studio/backend/assets/configs/model_defaults/llama/unsloth_llama-3-8b-bnb-4bit.yaml
@@ -34,6 +34,7 @@ lora:
- "down_proj"
use_rslora: false
use_loftq: false
+ use_dora: false
logging:
enable_wandb: false
diff --git a/studio/backend/assets/configs/model_defaults/llasa/unsloth_Llasa-3B.yaml b/studio/backend/assets/configs/model_defaults/llasa/unsloth_Llasa-3B.yaml
index 6bba9c9633..30e4440afb 100644
--- a/studio/backend/assets/configs/model_defaults/llasa/unsloth_Llasa-3B.yaml
+++ b/studio/backend/assets/configs/model_defaults/llasa/unsloth_Llasa-3B.yaml
@@ -30,6 +30,7 @@ lora:
- "v_proj"
use_rslora: false
use_loftq: false
+ use_dora: false
logging:
enable_wandb: false
diff --git a/studio/backend/assets/configs/model_defaults/mistral/unsloth_Magistral-Small-2509-unsloth-bnb-4bit.yaml b/studio/backend/assets/configs/model_defaults/mistral/unsloth_Magistral-Small-2509-unsloth-bnb-4bit.yaml
index f9833ce705..9bb0a93e63 100644
--- a/studio/backend/assets/configs/model_defaults/mistral/unsloth_Magistral-Small-2509-unsloth-bnb-4bit.yaml
+++ b/studio/backend/assets/configs/model_defaults/mistral/unsloth_Magistral-Small-2509-unsloth-bnb-4bit.yaml
@@ -35,6 +35,7 @@ lora:
- "down_proj"
use_rslora: false
use_loftq: false
+ use_dora: false
finetune_vision_layers: true
finetune_language_layers: true
finetune_attention_modules: true
diff --git a/studio/backend/assets/configs/model_defaults/mistral/unsloth_Ministral-3-3B-Instruct-2512.yaml b/studio/backend/assets/configs/model_defaults/mistral/unsloth_Ministral-3-3B-Instruct-2512.yaml
index 0ba857cd40..ded3607a14 100644
--- a/studio/backend/assets/configs/model_defaults/mistral/unsloth_Ministral-3-3B-Instruct-2512.yaml
+++ b/studio/backend/assets/configs/model_defaults/mistral/unsloth_Ministral-3-3B-Instruct-2512.yaml
@@ -35,6 +35,7 @@ lora:
- "down_proj"
use_rslora: false
use_loftq: false
+ use_dora: false
finetune_vision_layers: true
finetune_language_layers: true
finetune_attention_modules: true
diff --git a/studio/backend/assets/configs/model_defaults/mistral/unsloth_Mistral-Nemo-Base-2407-bnb-4bit.yaml b/studio/backend/assets/configs/model_defaults/mistral/unsloth_Mistral-Nemo-Base-2407-bnb-4bit.yaml
index 3476f2dd6d..2ac72f1c88 100644
--- a/studio/backend/assets/configs/model_defaults/mistral/unsloth_Mistral-Nemo-Base-2407-bnb-4bit.yaml
+++ b/studio/backend/assets/configs/model_defaults/mistral/unsloth_Mistral-Nemo-Base-2407-bnb-4bit.yaml
@@ -34,6 +34,7 @@ lora:
- "down_proj"
use_rslora: false
use_loftq: false
+ use_dora: false
logging:
enable_wandb: false
diff --git a/studio/backend/assets/configs/model_defaults/mistral/unsloth_Mistral-Small-Instruct-2409.yaml b/studio/backend/assets/configs/model_defaults/mistral/unsloth_Mistral-Small-Instruct-2409.yaml
index eda04d21f9..a087ced1f3 100644
--- a/studio/backend/assets/configs/model_defaults/mistral/unsloth_Mistral-Small-Instruct-2409.yaml
+++ b/studio/backend/assets/configs/model_defaults/mistral/unsloth_Mistral-Small-Instruct-2409.yaml
@@ -34,6 +34,7 @@ lora:
- "down_proj"
use_rslora: false
use_loftq: false
+ use_dora: false
logging:
enable_wandb: false
diff --git a/studio/backend/assets/configs/model_defaults/mistral/unsloth_Pixtral-12B-2409.yaml b/studio/backend/assets/configs/model_defaults/mistral/unsloth_Pixtral-12B-2409.yaml
index bcd0d20c8c..c9811f4f06 100644
--- a/studio/backend/assets/configs/model_defaults/mistral/unsloth_Pixtral-12B-2409.yaml
+++ b/studio/backend/assets/configs/model_defaults/mistral/unsloth_Pixtral-12B-2409.yaml
@@ -29,6 +29,7 @@ lora:
- "all-linear"
use_rslora: false
use_loftq: false
+ use_dora: false
finetune_vision_layers: true
finetune_language_layers: true
finetune_attention_modules: false
diff --git a/studio/backend/assets/configs/model_defaults/mistral/unsloth_mistral-7b-instruct-v0.3-bnb-4bit.yaml b/studio/backend/assets/configs/model_defaults/mistral/unsloth_mistral-7b-instruct-v0.3-bnb-4bit.yaml
index 34a033e32f..e3659d9fb0 100644
--- a/studio/backend/assets/configs/model_defaults/mistral/unsloth_mistral-7b-instruct-v0.3-bnb-4bit.yaml
+++ b/studio/backend/assets/configs/model_defaults/mistral/unsloth_mistral-7b-instruct-v0.3-bnb-4bit.yaml
@@ -34,6 +34,7 @@ lora:
- "down_proj"
use_rslora: false
use_loftq: false
+ use_dora: false
logging:
enable_wandb: false
diff --git a/studio/backend/assets/configs/model_defaults/mistral/unsloth_mistral-7b-v0.3-bnb-4bit.yaml b/studio/backend/assets/configs/model_defaults/mistral/unsloth_mistral-7b-v0.3-bnb-4bit.yaml
index 98105eaf38..ee17efc54d 100644
--- a/studio/backend/assets/configs/model_defaults/mistral/unsloth_mistral-7b-v0.3-bnb-4bit.yaml
+++ b/studio/backend/assets/configs/model_defaults/mistral/unsloth_mistral-7b-v0.3-bnb-4bit.yaml
@@ -33,6 +33,7 @@ lora:
- "down_proj"
use_rslora: false
use_loftq: false
+ use_dora: false
logging:
enable_wandb: false
diff --git a/studio/backend/assets/configs/model_defaults/other/OuteAI_Llama-OuteTTS-1.0-1B.yaml b/studio/backend/assets/configs/model_defaults/other/OuteAI_Llama-OuteTTS-1.0-1B.yaml
index 72b5b018e1..ef836b9b55 100644
--- a/studio/backend/assets/configs/model_defaults/other/OuteAI_Llama-OuteTTS-1.0-1B.yaml
+++ b/studio/backend/assets/configs/model_defaults/other/OuteAI_Llama-OuteTTS-1.0-1B.yaml
@@ -33,6 +33,7 @@ lora:
- "v_proj"
use_rslora: false
use_loftq: false
+ use_dora: false
logging:
enable_wandb: false
diff --git a/studio/backend/assets/configs/model_defaults/other/Spark-TTS-0.5B_LLM.yaml b/studio/backend/assets/configs/model_defaults/other/Spark-TTS-0.5B_LLM.yaml
index d20751b0c7..c80fad35a8 100644
--- a/studio/backend/assets/configs/model_defaults/other/Spark-TTS-0.5B_LLM.yaml
+++ b/studio/backend/assets/configs/model_defaults/other/Spark-TTS-0.5B_LLM.yaml
@@ -38,6 +38,7 @@ lora:
- "down_proj"
use_rslora: false
use_loftq: false
+ use_dora: false
logging:
enable_wandb: false
diff --git a/studio/backend/assets/configs/model_defaults/other/sesame_csm-1b.yaml b/studio/backend/assets/configs/model_defaults/other/sesame_csm-1b.yaml
index 8a80282a2a..034b5bd131 100644
--- a/studio/backend/assets/configs/model_defaults/other/sesame_csm-1b.yaml
+++ b/studio/backend/assets/configs/model_defaults/other/sesame_csm-1b.yaml
@@ -37,6 +37,7 @@ lora:
- "down_proj"
use_rslora: false
use_loftq: false
+ use_dora: false
logging:
enable_wandb: false
diff --git a/studio/backend/assets/configs/model_defaults/other/unsloth_GLM-4.7-Flash.yaml b/studio/backend/assets/configs/model_defaults/other/unsloth_GLM-4.7-Flash.yaml
index a973c2d4e4..d1a226be79 100644
--- a/studio/backend/assets/configs/model_defaults/other/unsloth_GLM-4.7-Flash.yaml
+++ b/studio/backend/assets/configs/model_defaults/other/unsloth_GLM-4.7-Flash.yaml
@@ -35,6 +35,7 @@ lora:
- "out_proj"
use_rslora: false
use_loftq: false
+ use_dora: false
logging:
enable_wandb: false
diff --git a/studio/backend/assets/configs/model_defaults/other/unsloth_LFM2-1.2B.yaml b/studio/backend/assets/configs/model_defaults/other/unsloth_LFM2-1.2B.yaml
index b0feafbd6e..1b8df5ced9 100644
--- a/studio/backend/assets/configs/model_defaults/other/unsloth_LFM2-1.2B.yaml
+++ b/studio/backend/assets/configs/model_defaults/other/unsloth_LFM2-1.2B.yaml
@@ -29,6 +29,7 @@ lora:
- "all-linear"
use_rslora: false
use_loftq: false
+ use_dora: false
logging:
enable_wandb: false
diff --git a/studio/backend/assets/configs/model_defaults/other/unsloth_Nemotron-3-Nano-30B-A3B.yaml b/studio/backend/assets/configs/model_defaults/other/unsloth_Nemotron-3-Nano-30B-A3B.yaml
index 2c44c91eab..cecab7f083 100644
--- a/studio/backend/assets/configs/model_defaults/other/unsloth_Nemotron-3-Nano-30B-A3B.yaml
+++ b/studio/backend/assets/configs/model_defaults/other/unsloth_Nemotron-3-Nano-30B-A3B.yaml
@@ -37,6 +37,7 @@ lora:
- "out_proj"
use_rslora: false
use_loftq: false
+ use_dora: false
logging:
enable_wandb: false
diff --git a/studio/backend/assets/configs/model_defaults/other/unsloth_PaddleOCR-VL.yaml b/studio/backend/assets/configs/model_defaults/other/unsloth_PaddleOCR-VL.yaml
index e1fbc08e4d..730be338cf 100644
--- a/studio/backend/assets/configs/model_defaults/other/unsloth_PaddleOCR-VL.yaml
+++ b/studio/backend/assets/configs/model_defaults/other/unsloth_PaddleOCR-VL.yaml
@@ -35,6 +35,7 @@ lora:
- "down_proj"
use_rslora: false
use_loftq: false
+ use_dora: false
finetune_vision_layers: true
finetune_language_layers: true
finetune_attention_modules: true
diff --git a/studio/backend/assets/configs/model_defaults/other/unsloth_answerdotai_ModernBERT-large.yaml b/studio/backend/assets/configs/model_defaults/other/unsloth_answerdotai_ModernBERT-large.yaml
index 2abdfd8ac3..a70ac0bd49 100644
--- a/studio/backend/assets/configs/model_defaults/other/unsloth_answerdotai_ModernBERT-large.yaml
+++ b/studio/backend/assets/configs/model_defaults/other/unsloth_answerdotai_ModernBERT-large.yaml
@@ -33,6 +33,7 @@ lora:
- "down_proj"
use_rslora: false
use_loftq: false
+ use_dora: false
logging:
enable_wandb: false
diff --git a/studio/backend/assets/configs/model_defaults/other/unsloth_orpheus-3b-0.1-ft.yaml b/studio/backend/assets/configs/model_defaults/other/unsloth_orpheus-3b-0.1-ft.yaml
index 5a3c4abb48..90ead037f6 100644
--- a/studio/backend/assets/configs/model_defaults/other/unsloth_orpheus-3b-0.1-ft.yaml
+++ b/studio/backend/assets/configs/model_defaults/other/unsloth_orpheus-3b-0.1-ft.yaml
@@ -38,6 +38,7 @@ lora:
- "down_proj"
use_rslora: false
use_loftq: false
+ use_dora: false
logging:
enable_wandb: false
diff --git a/studio/backend/assets/configs/model_defaults/other/unsloth_tinyllama-bnb-4bit.yaml b/studio/backend/assets/configs/model_defaults/other/unsloth_tinyllama-bnb-4bit.yaml
index a6ce27620f..a97c557c31 100644
--- a/studio/backend/assets/configs/model_defaults/other/unsloth_tinyllama-bnb-4bit.yaml
+++ b/studio/backend/assets/configs/model_defaults/other/unsloth_tinyllama-bnb-4bit.yaml
@@ -34,6 +34,7 @@ lora:
- "down_proj"
use_rslora: false
use_loftq: false
+ use_dora: false
logging:
enable_wandb: false
diff --git a/studio/backend/assets/configs/model_defaults/other/unsloth_whisper-large-v3.yaml b/studio/backend/assets/configs/model_defaults/other/unsloth_whisper-large-v3.yaml
index 050774a8cd..6855ed6a35 100644
--- a/studio/backend/assets/configs/model_defaults/other/unsloth_whisper-large-v3.yaml
+++ b/studio/backend/assets/configs/model_defaults/other/unsloth_whisper-large-v3.yaml
@@ -33,6 +33,7 @@ lora:
- "v_proj"
use_rslora: false
use_loftq: false
+ use_dora: false
logging:
enable_wandb: false
diff --git a/studio/backend/assets/configs/model_defaults/phi/unsloth_Phi-3-medium-4k-instruct.yaml b/studio/backend/assets/configs/model_defaults/phi/unsloth_Phi-3-medium-4k-instruct.yaml
index c574714d78..1933fed2ba 100644
--- a/studio/backend/assets/configs/model_defaults/phi/unsloth_Phi-3-medium-4k-instruct.yaml
+++ b/studio/backend/assets/configs/model_defaults/phi/unsloth_Phi-3-medium-4k-instruct.yaml
@@ -34,6 +34,7 @@ lora:
- "down_proj"
use_rslora: false
use_loftq: false
+ use_dora: false
logging:
enable_wandb: false
diff --git a/studio/backend/assets/configs/model_defaults/phi/unsloth_Phi-3.5-mini-instruct.yaml b/studio/backend/assets/configs/model_defaults/phi/unsloth_Phi-3.5-mini-instruct.yaml
index e803c842b3..fda4e64158 100644
--- a/studio/backend/assets/configs/model_defaults/phi/unsloth_Phi-3.5-mini-instruct.yaml
+++ b/studio/backend/assets/configs/model_defaults/phi/unsloth_Phi-3.5-mini-instruct.yaml
@@ -34,6 +34,7 @@ lora:
- "down_proj"
use_rslora: false
use_loftq: false
+ use_dora: false
logging:
enable_wandb: false
diff --git a/studio/backend/assets/configs/model_defaults/phi/unsloth_Phi-4.yaml b/studio/backend/assets/configs/model_defaults/phi/unsloth_Phi-4.yaml
index 4de3d9437d..c3910e3e5b 100644
--- a/studio/backend/assets/configs/model_defaults/phi/unsloth_Phi-4.yaml
+++ b/studio/backend/assets/configs/model_defaults/phi/unsloth_Phi-4.yaml
@@ -35,6 +35,7 @@ lora:
- "down_proj"
use_rslora: false
use_loftq: false
+ use_dora: false
logging:
enable_wandb: false
diff --git a/studio/backend/assets/configs/model_defaults/qwen/imdatta0_tiny_qwen3_moe_2.8B_0.7B.yaml b/studio/backend/assets/configs/model_defaults/qwen/imdatta0_tiny_qwen3_moe_2.8B_0.7B.yaml
index bb75b3ce52..765ffee938 100644
--- a/studio/backend/assets/configs/model_defaults/qwen/imdatta0_tiny_qwen3_moe_2.8B_0.7B.yaml
+++ b/studio/backend/assets/configs/model_defaults/qwen/imdatta0_tiny_qwen3_moe_2.8B_0.7B.yaml
@@ -36,6 +36,7 @@ lora:
- "gate_up_proj"
use_rslora: false
use_loftq: false
+ use_dora: false
logging:
enable_wandb: false
diff --git a/studio/backend/assets/configs/model_defaults/qwen/unsloth_Qwen2-7B.yaml b/studio/backend/assets/configs/model_defaults/qwen/unsloth_Qwen2-7B.yaml
index c305d328c2..39b30e9cee 100644
--- a/studio/backend/assets/configs/model_defaults/qwen/unsloth_Qwen2-7B.yaml
+++ b/studio/backend/assets/configs/model_defaults/qwen/unsloth_Qwen2-7B.yaml
@@ -34,6 +34,7 @@ lora:
- "down_proj"
use_rslora: false
use_loftq: false
+ use_dora: false
logging:
enable_wandb: false
diff --git a/studio/backend/assets/configs/model_defaults/qwen/unsloth_Qwen2-VL-7B-Instruct.yaml b/studio/backend/assets/configs/model_defaults/qwen/unsloth_Qwen2-VL-7B-Instruct.yaml
index 6cee3d0949..f97e525798 100644
--- a/studio/backend/assets/configs/model_defaults/qwen/unsloth_Qwen2-VL-7B-Instruct.yaml
+++ b/studio/backend/assets/configs/model_defaults/qwen/unsloth_Qwen2-VL-7B-Instruct.yaml
@@ -29,6 +29,7 @@ lora:
- "all-linear"
use_rslora: false
use_loftq: false
+ use_dora: false
finetune_vision_layers: true
finetune_language_layers: true
finetune_attention_modules: true
diff --git a/studio/backend/assets/configs/model_defaults/qwen/unsloth_Qwen2.5-1.5B-Instruct.yaml b/studio/backend/assets/configs/model_defaults/qwen/unsloth_Qwen2.5-1.5B-Instruct.yaml
index 20ba81df2c..e19b94ede2 100644
--- a/studio/backend/assets/configs/model_defaults/qwen/unsloth_Qwen2.5-1.5B-Instruct.yaml
+++ b/studio/backend/assets/configs/model_defaults/qwen/unsloth_Qwen2.5-1.5B-Instruct.yaml
@@ -34,6 +34,7 @@ lora:
- "down_proj"
use_rslora: false
use_loftq: false
+ use_dora: false
logging:
enable_wandb: false
diff --git a/studio/backend/assets/configs/model_defaults/qwen/unsloth_Qwen2.5-7B.yaml b/studio/backend/assets/configs/model_defaults/qwen/unsloth_Qwen2.5-7B.yaml
index 9930786c24..982f54b32f 100644
--- a/studio/backend/assets/configs/model_defaults/qwen/unsloth_Qwen2.5-7B.yaml
+++ b/studio/backend/assets/configs/model_defaults/qwen/unsloth_Qwen2.5-7B.yaml
@@ -34,6 +34,7 @@ lora:
- "down_proj"
use_rslora: false
use_loftq: false
+ use_dora: false
logging:
enable_wandb: false
diff --git a/studio/backend/assets/configs/model_defaults/qwen/unsloth_Qwen2.5-Coder-1.5B-Instruct.yaml b/studio/backend/assets/configs/model_defaults/qwen/unsloth_Qwen2.5-Coder-1.5B-Instruct.yaml
index 775c7ce08f..5242128004 100644
--- a/studio/backend/assets/configs/model_defaults/qwen/unsloth_Qwen2.5-Coder-1.5B-Instruct.yaml
+++ b/studio/backend/assets/configs/model_defaults/qwen/unsloth_Qwen2.5-Coder-1.5B-Instruct.yaml
@@ -34,6 +34,7 @@ lora:
- "down_proj"
use_rslora: false
use_loftq: false
+ use_dora: false
logging:
enable_wandb: false
diff --git a/studio/backend/assets/configs/model_defaults/qwen/unsloth_Qwen2.5-Coder-14B-Instruct.yaml b/studio/backend/assets/configs/model_defaults/qwen/unsloth_Qwen2.5-Coder-14B-Instruct.yaml
index 856db0c1b3..3559b636c6 100644
--- a/studio/backend/assets/configs/model_defaults/qwen/unsloth_Qwen2.5-Coder-14B-Instruct.yaml
+++ b/studio/backend/assets/configs/model_defaults/qwen/unsloth_Qwen2.5-Coder-14B-Instruct.yaml
@@ -35,6 +35,7 @@ lora:
- "down_proj"
use_rslora: false
use_loftq: false
+ use_dora: false
logging:
enable_wandb: false
diff --git a/studio/backend/assets/configs/model_defaults/qwen/unsloth_Qwen2.5-Coder-7B-Instruct-bnb-4bit.yaml b/studio/backend/assets/configs/model_defaults/qwen/unsloth_Qwen2.5-Coder-7B-Instruct-bnb-4bit.yaml
index 5900392547..3bc6d69afc 100644
--- a/studio/backend/assets/configs/model_defaults/qwen/unsloth_Qwen2.5-Coder-7B-Instruct-bnb-4bit.yaml
+++ b/studio/backend/assets/configs/model_defaults/qwen/unsloth_Qwen2.5-Coder-7B-Instruct-bnb-4bit.yaml
@@ -34,6 +34,7 @@ lora:
- "down_proj"
use_rslora: false
use_loftq: false
+ use_dora: false
logging:
enable_wandb: false
diff --git a/studio/backend/assets/configs/model_defaults/qwen/unsloth_Qwen2.5-VL-7B-Instruct-bnb-4bit.yaml b/studio/backend/assets/configs/model_defaults/qwen/unsloth_Qwen2.5-VL-7B-Instruct-bnb-4bit.yaml
index bd54b1d015..604b86dacd 100644
--- a/studio/backend/assets/configs/model_defaults/qwen/unsloth_Qwen2.5-VL-7B-Instruct-bnb-4bit.yaml
+++ b/studio/backend/assets/configs/model_defaults/qwen/unsloth_Qwen2.5-VL-7B-Instruct-bnb-4bit.yaml
@@ -29,6 +29,7 @@ lora:
- "all-linear"
use_rslora: false
use_loftq: false
+ use_dora: false
finetune_vision_layers: true
finetune_language_layers: true
finetune_attention_modules: true
diff --git a/studio/backend/assets/configs/model_defaults/qwen/unsloth_Qwen3-0.6B.yaml b/studio/backend/assets/configs/model_defaults/qwen/unsloth_Qwen3-0.6B.yaml
index 9feb6dcaae..daed4ebccb 100644
--- a/studio/backend/assets/configs/model_defaults/qwen/unsloth_Qwen3-0.6B.yaml
+++ b/studio/backend/assets/configs/model_defaults/qwen/unsloth_Qwen3-0.6B.yaml
@@ -35,6 +35,7 @@ lora:
- "down_proj"
use_rslora: false
use_loftq: false
+ use_dora: false
logging:
enable_wandb: false
diff --git a/studio/backend/assets/configs/model_defaults/qwen/unsloth_Qwen3-14B-Base-unsloth-bnb-4bit.yaml b/studio/backend/assets/configs/model_defaults/qwen/unsloth_Qwen3-14B-Base-unsloth-bnb-4bit.yaml
index a40eace253..05eef89b88 100644
--- a/studio/backend/assets/configs/model_defaults/qwen/unsloth_Qwen3-14B-Base-unsloth-bnb-4bit.yaml
+++ b/studio/backend/assets/configs/model_defaults/qwen/unsloth_Qwen3-14B-Base-unsloth-bnb-4bit.yaml
@@ -35,6 +35,7 @@ lora:
- "down_proj"
use_rslora: false
use_loftq: false
+ use_dora: false
logging:
enable_wandb: false
diff --git a/studio/backend/assets/configs/model_defaults/qwen/unsloth_Qwen3-14B.yaml b/studio/backend/assets/configs/model_defaults/qwen/unsloth_Qwen3-14B.yaml
index c130771c32..b4580e6d71 100644
--- a/studio/backend/assets/configs/model_defaults/qwen/unsloth_Qwen3-14B.yaml
+++ b/studio/backend/assets/configs/model_defaults/qwen/unsloth_Qwen3-14B.yaml
@@ -35,6 +35,7 @@ lora:
- "down_proj"
use_rslora: false
use_loftq: false
+ use_dora: false
logging:
enable_wandb: false
diff --git a/studio/backend/assets/configs/model_defaults/qwen/unsloth_Qwen3-30B-A3B-Instruct-2507.yaml b/studio/backend/assets/configs/model_defaults/qwen/unsloth_Qwen3-30B-A3B-Instruct-2507.yaml
index 2fb3a95c30..2eceb7d0de 100644
--- a/studio/backend/assets/configs/model_defaults/qwen/unsloth_Qwen3-30B-A3B-Instruct-2507.yaml
+++ b/studio/backend/assets/configs/model_defaults/qwen/unsloth_Qwen3-30B-A3B-Instruct-2507.yaml
@@ -36,6 +36,7 @@ lora:
- "gate_up_proj"
use_rslora: false
use_loftq: false
+ use_dora: false
logging:
enable_wandb: false
diff --git a/studio/backend/assets/configs/model_defaults/qwen/unsloth_Qwen3-32B.yaml b/studio/backend/assets/configs/model_defaults/qwen/unsloth_Qwen3-32B.yaml
index 152f4ae06a..032091880c 100644
--- a/studio/backend/assets/configs/model_defaults/qwen/unsloth_Qwen3-32B.yaml
+++ b/studio/backend/assets/configs/model_defaults/qwen/unsloth_Qwen3-32B.yaml
@@ -35,6 +35,7 @@ lora:
- "down_proj"
use_rslora: false
use_loftq: false
+ use_dora: false
logging:
enable_wandb: false
diff --git a/studio/backend/assets/configs/model_defaults/qwen/unsloth_Qwen3-4B-Instruct-2507.yaml b/studio/backend/assets/configs/model_defaults/qwen/unsloth_Qwen3-4B-Instruct-2507.yaml
index 94fe000708..e0e7f4ee3d 100644
--- a/studio/backend/assets/configs/model_defaults/qwen/unsloth_Qwen3-4B-Instruct-2507.yaml
+++ b/studio/backend/assets/configs/model_defaults/qwen/unsloth_Qwen3-4B-Instruct-2507.yaml
@@ -35,6 +35,7 @@ lora:
- "down_proj"
use_rslora: false
use_loftq: false
+ use_dora: false
logging:
enable_wandb: false
diff --git a/studio/backend/assets/configs/model_defaults/qwen/unsloth_Qwen3-4B-Thinking-2507.yaml b/studio/backend/assets/configs/model_defaults/qwen/unsloth_Qwen3-4B-Thinking-2507.yaml
index 3c325485d2..bb463849ed 100644
--- a/studio/backend/assets/configs/model_defaults/qwen/unsloth_Qwen3-4B-Thinking-2507.yaml
+++ b/studio/backend/assets/configs/model_defaults/qwen/unsloth_Qwen3-4B-Thinking-2507.yaml
@@ -35,6 +35,7 @@ lora:
- "down_proj"
use_rslora: false
use_loftq: false
+ use_dora: false
logging:
enable_wandb: false
diff --git a/studio/backend/assets/configs/model_defaults/qwen/unsloth_Qwen3-VL-8B-Instruct-unsloth-bnb-4bit.yaml b/studio/backend/assets/configs/model_defaults/qwen/unsloth_Qwen3-VL-8B-Instruct-unsloth-bnb-4bit.yaml
index 5b47c3bdd2..23e2b89dd0 100644
--- a/studio/backend/assets/configs/model_defaults/qwen/unsloth_Qwen3-VL-8B-Instruct-unsloth-bnb-4bit.yaml
+++ b/studio/backend/assets/configs/model_defaults/qwen/unsloth_Qwen3-VL-8B-Instruct-unsloth-bnb-4bit.yaml
@@ -29,6 +29,7 @@ lora:
- "all-linear"
use_rslora: false
use_loftq: false
+ use_dora: false
finetune_vision_layers: true
finetune_language_layers: true
finetune_attention_modules: true
diff --git a/studio/backend/assets/configs/vision_lora.yaml b/studio/backend/assets/configs/vision_lora.yaml
index 063a970316..a06f971523 100644
--- a/studio/backend/assets/configs/vision_lora.yaml
+++ b/studio/backend/assets/configs/vision_lora.yaml
@@ -30,6 +30,7 @@ lora:
vision_all_linear: true
use_rslora: false
use_loftq: false
+ use_dora: false
finetune_vision_layers: true
finetune_language_layers: true
finetune_attention_modules: true
diff --git a/studio/backend/auth/authentication.py b/studio/backend/auth/authentication.py
index dfb8fc513e..2e9520827e 100644
--- a/studio/backend/auth/authentication.py
+++ b/studio/backend/auth/authentication.py
@@ -11,11 +11,12 @@ import jwt
from .storage import (
API_KEY_PREFIX,
+ credential_generation,
get_jwt_secret,
get_user_and_secret,
load_jwt_secret,
save_refresh_token,
- validate_api_key,
+ validate_api_key_with_credential,
verify_refresh_token,
)
@@ -54,11 +55,14 @@ def create_access_token(
expires_delta: Optional[timedelta] = None,
*,
desktop: bool = False,
+ secret: Optional[str] = None,
) -> str:
"""
Create a signed JWT for the given subject (e.g. username).
- Valid across restarts: the signing secret is stored in SQLite.
+ Valid across restarts: the signing secret is stored in SQLite. Callers that
+ already verified a credential pass ``secret`` so a rotation landing mid-request
+ cannot sign the token with the credential that just replaced it.
"""
to_encode = {"sub": subject}
if desktop:
@@ -69,7 +73,7 @@ def create_access_token(
to_encode.update({"exp": expire})
return jwt.encode(
to_encode,
- _get_secret_for_subject(subject),
+ secret if secret is not None else _get_secret_for_subject(subject),
algorithm = ALGORITHM,
)
@@ -96,15 +100,28 @@ def is_desktop_access_token(token: str) -> bool:
return payload.get("sub") == subject and payload.get("desktop") is True
-def create_refresh_token(subject: str, *, desktop: bool = False) -> str:
+def create_refresh_token(
+ subject: str,
+ *,
+ desktop: bool = False,
+ secret: Optional[str] = None,
+) -> str:
"""
Create a random refresh token, store its hash in SQLite, and return it.
Refresh tokens are opaque (not JWTs); expire after REFRESH_TOKEN_EXPIRE_DAYS.
+ ``secret`` stamps the token with the credential version the caller verified,
+ so a rotation cannot leave a token minted from the replaced credential valid.
"""
token = secrets.token_urlsafe(48)
expires_at = datetime.now(timezone.utc) + timedelta(days = REFRESH_TOKEN_EXPIRE_DAYS)
- save_refresh_token(token, subject, expires_at.isoformat(), is_desktop = desktop)
+ save_refresh_token(
+ token,
+ subject,
+ expires_at.isoformat(),
+ is_desktop = desktop,
+ secret_gen = credential_generation(secret) if secret is not None else None,
+ )
return token
@@ -137,7 +154,22 @@ def reload_secret() -> None:
async def get_current_subject(credentials: HTTPAuthorizationCredentials = Depends(security)) -> str:
"""Validate JWT and require the password-change flow to be completed."""
- return await _get_current_subject(
+ subject, _generation = await _get_current_credential(
+ credentials,
+ allow_password_change = False,
+ )
+ return subject
+
+
+async def get_current_credential(
+ credentials: HTTPAuthorizationCredentials = Depends(security),
+) -> Tuple[str, Optional[str]]:
+ """As get_current_subject, but also returns the credential generation.
+
+ For routes that persist a new credential and must not do so on behalf of one
+ a concurrent reset has revoked.
+ """
+ return await _get_current_credential(
credentials,
allow_password_change = False,
)
@@ -158,27 +190,49 @@ async def get_current_subject_allow_password_change(
credentials: HTTPAuthorizationCredentials = Depends(security),
) -> str:
"""Validate JWT but allow access to the password-change endpoint."""
- return await _get_current_subject(
+ subject, _generation = await _get_current_credential(
credentials,
allow_password_change = True,
)
+ return subject
-async def _get_current_subject(
+# The literal the examples ship with; pasted unedited more often than a revoked key.
+API_KEY_PLACEHOLDER = f"{API_KEY_PREFIX}YOUR_KEY"
+
+
+def _invalid_api_key_detail(token: str) -> str:
+ """Why the key failed. Only the example placeholder is called out; every real
+ key gets one indistinguishable message, so this leaks no key existence."""
+ if token == API_KEY_PLACEHOLDER:
+ return (
+ "This is the placeholder key from the example. Create an API key in "
+ f"Unsloth Studio under Settings > API and use it in place of {API_KEY_PLACEHOLDER}."
+ )
+ return "Invalid or expired API key"
+
+
+async def _get_current_credential(
credentials: HTTPAuthorizationCredentials, *, allow_password_change: bool
-) -> str:
- """FastAPI dependency: validate the JWT and return the subject. Use on protected routes."""
+) -> Tuple[str, Optional[str]]:
+ """Validate the bearer and return ``(subject, credential generation)``.
+
+ The generation is the credential version this request actually authenticated
+ against. Routes that persist new credentials must bind their write to it, or
+ a reset landing mid-request would bless what it just revoked.
+ """
token = credentials.credentials
# --- API key path (sk-unsloth-...) ---
if token.startswith(API_KEY_PREFIX):
- username = validate_api_key(token)
- if username is None:
+ verified = validate_api_key_with_credential(token)
+ if verified is None:
raise HTTPException(
status_code = status.HTTP_401_UNAUTHORIZED,
- detail = "Invalid or expired API key",
+ detail = _invalid_api_key_detail(token),
)
- return username
+ username, secret = verified
+ return username, credential_generation(secret)
# --- JWT path ---
subject = _decode_subject_without_verification(token)
@@ -209,7 +263,7 @@ async def _get_current_subject(
status_code = status.HTTP_403_FORBIDDEN,
detail = "Password change required",
)
- return subject
+ return subject, credential_generation(jwt_secret)
except jwt.InvalidTokenError:
raise HTTPException(
status_code = status.HTTP_401_UNAUTHORIZED,
diff --git a/studio/backend/auth/storage.py b/studio/backend/auth/storage.py
index 39fa691304..6cf4d44834 100644
--- a/studio/backend/auth/storage.py
+++ b/studio/backend/auth/storage.py
@@ -9,6 +9,7 @@ import ipaddress
import os
import secrets
import sqlite3
+import tempfile
import threading
from datetime import datetime, timezone
from typing import Optional, Tuple
@@ -30,6 +31,97 @@ _BOOTSTRAP_PW_PATH = DB_PATH.parent / ".bootstrap_password"
_bootstrap_password: Optional[str] = None
+def _bootstrap_file_bytes(password: str) -> bytes:
+ """Exact on-disk form: the secret plus one LF.
+
+ Bytes, not text: text mode writes CRLF on Windows, and `$(cat ...)` strips
+ the LF but leaves the CR attached to the credential.
+ """
+ return (password + "\n").encode("utf-8")
+
+
+def _persist_bootstrap_password(password: str) -> None:
+ """Atomically write the bootstrap password 0600, LF terminated on every OS.
+
+ A partial write would destroy the only plaintext recovery credential.
+ """
+ fd, tmp_name = tempfile.mkstemp(
+ prefix = f".{_BOOTSTRAP_PW_PATH.name}.", dir = _BOOTSTRAP_PW_PATH.parent
+ )
+ try:
+ with os.fdopen(fd, "wb") as f:
+ f.write(_bootstrap_file_bytes(password))
+ try:
+ os.chmod(tmp_name, 0o600)
+ except OSError:
+ pass
+ os.replace(tmp_name, _BOOTSTRAP_PW_PATH)
+ except BaseException:
+ try:
+ os.unlink(tmp_name)
+ except OSError:
+ pass
+ raise
+
+
+def _normalise_bootstrap_file(raw: bytes, password: str) -> None:
+ """Append the LF a pre-newline release left off.
+
+ Append-only, and only when the file is exactly the credential:
+ clear_bootstrap_password() may unlink or (when unlink fails, notably on
+ Windows while this descriptor is open) truncate through another descriptor
+ after we read, so a rewrite could restore revoked plaintext. An append
+ cannot: worst case is a lone "\\n" over a cleared file, which strips back to
+ no bootstrap password. Pre-newline releases wrote no terminator at all, so
+ that is the only shape in the wild; anything else reads fine, since every
+ reader strips, and is left alone.
+ """
+ if raw != password.encode("utf-8"):
+ return
+
+ # O_BINARY: without it Windows opens in text mode and turns the LF straight
+ # back into CRLF, the bug being fixed.
+ fd = os.open(
+ _BOOTSTRAP_PW_PATH,
+ os.O_WRONLY | os.O_APPEND | getattr(os, "O_BINARY", 0),
+ )
+ try:
+ os.write(fd, b"\n")
+ try:
+ os.fchmod(fd, 0o600)
+ except (AttributeError, OSError):
+ # fchmod only reached Windows in 3.13.
+ pass
+ finally:
+ os.close(fd)
+
+
+def _read_persisted_bootstrap_password() -> Optional[str]:
+ """Read the persisted password, normalising the file if it is malformed."""
+ if not _BOOTSTRAP_PW_PATH.is_file():
+ return None
+
+ # No caller handles a raise, so an unreadable file has to mean "no bootstrap
+ # password", not a dead backend. We write UTF-8, so undecodable bytes are
+ # damage whose plaintext is worthless anyway.
+ try:
+ raw = _BOOTSTRAP_PW_PATH.read_bytes()
+ password = raw.decode("utf-8").strip()
+ except (OSError, UnicodeDecodeError):
+ return None
+ if not password:
+ return None
+
+ # Older releases wrote no terminator; best-effort, a read-only auth dir must
+ # not fail startup.
+ if raw != _bootstrap_file_bytes(password):
+ try:
+ _normalise_bootstrap_file(raw, password)
+ except OSError:
+ pass
+ return password
+
+
def generate_bootstrap_password() -> str:
"""Generate a 4-word diceware passphrase and persist it to disk.
@@ -43,10 +135,10 @@ def generate_bootstrap_password() -> str:
return _bootstrap_password
# Persisted from a previous run?
- if _BOOTSTRAP_PW_PATH.is_file():
- _bootstrap_password = _BOOTSTRAP_PW_PATH.read_text().strip()
- if _bootstrap_password:
- return _bootstrap_password
+ persisted = _read_persisted_bootstrap_password()
+ if persisted:
+ _bootstrap_password = persisted
+ return _bootstrap_password
# First startup: generate a fresh passphrase.
import diceware
@@ -57,11 +149,7 @@ def generate_bootstrap_password() -> str:
# Persist so the same passphrase survives restarts until password change.
ensure_dir(_BOOTSTRAP_PW_PATH.parent)
- _BOOTSTRAP_PW_PATH.write_text(_bootstrap_password)
- try:
- os.chmod(_BOOTSTRAP_PW_PATH, 0o600)
- except OSError:
- pass
+ _persist_bootstrap_password(_bootstrap_password)
return _bootstrap_password
@@ -72,13 +160,14 @@ def get_bootstrap_password() -> Optional[str]:
def _load_bootstrap_password() -> Optional[str]:
- """Load an existing bootstrap password without creating one."""
+ """Load an existing bootstrap password without creating one.
+
+ Upgrades take this path, not generate_bootstrap_password()
+ (ensure_default_admin short-circuits once the admin row exists), so it has
+ to normalise too.
+ """
global _bootstrap_password
- _bootstrap_password = None
- if _BOOTSTRAP_PW_PATH.is_file():
- bootstrap_password = _BOOTSTRAP_PW_PATH.read_text().strip()
- if bootstrap_password:
- _bootstrap_password = bootstrap_password
+ _bootstrap_password = _read_persisted_bootstrap_password()
return _bootstrap_password
@@ -97,9 +186,9 @@ def clear_bootstrap_password() -> None:
# Removal failed (Windows AV, read-only auth dir). The hash is already
# committed, so don't fail the change -- but truncate the file so its
# stale plaintext can't be re-seeded by generate_bootstrap_password()
- # if a later reset-password deletes auth.db and re-validates it.
+ # if auth.db is ever recreated.
try:
- _BOOTSTRAP_PW_PATH.write_text("")
+ _BOOTSTRAP_PW_PATH.write_text("", encoding = "utf-8")
cleared = True
except OSError:
cleared = False
@@ -132,6 +221,31 @@ def _hash_token(token: str) -> str:
return hashlib.sha256(token.encode("utf-8")).hexdigest()
+class CredentialRotated(Exception):
+ """A password reset revoked the credential this request authenticated with."""
+
+
+def credential_generation(jwt_secret: str) -> str:
+ """Marker for the credential version a refresh token was issued under.
+
+ Every password change rotates ``jwt_secret``, so a token stamped with the
+ previous one is rejected even if it was inserted after the revoking DELETE.
+ """
+ return hashlib.sha256(jwt_secret.encode("utf-8")).hexdigest()
+
+
+def _current_secret(conn: sqlite3.Connection, username: str) -> Optional[str]:
+ row = conn.execute(
+ "SELECT jwt_secret FROM auth_user WHERE username = ?", (username,)
+ ).fetchone()
+ return row["jwt_secret"] if row else None
+
+
+def _current_generation(conn: sqlite3.Connection, username: str) -> Optional[str]:
+ secret = _current_secret(conn, username)
+ return credential_generation(secret) if secret is not None else None
+
+
def get_connection() -> sqlite3.Connection:
"""Get a connection to the auth database, creating tables if needed."""
ensure_dir(DB_PATH.parent)
@@ -175,7 +289,8 @@ def get_connection() -> sqlite3.Connection:
token_hash TEXT NOT NULL,
username TEXT NOT NULL,
expires_at TEXT NOT NULL,
- is_desktop INTEGER NOT NULL DEFAULT 0
+ is_desktop INTEGER NOT NULL DEFAULT 0,
+ secret_gen TEXT
);
"""
)
@@ -214,6 +329,8 @@ def get_connection() -> sqlite3.Connection:
refresh_columns = {row["name"] for row in conn.execute("PRAGMA table_info(refresh_tokens)")}
if "is_desktop" not in refresh_columns:
conn.execute("ALTER TABLE refresh_tokens ADD COLUMN is_desktop INTEGER NOT NULL DEFAULT 0")
+ if "secret_gen" not in refresh_columns:
+ conn.execute("ALTER TABLE refresh_tokens ADD COLUMN secret_gen TEXT")
conn.commit()
return conn
@@ -587,12 +704,22 @@ def update_password(
new_password: str,
*,
revoke_refresh_tokens: bool = False,
-) -> bool:
+ expect_password_hash: Optional[str] = None,
+) -> Optional[str]:
"""Update password, clear first-login requirement, rotate JWT secret.
+ Returns the new JWT secret, or None when nothing was updated. Callers that
+ mint tokens for the caller must sign with the returned secret: re-reading it
+ would pick up a reset that landed between this commit and the mint.
+
``revoke_refresh_tokens`` deletes the user's refresh tokens in the SAME
transaction: a separate delete could fail after the password commit and
leave a pre-change token still able to mint access tokens.
+
+ ``expect_password_hash`` makes the write conditional on the credential the
+ caller verified still being current, so a request that checked the old
+ password cannot overwrite a reset that landed while it was in flight.
+ Returns False when the credential moved underneath it.
"""
from .hashing import hash_password
@@ -600,21 +727,32 @@ def update_password(
jwt_secret = secrets.token_urlsafe(64)
conn = get_connection()
try:
- cursor = conn.execute(
- """
- UPDATE auth_user
- SET password_salt = ?, password_hash = ?, jwt_secret = ?, must_change_password = 0
- WHERE username = ?
- """,
- (salt, pwd_hash, jwt_secret, username),
- )
+ if expect_password_hash is None:
+ cursor = conn.execute(
+ """
+ UPDATE auth_user
+ SET password_salt = ?, password_hash = ?, jwt_secret = ?, must_change_password = 0
+ WHERE username = ?
+ """,
+ (salt, pwd_hash, jwt_secret, username),
+ )
+ else:
+ cursor = conn.execute(
+ """
+ UPDATE auth_user
+ SET password_salt = ?, password_hash = ?, jwt_secret = ?, must_change_password = 0
+ WHERE username = ? AND password_hash = ?
+ """,
+ (salt, pwd_hash, jwt_secret, username, expect_password_hash),
+ )
if revoke_refresh_tokens and cursor.rowcount > 0:
conn.execute("DELETE FROM refresh_tokens WHERE username = ?", (username,))
conn.commit()
if cursor.rowcount > 0:
clear_bootstrap_password()
clear_desktop_secret()
- return cursor.rowcount > 0
+ return jwt_secret
+ return None
finally:
conn.close()
@@ -625,35 +763,49 @@ def save_refresh_token(
expires_at: str,
*,
is_desktop: bool = False,
+ secret_gen: Optional[str] = None,
) -> None:
"""
Store a hashed refresh token with its associated username and expiry.
+
+ ``secret_gen`` binds the token to a credential version; it defaults to the
+ current one, and callers that already verified a credential must pass the
+ version they verified rather than let this re-read a rotated one.
"""
token_hash = _hash_token(token)
conn = get_connection()
try:
+ if secret_gen is None:
+ secret_gen = _current_generation(conn, username)
conn.execute(
"""
- INSERT INTO refresh_tokens (token_hash, username, expires_at, is_desktop)
- VALUES (?, ?, ?, ?)
+ INSERT INTO refresh_tokens (token_hash, username, expires_at, is_desktop, secret_gen)
+ VALUES (?, ?, ?, ?, ?)
""",
- (token_hash, username, expires_at, int(is_desktop)),
+ (token_hash, username, expires_at, int(is_desktop), secret_gen),
)
conn.commit()
finally:
conn.close()
-def consume_refresh_token(token: str) -> Optional[Tuple[str, bool]]:
+def consume_refresh_token(token: str) -> Optional[Tuple[str, bool, str]]:
"""Atomically validate-and-delete a refresh token for single-use rotation.
DELETE RETURNING fuses validate and delete into one statement so two
- concurrent refresh requests cannot both consume the same token.
+ concurrent refresh requests cannot both consume the same token. Returns
+ ``(username, is_desktop, jwt_secret)``; the caller must mint the replacement
+ tokens against that secret so a rotation landing mid-refresh cannot issue a
+ post-rotation session from a pre-rotation token.
"""
token_hash = _hash_token(token)
now = datetime.now(timezone.utc).isoformat()
conn = get_connection()
try:
+ # One transaction with the delete: an unstamped legacy row has no
+ # generation to compare, so reading the credential after committing would
+ # hand a reset's new secret to a token issued before it.
+ conn.execute("BEGIN IMMEDIATE")
conn.execute(
"DELETE FROM refresh_tokens WHERE expires_at < ?",
(now,),
@@ -662,15 +814,21 @@ def consume_refresh_token(token: str) -> Optional[Tuple[str, bool]]:
"""
DELETE FROM refresh_tokens
WHERE token_hash = ? AND expires_at >= ?
- RETURNING username, is_desktop
+ RETURNING username, is_desktop, secret_gen
""",
(token_hash, now),
)
row = cur.fetchone()
- conn.commit()
if row is None:
+ conn.commit()
return None
- return row["username"], bool(row["is_desktop"])
+ secret = _current_secret(conn, row["username"])
+ conn.commit()
+ if secret is None:
+ return None
+ if row["secret_gen"] is not None and row["secret_gen"] != credential_generation(secret):
+ return None
+ return row["username"], bool(row["is_desktop"]), secret
finally:
conn.close()
@@ -694,7 +852,7 @@ def verify_refresh_token(token: str) -> Optional[Tuple[str, bool]]:
cur = conn.execute(
"""
- SELECT id, username, expires_at, is_desktop FROM refresh_tokens
+ SELECT id, username, expires_at, is_desktop, secret_gen FROM refresh_tokens
WHERE token_hash = ?
""",
(token_hash,),
@@ -703,6 +861,13 @@ def verify_refresh_token(token: str) -> Optional[Tuple[str, bool]]:
if row is None:
return None
+ if row["secret_gen"] is not None and row["secret_gen"] != _current_generation(
+ conn, row["username"]
+ ):
+ conn.execute("DELETE FROM refresh_tokens WHERE id = ?", (row["id"],))
+ conn.commit()
+ return None
+
# Check expiry
expires_at = datetime.fromisoformat(row["expires_at"])
if datetime.now(timezone.utc) > expires_at:
@@ -747,30 +912,41 @@ def create_desktop_secret() -> str:
conn.close()
-def validate_desktop_secret(raw_secret: str) -> Optional[str]:
- """Return the real admin username when the desktop secret matches."""
+def validate_desktop_secret_with_credential(raw_secret: str) -> Optional[Tuple[str, str]]:
+ """Validate the desktop secret and return ``(username, jwt_secret)``.
+
+ Both reads share one transaction so the returned secret is the credential
+ version the desktop secret was checked against; a reset landing mid-request
+ then invalidates the tokens minted from it rather than blessing them.
+ """
if not raw_secret.startswith(DESKTOP_SECRET_PREFIX):
return None
- if get_user_and_secret(DEFAULT_ADMIN_USERNAME) is None:
- return None
secret_hash = _pbkdf2_desktop_secret(raw_secret)
conn = get_connection()
try:
- cur = conn.execute(
+ conn.execute("BEGIN")
+ row = conn.execute(
"SELECT value FROM app_secrets WHERE key = ?",
(_DESKTOP_SECRET_HASH_KEY,),
- )
- row = cur.fetchone()
- if row is None:
+ ).fetchone()
+ if row is None or not secrets.compare_digest(row["value"], secret_hash):
return None
- if not secrets.compare_digest(row["value"], secret_hash):
+ jwt_secret = _current_secret(conn, DEFAULT_ADMIN_USERNAME)
+ if jwt_secret is None:
return None
- return DEFAULT_ADMIN_USERNAME
+ return DEFAULT_ADMIN_USERNAME, jwt_secret
finally:
+ conn.rollback()
conn.close()
+def validate_desktop_secret(raw_secret: str) -> Optional[str]:
+ """Return the real admin username when the desktop secret matches."""
+ verified = validate_desktop_secret_with_credential(raw_secret)
+ return verified[0] if verified else None
+
+
def clear_desktop_secret() -> None:
"""Remove backend-side desktop auth state."""
conn = get_connection()
@@ -796,6 +972,7 @@ def create_api_key(
name: str,
expires_at: Optional[str] = None,
internal: bool = False,
+ expect_gen: Optional[str] = None,
) -> Tuple[str, dict]:
"""Create a new API key for *username*.
@@ -804,6 +981,10 @@ def create_api_key(
Pass ``internal=True`` for keys minted by workflows (e.g. data-recipe
runs) that should not appear in user-facing key listings.
+
+ ``expect_gen`` ties the insert to the credential generation the request
+ authenticated under, so a session revoked by a concurrent password reset
+ cannot mint a key that outlives it. Raises ``CredentialRotated`` if it moved.
"""
raw_key = API_KEY_PREFIX + secrets.token_hex(16)
key_hash = _pbkdf2_api_key(raw_key)
@@ -812,6 +993,12 @@ def create_api_key(
conn = get_connection()
try:
+ if expect_gen is not None:
+ conn.execute("BEGIN IMMEDIATE")
+ if _current_generation(conn, username) != expect_gen:
+ raise CredentialRotated(
+ "The credential this request authenticated with was revoked."
+ )
conn.execute(
"""
INSERT INTO api_keys (username, key_prefix, key_hash, name, created_at, expires_at, is_internal)
@@ -900,15 +1087,25 @@ def revoke_internal_api_key(key_id: int) -> bool:
def validate_api_key(raw_key: str) -> Optional[str]:
- """Validate *raw_key* and return the owning username, or ``None``.
+ """Validate *raw_key* and return the owning username, or ``None``."""
+ verified = validate_api_key_with_credential(raw_key)
+ return verified[0] if verified else None
- Also updates ``last_used_at`` on success.
+
+def validate_api_key_with_credential(raw_key: str) -> Optional[Tuple[str, str]]:
+ """Validate *raw_key* and return ``(username, jwt_secret)``, or ``None``.
+
+ Also updates ``last_used_at`` on success. The key check and the credential
+ read share one write transaction, so the returned version is the one the key
+ was actually valid under: a reset committing right after cannot have its new
+ generation handed to a request the key it revoked authenticated.
"""
cache_id = _api_key_cache_id(raw_key)
cached_hash = _api_key_hash_cache.get(cache_id)
key_hash = cached_hash if cached_hash is not None else _pbkdf2_api_key(raw_key)
conn = get_connection()
try:
+ conn.execute("BEGIN IMMEDIATE")
cur = conn.execute(
"SELECT id, username, is_active, expires_at FROM api_keys WHERE key_hash = ?",
(key_hash,),
@@ -928,11 +1125,15 @@ def validate_api_key(raw_key: str) -> Optional[str]:
expires = datetime.fromisoformat(row["expires_at"])
if datetime.now(timezone.utc) > expires:
return None
+ secret = _current_secret(conn, row["username"])
+ if secret is None:
+ return None
conn.execute(
"UPDATE api_keys SET last_used_at = ? WHERE id = ?",
(datetime.now(timezone.utc).isoformat(), row["id"]),
)
conn.commit()
- return row["username"]
+ return row["username"], secret
finally:
+ conn.rollback()
conn.close()
diff --git a/studio/backend/auth/terminal_prompt.py b/studio/backend/auth/terminal_prompt.py
index e855f4078b..925404f47d 100644
--- a/studio/backend/auth/terminal_prompt.py
+++ b/studio/backend/auth/terminal_prompt.py
@@ -236,6 +236,10 @@ def prompt_for_password_change(
out.write(f"Password must be at least {min_length} characters; try again.\n")
out.flush()
continue
+ if any(ch.isspace() for ch in new_password):
+ out.write("Password cannot contain spaces; try again.\n")
+ out.flush()
+ continue
if is_current_password(new_password):
out.write(
"New password must differ from the current bootstrap password; try again.\n"
diff --git a/studio/backend/cloudflare_tunnel.py b/studio/backend/cloudflare_tunnel.py
index b1ddc74c32..f7967e2faa 100644
--- a/studio/backend/cloudflare_tunnel.py
+++ b/studio/backend/cloudflare_tunnel.py
@@ -20,6 +20,7 @@ import shutil
import subprocess
import sys
import threading
+import time
from pathlib import Path
from typing import Optional, Tuple
@@ -40,6 +41,22 @@ _RELEASE_BASE = "https://github.com/cloudflare/cloudflared/releases/latest/downl
_READY_TIMEOUT = 15.0 # seconds to wait for the URL + a registered edge connection
_DOWNLOAD_TIMEOUT = 60 # urlopen timeout for the one-time binary download
+# A registered edge connection does not mean the hostname resolves yet, so the
+# URL is fetched once before it is advertised.
+_PUBLIC_PROBE_PATH = "/api/health"
+_PUBLIC_PROBE_MARKER = "Unsloth UI Backend"
+# One deadline for DNS propagation + the health probe, bounding the startup stall.
+_PUBLIC_PROBE_TIMEOUT = 45.0
+_PUBLIC_PROBE_ATTEMPT_TIMEOUT = 5.0
+_PUBLIC_PROBE_RETRY_DELAY = 1.0
+
+# Wait for the hostname via DoH first: an early OS lookup negative-caches the
+# NXDOMAIN for up to 30 min.
+_DNS_POLL_DELAY = 2.0
+# Retry transient DoH failures, but give up fast when DoH is blocked outright.
+_DNS_MAX_DOH_ERRORS = 3
+_DOH_URL = "https://cloudflare-dns.com/dns-query?name={host}&type=A"
+
def _windows_hidden_kwargs() -> dict:
"""Suppress a child console window on Windows; no-op elsewhere."""
@@ -191,6 +208,59 @@ def ensure_cloudflared() -> Optional[str]:
return None
+def _wait_for_dns(host: str, deadline: float) -> None:
+ import json
+ import urllib.request
+
+ errors = 0
+ while True:
+ answered = False
+ try:
+ req = urllib.request.Request(
+ _DOH_URL.format(host = host),
+ headers = {"Accept": "application/dns-json", "User-Agent": "unsloth-studio"},
+ )
+ with urllib.request.urlopen(req, timeout = 5) as response:
+ answered = bool(json.loads(response.read(65536)).get("Answer"))
+ errors = 0
+ except Exception:
+ errors += 1
+ if errors >= _DNS_MAX_DOH_ERRORS:
+ return
+ if answered:
+ return
+ remaining = deadline - time.monotonic()
+ if remaining <= 0:
+ return
+ time.sleep(min(_DNS_POLL_DELAY, remaining))
+
+
+def verify_public_url(url: str, timeout: float = _PUBLIC_PROBE_TIMEOUT) -> bool:
+ import json
+ import urllib.request
+ from urllib.parse import urlsplit
+
+ deadline = time.monotonic() + timeout
+ host = urlsplit(url).hostname
+ if host:
+ _wait_for_dns(host, deadline)
+
+ probe_url = f"{url.rstrip('/')}{_PUBLIC_PROBE_PATH}"
+ while True:
+ try:
+ req = urllib.request.Request(probe_url, headers = {"User-Agent": "unsloth-studio"})
+ with urllib.request.urlopen(req, timeout = _PUBLIC_PROBE_ATTEMPT_TIMEOUT) as response:
+ body = response.read(4096)
+ if json.loads(body).get("service") == _PUBLIC_PROBE_MARKER:
+ return True
+ except Exception:
+ pass
+ remaining = deadline - time.monotonic()
+ if remaining <= 0:
+ return False
+ time.sleep(min(_PUBLIC_PROBE_RETRY_DELAY, remaining))
+
+
class CloudflareTunnel:
"""A cloudflared quick tunnel to http://localhost:. Best-effort throughout.
@@ -240,6 +310,7 @@ class CloudflareTunnel:
stderr = subprocess.STDOUT,
stdin = subprocess.DEVNULL,
text = True,
+ encoding = "utf-8",
errors = "replace",
bufsize = 1,
**_windows_hidden_kwargs(),
@@ -322,11 +393,12 @@ def start_studio_tunnel(port: int, timeout: float = _READY_TIMEOUT) -> Optional[
"""Start a quick tunnel and return its public URL once it is actually
serving, or None (best-effort).
- Waits for cloudflared to both mint the URL and register an edge connection
- before returning, so the caller never advertises a URL that yields Cloudflare
- error 1033 (HTTP 530). If a URL is minted but no connection registers within
- the window (e.g. quic is blocked on this network), retries once forcing the
- http2 protocol. On any failure the tunnel is stopped and None is returned.
+ Waits for cloudflared to both mint the URL and register an edge connection,
+ then fetches /api/health over the public URL, so the caller never advertises
+ a link that yields Cloudflare error 1033 (HTTP 530) or an unresolvable host.
+ If a URL is minted but no connection registers within the window (e.g. quic
+ is blocked on this network), retries once forcing the http2 protocol. On any
+ failure the tunnel is stopped and None is returned.
"""
global _active_tunnel, _shutdown_requested
binary = ensure_cloudflared()
@@ -349,9 +421,13 @@ def start_studio_tunnel(port: int, timeout: float = _READY_TIMEOUT) -> Optional[
prior, _active_tunnel = _active_tunnel, tunnel
if prior is not None:
prior.stop()
+ registered = False
try:
tunnel.start()
url = tunnel.wait_for_ready(timeout)
+ registered = url is not None
+ if url and not verify_public_url(url):
+ url = None
except Exception:
url = None
if url:
@@ -371,6 +447,9 @@ def start_studio_tunnel(port: int, timeout: float = _READY_TIMEOUT) -> Optional[
# http2 will not help, so do not burn another window on it.
if not saw_url:
return None
+ # probe failure after registering is DNS propagation; http2 would not help
+ if registered:
+ return None
return None
diff --git a/studio/backend/colab.py b/studio/backend/colab.py
index 1762469bcf..bf4a6a44b5 100644
--- a/studio/backend/colab.py
+++ b/studio/backend/colab.py
@@ -1,9 +1,7 @@
# SPDX-License-Identifier: AGPL-3.0-only
# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0
-"""
-Colab helpers for Unsloth Studio. Uses Colab's built-in proxy.
-"""
+"""Colab helpers for Unsloth Studio. Uses Colab's built-in proxy."""
from pathlib import Path
import sys
@@ -22,11 +20,9 @@ logger = get_logger(__name__)
def get_colab_url(port: int = 8888) -> str:
- """
- Get the Colab proxy URL for a port.
+ """Get the Colab proxy URL for a port.
- Retries up to 3 times, validating the result is a real HTTPS Colab URL.
- Falls back to http://localhost:{port} only when all attempts fail.
+ Retries 3x validating a real HTTPS Colab URL; falls back to localhost on failure.
"""
import time as _time
@@ -55,28 +51,243 @@ def get_colab_url(port: int = 8888) -> str:
return fallback
-def show_link(port: int = 8888, *, _url: "str | None" = None):
- """Display a styled clickable link to the UI.
-
- *_url* is an optional pre-fetched proxy URL; pass it to avoid a second eval_js round-trip.
- """
- from IPython.display import display, HTML
-
- url = _url if _url is not None else get_colab_url(port)
-
- # Truncated display URL; try/except so an odd URL shape still renders the link.
+def _short_colab_url(url: str, port: int) -> str:
+ """Truncated display form of a Colab proxy URL; falls back to the full URL."""
try:
port_prefix = f"{port}-"
idx = url.index(port_prefix)
next_dash = url.index("-", idx + len(port_prefix))
- short_url = url[: next_dash + 1] + "..."
+ return url[: next_dash + 1] + "..."
except (ValueError, IndexError):
- short_url = url
+ return url
- # Plain-text line so the URL shows even if HTML display fails.
- logger.info(f"🌐 Unsloth Studio URL: {url}")
- html = f"""
+def _is_colab_proxy_url(url: str, port: int) -> bool:
+ """True when *url* looks like a real Colab kernel proxy, not a localhost fallback."""
+ return bool(url and isinstance(url, str) and url.startswith("https://") and str(port) in url)
+
+
+def _is_colab_runtime() -> bool:
+ """True on a hosted Colab notebook kernel.
+
+ Reuses the backend's main Colab detector (``/content`` + Colab env / ``google.colab``)
+ instead of a single env var, which is not always present on hosted runtimes.
+ """
+ try:
+ from main import _IS_COLAB
+ return bool(_IS_COLAB)
+ except Exception:
+ return False
+
+
+def _colab_login_credentials_path() -> Path:
+ from auth.storage import DB_PATH
+ return DB_PATH.parent / ".colab_notebook_login"
+
+
+def _store_colab_login_credentials(username: str, password: str) -> None:
+ """Persist Colab admin credentials for notebook re-runs after interrupt."""
+ path = _colab_login_credentials_path()
+ try:
+ path.parent.mkdir(parents = True, exist_ok = True)
+ path.write_text(f"{username}\n{password}\n", encoding = "utf-8")
+ try:
+ import os
+ os.chmod(path, 0o600)
+ except OSError:
+ pass
+ except OSError as e:
+ logger.info(f"Could not persist Colab login credentials ({e}).")
+
+
+def _load_colab_login_credentials() -> "tuple[str, str] | None":
+ """Return stored Colab admin credentials from a previous ``start()`` run, if any."""
+ path = _colab_login_credentials_path()
+ try:
+ if not path.is_file():
+ return None
+ lines = path.read_text(encoding = "utf-8").splitlines()
+ if len(lines) >= 2 and lines[0] and lines[1]:
+ return lines[0], lines[1]
+ except (OSError, UnicodeDecodeError) as e:
+ logger.info(f"Could not load Colab login credentials ({e}).")
+ return None
+
+
+def _clear_colab_login_credentials() -> None:
+ """Drop the cached Colab credentials once they no longer authenticate."""
+ path = _colab_login_credentials_path()
+ try:
+ path.unlink(missing_ok = True)
+ except OSError as e:
+ logger.info(f"Could not clear Colab login credentials ({e}).")
+
+
+def _colab_credentials_still_valid(username: str, password: str) -> bool:
+ """True when *password* still matches the stored admin hash.
+
+ Guards against redisplaying a cached first-run password after the user has
+ changed the admin password through the app, which would print credentials
+ that no longer authenticate to the current Cloudflare tunnel.
+ """
+ try:
+ from auth.storage import get_user_and_secret
+ from auth.hashing import verify_password
+ except Exception as e:
+ logger.info(f"Could not load auth to validate cached Colab credentials ({e}).")
+ return False
+ try:
+ row = get_user_and_secret(username)
+ if not row:
+ return False
+ salt, pwd_hash = row[0], row[1]
+ return bool(verify_password(password, salt, pwd_hash))
+ except Exception as e:
+ logger.info(f"Could not validate cached Colab credentials ({e}).")
+ return False
+
+
+def _colab_wants_cloudflare(cloudflare: "bool | None") -> bool:
+ """Resolve whether to open a Cloudflare tunnel.
+
+ ``None`` auto-enables on real Colab (the in-cell proxy embed is often blank);
+ pass ``False`` to opt out.
+ """
+ if cloudflare is not None:
+ return cloudflare
+ return _is_colab_runtime()
+
+
+def _finalize_colab_admin_password() -> "tuple[str, str] | None":
+ """Clear the bootstrap-password gate on Colab so Cloudflare tunnels can start.
+
+ Returns ``(username, password)`` for display in the notebook. On first run the
+ random admin password is finalized; on later runs (e.g. after interrupt) the
+ stored credentials are re-displayed so the Cloudflare link stays usable.
+ Anyone who can read this cell already controls the runtime.
+ """
+ if not _is_colab_runtime():
+ return None
+ try:
+ from auth.storage import (
+ DEFAULT_ADMIN_USERNAME,
+ ensure_default_admin,
+ generate_bootstrap_password,
+ get_bootstrap_password,
+ requires_password_change,
+ update_password,
+ )
+ except Exception as e:
+ logger.warning(
+ f"Could not load auth for Colab setup ({e}); Cloudflare link may be blocked."
+ )
+ return None
+
+ try:
+ ensure_default_admin()
+ username = DEFAULT_ADMIN_USERNAME
+ if not requires_password_change(username):
+ creds = _load_colab_login_credentials()
+ if creds is not None and _colab_credentials_still_valid(username, creds[1]):
+ return creds
+ # The admin password was changed through the app after the first run,
+ # so the cached copy is stale; drop it instead of printing dead credentials.
+ _clear_colab_login_credentials()
+ return None
+ password = get_bootstrap_password() or generate_bootstrap_password()
+ if not update_password(username, password):
+ logger.warning(
+ "Could not finalize Colab admin password; Cloudflare link may be blocked."
+ )
+ return None
+ _store_colab_login_credentials(username, password)
+ return username, password
+ except Exception as e:
+ logger.warning(
+ f"Could not finalize Colab admin password ({e}); Cloudflare link may be blocked."
+ )
+ return None
+
+
+def _colab_login_html(username: str, password: str) -> str:
+ """Notebook card with Colab admin credentials (shown once after auto-finalize)."""
+ return f"""
+
+
+ Unsloth Studio Login (Colab)
+
+
+ Log in as {username} with this password. This cell is visible only in
+ your notebook session.
+
+
+ Password: {password}
+
+
+ """
+
+
+def _show_colab_login_credentials(username: str, password: str) -> None:
+ """Display Colab admin credentials in the notebook output."""
+ from IPython.display import HTML, display
+
+ logger.info(f"🔐 Unsloth Studio login — user: {username}")
+ display(HTML(_colab_login_html(username, password)))
+
+
+def _ready_card_html(
+ url: str,
+ port: int,
+ *,
+ has_cloudflare_link: bool = False,
+ cloudflare_requested: bool = False,
+) -> str:
+ """Branded ready card for the in-notebook Studio view.
+
+ Colab ``*.prod.colab.dev`` proxy URLs are session-scoped and 404 when opened as a
+ top-level tab or on another device, so never ``window.open`` them. On real Colab the
+ Cloudflare link is the supported entry point because in-cell proxy embeds often stay blank.
+ """
+ short_url = _short_colab_url(url, port)
+ if _is_colab_runtime() or _is_colab_proxy_url(url, port):
+ if has_cloudflare_link:
+ embed_note = (
+ "Open Studio with the Cloudflare link above. In-cell proxy previews on "
+ "current Colab often stay blank, so the tunnel link is the supported path."
+ )
+ elif cloudflare_requested:
+ embed_note = (
+ "Could not open a Cloudflare tunnel, so Studio may be unreachable on Colab. "
+ "Check the logs above and re-run this cell. Pass "
+ ''
+ "cloudflare=True after fixing any tunnel errors."
+ )
+ else:
+ embed_note = (
+ "Colab proxy links cannot be opened in a new tab (they 404 outside this "
+ 'notebook). Re-run with start(cloudflare=True) for a working link.'
+ )
+ return f"""
+
+
+
+ Unsloth Studio is Ready!
+
+
+ {embed_note}
+
+
+ {short_url}
+
+
+ """
+
+ return f"""
"""
- display(HTML(html))
+
+
+def show_link(
+ port: int = 8888,
+ *,
+ _url: "str | None" = None,
+ has_cloudflare_link: bool = False,
+ cloudflare_requested: bool = False,
+):
+ """Display a styled ready card for the UI.
+
+ Colab proxy URLs are informational only (no new-tab open; they 404 outside the cell);
+ non-proxy URLs keep a clickable open button. *_url* is an optional pre-fetched proxy
+ URL to avoid a second eval_js round-trip.
+ """
+ from IPython.display import display, HTML
+
+ url = _url if _url is not None else get_colab_url(port)
+ logger.info(f"🌐 Unsloth Studio URL: {url}")
+ display(
+ HTML(
+ _ready_card_html(
+ url,
+ port,
+ has_cloudflare_link = has_cloudflare_link,
+ cloudflare_requested = cloudflare_requested,
+ )
+ )
+ )
+
+
+def _warn_colab_cloudflare_missing(*, use_cloudflare: bool, cloudflare_url: "str | None") -> None:
+ """Log a prominent warning when Colab expected a tunnel but none was opened."""
+ if not use_cloudflare or cloudflare_url or not _is_colab_runtime():
+ return
+ logger.warning(
+ "Colab Cloudflare tunnel unavailable — Studio is unlikely to be reachable in this "
+ "notebook. Check the logs above for tunnel or auth errors, then re-run start()."
+ )
def _bootstrap_password_pending() -> bool:
"""True while the default admin still owes a bootstrap-password change.
- While pending, main.py injects that password into same-origin GETs, and a public
- tunnel GET (no Origin) reads as same-origin, so sharing the link would leak admin
- access. Fails safe to pending if the state cannot be read.
+ While pending, a public tunnel GET (no Origin) reads as same-origin and gets the
+ injected password, so sharing the link would leak admin access. Fails safe to pending.
"""
try:
from auth.storage import requires_password_change, DEFAULT_ADMIN_USERNAME
@@ -121,9 +369,8 @@ def _bootstrap_password_pending() -> bool:
def start_cloudflare_tunnel(port: int) -> "str | None":
"""Open a shareable Cloudflare quick tunnel to localhost:*port*, or None.
- run_server suppresses the tunnel on Colab by design, so we start it directly.
- Refused while the bootstrap password is pending; any failure collapses to None
- and the Colab proxy still works.
+ run_server suppresses the tunnel on Colab, so we start it directly. Refused while the
+ bootstrap password is pending; any failure collapses to None (Colab proxy still works).
"""
if _bootstrap_password_pending():
logger.warning(
@@ -152,9 +399,9 @@ def start_cloudflare_tunnel(port: int) -> "str | None":
def _publish_cloudflare_url(cloudflare_url: "str | None") -> None:
"""Publish a directly-started tunnel URL onto app.state so /api/health advertises it.
- run_server only sets this when it opens the tunnel itself, which it skips on Colab,
- so we set it here. Otherwise the frontend's API examples fall back to an
- unreachable server_url. Best-effort.
+ run_server sets this only when it opens the tunnel itself (skipped on Colab), so we
+ set it here; otherwise the frontend's API examples fall back to an unreachable
+ server_url. Best-effort.
"""
if not cloudflare_url:
return
@@ -183,8 +430,7 @@ def _stop_cloudflare_tunnel() -> None:
def _is_studio_healthy(port: int, timeout: float = 2.0) -> bool:
"""True only if Unsloth Studio (not some other app) answers /api/health on *port*.
- The service-marker check stops the reuse path reusing or tunneling a foreign
- process that merely serves /api/health.
+ The service-marker check stops the reuse path reusing or tunneling a foreign process.
"""
import json, urllib.request
try:
@@ -194,8 +440,29 @@ def _is_studio_healthy(port: int, timeout: float = 2.0) -> bool:
return False
-def _shareable_link_html(cloudflare_url: str) -> str:
- """Branded card for the shareable Cloudflare link, styled like the show_link banner."""
+def _shareable_link_html(
+ cloudflare_url: str,
+ password: "str | None" = None,
+ username: "str | None" = None,
+) -> str:
+ """Branded card for the shareable Cloudflare link, styled like the show_link banner.
+
+ *password* renders under the link so the credential sits in the card with the button
+ it unlocks. The username is always the default admin, so it reads inline.
+ """
+ login_block = ""
+ if password:
+ login_block = f"""
+
+ Password
+
+
{password}
+
+ Log in as {username} with this password. Shown only in your
+ notebook session, and never included in the shared link.
+
"""
return f"""
@@ -213,40 +480,55 @@ def _shareable_link_html(cloudflare_url: str) -> str:
Open Unsloth Studio
- This Cloudflare HTTPS link works from any device — share it with anyone. The Colab view below only works in this tab.
+ This Cloudflare HTTPS link works from any device, so you can share it with anyone.
"""
-def _show_and_embed(port: int, *, cloudflare_url: "str | None" = None):
- """Render the Unsloth header + iframe for *port*, with a shareable-link card above
- when *cloudflare_url* is set. Falls back to serve_kernel_port_as_iframe."""
- url = get_colab_url(port)
- logger.info(f"🌐 Unsloth Studio URL: {url}")
- if cloudflare_url:
- logger.info(f"🔗 Shareable Cloudflare link: {cloudflare_url}")
+# Height for serve_kernel_port_as_iframe (~82vh on a 1080p screen, clamped).
+_COLAB_IFRAME_HEIGHT = 900
+
+def _embed_kernel_port_iframe(port: int) -> bool:
+ """Embed Studio via Colab's native kernel-port iframe helper.
+
+ Only trusted on a real Colab runtime: colabtools can import ``google.colab`` and
+ queue browser-side JS without appending an iframe, so callers outside Colab must use
+ the HTML iframe path instead.
+ """
+ if not _is_colab_runtime():
+ return False
+ try:
+ from google.colab import output as colab_output
+ except ImportError:
+ return False
+ try:
+ colab_output.serve_kernel_port_as_iframe(
+ port,
+ height = _COLAB_IFRAME_HEIGHT,
+ width = "100%",
+ )
+ return True
+ except Exception as e:
+ logger.info(f"serve_kernel_port_as_iframe failed ({e}); trying HTML iframe.")
+ return False
+
+
+def _embed_html_iframe(url: str, port: int) -> bool:
+ """Fallback embed: raw HTML iframe when the Colab helper is unavailable."""
try:
from IPython.display import HTML, display
+ except ImportError:
+ return False
- iframe_id = f"unsloth-studio-{port}"
-
- # Truncated header URL — best-effort, falls back to full URL.
- try:
- port_prefix = f"{port}-"
- idx = url.index(port_prefix)
- next_dash = url.index("-", idx + len(port_prefix))
- short_url = url[: next_dash + 1] + "..."
- except (ValueError, IndexError):
- short_url = url
-
- if cloudflare_url:
- display(HTML(_shareable_link_html(cloudflare_url)))
-
+ short_url = _short_colab_url(url, port)
+ iframe_id = f"unsloth-studio-{port}"
+ try:
display(
HTML(f"""
""")
)
- except Exception:
- # Fallback: Colab's built-in helper.
+ return True
+ except Exception as e:
+ logger.info(f"HTML iframe embed failed ({e}).")
+ return False
+
+
+def _show_and_embed(
+ port: int,
+ *,
+ cloudflare_url: "str | None" = None,
+ colab_login: "tuple[str, str] | None" = None,
+ cloudflare_requested: bool = False,
+):
+ """Render the Unsloth ready card + iframe for *port*.
+
+ Prefer Colab's ``serve_kernel_port_as_iframe`` on real Colab; raw HTML iframe is the
+ fallback. Cloudflare cards stay clickable.
+ """
+ url = get_colab_url(port)
+ logger.info(f"🌐 Unsloth Studio URL: {url}")
+ if cloudflare_url:
+ logger.info(f"🔗 Shareable Cloudflare link: {cloudflare_url}")
+
+ _warn_colab_cloudflare_missing(
+ use_cloudflare = cloudflare_requested,
+ cloudflare_url = cloudflare_url,
+ )
+
+ # Fold the credentials into the link card rather than a second card below it.
+ credentials_shown = False
+ if cloudflare_url:
try:
- from google.colab import output as colab_output
- colab_output.serve_kernel_port_as_iframe(port, height = 900, width = "100%")
- except ImportError:
- pass
+ from IPython.display import HTML, display
+
+ username, password = colab_login if colab_login else (None, None)
+ display(HTML(_shareable_link_html(cloudflare_url, password, username)))
+ credentials_shown = bool(colab_login)
+ except Exception as e:
+ logger.info(f"Could not render Cloudflare link card ({e}).")
+
+ if colab_login and not credentials_shown:
+ try:
+ _show_colab_login_credentials(*colab_login)
+ except Exception as e:
+ logger.info(f"Could not render Colab login card ({e}).")
+
+ # With a tunnel up the embed below is skipped, so the ready card would only restate
+ # the link card and print a proxy URL that 404s outside this tab.
+ skip_ready_card = _is_colab_runtime() and bool(cloudflare_url)
+ if not skip_ready_card:
+ try:
+ show_link(
+ port,
+ _url = url,
+ has_cloudflare_link = bool(cloudflare_url),
+ cloudflare_requested = cloudflare_requested,
+ )
+ except Exception as e:
+ logger.info(f"Could not render Unsloth link card ({e}).")
+
+ # On Colab with a working tunnel, skip the in-cell proxy embed (often blank).
+ if _is_colab_runtime() and cloudflare_url:
+ return
+
+ # Real Colab: kernel helper needs only the port (works when eval_js failed).
+ if _is_colab_runtime():
+ if _embed_kernel_port_iframe(port):
+ return
+ _embed_html_iframe(url, port)
-def start(port: int = 8888, *, cloudflare: bool = False):
+def start(port: int = 8888, *, cloudflare: "bool | None" = None):
"""Start Unsloth Studio in Colab and display the URL.
Args:
port: Port to bind/serve on.
- cloudflare: Opt in to a shareable Cloudflare HTTPS link reachable from any
- device (default OFF). It exposes Unsloth's login page beyond Colab, so it
- stays an explicit opt-in; the default shows only the in-tab proxy iframe.
+ cloudflare: Shareable Cloudflare HTTPS link. ``None`` (default) auto-enables on
+ real Colab because the in-cell proxy embed is often blank; pass ``False`` to
+ skip the tunnel or ``True`` to force it on other runtimes.
Usage:
- start() # Colab-proxy iframe only (default)
- start(cloudflare=True) # also open a shareable Cloudflare link
+ start() # Cloudflare link on Colab (auto); proxy iframe elsewhere
+ start(cloudflare=False) # Colab proxy iframe only (often blank on current Colab)
+ start(cloudflare=True) # force Cloudflare link on any runtime
"""
import time
logger.info("🦥 Starting Unsloth Studio...")
+ use_cloudflare = _colab_wants_cloudflare(cloudflare)
- # Fast path: Unsloth already running (cell re-run). Re-launching would collide on
- # the port, so just re-show the link and iframe.
+ # Fast path: already running (cell re-run); re-show link/iframe instead of rebinding the port.
if _is_studio_healthy(port):
logger.info(f" Unsloth is already running on port {port} — reusing existing server.")
# try/finally: tear the tunnel down even if interrupted mid-start/render.
try:
- cf_url = start_cloudflare_tunnel(port) if cloudflare else None
+ colab_login = _finalize_colab_admin_password() if use_cloudflare else None
+ cf_url = start_cloudflare_tunnel(port) if use_cloudflare else None
_publish_cloudflare_url(cf_url)
- _show_and_embed(port, cloudflare_url = cf_url)
+ _show_and_embed(
+ port,
+ cloudflare_url = cf_url,
+ colab_login = colab_login,
+ cloudflare_requested = use_cloudflare,
+ )
for _ in range(10000):
time.sleep(300)
print("=", end = "", flush = True)
@@ -313,7 +664,6 @@ def start(port: int = 8888, *, cloudflare: bool = False):
logger.info(" Loading backend...")
from run import run_server
- # Auto-detect frontend path
repo_root = Path(__file__).parent.parent
frontend_path = repo_root / "frontend" / "dist"
@@ -323,8 +673,7 @@ def start(port: int = 8888, *, cloudflare: bool = False):
logger.info(" Starting server...")
try:
- # cloudflare=False: this helper owns the tunnel (Colab's own
- # start(cloudflare=...) drives it), so pin it off explicitly.
+ # cloudflare=False: this helper owns the tunnel (via start(cloudflare=...)), so pin it off.
app = run_server(
host = "0.0.0.0",
port = port,
@@ -339,14 +688,12 @@ def start(port: int = 8888, *, cloudflare: bool = False):
logger.error(f"❌ Unsloth Studio failed to start: {exc}")
return
- # run_server auto-increments the port if in use; read back the bound port so the
- # proxy URL and iframe point at the right place.
+ # run_server may auto-increment the port; read back the bound port for the proxy URL/iframe.
actual_port: int = getattr(getattr(app, "state", None), "server_port", None) or port
logger.info(f" Server started on port {actual_port}!")
- # Poll health endpoint before showing the link — avoids the race where ready_event
- # fires but the process hasn't finished binding.
+ # Poll health before showing the link: avoids the race where ready_event fires pre-bind.
import urllib.request
server_ready = False
@@ -365,12 +712,17 @@ def start(port: int = 8888, *, cloudflare: bool = False):
)
return
- # Open the tunnel now the server is healthy, publish its URL for /api/health, and
- # tear it down on interrupt (try/finally) rather than orphan the process.
+ # Server healthy: finalize Colab auth, open the tunnel, publish URL, tear down on interrupt.
try:
- cf_url = start_cloudflare_tunnel(actual_port) if cloudflare else None
+ colab_login = _finalize_colab_admin_password() if use_cloudflare else None
+ cf_url = start_cloudflare_tunnel(actual_port) if use_cloudflare else None
_publish_cloudflare_url(cf_url)
- _show_and_embed(actual_port, cloudflare_url = cf_url)
+ _show_and_embed(
+ actual_port,
+ cloudflare_url = cf_url,
+ colab_login = colab_login,
+ cloudflare_requested = use_cloudflare,
+ )
# Keep kernel alive so the daemon server thread runs.
for _ in range(10000):
diff --git a/studio/backend/core/data_recipe/jobs/manager.py b/studio/backend/core/data_recipe/jobs/manager.py
index 0e0044702e..135c9fccf6 100644
--- a/studio/backend/core/data_recipe/jobs/manager.py
+++ b/studio/backend/core/data_recipe/jobs/manager.py
@@ -27,7 +27,6 @@ from .constants import (
)
from .parse import apply_update, coerce_event, parse_log_message
from .types import Job
-from .worker import run_job_process
from loggers import get_logger
logger = get_logger(__name__)
@@ -169,12 +168,18 @@ class JobManager:
native_path_secret_removed_for_child_start,
run_without_native_path_secret,
)
+ from utils.hf_cache_settings import child_environment_for_spawn, get_hf_cache_paths
- with native_path_secret_removed_for_child_start():
+ cache_env = get_hf_cache_paths().child_env({})
+
+ with (
+ child_environment_for_spawn(cache_env),
+ native_path_secret_removed_for_child_start(),
+ ):
mp_q = _CTX.Queue()
proc = _CTX.Process(
target = run_without_native_path_secret,
- args = (run_job_process,),
+ args = ("core.data_recipe.jobs.worker", "run_job_process", cache_env),
kwargs = {"event_queue": mp_q, "recipe": recipe, "run": run_payload},
daemon = True,
)
diff --git a/studio/backend/core/data_recipe/local_callable_validators.py b/studio/backend/core/data_recipe/local_callable_validators.py
index ffc81669ae..143895d781 100644
--- a/studio/backend/core/data_recipe/local_callable_validators.py
+++ b/studio/backend/core/data_recipe/local_callable_validators.py
@@ -257,6 +257,8 @@ def _run_oxc_batch(
cwd = str(_OXC_TOOL_DIR),
input = json.dumps(payload),
text = True,
+ encoding = "utf-8",
+ errors = "replace",
capture_output = True,
check = False,
env = env,
diff --git a/studio/backend/core/export/export.py b/studio/backend/core/export/export.py
index c8be50b08b..4979ebd48d 100644
--- a/studio/backend/core/export/export.py
+++ b/studio/backend/core/export/export.py
@@ -81,6 +81,82 @@ _PYTORCH_MISSING_MESSAGE = (
_LLAMA_CPP_SCRIPTS_WARNING_EMITTED = False
+def _multi_gpu_device_map_kwargs() -> dict:
+ """``device_map`` kwargs for sharding a checkpoint across every visible GPU.
+
+ unsloth's ``from_pretrained`` defaults to ``device_map="sequential"``, which stacks
+ the whole model on GPU0 and OOMs multi-GPU hosts whose other GPUs sit empty (#7053).
+ Returns ``{"device_map": "balanced"}`` only on a real multi-GPU CUDA/ROCm host
+ (mirroring the inference loader's ``get_device_map``), else empty so single-GPU, CPU
+ and MLX loads keep the loader default."""
+ if _IS_MLX:
+ return {}
+ try:
+ from utils.hardware import get_device_map, get_parent_visible_gpu_ids
+
+ visible = get_parent_visible_gpu_ids()
+ if len(visible) > 1:
+ device_map = get_device_map(visible)
+ elif not visible:
+ # UUID/MIG masks resolve to no numeric ids; get_device_map(None) falls back
+ # to the visible-GPU count, so a multi-GPU UUID/MIG host still shards.
+ device_map = get_device_map(None)
+ else:
+ return {}
+ if device_map == "balanced":
+ return {"device_map": device_map}
+ except Exception as exc:
+ logger.debug(f"multi-GPU device_map resolution failed; using loader default: {exc}")
+ return {}
+
+
+def _is_oom_error(exc: BaseException) -> bool:
+ """True for an accelerator OOM, however it is spelled.
+
+ accelerate and transformers re-raise it as a plain ``RuntimeError`` on several paths
+ and ROCm/XPU use their own classes, so match the message too.
+ """
+ if torch is not None:
+ oom_types = tuple(
+ t
+ for t in (
+ getattr(torch, "OutOfMemoryError", None),
+ getattr(getattr(torch, "cuda", None), "OutOfMemoryError", None),
+ getattr(getattr(torch, "xpu", None), "OutOfMemoryError", None),
+ )
+ if isinstance(t, type)
+ )
+ if oom_types and isinstance(exc, oom_types):
+ return True
+ return "out of memory" in f"{type(exc).__name__}: {exc}".lower()
+
+
+def _is_cpu_spill_rejection(exc: BaseException) -> bool:
+ """bitsandbytes refuses a map that spills to CPU/disk with a plain ``ValueError``.
+
+ Busy secondary GPUs can make ``balanced`` spill to CPU even where the old sequential
+ load fit on GPU0, and that message says nothing about memory, so the retry has to
+ match it explicitly. See transformers ``quantizers/quantizer_bnb_4bit.py``.
+ """
+ return "dispatched on the cpu or the disk" in str(exc).lower()
+
+
+class _CpuSpillRetry(Exception):
+ """A multi-GPU load that succeeded but left modules offloaded to CPU/disk."""
+
+
+def _cpu_offloaded_modules(model) -> int:
+ """Count the modules a load parked on CPU or disk.
+
+ Only bitsandbytes refuses such a map; a full-precision load accepts it, leaves the
+ parameters on meta and dies much later in safetensors with "Cannot copy out of meta
+ tensor". Nothing raises at load time, so inspect the map directly. PEFT re-dispatches
+ when attaching an adapter, so in practice this catches merged checkpoints.
+ """
+ device_map = getattr(model, "hf_device_map", None) or {}
+ return sum(1 for target in device_map.values() if str(target) in ("cpu", "disk"))
+
+
def _supports_kwarg(fn, name):
"""True if `fn` accepts keyword `name` directly or via **kwargs."""
import inspect
@@ -165,7 +241,7 @@ def _offline_window_if(local_files_only):
def _is_wsl():
"""Detect if running under Windows Subsystem for Linux."""
try:
- return "microsoft" in open("/proc/version").read().lower()
+ return "microsoft" in open("/proc/version", encoding = "utf-8").read().lower()
except Exception:
return False
@@ -271,6 +347,7 @@ class ExportBackend:
load_in_4bit: bool = True,
trust_remote_code: bool = False,
hf_token: Optional[str] = None,
+ _device_map_override: Optional[dict] = None,
) -> Tuple[bool, str]:
"""
Load a checkpoint for export.
@@ -303,6 +380,14 @@ class ExportBackend:
# Skip the Hub when offline so a no-internet export uses the local cache.
local_files_only = _hf_offline()
+ # Shard across every visible GPU instead of stacking on GPU0 (#7053); {} on
+ # single-GPU/CPU/MLX. _device_map_override is the single-device retry below.
+ _device_map_kw = (
+ _multi_gpu_device_map_kwargs()
+ if _device_map_override is None
+ else _device_map_override
+ )
+
# Run the type-detection probes in the forced-offline window (else a gated
# base 404s); it covers is_vision_model's Hub reads + the transformers-5
# subprocess, and local_files_only makes detect_audio_type's requests.get skip.
@@ -328,6 +413,7 @@ class ExportBackend:
trust_remote_code = trust_remote_code,
token = token,
local_files_only = local_files_only,
+ **_device_map_kw,
)
elif self._audio_type == "whisper":
@@ -343,6 +429,7 @@ class ExportBackend:
trust_remote_code = trust_remote_code,
token = token,
local_files_only = local_files_only,
+ **_device_map_kw,
)
elif self._audio_type == "snac":
@@ -355,6 +442,7 @@ class ExportBackend:
trust_remote_code = trust_remote_code,
token = token,
local_files_only = local_files_only,
+ **_device_map_kw,
)
elif self._audio_type == "bicodec":
@@ -368,6 +456,7 @@ class ExportBackend:
trust_remote_code = trust_remote_code,
token = token,
local_files_only = local_files_only,
+ **_device_map_kw,
)
elif self._audio_type == "dac":
@@ -380,6 +469,7 @@ class ExportBackend:
trust_remote_code = trust_remote_code,
token = token,
local_files_only = local_files_only,
+ **_device_map_kw,
)
elif self.is_vision:
@@ -392,6 +482,7 @@ class ExportBackend:
trust_remote_code = trust_remote_code,
token = token,
local_files_only = local_files_only,
+ **_device_map_kw,
)
tokenizer = processor # vision: processor acts as tokenizer
@@ -405,8 +496,16 @@ class ExportBackend:
trust_remote_code = trust_remote_code,
token = token,
local_files_only = local_files_only,
+ **_device_map_kw,
)
+ # Only when we asked for the multi-GPU map: a single-GPU host has no second
+ # placement to retry on, so leave its behaviour untouched.
+ _offloaded = _cpu_offloaded_modules(model) if _device_map_kw else 0
+ if _device_map_override is None and _offloaded:
+ del model
+ raise _CpuSpillRetry(f"{_offloaded} module(s) offloaded to CPU/disk")
+
if _IS_MLX:
# MLX doesn't use PeftModel — detect LoRA via adapter_config.json
self.is_peft = adapter_config.exists()
@@ -429,11 +528,41 @@ class ExportBackend:
return True, f"Loaded {model_type} model{peft_info} successfully"
except Exception as e:
- logger.error(f"Error loading checkpoint: {e}")
- import traceback
+ # Sharding is an optimisation, never a requirement. "balanced" budgets from the
+ # free memory read BEFORE this process opens a CUDA context on each GPU, so when
+ # a training or chat job already owns the others the shard can OOM, or spill to
+ # CPU and be refused by bitsandbytes, where the old single-device load succeeded.
+ # Fall back once before giving up.
+ if (
+ _device_map_override is None
+ and (
+ isinstance(e, _CpuSpillRetry) or _is_oom_error(e) or _is_cpu_spill_rejection(e)
+ )
+ and _multi_gpu_device_map_kwargs()
+ ):
+ # Retry outside this block: the live traceback pins the half-built model's
+ # frames, so an in-block retry inherits the exhausted device.
+ retry_reason = str(e)
+ else:
+ logger.error(f"Error loading checkpoint: {e}")
+ import traceback
- logger.error(traceback.format_exc())
- return False, f"Failed to load checkpoint: {str(e)}"
+ logger.error(traceback.format_exc())
+ return False, f"Failed to load checkpoint: {str(e)}"
+
+ logger.warning(
+ f"Multi-GPU export load unusable ({retry_reason}); retrying on "
+ f"the single-device loader default."
+ )
+ self.cleanup_memory()
+ return self.load_checkpoint(
+ checkpoint_path,
+ max_seq_length = max_seq_length,
+ load_in_4bit = load_in_4bit,
+ trust_remote_code = trust_remote_code,
+ hf_token = hf_token,
+ _device_map_override = {},
+ )
def _write_export_metadata(self, save_directory: str):
"""Write export_metadata.json with base model info for Chat page discovery."""
@@ -445,7 +574,7 @@ class ExportBackend:
)
metadata = {"base_model": base_model}
metadata_path = os.path.join(save_directory, "export_metadata.json")
- with open(metadata_path, "w") as f:
+ with open(metadata_path, "w", encoding = "utf-8") as f:
json.dump(metadata, f, indent = 2)
logger.info(f"Wrote export metadata to {metadata_path}")
except Exception as e:
@@ -1048,6 +1177,21 @@ class ExportBackend:
"Use the safetensors adapter instead.",
None,
)
+ # llama.cpp's convert_lora_to_gguf.py has no concept of DoRA's
+ # lora_magnitude_vector tensors: it only reads the standard
+ # lora_A/lora_B delta, so exporting a DoRA adapter would silently
+ # drop the magnitude rescaling and produce a GGUF LoRA file that
+ # loads fine but no longer matches the trained model.
+ _peft_config = getattr(self.current_model, "peft_config", {}).get("default")
+ if getattr(_peft_config, "use_dora", False):
+ return (
+ False,
+ "GGUF LoRA export is not supported for DoRA adapters: the GGUF LoRA "
+ "format has no way to represent DoRA's magnitude vectors, so the "
+ "exported file would silently lose the DoRA behavior. Use the "
+ "safetensors adapter instead, or merge to a full GGUF model.",
+ None,
+ )
outtype = str(gguf_outtype).lower()
if outtype not in _GGUF_LORA_OUTTYPES:
return (
diff --git a/studio/backend/core/export/orchestrator.py b/studio/backend/core/export/orchestrator.py
index 6d1a928f2e..aaf48615f0 100644
--- a/studio/backend/core/export/orchestrator.py
+++ b/studio/backend/core/export/orchestrator.py
@@ -230,16 +230,20 @@ class ExportOrchestrator:
native_path_secret_removed_for_child_start,
run_without_native_path_secret,
)
+ from utils.hf_cache_settings import child_environment_for_spawn, get_hf_cache_paths
- from .worker import run_export_process
+ cache_env = get_hf_cache_paths().child_env({})
- with native_path_secret_removed_for_child_start():
+ with (
+ child_environment_for_spawn(cache_env),
+ native_path_secret_removed_for_child_start(),
+ ):
self._cmd_queue = _CTX.Queue()
self._resp_queue = _CTX.Queue()
self._proc = _CTX.Process(
target = run_without_native_path_secret,
- args = (run_export_process,),
+ args = ("core.export.worker", "run_export_process", cache_env),
kwargs = {
"cmd_queue": self._cmd_queue,
"resp_queue": self._resp_queue,
diff --git a/studio/backend/core/inference/_vulkan_probe.py b/studio/backend/core/inference/_vulkan_probe.py
index 706346daad..4bfefc21ce 100644
--- a/studio/backend/core/inference/_vulkan_probe.py
+++ b/studio/backend/core/inference/_vulkan_probe.py
@@ -6,12 +6,14 @@
Run in a short-lived subprocess (``python _vulkan_probe.py ``) so the
Vulkan instance never lives in the long-running backend process. Loads the
bundled ggml Vulkan backend from ```` and prints one
-``\\t\\t\\t`` line per device to stdout.
-Indices are ggml's own Vulkan device ordinals, which need not match nvidia-smi
-order. ``is_igpu`` (from ggml's device type) is ``1`` for an integrated GPU
-sharing system RAM. ``total_bytes`` is the device-local heap; the reader uses
-it to reserve absolute headroom on a discrete card (parity with the CUDA/ROCm
-fit) and ignores it for an iGPU, whose "VRAM" is shared system RAM.
+``\\t\\t\\t\\t`` line per device to
+stdout. Indices are ggml's own Vulkan device ordinals, which need not match
+nvidia-smi order. ``is_igpu`` (from ggml's device type) is ``1`` for an
+integrated GPU sharing system RAM. ``total_bytes`` is the device-local heap;
+the reader uses it to reserve absolute headroom on a discrete card (parity
+with the CUDA/ROCm fit) and ignores it for an iGPU, whose "VRAM" is shared
+system RAM. ``name`` is ggml's device description (the marketing name, e.g.
+"AMD Radeon RX 9070 XT"); empty when the registry lookup fails.
Uses only the standard library so it stays runnable as a bare script.
"""
@@ -24,15 +26,30 @@ import sys
_GGML_BACKEND_DEVICE_TYPE_IGPU = 2
-def _igpu_flags(base, lib, count: int) -> list[bool]:
- """Per-device integrated-GPU flags via ggml's backend registry.
+def _igpu_flags_and_names(base, lib, count: int) -> tuple[list[bool], list[str]]:
+ """Per-device integrated-GPU flags and descriptions via ggml's backend registry.
The Vulkan reg enumerates devices in the same order as
``ggml_backend_vk_get_device_memory`` (each context uses ``ctx->device =
- i``), so reg index == device ordinal. Returns all-False on any failure so
- the reader never over-caps a discrete card.
+ i``), so reg index == device ordinal. Returns all-False / empty-name on any
+ failure so the reader never over-caps a discrete card and the memory
+ readings still get through.
"""
flags = [False] * count
+ names = [""] * count
+
+ # The name lookup is bound OUTSIDE the type-detection try: a ggml-base
+ # without ggml_backend_dev_description (older/custom build) must degrade to
+ # unnamed devices, not abort before the iGPU flags are read (which would
+ # count an iGPU's shared RAM as VRAM).
+ describe = None
+ try:
+ base.ggml_backend_dev_description.restype = ctypes.c_char_p
+ base.ggml_backend_dev_description.argtypes = [ctypes.c_void_p]
+ describe = base.ggml_backend_dev_description
+ except Exception:
+ pass
+
try:
lib.ggml_backend_vk_reg.restype = ctypes.c_void_p
lib.ggml_backend_vk_reg.argtypes = []
@@ -45,17 +62,31 @@ def _igpu_flags(base, lib, count: int) -> list[bool]:
reg = lib.ggml_backend_vk_reg()
if not reg:
- return flags
+ return flags, names
dev_count = base.ggml_backend_reg_dev_count(reg)
for i in range(min(count, dev_count)):
dev = base.ggml_backend_reg_dev_get(reg, i)
if dev:
flags[i] = base.ggml_backend_dev_type(dev) == _GGML_BACKEND_DEVICE_TYPE_IGPU
+ if describe is not None:
+ try:
+ desc = describe(dev)
+ if desc:
+ # Tabs/newlines would corrupt the line protocol;
+ # spaces are safe.
+ names[i] = (
+ desc.decode("utf-8", errors = "replace")
+ .replace("\t", " ")
+ .replace("\n", " ")
+ .strip()
+ )
+ except Exception:
+ pass
except Exception:
- # Best-effort: any failure degrades to "discrete" so the memory
- # readings still get through instead of crashing the probe.
+ # Best-effort: any failure degrades to "discrete"/"unnamed" so the
+ # memory readings still get through instead of crashing the probe.
pass
- return flags
+ return flags, names
def main() -> int:
@@ -63,6 +94,14 @@ def main() -> int:
return 0
bindir = sys.argv[1]
+ # Device names can be non-ASCII (localized drivers); the platform-default
+ # stdout encoding (e.g. cp1252) would raise on them and lose the whole
+ # inventory. The reader decodes UTF-8 with the same error mode.
+ try:
+ sys.stdout.reconfigure(encoding = "utf-8", errors = "replace")
+ except Exception:
+ pass
+
# Hold add_dll_directory's handle for the rest of main() (the documented
# idiom) so bindir stays on the search path while the sibling ggml DLLs
# resolve below.
@@ -96,12 +135,12 @@ def main() -> int:
]
count = lib.ggml_backend_vk_get_device_count()
- igpu = _igpu_flags(base, lib, count)
+ igpu, names = _igpu_flags_and_names(base, lib, count)
rows = []
for i in range(count):
free, total = ctypes.c_size_t(0), ctypes.c_size_t(0)
lib.ggml_backend_vk_get_device_memory(i, ctypes.byref(free), ctypes.byref(total))
- rows.append("%d\t%d\t%d\t%d" % (i, free.value, int(igpu[i]), total.value))
+ rows.append("%d\t%d\t%d\t%d\t%s" % (i, free.value, int(igpu[i]), total.value, names[i]))
sys.stdout.write("\n".join(rows))
return 0
diff --git a/studio/backend/core/inference/anthropic_compat.py b/studio/backend/core/inference/anthropic_compat.py
index 34445cc58e..a32e372d73 100644
--- a/studio/backend/core/inference/anthropic_compat.py
+++ b/studio/backend/core/inference/anthropic_compat.py
@@ -172,6 +172,136 @@ def anthropic_messages_to_openai(
return result
+_ANTHROPIC_SCHEMA_CLIENT_TOOL_PARAMETERS = {
+ "bash": {
+ "type": "object",
+ "properties": {
+ "command": {"type": "string"},
+ "restart": {"type": "boolean"},
+ },
+ "anyOf": [
+ {"required": ["command"]},
+ {"properties": {"restart": {"const": True}}, "required": ["restart"]},
+ ],
+ },
+ "text_editor": {
+ "type": "object",
+ "properties": {
+ "command": {
+ "type": "string",
+ "enum": ["view", "str_replace", "create", "insert"],
+ },
+ "path": {"type": "string"},
+ "view_range": {
+ "type": "array",
+ "items": {"type": "integer"},
+ "minItems": 2,
+ "maxItems": 2,
+ },
+ "old_str": {"type": "string"},
+ "new_str": {"type": "string"},
+ "file_text": {"type": "string"},
+ "insert_line": {"type": "integer"},
+ "insert_text": {"type": "string"},
+ },
+ "required": ["command", "path"],
+ },
+ "computer": {
+ "type": "object",
+ "properties": {
+ "action": {"type": "string"},
+ "coordinate": {
+ "type": "array",
+ "items": {"type": "integer"},
+ "minItems": 2,
+ "maxItems": 2,
+ },
+ "text": {"type": "string"},
+ "duration": {"type": "number"},
+ "scroll_direction": {"type": "string"},
+ "scroll_amount": {"type": "integer"},
+ "start_coordinate": {
+ "type": "array",
+ "items": {"type": "integer"},
+ "minItems": 2,
+ "maxItems": 2,
+ },
+ "key": {"type": "string"},
+ },
+ "required": ["action"],
+ "additionalProperties": True,
+ },
+ "memory": {
+ "type": "object",
+ "properties": {
+ "command": {
+ "type": "string",
+ "enum": ["view", "create", "str_replace", "insert", "delete", "rename"],
+ },
+ "path": {"type": "string"},
+ "view_range": {
+ "type": "array",
+ "items": {"type": "integer"},
+ "minItems": 2,
+ "maxItems": 2,
+ },
+ "file_text": {"type": "string"},
+ "old_str": {"type": "string"},
+ "new_str": {"type": "string"},
+ "insert_line": {"type": "integer"},
+ "insert_text": {"type": "string"},
+ "old_path": {"type": "string"},
+ "new_path": {"type": "string"},
+ },
+ "required": ["command"],
+ },
+}
+
+_ANTHROPIC_SCHEMA_CLIENT_TOOL_DESCRIPTIONS = {
+ "bash": "Run a command in the caller-owned persistent bash session, or restart it.",
+ "text_editor": "View, create, or edit files in the caller-owned filesystem.",
+ "computer": "Interact with the caller-owned computer using an action and its parameters.",
+ "memory": "Store and retrieve files in the caller-owned persistent memory directory.",
+}
+
+
+def anthropic_schema_client_tool_kind(tool) -> Optional[str]:
+ """Return the kind of a schema-less Anthropic client tool, if recognized."""
+ td = tool if isinstance(tool, dict) else tool.model_dump()
+ if td.get("input_schema") is not None:
+ return None
+ type_ = td.get("type")
+ if not isinstance(type_, str):
+ return None
+ kind, separator, version = type_.rpartition("_")
+ if (
+ separator
+ and kind in _ANTHROPIC_SCHEMA_CLIENT_TOOL_PARAMETERS
+ and len(version) == 8
+ and version.isdigit()
+ ):
+ return kind
+ return None
+
+
+def _anthropic_schema_client_tool_parameters(td: dict, kind: str) -> dict:
+ parameters = _ANTHROPIC_SCHEMA_CLIENT_TOOL_PARAMETERS[kind]
+ if kind != "text_editor":
+ return parameters
+
+ version = td["type"].rpartition("_")[2]
+ commands = list(parameters["properties"]["command"]["enum"])
+ if version < "20250429":
+ commands.append("undo_edit")
+ return {
+ **parameters,
+ "properties": {
+ **parameters["properties"],
+ "command": {**parameters["properties"]["command"], "enum": commands},
+ },
+ }
+
+
def anthropic_tools_to_openai(tools: list) -> list[dict]:
"""Convert Anthropic client tools to OpenAI function-tool format."""
result = []
@@ -179,6 +309,9 @@ def anthropic_tools_to_openai(tools: list) -> list[dict]:
td = t if isinstance(t, dict) else t.model_dump()
name = td.get("name")
input_schema = td.get("input_schema")
+ schema_client_kind = anthropic_schema_client_tool_kind(td)
+ if schema_client_kind is not None:
+ input_schema = _anthropic_schema_client_tool_parameters(td, schema_client_kind)
if not name or input_schema is None:
continue
result.append(
@@ -186,7 +319,8 @@ def anthropic_tools_to_openai(tools: list) -> list[dict]:
"type": "function",
"function": {
"name": name,
- "description": td.get("description", ""),
+ "description": td.get("description")
+ or _ANTHROPIC_SCHEMA_CLIENT_TOOL_DESCRIPTIONS.get(schema_client_kind, ""),
"parameters": input_schema,
},
}
diff --git a/studio/backend/core/inference/api_monitor.py b/studio/backend/core/inference/api_monitor.py
index f76a38576f..b637ba56d1 100644
--- a/studio/backend/core/inference/api_monitor.py
+++ b/studio/backend/core/inference/api_monitor.py
@@ -5,6 +5,7 @@
from __future__ import annotations
+import os
import threading
import time
import uuid
@@ -18,6 +19,14 @@ _MAX_PROMPT_CHARS = 12000
_MAX_REPLY_CHARS = 12000
_PREVIEW_CHARS = 360
+# Opt-in startup kill switch for Studio's in-memory API monitor.
+_DISABLE_ENV = "UNSLOTH_STUDIO_DISABLE_API_MONITOR"
+_TRUE_VALUES = frozenset({"1", "true", "yes", "on"})
+
+
+def _api_monitor_disabled() -> bool:
+ return os.environ.get(_DISABLE_ENV, "").strip().lower() in _TRUE_VALUES
+
def _trim(text: Optional[str], limit: int) -> str:
if not text:
@@ -52,6 +61,13 @@ class ApiMonitorEntry:
total_tokens: Optional[int] = None
total_tokens_authoritative: bool = False
error: Optional[str] = None
+ # "request" (HTTP call) or "lifecycle" (model load/unload: event/reason, not a prompt; shared).
+ kind: str = "request"
+ event: Optional[str] = None
+ reason: Optional[str] = None
+ shared: bool = False
+ # 0-100 for a running download row; None when not applicable.
+ progress: Optional[float] = None
def snapshot(self, *, include_details: bool = True) -> dict[str, Any]:
duration_ms = None
@@ -85,6 +101,10 @@ class ApiMonitorEntry:
"completion_tokens": self.completion_tokens,
"total_tokens": self.total_tokens,
"error": self.error,
+ "kind": self.kind,
+ "event": self.event,
+ "reason": self.reason,
+ "progress": self.progress,
}
if include_details:
payload["prompt"] = self.prompt
@@ -93,10 +113,16 @@ class ApiMonitorEntry:
class ApiMonitor:
- def __init__(self, max_entries: int = _MAX_ENTRIES):
+ def __init__(
+ self,
+ max_entries: int = _MAX_ENTRIES,
+ *,
+ enabled: bool = True,
+ ):
self._entries: deque[ApiMonitorEntry] = deque()
self._max_entries = max(0, max_entries)
self._lock = threading.Lock()
+ self._enabled = enabled
def start(
self,
@@ -108,6 +134,8 @@ class ApiMonitor:
context_length: Optional[int] = None,
subject: Optional[str] = None,
) -> str:
+ if not self._enabled:
+ return ""
now = time.time()
entry = ApiMonitorEntry(
id = f"apireq_{uuid.uuid4().hex[:12]}",
@@ -127,6 +155,75 @@ class ApiMonitor:
self._trim_terminal_locked()
return entry.id
+ def record_lifecycle(
+ self,
+ *,
+ event: str,
+ model: str,
+ reason: Optional[str] = None,
+ running: bool = False,
+ ) -> str:
+ """Record a model load/unload alongside the request traffic that caused it.
+
+ ``running=True`` opens the row for the caller to close with :meth:`finish` /
+ :meth:`fail`; an unload is terminal on arrival. Rows are shared (visible to
+ every subject) and share the request retention budget.
+ """
+ if not self._enabled:
+ return ""
+ now = time.time()
+ entry = ApiMonitorEntry(
+ id = f"apievt_{uuid.uuid4().hex[:12]}",
+ endpoint = f"model.{event}",
+ method = "",
+ model = model or "default",
+ prompt = "",
+ status = "running" if running else "completed",
+ started_at = now,
+ updated_at = now,
+ started_monotonic = time.monotonic(),
+ finished_at = None if running else now,
+ finished_monotonic = None if running else time.monotonic(),
+ kind = "lifecycle",
+ event = event,
+ reason = reason,
+ shared = True,
+ )
+ with self._lock:
+ self._entries.appendleft(entry)
+ self._trim_terminal_locked()
+ return entry.id
+
+ def relabel(self, entry_id: Optional[str], model: str) -> None:
+ """Rename an open lifecycle row once the load resolves its real id: up front
+ the caller only has the load path, which may be an HF snapshot dir."""
+ if not entry_id or not model:
+ return
+ with self._lock:
+ entry = self._find_locked(entry_id)
+ if entry is not None:
+ entry.model = model
+ entry.updated_at = time.time()
+
+ def set_progress(self, entry_id: Optional[str], progress: Optional[float]) -> None:
+ """Update an open download row's percentage (clamped to 0-100)."""
+ if not entry_id or progress is None:
+ return
+ with self._lock:
+ entry = self._find_locked(entry_id)
+ if entry is not None and entry.status == "running":
+ entry.progress = min(100.0, max(0.0, float(progress)))
+ entry.updated_at = time.time()
+
+ def discard(self, entry_id: Optional[str]) -> None:
+ """Drop a row that turned out not to be an event (an already-satisfied load)."""
+ if not entry_id:
+ return
+ with self._lock:
+ entry = self._find_locked(entry_id)
+ if entry is not None:
+ self._entries.remove(entry)
+
def append_reply(self, entry_id: Optional[str], text: str) -> None:
if not entry_id or not text:
return
@@ -212,6 +309,18 @@ class ApiMonitor:
self._entries.appendleft(entry)
self._trim_terminal_locked()
+ def fail_open(self, entry_id: Optional[str], error: str) -> None:
+ """Fail only a still-open row: unlike :meth:`fail`, a catch-all in a
+ ``finally`` cannot stamp an error onto a request that already succeeded."""
+ if not entry_id:
+ return
+ with self._lock:
+ entry = self._find_locked(entry_id)
+ if entry is None or entry.finished_at is not None:
+ return
+ # Same lock as the check, so a finish() cannot land in between.
+ self._fail_locked(entry, error)
+
def fail(self, entry_id: Optional[str], error: str) -> None:
if not entry_id:
return
@@ -224,15 +333,18 @@ class ApiMonitor:
if error:
entry.error = _trim(error, 1000)
return
- now = time.time()
- entry.status = "error"
- entry.error = _trim(error, 1000)
- entry.updated_at = now
- entry.finished_at = now
- entry.finished_monotonic = time.monotonic()
- self._entries.remove(entry)
- self._entries.appendleft(entry)
- self._trim_terminal_locked()
+ self._fail_locked(entry, error)
+
+ def _fail_locked(self, entry: ApiMonitorEntry, error: str) -> None:
+ now = time.time()
+ entry.status = "error"
+ entry.error = _trim(error, 1000)
+ entry.updated_at = now
+ entry.finished_at = now
+ entry.finished_monotonic = time.monotonic()
+ self._entries.remove(entry)
+ self._entries.appendleft(entry)
+ self._trim_terminal_locked()
def snapshot(
self,
@@ -244,7 +356,7 @@ class ApiMonitor:
return [
entry.snapshot(include_details = include_details)
for entry in self._entries
- if subject is None or entry.subject == subject
+ if self._visible(entry, subject)
]
def get(
@@ -257,22 +369,29 @@ class ApiMonitor:
entry = self._find_locked(entry_id)
if entry is None:
return None
- if subject is not None and entry.subject != subject:
+ if not self._visible(entry, subject):
return None
return entry.snapshot(include_details = True)
def active_count(self, *, subject: Optional[str] = None) -> int:
+ # Lifecycle rows show as "running" while loading but are not in-flight API requests.
with self._lock:
return sum(
1
for entry in self._entries
- if entry.status == "running" and (subject is None or entry.subject == subject)
+ if entry.status == "running"
+ and entry.kind != "lifecycle"
+ and (subject is None or entry.subject == subject)
)
def clear(self) -> None:
with self._lock:
self._entries.clear()
+ @staticmethod
+ def _visible(entry: ApiMonitorEntry, subject: Optional[str]) -> bool:
+ return subject is None or entry.subject == subject or entry.shared
+
def _find_locked(self, entry_id: str) -> Optional[ApiMonitorEntry]:
for entry in self._entries:
if entry.id == entry_id:
@@ -292,4 +411,4 @@ class ApiMonitor:
self._entries = kept
-api_monitor = ApiMonitor()
+api_monitor = ApiMonitor(enabled = not _api_monitor_disabled())
diff --git a/studio/backend/core/inference/audio_codecs.py b/studio/backend/core/inference/audio_codecs.py
index 93c7da72cb..b59f2bcce0 100644
--- a/studio/backend/core/inference/audio_codecs.py
+++ b/studio/backend/core/inference/audio_codecs.py
@@ -76,8 +76,14 @@ class AudioCodecManager:
if self._snac_model is not None:
return
from snac import SNAC
+ from utils.hf_cache_settings import active_hf_hub_cache
- self._snac_model = SNAC.from_pretrained("hubertsiuzdak/snac_24khz").to(device).eval()
+ # Route weights to the selected cache; this can run in the main process.
+ self._snac_model = (
+ SNAC.from_pretrained("hubertsiuzdak/snac_24khz", cache_dir = active_hf_hub_cache())
+ .to(device)
+ .eval()
+ )
logger.info("Loaded SNAC codec (24kHz)")
def _load_bicodec(
diff --git a/studio/backend/core/inference/bypass_site/sitecustomize.py b/studio/backend/core/inference/bypass_site/sitecustomize.py
deleted file mode 100644
index 36ff0d197a..0000000000
--- a/studio/backend/core/inference/bypass_site/sitecustomize.py
+++ /dev/null
@@ -1,8 +0,0 @@
-"""Load path remapping without sandbox network guards."""
-
-import runpy
-from pathlib import Path
-
-
-_SHIM = Path(__file__).resolve().parents[1] / "sandbox_site" / "sitecustomize.py"
-runpy.run_path(str(_SHIM))
diff --git a/studio/backend/core/inference/chat_template_helpers.py b/studio/backend/core/inference/chat_template_helpers.py
index 528c059fbc..3a8463855b 100644
--- a/studio/backend/core/inference/chat_template_helpers.py
+++ b/studio/backend/core/inference/chat_template_helpers.py
@@ -326,6 +326,58 @@ def _normalize_tool_call_arguments(messages: list) -> list:
return out if mutated else messages
+def _take_tool_result(pending: list, call_id) -> Optional[dict]:
+ if call_id:
+ for i, result in enumerate(pending):
+ if result.get("tool_call_id") == call_id:
+ return pending.pop(i)
+ for i, result in enumerate(pending):
+ if not result.get("tool_call_id"):
+ return pending.pop(i)
+ return None
+
+
+def _split_parallel_tool_calls(messages: list) -> list:
+ """Llama 3.x templates render one call per message, so split parallel calls
+ into consecutive single-call messages, each followed by its own result."""
+ if not any(isinstance(m, dict) and len(m.get("tool_calls") or ()) > 1 for m in messages):
+ return messages
+
+ out: list = []
+ i = 0
+ total = len(messages)
+ while i < total:
+ msg = messages[i]
+ calls = msg.get("tool_calls") if isinstance(msg, dict) else None
+ if not calls or len(calls) <= 1:
+ out.append(msg)
+ i += 1
+ continue
+
+ # Tool results right after this message answer its calls.
+ j = i + 1
+ pending: list = []
+ while (
+ j < total
+ and isinstance(messages[j], dict)
+ and messages[j].get("role") in ("tool", "ipython")
+ ):
+ pending.append(messages[j])
+ j += 1
+
+ for idx, call in enumerate(calls):
+ piece = {**msg, "tool_calls": [call]}
+ if idx:
+ piece["content"] = ""
+ out.append(piece)
+ result = _take_tool_result(pending, call.get("id") if isinstance(call, dict) else None)
+ if result is not None:
+ out.append(result)
+ out.extend(pending)
+ i = j
+ return out
+
+
def apply_chat_template_for_generation(
tokenizer,
messages: list,
@@ -378,13 +430,21 @@ def apply_chat_template_for_generation(
try:
return _render(messages)
except Exception:
- # Strict tool templates reject the JSON-string ``arguments`` form via
- # TypeError or a broad Jinja raise_exception, so retry with dicts coerced.
- # Original messages render first, so working templates stay byte-identical.
+ # Retry with repairs applied cumulatively. Originals render first, so
+ # working templates stay byte-identical.
+ candidates: list = []
normalized = _normalize_tool_call_arguments(messages)
- if normalized is messages:
- raise
- return _render(normalized)
+ if normalized is not messages:
+ candidates.append(normalized)
+ split = _split_parallel_tool_calls(normalized)
+ if split is not normalized:
+ candidates.append(split)
+ for candidate in candidates:
+ try:
+ return _render(candidate)
+ except Exception:
+ continue
+ raise
def render_native_template(
diff --git a/studio/backend/core/inference/inference.py b/studio/backend/core/inference/inference.py
index 8d262bbb0f..e78bf1be8d 100644
--- a/studio/backend/core/inference/inference.py
+++ b/studio/backend/core/inference/inference.py
@@ -8,6 +8,7 @@ from unsloth.chat_templates import get_chat_template
from transformers import TextIteratorStreamer, TextStreamer
from peft import PeftModel, PeftModelForCausalLM
+import contextlib
import json
import sys
import torch
@@ -566,7 +567,7 @@ class InferenceBackend:
_meta_path = Path(config.path) / "export_metadata.json"
try:
if _meta_path.exists():
- _meta = json.loads(_meta_path.read_text())
+ _meta = json.loads(_meta_path.read_text(encoding = "utf-8-sig"))
if _meta.get("base_model"):
processor_source = _meta["base_model"]
except Exception:
@@ -1942,8 +1943,30 @@ class InferenceBackend:
+ text
+ "<|text_end|>\n<|audio_start|><|global_features_start|>\n"
)
+
with torch.inference_mode():
- with torch.amp.autocast("cuda", dtype = model.dtype):
+ # Derive the autocast device from the loaded model, not from the
+ # global backend: a CPU-fallback DAC on an XPU/CUDA host must not
+ # open a GPU autocast context around CPU tensors.
+ device_type = (
+ model.device.type
+ if hasattr(model.device, "type")
+ else str(model.device).split(":", 1)[0]
+ )
+ # Clamp to autocast-supported backends so exotic devices
+ # (e.g. "meta" during accelerate offloaded loading) do not raise.
+ # MPS is autocast-supported since torch 2.3, keep it in the set.
+ if device_type not in ("cuda", "xpu", "mps", "cpu"):
+ device_type = "cpu"
+ # CPU and XPU autocast only accept bfloat16/float16. For a
+ # float32 model, skip autocast entirely to avoid raising or
+ # producing a warning on every generate call.
+ autocast_dtype_supported = model.dtype in (torch.bfloat16, torch.float16)
+ if device_type in ("cpu", "xpu") and not autocast_dtype_supported:
+ autocast_ctx = contextlib.nullcontext()
+ else:
+ autocast_ctx = torch.amp.autocast(device_type, dtype = model.dtype)
+ with autocast_ctx:
inputs = tokenizer([prompt], return_tensors = "pt").to(model.device)
generated = model.generate(
**inputs,
@@ -2258,8 +2281,13 @@ class InferenceBackend:
except Exception as e:
logger.warning(f"Could not fully reset model state for {model_name}: {e}")
- def reset_generation_state(self):
- """Reset any cached generation state to prevent hanging after errors"""
+ def reset_generation_state(self, caller_cancel_event = None):
+ """Reset any cached generation state to prevent hanging after errors
+
+ ``caller_cancel_event`` is accepted for signature parity with the
+ orchestrator, which uses it to drop a reset from a request that never
+ started. Nothing here cancels a live generation, so it is unused.
+ """
try:
# Clear cached state for ALL loaded models
for model_name in self.models.keys():
diff --git a/studio/backend/core/inference/llama_admission.py b/studio/backend/core/inference/llama_admission.py
index b6a939c87b..7bf0dd7429 100644
--- a/studio/backend/core/inference/llama_admission.py
+++ b/studio/backend/core/inference/llama_admission.py
@@ -13,37 +13,159 @@ from __future__ import annotations
import asyncio
import os
+import sys
import threading
from collections import deque
from dataclasses import dataclass
from typing import Deque, Optional
-ADMISSION_CONTROL_ENV = "UNSLOTH_OPENAI_COMPAT_ADMISSION_CONTROL"
-ADMISSION_QUEUE_TIMEOUT_ENV = "UNSLOTH_OPENAI_COMPAT_ADMISSION_QUEUE_TIMEOUT"
-ADMISSION_KEEPALIVE_INTERVAL_ENV = "UNSLOTH_OPENAI_COMPAT_ADMISSION_KEEPALIVE_INTERVAL"
-ADMISSION_MAX_QUEUE_ENV = "UNSLOTH_OPENAI_COMPAT_ADMISSION_MAX_QUEUE"
+# dataclass(slots = True) halves per-instance overhead. Measured as perf-neutral
+# here, not a speed win: it costs a little on construction and gains it back on
+# access. It is 3.10+ and this package declares >=3.9, so gate it rather than
+# dropping it outright. Empty on 3.9 means a plain dataclass.
+_SLOTS = {"slots": True} if sys.version_info >= (3, 10) else {}
+
+
+ADMISSION_CONTROL_ENV = "UNSLOTH_LLAMA_ADMISSION_CONTROL"
+ADMISSION_QUEUE_TIMEOUT_ENV = "UNSLOTH_LLAMA_ADMISSION_QUEUE_TIMEOUT"
+ADMISSION_KEEPALIVE_INTERVAL_ENV = "UNSLOTH_LLAMA_ADMISSION_KEEPALIVE_INTERVAL"
+ADMISSION_MAX_QUEUE_ENV = "UNSLOTH_LLAMA_ADMISSION_MAX_QUEUE"
+ADMISSION_QUEUE_PER_SLOT_ENV = "UNSLOTH_LLAMA_ADMISSION_QUEUE_PER_SLOT"
+
+# The UNSLOTH_OPENAI_COMPAT_* spellings predate this queue being shared with the
+# Anthropic /v1/messages route (same llama-server slots). Still honored; the
+# neutral name above wins when both are set.
+_LEGACY_ENV = {
+ ADMISSION_CONTROL_ENV: "UNSLOTH_OPENAI_COMPAT_ADMISSION_CONTROL",
+ ADMISSION_QUEUE_TIMEOUT_ENV: "UNSLOTH_OPENAI_COMPAT_ADMISSION_QUEUE_TIMEOUT",
+ ADMISSION_KEEPALIVE_INTERVAL_ENV: "UNSLOTH_OPENAI_COMPAT_ADMISSION_KEEPALIVE_INTERVAL",
+ ADMISSION_MAX_QUEUE_ENV: "UNSLOTH_OPENAI_COMPAT_ADMISSION_MAX_QUEUE",
+}
DEFAULT_ADMISSION_ENABLED = True
+# None: a queued request waits for its slot indefinitely rather than timing out.
DEFAULT_ADMISSION_QUEUE_TIMEOUT_S = None
DEFAULT_ADMISSION_KEEPALIVE_INTERVAL_S = 5.0
-DEFAULT_ADMISSION_MAX_QUEUE = 64
+# None: no absolute cap, the wait line is sized from the pool instead.
+DEFAULT_ADMISSION_MAX_QUEUE = None
+# Wait line = 16 x the serving slots, so it tracks --parallel (4 slots -> 64
+# waiters, 8 -> 128). Purely a memory guard; waiting itself is never timed out.
+DEFAULT_ADMISSION_QUEUE_PER_SLOT = 16
+# Floor for the scaled line, so a 1-slot backend (plain `unsloth studio`, or any
+# load downshifted to fit VRAM) keeps the depth it had before scaling existed
+# rather than dropping to 16 and rejecting callers that used to queue.
+DEFAULT_ADMISSION_MIN_QUEUE = 64
-@dataclass(frozen = True)
+def _executor_workers() -> int:
+ """Threads asyncio's default executor runs to_thread work on.
+
+ Mirrors ThreadPoolExecutor's own default sizing, which is what
+ ``run_in_executor(None, ...)`` builds. 3.13 sizes it from
+ ``process_cpu_count()``, which honours CPU affinity and cgroup quotas;
+ ``cpu_count()`` would budget from the whole host inside a one-core container.
+ """
+ cpus = getattr(os, "process_cpu_count", os.cpu_count)() or 1
+ return min(32, cpus + 4)
+
+
+def _executor_reserve(workers: int) -> int:
+ """Threads kept clear of parked approvals, for generation steps, stream
+ teardown and unrelated to_thread work. Scaled rather than flat: a flat count
+ would leave a 5-worker executor (one usable CPU) no budget at all.
+ """
+ return max(2, workers // 8)
+
+
+def _max_parked(capacity: int) -> int:
+ """How many holders may sit on an approval prompt with their slot given back.
+
+ A pending prompt parks an executor thread (the loop blocks inside
+ to_thread(next, gen)) whether or not it parked its slot, the pool already
+ permits `capacity` of those, and every park admits one more, so budget only
+ what the executor has left over. Zero on a backend whose --parallel alone
+ fills it: the prompt then holds its slot, as it did before parking existed.
+ """
+ workers = _executor_workers()
+ spare = workers - _executor_reserve(workers) - max(0, capacity)
+ # A quarter of the executor, floored at two while `spare` allows: a quarter of
+ # five is one, and one park cannot cover the two simultaneous prompts #7455
+ # exists for.
+ return max(0, min(max(2, workers // 4), spare))
+
+
+# Process-wide, not per queue: there is one executor, and base_url takes a fresh
+# port on every load, so a per-queue budget would hand the same allowance to each
+# backend and to every reload, blind to the approvals parked on the old queue.
+_PARK_LOCK = threading.Lock()
+_parked_total = 0
+
+
+def _claim_park(limit: int) -> bool:
+ global _parked_total
+ with _PARK_LOCK:
+ if _parked_total >= limit:
+ return False
+ _parked_total += 1
+ return True
+
+
+def _drop_park() -> None:
+ global _parked_total
+ with _PARK_LOCK:
+ _parked_total = max(0, _parked_total - 1)
+
+
+def _live_capacity(current: "LlamaAdmissionQueue") -> int:
+ """Slots across every backend still serving requests.
+
+ One queue's capacity is the wrong denominator for a budget sized against the
+ one executor: a reload drains the old queue alongside the new one, and
+ prompts on both park threads. Idle queues hold nothing and are about to be
+ evicted.
+ """
+ with _QUEUES_LOCK:
+ queues = list(_QUEUES.values())
+ # is_idle takes each queue's own lock, so never while holding _QUEUES_LOCK.
+ total = sum(queue._capacity for queue in queues if queue is current or not queue.is_idle())
+ return total if any(queue is current for queue in queues) else total + current._capacity
+
+
+@dataclass(frozen = True, **_SLOTS)
class LlamaAdmissionConfig:
enabled: bool = DEFAULT_ADMISSION_ENABLED
queue_timeout_s: Optional[float] = DEFAULT_ADMISSION_QUEUE_TIMEOUT_S
keepalive_interval_s: float = DEFAULT_ADMISSION_KEEPALIVE_INTERVAL_S
max_queue: Optional[int] = DEFAULT_ADMISSION_MAX_QUEUE
+ queue_per_slot: Optional[int] = DEFAULT_ADMISSION_QUEUE_PER_SLOT
+ # Unconditional floor on the scaled line. The env path clears it when the
+ # operator sets QUEUE_PER_SLOT, so only the default multiplier is floored.
+ min_queue: Optional[int] = DEFAULT_ADMISSION_MIN_QUEUE
+
+ def queue_limit(self, capacity: int) -> Optional[int]:
+ """How many callers may line up for a pool of ``capacity`` slots.
+
+ An explicit ``max_queue`` wins; otherwise the line scales with the slots
+ so it follows ``--parallel``. The default multiplier is floored, so a
+ 1-slot backend does not end up shallower than it was before scaling. None
+ (or any non-positive setting) means an unbounded line.
+ """
+ if self.max_queue is not None:
+ return self.max_queue if self.max_queue > 0 else None
+ if not self.queue_per_slot or self.queue_per_slot <= 0:
+ return None
+ scaled = self.queue_per_slot * max(1, capacity)
+ return max(self.min_queue, scaled) if self.min_queue else scaled
-@dataclass(frozen = True)
+@dataclass(frozen = True, **_SLOTS)
class LlamaAdmissionSnapshot:
key: str
capacity: int
active: int
queued: int
+ free: int = 0
class LlamaAdmissionError(Exception):
@@ -69,8 +191,17 @@ class LlamaAdmissionCancelled(LlamaAdmissionError):
pass
-def _bool_env(name: str, default: bool) -> bool:
+def _raw_env(name: str) -> Optional[str]:
+ """Value for a canonical name, falling back to its legacy spelling."""
value = os.environ.get(name)
+ if value is None or not value.strip():
+ legacy = _LEGACY_ENV.get(name)
+ value = os.environ.get(legacy) if legacy else None
+ return value
+
+
+def _bool_env(name: str, default: bool) -> bool:
+ value = _raw_env(name)
if value is None or not value.strip():
return default
value = value.strip().lower()
@@ -82,7 +213,7 @@ def _bool_env(name: str, default: bool) -> bool:
def _optional_positive_float_env(name: str, default: Optional[float]) -> Optional[float]:
- value = os.environ.get(name)
+ value = _raw_env(name)
if value is None or not value.strip():
return default
try:
@@ -93,7 +224,7 @@ def _optional_positive_float_env(name: str, default: Optional[float]) -> Optiona
def _positive_float_env(name: str, default: float) -> float:
- value = os.environ.get(name)
+ value = _raw_env(name)
if value is None or not value.strip():
return default
try:
@@ -103,19 +234,38 @@ def _positive_float_env(name: str, default: float) -> float:
return parsed if parsed > 0 else default
-def _optional_positive_int_env(name: str, default: Optional[int]) -> Optional[int]:
- value = os.environ.get(name)
- if value is None or not value.strip():
- return default
+def _queue_limits_from_env() -> tuple[Optional[int], Optional[int], Optional[int]]:
+ """(max_queue, queue_per_slot, min_queue) from the environment.
+
+ An absolute MAX_QUEUE wins outright; MAX_QUEUE=0 asks for an unbounded line.
+ Unset leaves the per-slot multiplier in charge (itself 0 for unbounded). The
+ floor applies only to the default multiplier: setting QUEUE_PER_SLOT means
+ the operator wants that exact depth, however shallow.
+ """
+ # Explicit means it parsed, not just that something was set: a typo falls back
+ # to the default multiplier, so it has to keep the default's floor too.
+ raw_per_slot = _raw_env(ADMISSION_QUEUE_PER_SLOT_ENV)
try:
- parsed = int(value.strip())
+ per_slot = int((raw_per_slot or "").strip())
except ValueError:
- return default
- return parsed if parsed > 0 else None
+ per_slot, min_queue = DEFAULT_ADMISSION_QUEUE_PER_SLOT, DEFAULT_ADMISSION_MIN_QUEUE
+ else:
+ per_slot, min_queue = (per_slot if per_slot > 0 else None), None
+ raw = _raw_env(ADMISSION_MAX_QUEUE_ENV)
+ if raw is None or not raw.strip():
+ return None, per_slot, min_queue
+ try:
+ parsed = int(raw.strip())
+ except ValueError:
+ return None, per_slot, min_queue
+ return (parsed, None, None) if parsed > 0 else (None, None, None)
def llama_admission_config_from_env() -> LlamaAdmissionConfig:
+ max_queue, queue_per_slot, min_queue = _queue_limits_from_env()
return LlamaAdmissionConfig(
+ queue_per_slot = queue_per_slot,
+ min_queue = min_queue,
enabled = _bool_env(ADMISSION_CONTROL_ENV, DEFAULT_ADMISSION_ENABLED),
queue_timeout_s = _optional_positive_float_env(
ADMISSION_QUEUE_TIMEOUT_ENV,
@@ -125,14 +275,11 @@ def llama_admission_config_from_env() -> LlamaAdmissionConfig:
ADMISSION_KEEPALIVE_INTERVAL_ENV,
DEFAULT_ADMISSION_KEEPALIVE_INTERVAL_S,
),
- max_queue = _optional_positive_int_env(
- ADMISSION_MAX_QUEUE_ENV,
- DEFAULT_ADMISSION_MAX_QUEUE,
- ),
+ max_queue = max_queue,
)
-@dataclass
+@dataclass(**_SLOTS)
class _Waiter:
loop: asyncio.AbstractEventLoop
future: asyncio.Future
@@ -141,20 +288,130 @@ class _Waiter:
class LlamaAdmissionLease:
- def __init__(self, queue: Optional["LlamaAdmissionQueue"]):
+ __slots__ = ("_queue", "_slot", "_released", "_release_lock", "_parked", "_budgeted")
+
+ def __init__(
+ self,
+ queue: Optional["LlamaAdmissionQueue"],
+ slot: Optional[int] = None,
+ ):
self._queue = queue
+ self._slot = slot
self._released = False
self._release_lock = threading.Lock()
+ self._parked = False
+ self._budgeted = False
+
+ @property
+ def slot(self) -> Optional[int]:
+ """Pool slot this lease holds, or None when admission is disabled."""
+ return self._slot
+
+ def park(self) -> bool:
+ """Hand the slot back while this holder waits on something off the GPU.
+
+ A run stopped on a tool approval prompt is not decoding, so holding its
+ slot would let unanswered prompts fill the pool while llama-server idles.
+ The lease itself stays valid: releasing it after a park is still correct.
+
+ False when the park budget is spent and nothing was given back: the
+ caller keeps its slot across the prompt, as it did before parking
+ existed. Slower for whoever is behind it, but each freed slot admits
+ another run that can park too, on the executor the generators run on.
+ """
+ queue = self._queue
+ with self._release_lock:
+ if queue is None or self._released or self._parked:
+ return False
+ # Under the lease lock so the decision and the handover cannot split.
+ # Nothing takes the queue lock then a lease lock, so this order is
+ # the only one in play.
+ if not queue.try_park(self._slot):
+ return False
+ self._parked = True
+ self._budgeted = True
+ self._slot = None
+ return True
+
+ def _drop_budget(self) -> None:
+ """Give the executor budget back now the prompt wait is over.
+
+ Separate from the queue's parked count, which lasts until the slot is
+ back: the executor thread is free the moment the answer arrives. Holding
+ the budget until the resume lands would refuse someone else's park for a
+ finished wait, and that someone holds the slot the resumer wants.
+ """
+ with self._release_lock:
+ if not self._budgeted:
+ return
+ self._budgeted = False
+ _drop_park()
+
+ def unpark(self) -> None:
+ """Drop the parked state without reclaiming a slot.
+
+ For a holder that is tearing down: it will not decode again. Resuming
+ holders must use ``unpark_async``, which waits for a slot instead of
+ going back to llama-server past the admission limit.
+ """
+ with self._release_lock:
+ if not self._parked:
+ return
+ self._parked = False
+ self._drop_budget()
+ if self._queue is not None:
+ self._queue.unpark()
+
+ async def unpark_async(
+ self,
+ *,
+ cancel_event = None,
+ poll_s: float = 0.02,
+ ) -> None:
+ """Take a slot back, waiting until the pool has room.
+
+ ``park`` gave the slot to a waiter, so by the time the user answers the
+ prompt someone else may be decoding in it. Resuming regardless put two
+ holders on a one-slot server. Gives up if the caller is cancelled, since
+ the holder is then leaving anyway and must not be stuck here.
+ """
+ queue = self._queue
+ if queue is None or not self._parked:
+ return
+ # Before the wait, not after: the prompt is answered, so this holder is
+ # already off the executor and must not keep anyone else off it.
+ self._drop_budget()
+ slot = await queue.acquire_parked_slot(cancel_event = cancel_event, poll_s = poll_s)
+ stranded = None
+ with self._release_lock:
+ # release() may have run during the wait; it clears the flag and does
+ # the unpark itself, so only the caller that clears it here repeats one.
+ parked, self._parked = self._parked, False
+ if self._released:
+ # Released while waiting: this lease will never hand the slot
+ # back, so return it here rather than strand it for good.
+ stranded = slot
+ else:
+ self._slot = slot
+ if parked:
+ queue.unpark()
+ if stranded is not None:
+ queue.release(stranded)
def release(self) -> None:
queue = None
+ parked = False
with self._release_lock:
if self._released:
return
self._released = True
queue = self._queue
+ parked, self._parked = self._parked, False
+ self._drop_budget()
if queue is not None:
- queue.release()
+ if parked:
+ queue.unpark()
+ queue.release(self._slot)
async def __aenter__(self) -> "LlamaAdmissionLease":
return self
@@ -164,6 +421,8 @@ class LlamaAdmissionLease:
class LlamaAdmissionReservation:
+ __slots__ = ("_queue", "_lease", "_waiter", "snapshot")
+
def __init__(
self,
*,
@@ -195,6 +454,13 @@ class LlamaAdmissionReservation:
return self._lease
async def wait(self, timeout_s: float) -> Optional[LlamaAdmissionLease]:
+ """Wait up to ``timeout_s`` for a slot.
+
+ A timeout leaves this reservation queued so the caller can poll again.
+ Any exit that abandons the wait for good must call ``cancel()``, or the
+ slot granted later is delivered to a future nobody reads and is never
+ released.
+ """
lease = self.lease_nowait()
if lease is not None:
return lease
@@ -229,12 +495,74 @@ class LlamaAdmissionReservation:
class LlamaAdmissionQueue:
+ """A fixed pool of generation slots for one llama-server, plus a FIFO wait line.
+
+ The pool mirrors llama-server's own ``--parallel`` slots: ``capacity`` slot ids
+ are each either free or held by exactly one caller. A caller that finds every
+ slot busy waits in arrival order and is handed the next slot to free, so no
+ caller is starved. This bounds only the callers that reserve: chat completions
+ and messages do, while /v1/completions, Studio's own chat endpoint and RAG
+ captioning all reach llama-server directly, so it is not a global cap.
+ Waiting is unbounded in time by default (``queue_timeout_s``
+ None); the wait line itself is bounded, and only how many may line up before
+ new arrivals are rejected. By default that is ``16 x slots`` floored at 64,
+ not unlimited: an unbounded line takes ``max_queue`` or ``queue_per_slot``
+ set to 0. See ``LlamaAdmissionConfig.queue_limit``.
+ """
+
+ __slots__ = (
+ "key",
+ "_lock",
+ "_capacity",
+ "_free",
+ "_in_use",
+ "_held",
+ "_waiters",
+ "_parked",
+ "_unpark_tickets",
+ "_unpark_seq",
+ )
+
def __init__(self, key: str):
self.key = key
self._lock = threading.Lock()
- self._active = 0
self._capacity = 1
+ self._free: list[int] = [0]
+ # Held slots as a bitmask: one int instead of a set, so the pool costs the
+ # same whether it is idle or saturated. _held is its popcount, kept as a
+ # counter because int.bit_count() is 3.10+ and this package targets 3.9.
+ self._in_use = 0
+ self._held = 0
self._waiters: Deque[_Waiter] = deque()
+ # Holders parked on a tool approval prompt. They hold no slot, so this only
+ # keeps the queue off the idle-eviction list while they are away.
+ self._parked = 0
+ # FIFO tickets for holders resuming from a park (see acquire_parked_slot). A
+ # bare count deadlocked: every approved holder blocked every other one.
+ self._unpark_tickets: Deque[int] = deque()
+ self._unpark_seq = 0
+
+ def _resize_pool_locked(self, capacity: int) -> None:
+ # Slots past a shrunk capacity retire when their holder releases them.
+ if capacity == self._capacity:
+ return
+ self._capacity = capacity
+ self._free = [slot for slot in range(capacity) if not self._in_use >> slot & 1]
+
+ def _can_admit_locked(self, reserved: int) -> bool:
+ # Slots still held above a shrunk capacity keep occupying the backend, so
+ # count every held slot against the ceiling, not just the ids below it.
+ # ``reserved`` holds slots back for approved holders waiting to resume;
+ # without it a stream of new arrivals took the next slot, forever.
+ return bool(self._free) and (self._held + reserved) < self._capacity
+
+ def _take_slot_locked(self, reserved: int) -> Optional[int]:
+ if not self._can_admit_locked(reserved):
+ return None
+ slot = self._free.pop()
+ self._in_use |= 1 << slot
+ self._held += 1
+ return slot
def reserve(self, *, capacity: int, config: LlamaAdmissionConfig) -> LlamaAdmissionReservation:
capacity = max(1, int(capacity or 1))
@@ -242,22 +570,25 @@ class LlamaAdmissionQueue:
return LlamaAdmissionReservation(
queue = None,
lease = LlamaAdmissionLease(None),
- snapshot = LlamaAdmissionSnapshot(self.key, capacity, 0, 0),
+ snapshot = LlamaAdmissionSnapshot(self.key, capacity, 0, 0, capacity),
)
loop = asyncio.get_running_loop()
with self._lock:
- self._capacity = capacity
- self._prune_waiters_locked()
+ self._resize_pool_locked(capacity)
self._grant_waiters_locked()
- if self._active < self._capacity and not self._waiters:
- self._active += 1
- return LlamaAdmissionReservation(
- queue = self,
- lease = LlamaAdmissionLease(self),
- snapshot = self._snapshot_locked(),
- )
- if config.max_queue is not None and len(self._waiters) >= config.max_queue:
+ if not self._waiters:
+ slot = self._take_slot_locked(len(self._unpark_tickets))
+ if slot is not None:
+ # No snapshot here: callers read it through snapshot_now(),
+ # which re-reads the queue, so building one per admitted
+ # request would be pure allocation on the hot path.
+ return LlamaAdmissionReservation(
+ queue = self,
+ lease = LlamaAdmissionLease(self, slot),
+ )
+ limit = config.queue_limit(self._capacity)
+ if limit is not None and self._live_waiters_locked() >= limit:
raise LlamaAdmissionQueueFull(
"llama-server generation queue is full",
snapshot = self._snapshot_locked(),
@@ -270,15 +601,82 @@ class LlamaAdmissionQueue:
return LlamaAdmissionReservation(
queue = self,
waiter = waiter,
- snapshot = self._snapshot_locked(),
)
- def release(self) -> None:
+ def _release_slot_locked(self, slot: Optional[int]) -> None:
+ # A slot id at or past a shrunk capacity retires instead of returning.
+ if slot is None or not self._in_use >> slot & 1:
+ return
+ self._in_use &= ~(1 << slot)
+ self._held -= 1
+ if slot < self._capacity:
+ self._free.append(slot)
+
+ def release(self, slot: Optional[int]) -> None:
with self._lock:
- if self._active > 0:
- self._active -= 1
+ self._release_slot_locked(slot)
self._grant_waiters_locked()
+ def try_park(self, slot: Optional[int]) -> bool:
+ """Return a parked holder's slot to the pool. See ``LlamaAdmissionLease.park``.
+
+ False leaves the slot with its holder, so a refused park costs nothing to
+ undo. The per-queue count is only what ``is_idle`` reads; the budget and
+ the capacity it is sized from are both process-wide.
+ """
+ if not _claim_park(_max_parked(_live_capacity(self))):
+ return False
+ with self._lock:
+ self._parked += 1
+ self._release_slot_locked(slot)
+ self._grant_waiters_locked()
+ return True
+
+ def unpark(self) -> None:
+ with self._lock:
+ if self._parked > 0:
+ self._parked -= 1
+
+ async def acquire_parked_slot(
+ self,
+ *,
+ cancel_event = None,
+ poll_s: float = 0.02,
+ ) -> Optional[int]:
+ """Wait for a slot for a holder resuming from a park, None if cancelled.
+
+ Ordered by ticket rather than counted, so approvals resume in the order
+ they came back: counting them made every approved holder block every
+ other one, and with nothing decoding that never resolved.
+ """
+ with self._lock:
+ self._unpark_seq += 1
+ ticket = self._unpark_seq
+ self._unpark_tickets.append(ticket)
+ try:
+ while True:
+ with self._lock:
+ ahead = 0
+ for queued in self._unpark_tickets:
+ if queued == ticket:
+ break
+ ahead += 1
+ # Only the approvals ahead of this one hold slots back from it.
+ slot = self._take_slot_locked(ahead)
+ if slot is not None:
+ return slot
+ if cancel_event is not None and cancel_event.is_set():
+ return None
+ await asyncio.sleep(poll_s)
+ finally:
+ with self._lock:
+ try:
+ self._unpark_tickets.remove(ticket)
+ except ValueError:
+ pass
+ # This ticket was holding a slot back from the wait line.
+ self._grant_waiters_locked()
+
def cancel(self, waiter: _Waiter) -> None:
lease_to_release = None
with self._lock:
@@ -291,7 +689,13 @@ class LlamaAdmissionQueue:
lease_to_release = waiter.granted_lease
waiter.granted_lease = None
if not waiter.future.done():
- waiter.loop.call_soon_threadsafe(waiter.future.cancel)
+ try:
+ waiter.loop.call_soon_threadsafe(waiter.future.cancel)
+ except RuntimeError:
+ # Loop gone. Routes call cancel() from finally blocks, so
+ # raising here would both mask their exception and skip the
+ # release below, stranding the slot for the process lifetime.
+ pass
if lease_to_release is not None:
lease_to_release.release()
@@ -303,20 +707,32 @@ class LlamaAdmissionQueue:
def is_idle(self) -> bool:
with self._lock:
self._prune_waiters_locked()
- return self._active == 0 and not self._waiters
+ # A parked holder owns no slot but is coming back to this queue, so
+ # evicting it here would resume it against a fresh 1-slot pool.
+ return self._in_use == 0 and not self._waiters and not self._parked
def _grant_waiters_locked(self) -> None:
- self._prune_waiters_locked()
- while self._waiters and self._active < self._capacity:
+ # Dead waiters are skipped as they are popped, so no prune is needed here.
+ while self._waiters and self._can_admit_locked(len(self._unpark_tickets)):
waiter = self._waiters.popleft()
if waiter.cancelled or waiter.future.done():
continue
- self._active += 1
- lease = LlamaAdmissionLease(self)
+ slot = self._take_slot_locked(len(self._unpark_tickets))
+ lease = LlamaAdmissionLease(self, slot)
waiter.granted_lease = lease
- waiter.loop.call_soon_threadsafe(self._deliver_lease, waiter, lease)
+ try:
+ waiter.loop.call_soon_threadsafe(self._deliver_lease, waiter, lease)
+ except RuntimeError:
+ # Waiter's loop is gone. Reclaim the slot; leaving the bit set
+ # would strand it, since _free is rebuilt from the bitmask.
+ waiter.granted_lease = None
+ self._release_slot_locked(slot)
def _deliver_lease(self, waiter: _Waiter, lease: LlamaAdmissionLease) -> None:
+ # Runs on the waiter's own loop thread, which is also the only thread that
+ # cancels that reservation, so waiter state is safe to touch unlocked here.
+ # release() may be called from any thread, but only reaches this via
+ # call_soon_threadsafe. Cancelling off-loop would need this under _lock.
if waiter.cancelled or waiter.future.done():
waiter.granted_lease = None
if not waiter.future.done():
@@ -331,16 +747,32 @@ class LlamaAdmissionQueue:
lease.release()
def _prune_waiters_locked(self) -> None:
+ # Rebuilding the deque on every reserve/release dominated the hot path, so
+ # only pay it when a waiter actually died out of band (an externally
+ # cancelled future); cancel() already drops its own waiter eagerly.
+ for waiter in self._waiters:
+ if waiter.cancelled or waiter.future.done():
+ break
+ else:
+ return
self._waiters = deque(
waiter for waiter in self._waiters if not waiter.cancelled and not waiter.future.done()
)
+ def _live_waiters_locked(self) -> int:
+ self._prune_waiters_locked()
+ return len(self._waiters)
+
def _snapshot_locked(self) -> LlamaAdmissionSnapshot:
return LlamaAdmissionSnapshot(
key = self.key,
capacity = self._capacity,
- active = self._active,
+ active = self._held,
queued = len(self._waiters),
+ # What another caller could actually take, so the admission log never
+ # shows free slots next to queued requests: after a shrink, ids below
+ # the new capacity can be free while holdovers still fill the ceiling.
+ free = min(len(self._free), max(0, self._capacity - self._held)),
)
@@ -364,5 +796,10 @@ def get_llama_admission_queue(key: str) -> LlamaAdmissionQueue:
def reset_llama_admission_queues() -> None:
+ global _parked_total
with _QUEUES_LOCK:
_QUEUES.clear()
+ # The budget outlives the queues it was claimed against, so dropping them
+ # without it leaks the count and shrinks the budget for good.
+ with _PARK_LOCK:
+ _parked_total = 0
diff --git a/studio/backend/core/inference/llama_cpp.py b/studio/backend/core/inference/llama_cpp.py
index 2c7433f7a4..712caf43e5 100644
--- a/studio/backend/core/inference/llama_cpp.py
+++ b/studio/backend/core/inference/llama_cpp.py
@@ -33,6 +33,7 @@ from typing import (
List,
Literal,
Mapping,
+ MutableMapping,
Optional,
Union,
)
@@ -42,6 +43,7 @@ import httpx
from core.inference.llama_server_args import (
_LAYER_OFFLOAD_FLAGS,
_effective_tensor_parallel,
+ _flag_name,
_tensor_parallel_matches_loaded,
extra_args_disable_mmproj,
parse_cache_override,
@@ -83,6 +85,7 @@ from core.tool_healing import (
strip_outside_think,
)
from utils.native_path_leases import child_env_without_native_path_secret
+from utils.child_stdio import utf8_child_env
from utils.hf_xet_fallback import hf_hub_download_with_xet_fallback
from utils.subprocess_compat import (
windows_hidden_subprocess_kwargs as _windows_hidden_subprocess_kwargs,
@@ -90,13 +93,17 @@ from utils.subprocess_compat import (
from utils.process_lifetime import child_popen_kwargs as _child_popen_kwargs
from core.inference.tool_call_parser import (
MAX_ACT_REPROMPTS as _MAX_REPROMPTS,
+ NUDGE_TOOL_CALLS_STATUS as _NUDGE_TOOL_CALLS_STATUS,
REPROMPT_MAX_CHARS as _REPROMPT_MAX_CHARS,
+ is_reprompt_repeat as _is_reprompt_repeat,
+ is_reprompt_restatement as _is_reprompt_restatement,
is_short_intent_without_action as _is_short_intent_without_action,
reprompt_to_act_message as _reprompt_to_act_message,
)
from core.inference.tool_loop_controller import (
ToolLoopController,
append_deferred_nudges,
+ awaiting_approval_status,
tool_event_provenance,
)
from state.tool_approvals import (
@@ -126,6 +133,15 @@ LLAMA_SERVER_NOT_FOUND_DETAIL = (
"then try again. (Advanced: set LLAMA_SERVER_PATH to an existing binary.)"
)
+# Shared by the route, pre-teardown and post-metadata rejections (#7205).
+_VULKAN_DIFFUSION_GPU_IDS_ERROR = (
+ "GPU selection (gpu_ids) is not supported for a DiffusionGemma "
+ "GGUF on a Vulkan llama.cpp build: the diffusion runner selects "
+ "its device by CUDA physical index, which has no defined mapping "
+ "to ggml Vulkan device ordinals. Omit gpu_ids to use the default "
+ "device."
+)
+
# llama-server can serve HTTP 200 while running a model entirely on CPU when a
# GPU backend fails to init (#5807 / #5106 / #5830). Classify the startup log so
@@ -236,7 +252,7 @@ def _wsl_system_rocm_lib_dirs() -> "list[str]":
with open("/proc/version", encoding = "utf-8", errors = "replace") as fh:
if "microsoft" not in fh.read().lower():
return []
- except OSError:
+ except (OSError, UnicodeDecodeError):
return []
out: "list[str]" = []
for d in ("/opt/rocm/lib", "/opt/rocm/lib64"):
@@ -247,12 +263,97 @@ def _wsl_system_rocm_lib_dirs() -> "list[str]":
return out
+def _bundled_hip_present(binary_dir: str) -> bool:
+ """True when a prebuilt bundle ships its own HIP backend library."""
+ if not binary_dir:
+ return False
+ try:
+ # Glob the version suffix (libggml-hip.so, .so.0, .so.0.11.1) the same
+ # way the installer's runtime health check matches libggml-hip.so*.
+ return any(Path(str(binary_dir)).glob("libggml-hip.so*"))
+ except OSError:
+ return False
+
+
+def _native_linux_system_rocm_lib_dirs(binary_dir: str = "") -> "list[str]":
+ """System ROCm lib dir(s) to prepend before a prebuilt's bundled HIP, on native Linux.
+
+ The bundled bare-metal HIP runtime can mismatch the host amdkfd driver and crash
+ in hsa_init(); prepending the whole system ROCm lib dir loads a driver-matched,
+ version-consistent stack (libhsa-runtime64 / libamdhip64 / librocblas) ahead of it.
+ The whole dir is deliberate: mixing the bundle's rocBLAS with a different-version
+ system HIP/ROCR risks missing symbols. UNSLOTH_LLAMA_NO_SYSTEM_ROCM=1 keeps the pure
+ bundle (for a host whose system ROCm lacks this arch); no-op on WSL / non-Linux.
+ """
+ if os.environ.get("UNSLOTH_LLAMA_NO_SYSTEM_ROCM") == "1":
+ return []
+ if sys.platform != "linux" or os.path.exists("/dev/dxg"):
+ return []
+ if not os.path.exists("/dev/kfd"):
+ return []
+ if not _bundled_hip_present(binary_dir):
+ return []
+ # Env-configured ROCm root first; /opt/rocm only as a fallback so a stale
+ # /opt/rocm doesn't shadow the driver-matching install these vars point at.
+ candidates = []
+ for var in ("HIP_PATH", "HIP_PATH_57", "ROCM_PATH"):
+ val = os.environ.get(var)
+ if val:
+ candidates.append(val)
+ candidates.append("/opt/rocm")
+ out: "list[str]" = []
+ seen: "set[str]" = set()
+ for base in candidates:
+ for lib_sub in ("lib", "lib64"):
+ d = os.path.join(base, lib_sub)
+ if d in seen:
+ continue
+ seen.add(d)
+ if os.path.exists(os.path.join(d, "libhsa-runtime64.so")) or os.path.exists(
+ os.path.join(d, "libhsa-runtime64.so.1")
+ ):
+ out.append(d)
+ # ROCm keeps LLVM's versioned runtime under /lib/llvm, so a
+ # lib64 host still finds it under lib. Probe both and keep them
+ # ahead of the bundle, else system libamd_comgr binds to the
+ # bundle's incompatible libLLVM.so.*.
+ for _sub in (lib_sub, "lib"):
+ llvm_lib = os.path.join(base, _sub, "llvm", "lib")
+ if llvm_lib not in seen and os.path.isdir(llvm_lib):
+ seen.add(llvm_lib)
+ out.append(llvm_lib)
+ return out
+
+
# Plan-without-action re-prompt state now lives in tool_call_parser (imported above).
# Default max_tokens to the effective context when known. The floor is high
# enough for reasoning-heavy GGUFs and max_tokens-omitting API clients.
_DEFAULT_MAX_TOKENS_FLOOR = 32768
_DEFAULT_FIRST_TOKEN_TIMEOUT_S = 1200.0 # 20 min
+# A transport error can arrive before the child is reapable; a request path cannot
+# afford the 5s the background MTP reload spends on the same race.
+_RESPAWN_REAP_GRACE_S = 1.0
+
+
+def _finalize_reasoning_only_cumulative(
+ cumulative: str, reasoning_text: str, finish_reason: Optional[str], promote_reasoning_only: bool
+) -> str:
+ """Close a live thinking block and promote it only after a clean stop.
+
+ Local inference streams cumulative snapshots. Replacing ``...`` with
+ bare reasoning at EOF makes the final snapshot shorter, so suffix-based
+ route consumers drop the intended fallback. Keep the snapshot append-only.
+ A length-truncated thought is not a final answer, so close it without
+ promotion and let the client surface the ``length`` terminal state. Raw
+ consumers that do not split reasoning from visible content can disable the
+ fallback to avoid returning the same reasoning twice.
+ """
+ visible_fallback = (
+ reasoning_text if promote_reasoning_only and finish_reason != "length" else ""
+ )
+ return cumulative + "" + visible_fallback
+
# Only large streamed tool payloads get an early provisional card; render_html
# is exempt because it needs immediate artifact feedback.
@@ -262,12 +363,32 @@ _DEFAULT_STREAM_STALL_TIMEOUT_S = 120.0 # 2 min
# loop). Structured delta.tool_calls are grammar-bounded by llama-server; text
# parsed from content is not, so one runaway turn could fan out unbounded.
_MAX_TOOL_CALLS_PER_TURN = 8
-_FORCED_REPEAT_PLAN_SIGNAL = re.compile(
- r"\b(?:i\s+will|i'll|let\s+me|going\s+to|need\s+to|call|use|run|search|fetch|render)\b",
+# Obligation phrasing INTENT_SIGNAL leaves alone ("I need to call ..."), paired with
+# an action verb. Sentence-anchored: mid-sentence the same words are prose that names
+# a tool ("The API I should invoke is foo() because ..."), and suppressing that loses
+# a real answer. "should"/"must" sit outside the need|have|ought group because they
+# take a bare infinitive. "invoke"/"query" stay out of the verb list: they read as
+# technical prose far more often than as a stall.
+_FORCED_PLAN_INTENT = re.compile(
+ r"(?:^|[.!?]\s+)\s*"
+ r"(?:i\s+(?:(?:need|have|ought)\s+to|should|must)|need\s+to|going\s+to|must|should)"
+ r"\s+(?:\w+\s+){0,2}?(?:call|use|run|search|fetch|render)\b",
+ re.I | re.M,
+)
+# "the answer is not in the context" announces a *missing* answer, so the negated
+# forms are excluded or the plan behind them would ship as the final response.
+_FINAL_ANSWER_SIGNAL = re.compile(
+ r"\b(?:final\s+answer|answer\s*:|here\s+is|here's|in\s+summary|result\s*:"
+ r"|(?:the\s+)?answer\s+is(?!\s+(?:not|unavailable|unknown|unclear|missing)\b))\b",
re.I,
)
-_FINAL_ANSWER_SIGNAL = re.compile(
- r"\b(?:final\s+answer|answer\s*:|here\s+is|here's|in\s+summary|result\s*:)\b",
+# A plan that pivots ("I should call web_search, but Tokyo is the capital") has an
+# answer attached, so the turn must survive. Leaking a plan sentence is cosmetic;
+# dropping an answer is not, so the doubtful case keeps the output. The pivot has to
+# carry text of its own: "I should call web_search, though." answers nothing.
+_ANSWER_PIVOT = re.compile(
+ r"\b(?:but|however|although|though|that\s+said|in\s+the\s+meantime|meanwhile)\b"
+ r"[\W_]*(?:\w+[\W_]+){1,}\w",
re.I,
)
@@ -359,14 +480,28 @@ def _held_rehearsal_tail_len(text: str, active_tools: list[dict]) -> int:
return len(tail) if tail and _is_rehearsal_prefix(tail, active_tools) else 0
-def _should_suppress_forced_no_tool_output(text: str) -> bool:
- """Suppress only repeated forced-turn planning text, not final answers."""
+def _should_suppress_forced_no_tool_output(text: str, previous: str = "") -> bool:
+ """Suppress only repeated forced-turn planning text, not final answers.
+
+ ``previous`` is the stall text that triggered the nudge, so a retry that
+ moved on can be told from one that just said the same thing again.
+ """
stripped = text.strip()
if not stripped or len(stripped) >= _REPROMPT_MAX_CHARS:
return False
if _FINAL_ANSWER_SIGNAL.search(stripped):
return False
- return _FORCED_REPEAT_PLAN_SIGNAL.search(stripped) is not None
+ plan = _FORCED_PLAN_INTENT.search(stripped)
+ if plan is not None:
+ # Only the plan itself is safe to drop; anything the turn pivots to after it
+ # is the answer the user is waiting for.
+ return _ANSWER_PIVOT.search(stripped[plan.end() :]) is None
+ if not _is_short_intent_without_action(stripped):
+ return False
+ # INTENT_SIGNAL also fires on lead-ins to a real answer ("Now I have the results.
+ # The capital is Tokyo."), so a bare intent match is a stall only when the retry
+ # adds nothing. No ``previous`` keeps the standalone "is this a stall?" contract.
+ return not previous or _is_reprompt_restatement(stripped, previous)
# ── Pre-compiled patterns for GGUF shard detection ───────────
@@ -483,11 +618,11 @@ def _load_swa_cache() -> dict:
if _SWA_CACHE is not None:
return _SWA_CACHE
try:
- with open(_swa_cache_path()) as f:
+ with open(_swa_cache_path(), encoding = "utf-8-sig") as f:
_SWA_CACHE = json.load(f)
if not isinstance(_SWA_CACHE, dict):
_SWA_CACHE = {}
- except (FileNotFoundError, json.JSONDecodeError, OSError):
+ except (FileNotFoundError, json.JSONDecodeError, OSError, UnicodeDecodeError):
_SWA_CACHE = {}
return _SWA_CACHE
@@ -497,10 +632,10 @@ def _save_swa_cache(cache: dict) -> None:
path = _swa_cache_path()
path.parent.mkdir(parents = True, exist_ok = True)
tmp = path.with_suffix(".json.tmp")
- with open(tmp, "w") as f:
+ with open(tmp, "w", encoding = "utf-8") as f:
json.dump(cache, f, indent = 2, sort_keys = True)
tmp.replace(path)
- except OSError:
+ except (OSError, UnicodeDecodeError):
pass
@@ -526,8 +661,15 @@ def _swa_entry_from_layer_types(lt) -> Optional[object]:
def _fetch_swa_entry_from_hf(repo_id: str) -> Optional[object]:
try:
from huggingface_hub import hf_hub_download
- cfg_path = hf_hub_download(repo_id, "config.json", repo_type = "model")
- with open(cfg_path) as f:
+ from utils.hf_cache_settings import active_hf_hub_cache
+
+ cfg_path = hf_hub_download(
+ repo_id,
+ "config.json",
+ repo_type = "model",
+ cache_dir = active_hf_hub_cache(),
+ )
+ with open(cfg_path, encoding = "utf-8-sig") as f:
cfg = json.load(f)
except Exception:
return None
@@ -928,6 +1070,7 @@ def _cached_hf_snapshot_file(
filename: str,
*,
expected_size: Optional[int] = None,
+ cache_dir: Optional[str] = None,
) -> Optional[str]:
"""Return a cached snapshot file even when HF's current-ref probe misses it."""
if not filename:
@@ -936,8 +1079,22 @@ def _cached_hf_snapshot_file(
if not parts or any(part in (".", "..") for part in parts):
return None
try:
- from utils.models.model_config import _iter_hf_cache_snapshots
- for snap in _iter_hf_cache_snapshots(repo_id):
+ if cache_dir is None:
+ from utils.models.model_config import _iter_hf_cache_snapshots
+ snapshots = _iter_hf_cache_snapshots(repo_id)
+ else:
+ from hub.utils.hf_cache_state import iter_active_repo_cache_dirs
+ snapshots = (
+ snapshot
+ for repo_dir in iter_active_repo_cache_dirs(
+ "model",
+ repo_id,
+ root = Path(cache_dir),
+ )
+ for snapshot in (repo_dir / "snapshots").glob("*")
+ if snapshot.is_dir()
+ )
+ for snap in snapshots:
candidate = snap.joinpath(*parts)
if not candidate.is_file():
continue
@@ -1179,6 +1336,16 @@ def _snapshot_dir_of(path: str) -> Optional[Path]:
return None
+def _hub_cache_dir_for_snapshot_path(path: Optional[str]) -> Optional[str]:
+ """Return the HF Hub cache root that owns a snapshot-contained path."""
+ if not path:
+ return None
+ snapshot = _snapshot_dir_of(path)
+ if snapshot is None or snapshot.parent.name != "snapshots":
+ return None
+ return str(snapshot.parent.parent.parent)
+
+
def _companion_snapshot_sibling(
near_path: str, pick: Callable[[list[str]], Optional[str]]
) -> Optional[str]:
@@ -1390,6 +1557,21 @@ def _kv_bytes_per_elem(cache_type: Optional[str]) -> float:
}.get((cache_type or "f16").strip().lower(), 2.0)
+def _pad_kv_cells(cells: int) -> int:
+ return ((cells + 255) // 256) * 256
+
+
+def _kv_cache_cell_layout(n_ctx: int, n_parallel: int, kv_unified: bool) -> tuple[int, int, int]:
+ """Return llama.cpp's slot count, stream count, and cells per stream."""
+ slots = max(1, n_parallel)
+ padded_ctx = _pad_kv_cells(n_ctx)
+ streams = 1 if kv_unified else slots
+ if padded_ctx <= 0:
+ return slots, streams, 0
+ cells_per_stream = padded_ctx if kv_unified else _pad_kv_cells(padded_ctx // slots)
+ return slots, streams, cells_per_stream
+
+
def _env_main_cache_type_for_budget(env: Optional[Mapping[str, str]] = None) -> Optional[str]:
"""Heavier of the inherited LLAMA_ARG_CACHE_TYPE_K/_V env types when it
exceeds the f16 default, else None. Unsloth emits --cache-type only for the
@@ -1422,6 +1604,39 @@ def _extra_args_main_cache_type_for_budget(extra_args: Optional[Iterable[str]])
return max(candidates, key = _kv_bytes_per_elem)
+def _effective_main_cache_types(
+ args: Optional[Iterable[str]], env: Optional[Mapping[str, str]] = None
+) -> tuple[str, str]:
+ """Effective main K/V cache types after environment and CLI precedence."""
+ source_env = os.environ if env is None else env
+ env_k = (source_env.get("LLAMA_ARG_CACHE_TYPE_K") or "f16").strip().lower()
+ env_v = (source_env.get("LLAMA_ARG_CACHE_TYPE_V") or "f16").strip().lower()
+ arg_k, arg_v = parse_cache_override_per_axis(args)
+ return (
+ (arg_k or env_k).strip().lower(),
+ (arg_v or env_v).strip().lower(),
+ )
+
+
+def _planned_main_cache_types(
+ cache_type_kv: Optional[str],
+ extra_args: Optional[Iterable[str]],
+ env: Optional[Mapping[str, str]] = None,
+) -> tuple[str, str]:
+ """Main K/V types the loader's managed flags and user extras will produce."""
+ args = list(extra_args or ())
+ emitted_type = _extra_args_main_cache_type_for_budget(args) or cache_type_kv
+ if emitted_type:
+ args = [
+ "--cache-type-k",
+ emitted_type,
+ "--cache-type-v",
+ emitted_type,
+ *args,
+ ]
+ return _effective_main_cache_types(args, env)
+
+
def _auto_mode_drops_mtp(
req_mode: Optional[str],
size_b: Optional[float],
@@ -1465,26 +1680,90 @@ def _extra_args_set_spec_type(extra_args: Optional[Iterable[str]]) -> bool:
# set keeps detection and stripping from drifting.
_GPU_OFFLOAD_OVERRIDE_FLAGS = _LAYER_OFFLOAD_FLAGS
_THREAD_OVERRIDE_FLAGS = frozenset({"-t", "--threads"})
-
-
-def _extra_arg_flag_name(token: str) -> Optional[str]:
- if not token.startswith("-") or token in {"-", "--"}:
- return None
- if len(token) >= 2 and (token[1].isdigit() or token[1] == "."):
- return None
- return token.split("=", 1)[0]
+# common_params defaults in the bundled llama.cpp runtime.
+_DEFAULT_LLAMA_N_BATCH = 2048
+_DEFAULT_LLAMA_N_UBATCH = 512
+_LLAMA_ARG_TRUE_VALUES = frozenset({"on", "enabled", "true", "1"})
+_LLAMA_ARG_FALSE_VALUES = frozenset({"off", "disabled", "false", "0"})
+_LLAMA_ARG_AUTO_VALUES = frozenset({"auto", "-1"})
+_LLAMA_ARG_TRUE_OR_AUTO_VALUES = _LLAMA_ARG_TRUE_VALUES | _LLAMA_ARG_AUTO_VALUES
+_LLAMA_ARG_TRUE_FALSE_AUTO_VALUES = _LLAMA_ARG_TRUE_OR_AUTO_VALUES | _LLAMA_ARG_FALSE_VALUES
def _extra_args_set_any_flag(extra_args: Optional[Iterable[str]], flags: Collection[str]) -> bool:
if not extra_args:
return False
for raw in extra_args:
- flag = _extra_arg_flag_name(str(raw))
+ flag = _flag_name(str(raw))
if flag in flags:
return True
return False
+def _swa_full_from_args_or_env(
+ extra_args: Optional[Iterable[str]], env: Optional[Mapping[str, str]] = None
+) -> bool:
+ """Whether llama.cpp receives the enable-only full-size SWA option."""
+ if _extra_args_set_any_flag(extra_args, {"--swa-full"}):
+ return True
+ value = (os.environ if env is None else env).get("LLAMA_ARG_SWA_FULL")
+ return value in _LLAMA_ARG_TRUE_VALUES
+
+
+def _kv_unified_from_args(
+ extra_args: Optional[Iterable[str]],
+ default: bool = False,
+ env: Optional[Mapping[str, str]] = None,
+) -> bool:
+ """Resolve llama.cpp's environment and last-wins unified KV flags."""
+ enabled = False
+ value = (os.environ if env is None else env).get("LLAMA_ARG_KV_UNIFIED")
+ if value in _LLAMA_ARG_TRUE_VALUES:
+ enabled = True
+ elif value in _LLAMA_ARG_FALSE_VALUES:
+ enabled = False
+ if default:
+ # Studio's managed --kv-unified flag is appended after environment
+ # parsing and before user extras.
+ enabled = True
+ for raw in extra_args or ():
+ flag = _flag_name(str(raw))
+ if flag in {"-kvu", "--kv-unified"}:
+ enabled = True
+ elif flag in {"-no-kvu", "--no-kv-unified"}:
+ enabled = False
+ return enabled
+
+
+def _flash_attn_enabled_from_args(
+ args: Optional[Iterable[str]],
+ default: bool = True,
+ env: Optional[Mapping[str, str]] = None,
+) -> bool:
+ """Resolve llama.cpp's environment and last-wins flash-attention settings."""
+ enabled = default
+ # llama.cpp applies LLAMA_ARG_FLASH_ATTN before parsing argv (arg.cpp set_env),
+ # so the CLI still wins. --flash-attn has no args_neg, so no LLAMA_ARG_NO_ twin.
+ value = (os.environ if env is None else env).get("LLAMA_ARG_FLASH_ATTN")
+ if value in _LLAMA_ARG_FALSE_VALUES:
+ enabled = False
+ elif value in _LLAMA_ARG_TRUE_OR_AUTO_VALUES:
+ enabled = True
+ values = [str(arg) for arg in args] if args else []
+ for i, raw in enumerate(values):
+ if _flag_name(raw) not in {"-fa", "--flash-attn"}:
+ continue
+ _, eq, inline = raw.partition("=")
+ value = inline if eq else "on"
+ if not eq and i + 1 < len(values) and values[i + 1] in _LLAMA_ARG_TRUE_FALSE_AUTO_VALUES:
+ value = values[i + 1]
+ if value in _LLAMA_ARG_FALSE_VALUES:
+ enabled = False
+ elif value in _LLAMA_ARG_TRUE_OR_AUTO_VALUES:
+ enabled = True
+ return enabled
+
+
def _effective_spec_type(
extra_args: Optional[Iterable[str]], env: Optional[Mapping[str, str]] = None
) -> Optional[str]:
@@ -1496,7 +1775,8 @@ def _effective_spec_type(
cli_present = False
cli_value: Optional[str] = None
for i, raw in enumerate(args):
- flag, eq, inline = raw.partition("=")
+ flag = _flag_name(raw)
+ _, eq, inline = raw.partition("=")
if flag == "--spec-default":
cli_present = True
cli_value = "default"
@@ -1540,7 +1820,8 @@ def _extra_args_spec_draft_n_max(extra_args: Optional[Iterable[str]]) -> Optiona
args = [str(a) for a in extra_args]
found: Optional[int] = None
for i, raw in enumerate(args):
- flag, eq, inline = raw.partition("=")
+ flag = _flag_name(raw)
+ _, eq, inline = raw.partition("=")
if flag not in ("--spec-draft-n-max", "--draft-max"):
continue
value = inline if eq else (args[i + 1] if i + 1 < len(args) else "")
@@ -1570,7 +1851,8 @@ def _extra_args_mtp_draft_path(
args = [str(a) for a in extra_args] if extra_args else []
found: Optional[str] = None
for i, raw in enumerate(args):
- flag, eq, inline = raw.partition("=")
+ flag = _flag_name(raw)
+ _, eq, inline = raw.partition("=")
if flag not in flags:
continue
value = inline if eq else (args[i + 1] if i + 1 < len(args) else "")
@@ -1594,7 +1876,8 @@ def _extra_args_draft_cache_types(
k_type: Optional[str] = None
v_type: Optional[str] = None
for i, raw in enumerate(args):
- flag, eq, inline = raw.partition("=")
+ flag = _flag_name(raw)
+ _, eq, inline = raw.partition("=")
if flag not in k_flags and flag not in v_flags:
continue
value = inline if eq else (args[i + 1] if i + 1 < len(args) else "")
@@ -1626,7 +1909,8 @@ def _extra_args_draft_offloaded_to_cpu(
last_ngl: Optional[str] = None
last_dev: Optional[str] = None
for i, raw in enumerate(args):
- flag, eq, inline = raw.partition("=")
+ flag = _flag_name(raw)
+ _, eq, inline = raw.partition("=")
value = inline if eq else (args[i + 1] if i + 1 < len(args) else "")
if flag in ngl_flags:
last_ngl = value
@@ -1648,31 +1932,61 @@ def _extra_args_draft_offloaded_to_cpu(
def _extra_args_n_ubatch(
- extra_args: Optional[Iterable[str]], env: Optional[Mapping[str, str]] = None
+ extra_args: Optional[Iterable[str]],
+ env: Optional[Mapping[str, str]] = None,
+ n_ctx: Optional[int] = None,
) -> Optional[int]:
- """Physical micro-batch from extras (--ubatch-size/-ub) else the LLAMA_ARG_UBATCH
- env, else None. It sizes the compute-graph buffer, so an override must reach
- the VRAM reserve."""
+ """Effective ubatch after llama.cpp normalizes it, or None at defaults."""
+ values = {
+ "batch": _DEFAULT_LLAMA_N_BATCH,
+ "ubatch": _DEFAULT_LLAMA_N_UBATCH,
+ }
+ source_env = os.environ if env is None else env
+ overridden = False
+ for key, env_name in (
+ ("batch", "LLAMA_ARG_BATCH"),
+ ("ubatch", "LLAMA_ARG_UBATCH"),
+ ):
+ raw = source_env.get(env_name)
+ if raw:
+ try:
+ values[key] = int(raw)
+ overridden = True
+ except (TypeError, ValueError):
+ pass
+
args = [str(a) for a in extra_args] if extra_args else []
- found: Optional[int] = None
+ flags = {
+ "-b": "batch",
+ "--batch-size": "batch",
+ "-ub": "ubatch",
+ "--ubatch-size": "ubatch",
+ }
for i, raw in enumerate(args):
- flag, eq, inline = raw.partition("=")
- if flag not in ("--ubatch-size", "-ub"):
+ flag = _flag_name(raw)
+ _, eq, inline = raw.partition("=")
+ key = flags.get(flag)
+ if key is None:
continue
value = inline if eq else (args[i + 1] if i + 1 < len(args) else "")
try:
- found = int(value)
+ values[key] = int(value)
+ overridden = True
except (TypeError, ValueError):
continue
- if found is not None:
- return found
- raw = (os.environ if env is None else env).get("LLAMA_ARG_UBATCH")
- if raw:
- try:
- return int(raw)
- except (TypeError, ValueError):
- pass
- return None
+ if not overridden:
+ return None
+
+ # common_params stores signed values, then llama_context_params converts
+ # them to uint32_t. A zero ubatch means "use batch"; the context then caps
+ # ubatch at batch size.
+ batch = values["batch"] & 0xFFFFFFFF
+ raw_ubatch = values["ubatch"]
+ ubatch = batch if raw_ubatch == 0 else raw_ubatch & 0xFFFFFFFF
+ effective = min(batch, ubatch)
+ if n_ctx is not None and n_ctx > 0:
+ effective = min(effective, n_ctx)
+ return effective
def _build_ngram_mod_flags(
@@ -1916,6 +2230,8 @@ class LlamaCppBackend:
self._effective_context_length: Optional[int] = None
self._max_context_length: Optional[int] = None
self._effective_parallel_slots: int = 1
+ # --parallel the last load asked for, before any fit-time reduction.
+ self._requested_n_parallel: int = 1
self._chat_template: Optional[str] = None
self._chat_template_override: Optional[str] = None
self._supports_reasoning: bool = False
@@ -1938,6 +2254,10 @@ class LlamaCppBackend:
self._tensor_split: Optional[List[float]] = None
# User-picked physical GPU indices (None = automatic selection).
self._gpu_ids: Optional[List[int]] = None
+ # RAW requested GPU pin, before the fit narrowed it. self._gpu_ids records the
+ # EFFECTIVE (fit-narrowed) pin for /status; dedupe compares this raw value so a
+ # [0, 1] narrowed to [0] and re-sent as [0, 1] still matches (#7239).
+ self._requested_gpu_ids: Optional[List[int]] = None
# Layer load kept multi-GPU only to honor a downgraded tensor request, so a
# later explicit tensor-off reloads instead of deduping to it (#6659).
self._layer_preserves_tensor_intent: bool = False
@@ -1989,6 +2309,9 @@ class LlamaCppBackend:
# Serialises mid-session respawns so many generations hitting a killed
# server trigger at most one reload (see _respawn_if_dead).
self._respawn_lock = threading.Lock()
+ # Bumped by every unload. load_model clears _cancel_event, so a respawn that
+ # raced an unload needs a signal that survives the clear (see _respawn_if_dead).
+ self._unload_epoch = 0
# Set by the in-app updater while it swaps prebuilt binaries; load_model()
# rejects fast so no server starts from a half-swapped binary.
self._llama_update_in_progress = False
@@ -2024,6 +2347,14 @@ class LlamaCppBackend:
# save can tell whether the model files were swapped on disk since load.
self._slot_loaded_identity: Optional[tuple] = None
self._prompt_cache_disabled: bool = False
+ self._swa_full: bool = False
+ self._kv_cache_unified: bool = False
+ self._n_ubatch: int = self._DEFAULT_N_UBATCH
+ self._flash_attn_enabled: bool = True
+ self._effective_cache_types: tuple[str, str] = ("f16", "f16")
+ # Total KV allocation context across all slots. _effective_context_length
+ # becomes the per-slot request limit after /props reconciliation.
+ self._kv_cache_context_total: Optional[int] = None
# True once a probe has completed; cleared on transient failure.
self._is_audio: bool = False
self._audio_type: Optional[str] = None
@@ -2076,6 +2407,11 @@ class LlamaCppBackend:
"""True when the loaded GGUF is a block-diffusion model (DiffusionGemma)."""
return self._is_diffusion
+ @property
+ def swa_full(self) -> bool:
+ """Whether the active llama-server received full-size SWA mode."""
+ return self._swa_full
+
@property
def hf_variant(self) -> Optional[str]:
return self._hf_variant
@@ -2132,6 +2468,17 @@ class LlamaCppBackend:
slots = 1
return max(1, slots)
+ @property
+ def requested_parallel_slots(self) -> int:
+ """--parallel the last load asked for, before any fit-time reduction.
+ The reload dedupe compares requested-vs-requested (like requested_n_ctx);
+ the effective count would reload forever after a fitter reduction."""
+ try:
+ slots = int(getattr(self, "_requested_n_parallel", 1))
+ except (TypeError, ValueError):
+ slots = 1
+ return max(1, slots)
+
@property
def max_context_length(self) -> Optional[int]:
"""Return the largest context that fits on this hardware at load time.
@@ -2157,6 +2504,8 @@ class LlamaCppBackend:
def _reset_effective_parallel_slots(self) -> None:
self._effective_parallel_slots = 1
+ # Cleared with the effective count so a stale value can't skew the dedupe.
+ self._requested_n_parallel = 1
@staticmethod
def _read_rss_bytes(pid: int) -> Optional[int]:
@@ -2409,6 +2758,46 @@ class LlamaCppBackend:
"""User-picked physical GPU indices, or None for automatic selection."""
return self._gpu_ids
+ @property
+ def requested_gpu_ids(self) -> Optional[List[int]]:
+ """RAW requested GPU pin (before the fit narrowed it), or None for auto.
+ gpu_ids echoes the EFFECTIVE pin for /status."""
+ return self._requested_gpu_ids
+
+ def matches_gpu_ids(self, gpu_ids: Optional[List[int]]) -> bool:
+ """Whether a requested pin is already satisfied by the active runner.
+
+ A regular GGUF load may narrow the requested placement pool to the
+ smallest fitting subset. Accept both the original request and the
+ effective status-echoed subset so either can round-trip without a
+ needless reload. Diffusion drives one device and keeps its existing
+ lowest-device normalization.
+ """
+ if self._is_diffusion:
+ requested = [sorted(int(x) for x in gpu_ids)[0]] if gpu_ids else None
+ return requested == (self._gpu_ids or None)
+
+ requested = sorted(int(x) for x in gpu_ids) if gpu_ids else None
+ raw = self._requested_gpu_ids or None
+ effective = self._gpu_ids or None
+ return requested == raw or requested == effective
+
+ def _record_matching_gpu_request(self, gpu_ids: Optional[List[int]]) -> None:
+ """Adopt the caller's explicit pool after a full already-loaded match.
+
+ Matching an effective subset avoids a reload, but the incoming request
+ is still the user's latest placement intent. Record it so status and a
+ later reload do not restore GPUs the user just removed.
+ """
+ if self._is_diffusion:
+ self._requested_gpu_ids = [sorted(int(x) for x in gpu_ids)[0]] if gpu_ids else None
+ else:
+ self._requested_gpu_ids = sorted(int(x) for x in gpu_ids) if gpu_ids else None
+ if self._last_load_kwargs is not None:
+ self._last_load_kwargs["gpu_ids"] = (
+ list(self._requested_gpu_ids) if self._requested_gpu_ids else None
+ )
+
@property
def n_layers(self) -> Optional[int]:
"""Model layer count (GGUF block_count), or None if unknown."""
@@ -2652,6 +3041,7 @@ class LlamaCppBackend:
"found": False,
"mtp_token": None,
"supports_mtp": False,
+ "mtp_probe_inconclusive": True,
"ngram_mod_flavor": None,
"supports_ngram_mod": False,
"spec_draft_n_max_flag": None,
@@ -2684,17 +3074,22 @@ class LlamaCppBackend:
supports_no_cache_prompt = False
supports_metrics = False
supports_slot_save = False
+ saw_spec_type = False
+ probe_ok = False
+ help_text = ""
try:
probe_env = cls._llama_server_env_for_binary(bin_path)
result = subprocess.run(
[bin_path, "--help"],
capture_output = True,
text = True,
+ encoding = "utf-8",
errors = "replace",
timeout = 10,
check = False,
env = probe_env,
)
+ probe_ok = result.returncode == 0
help_text = (result.stdout or "") + "\n" + (result.stderr or "")
# Split into per-flag blocks (each --flag line + its indented
# continuation), so the "argument has been removed" description
@@ -2739,17 +3134,19 @@ class LlamaCppBackend:
return False
return "argument has been removed" not in desc
- # MTP token from the --spec-type line.
- spec_line = ""
- for line in help_text.splitlines():
- if "--spec-type" in line:
- spec_line = line
- break
- # PR #22673 used draft-mtp; later renamed to mtp.
- if "draft-mtp" in spec_line:
- mtp_token = "draft-mtp"
- elif re.search(r"[|,\[]mtp[|,\]]", spec_line):
- mtp_token = "mtp"
+ # MTP token from the full --spec-type help block (decl + indented
+ # continuation). First-line-only probing missed builds putting the
+ # enum on the next line (#7302). Prefer draft-mtp (PR #22673) over mtp.
+ spec_help = blocks.get("--spec-type") or ""
+ if not spec_help:
+ # Fallback: join --spec-type lines, avoiding incidental "mtp" in --help.
+ spec_help = "\n".join(
+ line for line in help_text.splitlines() if "--spec-type" in line
+ )
+ mtp_token = cls._mtp_token_from_spec_help(spec_help)
+ # Only a resolved --spec-type block confirms missing MTP; empty/crash
+ # leaves saw_spec_type False so supports_mtp fails open.
+ saw_spec_type = bool(spec_help.strip()) and "--spec-type" in spec_help
# ngram-mod flag flavor. Post-rename builds advertise both new
# args (real) and legacy ones (stubs); pre-rename builds only
@@ -2785,11 +3182,29 @@ class LlamaCppBackend:
supports_slot_save = _is_real("--slot-save-path")
except (OSError, subprocess.SubprocessError) as exc:
logger.debug(f"llama-server --help probe failed: {exc}")
+ saw_spec_type = False
+ probe_ok = False
+ help_text = ""
+
+ help_nonempty = bool(help_text.strip())
+ # Confirmed only when a successful --help lists a --spec-type block with
+ # mtp/draft-mtp; nonempty --help without it is a definitive pre-spec
+ # binary; failed/empty probes stay inconclusive (#7302).
+ if saw_spec_type and probe_ok:
+ supports_mtp = mtp_token is not None
+ mtp_probe_inconclusive = False
+ elif help_nonempty and probe_ok:
+ supports_mtp = False
+ mtp_probe_inconclusive = False
+ else:
+ supports_mtp = False
+ mtp_probe_inconclusive = True
info = {
"found": True,
"mtp_token": mtp_token,
- "supports_mtp": mtp_token is not None,
+ "supports_mtp": supports_mtp,
+ "mtp_probe_inconclusive": mtp_probe_inconclusive,
"ngram_mod_flavor": ngram_mod_flavor,
"supports_ngram_mod": ngram_mod_flavor is not None,
"spec_draft_n_max_flag": spec_draft_n_max_flag,
@@ -2805,6 +3220,21 @@ class LlamaCppBackend:
cls._capability_cache[cache_key] = info
return info
+ @staticmethod
+ def _mtp_token_from_spec_help(spec_help: str) -> Optional[str]:
+ """Extract ``draft-mtp`` / ``mtp`` from a ``--spec-type`` help snippet.
+
+ Prefers ``draft-mtp`` (llama.cpp PR #22673) over the later bare ``mtp``
+ rename. Returns ``None`` when neither token appears as an enum value.
+ """
+ text = spec_help or ""
+ if "draft-mtp" in text:
+ return "draft-mtp"
+ # Bare `mtp` enum token (`|mtp|`, `,mtp,`, ...), not a substring.
+ if re.search(r"(?physical mapping."""
try:
import torch
- is_rocm = getattr(torch.version, "hip", None) is not None
+
+ # Same ROCm detection as _emit_child_gpu_visibility: AMD SDK wheels
+ # leave version.hip unset but encode "rocm" in __version__. The two
+ # must agree, else an inherited ROCR mask reads back as "no mask",
+ # ordinal 0 is labelled physical 0, and the child's new ROCR pin
+ # re-exposes the GPU the inherited mask was hiding.
+ is_rocm = (
+ getattr(torch.version, "hip", None) is not None
+ or "rocm" in getattr(torch, "__version__", "").lower()
+ )
except Exception:
is_rocm = False
if is_rocm:
hip_v = os.environ.get("HIP_VISIBLE_DEVICES")
- rocr_v = os.environ.get("ROCR_VISIBLE_DEVICES")
+ # ROCR_VISIBLE_DEVICES is a Linux ROCr variable; Windows HIP has no
+ # ROCr layer, so a stray ROCR var there does not mask the runtime and
+ # must not be read as the ordinal->physical mapping (mirrors the
+ # Windows gate in _emit_child_gpu_visibility).
+ rocr_v = None if sys.platform == "win32" else os.environ.get("ROCR_VISIBLE_DEVICES")
cvd = (
hip_v
if hip_v is not None
@@ -2882,20 +3325,53 @@ class LlamaCppBackend:
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."""
+ def _emit_child_gpu_visibility(
+ env: dict,
+ pinned: str,
+ *,
+ prefer_rocr: bool = False,
+ ) -> None:
+ """Write the child's GPU visibility mask: CUDA, plus a ROCm mirror on AMD
+ (masking only CUDA_VISIBLE_DEVICES leaves an AMD child seeing every GPU).
+
+ Default: HIP_VISIBLE_DEVICES, clearing any inherited ROCR mask so the two
+ can't stack (ROCR re-indexes from 0, then a non-zero HIP pin points out of
+ range, HIP sees 0 devices, and llama.cpp falls back to CPU).
+
+ prefer_rocr masks at the ROCr/HSA layer instead (clearing HIP). A HIP mask
+ filters only AFTER the HSA runtime enumerates every agent, and that
+ enumeration segfaults at startup on a GPU the build has no kernels for
+ (e.g. a gfx1036 iGPU under a gfx103X prebuilt: that bundle maps only
+ gfx1030/1031/1032/1034), before llama-server logs a line. ROCR drops the
+ device at the driver layer, consuming physical ids.
+ The CPU-only sentinel ("-1") has no portable ROCR spelling, so it keeps
+ the HIP mask. Windows keeps the HIP mask too: ROCR_VISIBLE_DEVICES is a
+ Linux ROCr variable (Windows HIP has no ROCr layer), so the ROCR pin
+ would be dead there while the cleared HIP mask stops selecting."""
env["CUDA_VISIBLE_DEVICES"] = pinned
try:
import torch as _torch
- if getattr(_torch.version, "hip", None) is not None:
- env["HIP_VISIBLE_DEVICES"] = pinned
- env.pop("ROCR_VISIBLE_DEVICES", None)
+
+ # torch.version.hip is set on ROCm, None on CUDA; AMD SDK wheels may
+ # leave it unset but encode "rocm" in __version__ (mirrors detect_hardware).
+ if (
+ getattr(_torch.version, "hip", None) is not None
+ or "rocm" in getattr(_torch, "__version__", "").lower()
+ ):
+ if prefer_rocr and pinned != "-1" and sys.platform != "win32":
+ env["ROCR_VISIBLE_DEVICES"] = pinned
+ env.pop("HIP_VISIBLE_DEVICES", None)
+ # ROCR re-indexes the visible agents from 0, and with HIP
+ # cleared HIP honours CUDA_VISIBLE_DEVICES -- so it must carry
+ # the post-ROCR ordinals (0..N-1), not the physical ids, else a
+ # non-zero pick points out of range and HIP sees 0 devices (the
+ # same stacking the default path avoids by clearing ROCR).
+ env["CUDA_VISIBLE_DEVICES"] = ",".join(
+ str(i) for i in range(len(pinned.split(",")))
+ )
+ else:
+ env["HIP_VISIBLE_DEVICES"] = pinned
+ env.pop("ROCR_VISIBLE_DEVICES", None)
except Exception as e:
logger.debug("Failed to set ROCm visibility env vars for child: %s", e)
@@ -2930,11 +3406,25 @@ class LlamaCppBackend:
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))
+ # Re-emit at the layer that produced the mapping. A parent masked only
+ # via ROCR_VISIBLE_DEVICES hides agents at the driver layer, and the
+ # default HIP re-emission clears that mask -- HSA then enumerates every
+ # agent again and can segfault at startup on an unsupported GPU the
+ # parent was hiding (the crash prefer_rocr exists to avoid). Linux-only,
+ # mirroring _resolve_visible_physical_ids: on Windows a stray ROCR var
+ # is dead and was not the mapping's source.
+ prefer_rocr = (
+ sys.platform != "win32"
+ and env.get("HIP_VISIBLE_DEVICES") is None
+ and env.get("ROCR_VISIBLE_DEVICES") is not None
+ )
+ LlamaCppBackend._emit_child_gpu_visibility(
+ env, ",".join(str(i) for i in order), prefer_rocr = prefer_rocr
+ )
@staticmethod
def _amd_apu_wants_unified_memory(gpu_indices = None) -> bool:
- """True only for AMD unified-memory APUs (gfx1150/gfx1151), where
+ """True only for AMD unified-memory APUs (gfx1150/gfx1151/gfx1152), where
GGML_CUDA_ENABLE_UNIFIED_MEMORY lets llama.cpp use shared system RAM (it
hurts discrete GPUs). gpu_indices (PHYSICAL ids) scopes the check to the
selected GPUs, so a dGPU on a mixed host is not treated as unified-memory;
@@ -2964,7 +3454,9 @@ class LlamaCppBackend:
)
arch_by_id[pid] = _arch.split(":")[0].strip().lower()
for _i in list(gpu_indices) if gpu_indices is not None else list(arch_by_id):
- if arch_by_id.get(_i) in {"gfx1150", "gfx1151"}:
+ # gfx1152 is Krackan Point (Radeon 860M/840M), the third RDNA 3.5
+ # APU: same shared GPU/system-RAM pool as Strix Point/Halo.
+ if arch_by_id.get(_i) in {"gfx1150", "gfx1151", "gfx1152"}:
return True
except Exception:
return False
@@ -3164,6 +3656,8 @@ class LlamaCppBackend:
],
capture_output = True,
text = True,
+ encoding = "utf-8",
+ errors = "replace",
timeout = 10,
env = child_env_without_native_path_secret(),
**_windows_hidden_subprocess_kwargs(),
@@ -3235,18 +3729,17 @@ class LlamaCppBackend:
return []
@staticmethod
- def _get_gpu_free_memory_vulkan(binary: Optional[str] = None) -> list[tuple[int, int, int]]:
- """Query free (and total) VRAM per device via the bundled ggml Vulkan backend.
+ def _run_vulkan_probe(binary: Optional[str] = None) -> list[dict]:
+ """Run ``_vulkan_probe.py`` and parse its per-device lines.
- Loads ``libggml-vulkan`` in a short-lived subprocess (no Vulkan instance
- in this process) and returns (device_index, free_mib, total_mib) sorted
- by index. The index is ggml's compact Vulkan ordinal -- the one the
- registry names ``Vulkan`` and load_model pins with ``--device``,
- NOT the raw ``GGML_VK_VISIBLE_DEVICES`` space. A user-set
- ``GGML_VK_VISIBLE_DEVICES`` is honored by ggml (passed through), so the
- list already reflects it. iGPUs leave a host-RAM margin (see
- ``_apply_igpu_host_reserve_mib``) and report total 0; discrete cards pass
- their real total through. [] when no Vulkan build or device is reachable.
+ Returns raw (uncapped) rows sorted by index:
+ ``{"index", "free_mib", "total_mib", "is_igpu", "name"}``. The index is
+ ggml's compact Vulkan ordinal -- the one the registry names
+ ``Vulkan`` and load_model pins with ``--device``, NOT the raw
+ ``GGML_VK_VISIBLE_DEVICES`` space. A user-set ``GGML_VK_VISIBLE_DEVICES``
+ is honored by ggml (passed through), so the list already reflects it.
+ ``name`` is ggml's device description; "" from an older 4-column probe.
+ [] when no Vulkan build or device is reachable.
"""
binary = binary or LlamaCppBackend._find_llama_server_binary()
if not binary:
@@ -3271,12 +3764,15 @@ class LlamaCppBackend:
)
probe_script = Path(__file__).with_name("_vulkan_probe.py")
try:
+ # UTF-8 to match the probe's stdout reconfigure: device names can be
+ # non-ASCII, and the platform-default decode (cp1252) could throw.
result = subprocess.run(
[sys.executable, str(probe_script), str(binary_dir)],
capture_output = True,
- text = True,
+ encoding = "utf-8",
+ errors = "replace",
timeout = 15,
- env = env,
+ env = utf8_child_env(env),
**_windows_hidden_subprocess_kwargs(),
)
if result.returncode != 0:
@@ -3288,21 +3784,56 @@ class LlamaCppBackend:
logger.debug(f"vulkan GPU probe failed: {e}")
return []
- gpus: list[tuple[int, int, int]] = []
+ rows: list[dict] = []
for line in result.stdout.strip().splitlines():
parts = line.split("\t")
- if len(parts) != 4:
+ # 4 columns from an older probe (no name); 5 with the name column.
+ if len(parts) not in (4, 5):
continue
try:
- idx = int(parts[0])
- free_mib = int(parts[1]) // (1024 * 1024)
- is_igpu = parts[2] == "1"
- # iGPU "total" is shared RAM, not a VRAM budget -> keep 0 so the
- # fit stays on free*frac (the host reserve below is its
- # headroom); a discrete card passes its real total through.
- total_mib = 0 if is_igpu else int(parts[3]) // (1024 * 1024)
+ rows.append(
+ {
+ "index": int(parts[0]),
+ "free_mib": int(parts[1]) // (1024 * 1024),
+ "is_igpu": parts[2] == "1",
+ "total_mib": int(parts[3]) // (1024 * 1024),
+ "name": parts[4].strip() if len(parts) == 5 else "",
+ }
+ )
except ValueError:
continue
+ rows.sort(key = lambda r: r["index"])
+ return rows
+
+ @staticmethod
+ def vulkan_device_inventory(binary: Optional[str] = None) -> list[dict]:
+ """UI-facing Vulkan device list: the devices llama-server will actually
+ use, with real totals (an iGPU keeps its shared-RAM total here -- the
+ caller labels it, unlike the fit which zeroes it). Same rows as
+ ``_run_vulkan_probe``; names fall back to ``Vulkan``.
+ """
+ rows = LlamaCppBackend._run_vulkan_probe(binary)
+ for row in rows:
+ if not row["name"]:
+ row["name"] = f"Vulkan{row['index']}"
+ return rows
+
+ @staticmethod
+ def _get_gpu_free_memory_vulkan(binary: Optional[str] = None) -> list[tuple[int, int, int]]:
+ """Query free (and total) VRAM per device via the bundled ggml Vulkan backend.
+
+ Fit-oriented view of ``_run_vulkan_probe``: returns (device_index,
+ free_mib, total_mib) sorted by index. iGPUs leave a host-RAM margin (see
+ ``_apply_igpu_host_reserve_mib``) and report total 0; discrete cards pass
+ their real total through. [] when no Vulkan build or device is reachable.
+ """
+ gpus: list[tuple[int, int, int]] = []
+ for row in LlamaCppBackend._run_vulkan_probe(binary):
+ idx, free_mib, is_igpu = row["index"], row["free_mib"], row["is_igpu"]
+ # iGPU "total" is shared RAM, not a VRAM budget -> keep 0 so the
+ # fit stays on free*frac (the host reserve below is its
+ # headroom); a discrete card passes its real total through.
+ total_mib = 0 if is_igpu else row["total_mib"]
capped = _apply_igpu_host_reserve_mib(free_mib, is_igpu)
if capped < free_mib:
logger.info(
@@ -3311,7 +3842,6 @@ class LlamaCppBackend:
f"({free_mib}->{capped}MiB usable)"
)
gpus.append((idx, capped, total_mib))
- gpus.sort(key = lambda g: g[0])
if gpus:
logger.info(
"Vulkan GPU memory detected: "
@@ -3330,7 +3860,7 @@ class LlamaCppBackend:
except Exception:
pass
try:
- with open("/proc/meminfo") as f:
+ with open("/proc/meminfo", encoding = "utf-8") as f:
for line in f:
if line.startswith("MemAvailable:"):
return int(line.split()[1]) // 1024 # kB -> MiB
@@ -3448,6 +3978,14 @@ class LlamaCppBackend:
# aborts a --split-mode tensor load, so it's dropped for the tensor attempt.
_TENSOR_PARALLEL_KV_TYPES = frozenset({"f16", "bf16", "f32"})
+ # V cache types that llama.cpp can run WITHOUT flash attention. Only the V
+ # axis has the dependency: a quantized V cache (q8_0/q4_0/q4_1/q5_0/q5_1/
+ # iq4_nl) aborts init with "V cache quantization requires flash_attn", while
+ # a quantized K cache runs fine without FA. So the flash-attn-off crash-
+ # recovery fallback must reset a quantized V cache to f16 before it can
+ # launch (and leaves K alone). These three are the only non-quantized types.
+ _NON_QUANTIZED_KV_TYPES = frozenset({"f16", "bf16", "f32"})
+
# Main-model placement settings that Manual mode owns. They must not leak
# from Studio's parent environment into llama-server and silently override
# the command assembled from the current request. Draft-model placement is
@@ -3592,6 +4130,9 @@ class LlamaCppBackend:
lib_dirs.extend(_wsl_system_rocm_lib_dirs())
if lib_dirs:
env.setdefault("HSA_ENABLE_DXG_DETECTION", "1")
+ # Native Linux AMD: system ROCm libs before the bundle's HIP runtime,
+ # which can be incompatible with the host amdkfd driver.
+ lib_dirs.extend(_native_linux_system_rocm_lib_dirs(binary_dir))
lib_dirs.append(binary_dir)
_arch = platform.machine() # x86_64, aarch64, etc.
@@ -3742,6 +4283,32 @@ class LlamaCppBackend:
is non-None here."""
return self._embedding_length // self._n_heads if self._n_heads else 128 # type: ignore[operator]
+ def _max_kv_value_width(
+ self,
+ default_len: int,
+ swa_len: Optional[int] = None,
+ ) -> int:
+ """llama.cpp's hparams.n_embd_v_gqa_max() over every model layer."""
+ n_layers = self._n_layers or 1
+ n_kv = self._n_kv_heads or self._n_heads or 1
+ if self._sliding_window_pattern is None:
+ max_len = max(default_len, swa_len or default_len)
+ return max(
+ self._kv_heads_for_layer(layer_idx, n_kv) * max_len for layer_idx in range(n_layers)
+ )
+ return max(
+ self._kv_heads_for_layer(layer_idx, n_kv)
+ * (
+ (swa_len or default_len)
+ if (
+ layer_idx < len(self._sliding_window_pattern)
+ and self._sliding_window_pattern[layer_idx]
+ )
+ else default_len
+ )
+ for layer_idx in range(n_layers)
+ )
+
def _estimate_kv_cache_bytes(
self,
n_ctx: int,
@@ -3750,22 +4317,26 @@ class LlamaCppBackend:
swa_full: bool = False,
n_parallel: int = 1,
kv_unified: bool = True,
+ n_ubatch: Optional[int] = None,
ctx_checkpoints: int = 0,
+ flash_attn: bool = True,
) -> int:
"""Estimate KV cache VRAM for a given context length.
5-path architecture-aware estimation:
1. MLA -- compressed KV latent + RoPE, K-only (no separate V)
2. Hybrid -- only attention layers need KV (Mamba layers don't)
- 3. SWA -- sliding-window layers cache min(ctx, window) tokens
+ 3. SWA -- sliding-window layers use compact or full cache cells
4. GQA -- standard full KV with explicit key/value dimensions
5. Legacy -- fallback using embed // n_heads
Server-flag knobs (mirror llama-server's CLI):
swa_full -- --swa-full: SWA layers cache full n_ctx (path 3->4).
- n_parallel -- --parallel slots: non-SWA constant, SWA scale linearly.
- kv_unified -- --kv-unified: memory no-op (API forward-compat).
+ n_parallel -- --parallel slots: controls per-slot stream padding.
+ kv_unified -- --kv-unified: one shared stream vs one per slot.
+ n_ubatch -- --ubatch-size: SWA cache's processing headroom.
ctx_checkpoints -- --ctx-checkpoints: N SWA snapshots per slot.
+ flash_attn -- False pads variable-width V tensors to the model max.
Returns 0 if metadata is insufficient.
"""
@@ -3780,9 +4351,17 @@ class LlamaCppBackend:
n_kv = self._n_kv_heads or self._n_heads or 1 # type: ignore[assignment]
# Bytes per element depends on KV cache quantization
- bpe = _kv_bytes_per_elem(cache_type_kv)
+ bpe_k = _kv_bytes_per_elem(cache_type_kv)
+ # The automatic FA-off retry rewrites an invalid quantized V cache to
+ # f16. Pricing that viable retry here avoids under-reserving it.
+ bpe_v = bpe_k if flash_attn else max(bpe_k, _kv_bytes_per_elem("f16"))
- slots = max(1, n_parallel)
+ slots, streams, cells_per_stream = _kv_cache_cell_layout(n_ctx, n_parallel, kv_unified)
+ total_cells = cells_per_stream * streams
+ ubatch = max(
+ 0,
+ int(self._DEFAULT_N_UBATCH if n_ubatch is None else n_ubatch),
+ )
# Path 1: MLA (DeepSeek-V2/V3, GLM-4.7, GLM-5, Kimi-K2.5)
# One compressed KV latent per token/layer (shared across heads); V is
@@ -3793,7 +4372,7 @@ class LlamaCppBackend:
n_kv_mla = self._n_kv_heads or 1
rope_dim = self._key_length_mla or 64
key_len = self._kv_key_length or (self._kv_lora_rank + rope_dim)
- return int(n_layers_kv * n_ctx * n_kv_mla * key_len * bpe)
+ return int(n_layers_kv * total_cells * n_kv_mla * key_len * bpe_k)
key_len = self._kv_key_length
val_len = self._kv_value_length
@@ -3804,16 +4383,18 @@ class LlamaCppBackend:
fai = self._full_attention_interval
n_attn = -(-n_layers // fai) if fai > 0 else n_layers # ceiling division
if key_len is not None and val_len is not None:
- return int(n_attn * n_ctx * n_kv * (key_len + val_len) * bpe)
+ v_width = n_kv * val_len if flash_attn else self._max_kv_value_width(val_len)
+ return int(n_attn * total_cells * (n_kv * key_len * bpe_k + v_width * bpe_v))
head_dim = self._legacy_head_dim()
- return int(n_attn * n_ctx * n_kv * 2 * head_dim * bpe)
+ return int(n_attn * total_cells * n_kv * 2 * head_dim * bpe_k)
# Path 3: Sliding window (Gemma 2/3/3n/4, gpt-oss, Cohere2 ...). Pattern
# from the resolver; if absent, falls through to the legacy 1/4-global
# heuristic. --parallel N accounting (verified against llama-server):
- # non-SWA cells = n_ctx split across slots (CONSTANT); SWA per-slot cells
- # = 2*sliding_window (capped at n_ctx/per_slot_ctx) -> LINEAR in slots.
- # --swa-full forces full n_ctx for SWA; --ctx-checkpoints N adds snapshots.
+ # non-SWA cells total n_ctx across streams. Compact SWA adds one processing
+ # micro-batch to the window allowance and pads to 256 cells; unified mode
+ # holds all slots in one stream, while non-unified mode has one stream per
+ # slot. --swa-full expands SWA to each stream's full context.
if (
self._sliding_window is not None
and self._sliding_window > 0
@@ -3821,15 +4402,19 @@ class LlamaCppBackend:
and val_len is not None
):
swa = self._sliding_window
- per_slot_ctx = max(1, n_ctx // slots)
- # --swa-full caches full per_slot_ctx (constant n_ctx total); else SWA
- # caches 2*sliding_window per slot, clamped at per-slot ctx.
- swa_cells_per_slot = per_slot_ctx if swa_full else min(n_ctx, 2 * swa, per_slot_ctx)
+ if swa_full:
+ swa_cells_total = total_cells
+ else:
+ swa_limit = swa * (slots if kv_unified else 1) + ubatch
+ swa_cells_per_stream = min(cells_per_stream, swa_limit)
+ swa_cells_per_stream = _pad_kv_cells(swa_cells_per_stream)
+ swa_cells_total = swa_cells_per_stream * streams
key_len_swa = self._kv_key_length_swa or key_len
val_len_swa = self._kv_value_length_swa or val_len
+ padded_v_width = None if flash_attn else self._max_kv_value_width(val_len, val_len_swa)
if self._sliding_window_pattern is not None:
- global_bytes = 0.0 # constant across slots
- swa_bytes_per_slot = 0.0 # multiplied by slots
+ global_bytes = 0.0
+ swa_bytes = 0.0
checkpoint_extra_per_slot = 0.0
# Only layers that allocate their own KV; trailing shared layers
# reuse earlier caches.
@@ -3839,41 +4424,48 @@ class LlamaCppBackend:
layer_idx < len(self._sliding_window_pattern)
and self._sliding_window_pattern[layer_idx]
)
+ layer_key_bytes = layer_n_kv * (key_len_swa if is_swa else key_len) * bpe_k
+ layer_value_bytes = (
+ layer_n_kv * (val_len_swa if is_swa else val_len)
+ if padded_v_width is None
+ else padded_v_width
+ ) * bpe_v
+ layer_kv_bytes = layer_key_bytes + layer_value_bytes
if is_swa:
- swa_bytes_per_slot += (
- swa_cells_per_slot * layer_n_kv * (key_len_swa + val_len_swa) * bpe
- )
+ swa_bytes += swa_cells_total * layer_kv_bytes
if ctx_checkpoints > 0 and not swa_full:
- checkpoint_extra_per_slot += (
- ctx_checkpoints
- * swa
- * layer_n_kv
- * (key_len_swa + val_len_swa)
- * bpe
- )
+ checkpoint_extra_per_slot += ctx_checkpoints * swa * layer_kv_bytes
else:
- global_bytes += n_ctx * layer_n_kv * (key_len + val_len) * bpe
- return int(global_bytes + slots * (swa_bytes_per_slot + checkpoint_extra_per_slot))
+ global_bytes += total_cells * layer_kv_bytes
+ return int(global_bytes + swa_bytes + slots * checkpoint_extra_per_slot)
n_global = max(1, n_layers_kv // 4)
n_swa = n_layers_kv - n_global
- kv_per_token = n_kv * (key_len + val_len) * bpe
- kv_per_token_swa = n_kv * (key_len_swa + val_len_swa) * bpe
- global_bytes = n_global * n_ctx * kv_per_token
- swa_bytes_per_slot = n_swa * swa_cells_per_slot * kv_per_token_swa
+ global_v_width = n_kv * val_len if padded_v_width is None else padded_v_width
+ swa_v_width = n_kv * val_len_swa if padded_v_width is None else padded_v_width
+ kv_per_token = n_kv * key_len * bpe_k + global_v_width * bpe_v
+ kv_per_token_swa = n_kv * key_len_swa * bpe_k + swa_v_width * bpe_v
+ global_bytes = n_global * total_cells * kv_per_token
+ swa_bytes = n_swa * swa_cells_total * kv_per_token_swa
checkpoint_extra_per_slot = (
ctx_checkpoints * n_swa * swa * kv_per_token_swa
if ctx_checkpoints > 0 and not swa_full
else 0.0
)
- return int(global_bytes + slots * (swa_bytes_per_slot + checkpoint_extra_per_slot))
+ return int(global_bytes + swa_bytes + slots * checkpoint_extra_per_slot)
# Path 4: Standard GQA with explicit key/value dimensions
if key_len is not None and val_len is not None:
- return int(n_layers_kv * n_ctx * n_kv * (key_len + val_len) * bpe)
+ padded_v_width = None if flash_attn else self._max_kv_value_width(val_len)
+ bytes_per_cell = 0.0
+ for layer_idx in range(n_layers_kv):
+ layer_n_kv = self._kv_heads_for_layer(layer_idx, n_kv)
+ v_width = layer_n_kv * val_len if padded_v_width is None else padded_v_width
+ bytes_per_cell += layer_n_kv * key_len * bpe_k + v_width * bpe_v
+ return int(total_cells * bytes_per_cell)
# Path 5: Legacy fallback (old GGUFs without explicit dimensions)
head_dim = self._legacy_head_dim()
- return int(2 * n_kv * head_dim * n_layers_kv * n_ctx * bpe)
+ return int(2 * n_kv * head_dim * n_layers_kv * total_cells * bpe_k)
def _draft_backend_for(self, drafter_path: str) -> Optional["LlamaCppBackend"]:
"""Lightweight backend with a drafter GGUF's metadata, to size its own KV
@@ -3921,6 +4513,10 @@ class LlamaCppBackend:
draft_cache_type_k: Optional[str] = None,
draft_cache_type_v: Optional[str] = None,
n_parallel: int = 1,
+ swa_full: bool = False,
+ kv_unified: bool = True,
+ n_ubatch: Optional[int] = None,
+ flash_attn: bool = True,
) -> Optional[int]:
"""Draft KV cache bytes at n_ctx, sized from GGUF dims (K and V types are
independent). Separate drafter (Gemma): its own KV via _estimate_kv_cache_bytes
@@ -3934,12 +4530,23 @@ class LlamaCppBackend:
db = self._draft_backend_for(drafter_path)
if db is None or not db._can_estimate_kv():
return None
+ # Gemma 4 assistant layers share the target context's final global
+ # and SWA KV tensors, so only the drafter weights add memory.
+ if getattr(db, "_architecture", None) == "gemma4-assistant":
+ return 0
heavier = draft_cache_type_k if bpe_k >= bpe_v else draft_cache_type_v
- # The drafter is served under the same --parallel slot count as the
- # main model, so price its KV per slot too: a sliding-window drafter
- # (Gemma) grows KV with slots and would otherwise be under-reserved.
- kv = db._estimate_kv_cache_bytes(n_ctx, heavier, n_parallel = n_parallel)
- return kv or None
+ # The drafter uses the main model's slot and stream layout, so its
+ # compact SWA and per-stream padding must follow the same settings.
+ kv = db._estimate_kv_cache_bytes(
+ n_ctx,
+ heavier,
+ n_parallel = n_parallel,
+ swa_full = swa_full,
+ kv_unified = kv_unified,
+ n_ubatch = n_ubatch,
+ flash_attn = flash_attn,
+ )
+ return kv if kv > 0 else None
nextn = self._nextn_predict_layers or 0
n_kv = self._n_kv_heads or self._n_heads
k_len = self._kv_key_length
@@ -3953,7 +4560,14 @@ class LlamaCppBackend:
f16_bpe = _kv_bytes_per_elem("f16")
bpe_k = max(bpe_k, f16_bpe)
bpe_v = max(bpe_v, f16_bpe)
- return int(nextn * n_kv * (k_len * bpe_k + v_len * bpe_v) * n_ctx)
+ _, streams, cells_per_stream = _kv_cache_cell_layout(n_ctx, n_parallel, kv_unified)
+ v_width = n_kv * v_len
+ if not flash_attn:
+ v_width = self._max_kv_value_width(
+ v_len,
+ self._kv_value_length_swa,
+ )
+ return int(nextn * (n_kv * k_len * bpe_k + v_width * bpe_v) * cells_per_stream * streams)
def _estimate_mtp_overhead_bytes(
self,
@@ -3966,6 +4580,10 @@ class LlamaCppBackend:
draft_weights_bytes: int = 0,
n_parallel: int = 1,
mtp_keeps_target_ctx: bool = True,
+ swa_full: bool = False,
+ kv_unified: bool = True,
+ n_ubatch: Optional[int] = None,
+ flash_attn: bool = True,
) -> Optional[int]:
"""MTP draft reserve at ``n_ctx`` = draft KV (grows with ctx) + separate-
drafter weights + (MTP + MLA only) a duplicated target KV context. The
@@ -3981,6 +4599,10 @@ class LlamaCppBackend:
draft_cache_type_k = draft_cache_type_k,
draft_cache_type_v = draft_cache_type_v,
n_parallel = n_parallel,
+ swa_full = swa_full,
+ kv_unified = kv_unified,
+ n_ubatch = n_ubatch,
+ flash_attn = flash_attn,
)
weights = max(0, draft_weights_bytes)
# MLA models (GLM-5.x, DeepSeek, Kimi-K2) under MTP keep a *second* full copy
@@ -3996,7 +4618,15 @@ class LlamaCppBackend:
# rather than duplicating the target, so they must not be charged for it.
target_ctx_copy = 0
if mtp_keeps_target_ctx and self._kv_lora_rank is not None:
- target_ctx_copy = self._estimate_kv_cache_bytes(n_ctx, "f16", n_parallel = n_parallel)
+ target_ctx_copy = self._estimate_kv_cache_bytes(
+ n_ctx,
+ "f16",
+ n_parallel = n_parallel,
+ swa_full = swa_full,
+ kv_unified = kv_unified,
+ n_ubatch = n_ubatch,
+ flash_attn = flash_attn,
+ )
if draft_kv is None:
# KV unsized (exotic/remote drafter): still reserve known weights + any
# MLA target copy so a large config can't launch over budget (the small
@@ -4006,7 +4636,7 @@ class LlamaCppBackend:
return total if total > 0 else None
return draft_kv + weights + target_ctx_copy
- _DEFAULT_N_UBATCH = 512 # llama.cpp --ubatch default; Unsloth does not override it
+ _DEFAULT_N_UBATCH = _DEFAULT_LLAMA_N_UBATCH
_COMPUTE_BUFFER_SAFETY = 1.15 # upper-bound margin on the compute-buffer estimate
# Soft VRAM the modeled terms omit; charged to the fit budget on tight tiers (#6682).
_CUDA_CONTEXT_RESERVE_BYTES = 320 * 1024 * 1024 # CUDA ctx + cuBLAS workspace (~330 MiB)
@@ -4064,7 +4694,10 @@ class LlamaCppBackend:
n_embd = self._embedding_length or 0
if n_vocab <= 0 or n_embd <= 0:
return 0
- ub = max(1, int(n_ubatch if n_ubatch else self._DEFAULT_N_UBATCH))
+ ub = max(
+ 1,
+ int(self._DEFAULT_N_UBATCH if n_ubatch is None else n_ubatch),
+ )
par = max(1, int(n_parallel))
out_buffer = n_vocab * ub * 4 # f32 output/logits buffer
act_scratch = 4 * n_embd * ub * 4 # a few resident hidden-width buffers
@@ -4096,7 +4729,10 @@ class LlamaCppBackend:
n_embd = self._embedding_length or 0
if n_embd <= 0 or n_ctx <= 0:
return 0
- ub = max(1, int(n_ubatch if n_ubatch else self._DEFAULT_N_UBATCH))
+ ub = max(
+ 1,
+ int(self._DEFAULT_N_UBATCH if n_ubatch is None else n_ubatch),
+ )
if getattr(self, "_architecture", None) == "deepseek4":
# DSV4 indexer/CSA buffer (see constants): flat + linear, ub-scaled. Fires
# for any KV type -- the indexer scratch is present even with an f16 cache.
@@ -4144,6 +4780,9 @@ class LlamaCppBackend:
per_device_overhead_bytes: int,
min_gpus: int,
n_ubatch: Optional[int] = None,
+ swa_full: bool = False,
+ kv_unified: bool = True,
+ flash_attn: bool = True,
) -> tuple[Optional[list[int]], bool, int]:
"""Largest serving-slot count in [1, n_parallel) whose fully-on-GPU footprint fits,
so Unsloth keeps the model on GPU (-ngl -1) instead of --fit on, which offloads layers
@@ -4162,7 +4801,15 @@ class LlamaCppBackend:
total = (
base_footprint_bytes
+ cb
- + self._estimate_kv_cache_bytes(effective_ctx, cache_type_kv, n_parallel = slots)
+ + self._estimate_kv_cache_bytes(
+ effective_ctx,
+ cache_type_kv,
+ n_parallel = slots,
+ swa_full = swa_full,
+ kv_unified = kv_unified,
+ n_ubatch = n_ubatch,
+ flash_attn = flash_attn,
+ )
)
gpu_indices, use_fit = self._select_gpus(
total,
@@ -4187,7 +4834,9 @@ class LlamaCppBackend:
swa_full: bool = False,
n_parallel: int = 1,
kv_unified: bool = True,
+ n_ubatch: Optional[int] = None,
ctx_checkpoints: int = 0,
+ flash_attn: bool = True,
kv_on_gpu: bool = True,
mtp_engaged: bool = False,
mtp_overhead_fn: Optional[Callable[[int], int]] = None,
@@ -4224,7 +4873,9 @@ class LlamaCppBackend:
swa_full = swa_full,
n_parallel = n_parallel,
kv_unified = kv_unified,
+ n_ubatch = n_ubatch,
ctx_checkpoints = ctx_checkpoints,
+ flash_attn = flash_attn,
)
# byte-accurate mtp_overhead_fn supersedes the flat fraction (the fallback
@@ -4434,6 +5085,21 @@ class LlamaCppBackend:
LlamaCppBackend._gguf_skip_value(f, atype)
return None
+ @classmethod
+ def _gguf_path_is_diffusion(cls, gguf_path: str, model_identifier: str) -> bool:
+ """Classify a downloaded GGUF without mutating the active backend."""
+ probe = object.__new__(cls)
+ probe._model_identifier = model_identifier
+ probe._read_gguf_metadata(gguf_path)
+ return probe._is_diffusion
+
+ def _reject_vulkan_diffusion_gpu_ids_before_teardown(
+ self, gguf_path: str, model_identifier: str
+ ) -> None:
+ """Reject Vulkan + gpu_ids for diffusion GGUFs before Phase 1 teardown."""
+ if self._gguf_path_is_diffusion(gguf_path, model_identifier):
+ raise ValueError(_VULKAN_DIFFUSION_GPU_IDS_ERROR)
+
def _read_gguf_metadata(self, gguf_path: str) -> None:
"""Read context_length, architecture params, and chat_template from a GGUF header.
@@ -4846,7 +5512,7 @@ class LlamaCppBackend:
self._llama_log_path = log_dir / f"diffusion-{int(time.time())}-port-{self._port}.log"
self._llama_log_fh = open(self._llama_log_path, "w", encoding = "utf-8", buffering = 1)
logger.info(f"diffusion runner stdout/stderr -> {self._llama_log_path}")
- except OSError as e:
+ except (OSError, UnicodeDecodeError) as e:
logger.debug(f"Could not open diffusion runner log file: {e}")
# The shim (and its visual server) die with this backend process, so a
@@ -4856,7 +5522,9 @@ class LlamaCppBackend:
stdout = subprocess.PIPE,
stderr = subprocess.STDOUT,
text = True,
- env = env,
+ encoding = "utf-8",
+ errors = "replace",
+ env = utf8_child_env(env),
**_windows_hidden_subprocess_kwargs(),
**_child_popen_kwargs(),
)
@@ -4872,6 +5540,12 @@ class LlamaCppBackend:
self._is_audio = False # clear any prior TTS/audio model's routing flag
self._model_identifier = model_identifier
self._cache_type_kv = None
+ self._swa_full = False
+ self._kv_cache_unified = False
+ self._n_ubatch = self._DEFAULT_N_UBATCH
+ self._flash_attn_enabled = True
+ self._effective_cache_types = ("f16", "f16")
+ self._kv_cache_context_total = None
self._gpu_offload_active = True
# Diffusion doesn't use the llama.cpp GPU-memory knobs; reset them to
# defaults (the picked device is still recorded below) so /load, /status
@@ -4885,11 +5559,14 @@ class LlamaCppBackend:
# 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.
+ # The single-device runner records only the lowest selected GPU (chosen
+ # above), not the whole pick, and clears any explicit pin from a prior
+ # chat load; a multi-GPU list would misreport placement and mis-dedup.
self._gpu_ids = [sorted(gpu_ids)[0]] if gpu_ids else None
+ # The frontend prefers requested_gpu_ids when hydrating the picker.
+ # Diffusion uses only one device, so echo the collapsed effective pin,
+ # not unused members of the original request.
+ self._requested_gpu_ids = list(self._gpu_ids) if self._gpu_ids else None
if hf_variant:
self._hf_variant = hf_variant
elif gguf_path:
@@ -4955,6 +5632,9 @@ class LlamaCppBackend:
touching the shared one; defaults to the shared event.
"""
cancel_event = cancel_event if cancel_event is not None else self._cancel_event
+ from utils.hf_cache_settings import get_hf_cache_paths
+
+ download_cache_dir = str(get_hf_cache_paths().hub_cache)
try:
import huggingface_hub # noqa: F401 -- presence check only
except ImportError:
@@ -5050,7 +5730,11 @@ class LlamaCppBackend:
if not p.size:
continue
try:
- cached_path = try_to_load_from_cache(hf_repo, p.path)
+ cached_path = try_to_load_from_cache(
+ hf_repo,
+ p.path,
+ cache_dir = download_cache_dir,
+ )
except Exception:
cached_path = None
if (
@@ -5061,6 +5745,7 @@ class LlamaCppBackend:
hf_repo,
p.path,
expected_size = p.size,
+ cache_dir = download_cache_dir,
)
if isinstance(cached_path, str) and os.path.exists(cached_path):
try:
@@ -5074,12 +5759,8 @@ class LlamaCppBackend:
total_download_bytes = max(0, total_bytes - already_cached_bytes)
if total_download_bytes > 0:
- cache_dir = os.environ.get(
- "HF_HUB_CACHE",
- str(Path.home() / ".cache" / "huggingface" / "hub"),
- )
- Path(cache_dir).mkdir(parents = True, exist_ok = True)
- free_bytes = shutil.disk_usage(cache_dir).free
+ Path(download_cache_dir).mkdir(parents = True, exist_ok = True)
+ free_bytes = shutil.disk_usage(download_cache_dir).free
total_gb = total_download_bytes / (1024**3)
free_gb = free_bytes / (1024**3)
@@ -5097,7 +5778,7 @@ class LlamaCppBackend:
# surface the disk shortfall for the requested variant.
raise RuntimeError(
f"Not enough disk space to download {gguf_filename}. "
- f"Only {free_gb:.1f} GB free in {cache_dir}"
+ f"Only {free_gb:.1f} GB free in {download_cache_dir}"
)
smaller = self._find_smallest_fitting_variant(
hf_repo,
@@ -5128,7 +5809,7 @@ class LlamaCppBackend:
else:
raise RuntimeError(
f"Not enough disk space to download any variant. "
- f"Only {free_gb:.1f} GB free in {cache_dir}"
+ f"Only {free_gb:.1f} GB free in {download_cache_dir}"
)
except RuntimeError:
raise
@@ -5151,6 +5832,7 @@ class LlamaCppBackend:
cancel_event = cancel_event,
on_status = lambda m: logger.info(m),
force_download = force,
+ cache_dir = download_cache_dir,
)
for shard in gguf_extra_shards:
if cancel_event.is_set():
@@ -5162,6 +5844,7 @@ class LlamaCppBackend:
hf_token,
cancel_event = cancel_event,
force_download = force,
+ cache_dir = download_cache_dir,
)
except Exception as e:
if isinstance(e, RuntimeError) and "Cancelled" in str(e):
@@ -5207,6 +5890,12 @@ class LlamaCppBackend:
logger.info("Reusing cached %s: %s", label, cached)
return cached
+ from utils.hf_cache_settings import get_hf_cache_paths
+
+ companion_cache_dir = _hub_cache_dir_for_snapshot_path(near_path) or str(
+ get_hf_cache_paths().hub_cache
+ )
+
if _hub_download_in_flight(hf_repo):
logger.info("Skipping %s download while a hub download is active", label)
return None
@@ -5241,7 +5930,7 @@ class LlamaCppBackend:
if target is None:
try:
from utils.models.model_config import _iter_hf_cache_snapshots
- for snap in _iter_hf_cache_snapshots(hf_repo):
+ for snap in _iter_hf_cache_snapshots(hf_repo, companion_cache_dir):
rel_files = _gguf_snapshot_files(snap)
target = pick(rel_files)
if target is not None:
@@ -5259,7 +5948,11 @@ class LlamaCppBackend:
# hf_hub_download with hf_repo would miss the canonical file and silently
# drop the companion. _cached_hf_snapshot_file scans every case variant.
if _hf_env_offline():
- cached = _cached_hf_snapshot_file(hf_repo, target)
+ cached = _cached_hf_snapshot_file(
+ hf_repo,
+ target,
+ cache_dir = companion_cache_dir,
+ )
if cached:
logger.info("Resolved %s from local HF cache: %s", label, cached)
return cached
@@ -5272,6 +5965,7 @@ class LlamaCppBackend:
target,
hf_token,
cancel_event = cancel_event,
+ cache_dir = companion_cache_dir,
)
except Exception as e:
logger.warning(f"Could not download {label}: {e}")
@@ -5302,7 +5996,12 @@ class LlamaCppBackend:
near_path = near_path,
)
- def _cached_repo_mtp_drafter(self, hf_repo: str) -> Optional[str]:
+ def _cached_repo_mtp_drafter(
+ self,
+ hf_repo: str,
+ *,
+ cache_dir: Optional[str] = None,
+ ) -> Optional[str]:
"""A drafter already in this repo's local HF cache, reused offline when a
fresh copy can't be fetched. Prefers a repo-root ``mtp-*.gguf`` across all
cached snapshots; else an existing ``MTP/`` copy (any precision -- the
@@ -5312,7 +6011,12 @@ class LlamaCppBackend:
roots: list[Path] = []
subdirs: list[Path] = []
- for snap in _iter_hf_cache_snapshots(hf_repo): # newest first
+ snapshots = (
+ _iter_hf_cache_snapshots(hf_repo)
+ if cache_dir is None
+ else _iter_hf_cache_snapshots(hf_repo, cache_dir)
+ )
+ for snap in snapshots: # newest first
for f in sorted(_gguf_snapshot_files(snap)):
if _is_companion_gguf_path(f) and "mmproj" not in f.lower():
(roots if "/" not in f else subdirs).append(snap / f)
@@ -5365,7 +6069,10 @@ class LlamaCppBackend:
# current cached file and refetch a changed one, so skip the probe here
# rather than pair new weights with a stale draft.
if _hf_env_offline():
- cached = self._cached_repo_mtp_drafter(hf_repo)
+ cached = self._cached_repo_mtp_drafter(
+ hf_repo,
+ cache_dir = _hub_cache_dir_for_snapshot_path(near_path),
+ )
if cached:
logger.info(f"Reusing cached MTP drafter (offline): {cached}")
return cached
@@ -5580,6 +6287,9 @@ class LlamaCppBackend:
total_by_idx: Optional[dict[int, int]] = None,
n_ubatch: Optional[int] = None,
soft_overhead_bytes: int = 0,
+ swa_full: bool = False,
+ kv_unified: bool = True,
+ flash_attn: bool = True,
) -> tuple[int, int, list[int], Optional[list[int]]]:
"""Plan a ``--split-mode tensor`` load. Pure: no model or GPU needed.
@@ -5667,6 +6377,17 @@ class LlamaCppBackend:
def _mtp_at(ctx: int) -> int:
return mtp_overhead_fn(ctx) if mtp_overhead_fn is not None else 0
+ def _kv_at(ctx: int) -> int:
+ return self._estimate_kv_cache_bytes(
+ ctx,
+ cache_type_kv,
+ n_parallel = n_parallel,
+ swa_full = swa_full,
+ kv_unified = kv_unified,
+ n_ubatch = n_ubatch,
+ flash_attn = flash_attn,
+ )
+
# Context-linear compute buffer, summed over the split. Tensor mode
# replicates the compute graph on EVERY device (measured: the per-device
# buffer grows a flat n_ubatch*2 bytes/token, ~1024 B/tok on Qwen3.5-9B at
@@ -5692,31 +6413,21 @@ class LlamaCppBackend:
# Weights + buffers exceed the pool -> floor; the load then
# falls back to layer split.
return ctx_floor
- if mtp_overhead_fn is not None:
- # kv(ctx)+mtp(ctx)+compute(ctx) is not single-linear, so binary search.
- def _consumer(c: int) -> int:
- return (
- self._estimate_kv_cache_bytes(c, cache_type_kv, n_parallel = n_parallel)
- + _mtp_at(c)
- + _cc_ctx(c)
- )
- if _consumer(ctx) <= kv_budget_b:
- return ctx
- lo, hi, best = ctx_floor, ctx, ctx_floor
- while lo <= hi:
- mid = (lo + hi) // 2
- if _consumer(mid) <= kv_budget_b:
- best = mid
- lo = mid + 1
- else:
- hi = mid - 1
- return best
- kv_at = self._estimate_kv_cache_bytes(ctx, cache_type_kv, n_parallel = n_parallel)
- total_at = kv_at + _cc_ctx(ctx) # both ~linear through the origin
- if total_at <= kv_budget_b:
+ def _consumer(c: int) -> int:
+ return _kv_at(c) + _mtp_at(c) + _cc_ctx(c)
+
+ if _consumer(ctx) <= kv_budget_b:
return ctx
- return max(ctx_floor, int(ctx * kv_budget_b / total_at))
+ lo, hi, best = ctx_floor, ctx, ctx_floor
+ while lo <= hi:
+ mid = (lo + hi) // 2
+ if _consumer(mid) <= kv_budget_b:
+ best = mid
+ lo = mid + 1
+ else:
+ hi = mid - 1
+ return best
# KV size unknown -> can't prove a safe cap; floor.
return min(4096, ctx) if ctx > 0 else 4096
@@ -5728,11 +6439,7 @@ class LlamaCppBackend:
effective_ctx = min(_fit_ctx(target_ctx), max_available_ctx)
min_usable_mib = min(usable_by_idx.values())
- kv_bytes = (
- self._estimate_kv_cache_bytes(effective_ctx, cache_type_kv, n_parallel = n_parallel)
- if (self._can_estimate_kv() and effective_ctx > 0)
- else 0
- )
+ kv_bytes = _kv_at(effective_ctx) if (self._can_estimate_kv() and effective_ctx > 0) else 0
# The MTP reserve also has to fit the even split (mirror the pooled budget):
# byte-accurate per-ctx (0 when no fn) plus the same flat cushion as above.
mtp_bytes = (_mtp_at(effective_ctx) if effective_ctx > 0 else 0) + flat_mtp_bytes
@@ -5783,6 +6490,24 @@ class LlamaCppBackend:
and ("unknown" in text or "unsupported" in text or "not supported" in text)
)
+ @staticmethod
+ def _mmproj_retry_failure_message(*, projector_confirmed: bool, detail: str) -> str:
+ """User-facing error when the text-only --mmproj strip retry also fails.
+
+ Confirmed projector-format mismatches keep the historical wording.
+ Bare signal crashes (common on some ROCm/driver paths) must not be
+ reported as "Vision projector incompatible" — that misled #7302.
+ """
+ if projector_confirmed:
+ return (
+ "Vision projector incompatible with this llama.cpp "
+ "build, and the text-only retry also failed: " + detail
+ )
+ return (
+ "Vision model failed to start (llama-server crashed with "
+ "--mmproj), and the text-only retry also failed: " + detail
+ )
+
@staticmethod
def _output_has_nonprojector_diagnostic(output: str) -> bool:
"""True when the output already names a concrete non-projector cause (out
@@ -5851,28 +6576,98 @@ class LlamaCppBackend:
def explicit(i):
nxt = out[i + 1] if i + 1 < len(out) else None
- return nxt if nxt in ("on", "auto", "off") else None
+ return nxt if nxt in _LLAMA_ARG_TRUE_FALSE_AUTO_VALUES else None
effective = None
for i, tok in enumerate(out):
- if tok.startswith(("--flash-attn=", "-fa=")):
+ name = _flag_name(tok)
+ if name in ("--flash-attn", "-fa") and "=" in tok:
effective = tok.partition("=")[2]
- elif tok in ("--flash-attn", "-fa"):
+ elif name in ("--flash-attn", "-fa"):
effective = explicit(i) or "on"
- if effective not in ("on", "auto"):
+ if effective not in _LLAMA_ARG_TRUE_OR_AUTO_VALUES:
return None
for i, tok in enumerate(out):
- if tok.startswith(("--flash-attn=", "-fa=")):
+ name = _flag_name(tok)
+ if name in ("--flash-attn", "-fa") and "=" in tok:
flag, _, value = tok.partition("=")
- if value in ("on", "auto"):
+ if value in _LLAMA_ARG_TRUE_OR_AUTO_VALUES:
out[i] = f"{flag}=off"
- elif tok in ("--flash-attn", "-fa"):
- if explicit(i) in ("on", "auto"):
+ elif name in ("--flash-attn", "-fa"):
+ if explicit(i) in _LLAMA_ARG_TRUE_OR_AUTO_VALUES:
out[i + 1] = "off"
elif explicit(i) is None: # bare flag (reads as on) -> explicit off
out[i] = f"{tok}=off"
+
+ # A quantized V cache requires flash attention in llama.cpp: the init
+ # aborts with "V cache quantization requires flash_attn". A quantized K
+ # cache has no such requirement and runs fine without FA, so it is left
+ # untouched -- resetting it would needlessly enlarge the K cache and can
+ # OOM a memory-constrained config. Studio launches with FA on, so a
+ # quantized --cache-type-v is legal at launch but would make THIS FA-off
+ # retry crash on init instead of recovering. Reset a quantized V cache --
+ # main and draft (the draft context shares the global --flash-attn flag,
+ # so its V cache aborts too) -- to f16 (the llama.cpp default);
+ # non-quantized types -- f16/bf16/f32 -- run fine without FA and are left
+ # untouched. The value is rewritten in place so the list length is
+ # preserved for downstream slices, matching the flash-attn flip above.
+ _v_cache_flags = (
+ "--cache-type-v",
+ "-ctv",
+ "--cache-type-v-draft",
+ "--spec-draft-type-v",
+ "-ctvd",
+ )
+ _cache_reset = False
+ for i, tok in enumerate(out):
+ # llama.cpp rewrites '_' to '-' for any argv token starting with
+ # '--' before matching, so a legal pass-through spelling such as
+ # --cache_type_v parses as --cache-type-v and still enables a
+ # quantized V cache. Canonicalize the flag name the same way so the
+ # reset recognizes the underscore aliases too; short flags (-ctv)
+ # and the type value are left untouched.
+ name = _flag_name(tok)
+ if name not in _v_cache_flags:
+ continue
+ if "=" in tok:
+ flag, _, value = tok.partition("=")
+ if value.strip().lower() not in LlamaCppBackend._NON_QUANTIZED_KV_TYPES:
+ out[i] = f"{flag}=f16"
+ _cache_reset = True
+ elif i + 1 < len(out):
+ if out[i + 1].strip().lower() not in LlamaCppBackend._NON_QUANTIZED_KV_TYPES:
+ out[i + 1] = "f16"
+ _cache_reset = True
+ if _cache_reset:
+ logger.info(
+ "V cache dtype reset to f16 because flash attention was disabled "
+ "by the crash-recovery fallback (quantized V cache requires flash "
+ "attention in llama.cpp; the K cache is left untouched)."
+ )
return out
+ @staticmethod
+ def _drop_env_quantized_v_cache(env: MutableMapping[str, str]) -> bool:
+ """Drop an inherited quantized V-cache env var (main or draft) in place
+ before a flash-attn-off retry, returning True if anything was removed.
+
+ The argv rewrite in ``_with_flash_attn_off`` only reaches flags on the
+ command line. Studio deliberately lets an env-only cache type reach the
+ child untouched (an asymmetric K/V env must survive), so a quantized V
+ cache set purely through ``LLAMA_ARG_CACHE_TYPE_V`` (or the draft
+ ``LLAMA_ARG_SPEC_DRAFT_CACHE_TYPE_V``) would still abort the FA-off retry
+ with "V cache quantization requires flash_attn". Dropping it lets
+ llama.cpp fall back to the f16 default. Only V is dropped: a quantized K
+ cache runs fine without flash attention, so its env var is preserved.
+ """
+ dropped = False
+ for var in ("LLAMA_ARG_CACHE_TYPE_V", "LLAMA_ARG_SPEC_DRAFT_CACHE_TYPE_V"):
+ value = (env.get(var) or "").strip().lower()
+ if value and value not in LlamaCppBackend._NON_QUANTIZED_KV_TYPES:
+ env.pop(var, None)
+ dropped = True
+ return dropped
+
@staticmethod
def _strip_mmproj_args(cmd: list[str]) -> list[str]:
"""Return cmd without the '--mmproj ' pair (text-only retry).
@@ -5929,7 +6724,7 @@ class LlamaCppBackend:
buffering = 1,
)
logger.info(f"llama-server stdout/stderr -> {self._llama_log_path}")
- except OSError as e:
+ except (OSError, UnicodeDecodeError) as e:
# Best-effort; never block the load on logging.
logger.debug(f"Could not open llama-server log file: {e}")
self._llama_log_path = None
@@ -5943,6 +6738,8 @@ class LlamaCppBackend:
stdout = subprocess.PIPE,
stderr = subprocess.STDOUT,
text = True,
+ encoding = "utf-8",
+ errors = "replace",
env = env,
**_windows_hidden_subprocess_kwargs(),
**_child_popen_kwargs(),
@@ -5984,6 +6781,8 @@ class LlamaCppBackend:
gpu_layers: int = -1,
n_cpu_moe: int = 0,
tensor_split: Optional[List[float]] = None,
+ # Explicit GPU placement pool (issue #7164). None/[] = auto-select;
+ # the fitter may pin the smallest subset of this pool that fits.
gpu_ids: Optional[List[int]] = None,
n_threads: Optional[int] = None,
n_gpu_layers: Optional[int] = None, # caller compat, unused
@@ -6059,6 +6858,7 @@ class LlamaCppBackend:
chat_template_override = chat_template_override,
extra_args = extra_args,
is_vision = is_vision,
+ n_parallel = n_parallel,
preserve_multi_gpu_on_layer = preserve_multi_gpu_on_layer,
):
logger.info(
@@ -6081,15 +6881,80 @@ class LlamaCppBackend:
self._cancel_event.clear()
- # ── Phase 1: kill old process (under lock, fast) ──────────
- with self._lock:
- self._kill_process()
-
# Resolve llama-server now but defer a not-found error: a block-diffusion
# GGUF uses the diffusion runner, and its arch is only known after the header.
binary = self._find_llama_server_binary()
is_vulkan_backend = self._is_vulkan_backend(binary)
+ # Without --kv-unified an explicit --parallel N splits -c into windows of -c/N, so on a
+ # build lacking the flag the default of 4 would quarter every context window for a
+ # feature it cannot serve: fall back to one slot. Ahead of the KV estimates so the
+ # fit matches what launches.
+ if (
+ n_parallel > 1
+ and binary
+ and not self.probe_server_capabilities(binary).get("supports_kv_unified")
+ ):
+ logger.warning(
+ "llama-server at %s has no --kv-unified, so %d parallel slots would "
+ "split the context window %d ways. Using 1 slot instead; update "
+ "llama.cpp to run chats in parallel.",
+ binary,
+ n_parallel,
+ n_parallel,
+ )
+ n_parallel = 1
+
+ # ── Vulkan-ordinal preflight (BEFORE the Phase 1 kill) ────────
+ # An explicit Vulkan pin the ggml probe never enumerated cannot be honored.
+ # Validate it ABOVE the kill so an invalid selection leaves the live model
+ # untouched: CUDA ids are range-checked at the route, but Vulkan ordinals are
+ # not, so a stale gpu_ids=[99] used to kill the server then 400, leaving
+ # nothing running (#7239). _get_gpu_memory needs only the binary (safe pre-
+ # download) and reuses the later fit's issubset logic. Guarded on a found
+ # Vulkan build + a pin so a deferred not-found stays deferred for diffusion.
+ if is_vulkan_backend and gpu_ids and binary:
+ _pf_wanted = {int(x) for x in gpu_ids}
+ _pf_probed = {g[0] for g in self._get_gpu_memory(binary)}
+ if not _pf_wanted.issubset(_pf_probed):
+ raise ValueError(
+ f"Requested Vulkan GPU ordinal(s) {sorted(_pf_wanted)} not "
+ f"present. Available Vulkan devices: {sorted(_pf_probed)}."
+ )
+
+ # Classify before killing the healthy server (#7205); Phase 2 reuses this path.
+ _preflight_model_path = None
+ if is_vulkan_backend and gpu_ids and hf_repo:
+ _resolved_repo = _resolve_repo_id_casing(hf_repo)
+ if _resolved_repo != hf_repo:
+ logger.info(
+ "Using cached repo_id casing '%s' for requested '%s'",
+ _resolved_repo,
+ hf_repo,
+ )
+ hf_repo = _resolved_repo
+ with _hf_offline_if_dns_dead():
+ _preflight_model_path = self._download_gguf(
+ hf_repo = hf_repo,
+ hf_variant = hf_variant,
+ hf_token = hf_token,
+ )
+ self._reject_vulkan_diffusion_gpu_ids_before_teardown(
+ _preflight_model_path,
+ model_identifier,
+ )
+ elif is_vulkan_backend and gpu_ids and gguf_path and not hf_repo:
+ if not Path(gguf_path).is_file():
+ raise FileNotFoundError(f"GGUF file not found: {gguf_path}")
+ self._reject_vulkan_diffusion_gpu_ids_before_teardown(
+ gguf_path,
+ model_identifier,
+ )
+
+ # ── Phase 1: kill old process (under lock, fast) ──────────
+ with self._lock:
+ self._kill_process()
+
# ── Phase 2: download (NO lock held, so cancel can proceed) ──
# mtp_draft_path arrives set for local Gemma loads (detected
# sibling); for -hf loads it's None here and resolved just below.
@@ -6111,7 +6976,7 @@ class LlamaCppBackend:
)
hf_repo = _resolved_repo
with _hf_offline_if_dns_dead():
- model_path = self._download_gguf(
+ model_path = _preflight_model_path or self._download_gguf(
hf_repo = hf_repo,
hf_variant = hf_variant,
hf_token = hf_token,
@@ -6164,6 +7029,20 @@ class LlamaCppBackend:
# Not a tensor/layer GGUF: clear any preserved-fallback flag from a
# prior load (this path skips the command builder that clears it).
self._layer_preserves_tensor_intent = False
+ # On a Vulkan build gpu_ids are ggml Vulkan ordinals, but the diffusion
+ # runner selects its device by CUDA physical index (_diffusion_gpu_arg
+ # forwards gpu_ids[0] as a CUDA/DG_GPU token) with no mapping to them.
+ # The route rejects a CONFIRMED-diffusion pick up front; an uncached GGUF
+ # only classified as diffusion post-download still reaches here with a
+ # pin, so drop it and serve on the default device (like an unpinned load).
+ if gpu_ids and is_vulkan_backend:
+ logger.warning(
+ "Ignoring gpu_ids %s for diffusion GGUF on a Vulkan build: "
+ "the diffusion runner cannot map ggml Vulkan ordinals; "
+ "serving on the default device.",
+ gpu_ids,
+ )
+ gpu_ids = None
with self._lock:
if self._cancel_event.is_set():
logger.info("Load cancelled before diffusion server start")
@@ -6195,6 +7074,8 @@ class LlamaCppBackend:
# same message remote validation already shows.
raise LlamaServerNotFoundError(LLAMA_SERVER_NOT_FOUND_DETAIL)
+ server_caps = self.probe_server_capabilities(binary)
+
# Outside ``self._lock`` so /unload, /cancel, /status aren't
# blocked. ``unload_model`` also records the kill, so the
# frontend /unload+/load Apply path engages the wait here even
@@ -6215,6 +7096,18 @@ class LlamaCppBackend:
# state to publish.
ctx_override = parse_ctx_override(extra_args)
requested_ctx = resolve_requested_ctx(extra_args, n_ctx)
+ swa_full = _swa_full_from_args_or_env(extra_args)
+ _effective_ubatch = _extra_args_n_ubatch(
+ extra_args,
+ n_ctx = (requested_ctx if requested_ctx > 0 else self._context_length),
+ )
+ planned_kv_unified = _kv_unified_from_args(
+ extra_args,
+ default = n_parallel > 1 and server_caps.get("supports_kv_unified", False),
+ )
+ # A hard-crash recovery may relaunch this same plan with FA off.
+ # Size that larger cache up front so the recovery cannot OOM.
+ planned_flash_attn = False
cache_override = parse_cache_override(extra_args)
# Budget the heavier of asymmetric --cache-type-k/-v extras (they
# win per axis at launch, appended last); resolve_cache_type_kv only
@@ -6381,6 +7274,12 @@ class LlamaCppBackend:
# Layer-fallback min GPUs; raised below on a tensor downgrade. Bound
# before the try so the --fit-on except path still has it (no UnboundLocal).
_layer_min_gpus = 1
+ # An explicit Vulkan ordinal absent from the ggml probe cannot be
+ # honored; flag it in the fit and reject after the try (raising inside
+ # would be swallowed into the --fit-on fallback). Bound before the try.
+ _vulkan_explicit_unmatched = False
+ _vulkan_requested_ids: list[int] = []
+ _vulkan_available_ordinals: list[int] = []
try:
gguf_size = self._get_gguf_size_bytes(model_path)
# Include GPU-loaded mmproj in the fit budget (#5825).
@@ -6393,6 +7292,28 @@ class LlamaCppBackend:
# Pass binary so a Vulkan build probes ggml's Vulkan ordinals.
_gpu_mem = self._get_gpu_memory(binary)
gpus = [(idx, free) for idx, free, _t in _gpu_mem]
+ # Restrict the fit (and thus the layer plan + pin env) to the
+ # selected GPUs; fail-open if none match so a stale UI choice
+ # can't strand the load on CPU (issue #7164).
+ if gpu_ids:
+ # A Vulkan build indexes by ggml ordinal. An explicit ordinal
+ # absent from the probe can't be pinned, so reject after the try
+ # rather than fail-open onto a device the user didn't pick.
+ _wanted_ids = {int(x) for x in gpu_ids}
+ # Reject if ANY requested ordinal is absent, not only when none
+ # match: [0, 99] against {0, 1} silently drops 99. Comparing the
+ # full requested set (before filter narrows) still lets the fitter
+ # pick a valid subset later -- that is narrowing, not absence.
+ _probed_ordinals = {g[0] for g in gpus}
+ if is_vulkan_backend and not _wanted_ids.issubset(_probed_ordinals):
+ _vulkan_explicit_unmatched = True
+ _vulkan_requested_ids = sorted(_wanted_ids)
+ _vulkan_available_ordinals = sorted(_probed_ordinals)
+ # Restrict the probed pool to the selection; fail-open (keep the
+ # full pool) if none match so a stale UI choice can't strand the
+ # load on CPU (issue #7164).
+ _sel_gpus = [g for g in gpus if g[0] in _wanted_ids]
+ gpus = _sel_gpus if _sel_gpus else gpus
total_by_idx = {idx: total for idx, _f, total in _gpu_mem}
# GPU picker: restrict every mode to the chosen devices, so
# auto selection only considers them and manual mask to
@@ -6617,6 +7538,10 @@ class LlamaCppBackend:
draft_cache_type_k = _mtp_draft_ck,
draft_cache_type_v = _mtp_draft_cv,
n_parallel = n_parallel,
+ swa_full = swa_full,
+ kv_unified = planned_kv_unified,
+ n_ubatch = _effective_ubatch,
+ flash_attn = planned_flash_attn,
)
if (
self._estimate_mtp_overhead_bytes(
@@ -6628,6 +7553,10 @@ class LlamaCppBackend:
draft_weights_bytes = _mtp_draft_weights,
n_parallel = n_parallel,
mtp_keeps_target_ctx = _engaged_is_mtp,
+ swa_full = swa_full,
+ kv_unified = planned_kv_unified,
+ n_ubatch = _effective_ubatch,
+ flash_attn = planned_flash_attn,
)
is not None
):
@@ -6644,6 +7573,10 @@ class LlamaCppBackend:
_w: int = _mtp_draft_weights,
_np: int = n_parallel,
_mtp: bool = _engaged_is_mtp,
+ _swa_full: bool = swa_full,
+ _kv_unified: bool = planned_kv_unified,
+ _n_ubatch: Optional[int] = _effective_ubatch,
+ _flash_attn: bool = planned_flash_attn,
) -> int:
v = self._estimate_mtp_overhead_bytes(
ctx,
@@ -6654,15 +7587,26 @@ class LlamaCppBackend:
draft_weights_bytes = _w,
n_parallel = _np,
mtp_keeps_target_ctx = _mtp,
+ swa_full = _swa_full,
+ kv_unified = _kv_unified,
+ n_ubatch = _n_ubatch,
+ flash_attn = _flash_attn,
)
return v if v is not None else 0
def _mtp_bytes(ctx: int) -> int:
return mtp_overhead_fn(ctx) if mtp_overhead_fn is not None else 0
- # Effective micro-batch (a user --ubatch override scales the
- # compute buffer); None -> the 512 default in the estimate.
- _effective_ubatch = _extra_args_n_ubatch(extra_args)
+ def _kv_bytes(ctx: int) -> int:
+ return self._estimate_kv_cache_bytes(
+ ctx,
+ cache_type_kv,
+ n_parallel = n_parallel,
+ swa_full = swa_full,
+ kv_unified = planned_kv_unified,
+ n_ubatch = _effective_ubatch,
+ flash_attn = planned_flash_attn,
+ )
def _cc_bytes(ctx: int, n_gpus: int = 1) -> int:
# Context-linear compute-buffer growth (flash-attn KQ mask +
@@ -6902,6 +7846,9 @@ class LlamaCppBackend:
total_by_idx = total_by_idx,
n_ubatch = _effective_ubatch,
soft_overhead_bytes = _soft_overhead,
+ swa_full = swa_full,
+ kv_unified = planned_kv_unified,
+ flash_attn = planned_flash_attn,
)
use_fit = False
elif gpus and self._can_estimate_kv() and effective_ctx > 0:
@@ -6934,16 +7881,18 @@ class LlamaCppBackend:
pool_budget,
_ms,
cache_type_kv,
+ swa_full = swa_full,
n_parallel = n_parallel,
+ kv_unified = planned_kv_unified,
+ n_ubatch = _effective_ubatch,
+ flash_attn = planned_flash_attn,
mtp_engaged = _mtp_reserves_gpu,
mtp_overhead_fn = mtp_overhead_fn,
compute_ctx_bytes_fn = _cc_sub,
budget_frac = 1.0,
total_mib = None,
)
- kv = self._estimate_kv_cache_bytes(
- capped, cache_type_kv, n_parallel = n_parallel
- )
+ kv = _kv_bytes(capped)
footprint_mib = (
_ms + kv + _mtp_bytes(capped) + _cc_sub(capped)
) / (1024 * 1024)
@@ -6963,9 +7912,7 @@ class LlamaCppBackend:
# on and let llama-server flex -ngl (CPU offload).
requested_total = (
model_size_fit
- + self._estimate_kv_cache_bytes(
- effective_ctx, cache_type_kv, n_parallel = n_parallel
- )
+ + _kv_bytes(effective_ctx)
+ _mtp_bytes(effective_ctx)
+ _cc_bytes(effective_ctx)
)
@@ -7017,16 +7964,18 @@ class LlamaCppBackend:
pool_budget,
_ms,
cache_type_kv,
+ swa_full = swa_full,
n_parallel = n_parallel,
+ kv_unified = planned_kv_unified,
+ n_ubatch = _effective_ubatch,
+ flash_attn = planned_flash_attn,
mtp_engaged = _mtp_reserves_gpu,
mtp_overhead_fn = mtp_overhead_fn,
compute_ctx_bytes_fn = _cc_sub,
budget_frac = 1.0,
total_mib = None,
)
- kv = self._estimate_kv_cache_bytes(
- capped, cache_type_kv, n_parallel = n_parallel
- )
+ kv = _kv_bytes(capped)
footprint_mib = (
_ms + kv + _mtp_bytes(capped) + _cc_sub(capped)
) / (1024 * 1024)
@@ -7043,11 +7992,7 @@ class LlamaCppBackend:
if effective_ctx > 0:
for n_gpus in range(_auto_min_gpus, len(ranked) + 1):
subset = ranked[:n_gpus]
- kv = self._estimate_kv_cache_bytes(
- effective_ctx,
- cache_type_kv,
- n_parallel = n_parallel,
- )
+ kv = _kv_bytes(effective_ctx)
footprint_mib = (
_subset_model_size(n_gpus)
+ kv
@@ -7104,7 +8049,11 @@ class LlamaCppBackend:
_apple_fit_budget_mib,
model_size_fit,
cache_type_kv,
+ swa_full = swa_full,
n_parallel = n_parallel,
+ kv_unified = planned_kv_unified,
+ n_ubatch = _effective_ubatch,
+ flash_attn = planned_flash_attn,
mtp_engaged = _mtp_reserves_gpu,
mtp_overhead_fn = mtp_overhead_fn,
compute_ctx_bytes_fn = _cc_bytes,
@@ -7112,12 +8061,7 @@ class LlamaCppBackend:
total_mib = None,
)
_cap_footprint_mib = (
- model_size_fit
- + self._estimate_kv_cache_bytes(
- cap, cache_type_kv, n_parallel = n_parallel
- )
- + _mtp_bytes(cap)
- + _cc_bytes(cap)
+ model_size_fit + _kv_bytes(cap) + _mtp_bytes(cap) + _cc_bytes(cap)
) / (1024 * 1024)
# Fit returns the request unchanged when it fits OR weights
# exceed budget; only the latter over-commits, so floor to 4096.
@@ -7164,6 +8108,9 @@ class LlamaCppBackend:
_pipeline_overhead_bytes + _cc_bytes(effective_ctx),
_layer_min_gpus,
_effective_ubatch,
+ swa_full = swa_full,
+ kv_unified = planned_kv_unified,
+ flash_attn = planned_flash_attn,
)
if not _uf_slots:
logger.info(
@@ -7188,9 +8135,7 @@ class LlamaCppBackend:
_mtp_note = ""
if effective_ctx < original_ctx:
- kv_est = self._estimate_kv_cache_bytes(
- effective_ctx, cache_type_kv, n_parallel = n_parallel
- )
+ kv_est = _kv_bytes(effective_ctx)
logger.info(
f"Context auto-reduced: {original_ctx} -> {effective_ctx} "
f"(model: {model_size / (1024**3):.1f} GB, "
@@ -7199,9 +8144,7 @@ class LlamaCppBackend:
+ ")"
)
- kv_cache_bytes = self._estimate_kv_cache_bytes(
- effective_ctx, cache_type_kv, n_parallel = n_parallel
- )
+ kv_cache_bytes = _kv_bytes(effective_ctx)
mmproj_note = (
f"mmproj: {mmproj_size / (1024**3):.1f} GB, " if mmproj_size else ""
)
@@ -7219,6 +8162,17 @@ class LlamaCppBackend:
tp_tensor_split = None
effective_ctx = requested_ctx # fall back to original
+ # An unenumerated explicit Vulkan ordinal can't be pinned; fail loudly
+ # instead of fitting onto an unselected device. Clear the raw selection
+ # the early state-publish recorded so it never leaks into gpu_ids (#7239).
+ if _vulkan_explicit_unmatched:
+ self._gpu_ids = None
+ self._requested_gpu_ids = None
+ raise ValueError(
+ f"Requested Vulkan GPU ordinal(s) {_vulkan_requested_ids} not "
+ f"present. Available Vulkan devices: {_vulkan_available_ordinals}."
+ )
+
# GPU picker: when no narrower subset was chosen (manual, or
# a failed/file-size selection), pin the whole picked set so the
# model can't spill onto an unpicked GPU.
@@ -7355,7 +8309,6 @@ class LlamaCppBackend:
cmd.extend(["-ngl", "-1", "--fit", "off"])
fully_gpu_offloaded = True
- server_caps = self.probe_server_capabilities(binary)
# Expose Prometheus /metrics for the engine-stats logger, only
# when the binary advertises it (older/custom binaries may not).
if server_caps.get("supports_metrics"):
@@ -7427,6 +8380,11 @@ class LlamaCppBackend:
"iq4_nl",
"f32",
}
+ # Normalize like the budget does (_planned_main_cache_types): a
+ # case-sensitive match drops "Q8_0", emitting no flag, so llama.cpp
+ # runs f16 while the estimate priced q8_0. Emit the normalized
+ # spelling; kv_cache_type_from_str is case-sensitive.
+ cache_type_kv = cache_type_kv.strip().lower() if cache_type_kv else cache_type_kv
if (
cache_type_kv
and cache_type_kv in _valid_cache_types
@@ -7582,11 +8540,45 @@ class LlamaCppBackend:
", ".join(unsupported_cache_flags),
)
- # Vulkan pins via --device (a cmd arg, unlike the env-based
- # CUDA/ROCm pin below), emitted BEFORE user extras so llama.cpp's
- # last-wins parsing lets a user --device override Unsloth's pick.
- if is_vulkan_backend and gpu_indices is not None:
- cmd += LlamaCppBackend._vulkan_pin_args(gpu_indices)
+ # Vulkan pins via --device (a cmd arg), before user extras so a user
+ # --device wins. Fall back to raw ids when the fit did not narrow.
+ _vulkan_pin_ids = gpu_indices if gpu_indices is not None else (gpu_ids or None)
+
+ # Record the pin actually applied (fit-narrowed gpu_indices, else the raw
+ # request) for the keep-warm loop, dedupe, and /status, so an explicit
+ # [0, 1] narrowed to [0] records [0] and /status never echoes an ordinal
+ # the child never saw. Auto selection (no gpu_ids) stays None (#7239).
+ if is_vulkan_backend:
+ # Only record an EXPLICIT Vulkan pin: an auto pick still narrows +
+ # pins below, but recording it would misreport an explicit pin and
+ # make dedupe miss the loaded server; mirrors the CUDA/ROCm branch.
+ self._gpu_ids = (
+ sorted(int(x) for x in _vulkan_pin_ids)
+ if (gpu_ids and _vulkan_pin_ids)
+ else None
+ )
+ elif gpu_ids:
+ # Physical pin: the fit-selected subset when the fit ran, else the raw
+ # user selection so an explicit choice is honoured even when the fit
+ # could not size the model.
+ _effective_pin_ids = (
+ [int(x) for x in gpu_indices]
+ if gpu_indices is not None
+ else [int(x) for x in gpu_ids]
+ )
+ self._gpu_ids = (
+ sorted(int(x) for x in _effective_pin_ids) if _effective_pin_ids else None
+ )
+ else:
+ self._gpu_ids = None
+
+ # Also record the RAW requested pin (before the fit narrowed it). Load
+ # dedupe compares this so a [0, 1] narrowed to [0] and re-sent as [0, 1]
+ # still matches, while /status keeps echoing the effective pin (#7239).
+ self._requested_gpu_ids = sorted(int(x) for x in gpu_ids) if gpu_ids else None
+
+ if is_vulkan_backend and _vulkan_pin_ids is not None:
+ cmd += LlamaCppBackend._vulkan_pin_args(_vulkan_pin_ids)
# User pass-through args go last so llama.cpp's last-wins parsing
# lets the user override Unsloth's auto-set flags. Already
@@ -7595,6 +8587,8 @@ class LlamaCppBackend:
cmd.extend(str(a) for a in extra_args)
logger.info(f"Appending user extra args to llama-server: {list(extra_args)}")
+ kv_cache_unified = _kv_unified_from_args(cmd)
+
logger.info(f"Starting llama-server: {' '.join(self._redacted_cmd_for_log(cmd))}")
# Library paths so llama-server finds its shared libs and CUDA DLLs.
@@ -7655,10 +8649,10 @@ class LlamaCppBackend:
f"Data-center GPU detected: applied DC llama.cpp env tuning (multi_gpu={multi_gpu})"
)
- # Pin to selected GPU(s). On ROCm, narrowing only
- # CUDA_VISIBLE_DEVICES leaves an AMD child seeing the full set, so
- # set HIP_VISIBLE_DEVICES too. Vulkan is pinned via --device
- # (above), not here.
+ # Pin to selected GPU(s) (issue #7164; resolved above into gpu_indices).
+ # On ROCm, narrowing only CUDA_VISIBLE_DEVICES leaves the AMD child
+ # seeing the full set, so set HIP_VISIBLE_DEVICES too. Vulkan is pinned
+ # via --device (above), not here.
# A deliberate zero-offload load with no GPU companions runs
# entirely on CPU, yet a visible CUDA device still costs the child
# ~0.5 GB (context + compute scratch) that the CPU-only
@@ -7684,7 +8678,12 @@ class LlamaCppBackend:
# 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))
+ # Mask on AMD at the ROCr/HSA layer: HIP-only masking still
+ # enumerates every agent first, which segfaults on a deselected
+ # unsupported GPU (e.g. gfx1036 iGPU under a gfx103X prebuilt).
+ self._emit_child_gpu_visibility(
+ env, ",".join(str(i) for i in gpu_indices), prefer_rocr = True
+ )
elif manual_tensor_split_emitted and not is_vulkan_backend:
# A manual per-GPU ratio across ALL GPUs (no explicit pick, so
# no CUDA_VISIBLE_DEVICES mask above): the UI built the
@@ -7747,7 +8746,7 @@ class LlamaCppBackend:
buffering = 1,
)
logger.info(f"llama-server stdout/stderr -> {self._llama_log_path}")
- except OSError as e:
+ except (OSError, UnicodeDecodeError) as e:
# Best-effort; never block the load on logging.
logger.debug(f"Could not open llama-server log file: {e}")
self._llama_log_path = None
@@ -7757,6 +8756,8 @@ class LlamaCppBackend:
stdout = subprocess.PIPE,
stderr = subprocess.STDOUT,
text = True,
+ encoding = "utf-8",
+ errors = "replace",
env = env,
**_windows_hidden_subprocess_kwargs(),
**_child_popen_kwargs(),
@@ -7889,6 +8890,13 @@ class LlamaCppBackend:
_fa_rc,
)
self._kill_process()
+ # The argv rewrite can't reach an env-only quantized V
+ # cache; drop it so the FA-off child doesn't abort on it.
+ if self._drop_env_quantized_v_cache(env):
+ logger.info(
+ "Dropped inherited quantized V-cache env for the "
+ "--flash-attn off retry (requires flash attention)."
+ )
cmd = _fa_cmd
healthy = _spawn_and_wait(_fa_cmd, label = "-noflash")
@@ -7934,6 +8942,13 @@ class LlamaCppBackend:
_probe_rc,
)
self._kill_process()
+ # The argv rewrite can't reach an env-only quantized V
+ # cache; drop it so the FA-off child doesn't abort on it.
+ if self._drop_env_quantized_v_cache(env):
+ logger.info(
+ "Dropped inherited quantized V-cache env for the "
+ "--flash-attn off retry (requires flash attention)."
+ )
cmd = _fa_cmd
healthy = (
_spawn_and_wait(_fa_cmd, label = "-noflash-mtp")
@@ -8016,23 +9031,29 @@ class LlamaCppBackend:
self._kill_process()
# The #6415 split-axis abort is latched earlier (first spawn).
# Skip if a cancel/unload is pending (mirrors the MTP guard).
+ _projector_msg = self._is_projector_incompatibility(out)
+ _signal_mmproj_guess = self._is_signal_crash(
+ _crash_rc
+ ) and not self._output_has_nonprojector_diagnostic(out)
if (
launched_with_mmproj
and not self._cancel_event.is_set()
- and (
- self._is_projector_incompatibility(out)
- or (
- self._is_signal_crash(_crash_rc)
- and not self._output_has_nonprojector_diagnostic(out)
- )
- )
+ and (_projector_msg or _signal_mmproj_guess)
):
- logger.warning(
- "llama-server could not load this model's vision "
- "projector (--mmproj). The installed llama.cpp build is "
- "likely too old for it. Loading text-only for this "
- "session; run 'unsloth studio update' to enable vision."
- )
+ if _projector_msg:
+ logger.warning(
+ "llama-server could not load this model's vision "
+ "projector (--mmproj). The installed llama.cpp build is "
+ "likely too old for it. Loading text-only for this "
+ "session; run 'unsloth studio update' to enable vision."
+ )
+ else:
+ logger.warning(
+ "llama-server crashed while loading this model's vision "
+ "projector (--mmproj). Retrying text-only for this "
+ "session; if this persists, run 'unsloth studio update' "
+ "or check GPU/driver logs."
+ )
cmd = self._strip_mmproj_args(_last_spawn_cmd)
# This retry bypasses _spawn_and_wait, so refresh the
# launched-argv snapshot itself -- the zero-offload
@@ -8046,14 +9067,30 @@ class LlamaCppBackend:
# an OS-killed text-only retry still gets the OOM message.
_retry_rc = self._process.poll() if self._process is not None else None
self._kill_process()
+ # If the text-only retry ALSO hard-crashed (a signal, not
+ # OOM/timeout), the vision projector was never the cause:
+ # llama-server is faulting during GPU/driver init. Say so
+ # -- with the ROCm fix -- instead of blaming the mmproj.
+ if self._is_signal_crash(_retry_rc):
+ raise RuntimeError(
+ "llama-server crashed at startup on both the vision "
+ "and text-only attempts -- a GPU driver/runtime "
+ "initialization crash, not a model or vision-projector "
+ "problem. This often means an unsupported secondary "
+ "GPU; on AMD/ROCm, hide it with ROCR_VISIBLE_DEVICES "
+ "(e.g. ROCR_VISIBLE_DEVICES=0 exposes only the first "
+ "GPU) before launching Unsloth Studio."
+ )
+ _retry_detail = self._classify_llama_start_failure(
+ "\n".join(self._stdout_lines[-50:]),
+ gguf_path,
+ self._model_identifier,
+ _retry_rc,
+ )
raise RuntimeError(
- "Vision projector incompatible with this llama.cpp "
- "build, and the text-only retry also failed: "
- + self._classify_llama_start_failure(
- "\n".join(self._stdout_lines[-50:]),
- gguf_path,
- self._model_identifier,
- _retry_rc,
+ self._mmproj_retry_failure_message(
+ projector_confirmed = _projector_msg,
+ detail = _retry_detail,
)
)
else:
@@ -8068,6 +9105,21 @@ class LlamaCppBackend:
self._healthy = True
self._commit_effective_parallel_slots(n_parallel)
+ self._swa_full = swa_full
+ self._kv_cache_unified = kv_cache_unified
+ self._n_ubatch = max(
+ 0,
+ int(self._DEFAULT_N_UBATCH if _effective_ubatch is None else _effective_ubatch),
+ )
+ self._flash_attn_enabled = (
+ _flash_attn_enabled_from_args(_last_spawn_cmd, env = env)
+ and self._architecture != "grok"
+ )
+ self._effective_cache_types = _effective_main_cache_types(
+ _last_spawn_cmd,
+ env,
+ )
+ self._kv_cache_context_total = effective_ctx if effective_ctx > 0 else None
# Server is up: adopt the real per-request context it allocated
# -- the length --fit chose, or a --parallel slot split -- so the
@@ -8075,6 +9127,11 @@ class LlamaCppBackend:
# before the spawn above always failed; the seeded value was the
# requested/native length.)
self._reconcile_effective_ctx_with_server()
+ if self._kv_cache_context_total is not None:
+ self._n_ubatch = min(
+ self._n_ubatch,
+ self._kv_cache_context_total,
+ )
# Commit caller intent only after _healthy=True so a failed start
# can't poison the next inheritance check. None keeps prior, []
@@ -8084,6 +9141,8 @@ class LlamaCppBackend:
self._extra_args = list(extra_args)
self._extra_args_source = (model_identifier, hf_variant)
self._requested_n_ctx = int(n_ctx)
+ # Local n_parallel may have been reduced above; the snapshot has the ask.
+ self._requested_n_parallel = max(1, int(_pending_load_kwargs["n_parallel"]))
# Commit the known-good snapshot + whether MTP+tensor is live, then
# watch this load for a mid-generation crash.
self._last_load_kwargs = _pending_load_kwargs
@@ -8283,18 +9342,29 @@ class LlamaCppBackend:
caps = self.probe_server_capabilities(binary)
mtp_token = caps.get("mtp_token") if caps else None
if not mtp_token:
- logger.warning(
- "Requested MTP speculative decoding but "
- "llama-server lacks --spec-type mtp/draft-mtp; "
- "run `unsloth studio update`. Loading without "
- "speculative decoding."
- )
+ inconclusive = bool(caps.get("mtp_probe_inconclusive")) if caps else True
+ if inconclusive:
+ logger.info(
+ "Requested MTP speculative decoding but llama-server MTP "
+ "capability probe was inconclusive; loading without "
+ "speculative decoding."
+ )
+ else:
+ logger.warning(
+ "Requested MTP speculative decoding but "
+ "llama-server lacks --spec-type mtp/draft-mtp; "
+ "run `unsloth studio update`. Loading without "
+ "speculative decoding."
+ )
# Override an inherited LLAMA_ARG_SPEC_TYPE=draft-mtp (CLI wins
# over env) so the child matches the binary-capability gate and
# the no-MTP budget, like the sibling no-head/non-MTP fallbacks.
flags.append("--spec-default")
self._speculative_type = "default"
- self._spec_fallback_reason = "binary_no_mtp"
+ if inconclusive:
+ self._spec_fallback_reason = None
+ else:
+ self._spec_fallback_reason = "binary_no_mtp"
return False
draft_n_max = _resolved_draft_n_max()
n_max_flag = caps.get("spec_draft_n_max_flag") or "--spec-draft-n-max"
@@ -8485,6 +9555,7 @@ class LlamaCppBackend:
tensor_split: Optional[List[float]] = None,
gpu_ids: Optional[List[int]] = None,
mtp_draft_path: Optional[str] = None,
+ n_parallel: int = 1,
preserve_multi_gpu_on_layer: bool = False,
) -> bool:
"""True iff the live server already satisfies these load kwargs.
@@ -8520,7 +9591,6 @@ class LlamaCppBackend:
if _norm(self._cache_type_kv) != _norm(cache_type_kv):
return False
-
# Reconcile a user --split-mode in extras AND an inherited tensor
# LLAMA_ARG_SPLIT_MODE env, but only against a server that actually
# launched tensor: if load_model downgraded to layer split it scrubbed
@@ -8544,9 +9614,16 @@ class LlamaCppBackend:
# layer/MoE/split knobs), so a standing manual preference in the
# request must not force a needless reload -- only the GPU pick matters.
if not self._is_diffusion:
+ requested_extra_args = extra_args if extra_args is not None else self._extra_args
+ if self._swa_full != _swa_full_from_args_or_env(requested_extra_args):
+ return False
# A GPU-memory-mode flip (Unsloth / manual) must always reload.
if self._gpu_memory_mode != gpu_memory_mode:
return False
+ # Requested-vs-requested (like n_ctx): comparing the effective count
+ # would reload forever whenever the fitter launched fewer slots.
+ if self._requested_n_parallel != max(1, int(n_parallel)):
+ return False
# Manual: a layer-count change always reloads (covers Auto(-1) <-> a
# pinned count); MoE/split only matter with an explicit offload.
if gpu_memory_mode == "manual" and (
@@ -8560,16 +9637,10 @@ class LlamaCppBackend:
)
):
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:
+ # A changed GPU pick must reload. Regular GGUF accepts either the raw
+ # requested placement pool or the effective status-echoed subset;
+ # diffusion compares its normalized single-device pick.
+ if not self.matches_gpu_ids(gpu_ids):
return False
# Compare on the canonical requested mode. With --spec-type in
@@ -8627,6 +9698,7 @@ class LlamaCppBackend:
current = list(self._extra_args) if self._extra_args is not None else []
if list(extra_args) != current:
return False
+ self._record_matching_gpu_request(gpu_ids)
return True
def _classify_gpu_offload(
@@ -8676,7 +9748,8 @@ class LlamaCppBackend:
last_draft: Optional[str] = None
args = [str(arg) for arg in cmd]
for index, raw in enumerate(args):
- flag, equals, inline = raw.partition("=")
+ flag = _flag_name(raw)
+ _, equals, inline = raw.partition("=")
if flag not in main_flags and flag not in draft_flags:
continue
value = inline if equals else (args[index + 1] if index + 1 < len(args) else "")
@@ -8722,6 +9795,7 @@ class LlamaCppBackend:
"""Terminate the subprocess and cancel any in-flight download."""
self._cancel_event.set()
with self._lock:
+ self._unload_epoch += 1
self._kill_process()
logger.info(f"Unloaded GGUF model: {self._model_identifier}")
self._model_identifier = None
@@ -8748,6 +9822,12 @@ class LlamaCppBackend:
self._slot_save_binary = None
self._slot_loaded_identity = None
self._prompt_cache_disabled = False
+ self._swa_full = False
+ self._kv_cache_unified = False
+ self._n_ubatch = self._DEFAULT_N_UBATCH
+ self._flash_attn_enabled = True
+ self._effective_cache_types = ("f16", "f16")
+ self._kv_cache_context_total = None
self._chat_template = None
self._chat_template_override = None
self._supports_reasoning = False
@@ -8758,12 +9838,15 @@ class LlamaCppBackend:
self._supports_preserve_thinking = False
self._supports_tools = False
self._cache_type_kv = None
+ # GPU-pin state describes the active runner only; clear it so an explicit
+ # pin never leaks into the next (or diffusion) runner.
+ self._gpu_ids = None
+ self._requested_gpu_ids = None
self._tensor_parallel = False
self._gpu_memory_mode = "auto"
self._gpu_layers = -1
self._n_cpu_moe = 0
self._tensor_split = None
- self._gpu_ids = None
self._layer_preserves_tensor_intent = False
self._speculative_type = None
self._requested_spec_mode = None
@@ -8876,7 +9959,7 @@ class LlamaCppBackend:
return
try:
path.parent.mkdir(parents = True, exist_ok = True)
- path.write_text(f"{pid}:{cls._pid_start_identity(pid)}")
+ path.write_text(f"{pid}:{cls._pid_start_identity(pid)}", encoding = "utf-8")
except Exception as e:
logger.debug(f"Could not write llama-server pidfile: {e}")
@@ -9010,7 +10093,7 @@ class LlamaCppBackend:
pid = -1
identity = ""
try:
- pid_str, _, identity = path.read_text().strip().partition(":")
+ pid_str, _, identity = path.read_text(encoding = "utf-8").strip().partition(":")
pid = int(pid_str)
except Exception:
pid = -1
@@ -9177,6 +10260,8 @@ class LlamaCppBackend:
["pgrep", "-a", "-f", "llama-server"],
capture_output = True,
text = True,
+ encoding = "utf-8",
+ errors = "replace",
timeout = 5,
env = child_env_without_native_path_secret(),
)
@@ -9288,8 +10373,12 @@ class LlamaCppBackend:
tuple(sidecars),
self._requested_n_ctx,
self._effective_context_length,
- getattr(self, "_cache_type_kv", None),
+ self._effective_cache_types,
self.effective_parallel_slots,
+ self._swa_full,
+ self._kv_cache_unified,
+ self._n_ubatch,
+ self._flash_attn_enabled,
)
def _gguf_file_identity(self, path) -> Optional[tuple]:
@@ -9320,7 +10409,8 @@ class LlamaCppBackend:
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("=")
+ flag = _flag_name(arg)
+ _, sep, inline = arg.partition("=")
if flag not in self._SIDECAR_WEIGHT_FLAGS:
continue
operand = inline if sep else (args[i + 1] if i + 1 < len(args) else "")
@@ -9360,7 +10450,7 @@ class LlamaCppBackend:
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"}
+ return env in _LLAMA_ARG_FALSE_VALUES
def save_slots_for_resume(
self, should_abort: Optional[Callable[[], bool]] = None
@@ -9372,6 +10462,17 @@ class LlamaCppBackend:
or self._prompt_cache_off()
):
return None
+ # Same predicate as the estimator's SWA path: a window alone is not enough.
+ # phi3 GGUFs carry attention.sliding_window but no key/value length, and
+ # llama.cpp forces them back to a non-SWA cache, so their slots do restore.
+ if (
+ (self._sliding_window or 0) > 0
+ and self._kv_key_length is not None
+ and self._kv_value_length is not None
+ and not self._swa_full
+ ):
+ logger.debug("Skipping slot save: compact SWA cache cannot be reused after restart")
+ return None
save_dir = Path(self._slot_save_dir)
gguf_stat = self._gguf_file_identity(self._gguf_path)
if gguf_stat is None:
@@ -9388,9 +10489,16 @@ class LlamaCppBackend:
return None
try:
estimate = self._estimate_kv_cache_bytes(
- self._effective_context_length or self._context_length or 0,
- self._cache_type_kv,
+ self._kv_cache_context_total
+ or self._effective_context_length
+ or self._context_length
+ or 0,
+ max(self._effective_cache_types, key = _kv_bytes_per_elem),
n_parallel = self.effective_parallel_slots,
+ swa_full = self._swa_full,
+ kv_unified = self._kv_cache_unified,
+ n_ubatch = self._n_ubatch,
+ flash_attn = self._flash_attn_enabled,
)
# Skip before writing anything when the estimate alone blows the cap,
# rather than fully writing a slot and discarding it afterwards.
@@ -9518,15 +10626,18 @@ class LlamaCppBackend:
return False
if not self._mtp_runtime_fallback_active:
return False
- if not self._last_load_kwargs or self._process is None:
+ # Read before claiming: a raise after the claim strands the flag, and nothing
+ # else clears it, blocking every later respawn.
+ kwargs = self._last_load_kwargs
+ proc = self._process
+ if not kwargs or proc is None:
return False
# Single-flight: the first failure claims the reload.
with self._mtp_runtime_fallback_lock:
if self._mtp_runtime_fallback_in_progress:
return False
self._mtp_runtime_fallback_in_progress = True
- snapshot = dict(self._last_load_kwargs)
- proc = self._process
+ snapshot = dict(kwargs)
def _recover():
try:
@@ -9574,7 +10685,14 @@ class LlamaCppBackend:
with self._mtp_runtime_fallback_lock:
self._mtp_runtime_fallback_in_progress = False
- threading.Thread(target = _recover, daemon = True, name = "mtp-crash-reload").start()
+ try:
+ threading.Thread(target = _recover, daemon = True, name = "mtp-crash-reload").start()
+ except RuntimeError as exc:
+ # Release the claim: a reload that never started would block respawn forever.
+ with self._mtp_runtime_fallback_lock:
+ self._mtp_runtime_fallback_in_progress = False
+ logger.error(f"Could not start the MTP-crash reload: {exc}")
+ return False
return True
def _start_mtp_crash_watchdog(self) -> None:
@@ -9736,6 +10854,8 @@ class LlamaCppBackend:
actual_n_ctx = self._query_server_n_ctx()
if not actual_n_ctx or actual_n_ctx <= 0:
return
+ slots = 1 if self._kv_cache_unified else self.effective_parallel_slots
+ self._kv_cache_context_total = actual_n_ctx * slots
if self._effective_context_length and actual_n_ctx < self._effective_context_length:
logger.warning(
"llama-server allocated a smaller per-request context than "
@@ -10046,6 +11166,21 @@ class LlamaCppBackend:
finally:
_cancel_closed.set()
+ def _server_socket_is_open(self, timeout_s: float = 0.15) -> bool:
+ """True if anything still accepts on the server port.
+
+ The listening socket dies with the process, so this tells a live server
+ from a dead one without waiting for the child to become reapable.
+ """
+ port = self._port
+ if not port:
+ return False
+ try:
+ with socket.create_connection(("127.0.0.1", port), timeout = timeout_s):
+ return True
+ except OSError:
+ return False
+
def _respawn_if_dead(self) -> bool:
"""Relaunch the llama-server if its process has exited.
@@ -10055,28 +11190,114 @@ class LlamaCppBackend:
recover, returning True once healthy. Serialised on ``_respawn_lock`` so
many generations hitting the dead server trigger at most one reload.
"""
+ # Read outside the lock so a queued caller can tell the replacement from the child
+ # its own error came from; otherwise each burns the grace wait below, and that
+ # sleep is held under the lock, so the waits serialise.
+ served_by = self._process
with self._respawn_lock:
proc = self._process
if proc is None:
return False
- if proc.poll() is None:
- # Process is alive: either a concurrent caller already respawned
- # it (healthy), or this connection error wasn't a dead server.
+ if self._cancel_event.is_set():
+ # unload_model sets this before it kills, so the child can still be
+ # accepting. Reporting it healthy would aim the retry at a server
+ # that is deliberately going away.
+ return False
+ if proc is not served_by:
+ # Replaced while we queued: this child never served our request.
return self._healthy
- kwargs = self._last_load_kwargs
- if not kwargs:
- return False
- logger.warning(
- f"llama-server for '{self._model_identifier}' exited "
- f"(code {proc.returncode}); respawning to recover the session"
- )
- with self._lock:
- self._healthy = False
+ if proc.poll() is None:
+ # Still serving, so the error was transient. Charging it the grace below
+ # would cost a second per caller, serialised under this lock.
+ if self._server_socket_is_open():
+ return self._healthy
+ # A closing server can beat its own exit status: calling it alive returns
+ # the stale _healthy and spends the retry on the corpse.
+ deadline = time.monotonic() + _RESPAWN_REAP_GRACE_S
+ while proc.poll() is None and time.monotonic() < deadline:
+ time.sleep(0.05)
+ if proc.poll() is None:
+ # Alive: either a concurrent caller already respawned it (healthy), or
+ # this connection error wasn't a dead server.
+ return self._healthy
+ with self._mtp_runtime_fallback_lock:
+ if self._mtp_runtime_fallback_in_progress:
+ # An MTP-free reload owns this corpse; replaying the old kwargs
+ # restarts the crashing config and aborts that reload.
+ logger.info("Respawn skipped: an MTP-free reload is already recovering.")
+ return False
+ # The RLock lets the load_model below re-enter it.
+ with self._serial_load_lock:
+ if self._process is not proc:
+ logger.info("Respawn skipped: a newer load is already active.")
+ return self._healthy
+ # Snapshot under _lock, the one unload_model holds, so a teardown is
+ # either wholly before us (flag set) or wholly after (epoch bumped).
+ # _serial_load_lock alone would not exclude it: unload never takes it.
+ with self._lock:
+ if self._cancel_event.is_set():
+ logger.info("Respawn skipped: the model was unloaded.")
+ return False
+ kwargs = dict(self._last_load_kwargs or {})
+ if not kwargs:
+ return False
+ epoch = self._unload_epoch
+ self._healthy = False
+ logger.warning(
+ f"llama-server for '{self._model_identifier}' exited "
+ f"(code {proc.returncode}); respawning to recover the session"
+ )
+ try:
+ started = bool(self.load_model(**kwargs))
+ except Exception as exc:
+ logger.error(f"Failed to respawn llama-server: {exc}")
+ return False
+ if started and self._unload_epoch != epoch:
+ # An unload landed mid-reload. load_model cleared _cancel_event on
+ # the way in, so the epoch is the only surviving evidence; undo the
+ # replacement rather than leave a model the user stopped running.
+ logger.info("Respawn undone: the model was unloaded during the reload.")
+ self.unload_model()
+ return False
+ return started
+
+ @contextlib.contextmanager
+ def _open_chat_stream_with_respawn_retry(self, payload: dict, cancel_event):
+ """Open a chat stream, respawning a dead llama-server once before streaming.
+
+ Retry only when opening the response fails: once it is open a consumer may
+ already have emitted content or tool events, so a replay could duplicate
+ output and side effects. ``base_url`` is resolved per attempt because a
+ respawn may pick a new port. The budget is one retry per model request, not
+ per chat turn, so a long tool loop never discards a completed tool.
+
+ A child dying after the accept but before the headers surfaces as
+ ReadError/WriteError/RemoteProtocolError rather than ConnectError, and which
+ one differs per OS. llama-server flushes its 200 at slot start, so that window
+ is an upload still in flight or a request behind busy slots; a death during
+ decode arrives with the response open and is not replayed. Timeouts are
+ excluded: the server is slow, not dead, and a replay would spend the
+ first-token budget twice.
+ """
+ for attempt in range(2):
+ response_opened = False
try:
- return bool(self.load_model(**kwargs))
- except Exception as exc:
- logger.error(f"Failed to respawn llama-server: {exc}")
- return False
+ url = f"{self.base_url}/v1/chat/completions"
+ with self._open_stream(url, payload, cancel_event) as opened:
+ response_opened = True
+ yield opened
+ return
+ except (httpx.NetworkError, httpx.RemoteProtocolError) as exc:
+ if response_opened:
+ raise
+ if self._maybe_recover_from_mtp_crash(exc):
+ raise RuntimeError("Lost connection to llama-server") from exc
+ if attempt == 0 and self._respawn_if_dead():
+ logger.warning(
+ "llama-server was unreachable; respawned it and retrying the generation"
+ )
+ continue
+ raise
def generate_chat_completion(
self,
@@ -10095,6 +11316,7 @@ class LlamaCppBackend:
reasoning_effort: Optional[str] = None,
preserve_thinking: Optional[bool] = None,
seed: Optional[int] = None,
+ promote_reasoning_only: bool = True,
_allow_respawn_retry: bool = True,
) -> Generator[Union[str, dict], None, None]:
"""
@@ -10177,7 +11399,12 @@ class LlamaCppBackend:
# model put its whole reply in reasoning
# (e.g. Qwen3 always-think). Show it as
# the main response, not a thinking block.
- cumulative = reasoning_text
+ cumulative = _finalize_reasoning_only_cumulative(
+ cumulative,
+ reasoning_text,
+ _metadata_finish_reason,
+ promote_reasoning_only,
+ )
yield cumulative
_stream_done = True
break # exit inner while
@@ -10274,6 +11501,7 @@ class LlamaCppBackend:
reasoning_effort = reasoning_effort,
preserve_thinking = preserve_thinking,
seed = seed,
+ promote_reasoning_only = promote_reasoning_only,
_allow_respawn_retry = False,
)
return
@@ -10315,6 +11543,7 @@ class LlamaCppBackend:
confirm_tool_calls: bool = False,
bypass_permissions: bool = False,
permission_mode: Optional[str] = None,
+ promote_reasoning_only: bool = True,
) -> Generator[dict, None, None]:
"""
Agentic loop: let the model call tools, execute them, and continue.
@@ -10333,17 +11562,22 @@ class LlamaCppBackend:
from core.inference.tools import (
build_rag_autoinject,
execute_tool,
+ has_text_only_provisional_card,
is_always_safe_tool,
- is_potentially_unsafe_tool_call,
+ is_high_risk_tool_call,
)
- # Normalize the mode: "full" and bypass_permissions are the same
- # switch, whichever arrives first wins toward the permissive side.
- # "off" keeps the sandbox but never prompts.
+ # "full" and bypass_permissions are the same switch, whichever arrives
+ # first wins. "off" keeps the sandbox but never prompts. Unset defaults to
+ # "auto"; unknown falls back to the stricter "ask". An explicit
+ # confirm_tool_calls=True with no mode is already resolved to "ask" at the
+ # request layer, so it never arrives here as an ambiguous unset.
if permission_mode == "full":
bypass_permissions = True
elif bypass_permissions:
permission_mode = "full"
+ elif permission_mode is None:
+ permission_mode = "auto"
elif permission_mode not in ("ask", "auto", "off"):
permission_mode = "ask"
@@ -10366,7 +11600,6 @@ class LlamaCppBackend:
yield _ev
conversation.extend(_auto["messages"])
- url = f"{self.base_url}/v1/chat/completions"
_accumulated_completion_tokens = 0
_accumulated_predicted_ms = 0.0
_accumulated_predicted_n = 0
@@ -10527,6 +11760,10 @@ class LlamaCppBackend:
# direct answer ("4", "Hello!") won't match. Pattern shared with the
# safetensors loop (tool_call_parser.INTENT_SIGNAL).
_reprompt_count = 0
+ # Budgeted apart from _reprompt_count so a pre-tool nudge can't spend it.
+ _post_tool_reprompts = 0
+ # Text that triggered the last nudge; if the retry restates it, stop.
+ _last_reprompt_text = ""
# Gates ``max_tool_iterations`` on real tool turns (not the enlarged range) so reserved
# re-prompt slots don't extend the budget. Mirrors the safetensors guard.
_tool_iters_done = 0
@@ -10534,7 +11771,7 @@ class LlamaCppBackend:
# Reserve extra iterations for re-prompts so they don't consume the
# caller's tool-call budget; only when tool iterations are allowed.
- _extra = _MAX_REPROMPTS if max_tool_iterations > 0 else 0
+ _extra = _MAX_REPROMPTS + 1 if max_tool_iterations > 0 else 0
for iteration in range(max_tool_iterations + _extra):
if cancel_event is not None and cancel_event.is_set():
return
@@ -10600,6 +11837,7 @@ class LlamaCppBackend:
# Time each reasoning pass so final answers can replace tool timing.
_reasoning_started_at = None
_reasoning_summary_emitted = False
+ _deferred_reasoning_summary = None
cumulative_display = "" # Cumulative yielded text (with )
in_thinking = False
has_content_tokens = False
@@ -10626,7 +11864,7 @@ class LlamaCppBackend:
_text_args_name = ""
_confirm_gated_iteration = bool(confirm_tool_calls) and not bypass_permissions
- with self._open_stream(url, payload, cancel_event) as (
+ with self._open_chat_stream_with_respawn_retry(payload, cancel_event) as (
response,
first_token_deadline,
):
@@ -10657,7 +11895,12 @@ class LlamaCppBackend:
),
}
else:
- cumulative_display = reasoning_accum
+ cumulative_display = _finalize_reasoning_only_cumulative(
+ cumulative_display,
+ reasoning_accum,
+ _iter_finish_reason,
+ promote_reasoning_only,
+ )
if not _suppress_visible_output:
yield {
"type": "content",
@@ -10751,6 +11994,9 @@ class LlamaCppBackend:
permission_mode == "auto"
and is_always_safe_tool(current_name)
)
+ # A text-preview card still streams while gated;
+ # hiding it blanks the chat.
+ and not has_text_only_provisional_card(current_name)
)
# Keep small-argument tools on the normal path.
_args_len = len(
@@ -10843,7 +12089,11 @@ class LlamaCppBackend:
and not _reasoning_summary_emitted
):
_reasoning_summary_emitted = True
- yield _reasoning_summary_event(_reasoning_started_at)
+ _summary = _reasoning_summary_event(_reasoning_started_at)
+ if _suppress_visible_output:
+ _deferred_reasoning_summary = _summary
+ else:
+ yield _summary
has_content_tokens = True
content_accum += token
@@ -10852,20 +12102,27 @@ class LlamaCppBackend:
# TEXT call to a provisional card. Gated on an enabled-name
# sniff + size floor so prose/small calls spawn no pane; id
# matches the first call so the final tool_start reconciles.
- if (
- not has_structured_tc
- and not _confirm_gated_iteration
- and _text_args_call_start >= 0
- ):
+ if not has_structured_tc and _text_args_call_start >= 0:
if not _text_args_id:
_call_text = content_accum[_text_args_call_start:]
_sniffed = _sniff_text_tool_name(
_call_text, _enabled_tool_names
)
- if _sniffed and (
- _sniffed == "render_html"
- or len(_call_text)
- >= _PROVISIONAL_ARGS_MIN_CHARS
+ # Structured-path rule: gated calls
+ # stream only from a text-preview card.
+ if (
+ _sniffed
+ and not (
+ _confirm_gated_iteration
+ and not has_text_only_provisional_card(
+ _sniffed
+ )
+ )
+ and (
+ _sniffed == "render_html"
+ or len(_call_text)
+ >= _PROVISIONAL_ARGS_MIN_CHARS
+ )
):
_text_args_id = "call_0"
_text_args_name = _sniffed
@@ -11120,8 +12377,17 @@ class LlamaCppBackend:
# route's extractor closes the streamed ).
if _reasoning_started_at is not None and not _reasoning_summary_emitted:
_reasoning_summary_emitted = True
- yield _reasoning_summary_event(_reasoning_started_at)
- cumulative_display = reasoning_accum
+ _summary = _reasoning_summary_event(_reasoning_started_at)
+ if _suppress_visible_output:
+ _deferred_reasoning_summary = _summary
+ else:
+ yield _summary
+ cumulative_display = _finalize_reasoning_only_cumulative(
+ cumulative_display,
+ reasoning_accum,
+ _iter_finish_reason,
+ promote_reasoning_only,
+ )
if not _suppress_visible_output:
yield {
"type": "content",
@@ -11150,12 +12416,10 @@ class LlamaCppBackend:
)
if not _safety_tc:
# ── Re-prompt on plan-without-action ──
- # If the model described its intent (forward-looking
- # language) without calling a tool, nudge it to act.
- # Fires at most once per request, only on short
- # responses with intent signals -- "4" or "Hello!"
- # won't trigger it. Use content if available, else
- # fall back to reasoning text (reasoning-only stalls).
+ # Intent described without a tool call: nudge it to act. Up
+ # to _MAX_REPROMPTS times, only on short responses with intent
+ # signals -- "4" or "Hello!" won't trigger it. Uses content,
+ # else reasoning text (reasoning-only stalls).
_stripped = content_accum.strip()
if not _stripped:
_stripped = reasoning_accum.strip()
@@ -11165,18 +12429,33 @@ class LlamaCppBackend:
r"(?i)\brender[_\s-]?html\b",
_stripped,
)
+ # A post-tool stall still deserves a nudge, but each retry
+ # re-runs tools, so allow only one. RAG autoinject never lands
+ # in history, so _auto keeps a doc-grounded turn from reading
+ # as pre-tool (mirrors safetensors rag_autoinjected).
+ _already_acted = bool(_auto) or any(
+ record.executed for record in tool_controller.history
+ )
+ if _already_acted:
+ _reprompt_used, _reprompt_cap = _post_tool_reprompts, 1
+ else:
+ _reprompt_used, _reprompt_cap = _reprompt_count, _MAX_REPROMPTS
# None keeps the default-on re-prompt; False disables it.
if (
auto_heal_tool_calls
and (nudge_tool_calls is None or nudge_tool_calls)
and active_tools
and not _render_html_already_done_intent
- and _reprompt_count < _MAX_REPROMPTS
+ and _reprompt_used < _reprompt_cap
+ and not _is_reprompt_repeat(_stripped, _last_reprompt_text)
and _is_short_intent_without_action(_stripped)
):
_reprompt_count += 1
+ if _already_acted:
+ _post_tool_reprompts += 1
+ _last_reprompt_text = _stripped
logger.info(
- f"Re-prompt {_reprompt_count}/{_MAX_REPROMPTS}: "
+ f"Re-prompt {_reprompt_used + 1}/{_reprompt_cap}: "
f"model responded without calling tools "
f"({len(_stripped)} chars)"
)
@@ -11206,12 +12485,18 @@ class LlamaCppBackend:
_it_r = _iter_timings or {}
_accumulated_predicted_ms += _it_r.get("predicted_ms", 0)
_accumulated_predicted_n += _it_r.get("predicted_n", 0)
+ # Blank first (the route resets its text cursor only on an
+ # empty status), then the badge so the retry is not a hang.
yield {"type": "status", "text": ""}
+ yield {"type": "status", "text": _NUDGE_TOOL_CALLS_STATUS}
continue
if _forced_tool_call_pending:
_forced_tool_call_pending = False
- if not _should_suppress_forced_no_tool_output(_stripped):
+ if not _should_suppress_forced_no_tool_output(
+ _stripped,
+ _last_reprompt_text,
+ ):
if cumulative_display:
forced_visible_text = _strip_tool_markup(
cumulative_display,
@@ -11229,6 +12514,8 @@ class LlamaCppBackend:
"type": "content",
"text": forced_visible_text,
}
+ if _deferred_reasoning_summary is not None:
+ yield _deferred_reasoning_summary
elif not _suppress_visible_output:
# Turn ended as a plain answer (no [ARGS] followed): the held
# rehearsal tail is real prose, release it.
@@ -11428,18 +12715,16 @@ class LlamaCppBackend:
decision.as_assistant_tool_call()
)
- # Bypass wins over the confirm gate at the loop level too,
- # so a direct internal caller with both flags never prompts.
- # In "auto" mode only calls detected as potentially unsafe
- # pause; read-only calls run straight through. "off" never
- # prompts (sandbox stays on).
+ # Bypass wins here too, so a direct internal caller with both
+ # flags never prompts. "auto" pauses only high-risk calls;
+ # "off" never prompts (sandbox stays on).
needs_confirm = (
bool(confirm_tool_calls)
and not bypass_permissions
and permission_mode != "off"
)
if needs_confirm and permission_mode == "auto":
- needs_confirm = is_potentially_unsafe_tool_call(
+ needs_confirm = is_high_risk_tool_call(
decision.tool_name, decision.arguments
)
approval_id = new_approval_id() if needs_confirm else ""
@@ -11451,18 +12736,31 @@ class LlamaCppBackend:
start_event["awaiting_confirmation"] = needs_confirm
try:
- yield {"type": "status", "text": decision.status_text}
+ # Gated calls are not running yet; a "Running ..." badge
+ # counting up while it waits on a human reads as a hang.
+ yield {
+ "type": "status",
+ "text": (
+ awaiting_approval_status(decision.tool_name)
+ if needs_confirm
+ else decision.status_text
+ ),
+ }
yield start_event
- if (
- decision_slot is not None
- and wait_tool_decision(
+ _decision = (
+ wait_tool_decision(
decision_slot,
approval_id,
cancel_event = cancel_event,
)
- == "deny"
- ):
+ if decision_slot is not None
+ else None
+ )
+ if _decision is not None and _decision != "deny":
+ # Approved: now it really is running.
+ yield {"type": "status", "text": decision.status_text}
+ if _decision == "deny":
decision_slot = None
resolved_provisional_tool_call_ids.add(decision.tool_call_id)
yield {
@@ -11528,6 +12826,10 @@ class LlamaCppBackend:
_kb_search_count += 1
completion = tool_controller.record_result(decision, result)
resolved_provisional_tool_call_ids.add(decision.tool_call_id)
+ # A real execution opens the post-tool phase; carrying the pre-tool
+ # stall text over would read the same sentence as a repeat and
+ # swallow the one post-tool nudge.
+ _last_reprompt_text = ""
# A tool ran this turn, so it counts against the caller's budget.
_turn_executed_real_tool = True
yield completion.tool_end_event()
@@ -11653,7 +12955,7 @@ class LlamaCppBackend:
_stream_done = False
try:
- with self._open_stream(url, stream_payload, cancel_event) as (
+ with self._open_chat_stream_with_respawn_retry(stream_payload, cancel_event) as (
response,
first_token_deadline,
):
@@ -11685,7 +12987,12 @@ class LlamaCppBackend:
"text": _strip_tool_markup(cumulative, final = True),
}
else:
- cumulative = reasoning_text
+ cumulative = _finalize_reasoning_only_cumulative(
+ cumulative,
+ reasoning_text,
+ _metadata_finish_reason,
+ promote_reasoning_only,
+ )
yield {"type": "content", "text": cumulative}
_stream_done = True
break # exit inner while
@@ -12025,10 +13332,15 @@ class LlamaCppBackend:
min_p: float = 0.0,
max_new_tokens: int = 2048,
repetition_penalty: float = 1.1,
+ cancel_event: Optional[threading.Event] = None,
) -> tuple:
"""
Generate TTS audio via llama-server /completion + codec decode.
Returns (wav_bytes, sample_rate).
+
+ ``cancel_event`` lets a Stop or a forced model swap end the request: the
+ decode is one blocking POST, so a watcher closes the client out from under
+ it rather than polling. Raises RuntimeError once cancelled.
"""
if audio_type not in self._TTS_PROMPTS:
raise RuntimeError(f"GGUF TTS does not support '{audio_type}' codec.")
@@ -12050,15 +13362,47 @@ class LlamaCppBackend:
if need_ids:
payload["n_probs"] = 1
+ if cancel_event is not None and cancel_event.is_set():
+ raise RuntimeError("Audio generation cancelled")
+
with httpx.Client(
timeout = httpx.Timeout(300, connect = 10),
headers = self._auth_headers,
trust_env = False,
) as client:
- resp = client.post(f"{self.base_url}/completion", json = payload)
+ finished = threading.Event()
+ watcher: Optional[threading.Thread] = None
+ if cancel_event is not None:
+
+ def _close_when_cancelled() -> None:
+ while not finished.wait(0.05):
+ if cancel_event.is_set():
+ # Closing mid-request makes the blocking post raise
+ # httpx.RequestError, the only way out of it.
+ with contextlib.suppress(Exception):
+ client.close()
+ return
+
+ watcher = threading.Thread(target = _close_when_cancelled, daemon = True)
+ watcher.start()
+ try:
+ resp = client.post(f"{self.base_url}/completion", json = payload)
+ except httpx.RequestError:
+ if cancel_event is not None and cancel_event.is_set():
+ raise RuntimeError("Audio generation cancelled") from None
+ raise
+ finally:
+ finished.set()
+ if watcher is not None:
+ watcher.join(timeout = 0.5)
if resp.status_code != 200:
raise RuntimeError(f"llama-server returned {resp.status_code}: {resp.text}")
+ # The codec decode below is GPU work with no interruption point, so check here:
+ # cancelling after this only wastes the decode it cannot stop.
+ if cancel_event is not None and cancel_event.is_set():
+ raise RuntimeError("Audio generation cancelled")
+
data = resp.json()
token_ids = (
[p["id"] for p in data.get("completion_probabilities", []) if "id" in p]
diff --git a/studio/backend/core/inference/llama_keepwarm.py b/studio/backend/core/inference/llama_keepwarm.py
index 3380ebf5f5..05b1271b27 100644
--- a/studio/backend/core/inference/llama_keepwarm.py
+++ b/studio/backend/core/inference/llama_keepwarm.py
@@ -345,6 +345,22 @@ def _loaded_identity(backend):
return (backend.model_identifier, getattr(backend, "hf_variant", None), advertised)
+def _note_idle_unload_event(freed) -> None:
+ """Monitor row for an idle auto-unload. Best-effort; uses the stash's
+ advertised repo id so the row never shows the on-disk load path."""
+ try:
+ from core.inference.api_monitor import api_monitor
+ from core.inference.model_ids import public_model_id
+
+ identifier, variant, advertised = (list(freed) + [None, None, None])[:3]
+ label = public_model_id(advertised or identifier) or "model"
+ if variant and ":" not in label:
+ label = f"{label}:{variant}"
+ api_monitor.record_lifecycle(event = "unload", model = label, reason = "idle")
+ except Exception as exc:
+ logger.debug("idle unload monitor event failed: %s", exc)
+
+
async def idle_unload_loop(poll_seconds: float = 15.0) -> None:
"""Unload the loaded GGUF once idle past the configured TTL. Inert when off."""
from utils.openai_auto_switch_settings import (
@@ -407,6 +423,8 @@ async def idle_unload_loop(poll_seconds: float = 15.0) -> None:
elif manifest:
_delete_resume_files(manifest)
logger.info("Idle auto-unload: freed GGUF after %ss idle", ttl)
+ # An idle unload stashes for reload and skips note_model_unloaded.
+ _note_idle_unload_event(freed)
seen_model = None
except Exception as exc:
logger.debug("idle_unload_loop iteration failed: %s", exc)
diff --git a/studio/backend/core/inference/llama_server_args.py b/studio/backend/core/inference/llama_server_args.py
index 7b42d2f40d..7391e62516 100644
--- a/studio/backend/core/inference/llama_server_args.py
+++ b/studio/backend/core/inference/llama_server_args.py
@@ -16,11 +16,18 @@ from __future__ import annotations
import os
from typing import Iterable, Mapping, Optional
+# Valid llama-server --parallel range, shared with LoadRequest.n_parallel.
+# Mirrored by callers that cannot import this: run.py and unsloth_cli/commands/
+# studio.py (_PARALLEL_MIN/MAX), per-model-config.ts (N_PARALLEL_MIN/MAX);
+# test_parallel_slots_per_load.py pins them together.
+PARALLEL_MIN = 1
+PARALLEL_MAX = 64
+
# Each group = every alias (short + long) of one hard-denied flag.
# Extend the matching group when llama.cpp adds a new alias.
_DENYLIST_GROUPS: tuple[frozenset[str], ...] = (
- # Parallel slots: owned by typer --parallel; a pass-through would desync
- # app.state.llama_parallel_slots from llama-server.
+ # Parallel slots: owned by typer --parallel and LoadRequest.n_parallel; a
+ # pass-through would desync the slot bookkeeping from llama-server.
frozenset({"-np", "--parallel", "--n-parallel"}),
# Model identity: Unsloth resolves it from LoadRequest; a second -m would
# load a different model than Unsloth thinks it loaded.
@@ -80,9 +87,10 @@ _DENYLIST: frozenset[str] = frozenset().union(*_DENYLIST_GROUPS)
def _flag_name(token: str) -> Optional[str]:
"""Flag name for ``token``, or None if it isn't a flag.
- Peels `--key=value` to `--key`, treats `-1`/`-0.5` as values (shorts
- always start with a letter), and normalises attached `-np8` / `-np-1` /
- `-np8x` to `-np`. Mirrors the CLI's `_expand_attached_np_short`.
+ Peels `--key=value` to `--key`, normalises long-option underscores like
+ llama.cpp, treats `-1`/`-0.5` as values (shorts always start with a letter),
+ and normalises attached `-np8` / `-np-1` / `-np8x` to `-np`. Mirrors the
+ CLI's `_expand_attached_np_short`.
"""
token = token.strip()
if not token.startswith("-") or token in {"-", "--"}:
@@ -90,6 +98,8 @@ def _flag_name(token: str) -> Optional[str]:
if len(token) >= 2 and (token[1].isdigit() or token[1] == "."):
return None
name = token.split("=", 1)[0]
+ if name.startswith("--"):
+ name = name.replace("_", "-")
if len(name) > 3 and name.startswith("-np"):
suffix = name[3:]
if suffix[0].isdigit() or (
@@ -118,6 +128,7 @@ def validate_extra_args(args: Optional[Iterable[str]]) -> list[str]:
parse_ctx_override(out)
parse_cache_override(out)
parse_split_mode_override(out)
+ parse_gpu_layers_override(out)
return out
@@ -193,9 +204,8 @@ _SPLIT_SHADOWING_FLAGS: frozenset[str] = _SPLIT_MODE_FLAGS | _TENSOR_SPLIT_FLAGS
# 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"}
-)
+_GPU_LAYER_FLAGS: frozenset[str] = frozenset({"-ngl", "--gpu-layers", "--n-gpu-layers"})
+_LAYER_OFFLOAD_FLAGS: frozenset[str] = _GPU_LAYER_FLAGS | frozenset({"-fit", "--fit"})
_MOE_OFFLOAD_FLAGS: frozenset[str] = frozenset({"-ncmoe", "--n-cpu-moe", "-cmoe", "--cpu-moe"})
_OFFLOAD_SHADOWING_FLAGS: frozenset[str] = _LAYER_OFFLOAD_FLAGS | _MOE_OFFLOAD_FLAGS
@@ -306,6 +316,26 @@ def parse_cache_override(args: Optional[Iterable[str]]) -> Optional[str]:
return _last_flag_value(args, _CACHE_FLAGS)
+def parse_gpu_layers_override(args: Optional[Iterable[str]]) -> Optional[int]:
+ """Return the last user-supplied GPU layer count from extras.
+
+ Manual GPU memory mode strips llama.cpp offload flags because the
+ first-class load fields own them. Callers use this parser first to preserve
+ an explicit ``-ngl`` / ``--gpu-layers`` / ``--n-gpu-layers`` value when
+ translating the extras into those fields.
+ """
+ raw_value = _last_flag_value(args, _GPU_LAYER_FLAGS)
+ if raw_value is None:
+ return None
+ try:
+ value = int(raw_value)
+ except ValueError as exc:
+ raise ValueError("llama-server GPU layers flag requires an integer value") from exc
+ if value < -1:
+ raise ValueError("llama-server GPU layers flag requires an integer value of at least -1")
+ return value
+
+
def parse_cache_override_per_axis(
args: Optional[Iterable[str]],
) -> tuple[Optional[str], Optional[str]]:
diff --git a/studio/backend/core/inference/local_model_resolver.py b/studio/backend/core/inference/local_model_resolver.py
index 64ab38ec75..5d2a9e9c87 100644
--- a/studio/backend/core/inference/local_model_resolver.py
+++ b/studio/backend/core/inference/local_model_resolver.py
@@ -34,6 +34,15 @@ class _LocalGgufEntry:
_CACHE_TTL_S = 5.0
_lock = threading.Lock()
_scan: tuple[float, dict[str, _LocalGgufEntry]] = (0.0, {})
+# Not _lock: that is held for the whole scan, so the request path would wait on it.
+_warm_lock = threading.Lock()
+# Repos that finished downloading but are not in the published index yet: nothing
+# else covers them until the next scan, and the request path must not call them absent.
+_just_downloaded: set[str] = set()
+_warming = False
+_last_scan_s = 0.0
+# Rescan at most a tenth of the time: on the TTL alone a slow scan would run continuously.
+_WARM_DUTY = 10.0
def _is_abs_path_id(value: str) -> bool:
@@ -103,17 +112,26 @@ def _local_gguf_entry(loader_id: str, info) -> Optional[_LocalGgufEntry]:
load_dir = _resolve_load_dir(p)
variants, _ = list_local_gguf_variants(str(load_dir))
quants = tuple(v.quant for v in variants if getattr(v, "quant", None))
- return _LocalGgufEntry(loader_id, str(load_dir), quants) if quants else None
+ if not quants:
+ return None
+ # That call orders by descending size, so the head is the biggest quant (often
+ # F16). Downstream reads [0], and a bare id must mean whichever quant a plain
+ # load would take: answering with the largest can evict a model and then OOM.
+ from core.inference.openai_auto_download import preferred_quant
+
+ best = preferred_quant(quants)
+ if best and quants[0] != best:
+ quants = (best, *(q for q in quants if q != best))
+ return _LocalGgufEntry(loader_id, str(load_dir), quants)
except Exception:
return None
-def info_has_local_gguf(info) -> bool:
- """True when *info* (a LocalModelInfo) points to on-disk GGUF weights the
- auto-switch path can load. Read from the files, not ``info.model_format``: the
- HF-cache scanner leaves model_format unset for GGUF snapshots, so a
- model_format filter would drop every cached GGUF. Lets /v1/models advertise
- exactly what /v1 can serve."""
+def local_gguf_quants(info) -> Optional[tuple[str, ...]]:
+ """On-disk quant labels for *info*, or None when it is not a servable local
+ GGUF. Read from the files, not ``info.model_format``: the HF-cache scanner
+ leaves that unset for GGUF snapshots, so filtering on it drops every cached
+ GGUF. One scan tells /v1/models what it can serve and which quant to name."""
from pathlib import Path
path = getattr(info, "path", None)
@@ -123,8 +141,14 @@ def info_has_local_gguf(info) -> bool:
if isinstance(path, str) and any(
seg in (".studio_links", "ollama_links") for seg in Path(path).parts
):
- return False
- return _local_gguf_entry(getattr(info, "id", "") or "", info) is not None
+ return None
+ entry = _local_gguf_entry(getattr(info, "id", "") or "", info)
+ return entry.variants if entry is not None else None
+
+
+def info_has_local_gguf(info) -> bool:
+ """True when *info* points to on-disk GGUF weights the auto-switch path can load."""
+ return local_gguf_quants(info) is not None
def _build_index() -> dict[str, _LocalGgufEntry]:
@@ -146,10 +170,17 @@ def _build_index() -> dict[str, _LocalGgufEntry]:
_is_hidden_model,
)
from utils.paths import legacy_hf_cache_dir, hf_default_cache_dir, lmstudio_model_dirs
+ from utils.hf_cache_settings import known_hf_hub_caches
+ from core.inference.model_ids import public_model_id
index: dict[str, _LocalGgufEntry] = {}
seen_hf: set[str] = set()
+ try:
+ active_root = str(Path(_resolve_hf_cache_dir()).resolve())
+ except Exception:
+ active_root = None
+
def _scan_hf_once(directory) -> list:
if directory is None:
return []
@@ -161,7 +192,13 @@ def _build_index() -> dict[str, _LocalGgufEntry]:
if rp in seen_hf:
return []
seen_hf.add(rp)
- return _scan_hf_cache(directory)
+ # Only the active cache loads by repo id. Say so, or an inactive repo is
+ # indexed under an id it cannot load by, and its snapshot basename (what
+ # /v1/models advertises once loaded by path) is never a key at all.
+ # No format classification here: nothing on this path reads model_format,
+ # and its recursive walk would duplicate the one _local_gguf_entry already
+ # does per snapshot, on the request path.
+ return _scan_hf_cache(directory, active_cache = rp == active_root, classify_format = False)
except Exception as exc: # a missing/malformed root must skip, never crash the index
logger.debug("auto-switch: skipping HF cache dir %r: %s", directory, exc)
return []
@@ -174,7 +211,12 @@ def _build_index() -> dict[str, _LocalGgufEntry]:
except Exception as exc:
logger.debug("auto-switch: ./models scan failed: %s", exc)
try:
- for hf_dir in (_resolve_hf_cache_dir(), legacy_hf_cache_dir(), hf_default_cache_dir()):
+ for hf_dir in (
+ *known_hf_hub_caches(),
+ _resolve_hf_cache_dir(),
+ legacy_hf_cache_dir(),
+ hf_default_cache_dir(),
+ ):
found += _scan_hf_once(hf_dir)
except Exception as exc:
logger.debug("auto-switch: HF cache scan failed: %s", exc)
@@ -214,12 +256,91 @@ def _build_index() -> dict[str, _LocalGgufEntry]:
continue
# Index every alias (including the path) so a client can resolve by any of
# them, even though only the non-path loader_id is advertised.
- for key in (raw_id, getattr(info, "model_id", None), getattr(info, "display_name", None)):
+ for key in (
+ raw_id,
+ getattr(info, "model_id", None),
+ getattr(info, "display_name", None),
+ public_model_id(raw_id),
+ ):
if key:
index.setdefault(key.strip().lower(), entry)
+ # Other revisions of the same repo resolve to their own weights, so a pin on
+ # one keeps working after Hugging Face writes a newer snapshot.
+ for name, sibling_entry in _sibling_revision_entries(raw_id, loader_id):
+ index.setdefault(name.strip().lower(), sibling_entry)
return index
+def _sibling_revision_entries(raw_id: str, loader_id: str):
+ """Yield ``(revision_name, entry)`` for the repo's OTHER cached revisions.
+
+ An inactive-cache repo carries its snapshot path as the id, and /v1/models
+ advertises only that directory's basename once loaded, so anything durable
+ pinned to it (a subagent config) holds one revision hash. Hugging Face writes a
+ new snapshot dir on every update, and the scan emits a single entry per repo
+ pointed at the newest one, so that pin would otherwise stop resolving and drop
+ through to whatever model is loaded.
+
+ Each revision gets an entry for its OWN directory rather than an alias onto the
+ scanned one: aliasing would redirect a pin that names an older complete revision
+ onto a newer half-downloaded snapshot and break a request that works today.
+ Incomplete revisions are skipped for the same reason.
+
+ Sibling names are only revisions inside a real cache repo
+ (``/models--org--name/snapshots/``). A scan folder that merely happens
+ to be called ``snapshots`` holds unrelated models, and treating those as
+ revisions would silently serve one model in place of another.
+ """
+ from pathlib import Path
+ from types import SimpleNamespace
+
+ snapshots = Path(raw_id).parent
+ if snapshots.name != "snapshots" or not snapshots.parent.name.startswith("models--"):
+ return
+ from routes.models import snapshot_variants_all_complete
+
+ try:
+ siblings = [p for p in snapshots.iterdir() if p.is_dir() and p.name != Path(raw_id).name]
+ except OSError:
+ return
+ for sibling in siblings:
+ if not snapshot_variants_all_complete(str(sibling)):
+ continue
+ entry = _local_gguf_entry(loader_id, SimpleNamespace(path = str(sibling)))
+ if entry is not None:
+ yield sibling.name, entry
+
+
+def note_downloaded(repo_id: Optional[str]) -> None:
+ """Record a repo as present ahead of the scan that will index it."""
+ if not repo_id:
+ return
+ with _lock:
+ _just_downloaded.add(repo_id.strip().lower())
+
+
+def recently_downloaded(repo_id: str) -> bool:
+ """Whether *repo_id* finished downloading since the last completed scan."""
+ if not isinstance(repo_id, str) or not repo_id.strip():
+ return False
+ return repo_id.strip().lower() in _just_downloaded
+
+
+def invalidate_index() -> None:
+ """Mark the cached scan stale so the next resolve sees a just-finished download
+ instead of waiting out the TTL.
+
+ Keeps the entries: the request path reads this cache without scanning, so
+ emptying it would leave it with no evidence about any local model until the
+ rebuild lands, and a bare request for one would be answered by whatever is
+ resident. Only a completed download invalidates, and that only adds, so the
+ retained entries stay true.
+ """
+ global _scan
+ with _lock:
+ _scan = (0.0, _scan[1])
+
+
def _index() -> dict[str, _LocalGgufEntry]:
global _scan
# Build under the lock so concurrent callers with an expired cache don't all
@@ -234,23 +355,74 @@ def _index() -> dict[str, _LocalGgufEntry]:
# an install with many local models can itself exceed the TTL, which would
# store the cache already expired and make every request rebuild the index.
_scan = (time.monotonic(), fresh)
+ # The scan supersedes the notes: whatever landed is in the index now.
+ _just_downloaded.clear()
return fresh
-def resolve_local_gguf(requested: str) -> Optional[tuple[str, Optional[str], str]]:
+def index_is_built() -> bool:
+ """Whether a scan has ever completed, freshness aside.
+
+ Lock-free on purpose: ``_lock`` is held for the whole scan, so taking it would
+ park the request path on the scan it is trying to stay off. Safe because
+ ``_scan`` is only ever rebound, never mutated.
+ """
+ return bool(_scan[0])
+
+
+def warm_index_soon() -> None:
+ """(Re)build the index off the request path when it is missing or past its TTL.
+
+ The only refresh for callers using ``allow_scan=False``. Covers a stale index,
+ not just an absent one: a model downloaded through the Hub UI or dropped into a
+ scan folder has no invalidation hook and would otherwise stay invisible to them
+ for the life of the process. Never blocks, and never touches ``_lock``.
+ """
+ global _warming
+ if time.monotonic() - _scan[0] < max(_CACHE_TTL_S, _last_scan_s * _WARM_DUTY):
+ return
+ with _warm_lock:
+ if _warming:
+ return
+ _warming = True
+
+ def _run() -> None:
+ global _warming, _last_scan_s
+ started = time.monotonic()
+ try:
+ _index()
+ except Exception:
+ pass
+ finally:
+ _last_scan_s = time.monotonic() - started
+ with _warm_lock:
+ _warming = False
+
+ threading.Thread(target = _run, name = "local-model-index-warm", daemon = True).start()
+
+
+def resolve_local_gguf(
+ requested: str, *, allow_scan: bool = True
+) -> Optional[tuple[str, Optional[str], str]]:
"""Return ``(load_path, gguf_variant, loader_id)`` for a local match, else None.
``load_path`` is the concrete on-disk path to hand /load (so it never fetches
a remote), ``loader_id`` is the advertised id used as the launch-override key.
``requested`` is ``repo`` or ``repo:VARIANT``. An exact id match wins first
(so ids containing a colon still resolve); else the last ``:VARIANT`` is split
- off and resolves only when that quant is on disk.
+ off and resolves only when that quant is on disk, unless it names no quant at
+ all (an Ollama-style ":latest"), which means the repo.
+
+ ``allow_scan=False`` answers from the last built index and never rebuilds, for
+ the request path: the scan walks several model dirs and HF caches, takes seconds
+ on a large install, and holds a lock everyone queues behind. Stale is fine there,
+ since disk barely moves and a finished download calls :func:`invalidate_index`.
"""
if not isinstance(requested, str) or not requested.strip():
return None
requested = requested.strip()
try:
- index = _index()
+ index = _index() if allow_scan else _scan[1]
entry = index.get(requested.lower())
if entry is not None:
variant = entry.variants[0] if entry.variants else None
@@ -266,8 +438,44 @@ def resolve_local_gguf(requested: str) -> Optional[tuple[str, Optional[str], str
for v in entry.variants:
if v.lower() == wanted:
return entry.load_path, v, entry.loader_id
- return None
+ from core.inference.openai_auto_download import looks_like_quant
+
+ if looks_like_quant(variant):
+ return None
+ # ":latest" or ":8b" names no file, so it means the repo; a real quant that
+ # is not on disk still misses, or a swap would serve the wrong weights.
+ return entry.load_path, (entry.variants[0] if entry.variants else None), entry.loader_id
except Exception:
# Best-effort: any resolver failure falls through to the loaded model,
# so a malformed name can never turn a servable request into a 500.
return None
+
+
+MISS_MODEL_NOT_FOUND = "model_not_found"
+MISS_VARIANT_NOT_FOUND = "variant_not_found"
+
+
+def describe_local_miss(requested: str) -> tuple[str, tuple[str, ...]]:
+ """Why :func:`resolve_local_gguf` missed, so an error can say "wrong quant"
+ instead of "no such model".
+
+ ``(MISS_VARIANT_NOT_FOUND, )`` when the repo is downloaded but the
+ requested ``:VARIANT`` is not, else ``(MISS_MODEL_NOT_FOUND, ())``. Fail-safe: a
+ scan failure reports the generic miss rather than raising into the handler.
+ """
+ if not isinstance(requested, str) or not requested.strip():
+ return MISS_MODEL_NOT_FOUND, ()
+ base, sep, variant = requested.strip().rpartition(":")
+ from core.inference.openai_auto_download import looks_like_quant
+
+ # Split like the resolver or the two disagree: a tag naming no quant means the
+ # repo there, so reporting a missing quant for it would name one nobody asked for.
+ if not sep or not looks_like_quant(variant):
+ return MISS_MODEL_NOT_FOUND, ()
+ try:
+ entry = _index().get(base.strip().lower())
+ except Exception:
+ return MISS_MODEL_NOT_FOUND, ()
+ if entry is None or not entry.variants:
+ return MISS_MODEL_NOT_FOUND, ()
+ return MISS_VARIANT_NOT_FOUND, entry.variants
diff --git a/studio/backend/core/inference/mcp_client.py b/studio/backend/core/inference/mcp_client.py
index 0256df944e..98112c6d5b 100644
--- a/studio/backend/core/inference/mcp_client.py
+++ b/studio/backend/core/inference/mcp_client.py
@@ -971,7 +971,12 @@ def _call_stdio_tool(
raise RuntimeError("MCP server connection is not available")
else:
rem = _remaining()
- coro = _race_tool_call(session.client.call_tool(name, args), rem, cancel_event)
+ # raise_on_error=False for the same reason as the one-shot path.
+ coro = _race_tool_call(
+ session.client.call_tool(name, args, raise_on_error = False),
+ rem,
+ cancel_event,
+ )
return session.run(coro, rem)
except (_MCPCancelled, asyncio.TimeoutError):
# _race_tool_call cancels the pending call but cancellation is
diff --git a/studio/backend/core/inference/mlx_inference.py b/studio/backend/core/inference/mlx_inference.py
index e78c93b6f3..2b300a32b1 100644
--- a/studio/backend/core/inference/mlx_inference.py
+++ b/studio/backend/core/inference/mlx_inference.py
@@ -181,19 +181,27 @@ def _vlm_messages_have_tool_history(messages):
)
-def _build_generation_stats(prompt_n, prompt_tps, gen_n, gen_tps):
+def _build_generation_stats(
+ prompt_n,
+ prompt_tps,
+ gen_n,
+ gen_tps,
+ cached_n = 0,
+):
"""Map mlx stream stats onto the usage/timings shape llama-server emits."""
prompt_n = int(prompt_n or 0)
gen_n = int(gen_n or 0)
+ cached_n = int(cached_n or 0)
prompt_tps = float(prompt_tps or 0.0)
gen_tps = float(gen_tps or 0.0)
prompt_ms = (prompt_n / prompt_tps * 1000.0) if prompt_tps > 0 else 0.0
predicted_ms = (gen_n / gen_tps * 1000.0) if gen_tps > 0 else 0.0
+ total_prompt_n = prompt_n + cached_n
return {
"usage": {
- "prompt_tokens": prompt_n,
+ "prompt_tokens": total_prompt_n,
"completion_tokens": gen_n,
- "total_tokens": prompt_n + gen_n,
+ "total_tokens": total_prompt_n + gen_n,
},
"timings": {
"prompt_n": prompt_n,
@@ -204,11 +212,123 @@ def _build_generation_stats(prompt_n, prompt_tps, gen_n, gen_tps):
"predicted_ms": predicted_ms,
"predicted_per_token_ms": (predicted_ms / gen_n) if gen_n > 0 else 0.0,
"predicted_per_second": gen_tps,
- "cache_n": 0,
+ "cache_n": cached_n,
},
}
+PROMPT_CACHE_ENTRIES = 6
+PROMPT_CACHE_MEMORY_FRACTION = 0.15
+PROMPT_CACHE_FALLBACK_BYTES = 2 * 1024**3
+
+
+def _mlx_prompt_cache_api():
+ try:
+ from mlx_lm.models.cache import (
+ LRUPromptCache,
+ can_trim_prompt_cache,
+ make_prompt_cache,
+ trim_prompt_cache,
+ )
+ except ImportError:
+ return None
+ return LRUPromptCache, make_prompt_cache, can_trim_prompt_cache, trim_prompt_cache
+
+
+def _prompt_cache_max_bytes(recommended_gb = None):
+ override = os.environ.get("UNSLOTH_MLX_PROMPT_CACHE_BYTES")
+ if override:
+ try:
+ return max(int(override), 0)
+ except ValueError:
+ logger.warning("Ignoring non-integer UNSLOTH_MLX_PROMPT_CACHE_BYTES=%r", override)
+ if recommended_gb:
+ return int(recommended_gb * 1e9 * PROMPT_CACHE_MEMORY_FRACTION)
+ return PROMPT_CACHE_FALLBACK_BYTES
+
+
+def _flatten_kv_entries(cache):
+ for entry in cache:
+ nested = getattr(entry, "caches", None)
+ if nested is None:
+ yield entry
+ else:
+ yield from _flatten_kv_entries(nested)
+
+
+def _kv_prefix_coverage(cache):
+ covered = None
+ for entry in _flatten_kv_entries(cache):
+ offset = getattr(entry, "offset", None)
+ if offset is None:
+ return None
+ if getattr(entry, "start_position", 0):
+ return None
+ window = getattr(entry, "max_size", None)
+ if window is not None and offset > window:
+ return None
+ if covered is None:
+ covered = offset
+ elif covered != offset:
+ return None
+ return covered
+
+
+class _MLXPromptCacheHistory:
+ def __init__(self, max_entries, max_bytes):
+ api = _mlx_prompt_cache_api()
+ if api is None:
+ raise RuntimeError("mlx-lm is too old for LRUPromptCache")
+ lru_cls, make, can_trim, trim = api
+ self._make_prompt_cache = make
+ self._can_trim = can_trim
+ self._trim = trim
+ self._max_bytes = max_bytes
+ self._lru = lru_cls(max_size = max_entries, max_bytes = max_bytes)
+
+ def fetch(self, model, key, tokens):
+ cache, rest = self._lru.fetch_nearest_cache(key, list(tokens))
+ if cache is not None:
+ if rest:
+ return cache, list(rest)
+ if self._can_trim(cache) and self._trim(cache, 1) == 1:
+ return cache, list(tokens[-1:])
+ if len(tokens) > 1:
+ head = list(tokens[:-1])
+ cache, rest = self._lru.fetch_nearest_cache(key, head)
+ if cache is not None:
+ covered = len(head) - len(rest)
+ return cache, list(tokens[covered:])
+ return self._make_prompt_cache(model), list(tokens)
+
+ def insert(self, key, tokens, cache):
+ # An over-budget entry evicts itself and every other conversation.
+ nbytes = sum(getattr(entry, "nbytes", 0) for entry in cache)
+ if nbytes > self._max_bytes:
+ logger.debug(
+ "MLX prompt cache: skipping %.2f GB entry over the %.2f GB budget",
+ nbytes / 1e9,
+ self._max_bytes / 1e9,
+ )
+ return
+ covered = _kv_prefix_coverage(cache)
+ if covered is None:
+ logger.debug("MLX prompt cache: skipping cache with unverifiable prefix coverage")
+ return
+ tokens = list(tokens)
+ if covered > len(tokens):
+ logger.debug(
+ "MLX prompt cache: cache covers %d tokens but only %d were tracked",
+ covered,
+ len(tokens),
+ )
+ return
+ tokens = tokens[:covered]
+ if not tokens:
+ return
+ self._lru.insert_cache(key, tokens, cache)
+
+
def _mlx_distributed_rank_size(group = None):
"""Return ``(rank, world_size)`` for an optional MLX distributed group."""
if group is None:
@@ -313,6 +433,55 @@ class MLXInferenceBackend:
# Recorded for unload to release pinned memory back to the OS.
self._memory_limits_applied = {}
+ self._prompt_cache_history = None
+ self._prompt_cache_unavailable = False
+
+ def _prompt_cache(self):
+ if self._prompt_cache_history is not None or self._prompt_cache_unavailable:
+ return self._prompt_cache_history
+ max_bytes = _prompt_cache_max_bytes(self._memory_limits_applied.get("recommended_gb"))
+ if max_bytes <= 0:
+ self._prompt_cache_unavailable = True
+ logger.info("MLX prompt cache disabled by budget")
+ return None
+ try:
+ self._prompt_cache_history = _MLXPromptCacheHistory(
+ PROMPT_CACHE_ENTRIES,
+ max_bytes,
+ )
+ except Exception as exc:
+ self._prompt_cache_unavailable = True
+ logger.info("MLX prompt cache unavailable (%s); prefilling every request", exc)
+ return None
+ logger.info(
+ "MLX prompt cache: %d entries, %.2f GB budget",
+ PROMPT_CACHE_ENTRIES,
+ max_bytes / 1e9,
+ )
+ return self._prompt_cache_history
+
+ def _clear_prompt_cache(self):
+ self._prompt_cache_history = None
+ self._prompt_cache_unavailable = False
+
+ def _prepare_prompt_cache(self, prompt, adapter_state):
+ history = self._prompt_cache()
+ if history is None:
+ return prompt, None, None, None, 0
+ try:
+ tokenizer = self._tokenizer
+ bos = getattr(tokenizer, "bos_token", None)
+ add_special_tokens = bos is None or not prompt.startswith(bos)
+ tokens = list(tokenizer.encode(prompt, add_special_tokens = add_special_tokens))
+ if not tokens:
+ return prompt, None, None, None, 0
+ key = f"{self.active_model_name}|{adapter_state!r}"
+ cache, rest = history.fetch(self._model, key, tokens)
+ except Exception as exc:
+ logger.debug("MLX prompt cache lookup failed: %s", exc)
+ return prompt, None, None, None, 0
+ return rest, cache, key, tokens, len(tokens) - len(rest)
+
def _configure_memory_limits(self):
"""Apply Metal memory caps before loading a model.
@@ -535,6 +704,7 @@ class MLXInferenceBackend:
self._distributed_world_size = 1
if self.active_model_name == model_name:
self.active_model_name = None
+ self._clear_prompt_cache()
gc.collect()
mx.clear_cache()
@@ -731,24 +901,34 @@ class MLXInferenceBackend:
# prefix on every native-protocol snapshot just as the normal
# decoding path does below.
normalized_output = think_prefix
- logger.info(
- "Generating: prompt_len=%d, max_tokens=%d, model=%s, tokenizer=%s",
- len(prompt),
- max_new_tokens,
- type(self._model).__name__,
- type(self._tokenizer).__name__,
- )
with self._generation_lock, _temporary_mlx_adapter_state(self._model, _adapter_state):
+ (
+ gen_prompt,
+ prompt_cache,
+ cache_key,
+ prompt_tokens,
+ cached_n,
+ ) = self._prepare_prompt_cache(prompt, _adapter_state)
+ logger.info(
+ "Generating: prompt_len=%d, cached=%d, max_tokens=%d, model=%s, tokenizer=%s",
+ len(prompt),
+ cached_n,
+ max_new_tokens,
+ type(self._model).__name__,
+ type(self._tokenizer).__name__,
+ )
final_response = None
try:
# Enter request-scoped model state before yielding any response.
if think_prefix:
yield think_prefix
gen_kwargs = dict(
- prompt = prompt,
+ prompt = gen_prompt,
max_tokens = max_new_tokens,
sampler = sampler,
)
+ if prompt_cache is not None:
+ gen_kwargs["prompt_cache"] = prompt_cache
if logits_processors is not None:
gen_kwargs["logits_processors"] = logits_processors
for response in stream_generate(
@@ -757,6 +937,7 @@ class MLXInferenceBackend:
**gen_kwargs,
):
final_response = response
+ token_ids.append(response.token)
if preserve_native_channels:
piece = getattr(response, "text", None) or ""
delta = normalizer.feed(piece)
@@ -764,7 +945,6 @@ class MLXInferenceBackend:
normalized_output += delta
yield normalized_output
else:
- token_ids.append(response.token)
cumulative = self._tokenizer.decode(
token_ids,
skip_special_tokens = True,
@@ -773,6 +953,13 @@ class MLXInferenceBackend:
if cancel_event and cancel_event.is_set():
break
+ if prompt_cache is not None and prompt_tokens is not None:
+ history = self._prompt_cache_history
+ if history is not None:
+ try:
+ history.insert(cache_key, prompt_tokens + token_ids, prompt_cache)
+ except Exception as exc:
+ logger.debug("MLX prompt cache insert failed: %s", exc)
except Exception as e:
import traceback
logger.error("stream_generate failed:\n%s", traceback.format_exc())
@@ -785,6 +972,7 @@ class MLXInferenceBackend:
getattr(final_response, "prompt_tps", 0.0),
getattr(final_response, "generation_tokens", 0),
getattr(final_response, "generation_tps", 0.0),
+ cached_n,
)
if normalizer is not None:
cancelled = cancel_event is not None and cancel_event.is_set()
@@ -1001,7 +1189,8 @@ class MLXInferenceBackend:
**gen_kwargs,
)
- def reset_generation_state(self):
+ def reset_generation_state(self, caller_cancel_event = None):
+ # caller_cancel_event: signature parity with the orchestrator; unused here.
import mlx.core as mx
import gc
diff --git a/studio/backend/core/inference/model_ids.py b/studio/backend/core/inference/model_ids.py
index 548cc60f94..3886307ae2 100644
--- a/studio/backend/core/inference/model_ids.py
+++ b/studio/backend/core/inference/model_ids.py
@@ -39,10 +39,29 @@ def _looks_like_path(identifier: str) -> bool:
return False
+def hf_cache_repo_id(path: Optional[str]) -> Optional[str]:
+ """``.../models--org--name/snapshots/`` -> ``org/name``, else None.
+
+ A model loaded from the HF cache is identified by its snapshot dir, whose
+ basename is a commit hash; recover the repo id so callers don't show that.
+ """
+ if not path:
+ return None
+ parts = str(path).replace("\\", "/").split("/")
+ for index, part in enumerate(parts):
+ # Only inside the real cache layout: a "models--" name alone is not a repo id.
+ if part.startswith("models--") and parts[index + 1 : index + 2] == ["snapshots"]:
+ return part[len("models--") :].replace("--", "/")
+ return None
+
+
def public_model_id(identifier: Optional[str]) -> Optional[str]:
"""Return a clean, path-free public id for *identifier*.
- - Local GGUF path -> the file stem with ``.gguf`` stripped, e.g.
+ - HF cache path -> the repo id it came from, e.g.
+ ``~/.cache/huggingface/hub/models--unsloth--X-GGUF/snapshots/`` ->
+ ``unsloth/X-GGUF``.
+ - Other local GGUF path -> the file stem with ``.gguf`` stripped, e.g.
``/srv/models/Qwen3-30B-A3B-Q4_K_M.gguf`` -> ``Qwen3-30B-A3B-Q4_K_M``.
- HF repo id (``org/model``) and already-clean names -> returned unchanged.
- ``None`` / empty -> returned unchanged.
@@ -51,6 +70,9 @@ def public_model_id(identifier: Optional[str]) -> Optional[str]:
return identifier
if not _looks_like_path(identifier):
return identifier
+ repo_id = hf_cache_repo_id(identifier)
+ if repo_id:
+ return repo_id
name = os.path.basename(identifier.replace("\\", "/").rstrip("/"))
if name.lower().endswith(_GGUF_SUFFIX):
name = name[: -len(_GGUF_SUFFIX)]
diff --git a/studio/backend/core/inference/openai_auto_download.py b/studio/backend/core/inference/openai_auto_download.py
new file mode 100644
index 0000000000..cad5e40d14
--- /dev/null
+++ b/studio/backend/core/inference/openai_auto_download.py
@@ -0,0 +1,812 @@
+# SPDX-License-Identifier: AGPL-3.0-only
+# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0
+
+"""Opt-in: fetch a GGUF a /v1 request names but this server doesn't have.
+
+Auto-switch only loads models already on disk. With
+``openai_api_auto_download_model`` on, a miss that looks like a real Hub repo is
+fetched in the background and the request is told to retry rather than held
+open: a quant is routinely tens of GB, far longer than any client (or the
+Cloudflare edge on ``--secure``) will wait, and the inference lifecycle gate must
+not be held meanwhile. The resident model keeps serving, and the retry that lands
+after the download goes through the ordinary auto-switch path.
+
+Admission is deliberately narrow, since a request only needs an API key:
+- ``namespace/name`` only, and only when the Hub confirms GGUF weights. A
+ namespace is not evidence of intent (LiteLLM and OpenRouter address every
+ provider that way), so ``gpt-4`` and ``anthropic/claude-3.5-sonnet`` alike
+ fall through to the resident model as before.
+- GGUF only, decided from the remote file list, not the repo name: GGUF runs
+ under llama.cpp, which never imports repo Python.
+- ``auto_map`` is refused, so ``trust_remote_code`` is only ever granted
+ deliberately in the UI, never by an API call.
+- One download at a time, so a key holder cannot fan out fetches.
+"""
+
+from __future__ import annotations
+
+import asyncio
+import shutil
+import threading
+import time
+from dataclasses import dataclass
+from typing import Optional
+
+from loggers import get_logger
+
+logger = get_logger(__name__)
+
+# Keep the Hub probe short so a slow Hub can't stall the request path.
+_MODEL_INFO_TIMEOUT_S = 8.0
+# auth_check and hf_hub_download take no timeout of their own and run while the
+# provisional slot is held, so an unresponsive Hub would pin the single flight. The
+# code probe fetches up to three configs, so it gets more room than the auth call.
+_CODE_PROBE_TIMEOUT_S = 20.0
+# Headroom left free after the download, so filling the disk can't wedge the box.
+_DISK_RESERVE_BYTES = 5 * 1024**3
+_WATCH_POLL_S = 2.0
+# A stalled watcher must not pin the single-flight slot forever.
+_MAX_WATCH_S = 24 * 60 * 60
+# Past the watch window the row is resolved, so poll only to see whether the
+# worker is still alive and still owns the slot.
+_TIMED_OUT_POLL_S = 60.0
+_RETRY_AFTER_S = 30
+# Long enough for a client honouring Retry-After to come back and be told, short
+# enough that one that never returns cannot hold the slot.
+_FAILED_HOLD_S = 3 * _RETRY_AFTER_S
+_MAX_LISTED_VARIANTS = 8
+
+
+@dataclass(frozen = True)
+class AutoDownloadRefusal:
+ """Why this request cannot be served yet; the route raises it in the
+ surface's own error envelope."""
+
+ status: int
+ code: str
+ message: str
+ retry_after: Optional[int] = None
+
+
+@dataclass
+class _Active:
+ repo_id: str
+ # None while the Hub probe is still deciding which quant to fetch.
+ variant: Optional[str] = None
+ expected_bytes: int = 0
+ monitor_id: Optional[str] = None
+ started_at: float = 0.0
+ # Set when the worker failed. Held until a retry surfaces it: Retry-After is far
+ # longer than the watcher poll, so the client would restart the same failing download.
+ error: Optional[str] = None
+ failed_at: float = 0.0
+
+
+_lock = threading.Lock()
+_active: Optional[_Active] = None
+
+# Repos the Hub says are not servable, so a "vendor/model" miss doesn't re-probe every request.
+_NOT_SERVABLE_TTL_S = 10 * 60
+_NOT_SERVABLE_MAX = 256
+_cache_lock = threading.Lock()
+_not_servable: dict[str, float] = {}
+
+
+def _public_label(repo_id: str, variant: Optional[str]) -> str:
+ return f"{repo_id}:{variant}" if variant else repo_id
+
+
+def split_model_ref(requested: str) -> tuple[str, Optional[str]]:
+ """``org/repo:QUANT`` -> ``("org/repo", "QUANT")``; no suffix -> variant None.
+
+ Splits on the last colon. A slash-bearing suffix is only a variant when a real
+ Hub repo precedes it: "build/llama-13b" is a subdirectory GGUF key the catalog
+ advertises, while "C:/models/x.gguf" leaves a drive letter that is no repo id.
+ """
+ text = (requested or "").strip()
+ base, sep, suffix = text.rpartition(":")
+ if not sep or not base or not suffix:
+ return text, None
+ stripped = base.strip()
+ if "/" in suffix:
+ from hub.utils.paths import is_valid_repo_id
+ if "/" not in stripped or not is_valid_repo_id(stripped):
+ return text, None
+ return stripped, suffix.strip()
+
+
+def is_downloadable_ref(requested: str) -> bool:
+ """Whether *requested* is shaped like a Hub repo we may fetch.
+
+ Requires an explicit namespace: keeps ``gpt-4`` and other foreign ids falling
+ through, and stops ModelConfig.from_identifier's bare-name ``unsloth/``
+ prefixing from turning an unrelated label into a real repo.
+ """
+ from hub.utils.paths import is_valid_repo_id
+
+ repo_id, variant = split_model_ref(requested)
+ if "/" not in repo_id or not is_valid_repo_id(repo_id):
+ return False
+ if variant is not None:
+ from hub.utils.paths import is_valid_gguf_variant
+ return is_valid_gguf_variant(variant)
+ return True
+
+
+def looks_like_quant(variant: Optional[str]) -> bool:
+ """Whether a ``:suffix`` names a GGUF quant rather than a foreign tag.
+
+ Neither a namespace nor a colon proves a request was meant for this server
+ (``vendor/model`` is LiteLLM/OpenRouter, ``name:latest`` is Ollama). A real
+ quant label does.
+ """
+ import re
+
+ from utils.models.model_config import _GGUF_KNOWN_QUANT_RE
+
+ if not variant:
+ return False
+ # _extract_quant_label can append a bpw modifier (IQ4_XS-3.53bpw); still a quant.
+ label = re.sub(r"-[0-9]+(?:\.[0-9]+)?bpw$", "", variant.strip(), flags = re.IGNORECASE)
+ return _GGUF_KNOWN_QUANT_RE.fullmatch(label) is not None
+
+
+def _hub_token(hf_token: Optional[str]):
+ """The caller's token, or an explicit False. None makes huggingface_hub fall
+ back to a cached login (here the server owner's); only False is anonymous."""
+ return hf_token or False
+
+
+def _servable_key(repo_id: str, hf_token: Optional[str]) -> str:
+ """Cache key, per credential.
+
+ The Hub 404s a private repo the caller cannot see, so a tokenless verdict says
+ nothing about a caller who has one. Digested, so no token is held here.
+ """
+ import hashlib
+
+ seen_as = hashlib.sha256(hf_token.encode()).hexdigest()[:16] if hf_token else "anon"
+ return f"{repo_id.lower()}\n{seen_as}"
+
+
+def _mark_not_servable(repo_id: str, hf_token: Optional[str]) -> None:
+ with _cache_lock:
+ if len(_not_servable) >= _NOT_SERVABLE_MAX:
+ _not_servable.clear()
+ _not_servable[_servable_key(repo_id, hf_token)] = time.monotonic() + _NOT_SERVABLE_TTL_S
+
+
+def _is_not_servable(repo_id: str, hf_token: Optional[str]) -> bool:
+ key = _servable_key(repo_id, hf_token)
+ with _cache_lock:
+ expires = _not_servable.get(key)
+ if expires is None:
+ return False
+ if expires <= time.monotonic():
+ del _not_servable[key]
+ return False
+ return True
+
+
+def _gated_refusal(repo_id: str) -> AutoDownloadRefusal:
+ return AutoDownloadRefusal(
+ status = 403,
+ code = "model_access_denied",
+ message = (
+ f"'{repo_id}' is gated on Hugging Face. Accept its licence, then retry with "
+ "your own token in the X-Unsloth-HF-Token header: automatic download never "
+ "uses this server's Hugging Face identity."
+ ),
+ )
+
+
+async def _bounded_probe(fn, *args, timeout: float, default):
+ """Run a blocking Hub probe off the loop, bounding only the wait.
+
+ The thread is left to finish (a blocking socket read cannot be cancelled); the
+ caller takes *default*, chosen per call site so a timeout errs the safe way.
+ """
+ try:
+ return await asyncio.wait_for(asyncio.to_thread(fn, *args), timeout)
+ except (TimeoutError, asyncio.TimeoutError):
+ logger.debug("hub probe %s timed out after %ss", getattr(fn, "__name__", fn), timeout)
+ return default
+
+
+def _auth_denied(repo_id: str, hf_token: Optional[str]) -> bool:
+ """Whether this token lacks file access to a gated repo. False when the
+ check is inconclusive: the download's own auth is the real gate."""
+ from hub.utils.hf_errors import hf_error_status
+
+ try:
+ from huggingface_hub import auth_check
+ auth_check(repo_id, token = _hub_token(hf_token))
+ except Exception as exc:
+ return hf_error_status(exc) in (401, 403)
+ return False
+
+
+def _gguf_variants(siblings) -> dict[str, int]:
+ """Quant label -> bytes the download will actually fetch.
+
+ Mirrors list_gguf_variants for the selectable labels: companions (mmproj/MTP)
+ and big-endian builds are not quants, and sharded quants sum across shards.
+ Bytes come from the download plan, which folds companions back into every
+ quant, so the disk reserve is measured against what the worker fetches.
+ """
+ from hub.utils.gguf import extract_quant_label as canonical_quant_label
+ from hub.utils.gguf_plan import build_gguf_variant_plans
+ from utils.models.model_config import (
+ _extract_quant_label,
+ _is_big_endian_gguf_path,
+ _is_mmproj,
+ _is_mtp_drafter,
+ )
+
+ siblings = list(siblings or [])
+ plans = build_gguf_variant_plans(siblings)
+ sizes: dict[str, int] = {}
+ for sibling in siblings:
+ name = getattr(sibling, "rfilename", "") or ""
+ if not name.lower().endswith(".gguf"):
+ continue
+ quant = _extract_quant_label(name)
+ if not looks_like_quant(quant):
+ # With no recognized quant token the extractors part ways: this one takes
+ # the last hyphenated segment ("7b" of llama-7b) while the plan and worker
+ # key the whole stem, so advertising ours dispatches an unresolvable variant.
+ quant = canonical_quant_label(name) or quant
+ if _is_mmproj(name) or _is_mtp_drafter(name) or _is_big_endian_gguf_path(name, quant):
+ continue
+ plan = plans.get(quant.lower())
+ if plan is not None:
+ sizes[quant] = plan.download_size_bytes
+ else:
+ sizes[quant] = sizes.get(quant, 0) + int(getattr(sibling, "size", 0) or 0)
+ return sizes
+
+
+def _remaining_bytes(repo_id: str, plan, expected_bytes: int) -> int:
+ """Bytes still to fetch: a resumed quant or a companion shared with another
+ quant is already on disk, and charging for it can 507 a download that fits."""
+ try:
+ from hub.utils.download_registry import existing_blob_bytes
+
+ hashes = frozenset(
+ file.sha256 for file in getattr(plan, "expected_files", ()) or () if file.sha256
+ )
+ if not hashes:
+ return expected_bytes
+ return max(0, expected_bytes - existing_blob_bytes("model", repo_id, hashes))
+ except Exception:
+ return expected_bytes
+
+
+def _enough_disk(need_bytes: int) -> tuple[bool, int]:
+ """(fits, free_bytes). Fail-open on an unreadable cache root: the download
+ worker runs its own preflight, this only adds the reserve margin."""
+ try:
+ from hub.utils.hf_cache_state import hf_cache_root
+
+ root = hf_cache_root(create = True)
+ if root is None:
+ return True, 0
+ free = shutil.disk_usage(root).free
+ except Exception:
+ return True, 0
+ return free >= need_bytes + _DISK_RESERVE_BYTES, free
+
+
+def _gb(num_bytes: int) -> str:
+ return f"{num_bytes / 1024**3:.1f} GB"
+
+
+async def _job_state(repo_id: str, variant: Optional[str]) -> tuple[str, Optional[str]]:
+ from hub.services.models import downloads
+ try:
+ status = await downloads.get_download_status_response(repo_id, variant or "")
+ return status.state, status.error
+ except Exception as exc:
+ # "unknown", not "idle": idle ends the watch, and a failed probe proves nothing.
+ logger.debug("auto-download: status probe failed for %r: %s", repo_id, exc)
+ return "unknown", None
+
+
+async def _progress_percent(
+ repo_id: str, variant: Optional[str], expected_bytes: int, hf_token: Optional[str]
+) -> Optional[float]:
+ """0-100, or None. The hub service reports a 0-1 fraction, so scale it."""
+ from hub.services.models import downloads
+ try:
+ payload = await downloads.get_gguf_download_progress_response(
+ repo_id, variant or "", expected_bytes, hf_token
+ )
+ fraction = payload.get("progress")
+ if not isinstance(fraction, (int, float)):
+ return None
+ return min(100.0, max(0.0, float(fraction) * 100.0))
+ except Exception:
+ return None
+
+
+def _release(active: Optional[_Active]) -> None:
+ """Free the single-flight slot, but only while *active* still owns it.
+
+ Keying on ``repo_id`` alone let a stale operation clear a newer one: variant A
+ errors, an adopting request frees the slot, a retry starts B, then A's watcher
+ matches the repo and clears B, admitting a second download alongside it.
+ """
+ global _active
+ if active is None:
+ return
+ with _lock:
+ if _active is active:
+ _active = None
+
+
+async def _watch(active: _Active, hf_token: Optional[str]) -> None:
+ """Poll a dispatched job so the monitor row resolves and the resolver cache
+ is dropped the moment the weights land."""
+ from core.inference import api_monitor as monitor_module
+
+ api_monitor = monitor_module.api_monitor
+ deadline = time.monotonic() + _MAX_WATCH_S
+ timed_out = False
+ try:
+ while True:
+ await asyncio.sleep(_TIMED_OUT_POLL_S if timed_out else _WATCH_POLL_S)
+ state, error = await _job_state(active.repo_id, active.variant)
+ if state in ("running", "cancelling", "unknown"):
+ if timed_out:
+ # A running worker still owns the slot: releasing on the clock alone
+ # would admit a second multi-GB download beside it. "unknown" cannot
+ # confirm it is alive, so release then, or a broken probe wedges us.
+ if state == "unknown":
+ return
+ continue
+ if time.monotonic() >= deadline:
+ api_monitor.fail_open(active.monitor_id, "Download timed out")
+ timed_out = True
+ continue
+ # Only "running" has progress; the others are still in flight, so keep the slot.
+ if state == "running":
+ api_monitor.set_progress(
+ active.monitor_id,
+ await _progress_percent(
+ active.repo_id, active.variant, active.expected_bytes, hf_token
+ ),
+ )
+ continue
+ if state == "cancelled":
+ api_monitor.finish(active.monitor_id, status = "cancelled")
+ return
+ if state == "complete":
+ # No invalidate here: finalize_worker_exit already dropped the cache and
+ # warmed it; a second would mark that fresh scan stale and push a
+ # synchronous rescan onto the client's retry.
+ api_monitor.finish(active.monitor_id, status = "completed")
+ elif state == "idle":
+ # The job vanished without a terminal state (worker killed).
+ api_monitor.fail_open(active.monitor_id, "Download did not complete")
+ else:
+ api_monitor.fail_open(active.monitor_id, error or f"Download {state}")
+ # Keep the slot so the next retry is told it failed instead of
+ # silently restarting the same download.
+ active.error = error or f"Download {state}"
+ active.failed_at = time.monotonic()
+ return
+ return
+ except asyncio.CancelledError:
+ raise
+ except Exception as exc:
+ logger.warning("auto-download: watcher failed for %r: %s", active.repo_id, exc)
+ api_monitor.fail_open(active.monitor_id, "Download tracking failed")
+ finally:
+ if not active.failed_at:
+ _release(active)
+
+
+def _downloading_refusal(label: str, percent: Optional[float]) -> AutoDownloadRefusal:
+ progress = f" ({percent:.0f}% done)" if percent is not None else ""
+ return AutoDownloadRefusal(
+ status = 503,
+ code = "model_downloading",
+ message = (f"Downloading '{label}'{progress}. Retry shortly. Track it in Unsloth Studio."),
+ retry_after = _RETRY_AFTER_S,
+ )
+
+
+async def _is_downloadable_model(repo_id: str, hf_token: Optional[str]) -> bool:
+ """Whether the Hub has this repo with GGUF weights we could fetch.
+
+ Only asked while another download holds the slot, to tell a second download
+ apart from an ordinary foreign label. Any failure answers False: refusing
+ would strand normal traffic for the length of the download.
+ """
+ if _is_not_servable(repo_id, hf_token):
+ return False
+
+ def _probe():
+ from huggingface_hub import HfApi
+ return HfApi(token = _hub_token(hf_token)).model_info(repo_id, timeout = _MODEL_INFO_TIMEOUT_S)
+
+ try:
+ info = await asyncio.to_thread(_probe)
+ except Exception:
+ return False
+ # The same filter admission uses, not a bare .gguf test: mmproj, MTP drafters and
+ # big-endian builds are companions, not quants. Answering otherwise would hold an
+ # ordinary foreign label at model_download_busy for an unrelated download.
+ servable = bool(_gguf_variants(getattr(info, "siblings", None)))
+ if not servable:
+ _mark_not_servable(repo_id, hf_token)
+ return servable
+
+
+async def maybe_auto_download(
+ requested_model: str,
+ *,
+ hf_token: Optional[str] = None,
+ require_vision: bool = False,
+) -> Optional[AutoDownloadRefusal]:
+ """Start (or report on) a background fetch of *requested_model*.
+
+ Returns None when the request should carry on unchanged, or a refusal the
+ caller must raise. Only called after the local resolver has already missed.
+
+ ``require_vision`` refuses a target with no mmproj companion rather than spend
+ gigabytes on weights that cannot answer the request; the local capability guard
+ only ever sees an already-downloaded model.
+ """
+ global _active
+
+ repo_id, wanted_variant = split_model_ref(requested_model)
+ if not is_downloadable_ref(requested_model):
+ return None
+ if _is_not_servable(repo_id, hf_token) and not looks_like_quant(wanted_variant):
+ return None
+
+ # Settle the single-flight slot before the network, so retries during a download stay cheap.
+ busy: Optional[_Active] = None
+ with _lock:
+ current = _active
+ if current is not None and current.failed_at:
+ # A held failure only owns the slot until someone is told about it.
+ if current.repo_id != repo_id and time.monotonic() - current.failed_at > _FAILED_HOLD_S:
+ _active = current = None
+ if current is not None and current.repo_id == repo_id:
+ adopted = current
+ elif current is not None:
+ adopted = None
+ busy = current
+ else:
+ adopted = None
+ provisional = _Active(repo_id = repo_id, started_at = time.time())
+ _active = provisional
+
+ if busy is not None:
+ # Refusing before the probe blocks ordinary drop-in traffic: a namespaced label
+ # that is no downloadable GGUF repo (LiteLLM/OpenRouter style) would be told to
+ # wait out a multi-hour download. Only a downloadable label is a 2nd download.
+ if not await _is_downloadable_model(repo_id, hf_token):
+ return None
+ return AutoDownloadRefusal(
+ status = 503,
+ code = "model_download_busy",
+ message = (
+ f"Already downloading '{_public_label(busy.repo_id, busy.variant)}'. "
+ f"Retry '{requested_model}' once it finishes."
+ ),
+ retry_after = _RETRY_AFTER_S,
+ )
+
+ if adopted is not None:
+ if adopted.variant is None:
+ # Still probing: no job yet, and a stale whole-repo error would free the probe's slot.
+ return _downloading_refusal(adopted.repo_id, None)
+ state, error = await _job_state(adopted.repo_id, adopted.variant)
+ if state in ("running", "cancelling", "unknown"):
+ return _downloading_refusal(
+ _public_label(adopted.repo_id, adopted.variant),
+ await _progress_percent(
+ adopted.repo_id, adopted.variant, adopted.expected_bytes, hf_token
+ ),
+ )
+ if state == "error" or adopted.error:
+ error = error or adopted.error
+ # Surface once, then free the slot so a retry can start over.
+ _release(adopted)
+ return AutoDownloadRefusal(
+ status = 502,
+ code = "model_download_failed",
+ message = f"Downloading '{requested_model}' failed: {error or 'unknown error'}",
+ )
+ # complete/idle/cancelled: the watcher is about to free the slot, so retry once more.
+ return _downloading_refusal(
+ _public_label(adopted.repo_id, adopted.variant),
+ 100.0 if state == "complete" else None,
+ )
+
+ try:
+ return await _admit_and_start(
+ repo_id, wanted_variant, requested_model, hf_token, provisional, require_vision
+ )
+ except BaseException:
+ # Not `except Exception`: a cancel mid-probe would otherwise wedge the provisional slot.
+ _release(provisional)
+ raise
+
+
+async def _admit_and_start(
+ repo_id: str,
+ wanted_variant: Optional[str],
+ requested_model: str,
+ hf_token: Optional[str],
+ active: _Active,
+ require_vision: bool = False,
+) -> Optional[AutoDownloadRefusal]:
+ from hub.utils.hf_errors import hf_error_status
+
+ def _probe():
+ from huggingface_hub import HfApi
+ return HfApi(token = _hub_token(hf_token)).model_info(
+ repo_id, files_metadata = True, timeout = _MODEL_INFO_TIMEOUT_S
+ )
+
+ try:
+ info = await asyncio.to_thread(_probe)
+ except Exception as exc:
+ _release(active)
+ status = hf_error_status(exc)
+ if status == 401:
+ return AutoDownloadRefusal(
+ status = 401,
+ code = "model_access_denied",
+ message = (
+ f"Hugging Face rejected the token sent for '{repo_id}'. Replace the "
+ "X-Unsloth-HF-Token header with a valid token; retrying will not help."
+ ),
+ )
+ if status == 403:
+ return _gated_refusal(repo_id)
+ if status == 404:
+ _mark_not_servable(repo_id, hf_token)
+ # Unknown to the Hub reads as a foreign label; only an explicit quant makes it ours.
+ if not looks_like_quant(wanted_variant):
+ return None
+ # A private repo reads as absent without a token; don't confirm either way.
+ return AutoDownloadRefusal(
+ status = 404,
+ code = "model_not_found",
+ message = (
+ f"'{repo_id}' was not found on Hugging Face, or is not accessible. "
+ "If it is private, send a token in the X-Unsloth-HF-Token header."
+ ),
+ )
+ logger.warning("auto-download: Hub lookup failed for %r: %s", repo_id, exc)
+ return AutoDownloadRefusal(
+ status = 503,
+ code = "model_lookup_failed",
+ message = f"Could not reach Hugging Face to look up '{repo_id}'. Retry shortly.",
+ retry_after = _RETRY_AFTER_S,
+ )
+
+ # Inconclusive on timeout: the download's own auth is the real gate.
+ if getattr(info, "gated", False) and await _bounded_probe(
+ _auth_denied, repo_id, hf_token, timeout = _MODEL_INFO_TIMEOUT_S, default = False
+ ):
+ # Metadata for a gated repo is not file access; unchecked, the config read below lies.
+ _release(active)
+ return _gated_refusal(repo_id)
+
+ variants = _gguf_variants(getattr(info, "siblings", None))
+ if not variants:
+ _release(active)
+ _mark_not_servable(repo_id, hf_token)
+ if not looks_like_quant(wanted_variant):
+ return None
+ return AutoDownloadRefusal(
+ status = 400,
+ code = "model_not_supported",
+ message = (
+ f"'{repo_id}' has no GGUF weights. Automatic download serves GGUF only; "
+ "load other formats from Unsloth Studio."
+ ),
+ )
+
+ # trust_remote_code gate: _config_has_auto_map is tri-state, so refuse on True and on None.
+ from utils.security.consent import _config_has_auto_map
+
+ # _hub_token, not the raw token: None lets huggingface_hub fall back to a cached
+ # server login, so a caller-named repo would be probed with this server's identity.
+ # Defaults to None on timeout, which refuses: unchecked is not cleared.
+ has_auto_map = await _bounded_probe(
+ _config_has_auto_map,
+ repo_id,
+ _hub_token(hf_token),
+ timeout = _CODE_PROBE_TIMEOUT_S,
+ default = None,
+ )
+ if has_auto_map is not False:
+ _release(active)
+ unknown = has_auto_map is None
+ return AutoDownloadRefusal(
+ status = 403,
+ code = "remote_code_consent_required",
+ message = (
+ f"'{repo_id}' "
+ + (
+ "could not be checked for custom code"
+ if unknown
+ else "ships custom code that runs on load"
+ )
+ + ". Load it once in Unsloth Studio to review and approve it, then retry."
+ ),
+ )
+
+ variant = _match_variant(wanted_variant, variants)
+ if variant is None:
+ _release(active)
+ listed = sorted(variants)
+ shown = ", ".join(listed[:_MAX_LISTED_VARIANTS])
+ extra = len(listed) - _MAX_LISTED_VARIANTS
+ return AutoDownloadRefusal(
+ status = 404,
+ code = "model_not_found",
+ message = (
+ f"'{repo_id}' has no quant '{wanted_variant}'. Available quants: "
+ f"{shown}{f' and {extra} more' if extra > 0 else ''}."
+ ),
+ )
+
+ expected_bytes = variants[variant]
+ from hub.utils.gguf_plan import build_gguf_variant_plans
+
+ plan = build_gguf_variant_plans(list(getattr(info, "siblings", None) or [])).get(
+ variant.lower()
+ )
+ if require_vision and not (plan and plan.mmproj_filenames):
+ _release(active)
+ return AutoDownloadRefusal(
+ status = 400,
+ code = "invalid_value",
+ message = (
+ f"'{_public_label(repo_id, variant)}' ships no mmproj companion, so it "
+ "cannot answer the image or audio input in this request. It was not "
+ "downloaded."
+ ),
+ )
+
+ need_bytes = _remaining_bytes(repo_id, plan, expected_bytes)
+ fits, free = _enough_disk(need_bytes)
+ if not fits:
+ _release(active)
+ return AutoDownloadRefusal(
+ status = 507,
+ code = "insufficient_disk_space",
+ message = (
+ f"'{_public_label(repo_id, variant)}' needs {_gb(need_bytes)} plus "
+ f"{_gb(_DISK_RESERVE_BYTES)} headroom, but only {_gb(free)} is free."
+ ),
+ )
+
+ return await _dispatch(repo_id, variant, expected_bytes, requested_model, hf_token, active)
+
+
+def preferred_quant(labels) -> Optional[str]:
+ """The quant a plain load would pick from *labels*, or None.
+
+ The one ranking for "which quant did they mean": local resolution, remote
+ admission and /v1/models must agree, or a bare id means a different quant
+ depending on which of them answered it.
+ """
+ from utils.models.model_config import _pick_best_gguf
+
+ # _pick_best_gguf ranks filenames and matches upper-case tokens, so feed "