From ec46a5556836de2ee3e4daaf4c54c422b6c6b188 Mon Sep 17 00:00:00 2001 From: Daniel Han Date: Wed, 10 Jun 2026 11:24:39 -0700 Subject: [PATCH 01/42] Bump install.sh / install.ps1 pin to unsloth>=2026.6.2 (#6165) --- install.ps1 | 10 +++++----- install.sh | 10 +++++----- 2 files changed, 10 insertions(+), 10 deletions(-) diff --git a/install.ps1 b/install.ps1 index 16337f4b9c..cbc7b206c8 100644 --- a/install.ps1 +++ b/install.ps1 @@ -1830,7 +1830,7 @@ shell.Run cmd, 0, False if ($SkipTorch) { # No-torch: install unsloth + unsloth-zoo with --no-deps, then # runtime deps (typer, safetensors, transformers, etc.) with --no-deps. - $baseInstallExit = Invoke-InstallCommand { uv pip install --python $VenvPython --no-deps --reinstall-package unsloth --reinstall-package unsloth-zoo "unsloth>=2026.6.1" unsloth-zoo } + $baseInstallExit = Invoke-InstallCommand { uv pip install --python $VenvPython --no-deps --reinstall-package unsloth --reinstall-package unsloth-zoo "unsloth>=2026.6.2" unsloth-zoo } if ($baseInstallExit -eq 0) { # Resolve pydantic WITH deps so pip pins pydantic-core # to the matching version (no-torch-runtime.txt below @@ -1844,7 +1844,7 @@ shell.Run cmd, 0, False } } } else { - $baseInstallExit = Invoke-InstallCommand { uv pip install --python $VenvPython --reinstall-package unsloth --reinstall-package unsloth-zoo "unsloth>=2026.6.1" unsloth-zoo } + $baseInstallExit = Invoke-InstallCommand { uv pip install --python $VenvPython --reinstall-package unsloth --reinstall-package unsloth-zoo "unsloth>=2026.6.2" unsloth-zoo } } if ($baseInstallExit -ne 0) { Write-Host "[ERROR] Failed to install unsloth (exit code $baseInstallExit)" -ForegroundColor Red @@ -1891,7 +1891,7 @@ shell.Run cmd, 0, False if ($SkipTorch) { # No-torch: install unsloth + unsloth-zoo with --no-deps, then # runtime deps (typer, safetensors, transformers, etc.) with --no-deps. - $baseInstallExit = Invoke-InstallCommand { uv pip install --python $VenvPython --no-deps --upgrade-package unsloth --upgrade-package unsloth-zoo "unsloth>=2026.6.1" unsloth-zoo } + $baseInstallExit = Invoke-InstallCommand { uv pip install --python $VenvPython --no-deps --upgrade-package unsloth --upgrade-package unsloth-zoo "unsloth>=2026.6.2" unsloth-zoo } if ($baseInstallExit -eq 0) { # Same pydantic-with-deps trick as the migrated branch. $baseInstallExit = Invoke-InstallCommand { uv pip install --python $VenvPython pydantic } @@ -1903,7 +1903,7 @@ shell.Run cmd, 0, False } } } elseif ($StudioLocalInstall) { - $baseInstallExit = Invoke-InstallCommand { uv pip install --python $VenvPython --upgrade-package unsloth "unsloth>=2026.6.1" unsloth-zoo } + $baseInstallExit = Invoke-InstallCommand { uv pip install --python $VenvPython --upgrade-package unsloth "unsloth>=2026.6.2" unsloth-zoo } } else { $baseInstallExit = Invoke-InstallCommand { uv pip install --python $VenvPython --upgrade-package unsloth -- "$PackageName" } } @@ -1931,7 +1931,7 @@ shell.Run cmd, 0, False Write-TauriLog "STEP" "Installing unsloth" substep "installing unsloth (this may take a few minutes)..." if ($StudioLocalInstall) { - $baseInstallExit = Invoke-InstallCommand { uv pip install --python $VenvPython unsloth-zoo "unsloth>=2026.6.1" --torch-backend=auto } + $baseInstallExit = Invoke-InstallCommand { uv pip install --python $VenvPython unsloth-zoo "unsloth>=2026.6.2" --torch-backend=auto } if ($baseInstallExit -ne 0) { Write-Host "[ERROR] Failed to install unsloth (exit code $baseInstallExit)" -ForegroundColor Red return (Exit-InstallFailure "Failed to install unsloth (exit code $baseInstallExit)" $baseInstallExit) diff --git a/install.sh b/install.sh index af74410397..befcaaa186 100755 --- a/install.sh +++ b/install.sh @@ -2333,7 +2333,7 @@ if [ "$_MIGRATED" = true ]; then # to prevent transitive torch resolution. run_install_cmd "install unsloth (migrated no-torch)" uv pip install --python "$_VENV_PY" --no-deps \ --reinstall-package unsloth --reinstall-package unsloth-zoo \ - "unsloth>=2026.6.1" unsloth-zoo + "unsloth>=2026.6.2" unsloth-zoo # Resolve pydantic WITH deps so pip pins pydantic-core to the # matching version (no-torch-runtime.txt below is --no-deps). # All transitive deps are torch-free. @@ -2346,7 +2346,7 @@ if [ "$_MIGRATED" = true ]; then else run_install_cmd "install unsloth (migrated)" uv pip install --python "$_VENV_PY" \ --reinstall-package unsloth --reinstall-package unsloth-zoo \ - "unsloth>=2026.6.1" unsloth-zoo + "unsloth>=2026.6.2" unsloth-zoo fi if [ "$STUDIO_LOCAL_INSTALL" = true ]; then substep "overlaying local repo (editable)..." @@ -2550,7 +2550,7 @@ elif [ -n "$TORCH_INDEX_URL" ]; then # runtime deps (typer, safetensors, transformers, etc.) with --no-deps. run_install_cmd "install unsloth (no-torch)" uv pip install --python "$_VENV_PY" --no-deps \ --upgrade-package unsloth --upgrade-package unsloth-zoo \ - "unsloth>=2026.6.1" unsloth-zoo + "unsloth>=2026.6.2" unsloth-zoo # Same pydantic-with-deps trick as the migrated branch. run_install_cmd "install pydantic (with deps for compatible core)" \ uv pip install --python "$_VENV_PY" pydantic @@ -2568,7 +2568,7 @@ elif [ -n "$TORCH_INDEX_URL" ]; then fi elif [ "$STUDIO_LOCAL_INSTALL" = true ]; then run_install_cmd "install unsloth (local)" uv pip install --python "$_VENV_PY" \ - --upgrade-package unsloth "unsloth>=2026.6.1" unsloth-zoo + --upgrade-package unsloth "unsloth>=2026.6.2" unsloth-zoo substep "overlaying local repo (editable)..." run_install_cmd "overlay local repo" uv pip install --python "$_VENV_PY" -e "$_REPO_ROOT" --no-deps substep "overlaying unsloth-zoo from git main..." @@ -2600,7 +2600,7 @@ else tauri_log "STEP" "Installing Unsloth" substep "installing unsloth (this may take a few minutes)..." if [ "$STUDIO_LOCAL_INSTALL" = true ]; then - run_install_cmd "install unsloth (auto torch backend)" uv pip install --python "$_VENV_PY" unsloth-zoo "unsloth>=2026.6.1" --torch-backend=auto + run_install_cmd "install unsloth (auto torch backend)" uv pip install --python "$_VENV_PY" unsloth-zoo "unsloth>=2026.6.2" --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..." From bf2cd745b15922f0fed4649cbd67d563fbc2ec25 Mon Sep 17 00:00:00 2001 From: Leo Borcherding Date: Wed, 10 Jun 2026 23:00:11 -0500 Subject: [PATCH 02/42] Fix installer selecting ROCm torch on NVIDIA Linux hosts (#6174) * fix: prevent ROCm torch from installing on NVIDIA Linux hosts NVIDIA's open kernel module (driver 560+) registers GPU topology nodes in the KFD sysfs hierarchy with non-zero gpu_id values. The _has_amd_rocm_gpu (install.sh) and _has_rocm_gpu (install_python_stack.py) sysfs fallbacks previously treated any non-zero gpu_id as proof of an AMD GPU, so an NVIDIA-only host with the open kernel driver was misrouted to the ROCm install path, replacing the correctly-installed CUDA torch with ROCm wheels. Fixes: 1. install.sh _has_amd_rocm_gpu sysfs fallback: require vendor_id 4098 (AMD 0x1002) in the KFD node properties file before declaring an AMD GPU present. NVIDIA KFD nodes carry vendor_id 4318 (0x10DE) and are now skipped. 2. install_python_stack.py _has_rocm_gpu sysfs fallback: same vendor_id guard. Also preserves the existing fallback for older kernels that don't ship a properties file (trusts gpu_id alone there). 3. install.sh now exports UNSLOTH_TORCH_BACKEND ("cuda"/"rocm"/"cpu") immediately after get_torch_index_url() resolves the wheel family. install_python_stack.py reads this as _TORCH_BACKEND and short-circuits _ensure_rocm_torch() entirely on cuda/cpu hosts, providing a second layer of defense that is independent of subprocess GPU detection. Tests: 9 new cases in TestHasRocmGpuKfdVendorGuard, TestEnsureRocmTorch, and TestInstallShStructure cover all three changes. Full test_rocm_support.py suite: 289 passed, 2 skipped, 0 failed. Closes #6172 * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * fix: show actual torch backend in progress step labels The 'ROCm torch check' and 'ROCm torch (final)' step labels were hardcoded regardless of whether the installer was targeting CUDA, ROCm, or CPU. On NVIDIA hosts they showed 'ROCm' even though no ROCm wheels were being installed, which was misleading. Add _torch_step_label(suffix) which reads UNSLOTH_TORCH_BACKEND (set by install.sh) and formats the label as e.g. 'torch check (cuda)' or 'torch final (rocm)'. Falls back to live GPU detection for standalone studio update runs that bypass install.sh. * fix: make KFD sysfs vendor check conservative -- skip if no properties file The previous implementation fell through to `return True` when the KFD node's properties file was missing (OSError), intending to support older kernels. But NVIDIA open driver KFD nodes can also lack a properties file on some kernel versions, so the fallback still produced a false positive. Change the `except OSError: pass` to `continue` so any node without a readable properties file is skipped rather than trusted. KFD properties files exist on every kernel version that actually exposes /sys/class/kfd, so this does not regress real AMD GPU detection -- if the directory exists at all, properties files will be present for genuine GPU nodes. * fix: bulletproof NVIDIA vs AMD GPU detection Four changes that together ensure ROCm torch can never be installed on an NVIDIA host regardless of which detection path fires: 1. _has_rocm_gpu() (Python): NVIDIA guard at the top -- returns False immediately when _has_usable_nvidia_gpu() is True, blocking rocminfo, amd-smi, and KFD sysfs from producing a false positive even when ROCm tools are co-installed alongside the NVIDIA driver. 2. _has_amd_rocm_gpu() (install.sh): same NVIDIA guard -- calls _has_usable_nvidia_gpu first and returns 1 if it succeeds. 3. _has_usable_nvidia_gpu() (Python): adds /proc/driver/nvidia/gpus/ sysfs fallback. The NVIDIA driver populates this directory on Linux regardless of nvidia-smi state, so a subprocess PATH gap, timeout, or driver initialisation race can no longer silence NVIDIA detection. 4. _has_usable_nvidia_gpu() (install.sh): same /proc/driver/nvidia/gpus fallback, tried after nvidia-smi -L rather than instead of it. Together: NVIDIA wins at every decision point. If nvidia-smi works, it confirms NVIDIA. If it fails, /proc/driver/nvidia confirms NVIDIA. If somehow both fail, _has_rocm_gpu still checks NVIDIA first before any AMD path runs. * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * fix: two KFD/proc-only corner cases from Codex review 1. KFD awk state not reset per node file (Ryzen+NVIDIA false positive): The awk glob processes all topology node properties files in one pass. Without FNR==1 reset, a Ryzen+NVIDIA host where an AMD CPU-agent node sets amd=1 (vendor_id 4098, gpu_id 0) can combine with a later NVIDIA node setting gpu=1 (gpu_id > 0), triggering found=1 before vendor_id 4318 is seen. Added FNR==1{ gpu=0; amd=0 } to reset per file. 2. proc-only NVIDIA not reaching CUDA wheel selection: _has_usable_nvidia_gpu returning true via /proc/driver/nvidia fallback left _smi empty, so get_torch_index_url entered the AMD/CPU branch and selected CPU wheels despite NVIDIA being confirmed. Introduced _nvidia_detected flag (separate from _smi) so the AMD branch is skipped whenever NVIDIA is confirmed by any path, while _cuda_ver reads from _smi when available (with the existing cu126 fallback when _smi is absent). * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci --------- Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com> --- install.sh | 49 ++++- studio/install_python_stack.py | 115 +++++++++-- tests/studio/install/test_rocm_support.py | 233 +++++++++++++++++++++- 3 files changed, 366 insertions(+), 31 deletions(-) diff --git a/install.sh b/install.sh index befcaaa186..44d1c39d02 100755 --- a/install.sh +++ b/install.sh @@ -1717,8 +1717,15 @@ _ensure_rocm_probe_env() { # Returns 0 if an AMD GPU is present. Checks rocminfo, amd-smi, then sysfs # KFD topology (env-var-independent fallback for when HIP/ROCR_VISIBLE_DEVICES hides devices). +# Always returns 1 (false) when an NVIDIA GPU is present: blocks every +# detection path (rocminfo, amd-smi, KFD sysfs) from producing a false +# positive on NVIDIA-only or NVIDIA-primary hosts, even when ROCm tools +# are co-installed. _has_amd_rocm_gpu() { _ensure_rocm_probe_env + if _has_usable_nvidia_gpu; then + return 1 + fi if command -v rocminfo >/dev/null 2>&1 && \ rocminfo 2>/dev/null | awk '/Name:[[:space:]]*gfx[1-9][0-9]/{found=1} END{exit !found}'; then return 0 @@ -1726,27 +1733,42 @@ _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 '/gpu_id/{ if ($2+0 > 0) found=1 } END{ exit !found }' \ + 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 }' \ /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. return 0 fi return 1 } # ── NVIDIA usable-GPU helper ── -# Returns 0 (true) only if nvidia-smi is present AND actually lists a GPU. -# Prevents AMD-only hosts with a stale nvidia-smi on PATH from being routed -# into the CUDA branch. +# Returns 0 (true) if an NVIDIA GPU is present and usable. +# Primary probe: nvidia-smi -L. Fallback: /proc/driver/nvidia/gpus/ sysfs, +# which the NVIDIA driver populates on Linux regardless of nvidia-smi state +# -- handles PATH gaps, subprocess timeouts, and driver init races that +# could otherwise cause nvidia-smi to fail and silence NVIDIA detection. _has_usable_nvidia_gpu() { _nvsmi="" if command -v nvidia-smi >/dev/null 2>&1; then _nvsmi="nvidia-smi" elif [ -x "/usr/bin/nvidia-smi" ]; then _nvsmi="/usr/bin/nvidia-smi" - else - return 1 fi - "$_nvsmi" -L 2>/dev/null | awk '/^GPU[[:space:]]+[0-9]+:/{found=1} END{exit !found}' + if [ -n "$_nvsmi" ]; then + if "$_nvsmi" -L 2>/dev/null | awk '/^GPU[[:space:]]+[0-9]+:/{found=1} END{exit !found}'; then + return 0 + fi + fi + # Fallback: NVIDIA driver exposes one subdir per GPU under this path. + if [ -d /proc/driver/nvidia/gpus ] && \ + [ -n "$(ls -A /proc/driver/nvidia/gpus 2>/dev/null)" ]; then + return 0 + fi + return 1 } # ── Detect GPU and choose PyTorch index URL ── @@ -1763,14 +1785,16 @@ get_torch_index_url() { # packages) is not sufficient: otherwise an AMD-only host would # silently install CUDA wheels. _smi="" + _nvidia_detected=0 if _has_usable_nvidia_gpu; then + _nvidia_detected=1 if command -v nvidia-smi >/dev/null 2>&1; then _smi="nvidia-smi" elif [ -x "/usr/bin/nvidia-smi" ]; then _smi="/usr/bin/nvidia-smi" fi fi - if [ -z "$_smi" ]; then + if [ "$_nvidia_detected" -eq 0 ]; then # No NVIDIA GPU -- check for AMD ROCm GPU. # PyTorch only publishes ROCm wheels for linux-x86_64; skip the # ROCm branch entirely on aarch64 / arm64 / other architectures @@ -2098,6 +2122,15 @@ _maybe_bootstrap_rocm_wsl || true TORCH_INDEX_URL=$(get_torch_index_url) +# 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. +case "$TORCH_INDEX_URL" in + */rocm*|*/gfx*) export UNSLOTH_TORCH_BACKEND="rocm" ;; + */cpu) export UNSLOTH_TORCH_BACKEND="cpu" ;; + *) export UNSLOTH_TORCH_BACKEND="cuda" ;; +esac + # rocm7.2 ships torch 2.11.0 -- adjust the constraint to allow it. # All other ROCm tags and CUDA stay within <2.11.0. case "$TORCH_INDEX_URL" in diff --git a/studio/install_python_stack.py b/studio/install_python_stack.py index 0166bfd505..0460580922 100644 --- a/studio/install_python_stack.py +++ b/studio/install_python_stack.py @@ -557,7 +557,15 @@ def _persist_bnb_rocm_version(version: str) -> bool: def _has_rocm_gpu() -> bool: - """Return True only if an actual AMD GPU is visible (not just ROCm tools installed).""" + """Return True only if an actual AMD GPU is visible (not just ROCm tools installed). + + Always returns False when an NVIDIA GPU is present -- NVIDIA takes + priority on mixed hosts and prevents every detection path below + (rocminfo, amd-smi, KFD sysfs) from producing a false positive even + if ROCm tools are installed alongside the NVIDIA driver. + """ + if _has_usable_nvidia_gpu(): + return False for cmd, check_fn in ( # rocminfo: look for a real gfx GPU id (3-4 chars, nonzero first digit). # gfx000 is the CPU agent; ROCm 6.1+ also emits generic ISA lines like @@ -598,6 +606,13 @@ def _has_rocm_gpu() -> bool: # runtime-only detection. On minimal package-managed installs (no # rocminfo / no amd-smi tools), the kernel exposes AMD GPUs via # /sys/class/kfd so `studio update` can still detect and repair. + # + # Guard: reject any KFD node whose properties file reports a non-AMD + # vendor. With the NVIDIA open kernel module (driver 560+), NVIDIA GPUs + # can register KFD topology nodes with a non-zero gpu_id; those nodes + # have vendor_id 4318 (0x10DE) rather than the AMD value 4098 (0x1002). + # Without this check the fallback returns True on NVIDIA-only systems, + # causing _ensure_rocm_torch to install ROCm wheels on NVIDIA hardware. if sys.platform != "win32": try: kfd_nodes = "/sys/class/kfd/kfd/topology/nodes" @@ -609,29 +624,61 @@ def _has_rocm_gpu() -> bool: gpu_id = fh.read().strip() except OSError: continue - if gpu_id and gpu_id != "0": # gpu_id 0 = CPU node - return True + if not gpu_id or gpu_id == "0": # gpu_id 0 = CPU node + continue + # Require AMD vendor_id 4098 (0x1002) in the properties file. + # KFD properties files exist on every kernel that exposes + # /sys/class/kfd, so absence of the file means we cannot + # confirm AMD ownership -- skip the node rather than risk a + # false positive (e.g. NVIDIA open driver KFD nodes that + # lack a properties file on some kernel versions). + props_path = os.path.join(kfd_nodes, entry, "properties") + try: + with open(props_path) as fh: + props = fh.read() + except OSError: + continue # can't confirm vendor -- skip + if not re.search(r"\bvendor_id\s+4098\b", props): + continue + return True except OSError: pass return False def _has_usable_nvidia_gpu() -> bool: - """Return True only when nvidia-smi exists AND reports at least one GPU.""" + """Return True when an NVIDIA GPU is present and usable. + + Primary probe: nvidia-smi -L (subprocess). + Fallback: /proc/driver/nvidia/gpus/ sysfs (Linux only) -- handles the + case where nvidia-smi is present but the subprocess fails (PATH gap, + timeout, driver initialisation race). If either probe confirms an + NVIDIA GPU the function returns True so _has_rocm_gpu() is blocked. + """ exe = shutil.which("nvidia-smi") - if not exe: - return False - try: - result = subprocess.run( - [exe, "-L"], - stdout = subprocess.PIPE, - stderr = subprocess.DEVNULL, - text = True, - timeout = 10, - ) - except Exception: - return False - return result.returncode == 0 and "GPU " in result.stdout + if exe: + try: + result = subprocess.run( + [exe, "-L"], + stdout = subprocess.PIPE, + stderr = subprocess.DEVNULL, + text = True, + timeout = 10, + ) + if result.returncode == 0 and "GPU " in result.stdout: + return True + except Exception: + pass + # Fallback: the NVIDIA driver exposes one subdirectory per GPU under + # /proc/driver/nvidia/gpus/ on Linux regardless of nvidia-smi state. + if sys.platform != "win32": + try: + gpu_dir = "/proc/driver/nvidia/gpus" + if os.path.isdir(gpu_dir) and os.listdir(gpu_dir): + return True + except OSError: + pass + return False def _detect_amd_gfx_codes() -> list[str]: @@ -749,6 +796,13 @@ def _ensure_rocm_torch() -> None: Uses pip_install() to respect uv, constraints, and --python targeting. """ global _rocm_windows_torch_installed + # install.sh sets UNSLOTH_TORCH_BACKEND to the resolved wheel family + # ("cuda", "rocm", "cpu"). Skip ROCm operations entirely when install.sh + # already selected a non-ROCm backend -- this is the authoritative signal + # and avoids re-running GPU detection in a subprocess that may see a + # different environment (different PATH, CUDA_VISIBLE_DEVICES, etc.). + if _TORCH_BACKEND in ("cuda", "cpu"): + return # setup.ps1 sets this after installing AMD wheels; skip the probe only when # torch is actually importable as ROCm. If the venv was wiped between runs, # the stale env-var would suppress a needed reinstall. @@ -1088,6 +1142,29 @@ def _infer_no_torch() -> bool: NO_TORCH = _infer_no_torch() +# UNSLOTH_TORCH_BACKEND is set by install.sh after get_torch_index_url() so +# that this script knows which torch variant was selected without re-running +# GPU detection. Values: "cuda", "rocm", or "cpu". Empty means unknown +# (standalone `unsloth studio update` runs, where we re-detect normally). +_TORCH_BACKEND: str = os.environ.get("UNSLOTH_TORCH_BACKEND", "").lower() + + +def _torch_step_label(suffix: str) -> str: + """Return a progress label like 'torch check (cuda)' using the known backend. + + Falls back to GPU detection when UNSLOTH_TORCH_BACKEND is not set (e.g. + standalone `unsloth studio update` runs that bypass install.sh). + """ + backend = _TORCH_BACKEND + if not backend: + if _has_usable_nvidia_gpu(): + backend = "cuda" + elif _has_rocm_gpu(): + backend = "rocm" + else: + backend = "cpu" + return f"torch {suffix} ({backend})" + # -- Verbosity control ---------------------------------------------------------- # By default the installer shows a minimal in-place one-line progress bar. @@ -1770,7 +1847,7 @@ def install_python_stack() -> int: # venv got CPU-only torch (common when pip resolves torch from PyPI). # Must follow base packages so torch is present for inspection. if not IS_MACOS and not NO_TORCH: - _progress("ROCm torch check") + _progress(_torch_step_label("check")) _ensure_rocm_torch() # Windows + AMD GPU: warn if ROCm torch was not installed (wrong Python @@ -1955,7 +2032,7 @@ def install_python_stack() -> int: # Running the repair last ensures ROCm torch is in place at runtime, # whichever intermediate step clobbered it. if not IS_WINDOWS and not IS_MACOS and not NO_TORCH: - _progress("ROCm torch (final)") + _progress(_torch_step_label("final")) _ensure_rocm_torch() # 14. Final check (silent; third-party conflicts are expected) diff --git a/tests/studio/install/test_rocm_support.py b/tests/studio/install/test_rocm_support.py index 973c0a1e81..6bc94e3d0f 100644 --- a/tests/studio/install/test_rocm_support.py +++ b/tests/studio/install/test_rocm_support.py @@ -719,6 +719,143 @@ class TestEnsureRocmTorch: _ensure_rocm_torch() mock_pip.assert_not_called() + @patch.object(stack_mod, "pip_install") + @patch.object(stack_mod, "_has_rocm_gpu", return_value = True) + @patch.object(stack_mod, "_has_usable_nvidia_gpu", return_value = True) + def test_torch_backend_cuda_env_skips_entirely(self, mock_nvidia, mock_gpu, mock_pip): + """UNSLOTH_TORCH_BACKEND=cuda must short-circuit before any GPU probe.""" + with patch.dict(os.environ, {"UNSLOTH_TORCH_BACKEND": "cuda"}): + # Reload _TORCH_BACKEND from the patched environment. + with patch.object(stack_mod, "_TORCH_BACKEND", "cuda"): + _ensure_rocm_torch() + mock_pip.assert_not_called() + + @patch.object(stack_mod, "pip_install") + @patch.object(stack_mod, "_has_rocm_gpu", return_value = True) + @patch.object(stack_mod, "_has_usable_nvidia_gpu", return_value = True) + def test_torch_backend_cpu_env_skips_entirely(self, mock_nvidia, mock_gpu, mock_pip): + """UNSLOTH_TORCH_BACKEND=cpu must short-circuit before any GPU probe.""" + with patch.dict(os.environ, {"UNSLOTH_TORCH_BACKEND": "cpu"}): + with patch.object(stack_mod, "_TORCH_BACKEND", "cpu"): + _ensure_rocm_torch() + mock_pip.assert_not_called() + + +# TEST: install_python_stack.py -- _has_rocm_gpu KFD sysfs vendor_id guard + + +class TestHasRocmGpuKfdVendorGuard: + """Verify that the KFD sysfs fallback rejects non-AMD (NVIDIA) KFD nodes. + + These tests are source-level: they verify the regex and logic present in + the _has_rocm_gpu implementation rather than running the sysfs traversal + (which requires Linux path conventions). + """ + + def _src(self) -> str: + """Return the source of _has_rocm_gpu from install_python_stack.py.""" + import inspect + return inspect.getsource(stack_mod._has_rocm_gpu) + + def test_vendor_id_check_present(self): + """_has_rocm_gpu sysfs fallback must check vendor_id 4098 (AMD 0x1002).""" + src = self._src() + assert "vendor_id" in src, ( + "_has_rocm_gpu KFD sysfs fallback must read the properties file " + "to check vendor_id and exclude NVIDIA KFD nodes" + ) + assert "4098" in src, ( + "_has_rocm_gpu must require AMD vendor_id 4098 (0x1002) in the " + "KFD node properties to avoid false positives on NVIDIA systems" + ) + + def test_vendor_regex_pattern_anchored(self): + """The vendor_id regex must use a word boundary to avoid partial matches.""" + import re as _re + + src = self._src() + # The pattern should have a word boundary before and after the number + # so "vendor_id 41098" doesn't match "vendor_id 4098". + assert ( + _re.search(r"\\b.*vendor_id.*\\b", src) or "\\bvendor_id" in src + ), "_has_rocm_gpu vendor_id check should use word boundary anchors" + + def test_sysfs_fallback_guarded_by_non_win32(self): + """KFD sysfs fallback must be Linux-only (guarded by sys.platform != 'win32').""" + src = self._src() + assert "win32" in src, "_has_rocm_gpu sysfs fallback must be guarded by sys.platform check" + + def test_cpu_node_excluded(self): + """gpu_id == '0' must be excluded (CPU topology nodes).""" + src = self._src() + assert ( + '!= "0"' in src or "== '0'" in src or "!= '0'" in src or '"0"' in src + ), "_has_rocm_gpu must skip gpu_id 0 nodes (CPU nodes)" + + def test_install_sh_has_vendor_check(self): + """_has_amd_rocm_gpu in install.sh sysfs fallback must also check vendor_id 4098.""" + sh_path = PACKAGE_ROOT / "install.sh" + source = sh_path.read_text(encoding = "utf-8") + func_start = source.find("_has_amd_rocm_gpu()") + func_end = source.find("\n}", func_start) + func_body = source[func_start:func_end] + assert "vendor_id" in func_body, "_has_amd_rocm_gpu sysfs fallback must check vendor_id" + assert "4098" in func_body, "_has_amd_rocm_gpu must require AMD vendor_id 4098 (0x1002)" + + def test_has_rocm_gpu_returns_false_when_nvidia_present(self): + """_has_rocm_gpu must return False immediately when _has_usable_nvidia_gpu is True. + + This is the primary guard: even if rocminfo, amd-smi, or KFD sysfs + produce a false positive, an NVIDIA GPU always wins. + """ + with patch.object(stack_mod, "_has_usable_nvidia_gpu", return_value = True): + with patch("shutil.which", return_value = "/usr/bin/rocminfo"): + # Simulate rocminfo claiming an AMD GPU is present + mock_result = MagicMock() + mock_result.returncode = 0 + mock_result.stdout = "Name: gfx1100\n" + with patch("subprocess.run", return_value = mock_result): + assert not stack_mod._has_rocm_gpu(), ( + "_has_rocm_gpu must return False when NVIDIA GPU is detected, " + "regardless of what rocminfo reports" + ) + + def test_install_sh_has_rocm_gpu_nvidia_guard(self): + """_has_amd_rocm_gpu in install.sh must call _has_usable_nvidia_gpu and return 1 if true.""" + sh_path = PACKAGE_ROOT / "install.sh" + source = sh_path.read_text(encoding = "utf-8") + func_start = source.find("_has_amd_rocm_gpu()") + func_end = source.find("\n}", func_start) + func_body = source[func_start:func_end] + assert ( + "_has_usable_nvidia_gpu" in func_body + ), "_has_amd_rocm_gpu must call _has_usable_nvidia_gpu to block NVIDIA hosts" + assert ( + "return 1" in func_body + ), "_has_amd_rocm_gpu must return 1 (false) when NVIDIA GPU is detected" + + def test_has_usable_nvidia_gpu_proc_fallback_present(self): + """`_has_usable_nvidia_gpu` must have a /proc/driver/nvidia fallback.""" + import inspect + + src = inspect.getsource(stack_mod._has_usable_nvidia_gpu) + assert "/proc/driver/nvidia" in src, ( + "_has_usable_nvidia_gpu must fall back to /proc/driver/nvidia/gpus when " + "nvidia-smi subprocess fails, to handle PATH gaps and driver init races" + ) + + def test_install_sh_has_usable_nvidia_gpu_proc_fallback(self): + """_has_usable_nvidia_gpu in install.sh must also have a /proc/driver/nvidia fallback.""" + sh_path = PACKAGE_ROOT / "install.sh" + source = sh_path.read_text(encoding = "utf-8") + func_start = source.find("_has_usable_nvidia_gpu()") + func_end = source.find("\n}", func_start) + func_body = source[func_start:func_end] + assert "/proc/driver/nvidia" in func_body, ( + "_has_usable_nvidia_gpu in install.sh must fall back to " + "/proc/driver/nvidia/gpus when nvidia-smi fails" + ) + # TEST: install_python_stack.py -- _ROCM_TORCH_INDEX mapping @@ -927,16 +1064,21 @@ class TestInstallShStructure: source = sh_path.read_text(encoding = "utf-8") body = _extract_sh_function_body(source, "get_torch_index_url") nvidia_call = body.find("_has_usable_nvidia_gpu") - no_nvidia_branch = body.find('if [ -z "$_smi" ]') + # Gate changed from [ -z "$_smi" ] to [ "$_nvidia_detected" -eq 0 ] to + # handle proc-only NVIDIA hosts where nvidia-smi is absent but _has_usable_nvidia_gpu + # returns true via /proc/driver/nvidia/gpus. + no_nvidia_branch = body.find('if [ "$_nvidia_detected" -eq 0 ]') + if no_nvidia_branch < 0: + no_nvidia_branch = body.find('if [ -z "$_smi" ]') rocm_call = body.find("_has_amd_rocm_gpu") assert nvidia_call >= 0, "get_torch_index_url should call _has_usable_nvidia_gpu" - assert no_nvidia_branch >= 0, "get_torch_index_url should gate ROCm on no-nvidia-smi" + assert no_nvidia_branch >= 0, "get_torch_index_url should gate ROCm on no-nvidia branch" assert ( rocm_call > no_nvidia_branch - ), "ROCm detection should sit inside the 'no nvidia-smi' branch" + ), "ROCm detection should sit inside the 'no NVIDIA' branch" assert ( nvidia_call < no_nvidia_branch - ), "NVIDIA detection should run before the no-nvidia-smi branch" + ), "NVIDIA detection should run before the no-NVIDIA branch" def test_bitsandbytes_amd_install(self): """install.sh should install bitsandbytes for AMD when ROCm detected.""" @@ -1018,6 +1160,89 @@ class TestInstallShStructure: rocm_pos = func_body.find("amd-smi") assert darwin_pos < rocm_pos, "macOS check should come before ROCm detection" + def test_unsloth_torch_backend_exported_after_get_torch_index_url(self): + """install.sh must export UNSLOTH_TORCH_BACKEND after TORCH_INDEX_URL is set. + + This lets install_python_stack.py skip ROCm torch operations on CUDA + and CPU hosts without re-running GPU detection in a subprocess. + """ + sh_path = PACKAGE_ROOT / "install.sh" + source = sh_path.read_text(encoding = "utf-8") + torch_url_pos = source.find("TORCH_INDEX_URL=$(get_torch_index_url)") + backend_pos = source.find("UNSLOTH_TORCH_BACKEND") + assert backend_pos > 0, "UNSLOTH_TORCH_BACKEND must be set in install.sh" + assert ( + backend_pos > torch_url_pos + ), "UNSLOTH_TORCH_BACKEND must be set AFTER TORCH_INDEX_URL is resolved" + # Verify all three cases are covered + assert '"cuda"' in source[backend_pos : backend_pos + 500] + assert '"rocm"' in source[backend_pos : backend_pos + 500] + assert '"cpu"' in source[backend_pos : backend_pos + 500] + # Must be exported so subprocesses (setup.sh, install_python_stack.py) see it + assert "export UNSLOTH_TORCH_BACKEND" in source + + def test_kfd_sysfs_amd_vendor_check_in_has_amd_rocm_gpu(self): + """_has_amd_rocm_gpu sysfs fallback must require AMD vendor_id 4098. + + NVIDIA open kernel module (560+) registers KFD nodes with vendor_id + 4318 (0x10DE). Without the vendor check, _has_amd_rocm_gpu returns 0 + (true) on NVIDIA-only hosts that have the nvidia-open driver, causing + get_torch_index_url to select a ROCm wheel index. + """ + sh_path = PACKAGE_ROOT / "install.sh" + source = sh_path.read_text(encoding = "utf-8") + func_start = source.find("_has_amd_rocm_gpu()") + func_end = source.find("\n}", func_start) + func_body = source[func_start:func_end] + assert ( + "vendor_id" in func_body + ), "_has_amd_rocm_gpu sysfs fallback must check vendor_id to exclude NVIDIA KFD nodes" + assert ( + "4098" in func_body + ), "_has_amd_rocm_gpu sysfs fallback must require AMD vendor_id 4098 (0x1002)" + + def test_kfd_awk_resets_state_per_file(self): + """KFD sysfs awk must reset gpu/amd state per file (FNR==1). + + Without the reset, a Ryzen+NVIDIA host where node 0 is an AMD CPU + agent (vendor_id 4098, gpu_id 0) and node 1 is an NVIDIA GPU + (gpu_id > 0, vendor_id 4318) can produce a false positive: node 0 + sets amd=1, node 1 sets gpu=1, and the combined state triggers found=1 + before vendor_id 4318 is seen on node 1. + """ + sh_path = PACKAGE_ROOT / "install.sh" + source = sh_path.read_text(encoding = "utf-8") + func_start = source.find("_has_amd_rocm_gpu()") + func_end = source.find("\n}", func_start) + func_body = source[func_start:func_end] + assert "FNR==1" in func_body, ( + "_has_amd_rocm_gpu KFD awk must reset state per file with FNR==1 " + "to avoid false positives on Ryzen+NVIDIA hosts with multiple KFD nodes" + ) + + def test_get_torch_index_url_uses_nvidia_detected_flag(self): + """get_torch_index_url must track NVIDIA detection independently of _smi. + + When _has_usable_nvidia_gpu returns true via /proc/driver/nvidia fallback + but nvidia-smi is not on PATH, _smi stays empty. Without a separate + _nvidia_detected flag, the function falls into the AMD/CPU branch even + though NVIDIA was confirmed, silently installing CPU wheels instead of CUDA. + """ + sh_path = PACKAGE_ROOT / "install.sh" + source = sh_path.read_text(encoding = "utf-8") + func_start = source.find("get_torch_index_url()") + func_end = source.find("\n}", func_start) + func_body = source[func_start:func_end] + assert "_nvidia_detected" in func_body, ( + "get_torch_index_url must use a _nvidia_detected flag (separate from " + "_smi) so that proc-only NVIDIA detection still selects CUDA wheels" + ) + # The AMD/ROCm branch must be gated on _nvidia_detected being 0, not on + # _smi being empty. + assert ( + '_nvidia_detected" -eq 0' in func_body or "_nvidia_detected" in func_body + ), "get_torch_index_url AMD branch must be skipped when _nvidia_detected=1" + # TEST: Live regression on current host (NVIDIA B200 expected) From 9e2d17cc2d298b088876110b4192fc1ca28b9383 Mon Sep 17 00:00:00 2001 From: Daniel Han Date: Wed, 10 Jun 2026 21:51:16 -0700 Subject: [PATCH 03/42] Tests: follow Compare chat into the More submenu in the extra UI driver (#6177) PR #6153 taught the chat UI driver to find Compare chat inside the decluttered plus menu (#6140) but the extra UI driver kept looking for the old top-level item, so all three Studio UI CI workflows have been red on main since (Compare nav not found). Mirror the same hover-then-click More fallback here. --- tests/studio/playwright_extra_ui.py | 15 +++++++++++++++ 1 file changed, 15 insertions(+) diff --git a/tests/studio/playwright_extra_ui.py b/tests/studio/playwright_extra_ui.py index 20c7bda87c..26c3c244ca 100644 --- a/tests/studio/playwright_extra_ui.py +++ b/tests/studio/playwright_extra_ui.py @@ -332,6 +332,21 @@ with sync_playwright() as p: plus_btn.click(force = True) page.wait_for_timeout(400) compare_item = page.get_by_role("menuitem", name = re.compile(r"Compare chat", re.I)).first + if compare_item.count() == 0: + # Compare chat moved into the "More" submenu; hover, then click fallback. + more_trigger = page.get_by_role("menuitem", name = re.compile(r"^More$", re.I)).first + if more_trigger.count() > 0: + more_trigger.hover() + page.wait_for_timeout(400) + compare_item = page.get_by_role( + "menuitem", name = re.compile(r"Compare chat", re.I) + ).first + if compare_item.count() == 0: + more_trigger.click(force = True) + page.wait_for_timeout(400) + compare_item = page.get_by_role( + "menuitem", name = re.compile(r"Compare chat", re.I) + ).first if compare_item.count() > 0: compare_item.click(force = True) compare_opened = True From 8961cf154fd518fad5d1479784e5d920142066fa Mon Sep 17 00:00:00 2001 From: Michael Han <107991372+shimmyshimmer@users.noreply.github.com> Date: Wed, 10 Jun 2026 22:21:03 -0700 Subject: [PATCH 04/42] Give dark mode toasts a shadow and background separation (#6186) --- studio/frontend/src/index.css | 7 ++++--- 1 file changed, 4 insertions(+), 3 deletions(-) diff --git a/studio/frontend/src/index.css b/studio/frontend/src/index.css index 797e17420e..cb54d95e68 100644 --- a/studio/frontend/src/index.css +++ b/studio/frontend/src/index.css @@ -1751,10 +1751,11 @@ margin: 0 !important; } -/* Boost shadow on dark surfaces; mirrors .shadow-border / .menu-soft-surface pattern. */ +/* Composer shadow on the dark background color so toasts + do not merge into card-colored surfaces behind them. */ .dark [data-sonner-toast][data-styled='true'] { - background-color: var(--card) !important; - box-shadow: none !important; + background-color: var(--background) !important; + box-shadow: 0 2px 8px -2px rgba(0, 0, 0, 0.16) !important; } /* Selectable toast text; non-selectable toast buttons. */ From 53af33798dd72681c4826b4f6b32d75767e2027a Mon Sep 17 00:00:00 2001 From: oobabooga Date: Thu, 11 Jun 2026 02:55:43 -0300 Subject: [PATCH 05/42] Studio: forward `preserve_thinking` + `reasoning_effort` on the OpenAI passthrough (#6171) * Studio: forward preserve_thinking + reasoning_effort on the OpenAI passthrough * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * Provide _request_reasoning_kwargs on the responses passthrough test backend mock The OpenAI passthrough body builder now asks the active backend for capability-gated reasoning kwargs. The responses stream adapter test fakes the backend with a bare SimpleNamespace, so give it the same method a non-reasoning template would expose (returns None, keeping chat_template_kwargs out of the captured body). * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * Shorten the backend mock comment --------- Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com> Co-authored-by: Daniel Han --- studio/backend/routes/inference.py | 37 ++++++-- .../tests/test_openai_tool_passthrough.py | 94 +++++++++++++++++++ .../tests/test_responses_tool_passthrough.py | 4 + 3 files changed, 125 insertions(+), 10 deletions(-) diff --git a/studio/backend/routes/inference.py b/studio/backend/routes/inference.py index 70a5a3efe7..4cc76e7c2a 100644 --- a/studio/backend/routes/inference.py +++ b/studio/backend/routes/inference.py @@ -5037,7 +5037,9 @@ async def _responses_stream( detail = "Image provided but current GGUF model does not support vision.", ) - body = _build_openai_passthrough_body(chat_req, backend_ctx = llama_backend.context_length) + body = _build_openai_passthrough_body( + chat_req, backend_ctx = llama_backend.context_length, llama_backend = llama_backend + ) body["stream_options"] = {"include_usage": True} target_url = f"{llama_backend.base_url}/v1/chat/completions" @@ -6743,7 +6745,11 @@ def _extract_response_format(payload): return rf if isinstance(rf, dict) else None -def _build_openai_passthrough_body(payload, backend_ctx = None) -> dict: +def _build_openai_passthrough_body( + payload, + backend_ctx = None, + llama_backend = None, +) -> dict: """Assemble the llama-server request body from a ChatCompletionRequest. Only known OpenAI / llama-server fields are forwarded, so Studio-specific @@ -6754,12 +6760,19 @@ def _build_openai_passthrough_body(payload, backend_ctx = None) -> dict: system_prompt, _, _ = _extract_content_parts(payload.messages) messages = _set_or_prepend_system_message(messages, system_prompt) tool_choice = payload.tool_choice if payload.tool_choice is not None else "auto" - # When the caller asked for a specific reasoning mode, forward it via - # chat_template_kwargs so the Jinja template renders with (or without) the - # reasoning preamble. - tpl_kwargs = None - if payload.enable_thinking is not None: - tpl_kwargs = {"enable_thinking": bool(payload.enable_thinking)} + # Forward per-request reasoning fields (enable_thinking / reasoning_effort / + # preserve_thinking) via chat_template_kwargs so the Jinja template renders + # in the caller's mode, gated on the active template's capabilities exactly + # like the non-passthrough paths. + tpl_kwargs = ( + llama_backend._request_reasoning_kwargs( + payload.enable_thinking, + payload.reasoning_effort, + payload.preserve_thinking, + ) + if llama_backend is not None + else None + ) return _build_passthrough_payload( messages, payload.tools, @@ -6794,7 +6807,9 @@ async def _openai_passthrough_stream( the client sees a standard OpenAI response. """ target_url = f"{llama_backend.base_url}/v1/chat/completions" - body = _build_openai_passthrough_body(payload, backend_ctx = llama_backend.context_length) + body = _build_openai_passthrough_body( + payload, backend_ctx = llama_backend.context_length, llama_backend = llama_backend + ) _cancel_keys = (payload.cancel_id, payload.session_id, completion_id) _tracker = _TrackedCancel(cancel_event, *_cancel_keys) @@ -6939,7 +6954,9 @@ async def _openai_passthrough_non_streaming(llama_backend, payload, model_name): ``tool_calls``, and accurate ``usage`` token counts. """ target_url = f"{llama_backend.base_url}/v1/chat/completions" - body = _build_openai_passthrough_body(payload, backend_ctx = llama_backend.context_length) + body = _build_openai_passthrough_body( + payload, backend_ctx = llama_backend.context_length, llama_backend = llama_backend + ) try: async with httpx.AsyncClient() as client: diff --git a/studio/backend/tests/test_openai_tool_passthrough.py b/studio/backend/tests/test_openai_tool_passthrough.py index db4c07cdc8..d8a965c1f2 100644 --- a/studio/backend/tests/test_openai_tool_passthrough.py +++ b/studio/backend/tests/test_openai_tool_passthrough.py @@ -654,6 +654,100 @@ class TestBuildPassthroughPayloadToolChoice: ] +# ===================================================================== +# Passthrough reasoning kwargs — enable_thinking / reasoning_effort / +# preserve_thinking must reach llama-server via chat_template_kwargs, +# gated on template capabilities like the non-passthrough paths. +# ===================================================================== + + +def _reasoning_backend( + supports_reasoning = True, + reasoning_style = "enable_thinking", + reasoning_always_on = False, + supports_preserve_thinking = False, +): + """Bare LlamaCppBackend with just the reasoning capability flags set, + so _build_openai_passthrough_body exercises the real + _request_reasoning_kwargs gating.""" + from core.inference.llama_cpp import LlamaCppBackend + + backend = LlamaCppBackend.__new__(LlamaCppBackend) + backend._supports_reasoning = supports_reasoning + backend._reasoning_style = reasoning_style + backend._reasoning_always_on = reasoning_always_on + backend._supports_preserve_thinking = supports_preserve_thinking + return backend + + +class TestPassthroughReasoningKwargs: + def _payload(self, **fields): + return ChatCompletionRequest( + model = "default", + messages = [{"role": "user", "content": "hi"}], + **fields, + ) + + def test_enable_thinking_forwarded(self): + body = _build_openai_passthrough_body( + self._payload(enable_thinking = False), + backend_ctx = 4096, + llama_backend = _reasoning_backend(), + ) + assert body["chat_template_kwargs"] == {"enable_thinking": False} + + def test_preserve_thinking_forwarded_when_template_supports_it(self): + body = _build_openai_passthrough_body( + self._payload(enable_thinking = True, preserve_thinking = True), + backend_ctx = 4096, + llama_backend = _reasoning_backend(supports_preserve_thinking = True), + ) + assert body["chat_template_kwargs"] == { + "enable_thinking": True, + "preserve_thinking": True, + } + + def test_preserve_thinking_dropped_when_template_lacks_it(self): + body = _build_openai_passthrough_body( + self._payload(preserve_thinking = True), + backend_ctx = 4096, + llama_backend = _reasoning_backend(supports_preserve_thinking = False), + ) + assert "chat_template_kwargs" not in body + + def test_reasoning_effort_forwarded_for_effort_style_models(self): + body = _build_openai_passthrough_body( + self._payload(reasoning_effort = "high"), + backend_ctx = 4096, + llama_backend = _reasoning_backend(reasoning_style = "reasoning_effort"), + ) + assert body["chat_template_kwargs"] == {"reasoning_effort": "high"} + + def test_enable_thinking_maps_to_effort_for_effort_style_models(self): + body = _build_openai_passthrough_body( + self._payload(enable_thinking = False), + backend_ctx = 4096, + llama_backend = _reasoning_backend(reasoning_style = "reasoning_effort"), + ) + assert body["chat_template_kwargs"] == {"reasoning_effort": "low"} + + def test_always_on_reasoning_skips_thinking_kwargs(self): + body = _build_openai_passthrough_body( + self._payload(enable_thinking = False), + backend_ctx = 4096, + llama_backend = _reasoning_backend(reasoning_always_on = True), + ) + assert "chat_template_kwargs" not in body + + def test_no_reasoning_fields_omits_chat_template_kwargs(self): + body = _build_openai_passthrough_body( + self._payload(), + backend_ctx = 4096, + llama_backend = _reasoning_backend(supports_preserve_thinking = True), + ) + assert "chat_template_kwargs" not in body + + # ===================================================================== # OpenAI API compatibility helpers — verified spec edge cases # ===================================================================== diff --git a/studio/backend/tests/test_responses_tool_passthrough.py b/studio/backend/tests/test_responses_tool_passthrough.py index f0f2714214..beeee66b77 100644 --- a/studio/backend/tests/test_responses_tool_passthrough.py +++ b/studio/backend/tests/test_responses_tool_passthrough.py @@ -533,6 +533,10 @@ class TestResponsesStreamAdapter: is_vision = False, context_length = 4096, base_url = "http://llama.test", + # Non-reasoning template: the real backend returns None here. + _request_reasoning_kwargs = ( + lambda enable_thinking = None, reasoning_effort = None, preserve_thinking = None: None + ), ), ) From c8a7ba46438fa2cdfb0b170377af7260bb679439 Mon Sep 17 00:00:00 2001 From: Michael Han <107991372+shimmyshimmer@users.noreply.github.com> Date: Wed, 10 Jun 2026 23:55:20 -0700 Subject: [PATCH 06/42] Tone down the green glow on the login and password setup pages (#6185) * Make the auth page background glow much more subtle * Keep light mode rays at 0.35 opacity, dark mode at 0.15 --- studio/frontend/src/features/auth/change-password-page.tsx | 4 ++-- studio/frontend/src/features/auth/login-page.tsx | 4 ++-- 2 files changed, 4 insertions(+), 4 deletions(-) diff --git a/studio/frontend/src/features/auth/change-password-page.tsx b/studio/frontend/src/features/auth/change-password-page.tsx index 96ee8dd1a8..30f840faca 100644 --- a/studio/frontend/src/features/auth/change-password-page.tsx +++ b/studio/frontend/src/features/auth/change-password-page.tsx @@ -10,11 +10,11 @@ export function ChangePasswordPage() {
diff --git a/studio/frontend/src/features/auth/login-page.tsx b/studio/frontend/src/features/auth/login-page.tsx index eb33a6a71b..638df7d362 100644 --- a/studio/frontend/src/features/auth/login-page.tsx +++ b/studio/frontend/src/features/auth/login-page.tsx @@ -10,11 +10,11 @@ export function LoginPage() {
From 27d43a31f452b6761a7511a24e47df8389704402 Mon Sep 17 00:00:00 2001 From: Daniel Han Date: Thu, 11 Jun 2026 00:00:50 -0700 Subject: [PATCH 07/42] MLX CI: drop removed --simple-policy and stale ggml-org pin from the prebuilt step (#6189) #5963 folded the manifest resolver into the simple-path resolver, removed the installer's --simple-policy flag, and moved macOS prebuilts to the unslothai/llama.cpp fork. The MLX CI prebuilt step still passed the deleted flag, so argparse exits 2 and the workflow has been red on every main push since. Point the step at the fork with no extra flags, exactly like studio/setup.sh on macOS. --- .github/workflows/mlx-ci.yml | 51 +++++++++--------------------------- 1 file changed, 12 insertions(+), 39 deletions(-) diff --git a/.github/workflows/mlx-ci.yml b/.github/workflows/mlx-ci.yml index 75940832a0..221e86f235 100644 --- a/.github/workflows/mlx-ci.yml +++ b/.github/workflows/mlx-ci.yml @@ -231,33 +231,14 @@ jobs: tests/studio/test_is_mlx_dispatch_gate.py \ tests/studio/test_mlx_training_worker_behaviors.py - # Studio prebuilt llama.cpp install + GGUF inference. Drives the - # exact path Studio's setup.sh takes on macOS: invokes - # studio/install_llama_prebuilt.py with --published-repo - # ggml-org/llama.cpp and --published-release-tag b9049 (the - # latest llama.cpp release at the time this step was added; bump - # via UNSLOTH_LLAMA_TAG / DEFAULT_LLAMA_TAG when refreshing). - # The installer downloads llama-b9049-bin-macos-arm64.tar.gz, - # which is the universal Apple Silicon (arm64) build -- the - # same artifact works on M1/M2/M3/M4 because llama.cpp compiles - # against the ARMv8.2 baseline. - # - # The b9049 release also publishes: - # - llama-b9049-bin-macos-arm64-kleidiai.tar.gz - # KleidiAI dispatches at runtime; on M1 it falls back where - # ISA features (e.g. I8MM) are missing, so this asset also - # runs on M1 -- Studio just doesn't choose it by default. - # - llama-b9049-bin-macos-x64.tar.gz - # Intel-only; would only run on M1 via Rosetta 2 emulation, - # which we explicitly avoid. - # - iOS XCFramework - # iOS-app build artifact, unrelated to a macOS desktop CI. - # - # After install, downloads a small published GGUF - # (unsloth/gemma-3-270m-it-GGUF, Q4_K_M) from HuggingFace and - # runs the prebuilt llama-cli on it. Asserts the prompt echo - # appears in stdout. If the install fails OR the binary exits - # non-zero, that's an Unsloth/Studio bug. + # Studio prebuilt llama.cpp install + GGUF inference. Mirrors the + # path Studio's setup.sh takes on macOS since #5963: plan against + # the unslothai/llama.cpp fork's latest release, which ships the + # bin-macos-arm64 bundle plus the llama-prebuilt-manifest.json the + # default policy reads. After install, downloads a small published + # GGUF (unsloth/gemma-3-270m-it-GGUF, Q4_K_M) and validates + # llama-server /completion end to end. An install failure or a + # non-zero binary exit is an Unsloth/Studio bug. - name: Studio prebuilt llama.cpp install + GGUF inference (Mac M1) env: HF_TOKEN: ${{ secrets.HF_TOKEN }} @@ -272,20 +253,12 @@ jobs: set -euo pipefail INSTALL_DIR="$HOME/.unsloth-studio-prebuilt-test/llama.cpp" rm -rf "$INSTALL_DIR" - # --simple-policy is required when --published-repo points - # at upstream ggml-org/llama.cpp; that repo doesn't ship the - # llama-prebuilt-manifest.json asset Studio's default policy - # expects, so the simple platform-specific policy maps - # Darwin+arm64 -> bin-macos-arm64 directly. studio/setup.sh - # passes both --published-repo ggml-org/llama.cpp AND - # --simple-policy automatically on macOS, so this CI step - # exercises the same code path users hit when they run - # `curl -fsSL https://unsloth.ai/install.sh | sh`. + # Mirror studio/setup.sh on macOS (the install.sh user path): + # it plans against the unslothai/llama.cpp fork's latest + # release with no policy or tag flags. python studio/install_llama_prebuilt.py \ --install-dir "$INSTALL_DIR" \ - --published-repo ggml-org/llama.cpp \ - --published-release-tag b9049 \ - --simple-policy + --published-repo unslothai/llama.cpp # Studio bundles only llama-server + llama-quantize from the # prebuilt (not llama-cli) -- inference goes through From 898d3dd0b5937625f1724665e3892c7a905efbe7 Mon Sep 17 00:00:00 2001 From: Daniel Han Date: Thu, 11 Jun 2026 02:45:12 -0700 Subject: [PATCH 08/42] Studio: offer the in-app llama.cpp update for source-build (markerless) installs (#6188) * Studio: offer the in-app llama.cpp update for source-build (markerless) installs Source-build installs have no UNSLOTH_PREBUILT_INFO.json marker, so freshness reported supported=False and the Update button never showed (notably on macOS, where the fork shipped no prebuilt before b9585 and setup fell back to a source build). When an install has no marker but an official prebuilt now exists for the host, surface the update and let one click swap it in place. - install_llama_prebuilt.py: published_repo_for_host() (the setup.sh host->repo rule in Python) and a --resolve-prebuilt mode that reports whether a prebuilt exists for this host without downloading. - llama_cpp_update.py: markerless branch in get_update_status/start_update, version-suppressed so source builds already newer than latest are not nagged; fail-open throughout. * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * Studio: run llama update detection off the event loop, expose source_build The markerless source-build check probes the host and reads GitHub, so run get_update_status and start_update in a worker thread to keep the API responsive. Expose source_build in the status response so the banner can label the source-build switch. * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * Keep the llama-route auth stub out of sys.modules for the rest of the suite test_llama_route.py replaced sys.modules['auth.authentication'] with a bare stub at collection time and never restored it, so every later test importing create_access_token got the stub: 17 failures across test_desktop_auth, test_middleware, test_openai_tool_passthrough and test_rag_preview on all four Backend CI Python versions. Import the real module when its deps are available and only stub in minimal envs, popping the stubs after the standalone route load either way. * Studio: address review on the source-build update path - published_repo_for_host: route CPU-only Windows to ggml-org too (mirrors setup.ps1; the fork ships no win-cpu bundle), macOS always the fork. - markerless detection compares/display the upstream llama_tag, not a possible fork wrapper release_tag, so a source build is not wrongly judged newer. - do not offer when there is no resolvable install root (a pinned LLAMA_SERVER_PATH outside a managed dir): an apply would not take effect. * Ignore version probes in the update tests' subprocess capture The status polls in these tests trigger the new source-build detection, which shells out to llama-server --version through the same patched subprocess.run. On slow runners that probe lands after the installer call and clobbers the single captured argv, failing the flag assertions (seen on the 3.10/3.11 Backend CI jobs). Skip probe calls in all three fakes so only the installer invocation is captured. * Skip markerless re-detection while the update job is swapping the tree On a source-build install the frontend polls update-status every 3s during an apply, and each poll ran _source_build_status, which execs the very llama-server binary the job is concurrently replacing. On Windows that exec can hold the exe long enough to fail the installer's os.replace; everywhere it is a per-poll subprocess spawn for a status the poller does not read (it only consumes job progress). Gate the markerless branch on the job not running; the marked path is probe-free and still returns the live job state. * Studio: tighten source-build update root, repo routing, and downgrade guard Only manage a markerless install when the active binary lives under a resolvable llama.cpp root (marker dir, UNSLOTH_LLAMA_CPP_PATH it sits in, or a llama.cpp ancestor); a pinned LLAMA_SERVER_PATH or a PATH/system binary is left alone so an apply cannot install where it would not take effect. Gate start_update on the same suppression as detection so a direct POST cannot downgrade a source build newer than the latest prebuilt. Route Linux hosts with AMD tooling (rocminfo/amd-smi/hipconfig/ hipinfo) to the fork in --resolve-prebuilt, matching setup.sh, so a HIP source build is not offered an upstream CPU prebuilt. * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * Studio: cover inactive env root and pinned llama.cpp checkout in update root tests --------- Co-authored-by: danielhanchen Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com> --- studio/backend/routes/llama.py | 11 +- .../tests/test_install_resolve_prebuilt.py | 173 ++++++++++ studio/backend/tests/test_llama_cpp_update.py | 325 +++++++++++++++++- studio/backend/tests/test_llama_route.py | 111 ++++++ studio/backend/utils/llama_cpp_update.py | 209 ++++++++++- studio/install_llama_prebuilt.py | 65 ++++ 6 files changed, 873 insertions(+), 21 deletions(-) create mode 100644 studio/backend/tests/test_install_resolve_prebuilt.py create mode 100644 studio/backend/tests/test_llama_route.py diff --git a/studio/backend/routes/llama.py b/studio/backend/routes/llama.py index 11cc4c798e..b30c0f5aea 100644 --- a/studio/backend/routes/llama.py +++ b/studio/backend/routes/llama.py @@ -13,6 +13,7 @@ never blocks on a missing marker / offline GitHub. from __future__ import annotations +import asyncio from typing import Optional from fastapi import APIRouter, Depends, Query @@ -48,6 +49,9 @@ class LlamaUpdateStatusResponse(BaseModel): published_repo: Optional[str] = None installed_at_utc: Optional[str] = None age_days: Optional[int] = None + source_build: bool = Field( + False, description = "True when there is no marker (source build) but a prebuilt is offered." + ) job: LlamaUpdateJob = Field(default_factory = LlamaUpdateJob) @@ -65,11 +69,14 @@ async def llama_update_status( ), current_subject: str = Depends(get_current_subject), ) -> LlamaUpdateStatusResponse: - return LlamaUpdateStatusResponse(**get_update_status(force_refresh = force_refresh)) + # Off the event loop: detection may probe the host and read GitHub. + status = await asyncio.to_thread(get_update_status, force_refresh = force_refresh) + return LlamaUpdateStatusResponse(**status) @router.post("/update", response_model = LlamaUpdateActionResponse) async def llama_update( current_subject: str = Depends(get_current_subject), ) -> LlamaUpdateActionResponse: - return LlamaUpdateActionResponse(**start_update()) + action = await asyncio.to_thread(start_update) + return LlamaUpdateActionResponse(**action) diff --git a/studio/backend/tests/test_install_resolve_prebuilt.py b/studio/backend/tests/test_install_resolve_prebuilt.py new file mode 100644 index 0000000000..ede5629664 --- /dev/null +++ b/studio/backend/tests/test_install_resolve_prebuilt.py @@ -0,0 +1,173 @@ +# SPDX-License-Identifier: AGPL-3.0-only +# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0 + +"""install_llama_prebuilt.py: host->repo mapping and the --resolve-prebuilt mode. + +These back the in-app update for source-build (markerless) installs: the backend +asks the installer whether an official prebuilt exists for this host without +downloading. Network and host detection are stubbed; no GPU or internet needed. +""" + +from __future__ import annotations + +import importlib +import json +import sys +from pathlib import Path +from types import SimpleNamespace + +import pytest + +_studio = Path(__file__).resolve().parent.parent.parent +if str(_studio) not in sys.path: + sys.path.insert(0, str(_studio)) + +ilp = importlib.import_module("install_llama_prebuilt") + +if not hasattr(ilp, "published_repo_for_host") or not hasattr( + ilp, "resolve_simple_install_release_plans" +): + pytest.skip("PR symbols not present - check branch", allow_module_level = True) + +FORK = ilp.DEFAULT_PUBLISHED_REPO # unslothai/llama.cpp +UPSTREAM = ilp.UPSTREAM_REPO # ggml-org/llama.cpp + + +def _host(**kw): + base = dict( + system = "Linux", + machine = "x86_64", + is_windows = False, + is_linux = False, + is_macos = False, + is_x86_64 = False, + is_arm64 = False, + nvidia_smi = None, + driver_cuda_version = None, + compute_caps = [], + visible_cuda_devices = None, + has_physical_nvidia = False, + has_usable_nvidia = False, + has_rocm = False, + rocm_gfx_target = None, + macos_version = None, + ) + base.update(kw) + return ilp.HostInfo(**base) + + +def test_published_repo_for_host(): + # CPU-only Linux (x64 and arm64) -> ggml-org upstream. + assert ilp.published_repo_for_host(_host(is_linux = True, is_x86_64 = True)) == UPSTREAM + assert ( + ilp.published_repo_for_host(_host(is_linux = True, is_arm64 = True, machine = "aarch64")) + == UPSTREAM + ) + # GPU Linux -> fork. + assert ( + ilp.published_repo_for_host(_host(is_linux = True, is_x86_64 = True, has_usable_nvidia = True)) + == FORK + ) + assert ilp.published_repo_for_host(_host(is_linux = True, is_x86_64 = True, has_rocm = True)) == FORK + # CPU-only Windows -> ggml-org (setup.ps1: the fork ships no win-cpu bundle). + assert ( + ilp.published_repo_for_host(_host(system = "Windows", is_windows = True, is_x86_64 = True)) + == UPSTREAM + ) + # GPU Windows -> fork. + assert ( + ilp.published_repo_for_host( + _host(system = "Windows", is_windows = True, is_x86_64 = True, has_usable_nvidia = True) + ) + == FORK + ) + # macOS -> fork regardless of GPU (ggml-org macOS bundles need too-new macOS). + assert ( + ilp.published_repo_for_host( + _host(system = "Darwin", is_macos = True, is_arm64 = True, machine = "arm64") + ) + == FORK + ) + # Linux with AMD tooling but no probed GPU -> fork (setup.sh routes on tooling). + assert ( + ilp.published_repo_for_host( + _host(is_linux = True, is_x86_64 = True), linux_amd_tooling_present = True + ) + == FORK + ) + # The tooling hint is Linux-only: Windows CPU stays on ggml-org. + assert ( + ilp.published_repo_for_host( + _host(system = "Windows", is_windows = True, is_x86_64 = True), + linux_amd_tooling_present = True, + ) + == UPSTREAM + ) + + +def _run_resolve(monkeypatch, capsys, plans_or_exc): + monkeypatch.setattr( + ilp, + "detect_host", + lambda: _host(system = "Darwin", is_macos = True, is_arm64 = True, machine = "arm64"), + ) + + def _resolver(tag, host, repo, published_release_tag): + if isinstance(plans_or_exc, Exception): + raise plans_or_exc + return ("b9585", plans_or_exc) + + monkeypatch.setattr(ilp, "resolve_simple_install_release_plans", _resolver) + monkeypatch.setattr( + sys, + "argv", + ["install_llama_prebuilt.py", "--resolve-prebuilt", "latest", "--output-format", "json"], + ) + rc = ilp.main() + assert rc == ilp.EXIT_SUCCESS + return json.loads(capsys.readouterr().out.strip().splitlines()[-1]) + + +def test_resolve_prebuilt_available(monkeypatch, capsys): + plan = SimpleNamespace( + release_tag = "b9585", + llama_tag = "b9585", + attempts = [ + SimpleNamespace(name = "llama-b9585-bin-macos-arm64.tar.gz", install_kind = "macos-arm64") + ], + ) + out = _run_resolve(monkeypatch, capsys, [plan]) + assert out["prebuilt_available"] is True + assert out["repo"] == FORK + assert out["release_tag"] == "b9585" + assert out["asset"] == "llama-b9585-bin-macos-arm64.tar.gz" + assert out["install_kind"] == "macos-arm64" + + +def test_resolve_prebuilt_unavailable(monkeypatch, capsys): + out = _run_resolve(monkeypatch, capsys, ilp.PrebuiltFallback("no macOS asset")) + assert out["prebuilt_available"] is False + assert out["repo"] == FORK + + +def test_resolve_prebuilt_linux_amd_tooling_routes_to_fork(monkeypatch, capsys): + # CPU-probed Linux host but rocminfo on PATH: the dispatch must route to the + # fork so a HIP source build is not offered an upstream CPU prebuilt. + monkeypatch.setattr(ilp, "detect_host", lambda: _host(is_linux = True, is_x86_64 = True)) + monkeypatch.setattr(ilp.shutil, "which", lambda tool: tool == "rocminfo") + seen = {} + + def _resolver(tag, host, repo, published_release_tag): + seen["repo"] = repo + raise ilp.PrebuiltFallback("no asset") + + monkeypatch.setattr(ilp, "resolve_simple_install_release_plans", _resolver) + monkeypatch.setattr( + sys, + "argv", + ["install_llama_prebuilt.py", "--resolve-prebuilt", "latest", "--output-format", "json"], + ) + assert ilp.main() == ilp.EXIT_SUCCESS + out = json.loads(capsys.readouterr().out.strip().splitlines()[-1]) + assert seen["repo"] == FORK + assert out["repo"] == FORK diff --git a/studio/backend/tests/test_llama_cpp_update.py b/studio/backend/tests/test_llama_cpp_update.py index d2c854cdee..05a107377c 100644 --- a/studio/backend/tests/test_llama_cpp_update.py +++ b/studio/backend/tests/test_llama_cpp_update.py @@ -60,24 +60,151 @@ def _write_install( def _clean_state(monkeypatch): freshness.reset_caches() upd._reset_job_for_tests() + upd._resolve_memo.clear() + # Deterministic markerless paths: no host-pinned binary, no custom dir. + monkeypatch.delenv("LLAMA_SERVER_PATH", raising = False) + monkeypatch.delenv("UNSLOTH_LLAMA_CPP_PATH", raising = False) # Never hit the network in these tests. monkeypatch.setattr(freshness, "_fetch_latest_release_tag", lambda repo, timeout = 5.0: None) yield freshness.reset_caches() upd._reset_job_for_tests() + upd._resolve_memo.clear() -def test_status_no_marker(monkeypatch, tmp_path): +def _no_prebuilt(monkeypatch): + """Stub the host prebuilt probe to 'none available' (no source-build offer).""" + monkeypatch.setattr(upd, "_resolve_prebuilt_for_host", lambda *, force_refresh = False: None) + + +def _prebuilt( + monkeypatch, + *, + repo = "unslothai/llama.cpp", + release_tag = "b9585", + llama_tag = None, + asset = None, +): + """Stub the host prebuilt probe to report an available prebuilt.""" + payload = { + "prebuilt_available": True, + "repo": repo, + "release_tag": release_tag, + "llama_tag": llama_tag or release_tag, + "asset": asset or f"llama-{release_tag}-bin-macos-arm64.tar.gz", + "install_kind": "macos-arm64", + } + monkeypatch.setattr(upd, "_resolve_prebuilt_for_host", lambda *, force_refresh = False: payload) + + +def test_status_no_marker_no_prebuilt(monkeypatch, tmp_path): + # No marker AND no prebuilt available for the host -> unsupported (the genuine + # source-build-with-nothing-to-offer case). binary = tmp_path / "build" / "bin" / "llama-server" binary.parent.mkdir(parents = True) binary.write_text("stub") # no marker file alongside monkeypatch.setattr(upd, "_find_binary", lambda: str(binary)) + _no_prebuilt(monkeypatch) st = upd.get_update_status() assert st["supported"] is False assert st["update_available"] is False assert st["installed_tag"] is None +def test_status_source_build_offers_prebuilt(monkeypatch, tmp_path): + # Markerless source build with a prebuilt now available for the host: surface + # the update. Unknown installed version (source build) is treated as behind. + binary = tmp_path / "llama.cpp" / "build" / "bin" / "llama-server" + binary.parent.mkdir(parents = True) + binary.write_text("stub") + monkeypatch.setattr(upd, "_find_binary", lambda: str(binary)) + _prebuilt(monkeypatch, release_tag = "b9585") + monkeypatch.setattr(upd, "_installed_build_number", lambda b: None) + st = upd.get_update_status() + assert st["supported"] is True + assert st["update_available"] is True + assert st["source_build"] is True + assert st["latest_tag"] == "b9585" + assert st["published_repo"] == "unslothai/llama.cpp" + + +def test_status_source_build_compares_llama_tag(monkeypatch, tmp_path): + # release_tag may be a fork wrapper (v1.0); compare/display the upstream + # llama_tag (b9457) so a source build is not wrongly judged newer. + binary = tmp_path / "llama.cpp" / "build" / "bin" / "llama-server" + binary.parent.mkdir(parents = True) + binary.write_text("stub") + monkeypatch.setattr(upd, "_find_binary", lambda: str(binary)) + _prebuilt(monkeypatch, release_tag = "v1.0", llama_tag = "b9457") + monkeypatch.setattr(upd, "_installed_build_number", lambda b: 9000) + st = upd.get_update_status() + assert st["latest_tag"] == "b9457" # not the wrapper tag + assert st["update_available"] is True # 9000 < 9457 + + +def test_status_source_build_pinned_binary_not_offered(monkeypatch, tmp_path): + # LLAMA_SERVER_PATH pins a custom binary outside any llama.cpp dir; an apply + # could not take effect, so the button must not surface. + binary = tmp_path / "custom" / "llama-server" + binary.parent.mkdir(parents = True) + binary.write_text("stub") + monkeypatch.setenv("LLAMA_SERVER_PATH", str(binary)) + monkeypatch.setattr(upd, "_find_binary", lambda: str(binary)) + _prebuilt(monkeypatch) + st = upd.get_update_status() + assert st["supported"] is False + assert st["update_available"] is False + + +def test_llama_install_root_pinned_returns_none(monkeypatch, tmp_path): + binary = tmp_path / "custom" / "llama-server" + binary.parent.mkdir(parents = True) + binary.write_text("stub") + monkeypatch.setenv("LLAMA_SERVER_PATH", str(binary)) + assert upd._llama_install_root(str(binary)) is None + + +def test_status_source_build_suppressed_when_newer(monkeypatch, tmp_path): + # A source build already newer than the latest prebuilt is not nagged. + binary = tmp_path / "llama.cpp" / "build" / "bin" / "llama-server" + binary.parent.mkdir(parents = True) + binary.write_text("stub") + monkeypatch.setattr(upd, "_find_binary", lambda: str(binary)) + _prebuilt(monkeypatch, release_tag = "b9518") + monkeypatch.setattr(upd, "_installed_build_number", lambda b: 9600) + st = upd.get_update_status() + assert st["supported"] is True + assert st["update_available"] is False + assert st["installed_tag"] == "b9600" + + +def test_status_source_build_skips_probe_while_job_runs(monkeypatch, tmp_path): + # While the updater swaps the tree, status polls must not exec the binary + # being replaced (on Windows that exec can fail the installer's os.replace); + # the 3s poller only consumes job progress. + binary = tmp_path / "build" / "bin" / "llama-server" + binary.parent.mkdir(parents = True) + binary.write_text("stub") + monkeypatch.setattr(upd, "_find_binary", lambda: str(binary)) + probes = {"resolve": 0, "version": 0} + + def _count_resolve(*, force_refresh = False): + probes["resolve"] += 1 + return None + + def _count_version(b): + probes["version"] += 1 + return None + + monkeypatch.setattr(upd, "_resolve_prebuilt_for_host", _count_resolve) + monkeypatch.setattr(upd, "_installed_build_number", _count_version) + with upd._job_lock: + upd._job["state"] = upd._JOB_RUNNING + st = upd.get_update_status() + assert st["job"]["state"] == "running" + assert probes == {"resolve": 0, "version": 0} + + def test_status_update_available(monkeypatch, tmp_path): binary = _write_install(tmp_path, "b9493") monkeypatch.setattr(upd, "_find_binary", lambda: binary) @@ -99,13 +226,62 @@ def test_status_up_to_date(monkeypatch, tmp_path): assert st["update_available"] is False -def test_start_update_no_marker_refuses(monkeypatch, tmp_path): +def test_start_update_no_marker_no_prebuilt_refuses(monkeypatch, tmp_path): binary = tmp_path / "llama-server" binary.write_text("stub") # no marker monkeypatch.setattr(upd, "_find_binary", lambda: str(binary)) + monkeypatch.setattr(upd, "_installer_script", lambda: tmp_path / "install_llama_prebuilt.py") + _no_prebuilt(monkeypatch) res = upd.start_update() assert res["started"] is False - assert res["reason"] == "no_prebuilt_marker" + assert res["reason"] == "no_prebuilt_available" + + +def test_start_update_source_build_installs_prebuilt(monkeypatch, tmp_path): + # Markerless install + available prebuilt: install in place into the resolved + # root, with the asset-derived ROCm forwarding and the resolved repo. + install_dir = tmp_path / "llama.cpp" + binary = install_dir / "build" / "bin" / "llama-server" + binary.parent.mkdir(parents = True) + binary.write_text("stub") # no marker + monkeypatch.delenv("UNSLOTH_LLAMA_CPP_PATH", raising = False) + monkeypatch.setattr(upd, "_find_binary", lambda: str(binary)) + monkeypatch.setattr(upd, "_installer_script", lambda: tmp_path / "install_llama_prebuilt.py") + _prebuilt( + monkeypatch, repo = "unslothai/llama.cpp", asset = "app-b9585-linux-x64-rocm-gfx110X.tar.gz" + ) + + captured = {} + + class _Proc: + returncode = 0 + stdout = "installed" + stderr = "" + + def _fake_run(cmd, **kwargs): + cmd = list(cmd) + # Status polls probe `llama-server --version`; keep the installer argv. + if "--version" in cmd: + return _Proc() + captured["cmd"] = cmd + _write_install(install_dir, "b9585") # installer writes the marker + return _Proc() + + monkeypatch.setattr(upd.subprocess, "run", _fake_run) + + res = upd.start_update() + assert res["started"] is True, res + deadline = time.time() + 10 + while time.time() < deadline: + if upd.get_update_status()["job"]["state"] in ("success", "error"): + break + time.sleep(0.05) + cmd = captured["cmd"] + assert "--install-dir" in cmd and str(install_dir) in cmd + assert "--published-repo" in cmd and "unslothai/llama.cpp" in cmd + assert "--llama-tag" in cmd and "latest" in cmd + assert cmd[cmd.index("--rocm-gfx") + 1] == "gfx110x" + assert "--simple-policy" not in cmd and "--cpu-fallback" not in cmd def test_start_update_happy_path(monkeypatch, tmp_path): @@ -123,6 +299,10 @@ def test_start_update_happy_path(monkeypatch, tmp_path): stderr = "" def _fake_run(cmd, **kwargs): + cmd = list(cmd) + # Status polls probe `llama-server --version`; keep the installer argv. + if "--version" in cmd: + return _Proc() captured["cmd"] = cmd # Simulate the installer writing a new marker with the latest tag. _write_install(install_dir, "b9518") @@ -229,7 +409,11 @@ def _capture_install_cmd( stderr = "" def _fake_run(cmd, **kwargs): - captured["cmd"] = list(cmd) + cmd = list(cmd) + # Status polls probe `llama-server --version`; keep the installer argv. + if "--version" in cmd: + return _Proc() + captured["cmd"] = cmd _write_install(install_dir, latest, repo = repo, asset = asset) return _Proc() @@ -442,3 +626,136 @@ def test_update_fails_open_when_backend_unavailable(monkeypatch, tmp_path): break time.sleep(0.05) assert job["state"] == "success", job + + +# --- markerless helper units --- + + +def test_resolve_prebuilt_parses_and_caches(monkeypatch, tmp_path): + monkeypatch.setattr(upd, "_installer_script", lambda: tmp_path / "install_llama_prebuilt.py") + calls = {"n": 0} + + class _Proc: + returncode = 0 + # stderr noise plus the JSON line on stdout (installer logs to stderr). + stdout = ( + '{"prebuilt_available": true, "repo": "unslothai/llama.cpp", "release_tag": "b9585"}' + ) + stderr = "[llama-prebuilt] some log\n" + + def _fake_run(cmd, **kwargs): + calls["n"] += 1 + assert "--resolve-prebuilt" in cmd + return _Proc() + + monkeypatch.setattr(upd.subprocess, "run", _fake_run) + res = upd._resolve_prebuilt_for_host() + assert res["prebuilt_available"] is True and res["release_tag"] == "b9585" + # Second call is memoized (no second subprocess). + upd._resolve_prebuilt_for_host() + assert calls["n"] == 1 + + +def test_resolve_prebuilt_fails_open(monkeypatch, tmp_path): + monkeypatch.setattr(upd, "_installer_script", lambda: tmp_path / "install_llama_prebuilt.py") + + def _boom(cmd, **kwargs): + raise OSError("subprocess failed") + + monkeypatch.setattr(upd.subprocess, "run", _boom) + assert upd._resolve_prebuilt_for_host() is None + # Failures are not cached: a later success is observed. + + class _Proc: + returncode = 0 + stdout = '{"prebuilt_available": false}' + stderr = "" + + monkeypatch.setattr(upd.subprocess, "run", lambda cmd, **kw: _Proc()) + assert upd._resolve_prebuilt_for_host() == {"prebuilt_available": False} + + +def test_installed_build_number(monkeypatch): + def _ver(text): + class _Proc: + returncode = 0 + stdout = "" + stderr = text + + monkeypatch.setattr(upd.subprocess, "run", lambda cmd, **kw: _Proc()) + return upd._installed_build_number("/bin/llama-server") + + assert _ver("version: 9585 (abc1234)\nbuilt with clang\n") == 9585 + assert _ver("version: 1 (deadbee)\n") is None # source build without tags + assert _ver("no version here") is None + assert upd._installed_build_number(None) is None + + +def test_llama_install_root_finds_llama_cpp_ancestor(monkeypatch, tmp_path): + root = tmp_path / "llama.cpp" + binary = root / "build" / "bin" / "llama-server" + binary.parent.mkdir(parents = True) + binary.write_text("stub") + monkeypatch.delenv("UNSLOTH_LLAMA_CPP_PATH", raising = False) + assert upd._llama_install_root(str(binary)) == root + + +def test_llama_install_root_unmanaged_path_returns_none(monkeypatch, tmp_path): + # A binary on PATH (no marker, no env pin, no llama.cpp ancestor) is foreign: + # installing elsewhere would not replace it, so report no manageable root. + binary = tmp_path / "usr" / "local" / "bin" / "llama-server" + binary.parent.mkdir(parents = True) + binary.write_text("stub") + monkeypatch.delenv("UNSLOTH_LLAMA_CPP_PATH", raising = False) + assert upd._llama_install_root(str(binary)) is None + + +def test_llama_install_root_unsloth_env_dir(monkeypatch, tmp_path): + # UNSLOTH_LLAMA_CPP_PATH dir holding the active binary is the managed root. + root = tmp_path / "vendor" / "llama" + binary = root / "llama-server" + binary.parent.mkdir(parents = True) + binary.write_text("stub") + monkeypatch.setenv("UNSLOTH_LLAMA_CPP_PATH", str(root)) + assert upd._llama_install_root(str(binary)) == root + + +def test_llama_install_root_ignores_inactive_env_root(monkeypatch, tmp_path): + # UNSLOTH_LLAMA_CPP_PATH set but the active binary is not under it: do not + # target the stale env root, resolve from the binary's own llama.cpp tree. + inactive = tmp_path / "custom-empty" + inactive.mkdir() + active = tmp_path / "llama.cpp" + binary = active / "build" / "bin" / "llama-server" + binary.parent.mkdir(parents = True) + binary.write_text("stub") + monkeypatch.setenv("UNSLOTH_LLAMA_CPP_PATH", str(inactive)) + assert upd._llama_install_root(str(binary)) == active + + +def test_llama_install_root_refuses_pinned_checkout_under_llama_cpp(monkeypatch, tmp_path): + # The LLAMA_SERVER_PATH pin guard must run before the ancestor scan, or a + # user's own llama.cpp checkout could be handed to the installer. + root = tmp_path / "my-project" / "llama.cpp" + binary = root / "build" / "bin" / "llama-server" + binary.parent.mkdir(parents = True) + binary.write_text("stub") + monkeypatch.setenv("LLAMA_SERVER_PATH", str(binary)) + monkeypatch.delenv("UNSLOTH_LLAMA_CPP_PATH", raising = False) + assert upd._llama_install_root(str(binary)) is None + + +def test_start_update_source_build_refuses_when_newer(monkeypatch, tmp_path): + # A direct POST on a source build already newer than the prebuilt must not + # downgrade it; start_update mirrors the detection suppression. + install_dir = tmp_path / "llama.cpp" + binary = install_dir / "build" / "bin" / "llama-server" + binary.parent.mkdir(parents = True) + binary.write_text("stub") # no marker + monkeypatch.setattr(upd, "_find_binary", lambda: str(binary)) + monkeypatch.setattr(upd, "_installer_script", lambda: tmp_path / "install_llama_prebuilt.py") + _prebuilt(monkeypatch, release_tag = "b9518") + monkeypatch.setattr(upd, "_installed_build_number", lambda b: 9600) + res = upd.start_update() + assert res["started"] is False + assert res["reason"] == "up_to_date" diff --git a/studio/backend/tests/test_llama_route.py b/studio/backend/tests/test_llama_route.py new file mode 100644 index 0000000000..bf0c4b731f --- /dev/null +++ b/studio/backend/tests/test_llama_route.py @@ -0,0 +1,111 @@ +# SPDX-License-Identifier: AGPL-3.0-only +# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0 + +"""routes/llama.py: the source_build field is exposed and the handlers run the +(now subprocess-touching) detection off the event loop via a worker thread. + +The route file is loaded standalone with a stubbed auth dependency so the test +does not pull the whole routes package (matplotlib-heavy training router) and +works in a minimal env. +""" + +from __future__ import annotations + +import asyncio +import importlib.util +import sys +import threading +import types +from pathlib import Path + +import pytest + +_BACKEND = Path(__file__).resolve().parents[1] +if str(_BACKEND) not in sys.path: + sys.path.insert(0, str(_BACKEND)) + +pytest.importorskip("fastapi") + + +def _load_route(): + # Prefer the real auth module; stub it only in minimal envs where its + # deps are absent. Stubs are popped after the load so they never leak + # into sys.modules for the rest of the suite. + stubbed = [] + try: + import auth.authentication # noqa: F401 + except Exception: + auth_pkg = types.ModuleType("auth") + auth_pkg.__path__ = [] + auth_mod = types.ModuleType("auth.authentication") + auth_mod.get_current_subject = lambda: "test" + for name, stub in (("auth", auth_pkg), ("auth.authentication", auth_mod)): + if name not in sys.modules: + sys.modules[name] = stub + stubbed.append(name) + try: + spec = importlib.util.spec_from_file_location( + "llama_route_under_test", str(_BACKEND / "routes" / "llama.py") + ) + mod = importlib.util.module_from_spec(spec) + sys.modules["llama_route_under_test"] = mod # so pydantic resolves forward refs + spec.loader.exec_module(mod) + return mod + finally: + for name in stubbed: + sys.modules.pop(name, None) + + +rl = _load_route() + + +def test_status_response_exposes_source_build(): + payload = { + "supported": True, + "update_available": True, + "stale": False, + "installed_tag": None, + "latest_tag": "b9585", + "published_repo": "unslothai/llama.cpp", + "installed_at_utc": None, + "age_days": None, + "source_build": True, + "job": {"state": "idle"}, + } + model = rl.LlamaUpdateStatusResponse(**payload) + assert model.model_dump()["source_build"] is True + # Extra/unknown keys must not crash the response model. + rl.LlamaUpdateStatusResponse(**{**payload, "unexpected": 1}) + + +def test_status_handler_runs_off_event_loop(monkeypatch): + seen = {} + + def fake_status(force_refresh = False): + seen["thread"] = threading.current_thread() + return { + "supported": True, + "update_available": True, + "source_build": True, + "latest_tag": "b9585", + "job": {"state": "idle"}, + } + + monkeypatch.setattr(rl, "get_update_status", fake_status) + out = asyncio.run(rl.llama_update_status(force_refresh = False, current_subject = "t")) + assert out.source_build is True + # Detection ran in a worker thread, not the event-loop thread. + assert seen["thread"] is not threading.main_thread() + + +def test_update_handler_runs_off_event_loop(monkeypatch): + seen = {} + + def fake_start(): + seen["thread"] = threading.current_thread() + return {"started": True, "reason": None, "job": {"state": "running"}} + + monkeypatch.setattr(rl, "start_update", fake_start) + out = asyncio.run(rl.llama_update(current_subject = "t")) + assert out.started is True + assert seen["thread"] is not threading.main_thread() diff --git a/studio/backend/utils/llama_cpp_update.py b/studio/backend/utils/llama_cpp_update.py index 19946b7966..b138b29af3 100644 --- a/studio/backend/utils/llama_cpp_update.py +++ b/studio/backend/utils/llama_cpp_update.py @@ -21,6 +21,7 @@ Design notes: from __future__ import annotations +import json import os import re import subprocess @@ -108,6 +109,143 @@ def _installer_script() -> Optional[Path]: return None +# Markerless (source-build) installs have no UNSLOTH_PREBUILT_INFO.json, so we +# ask the installer whether an official prebuilt now exists for this host. Memo +# is 24h; only successful answers are cached so a network blip retries. +_RESOLVE_TTL_SECONDS = 24 * 60 * 60 +_resolve_memo: dict = {} + + +def _resolve_prebuilt_for_host(*, force_refresh: bool = False) -> Optional[dict]: + """Run install_llama_prebuilt.py --resolve-prebuilt (no download) and return + {prebuilt_available, repo, release_tag, llama_tag, asset, install_kind} or + None. Fail-open: any error -> None so a source build never blocks the app.""" + now = time.time() + if not force_refresh and _resolve_memo: + if now - _resolve_memo.get("at", 0.0) < _RESOLVE_TTL_SECONDS: + return _resolve_memo.get("value") + script = _installer_script() + if script is None: + return None + value: Optional[dict] = None + try: + proc = subprocess.run( + [ + sys.executable, + str(script), + "--resolve-prebuilt", + "latest", + "--output-format", + "json", + ], + capture_output = True, + text = True, + timeout = 60, + ) + out = (proc.stdout or "").strip() + if proc.returncode == 0 and out: + parsed = json.loads(out.splitlines()[-1]) + if isinstance(parsed, dict): + value = parsed + except Exception as exc: # pragma: no cover - subprocess/json defensive + logger.debug("llama update: resolve-prebuilt failed", error = str(exc)) + value = None + if value is not None: # cache real answers; let failures retry next poll + _resolve_memo.update(at = now, value = value) + return value + + +def _installed_build_number(binary: Optional[str]) -> Optional[int]: + """Best-effort build number from ``llama-server --version`` (e.g. + 'version: 9585 (abc)'). None when unparseable or <= 1: a source build with + no git tags reports 'version: 1', which we treat as unknown (offer update).""" + if not binary: + return None + try: + proc = subprocess.run([binary, "--version"], capture_output = True, text = True, timeout = 20) + except Exception: # pragma: no cover - defensive + return None + m = re.search(r"version:\s*(\d+)", (proc.stderr or "") + (proc.stdout or "")) + if not m: + return None + n = int(m.group(1)) + return n if n > 1 else None + + +def _is_under(path: Path, root: Path) -> bool: + try: + p, r = path.resolve(), root.resolve() + except (OSError, ValueError): + p, r = path, root + return p == r or r in p.parents + + +def _llama_install_root(binary: Optional[str]) -> Optional[Path]: + """The Studio-managed llama.cpp root the active binary lives under, or None + when the binary is unmanaged. Installing anywhere the active binary is not + would not replace what _find_llama_server_binary runs (which prefers a pinned + LLAMA_SERVER_PATH, then UNSLOTH_LLAMA_CPP_PATH, then a llama.cpp tree), so we + refuse rather than silently install into an inactive or foreign tree.""" + marked = _install_dir_for(binary) + if marked is not None: + return marked + if not binary: + return None + # LLAMA_SERVER_PATH is an explicit user pin that always wins in discovery; + # never auto-replace its tree (even a user's own llama.cpp checkout). + if os.environ.get("LLAMA_SERVER_PATH"): + return None + p = Path(binary) + env = os.environ.get("UNSLOTH_LLAMA_CPP_PATH") + if env and _is_under(p, Path(env)): + return Path(env) + for parent in p.parents: + if parent.name == "llama.cpp": + return parent + # PATH / system / custom install: not a managed tree, so do not offer. + return None + + +def _source_build_status(binary: str, *, force_refresh: bool) -> Optional[dict]: + """Update status for a markerless (source-build) install: offer the official + prebuilt when one exists for this host and is newer than the installed + binary. None -> caller falls through to the no-marker default (unsupported).""" + res = _resolve_prebuilt_for_host(force_refresh = force_refresh) + if not res or not res.get("prebuilt_available"): + return None + # llama_tag is the upstream build (bNNNN, what --version reports); release_tag + # can be a fork wrapper tag, so compare/display against llama_tag. + latest = res.get("llama_tag") or res.get("release_tag") + if not latest: + return None + # No resolvable install root (e.g. a pinned LLAMA_SERVER_PATH we cannot + # manage) means an apply would not take effect, so do not offer. + if _llama_install_root(binary) is None: + return None + installed_build = _installed_build_number(binary) + m = re.search(r"(\d+)", latest) + latest_build = int(m.group(1)) if m else None + # Suppress only when the source build is reliably newer/equal; unknown + # version (the involuntary source-build case) is treated as behind. + update_available = ( + installed_build is None or latest_build is None or installed_build < latest_build + ) + with _job_lock: + job = dict(_job) + return { + "supported": True, + "update_available": update_available, + "stale": False, + "installed_tag": (f"b{installed_build}" if installed_build else None), + "latest_tag": latest, + "published_repo": res.get("repo"), + "installed_at_utc": None, + "age_days": None, + "source_build": True, + "job": job, + } + + def get_update_status(*, force_refresh: bool = False) -> dict: """Report whether a newer prebuilt exists plus the current job state. @@ -115,6 +253,20 @@ def get_update_status(*, force_refresh: bool = False) -> dict: """ binary = _find_binary() marker = read_install_marker(binary) + + with _job_lock: + job_running = _job["state"] == _JOB_RUNNING + + # No marker = source build / custom path. Offer the official prebuilt if one + # now exists for this host (this is why macOS source builds showed no button). + # Skipped while the updater swaps the tree: each 3s poll would exec the + # half-replaced binary (on Windows that exec can make the installer's + # os.replace fail) and the poller only consumes job progress. + if marker is None and binary is not None and not job_running: + src = _source_build_status(binary, force_refresh = force_refresh) + if src is not None: + return src + repo = (marker or {}).get("published_repo") or DEFAULT_PUBLISHED_REPO if force_refresh and repo: @@ -143,6 +295,7 @@ def get_update_status(*, force_refresh: bool = False) -> dict: "published_repo": freshness.get("published_repo") or repo, "installed_at_utc": freshness.get("installed_at_utc"), "age_days": freshness.get("age_days"), + "source_build": False, "job": job, } @@ -254,19 +407,7 @@ def start_update() -> dict: """Kick off a background update. Idempotent: a second call while one is running returns the in-flight job rather than starting another.""" binary = _find_binary() - install_dir = _install_dir_for(binary) marker = read_install_marker(binary) - if install_dir is None or not marker: - return { - "started": False, - "reason": "no_prebuilt_marker", - "message": ( - "This llama.cpp install was not provisioned from an Unsloth " - "prebuilt (source build or custom path); in-app update is " - "unavailable." - ), - "job": get_update_status()["job"], - } script = _installer_script() if script is None: return { @@ -275,9 +416,47 @@ def start_update() -> dict: "message": "install_llama_prebuilt.py could not be located.", "job": get_update_status()["job"], } - repo = marker.get("published_repo") or DEFAULT_PUBLISHED_REPO - from_tag = marker.get("tag") or marker.get("release_tag") - asset = marker.get("asset") + + if marker: + install_dir = _install_dir_for(binary) + repo = marker.get("published_repo") or DEFAULT_PUBLISHED_REPO + from_tag = marker.get("tag") or marker.get("release_tag") + asset = marker.get("asset") + else: + # Source build / custom path: only proceed when the same detection logic + # would offer the update (prebuilt exists, install is behind, root is + # manageable), so a direct POST cannot downgrade a newer source build. + src = _source_build_status(binary, force_refresh = True) if binary else None + if src is None: + return { + "started": False, + "reason": "no_prebuilt_available", + "message": ( + "No official llama.cpp prebuilt is available for this host, " + "so the source build cannot be swapped automatically." + ), + "job": get_update_status()["job"], + } + if not src.get("update_available"): + return { + "started": False, + "reason": "up_to_date", + "message": "The installed llama.cpp build is already at or newer than the latest prebuilt.", + "job": get_update_status()["job"], + } + res = _resolve_prebuilt_for_host() + install_dir = _llama_install_root(binary) + repo = (res or {}).get("repo") or DEFAULT_PUBLISHED_REPO + from_tag = None + asset = (res or {}).get("asset") + + if install_dir is None: + return { + "started": False, + "reason": "no_install_dir", + "message": "Could not determine the llama.cpp install directory.", + "job": get_update_status()["job"], + } with _job_lock: if _job["state"] == _JOB_RUNNING: diff --git a/studio/install_llama_prebuilt.py b/studio/install_llama_prebuilt.py index 822a32ddb8..f9425f5c4f 100644 --- a/studio/install_llama_prebuilt.py +++ b/studio/install_llama_prebuilt.py @@ -3045,6 +3045,21 @@ def _apply_host_overrides( return host +def published_repo_for_host(host: HostInfo, *, linux_amd_tooling_present: bool = False) -> str: + """The release repo setup.sh / setup.ps1 pick for this host: macOS always the + fork (ggml-org macOS bundles need too-new macOS); else CPU-only Linux/Windows + -> ggml-org upstream (the fork ships no CPU bundle) and any usable GPU (NVIDIA + or ROCm) -> the fork. linux_amd_tooling_present mirrors setup.sh routing Linux + hosts that expose AMD tooling (rocminfo/amd-smi/hipconfig/hipinfo) to the fork + even when the probe cannot confirm an active GPU. Mirrors the shell routing.""" + if host.is_macos: + return DEFAULT_PUBLISHED_REPO + has_gpu = ( + host.has_usable_nvidia or host.has_rocm or (host.is_linux and linux_amd_tooling_present) + ) + return DEFAULT_PUBLISHED_REPO if has_gpu else UPSTREAM_REPO + + def pick_windows_cuda_runtime(host: HostInfo) -> str | None: if not host.driver_cuda_version: return None @@ -6690,6 +6705,16 @@ def parse_args() -> argparse.Namespace: const = "latest", help = ("Resolve the source-build fallback plan."), ) + resolve_group.add_argument( + "--resolve-prebuilt", + nargs = "?", + const = "latest", + help = ( + "Report whether an official prebuilt exists for this host without " + "downloading. Picks the host's published repo when --published-repo " + "is left at the default. Use --output-format json." + ), + ) parser.add_argument( "--output-format", choices = ("plain", "json"), @@ -6774,6 +6799,46 @@ def main() -> int: ) return EXIT_SUCCESS + if args.resolve_prebuilt is not None: + # Host-aware "is a prebuilt available" probe, no download. A default repo + # means "pick the repo for this host"; PrebuiltFallback == source build. + host = _apply_host_overrides( + detect_host(), + override_has_rocm = args.has_rocm, + override_rocm_gfx = args.rocm_gfx, + force_cpu = args.cpu_fallback, + ) + # setup.sh routes Linux hosts with AMD tooling to the fork even when no GPU + # is probed; mirror that so a HIP source build is not offered a CPU prebuilt. + amd_tooling = host.is_linux and any( + shutil.which(t) for t in ("rocminfo", "amd-smi", "hipconfig", "hipinfo") + ) + repo = ( + published_repo_for_host(host, linux_amd_tooling_present = amd_tooling) + if args.published_repo == DEFAULT_PUBLISHED_REPO + else args.published_repo + ) + try: + _requested, plans = resolve_simple_install_release_plans( + args.resolve_prebuilt, host, repo, args.published_release_tag or "" + ) + choice = plans[0].attempts[0] if plans and plans[0].attempts else None + if choice is None: + payload = {"prebuilt_available": False, "repo": repo} + else: + payload = { + "prebuilt_available": True, + "repo": repo, + "release_tag": plans[0].release_tag, + "llama_tag": plans[0].llama_tag, + "asset": choice.name, + "install_kind": choice.install_kind, + } + except PrebuiltFallback: + payload = {"prebuilt_available": False, "repo": repo} + emit_resolver_output(payload, output_format = args.output_format) + return EXIT_SUCCESS + if not args.install_dir: raise SystemExit( "install_llama_prebuilt.py: --install-dir is required unless --resolve-llama-tag, --resolve-install-tag, or --resolve-source-build is used" From fbaf5277870d79651c7787d3a25099d5cb0ff39d Mon Sep 17 00:00:00 2001 From: Daniel Han Date: Thu, 11 Jun 2026 04:02:56 -0700 Subject: [PATCH 09/42] Add FastDiffusionModel slow path for text-diffusion models (DiffusionGemma) (#6158) * Add FastDiffusionModel slow path for text-diffusion models (DiffusionGemma) Text-diffusion models (DiffusionGemma) use a block-diffusion generate loop and a novel backbone, so Unsloth's autoregressive kernel/compile patching does not apply. FastDiffusionModel loads the unmodified HuggingFace model (outputs stay bit-identical to transformers) and adds only the safe conveniences: 4bit/8bit loading, PEFT LoRA (attention + dense MLP; the fused 3D MoE experts are noted as a follow-up), the (model, tokenizer) API, and for_inference/for_training. FastModel.from_pretrained auto-routes diffusion model_types to this path, including a fallback that aliases the legacy "diffusion_gemma" config to the "diffusion_gemma4" classes current transformers ships. New file unsloth/models/diffusion.py; loader.py adds the dispatch and the slow-path-aware get_peft_model / for_inference / for_training. * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * FastDiffusionModel: use unsloth's standard dtype resolution Match the FastModel loader instead of the one-line bf16/fp16 pick: SUPPORTS_BFLOAT16 = is_bfloat16_supported(), default to bf16 when supported else fp16, downgrade an explicit bf16 on unsupported hardware with a warning, and assert a valid dtype. --------- Co-authored-by: danielhanchen Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com> --- unsloth/models/diffusion.py | 301 ++++++++++++++++++++++++++++++++++++ unsloth/models/loader.py | 51 ++++++ 2 files changed, 352 insertions(+) create mode 100644 unsloth/models/diffusion.py diff --git a/unsloth/models/diffusion.py b/unsloth/models/diffusion.py new file mode 100644 index 0000000000..12596b432e --- /dev/null +++ b/unsloth/models/diffusion.py @@ -0,0 +1,301 @@ +# Copyright 2023-present Daniel Han-Chen & the Unsloth team. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +""" +FastDiffusionModel: a transformers-only slow path for text-diffusion models (e.g. DiffusionGemma). + +These models use a block-diffusion sampling loop (custom generate) and a novel backbone, so we skip +Unsloth's autoregressive kernel/compile patching and load the unmodified HF model (outputs stay +bit-identical to transformers), keeping only the safe conveniences: 4bit/8bit loading, PEFT LoRA, the +(model, tokenizer) API, and for_inference/for_training. Extend DIFFUSION_MODEL_TYPES as more land. +""" + +import os +import torch +from transformers import AutoConfig, AutoProcessor, AutoTokenizer + +from ._utils import is_bfloat16_supported +from .llama import logger + +__all__ = ["FastDiffusionModel", "DIFFUSION_MODEL_TYPES", "is_diffusion_model_type"] + +# transformers model_type strings routed to this slow path +DIFFUSION_MODEL_TYPES = ("diffusion_gemma", "diffusion_gemma4") + +# Default LoRA targets: standard nn.Linear modules in the shared Gemma-4 backbone. The 128 MoE experts +# are fused 3D Parameters (gate_up_proj/down_proj), not nn.Linear, so PEFT LoRA cannot target them. +DIFFUSION_LORA_TARGETS = [ + "q_proj", + "k_proj", + "v_proj", + "o_proj", # attention + "gate_proj", + "up_proj", + "down_proj", # dense (non-expert) MLP +] + +# Vision tower uses a custom Linear with the same suffix names; exclude it so only the text path is wrapped. +DIFFUSION_LORA_EXCLUDE = r".*(vision_tower|embed_vision).*" + + +def is_diffusion_model_type(model_types): + """model_types: str or iterable -> True if any is a known diffusion model_type.""" + if isinstance(model_types, str): + model_types = (model_types,) + return any(mt in DIFFUSION_MODEL_TYPES for mt in model_types) + + +def _resolve_diffusion_model_class(config): + """Resolve the HF model class for a diffusion checkpoint from config.architectures.""" + import transformers + + archs = getattr(config, "architectures", None) or [] + for arch in archs: + cls = getattr(transformers, arch, None) + if cls is not None: + return cls + # Fallbacks across naming revisions. + for name in ( + "DiffusionGemmaForBlockDiffusion", + "DiffusionGemma4ModelForBlockDiffusion", + "DiffusionGemma4ForBlockDiffusion", + ): + cls = getattr(transformers, name, None) + if cls is not None: + return cls + raise RuntimeError( + f"Unsloth: could not resolve a diffusion model class from architectures={archs}. " + "Ensure you have the transformers build that ships the DiffusionGemma implementation." + ) + + +def _load_diffusion_config(model_name, token, trust_remote_code, revision, local_files_only): + """Load the config, aliasing the legacy ``diffusion_gemma`` model_type to the ``diffusion_gemma4`` + classes current transformers ships. AutoConfig raises on the legacy type; catch that, rewrite the + type/arch names in-memory, and rebuild.""" + try: + return AutoConfig.from_pretrained( + model_name, + token = token, + trust_remote_code = trust_remote_code, + revision = revision, + local_files_only = local_files_only, + ) + except ValueError as e: + if "diffusion_gemma" not in str(e): + raise + import json + from transformers.utils import cached_file + + cfg_path = cached_file( + model_name, + "config.json", + token = token, + revision = revision, + local_files_only = local_files_only, + ) + with open(cfg_path, encoding = "utf-8") as f: + cd = json.load(f) + cd["model_type"] = "diffusion_gemma4" + cd.setdefault("architectures", ["DiffusionGemma4ModelForBlockDiffusion"]) + if isinstance(cd.get("text_config"), dict): + cd["text_config"]["model_type"] = "diffusion_gemma4_text" + if isinstance(cd.get("vision_config"), dict): + cd["vision_config"]["model_type"] = "diffusion_gemma4_vision" + from transformers import DiffusionGemma4Config + + return DiffusionGemma4Config.from_dict(cd) + + +class FastDiffusionModel: + """transformers-only slow path for text-diffusion models.""" + + @staticmethod + def from_pretrained( + model_name = "google/diffusiongemma-26B-A4B-it", + max_seq_length = None, # API-compat; diffusion uses canvas_length + dtype = None, + load_in_4bit = False, + load_in_8bit = False, + load_in_16bit = False, + full_finetuning = False, + token = None, + device_map = "auto", + trust_remote_code = False, + attn_implementation = "eager", # exact match with the reference golden logits + revision = None, + return_tokenizer = True, + **kwargs, + ): + SUPPORTS_BFLOAT16 = is_bfloat16_supported() + if dtype is None: + dtype = torch.float16 if not SUPPORTS_BFLOAT16 else torch.bfloat16 + elif dtype == torch.bfloat16 and not SUPPORTS_BFLOAT16: + logger.warning_once("Device does not support bfloat16. Will change to float16.") + dtype = torch.float16 + assert dtype in (torch.float16, torch.bfloat16, torch.float32) + + # Honor an explicit local_files_only; else fall back to the offline env vars. + local_files_only = kwargs.pop("local_files_only", None) + if local_files_only is None: + local_files_only = ( + os.environ.get("HF_HUB_OFFLINE", "0") == "1" + or os.environ.get("TRANSFORMERS_OFFLINE", "0") == "1" + ) + config = _load_diffusion_config( + model_name, + token, + trust_remote_code, + revision, + local_files_only, + ) + model_type = getattr(config, "model_type", None) + if not is_diffusion_model_type(model_type): + raise RuntimeError( + f"Unsloth: FastDiffusionModel only supports diffusion model_types {DIFFUSION_MODEL_TYPES}, " + f"got '{model_type}'. Use FastModel/FastLanguageModel for autoregressive models." + ) + + model_cls = _resolve_diffusion_model_class(config) + + load_kwargs = dict( + dtype = dtype, + device_map = device_map, + token = token, + trust_remote_code = trust_remote_code, + attn_implementation = attn_implementation, + revision = revision, + local_files_only = local_files_only, + ) + + # Optional bitsandbytes quant. The MoE experts (3D Parameters) are not nn.Linear so bnb skips + # them; only attention + dense MLP Linears quantize, lm_head/embeddings stay full precision. + if load_in_4bit or load_in_8bit: + from transformers import BitsAndBytesConfig + if load_in_4bit: + qcfg = BitsAndBytesConfig( + load_in_4bit = True, + bnb_4bit_use_double_quant = True, + bnb_4bit_quant_type = "nf4", + bnb_4bit_compute_dtype = dtype, + llm_int8_skip_modules = [ + "lm_head", + "embed_tokens", + "experts", + "self_conditioning", + "router", + ], + ) + else: + qcfg = BitsAndBytesConfig(load_in_8bit = True) + load_kwargs["quantization_config"] = qcfg + + print(f"==(( Unsloth: FastDiffusionModel (slow / transformers-only path) ))==") + print(f" Model: {model_name} | class: {model_cls.__name__} | model_type: {model_type}") + print( + f" dtype: {dtype} | 4bit: {load_in_4bit} | 8bit: {load_in_8bit} | attn: {attn_implementation}" + ) + + model = model_cls.from_pretrained(model_name, **load_kwargs).eval() + # Mark before any early return so get_peft_model/for_* route to the slow path. + model._unsloth_slow_diffusion = True + + if not return_tokenizer: + return model, None + + # Prefer the processor (chat template + tokenizer); fall back to a bare tokenizer. Returned as + # "tokenizer" to match the Unsloth (model, tokenizer) contract. + try: + tokenizer = AutoProcessor.from_pretrained( + model_name, + token = token, + trust_remote_code = trust_remote_code, + revision = revision, + local_files_only = local_files_only, + ) + except Exception: + tokenizer = AutoTokenizer.from_pretrained( + model_name, + token = token, + trust_remote_code = trust_remote_code, + revision = revision, + local_files_only = local_files_only, + ) + + return model, tokenizer + + @staticmethod + def get_peft_model( + model, + r = 16, + target_modules = None, + lora_alpha = 16, + lora_dropout = 0.0, + bias = "none", + use_gradient_checkpointing = True, + random_state = 3407, + task_type = None, + **kwargs, + ): + """Attach a PEFT LoRA to the diffusion backbone (attention + dense MLP). No fused kernels.""" + from peft import LoraConfig, get_peft_model as peft_get_peft_model + + if target_modules is None: + target_modules = DIFFUSION_LORA_TARGETS + + lora_kwargs = dict( + r = r, + lora_alpha = lora_alpha, + lora_dropout = lora_dropout, + bias = bias, + target_modules = target_modules, + task_type = task_type, # None: diffusion has no standard CAUSAL_LM head + **{k: v for k, v in kwargs.items() if k in ("modules_to_save", "init_lora_weights")}, + ) + # Exclude the vision tower's custom (non-Linear) modules that share suffix names. + exclude = kwargs.get("exclude_modules", DIFFUSION_LORA_EXCLUDE) + try: + lora_config = LoraConfig(exclude_modules = exclude, **lora_kwargs) + except TypeError: + # Older PEFT without exclude_modules: scope the target to the text decoder by regex. + lora_kwargs["target_modules"] = ( + r".*model\.decoder\.layers\.\d+\.(self_attn\.[qkvo]_proj|mlp\.(gate|up|down)_proj)" + ) + lora_config = LoraConfig(**lora_kwargs) + if use_gradient_checkpointing: + model.gradient_checkpointing_enable() + if hasattr(model, "enable_input_require_grads"): + model.enable_input_require_grads() + + model = peft_get_peft_model(model, lora_config) + model._unsloth_slow_diffusion = True + try: + model.print_trainable_parameters() + except Exception: + pass + return model + + @staticmethod + def for_inference(model): + model.eval() + for _, m in model.named_modules(): + if hasattr(m, "gradient_checkpointing"): + m.gradient_checkpointing = False + return model + + @staticmethod + def for_training(model, use_gradient_checkpointing = True): + model.train() + if use_gradient_checkpointing and hasattr(model, "gradient_checkpointing_enable"): + model.gradient_checkpointing_enable() + return model diff --git a/unsloth/models/loader.py b/unsloth/models/loader.py index 0ee9e49ba1..4dc3046928 100644 --- a/unsloth/models/loader.py +++ b/unsloth/models/loader.py @@ -829,6 +829,7 @@ from ..kernels import ( post_patch_loss_function, ) from .vision import FastBaseModel +from .diffusion import FastDiffusionModel, is_diffusion_model_type from transformers import ( AutoModelForCausalLM, ) @@ -846,6 +847,25 @@ class FastModel(FastBaseModel): model = _prepare_model_for_qat(model, qat_scheme) return model + @staticmethod + def get_peft_model(model, *args, **kwargs): + # Route text-diffusion models (slow path) to the transformers-only PEFT helper. + if getattr(model, "_unsloth_slow_diffusion", False): + return FastDiffusionModel.get_peft_model(model, *args, **kwargs) + return FastBaseModel.get_peft_model(model, *args, **kwargs) + + @staticmethod + def for_inference(model): + if getattr(model, "_unsloth_slow_diffusion", False): + return FastDiffusionModel.for_inference(model) + return FastBaseModel.for_inference(model) + + @staticmethod + def for_training(model, use_gradient_checkpointing = True): + if getattr(model, "_unsloth_slow_diffusion", False): + return FastDiffusionModel.for_training(model, use_gradient_checkpointing) + return FastBaseModel.for_training(model, use_gradient_checkpointing) + @staticmethod def from_pretrained( model_name = "unsloth/Llama-3.2-11B-Vision-Instruct-bnb-4bit", @@ -1065,6 +1085,24 @@ class FastModel(FastBaseModel): local_files_only = True kwargs["local_files_only"] = True + # Text-diffusion slow-path dispatch, factored so both the normal route (below) and the + # legacy-config fallback (in the AutoConfig except handler) share one call site. + def _dispatch_diffusion(): + return FastDiffusionModel.from_pretrained( + model_name = model_name, + max_seq_length = max_seq_length, + dtype = dtype, + load_in_4bit = load_in_4bit, + load_in_8bit = load_in_8bit, + load_in_16bit = load_in_16bit, + full_finetuning = full_finetuning, + token = token, + device_map = device_map, + trust_remote_code = trust_remote_code, + revision = revision, + **kwargs, + ) + try: model_config = AutoConfig.from_pretrained( model_name, @@ -1078,6 +1116,12 @@ class FastModel(FastBaseModel): raise except Exception as error: autoconfig_error = str(error) + # Legacy text-diffusion configs use model_type "diffusion_gemma", which current + # transformers does not register by name (it ships "diffusion_gemma4"). AutoConfig + # raises before we can dispatch; route straight to the diffusion slow path, whose + # loader aliases the legacy type to the gemma4 classes. + if "diffusion_gemma" in autoconfig_error and is_diffusion_model_type("diffusion_gemma"): + return _dispatch_diffusion() if "architecture" in autoconfig_error: if "qwen3_5" in autoconfig_error: raise ImportError( @@ -1126,6 +1170,13 @@ class FastModel(FastBaseModel): ) model_types_all = ",".join(model_types) + "," + # ---- Text-diffusion models (e.g. DiffusionGemma) take a transformers-only slow path. ---- + # These use a custom block-diffusion `generate` and a novel backbone, so we skip Unsloth's + # autoregressive kernel/compile patching and load the unmodified HF model (bit-identical to + # naive transformers), keeping only 4bit/8bit + PEFT LoRA conveniences. + if is_diffusion_model_type(model_types): + return _dispatch_diffusion() + # Save model types and loading method lowered_model_name = model_name.lower() string = os.environ.get("UNSLOTH_MODEL_NAME", "") + model_types_all From 28fe978549096004b5112385776dab9850d35b17 Mon Sep 17 00:00:00 2001 From: Daniel Han Date: Thu, 11 Jun 2026 04:23:41 -0700 Subject: [PATCH 10/42] Lint CI: diff import-hoist check against the PR merge-base, not the base tip (#6190) The import-hoist safety step built its changed-file list with a two-dot diff against the base branch tip on a shallow clone. Once the base branch moves past a PR's branch point, that diff includes every file the base changed since, and the verifier compares newer base code (BEFORE) against the PR's older snapshot (AFTER). This time-reversed comparison reports the base branch's own refactors as blockers on PRs that never touched those files, forcing branch updates or admin merges. Resolve the true merge-base through the compare API, fetch that single commit by SHA (clone stays shallow), and use it for both the file list and --before. Verified against the PR 6137 failure: with the base tip the step flags studio files from a later main commit; with the merge-base the file list contains only the PR's own files and the verifier passes. --- .github/workflows/lint-ci.yml | 28 ++++++++++++++++++---------- 1 file changed, 18 insertions(+), 10 deletions(-) diff --git a/.github/workflows/lint-ci.yml b/.github/workflows/lint-ci.yml index 8cb3535075..bd859a6e9e 100644 --- a/.github/workflows/lint-ci.yml +++ b/.github/workflows/lint-ci.yml @@ -105,29 +105,37 @@ jobs: # exactly where a hoist refactor lives, and it skips brand-new # files whose re-export imports would otherwise look "unused". # - # actions/checkout uses fetch-depth: 1, so the base branch is not - # present locally. Fetch the single base commit with an explicit - # refspec so origin/ is reliably created (a bare - # `git fetch origin ` only updates FETCH_HEAD in some - # configs). Two-dot diff avoids needing a merge-base on a shallow - # clone. + # Diff against the true merge-base, not the base tip. A two-dot + # diff against the tip re-lints every file the base branch + # changed after the PR branched, comparing newer base code + # (BEFORE) against the PR's older snapshot (AFTER) - a + # time-reversed comparison that flags the base branch's own + # refactors as blockers on PRs that never touched those files. + # The compare API returns the merge-base without needing local + # history, and fetching that single commit by SHA keeps the + # shallow (fetch-depth: 1) clone. if: github.event_name == 'pull_request' + env: + GH_TOKEN: ${{ github.token }} run: | - git fetch --no-tags --depth=1 origin \ - "${{ github.base_ref }}:refs/remotes/origin/${{ github.base_ref }}" + MERGE_BASE=$(gh api \ + "repos/${{ github.repository }}/compare/${{ github.event.pull_request.base.sha }}...${{ github.event.pull_request.head.sha }}" \ + --jq .merge_base_commit.sha) + git fetch --no-tags --depth=1 origin "$MERGE_BASE" mapfile -t CHANGED < <( git diff --name-only --diff-filter=M \ - "origin/${{ github.base_ref }}" HEAD -- '*.py' \ + "$MERGE_BASE" HEAD -- '*.py' \ | grep -vE '(^|/)(unsloth_compiled_cache|node_modules|build|dist)/' || true ) if [ "${#CHANGED[@]}" -eq 0 ]; then echo "no in-place-modified Python files to check" exit 0 fi + printf 'merge base: %s\n' "$MERGE_BASE" printf 'checking %d file(s):\n' "${#CHANGED[@]}" printf ' %s\n' "${CHANGED[@]}" python scripts/verify_import_hoist.py \ - --before "origin/${{ github.base_ref }}" --after HEAD "${CHANGED[@]}" + --before "$MERGE_BASE" --after HEAD "${CHANGED[@]}" - name: No leftover debugger / pdb / breakpoint calls # Catches the "I'll just stick a breakpoint() here" mistake From 2db9fad4b5450daffce6c56dbf37292562ee602d Mon Sep 17 00:00:00 2001 From: Daniel Han Date: Thu, 11 Jun 2026 05:06:02 -0700 Subject: [PATCH 11/42] Installer: GPU detection follow-ups after #6174 (poisoned venv repair, llama.cpp routing, probe bounds) (#6183) * Installer: harden GPU detection follow-ups after #6174 Ports the NVIDIA-priority and /proc/driver/nvidia/gpus hardening from #6174 to the remaining pathways and adds recovery for already-poisoned venvs: - install_python_stack.py: add _ensure_cuda_torch so 'unsloth studio update' force-reinstalls CUDA torch when the venv carries a ROCm build on an NVIDIA Linux host (the pre-#6174 poisoning signature). Honors UNSLOTH_TORCH_BACKEND, UNSLOTH_ROCM_TORCH_INSTALLED, and CUDA_VISIBLE_DEVICES=-1/'' opt-outs; never touches healthy CUDA, deliberate CPU wheels, macOS, or Windows. - install_llama_prebuilt.py: detect_host gains the /proc NVIDIA fallback and skips ROCm probes when NVIDIA is usable; forwarded --rocm-gfx/--has-rocm overrides still win. - setup.sh: GPU summary classifies NVIDIA first through a timeout-bounded probe with the /proc fallback; AMD probes are bounded and gain a KFD vendor_id 4098 fallback; the llama.cpp source build only selects GGML_CUDA/GGML_HIP when the matching GPU is actually detected. - install.sh: bound both nvidia-smi calls with a 10s timeout (no behavior change when healthy or when the timeout binary is absent); classify the exported UNSLOTH_TORCH_BACKEND on the final index path segment so custom mirrors containing 'rocm'/'gfx' in their base path are not mislabeled. - install.ps1 + setup.ps1: NVIDIA probes now require a real 'GPU N:' row from nvidia-smi -L under a 10s bound instead of bare exit code 0; later CUDA version and compute_cap queries are bounded too. Tests: 3 new test files (50+ tests), suite at 788 passed. * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * Fix Resolve-CudaToolkit driver probe for extracted-function unit test tests/studio/test_resolve_cuda_toolkit.ps1 extracts Resolve-CudaToolkit alone into a child pwsh and stubs nvidia-smi with a .ps1 script. The bounded runner is not in scope there (and ProcessStartInfo cannot dispatch .ps1 stubs), so the DriverMaxCuda parse silently returned nothing and the major-mismatch scenarios failed. Fall back to direct invocation when Invoke-NvidiaSmiBounded is unavailable; production setup.ps1 always has it defined and keeps the 10s bound. * Treat CUDA_VISIBLE_DEVICES empty or -1 as hidden in NVIDIA-first guards The NVIDIA-first guards added in this branch only special-cased CUDA_VISIBLE_DEVICES=-1 at two setup.sh gates and ignored the empty-string form entirely, while the Python detector (install_llama_prebuilt.py) already treats both as hidden. On a mixed AMD+NVIDIA host steered to the AMD card via CUDA_VISIBLE_DEVICES, the guards suppressed the AMD probes, so setup.sh fell to a CPU llama.cpp build and install.sh picked CUDA wheels instead of ROCm. Move the policy into the helpers so every consumer agrees: - install.sh: new _cvd_hides_nvidia checked first in _has_usable_nvidia_gpu - studio/setup.sh: same via _setup_cvd_hides_nvidia; the two ad-hoc CUDA_VISIBLE_DEVICES=-1 gate conditions are now redundant and removed - studio/install_python_stack.py: _has_usable_nvidia_gpu returns False when CUDA_VISIBLE_DEVICES is set to or -1 (whitespace tolerated) Tests: 5 new sh scenarios (hidden via , -1, padded -1, visible device, and mixed host with hidden NVIDIA restoring the ROCm route) plus a pytest class covering all three implementations behaviourally. Addresses the review comment on the NVIDIA-first setup.sh block. * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * Retrigger CI after PyPI 503 outage during the previous run --------- Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com> --- install.ps1 | 55 +- install.sh | 51 +- studio/install_llama_prebuilt.py | 24 +- studio/install_python_stack.py | 155 ++++++ studio/setup.ps1 | 81 ++- studio/setup.sh | 166 ++++-- tests/sh/test_get_torch_index_url.sh | 51 +- tests/studio/install/test_cuda_repair.py | 248 +++++++++ .../install/test_gpu_detection_followups.py | 480 ++++++++++++++++++ tests/studio/install/test_probe_timeouts.py | 202 ++++++++ 10 files changed, 1439 insertions(+), 74 deletions(-) create mode 100644 tests/studio/install/test_cuda_repair.py create mode 100644 tests/studio/install/test_gpu_detection_followups.py create mode 100644 tests/studio/install/test_probe_timeouts.py diff --git a/install.ps1 b/install.ps1 index cbc7b206c8..2938d4334f 100644 --- a/install.ps1 +++ b/install.ps1 @@ -1406,14 +1406,58 @@ shell.Run cmd, 0, False } } + # ── Helper: run nvidia-smi under a timeout ── + # A wedged NVIDIA driver can make nvidia-smi block during init or after a + # reset; WaitForExit bounds it (mirrors Invoke-AmdSmiNoElevate) so detection + # cannot hang the installer. No RunAsInvoker compat layer: nvidia-smi does + # not auto-elevate. Returns combined stdout+stderr; "" on timeout/failure. + function Invoke-NvidiaSmiBounded { + param( + [Parameter(Mandatory = $true, Position = 0)][string]$Exe, + [Parameter(Position = 1)][string[]]$SmiArgs = @(), + [int]$TimeoutSec = 10 + ) + try { + $psi = New-Object System.Diagnostics.ProcessStartInfo + $psi.FileName = $Exe + $psi.Arguments = ($SmiArgs -join ' ') + $psi.UseShellExecute = $false + $psi.RedirectStandardOutput = $true + $psi.RedirectStandardError = $true + $psi.CreateNoWindow = $true + $proc = [System.Diagnostics.Process]::Start($psi) + $outTask = $proc.StandardOutput.ReadToEndAsync() + $errTask = $proc.StandardError.ReadToEndAsync() + if (-not $proc.WaitForExit($TimeoutSec * 1000)) { + try { $proc.Kill() } catch {} + $global:LASTEXITCODE = 124 + return "" + } + $global:LASTEXITCODE = $proc.ExitCode + return ($outTask.Result + "`n" + $errTask.Result) + } catch { + $global:LASTEXITCODE = 1 + return "" + } + } + + # ── Helper: nvidia-smi -L lists at least one real GPU ── + # Exit code 0 alone is not enough: a stale/driverless nvidia-smi can exit 0 + # while listing no GPU, which would mark an AMD host NVIDIA and suppress + # ROCm detection. Require a "GPU :" data row. + function Test-NvidiaSmiHasGpu { + param([Parameter(Mandatory = $true)][string]$Exe) + $out = Invoke-NvidiaSmiBounded $Exe @('-L') + return ($LASTEXITCODE -eq 0 -and $out -match '(?m)^GPU\s+\d+:') + } + # ── Detect GPU (robust: PATH + hardcoded fallback paths, mirrors setup.ps1) ── $HasNvidiaSmi = $false $NvidiaSmiExe = $null try { $nvSmiCmd = Get-Command nvidia-smi -ErrorAction SilentlyContinue - if ($nvSmiCmd) { - & $nvSmiCmd.Source *> $null - if ($LASTEXITCODE -eq 0) { $HasNvidiaSmi = $true; $NvidiaSmiExe = $nvSmiCmd.Source } + if ($nvSmiCmd -and (Test-NvidiaSmiHasGpu $nvSmiCmd.Source)) { + $HasNvidiaSmi = $true; $NvidiaSmiExe = $nvSmiCmd.Source } } catch {} if (-not $HasNvidiaSmi) { @@ -1423,8 +1467,7 @@ shell.Run cmd, 0, False )) { if (Test-Path $p) { try { - & $p *> $null - if ($LASTEXITCODE -eq 0) { $HasNvidiaSmi = $true; $NvidiaSmiExe = $p; break } + if (Test-NvidiaSmiHasGpu $p) { $HasNvidiaSmi = $true; $NvidiaSmiExe = $p; break } } catch {} } } @@ -1694,7 +1737,7 @@ shell.Run cmd, 0, False $baseUrl = if ($env:UNSLOTH_PYTORCH_MIRROR) { $env:UNSLOTH_PYTORCH_MIRROR.TrimEnd('/') } else { "https://download.pytorch.org/whl" } if (-not $NvidiaSmiExe) { return "$baseUrl/cpu" } try { - $output = & $NvidiaSmiExe 2>&1 | Out-String + $output = Invoke-NvidiaSmiBounded $NvidiaSmiExe # Newer NVIDIA drivers (e.g. 610.x on Windows) print # "CUDA UMD Version: X.Y" instead of the legacy "CUDA Version: X.Y". # Accept both spellings so we don't fall through to the cu126 default. diff --git a/install.sh b/install.sh index 44d1c39d02..df9eca65c7 100755 --- a/install.sh +++ b/install.sh @@ -1745,13 +1745,42 @@ _has_amd_rocm_gpu() { return 1 } +# ── Bounded command runner ── +# Runs a command under a 10s timeout when the `timeout` binary is available, +# otherwise runs it unbounded. Keeps a wedged nvidia-smi (blocking during +# driver init or after a reset) from hanging the installer: a timed-out probe +# exits nonzero and is treated exactly like a failed probe. No-op semantics on +# hosts without `timeout` (e.g. macOS) or when the probe is healthy. +_run_bounded() { + if command -v timeout >/dev/null 2>&1; then + timeout 10 "$@" + else + "$@" + fi +} + +# Returns 0 (true) when CUDA_VISIBLE_DEVICES is set to "" or "-1", i.e. every +# NVIDIA device is deliberately hidden (mixed AMD+NVIDIA hosts steering work to +# the AMD card). Unset means all devices visible. nvidia-smi ignores this env +# var, so the probes below cannot see the distinction on their own. +_cvd_hides_nvidia() { + [ "${CUDA_VISIBLE_DEVICES+set}" = "set" ] || return 1 + _cvd_trim=$(printf '%s' "$CUDA_VISIBLE_DEVICES" | tr -d '[:space:]') + [ -z "$_cvd_trim" ] || [ "$_cvd_trim" = "-1" ] +} + # ── NVIDIA usable-GPU helper ── # Returns 0 (true) if an NVIDIA GPU is present and usable. # Primary probe: nvidia-smi -L. Fallback: /proc/driver/nvidia/gpus/ sysfs, # which the NVIDIA driver populates on Linux regardless of nvidia-smi state # -- handles PATH gaps, subprocess timeouts, and driver init races that # could otherwise cause nvidia-smi to fail and silence NVIDIA detection. +# A GPU hidden via CUDA_VISIBLE_DEVICES=""/-1 counts as NOT usable (matches +# install_llama_prebuilt.py has_usable_nvidia), so AMD/CPU routing still runs. _has_usable_nvidia_gpu() { + if _cvd_hides_nvidia; then + return 1 + fi _nvsmi="" if command -v nvidia-smi >/dev/null 2>&1; then _nvsmi="nvidia-smi" @@ -1759,7 +1788,7 @@ _has_usable_nvidia_gpu() { _nvsmi="/usr/bin/nvidia-smi" fi if [ -n "$_nvsmi" ]; then - if "$_nvsmi" -L 2>/dev/null | awk '/^GPU[[:space:]]+[0-9]+:/{found=1} END{exit !found}'; then + if _run_bounded "$_nvsmi" -L 2>/dev/null | awk '/^GPU[[:space:]]+[0-9]+:/{found=1} END{exit !found}'; then return 0 fi fi @@ -1871,7 +1900,11 @@ get_torch_index_url() { # of the legacy "CUDA Version: X.Y"; accept both with two BRE expressions # (POSIX sed does not support "?" without -E). The two patterns are # mutually exclusive per line, so head -1 picks the first emitted match. - _cuda_ver=$(LC_ALL=C $_smi 2>/dev/null \ + # Bound the call (a wedged nvidia-smi would otherwise hang here) and force + # the C locale for stable parsing. LC_ALL is exported inside this command + # substitution subshell so it reaches nvidia-smi through _run_bounded + # without depending on `env`; the export is scoped to the subshell. + _cuda_ver=$(export LC_ALL=C; _run_bounded "$_smi" 2>/dev/null \ | sed -n \ -e 's/.*CUDA UMD Version:[[:space:]]*\([0-9][0-9]*\.[0-9][0-9]*\).*/\1/p' \ -e 's/.*CUDA Version:[[:space:]]*\([0-9][0-9]*\.[0-9][0-9]*\).*/\1/p' \ @@ -2125,10 +2158,16 @@ TORCH_INDEX_URL=$(get_torch_index_url) # 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. -case "$TORCH_INDEX_URL" in - */rocm*|*/gfx*) export UNSLOTH_TORCH_BACKEND="rocm" ;; - */cpu) export UNSLOTH_TORCH_BACKEND="cpu" ;; - *) export UNSLOTH_TORCH_BACKEND="cuda" ;; +# Classify on the FINAL path segment only: a custom UNSLOTH_PYTORCH_MIRROR +# whose base path happens to contain "rocm" or "gfx" must not mislabel a +# cu*/cpu index as ROCm (radeon repo URLs end in rocm-rel-X.Y/, Strix +# overrides in gfxNNNN/, so the trailing slash is stripped first). +_torch_index_leaf="${TORCH_INDEX_URL%/}" +_torch_index_leaf="${_torch_index_leaf##*/}" +case "$_torch_index_leaf" in + rocm*|gfx*) export UNSLOTH_TORCH_BACKEND="rocm" ;; + cpu) export UNSLOTH_TORCH_BACKEND="cpu" ;; + *) export UNSLOTH_TORCH_BACKEND="cuda" ;; esac # rocm7.2 ships torch 2.11.0 -- adjust the constraint to allow it. diff --git a/studio/install_llama_prebuilt.py b/studio/install_llama_prebuilt.py index f9425f5c4f..22fab82923 100644 --- a/studio/install_llama_prebuilt.py +++ b/studio/install_llama_prebuilt.py @@ -2890,7 +2890,27 @@ def detect_host() -> HostInfo: except Exception: pass + # Linux /proc/driver/nvidia/gpus fallback: the NVIDIA driver exposes one + # subdir per GPU here regardless of nvidia-smi state, so a host whose + # nvidia-smi is absent from PATH, wedged, or failing is still recognised as + # NVIDIA. Mirrors the fallback added to install.sh / install_python_stack.py + # in PR 6174 so the prebuilt installer does not misroute such hosts to ROCm + # or CPU. driver_cuda_version / compute_caps stay unset here; downstream + # CUDA asset selection treats unknown SMs as "prefer portable" and an + # unknown driver runtime line as "no published CUDA match" (returns None, + # no crash), so planning falls back to a source build with GGML_CUDA=ON. + if is_linux and not has_physical_nvidia: + try: + proc_gpu_dir = "/proc/driver/nvidia/gpus" + if os.path.isdir(proc_gpu_dir) and os.listdir(proc_gpu_dir): + has_physical_nvidia = True + has_usable_nvidia = visible_device_tokens != [] + except OSError: + pass + # Detect AMD ROCm (HIP) -- require actual GPU, not just tools installed + # NVIDIA takes precedence: when an NVIDIA GPU is usable, skip ROCm probing + # entirely so co-installed ROCm tools cannot misroute the host (PR 6174). def _amd_smi_has_gpu(stdout: str) -> bool: """Check for 'GPU: ' data rows, not just a table header.""" @@ -2898,7 +2918,7 @@ def detect_host() -> HostInfo: has_rocm = False rocm_gfx_target: str | None = None - if is_linux: + if is_linux and not has_usable_nvidia: # WSL2 ROCDXG: the system rocminfo enumerates the GPU over /dev/dxg # only when HSA_ENABLE_DXG_DETECTION=1 (a no-op on bare metal), and # rocminfo can live only under /opt/rocm/bin (the profile.d PATH @@ -2937,7 +2957,7 @@ def detect_host() -> HostInfo: has_rocm = True rocm_gfx_target = _pick_rocm_gfx_target(_result.stdout) break - elif is_windows: + elif is_windows and not has_usable_nvidia: # Windows: prefer active probes that validate GPU presence. # hipinfo / amd-smi are often NOT on PATH -- the HIP SDK installer # sets HIP_PATH / ROCM_PATH but does not always add the bin dir to diff --git a/studio/install_python_stack.py b/studio/install_python_stack.py index 0460580922..e540aac305 100644 --- a/studio/install_python_stack.py +++ b/studio/install_python_stack.py @@ -90,6 +90,18 @@ _PYTORCH_WHL_BASE = ( os.environ.get("UNSLOTH_PYTORCH_MIRROR") or "https://download.pytorch.org/whl" ).rstrip("/") +# CUDA torch repair specs (see _ensure_cuda_torch). torchvision/torchaudio are +# pinned to the torch<2.11 family rather than left bare: the install uses an +# exclusive --index-url (no PyPI fallback), so a bare name could resolve a +# torchvision built against a different torch major (e.g. 0.27 for torch 2.12) +# and fail at runtime with an ABI mismatch. Same bounds as the _default ROCm +# spec above, which targets the same torch family. +_CUDA_TORCH_PKG_SPEC: tuple[str, str, str] = ( + "torch>=2.4,<2.11.0", + "torchvision>=0.19,<0.26.0", + "torchaudio>=2.4,<2.11.0", +) + # AMD Windows ROCm wheels (repo.amd.com/rocm/whl/{arch_family}/). # Override with UNSLOTH_ROCM_WINDOWS_MIRROR for air-gapped/mirror installs. _ROCM_WINDOWS_INDEX_BASE = ( @@ -654,7 +666,15 @@ def _has_usable_nvidia_gpu() -> bool: case where nvidia-smi is present but the subprocess fails (PATH gap, timeout, driver initialisation race). If either probe confirms an NVIDIA GPU the function returns True so _has_rocm_gpu() is blocked. + + CUDA_VISIBLE_DEVICES set to "" or "-1" hides every NVIDIA device (mixed + AMD+NVIDIA hosts steering work to the AMD card); neither probe honours + that env var, so check it first and report the GPU as not usable. Unset + means all devices visible. """ + cvd = os.environ.get("CUDA_VISIBLE_DEVICES") + if cvd is not None and cvd.strip() in ("", "-1"): + return False exe = shutil.which("nvidia-smi") if exe: try: @@ -786,6 +806,139 @@ def _install_bnb_windows_rocm() -> bool: return True +def _detect_cuda_torch_index_url() -> str: + """Return the pytorch.org CUDA wheel index URL for the host's NVIDIA driver. + + Mirrors install.sh::get_torch_index_url's CUDA ladder so `studio update` + repairs to the same wheel family a fresh `curl | sh` install would pick. + Probes nvidia-smi (PATH, then /usr/bin/nvidia-smi) and parses both the + legacy "CUDA Version:" and the newer "CUDA UMD Version:" spellings. + Defaults to cu126 when nvidia-smi is missing or the version is unreadable + (e.g. NVIDIA detected only via the /proc/driver/nvidia/gpus fallback). + """ + exe = shutil.which("nvidia-smi") + if not exe and os.path.isfile("/usr/bin/nvidia-smi"): + exe = "/usr/bin/nvidia-smi" + tag = "cu126" # default when the driver CUDA version cannot be read + if exe: + try: + result = subprocess.run( + [exe], + stdout = subprocess.PIPE, + stderr = subprocess.DEVNULL, + text = True, + timeout = 10, + ) + if result.returncode == 0: + m = re.search(r"CUDA(?: UMD)? Version:\s*(\d+)\.(\d+)", result.stdout) + if m: + major, minor = int(m.group(1)), int(m.group(2)) + if major >= 13: + tag = "cu130" + elif major == 12 and minor >= 8: + tag = "cu128" + elif major == 12 and minor >= 6: + tag = "cu126" + elif major >= 12: + tag = "cu124" + elif major >= 11: + tag = "cu118" + else: + tag = "cpu" # ancient driver: no usable CUDA wheels + except Exception: + pass + return f"{_PYTORCH_WHL_BASE}/{tag}" + + +def _ensure_cuda_torch() -> None: + """Repair a venv whose torch is a ROCm build on an NVIDIA host. + + Counterpart to _ensure_rocm_torch. A venv poisoned by the pre-fix KFD + gpu_id false positive (ROCm torch installed on an NVIDIA-only machine) + keeps that broken torch on `studio update`, because a torch+rocm wheel + satisfies the version constraint and nothing force-reinstalls it. This + detects that exact case and reinstalls CUDA torch. + + Only repairs when torch actually links against HIP/ROCm. Healthy CUDA + torch and deliberate CPU-only torch are left untouched. + """ + # Respect an explicit backend choice from install.sh: only "" (standalone + # `studio update`) or "cuda" should ever force CUDA wheels. "rocm"/"cpu" + # (or any unrecognised value) are deliberate and must not be overridden. + if _TORCH_BACKEND not in ("", "cuda"): + return + # No CUDA torch on macOS; Windows venv/torch lifecycle is owned by + # install.ps1 (and the KFD poisoning bug is Linux-only), so skip both. + if IS_MACOS or IS_WINDOWS or NO_TORCH: + return + # Never undo a deliberate ROCm install (setup.ps1 sets this marker). + if os.environ.get("UNSLOTH_ROCM_TORCH_INSTALLED") == "1": + return + # CUDA_VISIBLE_DEVICES="" / "-1" deliberately hides the NVIDIA GPU (for + # example a mixed AMD+NVIDIA host that runs ROCm torch on the AMD card); + # never force CUDA wheels over that choice. + _cvd = os.environ.get("CUDA_VISIBLE_DEVICES") + if _cvd is not None and _cvd.strip() in ("", "-1"): + return + # Only NVIDIA hosts should carry CUDA torch. _has_usable_nvidia_gpu() + # covers the /proc/driver/nvidia/gpus fallback when nvidia-smi is absent. + if not _has_usable_nvidia_gpu(): + return + + # Classify the installed torch: "hip" (ROCm build -- the poisoning + # signature), "cuda" (healthy), or "cpu" (deliberate CPU wheel). A + # non-zero exit means torch is missing or un-importable; the base install + # step handles that, so leave it alone. + try: + probe = subprocess.run( + [ + sys.executable, + "-c", + ( + "import torch; " + "hip = getattr(torch.version, 'hip', '') or ''; " + "cuda = getattr(torch.version, 'cuda', '') or ''; " + "ver = getattr(torch, '__version__', '').lower(); " + "print('hip' if (hip or 'rocm' in ver) else ('cuda' if cuda else 'cpu'))" + ), + ], + stdout = subprocess.PIPE, + stderr = subprocess.DEVNULL, + timeout = 90, + ) + except (OSError, subprocess.TimeoutExpired): + return + if probe.returncode != 0: + return + # Take the last non-empty stdout line: stray output from sitecustomize or + # an import hook must not mask the marker (fail-closed either way). + _marker_lines = [ + line.strip() for line in probe.stdout.decode(errors = "replace").splitlines() if line.strip() + ] + if not _marker_lines or _marker_lines[-1] != "hip": + return # healthy CUDA torch, or a deliberate CPU wheel -- leave as-is + + index_url = _detect_cuda_torch_index_url() + _torch_pkg, _vision_pkg, _audio_pkg = _CUDA_TORCH_PKG_SPEC + print( + f" torch is a ROCm build on an NVIDIA host -- reinstalling " + f"CUDA torch from {index_url}\n" + f" (set UNSLOTH_TORCH_BACKEND=rocm to keep a deliberate ROCm torch " + f"on a mixed AMD+NVIDIA host)" + ) + pip_install( + "CUDA torch repair", + "--force-reinstall", + "--no-cache-dir", + _torch_pkg, + _vision_pkg, + _audio_pkg, + "--index-url", + index_url, + constrain = False, + ) + + def _ensure_rocm_torch() -> None: """Reinstall torch with ROCm wheels when the venv received CPU-only torch. @@ -1848,6 +2001,7 @@ def install_python_stack() -> int: # Must follow base packages so torch is present for inspection. if not IS_MACOS and not NO_TORCH: _progress(_torch_step_label("check")) + _ensure_cuda_torch() _ensure_rocm_torch() # Windows + AMD GPU: warn if ROCm torch was not installed (wrong Python @@ -2033,6 +2187,7 @@ def install_python_stack() -> int: # whichever intermediate step clobbered it. if not IS_WINDOWS and not IS_MACOS and not NO_TORCH: _progress(_torch_step_label("final")) + _ensure_cuda_torch() _ensure_rocm_torch() # 14. Final check (silent; third-party conflicts are expected) diff --git a/studio/setup.ps1 b/studio/setup.ps1 index e707cfda7a..29350e3932 100644 --- a/studio/setup.ps1 +++ b/studio/setup.ps1 @@ -299,7 +299,10 @@ function Get-CudaComputeCapability { if (-not $smiExe) { return $null } try { - $raw = & $smiExe --query-gpu=compute_cap --format=csv,noheader 2>$null + # Bounded: a wedged nvidia-smi must not hang setup after the initial + # -L probe succeeded (the helper merges stderr after stdout, so the + # first line is still the compute_cap value). + $raw = Invoke-NvidiaSmiBounded $smiExe @('--query-gpu=compute_cap', '--format=csv,noheader') if ($LASTEXITCODE -ne 0 -or -not $raw) { return $null } # nvidia-smi may return multiple GPUs; take the first one @@ -363,10 +366,10 @@ function Get-PytorchCudaTag { if (-not $smiExe) { return "cu126" } try { - # 2>&1 | Out-String merges stderr into stdout then converts to a single - # string. Plain 2>$null doesn't fully suppress stderr in PS 5.1 -- - # ErrorRecord objects leak into $output and break the -match. - $output = & $smiExe 2>&1 | Out-String + # Bounded: a wedged nvidia-smi must not hang setup. The helper merges + # stderr into the returned string, matching the old 2>&1 | Out-String + # shape (plain 2>$null leaks ErrorRecord objects in PS 5.1). + $output = Invoke-NvidiaSmiBounded $smiExe # Newer NVIDIA drivers (e.g. 610.x on Windows) print # "CUDA UMD Version: X.Y" instead of the legacy "CUDA Version: X.Y". # Accept both spellings so we don't fall through to the cu126 default. @@ -667,16 +670,58 @@ try { # ============================================ # 1a. GPU detection # ============================================ +# ── Helper: run nvidia-smi under a timeout ── +# A wedged NVIDIA driver can make nvidia-smi block during init or after a reset; +# WaitForExit bounds it (mirrors Invoke-AmdSmiNoElevate below) so detection +# cannot hang setup. No RunAsInvoker compat layer: nvidia-smi does not +# auto-elevate. Returns combined stdout+stderr; "" on timeout/failure. +function Invoke-NvidiaSmiBounded { + param( + [Parameter(Mandatory = $true, Position = 0)][string]$Exe, + [Parameter(Position = 1)][string[]]$SmiArgs = @(), + [int]$TimeoutSec = 10 + ) + try { + $psi = New-Object System.Diagnostics.ProcessStartInfo + $psi.FileName = $Exe + $psi.Arguments = ($SmiArgs -join ' ') + $psi.UseShellExecute = $false + $psi.RedirectStandardOutput = $true + $psi.RedirectStandardError = $true + $psi.CreateNoWindow = $true + $proc = [System.Diagnostics.Process]::Start($psi) + $outTask = $proc.StandardOutput.ReadToEndAsync() + $errTask = $proc.StandardError.ReadToEndAsync() + if (-not $proc.WaitForExit($TimeoutSec * 1000)) { + try { $proc.Kill() } catch {} + $global:LASTEXITCODE = 124 + return "" + } + $global:LASTEXITCODE = $proc.ExitCode + return ($outTask.Result + "`n" + $errTask.Result) + } catch { + $global:LASTEXITCODE = 1 + return "" + } +} + +# ── Helper: nvidia-smi -L lists at least one real GPU ── +# Exit code 0 alone is not enough: a stale/driverless nvidia-smi can exit 0 +# while listing no GPU, which would mark an AMD host NVIDIA and suppress ROCm +# detection. Require a "GPU :" data row. +function Test-NvidiaSmiHasGpu { + param([Parameter(Mandatory = $true)][string]$Exe) + $out = Invoke-NvidiaSmiBounded $Exe @('-L') + return ($LASTEXITCODE -eq 0 -and $out -match '(?m)^GPU\s+\d+:') +} + $HasNvidiaSmi = $false $NvidiaSmiExe = $null # Absolute path -- survives Refresh-Environment try { $nvSmiCmd = Get-Command nvidia-smi -ErrorAction SilentlyContinue - if ($nvSmiCmd) { - & $nvSmiCmd.Source *> $null - if ($LASTEXITCODE -eq 0) { - $HasNvidiaSmi = $true - $NvidiaSmiExe = $nvSmiCmd.Source - } + if ($nvSmiCmd -and (Test-NvidiaSmiHasGpu $nvSmiCmd.Source)) { + $HasNvidiaSmi = $true + $NvidiaSmiExe = $nvSmiCmd.Source } } catch {} # Fallback: nvidia-smi may not be on PATH even though a GPU + driver exist. @@ -689,8 +734,7 @@ if (-not $HasNvidiaSmi) { foreach ($p in $nvSmiDefaults) { if (Test-Path $p) { try { - & $p *> $null - if ($LASTEXITCODE -eq 0) { + if (Test-NvidiaSmiHasGpu $p) { $HasNvidiaSmi = $true $NvidiaSmiExe = $p Write-Host " Found nvidia-smi at $(Split-Path $p -Parent)" -ForegroundColor Gray @@ -1151,7 +1195,16 @@ function Resolve-CudaToolkit { $DriverMaxCuda = $null try { - $smiOut = & $NvidiaSmiExe 2>&1 | Out-String + # Bounded: source-build toolkit resolution must not hang on a wedged smi. + # test_resolve_cuda_toolkit.ps1 extracts this function alone into a child + # pwsh (no Invoke-NvidiaSmiBounded in scope) and stubs nvidia-smi with a + # .ps1 script, so fall back to direct invocation when the bounded runner + # is unavailable; production setup.ps1 always has it defined. + $smiOut = if (Get-Command Invoke-NvidiaSmiBounded -ErrorAction SilentlyContinue) { + Invoke-NvidiaSmiBounded $NvidiaSmiExe + } else { + & $NvidiaSmiExe 2>&1 | Out-String + } # Newer drivers report "CUDA UMD Version: X.Y" instead of "CUDA Version: X.Y"; accept both. if ($smiOut -match "CUDA(?: UMD)? Version:\s+([\d]+)\.([\d]+)") { $DriverMaxCuda = "$($Matches[1]).$($Matches[2])" diff --git a/studio/setup.sh b/studio/setup.sh index 66c03da391..a8603bd0da 100755 --- a/studio/setup.sh +++ b/studio/setup.sh @@ -155,9 +155,60 @@ _nvcc_meets_llama_minimum() { echo "$_raw" } +# Run a GPU probe under a 10s timeout when `timeout` is available so a wedged +# NVIDIA driver cannot hang setup; fall back to a bare call where it is not. +_setup_run_smi() { + if command -v timeout >/dev/null 2>&1; then + timeout 10 "$@" + else + "$@" + fi +} + +# Returns 0 when CUDA_VISIBLE_DEVICES is set to "" or "-1", i.e. every NVIDIA +# device is deliberately hidden (mixed AMD+NVIDIA hosts steering work to the +# AMD card). Unset means all devices visible. nvidia-smi ignores this env var, +# so the probes below cannot see the distinction on their own. +_setup_cvd_hides_nvidia() { + [ "${CUDA_VISIBLE_DEVICES+set}" = "set" ] || return 1 + _setup_cvd_trim=$(printf '%s' "$CUDA_VISIBLE_DEVICES" | tr -d '[:space:]') + [ -z "$_setup_cvd_trim" ] || [ "$_setup_cvd_trim" = "-1" ] +} + +# Returns 0 when an NVIDIA GPU is present and usable. Primary probe is +# `nvidia-smi -L` (timeout-bounded). Fallback is /proc/driver/nvidia/gpus, +# which the driver populates per GPU regardless of nvidia-smi state -- handles +# PATH gaps and driver init races. Mirrors install.sh _has_usable_nvidia_gpu +# (PR 6174) so setup routes the same way as the torch installer. A GPU hidden +# via CUDA_VISIBLE_DEVICES=""/-1 counts as NOT usable (matches +# install_llama_prebuilt.py has_usable_nvidia), so the AMD probes still run +# and a mixed host steered to its AMD card keeps the ROCm route. +_setup_has_usable_nvidia_gpu() { + if _setup_cvd_hides_nvidia; then + return 1 + fi + _setup_nvsmi="" + if command -v nvidia-smi >/dev/null 2>&1; then + _setup_nvsmi="nvidia-smi" + elif [ -x "/usr/bin/nvidia-smi" ]; then + _setup_nvsmi="/usr/bin/nvidia-smi" + fi + if [ -n "$_setup_nvsmi" ]; then + if _setup_run_smi "$_setup_nvsmi" -L 2>/dev/null \ + | awk '/^GPU[[:space:]]+[0-9]+:/{found=1} END{exit !found}'; then + return 0 + fi + fi + if [ -d /proc/driver/nvidia/gpus ] && \ + [ -n "$(ls -A /proc/driver/nvidia/gpus 2>/dev/null)" ]; then + return 0 + fi + return 1 +} + _cuda_driver_max_version() { command -v nvidia-smi >/dev/null 2>&1 || return 0 - nvidia-smi 2>/dev/null \ + _setup_run_smi nvidia-smi 2>/dev/null \ | sed -nE 's/.*CUDA( UMD)? Version:[[:space:]]*([0-9]+)\.([0-9]+).*/\2.\3/p' \ | head -1 || true } @@ -815,25 +866,42 @@ _setup_amd_detected=false _setup_nvidia_usable=false _setup_gfx_all="" _setup_mkt="" -if command -v rocminfo >/dev/null 2>&1 && \ - rocminfo 2>/dev/null | awk '/Name:[[:space:]]*gfx[1-9][0-9]/{found=1} END{exit !found}'; then - _setup_amd_detected=true - _setup_gfx_all=$(rocminfo 2>/dev/null | grep -oE 'gfx[1-9][0-9a-z]{2,3}' || true) - _setup_mkt=$(rocminfo 2>/dev/null | awk -F': ' \ - '/Marketing Name:/{gsub(/^[[:space:]]+|[[:space:]]+$/,"", $2); if($2){print $2; exit}}' || true) -elif command -v amd-smi >/dev/null 2>&1 && \ - amd-smi list 2>/dev/null | awk '/^GPU[[:space:]]*[:\[][[:space:]]*[0-9]/{ found=1 } END{ exit !found }'; then - _setup_amd_detected=true - _setup_gfx_all=$(amd-smi list 2>/dev/null | grep -oE 'gfx[1-9][0-9a-z]{2,3}' || true) - [ -z "$_setup_gfx_all" ] && \ - _setup_gfx_all=$(amd-smi static --asic 2>/dev/null | grep -oE 'gfx[1-9][0-9a-z]{2,3}' || true) - _setup_mkt=$(amd-smi static --asic 2>/dev/null | awk -F'[:|]' \ - '/[Mm]arket.?[Nn]ame/{gsub(/^[[:space:]]+|[[:space:]]+$/,"", $2); if($2){print $2; exit}}' || true) +# NVIDIA priority: classify NVIDIA first and skip the AMD probes entirely on +# a usable-NVIDIA host (mirrors _has_rocm_gpu in install_python_stack.py). +# This also keeps a wedged rocminfo/amd-smi from hanging setup before the +# host is classified; the AMD probes themselves run under _setup_run_smi. +if _setup_has_usable_nvidia_gpu; then + _setup_nvidia_usable=true +fi +if [ "$_setup_nvidia_usable" != true ]; then + if command -v rocminfo >/dev/null 2>&1 && \ + _setup_run_smi rocminfo 2>/dev/null | awk '/Name:[[:space:]]*gfx[1-9][0-9]/{found=1} END{exit !found}'; then + _setup_amd_detected=true + _setup_gfx_all=$(_setup_run_smi rocminfo 2>/dev/null | grep -oE 'gfx[1-9][0-9a-z]{2,3}' || true) + _setup_mkt=$(_setup_run_smi rocminfo 2>/dev/null | awk -F': ' \ + '/Marketing Name:/{gsub(/^[[:space:]]+|[[:space:]]+$/,"", $2); if($2){print $2; exit}}' || true) + elif command -v amd-smi >/dev/null 2>&1 && \ + _setup_run_smi amd-smi list 2>/dev/null | awk '/^GPU[[:space:]]*[:\[][[:space:]]*[0-9]/{ found=1 } END{ exit !found }'; then + _setup_amd_detected=true + _setup_gfx_all=$(_setup_run_smi amd-smi list 2>/dev/null | grep -oE 'gfx[1-9][0-9a-z]{2,3}' || true) + [ -z "$_setup_gfx_all" ] && \ + _setup_gfx_all=$(_setup_run_smi amd-smi static --asic 2>/dev/null | grep -oE 'gfx[1-9][0-9a-z]{2,3}' || true) + _setup_mkt=$(_setup_run_smi amd-smi static --asic 2>/dev/null | awk -F'[:|]' \ + '/[Mm]arket.?[Nn]ame/{gsub(/^[[:space:]]+|[[:space:]]+$/,"", $2); if($2){print $2; exit}}' || true) + 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 }' \ + /sys/class/kfd/kfd/topology/nodes/*/properties 2>/dev/null; then + # KFD sysfs fallback, AMD vendor_id 4098 only (mirrors install.sh + # _has_amd_rocm_gpu): covers AMD hosts where rocminfo/amd-smi are + # missing but the kernel exposes the GPU, so the source-build gate + # below does not drop them to a CPU llama.cpp build. No gfx arch is + # available from this path; name-based inference handles it. + _setup_amd_detected=true + fi fi -if command -v nvidia-smi >/dev/null 2>&1 && \ - nvidia-smi -L 2>/dev/null | awk '/^GPU[[:space:]]+[0-9]+:/{found=1} END{exit !found}'; then - _setup_nvidia_usable=true +if [ "$_setup_nvidia_usable" = true ]; then step "gpu" "NVIDIA GPU detected" elif [ "$_setup_amd_detected" = true ]; then _setup_vis="${HIP_VISIBLE_DEVICES:-${ROCR_VISIBLE_DEVICES:-}}" @@ -918,15 +986,15 @@ _HOST_MACHINE="$(uname -m 2>/dev/null || true)" # use unslothai. _LINUX_HAS_GPU=false # Route to the fork only for a usable GPU. NVIDIA counts only when a device is -# actually enumerated (_setup_nvidia_usable, from the nvidia-smi -L probe above) -# AND not hidden via CUDA_VISIBLE_DEVICES=-1 -- mirroring install_llama_prebuilt.py's -# has_usable_nvidia. Mere nvidia-smi presence (CPU-only CUDA-toolkit containers, -# broken drivers) or a hidden GPU therefore takes the ggml-org CPU prebuilt -# instead of a slow source build. AMD is deliberately left on tooling presence, -# not usability: an unusable NVIDIA host has a good CPU prebuilt to fall back to, -# whereas tightening AMD would regress ROCm hosts exposing only hipconfig/hipinfo -# into an unnecessary CPU build. -if [ "$_setup_nvidia_usable" = true ] && [ "${CUDA_VISIBLE_DEVICES:-}" != "-1" ]; then +# actually enumerated and not hidden via CUDA_VISIBLE_DEVICES=""/-1 +# (_setup_nvidia_usable, from _setup_has_usable_nvidia_gpu above) -- mirroring +# install_llama_prebuilt.py's has_usable_nvidia. Mere nvidia-smi presence +# (CPU-only CUDA-toolkit containers, broken drivers) or a hidden GPU therefore +# takes the ggml-org CPU prebuilt instead of a slow source build. AMD is +# deliberately left on tooling presence, not usability: an unusable NVIDIA host +# has a good CPU prebuilt to fall back to, whereas tightening AMD would regress +# ROCm hosts exposing only hipconfig/hipinfo into an unnecessary CPU build. +if [ "$_setup_nvidia_usable" = true ]; then _LINUX_HAS_GPU=true else for _GPU_TOOL in rocminfo amd-smi hipconfig hipinfo; do @@ -1271,23 +1339,35 @@ else GPU_BACKEND="" NVCC_PATH="" - if command -v nvcc &>/dev/null; then - NVCC_PATH="$(command -v nvcc)" - GPU_BACKEND="cuda" - elif [ -x /usr/local/cuda/bin/nvcc ]; then - NVCC_PATH="/usr/local/cuda/bin/nvcc" - export PATH="/usr/local/cuda/bin:$PATH" - GPU_BACKEND="cuda" - elif ls /usr/local/cuda-*/bin/nvcc &>/dev/null 2>&1; then - # Pick the newest cuda-XX.X directory - NVCC_PATH="$(ls -d /usr/local/cuda-*/bin/nvcc 2>/dev/null | sort -V | tail -1)" - export PATH="$(dirname "$NVCC_PATH"):$PATH" - GPU_BACKEND="cuda" + # Gate the CUDA toolkit search on an actually-usable NVIDIA GPU + # (_setup_nvidia_usable, computed in the GPU summary block above; + # already false when hidden via CUDA_VISIBLE_DEVICES=""/-1). + # A CUDA toolkit alone (CPU-only build container, leftover packages) + # is not proof of a GPU: building with -DGGML_CUDA=ON there yields a + # binary that fails at runtime, so fall through to the CPU build. + if [ "$_setup_nvidia_usable" = true ]; then + if command -v nvcc &>/dev/null; then + NVCC_PATH="$(command -v nvcc)" + GPU_BACKEND="cuda" + elif [ -x /usr/local/cuda/bin/nvcc ]; then + NVCC_PATH="/usr/local/cuda/bin/nvcc" + export PATH="/usr/local/cuda/bin:$PATH" + GPU_BACKEND="cuda" + elif ls /usr/local/cuda-*/bin/nvcc &>/dev/null 2>&1; then + # Pick the newest cuda-XX.X directory + NVCC_PATH="$(ls -d /usr/local/cuda-*/bin/nvcc 2>/dev/null | sort -V | tail -1)" + export PATH="$(dirname "$NVCC_PATH"):$PATH" + GPU_BACKEND="cuda" + fi fi - # Check for ROCm (AMD) only if CUDA was not already selected + # Check for ROCm (AMD) only if CUDA was not already selected, and + # only when an AMD GPU was actually detected (_setup_amd_detected). + # hipcc presence alone (HIP SDK, no GPU) must not select a HIP build. + # NVIDIA-usable hosts never build HIP (defense in depth: the AMD + # probes above are already skipped when NVIDIA is usable). ROCM_HIPCC="" - if [ -z "$GPU_BACKEND" ]; then + if [ -z "$GPU_BACKEND" ] && [ "$_setup_nvidia_usable" != true ] && [ "$_setup_amd_detected" = true ]; then if command -v hipcc &>/dev/null; then ROCM_HIPCC="$(command -v hipcc)" GPU_BACKEND="rocm" @@ -1349,7 +1429,7 @@ else CUDA_ARCHS="" if command -v nvidia-smi &>/dev/null; then - _raw_caps=$(nvidia-smi --query-gpu=compute_cap --format=csv,noheader 2>/dev/null || true) + _raw_caps=$(_setup_run_smi nvidia-smi --query-gpu=compute_cap --format=csv,noheader 2>/dev/null || true) while IFS= read -r _cap; do _cap=$(echo "$_cap" | tr -d '[:space:]') if [[ "$_cap" =~ ^([0-9]+)\.([0-9]+)$ ]]; then @@ -1455,7 +1535,7 @@ else CMAKE_ARGS="$CMAKE_ARGS -DGPU_TARGETS=${GPU_TARGETS}" _BUILD_DESC="building (ROCm, ${GPU_TARGETS//;/+})" fi - elif [ -d /usr/local/cuda ] || nvidia-smi &>/dev/null; then + elif [ -d /usr/local/cuda ] || _setup_run_smi nvidia-smi &>/dev/null; then _BUILD_DESC="building (CPU, CUDA driver found but nvcc missing)" elif [ -d /opt/rocm ] || command -v rocm-smi &>/dev/null; then _BUILD_DESC="building (CPU, ROCm driver found but hipcc missing)" diff --git a/tests/sh/test_get_torch_index_url.sh b/tests/sh/test_get_torch_index_url.sh index e20fd1ca86..3dece1f69a 100755 --- a/tests/sh/test_get_torch_index_url.sh +++ b/tests/sh/test_get_torch_index_url.sh @@ -13,6 +13,10 @@ FAIL=0 _FUNC_FILE=$(mktemp) _FAKE_SMI_DIR=$(mktemp -d) { + sed -n '/^_run_bounded()/,/^}/p' "$INSTALL_SH" + echo "" + sed -n '/^_cvd_hides_nvidia()/,/^}/p' "$INSTALL_SH" + echo "" sed -n '/^_has_amd_rocm_gpu()/,/^}/p' "$INSTALL_SH" echo "" sed -n '/^_has_usable_nvidia_gpu()/,/^}/p' "$INSTALL_SH" @@ -107,7 +111,7 @@ MOCK # Build a minimal tools directory with symlinks to essential commands # (uname, grep, head, etc.) but WITHOUT nvidia-smi or amd-smi. _TOOLS_DIR=$(mktemp -d) -for _cmd in uname grep sed head sh bash cat awk printf; do +for _cmd in uname grep sed head sh bash cat awk printf tr; do _real=$(command -v "$_cmd" 2>/dev/null || true) [ -n "$_real" ] && ln -sf "$_real" "$_TOOLS_DIR/$_cmd" done @@ -116,12 +120,19 @@ done # $1 = directory with mock nvidia-smi (prepended to PATH), or "none" for no-GPU test run_func() { _mock_dir="$1" + # Default: strip CUDA_VISIBLE_DEVICES so the host environment cannot leak + # in; a second argument sets it explicitly (hidden-GPU scenarios). + if [ "$#" -ge 2 ]; then + _cvd_setup="export CUDA_VISIBLE_DEVICES='$2'" + else + _cvd_setup="unset CUDA_VISIBLE_DEVICES" + fi if [ "$_mock_dir" = "none" ]; then # Minimal PATH with only basic tools, no nvidia-smi anywhere - PATH="$_TOOLS_DIR" bash -c ". '$_FUNC_FILE'; get_torch_index_url" 2>/dev/null + PATH="$_TOOLS_DIR" bash -c "$_cvd_setup; . '$_FUNC_FILE'; get_torch_index_url" 2>/dev/null else # Put mock nvidia-smi dir first, then basic tools - PATH="$_mock_dir:$_TOOLS_DIR" bash -c ". '$_FUNC_FILE'; get_torch_index_url" 2>/dev/null + PATH="$_mock_dir:$_TOOLS_DIR" bash -c "$_cvd_setup; . '$_FUNC_FILE'; get_torch_index_url" 2>/dev/null fi } @@ -332,6 +343,40 @@ _result=$(run_func "$_dir") assert_eq "CUDA Version 13.7 -> cu130" "https://download.pytorch.org/whl/cu130" "$_result" rm -rf "$_dir" +# 34) CUDA_VISIBLE_DEVICES="" hides the NVIDIA GPU -> cpu (no AMD present) +_dir=$(make_mock_smi "12.8") +_result=$(run_func "$_dir" "") +assert_eq "CVD='' hides NVIDIA -> cpu" "https://download.pytorch.org/whl/cpu" "$_result" +rm -rf "$_dir" + +# 35) CUDA_VISIBLE_DEVICES=-1 hides the NVIDIA GPU -> cpu (no AMD present) +_dir=$(make_mock_smi "12.8") +_result=$(run_func "$_dir" "-1") +assert_eq "CVD=-1 hides NVIDIA -> cpu" "https://download.pytorch.org/whl/cpu" "$_result" +rm -rf "$_dir" + +# 36) Mixed AMD+NVIDIA host with NVIDIA hidden -> ROCm route is restored +_cuda_dir=$(make_mock_smi "12.6") +_amd_dir=$(make_mock_amd_smi "6.4") +_combined_dir=$(mktemp -d) +ln -sf "$_cuda_dir/nvidia-smi" "$_combined_dir/nvidia-smi" +ln -sf "$_amd_dir/amd-smi" "$_combined_dir/amd-smi" +_result=$(run_func "$_combined_dir" "-1") +assert_eq "CUDA+ROCm with CVD=-1 -> rocm6.4" "https://download.pytorch.org/whl/rocm6.4" "$_result" +rm -rf "$_cuda_dir" "$_amd_dir" "$_combined_dir" + +# 37) CUDA_VISIBLE_DEVICES=0 (a visible device) must NOT hide the GPU +_dir=$(make_mock_smi "12.8") +_result=$(run_func "$_dir" "0") +assert_eq "CVD=0 keeps NVIDIA -> cu128" "https://download.pytorch.org/whl/cu128" "$_result" +rm -rf "$_dir" + +# 38) Whitespace-padded "-1" still hides the GPU +_dir=$(make_mock_smi "12.8") +_result=$(run_func "$_dir" " -1 ") +assert_eq "CVD=' -1 ' hides NVIDIA -> cpu" "https://download.pytorch.org/whl/cpu" "$_result" +rm -rf "$_dir" + rm -f "$_FUNC_FILE" rm -rf "$_FAKE_SMI_DIR" rm -rf "$_TOOLS_DIR" diff --git a/tests/studio/install/test_cuda_repair.py b/tests/studio/install/test_cuda_repair.py new file mode 100644 index 0000000000..83c2d962e8 --- /dev/null +++ b/tests/studio/install/test_cuda_repair.py @@ -0,0 +1,248 @@ +"""Tests for CUDA torch repair on poisoned NVIDIA venvs. + +Verifies _ensure_cuda_torch (studio/install_python_stack.py) reinstalls CUDA +torch when a venv on an NVIDIA host carries a ROCm torch build (the pre-fix KFD +gpu_id false positive), without touching healthy CUDA, deliberate CPU wheels, +ROCm hosts, macOS, or Windows. All tests use mocks -- no GPU required. +""" + +import importlib.util +import sys +from pathlib import Path +from unittest.mock import MagicMock, patch + +import pytest + + +# ── Load module under test (mirrors test_rocm_support.py) ──────────────────── + +PACKAGE_ROOT = Path(__file__).resolve().parents[3] + +_STACK_PATH = PACKAGE_ROOT / "studio" / "install_python_stack.py" +_STACK_SPEC = importlib.util.spec_from_file_location("studio_install_python_stack", _STACK_PATH) +assert _STACK_SPEC is not None and _STACK_SPEC.loader is not None +stack_mod = importlib.util.module_from_spec(_STACK_SPEC) +sys.modules[_STACK_SPEC.name] = stack_mod +_STACK_SPEC.loader.exec_module(stack_mod) + +_ensure_cuda_torch = stack_mod._ensure_cuda_torch +_detect_cuda_torch_index_url = stack_mod._detect_cuda_torch_index_url + + +# ── Helpers ────────────────────────────────────────────────────────────────── + + +def _make_run( + torch_state = "hip", + cuda_version = "12.8", + torch_rc = 0, + smi_rc = 0, +): + """Build a subprocess.run side_effect. + + The torch-classify probe runs sys.executable and reads bytes stdout; the + nvidia-smi version probe runs the smi path with text=True. Distinguish by + the executable. + """ + + def _run(cmd, *args, **kwargs): + result = MagicMock() + exe = str(cmd[0]) if cmd else "" + if exe == sys.executable: + result.returncode = torch_rc + result.stdout = (torch_state + "\n").encode() + return result + # nvidia-smi version probe (text = True) + result.returncode = smi_rc + out = f"CUDA Version: {cuda_version}\n" if cuda_version else "No devices found\n" + result.stdout = out if kwargs.get("text") else out.encode() + return result + + return _run + + +def _run_cuda_repair( + *, + backend = "", + nvidia = True, + torch_state = "hip", + cuda_version = "12.8", + torch_rc = 0, + smi_rc = 0, + is_macos = False, + is_windows = False, + no_torch = False, + rocm_marker = False, + smi_path = "/usr/bin/nvidia-smi", + cvd = None, +): + """Invoke _ensure_cuda_torch under a fully mocked host; return the pip mock. + + cvd controls CUDA_VISIBLE_DEVICES: None removes it from the environment + (the host machine may export one), any string sets it explicitly. + """ + env = {} + if rocm_marker: + env["UNSLOTH_ROCM_TORCH_INSTALLED"] = "1" + if cvd is not None: + env["CUDA_VISIBLE_DEVICES"] = cvd + + def _which(name, *a, **k): + if name == "nvidia-smi": + return smi_path + return None + + with ( + patch.object(stack_mod, "_TORCH_BACKEND", backend), + patch.object(stack_mod, "IS_MACOS", is_macos), + patch.object(stack_mod, "IS_WINDOWS", is_windows), + patch.object(stack_mod, "NO_TORCH", no_torch), + patch.object(stack_mod, "_has_usable_nvidia_gpu", return_value = nvidia), + patch.object(stack_mod.shutil, "which", side_effect = _which), + patch.object(stack_mod.os.path, "isfile", return_value = bool(smi_path)), + patch.object(stack_mod, "pip_install") as mock_pip, + patch.object( + stack_mod.subprocess, + "run", + side_effect = _make_run(torch_state, cuda_version, torch_rc, smi_rc), + ), + patch.dict(stack_mod.os.environ, env, clear = False), + ): + if not rocm_marker: + stack_mod.os.environ.pop("UNSLOTH_ROCM_TORCH_INSTALLED", None) + if cvd is None: + stack_mod.os.environ.pop("CUDA_VISIBLE_DEVICES", None) + _ensure_cuda_torch() + return mock_pip + + +def _index_url(mock_pip) -> str: + """Return the --index-url value from the recorded pip_install call.""" + args = [str(a) for a in mock_pip.call_args.args] + return args[args.index("--index-url") + 1] + + +# ── Repair fires only on the poisoning signature ───────────────────────────── + + +class TestCudaRepairFires: + def test_hip_build_on_nvidia_triggers_repair(self): + mock_pip = _run_cuda_repair(torch_state = "hip", cuda_version = "12.8") + assert mock_pip.call_count == 1 + call_args = [str(a) for a in mock_pip.call_args.args] + assert "--force-reinstall" in call_args + assert "--no-cache-dir" in call_args + assert "cu128" in _index_url(mock_pip) + assert mock_pip.call_args.kwargs["constrain"] is False + + def test_rocm_in_version_string_triggers_repair(self): + # AMD SDK / Radeon wheels may not set torch.version.hip but encode + # rocm in __version__; the probe prints "hip" for both. + mock_pip = _run_cuda_repair(torch_state = "hip") + assert mock_pip.call_count == 1 + + +# ── No-op cases ────────────────────────────────────────────────────────────── + + +class TestCudaRepairSkips: + def test_healthy_cuda_torch_no_repair(self): + mock_pip = _run_cuda_repair(torch_state = "cuda") + mock_pip.assert_not_called() + + def test_deliberate_cpu_wheel_no_repair(self): + mock_pip = _run_cuda_repair(torch_state = "cpu") + mock_pip.assert_not_called() + + def test_backend_rocm_skips(self): + mock_pip = _run_cuda_repair(backend = "rocm", torch_state = "hip") + mock_pip.assert_not_called() + + def test_backend_cpu_skips(self): + mock_pip = _run_cuda_repair(backend = "cpu", torch_state = "hip") + mock_pip.assert_not_called() + + def test_unknown_backend_skips(self): + mock_pip = _run_cuda_repair(backend = "auto", torch_state = "hip") + mock_pip.assert_not_called() + + def test_no_nvidia_gpu_skips(self): + mock_pip = _run_cuda_repair(nvidia = False, torch_state = "hip") + mock_pip.assert_not_called() + + def test_torch_missing_skips(self): + # Non-zero probe exit = torch missing / un-importable. + mock_pip = _run_cuda_repair(torch_state = "hip", torch_rc = 1) + mock_pip.assert_not_called() + + def test_macos_skips(self): + mock_pip = _run_cuda_repair(is_macos = True, torch_state = "hip") + mock_pip.assert_not_called() + + def test_windows_skips(self): + mock_pip = _run_cuda_repair(is_windows = True, torch_state = "hip") + mock_pip.assert_not_called() + + def test_no_torch_mode_skips(self): + mock_pip = _run_cuda_repair(no_torch = True, torch_state = "hip") + mock_pip.assert_not_called() + + def test_rocm_install_marker_skips(self): + mock_pip = _run_cuda_repair(rocm_marker = True, torch_state = "hip") + mock_pip.assert_not_called() + + def test_cvd_minus_one_skips(self): + # CUDA_VISIBLE_DEVICES=-1 deliberately hides the NVIDIA GPU (mixed + # AMD+NVIDIA host running ROCm torch on the AMD card). + mock_pip = _run_cuda_repair(cvd = "-1", torch_state = "hip") + mock_pip.assert_not_called() + + def test_cvd_empty_skips(self): + mock_pip = _run_cuda_repair(cvd = "", torch_state = "hip") + mock_pip.assert_not_called() + + def test_cvd_explicit_device_still_repairs(self): + mock_pip = _run_cuda_repair(cvd = "0", torch_state = "hip") + assert mock_pip.call_count == 1 + + +# ── CUDA index ladder ──────────────────────────────────────────────────────── + + +class TestCudaIndexResolution: + def test_cuda_128_selects_cu128(self): + assert "cu128" in _index_url(_run_cuda_repair(cuda_version = "12.8")) + + def test_cuda_130_selects_cu130(self): + assert "cu130" in _index_url(_run_cuda_repair(cuda_version = "13.0")) + + def test_cuda_126_selects_cu126(self): + assert "cu126" in _index_url(_run_cuda_repair(cuda_version = "12.6")) + + def test_cuda_124_selects_cu124(self): + assert "cu124" in _index_url(_run_cuda_repair(cuda_version = "12.4")) + + def test_cuda_118_selects_cu118(self): + assert "cu118" in _index_url(_run_cuda_repair(cuda_version = "11.8")) + + def test_unreadable_version_defaults_cu126(self): + # nvidia-smi runs but prints no CUDA version line (or fails). + mock_pip = _run_cuda_repair(cuda_version = "", smi_rc = 1) + assert "cu126" in _index_url(mock_pip) + + def test_proc_fallback_no_smi_defaults_cu126(self): + # NVIDIA usable via /proc fallback, nvidia-smi absent entirely. + mock_pip = _run_cuda_repair(smi_path = None) + assert "cu126" in _index_url(mock_pip) + + def test_detect_index_url_uses_pytorch_base(self): + with ( + patch.object(stack_mod.shutil, "which", return_value = None), + patch.object(stack_mod.os.path, "isfile", return_value = False), + ): + url = _detect_cuda_torch_index_url() + assert url == f"{stack_mod._PYTORCH_WHL_BASE}/cu126" + + +if __name__ == "__main__": + sys.exit(pytest.main([__file__, "-q"])) diff --git a/tests/studio/install/test_gpu_detection_followups.py b/tests/studio/install/test_gpu_detection_followups.py new file mode 100644 index 0000000000..983969a4d3 --- /dev/null +++ b/tests/studio/install/test_gpu_detection_followups.py @@ -0,0 +1,480 @@ +"""Tests for the GPU-detection follow-ups to PR 6174. + +PR 6174 made NVIDIA take precedence and added a /proc/driver/nvidia/gpus +fallback in install.sh and studio/install_python_stack.py. These tests cover the +same hardening ported to the llama.cpp prebuilt installer +(studio/install_llama_prebuilt.py) and the Studio shell setup (studio/setup.sh): + + * detect_host() recognises NVIDIA via /proc/driver/nvidia/gpus when nvidia-smi + is unavailable, and skips ROCm probing when NVIDIA is usable. + * setup.sh routes through a timeout-bounded NVIDIA probe with a /proc fallback + and only selects a CUDA/ROCm source build when the matching GPU is detected. + +All tests use mocks or source-level assertions -- no GPU, network, or real +nvidia-smi/rocminfo invocation. +""" + +import importlib.util +import sys +from pathlib import Path +from unittest.mock import MagicMock, patch + +import pytest + + +PACKAGE_ROOT = Path(__file__).resolve().parents[3] + +# Load studio/install_llama_prebuilt.py the same way the sibling suite does. +_MODULE_PATH = PACKAGE_ROOT / "studio" / "install_llama_prebuilt.py" +_SPEC = importlib.util.spec_from_file_location( + "studio_install_llama_prebuilt_followups", _MODULE_PATH +) +assert _SPEC is not None and _SPEC.loader is not None +prebuilt_mod = importlib.util.module_from_spec(_SPEC) +sys.modules[_SPEC.name] = prebuilt_mod +_SPEC.loader.exec_module(prebuilt_mod) + +detect_host = prebuilt_mod.detect_host +_apply_host_overrides = prebuilt_mod._apply_host_overrides + +SETUP_SH = PACKAGE_ROOT / "studio" / "setup.sh" + + +def _make_run_capture(rocminfo_stdout: str = ""): + """Return a fake run_capture: rocminfo reports rocminfo_stdout, everything + else (nvidia-smi, amd-smi) returns empty so only the patched probes matter.""" + + def _run_capture(cmd, *args, **kwargs): + exe = str(cmd[0]) if cmd else "" + result = MagicMock() + if exe.endswith("rocminfo"): + result.returncode = 0 + result.stdout = rocminfo_stdout + else: + result.returncode = 1 + result.stdout = "" + result.stderr = "" + return result + + return _run_capture + + +def _run_detect_host( + *, + machine: str = "x86_64", + system: str = "Linux", + which_map: dict | None = None, + proc_dir_entries: list | None = None, + rocminfo_stdout: str = "", + env: dict | None = None, +): + """Drive detect_host() against a fully synthetic host.""" + which_map = which_map or {} + proc_dir_entries = proc_dir_entries if proc_dir_entries is not None else [] + + real_isdir = prebuilt_mod.os.path.isdir + real_listdir = prebuilt_mod.os.listdir + proc_path = "/proc/driver/nvidia/gpus" + + def fake_isdir(p): + if str(p) == proc_path: + return bool(proc_dir_entries) + return real_isdir(p) + + def fake_listdir(p): + if str(p) == proc_path: + if not proc_dir_entries: + raise OSError("no such dir") + return list(proc_dir_entries) + return real_listdir(p) + + patches = [ + patch.object(prebuilt_mod.platform, "system", return_value = system), + patch.object(prebuilt_mod.platform, "machine", return_value = machine), + patch.object(prebuilt_mod.platform, "mac_ver", return_value = ("", ("", "", ""), "")), + patch.object(prebuilt_mod.shutil, "which", side_effect = lambda n: which_map.get(n)), + patch.object(prebuilt_mod, "run_capture", side_effect = _make_run_capture(rocminfo_stdout)), + patch.object(prebuilt_mod.os.path, "isdir", side_effect = fake_isdir), + patch.object(prebuilt_mod.os, "listdir", side_effect = fake_listdir), + patch.object(prebuilt_mod.os, "access", return_value = False), + patch.dict(prebuilt_mod.os.environ, env or {}, clear = False), + ] + for p in patches: + p.start() + try: + # Ensure CUDA_VISIBLE_DEVICES does not leak in from the test host unless + # the scenario sets it explicitly. + if env is None or "CUDA_VISIBLE_DEVICES" not in env: + prebuilt_mod.os.environ.pop("CUDA_VISIBLE_DEVICES", None) + return detect_host() + finally: + for p in patches: + p.stop() + + +# ── install_llama_prebuilt.detect_host(): /proc NVIDIA fallback ────────────── + + +class TestDetectHostProcFallback: + def test_proc_fallback_marks_physical_nvidia_when_smi_absent(self): + """No nvidia-smi, but /proc/driver/nvidia/gpus is populated -> NVIDIA.""" + host = _run_detect_host( + which_map = {}, # nvidia-smi resolves to None + proc_dir_entries = ["0000:01:00.0"], + ) + assert host.has_physical_nvidia is True + + def test_proc_fallback_has_usable_nvidia_when_devices_visible(self): + """Default CUDA_VISIBLE_DEVICES (unset) -> visible tokens non-empty -> usable.""" + host = _run_detect_host( + which_map = {}, + proc_dir_entries = ["0000:01:00.0"], + ) + assert host.has_usable_nvidia is True + + def test_proc_fallback_not_usable_when_devices_hidden(self): + """CUDA_VISIBLE_DEVICES='' hides all GPUs -> physical yes, usable no.""" + host = _run_detect_host( + which_map = {}, + proc_dir_entries = ["0000:01:00.0"], + env = {"CUDA_VISIBLE_DEVICES": ""}, + ) + assert host.has_physical_nvidia is True + assert host.has_usable_nvidia is False + + def test_empty_proc_dir_does_not_mark_nvidia(self): + """A driver dir that exists but is empty must not assert a GPU.""" + host = _run_detect_host(which_map = {}, proc_dir_entries = []) + assert host.has_physical_nvidia is False + + def test_proc_fallback_is_linux_only(self): + """The /proc fallback must not run on Windows (path is Linux-only).""" + host = _run_detect_host( + system = "Windows", + machine = "amd64", + which_map = {}, + proc_dir_entries = ["0000:01:00.0"], + ) + assert host.has_physical_nvidia is False + + +# ── install_llama_prebuilt.detect_host(): NVIDIA precedence over ROCm ──────── + + +class TestDetectHostNvidiaPrecedence: + def test_rocm_probe_skipped_when_proc_nvidia_present(self): + """rocminfo reports gfx1100, but a proc-detected NVIDIA GPU wins.""" + host = _run_detect_host( + which_map = {"rocminfo": "/usr/bin/rocminfo"}, + proc_dir_entries = ["0000:01:00.0"], + rocminfo_stdout = " Name: gfx1100\n", + ) + assert host.has_usable_nvidia is True + assert host.has_rocm is False + + def test_rocm_detected_when_no_nvidia(self): + """With no NVIDIA signal at all, rocminfo gfx1100 -> has_rocm True.""" + host = _run_detect_host( + which_map = {"rocminfo": "/usr/bin/rocminfo"}, + proc_dir_entries = [], + rocminfo_stdout = " Name: gfx1100\n", + ) + assert host.has_usable_nvidia is False + assert host.has_rocm is True + + +# ── _apply_host_overrides: forwarded --rocm-gfx / --has-rocm still win ─────── + + +class TestOverridesStillWin: + def test_forwarded_gfx_forces_rocm_on_non_nvidia_host(self): + host = _run_detect_host(which_map = {}, proc_dir_entries = []) + assert host.has_rocm is False + overridden = _apply_host_overrides(host, override_rocm_gfx = "gfx1100") + assert overridden.has_rocm is True + assert overridden.rocm_gfx_target == "gfx1100" + + def test_override_has_rocm_forces_rocm(self): + host = _run_detect_host(which_map = {}, proc_dir_entries = []) + overridden = _apply_host_overrides(host, override_has_rocm = True) + assert overridden.has_rocm is True + + def test_force_cpu_drops_nvidia_attributes(self): + host = _run_detect_host(which_map = {}, proc_dir_entries = ["0000:01:00.0"]) + assert host.has_usable_nvidia is True + overridden = _apply_host_overrides(host, force_cpu = True) + assert overridden.has_usable_nvidia is False + assert overridden.has_physical_nvidia is False + assert overridden.has_rocm is False + + +# ── setup.sh source-level guarantees ──────────────────────────────────────── + + +class TestSetupShHardening: + @pytest.fixture(scope = "class") + def setup_src(self) -> str: + return SETUP_SH.read_text(encoding = "utf-8") + + def test_has_usable_nvidia_helper_exists(self, setup_src): + assert "_setup_has_usable_nvidia_gpu()" in setup_src + + def test_helper_uses_proc_fallback(self, setup_src): + start = setup_src.find("_setup_has_usable_nvidia_gpu()") + end = setup_src.find("\n}", start) + body = setup_src[start:end] + assert ( + "/proc/driver/nvidia/gpus" in body + ), "_setup_has_usable_nvidia_gpu must fall back to /proc/driver/nvidia/gpus" + + def test_gpu_summary_uses_helper(self, setup_src): + assert "if _setup_has_usable_nvidia_gpu; then" in setup_src + + def test_timeout_wrapper_exists(self, setup_src): + start = setup_src.find("_setup_run_smi()") + assert start >= 0, "_setup_run_smi timeout wrapper must exist" + end = setup_src.find("\n}", start) + body = setup_src[start:end] + assert "timeout 10" in body + assert "command -v timeout" in body + + def test_cuda_source_build_gated_on_usable_nvidia(self, setup_src): + """The nvcc source-build search must be gated on _setup_nvidia_usable. + + The hidden-GPU policy (CUDA_VISIBLE_DEVICES=""/-1) lives inside + _setup_has_usable_nvidia_gpu, so the gate itself only needs the flag. + """ + anchor = setup_src.find('NVCC_PATH=""\n') + assert anchor >= 0 + window = setup_src[anchor : anchor + 700] + assert ( + 'if [ "$_setup_nvidia_usable" = true ]' in window + ), "CUDA toolkit search must require a usable NVIDIA GPU, not just nvcc" + + def test_nvidia_helper_honours_hidden_cvd(self, setup_src): + """_setup_has_usable_nvidia_gpu must consult the hidden-CVD helper so + CUDA_VISIBLE_DEVICES=""/-1 suppresses NVIDIA before the AMD probes are + gated (mixed hosts steered to the AMD card keep the ROCm route).""" + assert "_setup_cvd_hides_nvidia()" in setup_src + start = setup_src.find("_setup_has_usable_nvidia_gpu() {") + end = setup_src.find("\n}", start) + body = setup_src[start:end] + assert "_setup_cvd_hides_nvidia" in body + + def test_rocm_source_build_gated_on_amd_detected(self, setup_src): + """The hipcc source-build search must be gated on _setup_amd_detected.""" + anchor = setup_src.find('ROCM_HIPCC=""') + assert anchor >= 0 + window = setup_src[anchor : anchor + 400] + assert ( + '[ "$_setup_amd_detected" = true ]' in window + ), "ROCm toolkit search must require a detected AMD GPU, not just hipcc" + + def test_compute_cap_probe_timeout_wrapped(self, setup_src): + assert "_setup_run_smi nvidia-smi --query-gpu=compute_cap" in setup_src + + def test_driver_version_probe_timeout_wrapped(self, setup_src): + start = setup_src.find("_cuda_driver_max_version()") + end = setup_src.find("\n}", start) + body = setup_src[start:end] + assert "_setup_run_smi nvidia-smi" in body + + +# TEST: install.sh -- UNSLOTH_TORCH_BACKEND classified on the final path segment + + +class TestBackendExportLeafClassification: + """A custom UNSLOTH_PYTORCH_MIRROR whose base path contains "rocm" or + "gfx" must not mislabel a cu*/cpu index as ROCm; classification uses the + final path segment of TORCH_INDEX_URL only.""" + + @pytest.fixture(scope = "class") + def install_src(self) -> str: + return (PACKAGE_ROOT / "install.sh").read_text(encoding = "utf-8") + + def test_export_block_uses_leaf(self, install_src): + anchor = install_src.find("_torch_index_leaf=") + assert anchor >= 0, "backend export must classify on the final path segment" + window = install_src[anchor : anchor + 500] + assert 'export UNSLOTH_TORCH_BACKEND="rocm"' in window + assert 'export UNSLOTH_TORCH_BACKEND="cpu"' in window + assert 'export UNSLOTH_TORCH_BACKEND="cuda"' in window + + def test_leaf_classification_behaviour(self, tmp_path): + import subprocess as sp + + script = tmp_path / "leaf.sh" + src = (PACKAGE_ROOT / "install.sh").read_text(encoding = "utf-8") + anchor = src.find("_torch_index_leaf=") + block = src[anchor : src.find("esac", anchor) + 4] + # Drive the extracted block with adversarial mirror URLs. + script.write_text( + "#!/bin/sh\n" + 'TORCH_INDEX_URL="$1"\n' + block + "\n" + 'printf "%s" "$UNSLOTH_TORCH_BACKEND"\n' + ) + cases = { + "https://download.pytorch.org/whl/cu128": "cuda", + "https://download.pytorch.org/whl/cpu": "cpu", + "https://download.pytorch.org/whl/rocm6.4": "rocm", + "https://repo.radeon.com/rocm/manylinux/rocm-rel-7.2.1/": "rocm", + "https://repo.amd.com/rocm/whl/gfx1151/": "rocm", + "https://mirror.local/rocm-cache/cu128": "cuda", + "https://mirror.local/gfx-cache/cpu": "cpu", + } + for url, expected in cases.items(): + out = sp.run( + ["sh", str(script), url], capture_output = True, text = True, timeout = 30 + ).stdout.strip() + assert out == expected, f"{url} classified as {out!r}, expected {expected!r}" + + +# TEST: CUDA_VISIBLE_DEVICES=""/-1 hides NVIDIA in every usable-GPU helper + + +_STACK_PATH = PACKAGE_ROOT / "studio" / "install_python_stack.py" +_STACK_SPEC = importlib.util.spec_from_file_location( + "studio_install_python_stack_followups", _STACK_PATH +) +assert _STACK_SPEC is not None and _STACK_SPEC.loader is not None +stack_mod = importlib.util.module_from_spec(_STACK_SPEC) +sys.modules[_STACK_SPEC.name] = stack_mod +_STACK_SPEC.loader.exec_module(stack_mod) + + +def _stack_nvidia_usable(cvd): + """Drive install_python_stack._has_usable_nvidia_gpu with a mocked + nvidia-smi that always reports a GPU; cvd = None removes the env var.""" + + def fake_run(cmd, *args, **kwargs): + result = MagicMock() + result.returncode = 0 + result.stdout = "GPU 0: NVIDIA Fake (UUID: GPU-x)\n" + return result + + env = {} if cvd is None else {"CUDA_VISIBLE_DEVICES": cvd} + with ( + patch.object( + stack_mod.shutil, + "which", + side_effect = lambda n: "/usr/bin/nvidia-smi" if n == "nvidia-smi" else None, + ), + patch.object(stack_mod.subprocess, "run", side_effect = fake_run), + patch.dict(stack_mod.os.environ, env, clear = False), + ): + if cvd is None: + stack_mod.os.environ.pop("CUDA_VISIBLE_DEVICES", None) + return stack_mod._has_usable_nvidia_gpu() + + +class TestHiddenCvdNotUsable: + """CUDA_VISIBLE_DEVICES set to "" or "-1" deliberately hides every NVIDIA + device (mixed AMD+NVIDIA hosts steering work to the AMD card). All three + _has_usable_nvidia_gpu implementations (install_python_stack.py, install.sh, + setup.sh) must report the GPU as not usable so the AMD/CPU routes run, + matching install_llama_prebuilt.py's has_usable_nvidia.""" + + def test_python_unset_cvd_is_usable(self): + assert _stack_nvidia_usable(None) is True + + def test_python_empty_cvd_not_usable(self): + assert _stack_nvidia_usable("") is False + + def test_python_minus_one_not_usable(self): + assert _stack_nvidia_usable("-1") is False + + def test_python_padded_minus_one_not_usable(self): + assert _stack_nvidia_usable(" -1 ") is False + + def test_python_explicit_device_is_usable(self): + assert _stack_nvidia_usable("0") is True + + def test_python_device_list_is_usable(self): + assert _stack_nvidia_usable("0,1") is True + + def test_hidden_nvidia_restores_rocm_detection(self): + """Mixed host, NVIDIA hidden via CVD=-1, rocminfo reports gfx1100: + _has_rocm_gpu must proceed past the NVIDIA guard and return True + (before this fix the guard ignored CVD and blocked ROCm).""" + + def fake_run(cmd, *args, **kwargs): + result = MagicMock() + result.returncode = 0 + exe = str(cmd[0]) + if exe.endswith("rocminfo"): + result.stdout = " Name: gfx1100\n" + else: + result.stdout = "GPU 0: NVIDIA Fake (UUID: GPU-x)\n" + return result + + which_map = { + "rocminfo": "/usr/bin/rocminfo", + "nvidia-smi": "/usr/bin/nvidia-smi", + } + with ( + patch.object(stack_mod.shutil, "which", side_effect = which_map.get), + patch.object(stack_mod.subprocess, "run", side_effect = fake_run), + patch.dict(stack_mod.os.environ, {"CUDA_VISIBLE_DEVICES": "-1"}, clear = False), + ): + assert stack_mod._has_rocm_gpu() is True + + @staticmethod + def _run_sh_helper(tmp_path, src: str, fn_names: list, cvd): + """Extract shell functions, run the usable-GPU one against a fake + nvidia-smi, and return "usable"/"not_usable".""" + import os as _os + import subprocess as sp + + blocks = [] + for name in fn_names: + start = src.find(f"{name}() {{") + assert start >= 0, f"{name} missing" + end = src.find("\n}", start) + 2 + blocks.append(src[start:end]) + fake_bin = tmp_path / "bin" + fake_bin.mkdir(exist_ok = True) + smi = fake_bin / "nvidia-smi" + smi.write_text("#!/bin/sh\necho 'GPU 0: NVIDIA Fake (UUID: GPU-x)'\n") + smi.chmod(0o755) + script = tmp_path / "probe.sh" + script.write_text( + "#!/bin/sh\n" + "\n".join(blocks) + "\n" + f"if {fn_names[-1]}; then echo usable; else echo not_usable; fi\n" + ) + env = dict(_os.environ) + env["PATH"] = f"{fake_bin}:{env['PATH']}" + if cvd is None: + env.pop("CUDA_VISIBLE_DEVICES", None) + else: + env["CUDA_VISIBLE_DEVICES"] = cvd + return sp.run( + ["sh", str(script)], capture_output = True, text = True, timeout = 30, env = env + ).stdout.strip() + + @pytest.mark.parametrize( + "cvd, expected", + [(None, "usable"), ("", "not_usable"), ("-1", "not_usable"), ("0", "usable")], + ) + def test_install_sh_helper_cvd(self, tmp_path, cvd, expected): + src = (PACKAGE_ROOT / "install.sh").read_text(encoding = "utf-8") + out = self._run_sh_helper( + tmp_path, + src, + ["_run_bounded", "_cvd_hides_nvidia", "_has_usable_nvidia_gpu"], + cvd, + ) + assert out == expected + + @pytest.mark.parametrize( + "cvd, expected", + [(None, "usable"), ("", "not_usable"), ("-1", "not_usable"), ("0", "usable")], + ) + def test_setup_sh_helper_cvd(self, tmp_path, cvd, expected): + src = SETUP_SH.read_text(encoding = "utf-8") + out = self._run_sh_helper( + tmp_path, + src, + ["_setup_run_smi", "_setup_cvd_hides_nvidia", "_setup_has_usable_nvidia_gpu"], + cvd, + ) + assert out == expected diff --git a/tests/studio/install/test_probe_timeouts.py b/tests/studio/install/test_probe_timeouts.py new file mode 100644 index 0000000000..acea0ed34d --- /dev/null +++ b/tests/studio/install/test_probe_timeouts.py @@ -0,0 +1,202 @@ +"""Tests that NVIDIA probes in the installers are bounded by a timeout. + +Covers audit findings 5 and 6: a wedged nvidia-smi must not hang the installer, +and the Windows probe must require a real GPU listing (not just exit code 0). + +Source-level assertions verify the guards are present in install.sh / install.ps1 +/ setup.ps1; one behavioral shell test confirms the bash helper actually returns +within the timeout when nvidia-smi hangs. +""" + +import os +import shutil +import stat +import subprocess +import sys +import tempfile +from pathlib import Path + +import pytest + + +PACKAGE_ROOT = Path(__file__).resolve().parents[3] +INSTALL_SH = PACKAGE_ROOT / "install.sh" +INSTALL_PS1 = PACKAGE_ROOT / "install.ps1" +SETUP_PS1 = PACKAGE_ROOT / "studio" / "setup.ps1" + + +def _extract_sh_function_body(source: str, name: str) -> str: + """Return a shell function body from `source` by brace matching.""" + needle = f"{name}() {{" + start = source.find(needle) + if start < 0: + return "" + depth = 0 + i = start + len(needle) - 1 + n = len(source) + while i < n: + ch = source[i] + if ch == "{": + depth += 1 + elif ch == "}": + depth -= 1 + if depth == 0: + return source[start : i + 1] + i += 1 + return source[start:] + + +# ── install.sh: _run_bounded helper and its use at every nvidia-smi call ── + + +class TestInstallShBoundedProbe: + def _src(self) -> str: + return INSTALL_SH.read_text(encoding = "utf-8") + + def test_run_bounded_helper_defined(self): + body = _extract_sh_function_body(self._src(), "_run_bounded") + assert body, "install.sh must define a _run_bounded helper" + assert ( + "command -v timeout" in body + ), "_run_bounded must check for the `timeout` binary before using it" + assert "timeout 10" in body, "_run_bounded must apply a 10s timeout" + # Must fall back to running unbounded when `timeout` is unavailable + # (e.g. macOS) so semantics are unchanged there. + assert ( + "else" in body and '"$@"' in body + ), "_run_bounded must run the command unbounded when `timeout` is absent" + + def test_nvidia_smi_dash_l_probe_is_bounded(self): + body = _extract_sh_function_body(self._src(), "_has_usable_nvidia_gpu") + assert body, "install.sh must define _has_usable_nvidia_gpu" + # The -L probe must go through the bounded runner, not call nvidia-smi raw. + assert ( + '_run_bounded "$_nvsmi" -L' in body + ), "_has_usable_nvidia_gpu must run nvidia-smi -L through _run_bounded" + # The /proc fallback from PR 6174 must still be present. + assert "/proc/driver/nvidia" in body + + def test_cuda_version_parse_is_bounded(self): + body = _extract_sh_function_body(self._src(), "get_torch_index_url") + assert body, "install.sh must define get_torch_index_url" + assert ( + "_run_bounded" in body + ), "get_torch_index_url CUDA-version parse must run nvidia-smi through _run_bounded" + # The locale must be forced without depending on `env` being on PATH. + assert "LC_ALL=C" in body + # _nvidia_detected gating from PR 6174 must remain. + assert "_nvidia_detected" in body + + def test_no_unbounded_nvidia_smi_invocation_remains(self): + """Every nvidia-smi *execution* in install.sh goes through _run_bounded. + + `command -v nvidia-smi` and `-x /usr/bin/nvidia-smi` are resolution + checks, not executions, and are allowed. An execution looks like + `"$_nvsmi" ...` / `$_smi ...` / `nvidia-smi -L`. + """ + body_nvidia = _extract_sh_function_body(self._src(), "_has_usable_nvidia_gpu") + body_torch = _extract_sh_function_body(self._src(), "get_torch_index_url") + # In _has_usable_nvidia_gpu the only execution of $_nvsmi must be bounded. + assert '"$_nvsmi" -L' not in body_nvidia.replace( + '_run_bounded "$_nvsmi" -L', "" + ), "found an unbounded nvidia-smi -L execution in _has_usable_nvidia_gpu" + # In get_torch_index_url the $_smi execution must be bounded. + assert ( + "LC_ALL=C $_smi" not in body_torch + ), "found an unbounded LC_ALL=C $_smi execution in get_torch_index_url" + + +# ── install.ps1 / setup.ps1: bounded, GPU-row-validated Windows probe ── + + +class TestPowerShellBoundedProbe: + @pytest.mark.parametrize("path", [INSTALL_PS1, SETUP_PS1]) + def test_bounded_helper_present(self, path): + src = path.read_text(encoding = "utf-8") + assert ( + "function Invoke-NvidiaSmiBounded" in src + ), f"{path.name} must define Invoke-NvidiaSmiBounded" + assert ( + "WaitForExit($TimeoutSec * 1000)" in src + ), f"{path.name} bounded probe must use WaitForExit with a timeout" + # Kill + sentinel on timeout, mirroring Invoke-AmdSmiNoElevate. + assert ( + "$proc.Kill()" in src and "124" in src + ), f"{path.name} must kill nvidia-smi and signal a timeout exit code" + + @pytest.mark.parametrize("path", [INSTALL_PS1, SETUP_PS1]) + def test_probe_requires_gpu_row(self, path): + src = path.read_text(encoding = "utf-8") + assert ( + "function Test-NvidiaSmiHasGpu" in src + ), f"{path.name} must define Test-NvidiaSmiHasGpu" + assert "@('-L')" in src, f"{path.name} must probe nvidia-smi with -L" + assert ( + "^GPU\\s+\\d+:" in src + ), f"{path.name} must require a 'GPU :' data row, not just exit code 0" + + @pytest.mark.parametrize("path", [INSTALL_PS1, SETUP_PS1]) + def test_detection_uses_validated_probe(self, path): + src = path.read_text(encoding = "utf-8") + # The exit-code-only pattern must be gone from the detection block. + assert ( + "& $nvSmiCmd.Source *> $null" not in src + ), f"{path.name} must not use the exit-code-only nvidia-smi probe" + assert ( + "Test-NvidiaSmiHasGpu $nvSmiCmd.Source" in src + ), f"{path.name} PATH probe must use Test-NvidiaSmiHasGpu" + assert ( + "Test-NvidiaSmiHasGpu $p" in src + ), f"{path.name} hardcoded-path fallback must use Test-NvidiaSmiHasGpu" + + +# ── Behavioral: a hanging nvidia-smi must not hang _has_usable_nvidia_gpu ── + + +def _have_timeout() -> bool: + return shutil.which("timeout") is not None + + +@pytest.mark.skipif(not _have_timeout(), reason = "`timeout` binary not available") +def test_has_usable_nvidia_gpu_returns_under_timeout(): + """Extract _run_bounded + _has_usable_nvidia_gpu, point them at a fake + nvidia-smi that sleeps 30s, and assert the probe returns well under that. + """ + src = INSTALL_SH.read_text(encoding = "utf-8") + helper = _extract_sh_function_body(src, "_run_bounded") + fn = _extract_sh_function_body(src, "_has_usable_nvidia_gpu") + assert helper and fn + + workdir = tempfile.mkdtemp(prefix = "pr6174_timeout_", dir = str(PACKAGE_ROOT.parent)) + try: + fake_dir = Path(workdir, "bin") + fake_dir.mkdir() + fake_smi = fake_dir / "nvidia-smi" + fake_smi.write_text("#!/bin/sh\nsleep 30\n") + fake_smi.chmod(fake_smi.stat().st_mode | stat.S_IEXEC | stat.S_IXGRP | stat.S_IXOTH) + + # Build a minimal PATH that includes the fake nvidia-smi plus the real + # `timeout`/`awk`/`ls` it needs. Use the fake dir first so it wins. + real_bins = {Path(shutil.which(c)).parent for c in ("timeout", "awk", "ls", "sh")} + path_env = os.pathsep.join([str(fake_dir)] + [str(p) for p in real_bins]) + + # Force the /proc fallback off so the result depends only on the probe, + # and so a host with real NVIDIA does not mask the timeout behaviour. + script = ( + f"{helper}\n{fn}\n" + "if _has_usable_nvidia_gpu; then echo DETECTED; else echo NONE; fi\n" + ) + proc = subprocess.run( + ["sh", "-c", script], + env = {"PATH": path_env}, + stdout = subprocess.PIPE, + stderr = subprocess.DEVNULL, + text = True, + timeout = 20, # generous: the internal timeout is 10s, sleep is 30s + ) + # The probe must have returned (not hung). On this CI host /proc/driver/ + # nvidia/gpus is absent, so a timed-out smi yields NONE; on a real NVIDIA + # host the /proc fallback yields DETECTED. Either way it must not hang. + assert proc.stdout.strip() in {"NONE", "DETECTED"} + finally: + shutil.rmtree(workdir, ignore_errors = True) From 49024462f39e79a3b477306f46ea2f9f1dd18e73 Mon Sep 17 00:00:00 2001 From: Lee Jackson <130007945+Imagineer99@users.noreply.github.com> Date: Thu, 11 Jun 2026 13:13:27 +0100 Subject: [PATCH 12/42] Studio: ignore unsupported env proxy during Studio startup (#6102) * fix: ignore unsupported env proxy during Studio startup * fix: handle missing socksio env proxy at startup * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * Match printf logging style and inline the proxy predicate for PR #6102 * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci --------- Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com> Co-authored-by: Daniel Han --- .../core/inference/external_provider.py | 21 ++++- .../tests/test_external_provider_proxy_env.py | 80 +++++++++++++++++++ 2 files changed, 98 insertions(+), 3 deletions(-) create mode 100644 studio/backend/tests/test_external_provider_proxy_env.py diff --git a/studio/backend/core/inference/external_provider.py b/studio/backend/core/inference/external_provider.py index a7d44b992d..cae001c34d 100644 --- a/studio/backend/core/inference/external_provider.py +++ b/studio/backend/core/inference/external_provider.py @@ -469,9 +469,24 @@ def _apply_mistral_reasoning_controls( # Shared client reused across all requests for HTTP connection pooling. -# Auth headers and timeouts are per-request, so one client handles every -# provider without storing credentials. -_http_client = httpx.AsyncClient() +# Auth headers and timeouts are passed per-request, so a single client +# handles every provider without storing credentials. +def _create_shared_http_client() -> httpx.AsyncClient: + # Unsupported env proxy schemes (socks:// etc) raise at construction and + # would crash Studio startup (#6090); retry ignoring env proxies instead. + try: + return httpx.AsyncClient() + except (ImportError, ValueError) as exc: + exc_str = str(exc) + if "Unknown scheme for proxy URL" not in exc_str and "socksio" not in exc_str: + raise + logger.warning( + "Ignoring unsupported environment proxy for the shared HTTP client: %s", exc_str + ) + return httpx.AsyncClient(trust_env = False) + + +_http_client = _create_shared_http_client() # Cap per-image fetch well below Gemini's ~20 MB total request budget. diff --git a/studio/backend/tests/test_external_provider_proxy_env.py b/studio/backend/tests/test_external_provider_proxy_env.py new file mode 100644 index 0000000000..f17b655908 --- /dev/null +++ b/studio/backend/tests/test_external_provider_proxy_env.py @@ -0,0 +1,80 @@ +# SPDX-License-Identifier: AGPL-3.0-only +# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0 + +from pathlib import Path +import importlib.util +import sys + +_BACKEND_DIR = str(Path(__file__).resolve().parent.parent) +if _BACKEND_DIR not in sys.path: + sys.path.insert(0, _BACKEND_DIR) + +_EXTERNAL_PROVIDER_PATH = ( + Path(__file__).resolve().parent.parent / "core/inference/external_provider.py" +) + + +def _load_external_provider_module(): + spec = importlib.util.spec_from_file_location( + "external_provider_under_test", + _EXTERNAL_PROVIDER_PATH, + ) + module = importlib.util.module_from_spec(spec) + spec.loader.exec_module(module) + return module + + +def test_shared_http_client_ignores_unsupported_proxy_scheme(monkeypatch): + ep_mod = _load_external_provider_module() + calls = [] + + class FakeAsyncClient: + def __init__(self, **kwargs): + calls.append(kwargs) + if kwargs.get("trust_env") is not False: + raise ValueError("Unknown scheme for proxy URL URL('socks4://127.0.0.1:12345')") + + monkeypatch.setattr(ep_mod.httpx, "AsyncClient", FakeAsyncClient) + + client = ep_mod._create_shared_http_client() + + assert isinstance(client, FakeAsyncClient) + assert calls == [{}, {"trust_env": False}] + + +def test_shared_http_client_ignores_missing_socksio(monkeypatch): + ep_mod = _load_external_provider_module() + calls = [] + + class FakeAsyncClient: + def __init__(self, **kwargs): + calls.append(kwargs) + if kwargs.get("trust_env") is not False: + raise ImportError( + "Using SOCKS proxy, but the 'socksio' package is not installed. " + "Make sure to install httpx using `pip install httpx[socks]`." + ) + + monkeypatch.setattr(ep_mod.httpx, "AsyncClient", FakeAsyncClient) + + client = ep_mod._create_shared_http_client() + + assert isinstance(client, FakeAsyncClient) + assert calls == [{}, {"trust_env": False}] + + +def test_shared_http_client_reraises_other_value_errors(monkeypatch): + ep_mod = _load_external_provider_module() + + class FakeAsyncClient: + def __init__(self, **kwargs): + raise ValueError("different httpx setup error") + + monkeypatch.setattr(ep_mod.httpx, "AsyncClient", FakeAsyncClient) + + try: + ep_mod._create_shared_http_client() + except ValueError as exc: + assert str(exc) == "different httpx setup error" + else: + raise AssertionError("expected ValueError") From d0ffe2638321ff5f6abbcfe8d0567e1170cc3a74 Mon Sep 17 00:00:00 2001 From: Daniel Han Date: Thu, 11 Jun 2026 05:18:12 -0700 Subject: [PATCH 13/42] Keep audio feature extractors right padded when loading processors (#6157) * Keep audio feature extractors right padded when loading processors FastBaseModel.from_pretrained passes padding_side=left to AutoProcessor.from_pretrained for generation, and ProcessorMixin forwards the kwarg to every sub-component, including audio feature extractors. Stock transformers right-pads audio: frame-validity masks assume trailing padding. The leaked left padding shifts mel content to the end of the 30s window for Whisper and gives Gemma 4 one extra valid mel frame on clip lengths off the hop boundary, desyncing audio features from placeholder tokens and crashing training on transformers 5.5.0 to 5.9.x with 'Audio features and audio tokens do not match'. Reset the feature extractor to right padding at the single processor finalization point. Text tokenizer padding stays left. Verified on transformers 5.5.0: the loaded processor now matches a fresh stock AutoProcessor exactly, and a 9-combination Gemma 4 audio forward probe goes from 4 failures to none. * Add explicit None check before reading feature extractor attributes * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * Tighten comments --------- Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com> --- unsloth/models/vision.py | 9 +++++++++ 1 file changed, 9 insertions(+) diff --git a/unsloth/models/vision.py b/unsloth/models/vision.py index 34336aa80f..58814699b3 100644 --- a/unsloth/models/vision.py +++ b/unsloth/models/vision.py @@ -1280,6 +1280,15 @@ class FastBaseModel: tokenizer.padding_side = "left" # Force inference if hasattr(tokenizer, "tokenizer"): tokenizer.tokenizer.padding_side = "left" # Force inference + # Audio feature extractors must stay right padded: left (a text setting, + # forwarded by from_pretrained) shifts Whisper mels and desyncs Gemma 4 + # audio token counts (crash on transformers < 5.10). + feature_extractor = getattr(tokenizer, "feature_extractor", None) + if ( + feature_extractor is not None + and getattr(feature_extractor, "padding_side", None) == "left" + ): + feature_extractor.padding_side = "right" m = model while hasattr(m, "model"): m.max_seq_length = max_seq_length From 36bc0394d93ad5168ade756ad8021e643caca0f5 Mon Sep 17 00:00:00 2001 From: Daniel Han Date: Thu, 11 Jun 2026 05:22:38 -0700 Subject: [PATCH 14/42] Studio: fetch the release source asset for exact (mix) source builds (#6195) * Studio: fetch the release source asset for exact (mix) source builds The source-build fallback rebuilt the codeload/archive URL from the source repo and commit. A mix build's merged tree is never pushed to any repo (it ships only as the release's llama.cpp-source-commit-.tar.gz asset), so codeload 404s on the merge commit and an uncovered host could not build from source. When an exact-source asset exists, fetch it directly from the release and keep codeload as the fallback for vanilla builds whose commit is real. * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci --------- Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com> --- studio/install_llama_prebuilt.py | 30 ++++- .../test_install_llama_prebuilt_logic.py | 104 ++++++++++++++++++ 2 files changed, 133 insertions(+), 1 deletion(-) diff --git a/studio/install_llama_prebuilt.py b/studio/install_llama_prebuilt.py index 22fab82923..6c3be0e09f 100644 --- a/studio/install_llama_prebuilt.py +++ b/studio/install_llama_prebuilt.py @@ -1095,6 +1095,20 @@ def commit_source_archive_urls(repo: str, source_commit: str) -> list[str]: ] +def release_asset_download_url( + repo: str | None, release_tag: str | None, asset_name: str | None +) -> str | None: + """Direct download URL for a release asset, or None if any part is missing. + A mix build's merged commit is never pushed, so its source tree is only + reachable as this asset (codeload would 404 on the merge commit).""" + if not repo or not release_tag or not asset_name: + return None + return ( + f"https://github.com/{repo}/releases/download/" + f"{urllib.parse.quote(release_tag, safe = '')}/{urllib.parse.quote(asset_name, safe = '')}" + ) + + def github_release_assets(repo: str, tag: str) -> dict[str, str]: payload = fetch_json( f"https://api.github.com/repos/{repo}/releases/tags/{urllib.parse.quote(tag, safe = '')}" @@ -4495,13 +4509,17 @@ def hydrate_source_tree( expected_sha256: str | None, source_label: str | None = None, exact_source: bool = False, + asset_url: str | None = None, ) -> None: archive_path = work_dir / f"llama.cpp-source-{source_ref}.tar.gz" - source_urls = ( + repo_urls = ( commit_source_archive_urls(source_repo, source_ref) if exact_source else upstream_source_archive_urls(source_ref) ) + # Prefer the published release asset (the only copy of a mix build's merged + # tree); fall back to codeload/archive for vanilla builds whose commit is real. + source_urls = ([asset_url] if asset_url else []) + repo_urls label = source_label or f"llama.cpp source tree for {source_ref}" extract_dir = Path(tempfile.mkdtemp(prefix = "source-extract-", dir = work_dir)) @@ -6404,6 +6422,15 @@ def validate_prebuilt_choice( source_repo, source_ref, source_archive, exact_source = preferred_source_archive( approved_checksums, llama_tag ) + # For an exact (mix) source the merge commit lives only in the release asset, + # not in any repo, so fetch the asset directly; codeload stays the fallback. + asset_url = ( + release_asset_download_url( + approved_checksums.repo, approved_checksums.release_tag, source_archive.asset_name + ) + if exact_source and source_archive is not None + else None + ) if exact_source: log(f"hydrating exact llama.cpp source for {source_repo}@{source_ref} into {install_dir}") else: @@ -6420,6 +6447,7 @@ def validate_prebuilt_choice( else f"llama.cpp source tree for {llama_tag}" ), exact_source = exact_source, + asset_url = asset_url, ) log(f"overlaying prebuilt bundle {choice.name} into {install_dir}") server_path, quantize_path = install_from_archives(choice, host, install_dir, work_dir) diff --git a/tests/studio/install/test_install_llama_prebuilt_logic.py b/tests/studio/install/test_install_llama_prebuilt_logic.py index 057612d27b..7ca87aa8c4 100644 --- a/tests/studio/install/test_install_llama_prebuilt_logic.py +++ b/tests/studio/install/test_install_llama_prebuilt_logic.py @@ -206,6 +206,110 @@ def test_hydrate_source_tree_extracts_upstream_archive_contents( assert not (install_dir / f"llama.cpp-{upstream_tag}").exists() +def test_release_asset_download_url(): + fn = INSTALL_LLAMA_PREBUILT.release_asset_download_url + assert fn( + "unslothai/llama.cpp", "b9000-mix-abc1234", "llama.cpp-source-commit-deadbeef.tar.gz" + ) == ( + "https://github.com/unslothai/llama.cpp/releases/download/" + "b9000-mix-abc1234/llama.cpp-source-commit-deadbeef.tar.gz" + ) + # Any missing component -> None (no asset url, caller falls back to codeload). + assert fn(None, "b9000", "x.tar.gz") is None + assert fn("unslothai/llama.cpp", None, "x.tar.gz") is None + assert fn("unslothai/llama.cpp", "b9000", None) is None + + +def _mk_source_tarball(path: Path, tag: str) -> None: + with tarfile.open(path, "w:gz") as archive: + add_bytes_to_tar( + archive, f"llama.cpp-{tag}/CMakeLists.txt", b"cmake_minimum_required(VERSION 3.14)\n" + ) + add_bytes_to_tar( + archive, + f"llama.cpp-{tag}/convert_hf_to_gguf.py", + b"#!/usr/bin/env python3\nimport gguf\n", + ) + add_bytes_to_tar(archive, f"llama.cpp-{tag}/gguf-py/gguf/__init__.py", b"__all__ = []\n") + + +def test_hydrate_source_tree_prefers_release_asset_for_mix( + tmp_path: Path, monkeypatch: pytest.MonkeyPatch +): + # A mix build's merge commit is in no repo, so the codeload/archive URLs 404. + # hydrate must fetch the release asset and never touch codeload. + commit = "a" * 40 + archive_path = tmp_path / "merged-source.tar.gz" + _mk_source_tarball(archive_path, f"b9000-mix-{commit[:7]}") + asset_url = INSTALL_LLAMA_PREBUILT.release_asset_download_url( + "unslothai/llama.cpp", "b9000-mix-abc1234", f"llama.cpp-source-commit-{commit}.tar.gz" + ) + codeload_urls = set( + INSTALL_LLAMA_PREBUILT.commit_source_archive_urls("unslothai/llama.cpp", commit) + ) + seen = [] + + def fake_download_file(url: str, destination: Path) -> None: + seen.append(url) + if url in codeload_urls: + raise AssertionError("codeload was hit even though the release asset was available") + assert url == asset_url + destination.write_bytes(archive_path.read_bytes()) + + monkeypatch.setattr(INSTALL_LLAMA_PREBUILT, "download_file", fake_download_file) + + install_dir = tmp_path / "install" + work_dir = tmp_path / "work" + work_dir.mkdir() + hydrate_source_tree( + commit, + install_dir, + work_dir, + source_repo = "unslothai/llama.cpp", + expected_sha256 = sha256_file(archive_path), + exact_source = True, + asset_url = asset_url, + ) + assert seen == [asset_url] + assert (install_dir / "CMakeLists.txt").exists() + assert (install_dir / "convert_hf_to_gguf.py").exists() + + +def test_hydrate_source_tree_falls_back_to_codeload_when_asset_missing( + tmp_path: Path, monkeypatch: pytest.MonkeyPatch +): + # If the release asset 404s, fall back to codeload/archive (vanilla path). + commit = "b" * 40 + archive_path = tmp_path / "vanilla-source.tar.gz" + _mk_source_tarball(archive_path, f"commit-{commit[:7]}") + asset_url = INSTALL_LLAMA_PREBUILT.release_asset_download_url( + "unslothai/llama.cpp", "b9000", f"llama.cpp-source-commit-{commit}.tar.gz" + ) + codeload_urls = INSTALL_LLAMA_PREBUILT.commit_source_archive_urls("unslothai/llama.cpp", commit) + + def fake_download_file(url: str, destination: Path) -> None: + if url == asset_url: + raise RuntimeError("404 Not Found") + assert url in codeload_urls + destination.write_bytes(archive_path.read_bytes()) + + monkeypatch.setattr(INSTALL_LLAMA_PREBUILT, "download_file", fake_download_file) + + install_dir = tmp_path / "install" + work_dir = tmp_path / "work" + work_dir.mkdir() + hydrate_source_tree( + commit, + install_dir, + work_dir, + source_repo = "unslothai/llama.cpp", + expected_sha256 = sha256_file(archive_path), + exact_source = True, + asset_url = asset_url, + ) + assert (install_dir / "CMakeLists.txt").exists() + + def test_validate_prebuilt_choice_creates_repo_shaped_linux_install( tmp_path: Path, monkeypatch: pytest.MonkeyPatch ): From 8628439ede1a71203285db36286e85ad98662f03 Mon Sep 17 00:00:00 2001 From: Dariton4000 <184837391+Dariton4000@users.noreply.github.com> Date: Thu, 11 Jun 2026 14:27:13 +0200 Subject: [PATCH 15/42] fix: running unsloth inside of System32 now refuses to run and quits with an error (#5934) * -Added catch for unsloth running inside of system32 (for windows) - changed import statement to import os as _os instead of only os.path as _osp - replaced __osp mention with __os.path * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * Adressed geminis codereview. Did not change Error message as "System32" is precise enough and the _system32 variable is lowercase only * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci --------- Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com> Co-authored-by: Etherll <61019402+Etherll@users.noreply.github.com> --- unsloth_cli/__init__.py | 20 +++++++++++++++++--- 1 file changed, 17 insertions(+), 3 deletions(-) diff --git a/unsloth_cli/__init__.py b/unsloth_cli/__init__.py index c8ec4c66c0..e1047a6515 100644 --- a/unsloth_cli/__init__.py +++ b/unsloth_cli/__init__.py @@ -1,7 +1,7 @@ # SPDX-License-Identifier: AGPL-3.0-only # Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0 -import os.path as _osp +import os as _os import sys as _sys import typer @@ -20,7 +20,7 @@ from unsloth_cli.commands.studio import ( # Canonicalise `-np` only under the `unsloth` console-script; # third-party scripts that import unsloth_cli keep their argv intact. -_entry_base = _osp.basename(_sys.argv[0]).lower() if _sys.argv else "" +_entry_base = _os.path.basename(_sys.argv[0]).lower() if _sys.argv else "" if _entry_base in {"unsloth", "unsloth.exe"}: _expand_attached_np_short() del _entry_base @@ -53,7 +53,21 @@ def main( help = "Show version and exit.", ), ): - pass + if ( + _sys.platform == "win32" + ): # this block catches unsloth running inside of System32 or any subdirs, this WILL cause errors if not prevented. + _cwd = _os.path.normcase(_os.path.normpath(_os.getcwd())) + _system32 = _os.path.normcase( + _os.path.normpath(_os.path.join(_os.environ.get("WINDIR", r"C:\Windows"), "System32")) + ) + if _cwd == _system32 or _cwd.startswith(_system32 + _os.sep): + typer.secho( + "Refusing to run Unsloth inside System32 as it will lead to Errors.\n" + "cd to a normal working directory and try again.", + fg = "red", + err = True, + ) + raise typer.Exit(code = 1) app.command()(train) From c3604d01f7a009eba021a38aa1ab71b51b46e43d Mon Sep 17 00:00:00 2001 From: Daniel Han Date: Thu, 11 Jun 2026 05:27:46 -0700 Subject: [PATCH 16/42] Studio: enable MTP for sub-3B Gemma separate-drafter GGUFs (#6191) * Studio: enable MTP for sub-3B Gemma separate-drafter GGUFs The sub-3B auto-drop to ngram-mod was tuned for an embedded draft head (Qwen), whose per-token cost regresses below 3B. Gemma ships the head as a separate root mtp-*.gguf drafter, a tiny standalone model that is cheap enough to win below 3B: B200 Q4_K_XL bench, draft-mtp n=2 vs spec-off, gemma-4-E2B (2B) = 1.21x (accept ~0.65) while ngram-mod is 1.00x. Exempt a separate drafter from the sub-3B gate everywhere the threshold is applied: the resolver (_mtp_too_small), the auto-fit VRAM reserve, the drafter auto-download decision, and the reload-skip mirror via a has_separate_drafter flag on _auto_mode_drops_mtp. Embedded sub-3B heads (Qwen) still drop to ngram-mod. A drafter the binary cannot build (older prebuilt, or a CUDA kernel limit) still aborts the spawn and the load retries once without speculative decoding. Adds the full Qwen3.5 + Gemma-4 (regular and QAT) auto/off/forced resolver matrix, plus explicit sub-3B exemption tests. * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * Always compare the separate drafter in the reload-skip mirror The sub-3B wrapper around the drafter compare could skip it when the drafter was deleted out from under a running sub-3B server (detected None, stored set), leaving a stale launch. The resolved-path compare is cheap and already handles every case, so drop the _auto_mode_drops_mtp guard (and its now-unused imports) and always compare when the mode can use a drafter and the user does not own --spec-type. Addresses review feedback on #6191. --------- Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com> --- studio/backend/core/inference/llama_cpp.py | 60 ++++-- studio/backend/routes/inference.py | 26 +-- .../tests/test_llama_cpp_mtp_detection.py | 180 ++++++++++++++++++ 3 files changed, 230 insertions(+), 36 deletions(-) diff --git a/studio/backend/core/inference/llama_cpp.py b/studio/backend/core/inference/llama_cpp.py index 85ca32eaf5..7d67097d09 100644 --- a/studio/backend/core/inference/llama_cpp.py +++ b/studio/backend/core/inference/llama_cpp.py @@ -514,9 +514,18 @@ _CTX_FIT_VRAM_FRACTION = 0.90 _MTP_VRAM_RESERVE_FRAC = 0.05 -def _auto_mode_drops_mtp(req_mode: Optional[str], size_b: Optional[float]) -> bool: - """Auto mode drops MTP below _MTP_MIN_SIZE_B (draft-mtp regresses there); - forced mtp / mtp+ngram engage regardless of size.""" +def _auto_mode_drops_mtp( + req_mode: Optional[str], + size_b: Optional[float], + *, + has_separate_drafter: bool = False, +) -> bool: + """Auto mode drops MTP below _MTP_MIN_SIZE_B for an embedded draft head + (its per-token cost regresses there); a separate drafter (Gemma) is a tiny + standalone model that still speeds up below 3B, so it never drops. Forced + mtp / mtp+ngram engage regardless of size.""" + if has_separate_drafter: + return False return req_mode == "auto" and size_b is not None and size_b < _MTP_MIN_SIZE_B @@ -2910,19 +2919,14 @@ class LlamaCppBackend: # Auto-download the separate MTP drafter (e.g. Gemma) when # the requested spec mode can use it. Repos with the head # baked into the main GGUF (Qwen) have no mtp- sibling and - # this no-ops. Skipped when the user disabled MTP, drives - # --spec-type manually via extra_args, or in auto mode on a - # sub-3B model (e.g. Gemma E2B) where the resolver drops - # MTP anyway -- no point fetching a drafter it never uses. - # Forced mtp / mtp+ngram still download (user override). + # this no-ops, so the size gate stays out of it: a separate + # drafter speeds up even sub-3B (Gemma E2B), and the resolver + # below decides the final emission. Skipped only when the + # user disabled MTP or drives --spec-type manually. _spec_canon = _canonicalize_spec_mode(speculative_type) or "auto" - _auto_drops_mtp = _auto_mode_drops_mtp( - _spec_canon, _extract_model_size_b(model_identifier) - ) if ( not mtp_draft_path and _spec_canon in ("auto", "mtp", "mtp+ngram") - and not _auto_drops_mtp and not _extra_args_set_spec_type(extra_args) ): mtp_draft_path = self._download_mtp( @@ -3002,8 +3006,12 @@ class LlamaCppBackend: _mtp_canonical = _canonicalize_spec_mode(speculative_type) _mtp_effective = _mtp_canonical or "auto" _mtp_size_for_fit = _extract_model_size_b(model_identifier) + # Sub-3B drops MTP only for an embedded head; a separate + # drafter (Gemma) engages and needs its VRAM reserved. _mtp_sub_3b_for_fit = ( - _mtp_size_for_fit is not None and _mtp_size_for_fit < _MTP_MIN_SIZE_B + _mtp_size_for_fit is not None + and _mtp_size_for_fit < _MTP_MIN_SIZE_B + and not bool(mtp_draft_path) ) _mtp_will_engage = bool( not _extra_args_set_spec_type(extra_args) @@ -3771,18 +3779,26 @@ class LlamaCppBackend: https://github.com/ggml-org/llama.cpp/pull/18471 MTP guide: unsloth.ai/docs/models/qwen3.6#mtp-guide - Sub-3B dense MTP regresses vs spec-off: the draft head's per-token - cost exceeds the acceptance savings at this scale. Q4_K_XL clean - bench (each prompt once after an unrelated warmup) on B200 + x86 CPU: + Sub-3B dense MTP regresses vs spec-off when the head is baked into the + main GGUF (Qwen): the draft head's per-token cost exceeds the + acceptance savings at this scale. Q4_K_XL clean bench (each prompt once + after an unrelated warmup) on B200 + x86 CPU: 0.8B GPU: draft-mtp n=2 = 0.58x vs OFF; ngram-only = 1.10x 2B GPU: draft-mtp n=2 = 0.82x vs OFF; OFF or ngram = 1.00x 0.8B CPU: chained n=2 = 0.86x vs OFF; ngram-only = 1.19x 2B CPU: chained n=2 = 0.83x vs OFF; ngram-only = 1.01x 4B+ GPU/CPU: spec on is a net win (1.08x-1.46x). + A separate drafter (Gemma's root mtp-*.gguf) is a different, cheaper + mechanism that wins even below 3B, so it is exempt from the sub-3B drop + (``mtp_draft_path`` set -> not too small). B200 Q4_K_XL bench, draft-mtp + n=2 vs OFF: gemma-4-E2B (2B) = 1.21x, accept ~0.65 (vs ngram = 1.00x); + gemma-4-E4B (4B) and 12B engage as usual. Auto falls back to ngram-mod (zero-VRAM, near-zero idle cost on - diverse content); forced MTP on a model with no head/drafter defaults - back (mtp -> spec-default, mtp+ngram -> ngram-mod) since llama-server - aborts otherwise; sub-3B real-MTP engages with a warning. + diverse content) for an embedded sub-3B head; forced MTP on a model + with no head/drafter defaults back (mtp -> spec-default, mtp+ngram -> + ngram-mod) since llama-server aborts otherwise; a drafter the binary + cannot build (older prebuilt, or a CUDA kernel limit) aborts the spawn + and the load retries once without speculative decoding. """ flags: List[str] = [] # Reset; emit branches re-set on the resolved emission. @@ -3801,7 +3817,11 @@ class LlamaCppBackend: ) user_owns_spec_type = _extra_args_set_spec_type(extra_args) _mtp_size_b = _extract_model_size_b(model_identifier) - _mtp_too_small = _mtp_size_b is not None and _mtp_size_b < _MTP_MIN_SIZE_B + # The sub-3B regression is an embedded-head cost; a separate drafter + # (Gemma) is a cheap standalone model that wins below 3B, so exempt it. + _mtp_too_small = ( + _mtp_size_b is not None and _mtp_size_b < _MTP_MIN_SIZE_B and not bool(mtp_draft_path) + ) if user_owns_spec_type: # User --spec-type wins outright; suppress auto-emit to avoid a diff --git a/studio/backend/routes/inference.py b/studio/backend/routes/inference.py index 4cc76e7c2a..ef67f8bb59 100644 --- a/studio/backend/routes/inference.py +++ b/studio/backend/routes/inference.py @@ -417,7 +417,6 @@ try: LlamaCppBackend, _DEFAULT_MAX_TOKENS_FLOOR, _DEFAULT_T_MAX_PREDICT_MS, - _auto_mode_drops_mtp, _canonicalize_spec_mode, _extra_args_set_spec_type, _hf_offline_if_dns_dead, @@ -431,7 +430,6 @@ try: from utils.inference import load_inference_config from utils.models.model_config import ( detect_mtp_file, - extract_model_size_b, load_model_defaults, ) from utils.native_path_leases import ( @@ -450,7 +448,6 @@ except ImportError: LlamaCppBackend, _DEFAULT_MAX_TOKENS_FLOOR, _DEFAULT_T_MAX_PREDICT_MS, - _auto_mode_drops_mtp, _canonicalize_spec_mode, _extra_args_set_spec_type, _hf_offline_if_dns_dead, @@ -464,7 +461,6 @@ except ImportError: from utils.inference import load_inference_config from utils.models.model_config import ( detect_mtp_file, - extract_model_size_b, load_model_defaults, ) from utils.native_path_leases import ( @@ -1025,14 +1021,15 @@ def _request_matches_loaded_settings(request: LoadRequest, llama_backend: LlamaC else: if list(request.llama_extra_args) != backend_extra: return False - # A drafter that appeared next to the loaded weights since the last load - # changes the launch command (--model-draft) when the mode can use it; - # without this, a duplicate /load is deduped and MTP can't engage short - # of an unload. Runs last: it stats the filesystem (two dir scans against - # the resolved weight path -- covers local dirs and HF cache snapshots - # alike), so every pure-memory comparison above short-circuits first. - # Skipped when auto drops MTP anyway (sub-3B) or the user owns - # --spec-type, where a drafter changes nothing. Resolve both sides: the + # A separate drafter (Gemma's root mtp-*.gguf) appearing or disappearing + # next to the loaded weights changes the launch command (--model-draft), + # so a duplicate /load must reload rather than dedupe. Always compare the + # detected vs stored drafter when the mode can use one and the user does + # not own --spec-type: the resolved-path compare is cheap and handles all + # four cases (both None -> match; one None -> reload; equal -> match; + # different -> reload), including a drafter deleted out from under a + # running server. Runs last: it stats the filesystem, so every pure-memory + # comparison above short-circuits first. Resolve both sides since the # stored launch path may be a snapshot symlink while detect_mtp_file # returns the resolved blob. if req_mode in ("auto", "mtp", "mtp+ngram") and llama_backend.gguf_path: @@ -1041,10 +1038,7 @@ def _request_matches_loaded_settings(request: LoadRequest, llama_backend: LlamaC if request.llama_extra_args is not None else llama_backend.extra_args ) - size_b = extract_model_size_b(llama_backend.model_identifier or "") - if not _auto_mode_drops_mtp(req_mode, size_b) and not _extra_args_set_spec_type( - effective_extras - ): + if not _extra_args_set_spec_type(effective_extras): detected = detect_mtp_file(llama_backend.gguf_path) stored = llama_backend.mtp_draft_path try: diff --git a/studio/backend/tests/test_llama_cpp_mtp_detection.py b/studio/backend/tests/test_llama_cpp_mtp_detection.py index 25ce71a774..a7f30d2951 100644 --- a/studio/backend/tests/test_llama_cpp_mtp_detection.py +++ b/studio/backend/tests/test_llama_cpp_mtp_detection.py @@ -1219,3 +1219,183 @@ def test_forced_mtp_ngram_on_non_mtp_model_keeps_ngram(monkeypatch): assert parsed.get("--spec-type") == "ngram-mod" assert backend.speculative_type == "ngram-mod" assert backend.requested_spec_mode == "mtp+ngram" + + +# ── Full named-repo resolver matrix (the shipping Studio families) ───── +# +# Locks auto / off / forced-mtp routing for every Qwen3.5 (MTP + plain) and +# gemma-4 (regular + QAT) GGUF repo, including the giant MoEs that stay +# resolver-only (122B-A10B / 397B-A17B). Expectations are derived from the +# same signals load_model uses -- _extract_model_size_b (active>effective> +# total, so E2B->2, A3B->3, A10B->10, A17B->17), _is_mtp_model_name, and the +# separate-drafter flag -- so each row mirrors what the loader emits on a +# B200 (GPU default, n=2). gemma carries no -MTP marker; its MTP comes from +# the root mtp-*.gguf drafter, modelled here by passing mtp_draft_path. +# +# auto_spec: "draft-mtp" = head/drafter engaged (>=3B MTP, or any size with a +# separate drafter); "ngram-mod" = embedded sub-3B drop (zero-VRAM); None = +# non-MTP -> llama-server --spec-default. + +_GEMMA_DRAFTER = "/snap/mtp-gemma-4-it.gguf" # stand-in separate drafter + +_REAL_REPO_MATRIX = [ + # repo, drafter, auto_spec, auto_ngram_knobs + ("unsloth/Qwen3.5-0.8B-MTP-GGUF", None, "ngram-mod", True), + ("unsloth/Qwen3.5-2B-MTP-GGUF", None, "ngram-mod", True), + ("unsloth/Qwen3.5-4B-MTP-GGUF", None, "draft-mtp", False), + ("unsloth/Qwen3.5-9B-MTP-GGUF", None, "draft-mtp", False), + ("unsloth/Qwen3.5-27B-MTP-GGUF", None, "draft-mtp", False), + ("unsloth/Qwen3.5-35B-A3B-MTP-GGUF", None, "draft-mtp", False), + ("unsloth/Qwen3.5-122B-A10B-MTP-GGUF", None, "draft-mtp", False), + ("unsloth/Qwen3.5-397B-A17B-MTP-GGUF", None, "draft-mtp", False), + ("unsloth/Qwen3.5-0.8B-GGUF", None, None, False), + ("unsloth/Qwen3.5-2B-GGUF", None, None, False), + ("unsloth/Qwen3.5-4B-GGUF", None, None, False), + ("unsloth/Qwen3.5-9B-GGUF", None, None, False), + # E2B is 2B but ships a separate drafter -> exempt from the sub-3B drop. + ("unsloth/gemma-4-E2B-it-GGUF", _GEMMA_DRAFTER, "draft-mtp", False), + ("unsloth/gemma-4-E4B-it-GGUF", _GEMMA_DRAFTER, "draft-mtp", False), + ("unsloth/gemma-4-12b-it-GGUF", _GEMMA_DRAFTER, "draft-mtp", False), + ("unsloth/gemma-4-26B-A4B-it-GGUF", _GEMMA_DRAFTER, "draft-mtp", False), + ("unsloth/gemma-4-31B-it-GGUF", _GEMMA_DRAFTER, "draft-mtp", False), + ("unsloth/gemma-4-E2B-it-qat-GGUF", _GEMMA_DRAFTER, "draft-mtp", False), + ("unsloth/gemma-4-E4B-it-qat-GGUF", _GEMMA_DRAFTER, "draft-mtp", False), + ("unsloth/gemma-4-12b-it-qat-GGUF", _GEMMA_DRAFTER, "draft-mtp", False), + ("unsloth/gemma-4-26B-A4B-it-qat-GGUF", _GEMMA_DRAFTER, "draft-mtp", False), + ("unsloth/gemma-4-31B-it-qat-GGUF", _GEMMA_DRAFTER, "draft-mtp", False), +] + + +def _resolve_real(monkeypatch, repo, drafter, mode): + backend = _resolver_backend(monkeypatch) + flags = backend._build_speculative_flags( + speculative_type = mode, + spec_draft_n_max = None, + extra_args = None, + model_identifier = repo, + model_path = None, + gpus = True, # B200 default + binary = "/fake/llama-server", + mtp_draft_path = drafter, + ) + return backend, flags, _flags_dict(flags) + + +@pytest.mark.parametrize( + "repo, drafter, auto_spec, auto_ngram_knobs", + _REAL_REPO_MATRIX, + ids = [r[0].split("/")[-1] for r in _REAL_REPO_MATRIX], +) +def test_real_repo_auto_routing(monkeypatch, repo, drafter, auto_spec, auto_ngram_knobs): + # Auto is the default mode the dropdown ships with. + backend, flags, parsed = _resolve_real(monkeypatch, repo, drafter, "auto") + if auto_spec is None: + # Non-MTP: no draft-mtp, hand off to llama-server's own default. + assert "--spec-type" not in parsed + assert "--spec-default" in flags + assert backend.speculative_type == "default" + elif auto_spec == "draft-mtp": + assert parsed.get("--spec-type") == "draft-mtp" + assert parsed.get("--spec-draft-n-max") == "2" + assert backend.speculative_type == "draft-mtp" + # gemma ships a separate drafter; Qwen bakes the head into the GGUF. + assert ( + (parsed.get("--model-draft") == drafter) if drafter else ("--model-draft" not in parsed) + ) + else: # ngram-mod (sub-3B MTP drop) + assert parsed.get("--spec-type") == "ngram-mod" + assert "--model-draft" not in parsed # draft head dropped + assert backend.speculative_type == "ngram-mod" + if auto_ngram_knobs: + assert "--spec-ngram-mod-n-match" in parsed + assert backend.requested_spec_mode == "auto" + + +@pytest.mark.parametrize( + "repo, drafter", + [(r[0], r[1]) for r in _REAL_REPO_MATRIX], + ids = [r[0].split("/")[-1] for r in _REAL_REPO_MATRIX], +) +def test_real_repo_off_emits_nothing(monkeypatch, repo, drafter): + # Off must suppress speculative decoding for every family. + backend, flags, _ = _resolve_real(monkeypatch, repo, drafter, "off") + assert flags == [] + assert backend.speculative_type is None + assert backend.requested_spec_mode == "off" + + +@pytest.mark.parametrize( + "repo, drafter", + [(r[0], r[1]) for r in _REAL_REPO_MATRIX], + ids = [r[0].split("/")[-1] for r in _REAL_REPO_MATRIX], +) +def test_real_repo_forced_mtp_never_aborts(monkeypatch, repo, drafter): + # Forcing MTP on the dropdown: real MTP models (name marker or separate + # drafter) engage draft-mtp even below 3B; non-MTP models default back to + # --spec-default instead of emitting a draft-mtp llama-server will abort on. + backend, flags, parsed = _resolve_real(monkeypatch, repo, drafter, "mtp") + is_real_mtp = _is_mtp_model_name(repo) or bool(drafter) + if is_real_mtp: + assert parsed.get("--spec-type") == "draft-mtp" + assert backend.speculative_type == "draft-mtp" + assert ( + (parsed.get("--model-draft") == drafter) if drafter else ("--model-draft" not in parsed) + ) + else: + assert "--spec-type" not in parsed + assert "--spec-default" in flags + assert backend.speculative_type == "default" + assert backend.requested_spec_mode == "mtp" + + +# ── Sub-3B separate-drafter exemption (Gemma) ───────────────────────── +# +# The sub-3B MTP drop is an embedded-head cost (Qwen). A separate drafter +# (Gemma's root mtp-*.gguf) is a cheap standalone model that wins below 3B +# (B200 Q4_K_XL: gemma-4-E2B draft-mtp n=2 = 1.21x vs OFF), so it is exempt. + + +def test_sub3b_gemma_separate_drafter_engages_mtp(monkeypatch): + backend = _resolver_backend(monkeypatch) + flags = backend._build_speculative_flags( + speculative_type = "auto", + spec_draft_n_max = None, + extra_args = None, + model_identifier = "unsloth/gemma-4-E2B-it-GGUF", # 2B + model_path = None, + gpus = True, + binary = "/fake/llama-server", + mtp_draft_path = "/snap/mtp-gemma-4-E2B-it.gguf", # separate drafter + ) + parsed = _flags_dict(flags) + assert parsed.get("--spec-type") == "draft-mtp" + assert parsed.get("--model-draft") == "/snap/mtp-gemma-4-E2B-it.gguf" + assert "--spec-ngram-mod-n-match" not in parsed + assert backend.speculative_type == "draft-mtp" + + +def test_sub3b_qwen_embedded_head_still_drops_to_ngram(monkeypatch): + backend = _resolver_backend(monkeypatch) + flags = backend._build_speculative_flags( + speculative_type = "auto", + spec_draft_n_max = None, + extra_args = None, + model_identifier = "unsloth/Qwen3.5-2B-MTP-GGUF", # 2B, embedded head + model_path = None, + gpus = True, + binary = "/fake/llama-server", + mtp_draft_path = None, # no separate drafter + ) + parsed = _flags_dict(flags) + assert parsed.get("--spec-type") == "ngram-mod" + assert "--model-draft" not in parsed + assert backend.speculative_type == "ngram-mod" + + +def test_auto_mode_drops_mtp_exempts_separate_drafter(): + from core.inference.llama_cpp import _auto_mode_drops_mtp + + assert _auto_mode_drops_mtp("auto", 2.0) is True + assert _auto_mode_drops_mtp("auto", 2.0, has_separate_drafter = True) is False + assert _auto_mode_drops_mtp("auto", 4.0) is False + assert _auto_mode_drops_mtp("mtp", 2.0) is False # forced engages regardless From a22169941d0a7789b0ef16b314767c5bd2d463db Mon Sep 17 00:00:00 2001 From: Tai An Date: Thu, 11 Jun 2026 05:30:08 -0700 Subject: [PATCH 17/42] fix(studio): fall back to copy when os.replace is blocked during install activation (#6133) * fix(studio): fall back to copy when os.replace is blocked during install activation On Windows ARM64 the antivirus scanner can transiently hold a freshly extracted DLL open while MoveFileEx runs, so activating the staged llama.cpp prebuilt fails with [WinError 5] Access is denied. Attempt os.replace first, then fall back to a file-by-file copytree which bypasses the rename. * address review: keep os.replace for rollback, scope copy-fallback to staging The copy + rmtree fallback could silently corrupt a live install if the existing directory is busy. Restrict it to freshly extracted staging trees (renamed activate_staged_dir) and keep strict os.replace for the rollback move so a busy active install raises immediately. * fix(studio): scope copy-fallback to busy-lock errors, log it, and add tests --------- Co-authored-by: Lee Jackson <130007945+Imagineer99@users.noreply.github.com> Co-authored-by: Etherll <61019402+Etherll@users.noreply.github.com> --- studio/install_llama_prebuilt.py | 31 ++++++++++++- .../test_install_llama_prebuilt_logic.py | 44 +++++++++++++++++++ 2 files changed, 74 insertions(+), 1 deletion(-) diff --git a/studio/install_llama_prebuilt.py b/studio/install_llama_prebuilt.py index 6c3be0e09f..64ea255803 100644 --- a/studio/install_llama_prebuilt.py +++ b/studio/install_llama_prebuilt.py @@ -4870,6 +4870,35 @@ def confirm_install_tree(install_dir: Path, host: HostInfo) -> None: raise RuntimeError("activated install was missing expected files: " + ", ".join(missing)) +def activate_staged_dir(staging_dir: Path, dst: Path) -> None: + """Move a freshly extracted ``staging_dir`` onto ``dst``. + + ``os.replace`` is attempted first as the fast path. On Windows ARM64 the + antivirus scanner can transiently hold a freshly extracted DLL open at the + moment ``MoveFileEx`` runs, surfacing as ``[WinError 5] Access is denied``; + a file-by-file copy bypasses the rename entirely. + + This fallback is intentionally limited to staging trees we just extracted. + It must not be used to move an existing/active install aside: there an + ``os.replace`` failure means the directory is genuinely in use, and a + silent copy + ``rmtree`` could partially delete a live install. + + Only busy/lock errors (``is_busy_lock_error``) trigger the copy; anything + else (disk full, cross-device, missing path) re-raises so it cannot leave + a partially copied install behind. A copy is preferred over retrying the + rename because antivirus scans of large DLLs can outlast any reasonable + retry window. + """ + try: + os.replace(staging_dir, dst) + except OSError as exc: + if not is_busy_lock_error(exc): + raise + log(f"os.replace failed ({exc!r}); falling back to file-by-file copy of staging tree") + shutil.copytree(staging_dir, dst, dirs_exist_ok = True) + remove_tree(staging_dir) + + def activate_install_tree(staging_dir: Path, install_dir: Path, host: HostInfo) -> None: rollback_dir: Path | None = None failed_dir: Path | None = None @@ -4881,7 +4910,7 @@ def activate_install_tree(staging_dir: Path, install_dir: Path, host: HostInfo) log(f"moved existing install to rollback path {rollback_dir.name}") log(f"activating staged install {staging_dir} -> {install_dir}") - os.replace(staging_dir, install_dir) + activate_staged_dir(staging_dir, install_dir) log(f"activated staged install at {install_dir}") log(f"confirming activated install tree at {install_dir}") confirm_install_tree(install_dir, host) diff --git a/tests/studio/install/test_install_llama_prebuilt_logic.py b/tests/studio/install/test_install_llama_prebuilt_logic.py index 7ca87aa8c4..000c5a7ace 100644 --- a/tests/studio/install/test_install_llama_prebuilt_logic.py +++ b/tests/studio/install/test_install_llama_prebuilt_logic.py @@ -1,3 +1,4 @@ +import errno import importlib.util import io import json @@ -28,6 +29,7 @@ ApprovedReleaseChecksums = INSTALL_LLAMA_PREBUILT.ApprovedReleaseChecksums hydrate_source_tree = INSTALL_LLAMA_PREBUILT.hydrate_source_tree validate_prebuilt_choice = INSTALL_LLAMA_PREBUILT.validate_prebuilt_choice activate_install_tree = INSTALL_LLAMA_PREBUILT.activate_install_tree +activate_staged_dir = INSTALL_LLAMA_PREBUILT.activate_staged_dir create_install_staging_dir = INSTALL_LLAMA_PREBUILT.create_install_staging_dir sha256_file = INSTALL_LLAMA_PREBUILT.sha256_file source_archive_logical_name = INSTALL_LLAMA_PREBUILT.source_archive_logical_name @@ -666,6 +668,48 @@ def test_activate_install_tree_cleans_all_paths_when_rollback_restore_fails( assert "removing rollback path" in output +def test_activate_staged_dir_copies_when_replace_hits_busy_lock( + tmp_path: Path, monkeypatch: pytest.MonkeyPatch, capsys: pytest.CaptureFixture[str] +): + staging_dir = tmp_path / "llama.cpp.staging-test" + (staging_dir / "bin").mkdir(parents = True) + (staging_dir / "bin" / "ggml-base.dll").write_bytes(b"fake dll") + dst = tmp_path / "llama.cpp" + + def denied_replace(src, dst_arg): + raise PermissionError(errno.EACCES, "Access is denied", str(src)) + + monkeypatch.setattr(INSTALL_LLAMA_PREBUILT.os, "replace", denied_replace) + + activate_staged_dir(staging_dir, dst) + + assert (dst / "bin" / "ggml-base.dll").read_bytes() == b"fake dll" + assert not staging_dir.exists() + + captured = capsys.readouterr() + assert "falling back to file-by-file copy" in captured.out + captured.err + + +def test_activate_staged_dir_reraises_non_busy_errors( + tmp_path: Path, monkeypatch: pytest.MonkeyPatch +): + staging_dir = tmp_path / "llama.cpp.staging-test" + staging_dir.mkdir() + (staging_dir / "new.txt").write_text("new install\n") + dst = tmp_path / "llama.cpp" + + def out_of_space_replace(src, dst_arg): + raise OSError(errno.ENOSPC, "No space left on device", str(src)) + + monkeypatch.setattr(INSTALL_LLAMA_PREBUILT.os, "replace", out_of_space_replace) + + with pytest.raises(OSError, match = "No space left on device"): + activate_staged_dir(staging_dir, dst) + + assert not dst.exists() + assert (staging_dir / "new.txt").read_text() == "new install\n" + + def test_binary_env_linux_includes_binary_parent_in_ld_library_path( tmp_path: Path, monkeypatch: pytest.MonkeyPatch ): From ab9689c034ace6fef573ce017bbb30908826b114 Mon Sep 17 00:00:00 2001 From: Prathamesh Jadhav <55660103+lollinng@users.noreply.github.com> Date: Thu, 11 Jun 2026 18:00:54 +0530 Subject: [PATCH 18/42] Handle canonical inputs_embeds kwarg in unsloth_base_fast_generate (#3082) (#6015) The kwarg-dispatch in unsloth_base_fast_generate recognized input_ids, input, input_features and the misspelled input_embeds, but not HF's canonical inputs_embeds. So generate(inputs_embeds=...) fell through to the 'first kwarg' fallback, which picks whatever kwarg happens to come first (e.g. attention_mask) and uses it as input_ids -- giving the wrong tensor / batch size, or the KeyError reported in #3082 on older versions. Add an explicit inputs_embeds branch so embedding inputs (e.g. multimodal audio+text) are routed correctly regardless of kwarg order. Co-authored-by: Claude Opus 4.8 (1M context) Co-authored-by: Etherll <61019402+Etherll@users.noreply.github.com> --- unsloth/models/vision.py | 3 +++ 1 file changed, 3 insertions(+) diff --git a/unsloth/models/vision.py b/unsloth/models/vision.py index 58814699b3..a52ab359bd 100644 --- a/unsloth/models/vision.py +++ b/unsloth/models/vision.py @@ -283,6 +283,9 @@ def unsloth_base_fast_generate(self, *args, **kwargs): input_ids = kwargs["input"] elif "input_features" in kwargs: input_ids = kwargs["input_features"] + elif "inputs_embeds" in kwargs: + # canonical HF name for embedding inputs (e.g. multimodal generate) + input_ids = kwargs["inputs_embeds"] elif "input_embeds" in kwargs: input_ids = kwargs["input_embeds"] elif "inputs" in kwargs: From c84ba48dd64cb04ce16d171cdeadc6d13a77e6c4 Mon Sep 17 00:00:00 2001 From: Abhinav Date: Thu, 11 Jun 2026 18:08:35 +0530 Subject: [PATCH 19/42] fix: don't block the event loop with time.sleep in async load_checkpoint (#6135) The export route's load_checkpoint waits for the training subprocess to exit by calling time.sleep(0.5) in a loop (up to 30s) inside an async function. time.sleep blocks the whole event loop, so every other request to the server stalls for that duration. Use await asyncio.sleep(0.5), matching the async pattern already used elsewhere in this file (asyncio.to_thread, await asyncio.sleep). Co-authored-by: Wasim Yousef Said --- studio/backend/routes/export.py | 3 +-- 1 file changed, 1 insertion(+), 2 deletions(-) diff --git a/studio/backend/routes/export.py b/studio/backend/routes/export.py index abb89a13da..ef37daa158 100644 --- a/studio/backend/routes/export.py +++ b/studio/backend/routes/export.py @@ -78,8 +78,7 @@ async def load_checkpoint( for _ in range(60): # up to 30s if not trn.is_training_active(): break - import time - time.sleep(0.5) + await asyncio.sleep(0.5) else: logger.warning("Training subprocess did not exit within 30s, proceeding anyway") except Exception as e: From 22e2b63d2ebd30411bb9b92710952c0b1cd90ff8 Mon Sep 17 00:00:00 2001 From: ashzak Date: Thu, 11 Jun 2026 07:39:15 -0500 Subject: [PATCH 20/42] fix(studio): prevent UI freeze when switching tabs from heavy pages (#5978) * fix(studio): prevent UI freeze when switching tabs from heavy pages Change AnimatePresence mode from "wait" to "popLayout" to fix issue where switching tabs from Export (or other heavy pages) would cause the URL to update but the UI to freeze. With mode="wait", the exit animation must complete before the new component mounts. If the exiting page has expensive computations, this blocks the UI. mode="popLayout" allows the new route to mount immediately while the old one animates out. Fixes #5850 * fix: add relative positioning for popLayout mode AnimatePresence mode='popLayout' applies position: absolute to the exiting element, so the parent container needs position: relative to prevent layout jumps during transitions. --------- Co-authored-by: Lee Jackson <130007945+Imagineer99@users.noreply.github.com> --- studio/frontend/src/app/routes/__root.tsx | 9 +++++++-- 1 file changed, 7 insertions(+), 2 deletions(-) diff --git a/studio/frontend/src/app/routes/__root.tsx b/studio/frontend/src/app/routes/__root.tsx index c8e47902b9..c05b0e6451 100644 --- a/studio/frontend/src/app/routes/__root.tsx +++ b/studio/frontend/src/app/routes/__root.tsx @@ -140,9 +140,14 @@ function RootLayout() {
- + {/* Use mode="popLayout" instead of "wait" to prevent UI freezes when + switching from heavy pages (like Export with many checkpoints). + "popLayout" allows the new route to mount immediately while the + old one animates out, avoiding blocking on expensive exit renders. + See issue #5850. */} + Date: Thu, 11 Jun 2026 06:10:17 -0700 Subject: [PATCH 21/42] Studio: surface the llama.cpp update affordance when MTP is disabled (#6192) * Studio: surface the llama.cpp update affordance when MTP is disabled When a model asks for MTP (auto on an MTP model, or forced mtp / mtp+ngram) but it gets disabled, the load already degrades gracefully and serves without speculative decoding. Until now the UI gave no hint why, or that an update would fix it. Record why MTP was dropped on the backend (spec_fallback_reason): the probe found no mtp token (binary_no_mtp), the spawn aborted with an outdated-arch / context-build error such as a prebuilt that predates the Gemma drafter (binary_outdated), or the current build could not run it, e.g. a CUDA kernel limit (runtime_error). Expose it in the inference status. In the chat Speculative Decoding section, show a short note and, for the two update-fixable reasons, an inline Update llama.cpp button that reuses the existing update flow. A runtime_error gets the note without an update push, since a newer build may not fix it. Backend tests cover the reason being set / cleared. Frontend typechecks. * Address review: tighten the update hint to genuinely outdated binaries Reserve binary_outdated (which surfaces the Update llama.cpp affordance) for an unknown-architecture abort, which proves the prebuilt predates the model; classify the generic memory/context build failures as runtime_error, where an update may not help. Frontend: only append the "Update llama.cpp to enable it" sentence when an update is actually available, so the text never points at an action the UI is not offering. --- studio/backend/core/inference/llama_cpp.py | 24 +++++++- studio/backend/models/inference.py | 11 ++++ studio/backend/routes/inference.py | 1 + .../tests/test_llama_cpp_mtp_detection.py | 59 +++++++++++++++++++ .../src/features/chat/chat-settings-sheet.tsx | 50 +++++++++++++++- .../chat/hooks/use-chat-model-runtime.ts | 1 + .../chat/stores/chat-runtime-store.ts | 7 +++ .../frontend/src/features/chat/types/api.ts | 6 ++ 8 files changed, 157 insertions(+), 2 deletions(-) diff --git a/studio/backend/core/inference/llama_cpp.py b/studio/backend/core/inference/llama_cpp.py index 7d67097d09..e46d464784 100644 --- a/studio/backend/core/inference/llama_cpp.py +++ b/studio/backend/core/inference/llama_cpp.py @@ -676,6 +676,11 @@ class LlamaCppBackend: # Separate MTP drafter launched with the current model; reload-dedup # key so a drafter that appears next to the weights forces a reload. self._mtp_draft_path: Optional[str] = None + # Why MTP was disabled on the last load that asked for it (auto on an + # MTP model, or forced mtp / mtp+ngram), else None. Drives the "update + # llama.cpp" hint in the UI. "binary_no_mtp" / "binary_outdated" -> + # a newer prebuilt would help; "runtime_error" -> it may not. + self._spec_fallback_reason: Optional[str] = None self._hf_variant: Optional[str] = None self._is_vision: bool = False self._healthy = False @@ -794,6 +799,11 @@ class LlamaCppBackend: def mtp_draft_path(self) -> Optional[str]: return self._mtp_draft_path + @property + def spec_fallback_reason(self) -> Optional[str]: + """Why MTP was disabled on the last MTP-requesting load, else None.""" + return self._spec_fallback_reason + @property def extra_args(self) -> Optional[List[str]]: """Extra llama-server flags from the last load (a copy). None = @@ -3594,8 +3604,13 @@ class LlamaCppBackend: # failing (unknown arch / draft or context build); an # unrelated crash (e.g. OOM) gets a neutral message. _lo = "\n".join(self._stdout_lines).lower() + # Only an unknown architecture proves the prebuilt predates + # this MTP model (an update fixes it). The memory/context + # build failures are generic (VRAM / ctx pressure), where an + # update may not help, so classify those as runtime_error. + _arch_unsupported = "unknown model architecture" in _lo if ( - "unknown model architecture" in _lo + _arch_unsupported or "failed to measure draft model memory" in _lo or "failed to measure mtp context memory" in _lo or "failed to create llama_context" in _lo @@ -3605,10 +3620,14 @@ class LlamaCppBackend: "speculative decoding -- run `unsloth studio " "update` for MTP" ) + self._spec_fallback_reason = ( + "binary_outdated" if _arch_unsupported else "runtime_error" + ) else: _retry_reason = ( "retrying without speculative decoding in case MTP is the cause" ) + self._spec_fallback_reason = "runtime_error" _drafter = ( Path(launch_mtp_draft_path).name if launch_mtp_draft_path @@ -3804,6 +3823,7 @@ class LlamaCppBackend: # Reset; emit branches re-set on the resolved emission. self._spec_draft_n_max = None self._speculative_type = None + self._spec_fallback_reason = None # Canonical UI-facing requested mode (legacy values mapped via # _canonicalize_spec_mode). @@ -3853,6 +3873,7 @@ class LlamaCppBackend: "run `unsloth studio update`. Loading without " "speculative decoding." ) + 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" @@ -4118,6 +4139,7 @@ class LlamaCppBackend: self._gguf_path = None self._hf_repo = None self._mtp_draft_path = None + self._spec_fallback_reason = None self._hf_variant = None self._is_vision = False self._is_audio = False diff --git a/studio/backend/models/inference.py b/studio/backend/models/inference.py index c83c8ecc1b..f58a76d16e 100644 --- a/studio/backend/models/inference.py +++ b/studio/backend/models/inference.py @@ -346,6 +346,17 @@ class InferenceStatusResponse(BaseModel): "False -> recommend `unsloth studio update`." ), ) + spec_fallback_reason: Optional[str] = Field( + None, + description = ( + "Why MTP was disabled on the loaded model despite being requested " + "(auto on an MTP model, or forced mtp / mtp+ngram). " + "'binary_no_mtp' / 'binary_outdated' -> a newer prebuilt would " + "re-enable it (show the update affordance); 'runtime_error' -> the " + "current build could not run it. None when MTP engaged or was not " + "requested." + ), + ) llama_cpp_prebuilt_stale: bool = Field( False, description = ( diff --git a/studio/backend/routes/inference.py b/studio/backend/routes/inference.py index ef67f8bb59..0244956b72 100644 --- a/studio/backend/routes/inference.py +++ b/studio/backend/routes/inference.py @@ -1906,6 +1906,7 @@ async def get_status(current_subject: str = Depends(get_current_subject)): speculative_type = llama_backend.requested_spec_mode, spec_draft_n_max = llama_backend.spec_draft_n_max, llama_cpp_supports_mtp = _supports_mtp, + spec_fallback_reason = llama_backend.spec_fallback_reason, llama_cpp_prebuilt_stale = _stale, llama_cpp_installed_tag = _installed_tag, llama_cpp_latest_tag = _latest_tag, diff --git a/studio/backend/tests/test_llama_cpp_mtp_detection.py b/studio/backend/tests/test_llama_cpp_mtp_detection.py index a7f30d2951..b00cd7169e 100644 --- a/studio/backend/tests/test_llama_cpp_mtp_detection.py +++ b/studio/backend/tests/test_llama_cpp_mtp_detection.py @@ -1399,3 +1399,62 @@ def test_auto_mode_drops_mtp_exempts_separate_drafter(): assert _auto_mode_drops_mtp("auto", 2.0, has_separate_drafter = True) is False assert _auto_mode_drops_mtp("auto", 4.0) is False assert _auto_mode_drops_mtp("mtp", 2.0) is False # forced engages regardless + + +# ── spec_fallback_reason (drives the "update llama.cpp" UI hint) ─────── + + +def test_spec_fallback_reason_set_when_binary_lacks_mtp(monkeypatch): + # Outdated llama-server with no mtp token: a forced MTP request can't emit + # draft-mtp, so record the reason for the UI update affordance. + backend = _resolver_backend(monkeypatch, mtp_token = None) + backend._build_speculative_flags( + speculative_type = "mtp", + spec_draft_n_max = None, + extra_args = None, + model_identifier = _MTP_MODEL, + model_path = None, + gpus = True, + binary = "/fake/llama-server", + ) + assert backend.spec_fallback_reason == "binary_no_mtp" + + +def test_spec_fallback_reason_none_when_mtp_engages(monkeypatch): + backend = _resolver_backend(monkeypatch) + backend._build_speculative_flags( + speculative_type = "auto", + spec_draft_n_max = None, + extra_args = None, + model_identifier = _MTP_MODEL, + model_path = None, + gpus = True, + binary = "/fake/llama-server", + ) + assert backend.speculative_type == "draft-mtp" + assert backend.spec_fallback_reason is None + + +def test_spec_fallback_reason_reset_on_off(monkeypatch): + # A subsequent off load must clear a stale reason. + backend = _resolver_backend(monkeypatch, mtp_token = None) + backend._build_speculative_flags( + speculative_type = "mtp", + spec_draft_n_max = None, + extra_args = None, + model_identifier = _MTP_MODEL, + model_path = None, + gpus = True, + binary = "/fake/llama-server", + ) + assert backend.spec_fallback_reason == "binary_no_mtp" + backend._build_speculative_flags( + speculative_type = "off", + spec_draft_n_max = None, + extra_args = None, + model_identifier = _MTP_MODEL, + model_path = None, + gpus = True, + binary = "/fake/llama-server", + ) + assert backend.spec_fallback_reason is None diff --git a/studio/frontend/src/features/chat/chat-settings-sheet.tsx b/studio/frontend/src/features/chat/chat-settings-sheet.tsx index 0bd9d8eee6..03c183fe76 100644 --- a/studio/frontend/src/features/chat/chat-settings-sheet.tsx +++ b/studio/frontend/src/features/chat/chat-settings-sheet.tsx @@ -50,6 +50,7 @@ import { TooltipTrigger, } from "@/components/ui/tooltip"; import { useIsMobile } from "@/hooks/use-mobile"; +import { useLlamaUpdateCheck } from "@/hooks/use-llama-update-check"; import { cn } from "@/lib/utils"; import { ArrowDown01Icon, @@ -62,7 +63,7 @@ import { HugeiconsIcon } from "@hugeicons/react"; import { ChevronDown, ExternalLink } from "lucide-react"; import { Tooltip as TooltipPrimitive } from "radix-ui"; import { Fragment, type ReactNode } from "react"; -import { useEffect, useMemo, useRef, useState } from "react"; +import { useCallback, useEffect, useMemo, useRef, useState } from "react"; import { toast } from "@/lib/toast"; import { OpenAICodeExecSection } from "./components/openai-code-exec-section"; import { @@ -487,6 +488,27 @@ export function ChatSettingsPanel({ const loadedSpeculativeType = useChatRuntimeStore( (s) => s.loadedSpeculativeType, ); + const specFallbackReason = useChatRuntimeStore((s) => s.specFallbackReason); + // "binary_no_mtp" / "binary_outdated" mean a newer prebuilt would re-enable + // MTP; "runtime_error" means the current build cannot run it (no update push). + const mtpUpdatable = + specFallbackReason === "binary_no_mtp" || + specFallbackReason === "binary_outdated"; + const { + status: llamaUpdateStatus, + applying: llamaUpdating, + apply: applyLlamaUpdate, + } = useLlamaUpdateCheck({ enabled: mtpUpdatable }); + const handleMtpUpdate = useCallback(async () => { + const result = await applyLlamaUpdate(); + if (result.ok) { + toast.success( + `llama.cpp updated to ${result.tag ?? "the latest build"}. Reload your model to enable MTP.`, + ); + } else { + toast.error(`llama.cpp update failed: ${result.error ?? "unknown error"}`); + } + }, [applyLlamaUpdate]); const specDraftNMax = useChatRuntimeStore((s) => s.specDraftNMax); const setSpecDraftNMax = useChatRuntimeStore((s) => s.setSpecDraftNMax); const loadedSpecDraftNMax = useChatRuntimeStore( @@ -923,6 +945,32 @@ export function ChatSettingsPanel({
+ {specFallbackReason && + (speculativeType === "auto" || + speculativeType === "mtp" || + speculativeType === "mtp+ngram") && ( +
+

+ {specFallbackReason === "runtime_error" + ? "MTP could not start for this model on the installed llama.cpp build, so it is running without speculative decoding." + : "MTP is not available in the installed llama.cpp build, so this model is running without it." + + (llamaUpdateStatus?.update_available + ? " Update llama.cpp to enable it." + : "")} +

+ {mtpUpdatable && llamaUpdateStatus?.update_available && ( + + )} +
+ )} {(speculativeType === "mtp" || speculativeType === "mtp+ngram") && (
diff --git a/studio/frontend/src/features/chat/hooks/use-chat-model-runtime.ts b/studio/frontend/src/features/chat/hooks/use-chat-model-runtime.ts index 66ee2d04e4..4bd36e00d4 100644 --- a/studio/frontend/src/features/chat/hooks/use-chat-model-runtime.ts +++ b/studio/frontend/src/features/chat/hooks/use-chat-model-runtime.ts @@ -412,6 +412,7 @@ export function useChatModelRuntime() { statusRes.requires_trust_remote_code ?? false, defaultChatTemplate: nextDefaultChatTemplate, loadedIsMultimodal: isMultimodalResponse(statusRes), + specFallbackReason: statusRes.spec_fallback_reason ?? null, ...(prevState.loadedSpeculativeType === null && { speculativeType: currentSpecType, loadedSpeculativeType: currentSpecType, diff --git a/studio/frontend/src/features/chat/stores/chat-runtime-store.ts b/studio/frontend/src/features/chat/stores/chat-runtime-store.ts index aabf1b29a6..3ef63d2945 100644 --- a/studio/frontend/src/features/chat/stores/chat-runtime-store.ts +++ b/studio/frontend/src/features/chat/stores/chat-runtime-store.ts @@ -416,6 +416,11 @@ type ChatRuntimeStore = { loadedKvCacheDtype: string | null; speculativeType: string | null; loadedSpeculativeType: string | null; + /** + * Why MTP was disabled on the loaded model despite being requested, or null. + * Mirrors InferenceStatusResponse.spec_fallback_reason. + */ + specFallbackReason: string | null; /** User --spec-draft-n-max override (null = platform default). */ specDraftNMax: number | null; loadedSpecDraftNMax: number | null; @@ -765,6 +770,7 @@ export const useChatRuntimeStore = create((set, get) => ({ loadedKvCacheDtype: null, speculativeType: "auto", loadedSpeculativeType: null, + specFallbackReason: null, specDraftNMax: null, loadedSpecDraftNMax: null, loadedIsMultimodal: false, @@ -977,6 +983,7 @@ export const useChatRuntimeStore = create((set, get) => ({ loadedKvCacheDtype: null, speculativeType: "auto", loadedSpeculativeType: null, + specFallbackReason: null, specDraftNMax: null, loadedSpecDraftNMax: null, loadedIsMultimodal: false, diff --git a/studio/frontend/src/features/chat/types/api.ts b/studio/frontend/src/features/chat/types/api.ts index 49524c17b4..a0e5d5355e 100644 --- a/studio/frontend/src/features/chat/types/api.ts +++ b/studio/frontend/src/features/chat/types/api.ts @@ -174,6 +174,12 @@ export interface InferenceStatusResponse { /** Canonical UI-facing mode currently active. See LoadModelRequest. */ speculative_type?: string | null; spec_draft_n_max?: number | null; + /** + * Why MTP was disabled on the loaded model despite being requested. + * "binary_no_mtp" / "binary_outdated" -> updating llama.cpp would re-enable + * it; "runtime_error" -> the current build could not run it. Null otherwise. + */ + spec_fallback_reason?: string | null; } export interface AudioGenerationResponse { From 12a890e0cfc603df27352c7ec9796abecf2e184f Mon Sep 17 00:00:00 2001 From: Daniel Han Date: Thu, 11 Jun 2026 06:27:39 -0700 Subject: [PATCH 22/42] Fix bitsandbytes ROCm GPU arch and warp size detection on Windows (#6127) * Fix bitsandbytes ROCm GPU arch and warp size detection on Windows bitsandbytes resolves the ROCm GPU architecture (and warp size on 0.49.x) by shelling out to rocminfo / hipinfo.exe via PATH at import time. On Windows neither tool is normally on PATH (AMD torch wheels ship hipInfo.exe into the venv Scripts dir, only on PATH while activated), so every `import bitsandbytes` logs an ERROR and WARNING, ROCM_GPU_ARCH degrades to unknown, and the 0.49.x warp size defaults to 64, which is wrong on RDNA (wave 32) and silently disables pre-quantized 4-bit models via ALLOW_PREQUANTIZED_MODELS. Install a one-shot MetaPathFinder before unsloth_zoo is imported (the first bitsandbytes import on ROCm) that swaps get_rocm_gpu_arch and get_rocm_warpsize for torch-device-properties-first implementations right after bitsandbytes.cuda_specs executes, before cextension reads them. Falls back to running hipInfo.exe by absolute path (venv Scripts, conda Scripts, HIP SDK / AMD installer dirs). Repairs the constants in place when bitsandbytes was imported first. Strict no-op on non-Windows, non-ROCm builds, missing bitsandbytes, and versions that fix this upstream. Opt out with UNSLOTH_DISABLE_BNB_ROCM_FIX=1. Proposed upstream in bitsandbytes-foundation/bitsandbytes#1969; shipped here so all bitsandbytes versions are covered. Verified on gfx1151 Strix Halo, Windows 11, torch 2.11.0+rocm7.13.0 against bitsandbytes main, 0.49.2, and a torch-props-fixed variant. * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * Tighten comments in the bitsandbytes ROCm detection fix Comment and docstring pass only. AST comparison with docstrings stripped confirms every definition is identical to the version the 12 scenario suite ran against, and the suite plus the drift test pass unchanged on the edited files. * Keep the bitsandbytes cuda_specs finder installed for reload support Simulation testing caught a regression in the one-shot design: importlib.reload(bitsandbytes.cuda_specs) re-resolves the spec through sys.meta_path, so with the finder already removed the reload reinstalled the unpatched upstream detector and the Windows ROCm noise returned. Keep the finder on sys.meta_path permanently, matching the lifecycle of the existing causal_conv1d and vllm import blockers. The finder matches a single module name and patching stays idempotent via the sentinel flags, so repeat hits are no-ops. Validated on gfx1151 Windows 11: 22 simulation scenarios (conda and embedded layouts, Program Files scan ordering, paths with spaces and unicode, hanging probe timeout, lru-wrapped and C-function helper shapes, reload, failed-import retry, threads, spawn, dormant finder, Studio PATH coexistence, early fix-block ordering, bnb 0.45.5 / 0.47.0 / 0.49.2 / main / upstream-fixed) plus the original 12 scenario suite, CPU-torch and stale-HIP_PATH sandboxes, Python 3.10 to 3.13 gates, and a WSL Linux leg proving byte-identical Linux behavior with and without the fix, with and without rocminfo on PATH. --------- Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com> --- tests/test_import_fixes_drift.py | 51 +++++ unsloth/_gpu_init.py | 4 + unsloth/import_fixes.py | 334 +++++++++++++++++++++++++++++++ 3 files changed, 389 insertions(+) diff --git a/tests/test_import_fixes_drift.py b/tests/test_import_fixes_drift.py index 373830ae4a..6c75a4b886 100644 --- a/tests/test_import_fixes_drift.py +++ b/tests/test_import_fixes_drift.py @@ -584,3 +584,54 @@ def test_accelerate_utils_imports_module_present(): "accelerate.utils.imports.is_wandb_available is gone; " "disable_broken_wandb cannot patch the source module." ) + + +# =========================================================================== +# bitsandbytes -- ROCm arch / warp-size detection shape +# =========================================================================== + + +def test_bitsandbytes_rocm_detection_helpers_recognizable(): + """``fix_bitsandbytes_rocm_arch_detection`` swaps bnb's ROCm helpers + only when they shell out via subprocess and never consult torch device + props; a third shape is declined by design, silently restoring Windows + ROCm noise. Fail so the sniff gets updated. Reads source, no import.""" + spec = importlib.util.find_spec("bitsandbytes") + if spec is None: + pytest.skip("bitsandbytes not installed -- nothing to drift-check.") + cuda_specs_path = None + for location in spec.submodule_search_locations or []: + candidate = os.path.join(location, "cuda_specs.py") + if os.path.isfile(candidate): + cuda_specs_path = candidate + break + if cuda_specs_path is None: + pytest.skip("bitsandbytes has no cuda_specs.py (pre-ROCm version).") + + import ast + + with open(cuda_specs_path, "r", encoding = "utf-8") as f: + source = f.read() + helpers = [ + node + for node in ast.walk(ast.parse(source)) + if isinstance(node, ast.FunctionDef) + and node.name in ("get_rocm_gpu_arch", "get_rocm_warpsize") + ] + if not helpers: + pytest.skip("bitsandbytes cuda_specs has no ROCm detection helpers.") + for node in helpers: + segment = ast.get_source_segment(source, node) or "" + recognized = ( + "subprocess" in segment + or "get_device_properties" in segment + or "gcnArchName" in segment + ) + if not recognized: + pytest.fail( + f"DRIFT DETECTED: bitsandbytes.cuda_specs.{node.name} uses " + "neither subprocess nor torch device properties; " + "fix_bitsandbytes_rocm_arch_detection's shape sniff will " + "decline to patch it and Windows ROCm import-time noise / " + "wrong ROCM_GPU_ARCH may return." + ) diff --git a/unsloth/_gpu_init.py b/unsloth/_gpu_init.py index ede97ce68e..5da1e27a9d 100644 --- a/unsloth/_gpu_init.py +++ b/unsloth/_gpu_init.py @@ -30,6 +30,7 @@ from .import_fixes import ( disable_broken_causal_conv1d, disable_broken_vllm, configure_amdgpu_asic_id_table_path, + fix_bitsandbytes_rocm_arch_detection, torchvision_compatibility_check, fix_diffusers_warnings, fix_huggingface_hub, @@ -67,6 +68,8 @@ except Exception: # Configure libdrm ids table path early so ROCm can resolve AMD GPU names. configure_amdgpu_asic_id_table_path() +# Must precede `import unsloth_zoo` below, which imports bnb on ROCm. +fix_bitsandbytes_rocm_arch_detection() disable_broken_causal_conv1d() disable_broken_vllm() fix_message_factory_issue() @@ -75,6 +78,7 @@ torchvision_compatibility_check() fix_diffusers_warnings() fix_huggingface_hub() del configure_amdgpu_asic_id_table_path +del fix_bitsandbytes_rocm_arch_detection del disable_broken_causal_conv1d del disable_broken_vllm del fix_message_factory_issue diff --git a/unsloth/import_fixes.py b/unsloth/import_fixes.py index 0622f73965..695a6a577a 100644 --- a/unsloth/import_fixes.py +++ b/unsloth/import_fixes.py @@ -1823,6 +1823,340 @@ def configure_amdgpu_asic_id_table_path(): return None +# --------------------------------------------------------------------------- +# bitsandbytes Windows ROCm fix: cextension.py runs get_rocm_gpu_arch() +# (bnb >= 0.47) and get_rocm_warpsize() (0.49.x) at import, shelling out to +# rocminfo / hipinfo.exe via PATH. Neither is on PATH on Windows (AMD torch +# wheels put hipInfo.exe in venv Scripts), so every import logs ERROR + +# WARNING, ROCM_GPU_ARCH becomes "unknown", and warp size defaults to 64: +# wrong on RDNA (wave 32), breaking 4-bit blocksizes and +# ALLOW_PREQUANTIZED_MODELS. Upstream fix unmerged (bitsandbytes#1969), so a +# MetaPathFinder swaps both helpers for torch-device-props-first versions +# right after bitsandbytes.cuda_specs executes, before cextension reads +# them. Must run before `import unsloth_zoo` (imports bnb on ROCm). +# --------------------------------------------------------------------------- + +_BNB_CUDA_SPECS_MODULE = "bitsandbytes.cuda_specs" +_BNB_ROCM_FIX_FINDER_SENTINEL = "_unsloth_bnb_rocm_fix_finder" +_BNB_ROCM_FIX_FUNCTION_FLAG = "__unsloth_bnb_rocm_fix__" + + +def _torch_rocm_device_props(): + """Device-0 props on a ROCm torch build with a visible GPU, else None. + Never raises; bnb's own import initializes the device context anyway.""" + try: + import torch + + if not getattr(getattr(torch, "version", None), "hip", None): + return None + if not torch.cuda.is_available(): + return None + return torch.cuda.get_device_properties(0) + except Exception: + return None + + +def _iter_hipinfo_paths(): + """Yield existing hipInfo.exe paths: PATH, interpreter scripts dir (venv + and conda layouts), then HIP SDK / AMD installer locations.""" + import shutil + import sysconfig + + candidates = [] + try: + resolved = shutil.which("hipinfo.exe") + if resolved: + candidates.append(resolved) + except Exception: + pass + try: + scripts_dir = sysconfig.get_path("scripts") + if scripts_dir: + candidates.append(os.path.join(scripts_dir, "hipInfo.exe")) + except Exception: + pass + executable_dir = os.path.dirname(sys.executable or "") + if executable_dir: + candidates.append(os.path.join(executable_dir, "hipInfo.exe")) + candidates.append(os.path.join(executable_dir, "Scripts", "hipInfo.exe")) + for env_key in ("HIP_PATH", "ROCM_PATH"): + root = os.environ.get(env_key, "").strip() + if root: + candidates.append(os.path.join(root, "bin", "hipInfo.exe")) + rocm_root = os.path.join(os.environ.get("ProgramFiles", r"C:\Program Files"), "AMD", "ROCm") + try: + if os.path.isdir(rocm_root): + for version_dir in sorted(os.listdir(rocm_root), reverse = True): + candidates.append(os.path.join(rocm_root, version_dir, "bin", "hipInfo.exe")) + except Exception: + pass + + seen = set() + for candidate in candidates: + try: + key = os.path.normcase(os.path.normpath(candidate)) + if key in seen: + continue + seen.add(key) + if os.path.isfile(candidate): + yield candidate + except Exception: + continue + + +def _run_hipinfo(hipinfo_path): + """Run hipInfo.exe and return its stdout, or "" on any failure.""" + import subprocess + try: + result = subprocess.run( + [hipinfo_path], + capture_output = True, + text = True, + timeout = 15, + creationflags = getattr(subprocess, "CREATE_NO_WINDOW", 0), + ) + return result.stdout or "" + except Exception as e: + _log_rocm_detection(f"Unsloth: `{hipinfo_path}` failed: {e}") + return "" + + +def _unsloth_get_rocm_gpu_arch(): + """Replaces bnb's get_rocm_gpu_arch: torch device props first (no + subprocess), then hipInfo.exe by absolute path, then a quiet "unknown".""" + try: + import torch + if not getattr(getattr(torch, "version", None), "hip", None): + return "unknown" + except Exception: + return "unknown" + props = _torch_rocm_device_props() + if props is not None: + try: + # gcnArchName may carry feature flags, e.g. "gfx90a:sramecc+:xnack-" + arch = str(props.gcnArchName).split(":")[0].strip() + if arch.startswith("gfx"): + return arch + except Exception: + pass + for hipinfo_path in _iter_hipinfo_paths(): + match = re.search(r"gcnArchName:\s+gfx([a-zA-Z\d]+)", _run_hipinfo(hipinfo_path)) + if match: + return "gfx" + match.group(1) + _log_rocm_detection( + "Unsloth: Could not detect the ROCm GPU architecture - bitsandbytes will see `unknown`." + ) + return "unknown" + + +def _unsloth_get_rocm_warpsize(): + """Replaces bnb 0.49.x get_rocm_warpsize: upstream defaults to 64 when + rocminfo is missing, wrong on RDNA (wave 32).""" + try: + import torch + if not getattr(getattr(torch, "version", None), "hip", None): + return 32 # upstream behavior: NVIDIA warp size is always 32 + except Exception: + return 64 # upstream behavior: default to 64 on failure + props = _torch_rocm_device_props() + if props is not None: + # torch 2.11 ROCm exposes warp_size; some builds used warpSize. + for attribute_name in ("warp_size", "warpSize"): + warp_size = getattr(props, attribute_name, None) + if isinstance(warp_size, int) and warp_size in (32, 64): + return warp_size + for hipinfo_path in _iter_hipinfo_paths(): + match = re.search(r"^\s*warpSize:\s+(\d+)", _run_hipinfo(hipinfo_path), re.MULTILINE) + if match and int(match.group(1)) in (32, 64): + return int(match.group(1)) + _log_rocm_detection( + "Unsloth: Could not detect the ROCm warp size - defaulting to 64 " + "(bitsandbytes' own default)." + ) + return 64 + + +setattr(_unsloth_get_rocm_gpu_arch, _BNB_ROCM_FIX_FUNCTION_FLAG, True) +setattr(_unsloth_get_rocm_warpsize, _BNB_ROCM_FIX_FUNCTION_FLAG, True) + + +def _bnb_rocm_helper_is_broken(function): + """True only for upstream's subprocess-only detectors; co_names works + where getsource fails. Versions consulting torch props are untouched.""" + if function is None or not callable(function): + return False + if getattr(function, _BNB_ROCM_FIX_FUNCTION_FLAG, False): + return False # Already ours. + try: + function = inspect.unwrap(function) + except Exception: + pass + code = getattr(function, "__code__", None) + co_names = getattr(code, "co_names", ()) if code is not None else () + if not co_names: + return False # C function or opaque wrapper -- do not touch. + if "get_device_properties" in co_names or "gcnArchName" in co_names: + return False # Fixed upstream -- no-op. + return "subprocess" in co_names + + +def _patch_bnb_cuda_specs_module(module): + """Swap broken ROCm detection helpers on an executed cuda_specs module. + Returns True when the module ends up patched (now or previously).""" + patched = False + for attribute_name, replacement in ( + ("get_rocm_gpu_arch", _unsloth_get_rocm_gpu_arch), + ("get_rocm_warpsize", _unsloth_get_rocm_warpsize), + ): + original = getattr(module, attribute_name, None) + if getattr(original, _BNB_ROCM_FIX_FUNCTION_FLAG, False): + patched = True # Already ours. + continue + if not _bnb_rocm_helper_is_broken(original): + continue + setattr(module, attribute_name, replacement) + patched = True + logger.info( + f"Unsloth: Patched bitsandbytes.cuda_specs.{attribute_name} - " + f"avoids PATH-dependent subprocess GPU detection on Windows ROCm." + ) + return patched + + +class _BnbCudaSpecsPatchLoader(importlib.abc.Loader): + __slots__ = ("_loader",) + + def __init__(self, loader): + self._loader = loader + + def create_module(self, spec): + create_module = getattr(self._loader, "create_module", None) + if create_module is None: + return None + return create_module(spec) + + def exec_module(self, module): + self._loader.exec_module(module) + # Patch after the module body ran, before cextension calls it. The + # finder stays on sys.meta_path (same lifecycle as the blockers + # above) so importlib.reload(bitsandbytes.cuda_specs) re-patches. + try: + _patch_bnb_cuda_specs_module(module) + except Exception as e: + _log_rocm_detection(f"Unsloth: bitsandbytes ROCm detection patch failed: {e}") + + def __getattr__(self, name): + # Delegate get_source / get_filename etc. so introspection works. + return getattr(self._loader, name) + + +class _BnbCudaSpecsPatchFinder(importlib.abc.MetaPathFinder): + __slots__ = (_BNB_ROCM_FIX_FINDER_SENTINEL,) + + def __init__(self): + setattr(self, _BNB_ROCM_FIX_FINDER_SENTINEL, True) + + def find_spec( + self, + fullname, + path = None, + target = None, + ): + if fullname != _BNB_CUDA_SPECS_MODULE: + return None + # Delegate to remaining finders (editable installs, frozen apps) + # and wrap the loader that would actually be used. + spec = None + for finder in sys.meta_path: + if finder is self or getattr(finder, _BNB_ROCM_FIX_FINDER_SENTINEL, False): + continue + finder_find_spec = getattr(finder, "find_spec", None) + if finder_find_spec is None: + continue + try: + spec = finder_find_spec(fullname, path, target) + except Exception: + spec = None + if spec is not None: + break + if spec is None or spec.loader is None: + return None + if not hasattr(spec.loader, "exec_module"): + return None # Legacy loader -- let the stock machinery handle it. + spec.loader = _BnbCudaSpecsPatchLoader(spec.loader) + return spec + + +def _repair_imported_bitsandbytes_rocm_constants(): + """bnb imported before unsloth: noise already fired, but fix detectors + and cached constants, incl. by-value ROCM_WARP_SIZE_64 copies.""" + cuda_specs = sys.modules.get(_BNB_CUDA_SPECS_MODULE) + if cuda_specs is None: + return + if not _patch_bnb_cuda_specs_module(cuda_specs): + return + + try: + arch = cuda_specs.get_rocm_gpu_arch() + except Exception: + arch = "unknown" + warp_size_64 = None + get_rocm_warpsize = getattr(cuda_specs, "get_rocm_warpsize", None) + if callable(get_rocm_warpsize): + try: + warp_size_64 = get_rocm_warpsize() == 64 + except Exception: + warp_size_64 = None + + for module_name, module in list(sys.modules.items()): + if module is None or module is cuda_specs: + continue + if module_name != "bitsandbytes" and not module_name.startswith("bitsandbytes."): + continue + try: + if arch != "unknown" and getattr(module, "ROCM_GPU_ARCH", None) == "unknown": + module.ROCM_GPU_ARCH = arch + if warp_size_64 is not None and isinstance( + getattr(module, "ROCM_WARP_SIZE_64", None), bool + ): + module.ROCM_WARP_SIZE_64 = warp_size_64 + except Exception: + continue + logger.info("Unsloth: Repaired bitsandbytes ROCm arch / warp-size constants in place.") + + +def fix_bitsandbytes_rocm_arch_detection(): + """Fix bnb's import-time ROCm arch / warp-size detection on Windows + (see header above). No-op on non-Windows, non-ROCm, missing or + upstream-fixed bnb. Idempotent. Opt out: UNSLOTH_DISABLE_BNB_ROCM_FIX=1.""" + if os.environ.get("UNSLOTH_DISABLE_BNB_ROCM_FIX", "0") == "1": + return + if sys.platform != "win32": + return + if not _is_rocm_torch_build(): + return + + # Already imported: prevention impossible, repair in place instead. + if _BNB_CUDA_SPECS_MODULE in sys.modules: + try: + _repair_imported_bitsandbytes_rocm_constants() + except Exception: + pass + return + + try: + if importlib.util.find_spec("bitsandbytes") is None: + return + except Exception: + return + + for finder in sys.meta_path: + if getattr(finder, _BNB_ROCM_FIX_FINDER_SENTINEL, False): + return # Already installed -- idempotent. + sys.meta_path.insert(0, _BnbCudaSpecsPatchFinder()) + _log_rocm_detection("Unsloth: Installed the bitsandbytes ROCm arch detection patch hook.") + + def _is_causal_conv1d_name(module_name: str) -> bool: return module_name == _CAUSAL_CONV1D_PREFIX or module_name.startswith( _CAUSAL_CONV1D_PREFIX + "." From 73973435ad66084157efd843b25577dd65b999ae Mon Sep 17 00:00:00 2001 From: Daniel Han Date: Thu, 11 Jun 2026 06:52:38 -0700 Subject: [PATCH 23/42] Studio: require an installed ROCm DLL before forcing BNB_ROCM_VERSION; drop shadowing shutil imports in save.py (#6194) * Require a found ROCm DLL before forcing BNB_ROCM_VERSION in Studio paths main.py previously set BNB_ROCM_VERSION=72 whenever HIP_PATH or ROCM_PATH was set, and the training worker fell back to a blind 72 when DLL detection found nothing. On a Windows machine with the AMD HIP SDK installed but CUDA or CPU torch, that forces a ROCm backend onto a non-ROCm bitsandbytes wheel, which raises at import. Both paths now only write the override when a libbitsandbytes_rocm DLL actually exists (or a seeded value is already present), matching the strict gates in unsloth/import_fixes.py. Also removes four redundant local import shutil statements in unsloth/save.py that shadow the module-level import, the same pattern that caused the UnboundLocalError fixed in #6149. * Worker: gate the BNB override on a found ROCm DLL, preserving seeded marker Review follow-ups: track _found_rocm_bnb in the worker like main.py so a ROCm DLL with an unparsable name still gets the seeded or 72 fallback, and skip the env write entirely when no DLL exists so a seeded value keeps its sitecustomize marker and stays redetectable by later import fixes. --- studio/backend/core/training/worker.py | 30 ++++++++++++++--------- studio/backend/main.py | 15 ++++++------ tests/studio/install/test_rocm_support.py | 22 ++++++++++++++--- unsloth/save.py | 4 --- 4 files changed, 45 insertions(+), 26 deletions(-) diff --git a/studio/backend/core/training/worker.py b/studio/backend/core/training/worker.py index b50e17e5aa..fb5e9b5137 100644 --- a/studio/backend/core/training/worker.py +++ b/studio/backend/core/training/worker.py @@ -2143,14 +2143,14 @@ def run_training_process(*, event_queue: Any, stop_queue: Any, config: dict) -> # BNB picks a rocm DLL from torch.version.hip, but AMD's Windows BNB # wheel may ship a DLL whose suffix doesn't match. Detect the actual - # DLL name and override; "72" is a safe fallback. Values seeded by - # the installer are redetectable defaults, while caller overrides - # remain authoritative. + # DLL name and override. Values seeded by the installer are + # redetectable defaults, while caller overrides remain authoritative. if ( "BNB_ROCM_VERSION" not in os.environ or os.environ.get("UNSLOTH_BNB_ROCM_VERSION_SOURCE") == "sitecustomize" ): _bnb_rocm_ver = None + _found_rocm_bnb = False try: import glob as _glob import importlib.util as _ilu @@ -2163,6 +2163,7 @@ def run_training_process(*, event_queue: Any, stop_queue: Any, config: dict) -> for _dll in _glob.glob( os.path.join(_pkg_dir, "libbitsandbytes_rocm*.dll") ): + _found_rocm_bnb = True _m = _re.search( r"libbitsandbytes_rocm(\d+)\.dll", os.path.basename(_dll), @@ -2174,15 +2175,20 @@ def run_training_process(*, event_queue: Any, stop_queue: Any, config: dict) -> _bnb_rocm_ver = max(_all_vers, key = lambda v: int(v)) except Exception: pass - _bnb_rocm_ver = _bnb_rocm_ver or os.environ.get("BNB_ROCM_VERSION") or "72" - os.environ["BNB_ROCM_VERSION"] = _bnb_rocm_ver - os.environ["UNSLOTH_BNB_ROCM_VERSION_SOURCE"] = "detected" - logger.info( - "Windows ROCm: set BNB_ROCM_VERSION=%s " - "(detected from installed BNB wheel; " - "overrides torch.version.hip auto-detection)", - _bnb_rocm_ver, - ) + # Only when a ROCm bnb DLL actually exists (mirrors main.py): + # without one the seeded value and its marker stay untouched, + # so later import fixes can still redetect or opt out. DLL + # with unparsable name -> seeded value or "72". + if _found_rocm_bnb: + _bnb_rocm_ver = _bnb_rocm_ver or os.environ.get("BNB_ROCM_VERSION") or "72" + os.environ["BNB_ROCM_VERSION"] = _bnb_rocm_ver + os.environ["UNSLOTH_BNB_ROCM_VERSION_SOURCE"] = "detected" + logger.info( + "Windows ROCm: set BNB_ROCM_VERSION=%s " + "(detected from installed BNB wheel; " + "overrides torch.version.hip auto-detection)", + _bnb_rocm_ver, + ) # Parse HIP version for the kernel-fix gate below, falling back to # the rocm version embedded in torch.__version__ when version.hip is diff --git a/studio/backend/main.py b/studio/backend/main.py index dc3d802b78..2b34b11fef 100644 --- a/studio/backend/main.py +++ b/studio/backend/main.py @@ -83,9 +83,9 @@ if sys.platform == "win32": # ── Windows AMD ROCm: set BNB_ROCM_VERSION before any bitsandbytes import ─ # bitsandbytes derives the rocm.dll name from torch.version.hip, but the # wheel ships rocm72.dll, so the server crashes ("Configured ROCm binary not - # found") without this. Detect the shipped DLL and fall back to "72" (mirrors - # worker.py). Gate on the rocm bnb DLL / HIP_PATH rather than torch.version.hip - # to avoid importing torch on every Windows host. + # found") without this. Detect the shipped DLL (mirrors worker.py); gate on + # the rocm bnb DLL rather than torch.version.hip to avoid importing torch on + # every Windows host. # Values seeded by the installer's sitecustomize.py are redetectable # defaults; explicit caller values remain authoritative. if ( @@ -95,7 +95,6 @@ if sys.platform == "win32": import glob as _glob import logging as _logging - _hip_env = bool(os.environ.get("HIP_PATH") or os.environ.get("ROCM_PATH")) _bnb_rocm_ver = None _found_rocm_bnb = False try: @@ -118,11 +117,13 @@ if sys.platform == "win32": _bnb_rocm_ver = max(_all_vers_main, key = lambda v: int(v)) except Exception as _e: _logging.getLogger(__name__).warning( - "Windows ROCm: BNB DLL detection failed (%s); falling back to version '72'", + "Windows ROCm: BNB DLL detection failed (%s); leaving BNB_ROCM_VERSION as is", _e, ) - # rocm bnb DLL present, or HIP_PATH/ROCM_PATH set (DLL unparsable -> "72") - if _found_rocm_bnb or _hip_env: + # Only when a ROCm bnb DLL actually exists: HIP_PATH/ROCM_PATH alone + # (HIP SDK on a CUDA/CPU box) must not force a ROCm backend onto a + # non-ROCm bitsandbytes, which raises at import. DLL unparsable -> "72". + if _found_rocm_bnb: _bnb_rocm_ver_final = _bnb_rocm_ver or os.environ.get("BNB_ROCM_VERSION") or "72" os.environ["BNB_ROCM_VERSION"] = _bnb_rocm_ver_final os.environ["UNSLOTH_BNB_ROCM_VERSION_SOURCE"] = "detected" diff --git a/tests/studio/install/test_rocm_support.py b/tests/studio/install/test_rocm_support.py index 6bc94e3d0f..7d1ae535fd 100644 --- a/tests/studio/install/test_rocm_support.py +++ b/tests/studio/install/test_rocm_support.py @@ -2241,7 +2241,22 @@ class TestRuntimeBnbRocmSourceGuards: """A failed redetect must not downgrade a persisted suffix to '72'.""" for path in (self._MAIN_PATH, self._TRAINING_WORKER_PATH): source = path.read_text(encoding = "utf-8") - assert 'os.environ.get("BNB_ROCM_VERSION") or "72"' in source, path.name + assert ( + '_bnb_rocm_ver or os.environ.get("BNB_ROCM_VERSION") or "72"' in source + ), path.name + + def test_main_requires_found_rocm_dll(self): + """HIP_PATH/ROCM_PATH alone (HIP SDK on a CUDA/CPU box) must not force + a ROCm backend onto a non-ROCm bitsandbytes.""" + source = self._MAIN_PATH.read_text(encoding = "utf-8") + assert "if _found_rocm_bnb:" in source + assert "_hip_env" not in source + + def test_worker_requires_found_rocm_dll(self): + """No DLL found: the worker must not write any override or touch the + seeded marker (later import fixes must still see sitecustomize).""" + source = self._TRAINING_WORKER_PATH.read_text(encoding = "utf-8") + assert "if _found_rocm_bnb:" in source class TestDetectBnbRocmDllVer: @@ -2443,8 +2458,9 @@ class TestWorkerWindowsRocmPatches: assert "BNB_ROCM_VERSION" in source # Detection helper must be used assert "_detect_bnb_rocm_dll_ver" in source or "libbitsandbytes_rocm" in source - # "72" must appear as the safe fallback - assert '"72"' in source or "'72'" in source + # Falls back to the seeded value, never a blind "72" (which would + # force a ROCm backend onto a non-ROCm bitsandbytes wheel) + assert '_bnb_rocm_ver or os.environ.get("BNB_ROCM_VERSION")' in source def test_bnb_rocm_version_set_before_ml_imports(self): """BNB_ROCM_VERSION must appear in section 1f, before section 2 ML imports.""" diff --git a/unsloth/save.py b/unsloth/save.py index 4ff56cc477..629cbb9548 100644 --- a/unsloth/save.py +++ b/unsloth/save.py @@ -1092,7 +1092,6 @@ def unsloth_save_model( gc.collect() # Remove temporary location - import shutil shutil.rmtree(temporary_location, ignore_errors = True) @@ -1224,7 +1223,6 @@ def install_llama_cpp_old(version = -10): for i in range(30): print(f"**[WARNING]** Deleting llama.cpp directory... {30-i} seconds left.") time.sleep(1) - import shutil shutil.rmtree("llama.cpp", ignore_errors = True) @@ -2494,7 +2492,6 @@ def unsloth_push_to_hub_gguf( except Exception as e: if cleanup_temp: - import shutil for d in [save_directory, f"{save_directory}_gguf"]: try: shutil.rmtree(d) @@ -2681,7 +2678,6 @@ This model was finetuned and converted to GGUF format using [Unsloth](https://gi # Clean up temporary directory if cleanup_temp: print("Unsloth: Cleaning up temporary files...") - import shutil for d in [save_directory, f"{save_directory}_gguf"]: if os.path.exists(d): try: From 2e29363ad9fc943b2322765609dd3b66d55a25e1 Mon Sep 17 00:00:00 2001 From: Daniel Han Date: Thu, 11 Jun 2026 06:57:48 -0700 Subject: [PATCH 24/42] Studio CI: stop HF 429 rate limits from sinking the llama.cpp prebuilt path (#6199) * Stop HF 429 rate limits from sinking the llama.cpp prebuilt path in Studio CI The Windows Studio API smoke job failed when anonymous huggingface.co fetches of the tiny GGUF validation model (stories260K.gguf) hit HTTP 429 on the shared runner IP. The installer correctly refused the unvalidated prebuilt and fell back to a source build, which the prebuilt assert then flags. Three layers fix this: 1. Installer: auth_headers sends HF_TOKEN (or HUGGING_FACE_HUB_TOKEN) to huggingface.co hosts, mirroring the existing GH_TOKEN handling for the GitHub API rate limit. A redirect handler strips Authorization when a download is redirected off-host (CDN signed URLs reject foreign auth; urllib forwards headers on redirect, unlike requests/huggingface_hub). 2. Workflows: the HF_HOME prime steps also prefetch the validation model so the install's hf_hub_download resolves from the local cache even when the Hub is rate limiting; cache keys bumped v1 to v2 to repopulate. This also covers fork PRs, which cannot see secrets. 3. Workflows: every Install Studio / update step that already passes GH_TOKEN now also passes HF_TOKEN, so both the huggingface_hub path and the direct URL fallback are authenticated. Tests: tests/studio/install/test_hf_auth.py covers token-to-host routing, the cross-host redirect strip, and the download_bytes wiring (offline). Verified live: authenticated download of the validation model through the new opener (CDN redirect exercised, pinned sha matches) and an offline hf_hub_download cache hit against an HF_HOME primed by the new step. * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci --------- Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com> --- .github/workflows/studio-api-smoke.yml | 6 +- .github/workflows/studio-inference-smoke.yml | 13 +- .github/workflows/studio-mac-api-smoke.yml | 6 +- .../workflows/studio-mac-inference-smoke.yml | 8 +- .../workflows/studio-mac-install-matrix.yml | 1 + .github/workflows/studio-mac-ui-smoke.yml | 6 +- .github/workflows/studio-mac-update-smoke.yml | 3 + .github/workflows/studio-ui-smoke.yml | 6 +- .github/workflows/studio-update-smoke.yml | 3 + .../workflows/studio-windows-api-smoke.yml | 6 +- .../studio-windows-inference-smoke.yml | 13 +- .github/workflows/studio-windows-ui-smoke.yml | 6 +- .../workflows/studio-windows-update-smoke.yml | 3 + studio/install_llama_prebuilt.py | 35 +++++- tests/studio/install/test_hf_auth.py | 114 ++++++++++++++++++ 15 files changed, 205 insertions(+), 24 deletions(-) create mode 100644 tests/studio/install/test_hf_auth.py diff --git a/.github/workflows/studio-api-smoke.yml b/.github/workflows/studio-api-smoke.yml index 53514e2ce1..b196805cf7 100644 --- a/.github/workflows/studio-api-smoke.yml +++ b/.github/workflows/studio-api-smoke.yml @@ -77,7 +77,7 @@ jobs: path: hf-cache # Same key as studio-ui-smoke.yml so the two jobs share a # single GGUF download across CI. - key: ${{ runner.os }}-hf-${{ env.GGUF_REPO }}-${{ env.GGUF_VARIANT }}-v1 + key: ${{ runner.os }}-hf-${{ env.GGUF_REPO }}-${{ env.GGUF_VARIANT }}-v2 - name: Prime HF_HOME with the GGUF id: prime-hf @@ -88,17 +88,19 @@ jobs: python -m pip install --upgrade huggingface_hub mkdir -p hf-cache bash .github/scripts/hf-download-with-retry.sh "$GGUF_REPO" "$GGUF_FILE" + bash .github/scripts/hf-download-with-retry.sh ggml-org/models tinyllamas/stories260K.gguf - name: Save HF_HOME for ${{ env.GGUF_REPO }} if: always() && steps.prime-hf.outcome == 'success' uses: actions/cache/save@27d5ce7f107fe9357f9df03efb73ab90386fccae # v5.0.5 with: path: hf-cache - key: ${{ runner.os }}-hf-${{ env.GGUF_REPO }}-${{ env.GGUF_VARIANT }}-v1 + key: ${{ runner.os }}-hf-${{ env.GGUF_REPO }}-${{ env.GGUF_VARIANT }}-v2 - name: Install Studio (--local, --no-torch) env: GH_TOKEN: ${{ secrets.GITHUB_TOKEN }} + HF_TOKEN: ${{ secrets.HF_TOKEN }} run: | mkdir -p logs set -o pipefail diff --git a/.github/workflows/studio-inference-smoke.yml b/.github/workflows/studio-inference-smoke.yml index 133dd601d1..cffb33f71d 100644 --- a/.github/workflows/studio-inference-smoke.yml +++ b/.github/workflows/studio-inference-smoke.yml @@ -91,7 +91,7 @@ jobs: continue-on-error: true with: path: hf-cache - key: ${{ runner.os }}-hf-${{ env.GGUF_REPO }}-${{ env.GGUF_VARIANT }}-v1 + key: ${{ runner.os }}-hf-${{ env.GGUF_REPO }}-${{ env.GGUF_VARIANT }}-v2 - name: Prime HF_HOME with the GGUF id: prime-hf @@ -102,17 +102,19 @@ jobs: python -m pip install --upgrade huggingface_hub mkdir -p hf-cache bash .github/scripts/hf-download-with-retry.sh "$GGUF_REPO" "$GGUF_FILE" + bash .github/scripts/hf-download-with-retry.sh ggml-org/models tinyllamas/stories260K.gguf - name: Save HF_HOME for ${{ env.GGUF_REPO }} if: always() && steps.prime-hf.outcome == 'success' uses: actions/cache/save@27d5ce7f107fe9357f9df03efb73ab90386fccae # v5.0.5 with: path: hf-cache - key: ${{ runner.os }}-hf-${{ env.GGUF_REPO }}-${{ env.GGUF_VARIANT }}-v1 + key: ${{ runner.os }}-hf-${{ env.GGUF_REPO }}-${{ env.GGUF_VARIANT }}-v2 - name: Install Studio (--local, --no-torch) env: GH_TOKEN: ${{ secrets.GITHUB_TOKEN }} + HF_TOKEN: ${{ secrets.HF_TOKEN }} run: | mkdir -p logs set -o pipefail @@ -375,6 +377,7 @@ jobs: - name: Install Studio (--local, --no-torch) env: GH_TOKEN: ${{ secrets.GITHUB_TOKEN }} + HF_TOKEN: ${{ secrets.HF_TOKEN }} run: | mkdir -p logs set -o pipefail @@ -829,7 +832,7 @@ jobs: continue-on-error: true with: path: hf-cache - key: ${{ runner.os }}-hf-${{ env.GGUF_REPO }}-${{ env.GGUF_VARIANT }}-${{ env.MMPROJ_FILE }}-v1 + key: ${{ runner.os }}-hf-${{ env.GGUF_REPO }}-${{ env.GGUF_VARIANT }}-${{ env.MMPROJ_FILE }}-v2 - name: Prime HF_HOME with the GGUF + mmproj id: prime-hf @@ -841,17 +844,19 @@ jobs: mkdir -p hf-cache bash .github/scripts/hf-download-with-retry.sh "$GGUF_REPO" "$GGUF_FILE" bash .github/scripts/hf-download-with-retry.sh "$GGUF_REPO" "$MMPROJ_FILE" + bash .github/scripts/hf-download-with-retry.sh ggml-org/models tinyllamas/stories260K.gguf - name: Save HF_HOME for ${{ env.GGUF_REPO }} (model + mmproj) if: always() && steps.prime-hf.outcome == 'success' uses: actions/cache/save@27d5ce7f107fe9357f9df03efb73ab90386fccae # v5.0.5 with: path: hf-cache - key: ${{ runner.os }}-hf-${{ env.GGUF_REPO }}-${{ env.GGUF_VARIANT }}-${{ env.MMPROJ_FILE }}-v1 + key: ${{ runner.os }}-hf-${{ env.GGUF_REPO }}-${{ env.GGUF_VARIANT }}-${{ env.MMPROJ_FILE }}-v2 - name: Install Studio (--local, --no-torch) env: GH_TOKEN: ${{ secrets.GITHUB_TOKEN }} + HF_TOKEN: ${{ secrets.HF_TOKEN }} run: | mkdir -p logs set -o pipefail diff --git a/.github/workflows/studio-mac-api-smoke.yml b/.github/workflows/studio-mac-api-smoke.yml index 8839b559fd..412726538c 100644 --- a/.github/workflows/studio-mac-api-smoke.yml +++ b/.github/workflows/studio-mac-api-smoke.yml @@ -62,7 +62,7 @@ jobs: continue-on-error: true with: path: hf-cache - key: ${{ runner.os }}-hf-${{ env.GGUF_REPO }}-${{ env.GGUF_VARIANT }}-v1 + key: ${{ runner.os }}-hf-${{ env.GGUF_REPO }}-${{ env.GGUF_VARIANT }}-v2 - name: Prime HF_HOME with the GGUF id: prime-hf @@ -73,17 +73,19 @@ jobs: python -m pip install --upgrade huggingface_hub mkdir -p hf-cache bash .github/scripts/hf-download-with-retry.sh "$GGUF_REPO" "$GGUF_FILE" + bash .github/scripts/hf-download-with-retry.sh ggml-org/models tinyllamas/stories260K.gguf - name: Save HF_HOME for ${{ env.GGUF_REPO }} if: always() && steps.prime-hf.outcome == 'success' uses: actions/cache/save@27d5ce7f107fe9357f9df03efb73ab90386fccae # v5.0.5 with: path: hf-cache - key: ${{ runner.os }}-hf-${{ env.GGUF_REPO }}-${{ env.GGUF_VARIANT }}-v1 + key: ${{ runner.os }}-hf-${{ env.GGUF_REPO }}-${{ env.GGUF_VARIANT }}-v2 - name: Install Studio (--local, --no-torch) env: GH_TOKEN: ${{ secrets.GITHUB_TOKEN }} + HF_TOKEN: ${{ secrets.HF_TOKEN }} run: | mkdir -p logs set -o pipefail diff --git a/.github/workflows/studio-mac-inference-smoke.yml b/.github/workflows/studio-mac-inference-smoke.yml index e08ac6ca68..c794a34acd 100644 --- a/.github/workflows/studio-mac-inference-smoke.yml +++ b/.github/workflows/studio-mac-inference-smoke.yml @@ -85,7 +85,7 @@ jobs: continue-on-error: true with: path: hf-cache - key: ${{ runner.os }}-hf-${{ env.GGUF_REPO }}-${{ env.GGUF_VARIANT }}-v1 + key: ${{ runner.os }}-hf-${{ env.GGUF_REPO }}-${{ env.GGUF_VARIANT }}-v2 - name: Prime HF_HOME with the GGUF id: prime-hf @@ -96,6 +96,7 @@ jobs: python -m pip install --upgrade huggingface_hub mkdir -p hf-cache bash .github/scripts/hf-download-with-retry.sh "$GGUF_REPO" "$GGUF_FILE" + bash .github/scripts/hf-download-with-retry.sh ggml-org/models tinyllamas/stories260K.gguf # Save partial caches on cancel/timeout -- hf download resumes by # content hash. `outcome != skipped` keeps cache-hit a no-op. @@ -104,11 +105,12 @@ jobs: uses: actions/cache/save@27d5ce7f107fe9357f9df03efb73ab90386fccae # v5.0.5 with: path: hf-cache - key: ${{ runner.os }}-hf-${{ env.GGUF_REPO }}-${{ env.GGUF_VARIANT }}-v1 + key: ${{ runner.os }}-hf-${{ env.GGUF_REPO }}-${{ env.GGUF_VARIANT }}-v2 - name: Install Studio (--local, --no-torch) env: GH_TOKEN: ${{ secrets.GITHUB_TOKEN }} + HF_TOKEN: ${{ secrets.HF_TOKEN }} run: | mkdir -p logs set -o pipefail @@ -361,6 +363,7 @@ jobs: - name: Install Studio (--local, --no-torch) env: GH_TOKEN: ${{ secrets.GITHUB_TOKEN }} + HF_TOKEN: ${{ secrets.HF_TOKEN }} run: | mkdir -p logs set -o pipefail @@ -749,6 +752,7 @@ jobs: - name: Install Studio (--local, --no-torch) env: GH_TOKEN: ${{ secrets.GITHUB_TOKEN }} + HF_TOKEN: ${{ secrets.HF_TOKEN }} run: | mkdir -p logs set -o pipefail diff --git a/.github/workflows/studio-mac-install-matrix.yml b/.github/workflows/studio-mac-install-matrix.yml index 4e2722d1cd..da944d4b5c 100644 --- a/.github/workflows/studio-mac-install-matrix.yml +++ b/.github/workflows/studio-mac-install-matrix.yml @@ -63,6 +63,7 @@ jobs: - name: Install Studio (--local, --no-torch) env: GH_TOKEN: ${{ secrets.GITHUB_TOKEN }} + HF_TOKEN: ${{ secrets.HF_TOKEN }} run: | mkdir -p logs set -o pipefail diff --git a/.github/workflows/studio-mac-ui-smoke.yml b/.github/workflows/studio-mac-ui-smoke.yml index 0176b0a168..4f9f94b534 100644 --- a/.github/workflows/studio-mac-ui-smoke.yml +++ b/.github/workflows/studio-mac-ui-smoke.yml @@ -62,7 +62,7 @@ jobs: continue-on-error: true with: path: hf-cache - key: ${{ runner.os }}-hf-${{ env.GGUF_REPO }}-${{ env.GGUF_VARIANT }}-v1 + key: ${{ runner.os }}-hf-${{ env.GGUF_REPO }}-${{ env.GGUF_VARIANT }}-v2 - name: Prime HF_HOME with the GGUF id: prime-hf @@ -73,17 +73,19 @@ jobs: python -m pip install --upgrade huggingface_hub mkdir -p hf-cache bash .github/scripts/hf-download-with-retry.sh "$GGUF_REPO" "$GGUF_FILE" + bash .github/scripts/hf-download-with-retry.sh ggml-org/models tinyllamas/stories260K.gguf - name: Save HF_HOME for ${{ env.GGUF_REPO }} if: always() && steps.prime-hf.outcome == 'success' uses: actions/cache/save@27d5ce7f107fe9357f9df03efb73ab90386fccae # v5.0.5 with: path: hf-cache - key: ${{ runner.os }}-hf-${{ env.GGUF_REPO }}-${{ env.GGUF_VARIANT }}-v1 + key: ${{ runner.os }}-hf-${{ env.GGUF_REPO }}-${{ env.GGUF_VARIANT }}-v2 - name: Install Studio (--local, --no-torch) env: GH_TOKEN: ${{ secrets.GITHUB_TOKEN }} + HF_TOKEN: ${{ secrets.HF_TOKEN }} run: | mkdir -p logs set -o pipefail diff --git a/.github/workflows/studio-mac-update-smoke.yml b/.github/workflows/studio-mac-update-smoke.yml index 1eca227096..f554a16415 100644 --- a/.github/workflows/studio-mac-update-smoke.yml +++ b/.github/workflows/studio-mac-update-smoke.yml @@ -62,6 +62,7 @@ jobs: - name: Install Studio (--local, --no-torch) env: GH_TOKEN: ${{ secrets.GITHUB_TOKEN }} + HF_TOKEN: ${{ secrets.HF_TOKEN }} run: | mkdir -p logs set -o pipefail @@ -73,6 +74,7 @@ jobs: - name: First update should be a no-op (prebuilt already validated) env: GH_TOKEN: ${{ secrets.GITHUB_TOKEN }} + HF_TOKEN: ${{ secrets.HF_TOKEN }} run: | set -o pipefail unsloth studio update --local 2>&1 | tee logs/update.log @@ -91,6 +93,7 @@ jobs: - name: Second update must also be a no-op env: GH_TOKEN: ${{ secrets.GITHUB_TOKEN }} + HF_TOKEN: ${{ secrets.HF_TOKEN }} run: | set -o pipefail unsloth studio update --local 2>&1 | tee logs/update2.log diff --git a/.github/workflows/studio-ui-smoke.yml b/.github/workflows/studio-ui-smoke.yml index 455fe4b7e1..de106e201f 100644 --- a/.github/workflows/studio-ui-smoke.yml +++ b/.github/workflows/studio-ui-smoke.yml @@ -76,7 +76,7 @@ jobs: continue-on-error: true with: path: hf-cache - key: ${{ runner.os }}-hf-${{ env.GGUF_REPO }}-${{ env.GGUF_VARIANT }}-v1 + key: ${{ runner.os }}-hf-${{ env.GGUF_REPO }}-${{ env.GGUF_VARIANT }}-v2 - name: Prime HF_HOME with the GGUF id: prime-hf @@ -87,17 +87,19 @@ jobs: python -m pip install --upgrade huggingface_hub mkdir -p hf-cache bash .github/scripts/hf-download-with-retry.sh "$GGUF_REPO" "$GGUF_FILE" + bash .github/scripts/hf-download-with-retry.sh ggml-org/models tinyllamas/stories260K.gguf - name: Save HF_HOME for ${{ env.GGUF_REPO }} if: always() && steps.prime-hf.outcome == 'success' uses: actions/cache/save@27d5ce7f107fe9357f9df03efb73ab90386fccae # v5.0.5 with: path: hf-cache - key: ${{ runner.os }}-hf-${{ env.GGUF_REPO }}-${{ env.GGUF_VARIANT }}-v1 + key: ${{ runner.os }}-hf-${{ env.GGUF_REPO }}-${{ env.GGUF_VARIANT }}-v2 - name: Install Studio (--local, --no-torch) env: GH_TOKEN: ${{ secrets.GITHUB_TOKEN }} + HF_TOKEN: ${{ secrets.HF_TOKEN }} run: | mkdir -p logs set -o pipefail diff --git a/.github/workflows/studio-update-smoke.yml b/.github/workflows/studio-update-smoke.yml index 057aeacbd4..307bb51972 100644 --- a/.github/workflows/studio-update-smoke.yml +++ b/.github/workflows/studio-update-smoke.yml @@ -71,6 +71,7 @@ jobs: # prebuilt path falls back to source build. env: GH_TOKEN: ${{ secrets.GITHUB_TOKEN }} + HF_TOKEN: ${{ secrets.HF_TOKEN }} run: | mkdir -p logs set -o pipefail @@ -85,6 +86,7 @@ jobs: # idempotency regressed. env: GH_TOKEN: ${{ secrets.GITHUB_TOKEN }} + HF_TOKEN: ${{ secrets.HF_TOKEN }} run: | set -o pipefail unsloth studio update --local 2>&1 | tee logs/update.log @@ -107,6 +109,7 @@ jobs: # the first one. env: GH_TOKEN: ${{ secrets.GITHUB_TOKEN }} + HF_TOKEN: ${{ secrets.HF_TOKEN }} run: | set -o pipefail unsloth studio update --local 2>&1 | tee logs/update2.log diff --git a/.github/workflows/studio-windows-api-smoke.yml b/.github/workflows/studio-windows-api-smoke.yml index eee61516c7..78efe918ac 100644 --- a/.github/workflows/studio-windows-api-smoke.yml +++ b/.github/workflows/studio-windows-api-smoke.yml @@ -69,7 +69,7 @@ jobs: continue-on-error: true with: path: hf-cache - key: ${{ runner.os }}-hf-${{ env.GGUF_REPO }}-${{ env.GGUF_VARIANT }}-v1 + key: ${{ runner.os }}-hf-${{ env.GGUF_REPO }}-${{ env.GGUF_VARIANT }}-v2 - name: Prime HF_HOME with the GGUF id: prime-hf @@ -80,13 +80,14 @@ jobs: python -m pip install --upgrade huggingface_hub mkdir -p hf-cache bash .github/scripts/hf-download-with-retry.sh "$GGUF_REPO" "$GGUF_FILE" + bash .github/scripts/hf-download-with-retry.sh ggml-org/models tinyllamas/stories260K.gguf - name: Save HF_HOME for ${{ env.GGUF_REPO }} if: always() && steps.prime-hf.outcome == 'success' uses: actions/cache/save@27d5ce7f107fe9357f9df03efb73ab90386fccae # v5.0.5 with: path: hf-cache - key: ${{ runner.os }}-hf-${{ env.GGUF_REPO }}-${{ env.GGUF_VARIANT }}-v1 + key: ${{ runner.os }}-hf-${{ env.GGUF_REPO }}-${{ env.GGUF_VARIANT }}-v2 - name: Pre-install Windows tweaks (npm 11 + Defender exclusions) shell: pwsh @@ -123,6 +124,7 @@ jobs: shell: pwsh env: GH_TOKEN: ${{ secrets.GITHUB_TOKEN }} + HF_TOKEN: ${{ secrets.HF_TOKEN }} run: | New-Item -ItemType Directory -Force -Path logs | Out-Null # *>&1 captures Write-Host (Information stream) output; diff --git a/.github/workflows/studio-windows-inference-smoke.yml b/.github/workflows/studio-windows-inference-smoke.yml index 6148856016..a772a6d102 100644 --- a/.github/workflows/studio-windows-inference-smoke.yml +++ b/.github/workflows/studio-windows-inference-smoke.yml @@ -101,7 +101,7 @@ jobs: continue-on-error: true with: path: hf-cache - key: ${{ runner.os }}-hf-${{ env.GGUF_REPO }}-${{ env.GGUF_VARIANT }}-v1 + key: ${{ runner.os }}-hf-${{ env.GGUF_REPO }}-${{ env.GGUF_VARIANT }}-v2 - name: Prime HF_HOME with the GGUF id: prime-hf @@ -114,6 +114,7 @@ jobs: python -m pip install --upgrade huggingface_hub mkdir -p hf-cache bash .github/scripts/hf-download-with-retry.sh "$GGUF_REPO" "$GGUF_FILE" + bash .github/scripts/hf-download-with-retry.sh ggml-org/models tinyllamas/stories260K.gguf - name: Save HF_HOME cache for ${{ env.GGUF_REPO }} # Only write a fresh cache entry when we actually rebuilt the @@ -123,7 +124,7 @@ jobs: uses: actions/cache/save@27d5ce7f107fe9357f9df03efb73ab90386fccae # v5.0.5 with: path: hf-cache - key: ${{ runner.os }}-hf-${{ env.GGUF_REPO }}-${{ env.GGUF_VARIANT }}-v1 + key: ${{ runner.os }}-hf-${{ env.GGUF_REPO }}-${{ env.GGUF_VARIANT }}-v2 - name: Pre-install Windows tweaks (npm 11 + Defender exclusions) shell: pwsh @@ -160,6 +161,7 @@ jobs: shell: pwsh env: GH_TOKEN: ${{ secrets.GITHUB_TOKEN }} + HF_TOKEN: ${{ secrets.HF_TOKEN }} run: | New-Item -ItemType Directory -Force -Path logs | Out-Null # *>&1 captures Write-Host (Information stream) output; @@ -504,6 +506,7 @@ jobs: shell: pwsh env: GH_TOKEN: ${{ secrets.GITHUB_TOKEN }} + HF_TOKEN: ${{ secrets.HF_TOKEN }} run: | New-Item -ItemType Directory -Force -Path logs | Out-Null # *>&1 captures Write-Host (Information stream) output; @@ -879,7 +882,7 @@ jobs: continue-on-error: true with: path: hf-cache - key: ${{ runner.os }}-hf-${{ env.GGUF_REPO }}-${{ env.GGUF_VARIANT }}-${{ env.MMPROJ_FILE }}-v1 + key: ${{ runner.os }}-hf-${{ env.GGUF_REPO }}-${{ env.GGUF_VARIANT }}-${{ env.MMPROJ_FILE }}-v2 - name: Prime HF_HOME with the GGUF + mmproj id: prime-hf @@ -891,13 +894,14 @@ jobs: mkdir -p hf-cache bash .github/scripts/hf-download-with-retry.sh "$GGUF_REPO" "$GGUF_FILE" bash .github/scripts/hf-download-with-retry.sh "$GGUF_REPO" "$MMPROJ_FILE" + bash .github/scripts/hf-download-with-retry.sh ggml-org/models tinyllamas/stories260K.gguf - name: Save HF_HOME cache for ${{ env.GGUF_REPO }} (model + mmproj) if: always() && steps.prime-hf.outcome == 'success' uses: actions/cache/save@27d5ce7f107fe9357f9df03efb73ab90386fccae # v5.0.5 with: path: hf-cache - key: ${{ runner.os }}-hf-${{ env.GGUF_REPO }}-${{ env.GGUF_VARIANT }}-${{ env.MMPROJ_FILE }}-v1 + key: ${{ runner.os }}-hf-${{ env.GGUF_REPO }}-${{ env.GGUF_VARIANT }}-${{ env.MMPROJ_FILE }}-v2 - name: Pre-install Windows tweaks (npm 11 + Defender exclusions) shell: pwsh @@ -934,6 +938,7 @@ jobs: shell: pwsh env: GH_TOKEN: ${{ secrets.GITHUB_TOKEN }} + HF_TOKEN: ${{ secrets.HF_TOKEN }} run: | New-Item -ItemType Directory -Force -Path logs | Out-Null # *>&1 captures Write-Host (Information stream) output; diff --git a/.github/workflows/studio-windows-ui-smoke.yml b/.github/workflows/studio-windows-ui-smoke.yml index e3c8642122..40d8e530cd 100644 --- a/.github/workflows/studio-windows-ui-smoke.yml +++ b/.github/workflows/studio-windows-ui-smoke.yml @@ -85,7 +85,7 @@ jobs: continue-on-error: true with: path: hf-cache - key: ${{ runner.os }}-hf-${{ env.GGUF_REPO }}-${{ env.GGUF_VARIANT }}-v1 + key: ${{ runner.os }}-hf-${{ env.GGUF_REPO }}-${{ env.GGUF_VARIANT }}-v2 - name: Prime HF_HOME with the GGUF id: prime-hf @@ -96,13 +96,14 @@ jobs: python -m pip install --upgrade huggingface_hub mkdir -p hf-cache bash .github/scripts/hf-download-with-retry.sh "$GGUF_REPO" "$GGUF_FILE" + bash .github/scripts/hf-download-with-retry.sh ggml-org/models tinyllamas/stories260K.gguf - name: Save HF_HOME for ${{ env.GGUF_REPO }} if: always() && steps.prime-hf.outcome == 'success' uses: actions/cache/save@27d5ce7f107fe9357f9df03efb73ab90386fccae # v5.0.5 with: path: hf-cache - key: ${{ runner.os }}-hf-${{ env.GGUF_REPO }}-${{ env.GGUF_VARIANT }}-v1 + key: ${{ runner.os }}-hf-${{ env.GGUF_REPO }}-${{ env.GGUF_VARIANT }}-v2 - name: Pre-install Windows tweaks (npm 11 + Defender exclusions) shell: pwsh @@ -143,6 +144,7 @@ jobs: shell: pwsh env: GH_TOKEN: ${{ secrets.GITHUB_TOKEN }} + HF_TOKEN: ${{ secrets.HF_TOKEN }} run: | New-Item -ItemType Directory -Force -Path logs | Out-Null # *>&1 redirects ALL PowerShell streams (stdout, stderr, diff --git a/.github/workflows/studio-windows-update-smoke.yml b/.github/workflows/studio-windows-update-smoke.yml index aa3e35f052..4a4806cfb1 100644 --- a/.github/workflows/studio-windows-update-smoke.yml +++ b/.github/workflows/studio-windows-update-smoke.yml @@ -133,6 +133,7 @@ jobs: shell: pwsh env: GH_TOKEN: ${{ secrets.GITHUB_TOKEN }} + HF_TOKEN: ${{ secrets.HF_TOKEN }} run: | New-Item -ItemType Directory -Force -Path logs | Out-Null # *>&1 captures Write-Host (Information stream) output; @@ -179,6 +180,7 @@ jobs: - name: First update should be a no-op (prebuilt already validated) env: GH_TOKEN: ${{ secrets.GITHUB_TOKEN }} + HF_TOKEN: ${{ secrets.HF_TOKEN }} run: | set -o pipefail unsloth studio update --local 2>&1 | tee logs/update.log @@ -197,6 +199,7 @@ jobs: - name: Second update must also be a no-op env: GH_TOKEN: ${{ secrets.GITHUB_TOKEN }} + HF_TOKEN: ${{ secrets.HF_TOKEN }} run: | set -o pipefail unsloth studio update --local 2>&1 | tee logs/update2.log diff --git a/studio/install_llama_prebuilt.py b/studio/install_llama_prebuilt.py index 64ea255803..2b42a7a058 100644 --- a/studio/install_llama_prebuilt.py +++ b/studio/install_llama_prebuilt.py @@ -177,6 +177,7 @@ VALIDATION_MODEL_CACHE_FILENAME = "stories260K.gguf" INSTALL_LOCK_TIMEOUT_SECONDS = 300 INSTALL_STAGING_ROOT_NAME = ".staging" GITHUB_AUTH_HOSTS = {"api.github.com", "github.com"} +HF_AUTH_HOSTS = {"huggingface.co", "www.huggingface.co"} RETRYABLE_HTTP_STATUS = {408, 429, 500, 502, 503, 504} HTTP_FETCH_ATTEMPTS = 4 HTTP_FETCH_BASE_DELAY_SECONDS = 0.75 @@ -484,6 +485,10 @@ def should_send_github_auth(url: str | None) -> bool: return parsed_hostname(url) in GITHUB_AUTH_HOSTS +def should_send_hf_auth(url: str | None) -> bool: + return parsed_hostname(url) in HF_AUTH_HOSTS + + def auth_headers(url: str | None = None) -> dict[str, str]: headers = { "User-Agent": "unsloth-studio-llama-prebuilt", @@ -491,9 +496,35 @@ def auth_headers(url: str | None = None) -> dict[str, str]: token = os.environ.get("GH_TOKEN") or os.environ.get("GITHUB_TOKEN") if token and should_send_github_auth(url): headers["Authorization"] = f"Bearer {token}" + return headers + # Anonymous huggingface.co fetches share a per-IP rate limit that CI + # fleets exhaust (HTTP 429), sinking the prebuilt path into a source + # build. Authenticate when a token is available. + hf_token = os.environ.get("HF_TOKEN") or os.environ.get("HUGGING_FACE_HUB_TOKEN") + if hf_token and should_send_hf_auth(url): + headers["Authorization"] = f"Bearer {hf_token}" return headers +class _CrossHostAuthStrippingRedirectHandler(urllib.request.HTTPRedirectHandler): + """Drop Authorization when a redirect leaves the original host. + + huggingface.co redirects file downloads to CDN hosts whose signed URLs + can reject foreign Authorization headers; urllib forwards headers to + redirect targets by default (requests/huggingface_hub strip them). + """ + + def redirect_request(self, req, fp, code, msg, headers, newurl): + new_request = super().redirect_request(req, fp, code, msg, headers, newurl) + if new_request is not None and parsed_hostname(newurl) != parsed_hostname(req.full_url): + new_request.headers.pop("Authorization", None) + new_request.unredirected_hdrs.pop("Authorization", None) + return new_request + + +_URL_OPENER = urllib.request.build_opener(_CrossHostAuthStrippingRedirectHandler()) + + def github_api_headers(url: str | None = None) -> dict[str, str]: return { "Accept": "application/vnd.github+json", @@ -936,7 +967,7 @@ def download_bytes( for attempt in range(1, attempts + 1): try: request = urllib.request.Request(url, headers = headers or auth_headers(url)) - with urllib.request.urlopen(request, timeout = timeout) as response: + with _URL_OPENER.open(request, timeout = timeout) as response: total_bytes: int | None = None content_length = response.headers.get("Content-Length") if content_length and content_length.isdigit(): @@ -1015,7 +1046,7 @@ def download_file(url: str, destination: Path) -> None: delete = False, ) as handle: tmp_path = Path(handle.name) - with urllib.request.urlopen(request, timeout = 120) as response: + with _URL_OPENER.open(request, timeout = 120) as response: total_bytes: int | None = None content_length = response.headers.get("Content-Length") if content_length and content_length.isdigit(): diff --git a/tests/studio/install/test_hf_auth.py b/tests/studio/install/test_hf_auth.py new file mode 100644 index 0000000000..7c3296d4fd --- /dev/null +++ b/tests/studio/install/test_hf_auth.py @@ -0,0 +1,114 @@ +"""Tests for Hugging Face auth on the llama.cpp prebuilt installer's fetches. + +Anonymous huggingface.co downloads (tiny GGUF validation model) share a +per-IP rate limit that CI fleets exhaust (HTTP 429), forcing the prebuilt +path into a source build. auth_headers now sends HF_TOKEN to huggingface.co +hosts, and a redirect handler strips Authorization when the download is +redirected to a different host (CDN signed URLs). All tests are offline. +""" + +import importlib.util +import sys +import urllib.request +from pathlib import Path +from unittest.mock import MagicMock, patch + +import pytest + + +PACKAGE_ROOT = Path(__file__).resolve().parents[3] + +_MODULE_PATH = PACKAGE_ROOT / "studio" / "install_llama_prebuilt.py" +_SPEC = importlib.util.spec_from_file_location( + "studio_install_llama_prebuilt_hf_auth", _MODULE_PATH +) +assert _SPEC is not None and _SPEC.loader is not None +mod = importlib.util.module_from_spec(_SPEC) +sys.modules[_SPEC.name] = mod +_SPEC.loader.exec_module(mod) + +_TOKEN_VARS = ("GH_TOKEN", "GITHUB_TOKEN", "HF_TOKEN", "HUGGING_FACE_HUB_TOKEN") +HF_URL = "https://huggingface.co/ggml-org/models/resolve/main/tinyllamas/stories260K.gguf" +GH_URL = "https://api.github.com/repos/unslothai/llama.cpp/releases" + + +def _headers(url, env): + """auth_headers under a fully controlled token environment.""" + with patch.dict(mod.os.environ, env, clear = False): + for var in _TOKEN_VARS: + if var not in env: + mod.os.environ.pop(var, None) + return mod.auth_headers(url) + + +class TestAuthHeaderRouting: + def test_hf_token_sent_to_huggingface(self): + headers = _headers(HF_URL, {"HF_TOKEN": "hf_x"}) + assert headers.get("Authorization") == "Bearer hf_x" + + def test_hub_token_fallback(self): + headers = _headers(HF_URL, {"HUGGING_FACE_HUB_TOKEN": "hf_y"}) + assert headers.get("Authorization") == "Bearer hf_y" + + def test_hf_token_not_sent_to_github(self): + headers = _headers(GH_URL, {"HF_TOKEN": "hf_x"}) + assert "Authorization" not in headers + + def test_hf_token_not_sent_to_other_hosts(self): + headers = _headers("https://cdn-lfs.huggingface.co/x", {"HF_TOKEN": "hf_x"}) + assert "Authorization" not in headers + + def test_gh_token_not_sent_to_huggingface(self): + headers = _headers(HF_URL, {"GH_TOKEN": "gh_x"}) + assert "Authorization" not in headers + + def test_gh_token_still_wins_on_github(self): + headers = _headers(GH_URL, {"GH_TOKEN": "gh_x", "HF_TOKEN": "hf_x"}) + assert headers.get("Authorization") == "Bearer gh_x" + + def test_no_tokens_no_auth(self): + assert "Authorization" not in _headers(HF_URL, {}) + + def test_validation_model_url_is_hf(self): + assert mod.should_send_hf_auth(mod.TEST_MODEL_URL) is True + + +class TestCrossHostRedirectStripsAuth: + def _redirect(self, newurl): + req = urllib.request.Request(HF_URL, headers = {"Authorization": "Bearer hf_x"}) + handler = mod._CrossHostAuthStrippingRedirectHandler() + return handler.redirect_request(req, None, 302, "Found", {}, newurl) + + def test_cross_host_redirect_drops_authorization(self): + new_request = self._redirect("https://cdn-lfs.huggingface.co/signed/blob") + assert new_request is not None + assert "Authorization" not in new_request.headers + assert "Authorization" not in new_request.unredirected_hdrs + + def test_same_host_redirect_keeps_authorization(self): + new_request = self._redirect("https://huggingface.co/elsewhere/blob") + assert new_request is not None + assert new_request.headers.get("Authorization") == "Bearer hf_x" + + +class TestDownloadBytesWiring: + def test_download_bytes_sends_hf_auth(self): + response = MagicMock() + response.__enter__ = lambda s: s + response.__exit__ = lambda s, *a: False + response.headers.get.return_value = None + response.read.side_effect = [b"data", b""] + with ( + patch.object(mod._URL_OPENER, "open", return_value = response) as opened, + patch.dict(mod.os.environ, {"HF_TOKEN": "hf_x"}, clear = False), + ): + for var in ("GH_TOKEN", "GITHUB_TOKEN"): + mod.os.environ.pop(var, None) + data = mod.download_bytes(HF_URL) + assert data == b"data" + request = opened.call_args.args[0] + assert request.headers.get("Authorization") == "Bearer hf_x" + + +if __name__ == "__main__": + sys.exit(pytest.main([__file__, "-q"])) From de8d9b0691ec5ea776292ad2132060ae5ec41f18 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?=E0=A4=B9=E0=A4=BF=E0=A4=AE=E0=A4=BE=E0=A4=82=E0=A4=B6?= =?UTF-8?q?=E0=A5=81?= Date: Thu, 11 Jun 2026 19:33:47 +0530 Subject: [PATCH 25/42] center the search dialog and change the wrong borders (#6184) * center the search dialog and change the wrong borders * fix the mistake of 1 to l * Fix/adjust search dialog radius for PR #6184 --------- Co-authored-by: Lee Jackson <130007945+Imagineer99@users.noreply.github.com> Co-authored-by: wasimysaid --- .../src/features/chat/components/chat-search-dialog.tsx | 4 ++-- studio/frontend/src/index.css | 3 +-- 2 files changed, 3 insertions(+), 4 deletions(-) diff --git a/studio/frontend/src/features/chat/components/chat-search-dialog.tsx b/studio/frontend/src/features/chat/components/chat-search-dialog.tsx index 0a6ec0c884..b0fd3c5f27 100644 --- a/studio/frontend/src/features/chat/components/chat-search-dialog.tsx +++ b/studio/frontend/src/features/chat/components/chat-search-dialog.tsx @@ -64,10 +64,10 @@ export function ChatSearchDialog() { - +
Date: Thu, 11 Jun 2026 17:07:26 +0300 Subject: [PATCH 26/42] fix(studio): reuse venv Python in setup instead of re-probing system (#6033) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit * fix(studio): reuse venv Python in setup instead of re-probing system * Reuse venv Python for studio setup Pass the venv interpreter from install.ps1 to studio/setup.ps1 via UNSLOTH_SETUP_PYTHON and prefer it over probing the system. Added Resolve-ReusedSetupPython to accept the handed-off path (or derive the venv python when setup runs standalone), validate it (Python 3.11–3.13 and non-conda), and inject its Scripts dir onto PATH. When a reused interpreter is accepted, py.exe enumeration and further system probing are skipped. install.ps1 also sets the env var before running setup and removes it on cleanup to avoid leaving state behind. This prevents setup from being tripped by unsupported Python 3.14 or Windows Store stubs on PATH. * Harden setup Python detection for PR #6033: py -All, shared conda check, bare ~ guard --------- Co-authored-by: Daniel Han Co-authored-by: Lee Jackson <130007945+Imagineer99@users.noreply.github.com> --- install.ps1 | 17 +++++-- studio/setup.ps1 | 121 ++++++++++++++++++++++++++++++----------------- 2 files changed, 90 insertions(+), 48 deletions(-) diff --git a/install.ps1 b/install.ps1 index 2938d4334f..fd7d3f4a81 100644 --- a/install.ps1 +++ b/install.ps1 @@ -983,10 +983,13 @@ shell.Run cmd, 0, False function Find-CompatiblePython { # Try the Python Launcher first (most reliable on Windows) # py.exe resolves to the standard CPython install, not conda. - $pyLauncher = Get-Command py -CommandType Application -ErrorAction SilentlyContinue - if ($pyLauncher -and $pyLauncher.Source -notmatch $script:CondaSkipPattern) { - # Prefer the requested $PythonVersion, then newest-first fallback. - $minors = @($PythonVersion) + (@("3.13", "3.12", "3.11") | Where-Object { $_ -ne $PythonVersion }) + # Prefer the requested $PythonVersion, then newest-first fallback. + $minors = @($PythonVersion) + (@("3.13", "3.12", "3.11") | Where-Object { $_ -ne $PythonVersion }) + # Enumerate every py.exe on PATH with -All (Windows PowerShell 5.1 + # returns only the first launcher without it) and search each for a + # supported, non-conda interpreter. + foreach ($pyLauncher in @(Get-Command py -All -CommandType Application -ErrorAction SilentlyContinue)) { + if ($pyLauncher.Source -match $script:CondaSkipPattern) { continue } foreach ($minor in $minors) { try { $out = & $pyLauncher.Source "-$minor" --version 2>&1 | Out-String @@ -2089,6 +2092,11 @@ shell.Run cmd, 0, False $studioArgs = @('studio', 'setup') if ($script:UnslothVerbose) { $studioArgs += '--verbose' } $env:UNSLOTH_INSTALL_ROLLBACK_MANAGED = "1" + # Hand the venv interpreter to setup.ps1 so it reuses the Python we already + # resolved and built the venv with, instead of re-probing the system (which + # can trip over an unsupported `python` 3.14 or a Store stub on PATH even + # though the venv is fine). setup.ps1 Test-Path-guards this before use. + $env:UNSLOTH_SETUP_PYTHON = Join-Path $VenvDir "Scripts\python.exe" try { & $UnslothExe @studioArgs $setupExit = $LASTEXITCODE @@ -2099,6 +2107,7 @@ shell.Run cmd, 0, False Remove-Item Env:UNSLOTH_STUDIO_HOME -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 diff --git a/studio/setup.ps1 b/studio/setup.ps1 index 29350e3932..94817b561a 100644 --- a/studio/setup.ps1 +++ b/studio/setup.ps1 @@ -1552,10 +1552,48 @@ if ($IsPipInstall) { } } -# 1g. Python (>= 3.11 and < 3.14). Prefer the Studio venv that install.ps1 -# just created, then py.exe so a 3.14 ahead of 3.13 on PATH does not trip the gate. +# Conda CPython ships modified DLL search paths that break torch's c10.dll +# loading on Windows; a venv made from conda Python inherits its base_prefix, +# so check the executable path AND sys.base_prefix. +$CondaSkipPattern = '(?i)(conda|miniconda|anaconda|miniforge|mambaforge)' +function Test-IsConda { + param([string]$Exe) + if ($Exe -match $CondaSkipPattern) { return $true } + try { + $basePrefix = (& $Exe -c "import sys; print(sys.base_prefix)" 2>$null | Out-String).Trim() + if ($basePrefix -match $CondaSkipPattern) { return $true } + } catch { } + return $false +} + +# 1g. Python (>= 3.11 and < 3.14). Prefer the interpreter install.ps1 already +# resolved and built the venv with (UNSLOTH_SETUP_PYTHON), or the existing +# venv python, before re-probing a system where a 3.14 or a WindowsApps stub +# ahead on PATH would trip the gate. setup.ps1 only updates packages in that +# venv, so the handoff is safe to reuse once validated. +function Resolve-ReusedSetupPython { + if (-not [string]::IsNullOrWhiteSpace($env:UNSLOTH_SETUP_PYTHON) -and + (Test-Path -LiteralPath $env:UNSLOTH_SETUP_PYTHON)) { + return $env:UNSLOTH_SETUP_PYTHON + } + # Standalone `unsloth studio setup/update` (install.ps1 did not run): derive + # the venv python from the studio root, mirroring the resolver below. + $root = if (-not [string]::IsNullOrWhiteSpace($env:UNSLOTH_STUDIO_HOME)) { $env:UNSLOTH_STUDIO_HOME.Trim() } + elseif (-not [string]::IsNullOrWhiteSpace($env:STUDIO_HOME)) { $env:STUDIO_HOME.Trim() } + else { Join-Path $env:USERPROFILE ".unsloth\studio" } + if ($root -eq "~") { + # Join-Path with an empty child throws on Windows PowerShell 5.1. + $root = $env:USERPROFILE + } elseif ($root -like "~/*" -or $root -like "~\*") { + $root = Join-Path $env:USERPROFILE $root.Substring(1).TrimStart('/', '\') + } + $venvPy = Join-Path $root "unsloth_studio\Scripts\python.exe" + if (Test-Path -LiteralPath $venvPy) { return $venvPy } + return $null +} +$ReusedSetupPython = Resolve-ReusedSetupPython + $HasPython = $null -ne (Get-Command python -ErrorAction SilentlyContinue) -$PyLauncher = Get-Command py -CommandType Application -ErrorAction SilentlyContinue $PythonOk = $false $DetectedPyVer = $null @@ -1584,28 +1622,24 @@ function Add-PythonDirToProcessPath { } catch { } } -$_prereqStudioHome = $null -if (-not [string]::IsNullOrWhiteSpace($env:UNSLOTH_STUDIO_HOME)) { - $_prereqStudioHome = $env:UNSLOTH_STUDIO_HOME.Trim() -} elseif (-not [string]::IsNullOrWhiteSpace($env:STUDIO_HOME)) { - $_prereqStudioHome = $env:STUDIO_HOME.Trim() -} else { - $_prereqStudioHome = Join-Path $env:USERPROFILE ".unsloth\studio" -} -if ($_prereqStudioHome -eq "~" -or $_prereqStudioHome -like "~/*" -or $_prereqStudioHome -like "~\*") { - $_prereqStudioHome = (Join-Path $env:USERPROFILE $_prereqStudioHome.Substring(1).TrimStart('/','\')) -} -$_prereqVenvPython = Join-Path $_prereqStudioHome "unsloth_studio\Scripts\python.exe" -if (Test-Path -LiteralPath $_prereqVenvPython) { - $_venvPyVer = Get-CompatiblePythonVersion $_prereqVenvPython - if ($_venvPyVer) { - $DetectedPyVer = $_venvPyVer - Add-PythonDirToProcessPath $_prereqVenvPython +# Reuse the install.ps1 / venv interpreter before any system probe. +if ($ReusedSetupPython) { + $_reusedVer = Get-CompatiblePythonVersion $ReusedSetupPython + if ($_reusedVer -and -not (Test-IsConda $ReusedSetupPython)) { + $DetectedPyVer = $_reusedVer + Add-PythonDirToProcessPath $ReusedSetupPython $PythonOk = $true } } -if (-not $PythonOk -and $PyLauncher) { +# Fall back to every py.exe on PATH (all-users and per-user launchers can both +# register). -All is required: Windows PowerShell 5.1 returns only the first +# launcher without it, and the PowerShell 7 multi-match array breaks the call +# operator if used directly. +$PyLaunchers = if ($PythonOk) { @() } else { @(Get-Command py -All -CommandType Application -ErrorAction SilentlyContinue) } + +foreach ($PyLauncher in $PyLaunchers) { + if ($PyLauncher.Source -match $CondaSkipPattern) { continue } foreach ($minor in @("3.13", "3.12", "3.11")) { try { $out = & $PyLauncher.Source "-$minor" --version 2>&1 | Out-String @@ -1625,6 +1659,7 @@ if (-not $PythonOk -and $PyLauncher) { } } catch { } } + if ($PythonOk) { break } } if (-not $PythonOk -and $HasPython) { @@ -1857,36 +1892,33 @@ if (Test-Path $OxcValidatorDir) { Write-Host "" substep "setting up Python environment..." -# Find Python -- skip Anaconda/Miniconda distributions. -# Conda-bundled CPython ships modified DLL search paths that break -# torch's c10.dll loading on Windows. Standalone CPython (python.org, -# winget, uv) does not have this issue. -# Uses Get-Command -All to look past conda entries that shadow a valid -# standalone Python further down PATH, and probes py.exe (the Python -# Launcher) which reliably finds python.org installs. -# -# NOTE: A venv created from conda Python inherits conda's base_prefix -# even though the venv path itself does not contain "conda". We check -# both the executable path AND sys.base_prefix to catch this case. -$CondaSkipPattern = '(?i)(conda|miniconda|anaconda|miniforge|mambaforge)' +# Find Python -- skip Anaconda/Miniconda distributions ($CondaSkipPattern and +# Test-IsConda are defined above the 1g gate). Standalone CPython (python.org, +# winget, uv) does not have conda's torch c10.dll loading issue. $PythonCmd = $null -# Helper: check if a Python executable is conda-based by inspecting -# both the path and sys.base_prefix (catches venvs created from conda). -function Test-IsConda { - param([string]$Exe) - if ($Exe -match $CondaSkipPattern) { return $true } +# 0. Reuse the interpreter install.ps1 already resolved and built the venv with +# (UNSLOTH_SETUP_PYTHON, or the existing venv python) before probing the +# system -- it is already validated as supported and non-conda. +if ($ReusedSetupPython) { try { - $basePrefix = (& $Exe -c "import sys; print(sys.base_prefix)" 2>$null | Out-String).Trim() - if ($basePrefix -match $CondaSkipPattern) { return $true } + $out = & $ReusedSetupPython --version 2>&1 | Out-String + if ($out -match 'Python 3\.(\d+)') { + $pyMinor = [int]$Matches[1] + if ($pyMinor -ge 11 -and $pyMinor -le 13 -and -not (Test-IsConda $ReusedSetupPython)) { + $PythonCmd = $ReusedSetupPython + } + } } catch { } - return $false } # 1. Try the Python Launcher (py.exe) first -- most reliable on Windows. -# py.exe is installed by python.org and resolves to standalone CPython. -$pyLauncher = Get-Command py -CommandType Application -ErrorAction SilentlyContinue -if ($pyLauncher -and $pyLauncher.Source -notmatch $CondaSkipPattern) { +# Enumerate every launcher with -All (Windows PowerShell 5.1 returns only +# the first match without it) and search each for a supported, non-conda +# interpreter. +$PyLaunchersResolve = if ($PythonCmd) { @() } else { @(Get-Command py -All -CommandType Application -ErrorAction SilentlyContinue) } +foreach ($pyLauncher in $PyLaunchersResolve) { + if ($pyLauncher.Source -match $CondaSkipPattern) { continue } foreach ($minor in @("3.13", "3.12", "3.11")) { try { $out = & $pyLauncher.Source "-$minor" --version 2>&1 | Out-String @@ -1904,6 +1936,7 @@ if ($pyLauncher -and $pyLauncher.Source -notmatch $CondaSkipPattern) { } } catch { } } + if ($PythonCmd) { break } } # 2. Fall back to scanning python3.x / python3 / python on PATH. From 5a5ecc8fc0ed3c0199b0ac39ea57b65d944096fb Mon Sep 17 00:00:00 2001 From: Prathamesh Jadhav <55660103+lollinng@users.noreply.github.com> Date: Thu, 11 Jun 2026 19:38:34 +0530 Subject: [PATCH 27/42] Strip trailing whitespace in construct_chat_template (fixes #992) (#6008) * Strip trailing whitespace in construct_chat_template (fixes #992) construct_chat_template() only lstrip()s the template, so a template ending in a newline (e.g. the default Llama-3 template + a trailing '\n') keeps the trailing whitespace. That breaks the rfind-based repeated-example detection and, on current main, also the regex fallback, which then raises "Could not recover a two-example structure from chat_template". Surrounding whitespace is not significant to the template here (the left side is already stripped), so strip() both ends. Verified that the valid no-newline template still parses via the primary path and that the meaningful trailing token (e.g. <|eot_id|>) is preserved. Co-Authored-By: Claude Opus 4.8 (1M context) * Fix trailing whitespace handling via parse retry for PR #6008 --------- Co-authored-by: Claude Opus 4.8 (1M context) Co-authored-by: Daniel Han --- unsloth/chat_templates.py | 18 +++++++++++++++++- 1 file changed, 17 insertions(+), 1 deletion(-) diff --git a/unsloth/chat_templates.py b/unsloth/chat_templates.py index b96f8cb7a0..fe078c5025 100644 --- a/unsloth/chat_templates.py +++ b/unsloth/chat_templates.py @@ -2341,13 +2341,15 @@ extra_eos_tokens = None, You must use {INPUT}, {OUTPUT} twice, and {SYSTEM} is optional. """ - # Strip only the left + # Strip only the left: trailing whitespace can be part of the repeated example + # (e.g. "{OUTPUT}\n"). Accidental trailing whitespace (#992) is retried on failure. chat_template = chat_template.lstrip() assert(tokenizer is not None) if extra_eos_tokens is None: extra_eos_tokens = [] elif type(extra_eos_tokens) is str: extra_eos_tokens = [extra_eos_tokens,] + original_extra_eos_tokens = list(extra_eos_tokens) vocab = tokenizer.get_vocab() for extra_eos in extra_eos_tokens: @@ -2454,6 +2456,20 @@ extra_eos_tokens = None, f"{left_changed}" ) except: + # Accidental trailing whitespace (#992) desyncs the two-example detection, + # so retry once without it. Templates that parse as-is are never altered. + rstripped_chat_template = chat_template.rstrip() + if rstripped_chat_template != chat_template: + try: + return construct_chat_template( + tokenizer = tokenizer, + chat_template = rstripped_chat_template, + default_system_message = default_system_message, + extra_eos_tokens = original_extra_eos_tokens, + ) + except Exception: + pass + output_pos = chat_template.find("{OUTPUT}") input_pos = chat_template.find("{INPUT}") if output_pos == -1 or input_pos == -1: From 1a99980b469c757091fe00b654ca78804b17a808 Mon Sep 17 00:00:00 2001 From: Daniel Han Date: Thu, 11 Jun 2026 07:10:08 -0700 Subject: [PATCH 28/42] Studio: auto Cloudflare tunnel for 0.0.0.0 launches (#6204) * Studio: auto Cloudflare tunnel for 0.0.0.0 launches Binding Studio to 0.0.0.0 for remote access often leaves the raw http://: URL unreachable (https-vs-http, blocked high ports, closed cloud security groups). On a wildcard bind, auto-start a free cloudflared quick tunnel and show its https://*.trycloudflare.com URL in the startup banner: Secure link access via Cloudflare: https://.trycloudflare.com - new studio/backend/cloudflare_tunnel.py: find or download+cache the cloudflared binary (per-OS/arch GitHub release, safe .tgz extract), start the tunnel, parse the URL, tear it down. Stdlib only; best-effort and non-fatal throughout (a missing binary or offline box never blocks or slows startup). - run_server starts the tunnel for 0.0.0.0 only (skips loopback, api-only and Colab), prints the line in the banner, and _graceful_shutdown stops the child so it never orphans. - --cloudflare/--no-cloudflare flag (default on) on `unsloth studio` and `unsloth studio run`, forwarded through the re-exec into run_server. - tests for the helper, the CLI flag forwarding, and the run.py defaults. * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * Studio cloudflare: send a User-Agent on the cloudflared download GitHub's CDN can 403 the default Python-urllib User-Agent on release asset downloads. Set an explicit UA and pin it with a test. * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * Studio cloudflare: address review (opt-out for subcommands, tunnel teardown) - reject --no-cloudflare placed before a subcommand (it would not reach the subcommand), mirroring the --parallel guard - register the tunnel before waiting for its URL so a shutdown during the wait stops cloudflared instead of orphaning it - tear the server + children down if `unsloth studio run` startup aborts (health timeout, model-load error, Ctrl+C) before the wait loop * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci --------- Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com> --- studio/backend/cloudflare_tunnel.py | 302 +++++++++++++ studio/backend/run.py | 52 +++ .../backend/tests/test_cloudflare_tunnel.py | 425 ++++++++++++++++++ studio/backend/utils/paths/storage_roots.py | 5 + unsloth_cli/commands/studio.py | 96 ++-- .../tests/test_studio_cloudflare_flag.py | 292 ++++++++++++ 6 files changed, 1147 insertions(+), 25 deletions(-) create mode 100644 studio/backend/cloudflare_tunnel.py create mode 100644 studio/backend/tests/test_cloudflare_tunnel.py create mode 100644 unsloth_cli/tests/test_studio_cloudflare_flag.py diff --git a/studio/backend/cloudflare_tunnel.py b/studio/backend/cloudflare_tunnel.py new file mode 100644 index 0000000000..7c088b6e1e --- /dev/null +++ b/studio/backend/cloudflare_tunnel.py @@ -0,0 +1,302 @@ +# SPDX-License-Identifier: AGPL-3.0-only +# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0 + +"""Free Cloudflare quick tunnel for Studio's 0.0.0.0 launches. + +The raw http://: is often unreachable (https-vs-http, blocked ports, +closed security groups); a cloudflared quick tunnel gives a free +https://*.trycloudflare.com URL that works anywhere, with no account or domain. + +Best-effort throughout: any failure collapses to "no URL" and Studio keeps +running. Stdlib only (back-end imports are lazy) so it is safe to import early. +""" + +from __future__ import annotations + +import os +import platform +import re +import shutil +import subprocess +import sys +import threading +from pathlib import Path +from typing import Optional, Tuple + +# cloudflared logs the quick-tunnel URL; match only the URL so we do not depend +# on the surrounding wording, which Cloudflare may change. +_URL_RE = re.compile(r"https://[A-Za-z0-9-]+\.trycloudflare\.com") + +_RELEASE_BASE = "https://github.com/cloudflare/cloudflared/releases/latest/download" + +_URL_TIMEOUT = 15.0 # seconds to wait for the public URL before giving up +_DOWNLOAD_TIMEOUT = 60 # urlopen timeout for the one-time binary download + + +def _windows_hidden_kwargs() -> dict: + """Suppress a child console window on Windows; no-op elsewhere.""" + if sys.platform != "win32": + return {} + flags = getattr(subprocess, "CREATE_NO_WINDOW", 0) + return {"creationflags": flags} if flags else {} + + +def _asset_name() -> Optional[Tuple[str, bool]]: + """(release asset filename, is_tgz) for this OS/arch, or None if unsupported.""" + system = platform.system().lower() + machine = platform.machine().lower() + is_x64 = machine in ("x86_64", "amd64", "x64") + is_arm64 = machine in ("aarch64", "arm64") + is_x86 = machine in ("i386", "i686", "x86") + if system == "linux": + if is_x64: + return ("cloudflared-linux-amd64", False) + if is_arm64: + return ("cloudflared-linux-arm64", False) + elif system == "darwin": + if is_arm64: + return ("cloudflared-darwin-arm64.tgz", True) + if is_x64: + return ("cloudflared-darwin-amd64.tgz", True) + elif system == "windows": + if is_x64: + return ("cloudflared-windows-amd64.exe", False) + if is_x86: + return ("cloudflared-windows-386.exe", False) + return None + + +def _cache_path() -> Optional[Path]: + """studio_bin_root()/cloudflared(.exe), or None if the studio home is unresolvable.""" + try: + from utils.paths.storage_roots import studio_bin_root # lazy: backend-only import + except Exception: + return None + name = "cloudflared.exe" if sys.platform == "win32" else "cloudflared" + return studio_bin_root() / name + + +def find_cloudflared() -> Optional[str]: + """Locate an existing cloudflared: PATH first, then the Studio bin cache.""" + on_path = shutil.which("cloudflared") + if on_path: + return on_path + cached = _cache_path() + if cached is not None and cached.is_file() and os.access(cached, os.X_OK): + return str(cached) + return None + + +def _download(url: str, dest: Path) -> bool: + """Download url to dest via urllib (temp file + atomic rename). Best-effort -> bool.""" + import tempfile + import urllib.request + + tmp_path: Optional[Path] = None + try: + dest.parent.mkdir(parents = True, exist_ok = True) + with tempfile.NamedTemporaryFile( + prefix = dest.name + ".tmp-", dir = dest.parent, delete = False + ) as handle: + tmp_path = Path(handle.name) + # GitHub's CDN 403s the default Python-urllib User-Agent. + req = urllib.request.Request(url, headers = {"User-Agent": "unsloth-studio"}) + with urllib.request.urlopen(req, timeout = _DOWNLOAD_TIMEOUT) as response: + shutil.copyfileobj(response, handle) + if tmp_path.stat().st_size == 0: + raise RuntimeError("empty download") + os.replace(tmp_path, dest) + return True + except Exception: + if tmp_path is not None: + try: + tmp_path.unlink(missing_ok = True) + except Exception: + pass + return False + + +def _extract_tgz_member(tgz_path: Path, dest: Path) -> bool: + """Extract just the `cloudflared` member from a darwin .tgz to dest. + + Rejects absolute paths and `..` traversal so a hostile archive cannot write + outside dest. Best-effort -> bool. + """ + import tarfile + try: + with tarfile.open(tgz_path, "r:gz") as tar: + member = None + for m in tar.getmembers(): + if not m.isfile() or os.path.basename(m.name) != "cloudflared": + continue + if m.name.startswith("/") or ".." in Path(m.name).parts: + continue + member = m + break + if member is None: + return False + src = tar.extractfile(member) + if src is None: + return False + with src, open(dest, "wb") as out: + shutil.copyfileobj(src, out) + return True + except Exception: + return False + + +def ensure_cloudflared() -> Optional[str]: + """Return a cloudflared path, downloading + caching the binary once if missing.""" + existing = find_cloudflared() + if existing: + return existing + asset = _asset_name() + cached = _cache_path() + if asset is None or cached is None: + return None + name, is_tgz = asset + url = f"{_RELEASE_BASE}/{name}" + try: + cached.parent.mkdir(parents = True, exist_ok = True) + if is_tgz: + tgz = cached.with_suffix(".tgz") + if not _download(url, tgz) or not _extract_tgz_member(tgz, cached): + tgz.unlink(missing_ok = True) + return None + tgz.unlink(missing_ok = True) + elif not _download(url, cached): + return None + if sys.platform != "win32": + os.chmod(cached, 0o755) + return str(cached) + except Exception: + return None + + +class CloudflareTunnel: + """A cloudflared quick tunnel to http://localhost:. Best-effort throughout. + + Use localhost (not the wildcard bind) as the tunnel origin so cloudflared's + upstream stays local-only. + """ + + def __init__(self, port: int, binary: str): + self.port = port + self.binary = binary + self._proc: Optional[subprocess.Popen] = None + self._lock = threading.Lock() + self._url_event = threading.Event() + self.url: Optional[str] = None + self.error: Optional[str] = None + + def start(self) -> None: + cmd = [ + self.binary, + "tunnel", + "--url", + f"http://localhost:{self.port}", + "--no-autoupdate", + ] + proc = subprocess.Popen( + cmd, + stdout = subprocess.PIPE, + stderr = subprocess.STDOUT, + stdin = subprocess.DEVNULL, + text = True, + errors = "replace", + bufsize = 1, + **_windows_hidden_kwargs(), + ) + with self._lock: + self._proc = proc + threading.Thread( + target = self._reader, args = (proc,), name = "cloudflared-reader", daemon = True + ).start() + + def _reader(self, proc: subprocess.Popen) -> None: + # Drain cloudflared's output, capture the first trycloudflare URL, and + # keep draining so it never blocks on a full pipe. + try: + if proc.stdout is not None: + for line in proc.stdout: + if self.url is None: + match = _URL_RE.search(line) + if match: + self.url = match.group(0) + self._url_event.set() + except Exception: + pass + finally: + if self.url is None: + self.error = "cloudflared exited before emitting a tunnel URL" + self._url_event.set() + + def wait_for_url(self, timeout: float = _URL_TIMEOUT) -> Optional[str]: + self._url_event.wait(timeout) + return self.url + + def stop(self) -> None: + """Terminate the tunnel. Idempotent and safe to call from a signal handler.""" + with self._lock: + proc, self._proc = self._proc, None + if proc is None: + return + try: + if proc.poll() is None: + proc.terminate() + try: + proc.wait(timeout = 5) + except subprocess.TimeoutExpired: + proc.kill() + try: + proc.wait(timeout = 5) + except Exception: + pass + except Exception: + pass + + +# Single serving process per Studio launch, so one module-level tunnel handle is +# enough; the lock guards the start/stop/shutdown races. +_active_tunnel: Optional[CloudflareTunnel] = None +_active_lock = threading.Lock() + + +def start_studio_tunnel(port: int, timeout: float = _URL_TIMEOUT) -> Optional[str]: + """Start a quick tunnel and return its public URL, or None (best-effort). + + On any failure (no binary, no URL within timeout, early crash) the tunnel is + stopped and None is returned, so the caller prints a hint and continues. + """ + global _active_tunnel + binary = ensure_cloudflared() + if not binary: + return None + tunnel = CloudflareTunnel(port, binary) + # Register before start/wait so a shutdown during the URL wait can stop it. + with _active_lock: + prior, _active_tunnel = _active_tunnel, tunnel + if prior is not None: + prior.stop() + try: + tunnel.start() + url = tunnel.wait_for_url(timeout) + except Exception: + url = None + if url: + return url + # No URL (or crash): drop it unless a concurrent shutdown already replaced it. + with _active_lock: + if _active_tunnel is tunnel: + _active_tunnel = None + tunnel.stop() + return None + + +def stop_studio_tunnel() -> None: + """Terminate the active tunnel, if any. Idempotent.""" + global _active_tunnel + with _active_lock: + tunnel, _active_tunnel = _active_tunnel, None + if tunnel is not None: + tunnel.stop() diff --git a/studio/backend/run.py b/studio/backend/run.py index 4619cb849c..32dfe06a18 100644 --- a/studio/backend/run.py +++ b/studio/backend/run.py @@ -414,9 +414,26 @@ def _emit_startup_output(host: str, port: int, display_host: str) -> None: print_studio_stop_hint() elif wildcard_bind: _verify_global_reachability(display_host, port) + _print_cloudflare_line() print_studio_stop_hint() +def _print_cloudflare_line() -> None: + """Print the Cloudflare quick-tunnel URL for 0.0.0.0 binds, if one is up. + + Reads the module-level URL set by ``run_server``. Prints nothing when the + tunnel is disabled or failed -- failures are silently ignored. + """ + if not _cloudflare_url: + return + from startup_banner import stdout_supports_color + + accent = "\033[38;5;150;1m" + reset = "\033[0m" + line = f" Secure link access via Cloudflare: {_cloudflare_url}" + print(f"{accent}{line}{reset}" if stdout_supports_color() else line) + + def _get_pid_on_port(port: int) -> "tuple[int, str] | None": """Return (pid, process_name) listening on *port*, or None. @@ -584,6 +601,13 @@ def _graceful_shutdown(server = None): except Exception as e: logger.warning("Error shutting down llama-server: %s", e) + # 6. Stop the Cloudflare tunnel (if started). + try: + from cloudflare_tunnel import stop_studio_tunnel + stop_studio_tunnel() + except Exception as e: + logger.warning("Error stopping Cloudflare tunnel: %s", e) + logger.info("All subprocesses cleaned up") @@ -594,6 +618,10 @@ _server = None # Shutdown event -- wakes the main loop on signal. _shutdown_event = None +# trycloudflare.com URL for 0.0.0.0 binds (set by run_server, read by the banner); +# None when there is no tunnel (loopback, disabled, or a silently-ignored failure). +_cloudflare_url = None + _DEFAULT_FRONTEND_PATH = Path(__file__).resolve().parent.parent / "frontend" / "dist" @@ -769,6 +797,7 @@ def run_server( silent: bool = False, api_only: bool = False, llama_parallel_slots: int = 1, + cloudflare: bool = True, ): """ Start the FastAPI server. @@ -973,6 +1002,21 @@ def run_server( if api_only: print(f"TAURI_PORT={port}", flush = True) + # Free trycloudflare.com tunnel for 0.0.0.0 binds (the raw ip:port is often + # unreachable). Started pre-banner and even when silent so the CLI banner can + # read app.state.cloudflare_url; torn down by _graceful_shutdown. + global _cloudflare_url + _cloudflare_url = None + app.state.cloudflare_url = None + _cloudflare_enabled = cloudflare and host == "0.0.0.0" and not api_only and not _IS_COLAB + if _cloudflare_enabled: + try: # best-effort: any failure must not block startup + from cloudflare_tunnel import start_studio_tunnel + _cloudflare_url = start_studio_tunnel(port) + app.state.cloudflare_url = _cloudflare_url + except Exception as e: + logger.debug("Cloudflare tunnel skipped: %s", e) + if not silent: _emit_startup_output(host, port, display_host) @@ -1011,6 +1055,13 @@ if __name__ == "__main__": action = "store_true", help = "API server only, no frontend (for Tauri)", ) + parser.add_argument( + "--cloudflare", + action = argparse.BooleanOptionalAction, + default = True, + help = "Auto-create a free Cloudflare HTTPS tunnel when bound to 0.0.0.0 " + "(default on; --no-cloudflare to disable)", + ) # Mirror unsloth_cli/commands/studio.py's _PARALLEL_*. Default 1 is for direct # backend launches; `unsloth studio run` always passes its own value (4). _PARALLEL_MIN = 1 @@ -1037,6 +1088,7 @@ if __name__ == "__main__": silent = args.silent, api_only = args.api_only, llama_parallel_slots = args.parallel, + cloudflare = args.cloudflare, ) if args.frontend is not None: kwargs["frontend_path"] = Path(args.frontend) diff --git a/studio/backend/tests/test_cloudflare_tunnel.py b/studio/backend/tests/test_cloudflare_tunnel.py new file mode 100644 index 0000000000..8208f7f83f --- /dev/null +++ b/studio/backend/tests/test_cloudflare_tunnel.py @@ -0,0 +1,425 @@ +# SPDX-License-Identifier: AGPL-3.0-only +# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0 + +"""Tests for the Cloudflare quick-tunnel helper and run.py wiring. + +cloudflare_tunnel.py is stdlib-only (storage_roots is imported lazily), so it +loads via spec_from_file_location without the studio venv. run.py defaults are +checked by AST so we never import its heavy deps (uvicorn/structlog). +""" + +import ast +import importlib.util +import io +import sys +import tarfile +from pathlib import Path + +import pytest + +_BACKEND = Path(__file__).resolve().parent.parent +_CT_PY = _BACKEND / "cloudflare_tunnel.py" +_RUN_PY = _BACKEND / "run.py" + + +def _load_ct(): + spec = importlib.util.spec_from_file_location("cloudflare_tunnel", _CT_PY) + mod = importlib.util.module_from_spec(spec) + spec.loader.exec_module(mod) + return mod + + +ct = _load_ct() + + +# ── URL parsing ────────────────────────────────────────────────────── + + +def test_url_regex_extracts_and_ignores_noise(): + blob = ( + "2026-06-11T10:00:00Z INF Thank you for trying Cloudflare Tunnel.\n" + "2026-06-11T10:00:01Z INF Requesting new quick Tunnel on trycloudflare.com...\n" + "2026-06-11T10:00:01Z INF | https://setting-democracy-gathering.trycloudflare.com |\n" + "2026-06-11T10:00:02Z INF Registered tunnel connection https://not-the-url.example.com\n" + ) + m = ct._URL_RE.search(blob) + assert m is not None + assert m.group(0) == "https://setting-democracy-gathering.trycloudflare.com" + + +def test_url_regex_no_match_on_unrelated(): + assert ct._URL_RE.search("INF connecting to https://api.cloudflare.com/v4") is None + + +# ── asset mapping ──────────────────────────────────────────────────── + + +@pytest.mark.parametrize( + "system,machine,expected", + [ + ("Linux", "x86_64", ("cloudflared-linux-amd64", False)), + ("Linux", "aarch64", ("cloudflared-linux-arm64", False)), + ("Darwin", "arm64", ("cloudflared-darwin-arm64.tgz", True)), + ("Darwin", "x86_64", ("cloudflared-darwin-amd64.tgz", True)), + ("Windows", "AMD64", ("cloudflared-windows-amd64.exe", False)), + ("Windows", "x86", ("cloudflared-windows-386.exe", False)), + ("Linux", "mips", None), + ("Plan9", "x86_64", None), + ], +) +def test_asset_name(monkeypatch, system, machine, expected): + monkeypatch.setattr(ct.platform, "system", lambda: system) + monkeypatch.setattr(ct.platform, "machine", lambda: machine) + assert ct._asset_name() == expected + + +# ── binary discovery ───────────────────────────────────────────────── + + +def test_find_cloudflared_prefers_path(monkeypatch): + monkeypatch.setattr(ct.shutil, "which", lambda name: "/usr/local/bin/cloudflared") + assert ct.find_cloudflared() == "/usr/local/bin/cloudflared" + + +def test_find_cloudflared_falls_back_to_cache(monkeypatch, tmp_path): + cached = tmp_path / "cloudflared" + cached.write_text("#!/bin/sh\n") + cached.chmod(0o755) + monkeypatch.setattr(ct.shutil, "which", lambda name: None) + monkeypatch.setattr(ct, "_cache_path", lambda: cached) + assert ct.find_cloudflared() == str(cached) + + +def test_find_cloudflared_none_when_missing(monkeypatch, tmp_path): + monkeypatch.setattr(ct.shutil, "which", lambda name: None) + monkeypatch.setattr(ct, "_cache_path", lambda: tmp_path / "absent") + assert ct.find_cloudflared() is None + + +# ── ensure / download ──────────────────────────────────────────────── + + +def test_ensure_downloads_and_chmods_when_missing(monkeypatch, tmp_path): + cached = tmp_path / "cloudflared" + monkeypatch.setattr(ct, "find_cloudflared", lambda: None) + monkeypatch.setattr(ct, "_asset_name", lambda: ("cloudflared-linux-amd64", False)) + monkeypatch.setattr(ct, "_cache_path", lambda: cached) + + def fake_download(url, dest): + assert url.endswith("/cloudflared-linux-amd64") + dest.write_bytes(b"ELF-ish") + return True + + monkeypatch.setattr(ct, "_download", fake_download) + monkeypatch.setattr(ct.sys, "platform", "linux") + path = ct.ensure_cloudflared() + assert path == str(cached) + assert cached.exists() + assert cached.stat().st_mode & 0o111 # executable bit set + + +def test_ensure_returns_none_on_download_failure(monkeypatch, tmp_path): + monkeypatch.setattr(ct, "find_cloudflared", lambda: None) + monkeypatch.setattr(ct, "_asset_name", lambda: ("cloudflared-linux-amd64", False)) + monkeypatch.setattr(ct, "_cache_path", lambda: tmp_path / "cloudflared") + monkeypatch.setattr(ct, "_download", lambda url, dest: False) + assert ct.ensure_cloudflared() is None + + +def test_ensure_returns_none_for_unsupported_arch(monkeypatch, tmp_path): + monkeypatch.setattr(ct, "find_cloudflared", lambda: None) + monkeypatch.setattr(ct, "_asset_name", lambda: None) + monkeypatch.setattr(ct, "_cache_path", lambda: tmp_path / "cloudflared") + assert ct.ensure_cloudflared() is None + + +def test_download_sets_user_agent(monkeypatch, tmp_path): + import urllib.request + + captured = {} + + class _Resp: + _sent = False + + def __enter__(self): + return self + + def __exit__(self, *a): + return False + + def read(self, n = -1): + if self._sent: + return b"" + self._sent = True + return b"payload" + + def fake_urlopen(req, timeout = None): + captured["ua"] = req.get_header("User-agent") + return _Resp() + + monkeypatch.setattr(urllib.request, "urlopen", fake_urlopen) + dest = tmp_path / "cloudflared" + assert ct._download("https://github.com/cloudflare/cloudflared/x", dest) is True + assert captured["ua"] == "unsloth-studio" # GitHub CDN 403s the default UA + assert dest.read_bytes() == b"payload" + + +# ── cross-platform: Windows (.exe), macOS (.tgz) ───────────────────── + + +def test_cache_path_uses_exe_on_windows(monkeypatch, tmp_path): + import types + + fake_sr = types.ModuleType("utils.paths.storage_roots") + fake_sr.studio_bin_root = lambda: tmp_path + monkeypatch.setitem(sys.modules, "utils.paths.storage_roots", fake_sr) + monkeypatch.setattr(ct.sys, "platform", "win32") + assert ct._cache_path() == tmp_path / "cloudflared.exe" + + +def test_ensure_windows_downloads_exe(monkeypatch, tmp_path): + cached = tmp_path / "cloudflared.exe" + monkeypatch.setattr(ct, "find_cloudflared", lambda: None) + monkeypatch.setattr(ct, "_asset_name", lambda: ("cloudflared-windows-amd64.exe", False)) + monkeypatch.setattr(ct, "_cache_path", lambda: cached) + monkeypatch.setattr(ct.sys, "platform", "win32") + + def fake_download(url, dest): + assert url.endswith("/cloudflared-windows-amd64.exe") + dest.write_bytes(b"MZ") # PE header magic + return True + + monkeypatch.setattr(ct, "_download", fake_download) + # chmod is skipped on Windows; would raise on a path that does not exist yet. + monkeypatch.setattr(ct.os, "chmod", lambda *a, **k: pytest.fail("chmod called on win32")) + assert ct.ensure_cloudflared() == str(cached) + assert cached.read_bytes() == b"MZ" + + +def test_ensure_macos_extracts_tgz_and_chmods(monkeypatch, tmp_path): + cached = tmp_path / "cloudflared" + monkeypatch.setattr(ct, "find_cloudflared", lambda: None) + monkeypatch.setattr(ct, "_asset_name", lambda: ("cloudflared-darwin-arm64.tgz", True)) + monkeypatch.setattr(ct, "_cache_path", lambda: cached) + monkeypatch.setattr(ct.sys, "platform", "darwin") + + def fake_download(url, dest): + # dest is cached.with_suffix(".tgz"); write a real archive there. + assert url.endswith("/cloudflared-darwin-arm64.tgz") + with tarfile.open(dest, "w:gz") as tar: + data = b"mach-o" + info = tarfile.TarInfo(name = "cloudflared") + info.size = len(data) + tar.addfile(info, io.BytesIO(data)) + return True + + monkeypatch.setattr(ct, "_download", fake_download) + path = ct.ensure_cloudflared() + assert path == str(cached) + assert cached.read_bytes() == b"mach-o" + assert cached.stat().st_mode & 0o111 # chmod applied on posix + assert not cached.with_suffix(".tgz").exists() # temp archive cleaned up + + +# ── .tgz extraction (darwin) ───────────────────────────────────────── + + +def _make_tgz( + tmp_path, + member_name, + data = b"bin", +): + tgz = tmp_path / "cf.tgz" + with tarfile.open(tgz, "w:gz") as tar: + info = tarfile.TarInfo(name = member_name) + info.size = len(data) + tar.addfile(info, io.BytesIO(data)) + return tgz + + +def test_tgz_extraction_extracts_clean_member(tmp_path): + tgz = _make_tgz(tmp_path, "cloudflared") + dest = tmp_path / "out" + assert ct._extract_tgz_member(tgz, dest) is True + assert dest.read_bytes() == b"bin" + + +def test_tgz_extraction_rejects_traversal(tmp_path): + tgz = _make_tgz(tmp_path, "../cloudflared") + dest = tmp_path / "out" + assert ct._extract_tgz_member(tgz, dest) is False + assert not dest.exists() + + +def test_tgz_extraction_missing_member(tmp_path): + tgz = _make_tgz(tmp_path, "README") + dest = tmp_path / "out" + assert ct._extract_tgz_member(tgz, dest) is False + + +# ── tunnel lifecycle ───────────────────────────────────────────────── + + +class _FakePopen: + def __init__(self): + self.terminated = False + self.killed = False + self._alive = True + + def poll(self): + return None if self._alive else 0 + + def terminate(self): + self.terminated = True + self._alive = False + + def wait(self, timeout = None): + if self._alive: + raise ct.subprocess.TimeoutExpired(cmd = "cloudflared", timeout = timeout) + return 0 + + def kill(self): + self.killed = True + self._alive = False + + +def test_stop_terminates_process(): + t = ct.CloudflareTunnel(8080, "/bin/cloudflared") + fake = _FakePopen() + t._proc = fake + t.stop() + assert fake.terminated is True + assert t._proc is None + # second stop is a no-op (idempotent) + t.stop() + + +def test_wait_for_url_times_out_without_blocking(): + t = ct.CloudflareTunnel(8080, "/bin/cloudflared") + assert t.wait_for_url(timeout = 0.05) is None + + +def test_start_studio_tunnel_no_binary(monkeypatch): + monkeypatch.setattr(ct, "ensure_cloudflared", lambda: None) + assert ct.start_studio_tunnel(8080) is None + + +def test_start_studio_tunnel_registers_before_wait(monkeypatch): + # The tunnel must be visible to stop_studio_tunnel() during the URL wait, + # else a shutdown in that window orphans cloudflared. + seen = {} + + class _Stub: + def __init__(self, port, binary): + self.url = None + + def start(self): + pass + + def wait_for_url(self, timeout): + seen["active_during_wait"] = ct._active_tunnel is self + self.url = "https://x.trycloudflare.com" + return self.url + + def stop(self): + seen["stopped"] = True + + monkeypatch.setattr(ct, "ensure_cloudflared", lambda: "/bin/cloudflared") + monkeypatch.setattr(ct, "CloudflareTunnel", _Stub) + try: + assert ct.start_studio_tunnel(8080) == "https://x.trycloudflare.com" + assert seen["active_during_wait"] is True + finally: + ct.stop_studio_tunnel() + + +def test_start_studio_tunnel_clears_and_stops_on_no_url(monkeypatch): + seen = {} + + class _Stub: + def __init__(self, port, binary): + self.url = None + + def start(self): + pass + + def wait_for_url(self, timeout): + return None + + def stop(self): + seen["stopped"] = True + + monkeypatch.setattr(ct, "ensure_cloudflared", lambda: "/bin/cloudflared") + monkeypatch.setattr(ct, "CloudflareTunnel", _Stub) + assert ct.start_studio_tunnel(8080) is None + assert seen.get("stopped") is True + assert ct._active_tunnel is None + + +def test_start_studio_tunnel_returns_url(monkeypatch): + class _StubTunnel: + def __init__(self, port, binary): + self.url = None + + def start(self): + self.url = "https://stub-xyz.trycloudflare.com" + + def wait_for_url(self, timeout): + return self.url + + def stop(self): + pass + + monkeypatch.setattr(ct, "ensure_cloudflared", lambda: "/bin/cloudflared") + monkeypatch.setattr(ct, "CloudflareTunnel", _StubTunnel) + try: + assert ct.start_studio_tunnel(8080) == "https://stub-xyz.trycloudflare.com" + finally: + ct.stop_studio_tunnel() + + +# ── run.py source-level pins (AST / source, no heavy import) ───────── + + +def _func_param_defaults(source, func_name): + tree = ast.parse(source) + for node in ast.walk(tree): + if isinstance(node, (ast.FunctionDef, ast.AsyncFunctionDef)) and node.name == func_name: + args = node.args.args + defaults = node.args.defaults + offset = len(args) - len(defaults) + out = {} + for i, d in enumerate(defaults): + if isinstance(d, ast.Constant): + out[args[offset + i].arg] = d.value + return out + return {} + + +def _argparse_default(source, option): + tree = ast.parse(source) + for node in ast.walk(tree): + if isinstance(node, ast.Call) and isinstance(node.func, ast.Attribute): + if node.func.attr == "add_argument" and node.args: + a0 = node.args[0] + if isinstance(a0, ast.Constant) and a0.value == option: + for kw in node.keywords: + if kw.arg == "default" and isinstance(kw.value, ast.Constant): + return kw.value.value + return None + + +def test_run_server_cloudflare_default_true(): + defaults = _func_param_defaults(_RUN_PY.read_text(), "run_server") + assert defaults.get("cloudflare") is True + + +def test_argparse_cloudflare_default_true(): + assert _argparse_default(_RUN_PY.read_text(), "--cloudflare") is True + + +def test_run_server_gates_tunnel_on_wildcard(): + # Guard against accidentally widening the trigger beyond 0.0.0.0. + source = _RUN_PY.read_text() + assert "_cloudflare_enabled" in source + assert 'host == "0.0.0.0"' in source diff --git a/studio/backend/utils/paths/storage_roots.py b/studio/backend/utils/paths/storage_roots.py index 10597a689b..2d4f5ac243 100644 --- a/studio/backend/utils/paths/storage_roots.py +++ b/studio/backend/utils/paths/storage_roots.py @@ -60,6 +60,11 @@ def cache_root() -> Path: return studio_root() / "cache" +def studio_bin_root() -> Path: + """Dir for Studio-managed executables (the `unsloth` shim, downloaded tools like cloudflared).""" + return studio_root() / "bin" + + def assets_root() -> Path: return studio_root() / "assets" diff --git a/unsloth_cli/commands/studio.py b/unsloth_cli/commands/studio.py index b506f764d7..7493b87eee 100644 --- a/unsloth_cli/commands/studio.py +++ b/unsloth_cli/commands/studio.py @@ -597,6 +597,11 @@ def studio_default( f"defaults to {_PARALLEL_DEFAULT_RUN}." ), ), + cloudflare: bool = typer.Option( + True, + "--cloudflare/--no-cloudflare", + help = "Auto-create a free Cloudflare HTTPS tunnel when bound to 0.0.0.0 (default on).", + ), ): """Launch the Unsloth Studio server.""" # Runs before every subcommand (run/setup/update/...). @@ -614,6 +619,16 @@ def studio_default( err = True, ) raise typer.Exit(2) + # Same for --no-cloudflare: it would not reach the subcommand. + if not cloudflare: + typer.echo( + f"Error: --no-cloudflare on `unsloth studio` applies to the " + f"plain-server path only. For `unsloth studio " + f"{ctx.invoked_subcommand}`, put it after the subcommand: " + f"`unsloth studio {ctx.invoked_subcommand} --no-cloudflare ...`", + err = True, + ) + raise typer.Exit(2) return # Use the studio venv if it exists and we aren't already in it. @@ -648,6 +663,8 @@ def studio_default( args.append("--silent") if api_only: args.append("--api-only") + # Forward the explicit polarity (matches run.py's BooleanOptionalAction). + args.append("--cloudflare" if cloudflare else "--no-cloudflare") # On Windows os.execvp keeps the parent alive, so Ctrl+C # would orphan the child; use Popen+wait instead. if sys.platform == "win32": @@ -689,6 +706,7 @@ def studio_default( silent = silent, api_only = api_only, llama_parallel_slots = parallel, + cloudflare = cloudflare, ) if frontend is not None: run_kwargs["frontend_path"] = frontend @@ -861,6 +879,11 @@ def run( f"{_PARALLEL_DEFAULT_RUN} (pre-PR hardcoded value)." ), ), + cloudflare: bool = typer.Option( + True, + "--cloudflare/--no-cloudflare", + help = "Auto-create a free Cloudflare HTTPS tunnel when bound to 0.0.0.0 (default on).", + ), ): """Start Studio, load a model, print an API key -- one-liner server. @@ -980,6 +1003,8 @@ def run( # Typer claims --parallel outside ctx.args; without this the # child reverts to its default and silently drops the value. args.extend(["--parallel", str(parallel)]) + # Forward the explicit polarity (same rationale as --load-in-4bit above). + args.append("--cloudflare" if cloudflare else "--no-cloudflare") # llama-server pass-through extras → child ctx.args → load payload. if extra_llama_args: args.extend(extra_llama_args) @@ -997,7 +1022,13 @@ def run( # ── 2. Start server (always suppress built-in banner) ───────────── from studio.backend.run import run_server, _resolve_external_ip - run_kwargs = dict(host = host, port = port, silent = True, llama_parallel_slots = parallel) + run_kwargs = dict( + host = host, + port = port, + silent = True, + llama_parallel_slots = parallel, + cloudflare = cloudflare, + ) if frontend is not None: run_kwargs["frontend_path"] = frontend app = run_server(**run_kwargs) @@ -1011,32 +1042,41 @@ def run( set_tool_policy(enable_tools) - # 3. Wait for server health. - if not silent: - typer.echo("Starting Unsloth Studio...") - if not _wait_for_server(actual_port): - typer.echo("Error: server did not become healthy within 30 seconds.", err = True) - raise typer.Exit(1) + # Steps 3-5 can abort (health timeout, model-load error, or Ctrl+C during the + # slow load); tear the server and its children (llama-server, cloudflared) down + # on any abort so they never orphan. + from studio.backend.run import _graceful_shutdown, _server - # 4. Create API key in-process. - api_key = _create_api_key_inprocess(api_key_name) - - # 5. Load model via HTTP. - if not silent: - typer.echo(f"Loading model: {model}...") try: - result = _load_model_via_http( - port = actual_port, - api_key = api_key, - model = model, - gguf_variant = gguf_variant, - max_seq_length = max_seq_length, - load_in_4bit = load_in_4bit, - llama_extra_args = extra_llama_args, - ) - except RuntimeError as exc: - typer.echo(f"Error: {exc}", err = True) - raise typer.Exit(1) + # 3. Wait for server health. + if not silent: + typer.echo("Starting Unsloth Studio...") + if not _wait_for_server(actual_port): + typer.echo("Error: server did not become healthy within 30 seconds.", err = True) + raise typer.Exit(1) + + # 4. Create API key in-process. + api_key = _create_api_key_inprocess(api_key_name) + + # 5. Load model via HTTP. + if not silent: + typer.echo(f"Loading model: {model}...") + try: + result = _load_model_via_http( + port = actual_port, + api_key = api_key, + model = model, + gguf_variant = gguf_variant, + max_seq_length = max_seq_length, + load_in_4bit = load_in_4bit, + llama_extra_args = extra_llama_args, + ) + except RuntimeError as exc: + typer.echo(f"Error: {exc}", err = True) + raise typer.Exit(1) + except BaseException: + _graceful_shutdown(_server) + raise loaded_model = result.get("model", model) display_variant = f" ({gguf_variant})" if gguf_variant else "" @@ -1045,6 +1085,8 @@ def run( display_host = _resolve_external_ip() if host == "0.0.0.0" else host base_url = f"http://{display_host}:{actual_port}" sdk_base_url = f"{base_url}/v1" + # run_server started the tunnel during the silent run above (0.0.0.0 only). + _cf_url = getattr(app.state, "cloudflare_url", None) # Orange so the tool-policy notice stands out; printed under # --silent / --yes too so the policy is never invisible. @@ -1074,6 +1116,8 @@ def run( typer.echo("") typer.echo("=" * 56) typer.echo(f" Unsloth Studio running at {base_url}") + if _cf_url: + typer.echo(f" Secure link access via Cloudflare: {_cf_url}") typer.echo(f" Model loaded: {loaded_model}{display_variant}") typer.echo(f" API Key: {api_key}") typer.echo("") @@ -1107,6 +1151,8 @@ def run( else: # Silent still prints URL + API key + tool-status policy. typer.echo(f"URL: {base_url}") + if _cf_url: + typer.echo(f"Secure link access via Cloudflare: {_cf_url}") typer.echo(f"API Key: {api_key}") typer.secho(_tool_notice, fg = _tool_notice_fg, bold = True) diff --git a/unsloth_cli/tests/test_studio_cloudflare_flag.py b/unsloth_cli/tests/test_studio_cloudflare_flag.py new file mode 100644 index 0000000000..5c2a039547 --- /dev/null +++ b/unsloth_cli/tests/test_studio_cloudflare_flag.py @@ -0,0 +1,292 @@ +# SPDX-License-Identifier: AGPL-3.0-only +# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0 + +"""Tests for the `--cloudflare/--no-cloudflare` Studio flag. + +Pins the typer Option (default on) on both `unsloth studio` and +`unsloth studio run`, and that the chosen polarity reaches the re-exec'd +child and run_server. Modeled on test_studio_run_parallel_flag.py. +""" + +from __future__ import annotations + +import sys +from pathlib import Path + +import pytest +from typer.testing import CliRunner + + +_REPO_ROOT = Path(__file__).resolve().parents[2] +if str(_REPO_ROOT) not in sys.path: + sys.path.insert(0, str(_REPO_ROOT)) + + +def _studio(): + from unsloth_cli.commands import studio as _studio_mod + return _studio_mod + + +_BASE = ["--model", "unsloth/Qwen3-1.7B-GGUF"] + + +# ── option registration ────────────────────────────────────────────── + + +def test_run_exposes_cloudflare_option_default_on(): + import inspect + + sig = inspect.signature(_studio().run) + assert "cloudflare" in sig.parameters + opt = sig.parameters["cloudflare"].default + decls = set(getattr(opt, "param_decls", []) or []) + assert "--cloudflare/--no-cloudflare" in decls + assert getattr(opt, "default", None) is True + + +def test_studio_default_exposes_cloudflare_option_default_on(): + import inspect + + sig = inspect.signature(_studio().studio_default) + assert "cloudflare" in sig.parameters + opt = sig.parameters["cloudflare"].default + assert getattr(opt, "default", None) is True + + +# ── re-exec forwarding: `unsloth studio run` ───────────────────────── + + +class _ExecCaptured(SystemExit): + def __init__(self, argv): + super().__init__(0) + self.argv = list(argv) + + +def _install_run_reexec_capture(monkeypatch, *, platform = "linux"): + studio_mod = _studio() + captured = [] + + monkeypatch.setattr(sys, "prefix", "/nonexistent/outer/venv") + fake_venv = Path("/fake/studio/venv/unsloth_studio") + monkeypatch.setattr(studio_mod, "_studio_venv_python", lambda: fake_venv / "bin" / "python") + fake_bin = fake_venv / "bin" / "unsloth" + real_is_file = Path.is_file + monkeypatch.setattr( + Path, + "is_file", + lambda self: True if str(self) == str(fake_bin) else real_is_file(self), + ) + from unsloth_cli import _tool_policy as _tp_mod + + monkeypatch.setattr( + _tp_mod, + "resolve_tool_policy", + lambda host, flag, yes, silent: False if flag is None else bool(flag), + ) + monkeypatch.setattr(sys, "platform", platform) + + def fake_execvp(file, argv): + captured.append(list(argv)) + raise _ExecCaptured(argv) + + monkeypatch.setattr(studio_mod.os, "execvp", fake_execvp) + return captured + + +def _invoke_run(monkeypatch, args): + import typer as _typer + + studio_mod = _studio() + captured = _install_run_reexec_capture(monkeypatch) + app = _typer.Typer() + app.command( + context_settings = {"allow_extra_args": True, "ignore_unknown_options": True}, + )(studio_mod.run) + CliRunner().invoke(app, args, catch_exceptions = True) + return captured + + +@pytest.mark.parametrize( + "user_flag,expected,unexpected", + [ + (None, "--cloudflare", "--no-cloudflare"), # default on + ("--cloudflare", "--cloudflare", "--no-cloudflare"), + ("--no-cloudflare", "--no-cloudflare", "--cloudflare"), + ], +) +def test_run_reexec_forwards_cloudflare_polarity(monkeypatch, user_flag, expected, unexpected): + extras = [user_flag] if user_flag else [] + captured = _invoke_run(monkeypatch, _BASE + extras) + assert len(captured) == 1, captured + argv = captured[0] + assert expected in argv, f"expected {expected} in child argv; got {argv}" + assert unexpected not in argv, f"unexpected {unexpected} in child argv; got {argv}" + + +# ── re-exec forwarding: plain `unsloth studio` ─────────────────────── + + +def _invoke_studio_default( + monkeypatch, + args, + *, + platform = "linux", +): + import typer as _typer + + studio_mod = _studio() + captured = [] + + monkeypatch.setattr(sys, "prefix", "/nonexistent/outer/venv") + monkeypatch.setattr(studio_mod, "_ensure_studio_env_exported", lambda: None) + fake_venv = Path("/fake/studio/venv/unsloth_studio") + monkeypatch.setattr(studio_mod, "_studio_venv_python", lambda: fake_venv / "bin" / "python") + monkeypatch.setattr(studio_mod, "_find_run_py", lambda: Path("/fake/studio/run.py")) + monkeypatch.setattr(studio_mod, "_find_frontend_dist", lambda: None) + monkeypatch.setattr(sys, "platform", platform) + + def fake_execvp(file, argv): + captured.append(list(argv)) + raise _ExecCaptured(argv) + + monkeypatch.setattr(studio_mod.os, "execvp", fake_execvp) + + app = _typer.Typer() + app.command()(studio_mod.studio_default) + CliRunner().invoke(app, args, catch_exceptions = True) + return captured + + +@pytest.mark.parametrize( + "user_flag,expected,unexpected", + [ + (None, "--cloudflare", "--no-cloudflare"), + ("--no-cloudflare", "--no-cloudflare", "--cloudflare"), + ], +) +def test_studio_default_reexec_forwards_cloudflare(monkeypatch, user_flag, expected, unexpected): + extras = [user_flag] if user_flag else [] + captured = _invoke_studio_default(monkeypatch, ["-H", "0.0.0.0"] + extras) + assert len(captured) == 1, captured + argv = captured[0] + assert expected in argv, f"expected {expected}; got {argv}" + assert unexpected not in argv, f"unexpected {unexpected}; got {argv}" + + +# ── in-venv path forwards cloudflare into run_server ───────────────── + + +class _RunServerCaptured(SystemExit): + def __init__(self, kwargs): + super().__init__(0) + self.kwargs = dict(kwargs) + + +@pytest.mark.parametrize("user_flag,expected", [(None, True), ("--no-cloudflare", False)]) +def test_run_in_venv_passes_cloudflare_to_run_server(monkeypatch, user_flag, expected): + import types + + studio_mod = _studio() + fake_venv = Path("/fake/studio/venv/unsloth_studio") + monkeypatch.setattr(sys, "prefix", str(fake_venv)) + monkeypatch.setattr(studio_mod, "STUDIO_HOME", fake_venv.parent) + + from unsloth_cli import _tool_policy as _tp_mod + + monkeypatch.setattr( + _tp_mod, + "resolve_tool_policy", + lambda host, flag, yes, silent: False if flag is None else bool(flag), + ) + + captured: dict = {} + + def fake_run_server(**kwargs): + captured.update(kwargs) + raise _RunServerCaptured(kwargs) + + fake_backend_run = sys.modules.setdefault( + "studio.backend.run", types.ModuleType("studio.backend.run") + ) + fake_backend_run.run_server = fake_run_server + fake_backend_run._resolve_external_ip = lambda: "127.0.0.1" + + import typer as _typer + + app = _typer.Typer() + app.command( + context_settings = {"allow_extra_args": True, "ignore_unknown_options": True}, + )(studio_mod.run) + extras = [user_flag] if user_flag else [] + CliRunner().invoke(app, _BASE + extras, catch_exceptions = True) + + assert captured.get("cloudflare") is expected, captured + + +# ── parent-level --no-cloudflare with a subcommand is rejected ─────── + + +def test_studio_default_rejects_no_cloudflare_with_subcommand(monkeypatch): + # `unsloth studio --no-cloudflare run ...` would not reach the subcommand, + # so it must error (mirrors --parallel) rather than silently still tunnel. + import typer as _typer + + studio_mod = _studio() + app = _typer.Typer() + app.add_typer(studio_mod.studio_app, name = "studio") + result = CliRunner().invoke(app, ["studio", "--no-cloudflare", "run", "--model", "X"]) + assert result.exit_code == 2, result.output + combined = (result.output or "") + (getattr(result, "stderr", "") or "") + assert "--no-cloudflare" in combined, combined + + +# ── run() tears the server + tunnel down if startup aborts ─────────── + + +def test_run_in_venv_shuts_down_on_startup_abort(monkeypatch): + import types + + studio_mod = _studio() + fake_venv = Path("/fake/studio/venv/unsloth_studio") + monkeypatch.setattr(sys, "prefix", str(fake_venv)) + monkeypatch.setattr(studio_mod, "STUDIO_HOME", fake_venv.parent) + + from unsloth_cli import _tool_policy as _tp_mod + + monkeypatch.setattr( + _tp_mod, + "resolve_tool_policy", + lambda host, flag, yes, silent: False if flag is None else bool(flag), + ) + + class _App: + class state: + server_port = 8888 + + shutdown_calls = [] + backend = sys.modules.setdefault("studio.backend.run", types.ModuleType("studio.backend.run")) + backend.run_server = lambda **k: _App() + backend._resolve_external_ip = lambda: "1.2.3.4" + backend._server = object() + backend._shutdown_event = None + backend._graceful_shutdown = lambda server: shutdown_calls.append(server) + + # set_tool_policy is imported as `from state.tool_policy import set_tool_policy`. + state_mod = sys.modules.setdefault("state", types.ModuleType("state")) + tp_mod = sys.modules.setdefault("state.tool_policy", types.ModuleType("state.tool_policy")) + tp_mod.set_tool_policy = lambda *a, **k: None + state_mod.tool_policy = tp_mod + + # Force the health check to fail so startup aborts after run_server(). + monkeypatch.setattr(studio_mod, "_wait_for_server", lambda *a, **k: False) + + import typer as _typer + + app = _typer.Typer() + app.command( + context_settings = {"allow_extra_args": True, "ignore_unknown_options": True}, + )(studio_mod.run) + result = CliRunner().invoke(app, _BASE + ["-H", "0.0.0.0"], catch_exceptions = True) + + assert result.exit_code == 1, result.output + assert len(shutdown_calls) == 1, "startup abort must call _graceful_shutdown" From 14f679da8acd3f083a9f2a5e2fe84d2e38f06351 Mon Sep 17 00:00:00 2001 From: BardiaKoopah Date: Thu, 11 Jun 2026 07:25:24 -0700 Subject: [PATCH 29/42] fix(studio): surface NaN loss honestly instead of laundering to last finite value (#6016) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit When training produced a NaN or Inf loss event, the handler filtered the value to None but never updated progress.loss — clients kept seeing the last finite value as if everything were fine. Now: on non-finite loss, clear progress.loss to None and log a one-shot warning. Training continues (no phase=error, no _should_stop), matching the expected behavior for a non-fatal numerical event. Test: tests/test_training_nan_loss_handling.py with 6 cases covering finite, NaN, +/-Inf, idempotency of the one-shot warning, and recovery when a finite step follows a non-finite one. --- studio/backend/core/training/training.py | 19 +++- .../tests/test_training_nan_loss_handling.py | 106 ++++++++++++++++++ 2 files changed, 124 insertions(+), 1 deletion(-) create mode 100644 studio/backend/tests/test_training_nan_loss_handling.py diff --git a/studio/backend/core/training/training.py b/studio/backend/core/training/training.py index c5bc3dbad5..5bd9bc5338 100644 --- a/studio/backend/core/training/training.py +++ b/studio/backend/core/training/training.py @@ -534,8 +534,21 @@ class TrainingBackend: except (TypeError, ValueError): logger.debug("Could not convert loss to float: %s", _raw_loss) _safe_loss = None - if _safe_loss is not None and not math.isfinite(_safe_loss): + _loss_is_nonfinite = ( + _safe_loss is not None and not math.isfinite(_safe_loss) + ) + if _loss_is_nonfinite: + # Drop the value rather than laundering it back to the last + # finite loss; clients see loss=None at this step so the NaN + # is not hidden behind a stale value. Training continues. _safe_loss = None + if not getattr(self._progress, "_nonfinite_loss_warned", False): + self._progress._nonfinite_loss_warned = True + logger.warning( + "Training produced non-finite loss at step %s; " + "loss field will report null until it recovers.", + event.get("step", "?"), + ) try: _safe_lr = float(_raw_lr) if _raw_lr is not None else None except (TypeError, ValueError): @@ -545,6 +558,10 @@ class TrainingBackend: _safe_lr = None if _safe_loss is not None: self._progress.loss = _safe_loss + elif _loss_is_nonfinite: + # Clear stale finite loss so the API doesn't keep + # reporting the last good value while NaN is happening. + self._progress.loss = None if _safe_lr is not None: self._progress.learning_rate = _safe_lr self._progress.total_steps = event.get("total_steps", self._progress.total_steps) diff --git a/studio/backend/tests/test_training_nan_loss_handling.py b/studio/backend/tests/test_training_nan_loss_handling.py new file mode 100644 index 0000000000..0ab444f9bb --- /dev/null +++ b/studio/backend/tests/test_training_nan_loss_handling.py @@ -0,0 +1,106 @@ +# SPDX-License-Identifier: AGPL-3.0-only +# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. + +"""Pin Studio's behavior when a training event reports non-finite (NaN/Inf) loss. + +The training event handler used to filter NaN/Inf to None silently while +leaving the previous finite loss in progress.loss — so the API kept reporting +the stale value as if everything were fine. We now drop the stale value: +clients see loss=None at the affected step and a one-shot warning is logged. +Training continues; the run is not marked failed. +""" + +from __future__ import annotations + +import math +import os +import sys + +import pytest + +_BACKEND = os.path.join(os.path.dirname(__file__), "..") +if _BACKEND not in sys.path: + sys.path.insert(0, _BACKEND) + +from core.training.training import TrainingBackend + + +def _make_backend() -> TrainingBackend: + return TrainingBackend() + + +def _progress_event(step: int, loss: float, lr: float = 1e-4) -> dict: + return { + "type": "progress", + "step": step, + "loss": loss, + "learning_rate": lr, + "epoch": 0.0, + "total_steps": 100, + } + + +class TestNonfiniteLossSoftHandling: + def test_finite_loss_updates_progress_normally(self): + b = _make_backend() + b._handle_event(_progress_event(step=1, loss=0.97)) + assert b._progress.loss == pytest.approx(0.97) + assert b._progress.error is None + assert b._should_stop is False + assert getattr(b._progress, "_nonfinite_loss_warned", False) is False + + def test_nan_loss_clears_progress_loss(self): + b = _make_backend() + b._handle_event(_progress_event(step=1, loss=0.97)) + assert b._progress.loss == pytest.approx(0.97) + b._handle_event(_progress_event(step=2, loss=float("nan"))) + # Stale finite loss must NOT leak through + assert b._progress.loss is None + # Run is not marked failed + assert b._progress.error is None + assert b._should_stop is False + # Warning flag is set so we don't re-log on every subsequent NaN step + assert b._progress._nonfinite_loss_warned is True + + def test_inf_loss_clears_progress_loss(self): + b = _make_backend() + b._handle_event(_progress_event(step=1, loss=float("inf"))) + assert b._progress.loss is None + assert b._progress.error is None + assert b._should_stop is False + assert b._progress._nonfinite_loss_warned is True + + def test_negative_inf_loss_clears_progress_loss(self): + b = _make_backend() + b._handle_event(_progress_event(step=1, loss=float("-inf"))) + assert b._progress.loss is None + assert b._progress.error is None + assert b._should_stop is False + assert b._progress._nonfinite_loss_warned is True + + def test_repeated_nan_only_warns_once(self): + """Subsequent NaN events must not re-fire the warning flag setter. + The flag should already be True after the first NaN.""" + b = _make_backend() + b._handle_event(_progress_event(step=1, loss=0.97)) + b._handle_event(_progress_event(step=2, loss=float("nan"))) + assert b._progress._nonfinite_loss_warned is True + # Further NaN steps don't change anything we care about + b._handle_event(_progress_event(step=3, loss=float("nan"))) + b._handle_event(_progress_event(step=4, loss=float("nan"))) + assert b._progress._nonfinite_loss_warned is True + assert b._progress.loss is None + assert b._progress.error is None + assert b._should_stop is False + + def test_recovery_updates_loss_when_finite_again(self): + """If a NaN step is followed by a finite step, progress.loss must + reflect the new finite value (not stay stuck at None).""" + b = _make_backend() + b._handle_event(_progress_event(step=1, loss=0.97)) + b._handle_event(_progress_event(step=2, loss=float("nan"))) + assert b._progress.loss is None + b._handle_event(_progress_event(step=3, loss=0.85)) + assert b._progress.loss == pytest.approx(0.85) + # Warning flag stays set (we don't reset it on recovery) + assert b._progress._nonfinite_loss_warned is True From 0ac051177ab03c597af64d679b2063f05443a397 Mon Sep 17 00:00:00 2001 From: BardiaKoopah Date: Thu, 11 Jun 2026 07:29:05 -0700 Subject: [PATCH 30/42] fix(mlx): forward resume_from_checkpoint to MLXTrainer.train() (#6173) Studio's frontend exposes a Resume action and submits requests with resume_from_checkpoint set to a previous run's output_dir. The CUDA training paths in worker.py read this field from config and pass it to trainer.train() (see lines 2729-2787 and 3108-3229). The MLX path _run_mlx_training did neither: it never read config['resume_from_checkpoint'] and called trainer.train() with no args. The MLX trainer also did not accept the kwarg, so even threading it through would have been a no-op. With this PR + the unsloth-zoo companion PR adding the trainer-side support (saves optimizer_state + trainer_state, accepts and applies resume_from_checkpoint in MLXTrainer.train()), MLX Resume now works end-to-end. Verified on M2 16GB with Qwen3-0.6B + unsloth/LaTeX_OCR: loss at every post-resume step matches a fresh run bit for bit (2.168627977371216 == 2.168627977371216 at step 6, etc). Two lines: read the field near the other config.get() extractions in _run_mlx_training, pass it as a kwarg at the trainer.train() call site. Companion PR: unslothai/unsloth-zoo#751 --- studio/backend/core/training/worker.py | 8 +++++++- 1 file changed, 7 insertions(+), 1 deletion(-) diff --git a/studio/backend/core/training/worker.py b/studio/backend/core/training/worker.py index fb5e9b5137..427687451c 100644 --- a/studio/backend/core/training/worker.py +++ b/studio/backend/core/training/worker.py @@ -1389,6 +1389,12 @@ def _run_mlx_training(event_queue, stop_queue, config): # Force text-only for non-image datasets even on vision-capable models # (e.g. Qwen3.5-VL trained on plain alpaca text). _send("status", status_message = f"Loading {model_name}...") + # Pull through resume_from_checkpoint so MLXTrainer.train() can restore + # optimizer + step state and continue cleanly. Was previously dropped on + # the floor for the MLX path, so the Resume UI button silently restarted + # from step 0 (the CUDA path at lines 2729 / 3108 has been forwarding + # this all along). + resume_from_checkpoint = config.get("resume_from_checkpoint") or None is_dataset_image = bool(config.get("is_dataset_image", False)) training_type = config.get("training_type", "LoRA/QLoRA") use_lora = training_type == "LoRA/QLoRA" @@ -1852,7 +1858,7 @@ def _run_mlx_training(event_queue, stop_queue, config): # ── 11. Run training ── gc.collect() mx.synchronize() - trainer.train() + trainer.train(resume_from_checkpoint=resume_from_checkpoint) # ── 12. Save and finalize ── if trainer.stop_requested and not _stop_save[0]: From 7467064c6650b401665beb03d285cd6d19cb6776 Mon Sep 17 00:00:00 2001 From: Burak Emir Date: Thu, 11 Jun 2026 16:45:10 +0200 Subject: [PATCH 31/42] Bump hono to 4.12.21, fixes CVE-2026-47676 (#6014) --- studio/frontend/package-lock.json | 6 +++--- studio/frontend/package.json | 2 +- 2 files changed, 4 insertions(+), 4 deletions(-) diff --git a/studio/frontend/package-lock.json b/studio/frontend/package-lock.json index 690da8c0c4..a521552d79 100644 --- a/studio/frontend/package-lock.json +++ b/studio/frontend/package-lock.json @@ -10264,9 +10264,9 @@ } }, "node_modules/hono": { - "version": "4.12.18", - "resolved": "https://registry.npmjs.org/hono/-/hono-4.12.18.tgz", - "integrity": "sha512-RWzP96k/yv0PQfyXnWjs6zot20TqfpfsNXhOnev8d1InAxubW93L11/oNUc3tQqn2G0bSdAOBpX+2uDFHV7kdQ==", + "version": "4.12.21", + "resolved": "https://registry.npmjs.org/hono/-/hono-4.12.21.tgz", + "integrity": "sha512-uV63apnb0kyPtAUwoWgaGh9HyIFcv8lgmzPZSiTBQAFOFGIzka5EZ1dZocmGnn0XdX0+XTqJ6Tqv7selMuGLRQ==", "license": "MIT", "engines": { "node": ">=16.9.0" diff --git a/studio/frontend/package.json b/studio/frontend/package.json index 2217dd080b..c49b62ab50 100644 --- a/studio/frontend/package.json +++ b/studio/frontend/package.json @@ -86,7 +86,7 @@ "@tanstack/router-core": "1.169.2", "@tanstack/history": "1.161.6", "mermaid": "11.15.0", - "hono": "4.12.18", + "hono": "4.12.21", "qs": "6.15.2", "ip-address": "10.1.1", "brace-expansion@5.0.5": "5.0.6" From d24ee77f1736b9bdcf87296905e18ca0ef5147d4 Mon Sep 17 00:00:00 2001 From: Daniel Han Date: Thu, 11 Jun 2026 07:48:21 -0700 Subject: [PATCH 32/42] Fix Llama 3.1+ rope scaling dropped on the FastLanguageModel path (long inputs become gibberish past ~29K tokens) (#6197) * Fix config.rope_scaling being dropped by the replaced rotary embedding (#2405) On modern transformers, LlamaModel builds its rotary embedding from config using unsloth's replacement LlamaRotaryEmbedding class, whose config path computed vanilla inv_freq and ignored config.rope_scaling entirely. The llama3/linear/longrope dispatch in patch_llama_rope_scaling rewrites LlamaAttention.__init__, which no longer constructs rotary embeddings, so it never fires; the model-level rotary is then copied onto every attention layer. Result: Llama-3.1/3.2/3.3 ran with unscaled RoPE on the FastLanguageModel path and collapsed into repetition loops past roughly 29K tokens (PASS at 28867, FAIL at 31767 in needle retrieval). FastModel was unaffected because vision.py keeps transformers' own rotary. qwen2, qwen3, qwen3_moe, mistral and cohere assign the same base class, so any rope-scaled config of those families was equally exposed. The fix makes the base class config path compute inv_freq and attention_scaling via transformers' ROPE_INIT_FUNCTIONS (covers llama3, linear, dynamic, yarn, longrope), with an inline llama3 fallback reading factors from config for older transformers, degrading to prior behavior on any failure. attention_scaling is applied in _set_cos_sin_cache (1.0 default, exact no-op for unscaled paths) and persists across extend_rope_embedding. A type(self) guard prevents double-scaling via the legacy scaled subclasses. Adds tests/utils/test_rope_scaling_drift.py (AST tripwire + behavioral inv_freq/cos-cache/extension checks, validated to fail 4 of 5 on the unfixed code) and wires it into the existing consolidated CI HARD GATE step. Verified on GPU: 48K-token needle retrieval flips FAIL to PASS for FastLanguageModel in bf16 and 4bit, 20K stays PASS, scaled inv_freq matches transformers exactly, and the left-padded batch generation guard still gets exact solo-vs-batched token matches. * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * Address review: normalize object-style rope_scaling, vectorize llama3 fallback config.rope_scaling can be a config object rather than a dict on newer transformers; _rope_scaling_as_dict normalizes it (to_dict/dict/vars fallbacks) before any .get() access, with a regression test using a dataclass stand-in. The inline llama3 fallback now uses torch.where instead of a per-frequency Python loop; verified bit-for-bit equal to transformers ROPE_INIT_FUNCTIONS for factor 8 (Llama-3.1) and factor 32 (Llama-3.2). * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * Address review: CPU-safe rope guard tests, normalized config for delegation The rotary constructor builds per-device CUDA caches, so the behavioral tests that instantiate it cannot run on GPU-less CI. Restructured into three layers: the AST tripwire now also asserts the constructor stays wired to _compute_config_rope_inv_freq; the CPU layer tests that pure helper directly (llama3 dict, llama3 object, linear object, default type) with no instantiation; the instantiation and cache tests are gated behind a real CUDA probe (actual tensor allocation, so import-time CUDA spoofs cannot fool the gate). Verified: 9 passed with GPU; 5 passed 4 skipped with CUDA hidden; 5 failed 4 skipped on the unfixed code in CPU mode. Delegation to ROPE_INIT_FUNCTIONS now retries with a shallow config copy carrying the normalized rope_scaling dict when the original was an object the installed transformers cannot read; covered by a linear-object test, which has no inline fallback and passes only through that retry path. * Tighten comments in rope scaling fix and guard test Comment and docstring reduction only; verified code-identical with scripts/comment_tools.py check --strip-docstrings (AST signature match on both Python files). All guard tests unchanged: 20 passed with GPU, 5 passed 4 skipped with CUDA hidden. * Apply repo kwarg-spacing format --------- Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com> --- .github/workflows/consolidated-tests-ci.yml | 21 +- tests/utils/test_rope_scaling_drift.py | 312 ++++++++++++++++++++ unsloth/models/llama.py | 123 +++++++- 3 files changed, 436 insertions(+), 20 deletions(-) create mode 100644 tests/utils/test_rope_scaling_drift.py diff --git a/.github/workflows/consolidated-tests-ci.yml b/.github/workflows/consolidated-tests-ci.yml index 022c796ac4..f7c338d76b 100644 --- a/.github/workflows/consolidated-tests-ci.yml +++ b/.github/workflows/consolidated-tests-ci.yml @@ -333,19 +333,16 @@ jobs: run: | python -m pytest -v --tb=short tests/test_callback_signature_drift.py - - name: batched left-padding generation guard (HARD GATE) - # Guards _fast_prepare_inputs_for_generation against the bug class of - # issues #1066 / #3699: position_ids taken from cache_position (which - # counts left-pad tokens) or the 2D attention mask truncated to its - # last column. Both shipped in cc4c5d77 and were fixed by #2216 and - # #4100; nothing tested this path, so each regression reached users. - # Layer 1 in the file is stdlib-ast-only (survives unsloth import - # breakage), layer 2 calls the real function on CPU via the - # tests/conftest.py CUDA spoof. Validated to fail on the pre-#2216 - # and pre-#4100 code states; staging proof on GPU-less runners: - # danielhanchen/unsloth-staging-2 PR 170 (green, gate passed in all combos) / PR 172 (red, gate failed in all combos). + - name: generation correctness guards (HARD GATE) + # Deterministic CPU guards, each validated to fail on its pre-fix code: + # leftpad = batched left-padded generation (#1066/#3699, fixed by + # #2216 + #4100; staging proof: unsloth-staging-2 PRs 170/172); + # rope_scaling_drift = config.rope_scaling dropped by replaced rotary + # classes (#2405). AST checks run first so import breakage cannot mask them. run: | - python -m pytest -v --tb=short tests/utils/test_prepare_inputs_leftpad.py + python -m pytest -v --tb=short \ + tests/utils/test_prepare_inputs_leftpad.py \ + tests/utils/test_rope_scaling_drift.py - name: unsloth Bucket-A — CPU tests not in Repo tests (CPU) # CPU tests across 6 files under tests/saving/, tests/utils/, tests/python/ diff --git a/tests/utils/test_rope_scaling_drift.py b/tests/utils/test_rope_scaling_drift.py new file mode 100644 index 0000000000..fae5a7d8a8 --- /dev/null +++ b/tests/utils/test_rope_scaling_drift.py @@ -0,0 +1,312 @@ +"""Guard for config.rope_scaling being silently dropped (issue #2405). + +Unsloth's replacement rotary classes ignored rope_scaling when constructed +from a config (the modern-transformers path), so Llama-3.1 ran with unscaled +RoPE and collapsed into gibberish past ~32K tokens. + +Layers: (1) AST tripwire, stdlib only; (2) CPU checks of the pure helper +_compute_config_rope_inv_freq against transformers' ROPE_INIT_FUNCTIONS; +(3) CUDA checks instantiating the real class (skipped without a real device, +probed by allocating a tensor so import-time CUDA spoofs cannot fool the gate). +Layers 2 and 3 fail on the unfixed code. +""" + +import ast +import math +from pathlib import Path + +import pytest +import torch + + +def _has_real_cuda(): + try: + torch.zeros(1).to("cuda") + return True + except Exception: + return False + + +HAS_REAL_CUDA = _has_real_cuda() +requires_cuda = pytest.mark.skipif( + not HAS_REAL_CUDA, + reason = "LlamaRotaryEmbedding builds per-device CUDA caches in __init__", +) + +REPO_ROOT = Path(__file__).resolve().parents[2] +LLAMA_PY = REPO_ROOT / "unsloth" / "models" / "llama.py" + +CLASS_NAME = "LlamaRotaryEmbedding" + +# Llama-3.1-style rope_scaling. +LLAMA3_ROPE_SCALING = { + "rope_type": "llama3", + "factor": 8.0, + "low_freq_factor": 1.0, + "high_freq_factor": 4.0, + "original_max_position_embeddings": 8192, +} +ROPE_THETA = 500000.0 +HEAD_DIM = 128 +MAX_POS = 131072 + + +# --- Layer 1: AST structural tripwire (stdlib only, no unsloth import) --- + + +def _load_class_init(): + tree = ast.parse(LLAMA_PY.read_text()) + for node in ast.walk(tree): + if isinstance(node, ast.ClassDef) and node.name == CLASS_NAME: + for sub in node.body: + if isinstance(sub, ast.FunctionDef) and sub.name == "__init__": + return sub + raise AssertionError( + f"{CLASS_NAME}.__init__ not found in {LLAMA_PY}; if it was renamed or " + "moved, update this guard so RoPE scaling stays protected (issue #2405)" + ) + + +def _config_branch(init_fn): + """The `if config is not None:` block at the top of __init__.""" + for node in init_fn.body: + if isinstance(node, ast.If): + test = node.test + is_config_test = ( + isinstance(test, ast.Compare) + and isinstance(test.left, ast.Name) + and test.left.id == "config" + ) + if is_config_test: + return node + return None + + +def test_config_path_inspects_rope_scaling(): + init_fn = _load_class_init() + branch = _config_branch(init_fn) + assert branch is not None, ( + f"{CLASS_NAME}.__init__ no longer has an `if config is not None:` " + "branch; the config constructor path must read config.rope_scaling so " + "scaled models (llama3/linear/longrope) are not silently unscaled " + "(issue #2405)" + ) + + names = set() + for stmt in branch.body: + for sub in ast.walk(stmt): + if isinstance(sub, ast.Attribute): + names.add(sub.attr) + elif isinstance(sub, ast.Constant) and isinstance(sub.value, str): + names.add(sub.value) + assert "rope_scaling" in names, ( + f"{CLASS_NAME}.__init__ config path does not reference `rope_scaling`. " + "When a rotary class is built straight from a config (the path modern " + "transformers takes, since rotary moved to LlamaModel), the llama3 / " + "linear / longrope scaling must still be applied; otherwise long inputs " + "produce repeated-pattern gibberish (issue #2405)." + ) + + called = { + sub.func.id + for stmt in branch.body + for sub in ast.walk(stmt) + if isinstance(sub, ast.Call) and isinstance(sub.func, ast.Name) + } + assert "_compute_config_rope_inv_freq" in called, ( + f"{CLASS_NAME}.__init__ config path no longer calls " + "_compute_config_rope_inv_freq; the CPU behavioral tests below cover " + "that helper directly, so the constructor must stay wired to it or " + "scaled configs silently lose RoPE scaling again (issue #2405)." + ) + + +# --- Layer 2: CPU behavioral guard (pure helper, no instantiation) --- + + +def _make_config(rope_scaling): + from transformers import LlamaConfig + return LlamaConfig( + hidden_size = 256, + num_attention_heads = 2, + num_key_value_heads = 2, + head_dim = HEAD_DIM, + rope_theta = ROPE_THETA, + max_position_embeddings = MAX_POS, + rope_scaling = rope_scaling, + ) + + +def _unsloth_rotary(config): + from unsloth.models import llama as llama_mod + return llama_mod.LlamaRotaryEmbedding(config = config) + + +def _reference_inv_freq(config, rope_type): + from transformers.modeling_rope_utils import ROPE_INIT_FUNCTIONS + inv_freq, _attention_factor = ROPE_INIT_FUNCTIONS[rope_type](config, "cpu") + return inv_freq.float().cpu() + + +def _vanilla_inv_freq(): + return 1.0 / ( + ROPE_THETA ** (torch.arange(0, HEAD_DIM, 2, dtype = torch.int64).float() / HEAD_DIM) + ) + + +def _compute_helper(config, rope_scaling): + from unsloth.models.llama import _compute_config_rope_inv_freq + return _compute_config_rope_inv_freq(config, rope_scaling) + + +def test_llama3_scaling_applied_to_inv_freq(): + config = _make_config(LLAMA3_ROPE_SCALING) + got, attention_scaling = _compute_helper(config, config.rope_scaling) + expected = _reference_inv_freq(config, "llama3") + vanilla = _vanilla_inv_freq() + + # Guard against a vacuous test. + assert not torch.allclose( + expected, vanilla, rtol = 1e-4 + ), "test setup error: llama3-scaled inv_freq should differ from vanilla" + assert got is not None, ( + "_compute_config_rope_inv_freq returned None for a llama3 config; the " + "config path is dropping config.rope_scaling, so long-context inference " + "degrades into repeated-pattern gibberish (issue #2405)." + ) + got = got.float().cpu() + assert torch.allclose(got, expected, rtol = 1e-4, atol = 1e-6), ( + "inv_freq for a llama3 config does not match transformers' llama3 RoPE " + "scaling (issue #2405).\n" + f"got[:6]={got[:6].tolist()}\nexpected[:6]={expected[:6].tolist()}" + ) + + +def test_default_rope_type_matches_vanilla_inv_freq(): + config = _make_config(None) + got, attention_scaling = _compute_helper(config, {"rope_type": "default"}) + assert got is not None + vanilla = _vanilla_inv_freq() + assert torch.allclose(got.float().cpu(), vanilla, rtol = 1e-4, atol = 1e-6), ( + "default rope_type must equal the vanilla inv_freq; " + f"got[:6]={got[:6].tolist()} vanilla[:6]={vanilla[:6].tolist()}" + ) + + +def _cos_at_position(rot, position): + """cos row at one position, built like _set_cos_sin_cache but CPU-only.""" + inv_freq = rot.inv_freq.float().cpu() + t = torch.tensor([position], dtype = torch.float32) + t = rot._apply_time_scaling(t.clone()) if hasattr(rot, "_apply_time_scaling") else t + freqs = torch.outer(t, inv_freq) + emb = torch.cat((freqs, freqs), dim = -1) + return emb.cos().squeeze(0) + + +# --- Layer 3: CUDA behavioral guard (real instantiation needs a device) --- + + +@requires_cuda +def test_constructor_applies_llama3_scaling(): + config = _make_config(LLAMA3_ROPE_SCALING) + rot = _unsloth_rotary(config) + got = rot.inv_freq.float().cpu() + expected = _reference_inv_freq(config, "llama3") + assert torch.allclose( + got, expected, rtol = 1e-4, atol = 1e-6 + ), "LlamaRotaryEmbedding built from a llama3 config produced unscaled inv_freq (issue #2405)." + + +@requires_cuda +def test_constructor_unscaled_config_uses_vanilla_inv_freq(): + rot = _unsloth_rotary(_make_config(None)) + got = rot.inv_freq.float().cpu() + vanilla = _vanilla_inv_freq() + assert torch.allclose( + got, vanilla, rtol = 1e-4, atol = 1e-6 + ), "LlamaRotaryEmbedding with no rope_scaling must use the vanilla inv_freq" + + +@requires_cuda +def test_cos_cache_differs_between_scaled_and_unscaled_at_long_position(): + scaled = _unsloth_rotary(_make_config(LLAMA3_ROPE_SCALING)) + unscaled = _unsloth_rotary(_make_config(None)) + + pos = 10000 + cos_scaled = _cos_at_position(scaled, pos) + cos_unscaled = _cos_at_position(unscaled, pos) + assert not torch.allclose(cos_scaled, cos_unscaled, rtol = 1e-4, atol = 1e-5), ( + f"cos values at position {pos} are identical for a llama3-scaled and an " + "unscaled rotary embedding, which means scaling was dropped (issue " + "#2405). With correct llama3 scaling the low-frequency bands shrink by " + "up to 8x and must change the angles at long positions." + ) + + +@requires_cuda +def test_extended_cache_keeps_scaling_after_growth(): + scaled = _unsloth_rotary(_make_config(LLAMA3_ROPE_SCALING)) + # Grow past the initial cache size (mirrors long-context decode). + dummy = torch.zeros(1, dtype = torch.float32) + scaled.extend_rope_embedding(dummy, seq_len = 40960) + + config = _make_config(LLAMA3_ROPE_SCALING) + expected = _reference_inv_freq(config, "llama3") + got = scaled.inv_freq.float().cpu() + assert torch.allclose(got, expected, rtol = 1e-4, atol = 1e-6), ( + "growing the RoPE cache (extend_rope_embedding) must preserve llama3 " + "scaling of inv_freq; long-context decode loses scaling otherwise " + "(issue #2405)." + ) + + +def test_object_style_rope_scaling_does_not_crash(): + # Object-style rope_scaling must be normalized, not .get()'d directly. + from dataclasses import dataclass + + from unsloth.models.llama import _compute_config_rope_inv_freq + + @dataclass + class FakeRopeScalingConfig: + rope_type: str = "llama3" + factor: float = 8.0 + low_freq_factor: float = 1.0 + high_freq_factor: float = 4.0 + original_max_position_embeddings: int = 8192 + + config = _make_config(LLAMA3_ROPE_SCALING) + inv_freq, attention_scaling = _compute_config_rope_inv_freq(config, FakeRopeScalingConfig()) + assert inv_freq is not None, ( + "object-style (non-dict) config.rope_scaling must be normalized, not " + "dropped; otherwise scaled models silently lose RoPE scaling again " + "(issue #2405)." + ) + expected = _reference_inv_freq(config, "llama3") + assert torch.allclose(inv_freq.float().cpu(), expected, rtol = 1e-4, atol = 1e-6) + + +def test_object_style_rope_scaling_on_config_delegates_correctly(): + # 'linear' has no inline fallback; only the normalized-config retry passes this. + from dataclasses import dataclass + + from unsloth.models.llama import _compute_config_rope_inv_freq + + @dataclass + class FakeLinearRopeScalingConfig: + rope_type: str = "linear" + factor: float = 4.0 + + dict_config = _make_config({"rope_type": "linear", "factor": 4.0}) + expected = _reference_inv_freq(dict_config, "linear") + + object_config = _make_config({"rope_type": "linear", "factor": 4.0}) + object_config.rope_scaling = FakeLinearRopeScalingConfig() + inv_freq, attention_scaling = _compute_config_rope_inv_freq( + object_config, object_config.rope_scaling + ) + assert inv_freq is not None, ( + "linear rope_scaling exposed as a config object was silently dropped; " + "delegation must retry with a config copy carrying the normalized dict " + "(issue #2405)." + ) + assert torch.allclose(inv_freq.float().cpu(), expected, rtol = 1e-4, atol = 1e-6) diff --git a/unsloth/models/llama.py b/unsloth/models/llama.py index d5883e4346..2d31f71ab1 100644 --- a/unsloth/models/llama.py +++ b/unsloth/models/llama.py @@ -1622,6 +1622,93 @@ def _get_rope_theta(config, default = 10000.0): return default +def _rope_scaling_as_dict(rope_scaling): + """Normalize config.rope_scaling (dict or config object) to a dict; {} on failure.""" + if isinstance(rope_scaling, dict): + return rope_scaling + for converter in ("to_dict", "dict"): + fn = getattr(rope_scaling, converter, None) + if callable(fn): + try: + d = fn() + if isinstance(d, dict): + return d + except Exception: + pass + try: + return {k: v for k, v in vars(rope_scaling).items() if not k.startswith("_")} + except TypeError: + return {} + + +def _llama3_inv_freq_from_config( + config, + rope_scaling, + device = "cpu", +): + """llama3 inv_freq with factors from config; fallback when modeling_rope_utils is missing.""" + base = _get_rope_theta(config, default = 10000.0) + dim = getattr(config, "head_dim", None) + if dim is None: + dim = int(config.hidden_size // config.num_attention_heads) + inv_freq = 1.0 / ( + base ** (torch.arange(0, dim, 2, dtype = torch.int64, device = device).float() / dim) + ) + + scale_factor = rope_scaling.get("factor", 8.0) + low_freq_factor = rope_scaling.get("low_freq_factor", 1.0) + high_freq_factor = rope_scaling.get("high_freq_factor", 4.0) + old_context_len = rope_scaling.get("original_max_position_embeddings", 8192) + + low_freq_wavelen = old_context_len / low_freq_factor + high_freq_wavelen = old_context_len / high_freq_factor + assert low_freq_wavelen != high_freq_wavelen + + # Vectorized meta-llama bands: high freqs kept, low divided by factor, medium blended. + wavelen = 2 * math.pi / inv_freq + scaled = torch.where(wavelen > low_freq_wavelen, inv_freq / scale_factor, inv_freq) + smooth = (old_context_len / wavelen - low_freq_factor) / (high_freq_factor - low_freq_factor) + smoothed = (1 - smooth) * inv_freq / scale_factor + smooth * inv_freq + is_medium = (wavelen >= high_freq_wavelen) & (wavelen <= low_freq_wavelen) + return torch.where(is_medium, smoothed, scaled) + + +def _compute_config_rope_inv_freq(config, rope_scaling): + """(inv_freq, attention_scaling) per config.rope_scaling via transformers' + ROPE_INIT_FUNCTIONS, with an inline llama3 fallback; (None, 1.0) on failure.""" + original_rope_scaling = rope_scaling + rope_scaling = _rope_scaling_as_dict(rope_scaling) + rope_type = rope_scaling.get("rope_type", None) or rope_scaling.get("type", None) + try: + from transformers.modeling_rope_utils import ROPE_INIT_FUNCTIONS + + rope_init_fn = ROPE_INIT_FUNCTIONS[rope_type] + try: + inv_freq, attention_scaling = rope_init_fn(config, torch.device("cpu")) + except Exception: + # Object-style rope_scaling: retry with a config copy carrying the plain dict. + if isinstance(original_rope_scaling, dict): + raise + import copy as _copy + + config_copy = _copy.copy(config) + config_copy.rope_scaling = rope_scaling + inv_freq, attention_scaling = rope_init_fn(config_copy, torch.device("cpu")) + return inv_freq.to(dtype = torch.float32, device = "cpu"), float(attention_scaling) + except Exception as exception: + if rope_type == "llama3": + try: + return _llama3_inv_freq_from_config(config, rope_scaling), 1.0 + except Exception: + pass + logger.warning_once( + f"Unsloth: Could not apply RoPE scaling '{rope_type}' from config " + f"({type(exception).__name__}: {exception}); falling back to unscaled RoPE. " + "Long-context generation may degrade." + ) + return None, 1.0 + + # Solves https://github.com/unslothai/unsloth/issues/168 # Static KV Cache was introduced in 4.38.0, causing training to be much slower. # Inference can now be CUDAGraphed, but we shall retain the old rotary embeddings. @@ -1640,6 +1727,12 @@ class LlamaRotaryEmbedding(torch.nn.Module): config = None, # [TODO] Hack to pass in config - need to remove later ): super().__init__() + # cos/sin multiplier (1.0 except yarn / longrope); set before any cache build. + self.attention_scaling = 1.0 + # Base-class-from-config path (modern transformers): derive inv_freq like + # transformers so config.rope_scaling is not dropped (#2405). Scaled + # subclasses are excluded to avoid double-scaling. + config_inv_freq = None if config is not None: # [TODO] Hack to pass in config - need to remove later base = _get_rope_theta(config, default = base) @@ -1652,6 +1745,13 @@ class LlamaRotaryEmbedding(torch.nn.Module): device = DEVICE_TYPE_TORCH max_position_embeddings = config.max_position_embeddings + rope_scaling = getattr(config, "rope_scaling", None) + if rope_scaling is not None and type(self) is LlamaRotaryEmbedding: + config_inv_freq, self.attention_scaling = _compute_config_rope_inv_freq( + config, + rope_scaling, + ) + self.dim = dim self.max_position_embeddings = max_position_embeddings self.base = base @@ -1660,12 +1760,17 @@ class LlamaRotaryEmbedding(torch.nn.Module): self.multi_gpu_cos_cached = [None] * DEVICE_COUNT self.multi_gpu_sin_cached = [None] * DEVICE_COUNT - # Normal Llama-3 RoPE - inv_freq = 1.0 / ( - self.base - ** (torch.arange(0, self.dim, 2, dtype = torch.int64, device = "cpu").float() / self.dim) - ) - inv_freq = self._apply_inv_freq_scaling(inv_freq) + if config_inv_freq is not None: + inv_freq = config_inv_freq # already scaled; skip subclass scaling + else: + # Normal Llama-3 RoPE + inv_freq = 1.0 / ( + self.base + ** ( + torch.arange(0, self.dim, 2, dtype = torch.int64, device = "cpu").float() / self.dim + ) + ) + inv_freq = self._apply_inv_freq_scaling(inv_freq) self.register_buffer("inv_freq", inv_freq, persistent = False) # Build here to make `torch.jit.trace` work. @@ -1704,8 +1809,10 @@ class LlamaRotaryEmbedding(torch.nn.Module): freqs = torch.outer(t, self.inv_freq) # Different from paper, but it uses a different permutation in order to obtain the same calculation emb = torch.cat((freqs, freqs), dim = -1) - cos = emb.cos().to(dtype = dtype, device = device, non_blocking = True) - sin = emb.sin().to(dtype = dtype, device = device, non_blocking = True) + # Applied here so attention_scaling survives extend_rope_embedding rebuilds; + # default 1.0 keeps unscaled paths bit-identical. + cos = (emb.cos() * self.attention_scaling).to(dtype = dtype, device = device, non_blocking = True) + sin = (emb.sin() * self.attention_scaling).to(dtype = dtype, device = device, non_blocking = True) self.multi_gpu_cos_cached[device.index] = cos self.multi_gpu_sin_cached[device.index] = sin return cos, sin From bc85ecd145ed5ceae8dab662e3686485c1e6a559 Mon Sep 17 00:00:00 2001 From: Daniel Han Date: Thu, 11 Jun 2026 07:49:55 -0700 Subject: [PATCH 33/42] Studio: report the real llama-server context window and add an opt-in overflow policy for OpenAI-compatible serving (#6164) * Studio: report the real llama-server context window and add an opt-in overflow policy for OpenAI-compatible serving A community report showed OpenCode failing tool calls every few minutes against Studio's OpenAI-compatible API while the same GGUF was stable on LM Studio. Root cause: Studio advertises the requested context length, but llama-server can allocate less (memory-fit step on small GPUs, --parallel slot split), so clients budget against a window that does not exist. Their generations truncate mid tool call at the real wall (finish_reason=length with cut JSON arguments) and eventually the prompt itself exceeds the real window, returning a 400 that agentic clients treat as non-retryable. Changes: - After llama-server health, read default_generation_settings.n_ctx from /props and adopt it whenever it is below Studio's computed context, with a warning. The load response, status route, UI value, and the passthrough max_tokens ceiling all become honest automatically. - Expose context_length and max_context_length on /v1/models so clients can budget against the enforced window. - Accept empty role=tool content (commands with no output are routine in agentic loops; OpenAI and llama-server both accept it) instead of a 400. - Add context_overflow=truncate_middle (per request, or server-wide via UNSLOTH_CONTEXT_OVERFLOW=truncate_middle): on exceed_context_size_error the passthrough drops whole middle turn-groups (system prompt, first turn, and recent turns kept; tool calls stay paired with their results), clips oversized contents middle-out when group-dropping is not enough, clamps max_tokens to the generation headroom, and retries. Default stays 'error' with code=context_length_exceeded so clients running their own compaction keep full control. * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * Studio: allocate the requested context for real (kv-unified, fit-ctx floor) Two launch-flag gaps caused the advertised vs allocated divergence at the source: - llama-server enables --kv-unified only when the slot count is auto; Studio always passes --parallel N, which silently splits -c into per-slot windows of -c/N. Pass --kv-unified when N > 1 so a single request can use the full advertised window (same total KV memory, shared pool). - with --fit on the fit step may set ctx as low as 4096; pass --fit-ctx for explicit requests so fit offloads or fails into the existing --fit off retry instead of silently shrinking the window. Both flags are gated on --help capability probing so older builds keep the current behavior, where the /props readback remains the backstop. Verified live: -c 98304 --parallel 4 now serves per-slot n_ctx 98304 (was 24576), 48k-token requests pass through the passthrough, and the readback warning no longer fires. * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci --------- Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com> --- studio/backend/core/inference/llama_cpp.py | 78 +++++ studio/backend/models/inference.py | 16 +- studio/backend/routes/inference.py | 288 +++++++++++++++--- .../tests/test_context_overflow_truncation.py | 277 +++++++++++++++++ .../tests/test_llama_cpp_props_readback.py | 254 +++++++++++++++ .../tests/test_openai_tool_passthrough.py | 9 +- .../tests/test_tool_message_empty_content.py | 53 ++++ 7 files changed, 924 insertions(+), 51 deletions(-) create mode 100644 studio/backend/tests/test_context_overflow_truncation.py create mode 100644 studio/backend/tests/test_llama_cpp_props_readback.py create mode 100644 studio/backend/tests/test_tool_message_empty_content.py diff --git a/studio/backend/core/inference/llama_cpp.py b/studio/backend/core/inference/llama_cpp.py index e46d464784..2560f9e5ac 100644 --- a/studio/backend/core/inference/llama_cpp.py +++ b/studio/backend/core/inference/llama_cpp.py @@ -1153,6 +1153,8 @@ class LlamaCppBackend: "ngram_mod_flavor": None, "supports_ngram_mod": False, "spec_draft_n_max_flag": None, + "supports_kv_unified": False, + "supports_fit_ctx": False, } try: mtime = int(Path(bin_path).stat().st_mtime) @@ -1166,6 +1168,8 @@ class LlamaCppBackend: mtp_token: Optional[str] = None ngram_mod_flavor: Optional[str] = None spec_draft_n_max_flag: Optional[str] = None + supports_kv_unified = False + supports_fit_ctx = False try: result = subprocess.run( [bin_path, "--help"], @@ -1253,6 +1257,9 @@ class LlamaCppBackend: spec_draft_n_max_flag = "--spec-draft-n-max" elif _is_real("--draft-max"): spec_draft_n_max_flag = "--draft-max" + + supports_kv_unified = _is_real("--kv-unified") + supports_fit_ctx = _is_real("--fit-ctx") except (OSError, subprocess.SubprocessError) as exc: logger.debug(f"llama-server --help probe failed: {exc}") @@ -1263,6 +1270,8 @@ class LlamaCppBackend: "ngram_mod_flavor": ngram_mod_flavor, "supports_ngram_mod": ngram_mod_flavor is not None, "spec_draft_n_max_flag": spec_draft_n_max_flag, + "supports_kv_unified": supports_kv_unified, + "supports_fit_ctx": supports_fit_ctx, } cls._capability_cache[cache_key] = info return info @@ -3232,6 +3241,16 @@ class LlamaCppBackend: # Fits on selected GPU(s) -- offload all layers cmd.extend(["-ngl", "-1"]) + cmd.extend( + self._ctx_integrity_flags( + n_parallel, + use_fit, + requested_ctx, + effective_ctx, + self.probe_server_capabilities(binary), + ) + ) + # -1 = llama.cpp auto-detect (physical cores). Pass explicitly # so we don't inherit llama-server's internal default, which # has varied (hardware concurrency incl. hyperthreads on some @@ -3585,6 +3604,7 @@ class LlamaCppBackend: self._effective_context_length = ( effective_ctx if effective_ctx > 0 else self._context_length ) + self._reconcile_effective_ctx_with_server() self._max_context_length = ( max_available_ctx if max_available_ctx > 0 else self._effective_context_length ) @@ -4468,6 +4488,64 @@ class LlamaCppBackend: logger.error(f"llama-server health check timed out after {timeout}s") return False + @staticmethod + def _ctx_integrity_flags( + n_parallel: int, use_fit: bool, requested_ctx: int, effective_ctx: int, caps: dict + ) -> list[str]: + """Flags that keep the per-request window equal to the advertised ctx. + + Explicit ``--parallel`` disables llama-server's auto-slots + ``--kv-unified`` default, silently splitting ``-c`` into per-slot + windows of ``-c / N``; restore the shared pool so one request can use + the full context. With ``--fit on``, ``--fit-ctx`` floors the fit step + at an explicitly requested ctx (default floor is 4096) so it offloads + or fails instead of silently shrinking the window. + """ + flags: list[str] = [] + if n_parallel > 1 and caps.get("supports_kv_unified"): + flags.append("--kv-unified") + if use_fit and requested_ctx > 0 and effective_ctx > 0 and caps.get("supports_fit_ctx"): + flags.extend(["--fit-ctx", str(effective_ctx)]) + return flags + + def _query_server_n_ctx(self) -> Optional[int]: + """Per-slot context llama-server actually allocated, from ``/props``. + + The memory-fit step or ``--parallel`` slot split can leave this below + the requested ``-c``; requests are validated against this value. + """ + url = f"http://127.0.0.1:{self._port}/props" + try: + resp = httpx.get(url, timeout = 5.0) + if resp.status_code != 200: + return None + settings = resp.json().get("default_generation_settings") or {} + n_ctx = settings.get("n_ctx") + return int(n_ctx) if n_ctx else None + except Exception: + return None + + def _reconcile_effective_ctx_with_server(self) -> None: + """Adopt the server's real ``n_ctx`` when it is below Studio's value. + + Keeps ``context_length`` (load response, status route, passthrough + ``max_tokens`` ceiling) honest; clients sized to the requested value + would otherwise hit ``exceed_context_size_error`` 400s early. + """ + actual_n_ctx = self._query_server_n_ctx() + if not actual_n_ctx or actual_n_ctx <= 0: + return + if self._effective_context_length and actual_n_ctx < self._effective_context_length: + logger.warning( + "llama-server allocated a smaller per-request context than " + f"requested ({self._effective_context_length} -> {actual_n_ctx}; " + "memory fit or --parallel slot split); clients must treat " + f"{actual_n_ctx} as the real context window." + ) + self._effective_context_length = actual_n_ctx + elif not self._effective_context_length: + self._effective_context_length = actual_n_ctx + # ── Message building (OpenAI format) ────────────────────────── @staticmethod diff --git a/studio/backend/models/inference.py b/studio/backend/models/inference.py index f58a76d16e..f94d5366eb 100644 --- a/studio/backend/models/inference.py +++ b/studio/backend/models/inference.py @@ -544,8 +544,10 @@ class ChatMessage(BaseModel): if self.role == "tool": # tool_call_id resolution happens at ChatCompletionRequest scope. - if not self.content: - raise ValueError('role="tool" messages require non-empty "content".') + # OpenAI accepts empty tool results (commands with no output); + # normalize to "" instead of a 400 agentic clients treat as fatal. + if self.content is None or self.content == []: + self.content = "" elif self.role == "assistant": # Post-Stop sentinel: collapse content="" / [] to None. if (self.content == "" or self.content == []) and not self.tool_calls: @@ -692,6 +694,16 @@ class ChatCompletionRequest(BaseModel): True, description = "[x-unsloth] Auto-detect and fix malformed tool calls from model output.", ) + context_overflow: Optional[Literal["error", "truncate_middle"]] = Field( + None, + description = ( + "[x-unsloth] Passthrough behavior when the prompt exceeds the real " + "context window. 'error' (default) returns a 400 with " + "code=context_length_exceeded. 'truncate_middle' drops middle " + "turn-groups (system prompt, first turn, and recent turns kept; " + "tool calls stay paired with their results) and retries." + ), + ) max_tool_calls_per_message: Optional[int] = Field( 25, ge = 0, diff --git a/studio/backend/routes/inference.py b/studio/backend/routes/inference.py index 0244956b72..9dc9feb920 100644 --- a/studio/backend/routes/inference.py +++ b/studio/backend/routes/inference.py @@ -242,6 +242,169 @@ def _openai_passthrough_error(status_code, text) -> "HTTPException": ) +_OVERFLOW_TRUNCATE_MAX_RETRIES = 3 +# Truncated-prompt share of the real window; the rest is generation headroom +# so a near-full prompt cannot cut a tool call mid-JSON at the wall. +_OVERFLOW_PROMPT_TARGET_FRACTION = 0.75 + + +def _overflow_truncation_requested(payload) -> bool: + """True when the request (or the UNSLOTH_CONTEXT_OVERFLOW server default, + for clients that cannot send custom fields) opted into truncation.""" + requested = getattr(payload, "context_overflow", None) + if requested is not None: + return requested == "truncate_middle" + return os.environ.get("UNSLOTH_CONTEXT_OVERFLOW", "").strip().lower() == "truncate_middle" + + +def _parse_overflow_counts(err_text: str): + """(n_prompt_tokens, n_ctx) from an exceed_context_size_error body, or + None. Tolerates \\" around keys (body may be a re-wrapped JSON string).""" + m_prompt = _re.search(r'n_prompt_tokens\\?"?\s*:\s*(\d+)', err_text) + m_ctx = _re.search(r'n_ctx\\?"?\s*:\s*(\d+)', err_text) + if m_prompt and m_ctx: + return int(m_prompt.group(1)), int(m_ctx.group(1)) + return None + + +def _estimate_message_tokens(msg: dict) -> int: + try: + return max(1, len(json.dumps(msg, ensure_ascii = False)) // 4) + except Exception: + return 1 + + +def _truncate_middle_messages(messages: list, keep_ratio: float): + """Drop whole turn-groups from the middle of an OpenAI message list. + + Always kept: leading system message(s), the first group (task anchor), + and the trailing groups. A group is a user message, or an assistant + message plus its following tool results, so surviving tool_calls stay + paired with their results as chat templates require. + Returns (new_messages, dropped_message_count). + """ + if not messages or keep_ratio >= 1.0: + return messages, 0 + + head: list = [] + idx = 0 + while idx < len(messages) and messages[idx].get("role") in ("system", "developer"): + head.append(messages[idx]) + idx += 1 + + groups: list[list] = [] + for msg in messages[idx:]: + role = msg.get("role") + if role == "tool" and groups: + groups[-1].append(msg) + elif role == "tool": + groups.append([msg]) # orphan tool result; treat as its own group + else: + groups.append([msg]) + + # Anchor group plus the last 3 groups stay. + protected_tail = min(3, max(1, len(groups) - 1)) + if len(groups) <= 1 + protected_tail: + return messages, 0 + + total_est = sum(_estimate_message_tokens(m) for m in messages) + target_est = int(total_est * keep_ratio) + + anchor = groups[0] + middle = groups[1:-protected_tail] + tail = groups[-protected_tail:] + + current_est = total_est + kept_middle: list[list] = list(middle) + dropped = 0 + # Drop oldest-first until the estimate fits the target. + while kept_middle and current_est > target_est: + victim = kept_middle.pop(0) + dropped += len(victim) + current_est -= sum(_estimate_message_tokens(m) for m in victim) + + if dropped == 0: + return messages, 0 + + new_messages = head + anchor + for grp in kept_middle: + new_messages.extend(grp) + for grp in tail: + new_messages.extend(grp) + return new_messages, dropped + + +_CLIP_MARKER = "\n[... truncated by context_overflow=truncate_middle ...]\n" +# Generous head+tail first; cut harder if the estimate still misses the target. +_CLIP_KEEP_CHARS = (1500, 400) + + +def _clip_long_contents(messages: list, target_est: int) -> int: + """Clip oversized string contents middle-out until ``target_est`` is met. + + Tool results first, then earlier user turns, the final message last. + Message count and roles never change, so tool pairing holds even when + group-dropping could not free enough. Returns messages clipped. + """ + + def _candidates(): + tools = [m for m in messages if m.get("role") == "tool"] + users = [m for m in messages[:-1] if m.get("role") == "user"] + last = [messages[-1]] if messages else [] + return tools + users + last + + clipped = 0 + for keep in _CLIP_KEEP_CHARS: + for msg in _candidates(): + if sum(_estimate_message_tokens(m) for m in messages) <= target_est: + return clipped + content = msg.get("content") + if not isinstance(content, str) or len(content) <= 2 * keep + len(_CLIP_MARKER): + continue + msg["content"] = content[:keep] + _CLIP_MARKER + content[-keep:] + clipped += 1 + return clipped + + +def _apply_overflow_truncation(body: dict, err_text: str) -> bool: + """Shrink a passthrough body after an upstream context overflow: drop + middle turn-groups, clip still-oversized contents, clamp ``max_tokens`` + to the generation headroom. Returns False when nothing could shrink.""" + counts = _parse_overflow_counts(err_text) + messages = body.get("messages") or [] + total_est = sum(_estimate_message_tokens(m) for m in messages) + if counts: + n_prompt, n_ctx = counts + keep_ratio = min(0.95, (_OVERFLOW_PROMPT_TARGET_FRACTION * n_ctx) / max(1, n_prompt)) + # Scale the server-token target into char-estimate units. + target_est = int(total_est * keep_ratio) + else: + n_ctx = None + keep_ratio = 0.6 # no counts in the error; cut conservatively + target_est = int(total_est * keep_ratio) + + new_messages, dropped = _truncate_middle_messages(messages, keep_ratio) + if dropped: + body["messages"] = new_messages + clipped = 0 + if sum(_estimate_message_tokens(m) for m in body.get("messages") or []) > target_est: + clipped = _clip_long_contents(body.get("messages") or [], target_est) + if not dropped and not clipped: + return False + if n_ctx: + headroom = max(1024, int(n_ctx * (1.0 - _OVERFLOW_PROMPT_TARGET_FRACTION))) + cur_max = body.get("max_tokens") + body["max_tokens"] = min(cur_max, headroom) if cur_max else headroom + logger.warning( + "context_overflow=truncate_middle: dropped %d middle messages, clipped " + "%d contents (keep_ratio %.2f); retrying within the real window", + dropped, + clipped, + keep_ratio, + ) + return True + + def _anthropic_stream_error_event(exc): """Anthropic in-band SSE ``error`` event for a mid-stream failure, or ``None`` to fall through to a normal message_delta finish. Returns an event only for a @@ -4456,26 +4619,35 @@ def _openai_model_objects() -> list[dict]: # Check GGUF backend llama_backend = get_llama_cpp_backend() if llama_backend.is_loaded: - models.append( - { - "id": llama_backend.model_identifier, - "object": "model", - "created": _created, - "owned_by": "local", - } - ) + entry = { + "id": llama_backend.model_identifier, + "object": "model", + "created": _created, + "owned_by": "local", + } + # Extension fields: the real per-request window (post /props readback) + # so clients can budget/compact against the enforced limit. + if llama_backend.context_length: + entry["context_length"] = llama_backend.context_length + if llama_backend.max_context_length: + entry["max_context_length"] = llama_backend.max_context_length + models.append(entry) # Check Unsloth backend backend = get_inference_backend() if backend.active_model_name: - models.append( - { - "id": backend.active_model_name, - "object": "model", - "created": _created, - "owned_by": "local", - } + entry = { + "id": backend.active_model_name, + "object": "model", + "created": _created, + "owned_by": "local", + } + _sf_ctx = getattr(backend, "context_length", None) or getattr( + backend, "max_seq_length", None ) + if _sf_ctx: + entry["context_length"] = _sf_ctx + models.append(entry) return models @@ -6822,27 +6994,32 @@ async def _openai_passthrough_stream( limits = httpx.Limits(max_keepalive_connections = 0), ) resp = None - try: - req = client.build_request("POST", target_url, json = body) - resp = await client.send(req, stream = True) - except httpx.RequestError as e: - # llama-server subprocess crashed / starting / unreachable. - logger.error("openai passthrough stream: upstream unreachable: %s", e) - if resp is not None: + _truncate_budget = ( + _OVERFLOW_TRUNCATE_MAX_RETRIES if _overflow_truncation_requested(payload) else 0 + ) + while True: + try: + req = client.build_request("POST", target_url, json = body) + resp = await client.send(req, stream = True) + except httpx.RequestError as e: + # llama-server subprocess crashed / starting / unreachable. + logger.error("openai passthrough stream: upstream unreachable: %s", e) + if resp is not None: + try: + await resp.aclose() + except Exception: + pass try: - await resp.aclose() + await client.aclose() except Exception: pass - try: - await client.aclose() - except Exception: - pass - raise HTTPException( - status_code = 502, - detail = _friendly_error(e), - ) + raise HTTPException( + status_code = 502, + detail = _friendly_error(e), + ) - if resp.status_code != 200: + if resp.status_code == 200: + break err_bytes = await resp.aread() err_text = err_bytes.decode("utf-8", errors = "replace") logger.error( @@ -6855,6 +7032,14 @@ async def _openai_passthrough_stream( await resp.aclose() except Exception: pass + # Opt-in overflow policy: shrink and retry instead of a fatal 400. + if ( + _truncate_budget > 0 + and _classify_llama_generation_error(Exception(err_text)) + and _apply_overflow_truncation(body, err_text) + ): + _truncate_budget -= 1 + continue try: await client.aclose() except Exception: @@ -6953,20 +7138,33 @@ async def _openai_passthrough_non_streaming(llama_backend, payload, model_name): payload, backend_ctx = llama_backend.context_length, llama_backend = llama_backend ) - try: - async with httpx.AsyncClient() as client: - resp = await client.post(target_url, json = body, timeout = 600) - except httpx.RequestError as e: - # llama-server subprocess crashed / starting / unreachable. Surface the - # same friendly message the sync chat path emits so operators don't see - # a bare 500 with no diagnostic. - logger.error("openai passthrough non-streaming: upstream unreachable: %s", e) - raise HTTPException( - status_code = 502, - detail = _friendly_error(e), - ) + _truncate_budget = ( + _OVERFLOW_TRUNCATE_MAX_RETRIES if _overflow_truncation_requested(payload) else 0 + ) + while True: + try: + async with httpx.AsyncClient() as client: + resp = await client.post(target_url, json = body, timeout = 600) + except httpx.RequestError as e: + # llama-server subprocess crashed / starting / unreachable. Surface the + # same friendly message the sync chat path emits so operators don't see + # a bare 500 with no diagnostic. + logger.error("openai passthrough non-streaming: upstream unreachable: %s", e) + raise HTTPException( + status_code = 502, + detail = _friendly_error(e), + ) - if resp.status_code != 200: + if resp.status_code == 200: + break + # Opt-in overflow policy: shrink and retry instead of a fatal 400. + if ( + _truncate_budget > 0 + and _classify_llama_generation_error(Exception(resp.text)) + and _apply_overflow_truncation(body, resp.text) + ): + _truncate_budget -= 1 + continue raise _openai_passthrough_error(resp.status_code, resp.text) # The guided-decoding fence wraps each choice's JSON content in a diff --git a/studio/backend/tests/test_context_overflow_truncation.py b/studio/backend/tests/test_context_overflow_truncation.py new file mode 100644 index 0000000000..4f4c240934 --- /dev/null +++ b/studio/backend/tests/test_context_overflow_truncation.py @@ -0,0 +1,277 @@ +# SPDX-License-Identifier: AGPL-3.0-only +# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0 + +"""Tests for the opt-in ``context_overflow="truncate_middle"`` passthrough policy. + +On ``exceed_context_size_error`` the passthrough drops middle turn-groups and +retries inside the real window instead of surfacing a fatal 400. Truncation +keeps the system prompt, the first turn, and recent turns, and never orphans +a tool result from its tool_calls turn. Also covers ``/v1/models`` exposing +the real post-readback context window. +""" + +from __future__ import annotations + +import sys +from pathlib import Path + +import pytest + +_BACKEND_DIR = str(Path(__file__).resolve().parent.parent) +if _BACKEND_DIR not in sys.path: + sys.path.insert(0, _BACKEND_DIR) + +from routes.inference import ( + _apply_overflow_truncation, + _clip_long_contents, + _CLIP_MARKER, + _estimate_message_tokens, + _openai_model_objects, + _overflow_truncation_requested, + _parse_overflow_counts, + _truncate_middle_messages, +) +import routes.inference as routes_mod + + +# Nick's actual error body from the Discord report logs. +_NICK_ERROR = ( + '{"detail":"llama-server error: {\\"error\\":{\\"code\\":400,' + '\\"message\\":\\"request (70494 tokens) exceeds the available context size ' + '(67584 tokens), try increasing it\\",\\"type\\":\\"exceed_context_size_error\\",' + '\\"n_prompt_tokens\\":70494,\\"n_ctx\\":67584}}"}' +) + + +def _tool_turn(i: int, result_chars: int = 400) -> list[dict]: + """An assistant tool_calls turn paired with its tool result.""" + return [ + { + "role": "assistant", + "content": "", + "tool_calls": [ + { + "id": f"call_{i}", + "type": "function", + "function": {"name": "read", "arguments": f'{{"filePath":"/f{i}"}}'}, + } + ], + }, + {"role": "tool", "tool_call_id": f"call_{i}", "content": "x" * result_chars}, + ] + + +def _conversation(n_tool_turns: int = 12) -> list[dict]: + msgs = [ + {"role": "system", "content": "You are an agent." * 20}, + {"role": "user", "content": "Do the big task." * 20}, + ] + for i in range(n_tool_turns): + msgs.extend(_tool_turn(i)) + msgs.append({"role": "assistant", "content": "halfway summary"}) + msgs.append({"role": "user", "content": "keep going"}) + return msgs + + +# --------------------------------------------------------------------------- +# _parse_overflow_counts +# --------------------------------------------------------------------------- + + +def test_parse_overflow_counts_nick_error(): + assert _parse_overflow_counts(_NICK_ERROR) == (70494, 67584) + + +def test_parse_overflow_counts_missing_fields(): + assert _parse_overflow_counts('{"error":"something else"}') is None + + +# --------------------------------------------------------------------------- +# _truncate_middle_messages +# --------------------------------------------------------------------------- + + +def test_truncation_drops_middle_keeps_anchors(): + msgs = _conversation() + new, dropped = _truncate_middle_messages(msgs, keep_ratio = 0.5) + assert dropped > 0 + assert len(new) == len(msgs) - dropped + # System prompt and task anchor survive. + assert new[0]["role"] == "system" + assert new[1] == msgs[1] + # The most recent turns survive verbatim. + assert new[-1] == msgs[-1] + assert new[-2] == msgs[-2] + + +def test_truncation_never_orphans_tool_results(): + msgs = _conversation() + new, dropped = _truncate_middle_messages(msgs, keep_ratio = 0.4) + assert dropped > 0 + surviving_call_ids = { + tc["id"] for m in new if m.get("role") == "assistant" for tc in (m.get("tool_calls") or []) + } + for m in new: + if m.get("role") == "tool": + assert m["tool_call_id"] in surviving_call_ids + + +def test_truncation_reduces_estimated_size_toward_target(): + msgs = _conversation() + total = sum(_estimate_message_tokens(m) for m in msgs) + new, dropped = _truncate_middle_messages(msgs, keep_ratio = 0.5) + new_total = sum(_estimate_message_tokens(m) for m in new) + assert dropped > 0 + assert new_total < total + # Should land at or below the requested share, modulo one whole group. + biggest_group = max( + _estimate_message_tokens(a) + _estimate_message_tokens(b) + for a, b in zip(msgs[2:-2:2], msgs[3:-2:2]) + ) + assert new_total <= int(total * 0.5) + biggest_group + + +def test_truncation_noop_when_keep_ratio_full(): + msgs = _conversation() + new, dropped = _truncate_middle_messages(msgs, keep_ratio = 1.0) + assert dropped == 0 + assert new == msgs + + +def test_truncation_noop_when_only_protected_turns_remain(): + msgs = [ + {"role": "system", "content": "sys"}, + {"role": "user", "content": "task"}, + *_tool_turn(0), + {"role": "user", "content": "latest"}, + ] + new, dropped = _truncate_middle_messages(msgs, keep_ratio = 0.1) + assert dropped == 0 + assert new == msgs + + +# --------------------------------------------------------------------------- +# _apply_overflow_truncation +# --------------------------------------------------------------------------- + + +def test_apply_overflow_truncation_mutates_body_and_clamps_max_tokens(): + body = {"messages": _conversation(), "max_tokens": 32000} + assert _apply_overflow_truncation(body, _NICK_ERROR) is True + assert len(body["messages"]) < len(_conversation()) + # Generation headroom: max_tokens clamped to the non-prompt share of n_ctx. + assert body["max_tokens"] <= max(1024, int(67584 * 0.25)) + + +def test_apply_overflow_truncation_returns_false_when_nothing_droppable(): + body = { + "messages": [ + {"role": "system", "content": "sys"}, + {"role": "user", "content": "task"}, + {"role": "user", "content": "latest"}, + ], + "max_tokens": 32000, + } + assert _apply_overflow_truncation(body, _NICK_ERROR) is False + + +def test_apply_overflow_truncation_clips_giant_protected_tool_results(): + """One giant burst (few turn-groups, all protected) must still shrink: + stage 2 clips oversized tool contents instead of giving up.""" + msgs = [ + {"role": "system", "content": "sys"}, + {"role": "user", "content": "task"}, + *_tool_turn(0, result_chars = 60000), + *_tool_turn(1, result_chars = 60000), + ] + body = {"messages": msgs, "max_tokens": 32000} + n_before = len(msgs) + assert _apply_overflow_truncation(body, _NICK_ERROR) is True + # No message disappeared (pairing intact), but contents were clipped. + assert len(body["messages"]) == n_before + clipped = [m for m in body["messages"] if _CLIP_MARKER in str(m.get("content"))] + assert clipped, "expected at least one clipped tool result" + surviving_call_ids = { + tc["id"] + for m in body["messages"] + if m.get("role") == "assistant" + for tc in (m.get("tool_calls") or []) + } + for m in body["messages"]: + if m.get("role") == "tool": + assert m["tool_call_id"] in surviving_call_ids + + +def test_clip_long_contents_reaches_target_and_keeps_structure(): + msgs = [ + {"role": "system", "content": "sys"}, + {"role": "user", "content": "task"}, + *_tool_turn(0, result_chars = 40000), + {"role": "user", "content": "latest question"}, + ] + total = sum(_estimate_message_tokens(m) for m in msgs) + clipped = _clip_long_contents(msgs, target_est = total // 4) + assert clipped >= 1 + assert sum(_estimate_message_tokens(m) for m in msgs) <= total // 4 + # Roles and count unchanged; the short final user message untouched. + assert [m["role"] for m in msgs] == ["system", "user", "assistant", "tool", "user"] + assert msgs[-1]["content"] == "latest question" + + +def test_overflow_truncation_requested_reads_field(monkeypatch): + monkeypatch.delenv("UNSLOTH_CONTEXT_OVERFLOW", raising = False) + + class _P: + context_overflow = "truncate_middle" + + class _Q: + context_overflow = None + + assert _overflow_truncation_requested(_P()) is True + assert _overflow_truncation_requested(_Q()) is False + assert _overflow_truncation_requested(object()) is False + + +def test_overflow_truncation_server_default_env(monkeypatch): + """UNSLOTH_CONTEXT_OVERFLOW enables the policy for clients that cannot + send custom body fields; an explicit per-request 'error' still wins.""" + + class _Unset: + context_overflow = None + + class _ExplicitError: + context_overflow = "error" + + monkeypatch.setenv("UNSLOTH_CONTEXT_OVERFLOW", "truncate_middle") + assert _overflow_truncation_requested(_Unset()) is True + assert _overflow_truncation_requested(_ExplicitError()) is False + monkeypatch.setenv("UNSLOTH_CONTEXT_OVERFLOW", "error") + assert _overflow_truncation_requested(_Unset()) is False + + +# --------------------------------------------------------------------------- +# /v1/models context metadata +# --------------------------------------------------------------------------- + + +class _FakeLlamaBackend: + is_loaded = True + model_identifier = "unsloth/Qwen3.6-27B-GGUF" + context_length = 67584 + max_context_length = 262144 + + +class _FakeEmptyBackend: + active_model_name = None + + +def test_v1_models_exposes_real_context_window(monkeypatch): + monkeypatch.setattr(routes_mod, "get_llama_cpp_backend", lambda: _FakeLlamaBackend()) + monkeypatch.setattr(routes_mod, "get_inference_backend", lambda: _FakeEmptyBackend()) + models = _openai_model_objects() + assert len(models) == 1 + entry = models[0] + assert entry["id"] == "unsloth/Qwen3.6-27B-GGUF" + # The REAL (post /props readback) window, not the requested one. + assert entry["context_length"] == 67584 + assert entry["max_context_length"] == 262144 diff --git a/studio/backend/tests/test_llama_cpp_props_readback.py b/studio/backend/tests/test_llama_cpp_props_readback.py new file mode 100644 index 0000000000..d87c05f2c6 --- /dev/null +++ b/studio/backend/tests/test_llama_cpp_props_readback.py @@ -0,0 +1,254 @@ +# SPDX-License-Identifier: AGPL-3.0-only +# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0 + +"""Tests for the post-launch /props context readback. + +llama-server's memory-fit step or --parallel slot split can allocate less +context than the requested -c while Studio keeps advertising the requested +value; clients sized to it then die on exceed_context_size_error 400s. +``_reconcile_effective_ctx_with_server`` must adopt the server's real +``default_generation_settings.n_ctx`` whenever it is smaller. + +Stubbed httpx; no subprocess, GPU, or network. Cross-platform. +""" + +from __future__ import annotations + +import json +import sys +import types as _types +from pathlib import Path + +import pytest + +# --------------------------------------------------------------------------- +# Stub heavy/unavailable deps before importing the module under test. +# Mirrors test_llama_cpp_context_fit.py. +# --------------------------------------------------------------------------- + +_BACKEND_DIR = str(Path(__file__).resolve().parent.parent) +if _BACKEND_DIR not in sys.path: + sys.path.insert(0, _BACKEND_DIR) + +# Prefer the real modules so importing this file first cannot poison later +# test modules with stubs; only stub what the environment genuinely lacks. +try: + import loggers # noqa: F401 +except ImportError: + _loggers_stub = _types.ModuleType("loggers") + _loggers_stub.get_logger = lambda name: __import__("logging").getLogger(name) + sys.modules.setdefault("loggers", _loggers_stub) + +try: + import structlog # noqa: F401 +except ImportError: + sys.modules.setdefault("structlog", _types.ModuleType("structlog")) + +try: + import httpx # noqa: F401 +except ImportError: + _httpx_stub = _types.ModuleType("httpx") + for _exc_name in ( + "ConnectError", + "TimeoutException", + "ReadTimeout", + "ReadError", + "RemoteProtocolError", + "CloseError", + "WriteError", + "HTTPError", + ): + setattr(_httpx_stub, _exc_name, type(_exc_name, (Exception,), {})) + + class _FakeTimeout: + def __init__(self, *a, **kw): + pass + + _httpx_stub.Timeout = _FakeTimeout + _httpx_stub.Client = type( + "Client", + (), + { + "__init__": lambda self, **kw: None, + "__enter__": lambda self: self, + "__exit__": lambda self, *a: None, + }, + ) + _httpx_stub.get = lambda *a, **kw: (_ for _ in ()).throw(RuntimeError("unstubbed httpx.get")) + sys.modules.setdefault("httpx", _httpx_stub) + +from core.inference.llama_cpp import LlamaCppBackend +import core.inference.llama_cpp as llama_cpp_mod + + +# --------------------------------------------------------------------------- +# Helpers +# --------------------------------------------------------------------------- + + +class _FakeResponse: + def __init__( + self, + status_code = 200, + body = None, + ): + self.status_code = status_code + self._body = body or {} + + def json(self): + return self._body + + +def _make_backend(effective_ctx = 98304, port = 51234): + inst = LlamaCppBackend.__new__(LlamaCppBackend) + inst._port = port + inst._effective_context_length = effective_ctx + inst._context_length = 262144 + return inst + + +def _stub_props( + monkeypatch, + status_code = 200, + body = None, + exc = None, +): + def fake_get(url, timeout = None): + assert url.endswith("/props") + if exc is not None: + raise exc + return _FakeResponse(status_code, body) + + monkeypatch.setattr(llama_cpp_mod.httpx, "get", fake_get, raising = False) + + +# --------------------------------------------------------------------------- +# _query_server_n_ctx parsing +# --------------------------------------------------------------------------- + + +def test_query_n_ctx_reads_default_generation_settings(monkeypatch): + _stub_props( + monkeypatch, + body = {"default_generation_settings": {"n_ctx": 67584}}, + ) + assert _make_backend()._query_server_n_ctx() == 67584 + + +def test_query_n_ctx_non_200_returns_none(monkeypatch): + _stub_props(monkeypatch, status_code = 503) + assert _make_backend()._query_server_n_ctx() is None + + +def test_query_n_ctx_missing_key_returns_none(monkeypatch): + _stub_props(monkeypatch, body = {"default_generation_settings": {}}) + assert _make_backend()._query_server_n_ctx() is None + + +def test_query_n_ctx_swallows_transport_errors(monkeypatch): + _stub_props(monkeypatch, exc = RuntimeError("connection refused")) + assert _make_backend()._query_server_n_ctx() is None + + +# --------------------------------------------------------------------------- +# _reconcile_effective_ctx_with_server decisions +# --------------------------------------------------------------------------- + + +def test_fit_shrunk_ctx_overwrites_advertised_value(monkeypatch): + """The Nick repro: requested/advertised 98304, server really at 67584.""" + inst = _make_backend(effective_ctx = 98304) + _stub_props( + monkeypatch, + body = {"default_generation_settings": {"n_ctx": 67584}}, + ) + inst._reconcile_effective_ctx_with_server() + assert inst._effective_context_length == 67584 + assert inst.context_length == 67584 + + +def test_matching_ctx_is_left_alone(monkeypatch): + inst = _make_backend(effective_ctx = 98304) + _stub_props( + monkeypatch, + body = {"default_generation_settings": {"n_ctx": 98304}}, + ) + inst._reconcile_effective_ctx_with_server() + assert inst._effective_context_length == 98304 + + +def test_larger_server_ctx_does_not_inflate_advertised_value(monkeypatch): + """Never advertise more than the user asked for, even if the server could.""" + inst = _make_backend(effective_ctx = 32768) + _stub_props( + monkeypatch, + body = {"default_generation_settings": {"n_ctx": 65536}}, + ) + inst._reconcile_effective_ctx_with_server() + assert inst._effective_context_length == 32768 + + +def test_unset_effective_ctx_adopts_server_value(monkeypatch): + inst = _make_backend(effective_ctx = None) + inst._context_length = None + _stub_props( + monkeypatch, + body = {"default_generation_settings": {"n_ctx": 40960}}, + ) + inst._reconcile_effective_ctx_with_server() + assert inst._effective_context_length == 40960 + + +def test_props_failure_keeps_studio_value(monkeypatch): + """A flaky /props must never wipe the computed context.""" + inst = _make_backend(effective_ctx = 98304) + _stub_props(monkeypatch, exc = RuntimeError("boom")) + inst._reconcile_effective_ctx_with_server() + assert inst._effective_context_length == 98304 + + +# --------------------------------------------------------------------------- +# _ctx_integrity_flags: keep the per-request window equal to the advertised ctx +# --------------------------------------------------------------------------- + +_CAPS_ALL = {"supports_kv_unified": True, "supports_fit_ctx": True} +_CAPS_NONE = {"supports_kv_unified": False, "supports_fit_ctx": False} + + +def test_kv_unified_added_for_multi_slot(): + """Explicit --parallel N disables llama-server's auto-slots kv-unified + default, splitting -c into per-slot windows of -c/N; Studio must restore + the shared pool so one request can use the full advertised context.""" + flags = LlamaCppBackend._ctx_integrity_flags(4, False, 98304, 98304, _CAPS_ALL) + assert "--kv-unified" in flags + + +def test_kv_unified_skipped_for_single_slot_or_old_build(): + assert "--kv-unified" not in LlamaCppBackend._ctx_integrity_flags( + 1, False, 98304, 98304, _CAPS_ALL + ) + assert "--kv-unified" not in LlamaCppBackend._ctx_integrity_flags( + 4, False, 98304, 98304, _CAPS_NONE + ) + + +def test_fit_ctx_floors_explicit_request_under_fit(): + flags = LlamaCppBackend._ctx_integrity_flags(1, True, 98304, 98304, _CAPS_ALL) + assert flags[flags.index("--fit-ctx") + 1] == "98304" + + +def test_fit_ctx_skipped_without_fit_or_explicit_ctx_or_support(): + assert "--fit-ctx" not in LlamaCppBackend._ctx_integrity_flags( + 1, False, 98304, 98304, _CAPS_ALL + ) + assert "--fit-ctx" not in LlamaCppBackend._ctx_integrity_flags(1, True, 0, 262144, _CAPS_ALL) + assert "--fit-ctx" not in LlamaCppBackend._ctx_integrity_flags( + 1, True, 98304, 98304, _CAPS_NONE + ) + + +def test_probe_missing_binary_reports_new_capabilities_false(): + info = LlamaCppBackend.probe_server_capabilities(binary = "/nonexistent/llama-server") + assert info["found"] is False + assert info["supports_kv_unified"] is False + assert info["supports_fit_ctx"] is False diff --git a/studio/backend/tests/test_openai_tool_passthrough.py b/studio/backend/tests/test_openai_tool_passthrough.py index d8a965c1f2..1d994b46c0 100644 --- a/studio/backend/tests/test_openai_tool_passthrough.py +++ b/studio/backend/tests/test_openai_tool_passthrough.py @@ -166,10 +166,11 @@ class TestChatMessageToolRoles: with pytest.raises(ValidationError): ChatMessage(role = "user", content = []) - def test_tool_empty_content_rejected(self): - with pytest.raises(ValidationError) as exc_info: - ChatMessage(role = "tool", tool_call_id = "call_1", content = "") - assert "content" in str(exc_info.value) + def test_tool_empty_content_accepted(self): + # Empty tool output (mkdir, git add, ...) is routine in agentic loops; + # OpenAI and llama-server both accept it, so Studio must not 400. + msg = ChatMessage(role = "tool", tool_call_id = "call_1", content = "") + assert msg.content == "" def test_assistant_without_content_or_tool_calls_tolerated(self): # Stop-button leaves an empty assistant turn; tolerate for replay. diff --git a/studio/backend/tests/test_tool_message_empty_content.py b/studio/backend/tests/test_tool_message_empty_content.py new file mode 100644 index 0000000000..d63b16ce80 --- /dev/null +++ b/studio/backend/tests/test_tool_message_empty_content.py @@ -0,0 +1,53 @@ +# SPDX-License-Identifier: AGPL-3.0-only +# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0 + +"""Empty ``role="tool"`` content must be accepted on the OpenAI-compat surface. + +Agentic clients send ``content: ""`` when a command produced no output; +OpenAI and llama-server both accept it. Studio used to 400, which standard +clients treat as non-retryable and kill the session. The validator must +normalize empty/missing tool content to ``""`` instead of raising. +""" + +from __future__ import annotations + +import sys +from pathlib import Path + +import pytest + +_BACKEND_DIR = str(Path(__file__).resolve().parent.parent) +if _BACKEND_DIR not in sys.path: + sys.path.insert(0, _BACKEND_DIR) + +from models.inference import ChatMessage + + +def test_tool_message_empty_string_content_is_accepted(): + msg = ChatMessage(role = "tool", content = "", tool_call_id = "call_1") + assert msg.content == "" + + +def test_tool_message_none_content_normalizes_to_empty_string(): + msg = ChatMessage(role = "tool", content = None, tool_call_id = "call_1") + assert msg.content == "" + + +def test_tool_message_empty_list_content_normalizes_to_empty_string(): + msg = ChatMessage(role = "tool", content = [], tool_call_id = "call_1") + assert msg.content == "" + + +def test_tool_message_real_content_is_preserved(): + msg = ChatMessage(role = "tool", content = "ok", tool_call_id = "call_1") + assert msg.content == "ok" + + +def test_user_message_still_requires_content(): + with pytest.raises(ValueError): + ChatMessage(role = "user", content = None) + + +def test_assistant_empty_content_still_collapses_to_none(): + msg = ChatMessage(role = "assistant", content = "") + assert msg.content is None From 3733e0b2741e666ebdd677424c2741c348dc9cfc Mon Sep 17 00:00:00 2001 From: Daniel Han Date: Thu, 11 Jun 2026 07:50:13 -0700 Subject: [PATCH 34/42] fix(studio): surface live step with null loss through the SSE progress stream (#6206) * fix(studio): surface live step with null loss through the SSE progress stream The metric histories skip non-finite steps, so during a NaN stretch the SSE live loop and final complete event replayed the last finite step/loss pair. Follow the live progress step when it is ahead of the history tail and report its loss honestly (null until recovery). Completes the NaN honesty fix for the SSE consumer flagged in review. * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * Apply live-step handling to inactive streams and clear the UI loss on null for PR #6206 Fresh /progress connections after a finished run took the inactive branch which still replayed the last finite step and loss pair; apply the same live-step correction there. On the frontend, applyProgress kept the stale currentLoss when a payload advanced the step with a null loss; clear it so the display shows -- until the loss recovers. Widen the runtime state type to number | null, which the view layer already handles. --------- Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com> --- studio/backend/core/training/training.py | 4 +- studio/backend/core/training/worker.py | 2 +- studio/backend/routes/training.py | 31 ++++- .../tests/test_training_nan_loss_handling.py | 30 ++-- .../test_training_progress_stream_nan.py | 129 ++++++++++++++++++ .../training/stores/training-runtime-store.ts | 5 +- .../src/features/training/types/runtime.ts | 3 +- 7 files changed, 180 insertions(+), 24 deletions(-) create mode 100644 studio/backend/tests/test_training_progress_stream_nan.py diff --git a/studio/backend/core/training/training.py b/studio/backend/core/training/training.py index 5bd9bc5338..2764101d93 100644 --- a/studio/backend/core/training/training.py +++ b/studio/backend/core/training/training.py @@ -534,9 +534,7 @@ class TrainingBackend: except (TypeError, ValueError): logger.debug("Could not convert loss to float: %s", _raw_loss) _safe_loss = None - _loss_is_nonfinite = ( - _safe_loss is not None and not math.isfinite(_safe_loss) - ) + _loss_is_nonfinite = _safe_loss is not None and not math.isfinite(_safe_loss) if _loss_is_nonfinite: # Drop the value rather than laundering it back to the last # finite loss; clients see loss=None at this step so the NaN diff --git a/studio/backend/core/training/worker.py b/studio/backend/core/training/worker.py index 427687451c..18b25cb4fe 100644 --- a/studio/backend/core/training/worker.py +++ b/studio/backend/core/training/worker.py @@ -1858,7 +1858,7 @@ def _run_mlx_training(event_queue, stop_queue, config): # ── 11. Run training ── gc.collect() mx.synchronize() - trainer.train(resume_from_checkpoint=resume_from_checkpoint) + trainer.train(resume_from_checkpoint = resume_from_checkpoint) # ── 12. Save and finalize ── if trainer.stop_requested and not _stop_save[0]: diff --git a/studio/backend/routes/training.py b/studio/backend/routes/training.py index 96403e21a5..d7687ffdee 100644 --- a/studio/backend/routes/training.py +++ b/studio/backend/routes/training.py @@ -665,10 +665,17 @@ async def stream_training_progress( # If not active, send final state and exit if not is_active: - if backend.step_history: - final_step = backend.step_history[-1] + _live = (getattr(tp, "step", 0) or 0) if tp else 0 + if backend.step_history or _live > 0: + final_step = backend.step_history[-1] if backend.step_history else 0 final_loss = backend.loss_history[-1] if backend.loss_history else None final_lr = backend.lr_history[-1] if backend.lr_history else None + # Histories skip non-finite steps; report the live step with + # loss=None instead of the last finite pair. + if _live > final_step: + final_step = _live + final_loss = getattr(tp, "loss", None) + final_lr = getattr(tp, "learning_rate", final_lr) final_total_steps = getattr(tp, "total_steps", final_step) if tp else final_step final_epoch = getattr(tp, "epoch", None) if tp else None payload = build_progress( @@ -697,11 +704,18 @@ async def stream_training_progress( while backend.is_training_active(): try: - if backend.step_history: - current_step = backend.step_history[-1] + tp_inner = getattr(getattr(backend, "trainer", None), "training_progress", None) + live_step = (getattr(tp_inner, "step", 0) or 0) if tp_inner else 0 + if backend.step_history or live_step > 0: + current_step = backend.step_history[-1] if backend.step_history else 0 current_loss = backend.loss_history[-1] if backend.loss_history else None current_lr = backend.lr_history[-1] if backend.lr_history else None - tp_inner = getattr(getattr(backend, "trainer", None), "training_progress", None) + # Histories skip non-finite steps; follow the live progress + # step and report its loss (None until it recovers). + if live_step > current_step: + current_step = live_step + current_loss = getattr(tp_inner, "loss", None) + current_lr = getattr(tp_inner, "learning_rate", current_lr) current_total_steps = ( getattr(tp_inner, "total_steps", current_step) if tp_inner else current_step ) @@ -798,6 +812,13 @@ async def stream_training_progress( final_loss = backend.loss_history[-1] if backend.loss_history else None final_lr = backend.lr_history[-1] if backend.lr_history else None final_tp = getattr(getattr(backend, "trainer", None), "training_progress", None) + # If the run ended on a non-finite stretch, report the live step with + # loss=None instead of rolling back to the last finite pair. + _final_live_step = (getattr(final_tp, "step", 0) or 0) if final_tp else 0 + if _final_live_step > (final_step if final_step is not None else -1): + final_step = _final_live_step + final_loss = getattr(final_tp, "loss", None) + final_lr = getattr(final_tp, "learning_rate", final_lr) final_total_steps = getattr(final_tp, "total_steps", final_step) if final_tp else final_step final_epoch = getattr(final_tp, "epoch", None) if final_tp else None final_payload = build_progress( diff --git a/studio/backend/tests/test_training_nan_loss_handling.py b/studio/backend/tests/test_training_nan_loss_handling.py index 0ab444f9bb..a2dc78bee2 100644 --- a/studio/backend/tests/test_training_nan_loss_handling.py +++ b/studio/backend/tests/test_training_nan_loss_handling.py @@ -29,7 +29,11 @@ def _make_backend() -> TrainingBackend: return TrainingBackend() -def _progress_event(step: int, loss: float, lr: float = 1e-4) -> dict: +def _progress_event( + step: int, + loss: float, + lr: float = 1e-4, +) -> dict: return { "type": "progress", "step": step, @@ -43,7 +47,7 @@ def _progress_event(step: int, loss: float, lr: float = 1e-4) -> dict: class TestNonfiniteLossSoftHandling: def test_finite_loss_updates_progress_normally(self): b = _make_backend() - b._handle_event(_progress_event(step=1, loss=0.97)) + b._handle_event(_progress_event(step = 1, loss = 0.97)) assert b._progress.loss == pytest.approx(0.97) assert b._progress.error is None assert b._should_stop is False @@ -51,9 +55,9 @@ class TestNonfiniteLossSoftHandling: def test_nan_loss_clears_progress_loss(self): b = _make_backend() - b._handle_event(_progress_event(step=1, loss=0.97)) + b._handle_event(_progress_event(step = 1, loss = 0.97)) assert b._progress.loss == pytest.approx(0.97) - b._handle_event(_progress_event(step=2, loss=float("nan"))) + b._handle_event(_progress_event(step = 2, loss = float("nan"))) # Stale finite loss must NOT leak through assert b._progress.loss is None # Run is not marked failed @@ -64,7 +68,7 @@ class TestNonfiniteLossSoftHandling: def test_inf_loss_clears_progress_loss(self): b = _make_backend() - b._handle_event(_progress_event(step=1, loss=float("inf"))) + b._handle_event(_progress_event(step = 1, loss = float("inf"))) assert b._progress.loss is None assert b._progress.error is None assert b._should_stop is False @@ -72,7 +76,7 @@ class TestNonfiniteLossSoftHandling: def test_negative_inf_loss_clears_progress_loss(self): b = _make_backend() - b._handle_event(_progress_event(step=1, loss=float("-inf"))) + b._handle_event(_progress_event(step = 1, loss = float("-inf"))) assert b._progress.loss is None assert b._progress.error is None assert b._should_stop is False @@ -82,12 +86,12 @@ class TestNonfiniteLossSoftHandling: """Subsequent NaN events must not re-fire the warning flag setter. The flag should already be True after the first NaN.""" b = _make_backend() - b._handle_event(_progress_event(step=1, loss=0.97)) - b._handle_event(_progress_event(step=2, loss=float("nan"))) + b._handle_event(_progress_event(step = 1, loss = 0.97)) + b._handle_event(_progress_event(step = 2, loss = float("nan"))) assert b._progress._nonfinite_loss_warned is True # Further NaN steps don't change anything we care about - b._handle_event(_progress_event(step=3, loss=float("nan"))) - b._handle_event(_progress_event(step=4, loss=float("nan"))) + b._handle_event(_progress_event(step = 3, loss = float("nan"))) + b._handle_event(_progress_event(step = 4, loss = float("nan"))) assert b._progress._nonfinite_loss_warned is True assert b._progress.loss is None assert b._progress.error is None @@ -97,10 +101,10 @@ class TestNonfiniteLossSoftHandling: """If a NaN step is followed by a finite step, progress.loss must reflect the new finite value (not stay stuck at None).""" b = _make_backend() - b._handle_event(_progress_event(step=1, loss=0.97)) - b._handle_event(_progress_event(step=2, loss=float("nan"))) + b._handle_event(_progress_event(step = 1, loss = 0.97)) + b._handle_event(_progress_event(step = 2, loss = float("nan"))) assert b._progress.loss is None - b._handle_event(_progress_event(step=3, loss=0.85)) + b._handle_event(_progress_event(step = 3, loss = 0.85)) assert b._progress.loss == pytest.approx(0.85) # Warning flag stays set (we don't reset it on recovery) assert b._progress._nonfinite_loss_warned is True diff --git a/studio/backend/tests/test_training_progress_stream_nan.py b/studio/backend/tests/test_training_progress_stream_nan.py new file mode 100644 index 0000000000..899527a04d --- /dev/null +++ b/studio/backend/tests/test_training_progress_stream_nan.py @@ -0,0 +1,129 @@ +# SPDX-License-Identifier: AGPL-3.0-only +# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0 + +"""The SSE progress stream must follow the live progress step during +non-finite-loss stretches (loss reported as null) instead of replaying the +last finite step/loss pair from the metric histories, which skip NaN steps.""" + +import asyncio +import json +import sys +import types + +import pytest + +if "structlog" not in sys.modules: + + class _DummyLogger: + def __getattr__(self, _name): + return lambda *args, **kwargs: None + + sys.modules["structlog"] = types.SimpleNamespace( + BoundLogger = _DummyLogger, + get_logger = lambda *args, **kwargs: _DummyLogger(), + ) + +import routes.training as rt + + +class _Progress: + def __init__(self): + self.step = 5 + self.total_steps = 10 + self.loss = None # cleared by the NaN honesty fix in core training + self.learning_rate = 8e-5 + self.epoch = 0.1 + self.grad_norm = None + self.num_tokens = None + self.eval_loss = None + self.elapsed_seconds = None + self.eta_seconds = None + + +class _FakeBackend: + """Finite history stops at step 2; live progress is at step 5 with NaN + (loss=None). Active for a few polls, then done.""" + + def __init__(self, active_polls = 2): + self.current_job_id = "job-1" + self.step_history = [1, 2] + self.loss_history = [2.0, 1.5] + self.lr_history = [1e-4, 9e-5] + self.eval_enabled = False + self._active_calls = 0 + self._active_polls = active_polls + self.trainer = types.SimpleNamespace(training_progress = _Progress()) + + def is_training_active(self): + self._active_calls += 1 + return self._active_calls <= self._active_polls + + +class _FakeRequest: + headers = {} + + +def _collect_events(response, timeout = 15): + async def _drain(): + chunks = [] + async for chunk in response.body_iterator: + chunks.append(chunk) + return "".join(c.decode() if isinstance(c, bytes) else c for c in chunks) + + return asyncio.run(asyncio.wait_for(_drain(), timeout)) + + +def _progress_payloads(raw): + payloads = [] + for block in raw.split("\n\n"): + lines = block.strip().splitlines() + data = next((l[6:] for l in lines if l.startswith("data: ")), None) + if data: + payloads.append(json.loads(data)) + return payloads + + +def test_stream_reports_live_step_with_null_loss_during_nan(monkeypatch): + backend = _FakeBackend(active_polls = 2) + monkeypatch.setattr(rt, "get_training_backend", lambda: backend) + + response = asyncio.run(rt.stream_training_progress(_FakeRequest(), current_subject = "tester")) + raw = _collect_events(response) + payloads = _progress_payloads(raw) + assert payloads, f"no SSE payloads parsed from: {raw!r}" + + live = [p for p in payloads if p.get("step") == 5] + assert live, ( + "stream never advanced to the live progress step during the NaN " + f"stretch; steps seen: {[p.get('step') for p in payloads]}" + ) + assert live[0]["loss"] is None + # The stale finite pair must not be re-emitted as the latest progress. + stale = [p for p in payloads if p.get("step") == 2 and p.get("loss") == 1.5] + assert not stale + + +def test_inactive_stream_completes_with_live_step_and_null_loss(monkeypatch): + # Fresh connection after the run already ended during a NaN stretch: the + # immediate complete event must not replay the stale finite pair either. + backend = _FakeBackend(active_polls = 0) + monkeypatch.setattr(rt, "get_training_backend", lambda: backend) + + response = asyncio.run(rt.stream_training_progress(_FakeRequest(), current_subject = "tester")) + payloads = _progress_payloads(_collect_events(response)) + final = payloads[-1] + assert final["step"] == 5 + assert final["loss"] is None + + +def test_stream_uses_finite_history_when_progress_in_sync(monkeypatch): + backend = _FakeBackend(active_polls = 2) + # Live progress agrees with the history tail: normal finite behavior. + backend.trainer.training_progress.step = 2 + backend.trainer.training_progress.loss = 1.5 + monkeypatch.setattr(rt, "get_training_backend", lambda: backend) + + response = asyncio.run(rt.stream_training_progress(_FakeRequest(), current_subject = "tester")) + payloads = _progress_payloads(_collect_events(response)) + finite = [p for p in payloads if p.get("step") == 2] + assert finite and finite[0]["loss"] == 1.5 diff --git a/studio/frontend/src/features/training/stores/training-runtime-store.ts b/studio/frontend/src/features/training/stores/training-runtime-store.ts index f1966f03fb..97fbd32d57 100644 --- a/studio/frontend/src/features/training/stores/training-runtime-store.ts +++ b/studio/frontend/src/features/training/stores/training-runtime-store.ts @@ -274,7 +274,10 @@ export const useTrainingRuntimeStore = create()((set) => ( jobId: payload.job_id || state.jobId, currentStep: step, totalSteps: Math.max(payload.total_steps, state.totalSteps), - currentLoss: currentLoss ?? state.currentLoss, + // A null loss at a new step means the backend reported a non-finite + // loss; clear the display instead of keeping the stale value. + currentLoss: + currentLoss ?? (step > state.currentStep ? null : state.currentLoss), currentLearningRate: currentLearningRate ?? state.currentLearningRate, progressPercent: payload.progress_percent, currentEpoch: payload.epoch ?? state.currentEpoch, diff --git a/studio/frontend/src/features/training/types/runtime.ts b/studio/frontend/src/features/training/types/runtime.ts index d30c3f75c2..5fd8286d22 100644 --- a/studio/frontend/src/features/training/types/runtime.ts +++ b/studio/frontend/src/features/training/types/runtime.ts @@ -90,7 +90,8 @@ export interface TrainingRuntimeState { currentStep: number; totalSteps: number; currentEpoch: number; - currentLoss: number; + // null while the latest reported loss is non-finite + currentLoss: number | null; currentLearningRate: number; progressPercent: number; elapsedSeconds: number | null; From f64c3c8abaecf017061ff5d9e149542bee29a25b Mon Sep 17 00:00:00 2001 From: Nilay <118994073+NilayYadav@users.noreply.github.com> Date: Thu, 11 Jun 2026 20:39:34 +0530 Subject: [PATCH 35/42] Studio: add `unsloth chat` CLI command (#6170) * Studio: add `unsloth chat` CLI command Interactive chat REPL on the shared Studio backend: trained-model picker when no model is given, /think and /compare toggles (adapter toggle on CUDA, side-by-side base-model load on MLX), markdown streaming, and connect-if-running Studio server mode so models stay warm across sessions and are shared with the UI. * fix settings * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * fix error handling and compare base precision * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * Fix chat CLI backend imports and GGUF drafter loading * Hide split thinking tags in chat CLI streams --------- Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com> Co-authored-by: Lee Jackson <130007945+Imagineer99@users.noreply.github.com> Co-authored-by: imagineer99 --- pyproject.toml | 1 + unsloth_cli/__init__.py | 2 + unsloth_cli/_inference.py | 417 +++++++++++++++++++++++ unsloth_cli/commands/chat.py | 340 ++++++++++++++++++ unsloth_cli/commands/inference.py | 87 ++--- unsloth_cli/tests/test_inference_chat.py | 401 ++++++++++++++++++++++ 6 files changed, 1208 insertions(+), 40 deletions(-) create mode 100644 unsloth_cli/_inference.py create mode 100644 unsloth_cli/commands/chat.py create mode 100644 unsloth_cli/tests/test_inference_chat.py diff --git a/pyproject.toml b/pyproject.toml index 79c14c907c..9f6465d3a3 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -26,6 +26,7 @@ classifiers = [ ] dependencies = [ "typer", + "rich", "pydantic", "pyyaml", "nest-asyncio", diff --git a/unsloth_cli/__init__.py b/unsloth_cli/__init__.py index e1047a6515..08d7e17e8c 100644 --- a/unsloth_cli/__init__.py +++ b/unsloth_cli/__init__.py @@ -10,6 +10,7 @@ from importlib.metadata import version as package_version, PackageNotFoundError from unsloth_cli.commands.train import train from unsloth_cli.commands.inference import inference +from unsloth_cli.commands.chat import chat from unsloth_cli.commands.export import export, list_checkpoints from unsloth_cli.commands.studio import ( run as studio_run, @@ -72,6 +73,7 @@ def main( app.command()(train) app.command()(inference) +app.command()(chat) app.command()(export) app.command("list-checkpoints")(list_checkpoints) app.add_typer(studio_app, name = "studio", help = "Unsloth Studio commands.") diff --git a/unsloth_cli/_inference.py b/unsloth_cli/_inference.py new file mode 100644 index 0000000000..5e734ee2a1 --- /dev/null +++ b/unsloth_cli/_inference.py @@ -0,0 +1,417 @@ +# SPDX-License-Identifier: AGPL-3.0-only +# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0 + +"""Model loading and streaming shared by `inference` and `chat`.""" + +import os +import re +import sys +from pathlib import Path +from typing import Optional + +import typer + +_THINK_OPEN = "" +_THINK_BLOCK = re.compile(rf"{re.escape(_THINK_OPEN)}.*?", re.DOTALL) + + +def ensure_studio_backend_path() -> None: + backend_dir = str(Path(__file__).resolve().parents[1] / "studio" / "backend") + if backend_dir not in sys.path: + sys.path.insert(0, backend_dir) + + +def configure_quiet_logging() -> None: + import logging + + import structlog + + # The CLI never configures structlog, so without this every backend INFO + # line prints. LOG_LEVEL is exported so the worker subprocess inherits it. + level_name = os.environ.setdefault("LOG_LEVEL", "WARNING").upper() + level = getattr(logging, level_name, logging.WARNING) + structlog.configure(wrapper_class = structlog.make_filtering_bound_logger(level)) + os.environ.setdefault("HF_HUB_DISABLE_PROGRESS_BARS", "1") + + +def visible_text(text: str, show_thinking: bool) -> str: + if show_thinking: + return text + text = _THINK_BLOCK.sub("", text) + # Hold back an unclosed trailing so reasoning never leaks mid-stream. + open_idx = text.find(_THINK_OPEN) + if open_idx != -1: + text = text[:open_idx] + max_prefix = min(len(text), len(_THINK_OPEN) - 1) + for size in range(max_prefix, 0, -1): + if _THINK_OPEN.startswith(text[-size:]): + return text[:-size] + return text + + +def stream_to_stdout(stream, show_thinking: bool) -> str: + # Backends yield the full text-so-far on each step (llama.cpp ends with a + # metadata dict, skipped); print the growing tail, return the raw text. + raw = "" + shown = "" + for chunk in stream: + if not isinstance(chunk, str): + continue + raw = chunk + rendered = visible_text(chunk, show_thinking) + delta = rendered[len(shown) :] + if delta: + sys.stdout.write(delta) + sys.stdout.flush() + shown = rendered + sys.stdout.write("\n") + sys.stdout.flush() + return raw + + +def stream_markdown(stream, show_thinking: bool, *, console) -> str: + from rich.live import Live + from rich.markdown import Markdown + from rich.text import Text + + raw = "" + with Live(console = console, refresh_per_second = 12, vertical_overflow = "visible") as live: + for chunk in stream: + if not isinstance(chunk, str): + continue + raw = chunk + visible = visible_text(chunk, show_thinking) + live.update(Markdown(visible) if visible.strip() else Text("")) + return raw + + +def collect_stream(stream, show_thinking: bool) -> str: + raw = "" + for chunk in stream: + if isinstance(chunk, str): + raw = chunk + return visible_text(raw, show_thinking) + + +def render_columns( + left_label: str, + left_text: str, + right_label: str, + right_text: str, + *, + console = None, +) -> None: + from rich import box + from rich.console import Console + from rich.table import Table + + table = Table(box = box.MINIMAL, expand = True, padding = (0, 1), pad_edge = False) + table.add_column(left_label, header_style = "bold yellow", ratio = 1, overflow = "fold") + table.add_column(right_label, header_style = "bold magenta", ratio = 1, overflow = "fold") + table.add_row(left_text or "", right_text or "") + (console or Console()).print(table) + + +class ChatBackend: + """Uniform stream()/close() over the llama-server and Unsloth backends.""" + + def __init__(self, kind: str, backend) -> None: + self._kind = kind # "gguf" | "unsloth" + self._backend = backend + + def stream( + self, + messages: list, + *, + system_prompt: str, + temperature: float, + top_p: float, + top_k: int, + max_new_tokens: int, + repetition_penalty: float, + enable_thinking: bool, + use_adapter: Optional[bool] = None, + ): + if self._kind == "gguf": + # llama-server takes the system prompt as the first message. + msgs = list(messages) + if system_prompt: + msgs = [{"role": "system", "content": system_prompt}, *msgs] + return self._backend.generate_chat_completion( + messages = msgs, + temperature = temperature, + top_p = top_p, + top_k = top_k, + max_tokens = max_new_tokens, + repetition_penalty = repetition_penalty, + enable_thinking = enable_thinking, + ) + gen_kwargs = dict( + messages = messages, + system_prompt = system_prompt, + temperature = temperature, + top_p = top_p, + top_k = top_k, + max_new_tokens = max_new_tokens, + repetition_penalty = repetition_penalty, + enable_thinking = enable_thinking, + ) + if use_adapter is not None: + return self._backend.generate_with_adapter_control( + use_adapter = use_adapter, **gen_kwargs + ) + return self._backend.generate_chat_response(**gen_kwargs) + + def close(self) -> None: + # Shut the worker down directly: the graceful unload_model waits for + # an ack that compare mode can swallow, hanging exit for minutes. + try: + if self._kind == "gguf": + self._backend.unload_model() + else: + self._backend._shutdown_subprocess(timeout = 2.0) + except Exception: + pass + + +def resolve_model_config(model: str, *, hf_token: Optional[str]): + ensure_studio_backend_path() + from utils.models import ModelConfig + + model_config = ModelConfig.from_identifier(model_id = model, hf_token = hf_token) + if not model_config: + typer.echo("Could not resolve model config", err = True) + raise typer.Exit(code = 1) + return model_config + + +def _load_gguf_backend(model_config, *, hf_token, max_seq_length): + ensure_studio_backend_path() + from core.inference.llama_cpp import LlamaCppBackend + + llama_backend = LlamaCppBackend() + common = dict( + hf_variant = model_config.gguf_variant, + model_identifier = model_config.identifier, + is_vision = model_config.is_vision, + n_ctx = max_seq_length, + ) + if model_config.gguf_hf_repo: + loaded = llama_backend.load_model( + hf_repo = model_config.gguf_hf_repo, hf_token = hf_token, **common + ) + else: + loaded = llama_backend.load_model( + gguf_path = model_config.gguf_file, + mmproj_path = model_config.gguf_mmproj_file, + mtp_draft_path = model_config.gguf_mtp_file, + **common, + ) + if not loaded: + typer.echo("Model load failed", err = True) + raise typer.Exit(code = 1) + return ChatBackend("gguf", llama_backend) + + +def load_chat_backend( + model: str, + *, + hf_token: Optional[str], + max_seq_length: int, + load_in_4bit: bool, + model_config = None, + fresh_backend: bool = False, +): + """Load `model` in-process: GGUF via llama-server, else the orchestrator. + + fresh_backend uses a private orchestrator so a second model (compare's + base column) can run alongside the main one. + """ + if model_config is None: + model_config = resolve_model_config(model, hf_token = hf_token) + + typer.echo(f"Loading {model}", err = True) + + if model_config.is_gguf: + return _load_gguf_backend(model_config, hf_token = hf_token, max_seq_length = max_seq_length) + + if fresh_backend: + ensure_studio_backend_path() + from core.inference import InferenceOrchestrator + backend = InferenceOrchestrator() + else: + ensure_studio_backend_path() + from core.inference import get_inference_backend + backend = get_inference_backend() + if not backend.load_model( + config = model_config, + max_seq_length = max_seq_length, + load_in_4bit = load_in_4bit, + hf_token = hf_token, + ): + typer.echo("Model load failed", err = True) + raise typer.Exit(code = 1) + return ChatBackend("unsloth", backend) + + +def find_studio_server(timeout: float = 0.4) -> Optional[str]: + import urllib.request + base = os.environ.get("UNSLOTH_STUDIO_URL", "http://127.0.0.1:8888").rstrip("/") + try: + with urllib.request.urlopen(f"{base}/api/health", timeout = timeout): + return base + except Exception: + return None + + +def _studio_token() -> Optional[str]: + """Self-issue a JWT: the CLI runs as the same OS user as the server, so it + signs with the same stored secret the server validates against.""" + try: + import studio.backend.core # noqa: F401 puts studio/backend on sys.path + + from studio.backend.auth import storage + from studio.backend.auth.authentication import create_access_token + + row = storage.get_connection().execute("SELECT username FROM auth_user LIMIT 1").fetchone() + return create_access_token(row[0], desktop = True) if row else None + except Exception: + return None + + +class HttpChatBackend: + """Chat against a running Studio server over its OpenAI-compatible API. + + close() leaves the model loaded on purpose — the next session (or the + UI) starts instantly. + """ + + def __init__(self, base_url: str, token: str) -> None: + self._base = base_url + self._token = token + + def _request( + self, + method: str, + path: str, + payload = None, + timeout = None, + ): + import json + import urllib.request + + request = urllib.request.Request( + self._base + path, + data = None if payload is None else json.dumps(payload).encode(), + headers = { + "Authorization": f"Bearer {self._token}", + "Content-Type": "application/json", + }, + method = method, + ) + return urllib.request.urlopen(request, timeout = timeout) + + def ensure_loaded(self, model: str, *, hf_token, max_seq_length, load_in_4bit) -> None: + typer.echo(f"Loading {model} on the Studio server", err = True) + try: + self._request( + "POST", + "/api/inference/load", + { + "model_path": model, + "hf_token": hf_token, + "max_seq_length": max_seq_length, + "load_in_4bit": load_in_4bit, + }, + ).close() + except Exception as exc: + typer.echo(f"Model load failed: {exc}", err = True) + raise typer.Exit(code = 1) + + def stream( + self, + messages: list, + *, + system_prompt: str, + temperature: float, + top_p: float, + top_k: int, + max_new_tokens: int, + repetition_penalty: float, + enable_thinking: bool, + use_adapter: Optional[bool] = None, + ): + import json + + msgs = list(messages) + if system_prompt: + msgs = [{"role": "system", "content": system_prompt}, *msgs] + resp = self._request( + "POST", + "/v1/chat/completions", + { + "model": "default", + "messages": msgs, + "stream": True, + "temperature": temperature, + "top_p": top_p, + "top_k": top_k, + "max_tokens": max_new_tokens, + "repetition_penalty": repetition_penalty, + "enable_thinking": enable_thinking, + }, + ) + + def cumulative(): + # Accumulate SSE deltas into the full-text-so-far convention the + # stream helpers expect. + text = "" + with resp: + for raw_line in resp: + line = raw_line.decode("utf-8", "replace").strip() + if not line.startswith("data:"): + continue + data = line[len("data:") :].strip() + if data == "[DONE]": + break + try: + parsed = json.loads(data) + except ValueError: + continue + if "error" in parsed: + raise RuntimeError( + f"Server error: {parsed['error'].get('message', 'Unknown server error')}" + ) + try: + delta = parsed["choices"][0]["delta"].get("content") + except (KeyError, IndexError): + continue + if not delta: + continue + text += delta + # An emoji can arrive split across two deltas as lone + # surrogate halves: hold back a trailing half, merge pairs. + visible = text + if "\ud800" <= visible[-1] <= "\udbff": + visible = visible[:-1] + yield visible.encode("utf-16", "surrogatepass").decode("utf-16", "replace") + + return cumulative() + + def close(self) -> None: + pass + + +def connect_studio_server(model: str, *, hf_token, max_seq_length, load_in_4bit): + """Backend on a running Studio server, or None (caller loads locally).""" + base_url = find_studio_server() + if not base_url: + return None + token = _studio_token() + if not token: + return None + backend = HttpChatBackend(base_url, token) + backend.ensure_loaded( + model, hf_token = hf_token, max_seq_length = max_seq_length, load_in_4bit = load_in_4bit + ) + return backend diff --git a/unsloth_cli/commands/chat.py b/unsloth_cli/commands/chat.py new file mode 100644 index 0000000000..c62916bc75 --- /dev/null +++ b/unsloth_cli/commands/chat.py @@ -0,0 +1,340 @@ +# SPDX-License-Identifier: AGPL-3.0-only +# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0 + +from typing import Optional + +import typer +from rich.console import Console + +from unsloth_cli._inference import ( + collect_stream, + configure_quiet_logging, + connect_studio_server, + ensure_studio_backend_path, + load_chat_backend, + render_columns, + resolve_model_config, + stream_markdown, + visible_text, +) + +_HELP = ( + "Commands: /exit (quit), /reset (clear history), " + "/think (toggle reasoning), /compare (base vs tuned), /help" +) + + +def _you_prompt(colors: bool) -> str: + # The prompt must go through input(), not a separate print — readline + # redraws erase anything they didn't draw, eating the label. GNU readline + # wants colors wrapped in \001/\002; libedit (macOS) prints those + # literally, so it gets raw ANSI. + try: + import readline + except ImportError: + return "\n\x1b[1;36mYou: \x1b[0m" if colors else "\nYou: " + libedit = ( + "libedit" in (readline.__doc__ or "") or getattr(readline, "backend", "") == "editline" + ) + if not colors: + return "\nYou: " + if libedit: + return "\n\x1b[1;36mYou: \x1b[0m" + return "\n\001\x1b[1;36m\002You: \001\x1b[0m\002" + + +def _compare_blocked_reason(model_config) -> Optional[str]: + if model_config.is_gguf: + return ( + "GGUF models can't toggle adapters — load a LoRA fine-tune " + "(transformers backend) to compare base vs tuned." + ) + if not model_config.is_lora: + return ( + "this isn't a LoRA adapter — compare turns the adapter off for the " + "'base' column, so there's nothing to compare against." + ) + return None + + +def _get_base_load_in_4bit(model_config) -> bool: + """Determine load_in_4bit for base model based on tuned adapter precision.""" + if not model_config.is_lora or not model_config.path: + # Fallback to default if not a LoRA or no path + return True + + try: + import json + from pathlib import Path + + adapter_cfg_path = Path(model_config.path) / "adapter_config.json" + if not adapter_cfg_path.exists(): + return True + + with open(adapter_cfg_path) as f: + adapter_cfg = json.load(f) + + training_method = adapter_cfg.get("unsloth_training_method") + if training_method == "lora": + return False + elif training_method == "qlora": + return True + elif not training_method: + # Fallback: check base model name for -bnb-4bit suffix + if model_config.base_model and "-bnb-4bit" not in model_config.base_model.lower(): + return False + return True + return True + except Exception: + return True + + +def _compare_needs_second_model() -> bool: + # MLX can't toggle the adapter off, so compare loads the base separately. + # detect_hardware() would print into the chat (and import torch), so + # probe its MLX condition quietly: Apple Silicon with mlx installed. + try: + from studio.backend.utils.hardware import hardware as hw + + if hw.DEVICE is not None: + return hw.DEVICE == hw.DeviceType.MLX + if not hw.is_apple_silicon(): + return False + import mlx.core # noqa: F401 + + return True + except Exception: + return False + + +def _pick_trained_model(console) -> str: + ensure_studio_backend_path() + from utils.models import scan_trained_models + + trained = scan_trained_models() + if not trained: + typer.echo( + "No trained models found in your outputs folder. " + "Pass a model id or path: `unsloth chat `.", + err = True, + ) + raise typer.Exit(code = 1) + + console.print("Your trained models (newest first):", style = "bold") + for i, (display_name, _, model_type) in enumerate(trained, 1): + console.print(f" {i}. {display_name} ({model_type})", markup = False) + + while True: + try: + raw = input(f"Chat with [1-{len(trained)}, Enter = 1]: ").strip() + except (EOFError, KeyboardInterrupt): + raise typer.Exit(code = 1) + if not raw: + return trained[0][1] + if raw.isdigit() and 1 <= int(raw) <= len(trained): + return trained[int(raw) - 1][1] + console.print(f"Pick a number between 1 and {len(trained)}.", style = "yellow") + + +def chat( + model: Optional[str] = typer.Argument( + None, help = "HF model id or local path. Omit to pick one of your trained models." + ), + hf_token: Optional[str] = typer.Option( + None, "--hf-token", envvar = "HF_TOKEN", help = "Hugging Face token if needed." + ), + temperature: float = typer.Option(0.7, "--temperature"), + top_p: float = typer.Option(0.9, "--top-p"), + top_k: int = typer.Option(40, "--top-k"), + max_new_tokens: int = typer.Option(512, "--max-new-tokens"), + repetition_penalty: float = typer.Option(1.1, "--repetition-penalty"), + system_prompt: str = typer.Option( + "", "--system-prompt", help = "Optional system prompt for the conversation." + ), + max_seq_length: int = typer.Option(4096, "--max-seq-length"), + load_in_4bit: bool = typer.Option(True, "--load-in-4bit/--no-load-in-4bit"), + think: bool = typer.Option( + False, + "--think/--no-think", + help = "Start with the model's reasoning shown. Toggle live with /think.", + ), + compare: bool = typer.Option( + False, + "--compare/--no-compare", + help = "Answer each prompt twice — base vs fine-tuned — side by side. " + "Needs a LoRA adapter. Toggle live with /compare.", + ), + verbose: bool = typer.Option( + False, "--verbose", "-v", help = "Show backend and llama-server logs." + ), + no_server: bool = typer.Option( + False, + "--no-server", + help = "Load the model in-process even if a Studio server is running.", + ), +): + """Start an interactive chat with a model (loads once, stays warm).""" + if not verbose: + configure_quiet_logging() + + console = Console() + err = Console(stderr = True) + + if model is None: + model = _pick_trained_model(console) + + # Resolve first so --compare can be rejected before the slow load. + model_config = resolve_model_config(model, hf_token = hf_token) + compare_blocked = _compare_blocked_reason(model_config) + if compare and compare_blocked: + err.print(f"--compare unavailable: {compare_blocked}", style = "red", markup = False) + raise typer.Exit(code = 1) + + load_opts = dict(hf_token = hf_token, max_seq_length = max_seq_length, load_in_4bit = load_in_4bit) + + # Prefer a running Studio server: instant starts, model shared with the UI. + chat_backend = None if no_server else connect_studio_server(model, **load_opts) + server_mode = chat_backend is not None + if server_mode: + console.print( + "(Studio server connected — model stays warm after /exit)", + style = "bright_black", + ) + else: + chat_backend = load_chat_backend(model, model_config = model_config, **load_opts) + + name = model_config.display_name or model + show_thinking = think + compare_mode = compare + messages = [] + + # Compare's base column: server mode keeps the tuned model remote and + # loads the base locally; local MLX (no adapter toggle) does the same; + # local CUDA just toggles the adapter on the one loaded model. + dual_compare = compare_blocked is None and (server_mode or _compare_needs_second_model()) + base_backend = None + + def load_base_for_compare(): + nonlocal base_backend + if base_backend is not None: + return True + base_id = model_config.base_model + if not base_id: + console.print( + "(compare unavailable: this adapter doesn't record its base model)", + style = "yellow", + ) + return False + console.print( + f"(loading base model {base_id} for compare — keeps two models in memory)", + style = "bright_black", + markup = False, + ) + try: + # Use the same precision as the tuned model for fair comparison + base_load_opts = dict(load_opts) # Copy original options + base_load_opts["load_in_4bit"] = _get_base_load_in_4bit(model_config) + base_backend = load_chat_backend(base_id, fresh_backend = True, **base_load_opts) + except Exception as exc: + err.print(f"(base model load failed: {exc})", style = "red", markup = False) + return False + return True + + if compare and dual_compare and not load_base_for_compare(): + raise typer.Exit(code = 1) + + def generate(backend = None, use_adapter = None): + # Reads messages and show_thinking live, so /reset and /think apply. + return (backend or chat_backend).stream( + messages, + system_prompt = system_prompt, + temperature = temperature, + top_p = top_p, + top_k = top_k, + max_new_tokens = max_new_tokens, + repetition_penalty = repetition_penalty, + enable_thinking = show_thinking, + use_adapter = use_adapter, + ) + + console.print() + console.print(f"Chatting with {name}", style = "bold green", markup = False) + console.print(_HELP, style = "bright_black") + + # legacy_windows: pre-VT consoles print raw ANSI as ←[1;36m garbage. + you_prompt = _you_prompt(console.is_terminal and not console.legacy_windows) + assistant_label = "[bold magenta]Assistant:[/bold magenta]" + + try: + while True: + try: + user = input(you_prompt).strip() + except (EOFError, KeyboardInterrupt): + console.print() + break + + if not user: + continue + if user in ("/exit", "/quit"): + break + if user == "/reset": + messages = [] + console.print("(history cleared)", style = "bright_black") + continue + if user == "/think": + show_thinking = not show_thinking + state = "on" if show_thinking else "off" + console.print(f"(thinking {state})", style = "bright_black") + continue + if user == "/compare": + if compare_blocked: + console.print(f"(compare unavailable: {compare_blocked})", style = "yellow") + continue + if not compare_mode and dual_compare and not load_base_for_compare(): + continue + compare_mode = not compare_mode + state = "on" if compare_mode else "off" + console.print(f"(compare {state})", style = "bright_black") + continue + if user in ("/help", "/?"): + console.print(_HELP, style = "bright_black") + continue + + messages.append({"role": "user", "content": user}) + + try: + if compare_mode: + console.print("(comparing base vs tuned…)", style = "bright_black") + if dual_compare: + base_text = collect_stream(generate(backend = base_backend), show_thinking) + tuned_text = collect_stream(generate(), show_thinking) + else: + base_text = collect_stream(generate(use_adapter = False), show_thinking) + tuned_text = collect_stream(generate(use_adapter = True), show_thinking) + console.print() + render_columns( + "base", base_text, f"{name} (tuned)", tuned_text, console = console + ) + # History continues as the tuned model; base is just the reference. + answer = tuned_text + else: + console.print(assistant_label) + answer = stream_markdown(generate(), show_thinking, console = console) + except KeyboardInterrupt: + # Ctrl-C aborts this answer only; drop the unanswered turn. + console.print("\n(interrupted)", style = "bright_black") + messages.pop() + continue + except Exception as exc: + err.print(f"\n(error: {exc})", style = "red", markup = False) + messages.pop() + continue + + messages.append( + {"role": "assistant", "content": visible_text(answer, show_thinking = False)} + ) + finally: + chat_backend.close() + if base_backend is not None: + base_backend.close() + err.print("\nBye.", style = "bright_black") diff --git a/unsloth_cli/commands/inference.py b/unsloth_cli/commands/inference.py index 18de68c4d9..5dbc32c7d2 100644 --- a/unsloth_cli/commands/inference.py +++ b/unsloth_cli/commands/inference.py @@ -1,11 +1,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 -import sys from typing import Optional import typer +from unsloth_cli._inference import ( + configure_quiet_logging, + connect_studio_server, + load_chat_backend, + stream_to_stdout, +) + def inference( model: str = typer.Argument(..., help = "HF model id or local path."), @@ -25,45 +31,46 @@ def inference( ), max_seq_length: int = typer.Option(2048, "--max-seq-length"), load_in_4bit: bool = typer.Option(True, "--load-in-4bit/--no-load-in-4bit"), + think: bool = typer.Option( + False, + "--think/--no-think", + help = "Show the model's reasoning. Off by default so reasoning " + "models answer directly instead of spending the token budget thinking.", + ), + verbose: bool = typer.Option( + False, + "--verbose", + "-v", + help = "Show backend and llama-server logs (otherwise only the answer).", + ), + no_server: bool = typer.Option( + False, + "--no-server", + help = "Load the model in-process even if a Studio server is running.", + ), ): """Run a single inference using the specified model.""" - from studio.backend.core import ModelConfig, get_inference_backend + if not verbose: + configure_quiet_logging() - inference_backend = get_inference_backend() - model_config = ModelConfig.from_ui_selection( - dropdown_value = model, search_value = None, hf_token = hf_token, is_lora = False - ) - if not model_config: - typer.echo("Could not resolve model config", err = True) - raise typer.Exit(code = 1) - - if not inference_backend.load_model( - config = model_config, - max_seq_length = max_seq_length, - load_in_4bit = load_in_4bit, - hf_token = hf_token, - ): - typer.echo("Model load failed", err = True) - raise typer.Exit(code = 1) - - messages = [{"role": "user", "content": prompt}] - stream = inference_backend.generate_chat_response( - messages = messages, - system_prompt = system_prompt, - temperature = temperature, - top_p = top_p, - top_k = top_k, - max_new_tokens = max_new_tokens, - repetition_penalty = repetition_penalty, - ) - - typer.echo("Assistant:", nl = True) - previous = "" - for chunk in stream: - delta = chunk[len(previous) :] - if delta: - sys.stdout.write(delta) - sys.stdout.flush() - previous = chunk - sys.stdout.write("\n") - sys.stdout.flush() + # A running Studio server keeps the model warm between runs, which is + # exactly what a one-shot command wants. + load_opts = dict(hf_token = hf_token, max_seq_length = max_seq_length, load_in_4bit = load_in_4bit) + chat_backend = None if no_server else connect_studio_server(model, **load_opts) + if chat_backend is None: + chat_backend = load_chat_backend(model, **load_opts) + try: + stream = chat_backend.stream( + [{"role": "user", "content": prompt}], + system_prompt = system_prompt, + temperature = temperature, + top_p = top_p, + top_k = top_k, + max_new_tokens = max_new_tokens, + repetition_penalty = repetition_penalty, + enable_thinking = think, + ) + typer.echo("Assistant:") + stream_to_stdout(stream, show_thinking = think) + finally: + chat_backend.close() diff --git a/unsloth_cli/tests/test_inference_chat.py b/unsloth_cli/tests/test_inference_chat.py new file mode 100644 index 0000000000..629cb46405 --- /dev/null +++ b/unsloth_cli/tests/test_inference_chat.py @@ -0,0 +1,401 @@ +# SPDX-License-Identifier: AGPL-3.0-only +# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0 + +"""Tests for the `unsloth chat` / `unsloth inference` CLI — fakes only, no model loads.""" + +from __future__ import annotations + +import inspect +import sys +import types +from pathlib import Path + +_REPO_ROOT = Path(__file__).resolve().parents[2] +if str(_REPO_ROOT) not in sys.path: + sys.path.insert(0, str(_REPO_ROOT)) + + +import typer +from rich.console import Console +from typer.testing import CliRunner + +import unsloth_cli.commands.chat as chatmod +from unsloth_cli._inference import ( + ChatBackend, + HttpChatBackend, + collect_stream, + render_columns, + visible_text, +) + + +class _FakeConfig: + is_gguf = False + is_lora = True + display_name = "fake-model" + base_model = "fake/base" + path = None + + +def _chat_app(): + cli = typer.Typer() + cli.command()(chatmod.chat) + return cli + + +def test_visible_text_passthrough_when_shown(): + text = "reasoninganswer" + assert visible_text(text, show_thinking = True) == text + + +def test_visible_text_strips_closed_think_block(): + text = "step 1\nstep 2The answer is 42." + assert visible_text(text, show_thinking = False) == "The answer is 42." + + +def test_visible_text_holds_unclosed_think(): + # An open is held back so partial reasoning never leaks mid-stream. + assert visible_text("still thinking", show_thinking = False) == "" + assert visible_text("done.more thinking", show_thinking = False) == "done." + + +def test_visible_text_holds_partial_think_prefix(): + # Streams are cumulative, so the opening tag can arrive as "<", "". Hold possible tag prefixes until they are disambiguated. + assert visible_text("<", show_thinking = False) == "" + assert visible_text("rhel", "rhello"]) + assert collect_stream(stream, show_thinking = False) == "hello" + + +def test_render_columns_emits_both_answers_with_separator(capsys): + render_columns("base", "alpha", "tuned", "beta") + out = capsys.readouterr().out + assert "base" in out and "tuned" in out + assert "alpha" in out and "beta" in out + assert "│" in out + + +def test_you_prompt_matches_readline_backend(monkeypatch): + gnu = types.ModuleType("readline") + gnu.__doc__ = "Importing this module enables command line editing using GNU readline." + monkeypatch.setitem(sys.modules, "readline", gnu) + prompt = chatmod._you_prompt(colors = True) + assert "You: " in prompt and "\001" in prompt + + libedit = types.ModuleType("readline") + libedit.__doc__ = "Importing this module enables command line editing using libedit readline." + monkeypatch.setitem(sys.modules, "readline", libedit) + assert chatmod._you_prompt(colors = True) == "\n\x1b[1;36mYou: \x1b[0m" + assert chatmod._you_prompt(colors = False) == "\nYou: " + + # Windows: no readline module at all; the console's own line editing + # handles backspace, so plain ANSI color (no markers) is safe. + monkeypatch.setitem(sys.modules, "readline", None) + assert chatmod._you_prompt(colors = True) == "\n\x1b[1;36mYou: \x1b[0m" + assert chatmod._you_prompt(colors = False) == "\nYou: " + + +def test_chat_registered_on_app(): + from unsloth_cli import app + + # cmd.name is None until typer resolves it from the callback name. + names = {(cmd.name or cmd.callback.__name__) for cmd in app.registered_commands} + assert "chat" in names + + +def test_chat_exits_cleanly_on_slash_exit(monkeypatch): + closed = [] + + class _FakeChatBackend: + def stream(self, *a, **k): + return iter(["hello"]) + + def close(self): + closed.append(True) + + monkeypatch.setattr(chatmod, "resolve_model_config", lambda *a, **k: _FakeConfig()) + monkeypatch.setattr(chatmod, "load_chat_backend", lambda *a, **k: _FakeChatBackend()) + monkeypatch.setattr(chatmod, "_compare_needs_second_model", lambda: False) + monkeypatch.setattr(chatmod, "connect_studio_server", lambda *a, **k: None) + + runner = CliRunner() + for args in (["fake-model"], ["fake-model", "--compare"]): + closed.clear() + result = runner.invoke(_chat_app(), args, input = "hi\n/exit\n") + assert result.exit_code == 0, result.output + assert closed == [True] + assert "Bye." in result.output + # The prompt must go through input() (readline-safe), not a print. + assert "You: " in result.output + assert "You: You:" not in result.output + + +def test_pick_trained_model_lists_and_selects(monkeypatch): + fake_models = types.ModuleType("utils.models") + fake_models.scan_trained_models = lambda: [ + ("run-new", "outputs/run-new", "lora"), + ("run-old", "outputs/run-old", "merged"), + ] + monkeypatch.setitem(sys.modules, "utils.models", fake_models) + + monkeypatch.setattr("builtins.input", lambda prompt = "": "2") + assert chatmod._pick_trained_model(Console()) == "outputs/run-old" + + monkeypatch.setattr("builtins.input", lambda prompt = "": "") + assert chatmod._pick_trained_model(Console()) == "outputs/run-new" + + +def test_chat_no_arg_chats_with_picked_trained_model(monkeypatch): + class _FakeChatBackend: + def stream(self, *a, **k): + return iter(["hello"]) + + def close(self): + pass + + resolved = [] + monkeypatch.setattr(chatmod, "_pick_trained_model", lambda console: "outputs/run-42") + monkeypatch.setattr( + chatmod, + "resolve_model_config", + lambda model, **k: (resolved.append(model), _FakeConfig())[1], + ) + monkeypatch.setattr(chatmod, "load_chat_backend", lambda *a, **k: _FakeChatBackend()) + monkeypatch.setattr(chatmod, "_compare_needs_second_model", lambda: False) + monkeypatch.setattr(chatmod, "connect_studio_server", lambda *a, **k: None) + + result = CliRunner().invoke(_chat_app(), [], input = "/exit\n") + assert result.exit_code == 0, result.output + assert resolved == ["outputs/run-42"] + + +def test_find_studio_server_none_when_not_running(monkeypatch): + import urllib.request + + from unsloth_cli import _inference + + def refuse(*a, **k): + raise OSError("connection refused") + + monkeypatch.setattr(urllib.request, "urlopen", refuse) + assert _inference.find_studio_server() is None + + +class _FakeSSEResponse: + def __init__(self, lines): + self._lines = lines + + def __iter__(self): + return iter(self._lines) + + def __enter__(self): + return self + + def __exit__(self, *exc): + return False + + +def test_http_backend_streams_cumulative_text(monkeypatch): + backend = HttpChatBackend("http://localhost:8888", "token") + response = _FakeSSEResponse( + [ + b'data: {"choices":[{"delta":{"content":"He"}}]}\n', + b"\n", + b'data: {"choices":[{"delta":{"content":"llo"}}]}\n', + b"data: [DONE]\n", + ] + ) + monkeypatch.setattr(backend, "_request", lambda *a, **k: response) + + out = list(backend.stream([{"role": "user", "content": "hi"}], **_STREAM_KWARGS)) + assert out == ["He", "Hello"] + + +def test_http_backend_merges_emoji_split_across_deltas(monkeypatch): + backend = HttpChatBackend("http://localhost:8888", "token") + response = _FakeSSEResponse( + [ + b'data: {"choices":[{"delta":{"content":"hi "}}]}\n', + b'data: {"choices":[{"delta":{"content":"\\ud83d"}}]}\n', + b'data: {"choices":[{"delta":{"content":"\\ude0a"}}]}\n', + b"data: [DONE]\n", + ] + ) + monkeypatch.setattr(backend, "_request", lambda *a, **k: response) + + out = list(backend.stream([{"role": "user", "content": "hi"}], **_STREAM_KWARGS)) + # The lone high surrogate is held back, then merged with its other half. + assert out == ["hi ", "hi ", "hi 😊"] + + +def test_chat_prefers_running_studio_server(monkeypatch): + closed = [] + + class _FakeHttpBackend: + def stream(self, *a, **k): + return iter(["hello"]) + + def close(self): + closed.append("http") + + local_loads = [] + monkeypatch.setattr(chatmod, "resolve_model_config", lambda *a, **k: _FakeConfig()) + monkeypatch.setattr(chatmod, "connect_studio_server", lambda *a, **k: _FakeHttpBackend()) + monkeypatch.setattr(chatmod, "load_chat_backend", lambda *a, **k: local_loads.append(1)) + monkeypatch.setattr(chatmod, "_compare_needs_second_model", lambda: False) + + result = CliRunner().invoke(_chat_app(), ["fake-model"], input = "hi\n/exit\n") + + assert result.exit_code == 0, result.output + assert local_loads == [] + assert "stays warm" in result.output + assert closed == ["http"] + + +def test_chat_server_mode_compare_loads_base_locally(monkeypatch): + streamed, closed, base_loads = [], [], [] + + class _FakeHttpBackend: + def stream(self, *a, **k): + streamed.append("tuned") + return iter(["tuned-answer"]) + + def close(self): + closed.append("http") + + class _FakeBaseBackend: + def stream(self, *a, **k): + streamed.append("base") + return iter(["base-answer"]) + + def close(self): + closed.append("base") + + def fake_local_load(model, **kwargs): + base_loads.append((model, kwargs.get("fresh_backend", False))) + return _FakeBaseBackend() + + monkeypatch.setattr(chatmod, "resolve_model_config", lambda *a, **k: _FakeConfig()) + monkeypatch.setattr(chatmod, "connect_studio_server", lambda *a, **k: _FakeHttpBackend()) + monkeypatch.setattr(chatmod, "load_chat_backend", fake_local_load) + + result = CliRunner().invoke(_chat_app(), ["tuned-run"], input = "/compare\nhi\n/exit\n") + + assert result.exit_code == 0, result.output + assert "(compare on)" in result.output + # Only the base model loaded locally, on its own private backend. + assert base_loads == [("fake/base", True)] + assert streamed == ["base", "tuned"] + assert set(closed) == {"http", "base"} + + +def test_chat_compare_on_mlx_loads_base_model_side_by_side(monkeypatch): + loads, streamed, closed = [], [], [] + + class _FakeLocalBackend: + def __init__(self, role): + self.role = role + + def stream(self, *a, **k): + streamed.append((self.role, k.get("use_adapter"))) + return iter([f"{self.role}-answer"]) + + def close(self): + closed.append(self.role) + + def fake_load(model, **kwargs): + fresh = kwargs.get("fresh_backend", False) + loads.append((model, fresh)) + return _FakeLocalBackend("base" if fresh else "tuned") + + monkeypatch.setattr(chatmod, "resolve_model_config", lambda *a, **k: _FakeConfig()) + monkeypatch.setattr(chatmod, "load_chat_backend", fake_load) + monkeypatch.setattr(chatmod, "_compare_needs_second_model", lambda: True) + monkeypatch.setattr(chatmod, "connect_studio_server", lambda *a, **k: None) + + result = CliRunner().invoke(_chat_app(), ["tuned-run", "--compare"], input = "hi\n/exit\n") + + assert result.exit_code == 0, result.output + assert loads == [("tuned-run", False), ("fake/base", True)] + # Both models answered the turn, via plain generation (no adapter toggle). + assert ("base", None) in streamed and ("tuned", None) in streamed + assert set(closed) == {"tuned", "base"} From 004577c9cd4b156d02e9d10b1ede822c2822e22f Mon Sep 17 00:00:00 2001 From: Nilay <118994073+NilayYadav@users.noreply.github.com> Date: Thu, 11 Jun 2026 20:47:22 +0530 Subject: [PATCH 36/42] studio: show MCP "Import config" on the add-server form (#6030) * studio: import MCP servers from a config file * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * import config' on the add-server form * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * fix: defensively handle MCP config imports * fix: address MCP import review follow-ups * fix: preserve apostrophes in Windows MCP commands * fix: preserve apostrophe-wrapped Windows MCP args * fix: align Windows MCP parsing with list2cmdline * fix: preserve explicit MCP remote transport intent * fix: trim MCP remote URLs before transport checks --------- Co-authored-by: Roland Tannous Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com> Co-authored-by: Roland Tannous <115670425+rolandtannous@users.noreply.github.com> Co-authored-by: Lee Jackson <130007945+Imagineer99@users.noreply.github.com> Co-authored-by: imagineer99 --- studio/backend/core/inference/mcp_client.py | 81 ++++- .../core/inference/mcp_config_import.py | 169 +++++++++ studio/backend/models/mcp_servers.py | 11 + studio/backend/routes/mcp_servers.py | 43 +++ .../backend/tests/test_mcp_config_import.py | 341 ++++++++++++++++++ .../src/features/chat/api/mcp-servers-api.ts | 15 + .../features/chat/chat-mcp-servers-dialog.tsx | 88 ++++- 7 files changed, 739 insertions(+), 9 deletions(-) create mode 100644 studio/backend/core/inference/mcp_config_import.py create mode 100644 studio/backend/tests/test_mcp_config_import.py diff --git a/studio/backend/core/inference/mcp_client.py b/studio/backend/core/inference/mcp_client.py index 9a7fd1f074..a0d79bbdf4 100644 --- a/studio/backend/core/inference/mcp_client.py +++ b/studio/backend/core/inference/mcp_client.py @@ -25,16 +25,85 @@ def is_stdio(address: str) -> bool: return not address.strip().lower().startswith(("http://", "https://")) +def _split_windows_command_line(address: str) -> list[str]: + """Parse a Windows command line using the same backslash/quote rules that + subprocess.list2cmdline() writes. This keeps trailing backslashes before a + closing quote from being doubled in the resulting argv.""" + parts: list[str] = [] + current: list[str] = [] + in_quotes = False + backslashes = 0 + arg_started = False + i = 0 + + while i < len(address): + ch = address[i] + if ch == "\\": + backslashes += 1 + i += 1 + continue + if ch == '"': + current.extend("\\" * (backslashes // 2)) + if backslashes % 2: + current.append('"') + else: + in_quotes = not in_quotes + arg_started = True + backslashes = 0 + i += 1 + continue + if ch.isspace() and not in_quotes: + if backslashes: + current.extend("\\" * backslashes) + arg_started = True + backslashes = 0 + if arg_started or current: + parts.append("".join(current)) + current = [] + arg_started = False + i += 1 + while i < len(address) and address[i].isspace(): + i += 1 + continue + if backslashes: + current.extend("\\" * backslashes) + arg_started = True + backslashes = 0 + current.append(ch) + arg_started = True + i += 1 + + if backslashes: + current.extend("\\" * backslashes) + arg_started = True + if in_quotes: + raise ValueError("No closing quotation") + if arg_started or current: + parts.append("".join(current)) + return parts + + def parse_stdio_command(address: str) -> list[str]: """Split a stdio command line into argv. Shared by route validation and the transport so both agree on quoting (notably Windows backslash paths).""" posix = sys.platform != "win32" - parts = shlex.split(address, posix = posix) - if not posix: - # posix=False keeps backslash paths but also keeps surrounding quotes; - # strip a matched pair so argv reaches the subprocess clean. - parts = [p[1:-1] if len(p) >= 2 and p[0] == p[-1] and p[0] in "\"'" else p for p in parts] - return parts + if posix: + return shlex.split(address, posix = posix) + if address.lstrip().startswith("'"): + raise ValueError("Single-quoted executables are not supported on Windows") + return _split_windows_command_line(address) + + +def join_stdio_command(parts: list[str]) -> str: + """Inverse of parse_stdio_command: join argv into a single command string + that parse_stdio_command() splits back into ``parts`` on this platform. + Config files (issue #5936) carry structured command + args; storage holds + one string in the url field. Windows uses list2cmdline so spaced/backslash + paths round-trip through the posix=False quote-strip; posix uses shlex.""" + if sys.platform == "win32": + import subprocess + return subprocess.list2cmdline(parts) + return shlex.join(parts) def stdio_mcp_enabled() -> bool: diff --git a/studio/backend/core/inference/mcp_config_import.py b/studio/backend/core/inference/mcp_config_import.py new file mode 100644 index 0000000000..534ed84c1d --- /dev/null +++ b/studio/backend/core/inference/mcp_config_import.py @@ -0,0 +1,169 @@ +# SPDX-License-Identifier: AGPL-3.0-only +# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0 + +"""Parse a standard ``mcpServers`` JSON config (Claude Desktop / Cursor / Cline +/ VS Code) into entries the existing MCP storage understands. See issue #5936. + +A stdio entry (``command`` + ``args`` + ``env``) is joined into the single +command string the ``url`` field already stores; a remote entry (``url`` + +``headers``) maps straight through. Parsing never raises on a single bad entry: +it returns ``(entries, errors)`` so one malformed server can't sink the import. +""" + +from __future__ import annotations + +from dataclasses import dataclass +from typing import Optional + +from core.inference.mcp_client import join_stdio_command + +_SCALAR = (str, int, float, bool) +_UNSUPPORTED_STDIO_FIELDS = ("cwd", "envFile") +_UNSUPPORTED_TIMEOUT_FIELDS = ("timeout", "timeoutMs", "timeoutSeconds") +_HTTP_REMOTE_TYPES = ("http", "streamableHttp") + + +@dataclass +class ParsedMcpEntry: + display_name: str + url: str # joined command (stdio) or http(s) url (remote) + headers: Optional[dict[str, str]] # env vars (stdio) or http headers (remote) + is_stdio: bool + is_enabled: bool = True + use_oauth: bool = False + + +def _coerce_str_dict(value: dict) -> dict[str, str]: + return {str(k): str(v) for k, v in value.items()} + + +def _has_variable_reference(value: object) -> bool: + if isinstance(value, str): + return "${" in value + if isinstance(value, list): + return any(_has_variable_reference(item) for item in value) + if isinstance(value, dict): + return any(_has_variable_reference(item) for item in value.values()) + return False + + +def _has_null_value(value: object) -> bool: + return isinstance(value, dict) and any(item is None for item in value.values()) + + +def _enabled_from_spec(label: str, spec: dict) -> tuple[Optional[bool], Optional[str]]: + disabled = spec.get("disabled") + if disabled is None: + return True, None + if not isinstance(disabled, bool): + return None, f"{label}: 'disabled' must be true or false." + return not disabled, None + + +def _parse_entry(name: str, spec: object) -> tuple[Optional[ParsedMcpEntry], Optional[str]]: + label = str(name).strip() + if not label: + return None, "Server entry has an empty name." + if not isinstance(spec, dict): + return None, f"{label}: entry must be an object." + if _has_variable_reference(spec): + return None, f"{label}: VS Code variable references are not supported by import." + + is_enabled, error = _enabled_from_spec(label, spec) + if error: + return None, error + + has_command = bool(spec.get("command")) + has_url = bool(spec.get("url")) + if has_command and has_url: + return None, f"{label}: entry has both 'command' and 'url'; use one." + if not has_command and not has_url: + return None, f"{label}: entry needs a 'command' (stdio) or 'url' (remote)." + + if has_command: + command = spec["command"] + if not isinstance(command, str): + return None, f"{label}: 'command' must be a string." + entry_type = spec.get("type") + if entry_type is not None and entry_type != "stdio": + return None, f"{label}: stdio entry has unsupported type {entry_type!r}." + sandbox_enabled = spec.get("sandboxEnabled") + if sandbox_enabled is not None and not isinstance(sandbox_enabled, bool): + return None, f"{label}: 'sandboxEnabled' must be true or false." + if sandbox_enabled: + return None, f"{label}: sandboxed stdio servers cannot be preserved by import." + unsupported = [field for field in _UNSUPPORTED_STDIO_FIELDS if spec.get(field) is not None] + if unsupported: + return None, f"{label}: import cannot preserve {', '.join(unsupported)}." + if spec.get("oauth") is not None: + return None, f"{label}: 'oauth' is only supported for remote servers." + args = spec.get("args") or [] + if not isinstance(args, list) or not all(isinstance(a, _SCALAR) for a in args): + return None, f"{label}: 'args' must be a list of strings." + env = spec.get("env") + if env is not None and not isinstance(env, dict): + return None, f"{label}: 'env' must be an object." + if _has_null_value(env): + return None, f"{label}: null environment values are not supported by import." + url = join_stdio_command([command, *(str(a) for a in args)]) + headers = _coerce_str_dict(env) if env else None + return ParsedMcpEntry(label, url, headers, True, is_enabled = is_enabled), None + + url = spec["url"] + if not isinstance(url, str): + return None, f"{label}: 'url' must be a string." + url = url.strip() + entry_type = spec.get("type") + if entry_type is not None and entry_type not in (*_HTTP_REMOTE_TYPES, "sse"): + return None, f"{label}: remote entry has unsupported type {entry_type!r}." + unsupported_timeout = [ + field for field in _UNSUPPORTED_TIMEOUT_FIELDS if spec.get(field) is not None + ] + if unsupported_timeout: + return None, f"{label}: import cannot preserve {', '.join(unsupported_timeout)}." + url_infers_sse = url.rstrip("/").endswith("/sse") + if entry_type == "sse" and not url_infers_sse: + return None, f"{label}: explicit SSE transport cannot be preserved for this URL." + if entry_type in _HTTP_REMOTE_TYPES and url_infers_sse: + return None, f"{label}: explicit HTTP transport cannot be preserved for this URL." + oauth_raw = spec.get("oauth") + if oauth_raw is not None and not isinstance(oauth_raw, dict): + return None, f"{label}: 'oauth' must be an object." + headers_raw = spec.get("headers") + if headers_raw is not None and not isinstance(headers_raw, dict): + return None, f"{label}: 'headers' must be an object." + if _has_null_value(headers_raw): + return None, f"{label}: null header values are not supported by import." + headers = _coerce_str_dict(headers_raw) if headers_raw else None + return ParsedMcpEntry( + label, + url, + headers, + False, + is_enabled = is_enabled, + use_oauth = oauth_raw is not None, + ), None + + +def parse_mcp_config(config: object) -> tuple[list[ParsedMcpEntry], list[str]]: + """Parse a Claude-Desktop/Cursor/Cline/VS Code config. Accepts the + ``mcpServers`` key (primary) or ``servers`` (VS Code alias). Returns + ``(entries, errors)``; a bad entry adds an error rather than raising.""" + if not isinstance(config, dict): + return [], ["Config must be a JSON object."] + servers_key = "mcpServers" if "mcpServers" in config else "servers" + servers = config.get(servers_key) + if servers is None: + return [], ["Config has no 'mcpServers' (or 'servers') object."] + if not isinstance(servers, dict): + return [], [f"'{servers_key}' must be an object mapping name -> server."] + + entries: list[ParsedMcpEntry] = [] + errors: list[str] = [] + for name, spec in servers.items(): + entry, error = _parse_entry(name, spec) + if error: + errors.append(error) + elif entry: + entries.append(entry) + return entries, errors diff --git a/studio/backend/models/mcp_servers.py b/studio/backend/models/mcp_servers.py index c696eb0faa..606c2423bf 100644 --- a/studio/backend/models/mcp_servers.py +++ b/studio/backend/models/mcp_servers.py @@ -44,3 +44,14 @@ class McpServerProbeResult(BaseModel): ok: bool tool_count: int = 0 error: Optional[str] = None + + +class McpServerImportRequest(BaseModel): + # A standard mcpServers JSON config (Claude Desktop / Cursor / Cline / VS Code). + config: dict + + +class McpServerImportResult(BaseModel): + created: list[McpServerResponse] = Field(default_factory = list) + skipped: list[str] = Field(default_factory = list) # display names skipped as duplicates + errors: list[str] = Field(default_factory = list) diff --git a/studio/backend/routes/mcp_servers.py b/studio/backend/routes/mcp_servers.py index 098c64fe39..3001c6b7c9 100644 --- a/studio/backend/routes/mcp_servers.py +++ b/studio/backend/routes/mcp_servers.py @@ -18,8 +18,11 @@ from core.inference.mcp_client import ( probe_timeout, stdio_mcp_enabled, ) +from core.inference.mcp_config_import import parse_mcp_config from models.mcp_servers import ( McpServerCreate, + McpServerImportRequest, + McpServerImportResult, McpServerProbeResult, McpServerResponse, McpServerTestRequest, @@ -240,6 +243,46 @@ async def refresh_mcp_server_tools( return McpServerProbeResult(ok = True, tool_count = len(tools)) +@router.post("/import", response_model = McpServerImportResult) +async def import_mcp_servers( + payload: McpServerImportRequest, current_subject: str = Depends(get_current_subject) +): + """Bulk-register servers from a standard mcpServers JSON config (issue + #5936). Each entry rides the existing create path: _validate_url applies + the same stdio gate (a stdio entry becomes a per-entry error when stdio is + off; http still imports), and entries whose url already exists are skipped + so re-importing the same file is idempotent. One bad entry never 400s the + whole batch -- failures are reported per entry.""" + entries, errors = parse_mcp_config(payload.config) + created: list[McpServerResponse] = [] + skipped: list[str] = [] + seen_urls = {row["url"] for row in mcp_servers_db.list_servers()} + + for entry in entries: + try: + url = _validate_url(entry.url) + except HTTPException as exc: + errors.append(f"{entry.display_name}: {exc.detail}") + continue + if url in seen_urls: + skipped.append(entry.display_name) + continue + headers = _normalize_headers(entry.headers) + server_id = uuid.uuid4().hex[:16] + mcp_servers_db.create_server( + id = server_id, + display_name = entry.display_name, + url = url, + headers_json = json.dumps(headers) if headers else None, + is_enabled = entry.is_enabled, + use_oauth = entry.use_oauth and not is_stdio(url), + ) + seen_urls.add(url) + created.append(_row_to_response(mcp_servers_db.get_server(server_id))) + + return McpServerImportResult(created = created, skipped = skipped, errors = errors) + + @router.post("/test", response_model = McpServerProbeResult) async def test_mcp_server( payload: McpServerTestRequest, current_subject: str = Depends(get_current_subject) diff --git a/studio/backend/tests/test_mcp_config_import.py b/studio/backend/tests/test_mcp_config_import.py new file mode 100644 index 0000000000..4082733490 --- /dev/null +++ b/studio/backend/tests/test_mcp_config_import.py @@ -0,0 +1,341 @@ +"""Tests for MCP config-file import (issue #5936). + +Covers the round-trip-safe command join/split inverse (join_stdio_command ↔ +parse_stdio_command, on both posix and win32 using the issue's Windows +fixtures), the pure config parser (parse_mcp_config), and the POST /import +route (stdio gate on/off, url dedup, one bad entry not sinking the batch). + +Run from studio/backend: python -m pytest tests/test_mcp_config_import.py -q +""" + +import sys + +import pytest + +from core.inference import mcp_client +from core.inference.mcp_config_import import parse_mcp_config +from storage import mcp_servers_db + + +def _reset_db(tmp_path, monkeypatch): + monkeypatch.setenv("UNSLOTH_STUDIO_HOME", str(tmp_path)) + monkeypatch.setattr(mcp_servers_db, "_schema_ready", False) + + +def _enable(monkeypatch): + monkeypatch.setenv("UNSLOTH_STUDIO_ALLOW_STDIO_MCP", "1") + + +def _disable(monkeypatch): + monkeypatch.delenv("UNSLOTH_STUDIO_ALLOW_STDIO_MCP", raising = False) + + +# ── 1. join_stdio_command ↔ parse_stdio_command round-trip ────────── + + +@pytest.mark.parametrize( + "parts", + [ + ["npx", "-y", "@modelcontextprotocol/server-filesystem", "/tmp"], + ["python", "-m", "mod", "--name", "a b"], + ["uvx", "some-server", "--flag"], + ["/usr/local/bin/my-server"], + ["mcp-server-sqlite"], + ], +) +def test_join_parse_roundtrip_posix(monkeypatch, parts): + monkeypatch.setattr(sys, "platform", "linux") + joined = mcp_client.join_stdio_command(parts) + assert mcp_client.parse_stdio_command(joined) == parts + + +@pytest.mark.parametrize( + "parts", + [ + # Issue #5936's literal Windows examples: absolute .exe with a path, and + # backslash drive/dir args must survive the join→split round-trip intact. + [ + "C:\\Users\\user\\Documents\\Office-Word-MCP-Server\\.venv\\Scripts\\python.exe", + "C:\\Users\\user\\Documents\\Office-Word-MCP-Server\\word_mcp_server.py", + ], + [ + "node", + "C:\\Users\\user\\Documents\\DesktopCommanderMCP\\dist\\index.js", + "--no-onboarding", + ], + [ + "node", + "C:\\Users\\user\\AppData\\Roaming\\npm\\node_modules\\@modelcontextprotocol\\server-filesystem\\dist\\index.js", + "D:\\", + "O:\\", + ], + # A command path with spaces is the case that actually needs quoting. + ["C:\\Program Files\\node\\node.exe", "server.js"], + ["C:\\Program Files\\Foo\\", "server.js"], + ["C:\\Program Files\\Foo\\", '{"foo":"bar"}'], + ["'C:\\Program Files\\node\\node.exe'", "server.js"], + ["node", "O'Reilly"], + ["node", "C:\\Users\\O'Reilly\\server.js"], + ["node", "'draft'"], + ["node", "'open", "close'"], + ["node", ""], + ], +) +def test_join_parse_roundtrip_win32(monkeypatch, parts): + monkeypatch.setattr(sys, "platform", "win32") + joined = mcp_client.join_stdio_command(parts) + assert mcp_client.parse_stdio_command(joined) == parts + + +def test_parse_rejects_manual_single_quoted_windows_executable(monkeypatch): + monkeypatch.setattr(sys, "platform", "win32") + command = "'C:\\Program Files\\node\\node.exe' server.js" + with pytest.raises(ValueError): + mcp_client.parse_stdio_command(command) + + +def test_parse_windows_apostrophes_as_literals(monkeypatch): + monkeypatch.setattr(sys, "platform", "win32") + assert mcp_client.parse_stdio_command("node O'Reilly") == ["node", "O'Reilly"] + assert mcp_client.parse_stdio_command("node C:\\Users\\O'Reilly\\server.js") == [ + "node", + "C:\\Users\\O'Reilly\\server.js", + ] + assert mcp_client.parse_stdio_command("node 'draft'") == ["node", "'draft'"] + assert mcp_client.parse_stdio_command("node 'open close'") == ["node", "'open", "close'"] + + +def test_parse_rejects_unterminated_windows_double_quote(monkeypatch): + monkeypatch.setattr(sys, "platform", "win32") + with pytest.raises(ValueError): + mcp_client.parse_stdio_command('node "C:\\path with spaces') + + +# ── 2. parse_mcp_config ───────────────────────────────────────────── + + +def test_parse_stdio_entry(): + cfg = { + "mcpServers": { + "fs": { + "command": "npx", + "args": ["-y", "server", "/tmp"], + "env": {"K": "v"}, + } + } + } + entries, errors = parse_mcp_config(cfg) + assert errors == [] + assert len(entries) == 1 + entry = entries[0] + assert entry.display_name == "fs" + assert entry.is_stdio is True + assert entry.headers == {"K": "v"} + assert mcp_client.parse_stdio_command(entry.url) == ["npx", "-y", "server", "/tmp"] + + +def test_parse_remote_entry(): + cfg = { + "mcpServers": { + "remote": { + "url": "https://example.com/mcp", + "headers": {"Authorization": "Bearer x"}, + } + } + } + entries, errors = parse_mcp_config(cfg) + assert errors == [] + assert entries[0].url == "https://example.com/mcp" + assert entries[0].is_stdio is False + assert entries[0].headers == {"Authorization": "Bearer x"} + + +def test_parse_preserves_disabled_and_oauth(): + cfg = { + "servers": { + "remote": { + "type": "http", + "url": "https://example.com/mcp", + "oauth": {"clientId": "client"}, + "disabled": True, + } + } + } + entries, errors = parse_mcp_config(cfg) + assert errors == [] + assert entries[0].is_enabled is False + assert entries[0].use_oauth is True + + +def test_parse_accepts_cline_streamable_http_alias(): + cfg = { + "mcpServers": { + "remote": { + "type": "streamableHttp", + "url": "https://example.com/mcp", + } + } + } + entries, errors = parse_mcp_config(cfg) + assert errors == [] + assert entries[0].url == "https://example.com/mcp" + assert entries[0].is_stdio is False + + +@pytest.mark.parametrize( + "server", + [ + {"command": "node", "args": ["server.js"], "cwd": "/tmp/server"}, + {"command": "node", "args": ["server.js"], "envFile": ".env"}, + {"command": "node", "args": ["server.js"], "env": {"API_KEY": "${input:api-key}"}}, + {"command": "node", "args": ["${workspaceFolder}/server.js"]}, + {"command": "node", "args": ["server.js"], "env": {"HTTP_PROXY": None}}, + {"command": "node", "args": ["server.js"], "sandboxEnabled": True}, + {"url": "https://example.com/mcp", "headers": {"Authorization": "Bearer ${input:token}"}}, + {"url": "https://example.com/mcp", "headers": {"Authorization": None}}, + {"type": "http", "url": "https://example.com/sse"}, + {"type": "http", "url": "https://example.com/sse "}, + {"type": "streamableHttp", "url": "https://example.com/sse"}, + {"url": "https://example.com/mcp", "timeout": 120}, + {"url": "https://example.com/mcp", "timeoutMs": 120000}, + {"url": "https://example.com/mcp", "timeoutSeconds": 120}, + {"type": "sse", "url": "https://example.com/custom"}, + ], +) +def test_parse_rejects_unrepresentable_imports(server): + entries, errors = parse_mcp_config({"servers": {"bad": server}}) + assert entries == [] + assert len(errors) == 1 + + +def test_servers_alias_key(): + # VS Code uses "servers" instead of "mcpServers". + cfg = {"servers": {"fs": {"command": "node", "args": ["x.js"]}}} + entries, errors = parse_mcp_config(cfg) + assert errors == [] + assert len(entries) == 1 + + +def test_env_and_args_values_coerced_to_str(): + cfg = {"mcpServers": {"fs": {"command": "node", "args": [8080], "env": {"PORT": 8080}}}} + entries, errors = parse_mcp_config(cfg) + assert errors == [] + assert entries[0].headers == {"PORT": "8080"} + assert mcp_client.parse_stdio_command(entries[0].url) == ["node", "8080"] + + +def test_args_optional(): + cfg = {"mcpServers": {"sqlite": {"command": "mcp-server-sqlite"}}} + entries, errors = parse_mcp_config(cfg) + assert errors == [] + assert entries[0].url == "mcp-server-sqlite" + assert entries[0].headers is None + + +def test_bad_entry_does_not_sink_batch(): + cfg = { + "mcpServers": { + "good": {"command": "node", "args": ["x.js"]}, + "both": {"command": "node", "url": "https://x/mcp"}, + "neither": {"name": "oops"}, + "bad_args": {"command": "node", "args": "x.js"}, + "bad_env": {"command": "node", "env": ["NOT", "A", "DICT"]}, + } + } + entries, errors = parse_mcp_config(cfg) + assert {e.display_name for e in entries} == {"good"} + assert len(errors) == 4 + + +def test_not_a_dict(): + entries, errors = parse_mcp_config([]) + assert entries == [] + assert len(errors) == 1 + + +def test_missing_servers_key(): + entries, errors = parse_mcp_config({"foo": {}}) + assert entries == [] + assert len(errors) == 1 + + +def test_servers_alias_error_names_actual_key(): + entries, errors = parse_mcp_config({"servers": []}) + assert entries == [] + assert errors == ["'servers' must be an object mapping name -> server."] + + +# ── 3. POST /import route ─────────────────────────────────────────── + + +def test_import_route_creates_and_dedups(tmp_path, monkeypatch): + import asyncio + + from models.mcp_servers import McpServerImportRequest + import routes.mcp_servers as routes_mcp + + _reset_db(tmp_path, monkeypatch) + _enable(monkeypatch) + cfg = { + "mcpServers": { + "fs": { + "command": "npx", + "args": ["-y", "server", "/tmp"], + "env": {"API_KEY": "sk"}, + }, + "remote": {"url": "https://example.com/mcp"}, + "oauth": { + "type": "http", + "url": "https://auth.example.com/mcp", + "oauth": {"clientId": "client"}, + }, + "disabled": { + "url": "https://disabled.example.com/mcp", + "disabled": True, + }, + } + } + res = asyncio.run( + routes_mcp.import_mcp_servers(McpServerImportRequest(config = cfg), current_subject = "u") + ) + assert res.errors == [] + assert res.skipped == [] + assert {c.display_name for c in res.created} == {"fs", "remote", "oauth", "disabled"} + fs = next(c for c in res.created if c.display_name == "fs") + assert fs.headers == {"API_KEY": "sk"} + assert fs.use_oauth is False + assert fs.is_enabled is True + oauth = next(c for c in res.created if c.display_name == "oauth") + assert oauth.use_oauth is True + disabled = next(c for c in res.created if c.display_name == "disabled") + assert disabled.is_enabled is False + + # Re-importing the same config skips both by url. + res2 = asyncio.run( + routes_mcp.import_mcp_servers(McpServerImportRequest(config = cfg), current_subject = "u") + ) + assert res2.created == [] + assert set(res2.skipped) == {"fs", "remote", "oauth", "disabled"} + + +def test_import_route_gates_stdio_when_disabled(tmp_path, monkeypatch): + import asyncio + + from models.mcp_servers import McpServerImportRequest + import routes.mcp_servers as routes_mcp + + _reset_db(tmp_path, monkeypatch) + _disable(monkeypatch) + cfg = { + "mcpServers": { + "fs": {"command": "npx", "args": ["server"]}, + "remote": {"url": "https://example.com/mcp"}, + } + } + res = asyncio.run( + routes_mcp.import_mcp_servers(McpServerImportRequest(config = cfg), current_subject = "u") + ) + # Remote still imports; the stdio entry is rejected per-entry (gate off). + assert {c.display_name for c in res.created} == {"remote"} + assert any("fs" in err for err in res.errors) + assert len(mcp_servers_db.list_servers()) == 1 diff --git a/studio/frontend/src/features/chat/api/mcp-servers-api.ts b/studio/frontend/src/features/chat/api/mcp-servers-api.ts index be88d12664..0dc0719da4 100644 --- a/studio/frontend/src/features/chat/api/mcp-servers-api.ts +++ b/studio/frontend/src/features/chat/api/mcp-servers-api.ts @@ -21,6 +21,12 @@ export interface McpServerProbeResult { error: string | null; } +export interface McpServerImportResult { + created: McpServerConfig[]; + skipped: string[]; + errors: string[]; +} + function parseErrorText(status: number, body: unknown): string { if (body && typeof body === "object") { const { detail, message } = body as { detail?: unknown; message?: unknown }; @@ -114,3 +120,12 @@ export function testMcpServer(payload: { }, }); } + +// Bulk-import servers from a standard mcpServers JSON config (Claude Desktop, +// Cursor, Cline, VS Code). The backend skips duplicates and reports per-entry +// errors instead of failing the whole batch. +export function importMcpServers( + config: unknown, +): Promise { + return mcpRequest("/import", { method: "POST", body: { config } }); +} diff --git a/studio/frontend/src/features/chat/chat-mcp-servers-dialog.tsx b/studio/frontend/src/features/chat/chat-mcp-servers-dialog.tsx index f78784b436..17d6d1e9ec 100644 --- a/studio/frontend/src/features/chat/chat-mcp-servers-dialog.tsx +++ b/studio/frontend/src/features/chat/chat-mcp-servers-dialog.tsx @@ -1,11 +1,11 @@ // SPDX-License-Identifier: AGPL-3.0-only // Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0 -import { useCallback, useEffect, useState } from "react"; +import { type ChangeEvent, useCallback, useEffect, useRef, useState } from "react"; import { toast } from "sonner"; import { Delete02Icon, Edit03Icon, PlusSignIcon } from "@hugeicons/core-free-icons"; import { HugeiconsIcon } from "@hugeicons/react"; -import { RefreshCwIcon } from "lucide-react"; +import { RefreshCwIcon, UploadIcon } from "lucide-react"; import { Button } from "@/components/ui/button"; import { @@ -23,6 +23,7 @@ import { type McpServerConfig, createMcpServer, deleteMcpServer, + importMcpServers, listMcpServers, refreshMcpServerTools, testMcpServer, @@ -194,7 +195,9 @@ export function ChatMcpServersDialog({ const [form, setForm] = useState(EMPTY_FORM); const [saving, setSaving] = useState(false); const [testing, setTesting] = useState(false); + const [importing, setImporting] = useState(false); const [refreshingId, setRefreshingId] = useState(null); + const fileInputRef = useRef(null); const refresh = useCallback(async () => { setLoading(true); @@ -315,6 +318,49 @@ export function ChatMcpServersDialog({ } } + async function onImportFile(e: ChangeEvent) { + const file = e.target.files?.[0]; + e.target.value = ""; // let the user re-pick the same file later + if (!file) return; + let config: unknown; + try { + config = JSON.parse(await file.text()); + } catch { + toast.error("Invalid JSON file"); + return; + } + setImporting(true); + try { + const result = await importMcpServers(config); + const parts = [`${result.created.length} added`]; + if (result.skipped.length) parts.push(`${result.skipped.length} skipped`); + if (result.errors.length) { + parts.push( + `${result.errors.length} error${result.errors.length === 1 ? "" : "s"}`, + ); + } + const summary = parts.join(", "); + if (result.errors.length) { + toast.warning(summary, { + description: ( +
+ {result.errors.slice(0, 5).join("\n")} +
+ ), + }); + } else { + toast.success(summary); + } + await refresh(); + } catch (err) { + toast.error("Import failed", { + description: err instanceof Error ? err.message : String(err), + }); + } finally { + setImporting(false); + } + } + async function removeServer(server: McpServerConfig) { const ok = window.confirm(`Delete MCP server "${server.display_name}"?`); if (!ok) return; @@ -385,9 +431,35 @@ export function ChatMcpServersDialog({ Register remote (HTTP) or local (stdio command) MCP servers. + {showForm ? (
+ {view.kind === "create" && ( +
+ + Import servers from a config file. + + +
+ )}
) : (
-
+
+