diff --git a/.github/scripts/agent-guides-drive.sh b/.github/scripts/agent-guides-drive.sh index 2457f08407..b63ac94b93 100755 --- a/.github/scripts/agent-guides-drive.sh +++ b/.github/scripts/agent-guides-drive.sh @@ -36,6 +36,23 @@ AGENT="${2:?usage: agent-guides-drive.sh }" # Determinism (seed/temp) is applied at the server level by # serve-unsloth-run.sh --extra; agents inherit it through the API. TIMEOUT="${AGENT_INVOKE_TIMEOUT:-180}" +# opencode is the slow outlier. Unlike the print-mode agents (claude -p, codex +# exec) it runs a full turn AND a separate small_model call to name the session, +# so one connection reply takes ~8 min on a CPU-served 4B -- right at the shared +# 600s cap, so the cell flaked when a run drifted past a ~480s success. Give it +# headroom (still well under the 40-min job budget); the fast agents keep the +# tight cap that still catches a real headless-TTY hang. +case "$AGENT" in + opencode) + # Double it, but only for a bare-integer seconds value. A GNU timeout(1) + # duration suffix (s/m/h/d, including floats like 0.5s) is left unchanged so + # the arithmetic never sees a non-number; timeout(1) parses it directly. + case "$TIMEOUT" in + *[!0-9]*) ;; + *) TIMEOUT=$(( TIMEOUT * 2 )) ;; + esac + ;; +esac # Claude refuses --dangerously-skip-permissions outside a sandbox; the CI runner # IS the sandbox, so declare it (mirrors unslothai/scripts launcher.sh). Harmless diff --git a/.github/scripts/run-studio-permission-browser.sh b/.github/scripts/run-studio-permission-browser.sh new file mode 100755 index 0000000000..2007789035 --- /dev/null +++ b/.github/scripts/run-studio-permission-browser.sh @@ -0,0 +1,69 @@ +#!/usr/bin/env bash +# SPDX-License-Identifier: AGPL-3.0-only +# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. + +set -euo pipefail + +port="${1:?usage: $0 PORT BROWSER [CHANNEL]}" +browser="${2:?usage: $0 PORT BROWSER [CHANNEL]}" +channel="${3:-}" +slug="$browser${channel:+-$channel}" +artifact_dir="logs/playwright-permissions-$slug" +server_log="logs/studio-permissions-$slug.log" +studio_home="${UNSLOTH_STUDIO_HOME:-$HOME/.unsloth/studio}" +set -- +if [ -n "${STUDIO_PERMISSION_FRONTEND:-}" ]; then + set -- -f "$STUDIO_PERMISSION_FRONTEND" +fi + +mkdir -p "$artifact_dir" +unsloth studio reset-password +UNSLOTH_API_ONLY=1 unsloth studio -H 127.0.0.1 -p "$port" "$@" \ + >"$server_log" 2>&1 & +studio_pid=$! + +cleanup() { + kill "$studio_pid" 2>/dev/null || true + wait "$studio_pid" 2>/dev/null || true +} +trap cleanup EXIT + +healthy=0 +for _ in $(seq 1 180); do + if curl -fs "http://127.0.0.1:$port/api/health" >/dev/null; then + healthy=1 + break + fi + if ! kill -0 "$studio_pid" 2>/dev/null; then + tail -100 "$server_log" || true + exit 1 + fi + sleep 1 +done +if [ "$healthy" -ne 1 ]; then + tail -100 "$server_log" || true + exit 1 +fi + +old_password=$(cat "$studio_home/auth/.bootstrap_password") +new_password="CIPerm-$(python -c 'import secrets; print(secrets.token_urlsafe(16))')" +if [ "${GITHUB_ACTIONS:-}" = "true" ]; then + echo "::add-mask::$old_password" + echo "::add-mask::$new_password" +fi + +export BASE_URL="http://127.0.0.1:$port" +export STUDIO_OLD_PW="$old_password" +export STUDIO_NEW_PW="$new_password" +export STUDIO_UI_STRICT=1 +export STUDIO_UI_PERMISSION_ONLY=1 +export STUDIO_UI_WALL_TIMEOUT_S=240 +export STUDIO_PLAYWRIGHT_BROWSER="$browser" +export PW_ART_DIR="$artifact_dir" +if [ -n "$channel" ]; then + export STUDIO_PLAYWRIGHT_CHANNEL="$channel" +else + unset STUDIO_PLAYWRIGHT_CHANNEL || true +fi + +python tests/studio/playwright_chat_ui.py diff --git a/.github/workflows/studio-mac-ui-smoke.yml b/.github/workflows/studio-mac-ui-smoke.yml index 378e8ee5a6..7375e9bcbf 100644 --- a/.github/workflows/studio-mac-ui-smoke.yml +++ b/.github/workflows/studio-mac-ui-smoke.yml @@ -19,6 +19,7 @@ on: - 'install.sh' - 'pyproject.toml' - 'tests/studio/**' + - '.github/scripts/run-studio-permission-browser.sh' - '.github/workflows/studio-mac-ui-smoke.yml' push: branches: [main, pip] @@ -96,7 +97,7 @@ jobs: - name: Assert llama.cpp loads on this macOS run: bash .github/scripts/assert-llama-loads.sh - - name: Install Playwright + Chromium + - name: Install Playwright browsers # No --with-deps on Mac: that flag installs Linux apt packages. # GitHub-hosted macos-14 ships the system frameworks Chromium # needs already. @@ -112,7 +113,7 @@ jobs: # in-script retry recover from any residual flakes. run: | pip install 'playwright>=1.55,<1.58' - python -m playwright install chromium + python -m playwright install chromium webkit - name: Patch Playwright pipeTransport.js to tolerate malformed JSON # In Playwright 1.55-1.58, pipeTransport.js does @@ -244,6 +245,10 @@ jobs: kill "${STUDIO_PID}" 2>/dev/null || true sleep 2 + - name: Cross-browser permission controls + run: | + bash .github/scripts/run-studio-permission-browser.sh 18895 webkit + - name: Reset auth + boot Unsloth for extra UI tests (port 18897) run: | unsloth studio reset-password @@ -343,5 +348,7 @@ jobs: logs/studio_extra.log logs/install.log logs/playwright + logs/playwright-permissions-* logs/playwright_extra + logs/studio-permissions-*.log retention-days: 7 diff --git a/.github/workflows/studio-ui-smoke.yml b/.github/workflows/studio-ui-smoke.yml index b6d6d7d6e2..0ad55ebd6d 100644 --- a/.github/workflows/studio-ui-smoke.yml +++ b/.github/workflows/studio-ui-smoke.yml @@ -27,6 +27,7 @@ on: # The Playwright test files themselves -- a PR that ONLY edits # the test must still trigger UI CI. - 'tests/studio/**' + - '.github/scripts/run-studio-permission-browser.sh' - '.github/workflows/studio-ui-smoke.yml' push: branches: [main, pip] @@ -107,13 +108,10 @@ jobs: set -o pipefail bash install.sh --local --no-torch 2>&1 | tee logs/install.log - - name: Install Playwright + Chromium + - name: Install Playwright browsers run: | pip install 'playwright>=1.45' - # --with-deps installs the OS-level runtime libs Chromium - # needs (libnss3, libxkbcommon, etc.). About 30 s on a - # warm runner. - python -m playwright install --with-deps chromium + python -m playwright install --with-deps chromium firefox webkit - name: Reset auth + boot Unsloth run: | @@ -182,6 +180,12 @@ jobs: kill "${STUDIO_PID}" 2>/dev/null || true sleep 2 + - name: Cross-browser permission controls + run: | + bash .github/scripts/run-studio-permission-browser.sh 18893 firefox + bash .github/scripts/run-studio-permission-browser.sh 18893 webkit + bash .github/scripts/run-studio-permission-browser.sh 18893 chromium chrome + # The chat UI test ends by clicking the Shutdown menuitem, which # leaves the server dead. The extra UI test (Compare / Recipes / # Export / Unsloth / Settings) needs a fresh Unsloth, so we boot a @@ -233,6 +237,54 @@ jobs: kill "${STUDIO_EXTRA_PID}" 2>/dev/null || true sleep 2 + # Model-picker per-model-config regression (PR #7207 re-land of #6647). + # Fourth Unsloth on its own port; loads the tiny GGUF and drives the + # picker's run-settings surface: Context Length persists across a reload, + # Reset clears the stored override (never pins it), and the infra models + # (RAG embedder + llama.cpp probe) stay hidden from the picker. + - name: Reset auth + boot Unsloth for model-config tests (port 18898) + run: | + unsloth studio reset-password + mkdir -p logs + UNSLOTH_API_ONLY=1 unsloth studio -H 127.0.0.1 -p 18898 \ + > logs/studio_modelcfg.log 2>&1 & + echo "STUDIO_MODELCFG_PID=$!" >> "$GITHUB_ENV" + + - name: Wait for /api/health on 18898 + run: | + for i in $(seq 1 180); do + if curl -fs "http://127.0.0.1:18898/api/health" > /tmp/health4.json; then + jq -e '.status == "healthy"' /tmp/health4.json && break + fi + sleep 1 + done + jq -e '.status == "healthy"' /tmp/health4.json + + - name: Pass bootstrap pw for model-config test + run: | + NEW="CIModelCfg-$(python -c 'import secrets; print(secrets.token_urlsafe(16))')" + echo "::add-mask::$NEW" + echo "STUDIO_MODELCFG_NEW_PW=$NEW" >> "$GITHUB_ENV" + + - name: Drive model-picker per-model-config with Playwright + env: + BASE_URL: http://127.0.0.1:18898 + STUDIO_NEW_PW: ${{ env.STUDIO_MODELCFG_NEW_PW }} + PW_ART_DIR: logs/playwright_modelcfg + STUDIO_UI_STRICT: '1' + GGUF_REPO: ${{ env.GGUF_REPO }} + GGUF_VARIANT: ${{ env.GGUF_VARIANT }} + STUDIO_MODEL_HINT: gemma-3-270m + run: | + mkdir -p logs/playwright_modelcfg + python tests/studio/playwright_model_config.py + + - name: Stop fourth Unsloth + if: always() + run: | + kill "${STUDIO_MODELCFG_PID}" 2>/dev/null || true + sleep 2 + # IME + multilingual paste regression (issue #5318 / PR #5327). # Third Unsloth on its own port so a hang here cannot poison the # earlier UI tests. No GGUF -- the bug surface is the composer. @@ -293,10 +345,14 @@ jobs: path: | logs/studio.log logs/studio_extra.log + logs/studio_modelcfg.log logs/studio_ime.log logs/install.log logs/server-logs/ logs/playwright + logs/playwright-permissions-* logs/playwright_extra + logs/playwright_modelcfg logs/playwright_ime + logs/studio-permissions-*.log retention-days: 7 diff --git a/.github/workflows/studio-windows-ui-smoke.yml b/.github/workflows/studio-windows-ui-smoke.yml index 12d7475b53..f401f7be44 100644 --- a/.github/workflows/studio-windows-ui-smoke.yml +++ b/.github/workflows/studio-windows-ui-smoke.yml @@ -19,6 +19,7 @@ on: - 'install.ps1' - 'pyproject.toml' - 'tests/studio/**' + - '.github/scripts/run-studio-permission-browser.sh' - '.github/workflows/studio-windows-ui-smoke.yml' push: branches: [main, pip] @@ -345,6 +346,10 @@ jobs: kill "${STUDIO_PID}" 2>/dev/null || true sleep 2 + - name: Edge permission controls + run: | + bash .github/scripts/run-studio-permission-browser.sh 18895 chromium msedge + - name: Reset auth + boot Unsloth for extra UI tests (port 18897) run: | unsloth studio reset-password @@ -402,5 +407,7 @@ jobs: logs/studio_extra.log logs/install.log logs/playwright + logs/playwright-permissions-* logs/playwright_extra + logs/studio-permissions-*.log retention-days: 7 diff --git a/README.md b/README.md index 085c7718e5..6aa8f4f4c3 100644 --- a/README.md +++ b/README.md @@ -11,6 +11,7 @@ Unsloth Studio lets you run and train models locally.

Features • + News • Quickstart • Notebooks • Documentation @@ -47,15 +48,44 @@ Unsloth Studio (Beta) lets you run and train text, [audio](https://unsloth.ai/do * [Auto set inference settings](https://unsloth.ai/docs/new/studio/chat#auto-parameter-tuning) and customize chat templates. * We work directly with teams behind [gpt-oss](https://docs.unsloth.ai/new/gpt-oss-how-to-run-and-fine-tune#unsloth-fixes-for-gpt-oss), [Qwen3](https://www.reddit.com/r/LocalLLaMA/comments/1kaodxu/qwen3_unsloth_dynamic_ggufs_128k_context_bug_fixes/), [Llama 4](https://github.com/ggml-org/llama.cpp/pull/12889), [Mistral](https://huggingface.co/mistralai/Mistral-Medium-3.5-128B/discussions/18), [Gemma 1-3](https://news.ycombinator.com/item?id=39671146), and [Phi-4](https://unsloth.ai/blog/phi4), where we’ve fixed bugs that improve model accuracy. * Chat with images, audio, PDFs, code, DOCX and more. [Connect API providers](https://unsloth.ai/docs/integrations/connections) (OpenAI, Anthropic) or servers (vLLM, Ollama). +* [**Compare any two models**](https://unsloth.ai/docs/new/studio/chat#model-arena) side by side with the same prompt. +* **OpenAI/Anthropic-compatible APIs**: Serve local models through `/v1/chat/completions`, `/v1/responses` and `/v1/messages`. +* **Connect local models to agents**: Use `unsloth start` with Claude Code, Codex, Hermes and more. +* **Web/PDF search** can read PDF papers, manuals and other PDF results. +* **GGUF hardware controls**: Choose GPUs/layers, offload MoE experts, use multi-GPU or Tensor Parallelism. +* The opt-in **MCP control endpoint** lets AI clients manage models, training, recipes and exports. ### Training -* Train and RL **500+ models** up to **2x faster** with up to **70% less VRAM**, with no accuracy loss. -* Custom Triton and mathematical **kernels**. See some collabs we did with [PyTorch](https://unsloth.ai/docs/get-started/reinforcement-learning-rl-guide/fp8-reinforcement-learning) and [Hugging Face](https://unsloth.ai/docs/new/faster-moe). +* Train and RL **500+ models** up to **2x faster** with **70% less VRAM**; MoE up to **12x faster**. +* Train and run RL on [AMD GPUs](https://unsloth.ai/docs/basics/amd) across Windows, WSL and Linux. * **Data Recipes**: [Auto-create datasets](https://unsloth.ai/docs/new/studio/data-recipe) from **PDF, CSV, DOCX** etc. Edit data in a visual-node workflow. -* **[Reinforcement Learning](https://unsloth.ai/docs/get-started/reinforcement-learning-rl-guide)** (RL): The most efficient [RL](https://unsloth.ai/docs/get-started/reinforcement-learning-rl-guide) library, using **80% less VRAM** for GRPO, [FP8](https://unsloth.ai/docs/get-started/reinforcement-learning-rl-guide/fp8-reinforcement-learning) etc. -* Supports full fine-tuning, RL, pretraining, 4-bit, 16-bit and, FP8 training. +* **[Reinforcement Learning](https://unsloth.ai/docs/get-started/reinforcement-learning-rl-guide)** uses **80% less VRAM** for GRPO, FP8 and vision RL, with 7x longer contexts. +* [**Long-context training**](https://unsloth.ai/docs/new/3x-faster-training-packing): **3x faster**, 30% less VRAM and 500K+ context. +* Supports LoRA/QLoRA, full fine-tuning, RL, pretraining, 4-bit, 16-bit and FP8. +* Custom Triton and mathematical **kernels** built with PyTorch and Hugging Face. * **Observability**: Monitor training live, track loss and GPU usage and customize graphs. * [Multi-GPU](https://unsloth.ai/docs/basics/multi-gpu-training-with-unsloth) training is supported, with major improvements coming soon. +## 🚀 Unsloth Start + +[Unsloth Start](https://unsloth.ai/docs/integrations/unsloth-start) connects [Claude Code](https://unsloth.ai/docs/basics/claude-code), [Codex](https://unsloth.ai/docs/basics/codex) and other agents to local models with one command. + +Start Unsloth, load a model, open your project folder, then run: + +```bash +unsloth start claude +``` + +Replace `claude` with any supported agent: + +| Agent | Command | +| --- | --- | +| Claude Code | `unsloth start claude` | +| OpenAI Codex | `unsloth start codex` | +| Hermes Agent | `unsloth start hermes` | +| OpenClaw | `unsloth start openclaw` | +| OpenCode | `unsloth start opencode` | +| Pi Coding Agent | `unsloth start pi` | + ## 📥 Install Unsloth can be used in two ways: through **[Unsloth Studio](https://unsloth.ai/docs/new/studio/)**, the web UI, or through **Unsloth Core**, the code-based version. Each has different requirements. @@ -65,7 +95,8 @@ Unsloth Studio (Beta) works on **Windows, Linux, WSL** and **macOS**. * **CPU:** Supported for Chat and Data Recipes currently * **NVIDIA:** Training works on RTX 30/40/50, Blackwell, DGX Spark, Station and more * **macOS:** Training, MLX and GGUF inference are ALL supported. -* **AMD:** Chat + Data works. Train with [Unsloth Core](#unsloth-core-code-based). Unsloth Studio support is out soon. +* **AMD:** Training, RL, chat and deployment work on Windows, WSL and Linux. [Read the AMD guide](https://unsloth.ai/docs/basics/amd). +* **Vulkan:** GGUF inference is supported on [compatible GPUs, including Intel GPUs](https://github.com/unslothai/unsloth/pull/5819). * **Multi-GPU:** Available now, with a major upgrade on the way #### macOS, Linux, WSL: @@ -122,7 +153,7 @@ You can use the same Docker image as Unsloth Studio. #### AMD, Intel: For RTX 50x, B200, 6000 GPUs: `uv pip install unsloth --torch-backend=auto`. Read our guides for: [Blackwell](https://unsloth.ai/docs/blog/fine-tuning-llms-with-blackwell-rtx-50-series-and-unsloth) and [DGX Spark](https://unsloth.ai/docs/blog/fine-tuning-llms-with-nvidia-dgx-spark-and-unsloth).
-To install Unsloth on **AMD** and **Intel** GPUs, follow our [AMD Guide](https://unsloth.ai/docs/get-started/install/amd) and [Intel Guide](https://unsloth.ai/docs/get-started/install/intel). +To install Unsloth on **AMD** and **Intel** GPUs, follow our [AMD Guide](https://unsloth.ai/docs/basics/amd) and [Intel Guide](https://unsloth.ai/docs/get-started/install/intel). ## 📒 Free Notebooks @@ -148,13 +179,20 @@ Read our [guide](https://unsloth.ai/docs/get-started/fine-tuning-llms-guide). Ad - See detailed documentation for Unsloth [here](https://unsloth.ai/docs) ## 🦥 Unsloth News -- **Connections**: Connect any API provider (OpenAI, Anthropic) or server (vLLM, Ollama). [Guide](https://unsloth.ai/docs/integrations/connections) -- **MTP**: Run Qwen3.6 MTP in Unsloth. MTP settings are autoset specific to your hardware. [Guide](https://unsloth.ai/docs/models/qwen3.6#mtp-guide) -- **API inference endpoint**: Deploy and run local LLMs in Claude Code, Codex tools. [Guide](https://unsloth.ai/docs/basics/api) -- **Qwen3.6**: Qwen3.6-35B-A3B can now be trained and run in Unsloth Studio. [Blog](https://unsloth.ai/docs/models/qwen3.6) -- **Gemma 4**: Run and train Google’s new models directly in Unsloth. [Blog](https://unsloth.ai/docs/models/gemma-4) +- **AMD training**: Train, run RL, chat and deploy on AMD GPUs across Windows, WSL and Linux. [Guide](https://unsloth.ai/docs/basics/amd) +- **GGUF hardware controls**: Choose GPU/layer placement, offload MoE experts and use multi-GPU or Tensor Parallelism. [#6414](https://github.com/unslothai/unsloth/pull/6414) +- **Local models for any agent**: Use `unsloth start` with Claude Code, Codex, Hermes, OpenCode, OpenClaw, Pi and more through Unsloth's OpenAI- and Anthropic-compatible APIs. [Guide](https://unsloth.ai/docs/basics/api) +- **MCP control endpoint**: Let compatible clients manage models, training, recipes, checkpoints and exports. [#7191](https://github.com/unslothai/unsloth/pull/7191) +- **Local inference reliability**: Resume long chats faster, recover stalled downloads and reuse existing GGUF files. [#7204](https://github.com/unslothai/unsloth/pull/7204) • [#6858](https://github.com/unslothai/unsloth/pull/6858) • [#7209](https://github.com/unslothai/unsloth/pull/7209) +- **New models**: [Qwen-AgentWorld](https://huggingface.co/unsloth/Qwen-AgentWorld-35B-A3B-GGUF), [Ornith](https://huggingface.co/unsloth/models?search=ornith), [Kimi K2.7 Code](https://unsloth.ai/docs/models/kimi-k2.7-code) and [MiniMax M3](https://unsloth.ai/docs/models/minimax-m3) +- **GLM-5.2**: Run Z.ai's 744B-parameter, 1M-context open model locally with Unsloth Dynamic GGUFs. [Guide](https://unsloth.ai/docs/models/glm-5.2) +- **DeepSeek-V4**: Run DeepSeek-V4-Flash locally with corrected multi-turn and tool-calling behavior. [Guide](https://unsloth.ai/docs/models/deepseek-v4) +- **DiffusionGemma**: Run and fine-tune Google's diffusion language model with 1.8x faster inference in Unsloth Studio. [Guide](https://unsloth.ai/docs/models/diffusiongemma) +- **Qwen3.6**: Run and train Qwen3.6 with MTP for 1.4-2.2x faster inference and NVFP4 quants for supported GPUs. [Guide](https://unsloth.ai/docs/models/qwen3.6) +- **Gemma 4**: Run and train Gemma 4 text, image and audio models with QAT, MTP, GGUF and MLX support. [Guide](https://unsloth.ai/docs/models/gemma-4) +- **MCP servers**: Connect local models to files, apps, databases and external tools through Model Context Protocol. [Guide](https://unsloth.ai/docs/basics/mcp) +- **Connections**: Mix local models with API providers (OpenAI, Anthropic) or servers (vLLM, Ollama) in the same interface. [Guide](https://unsloth.ai/docs/integrations/connections) - **Introducing Unsloth Studio**: our new web UI for running and training LLMs. [Blog](https://unsloth.ai/docs/new/studio) -- **Qwen3.5** - 0.8B, 2B, 4B, 9B, 27B, 35-A3B, 112B-A10B are now supported. [Guide + notebooks](https://unsloth.ai/docs/models/qwen3.5/fine-tune) - Train **MoE LLMs 12x faster** with 35% less VRAM - DeepSeek, GLM, Qwen and gpt-oss. [Blog](https://unsloth.ai/docs/new/faster-moe) - **Embedding models**: Unsloth now supports ~1.8-3.3x faster embedding fine-tuning. [Blog](https://unsloth.ai/docs/new/embedding-finetuning) • [Notebooks](https://unsloth.ai/docs/get-started/unsloth-notebooks#embedding-models) - New **7x longer context RL** vs. all other setups, via our new batching algorithms. [Blog](https://unsloth.ai/docs/new/grpo-long-context) diff --git a/install.ps1 b/install.ps1 index df49414620..cceefd2647 100644 --- a/install.ps1 +++ b/install.ps1 @@ -53,7 +53,8 @@ function Install-UnslothStudio { param([string]$TorchIndexUrl) if ($SkipTorch) { return "none" } if ([string]::IsNullOrWhiteSpace($TorchIndexUrl)) { return "none" } - $leaf = ($TorchIndexUrl.TrimEnd('/') -split '/')[-1].ToLowerInvariant() + # Drop query/fragment first so a token-authenticated pin classifies by family. + $leaf = (($TorchIndexUrl -split '[?#]', 2)[0].TrimEnd('/') -split '/')[-1].ToLowerInvariant() if (@("cpu", "cu118", "cu124", "cu126", "cu128", "cu130") -contains $leaf) { return $leaf } if ($leaf -match '^rocm[0-9]+\.[0-9]+$') { return $leaf } return "auto" @@ -62,7 +63,8 @@ function Install-UnslothStudio { function Get-TauriGpuBranch { param([string]$TorchIndexFamily) if ($SkipTorch) { return "no_torch" } - if ($TorchIndexFamily -like "cu*") { return "cuda" } + # Require a digit after "cu" so /current or /custom isn't branded CUDA (parity ^cu[0-9]). + if ($TorchIndexFamily -match '^cu[0-9]') { return "cuda" } if ($TorchIndexFamily -like "rocm*") { return "rocm" } if ($TorchIndexFamily -eq "cpu") { return "cpu" } return "unknown" @@ -467,22 +469,35 @@ function Install-UnslothStudio { } } + # Redact index-URL credentials (userinfo + ?query= + #fragment) from captured installer + # output before printing on failure; uv/pip errors echo the failing --index-url verbatim. + # Mirrors the other installers. Verbose mode streams uncaptured, so it isn't redacted. + function Redact-InstallOutput { + param([string]$Text) + if (-not $Text) { return $Text } + $Text = $Text -replace '(https?://)[^/@\s`]+@', '$1@' + $Text = $Text -replace '([?&][^=\s&`]+)=[^&#\s`]+', '$1=' + # A #token=... fragment is as sensitive as a query; URL-anchored. + return $Text -replace '(https?://[^\s`#]+)#[^\s`]+', '$1#' + } + # Run native commands quietly by default to match install.sh behavior. # Full command output is shown only when --verbose / UNSLOTH_VERBOSE=1. function Invoke-InstallCommand { param( [Parameter(Mandatory = $true)][ScriptBlock]$Command ) - # Installer-pinned index installs (torch) must beat an inherited uv mirror - # (#6898): when the command pins an index, clear every uv index env var so - # it wins, then restore in finally. Other installs keep the user's mirror. + # Installer-pinned index installs (torch) must beat an inherited uv mirror (#6898): + # for --default-index, clear the uv index env vars (restore in finally) and set + # UV_NO_CONFIG=1 so a uv.toml/pyproject index can't outrank the CLI pin (uv 0.10). $savedUvIndex = $null if ($Command.ToString() -match '--default-index') { $savedUvIndex = @{} - foreach ($n in 'UV_DEFAULT_INDEX', 'UV_INDEX_URL', 'UV_INDEX', 'UV_EXTRA_INDEX_URL') { + foreach ($n in 'UV_DEFAULT_INDEX', 'UV_INDEX_URL', 'UV_INDEX', 'UV_EXTRA_INDEX_URL', 'UV_TORCH_BACKEND', 'UV_FIND_LINKS', 'UV_CONFIG_FILE', 'UV_NO_CONFIG') { $savedUvIndex[$n] = [Environment]::GetEnvironmentVariable($n) Remove-Item "Env:$n" -ErrorAction SilentlyContinue } + $env:UV_NO_CONFIG = '1' } $prevEap = $ErrorActionPreference $ErrorActionPreference = "Continue" @@ -493,17 +508,23 @@ function Install-UnslothStudio { # Merge stderr into stdout so progress/warning output stays visible # without flipping $? on successful native commands (PS 5.1 treats # stderr records as errors that set $? = $false even on exit code 0). - & $Command 2>&1 | Out-Host + # Redact per record: uv echoes index URLs (credentials and all) in + # its errors, and verbose mode must not bypass the quiet path's + # redaction. ForEach-Object/Out-Host leave $LASTEXITCODE untouched. + & $Command 2>&1 | ForEach-Object { Redact-InstallOutput "$_" } | Out-Host } else { $output = & $Command 2>&1 | Out-String if ($LASTEXITCODE -ne 0) { - Write-Host $output -ForegroundColor Red + Write-Host (Redact-InstallOutput $output) -ForegroundColor Red } } return [int]$LASTEXITCODE } finally { $ErrorActionPreference = $prevEap - if ($savedUvIndex) { foreach ($n in $savedUvIndex.Keys) { if ($null -ne $savedUvIndex[$n]) { Set-Item "Env:$n" $savedUvIndex[$n] } } } + if ($savedUvIndex) { + Remove-Item "Env:UV_NO_CONFIG" -ErrorAction SilentlyContinue + foreach ($n in $savedUvIndex.Keys) { if ($null -ne $savedUvIndex[$n]) { Set-Item "Env:$n" $savedUvIndex[$n] } } + } } } @@ -1960,10 +1981,31 @@ exit 0 # On an AMD GPU (no NVIDIA), surface the optional WSL-ROCm driver hint. if (-not $HasNvidiaSmi -and ($ROCmGfxArch -or $ROCmGpuLabel)) { Show-AmdWslDriverHint } + # Trim trailing slashes from the URL PATH only, preserving ?query / #fragment: a whole-URL + # TrimEnd corrupts a token ending in "/", a single strip leaves .../cu128// empty. Shared. + function Trim-IndexPathSlashes { + param([string]$Url) + $value = $Url.Trim() + $idx = $value.IndexOfAny([char[]]@('?', '#')) + if ($idx -lt 0) { + return $value.TrimEnd('/') + } + return $value.Substring(0, $idx).TrimEnd('/') + $value.Substring($idx) + } + # ── Choose the correct PyTorch index URL based on driver CUDA version ── # Mirrors Get-PytorchCudaTag in setup.ps1. function Get-TorchIndexUrl { $baseUrl = if ($env:UNSLOTH_PYTORCH_MIRROR) { $env:UNSLOTH_PYTORCH_MIRROR.TrimEnd('/') } else { "https://download.pytorch.org/whl" } + # Explicit pin -- skip ALL GPU probing (headless / CI / cross-install). + # UNSLOTH_TORCH_INDEX_URL wins (full URL, verbatim); _FAMILY is the leaf appended + # to the mirror base. Matches install.sh / install_python_stack.py. + if (-not [string]::IsNullOrWhiteSpace($env:UNSLOTH_TORCH_INDEX_URL)) { + return (Trim-IndexPathSlashes $env:UNSLOTH_TORCH_INDEX_URL) + } + if (-not [string]::IsNullOrWhiteSpace($env:UNSLOTH_TORCH_INDEX_FAMILY)) { + return "$baseUrl/$($env:UNSLOTH_TORCH_INDEX_FAMILY.Trim().Trim('/'))" + } if (-not $NvidiaSmiExe) { return "$baseUrl/cpu" } try { $output = Invoke-NvidiaSmiBounded $NvidiaSmiExe @@ -1984,6 +2026,25 @@ exit 0 return "$baseUrl/cu126" } + # Strip userinfo AND query/fragment so an authenticated pin never leaks. Shared with + # _strip_index_url_credentials (install.sh / py / setup.ps1). + function Remove-IndexUrlCredentials { + param([string]$Url) + $sep = $Url.IndexOf('://') + if ($sep -lt 0) { return $Url } + $scheme = $Url.Substring(0, $sep) + $rest = $Url.Substring($sep + 3) + # Drop query / fragment (may hold auth tokens). + $q = $rest.IndexOfAny([char[]]('?', '#')) + if ($q -ge 0) { $rest = $rest.Substring(0, $q) } + $slash = $rest.IndexOf('/') + $authority = if ($slash -ge 0) { $rest.Substring(0, $slash) } else { $rest } + $at = $authority.LastIndexOf('@') + $host_ = if ($at -ge 0) { $authority.Substring($at + 1) } else { $authority } + if ($slash -ge 0) { return "${scheme}://${host_}$($rest.Substring($slash))" } + return "${scheme}://${host_}" + } + # ── Torch flavor helpers (to repair a stale CPU / wrong-CUDA wheel) ── # torch.__version__ -> flavor tag (cuXXX / rocm / cpu); untagged wheel = cpu, # matching setup.ps1's stale-venv parse. @@ -2002,11 +2063,13 @@ exit 0 param([string]$TorchIndexUrl, [string]$ROCmIndexUrl) if (-not [string]::IsNullOrWhiteSpace($ROCmIndexUrl)) { return 'rocm' } if ([string]::IsNullOrWhiteSpace($TorchIndexUrl)) { return $null } - $leaf = ($TorchIndexUrl.TrimEnd('/') -split '/')[-1].ToLowerInvariant() + # Drop query/fragment first so .../cu128?token=x classifies as cu128 (else it reinstalls every run). + $leaf = (($TorchIndexUrl -split '[?#]', 2)[0].TrimEnd('/') -split '/')[-1].ToLowerInvariant() if ($leaf -match '^cu\d+$') { return $leaf } if ($leaf -eq 'cpu') { return 'cpu' } if ($leaf -match '^rocm') { return 'rocm' } - if ($leaf -match '^gfx') { return 'rocm' } + # gfx must be followed by a digit (an architecture leaf); gfx-private is custom. + if ($leaf -match '^gfx[0-9]') { return 'rocm' } return $null } @@ -2041,6 +2104,10 @@ exit 0 } catch { return $null } } + # An explicit pin is authoritative: the AMD ROCm reroute below must not rewrite it + # (e.g. a deliberate cpu pin on an AMD host). + $TorchIndexPinned = (-not [string]::IsNullOrWhiteSpace($env:UNSLOTH_TORCH_INDEX_URL)) -or ` + (-not [string]::IsNullOrWhiteSpace($env:UNSLOTH_TORCH_INDEX_FAMILY)) $TorchIndexUrl = Get-TorchIndexUrl # ── GPU arch → newest compatible Windows ROCm wheel release ── @@ -2052,13 +2119,19 @@ exit 0 # Override with UNSLOTH_ROCM_WINDOWS_MIRROR for air-gapped / mirror installs. $ROCmIndexUrl = $null $ROCmTorchFloor = $null - if (($HasROCm -or $ROCmGfxArch) -and $TorchIndexUrl -like "*/cpu" -and -not $SkipTorch) { + $PinnedRocmVisionSpec = $null + $PinnedRocmAudioSpec = $null + if (-not $TorchIndexPinned -and ($HasROCm -or $ROCmGfxArch) -and $TorchIndexUrl -like "*/cpu" -and -not $SkipTorch) { $amdIndexBase = if ($env:UNSLOTH_ROCM_WINDOWS_MIRROR) { $env:UNSLOTH_ROCM_WINDOWS_MIRROR.TrimEnd('/') } else { "https://repo.amd.com/rocm/whl" } $archFamilyMap = @{ "gfx1201" = "gfx120X-all"; "gfx1200" = "gfx120X-all" # RDNA 4 "gfx1151" = "gfx1151"; "gfx1150" = "gfx1150" # RDNA 3.5 (Strix Halo/Point) "gfx1103" = "gfx110X-all"; "gfx1102" = "gfx110X-all" # RDNA 3 "gfx1101" = "gfx110X-all"; "gfx1100" = "gfx110X-all" + "gfx1036" = "gfx103X-all"; "gfx1035" = "gfx103X-all" # RDNA 2 (RX 6000) + "gfx1034" = "gfx103X-all"; "gfx1033" = "gfx103X-all" + "gfx1032" = "gfx103X-all"; "gfx1031" = "gfx103X-all" + "gfx1030" = "gfx103X-all" "gfx90a" = "gfx90a"; "gfx908" = "gfx908" # MI200/MI100 } # gfx120X (RDNA 4) and gfx1151/gfx1150 (Strix) have a null-pointer bug in @@ -2102,6 +2175,32 @@ exit 0 } } + # A gfx*/rocm pin skips the auto-reroute above, but the generic CPU/CUDA install below + # would use torch>=2.4,<2.11 and pull a known-bad wheel on the gfx115x/gfx120x/rocm>=7.2 + # indexes (the _grouped_mm bug). Route a pinned ROCm index through the ROCm path. + if ($TorchIndexPinned -and -not $ROCmIndexUrl -and -not $SkipTorch) { + $_pinLeaf = (($TorchIndexUrl -split '[?#]', 2)[0].TrimEnd('/') -split '/')[-1].ToLower() + $_pinRocm211 = $false + # Anchor ($) so a suffixed custom leaf (rocm7.2-private) falls through to verbatim. + if ($_pinLeaf -match '^rocm(\d+)\.(\d+)$') { + # Only KNOWN-2.11 rocm (rocm7.2) gets the floor. Matches Test-RocmKnown211Version. + $_pinRocm211 = ([int]$Matches[1] -eq 7 -and [int]$Matches[2] -eq 2) + } + # Only the 2.11-allowlist gfx arches need the floor; others publish <2.11 and stay bare. + $_pinGfx211 = @('gfx120x-all', 'gfx1151', 'gfx1150') -contains $_pinLeaf + if ($_pinGfx211 -or $_pinRocm211) { + $ROCmIndexUrl = $TorchIndexUrl + $ROCmTorchFloor = "torch>=2.11.0,<2.12.0" + $PinnedRocmVisionSpec = "torchvision>=0.26.0,<0.27.0" + $PinnedRocmAudioSpec = "torchaudio>=2.11.0,<2.12.0" + substep "pinned ROCm index ($_pinLeaf) -- enforcing $ROCmTorchFloor" "Cyan" + } elseif ($_pinLeaf -match '^gfx[0-9]' -or $_pinLeaf -match '^rocm[0-9]+(\.[0-9]+)?$') { + # Other gfx / older rocm (<=7.1) ship torch <2.11; route via the ROCm path with + # bare specs. Only EXACT rocm/gfx* are families; a suffixed leaf is verbatim. + $ROCmIndexUrl = $TorchIndexUrl + } + } + if ($ROCmIndexUrl) { $TorchIndexFamily = "rocm" } else { @@ -2164,14 +2263,14 @@ exit 0 } if ($_Migrated) { - # Migrated env: force-reinstall unsloth+unsloth-zoo to ensure clean state - # in the new venv location, while preserving existing torch/CUDA + # Migrated env: force-reinstall unsloth+unsloth-zoo for a clean state, preserving + # existing torch/CUDA unless the flavor repair below re-lands it. Write-TauriLog "STEP" "Installing unsloth" substep "upgrading unsloth in migrated environment..." if ($SkipTorch) { # No-torch: install unsloth + unsloth-zoo with --no-deps, then # runtime deps (typer, safetensors, transformers, etc.) with --no-deps. - $baseInstallExit = Invoke-InstallCommandRetry -Label "install unsloth (migrated no-torch)" { uv pip install --python $VenvPython --no-deps --reinstall-package unsloth --reinstall-package unsloth-zoo "unsloth>=2026.7.3" "unsloth-zoo>=2026.7.3" } + $baseInstallExit = Invoke-InstallCommandRetry -Label "install unsloth (migrated no-torch)" { uv pip install --python $VenvPython --no-deps --reinstall-package unsloth --reinstall-package unsloth-zoo "unsloth>=2026.7.4" "unsloth-zoo>=2026.7.4" } if ($baseInstallExit -eq 0) { # Resolve pydantic WITH deps so pip pins pydantic-core # to the matching version (no-torch-runtime.txt below @@ -2185,7 +2284,7 @@ exit 0 } } } else { - $baseInstallExit = Invoke-InstallCommandRetry -Label "install unsloth (migrated)" { uv pip install --python $VenvPython --reinstall-package unsloth --reinstall-package unsloth-zoo "unsloth>=2026.7.3" "unsloth-zoo>=2026.7.3" } + $baseInstallExit = Invoke-InstallCommandRetry -Label "install unsloth (migrated)" { uv pip install --python $VenvPython --reinstall-package unsloth --reinstall-package unsloth-zoo "unsloth>=2026.7.4" "unsloth-zoo>=2026.7.4" } } if ($baseInstallExit -ne 0) { Write-Host "[ERROR] Failed to install unsloth (exit code $baseInstallExit)" -ForegroundColor Red @@ -2210,22 +2309,24 @@ exit 0 substep "skipping PyTorch (--no-torch flag set)." "Yellow" } elseif ($ROCmIndexUrl) { Write-TauriLog "STEP" "Installing PyTorch (AMD ROCm Windows)" - substep "installing PyTorch from $ROCmIndexUrl..." + substep "installing PyTorch from $(Remove-IndexUrlCredentials $ROCmIndexUrl)..." $torchSpec = if ($ROCmTorchFloor) { $ROCmTorchFloor } else { "torch" } # Pin the companions to match $torchSpec; bare names can resolve an # ABI-incompatible torchvision/torchaudio on AMD's per-arch index. - $visionSpec = if ($ROCmGfxArch -and $torchvisionFloorMap.ContainsKey($ROCmGfxArch)) { $torchvisionFloorMap[$ROCmGfxArch] } else { "torchvision" } - $audioSpec = if ($ROCmGfxArch -and $torchaudioFloorMap.ContainsKey($ROCmGfxArch)) { $torchaudioFloorMap[$ROCmGfxArch] } else { "torchaudio" } + $visionSpec = if ($PinnedRocmVisionSpec) { $PinnedRocmVisionSpec } elseif ($ROCmGfxArch -and $torchvisionFloorMap -and $torchvisionFloorMap.ContainsKey($ROCmGfxArch)) { $torchvisionFloorMap[$ROCmGfxArch] } else { "torchvision" } + $audioSpec = if ($PinnedRocmAudioSpec) { $PinnedRocmAudioSpec } elseif ($ROCmGfxArch -and $torchaudioFloorMap -and $torchaudioFloorMap.ContainsKey($ROCmGfxArch)) { $torchaudioFloorMap[$ROCmGfxArch] } else { "torchaudio" } $torchInstallExit = Invoke-InstallCommandRetry -Label "install PyTorch (AMD ROCm)" { uv pip install --python $VenvPython --force-reinstall --default-index $ROCmIndexUrl $torchSpec $visionSpec $audioSpec } if ($torchInstallExit -ne 0) { - # Transient AMD-index failure: fall back to a CPU base so the install - # still completes; Unsloth setup retries ROCm afterwards. + # Transient AMD-index failure: fall back to a CPU base (Unsloth setup retries + # ROCm). Use an explicit CPU index -- for a pinned ROCm index $TorchIndexUrl IS + # the ROCm mirror, so reusing it would just retry it. + $CpuFallbackIndexUrl = if ($env:UNSLOTH_PYTORCH_MIRROR) { "$($env:UNSLOTH_PYTORCH_MIRROR.TrimEnd('/'))/cpu" } else { "https://download.pytorch.org/whl/cpu" } substep "ROCm PyTorch install failed (exit $torchInstallExit); using a CPU base, Unsloth setup retries ROCm." "Yellow" # --force-reinstall: a failed ROCm install can leave an unpinned ROCm # torch (e.g. 2.10.0+rocm on gfx110X/gfx90a) that still satisfies the CPU # torch>= range, so without it uv would keep the ROCm build and only swap # the companions -- a mismatched venv the flavor-repair block won't fix. - $torchInstallExit = Invoke-InstallCommandRetry -Label "install PyTorch (CPU fallback)" { uv pip install --python $VenvPython --force-reinstall "torch>=2.4,<2.11.0" torchvision torchaudio --default-index $TorchIndexUrl } + $torchInstallExit = Invoke-InstallCommandRetry -Label "install PyTorch (CPU fallback)" { uv pip install --python $VenvPython --force-reinstall "torch>=2.4,<2.11.0" "torchvision>=0.19,<0.26.0" "torchaudio>=2.4,<2.11.0" --default-index $CpuFallbackIndexUrl } if ($torchInstallExit -ne 0) { Write-Host "[ERROR] Failed to install PyTorch (ROCm and CPU base both failed, exit code $torchInstallExit)" -ForegroundColor Red return (Exit-InstallFailure "Failed to install PyTorch (exit code $torchInstallExit)" $torchInstallExit) @@ -2238,8 +2339,14 @@ exit 0 } } else { Write-TauriLog "STEP" "Installing PyTorch" - substep "installing PyTorch ($TorchIndexUrl)..." - $torchInstallExit = Invoke-InstallCommandRetry -Label "install PyTorch" { uv pip install --python $VenvPython "torch>=2.4,<2.11.0" torchvision torchaudio --default-index $TorchIndexUrl } + substep "installing PyTorch ($(Remove-IndexUrlCredentials $TorchIndexUrl))..." + # Bound the companions to the capped torch on EVERY index, cu + # families included: torchaudio 2.11 dropped its exact torch pin from + # the wheel metadata, so a bare companion next to torch<2.11 can + # resolve a mismatched 2.11.0 build. Mirrors install.sh. + $_pinVisionSpec = "torchvision>=0.19,<0.26.0" + $_pinAudioSpec = "torchaudio>=2.4,<2.11.0" + $torchInstallExit = Invoke-InstallCommandRetry -Label "install PyTorch" { uv pip install --python $VenvPython "torch>=2.4,<2.11.0" $_pinVisionSpec $_pinAudioSpec --default-index $TorchIndexUrl } if ($torchInstallExit -ne 0) { Write-Host "[ERROR] Failed to install PyTorch (exit code $torchInstallExit)" -ForegroundColor Red return (Exit-InstallFailure "Failed to install PyTorch (exit code $torchInstallExit)" $torchInstallExit) @@ -2251,7 +2358,7 @@ exit 0 if ($SkipTorch) { # No-torch: install unsloth + unsloth-zoo with --no-deps, then # runtime deps (typer, safetensors, transformers, etc.) with --no-deps. - $baseInstallExit = Invoke-InstallCommandRetry -Label "install unsloth (no-torch)" { uv pip install --python $VenvPython --no-deps --upgrade-package unsloth --upgrade-package unsloth-zoo "unsloth>=2026.7.3" "unsloth-zoo>=2026.7.3" } + $baseInstallExit = Invoke-InstallCommandRetry -Label "install unsloth (no-torch)" { uv pip install --python $VenvPython --no-deps --upgrade-package unsloth --upgrade-package unsloth-zoo "unsloth>=2026.7.4" "unsloth-zoo>=2026.7.4" } if ($baseInstallExit -eq 0) { # Same pydantic-with-deps trick as the migrated branch. $baseInstallExit = Invoke-InstallCommandRetry -Label "install pydantic" { uv pip install --python $VenvPython pydantic } @@ -2263,7 +2370,7 @@ exit 0 } } } elseif ($StudioLocalInstall) { - $baseInstallExit = Invoke-InstallCommandRetry -Label "install unsloth (local)" { uv pip install --python $VenvPython --upgrade-package unsloth "unsloth>=2026.7.3" "unsloth-zoo>=2026.7.3" } + $baseInstallExit = Invoke-InstallCommandRetry -Label "install unsloth (local)" { uv pip install --python $VenvPython --upgrade-package unsloth "unsloth>=2026.7.4" "unsloth-zoo>=2026.7.4" } } else { $baseInstallExit = Invoke-InstallCommandRetry -Label "install unsloth" { uv pip install --python $VenvPython --upgrade-package unsloth -- "$PackageName" } } @@ -2291,7 +2398,7 @@ exit 0 Write-TauriLog "STEP" "Installing unsloth" substep "installing unsloth (this may take a few minutes)..." if ($StudioLocalInstall) { - $baseInstallExit = Invoke-InstallCommandRetry -Label "install unsloth (auto torch backend)" { uv pip install --python $VenvPython "unsloth-zoo>=2026.7.3" "unsloth>=2026.7.3" --torch-backend=auto } + $baseInstallExit = Invoke-InstallCommandRetry -Label "install unsloth (auto torch backend)" { uv pip install --python $VenvPython "unsloth-zoo>=2026.7.4" "unsloth>=2026.7.4" --torch-backend=auto } if ($baseInstallExit -ne 0) { Write-Host "[ERROR] Failed to install unsloth (exit code $baseInstallExit)" -ForegroundColor Red return (Exit-InstallFailure "Failed to install unsloth (exit code $baseInstallExit)" $baseInstallExit) @@ -2317,6 +2424,13 @@ exit 0 } } + $installedPackageVersion = (& $VenvPython -c "from importlib.metadata import version; import sys; print(version(sys.argv[1]))" $PackageName 2>$null | Out-String).Trim() + if ($LASTEXITCODE -eq 0 -and $installedPackageVersion) { + step $PackageName "$installedPackageVersion installed" + } else { + substep "[WARN] installed $PackageName version could not be determined" "Yellow" + } + # ── Enforce the installed torch flavor matches the detected GPU build ── # PEP 440 ignores the +cpu/+cuXXX/+rocm local label in a version range, so uv # keeps a stale torch==X+cpu against a CUDA index and setup.ps1 then loops on @@ -2335,8 +2449,8 @@ exit 0 $rocmSpec = if ($ROCmTorchFloor) { $ROCmTorchFloor } else { "torch" } # Pin companions like the fresh ROCm path (bare names can pull an # ABI-incompatible torchvision/torchaudio from the per-arch index). - $visionSpec = if ($ROCmGfxArch -and $torchvisionFloorMap.ContainsKey($ROCmGfxArch)) { $torchvisionFloorMap[$ROCmGfxArch] } else { "torchvision" } - $audioSpec = if ($ROCmGfxArch -and $torchaudioFloorMap.ContainsKey($ROCmGfxArch)) { $torchaudioFloorMap[$ROCmGfxArch] } else { "torchaudio" } + $visionSpec = if ($PinnedRocmVisionSpec) { $PinnedRocmVisionSpec } elseif ($ROCmGfxArch -and $torchvisionFloorMap -and $torchvisionFloorMap.ContainsKey($ROCmGfxArch)) { $torchvisionFloorMap[$ROCmGfxArch] } else { "torchvision" } + $audioSpec = if ($PinnedRocmAudioSpec) { $PinnedRocmAudioSpec } elseif ($ROCmGfxArch -and $torchaudioFloorMap -and $torchaudioFloorMap.ContainsKey($ROCmGfxArch)) { $torchaudioFloorMap[$ROCmGfxArch] } else { "torchaudio" } substep "PyTorch flavor mismatch (installed $installedTorchTag, need ROCm) -- reinstalling correct build..." "Yellow" $torchFixExit = Invoke-InstallCommand { uv pip install --python $VenvPython --force-reinstall --default-index $ROCmIndexUrl $rocmSpec $visionSpec $audioSpec } if ($torchFixExit -ne 0) { @@ -2347,7 +2461,7 @@ exit 0 } elseif ($expectedTorchTag -ne 'rocm') { # CUDA: stale +cpu (or wrong cuXXX) against a CUDA index -> reinstall triplet. substep "PyTorch flavor mismatch (installed $installedTorchTag, need $expectedTorchTag) -- reinstalling correct build..." "Yellow" - $torchFixExit = Invoke-InstallCommand { uv pip install --python $VenvPython "torch>=2.4,<2.11.0" torchvision torchaudio --default-index $TorchIndexUrl --reinstall-package torch --reinstall-package torchvision --reinstall-package torchaudio } + $torchFixExit = Invoke-InstallCommand { uv pip install --python $VenvPython "torch>=2.4,<2.11.0" "torchvision>=0.19,<0.26.0" "torchaudio>=2.4,<2.11.0" --default-index $TorchIndexUrl --reinstall-package torch --reinstall-package torchvision --reinstall-package torchaudio } if ($torchFixExit -ne 0) { Write-Host "[ERROR] Failed to reinstall PyTorch with the correct CUDA build (exit code $torchFixExit)" -ForegroundColor Red return (Exit-InstallFailure "Failed to reinstall PyTorch ($expectedTorchTag) (exit code $torchFixExit)" $torchFixExit) diff --git a/install.sh b/install.sh index 5972379d26..445bab616a 100755 --- a/install.sh +++ b/install.sh @@ -159,18 +159,58 @@ run_maybe_quiet() { fi } +# Trim trailing slashes from the URL PATH only, preserving ?query / #fragment: a whole-URL +# strip corrupts a token ending in "/", a single strip leaves .../cu128// empty. Shared. +_trim_index_path_slashes() { + _tips_v="$1" + case "$_tips_v" in + *[?#]*) + _tips_head="${_tips_v%%[?#]*}" + _tips_tail="${_tips_v#"$_tips_head"}" + ;; + *) + _tips_head="$_tips_v" + _tips_tail="" + ;; + esac + while [ -n "$_tips_head" ] && [ "${_tips_head%/}" != "$_tips_head" ]; do + _tips_head="${_tips_head%/}" + done + printf '%s%s' "$_tips_head" "$_tips_tail" +} + +# Redact index-URL credentials (userinfo + ?query= + #fragment) from captured installer +# output before printing on failure; uv/pip errors echo the failing --index-url verbatim. +# Mirrors the other installers. Verbose mode streams uncaptured, so it isn't redacted. +_redact_install_output() { + sed -E \ + -e 's#(https?://)[^/@[:space:]`]+@#\1@#g' \ + -e 's#([?&][^=[:space:]&`]+)=[^&#[:space:]`]+#\1=#g' \ + -e 's|(https?://[^[:space:]`#]+)#[^[:space:]`]+|\1#|g' \ + "$@" +} + run_install_cmd() { _label="$1" shift - # Installer-pinned index installs (torch) must beat an inherited uv mirror - # (#6898): when we pass --default-index, neutralize every uv index env var so - # the pinned index wins. Other installs keep the user's mirror. + # Installer-pinned index installs (torch) must beat an inherited uv mirror (#6898): + # for --default-index, neutralize the uv index/backend/config vars (UV_TORCH_BACKEND + # redirects torch; UV_NO_CONFIG=1 + dropping UV_CONFIG_FILE stops a uv.toml/pyproject + # index outranking the CLI pin, uv 0.10). case " $* " in - *" --default-index "*) set -- env -u UV_DEFAULT_INDEX -u UV_INDEX_URL -u UV_INDEX -u UV_EXTRA_INDEX_URL "$@" ;; + *" --default-index "*) set -- env -u UV_DEFAULT_INDEX -u UV_INDEX_URL -u UV_INDEX -u UV_EXTRA_INDEX_URL -u UV_TORCH_BACKEND -u UV_FIND_LINKS -u UV_CONFIG_FILE UV_NO_CONFIG=1 "$@" ;; esac if _is_verbose; then - "$@" && return 0 - _rc=$? + # Stream through the redactor: uv echoes index URLs (credentials and + # all) in its errors, and verbose mode previously bypassed the + # redaction the quiet path applies. The rc file preserves the + # command's exit code across the pipe without relying on pipefail + # (this script runs under plain sh). + _rcf=$(mktemp) + { "$@" 2>&1; printf '%s' "$?" > "$_rcf"; } | _redact_install_output + _rc=$(cat "$_rcf" 2>/dev/null || echo 1) + rm -f "$_rcf" + [ "${_rc:-1}" -eq 0 ] 2>/dev/null && return 0 step "error" "$_label failed (exit code $_rc)" "$C_ERR" >&2 return "$_rc" fi @@ -178,7 +218,7 @@ run_install_cmd() { "$@" >"$_log" 2>&1 && { rm -f "$_log"; return 0; } _rc=$? step "error" "$_label failed (exit code $_rc)" "$C_ERR" >&2 - cat "$_log" >&2 + _redact_install_output "$_log" >&2 rm -f "$_log" return $_rc } @@ -257,7 +297,7 @@ _install_bnb_rocm() { fi _bnb_rc=$? if _is_verbose; then - cat "$_bnb_log" >&2 + _redact_install_output "$_bnb_log" >&2 fi rm -f "$_bnb_log" step "warning" "$_label (pre-release) failed (exit code $_bnb_rc)" "$C_WARN" >&2 @@ -310,6 +350,11 @@ _tauri_torch_index_family() { return fi _diag_url="${1:-}" + # Strip query/fragment AND a trailing slash before classifying (like _torch_index_url_leaf): + # a token isn't echoed into [TAURI:DIAG], and .../cu128/?token=x still classifies as cu128. + _diag_url="${_diag_url%%\?*}" + _diag_url="${_diag_url%%#*}" + _diag_url="${_diag_url%/}" case "$_diag_url" in */cu118) echo "cu118" ;; */cu124) echo "cu124" ;; @@ -343,7 +388,8 @@ _tauri_gpu_branch() { return fi case "$_diag_family" in - cu*) echo "cuda" ;; + # Require a digit after cu so /current or /custom isn't branded CUDA (parity ^cu[0-9]). + cu[0-9]*) echo "cuda" ;; rocm*) if [ "$_diag_radeon" = true ]; then echo "rocm_radeon" @@ -472,11 +518,13 @@ _on_install_exit() { _restore_studio_venv_replacement fi [ -n "${_UV_OVERRIDE_TMPDIR:-}" ] && rm -rf "$_UV_OVERRIDE_TMPDIR" 2>/dev/null || true + [ -n "${_UNSLOTH_TORCH_OVERRIDES:-}" ] && rm -f "$_UNSLOTH_TORCH_OVERRIDES" 2>/dev/null || true exit "$_status" } -# Empty so an inherited value can never reach the trap's rm; only a temp dir -# this script creates below (Apple Silicon, spaced path) is ever removed. +# Empty so an inherited value never reaches the trap's rm; only temp paths this +# script creates below (spaced-path dir, torch-trio overrides) are removed. _UV_OVERRIDE_TMPDIR="" +_UNSLOTH_TORCH_OVERRIDES="" trap _on_install_exit EXIT # ── Helper: download a URL to a file (supports curl and wget) ── @@ -1573,6 +1621,12 @@ _has_usable_nvidia_gpu() { # the STUDIO_HOME mkdir/venv so the origin distro is untouched. _maybe_reroute_strixhalo_to_2404() { [ "${OS:-}" = "wsl" ] || return 0 + # An explicit index pin skips every GPU-driven reroute (same contract as + # the later Radeon/Strix guard): the pin is honored in THIS distro rather + # than probing the GPU and switching distributions. Whitespace-only + # overrides do not gate (parity with get_torch_index_url). + _rr_pin=$(printf '%s' "${UNSLOTH_TORCH_INDEX_URL:-}${UNSLOTH_TORCH_INDEX_FAMILY:-}" | tr -d '[:space:]') + [ -n "$_rr_pin" ] && return 0 [ "${SKIP_TORCH:-false}" = "false" ] || return 0 [ "${UNSLOTH_SKIP_ROCM_WSL_SETUP:-0}" = "1" ] && return 0 [ "${UNSLOTH_WSL_REROUTED:-0}" = "1" ] && return 0 @@ -1634,6 +1688,10 @@ _maybe_reroute_strixhalo_to_2404() { # Forward explicit ROCm-bootstrap consent (e.g. Tauri) so the child auto-enables the # GPU instead of falling back to the desktop-app prompt path. [ "${UNSLOTH_ROCM_WSL_AUTO:-0}" = "1" ] && _rr_exports="$_rr_exports; export UNSLOTH_ROCM_WSL_AUTO=1" + # Forward a pinned torch index into the rerouted distro; dropping it would + # silently revert the child install to auto-detection. + [ -n "${UNSLOTH_TORCH_INDEX_URL:-}" ] && _rr_exports="$_rr_exports; export UNSLOTH_TORCH_INDEX_URL=$(_rr_q "$UNSLOTH_TORCH_INDEX_URL")" + [ -n "${UNSLOTH_TORCH_INDEX_FAMILY:-}" ] && _rr_exports="$_rr_exports; export UNSLOTH_TORCH_INDEX_FAMILY=$(_rr_q "$UNSLOTH_TORCH_INDEX_FAMILY")" [ "$_SKIP_AUTOSTART" = true ] && _rr_exports="$_rr_exports; export UNSLOTH_SKIP_AUTOSTART=1" _rr_args="" [ "$PACKAGE_NAME" != "unsloth" ] && _rr_args="$_rr_args --package $(_rr_q "$PACKAGE_NAME")" @@ -1821,6 +1879,8 @@ tauri_log "STEP" "Creating virtual environment" mkdir -p "$STUDIO_HOME" _MIGRATED=false +# Empty so an inherited value can never masquerade as a probed torch version. +_PREV_TORCH_VER="" if [ -x "$VENV_DIR/bin/python" ]; then # why: matching guard to the .venv branch below -- in env-mode @@ -1838,6 +1898,12 @@ if [ -x "$VENV_DIR/bin/python" ]; then echo " Move it aside or choose an empty UNSLOTH_STUDIO_HOME." >&2 exit 1 fi + # Record the existing venv's torch BEFORE the replacement moves it aside: a re-run + # rebuilds the venv for clean state, but must keep the torch release the user + # already has (see _previous_torch_pin below). Last line only: sitecustomize or + # import-hook noise on stdout must not corrupt the version. + _PREV_TORCH_VER=$("$VENV_DIR/bin/python" -c \ + "import torch; print(torch.__version__)" 2>/dev/null | tail -n 1 || true) # New layout already exists — replace only after preserving rollback copy. substep "preserving existing environment for rollback..." _start_studio_venv_replacement "$VENV_DIR" @@ -1991,6 +2057,15 @@ if [ "$SKIP_TORCH" = false ] && [ "$OS" = "macos" ] && [ "$_ARCH" = "arm64" ]; t TORCH_CONSTRAINT="torch>=2.6,<2.11.0" fi fi +# Companion (torchvision/torchaudio) constraints, bounded to torch's window. +# torchaudio 2.11 dropped its exact torch pin, so a bare companion next to a +# <2.11-capped torch resolves torchaudio 2.11 (verified: cpu leaf installed +# torch 2.10.0+cpu with torchaudio 2.11.0+cpu). torchvision still exact-pins +# torch and self-corrects, but is bounded for symmetry. Widened alongside the +# cu* torch window below; the torch-2.11 AMD paths (rocm7.2 / per-gfx / Strix) +# pin their own trio. +TORCHVISION_CONSTRAINT="torchvision>=0.19,<0.26.0" +TORCHAUDIO_CONSTRAINT="torchaudio>=2.4,<2.11.0" # ── Resolve repo root (for --local installs) ── _REPO_ROOT="$(cd "$(dirname "$0" 2>/dev/null || echo ".")" && pwd)" @@ -2059,6 +2134,24 @@ _has_amd_rocm_gpu() { get_torch_index_url() { _base="${UNSLOTH_PYTORCH_MIRROR:-https://download.pytorch.org/whl}" _base="${_base%/}" + # Explicit override -- skip ALL GPU probing (headless / container / CI / cross-install). + # UNSLOTH_TORCH_INDEX_URL wins (full URL, verbatim); _FAMILY is the leaf (cpu, cu128, ...) + # appended to the mirror base. Trim whitespace so a whitespace-only value is unset. + _url="${UNSLOTH_TORCH_INDEX_URL:-}" + _url="${_url#"${_url%%[![:space:]]*}"}"; _url="${_url%"${_url##*[![:space:]]}"}" + if [ -n "$_url" ]; then + # Trim trailing PATH slashes (a multi-slash path 404s on strict pip proxies) while + # preserving a ?query/#fragment token (a whole-URL strip would eat a "/"-ending token). + _url=$(_trim_index_path_slashes "$_url") + echo "$_url"; return + fi + _family="${UNSLOTH_TORCH_INDEX_FAMILY:-}" + _family="${_family#"${_family%%[![:space:]]*}"}"; _family="${_family%"${_family##*[![:space:]]}"}" + if [ -n "$_family" ]; then + while [ "${_family#/}" != "$_family" ]; do _family="${_family#/}"; done + while [ "${_family%/}" != "$_family" ]; do _family="${_family%/}"; done + echo "$_base/$_family"; return + fi # macOS: always CPU (no CUDA support) case "$(uname -s)" in Darwin) echo "$_base/cpu"; return ;; esac # Try nvidia-smi -- require the binary to actually list a usable GPU. @@ -2187,16 +2280,155 @@ _torch_flavor_tag() { esac } +# Final path segment of a wheel index URL ($1), lowercased, query/fragment stripped first +# so a token-authenticated pin (.../cu128?token=x) classifies as cu128 (else it reinstalls +# every update). Classification only. Shared with the py / ps1 leaf extractors. +_torch_index_url_leaf() { + _tl_u="${1%%\?*}" + _tl_u="${_tl_u%%#*}" + # Strip ALL trailing slashes, not one: .../rocm7.2// must yield rocm7.2, not an empty leaf. + while [ -n "$_tl_u" ] && [ "${_tl_u%/}" != "$_tl_u" ]; do + _tl_u="${_tl_u%/}" + done + printf '%s' "${_tl_u##*/}" | tr '[:upper:]' '[:lower:]' +} + +# True (exit 0) when a lowercased leaf is an EXACT pip ROCm family: rocm[.] +# or a gfx ARCHITECTURE leaf (gfx followed by a digit: gfx90a, gfx1151, gfx120x-all). A leaf +# that merely starts with rocm/gfx (rocm7.2-private, gfx-private) is a custom verbatim pin. +# Matches the py / ps1 sides. +_is_pip_rocm_family_leaf() { + case "$1" in + gfx[0-9]*) return 0 ;; + rocm[0-9]*) + # Exact rocm[.]: both major and minor must be non-empty all-digits + # (rocm7., rocm7.2.1, rocm7.2-private are all custom pins, not a family). + _rocm_rest="${1#rocm}" + case "$_rocm_rest" in + *.*.*) return 1 ;; + *.*) + _rocm_minor="${_rocm_rest#*.}" + case "${_rocm_rest%%.*}" in "" | *[!0-9]*) return 1 ;; esac + case "$_rocm_minor" in "" | *[!0-9]*) return 1 ;; esac + ;; + *[!0-9]*) return 1 ;; + esac + return 0 + ;; + *) return 1 ;; + esac +} + +# Whether release base $1 (X.Y[.Z...]) falls inside constraint window $2 +# ("torch>=A.B[.C],="*",<"*) ;; + *) echo "no"; return ;; + esac + _trw_floor="${_trw_con#torch>=}"; _trw_floor="${_trw_floor%%,*}" + _trw_ceil="${_trw_con##*,<}" + _v_maj="${1%%.*}"; _v_rest="${1#*.}"; _v_min="${_v_rest%%.*}" + _f_maj="${_trw_floor%%.*}"; _f_rest="${_trw_floor#*.}"; _f_min="${_f_rest%%.*}" + _c_maj="${_trw_ceil%%.*}"; _c_rest="${_trw_ceil#*.}"; _c_min="${_c_rest%%.*}" + for _trw_n in "$_v_maj" "$_v_min" "$_f_maj" "$_f_min" "$_c_maj" "$_c_min"; do + case "$_trw_n" in ''|*[!0-9]*) echo "no"; return ;; esac + done + if [ "$_v_maj" -gt "$_f_maj" ] || { [ "$_v_maj" -eq "$_f_maj" ] && [ "$_v_min" -ge "$_f_min" ]; }; then + if [ "$_v_maj" -lt "$_c_maj" ] || { [ "$_v_maj" -eq "$_c_maj" ] && [ "$_v_min" -lt "$_c_min" ]; }; then + echo "yes" + return + fi + fi + echo "no" +} + +# Keep the previous venv's torch on a re-run: echo "torch==X.Y.Z" when the probed +# version ($1) is inside the active constraint window ($2), else "". The RELEASE is kept +# regardless of flavor tag; the pin installs from the freshly chosen index, so flavor +# follows the machine (cpu <-> cuda, cu126 -> cu130, PyPI bare -> +cu130) while the +# release follows the user. Gating on flavor was wrong: a PyPI torch reports a BARE +# version (on Linux the PyPI wheel IS CUDA), misclassified "cpu", so a healthy 2.10 on a +# cu130 host was moved to 2.11. Per-leaf floors still win (rocm7.2 / gfx >=2.11 for the +# Strix _grouped_mm fix, out-of-window manual installs) and are never pinned; the caller's +# _PREV_FALLBACK_CONSTRAINT installs the newest supported release when the index lacks the +# exact one. Opt out with UNSLOTH_TORCH_UPGRADE=1. +_previous_torch_pin() { + _ptp_ver="$1" + _ptp_con="$2" + [ -n "$_ptp_ver" ] || { echo ""; return; } + [ "${UNSLOTH_TORCH_UPGRADE:-0}" = "1" ] && { echo ""; return; } + _ptp_base="${_ptp_ver%%+*}" + # Base must be a plain numeric release (X.Y[.Z]); probe noise and + # nightly/dev/source builds (2.11.0.dev20250704, 2.9.0a0) must never + # become a pin -- no stable index carries them, so pinning would only + # print "keeping it" and then burn a doomed resolve before falling back. + case "$_ptp_base" in + *[!0-9.]* | *..* | .* | *.) echo ""; return ;; + [0-9]*.[0-9]*) ;; + *) echo ""; return ;; + esac + [ "$(_torch_release_in_window "$_ptp_base" "$_ptp_con")" = "yes" ] || { echo ""; return; } + echo "torch==$_ptp_base" +} + +# Install torch from TORCH_INDEX_URL honoring a kept-release pin: with _PREV_TORCH_PIN +# set, TORCH_CONSTRAINT is the exact previous release; fall back to the supported range +# if the index lacks it (pruned mirror) rather than failing. Used by every --default-index +# path (NVIDIA cu*, AMD rocm/gfx fallbacks, cpu/mac, ROCm repairs) so preservation is +# uniform. Extra args (e.g. --force-reinstall) are passed through to uv. +_install_torch_default_index() { + if [ -n "$_PREV_TORCH_PIN" ]; then + # Pair the companions with the kept torch minor: torchaudio no longer + # exact-pins torch in its metadata, so leaving it unconstrained resolves + # a newer mismatched build (a kept torch 2.9.0 pulled torchaudio 2.11.0). + _itdi_base="${_PREV_TORCH_PIN#torch==}" + _itdi_minor="${_itdi_base#*.}" + _itdi_minor="${_itdi_minor%%.*}" + _itdi_tv="torchvision" + _itdi_ta="torchaudio" + case "$_itdi_base" in + 2.*) + _itdi_tv="torchvision==0.$((_itdi_minor + 15)).*" + _itdi_ta="torchaudio==2.${_itdi_minor}.*" + ;; + esac + if ! run_install_cmd_retry "install PyTorch (kept release)" uv pip install --python "$_VENV_PY" "$TORCH_CONSTRAINT" "$_itdi_tv" "$_itdi_ta" \ + --default-index "$TORCH_INDEX_URL" "$@"; then + substep "[WARN] $_PREV_TORCH_PIN is not installable from $(_strip_index_url_credentials "$TORCH_INDEX_URL") -- installing the newest supported release instead" "$C_WARN" + TORCH_CONSTRAINT="$_PREV_FALLBACK_CONSTRAINT" + _PREV_TORCH_PIN="" + run_install_cmd_retry "install PyTorch" uv pip install --python "$_VENV_PY" "$TORCH_CONSTRAINT" "$TORCHVISION_CONSTRAINT" "$TORCHAUDIO_CONSTRAINT" \ + --default-index "$TORCH_INDEX_URL" "$@" + fi + else + run_install_cmd_retry "install PyTorch" uv pip install --python "$_VENV_PY" "$TORCH_CONSTRAINT" "$TORCHVISION_CONSTRAINT" "$TORCHAUDIO_CONSTRAINT" \ + --default-index "$TORCH_INDEX_URL" "$@" + fi +} + # Expected tag from the index leaf ($1): cuXXX / cpu / rocm (rocmX.Y and gfx* -> # rocm). Empty on an unknown leaf (odd mirror) so the repair safely no-ops. _expected_torch_flavor_tag() { - _u="${1%/}" - _leaf="${_u##*/}" + _leaf=$(_torch_index_url_leaf "$1") case "$_leaf" in - cu[0-9]*) echo "$_leaf" ;; - cpu) echo "cpu" ;; - rocm*|gfx*) echo "rocm" ;; - *) echo "" ;; + cu[0-9]*) + # Exact cu + digits only; a cu*-suffixed leaf (cu128-private) -> "" (custom), + # else a correct +cu128 wheel is force-reinstalled every run. + case "${_leaf#cu}" in + *[!0-9]*) echo "" ;; + *) echo "$_leaf" ;; + esac + ;; + cpu) echo "cpu" ;; + # Exact rocm/gfx families only; a custom rocm*-suffixed leaf -> "" (custom). + *) + if _is_pip_rocm_family_leaf "$_leaf"; then echo "rocm"; else echo ""; fi + ;; esac } @@ -2206,14 +2438,42 @@ _expected_torch_flavor_tag() { # fresh-install paths above already use -- so a stale wheel is auto-repairable. # Unknown/odd-mirror leaves -> no, so we warn rather than risk a wrong reinstall. _torch_index_repairable() { - _u="${1%/}" - _leaf="${_u##*/}" + _leaf=$(_torch_index_url_leaf "$1") case "$_leaf" in - cu[0-9]*|rocm[0-9]*|gfx*) echo "yes" ;; - *) echo "no" ;; + cu[0-9]*) echo "yes" ;; + # Only EXACT rocm/gfx families resolve via --default-index; a suffixed leaf is verbatim. + *) + if _is_pip_rocm_family_leaf "$_leaf"; then echo "yes"; else echo "no"; fi + ;; esac } +# Remove credentials from a wheel index URL ($1) so an authenticated pin never leaks: +# drops userinfo AND query/fragment; scheme/host/path stay exact. Shared with py / ps1. +_strip_index_url_credentials() { + _sic_url="$1" + case "$_sic_url" in + *://*) ;; + *) printf '%s' "$_sic_url"; return ;; + esac + _sic_scheme="${_sic_url%%://*}" + _sic_rest="${_sic_url#*://}" + # Drop query / fragment (may hold auth tokens). + _sic_rest="${_sic_rest%%\?*}" + _sic_rest="${_sic_rest%%#*}" + _sic_auth="${_sic_rest%%/*}" + # Drop user:pass@ userinfo if present. + case "$_sic_auth" in + *@*) _sic_host="${_sic_auth##*@}" ;; + *) _sic_host="$_sic_auth" ;; + esac + if [ "$_sic_auth" = "$_sic_rest" ]; then + printf '%s://%s' "$_sic_scheme" "$_sic_host" + else + printf '%s://%s/%s' "$_sic_scheme" "$_sic_host" "${_sic_rest#*/}" + fi +} + get_radeon_wheel_url() { # Only meaningful on Linux. Picks a repo.radeon.com base URL whose listing # contains torch wheels. Tries paths like rocm-rel-7.2.1/, rocm-rel-7.2/, @@ -2459,7 +2719,19 @@ _maybe_bootstrap_rocm_wsl() { [ -n "$_rw_tmp" ] && rm -f "$_rw_tmp" return 0 } -_maybe_bootstrap_rocm_wsl || true +# When the caller pins the wheel index (UNSLOTH_TORCH_INDEX_URL / _FAMILY), honour it +# everywhere: skip the WSL ROCm bootstrap and the Radeon/Strix reroute below (which would +# re-probe the GPU and overwrite the pin). Trim whitespace first (parity with +# get_torch_index_url): a whitespace-only override is unset there, so must not flip this true. +_torch_index_pinned=false +_ti_url_trim="${UNSLOTH_TORCH_INDEX_URL:-}" +_ti_url_trim="${_ti_url_trim#"${_ti_url_trim%%[![:space:]]*}"}"; _ti_url_trim="${_ti_url_trim%"${_ti_url_trim##*[![:space:]]}"}" +_ti_family_trim="${UNSLOTH_TORCH_INDEX_FAMILY:-}" +_ti_family_trim="${_ti_family_trim#"${_ti_family_trim%%[![:space:]]*}"}"; _ti_family_trim="${_ti_family_trim%"${_ti_family_trim##*[![:space:]]}"}" +if [ -n "$_ti_url_trim" ] || [ -n "$_ti_family_trim" ]; then + _torch_index_pinned=true +fi +[ "$_torch_index_pinned" = true ] || _maybe_bootstrap_rocm_wsl || true TORCH_INDEX_URL=$(get_torch_index_url) @@ -2470,24 +2742,74 @@ TORCH_INDEX_URL=$(get_torch_index_url) # whose base path happens to contain "rocm" or "gfx" must not mislabel a # cu*/cpu index as ROCm (radeon repo URLs end in rocm-rel-X.Y/, Strix # overrides in gfxNNNN/, so the trailing slash is stripped first). -_torch_index_leaf="${TORCH_INDEX_URL%/}" +# Lowercase the leaf so every gfx*/rocm*/cu* arm matches regardless of case (canonical AMD +# RDNA4 leaf is gfx120X-all). CUDA is branded only on a real cu[0-9]* leaf, so a mirror +# leaf (/current) does NOT commit a CUDA backend; an unknown leaf leaves the var unset so +# the stack probes the GPU. Query/fragment dropped first, then ALL trailing slashes (in +# lockstep with the shared _torch_index_url_leaf extractor). +_torch_index_leaf="${TORCH_INDEX_URL%%\?*}" +_torch_index_leaf="${_torch_index_leaf%%#*}" +# Strip ALL trailing slashes, not one: .../cu128// must yield cu128, not an empty leaf. +while [ -n "$_torch_index_leaf" ] && [ "${_torch_index_leaf%/}" != "$_torch_index_leaf" ]; do + _torch_index_leaf="${_torch_index_leaf%/}" +done _torch_index_leaf="${_torch_index_leaf##*/}" +_torch_index_leaf=$(printf '%s' "$_torch_index_leaf" | tr '[:upper:]' '[:lower:]') case "$_torch_index_leaf" in rocm*|gfx*) export UNSLOTH_TORCH_BACKEND="rocm" ;; cpu) export UNSLOTH_TORCH_BACKEND="cpu" ;; - *) export UNSLOTH_TORCH_BACKEND="cuda" ;; + cu[0-9]*) export UNSLOTH_TORCH_BACKEND="cuda" ;; + # Unknown leaf (odd mirror, /current): unset so a stale inherited value can't leak and + # the stack probes the GPU. + *) unset UNSLOTH_TORCH_BACKEND ;; esac -# rocm7.2 ships torch 2.11.0 -- adjust the constraint to allow it. -# All other ROCm tags and CUDA stay within <2.11.0. -case "$TORCH_INDEX_URL" in - */rocm7.2) TORCH_CONSTRAINT="torch>=2.11.0,<2.12.0" ;; +# Whether TORCH_INDEX_URL names an actual pip ROCm family (rocm* / gfx*), gating the +# ROCm-only side effects below (AMD bitsandbytes, ROCm-torch repair). Digit-gated so a leaf +# merely STARTING with "rocm" isn't force-repaired from the wrong path. +if _is_pip_rocm_family_leaf "$_torch_index_leaf"; then + _torch_index_is_rocm_family=true +else + _torch_index_is_rocm_family=false +fi + +# rocm7.2 and the per-gfx indexes with the _grouped_mm <2.11 bug (gfx120X-all, gfx1151, +# gfx1150) ship torch 2.11.0 -- raise the floor (also covers a pinned override that skipped +# the Strix reroute). Pin the companions too: the per-gfx index publishes them independently +# and a bare name can resolve a 2.12 ABI-mismatched wheel. Match on the FINAL leaf so a +# custom mirror with a gfx/rocm7.2 path segment but a cu*/cpu family isn't forced. +case "$_torch_index_leaf" in + rocm7.2|gfx120x-all|gfx1151|gfx1150) + TORCH_CONSTRAINT="torch>=2.11.0,<2.12.0" + TORCHVISION_CONSTRAINT="torchvision>=0.26.0,<0.27.0" + TORCHAUDIO_CONSTRAINT="torchaudio>=2.11.0,<2.12.0" + ;; + # CUDA cu12x/cu13x indexes ship torch 2.11.x: widen the ceiling to <2.12.0 (matches + # _CUDA_TORCH_PKG_SPEC) and widen the companions with it so the trio stays paired. + cu[0-9]*) + TORCH_CONSTRAINT="torch>=2.4,<2.12.0" + TORCHVISION_CONSTRAINT="torchvision>=0.19,<0.27.0" + TORCHAUDIO_CONSTRAINT="torchaudio>=2.4,<2.12.0" + ;; esac +# A pinned custom/unknown-leaf index (/simple, /current, /cu128-private) has no curated +# companion set, so bound torchvision/torchaudio to the same <2.11 range the Python path pins +# (else a mirror with newer companions resolves a 2.12 ABI-mismatched wheel). Known families +# keep their curated companions above (_expected_torch_flavor_tag returns "" only for custom). +if [ "$_torch_index_pinned" = true ] && \ + [ -z "$(_expected_torch_flavor_tag "$TORCH_INDEX_URL")" ]; then + TORCHVISION_CONSTRAINT="torchvision>=0.19,<0.26.0" + TORCHAUDIO_CONSTRAINT="torchaudio>=2.4,<2.11.0" +fi + # Auto-detect GPU for AMD ROCm based # get_torch_index_url must have chosen */rocm* # (gfx in rocminfo or amd-smi list). Then require rocminfo "Marketing Name:.*Radeon". +# Skipped when the index is pinned: an explicit override must not be rerouted to the +# Radeon/Strix repos by GPU probing. _amd_gpu_radeon=false +if [ "$_torch_index_pinned" = false ]; then case "$TORCH_INDEX_URL" in */rocm*) if _has_amd_rocm_gpu && command -v rocminfo >/dev/null 2>&1 && \ @@ -2564,10 +2886,31 @@ case "$TORCH_INDEX_URL" in done TORCH_INDEX_URL="${_amd_strix_base}/${_strix_gfx}/" TORCH_CONSTRAINT="torch>=2.11.0,<2.12.0" + # Pin companions to 2.11 (per-gfx index publishes them independently). + TORCHVISION_CONSTRAINT="torchvision>=0.26.0,<0.27.0" + TORCHAUDIO_CONSTRAINT="torchaudio>=2.11.0,<2.12.0" _amd_gpu_radeon=false fi ;; esac +fi # _torch_index_pinned guard (Radeon + Strix reroute) +# Re-run over an existing install: keep the previous venv's torch RELEASE; the fresh +# index above supplies the right flavor for this machine. Evaluated HERE, after every +# index/constraint decision including the Strix reroute, so the window checked is the +# final one and a raised floor (rocm7.2 / Strix gfx) rejects an older release. +# _PREV_FALLBACK_CONSTRAINT keeps the range so the install can fall back when the exact +# release is not on the chosen index (mirrors may prune old wheels). Skipped for --no-torch. +_PREV_TORCH_PIN="" +_PREV_FALLBACK_CONSTRAINT="$TORCH_CONSTRAINT" +if [ "$SKIP_TORCH" = false ]; then + _prev_pin=$(_previous_torch_pin "$_PREV_TORCH_VER" "$TORCH_CONSTRAINT") + if [ -n "$_prev_pin" ]; then + _PREV_TORCH_PIN="$_prev_pin" + TORCH_CONSTRAINT="$_prev_pin" + substep "existing install has torch $_PREV_TORCH_VER -- keeping it (set UNSLOTH_TORCH_UPGRADE=1 to get the newest release)" + fi +fi + _TAURI_TORCH_INDEX_FAMILY=$(_tauri_torch_index_family "$TORCH_INDEX_URL") if [ "$_amd_gpu_radeon" = true ] && [ "$SKIP_TORCH" = false ]; then _TAURI_TORCH_INDEX_FAMILY="radeon" @@ -2697,7 +3040,7 @@ case "$TORCH_INDEX_URL" in if [ "$_amd_gpu_radeon" = true ]; then substep "wheels: repo.radeon.com (Radeon)" else - substep "wheels: $TORCH_INDEX_URL" + substep "wheels: $(_strip_index_url_credentials "$TORCH_INDEX_URL")" fi ;; esac @@ -2705,9 +3048,46 @@ esac # ── Install unsloth directly into the venv (no activation needed) ── tauri_log "STEP" "Installing PyTorch" _VENV_PY="$VENV_DIR/bin/python" + +# A released unsloth wheel can pin an older torch (unsloth 2026.7.2 declares +# torch<2.11.0); a with-deps PyPI resolve then downgrades the whole trio, +# swapping the pinned +cuXXX/+rocm build for PyPI's default. The flavor guard +# below misses this (PyPI's torch 2.10 default is itself cu128-flavored), so +# freeze the trio via uv --overrides (overrides replace dependency requirements +# during resolution) while unsloth's other deps resolve normally. Sets +# _UNSLOTH_TORCH_OVERRIDES from the trio in the venv; every with-deps unsloth +# install (migrated and fresh) must call this before resolving and rm it after. +_build_unsloth_torch_overrides() { + _UNSLOTH_TORCH_OVERRIDES="" + [ "$SKIP_TORCH" = false ] || return 0 + _torch_trio_pins=$("$_VENV_PY" -c " +from importlib.metadata import version, PackageNotFoundError +for _p in ('torch', 'torchvision', 'torchaudio'): + try: + print(_p + '==' + version(_p)) + except PackageNotFoundError: + pass +" 2>/dev/null) || _torch_trio_pins="" + case "$_torch_trio_pins" in + torch==*) + _UNSLOTH_TORCH_OVERRIDES=$(mktemp) + printf '%s\n' "$_torch_trio_pins" > "$_UNSLOTH_TORCH_OVERRIDES" + # The CLI --overrides flag replaces any UV_OVERRIDE env file (same + # uv setting; macOS arm64 exports one here), so fold its pins in. + # awk, not cat: it drops inherited torch-trio lines (uv intersects + # duplicate overrides, so a conflicting pin would make resolution + # unsatisfiable) and newline-terminates the last line so an + # unterminated file cannot join two requirements into one. + for _ov_file in ${UV_OVERRIDE:-}; do + [ -f "$_ov_file" ] && awk '!/^[[:space:]]*torch(vision|audio)?([[:space:]<>=!~;@[]|$)/' "$_ov_file" >> "$_UNSLOTH_TORCH_OVERRIDES" + done + ;; + esac +} + if [ "$_MIGRATED" = true ]; then - # Migrated env: force-reinstall unsloth+unsloth-zoo to ensure clean state - # in the new venv location, while preserving existing torch/CUDA + # Migrated env: force-reinstall unsloth+unsloth-zoo for a clean state, preserving + # existing torch/CUDA unless the ROCm repair below fires. substep "upgrading unsloth in migrated environment..." if [ "$SKIP_TORCH" = true ]; then # No-torch: install unsloth + unsloth-zoo with --no-deps (current @@ -2716,7 +3096,7 @@ if [ "$_MIGRATED" = true ]; then # to prevent transitive torch resolution. run_install_cmd_retry "install unsloth (migrated no-torch)" uv pip install --python "$_VENV_PY" --no-deps \ --reinstall-package unsloth --reinstall-package unsloth-zoo \ - "unsloth>=2026.7.3" "unsloth-zoo>=2026.7.3" + "unsloth>=2026.7.4" "unsloth-zoo>=2026.7.4" # Resolve pydantic WITH deps so pip pins pydantic-core to the # matching version (no-torch-runtime.txt below is --no-deps). # All transitive deps are torch-free. @@ -2729,9 +3109,13 @@ if [ "$_MIGRATED" = true ]; then else # Pin mlx-lm away from 0.31.3 here too: a curl-piped migration has no # overrides file, so UV_OVERRIDE is unset and this positional is the only cover. + _build_unsloth_torch_overrides run_install_cmd_retry "install unsloth (migrated)" uv pip install --python "$_VENV_PY" \ + ${_UNSLOTH_TORCH_OVERRIDES:+--overrides "$_UNSLOTH_TORCH_OVERRIDES"} \ --reinstall-package unsloth --reinstall-package unsloth-zoo \ - "unsloth>=2026.7.3" "unsloth-zoo>=2026.7.3" ${_MLX_LM_EXCLUDE_ARG:-} + "unsloth>=2026.7.4" "unsloth-zoo>=2026.7.4" ${_MLX_LM_EXCLUDE_ARG:-} + [ -n "$_UNSLOTH_TORCH_OVERRIDES" ] && rm -f "$_UNSLOTH_TORCH_OVERRIDES" + _UNSLOTH_TORCH_OVERRIDES="" fi if [ "$STUDIO_LOCAL_INSTALL" = true ]; then substep "overlaying local repo (editable)..." @@ -2744,21 +3128,14 @@ if [ "$_MIGRATED" = true ]; then # AMD ROCm: install bitsandbytes even in migrated environments so # existing ROCm installs gain the AMD bitsandbytes build without a # fresh reinstall. - if [ "$SKIP_TORCH" = false ]; then - case "$TORCH_INDEX_URL" in - */rocm*|*/gfx*) - _install_bnb_rocm "install bitsandbytes (AMD)" "$_VENV_PY" - # Repair ROCm torch if overwritten during migrated install - _has_hip=$("$_VENV_PY" -c "import torch; print(getattr(torch.version,'hip','') or '')" 2>/dev/null || true) - if [ -z "$_has_hip" ]; then - substep "repairing ROCm torch (overwritten by dependency resolution)..." - run_install_cmd_retry "repair ROCm torch" uv pip install --python "$_VENV_PY" \ - "$TORCH_CONSTRAINT" torchvision torchaudio \ - --default-index "$TORCH_INDEX_URL" \ - --force-reinstall - fi - ;; - esac + if [ "$SKIP_TORCH" = false ] && [ "$_torch_index_is_rocm_family" = true ]; then + _install_bnb_rocm "install bitsandbytes (AMD)" "$_VENV_PY" + # Repair ROCm torch if overwritten during migrated install + _has_hip=$("$_VENV_PY" -c "import torch; print(getattr(torch.version,'hip','') or '')" 2>/dev/null || true) + if [ -z "$_has_hip" ]; then + substep "repairing ROCm torch (overwritten by dependency resolution)..." + _install_torch_default_index --force-reinstall + fi fi elif [ -n "$TORCH_INDEX_URL" ]; then # Fresh: Step 1 - install torch from explicit index (skip when --no-torch or Intel Mac) @@ -2820,7 +3197,42 @@ elif [ -n "$TORCH_INDEX_URL" ]; then _ta_ver=$(_extract_version "$_ta_whl" "torchaudio") _radeon_versions_match=false - if [ -n "$_torch_ver" ] && [ -n "$_tv_ver" ] && [ -n "$_ta_ver" ]; then + # Kept release (_PREV_TORCH_PIN) wins here too: pick its exact + # patch (else the newest patch of its minor) plus the paired + # vision/audio wheels. Any gap falls back to the newest-trio + # search below, mirroring _install_torch_default_index, so a + # rerun never drifts to another release nor below the kept one. + if [ -n "$_PREV_TORCH_PIN" ]; then + _prev_kept_base="${_PREV_TORCH_PIN#torch==}" + _prev_kept_minor="${_prev_kept_base#*.}" + _prev_kept_minor="${_prev_kept_minor%%.*}" + case "$_prev_kept_minor" in + ''|*[!0-9]*) ;; + *) + _kept_torch=$(_pick_radeon_wheel "torch" "${_prev_kept_base}" 2>/dev/null) || _kept_torch="" + [ -z "$_kept_torch" ] && { _kept_torch=$(_pick_radeon_wheel "torch" "2.${_prev_kept_minor}." 2>/dev/null) || _kept_torch=""; } + _kept_tv=$(_pick_radeon_wheel "torchvision" "0.$((_prev_kept_minor + 15))." 2>/dev/null) || _kept_tv="" + _kept_ta=$(_pick_radeon_wheel "torchaudio" "2.${_prev_kept_minor}." 2>/dev/null) || _kept_ta="" + if [ -n "$_kept_torch" ] && [ -n "$_kept_tv" ] && [ -n "$_kept_ta" ]; then + _torch_whl=$_kept_torch + _tv_whl=$_kept_tv + _ta_whl=$_kept_ta + _tri_whl="" + _radeon_versions_match=true + # Say so when the listing pruned the exact patch + # and a same-series build is installed instead. + case "$(printf '%s' "${_kept_torch##*/}" | sed 's/%2[Bb]/+/g')" in + "torch-${_prev_kept_base}"[+-]*) ;; + *) substep "kept release ${_prev_kept_base} is not in the Radeon listing -- installing the closest 2.${_prev_kept_minor} series build instead" ;; + esac + else + substep "[WARN] Radeon repo lacks a complete wheel set for kept $_PREV_TORCH_PIN -- installing the newest compatible set instead" "$C_WARN" + fi + ;; + esac + fi + if [ "$_radeon_versions_match" != true ] && \ + [ -n "$_torch_ver" ] && [ -n "$_tv_ver" ] && [ -n "$_ta_ver" ]; then _torch_minor=${_torch_ver#*.} _ta_minor=${_ta_ver#*.} _tv_minor=${_tv_ver#*.} @@ -2877,10 +3289,8 @@ elif [ -n "$TORCH_INDEX_URL" ]; then if [ -z "$_torch_whl" ] || [ -z "$_tv_whl" ] || [ -z "$_ta_whl" ] || \ [ "$_radeon_versions_match" != true ]; then - substep "[WARN] Radeon repo lacks a compatible wheel set for this Python; falling back to ROCm index ($TORCH_INDEX_URL)" "$C_WARN" - run_install_cmd_retry "install PyTorch" uv pip install --python "$_VENV_PY" \ - "$TORCH_CONSTRAINT" torchvision torchaudio \ - --default-index "$TORCH_INDEX_URL" + substep "[WARN] Radeon repo lacks a compatible wheel set for this Python; falling back to ROCm index ($(_strip_index_url_credentials "$TORCH_INDEX_URL"))" "$C_WARN" + _install_torch_default_index else substep "installing PyTorch from Radeon repo (${_RADEON_BASE_URL})..." # Pass explicit wheel URLs so the matched trio is @@ -2900,42 +3310,34 @@ elif [ -n "$TORCH_INDEX_URL" ]; then fi fi else - substep "[WARN] Radeon repo unavailable; falling back to ROCm index ($TORCH_INDEX_URL)" "$C_WARN" - run_install_cmd_retry "install PyTorch" uv pip install --python "$_VENV_PY" \ - "$TORCH_CONSTRAINT" torchvision torchaudio \ - --default-index "$TORCH_INDEX_URL" + substep "[WARN] Radeon repo unavailable; falling back to ROCm index ($(_strip_index_url_credentials "$TORCH_INDEX_URL"))" "$C_WARN" + _install_torch_default_index fi else substep "[WARN] Radeon GPU detected but could not detect full ROCm version; falling back to ROCm index" "$C_WARN" - run_install_cmd_retry "install PyTorch" uv pip install --python "$_VENV_PY" \ - "$TORCH_CONSTRAINT" torchvision torchaudio \ - --default-index "$TORCH_INDEX_URL" + _install_torch_default_index fi else - substep "installing PyTorch ($TORCH_INDEX_URL)..." - run_install_cmd_retry "install PyTorch" uv pip install --python "$_VENV_PY" "$TORCH_CONSTRAINT" torchvision torchaudio \ - --default-index "$TORCH_INDEX_URL" + substep "installing PyTorch ($(_strip_index_url_credentials "$TORCH_INDEX_URL"))..." + _install_torch_default_index fi # AMD ROCm: install bitsandbytes (once, after torch, for all ROCm paths). # Gate on SKIP_TORCH=false so a user running with --no-torch on a ROCm # host stays in GGUF-only mode rather than pulling in bitsandbytes, # which is only useful once torch is present for training. - if [ "$SKIP_TORCH" = false ]; then - case "$TORCH_INDEX_URL" in - */rocm*|*/gfx*) - _install_bnb_rocm "install bitsandbytes (AMD)" "$_VENV_PY" - ;; - esac + if [ "$SKIP_TORCH" = false ] && [ "$_torch_index_is_rocm_family" = true ]; then + _install_bnb_rocm "install bitsandbytes (AMD)" "$_VENV_PY" fi - # Fresh: Step 2 - install unsloth, preserving pre-installed torch + # Fresh: Step 2 - install unsloth, preserving the torch Step 1 installed tauri_log "STEP" "Installing Unsloth" substep "installing unsloth (this may take a few minutes)..." + _build_unsloth_torch_overrides if [ "$SKIP_TORCH" = true ]; then # No-torch: install unsloth + unsloth-zoo with --no-deps, then # runtime deps (typer, safetensors, transformers, etc.) with --no-deps. run_install_cmd_retry "install unsloth (no-torch)" uv pip install --python "$_VENV_PY" --no-deps \ --upgrade-package unsloth --upgrade-package unsloth-zoo \ - "unsloth>=2026.7.3" "unsloth-zoo>=2026.7.3" + "unsloth>=2026.7.4" "unsloth-zoo>=2026.7.4" # Same pydantic-with-deps trick as the migrated branch. run_install_cmd_retry "install pydantic (with deps for compatible core)" \ uv pip install --python "$_VENV_PY" pydantic @@ -2953,7 +3355,8 @@ elif [ -n "$TORCH_INDEX_URL" ]; then fi elif [ "$STUDIO_LOCAL_INSTALL" = true ]; then run_install_cmd_retry "install unsloth (local)" uv pip install --python "$_VENV_PY" \ - --upgrade-package unsloth "unsloth>=2026.7.3" "unsloth-zoo>=2026.7.3" + ${_UNSLOTH_TORCH_OVERRIDES:+--overrides "$_UNSLOTH_TORCH_OVERRIDES"} \ + --upgrade-package unsloth "unsloth>=2026.7.4" "unsloth-zoo>=2026.7.4" substep "overlaying local repo (editable)..." run_install_cmd "overlay local repo" uv pip install --python "$_VENV_PY" -e "$_REPO_ROOT" --no-deps substep "overlaying unsloth-zoo from git main..." @@ -2962,30 +3365,26 @@ elif [ -n "$TORCH_INDEX_URL" ]; then "unsloth-zoo @ git+https://github.com/unslothai/unsloth-zoo" else run_install_cmd_retry "install unsloth" uv pip install --python "$_VENV_PY" \ + ${_UNSLOTH_TORCH_OVERRIDES:+--overrides "$_UNSLOTH_TORCH_OVERRIDES"} \ --upgrade-package unsloth -- "$PACKAGE_NAME" ${_MLX_LM_EXCLUDE_ARG:-} fi + [ -n "$_UNSLOTH_TORCH_OVERRIDES" ] && rm -f "$_UNSLOTH_TORCH_OVERRIDES" + _UNSLOTH_TORCH_OVERRIDES="" # AMD ROCm: repair torch if the unsloth/unsloth-zoo install pulled in # CUDA torch from PyPI, overwriting the ROCm wheels installed in Step 1. - if [ "$SKIP_TORCH" = false ]; then - case "$TORCH_INDEX_URL" in - */rocm*|*/gfx*) - _has_hip=$("$_VENV_PY" -c "import torch; print(getattr(torch.version,'hip','') or '')" 2>/dev/null || true) - if [ -z "$_has_hip" ]; then - substep "repairing ROCm torch (overwritten by dependency resolution)..." - run_install_cmd_retry "repair ROCm torch" uv pip install --python "$_VENV_PY" \ - "$TORCH_CONSTRAINT" torchvision torchaudio \ - --default-index "$TORCH_INDEX_URL" \ - --force-reinstall - fi - ;; - esac + if [ "$SKIP_TORCH" = false ] && [ "$_torch_index_is_rocm_family" = true ]; then + _has_hip=$("$_VENV_PY" -c "import torch; print(getattr(torch.version,'hip','') or '')" 2>/dev/null || true) + if [ -z "$_has_hip" ]; then + substep "repairing ROCm torch (overwritten by dependency resolution)..." + _install_torch_default_index --force-reinstall + fi fi else # Fallback: GPU detection failed to produce a URL -- let uv resolve torch tauri_log "STEP" "Installing Unsloth" substep "installing unsloth (this may take a few minutes)..." if [ "$STUDIO_LOCAL_INSTALL" = true ]; then - run_install_cmd_retry "install unsloth (auto torch backend)" uv pip install --python "$_VENV_PY" "unsloth-zoo>=2026.7.3" "unsloth>=2026.7.3" --torch-backend=auto + run_install_cmd_retry "install unsloth (auto torch backend)" uv pip install --python "$_VENV_PY" "unsloth-zoo>=2026.7.4" "unsloth>=2026.7.4" --torch-backend=auto substep "overlaying local repo (editable)..." run_install_cmd "overlay local repo" uv pip install --python "$_VENV_PY" -e "$_REPO_ROOT" --no-deps substep "overlaying unsloth-zoo from git main..." @@ -2997,6 +3396,15 @@ else fi fi +_installed_package_version=$("$_VENV_PY" -c \ + 'from importlib.metadata import version; import sys; print(version(sys.argv[1]))' \ + "$PACKAGE_NAME" 2>/dev/null || true) +if [ -n "$_installed_package_version" ]; then + step "$PACKAGE_NAME" "$_installed_package_version installed" +else + substep "[WARN] installed $PACKAGE_NAME version could not be determined" "$C_WARN" +fi + # ── Enforce the installed torch flavor matches the detected GPU build ── # PEP 440 ignores the +cpu/+cuXXX/+rocm local label in a version range, so uv # keeps a stale torch==X+cpu against a GPU index and the venv silently trains on @@ -3014,9 +3422,7 @@ if [ "$SKIP_TORCH" = false ] && [ -n "${TORCH_INDEX_URL:-}" ]; then if [ -n "$_installed_torch_tag" ] && [ "$_installed_torch_tag" != "$_expected_torch_tag" ] \ && [ "$(_torch_index_repairable "$TORCH_INDEX_URL")" = "yes" ]; then substep "PyTorch flavor mismatch (installed $_installed_torch_tag, need $_expected_torch_tag) -- reinstalling correct build..." - run_install_cmd "reinstall PyTorch ($_expected_torch_tag)" uv pip install --python "$_VENV_PY" \ - "$TORCH_CONSTRAINT" torchvision torchaudio \ - --default-index "$TORCH_INDEX_URL" \ + _install_torch_default_index \ --reinstall-package torch --reinstall-package torchvision --reinstall-package torchaudio _installed_torch_ver=$("$_VENV_PY" -c "import torch; print(torch.__version__)" 2>/dev/null || true) _installed_torch_tag="" @@ -3027,7 +3433,7 @@ if [ "$SKIP_TORCH" = false ] && [ -n "${TORCH_INDEX_URL:-}" ]; then substep "[WARN] PyTorch is CPU-only but a $_expected_torch_tag GPU build was expected for this machine." "$C_WARN" substep "[WARN] Training and GPU inference will run on CPU until this is fixed." "$C_WARN" substep "[WARN] Re-run this installer, or reinstall the GPU build manually:" "$C_WARN" - substep "[WARN] uv pip install --python \"$_VENV_PY\" \"$TORCH_CONSTRAINT\" torchvision torchaudio --default-index $TORCH_INDEX_URL --reinstall-package torch --reinstall-package torchvision --reinstall-package torchaudio" "$C_WARN" + substep "[WARN] uv pip install --python \"$_VENV_PY\" \"$TORCH_CONSTRAINT\" \"$TORCHVISION_CONSTRAINT\" \"$TORCHAUDIO_CONSTRAINT\" --default-index $(_strip_index_url_credentials "$TORCH_INDEX_URL") --reinstall-package torch --reinstall-package torchvision --reinstall-package torchaudio" "$C_WARN" fi fi fi diff --git a/pyproject.toml b/pyproject.toml index 7b8fdd100d..071258eb8f 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -74,7 +74,7 @@ triton = [ ] huggingfacenotorch = [ - "unsloth_zoo>=2026.7.3", + "unsloth_zoo>=2026.7.4", "wheel>=0.42.0", "packaging", "numpy", @@ -95,7 +95,7 @@ huggingfacenotorch = [ ] huggingface = [ "unsloth[huggingfacenotorch]", - "unsloth_zoo>=2026.7.3", + "unsloth_zoo>=2026.7.4", "torchvision", "unsloth[triton]", ] @@ -580,7 +580,7 @@ colab-ampere-torch220 = [ "flash-attn>=2.6.3 ; ('linux' in sys_platform)", ] colab-new = [ - "unsloth_zoo>=2026.7.3", + "unsloth_zoo>=2026.7.4", "packaging", "tyro", "transformers>=4.51.3,!=4.52.0,!=4.52.1,!=4.52.2,!=4.52.3,!=4.53.0,!=4.54.0,!=4.55.0,!=4.55.1,!=4.57.0,!=4.57.4,!=4.57.5,!=5.0.0,!=5.1.0,<=5.5.0", diff --git a/studio/backend/core/inference/llama_cpp.py b/studio/backend/core/inference/llama_cpp.py index c9ab7eb83b..2c7433f7a4 100644 --- a/studio/backend/core/inference/llama_cpp.py +++ b/studio/backend/core/inference/llama_cpp.py @@ -11,6 +11,7 @@ import atexit import contextlib import functools import json +import math import os import re import struct @@ -22,12 +23,24 @@ import subprocess import sys import threading import time +import uuid from pathlib import Path -from typing import Callable, Collection, Generator, Iterable, List, Mapping, Optional, Union +from typing import ( + Callable, + Collection, + Generator, + Iterable, + List, + Literal, + Mapping, + Optional, + Union, +) import httpx from core.inference.llama_server_args import ( + _LAYER_OFFLOAD_FLAGS, _effective_tensor_parallel, _tensor_parallel_matches_loaded, extra_args_disable_mmproj, @@ -234,8 +247,7 @@ def _wsl_system_rocm_lib_dirs() -> "list[str]": return out -# Plan-without-action re-prompt state (intent signal, caps, message) now lives -# in tool_call_parser, imported above under its old aliases. +# Plan-without-action re-prompt state now lives in tool_call_parser (imported above). # Default max_tokens to the effective context when known. The floor is high # enough for reasoning-heavy GGUFs and max_tokens-omitting API clients. @@ -442,6 +454,23 @@ def _hf_offline_if_dns_dead(): os.environ.pop("TRANSFORMERS_OFFLINE", None) +try: + _SLOT_SAVE_MAX_BYTES = int(os.environ.get("UNSLOTH_SLOT_SAVE_MAX_BYTES") or (10 << 30)) +except ValueError: + _SLOT_SAVE_MAX_BYTES = 10 << 30 + +# The idle loop holds the lifecycle gate across a slot save, so a newly arriving +# request waits on the in-flight save's HTTP call. Bound it (was 120s) so a slow +# or stuck save can't stall the next request for minutes; best-effort save just +# falls back to a plain unload. Override with UNSLOTH_SLOT_SAVE_TIMEOUT (seconds). +try: + _SLOT_SAVE_HTTP_TIMEOUT = float(os.environ.get("UNSLOTH_SLOT_SAVE_TIMEOUT") or 30.0) +except ValueError: + _SLOT_SAVE_HTTP_TIMEOUT = 30.0 +if _SLOT_SAVE_HTTP_TIMEOUT <= 0: + _SLOT_SAVE_HTTP_TIMEOUT = 30.0 + + def _swa_cache_path() -> Path: home = os.environ.get("UNSLOTH_STUDIO_HOME") or os.environ.get("STUDIO_HOME") base = Path(home) if home else Path.home() / ".unsloth" / "studio" @@ -1431,7 +1460,10 @@ def _extra_args_set_spec_type(extra_args: Optional[Iterable[str]]) -> bool: return _extra_args_set_any_flag(extra_args, {"--spec-type", "--spec-default"}) -_GPU_OFFLOAD_OVERRIDE_FLAGS = frozenset({"-ngl", "--gpu-layers", "--n-gpu-layers", "-fit", "--fit"}) +# Layer-offload override detection. Single-sourced from llama_server_args, which +# also strips these (plus the MoE flags) from inherited extras; sharing the layer +# set keeps detection and stripping from drifting. +_GPU_OFFLOAD_OVERRIDE_FLAGS = _LAYER_OFFLOAD_FLAGS _THREAD_OVERRIDE_FLAGS = frozenset({"-t", "--threads"}) @@ -1895,6 +1927,17 @@ class LlamaCppBackend: self._cache_type_kv: Optional[str] = None # Whether --split-mode tensor was applied on the active load. self._tensor_parallel: bool = False + # GPU memory strategy applied on the active load ("auto"/"manual"). + self._gpu_memory_mode: str = "auto" + # Manual-mode load options (echoed back so the UI round-trips them). + self._gpu_layers: int = -1 + # MoE expert layers to keep on CPU (--n-cpu-moe); 0 = none. + self._n_cpu_moe: int = 0 + # Relative model share per GPU (--tensor-split), in GPU order; None = + # default (llama.cpp splits by free VRAM). + self._tensor_split: Optional[List[float]] = None + # User-picked physical GPU indices (None = automatic selection). + self._gpu_ids: Optional[List[int]] = None # Layer load kept multi-GPU only to honor a downgraded tensor request, so a # later explicit tensor-off reloads instead of deduping to it (#6659). self._layer_preserves_tensor_intent: bool = False @@ -1909,6 +1952,11 @@ class LlamaCppBackend: self._spec_draft_n_max: Optional[int] = None # KV-cache estimation fields (populated by _read_gguf_metadata) self._n_layers: Optional[int] = None + # MoE metadata (populated by _read_gguf_metadata): expert count (>0 = + # MoE) and leading dense-layer count (offsets --n-cpu-moe, which counts + # from layer 0). See the n_moe_layers property. + self._n_experts: Optional[int] = None + self._leading_dense_block_count: Optional[int] = None self._n_kv_heads: Optional[int] = None self._n_kv_heads_by_layer: Optional[list[int]] = None self._n_heads: Optional[int] = None @@ -1970,6 +2018,12 @@ class LlamaCppBackend: self._llama_log_path: Optional[Path] = None self._cancel_event = threading.Event() self._api_key: Optional[str] = None + self._slot_save_dir: Optional[str] = None + self._slot_save_binary: Optional[tuple[str, int]] = None + # (gguf_identity, launch_fingerprint) snapshotted at load, so a later slot + # save can tell whether the model files were swapped on disk since load. + self._slot_loaded_identity: Optional[tuple] = None + self._prompt_cache_disabled: bool = False # True once a probe has completed; cleared on transient failure. self._is_audio: bool = False self._audio_type: Optional[str] = None @@ -2329,6 +2383,79 @@ class LlamaCppBackend: """Whether --split-mode tensor is active on the loaded server.""" return self._tensor_parallel + @property + def gpu_memory_mode(self) -> str: + """Active GPU memory strategy: 'auto' or 'manual' (gpu_layers < 0 = Auto/--fit, >= 0 = pinned).""" + return self._gpu_memory_mode + + @property + def gpu_layers(self) -> int: + """Requested --gpu-layers for manual mode (-1 when not manual).""" + return self._gpu_layers + + @property + def n_cpu_moe(self) -> int: + """MoE expert layers manual mode kept on CPU (--n-cpu-moe); 0 = none.""" + return self._n_cpu_moe + + @property + def tensor_split(self) -> Optional[List[float]]: + """Manual-mode relative model share per GPU (--tensor-split); None = + default (split by free VRAM).""" + return self._tensor_split + + @property + def gpu_ids(self) -> Optional[List[int]]: + """User-picked physical GPU indices, or None for automatic selection.""" + return self._gpu_ids + + @property + def n_layers(self) -> Optional[int]: + """Model layer count (GGUF block_count), or None if unknown.""" + return self._n_layers + + @property + def n_moe_layers(self) -> int: + """Number of MoE expert layers (the --n-cpu-moe ceiling), 0 if not MoE. + + block_count minus the leading dense layers (which carry no experts): + --n-cpu-moe counts from layer 0, so those dense layers are no-ops. + """ + if not self._n_experts or not self._n_layers: + return 0 + return max(0, self._n_layers - (self._leading_dense_block_count or 0)) + + @staticmethod + def _resolve_cpu_moe_flag( + n_cpu_moe: int, n_moe_layers: int, leading_dense: int + ) -> Optional[int]: + """The --n-cpu-moe value (absolute first-N layers), or None to omit it. + + Clamps the requested count to the model's MoE layers, then offsets past + the leading dense layers (--n-cpu-moe counts from layer 0). Returns None + for nothing-to-offload (0 requested) or a non-MoE model. + """ + if n_cpu_moe <= 0 or n_moe_layers <= 0: + return None + return leading_dense + min(n_cpu_moe, n_moe_layers) + + @staticmethod + def _sanitize_tensor_split(tensor_split: Optional[List[float]]) -> List[float]: + """Per-GPU shares with negative and non-finite entries clamped to 0. + + A direct caller's negative entry would launch a placement different + from the ratio the UI showed, and inf would pass a plain ``> 0`` total + gate and emit ``--tensor-split inf,...``. Returns [] for input that + can't be read as floats (the length gate at the call site then drops + the split). + """ + try: + return [ + x if math.isfinite(x) and x > 0.0 else 0.0 for x in (float(v) for v in tensor_split) + ] + except (TypeError, ValueError, OverflowError): + return [] + @property def layer_preserves_tensor_intent(self) -> bool: """True when a downgraded tensor request kept this layer load multi-GPU.""" @@ -2530,10 +2657,12 @@ class LlamaCppBackend: "spec_draft_n_max_flag": None, "supports_kv_unified": False, "supports_fit_ctx": False, + "supports_fit_target": False, "supports_cache_ram": False, "supports_ctx_checkpoints": False, "supports_no_cache_prompt": False, "supports_metrics": False, + "supports_slot_save": False, } try: mtime = int(Path(bin_path).stat().st_mtime) @@ -2549,10 +2678,12 @@ class LlamaCppBackend: spec_draft_n_max_flag: Optional[str] = None supports_kv_unified = False supports_fit_ctx = False + supports_fit_target = False supports_cache_ram = False supports_ctx_checkpoints = False supports_no_cache_prompt = False supports_metrics = False + supports_slot_save = False try: probe_env = cls._llama_server_env_for_binary(bin_path) result = subprocess.run( @@ -2646,10 +2777,12 @@ class LlamaCppBackend: supports_kv_unified = _is_real("--kv-unified") supports_fit_ctx = _is_real("--fit-ctx") + supports_fit_target = _is_real("--fit-target") supports_cache_ram = _is_real("--cache-ram") supports_ctx_checkpoints = _is_real("--ctx-checkpoints") supports_no_cache_prompt = _is_real("--no-cache-prompt") supports_metrics = _is_real("--metrics") + supports_slot_save = _is_real("--slot-save-path") except (OSError, subprocess.SubprocessError) as exc: logger.debug(f"llama-server --help probe failed: {exc}") @@ -2662,10 +2795,12 @@ class LlamaCppBackend: "spec_draft_n_max_flag": spec_draft_n_max_flag, "supports_kv_unified": supports_kv_unified, "supports_fit_ctx": supports_fit_ctx, + "supports_fit_target": supports_fit_target, "supports_cache_ram": supports_cache_ram, "supports_ctx_checkpoints": supports_ctx_checkpoints, "supports_no_cache_prompt": supports_no_cache_prompt, "supports_metrics": supports_metrics, + "supports_slot_save": supports_slot_save, } cls._capability_cache[cache_key] = info return info @@ -2746,6 +2881,57 @@ class LlamaCppBackend: except ValueError: return None + @staticmethod + def _emit_child_gpu_visibility(env: dict, pinned: str) -> None: + """Write the child's GPU visibility mask (CUDA, plus the HIP mirror on + ROCm, where narrowing only CUDA_VISIBLE_DEVICES leaves an AMD child + seeing the full set). Do NOT also set ROCR_VISIBLE_DEVICES: ROCR and HIP + mask at different layers, so the same indices apply twice -- ROCR reduces + and re-indexes from 0, then a non-zero HIP pin points out of range, HIP + enumerates 0 devices, and llama.cpp falls back to CPU. The HIP mask alone + narrows correctly; clear any inherited ROCR mask so it can't double up.""" + env["CUDA_VISIBLE_DEVICES"] = pinned + try: + import torch as _torch + if getattr(_torch.version, "hip", None) is not None: + env["HIP_VISIBLE_DEVICES"] = pinned + env.pop("ROCR_VISIBLE_DEVICES", None) + except Exception as e: + logger.debug("Failed to set ROCm visibility env vars for child: %s", e) + + @staticmethod + def _pin_visible_gpu_order_for_split(env: dict) -> None: + """Pin the child's GPU enumeration to the picker's order for a manual + ``--tensor-split`` across the whole visible set. CUDA's default + FASTEST_FIRST enumeration applies the shares to the wrong cards on + heterogeneous hosts (#5025), and CUDA_DEVICE_ORDER only fixes the + numbering base: an inherited numeric visibility mask ALSO defines + enumeration order, so a reordered parent mask (CUDA_VISIBLE_DEVICES=3,1) + would still hand the shares to the wrong cards. The UI built the split + positionally over get_backend_visible_gpu_info's device list (ascending + physical via nvidia-smi, inherited mask order on the torch fallback), so + re-emit the same set in that report order -- not an assumed ascending + sort. The visible set itself never changes. No mask, an empty mask, or a + UUID/MIG mask (which resolves to None) is left alone -- the multi-GPU + controls are hidden for the latter.""" + env["CUDA_DEVICE_ORDER"] = "PCI_BUS_ID" + inherited = LlamaCppBackend._resolve_visible_physical_ids() + if not inherited: + return + order = None + try: + from utils.hardware import get_backend_visible_gpu_info + info = get_backend_visible_gpu_info() + if info.get("available") and info.get("index_kind") == "physical": + reported = [d["index"] for d in info.get("devices", [])] + if sorted(reported) == sorted(inherited): + order = reported + except Exception as e: + logger.debug("Could not read reported GPU order for split pin: %s", e) + if order is None: + order = sorted(inherited) + LlamaCppBackend._emit_child_gpu_visibility(env, ",".join(str(i) for i in order)) + @staticmethod def _amd_apu_wants_unified_memory(gpu_indices = None) -> bool: """True only for AMD unified-memory APUs (gfx1150/gfx1151), where @@ -3262,6 +3448,20 @@ class LlamaCppBackend: # aborts a --split-mode tensor load, so it's dropped for the tensor attempt. _TENSOR_PARALLEL_KV_TYPES = frozenset({"f16", "bf16", "f32"}) + # Main-model placement settings that Manual mode owns. They must not leak + # from Studio's parent environment into llama-server and silently override + # the command assembled from the current request. Draft-model placement is + # intentionally separate and remains available to speculative decoding. + _MANUAL_PLACEMENT_ENV_VARS = ( + "LLAMA_ARG_CPU_MOE", + "LLAMA_ARG_N_CPU_MOE", + "LLAMA_ARG_N_GPU_LAYERS", + "LLAMA_ARG_TENSOR_SPLIT", + "LLAMA_ARG_FIT", + "LLAMA_ARG_FIT_TARGET", + "LLAMA_ARG_FIT_CTX", + ) + # (binary, mtime, model) that aborted on --split-mode tensor this process (#6415 # geometry limit, e.g. MQA n_head_kv=1). Model-keyed so one model's abort doesn't # skip tensor for others; tensor is tried by default, recorded only on a real abort. @@ -3426,6 +3626,12 @@ class LlamaCppBackend: return env + @classmethod + def _clear_manual_placement_env(cls, env: dict[str, str]) -> None: + """Remove inherited main-model placement owned by Manual mode.""" + for name in cls._MANUAL_PLACEMENT_ENV_VARS: + env.pop(name, None) + @staticmethod def _select_gpus( model_size_bytes: int, @@ -4246,6 +4452,8 @@ class LlamaCppBackend: self._supports_preserve_thinking = False self._supports_tools = False self._n_layers = None + self._n_experts = None + self._leading_dense_block_count = None self._n_kv_heads = None self._n_kv_heads_by_layer = None self._n_heads = None @@ -4335,6 +4543,8 @@ class LlamaCppBackend: arch_keys = { f"{arch}.context_length": "context_length", f"{arch}.block_count": "n_layers", + f"{arch}.expert_count": "n_experts", + f"{arch}.leading_dense_block_count": "leading_dense_block_count", f"{arch}.attention.head_count_kv": "n_kv_heads", f"{arch}.attention.head_count": "n_heads", f"{arch}.embedding_length": "embedding_length", @@ -4523,6 +4733,28 @@ class LlamaCppBackend: return None + @staticmethod + def _diffusion_gpu_arg(gpu_ids: Optional[List[int]], *, cpu_only: bool = False) -> str: + """Device token passed to the diffusion visual-server child. + + The visual engine replaces its child's CUDA visibility mask with this + token, so an unpinned load must carry forward the first token from the + parent's mask rather than turning a parent-relative ordinal into a new + physical selection. + """ + if gpu_ids: + return str(sorted(gpu_ids)[0]) + if cpu_only: + return "" + if "DG_GPU" in os.environ: + return os.environ["DG_GPU"] + parent_mask = os.environ.get("CUDA_VISIBLE_DEVICES") + if parent_mask: + first = next((token.strip() for token in parent_mask.split(",") if token.strip()), "") + if first and first != "-1": + return first + return "0" + def _start_diffusion_server( self, *, @@ -4533,6 +4765,7 @@ class LlamaCppBackend: model_identifier: str, n_ctx: int, extra_args: Optional[List[str]], + gpu_ids: Optional[List[int]] = None, ) -> bool: """Launch the OpenAI-compat diffusion shim (which drives the on-device visual decoder) and wait for health. Presents the same /v1 + /health @@ -4558,7 +4791,11 @@ class LlamaCppBackend: # CUDA_VISIBLE_DEVICES="" to force CPU serving. Keep the visual-server child # CPU-masked (empty --gpu) so the shim does not re-expose GPU 0 via its default. cpu_only = self._effective_gpu_count() == 0 - gpu = "" if cpu_only else os.environ.get("DG_GPU", "0") + # Honor the GPU picker first: the diffusion runner takes a single device, + # so use the lowest selected GPU (matches the sorted set recorded below, so + # the device used == the echoed gpu_ids[0]). With no pick, fall back to the + # CPU-only mask, else DG_GPU / 0. + gpu = self._diffusion_gpu_arg(gpu_ids, cpu_only = cpu_only) cmd = list(shim_cmd) + [ "--gguf", @@ -4586,6 +4823,11 @@ class LlamaCppBackend: env.setdefault("UNSLOTH_ALLOW_CPU", "1") env["DG_VISUAL_BIN"] = visual_bin env["DG_GPU"] = gpu + if gpu_ids: + # The visual server remasks via CUDA_VISIBLE_DEVICES=; pin PCI + # order (as the llama-server path does) so the picked physical id maps + # to the GPU the picker showed, not CUDA's default fastest-first order. + env["CUDA_DEVICE_ORDER"] = "PCI_BUS_ID" # The file-override shim imports its sibling visual_engine; put its dir on PYTHONPATH. # (The zoo-package shim is an installed module and needs no PYTHONPATH change.) if extra_pythonpath: @@ -4631,6 +4873,23 @@ class LlamaCppBackend: self._model_identifier = model_identifier self._cache_type_kv = None self._gpu_offload_active = True + # Diffusion doesn't use the llama.cpp GPU-memory knobs; reset them to + # defaults (the picked device is still recorded below) so /load, /status + # and reload dedup don't report a previous GGUF's manual settings. + self._gpu_memory_mode = "auto" + self._gpu_layers = -1 + self._n_cpu_moe = 0 + self._tensor_split = None + # Diffusion is never tensor-parallel; clear any state left by a prior TP + # chat load (load_model phase 1 only kills the process, it doesn't run + # the unload reset) so /status doesn't misreport TP and an identical + # re-Apply doesn't reload against stale tensor-parallel state. + self._tensor_parallel = False + # Record only the single device the runner actually uses (the lowest + # selected GPU, chosen above) -- not the whole pick. The diffusion runner + # is single-device, so echoing a multi-GPU list would misreport placement + # in /status and let a re-Apply dedup against GPUs the runner never used. + self._gpu_ids = [sorted(gpu_ids)[0]] if gpu_ids else None if hf_variant: self._hf_variant = hf_variant elif gguf_path: @@ -5721,6 +5980,11 @@ class LlamaCppBackend: speculative_type: Optional[str] = None, spec_draft_n_max: Optional[int] = None, tensor_parallel: bool = False, + gpu_memory_mode: Literal["auto", "manual"] = "auto", + gpu_layers: int = -1, + n_cpu_moe: int = 0, + tensor_split: Optional[List[float]] = None, + gpu_ids: Optional[List[int]] = None, n_threads: Optional[int] = None, n_gpu_layers: Optional[int] = None, # caller compat, unused n_parallel: int = 1, @@ -5753,6 +6017,14 @@ class LlamaCppBackend: "speculative_type": speculative_type, "spec_draft_n_max": spec_draft_n_max, "tensor_parallel": tensor_parallel, + # GPU-memory placement: replayed on respawn so a server SIGKILL'd by + # GPU/RAM pressure reloads onto the same devices with the same + # offload, not the auto defaults. + "gpu_memory_mode": gpu_memory_mode, + "gpu_layers": gpu_layers, + "n_cpu_moe": n_cpu_moe, + "tensor_split": list(tensor_split) if tensor_split is not None else None, + "gpu_ids": list(gpu_ids) if gpu_ids is not None else None, "n_threads": n_threads, "n_gpu_layers": n_gpu_layers, "n_parallel": n_parallel, @@ -5779,6 +6051,11 @@ class LlamaCppBackend: speculative_type = speculative_type, spec_draft_n_max = spec_draft_n_max, tensor_parallel = tensor_parallel, + gpu_memory_mode = gpu_memory_mode, + gpu_layers = gpu_layers, + n_cpu_moe = n_cpu_moe, + tensor_split = tensor_split, + gpu_ids = gpu_ids, chat_template_override = chat_template_override, extra_args = extra_args, is_vision = is_vision, @@ -5899,6 +6176,7 @@ class LlamaCppBackend: model_identifier = model_identifier, n_ctx = n_ctx, extra_args = extra_args, + gpu_ids = gpu_ids, ) if not binary: @@ -5960,6 +6238,59 @@ class LlamaCppBackend: # use the same helper so a healthy env-driven tensor server matches. split_mode_override = parse_split_mode_override(extra_args) tensor_parallel = _effective_tensor_parallel(extra_args, tensor_parallel) + # gpu_layers=0 leaves nothing to split, yet --split-mode tensor or + # a per-GPU ratio still launches tensor mode -- and under the + # CPU-only mask below (no visible devices) that aborts the server + # instead of loading on CPU. Drop both here (nothing to split). + if gpu_memory_mode == "manual" and gpu_layers == 0: + if tensor_parallel or tensor_split: + logger.info( + "Manual gpu_layers=0: dropping tensor split/parallel " + "flags (nothing to split on the GPU)" + ) + tensor_parallel = False + tensor_split = None + # Record the requested strategy for /status and the load + # response. 'manual' has no fallback, so the request value is the + # value actually applied. + self._gpu_memory_mode = gpu_memory_mode + # The layer/MoE/split knobs apply only with an explicit offload + # (manual + gpu_layers >= 0); else record defaults so /status and + # /load don't report knobs the server never applied. + if gpu_memory_mode == "manual" and gpu_layers >= 0: + self._gpu_layers = gpu_layers + self._n_cpu_moe = n_cpu_moe + self._tensor_split = tensor_split + else: + self._gpu_layers = -1 + self._n_cpu_moe = 0 + self._tensor_split = None + self._gpu_ids = sorted(gpu_ids) if gpu_ids else None + # Manual offload skips the TP planner but still emits --split-mode + # tensor at launch; drop it when fewer than 2 GPUs are in use -- + # tensor split is a no-op there and aborts on some architectures. + # Done before the cache-drop below so a quantized KV survives. + if ( + tensor_parallel + and gpu_memory_mode == "manual" + and gpu_layers >= 0 + and self._effective_gpu_count(sorted(gpu_ids) if gpu_ids else None) < 2 + ): + logger.info( + "Tensor parallelism requested in manual mode but fewer " + "than 2 GPUs are in use; ignoring (needs >= 2)." + ) + tensor_parallel = False + # Drop TP for manual + Auto layers before the cache-drop below (like + # the <2-GPU guard above), so a requested quantized KV survives into + # the --fit load rather than being stripped for a tensor attempt. + if tensor_parallel and gpu_memory_mode == "manual" and gpu_layers < 0: + logger.info( + "Manual mode with Auto layers hands memory management to " + "llama.cpp --fit, which is incompatible with tensor " + "parallelism; ignoring the tensor split." + ) + tensor_parallel = False # Tensor mode aborts on a quantized KV cache, so drop it for the # tensor attempt (and strip any inherited/explicit --cache-type # that would re-impose it when appended last). Layer split does @@ -6040,6 +6371,12 @@ class LlamaCppBackend: "Vision-capable GGUF loaded without a usable mmproj; " "image input will be disabled for this session" ) + # Seed before the try: the except (GPU-selection failure -> + # --fit on) falls through to the launch which reads this, and the + # probe that assigns it may throw first. Captured before manual + # empty `gpus` so the speculative defaults stay GPU-aware and the + # CPU-fallback check still knows GPUs were present. + _detected_gpus: list[tuple[int, int]] = [] model_size = None # set in the fit try; used by the APU RAM guard # Layer-fallback min GPUs; raised below on a tensor downgrade. Bound # before the try so the --fit-on except path still has it (no UnboundLocal). @@ -6057,6 +6394,18 @@ class LlamaCppBackend: _gpu_mem = self._get_gpu_memory(binary) gpus = [(idx, free) for idx, free, _t in _gpu_mem] total_by_idx = {idx: total for idx, _f, total in _gpu_mem} + # GPU picker: restrict every mode to the chosen devices, so + # auto selection only considers them and manual mask to + # them (the env block below pins CUDA/HIP_VISIBLE_DEVICES). + if gpu_ids: + _picked = set(gpu_ids) + gpus = [g for g in gpus if g[0] in _picked] + + # GPUs the model will run on -- captured before manual + # empty `gpus` to bypass the planner. bool() drives the + # GPU-aware speculative defaults; the list feeds the + # CPU-fallback check. + _detected_gpus = list(gpus) def _gpu_usable(g, frac = _CTX_FIT_VRAM_FRACTION): # Per-GPU usable budget for ranking: free - (1-frac)*total. @@ -6088,6 +6437,44 @@ class LlamaCppBackend: # GPU/VRAM-fit logic below may shrink it on limited HW. max_available_ctx = self._context_length or effective_ctx + # Manual + Auto layers (the Manual default): hand memory + # management to llama.cpp's --fit. Emptying the probed GPU set + # no-ops the selection/TP planning below, leaving gpu_indices + # None (an explicit gpu_ids pick still pins below) and use_fit + # True. An explicit context is honored (--fit optimizes around + # it); 0 lets --fit size it. + if gpu_memory_mode == "manual" and gpu_layers < 0: + # Tensor parallelism was already dropped above (before the + # cache-drop), so a quantized KV survives into this --fit load. + gpus = [] + effective_ctx = requested_ctx if requested_ctx > 0 else 0 + original_ctx = effective_ctx + # --fit aborts under --split-mode tensor; a raw extras + # --split-mode/--tensor-split (appended last) would + # otherwise reach llama-server. Strip it like the TP + # downgrade does. + extra_args = strip_split_mode_only(extra_args) + elif gpu_memory_mode == "manual": + # Manual offload (--gpu-layers + --fit off): no automatic + # device masking (a gpu_ids pick still pins below) or + # context cap -- the user owns both. tensor_parallel is + # honored but skips the memory-based planner (gpus = []); + # the toggle just emits --split-mode tensor (split by free + # VRAM, or by the Split ratio if set). + gpus = [] + effective_ctx = ( + requested_ctx if requested_ctx > 0 else (self._context_length or 0) + ) + original_ctx = effective_ctx + # Strip the user --split-mode when the toggle owns the split + # (TP engaged -> Studio emits --split-mode tensor) or when the + # user asked for tensor (which aborts on a single GPU even if + # the manual <2-GPU guard downgraded TP). Otherwise keep their + # non-tensor mode (row/none/layer) -- the toggle can't express + # those. + if tensor_parallel or split_mode_override == "tensor": + extra_args = strip_split_mode_only(extra_args) + # Will MTP engage? If so, auto-fit reserves draft-model VRAM. # Mirrors _build_speculative_flags: forced mtp/mtp+ngram always # engage; auto only on an MTP model >= 3B; ngram/off never. A @@ -6175,7 +6562,10 @@ class LlamaCppBackend: _extra_n_max = _extra_args_spec_draft_n_max(extra_args) _mtp_eff_n_max = _extra_n_max if _extra_n_max is not None else spec_draft_n_max if _mtp_eff_n_max is None: - _mtp_eff_n_max = 2 if gpus else 3 + # _detected_gpus (not gpus) so manual -- which empty + # gpus to bypass the planner -- keep the GPU draft depth the + # launch flags also use, instead of the CPU default. + _mtp_eff_n_max = 2 if _detected_gpus else 3 # Separate-drafter weights live on GPU (an embedded head is # already in model_size). Size the drafter the launch loads, by # precedence: extras --model-draft (last-wins), else Unsloth's @@ -6313,7 +6703,8 @@ class LlamaCppBackend: # honor it, cap only if it fits no combination. Auto (native): # prefer fewer GPUs with reduced context (multi-GPU is slower). gpu_indices, use_fit = None, True - # Per-GPU weight proportions for tensor mode (None = even). + # Per-GPU weight proportions for tensor mode (None lets + # llama.cpp split by free VRAM). tp_tensor_split: Optional[list[int]] = None explicit_ctx = requested_ctx > 0 # Flat MTP reserve fraction: used only as the fallback when the @@ -6388,7 +6779,12 @@ class LlamaCppBackend: # GPUs below that reserve from the set up front (gpu_indices # becomes the CUDA_VISIBLE_DEVICES mask, fully excluding them). tp_gpus = gpus - if tensor_parallel: + # Manual mode owns the layer count and context, so it skips + # the memory-based planner; its toggle still emits + # --split-mode tensor below (split by free VRAM, or by the + # Split ratio if set). auto plans here. + plan_tp = tensor_parallel and gpu_memory_mode != "manual" + if plan_tp: # Deterministic per-device compute buffer (replicated on # every device in tensor mode); flat fallback when dims # are unavailable. _plan_tensor_parallel uses the same. @@ -6407,7 +6803,7 @@ class LlamaCppBackend: # free yet have no budget left. tp_gpus = [g for g in gpus if _gpu_usable(g) >= reserve_mib] - if tensor_parallel and len(tp_gpus) < 2: + if plan_tp and len(tp_gpus) < 2: # Tensor parallelism needs >= 2 usable GPUs. On a single # GPU --split-mode tensor is a no-op; with 0 GPUs (CPU-only # or probe failed) it must not reach llama-server; and a @@ -6823,6 +7219,12 @@ class LlamaCppBackend: tp_tensor_split = None effective_ctx = requested_ctx # fall back to original + # GPU picker: when no narrower subset was chosen (manual, or + # a failed/file-size selection), pin the whole picked set so the + # model can't spill onto an unpicked GPU. + if gpu_ids and gpu_indices is None: + gpu_indices = sorted(gpu_ids) + # Unified-memory APUs load weights into system RAM (under WSL the VM # cap, not the ROCm-reported VRAM, is the real ceiling); refuse an # oversize load the OS would otherwise kill mid-flight. Base model @@ -6859,8 +7261,6 @@ class LlamaCppBackend: model_path, "--port", str(self._port), - "-c", - str(effective_ctx) if effective_ctx > 0 else "0", "--parallel", str(n_parallel), "--flash-attn", @@ -6868,6 +7268,17 @@ class LlamaCppBackend: # Error out at n_ctx instead of silently rotating the KV cache; frontend catches it and points the user at "Context Length". "--no-context-shift", ] + # A positive context is always passed (in auto-fit, --fit then + # optimizes the gpu-layer offload around it). When auto-fit has + # no explicit context, omit -c so --fit sizes it to fit VRAM: + # "-c 0" would instead pin the FULL native context (llama.cpp's + # -c handler sets fit_params_min_ctx = UINT32_MAX on value 0, + # disabling --fit's reduction). See gpu_memory_mode. + auto_fit = gpu_memory_mode == "manual" and gpu_layers < 0 + if effective_ctx > 0: + cmd.extend(["-c", str(effective_ctx)]) + elif not auto_fit: + cmd.extend(["-c", "0"]) # Report a clean public model id (matching GET /v1/models) rather # than the raw -m path in llama-server's own /v1/models and the @@ -6879,7 +7290,63 @@ class LlamaCppBackend: cmd.extend(["--alias", _alias]) fully_gpu_offloaded = False - if use_fit: + # Set when a positional --tensor-split is emitted, so the env block + # can pin CUDA to PCI order even without a GPU subset (see below). + manual_tensor_split_emitted = False + if gpu_memory_mode == "manual" and gpu_layers >= 0: + # Pin the user's layer count and disable auto-fit. --fit off + # also means _ctx_integrity_flags must not add --fit-ctx. + use_fit = False + cmd.extend(["--gpu-layers", str(gpu_layers), "--fit", "off"]) + # Keep the first n_cpu_moe MoE layers' experts on CPU. + moe_flag = self._resolve_cpu_moe_flag( + n_cpu_moe, + self.n_moe_layers, + self._leading_dense_block_count or 0, + ) + if moe_flag is not None: + cmd.extend(["--n-cpu-moe", str(moe_flag)]) + elif n_cpu_moe: + # Requested on a dense model: nothing was emitted, so + # don't report a count llama-server never received. + self._n_cpu_moe = 0 + # Distribute the model across GPUs by the user's per-GPU shares + # (--tensor-split). Works with layer split and tensor + # parallelism; --fit off means no fit/tensor abort. Only emit + # when >1 GPU is in use AND the list length matches that count: + # the field is hidden (not cleared) when the picker narrows to + # one, and a direct caller can send a stale ratio for a different + # GPU set. Studio drops any mismatch to the free-VRAM default + # (llama.cpp would silently zero-pad a short list, or abort past + # its 16-device cap). + _split_gpus = self._effective_gpu_count(gpu_indices) + if tensor_split and _split_gpus > 1: + # An all-zero/non-positive sanitized split assigns nothing + # anywhere, so fall through to the free-VRAM default in + # that case. + _sanitized_split = self._sanitize_tensor_split(tensor_split) + _split_total = sum(_sanitized_split) + if len(_sanitized_split) == _split_gpus and _split_total > 0: + cmd.extend( + ["--tensor-split", ",".join(f"{x:g}" for x in _sanitized_split)] + ) + self._tensor_split = _sanitized_split + manual_tensor_split_emitted = True + else: + logger.warning( + "Dropping manual --tensor-split (%d entries for " + "%d GPUs, sanitized total %s); llama.cpp's " + "free-VRAM split applies instead", + len(tensor_split), + _split_gpus, + _split_total, + ) + self._tensor_split = None + elif tensor_split: + # Single effective GPU: the split is never emitted, so + # don't report it as active via /status and /load. + self._tensor_split = None + elif use_fit: cmd.extend(["--fit", "on"]) elif gpu_indices is not None: # Fits on selected GPU(s) -- force all layers on GPU. --fit off is @@ -6893,10 +7360,31 @@ class LlamaCppBackend: # when the binary advertises it (older/custom binaries may not). if server_caps.get("supports_metrics"): cmd.append("--metrics") + self._slot_save_dir = None + self._slot_save_binary = None + self._prompt_cache_disabled = False + if server_caps.get("supports_slot_save"): + try: + from utils.paths.storage_roots import ( # noqa: WPS433 + llama_slot_cache_root, + ) + + slot_dir = llama_slot_cache_root() + slot_dir.mkdir(parents = True, exist_ok = True) + # Saved KV encodes chat content; keep it from other local users. + with contextlib.suppress(OSError): + os.chmod(slot_dir, 0o700) + cmd.extend(["--slot-save-path", str(slot_dir)]) + self._slot_save_dir = str(slot_dir) + self._slot_save_binary = (binary, Path(binary).stat().st_mtime_ns) + except OSError: + self._slot_save_dir = None + self._slot_save_binary = None cmd.extend( self._ctx_integrity_flags( n_parallel, use_fit, + auto_fit, requested_ctx, effective_ctx, server_caps, @@ -6960,9 +7448,11 @@ class LlamaCppBackend: self._cache_type_kv = None # Tensor parallelism: split the model across GPUs by tensor - # rather than by layer. Multi-GPU only -- a no-op on a single - # GPU. Default (layer split) is left implicit by omitting the - # flag. See llama.cpp --split-mode. + # rather than by layer. The UI only offers it on multi-GPU; a + # direct single-GPU caller is redundant (supported archs no-op, + # unsupported ones abort and the /load path retries layer split). + # Default (layer split) is left implicit by omitting the flag. + # See llama.cpp --split-mode. if tensor_parallel: cmd.extend(["--split-mode", "tensor"]) if tp_tensor_split and len(tp_tensor_split) > 1: @@ -6994,7 +7484,7 @@ class LlamaCppBackend: extra_args = extra_args, model_identifier = model_identifier, model_path = model_path, - gpus = bool(gpus), + gpus = bool(_detected_gpus), binary = binary, mtp_draft_path = launch_mtp_draft_path, ) @@ -7073,8 +7563,9 @@ class LlamaCppBackend: else: self._api_key = None - # Windows + full offload: disable KV checkpoints (WDDM/PCI-E - # overhead). CPU/partial offload keeps prompt caching. #5692. + # Windows + full offload: drop the host-RAM KV checkpoints that cause + # WDDM/PCI-E overhead, but keep prompt caching (in-VRAM prefix reuse) so + # a repeated prompt is not re-prefilled on every request. #5692. if sys.platform == "win32" and full_offload_tuning_active: unsupported_cache_flags: list[str] = [] if server_caps.get("supports_cache_ram"): @@ -7085,10 +7576,6 @@ class LlamaCppBackend: cmd.extend(["--ctx-checkpoints", "0"]) else: unsupported_cache_flags.append("--ctx-checkpoints") - if server_caps.get("supports_no_cache_prompt"): - cmd.append("--no-cache-prompt") - else: - unsupported_cache_flags.append("--no-cache-prompt") if unsupported_cache_flags: logger.info( "Skipping unsupported Windows cache flags for llama-server: %s", @@ -7112,6 +7599,8 @@ class LlamaCppBackend: # Library paths so llama-server finds its shared libs and CUDA DLLs. env = self._llama_server_env_for_binary(binary) + if gpu_memory_mode == "manual": + self._clear_manual_placement_env(env) # Omitting --threads relies on llama.cpp's physical-core default, so # drop an inherited LLAMA_ARG_THREADS that would otherwise feed the # arg handler and silently force hardware_concurrency(). #5692 @@ -7170,28 +7659,39 @@ class LlamaCppBackend: # CUDA_VISIBLE_DEVICES leaves an AMD child seeing the full set, so # set HIP_VISIBLE_DEVICES too. Vulkan is pinned via --device # (above), not here. - if gpu_indices is not None and not is_vulkan_backend: - pinned = ",".join(str(i) for i in gpu_indices) - env["CUDA_VISIBLE_DEVICES"] = pinned - try: - import torch as _torch - if getattr(_torch.version, "hip", None) is not None: - env["HIP_VISIBLE_DEVICES"] = pinned - # Do NOT also set ROCR_VISIBLE_DEVICES to the same - # value. ROCR_VISIBLE_DEVICES filters at the HSA/ROCr - # layer and HIP_VISIBLE_DEVICES at the HIP layer, so - # setting both with the same physical indices applies - # the mask twice: ROCR reduces the visible set and - # re-indexes it from 0, then HIP indexes into the - # already-reduced set. A single non-zero pin (e.g. - # "1") then points out of range at the HIP layer, HIP - # enumerates 0 devices, and llama.cpp falls back to - # CPU ("ggml_cuda_init: no ROCm-capable device is - # detected"). The HIP mask alone narrows correctly; - # clear any inherited ROCR mask so it can't double up. - env.pop("ROCR_VISIBLE_DEVICES", None) - except Exception as e: - logger.debug("Failed to set ROCm visibility env vars for child: %s", e) + # A deliberate zero-offload load with no GPU companions runs + # entirely on CPU, yet a visible CUDA device still costs the child + # ~0.5 GB (context + compute scratch) that the CPU-only + # classification below reports as free. Hide the GPUs so the load + # is exactly what it claims: zero VRAM (verified: GPU stays at idle + # baseline and generation runs). Companion loads keep the normal + # masking, and a user device pin (in extras or an inherited + # LLAMA_ARG_DEVICE) keeps control of its own devices -- the child + # aborts on a pin it can't see. The draft-device forms count too: + # llama-server parses them even with no drafter loaded. + _cpu_only_zero_offload = ( + gpu_memory_mode == "manual" + and gpu_layers == 0 + and not is_vulkan_backend + and not self._zero_offload_keeps_gpu_visible(cmd, env) + ) + if _cpu_only_zero_offload: + self._emit_child_gpu_visibility(env, "-1") + elif gpu_indices is not None and not is_vulkan_backend: + # When the user picked GPUs by index, align CUDA's ordering + # with the PCI-bus order the picker enumerated (nvidia-smi), + # so "GPU 1" in the UI is GPU 1 to llama.cpp -- not CUDA's + # default FASTEST_FIRST order (#5025). + if gpu_ids: + env["CUDA_DEVICE_ORDER"] = "PCI_BUS_ID" + self._emit_child_gpu_visibility(env, ",".join(str(i) for i in gpu_indices)) + elif manual_tensor_split_emitted and not is_vulkan_backend: + # A manual per-GPU ratio across ALL GPUs (no explicit pick, so + # no CUDA_VISIBLE_DEVICES mask above): the UI built the + # --tensor-split list in ascending physical/PCI index order, + # so pin the child's enumeration to that order too. The whole + # visible set stays in use; only its ordering is fixed. + self._pin_visible_gpu_order_for_split(env) # Captured before any text-only fallback strips it from cmd. launched_with_mmproj = "--mmproj" in cmd @@ -7350,7 +7850,6 @@ class LlamaCppBackend: self._effective_context_length = ( effective_ctx if effective_ctx > 0 else self._context_length ) - self._reconcile_effective_ctx_with_server() self._max_context_length = ( max_available_ctx if max_available_ctx > 0 else self._effective_context_length ) @@ -7535,6 +8034,10 @@ class LlamaCppBackend: "session; run 'unsloth studio update' to enable vision." ) cmd = self._strip_mmproj_args(_last_spawn_cmd) + # This retry bypasses _spawn_and_wait, so refresh the + # launched-argv snapshot itself -- the zero-offload + # classification below must not see the stripped --mmproj. + _last_spawn_cmd = list(cmd) self._is_vision = False self._mmproj_has_audio = False self._start_llama_process(cmd, env) @@ -7566,6 +8069,13 @@ class LlamaCppBackend: self._healthy = True self._commit_effective_parallel_slots(n_parallel) + # Server is up: adopt the real per-request context it allocated + # -- the length --fit chose, or a --parallel slot split -- so the + # reported context_length matches reality. (Querying /props + # before the spawn above always failed; the seeded value was the + # requested/native length.) + self._reconcile_effective_ctx_with_server() + # Commit caller intent only after _healthy=True so a failed start # can't poison the next inheritance check. None keeps prior, [] # clears, list sets. Source records hf_variant for the route's @@ -7580,11 +8090,24 @@ class LlamaCppBackend: self._mtp_runtime_fallback_active = _mtp_active_for_launched_server self._start_mtp_crash_watchdog() - # Catch silent CPU fallback when GPU was intended (#5106). - self._gpu_offload_active = self._classify_gpu_offload( - gpu_indices is not None or use_fit, gpus or [] - ) - if self._gpu_offload_active is False: + # Catch silent CPU fallback when GPU was intended (#5106). Manual + # offload (no picker) leaves gpu_indices None and use_fit False, so + # include its GPU-layer intent; use the preserved probe since + # auto-layers/manual empty `gpus`. A deliberate zero-offload load + # classifies by its launched argv instead: the main model is + # CPU-only by construction and must read False (not None), or + # training needlessly unloads a server holding no VRAM. + _deliberate_cpu_only = gpu_memory_mode == "manual" and gpu_layers == 0 + if _deliberate_cpu_only: + self._gpu_offload_active = self._zero_offload_gpu_flag( + _last_spawn_cmd, _detected_gpus, env + ) + else: + self._gpu_offload_active = self._classify_gpu_offload( + gpu_indices is not None or use_fit or gpu_memory_mode == "manual", + _detected_gpus, + ) + if self._gpu_offload_active is False and not _deliberate_cpu_only: logger.warning( "llama-server appears to have loaded the model entirely " "on CPU even though Unsloth detected at least one GPU. " @@ -7627,6 +8150,15 @@ class LlamaCppBackend: if not self._healthy: return False + # Snapshot the files the server actually loaded. If a GGUF shard or a + # LoRA/control-vector sidecar is swapped on disk afterwards while the + # old weights stay mapped, save_slots_for_resume() compares against + # this and refuses to persist KV that a reload could misapply. + if self._slot_save_dir: + self._slot_loaded_identity = ( + self._gguf_file_identity(self._gguf_path), + self._slot_launch_fingerprint(), + ) return True def _build_speculative_flags( @@ -7947,6 +8479,11 @@ class LlamaCppBackend: gguf_path: Optional[str] = None, spec_draft_n_max: Optional[int] = None, tensor_parallel: bool = False, + gpu_memory_mode: Literal["auto", "manual"] = "auto", + gpu_layers: int = -1, + n_cpu_moe: int = 0, + tensor_split: Optional[List[float]] = None, + gpu_ids: Optional[List[int]] = None, mtp_draft_path: Optional[str] = None, preserve_multi_gpu_on_layer: bool = False, ) -> bool: @@ -8003,6 +8540,38 @@ class LlamaCppBackend: ): return False + # The diffusion runner is mode-agnostic (always "auto", ignores the + # layer/MoE/split knobs), so a standing manual preference in the + # request must not force a needless reload -- only the GPU pick matters. + if not self._is_diffusion: + # A GPU-memory-mode flip (Unsloth / manual) must always reload. + if self._gpu_memory_mode != gpu_memory_mode: + return False + # Manual: a layer-count change always reloads (covers Auto(-1) <-> a + # pinned count); MoE/split only matter with an explicit offload. + if gpu_memory_mode == "manual" and ( + self._gpu_layers != gpu_layers + or ( + gpu_layers >= 0 + and ( + self._n_cpu_moe != n_cpu_moe + or (self._tensor_split or None) != (tensor_split or None) + ) + ) + ): + return False + # A changed GPU pick must reload (compare order-insensitively; None/[] + # both mean automatic). The diffusion runner collapses a multi-GPU pick + # to its single lowest device, so self._gpu_ids holds just that device; + # normalize the request the same way, or a multi-GPU pick that resolves + # to the same device needlessly reloads. + if self._is_diffusion: + requested_gpu_pick = [sorted(gpu_ids)[0]] if gpu_ids else None + else: + requested_gpu_pick = sorted(gpu_ids) if gpu_ids else None + if (self._gpu_ids or None) != requested_gpu_pick: + return False + # Compare on the canonical requested mode. With --spec-type in # extra_args the backend stores None; mirror that here. if _extra_args_set_spec_type(extra_args): @@ -8071,6 +8640,78 @@ class LlamaCppBackend: return None return classify_gpu_offload_lines(self._stdout_lines) + @staticmethod + def _cmd_has_gpu_companion(cmd: list, env: Optional[Mapping[str, str]] = None) -> bool: + """True when the argv/env carries a GPU companion: any --mmproj form, or + a drafter (Studio's --model-draft, the extras aliases, or the + LLAMA_ARG_SPEC_DRAFT_* env) -- these offload to the GPU regardless of + the main ``--gpu-layers``. A drafter explicitly forced to CPU + (--spec-draft-ngl 0 / --spec-draft-device cpu) doesn't count.""" + if any(str(a).startswith("--mmproj") for a in cmd): + return True + if _extra_args_mtp_draft_path(cmd, env) is None: + return False + return not _extra_args_draft_offloaded_to_cpu(cmd, env) + + @staticmethod + def _zero_offload_keeps_gpu_visible(cmd: list, env: Optional[Mapping[str, str]] = None) -> bool: + """Whether a zero-layer launch still has a reason to use visible GPUs. + + Keep this shared by child masking and post-launch residency bookkeeping: + a device pin, surviving tensor mode, mmproj, or GPU drafter prevents the + launch from being a confirmed zero-VRAM server. + """ + return ( + LlamaCppBackend._cmd_has_gpu_device_pin(cmd, env) + or _effective_tensor_parallel(cmd, False, env) + or LlamaCppBackend._cmd_has_gpu_companion(cmd, env) + ) + + @staticmethod + def _cmd_has_gpu_device_pin(cmd: list, env: Optional[Mapping[str, str]] = None) -> bool: + """True when the effective main or draft ``--device`` pin names a GPU.""" + main_flags = {"--device", "-dev"} + draft_flags = {"--spec-draft-device", "-devd", "--device-draft"} + last_main: Optional[str] = None + last_draft: Optional[str] = None + args = [str(arg) for arg in cmd] + for index, raw in enumerate(args): + flag, equals, inline = raw.partition("=") + if flag not in main_flags and flag not in draft_flags: + continue + value = inline if equals else (args[index + 1] if index + 1 < len(args) else "") + if flag in main_flags: + last_main = value + else: + last_draft = value + if last_main is None: + last_main = (env or {}).get("LLAMA_ARG_DEVICE") + + def _names_gpu(value: Optional[str]) -> bool: + if value is None: + return False + devices = [item.strip().lower() for item in value.split(",") if item.strip()] + return not devices or any(item not in ("cpu", "none") for item in devices) + + return _names_gpu(last_main) or _names_gpu(last_draft) + + @staticmethod + def _zero_offload_gpu_flag( + spawn_cmd: list, + detected_gpus: list, + env: Optional[Mapping[str, str]] = None, + ) -> Optional[bool]: + """GPU-residency flag for a deliberate manual zero-offload load. The + main model is CPU-only by construction, but device pins, tensor mode, + mmproj, and GPU drafters can still make the server hold VRAM. The counted + offload classifier cannot see those allocations. This uses the same + predicate as the launch-time zero-VRAM mask; None means no GPU signal.""" + if not detected_gpus: + return None + if LlamaCppBackend._is_vulkan_backend(): + return True + return LlamaCppBackend._zero_offload_keeps_gpu_visible(spawn_cmd, env) + def load_cancelled(self) -> bool: """True if a load was cancelled (e.g. via unload/_cancel_event) and not yet consumed by the next load_model. Lets the tensor->layer fallback @@ -8103,6 +8744,10 @@ class LlamaCppBackend: self._effective_context_length = None self._max_context_length = None self._reset_effective_parallel_slots() + self._slot_save_dir = None + self._slot_save_binary = None + self._slot_loaded_identity = None + self._prompt_cache_disabled = False self._chat_template = None self._chat_template_override = None self._supports_reasoning = False @@ -8114,11 +8759,18 @@ class LlamaCppBackend: self._supports_tools = False self._cache_type_kv = None self._tensor_parallel = False + self._gpu_memory_mode = "auto" + self._gpu_layers = -1 + self._n_cpu_moe = 0 + self._tensor_split = None + self._gpu_ids = None self._layer_preserves_tensor_intent = False self._speculative_type = None self._requested_spec_mode = None self._spec_draft_n_max = None self._n_layers = None + self._n_experts = None + self._leading_dense_block_count = None self._n_kv_heads = None self._n_kv_heads_by_layer = None self._n_heads = None @@ -8181,6 +8833,10 @@ class LlamaCppBackend: # Clear healthy so a /load during the replacement's warm-up can't # short-circuit against the previous server's health (#5401). self._healthy = False + # Reset to unknown so the training guard treats the next (still + # loading) server as VRAM-resident rather than reading the killed + # server's stale zero-offload flag until the health probe reclassifies. + self._gpu_offload_active = None # Drives _wait_for_vram_settle in the next load_model; set in finally # so both in-process and frontend Apply paths record the kill. self._last_kill_monotonic = time.monotonic() @@ -8618,6 +9274,237 @@ class LlamaCppBackend: return False return True + def _slot_launch_fingerprint(self) -> tuple: + # KV validity keys on extra args, stat'd sidecar weights, effective ctx. + sidecars = [] + for path in self._sidecar_weight_files(): + try: + st = os.stat(path) + sidecars.append((path, st.st_size, st.st_mtime_ns)) + except OSError: + sidecars.append((path, None, None)) + return ( + tuple(self._extra_args or ()), + tuple(sidecars), + self._requested_n_ctx, + self._effective_context_length, + getattr(self, "_cache_type_kv", None), + self.effective_parallel_slots, + ) + + def _gguf_file_identity(self, path) -> Optional[tuple]: + # (size, mtime_ns) per shard: a split GGUF keys KV validity on every sibling. + p = Path(path) + paths = [p] + m = _SHARD_FULL_RE.match(p.name) + if m: + prefix, _first, total = m.groups() + paths = [ + p.with_name(f"{prefix}-{i:05d}-of-{total}{p.suffix}") + for i in range(1, int(total) + 1) + ] + try: + return tuple((sp.stat().st_size, sp.stat().st_mtime_ns) for sp in paths) + except OSError: + return None + + _SIDECAR_WEIGHT_FLAGS = ( + "--lora", + "--lora-scaled", + "--control-vector", + "--control-vector-scaled", + ) + + def _sidecar_weight_files(self) -> list[str]: + # llama.cpp: comma-separated paths, FNAME:SCALE on -scaled (older builds: FNAME SCALE). + args = [str(a).strip() for a in (self._extra_args or ())] + files: list[str] = [] + for i, arg in enumerate(args): + flag, sep, inline = arg.partition("=") + if flag not in self._SIDECAR_WEIGHT_FLAGS: + continue + operand = inline if sep else (args[i + 1] if i + 1 < len(args) else "") + if not operand: + continue + candidates = [operand] + pieces = [p for p in operand.split(",") if p] + if len(pieces) > 1: + candidates.extend(pieces) + if flag.endswith("-scaled"): + for item in list(candidates): + # ":" tail is a scale; rpartition spares drive letters. + head, colon, tail = item.rpartition(":") + if not (colon and head): + continue + try: + float(tail) + except ValueError: + continue + candidates.append(head) + for cand in candidates: + if cand not in files: + files.append(cand) + return files + + def _prompt_cache_off(self) -> bool: + # Caching off makes restores useless; last prompt-cache flag wins, env only when unset. + last = None + for arg in self._extra_args or (): + flag = arg.strip().split("=", 1)[0] + if flag in ("--cache-prompt", "--no-cache-prompt"): + last = flag + if last is not None: + return last == "--no-cache-prompt" + if self._prompt_cache_disabled: + return True + if os.environ.get("LLAMA_ARG_NO_CACHE_PROMPT") is not None: + return True + env = (os.environ.get("LLAMA_ARG_CACHE_PROMPT") or "").strip().lower() + return env in {"off", "disabled", "false", "0"} + + def save_slots_for_resume( + self, should_abort: Optional[Callable[[], bool]] = None + ) -> Optional[dict]: + if ( + not self.is_loaded + or not self._slot_save_dir + or not self._gguf_path + or self._prompt_cache_off() + ): + return None + save_dir = Path(self._slot_save_dir) + gguf_stat = self._gguf_file_identity(self._gguf_path) + if gguf_stat is None: + return None + launch = self._slot_launch_fingerprint() + # If the GGUF or a sidecar was swapped on disk while the original weights + # stayed mapped, the live KV belongs to the old weights but a reload would + # load the new file. Persisting it would let restore misapply stale KV. + if self._slot_loaded_identity is not None and self._slot_loaded_identity != ( + gguf_stat, + launch, + ): + logger.debug("Skipping slot save: model files changed on disk since load") + return None + try: + estimate = self._estimate_kv_cache_bytes( + self._effective_context_length or self._context_length or 0, + self._cache_type_kv, + n_parallel = self.effective_parallel_slots, + ) + # Skip before writing anything when the estimate alone blows the cap, + # rather than fully writing a slot and discarding it afterwards. + if estimate > _SLOT_SAVE_MAX_BYTES: + logger.debug( + "Skipping slot save: estimated %d bytes exceeds cap %d", + estimate, + _SLOT_SAVE_MAX_BYTES, + ) + return None + # A 0 estimate means metadata was insufficient, not a zero-byte cache: + # a slot can still be many GiB, so demand room for the whole cap before + # trusting the post-write check. + required = (estimate if estimate > 0 else _SLOT_SAVE_MAX_BYTES) + (1 << 30) + if shutil.disk_usage(save_dir).free < required: + logger.debug("Skipping slot save: insufficient free disk") + return None + except Exception: + pass + token = uuid.uuid4().hex[:8] + entries: list[dict] = [] + total_bytes = 0 + for slot in range(self.effective_parallel_slots): + # A request pending mid-save waits on the gate; stop wasting its time. + if should_abort is not None and should_abort(): + break + filename = f"resume-{token}-slot{slot}.bin" + path = save_dir / filename + try: + resp = httpx.post( + f"{self.base_url}/slots/{slot}", + params = {"action": "save"}, + json = {"filename": filename}, + headers = self._auth_headers, + timeout = _SLOT_SAVE_HTTP_TIMEOUT, + trust_env = False, + ) + except Exception as e: + logger.debug(f"slot {slot} save failed: {e}") + with contextlib.suppress(OSError): + path.unlink() + break + if resp.status_code != 200: + logger.debug(f"slot {slot} save returned HTTP {resp.status_code}") + with contextlib.suppress(OSError): + path.unlink() + continue + try: + body = resp.json() + if not isinstance(body, dict): + raise ValueError("slot save response was not a JSON object") + n_saved = int(body.get("n_saved") or 0) + except Exception as e: + # A 200 that still wrote a file but returns a malformed body must + # clean up like the transport/HTTP error paths above, or the file + # (which holds chat KV) is orphaned until the next startup sweep. + logger.debug(f"slot {slot} save returned an invalid response: {e}") + with contextlib.suppress(OSError): + path.unlink() + continue + if n_saved <= 0: + with contextlib.suppress(OSError): + path.unlink() + continue + # Account by the bytes actually on disk, not the server-reported + # count, so the cap holds even if a custom binary under-reports. + try: + n_written = path.stat().st_size + except OSError: + n_written = 0 + total_bytes += n_written + entries.append({"id": slot, "filename": filename, "n_saved": n_saved}) + if total_bytes > _SLOT_SAVE_MAX_BYTES: + break # already over the cap; the discard below cleans up + if not entries: + return None + if total_bytes > _SLOT_SAVE_MAX_BYTES: + logger.debug( + "Discarding slot save: %d bytes exceeds cap %d", + total_bytes, + _SLOT_SAVE_MAX_BYTES, + ) + for entry in entries: + with contextlib.suppress(OSError): + (save_dir / entry["filename"]).unlink() + return None + return { + "dir": self._slot_save_dir, + "binary": self._slot_save_binary, + "gguf": str(self._gguf_path), + "gguf_stat": gguf_stat, + "launch": launch, + "slots": entries, + } + + def restore_slots_for_resume(self, manifest: dict) -> None: + if not self.is_loaded or not self._slot_save_dir: + return + for entry in manifest.get("slots") or []: + try: + resp = httpx.post( + f"{self.base_url}/slots/{int(entry['id'])}", + params = {"action": "restore"}, + json = {"filename": str(entry["filename"])}, + headers = self._auth_headers, + timeout = _SLOT_SAVE_HTTP_TIMEOUT, + trust_env = False, + ) + except Exception as e: + logger.debug(f"slot restore failed: {e}") + break + if resp.status_code != 200: + logger.debug(f"slot {entry.get('id')} restore returned HTTP {resp.status_code}") + def _maybe_recover_from_mtp_crash(self, exc: Optional[BaseException] = None) -> bool: """Schedule one background reload without MTP after a mid-generation death. @@ -8785,7 +9672,12 @@ class LlamaCppBackend: @staticmethod def _ctx_integrity_flags( - n_parallel: int, use_fit: bool, requested_ctx: int, effective_ctx: int, caps: dict + n_parallel: int, + use_fit: bool, + auto_fit: bool, + requested_ctx: int, + effective_ctx: int, + caps: dict, ) -> list[str]: """Flags that keep the per-request window equal to the advertised ctx. @@ -8793,14 +9685,28 @@ class LlamaCppBackend: ``--kv-unified`` default, silently splitting ``-c`` into per-slot windows of ``-c / N``; restore the shared pool so one request can use the full context. With ``--fit on``, ``--fit-ctx`` floors the fit step - at an explicitly requested ctx (default floor is 4096) so it offloads - or fails instead of silently shrinking the window. + at an explicitly requested ctx so it offloads or fails instead of + silently shrinking the window. The 8192 auto-floor and the tighter + ``--fit-target`` margin apply only under Manual + Auto (``auto_fit``), + which omits ``-c``: on the legacy auto path ``-c 0`` already pins the + native window and ``--fit-ctx 8192`` would override it down to 8192. """ flags: list[str] = [] if n_parallel > 1 and caps.get("supports_kv_unified"): flags.append("--kv-unified") - if use_fit and requested_ctx > 0 and effective_ctx > 0 and caps.get("supports_fit_ctx"): - flags.extend(["--fit-ctx", str(effective_ctx)]) + if use_fit and caps.get("supports_fit_ctx"): + if requested_ctx > 0 and effective_ctx > 0: + # Floor the fit step at the explicitly requested ctx. + flags.extend(["--fit-ctx", str(effective_ctx)]) + elif auto_fit: + # Manual + Auto omits -c, so floor at 8192 so --fit doesn't + # shrink the window below a usable size. + flags.extend(["--fit-ctx", "8192"]) + if use_fit and auto_fit and caps.get("supports_fit_target"): + # llama.cpp's --fit leaves 1 GiB free per device by default; + # tighten that to 512 MiB so it packs more of the model onto + # the GPU before spilling to system RAM. + flags.extend(["--fit-target", "512"]) return flags def _query_server_n_ctx(self) -> Optional[int]: @@ -8997,6 +9903,75 @@ class LlamaCppBackend: except Exception: logger.debug("Could not close httpx client", exc_info = True) + @staticmethod + def _install_cancel_aware_read( + client: "httpx.Client", + cancel_event: threading.Event, + response: Optional["httpx.Response"] = None, + poll_s: float = 0.2, + ) -> None: + """Wrap the httpcore stream so the reader interrupts its own blocked recv() on cancel. + + A cross-thread socket shutdown wakes a parked recv() on POSIX but not on + Windows (Winsock), so read in short slices and poll cancel_event between them + (plain or TLS); slice timeouts are swallowed so a slow-but-alive stream survives. + httpcore snapshots request.extensions["timeout"]["read"] once at body start, so + given ``response`` we re-read the live value per call to honor the post-first-token + stall timeout instead of the long prefill timeout.""" + import httpcore + + def _live_read_timeout() -> Optional[float]: + if response is None: + return None + try: + ext = response.request.extensions.get("timeout") + if isinstance(ext, dict): + value = ext.get("read") + if isinstance(value, (int, float)): + return float(value) + except Exception: + pass + return None + + try: + pool = getattr(getattr(client, "_transport", None), "_pool", None) + for connection in list(getattr(pool, "_connections", []) or []): + inner = getattr(connection, "_connection", None) + stream = getattr(inner, "_network_stream", None) + if stream is None or getattr(stream, "_unsloth_cancel_wrapped", False): + continue + orig_read = stream.read + + def read( + max_bytes, + timeout = None, + _orig = orig_read, + ): + live = _live_read_timeout() + effective = live if live is not None else timeout + deadline = None if effective is None else time.monotonic() + effective + while True: + if cancel_event.is_set(): + raise httpcore.ReadError("stream cancelled by user") + if deadline is None: + step = poll_s + else: + remaining = deadline - time.monotonic() + if remaining <= 0: + raise httpcore.ReadTimeout("read operation timed out") + step = min(poll_s, remaining) + try: + return _orig(max_bytes, timeout = step) + except httpcore.ReadTimeout: + if deadline is not None and time.monotonic() >= deadline: + raise + continue # slow but alive: keep reading + + stream.read = read + stream._unsloth_cancel_wrapped = True + except Exception: + logger.debug("Could not install cancel-aware read", exc_info = True) + @staticmethod @contextlib.contextmanager def _stream_with_retry( @@ -9054,6 +10029,11 @@ class LlamaCppBackend: headers = headers, ) as response: _response_ref[0] = response + if cancel_event is not None: + # Portable mid-stream cancel: the reader polls cancel itself, so + # Stop interrupts a stalled read where the watcher's Windows socket + # shutdown does not. Pass response to honor the live stall timeout. + LlamaCppBackend._install_cancel_aware_read(client, cancel_event, response) if cancel_event is not None and cancel_event.is_set(): raise _LlamaStreamCancelled yield response diff --git a/studio/backend/core/inference/llama_keepwarm.py b/studio/backend/core/inference/llama_keepwarm.py index 86a8c8a404..3380ebf5f5 100644 --- a/studio/backend/core/inference/llama_keepwarm.py +++ b/studio/backend/core/inference/llama_keepwarm.py @@ -15,6 +15,7 @@ import asyncio import contextlib import threading import time +from pathlib import Path from loggers import get_logger @@ -30,6 +31,8 @@ _last_active = time.monotonic() # otherwise 503 against an empty backend can reload it (set on unload, cleared on # reload). Storing the quant means the reload restores the exact freed variant. _last_unloaded_model = None +# Slot KV manifest saved by the idle unload; whoever pops it owns deleting its files. +_kv_resume = None # Guards inflight bumps against the idle-check-then-unload race, and blocks new # inference from starting mid-swap. Process-wide, not per-loop: the backend slot is # shared across every event loop in the process, so a per-loop gate would let a @@ -161,11 +164,17 @@ def inference_lifecycle_gate(): return _unload_gate() -def note_model_loaded() -> None: - """Record a successful GGUF load: stamp activity and drop any reload stash so - a manual load clears it synchronously, not only on the next idle poll.""" +def note_model_loaded(backend = None) -> None: + """Stamp activity and synchronously drop any reload stash.""" _note_activity() + resume = take_kv_resume() _set_last_unloaded(None) + if resume is None: + return + if backend is not None: + restore_kv_resume(backend, resume) + else: + _delete_resume_files(resume) def note_model_unloaded() -> None: @@ -182,9 +191,81 @@ def get_last_unloaded_model(): def _set_last_unloaded(value) -> None: - global _last_unloaded_model + global _last_unloaded_model, _kv_resume + stale = None with _lock: _last_unloaded_model = value + if value is None and _kv_resume is not None: + stale, _kv_resume = _kv_resume, None + if stale: + _delete_resume_files(stale) + + +def _delete_resume_files(manifest) -> None: + try: + base = Path(manifest.get("dir") or "") + for entry in manifest.get("slots") or []: + with contextlib.suppress(OSError): + (base / str(entry.get("filename"))).unlink() + except Exception: + pass + + +def _set_kv_resume(value) -> None: + global _kv_resume + stale = None + with _lock: + if _kv_resume is not None and _kv_resume is not value: + stale = _kv_resume + _kv_resume = value + if stale: + _delete_resume_files(stale) + + +def take_kv_resume(): + global _kv_resume + with _lock: + manifest, _kv_resume = _kv_resume, None + return manifest + + +def purge_kv_resume() -> None: + resume = take_kv_resume() + if resume: + _delete_resume_files(resume) + + +def restore_kv_resume(backend, manifest) -> None: + try: + gguf = manifest.get("gguf") + binary = manifest.get("binary") + current = getattr(backend, "_gguf_path", None) + same_gguf = bool(gguf and current) and Path(current).resolve() == Path(gguf).resolve() + if same_gguf: + # Same path is not enough: shards may have been rewritten meanwhile. + identity = getattr(backend, "_gguf_file_identity", None) + same_gguf = callable(identity) and identity(current) == manifest.get("gguf_stat") + if same_gguf: + # Nor the same file: launch overrides can invalidate KV numerics. + fingerprint = getattr(backend, "_slot_launch_fingerprint", None) + same_gguf = callable(fingerprint) and manifest.get("launch") == fingerprint() + if same_gguf and binary and binary == getattr(backend, "_slot_save_binary", None): + logger.info("Restoring saved slot KV onto the reloaded model") + backend.restore_slots_for_resume(manifest) + except Exception as exc: + logger.debug("slot restore after reload failed: %s", exc) + finally: + _delete_resume_files(manifest) + + +def sweep_slot_save_dir() -> None: + try: + from utils.paths.storage_roots import llama_slot_cache_root + for path in llama_slot_cache_root().glob("resume-*.bin"): + with contextlib.suppress(OSError): + path.unlink() + except Exception: + pass class LlamaKeepWarmMiddleware: @@ -266,7 +347,10 @@ def _loaded_identity(backend): async def idle_unload_loop(poll_seconds: float = 15.0) -> None: """Unload the loaded GGUF once idle past the configured TTL. Inert when off.""" - from utils.openai_auto_switch_settings import get_auto_unload_idle_seconds + from utils.openai_auto_switch_settings import ( + get_auto_unload_idle_seconds, + get_auto_unload_keep_kv, + ) seen_model = None while True: @@ -281,17 +365,47 @@ async def idle_unload_loop(poll_seconds: float = 15.0) -> None: # Track by (id, variant): a (re)loaded model -- including the same repo # at a different quant -- counts as activity so it survives one TTL # before its first request (loads bypass the activity middleware). - current = _loaded_identity(backend) - if current != seen_model: - seen_model = current - if current is not None: - _note_activity() - _set_last_unloaded(None) # a model is loaded; drop stale stash async with _unload_gate(): + # Purging the stash mid-reload would race the restore. + current = _loaded_identity(backend) + if current != seen_model: + seen_model = current + if current is not None: + _note_activity() + _set_last_unloaded(None) # a model is loaded; drop stale stash if backend.is_loaded and _is_idle(ttl): freed = _loaded_identity(backend) - await asyncio.to_thread(backend.unload_model) + manifest = None + if get_auto_unload_keep_kv(): + try: + manifest = await asyncio.to_thread( + backend.save_slots_for_resume, + lambda: not _is_idle(ttl), + ) + except Exception as exc: + logger.debug("slot save before idle unload failed: %s", exc) + # Re-read settings: the save can outlive a settings change. + ttl = get_auto_unload_idle_seconds() + if ttl <= 0 or not _is_idle(ttl): + if manifest: + _delete_resume_files(manifest) + continue + if manifest and not get_auto_unload_keep_kv(): + _delete_resume_files(manifest) + manifest = None + try: + await asyncio.to_thread(backend.unload_model) + except Exception: + # Failed unload means nothing will stash the manifest. + if manifest: + _delete_resume_files(manifest) + raise _set_last_unloaded(freed) # let an alias request reload it + if manifest and freed: + _set_kv_resume({"identity": freed, **manifest}) + logger.info("Idle auto-unload: saved slot KV for restore on reload") + elif manifest: + _delete_resume_files(manifest) logger.info("Idle auto-unload: freed GGUF after %ss idle", ttl) seen_model = None except Exception as exc: diff --git a/studio/backend/core/inference/llama_server_args.py b/studio/backend/core/inference/llama_server_args.py index 70d0dc774d..7b42d2f40d 100644 --- a/studio/backend/core/inference/llama_server_args.py +++ b/studio/backend/core/inference/llama_server_args.py @@ -70,6 +70,8 @@ _DENYLIST_GROUPS: tuple[frozenset[str], ...] = ( # llama-server's own built-in tools flag would silently stack on top of # Unsloth's --enable-tools / --disable-tools policy resolver. frozenset({"--tools"}), + # Slot-state dir: Studio owns it for KV persistence across idle unload. + frozenset({"--slot-save-path"}), ) _DENYLIST: frozenset[str] = frozenset().union(*_DENYLIST_GROUPS) @@ -186,12 +188,25 @@ _SPLIT_MODE_FLAGS: frozenset[str] = frozenset({"-sm", "--split-mode"}) _TENSOR_SPLIT_FLAGS: frozenset[str] = frozenset({"-ts", "--tensor-split"}) _SPLIT_SHADOWING_FLAGS: frozenset[str] = _SPLIT_MODE_FLAGS | _TENSOR_SPLIT_FLAGS +# GPU-offload flags. Stripped only when the GPU Memory mode owns offload +# (manual emits --fit / --gpu-layers / --n-cpu-moe); in auto, a user's +# inherited -ngl is respected (the offload_overridden path), so this group is +# opt-in, not default. Layer flags are shared with llama_cpp's override +# detection; the MoE flags are strip-only (manual's --n-cpu-moe slider owns them). +_LAYER_OFFLOAD_FLAGS: frozenset[str] = frozenset( + {"-ngl", "--gpu-layers", "--n-gpu-layers", "-fit", "--fit"} +) +_MOE_OFFLOAD_FLAGS: frozenset[str] = frozenset({"-ncmoe", "--n-cpu-moe", "-cmoe", "--cpu-moe"}) +_OFFLOAD_SHADOWING_FLAGS: frozenset[str] = _LAYER_OFFLOAD_FLAGS | _MOE_OFFLOAD_FLAGS + _SHADOWING_FLAGS: frozenset[str] = ( _CONTEXT_FLAGS | _CACHE_FLAGS | _SPEC_FLAGS | _TEMPLATE_FLAGS | _SPLIT_SHADOWING_FLAGS ) # Shadowing flags that take no value -- strip the flag only, not the next token. -_BOOLEAN_SHADOWING_FLAGS: frozenset[str] = frozenset({"--spec-default", "--jinja", "--no-jinja"}) +_BOOLEAN_SHADOWING_FLAGS: frozenset[str] = frozenset( + {"--spec-default", "--jinja", "--no-jinja", "-cmoe", "--cpu-moe"} +) def parse_ctx_override(args: Optional[Iterable[str]]) -> Optional[int]: @@ -424,6 +439,8 @@ def strip_shadowing_flags( strip_spec: bool = True, strip_template: bool = True, strip_split_mode: bool = True, + strip_tensor_split: bool = False, + strip_offload: bool = False, ) -> list[str]: """Strip flags that shadow first-class Unsloth settings. @@ -432,6 +449,12 @@ def strip_shadowing_flags( (same for cache / spec / template / split-mode). Each ``strip_*`` toggle controls one group; the route only strips groups whose first-class field the caller actually supplied. + + ``strip_split_mode`` removes both ``--split-mode`` and the coupled + ``--tensor-split`` (the Tensor Parallelism toggle owns the whole split). + ``strip_tensor_split`` removes ``--tensor-split`` *alone*, so manual mode can + replace an inherited per-GPU ratio while leaving the user's ``--split-mode`` + row/none/layer choice intact. """ shadowing: set[str] = set() if strip_context: @@ -444,6 +467,10 @@ def strip_shadowing_flags( shadowing |= _TEMPLATE_FLAGS if strip_split_mode: shadowing |= _SPLIT_SHADOWING_FLAGS + if strip_tensor_split: + shadowing |= _TENSOR_SPLIT_FLAGS + if strip_offload: + shadowing |= _OFFLOAD_SHADOWING_FLAGS tokens = [str(a) for a in (args or [])] out: list[str] = [] diff --git a/studio/backend/core/inference/orchestrator.py b/studio/backend/core/inference/orchestrator.py index eaa474d9b8..75ef9c2399 100644 --- a/studio/backend/core/inference/orchestrator.py +++ b/studio/backend/core/inference/orchestrator.py @@ -54,9 +54,8 @@ class GenStreamError(str): """A stream chunk carrying a real backend/generation error, not model text. Subclasses str so existing display/logging consumers are unaffected, while - callers that must abort a distributed run on error (raise_on_streamed_error) - can distinguish a real error from model output whose visible text starts with - "Error:" by checking isinstance(chunk, GenStreamError). + callers can distinguish a real error from model output whose visible text + starts with "Error:" by checking isinstance(chunk, GenStreamError). """ __slots__ = ("public",) diff --git a/studio/backend/core/rag/store.py b/studio/backend/core/rag/store.py index f9128d1715..1165b6bb0e 100644 --- a/studio/backend/core/rag/store.py +++ b/studio/backend/core/rag/store.py @@ -158,6 +158,16 @@ def list_documents(conn: sqlite3.Connection, scope: str) -> list[dict]: return [dict(r) for r in rows] +def list_all_documents(conn: sqlite3.Connection) -> list[dict]: + """Every uploaded document across all scopes (KBs, threads, projects).""" + rows = conn.execute( + "SELECT id, scope, kb_id, thread_id, project_id, filename, sha256, status, error, " + "num_chunks, stored_path, created_at " + "FROM documents ORDER BY created_at DESC" + ).fetchall() + return [dict(r) for r in rows] + + def get_document(conn: sqlite3.Connection, document_id: str) -> dict | None: row = conn.execute("SELECT * FROM documents WHERE id=?", (document_id,)).fetchone() return dict(row) if row else None diff --git a/studio/backend/core/training/resume.py b/studio/backend/core/training/resume.py index bbd9a895ab..17183484c5 100644 --- a/studio/backend/core/training/resume.py +++ b/studio/backend/core/training/resume.py @@ -4,6 +4,8 @@ """Helpers for validating resumable training outputs.""" import json +import pickletools +import zipfile from pathlib import Path from typing import Optional @@ -33,21 +35,158 @@ def _checkpoint_step(path: Path) -> int: return -1 -def get_resume_checkpoint_path(path_value: str) -> Optional[str]: +_MODEL_FILES = ( + "adapter_model.safetensors", + "adapter_model.bin", + "model.safetensors", + "pytorch_model.bin", +) +_MODEL_INDEXES = ("model.safetensors.index.json", "pytorch_model.bin.index.json") + + +def _valid_state_file(path: Path, require_tensor: bool = True) -> bool: + try: + if not path.is_file() or path.stat().st_size == 0: + return False + if path.suffix == ".safetensors": + try: + from safetensors import SafetensorError, safe_open + except ImportError: + return False + try: + with safe_open(str(path), framework = "np") as state: + return bool(state.keys()) + except SafetensorError: + return False + if path.suffix in {".bin", ".pt"}: + with zipfile.ZipFile(path) as state: + infos = state.infolist() + names = [info.filename for info in infos] + data_name = next( + (name for name in names if name == "data.pkl" or name.endswith("/data.pkl")), + None, + ) + if data_name is None: + return False + data_prefix = data_name.removesuffix("data.pkl") + "data/" + operations = list(pickletools.genops(state.read(data_name))) + if not operations or operations[-1][0].name != "STOP": + return False + if not require_tensor: + return True + # Require a non-empty tensor record; a zero-byte one fails torch.load. + return any( + info.filename.startswith(data_prefix) + and not info.is_dir() + and info.file_size > 0 + for info in infos + ) + # Unrecognized state-file formats are not usable resume state. + return False + except (OSError, ValueError, zipfile.BadZipFile): + return False + + +def _checkpoint_state(path: Path) -> Optional[int]: + try: + state = json.loads((path / "trainer_state.json").read_text(encoding = "utf-8")) + step = state.get("global_step") if isinstance(state, dict) else None + except (OSError, UnicodeDecodeError, json.JSONDecodeError): + return None + if isinstance(step, bool) or not isinstance(step, int) or step < 0: + return None + directory_step = _checkpoint_step(path) + return step if directory_step < 0 or step == directory_step else None + + +_INDEX_SHARD_SUFFIX = { + "model.safetensors.index.json": ".safetensors", + "pytorch_model.bin.index.json": ".bin", +} + + +def _valid_indexed_shard(checkpoint: Path, shard: object, expected_suffix: str) -> bool: + # Shard must be a relative, in-format path contained in the checkpoint dir. + if not isinstance(shard, str) or not shard: + return False + if Path(shard).is_absolute() or Path(shard).suffix != expected_suffix: + return False + try: + root = checkpoint.resolve(strict = True) + candidate = (checkpoint / shard).resolve(strict = True) + candidate.relative_to(root) + except (OSError, ValueError): + return False + return _valid_state_file(candidate) + + +def _has_model_state(path: Path) -> bool: + if any(_valid_state_file(path / name) for name in _MODEL_FILES): + return True + for name in _MODEL_INDEXES: + try: + index = json.loads((path / name).read_text(encoding = "utf-8")) + shards = set(index["weight_map"].values()) + except ( + AttributeError, + OSError, + KeyError, + TypeError, + UnicodeDecodeError, + json.JSONDecodeError, + ): + continue + expected_suffix = _INDEX_SHARD_SUFFIX[name] + if shards and all(_valid_indexed_shard(path, shard, expected_suffix) for shard in shards): + return True + return False + + +def is_resume_checkpoint_valid( + path: Path, + expected_step: Optional[int] = None, + backend: Optional[str] = None, +) -> bool: + step = _checkpoint_state(path) if path.is_dir() else None + step_valid = step is not None and (expected_step is None or step == expected_step) + if backend == "mlx": + valid_bundle = _valid_state_file(path / "adapters.safetensors") and _valid_state_file( + path / "optimizer_state.safetensors" + ) + else: + valid_bundle = ( + _has_model_state(path) + # optimizer/scheduler state can be validly tensor-free (e.g. SGD without + # momentum); _has_model_state still requires real model tensors. + and _valid_state_file(path / "optimizer.pt", require_tensor = False) + and _valid_state_file(path / "scheduler.pt", require_tensor = False) + ) + if backend is None and not valid_bundle: + valid_bundle = _valid_state_file(path / "adapters.safetensors") and _valid_state_file( + path / "optimizer_state.safetensors" + ) + return step_valid and valid_bundle + + +def get_resume_checkpoint_path( + path_value: str, expected_step: Optional[int] = None +) -> Optional[str]: path = resolve_output_dir(path_value) if not _is_under_outputs(path) or not path.is_dir(): return None - if (path / "trainer_state.json").is_file(): + if is_resume_checkpoint_valid(path, expected_step): return str(path) - checkpoints = [ - child - for child in path.glob("checkpoint-*") - if child.is_dir() and (child / "trainer_state.json").is_file() - ] - if not checkpoints: - return None - return str(max(checkpoints, key = _checkpoint_step)) + checkpoints = sorted(path.glob("checkpoint-*"), key = _checkpoint_step, reverse = True) + return next( + ( + str(checkpoint) + for checkpoint in checkpoints + if _checkpoint_step(checkpoint) >= 0 + and is_resume_checkpoint_valid(checkpoint, expected_step) + ), + None, + ) def normalize_resume_output_dir(path_value: str) -> str: @@ -78,9 +217,17 @@ def _uses_s3_dataset(run: dict) -> bool: def can_resume_run(run: dict) -> bool: if run.get("resumed_later"): return False + # Set when a stop-and-save failed to write a current-step checkpoint. + if run.get("resume_blocked"): + return False if _uses_s3_dataset(run): return False + status = run.get("status") + if status == "error": + # A save-time crash can report final_step == total_steps with no artifacts; checkpoint state alone decides resumability. + return has_resume_state(run.get("output_dir")) + final_step = run.get("final_step") total_steps = run.get("total_steps") has_remaining_steps = ( @@ -89,8 +236,4 @@ def can_resume_run(run: dict) -> bool: or total_steps <= 0 or final_step < total_steps ) - return ( - run.get("status") == "stopped" - and has_remaining_steps - and has_resume_state(run.get("output_dir")) - ) + return status == "stopped" and has_remaining_steps and has_resume_state(run.get("output_dir")) diff --git a/studio/backend/core/training/trainer.py b/studio/backend/core/training/trainer.py index 26720865f4..8e419849cb 100644 --- a/studio/backend/core/training/trainer.py +++ b/studio/backend/core/training/trainer.py @@ -3425,15 +3425,19 @@ class UnslothTrainer: logger.info( f"CPT: using UnslothTrainer with embedding_learning_rate={embedding_lr}\n" ) + cpt_args = _UnslothTrainingArguments( + embedding_learning_rate = embedding_lr, + **config_args, + ) + if config_args.get("packing", False): + cpt_args.packing_strategy = "wrapped" + logger.info("CPT packing strategy: wrapped\n") trainer_kwargs = { "model": self.model, "tokenizer": sft_tokenizer, "train_dataset": dataset["dataset"], "data_collator": data_collator, - "args": _UnslothTrainingArguments( - embedding_learning_rate = embedding_lr, - **config_args, - ), + "args": cpt_args, } if eval_dataset is not None: trainer_kwargs["eval_dataset"] = eval_dataset diff --git a/studio/backend/core/training/training.py b/studio/backend/core/training/training.py index b407ba39a5..26d8c23b66 100644 --- a/studio/backend/core/training/training.py +++ b/studio/backend/core/training/training.py @@ -761,6 +761,7 @@ class TrainingBackend: # Left True after an abnormal death so _ensure_pump_alive spots a crash. self._pump_running: bool = False self._lock = threading.Lock() + self._run_intent_lock = threading.RLock() # Stop watchdog: after a stop is requested, escalates to force_terminate() # if the worker does not exit on its own within a bounded time. The watched @@ -773,6 +774,7 @@ class TrainingBackend: self._progress = TrainingProgress() self._should_stop = False self._cancel_requested = False # True only for stop(save=False) + self._cancel_cleanup_output_dir: Optional[str] = None # Throttled training-status logging to the server log (not one line/step). self._last_progress_log_ts: float = 0.0 @@ -792,6 +794,8 @@ class TrainingBackend: # Job metadata self.current_job_id: Optional[str] = None self._output_dir: Optional[str] = None + self._resume_source_run_id: Optional[str] = None + self._terminal_finalize_payload: Optional[dict] = None # DB persistence self._metric_buffer: list[dict] = [] @@ -819,6 +823,7 @@ class TrainingBackend: job_id: str, *, before_spawn = None, + resume_source_run_id: Optional[str] = None, **kwargs, ) -> bool: """Spawn a subprocess to run the full training pipeline. @@ -956,6 +961,7 @@ class TrainingBackend: self.current_job_id = job_id self._should_stop = False self._cancel_requested = False + self._cancel_cleanup_output_dir = None self._complete_seen.clear() self._progress = TrainingProgress( is_training = True, status_message = "Initializing training..." @@ -972,7 +978,10 @@ class TrainingBackend: self.eval_loss_history.clear() self.eval_step_history.clear() self.eval_enabled = False - self._output_dir = None + self._output_dir = config.get("output_dir") if resume_source_run_id else None + self._progress.output_dir = self._output_dir + self._resume_source_run_id = resume_source_run_id + self._terminal_finalize_payload = None self._metric_buffer.clear() self._run_finalized = False self._db_run_created = False @@ -990,6 +999,17 @@ class TrainingBackend: # in history during model loading and a fast terminal worker can't race the # pump into a duplicate create/finalize. From here the pump only finalizes. self._ensure_db_run_created() + if resume_source_run_id and not self._db_run_created: + if proc.is_alive(): + proc.terminate() + proc.join(timeout = 5.0) + if proc.is_alive(): + proc.kill() + proc.join(timeout = 2.0) + self._progress.is_training = False + self._progress.error = "Resume checkpoint is no longer available." + self._spawn_in_progress = False + return False # Assign handles and start the pump together under the lock so a concurrent # poll can't see a live _proc with no pump and spawn a duplicate. @@ -1011,28 +1031,75 @@ class TrainingBackend: def stop_training(self, save: bool = True) -> bool: """Send stop signal to the training subprocess.""" - self._should_stop = True - if not save: - self._cancel_requested = True - with self._lock: - if self._stop_queue is not None: - try: - self._stop_queue.put({"type": "stop", "save": save}) - except (OSError, ValueError): - pass - # Update progress immediately for responsive UI. - self._progress.status_message = ( - "Stopping training and saving checkpoint..." if save else "Cancelling training..." - ) - # Guarantee the run finalizes even if the worker wedges after saving. - self._start_stop_watchdog(cancel = not save) + with self._run_intent_lock: + with self._lock: + run_id = self.current_job_id + if not save and run_id: + persist_error: Optional[Exception] = None + for attempt in range(_DB_FINALIZE_RETRIES): + try: + from storage.studio_db import mark_run_cancel_requested + + self._ensure_db_run_created() + with self._lock: + terminal_payload = self._terminal_finalize_payload + if ( + terminal_payload + and terminal_payload.get("expected_job_id") == run_id + ): + return False + if not mark_run_cancel_requested(run_id): + if self._db_run_created: + return False + raise RuntimeError( + "Training run disappeared before cancellation persisted" + ) + if self.current_job_id != run_id: + return False + self._should_stop = self._cancel_requested = True + self._cancel_cleanup_output_dir = self._output_dir + self._output_dir = self._progress.output_dir = None + persist_error = None + break + except Exception as exc: + persist_error = exc + if attempt + 1 < _DB_FINALIZE_RETRIES: + time.sleep(_DB_FINALIZE_RETRY_S) + if persist_error is not None: + raise RuntimeError("Failed to persist Stop-without-Save") from persist_error + with self._lock: + if self.current_job_id != run_id: + return False + if save or not run_id: + self._should_stop = True + if not save and not run_id: + self._cancel_requested = True + self._cancel_cleanup_output_dir = self._output_dir + self._output_dir = self._progress.output_dir = None + if self._stop_queue is not None: + try: + self._stop_queue.put({"type": "stop", "save": save}) + except (OSError, ValueError): + pass + self._progress.status_message = ( + "Stopping training and saving checkpoint..." + if save + else "Cancelling training..." + ) + self._start_stop_watchdog(cancel = not save, expected_job_id = run_id) return True - def _start_stop_watchdog(self, cancel: bool) -> None: + def _start_stop_watchdog( + self, + cancel: bool, + expected_job_id: Optional[str] = None, + ) -> None: """Start a daemon that force-terminates the worker if a requested stop does not exit on its own. No-op if no worker is alive or a live watchdog already watches this proc (a stale watchdog on an old proc never blocks a new run's watcher).""" with self._lock: + if expected_job_id is not None and self.current_job_id != expected_job_id: + return proc = self._proc if proc is None or not proc.is_alive(): return @@ -1113,8 +1180,9 @@ class TrainingBackend: watched_job_id: Optional[str] = None, ) -> None: """Finalize parent state after a force-terminate so the UI leaves "Stopping..." - even if the worker is wedged in driver teardown; preserves output_dir so a saved - checkpoint is kept. No-ops if a new run already replaced the watched worker, so a + even if the worker is wedged in driver teardown; preserves output_dir on a save so + the checkpoint is kept, and clears it on a cancel (Stop without saving must not + offer resume/export). No-ops if a new run already replaced the watched worker, so a stale watchdog never marks a fresh run stopped or drops its handle. Supersession is checked on both the watched proc and job id: start_training sets @@ -1134,7 +1202,18 @@ class TrainingBackend: return # a new run is already starting up; leave its state alone run_id = self.current_job_id # == watched_job_id self._progress.is_training = False - self._progress.status_message = "Training stopped." + terminal_payload = self._terminal_finalize_kwargs() + status = terminal_payload["status"] + error_message = terminal_payload.get("error_message") + output_dir = terminal_payload["output_dir"] + clear_output_dir = terminal_payload["clear_output_dir"] + resume_blocked = bool(terminal_payload.get("resume_blocked")) + with self._lock: + if self.current_job_id != run_id: + return + self._progress.status_message = error_message or "Training stopped." + if error_message: + self._progress.error = error_message # Create the row if a start-time create failed (no-op otherwise; skips when the pump # is mid-create, in which case its create-then-finalize records the run instead). self._ensure_db_run_created() @@ -1148,7 +1227,8 @@ class TrainingBackend: batch: list = [] final_step = final_loss = duration = None loss_history: list = [] - output_dir = self._output_dir + if clear_output_dir: + self._output_dir = self._progress.output_dir = None if claim: self._run_finalized = True # claim this run's finalize batch = list(self._metric_buffer) @@ -1161,7 +1241,17 @@ class TrainingBackend: loss_history = list(self.loss_history) if claim: self._finish_stopped_run( - run_id, output_dir, batch, final_step, final_loss, duration, loss_history + run_id, + output_dir, + batch, + final_step, + final_loss, + duration, + loss_history, + status = status, + error_message = error_message, + clear_output_dir = clear_output_dir, + resume_blocked = resume_blocked, ) with self._lock: if target_proc is None or self._proc is target_proc: @@ -1176,6 +1266,10 @@ class TrainingBackend: final_loss: Optional[float], duration: Optional[float], loss_history: list, + status: str = "stopped", + error_message: Optional[str] = None, + clear_output_dir: bool = False, + resume_blocked: bool = False, ) -> None: """Record a force-stopped run finished by its captured id, from state snapshotted under the lock. insert_metrics_batch upserts and finish_run is an idempotent UPDATE, @@ -1194,14 +1288,16 @@ class TrainingBackend: sparkline = downsample(loss_history, 50) finish_run( id = run_id, - status = "stopped", + status = status, ended_at = datetime.now(timezone.utc).isoformat(), final_step = final_step, final_loss = final_loss, duration_seconds = duration, loss_sparkline = _json.dumps(sparkline), output_dir = output_dir, - error_message = None, + error_message = error_message, + clear_output_dir = clear_output_dir, + resume_blocked = resume_blocked, ) return except Exception: @@ -1231,7 +1327,7 @@ class TrainingBackend: logger.info("Force-terminating training subprocess (pid=%s)", proc.pid) proc.terminate() cancelled = self._cancel_requested - output_dir = self._output_dir + output_dir = self._cancel_cleanup_output_dir or self._output_dir if proc is not None: proc.join(timeout = 5.0) @@ -1595,17 +1691,60 @@ class TrainingBackend: ) self._ensure_db_run_created() - self._finalize_run_in_db( - status = "stopped" if self._should_stop else "error", - error_message = None - if self._should_stop - else "Training process terminated unexpectedly", - ) + terminal_payload = self._terminal_finalize_kwargs() + with self._lock: + if terminal_payload["clear_output_dir"]: + self._output_dir = self._progress.output_dir = None + if terminal_payload.get("error_message"): + self._progress.error = terminal_payload["error_message"] + self._progress.status_message = terminal_payload["error_message"] + self._finalize_run_in_db(**terminal_payload) except Exception: logger.exception("Training event pump: finalization after worker exit failed") self._pump_running = False return + def _has_current_resume_checkpoint(self, output_dir, step) -> bool: + # A valid checkpoint at the current step means the stop-and-save landed on + # disk even if the worker died before confirming it. + if not output_dir or not isinstance(step, int) or step <= 0: + return False + from core.training.resume import get_resume_checkpoint_path + return get_resume_checkpoint_path(output_dir, expected_step = step) is not None + + def _terminal_finalize_kwargs(self) -> dict: + with self._lock: + job_id = self.current_job_id + payload = self._terminal_finalize_payload + if payload and payload.get("expected_job_id") == job_id: + return dict(payload) + cancel, stopped = self._cancel_requested, self._should_stop + output_dir = None if cancel else self._output_dir + step = self._progress.step + existing_error = self._progress.error + status, error, blocked = ( + ("stopped", None, cancel) + if stopped + else ( + "error", + existing_error or "Training process terminated unexpectedly", + False, + ) + ) + # Block only when no valid current-step checkpoint actually landed. + if stopped and not cancel and not self._has_current_resume_checkpoint(output_dir, step): + status = "error" + error = "Stop and Save ended before a valid current-step checkpoint was written." + blocked = True + return { + "status": status, + "error_message": error, + "output_dir": output_dir, + "clear_output_dir": cancel, + "resume_blocked": blocked, + "expected_job_id": job_id, + } + def _handle_event(self, event: dict) -> None: """Apply a subprocess event to local state. @@ -1764,6 +1903,15 @@ class TrainingBackend: elif etype == "eval_configured": self.eval_enabled = True + elif etype == "output_dir": + event_output_dir = event.get("output_dir") + if self._cancel_requested: + self._cancel_cleanup_output_dir = event_output_dir + self._output_dir = self._progress.output_dir = None + else: + self._output_dir = event_output_dir + db_action = "persist_output_dir" + elif etype == "status": self._progress.status_message = event.get("message", "") self._progress.is_training = True @@ -1778,7 +1926,12 @@ class TrainingBackend: self._complete_seen.set() self._progress.is_training = False self._progress.is_completed = not stopped - self._output_dir = event.get("output_dir") + event_output_dir = event.get("output_dir") + if self._cancel_requested: + self._cancel_cleanup_output_dir = event_output_dir + self._output_dir = None + else: + self._output_dir = event_output_dir self._progress.output_dir = self._output_dir self._progress.status_message = msg if not self._db_run_created and self.current_job_id and self._db_config: @@ -1788,11 +1941,16 @@ class TrainingBackend: db_action_kwargs = { "status": "stopped" if stopped else "completed", "output_dir": self._output_dir, + "clear_output_dir": self._cancel_requested, + "expected_job_id": self.current_job_id, } + self._terminal_finalize_payload = dict(db_action_kwargs) elif etype == "error": self._progress.is_training = False self._progress.error = event.get("error", "Unknown error") + if self._cancel_requested: + self._output_dir = self._progress.output_dir = None logger.error("Training error: %s", event.get("error")) stack = event.get("stack", "") if stack: @@ -1801,29 +1959,36 @@ class TrainingBackend: db_action = "create_and_finalize" else: db_action = "finalize" + stop_save_failed = ( + self._should_stop + and not self._cancel_requested + and not self._has_current_resume_checkpoint( + self._output_dir, self._progress.step + ) + ) db_action_kwargs = { - "status": "stopped" if self._should_stop else "error", + "status": "stopped" + if self._should_stop + and not stop_save_failed + and not event.get("keep_error_status") + else "error", "error_message": event.get("error", "Unknown error"), + "output_dir": self._output_dir, + "clear_output_dir": self._cancel_requested, + "resume_blocked": stop_save_failed or bool(event.get("resume_blocked")), + "expected_job_id": self.current_job_id, } + self._terminal_finalize_payload = dict(db_action_kwargs) # --- DB I/O outside the lock --- if db_action == "create_run": - try: - from storage.studio_db import create_run - - create_run( - id = db_action_kwargs["job_id"], - model_name = db_action_kwargs["model_name"], - dataset_name = db_action_kwargs["dataset_name"], - config_json = db_action_kwargs["config_json"], - started_at = db_action_kwargs["started_at"], - total_steps = db_action_kwargs["total_steps"], - ) - self._db_run_created = True + self._ensure_db_run_created() + if self._db_run_created: if db_action_kwargs["total_steps"]: self._db_total_steps_set = True - except Exception: - logger.warning("Failed to create DB run record", exc_info = True) + self._persist_output_dir() + elif db_action == "persist_output_dir": + self._persist_output_dir() elif db_action == "create_and_finalize": self._ensure_db_run_created() self._finalize_run_in_db(**db_action_kwargs) @@ -1842,6 +2007,22 @@ class TrainingBackend: if etype == "progress": self._log_training_progress() + def _persist_output_dir(self) -> None: + with self._lock: + if ( + not self._output_dir + or not self.current_job_id + or not self._db_run_created + or self._cancel_requested + ): + return + run_id, output_dir = self.current_job_id, self._output_dir + try: + from storage.studio_db import update_run_output_dir + update_run_output_dir(run_id, output_dir) + except Exception: + logger.warning("Failed to persist output_dir", exc_info = True) + def _log_training_progress(self) -> None: """One throttled training-status line to the server log (the per-step stream still goes to the UI via SSE): first step, then at most every 30s, plus the @@ -1875,6 +2056,7 @@ class TrainingBackend: caller create at a time, and ``_db_run_created`` is published only after ``create_run`` commits, so a concurrent finalize never runs ``finish_run`` against a not-yet-inserted row (a zero-row UPDATE that would leave the run stuck as running).""" + self._run_intent_lock.acquire() with self._lock: if ( self._db_run_created @@ -1882,6 +2064,7 @@ class TrainingBackend: or not self.current_job_id or not self._db_config ): + self._run_intent_lock.release() return self._db_create_in_progress = True # only one caller creates job_id = self.current_job_id @@ -1898,6 +2081,12 @@ class TrainingBackend: or _s3_dataset_name(db_config.get("s3_dataset")) or "unknown" ) + with self._lock: + if self.current_job_id != job_id: + return + output_dir = self._output_dir + cancel_requested = self._cancel_requested + resumed_from_run_id = self._resume_source_run_id create_run( id = job_id, model_name = db_config["model_name"], @@ -1905,6 +2094,9 @@ class TrainingBackend: config_json = _json.dumps(db_config), started_at = started_at, total_steps = total_steps, + output_dir = output_dir, + cancel_requested = cancel_requested, + resumed_from_run_id = resumed_from_run_id, ) created = True except Exception: @@ -1919,12 +2111,15 @@ class TrainingBackend: if created: self._db_run_created = True # publish only after the insert commits self._db_create_in_progress = False + self._run_intent_lock.release() def _finalize_run_in_db( self, status: str, error_message: Optional[str] = None, output_dir: Optional[str] = None, + clear_output_dir: bool = False, + resume_blocked: bool = False, expected_job_id: Optional[str] = None, ) -> None: """Flush remaining metrics and mark a run finished in the DB. Claims the finalize @@ -1947,26 +2142,33 @@ class TrainingBackend: duration = self._progress.elapsed_seconds loss_history = list(self.loss_history) self._flush_metrics_to_db(run_id = run_id) - try: - from storage.studio_db import finish_run - from utils.downsample import downsample + for attempt in range(_DB_FINALIZE_RETRIES): + try: + from storage.studio_db import finish_run + from utils.downsample import downsample - sparkline = downsample(loss_history, 50) - finish_run( - id = run_id, - status = status, - ended_at = datetime.now(timezone.utc).isoformat(), - final_step = final_step, - final_loss = final_loss, - duration_seconds = duration, - loss_sparkline = _json.dumps(sparkline), - output_dir = output_dir, - error_message = error_message, - ) - except Exception: - with self._lock: - self._run_finalized = False # unclaim so a later flush can retry - logger.warning("Failed to finalize run in DB (status=%s)", status, exc_info = True) + finish_run( + id = run_id, + status = status, + ended_at = datetime.now(timezone.utc).isoformat(), + final_step = final_step, + final_loss = final_loss, + duration_seconds = duration, + loss_sparkline = _json.dumps(downsample(loss_history, 50)), + output_dir = output_dir, + error_message = error_message, + clear_output_dir = clear_output_dir, + resume_blocked = resume_blocked, + ) + return + except Exception: + if attempt + 1 < _DB_FINALIZE_RETRIES: + time.sleep(_DB_FINALIZE_RETRY_S) + continue + with self._lock: + if self.current_job_id == run_id: + self._run_finalized = False + logger.warning("Failed to finalize run in DB (status=%s)", status, exc_info = True) def _flush_metrics_to_db(self, run_id: Optional[str] = None) -> None: """Flush buffered metrics to the DB and update live progress. The target run id, diff --git a/studio/backend/core/training/worker.py b/studio/backend/core/training/worker.py index 111f4fdd0f..2df4fa58c6 100644 --- a/studio/backend/core/training/worker.py +++ b/studio/backend/core/training/worker.py @@ -1840,8 +1840,15 @@ def _run_mlx_training(event_queue, stop_queue, config): # Resolve to ~/.unsloth/studio/outputs/ so the export page finds it from utils.paths import ensure_dir - output_dir = _resolve_mlx_output_dir(config, model_name) + # Resume must land in the original run dir even when config lacks output_dir. + resume_dir = config.get("output_dir", "") or _output_dir_from_resume_checkpoint( + resume_from_checkpoint + ) + output_dir = _resolve_mlx_output_dir( + {**config, "output_dir": resume_dir} if resume_dir else config, model_name + ) ensure_dir(Path(output_dir)) + _emit_output_dir(event_queue, output_dir) # ── 6. Create trainer ── eval_steps_val = config.get("eval_steps", 0) or 0 @@ -2067,6 +2074,17 @@ def _run_mlx_training(event_queue, stop_queue, config): trainer.add_eval_callback(_on_eval) + _opt_ref = [None] + _orig_build_optimizer = getattr(trainer, "_build_optimizer", None) + + if callable(_orig_build_optimizer): + + def _capture_optimizer(total_steps): + _opt_ref[0] = _orig_build_optimizer(total_steps) + return _opt_ref[0] + + trainer._build_optimizer = _capture_optimizer + # ── 11. Run training ── gc.collect() mx.synchronize() @@ -2082,31 +2100,58 @@ def _run_mlx_training(event_queue, stop_queue, config): trainer.save_model = _save_model # ── 12. Save and finalize ── - if trainer.stop_requested: - if not _stop_save[0]: - # Cancel (save=False): skip saving. - _send("complete", output_dir = None, status_message = "Training cancelled") + def _finish_tracking() -> None: + # Runs on every save/finalize exit so TB/W&B never leak on early return. + if tb_writer is not None: + try: + tb_writer.close() + except Exception: + pass + if wandb_run is not None: + try: + wandb_run.finish() + except Exception: + pass + + def _stop_checkpoint_ok() -> bool: + if _write_mlx_stop_checkpoint(trainer, _opt_ref[0], output_dir): + return True + _send( + "error", + error = ( + "Failed to save a resumable checkpoint after stop. " + "Model files were saved, but this run cannot be resumed." + ), + # A user stop finalizes as 'stopped'; keep this failure's error status so history explains it. + keep_error_status = True, + # Older checkpoints are stale; resuming would roll back past this stop. + resume_blocked = True, + ) + return False + + try: + if trainer.stop_requested: + if not _stop_save[0]: + # Cancel (save=False): skip saving. + _send("complete", output_dir = None, status_message = "Training cancelled") + else: + _send("status", status_message = "Saving stopped model...") + mx.synchronize() + trainer.save_model(output_dir) + # Stop-and-save promises a resumable checkpoint, not just model files. + if not _stop_checkpoint_ok(): + return + _send("complete", output_dir = output_dir, status_message = "Training stopped") else: - _send("status", status_message = "Saving stopped model...") + _send("status", status_message = "Saving model...") mx.synchronize() trainer.save_model(output_dir) - _send("complete", output_dir = output_dir, status_message = "Training stopped") - else: - _send("status", status_message = "Saving model...") - mx.synchronize() - trainer.save_model(output_dir) - _send("complete", output_dir = output_dir, status_message = "Training completed") - - if tb_writer is not None: - try: - tb_writer.close() - except Exception: - pass - if wandb_run is not None: - try: - wandb_run.finish() - except Exception: - pass + # A save-stop can race the natural final save; it made the same promise. + if trainer.stop_requested and _stop_save[0] and not _stop_checkpoint_ok(): + return + _send("complete", output_dir = output_dir, status_message = "Training completed") + finally: + _finish_tracking() def _is_current_process_apple_silicon() -> bool: @@ -3177,6 +3222,7 @@ def run_training_process(*, event_queue: Any, stop_queue: Any, config: dict) -> ) output_dir = str(resolve_output_dir(output_dir)) ensure_dir(Path(output_dir)) + _emit_output_dir(event_queue, output_dir) tensorboard_dir = config.get("tensorboard_dir") if config.get("enable_tensorboard", False): @@ -3296,6 +3342,61 @@ def _send_status(event_queue: Any, message: str) -> None: ) +def _emit_output_dir(event_queue: Any, output_dir: str) -> None: + try: + event_queue.put({"type": "output_dir", "output_dir": output_dir, "ts": time.time()}) + except Exception: + pass + + +def _mlx_has_checkpoint_at_step(output_dir, step: int) -> bool: + if step <= 0: + return False + from core.training.resume import is_resume_checkpoint_valid + return is_resume_checkpoint_valid( + Path(output_dir) / f"checkpoint-{step}", expected_step = step, backend = "mlx" + ) + + +def _write_mlx_stop_checkpoint(trainer, optimizer, output_dir) -> bool: + """Write a full resume checkpoint for a stopped MLX run. + + Returns True when a checkpoint for the current training step exists. + """ + step = int(getattr(trainer, "_global_step", 0) or 0) + # A periodic save or a resumed run may already cover the current step. + if _mlx_has_checkpoint_at_step(output_dir, step): + return True + if step <= 0 or optimizer is None: + return False + ckpt_dir = Path(output_dir) / f"checkpoint-{step}" + if ckpt_dir.is_symlink(): + # Refuse a symlinked dir: it could redirect writes outside output_dir. + logger.error("Refusing to write MLX stop checkpoint through symlink: %s", ckpt_dir) + return False + try: + ckpt_dir.mkdir(parents = True, exist_ok = True) + from unsloth_zoo.mlx.utils import ( + save_optimizer_state, + save_trainable_adapters, + save_trainer_state, + ) + + save_trainable_adapters(trainer.model, str(ckpt_dir)) + save_optimizer_state(optimizer, str(ckpt_dir)) + save_trainer_state( + { + "global_step": step, + "train_loss_history": list(getattr(trainer, "_train_loss_history", [])), + }, + str(ckpt_dir), + ) + logger.info("Saved stop checkpoint to %s", ckpt_dir) + except Exception: + logger.exception("Failed to write stop checkpoint under %s", output_dir) + return _mlx_has_checkpoint_at_step(output_dir, step) + + def _run_embedding_training(event_queue: Any, stop_queue: Any, config: dict) -> None: """Self-contained embedding model training pipeline. @@ -3660,6 +3761,7 @@ def _run_embedding_training(event_queue: Any, stop_queue: Any, config: dict) -> config.get("project_name"), ) output_dir = str(resolve_output_dir(output_dir)) + _emit_output_dir(event_queue, output_dir) num_epochs = config.get("num_epochs", 2) batch_size = config.get("batch_size", 256) diff --git a/studio/backend/hub/routes/__init__.py b/studio/backend/hub/routes/__init__.py index e9579635b0..7c5cfb9b3c 100644 --- a/studio/backend/hub/routes/__init__.py +++ b/studio/backend/hub/routes/__init__.py @@ -5,8 +5,10 @@ from hub.routes.inventory import router as inventory_router from hub.routes.datasets import router as datasets_router +from hub.routes.token import router as token_router __all__ = [ "inventory_router", "datasets_router", + "token_router", ] diff --git a/studio/backend/hub/routes/inventory.py b/studio/backend/hub/routes/inventory.py index 4b6c179a2b..1ffadf0544 100644 --- a/studio/backend/hub/routes/inventory.py +++ b/studio/backend/hub/routes/inventory.py @@ -28,6 +28,7 @@ from hub.schemas.inventory import ( CachedModelsResponse, DeleteCachedModelResponse, GgufVariantsResponse, + HiddenModelsResponse, LocalModelListResponse, ModelsFolderResponse, RecommendedFoldersResponse, @@ -214,6 +215,16 @@ async def list_cached_models( return await cache_inventory.list_cached_models_response(hf_token) +@router.get("/hidden-models", response_model = HiddenModelsResponse) +async def list_hidden_models(current_subject: str = Depends(get_current_subject)): + import asyncio + + from routes.models import hidden_model_matchers + + needles, exact_ids, exact_paths = await asyncio.to_thread(hidden_model_matchers) + return HiddenModelsResponse(needles = needles, exact_ids = exact_ids, exact_paths = exact_paths) + + @router.delete( "/delete-cached", response_model = DeleteCachedModelResponse, diff --git a/studio/backend/hub/routes/token.py b/studio/backend/hub/routes/token.py new file mode 100644 index 0000000000..1b7ad733a2 --- /dev/null +++ b/studio/backend/hub/routes/token.py @@ -0,0 +1,44 @@ +# SPDX-License-Identifier: AGPL-3.0-only +# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0 + +"""Hugging Face token validation endpoint.""" + +from __future__ import annotations + +import asyncio +from typing import Literal, Optional + +from fastapi import APIRouter, Depends, Request +from pydantic import BaseModel + +from auth.authentication import get_current_subject +from hub.dependencies import get_hf_token +from utils.client_ip import client_ip +from utils.hf_token_validation import validate_hf_token + + +router = APIRouter() + + +class HfTokenValidationResponse(BaseModel): + status: Literal["missing", "valid", "invalid", "rate_limited", "unavailable"] + retry_after_seconds: Optional[int] = None + + +@router.post("/token/validate", response_model = HfTokenValidationResponse) +async def validate_token( + request: Request, + hf_token: Optional[str] = Depends(get_hf_token), + current_subject: str = Depends(get_current_subject), +): + if not hf_token: + return HfTokenValidationResponse(status = "missing") + result = await asyncio.to_thread( + validate_hf_token, + hf_token, + rate_key = f"{current_subject}:{client_ip(request)}", + ) + return HfTokenValidationResponse( + status = result.status, + retry_after_seconds = result.retry_after_seconds, + ) diff --git a/studio/backend/hub/schemas/inventory.py b/studio/backend/hub/schemas/inventory.py index ef95efe2f2..19d6da3e11 100644 --- a/studio/backend/hub/schemas/inventory.py +++ b/studio/backend/hub/schemas/inventory.py @@ -160,6 +160,7 @@ class CachedRepoBase(BaseModel): repo_id: str size_bytes: int = 0 cache_path: Optional[str] = None + last_modified: Optional[float] = None partial: bool = False partial_transport: Optional[str] = None inventory_id: Optional[str] = None @@ -189,6 +190,12 @@ class CachedModelsResponse(BaseModel): cached: List[CachedModelRepo] = Field(default_factory = list) +class HiddenModelsResponse(BaseModel): + needles: List[str] = Field(default_factory = list) + exact_ids: List[str] = Field(default_factory = list) + exact_paths: List[str] = Field(default_factory = list) + + class AddScanFolderRequest(BaseModel): """Request body for adding a custom scan folder.""" diff --git a/studio/backend/hub/services/models/cache_inventory.py b/studio/backend/hub/services/models/cache_inventory.py index 54a25482f2..c1b864bb63 100644 --- a/studio/backend/hub/services/models/cache_inventory.py +++ b/studio/backend/hub/services/models/cache_inventory.py @@ -31,6 +31,7 @@ from hub.services.models.common import ( _is_checkpoint_weight_name, _is_gguf_filename, _is_main_gguf_filename, + _is_mmproj_filename, _is_transformers_safetensors_weight_name, _local_inventory_id, _prefer_complete_larger, @@ -132,6 +133,39 @@ def _repo_has_gguf_files(repo_info) -> bool: return _repo_gguf_size_bytes(repo_info) > 0 +def _blob_mtime(file_obj) -> float: + ts = getattr(file_obj, "blob_last_modified", None) + if isinstance(ts, (int, float)) and ts > 0: + return float(ts) + blob_path = getattr(file_obj, "blob_path", None) + if blob_path: + try: + return float(Path(blob_path).stat().st_mtime) + except OSError: + pass + return 0.0 + + +def _repo_gguf_last_modified(repo_info) -> float: + latest = 0.0 + for revision in repo_info.revisions: + for f in revision.files: + if _is_main_gguf_filename(f.file_name): + latest = max(latest, _blob_mtime(f)) + return latest + + +def _repo_has_mmproj(repo_info) -> bool: + # An mmproj file only makes a repo vision-capable when it is an actual GGUF + # projector; a non-GGUF sidecar (e.g. mmproj_config.json) does not, and the + # runtime's projector detection is GGUF-only. + return any( + _is_gguf_filename(f.file_name) and _is_mmproj_filename(f.file_name) + for revision in repo_info.revisions + for f in revision.files + ) + + def _cached_repo_file_name(file_obj) -> str: file_path = getattr(file_obj, "file_path", None) if file_path: @@ -291,6 +325,7 @@ def _scan_cached_gguf() -> list[dict]: continue key = repo_id.lower() existing = seen_lower.get(key) + last_modified = _repo_gguf_last_modified(repo_info) row = { "repo_id": repo_id, "size_bytes": max(total_size, variant_state_size), @@ -300,6 +335,9 @@ def _scan_cached_gguf() -> list[dict]: # per-variant detail lives on GgufVariantDetail. "partial_transport": None, } + last_modified = max(last_modified, (existing or {}).get("last_modified", 0.0)) + if last_modified > 0: + row["last_modified"] = last_modified row.update( _cache_inventory_fields( repo_id, @@ -308,11 +346,20 @@ def _scan_cached_gguf() -> list[dict]: requires_variant = True, ) ) + if _repo_has_mmproj(repo_info): + row["capabilities"]["supports_vision"] = True # Visible infra variants remain management-only. if is_hidden_infra: row["capabilities"]["can_chat"] = False if _prefer_cache_row(row, existing): + if existing and existing["capabilities"].get("supports_vision"): + row["capabilities"]["supports_vision"] = True seen_lower[key] = row + else: + if last_modified > existing.get("last_modified", 0.0): + existing["last_modified"] = last_modified + if row["capabilities"].get("supports_vision"): + existing["capabilities"]["supports_vision"] = True except Exception as e: repo_label = getattr(repo_info, "repo_id", "") logger.warning(f"Skipping cached GGUF repo {repo_label}: {e}") @@ -340,13 +387,14 @@ class _CachedNonGgufPayload(NamedTuple): size_bytes: int has_runnable_weights: bool model_format: ModelFormat + last_modified: float def _repo_non_gguf_model_payload(repo_info) -> _CachedNonGgufPayload: - all_weight_blobs: dict[str, int] = {} - adapter_blobs: dict[str, int] = {} - safetensors_blobs: dict[str, int] = {} - checkpoint_blobs: dict[str, int] = {} + all_weight_blobs: dict[str, tuple[int, float]] = {} + adapter_blobs: dict[str, tuple[int, float]] = {} + safetensors_blobs: dict[str, tuple[int, float]] = {} + checkpoint_blobs: dict[str, tuple[int, float]] = {} has_config = False has_adapter_config = False has_adapter_weights = False @@ -354,12 +402,15 @@ def _repo_non_gguf_model_payload(repo_info) -> _CachedNonGgufPayload: has_transformers_safetensors = False has_checkpoint = False - def _record_blob(target: dict[str, int], file_obj, rev_id: str, file_name: str) -> None: + def _record_blob( + target: dict[str, tuple[int, float]], file_obj, rev_id: str, file_name: str + ) -> None: blob_path = getattr(file_obj, "blob_path", None) size = int(file_obj.size_on_disk or 0) key = str(blob_path) if blob_path else f"{rev_id}:{file_name}" - target[key] = size - all_weight_blobs[key] = size + value = (size, _blob_mtime(file_obj)) + target[key] = value + all_weight_blobs[key] = value for revision in repo_info.revisions: rev_id = getattr(revision, "commit_hash", None) or str(id(revision)) @@ -403,18 +454,19 @@ def _repo_non_gguf_model_payload(repo_info) -> _CachedNonGgufPayload: or "unknown" ) if model_format == "adapter": - size_bytes = sum(adapter_blobs.values()) + selected_blobs = adapter_blobs elif model_format == "safetensors": - size_bytes = sum(safetensors_blobs.values()) + selected_blobs = safetensors_blobs elif model_format == "checkpoint": - size_bytes = sum(checkpoint_blobs.values()) + selected_blobs = checkpoint_blobs else: - size_bytes = sum(all_weight_blobs.values()) + selected_blobs = all_weight_blobs return _CachedNonGgufPayload( - size_bytes = size_bytes, + size_bytes = sum(size for size, _mtime in selected_blobs.values()), has_runnable_weights = model_format != "unknown", model_format = model_format, + last_modified = max((mtime for _size, mtime in selected_blobs.values()), default = 0.0), ) @@ -544,6 +596,12 @@ def _scan_cached_models() -> list[dict]: ), **_cached_model_local_metadata(repo_path), } + last_modified = max( + payload.last_modified, + (existing or {}).get("last_modified", 0.0), + ) + if last_modified > 0: + row["last_modified"] = last_modified row.update( _cache_inventory_fields( repo_id, @@ -553,6 +611,8 @@ def _scan_cached_models() -> list[dict]: ) if _prefer_cache_row(row, existing): seen_lower[key] = row + elif last_modified > existing.get("last_modified", 0.0): + existing["last_modified"] = last_modified except Exception as e: repo_label = getattr(repo_info, "repo_id", "") logger.warning(f"Skipping cached model repo {repo_label}: {e}") diff --git a/studio/backend/main.py b/studio/backend/main.py index 4797764ce7..3f244dc22e 100644 --- a/studio/backend/main.py +++ b/studio/backend/main.py @@ -289,6 +289,7 @@ from fastapi import Depends, FastAPI, HTTPException, Query, Request from fastapi.middleware.cors import CORSMiddleware from fastapi.staticfiles import StaticFiles from fastapi.responses import FileResponse, HTMLResponse, Response +from starlette.middleware.gzip import GZipMiddleware from pathlib import Path from datetime import datetime @@ -312,7 +313,9 @@ from routes.preview import router as preview_router from hub.routes import ( inventory_router as hub_inventory_router, datasets_router as hub_datasets_router, + token_router as hub_token_router, ) +from picker.routes import templates_router as picker_templates_router from hub.schemas.downloads import TransportCapabilities from hub.utils.download_registry import ( get_download_transport_capabilities, @@ -547,8 +550,9 @@ async def lifespan(app: FastAPI): threading.Thread(target = _warm_rag_embedder, daemon = True, name = "rag-embedder-warm").start() # Idle auto-unload loop (no-op unless the OpenAI auto-unload TTL is set). - from core.inference.llama_keepwarm import idle_unload_loop + from core.inference.llama_keepwarm import idle_unload_loop, sweep_slot_save_dir + sweep_slot_save_dir() app.state.idle_unload_task = asyncio.create_task(idle_unload_loop()) # Initialize RSA key pair for API key encryption (external providers). @@ -761,6 +765,7 @@ _BODY_PROTECTED_PREFIXES = ( "/v1/completions", "/p/", "/api/inference", + "/api/picker", "/api/data-recipe", "/api/datasets", "/api/hub", @@ -992,6 +997,8 @@ app.include_router(rag_router, prefix = "/api/rag", tags = ["rag"]) app.include_router(training_history_router, prefix = "/api/train", tags = ["training-history"]) app.include_router(hub_inventory_router, prefix = "/api/hub", tags = ["hub"]) app.include_router(hub_datasets_router, prefix = "/api/hub/datasets", tags = ["hub"]) +app.include_router(picker_templates_router, prefix = "/api/picker", tags = ["picker"]) +app.include_router(hub_token_router, prefix = "/api/hub", tags = ["hub"]) # Re-wrap client-error responses on the /v1/* surface into OpenAI/Anthropic # error envelopes; non-/v1 paths keep FastAPI's default {"detail": ...} shape. @@ -1148,17 +1155,35 @@ def _get_cached_system_gpu_info(logger) -> dict[str, Any]: util = util_devices.get(idx, {}) total_vram = util.get("vram_total_gb") or dev.get("memory_total_gb") or 0 - used_vram = util.get("vram_used_gb") or 0 + # Keep None (usage unknown, e.g. Windows ROCm perf counter) so the UI + # shows unknown, not a fabricated 0 used / full free. + used_vram = util.get("vram_used_gb") enriched_dev = dict(dev) enriched_dev["vram_used_gb"] = used_vram - enriched_dev["vram_free_gb"] = round(total_vram - used_vram, 2) if total_vram else 0 + enriched_dev["vram_free_gb"] = ( + round(total_vram - used_vram, 2) if total_vram and used_vram is not None else None + ) enriched_dev["vram_utilization_pct"] = util.get("vram_utilization_pct") enriched_devices.append(enriched_dev) + # Whether GGUF loads accept an explicit gpu_ids pick: /load and + # /validate 400 picks on XPU hosts (no visibility mask speaks torch-xpu + # ordinals) and on Vulkan-only builds (--device pins ggml's own + # ordinals), so the picker must not offer them. + try: + from core.inference.llama_cpp import LlamaCppBackend + from utils.hardware import DeviceType, get_device + gpu_ids_supported = ( + get_device() != DeviceType.XPU and not LlamaCppBackend._is_vulkan_backend() + ) + except Exception as e: + logger.debug(f"Could not resolve gpu_ids support: {e}") + gpu_ids_supported = True gpu_info = { "available": visibility_info.get("available", False), "devices": enriched_devices, + "gguf_gpu_ids_supported": gpu_ids_supported, } _system_gpu_cache = (time.monotonic(), gpu_info) return gpu_info @@ -1488,6 +1513,34 @@ def _should_inject_bootstrap(request: Request) -> bool: return _is_local_bootstrap_request(request) +_IMMUTABLE_ASSET_CACHE_CONTROL = "public, max-age=31536000, immutable" + + +class ImmutableStaticFiles(StaticFiles): + """Serve Vite's content-hashed assets without browser revalidation.""" + + def file_response( + self, + full_path, + stat_result, + scope, + status_code = 200, + ): + response = super().file_response(full_path, stat_result, scope, status_code) + response.headers["Cache-Control"] = _IMMUTABLE_ASSET_CACHE_CONTROL + return response + + +class _AssetGZipMiddleware(GZipMiddleware): + """Serve range requests uncompressed; gzip + 206 mislabels Content-Range.""" + + async def __call__(self, scope, receive, send): + if scope["type"] == "http" and any(key == b"range" for key, _ in scope["headers"]): + await self.app(scope, receive, send) + return + await super().__call__(scope, receive, send) + + def setup_frontend(app: FastAPI, build_path: Path): """Mount frontend static files (optional)""" if not build_path.exists(): @@ -1495,7 +1548,12 @@ def setup_frontend(app: FastAPI, build_path: Path): assets_dir = build_path / "assets" if assets_dir.exists(): - app.mount("/assets", StaticFiles(directory = assets_dir), name = "assets") + assets_app = _AssetGZipMiddleware( + ImmutableStaticFiles(directory = assets_dir), + minimum_size = 1024, + compresslevel = 6, + ) + app.mount("/assets", assets_app, name = "assets") def _build_index_response(request: Request) -> Response: content = (build_path / "index.html").read_bytes() diff --git a/studio/backend/models/inference.py b/studio/backend/models/inference.py index f3ae0f70df..580a74dddf 100644 --- a/studio/backend/models/inference.py +++ b/studio/backend/models/inference.py @@ -18,6 +18,8 @@ from pydantic import ( model_validator, ) +from picker.schemas import MAX_CHAT_TEMPLATE_BYTES + class LoadRequest(BaseModel): """Request to load a model for inference""" @@ -54,8 +56,16 @@ class LoadRequest(BaseModel): @field_validator("chat_template_override") @classmethod def normalize_blank_chat_template_override(cls, value: Optional[str]) -> Optional[str]: - if value is not None and value.strip() == "": + if value is None: return None + # Char count is a lower bound on UTF-8 byte length: reject an oversized + # template before spending work encoding it. + if len(value) > MAX_CHAT_TEMPLATE_BYTES: + raise ValueError(f"Chat template exceeds the {MAX_CHAT_TEMPLATE_BYTES}-byte limit.") + if value.strip() == "": + return None + if len(value.encode("utf-8")) > MAX_CHAT_TEMPLATE_BYTES: + raise ValueError(f"Chat template exceeds the {MAX_CHAT_TEMPLATE_BYTES}-byte limit.") return value cache_type_kv: Optional[str] = Field( @@ -64,7 +74,7 @@ class LoadRequest(BaseModel): ) gpu_ids: Optional[List[int]] = Field( None, - description = "Physical GPU indices to use, for example [0, 1]. Omit or pass [] to use automatic selection. Explicit gpu_ids are unsupported when the parent CUDA_VISIBLE_DEVICES uses UUID/MIG entries. Not supported for GGUF models.", + description = "Physical GPU indices to use, for example [0, 1]. Omit or pass [] to use automatic selection. Explicit gpu_ids are unsupported when the parent CUDA_VISIBLE_DEVICES uses UUID/MIG entries. For GGUF models the picked devices are pinned via CUDA/HIP_VISIBLE_DEVICES.", ) speculative_type: Optional[str] = Field( None, @@ -100,6 +110,66 @@ class LoadRequest(BaseModel): "No effect on a single GPU. Ignored for non-GGUF models." ), ) + gpu_memory_mode: Literal["auto", "manual"] = Field( + "auto", + description = ( + "GPU memory strategy for GGUF models. 'auto' (default): Unsloth " + "selects GPUs and caps context to fit VRAM. 'manual': you own the " + "offload. Leave gpu_layers at -1 (Auto) to hand memory management to " + "llama.cpp's --fit (no device masking, no context auto-reduce, no " + "gpu-layer/tensor-split planning); set gpu_layers >= 0 to pin layers " + "and n_cpu_moe yourself (--fit off), with tensor_parallel still " + "applying (split by free VRAM unless tensor_split is set, no planner). " + "Ignored for non-GGUF." + ), + ) + gpu_layers: int = Field( + -1, + ge = -1, + description = ( + "Manual mode only: number of layers to offload to the GPU " + "(--gpu-layers, with --fit off). A value >= the model's layer count " + "offloads all of them. -1 = Auto: hand layer + context sizing to " + "llama.cpp's --fit. Ignored unless gpu_memory_mode is 'manual'." + ), + ) + n_cpu_moe: int = Field( + 0, + ge = 0, + description = ( + "Manual mode only: keep the first N MoE expert layers on the CPU " + "(--n-cpu-moe) to save VRAM on MoE models. 0 = none, N = number of " + "MoE layers offloaded (the backend offsets past any leading dense " + "layers). Ignored unless gpu_memory_mode is 'manual' with gpu_layers >= 0." + ), + ) + tensor_split: Optional[List[float]] = Field( + None, + description = ( + "Manual mode only: relative share of the model per GPU (--tensor-split), " + "in the order of the GPUs in use, e.g. [2, 1] for 2:1. Omit it to let " + "llama.cpp use its default, which splits by free VRAM. Any list given is " + "passed through as-is, so send [1, 1] to force an even split. Ignored " + "unless gpu_memory_mode is 'manual' with gpu_layers >= 0." + ), + ) + + @field_validator("tensor_split") + @classmethod + def _reject_degenerate_tensor_split(cls, value: Optional[List[float]]) -> Optional[List[float]]: + # A negative / non-finite / all-zero split is silently dropped at launch + # (stored as None) yet still compared raw in the reload dedupe, so an + # identical Apply reloads forever. Reject it up front; [] = no split. + if not value: + return value + import math + + if any((not math.isfinite(v)) or v < 0 for v in value): + raise ValueError("tensor_split entries must be finite and non-negative") + if sum(value) <= 0: + raise ValueError("tensor_split must have a positive total") + return value + llama_extra_args: Optional[List[str]] = Field( None, description = ( @@ -133,11 +203,26 @@ class ValidateModelRequest(BaseModel): max_seq_length: int = Field(0, ge = 0, le = 1048576) load_in_4bit: bool = Field(True) gpu_ids: Optional[List[int]] = Field(None) + gpu_memory_mode: Literal["auto", "manual"] = Field( + "auto", + description = ( + "GGUF GPU-memory strategy intended for the follow-up load. Manual " + "placement bypasses the training coexistence estimate: Auto layers " + "delegate fitting to llama.cpp, while explicit layers are user-owned." + ), + ) include_context_length: bool = Field( False, description = "Also read the native context length from the local GGUF header. " "Opt-in so the normal load preflight doesn't pay for a cache scan it doesn't need.", ) + include_chat_template: bool = Field( + False, + description = "Also read the embedded chat template from the local GGUF header, so a " + "native (picked / drag-drop) file's default template can be shown before it is loaded. " + "Opt-in and, like include_context_length, a metadata-only probe that skips the training " + "guard. Only the leased file's own embedded template is read, never sibling sidecars.", + ) class TransformersUpgradeInfo(BaseModel): @@ -188,6 +273,21 @@ class ValidateModelResponse(BaseModel): description = "Native training context length, read from the GGUF header when the file " "is already downloaded locally; None for non-GGUF, gated, or not-yet-downloaded models.", ) + layer_count: Optional[int] = Field( + None, + description = "Total layer count (GGUF block_count), the manual gpu-layers ceiling, read " + "from the header alongside context_length; None when not read.", + ) + moe_layer_count: Optional[int] = Field( + None, + description = "MoE expert-layer count (the manual --n-cpu-moe ceiling), read from the GGUF " + "header alongside context_length; 0 for dense models, None when not read.", + ) + chat_template: Optional[str] = Field( + None, + description = "Embedded GGUF chat template, read from the header when include_chat_template " + "is set (native lease-backed picks); None for non-GGUF, over-cap, or not-read templates.", + ) # Additive fields; the consuming consent dialog ships in a follow-up frontend PR. requires_transformers_upgrade: bool = Field( False, @@ -333,6 +433,34 @@ class LoadResponse(BaseModel): False, description = "Whether tensor-parallel split (--split-mode tensor) is active.", ) + gpu_memory_mode: Literal["auto", "manual"] = Field( + "auto", + description = "Active GPU memory strategy ('auto' or 'manual').", + ) + gpu_layers: int = Field( + -1, + description = "Manual mode: requested --gpu-layers value (-1 = Auto/--fit, or when not manual).", + ) + n_cpu_moe: int = Field( + 0, + description = "Manual mode: MoE expert layers pinned to CPU (--n-cpu-moe); 0 = none.", + ) + tensor_split: Optional[List[float]] = Field( + None, + description = "Manual mode: relative model share per GPU (--tensor-split); None = default (split by free VRAM).", + ) + n_layers: Optional[int] = Field( + None, + description = "Model's layer count (GGUF block_count), for the manual gpu-layers ceiling.", + ) + n_moe_layers: int = Field( + 0, + description = "Model's MoE expert-layer count (the n_cpu_moe ceiling); 0 if not an MoE model.", + ) + gpu_ids: Optional[List[int]] = Field( + None, + description = "Physical GPU indices the model is pinned to, or None for automatic selection.", + ) class UnloadResponse(BaseModel): @@ -461,6 +589,42 @@ class InferenceStatusResponse(BaseModel): False, description = "Whether tensor-parallel split (--split-mode tensor) is active.", ) + gpu_memory_mode: Literal["auto", "manual"] = Field( + "auto", + description = "Active GPU memory strategy ('auto' or 'manual').", + ) + gpu_layers: int = Field( + -1, + description = "Manual mode: requested --gpu-layers value (-1 = Auto/--fit, or when not manual).", + ) + n_cpu_moe: int = Field( + 0, + description = "Manual mode: MoE expert layers pinned to CPU (--n-cpu-moe); 0 = none.", + ) + tensor_split: Optional[List[float]] = Field( + None, + description = "Manual mode: relative model share per GPU (--tensor-split); None = default (split by free VRAM).", + ) + requested_context_length: Optional[int] = Field( + None, + description = ( + "The n_ctx the active GGUF load was invoked with (0 = Auto). Lets the " + "UI re-seed a Manual + Auto-layers context pin on hydration, where " + "context_length only exposes the resolved value. None for non-GGUF." + ), + ) + n_layers: Optional[int] = Field( + None, + description = "Model's layer count (GGUF block_count), for the manual gpu-layers ceiling.", + ) + n_moe_layers: int = Field( + 0, + description = "Model's MoE expert-layer count (the n_cpu_moe ceiling); 0 if not an MoE model.", + ) + gpu_ids: Optional[List[int]] = Field( + None, + description = "Physical GPU indices the model is pinned to, or None for automatic selection.", + ) llama_cpp_supports_mtp: bool = Field( True, description = ( diff --git a/studio/backend/models/training.py b/studio/backend/models/training.py index 0b50f63b95..0f88b78f9f 100644 --- a/studio/backend/models/training.py +++ b/studio/backend/models/training.py @@ -505,6 +505,13 @@ class TrainingStartRequest(BaseModel): description = "S3 bucket configuration for loading datasets from AWS S3. Requires boto3 to be installed.", ) + @field_validator("target_modules", mode = "before") + @classmethod + def _normalize_target_modules(cls, value: Any) -> Any: + # Sanitized non-LoRA history stores the unused value as null; treat it as a + # fresh request's omitted/default empty list on resume. + return [] if value is None else value + @model_validator(mode = "after") def _validate_streaming_splits(self) -> "TrainingStartRequest": # Streaming load_dataset does not accept HF slice syntax (e.g. "train[:50%]" diff --git a/studio/backend/picker/__init__.py b/studio/backend/picker/__init__.py new file mode 100644 index 0000000000..32014236c6 --- /dev/null +++ b/studio/backend/picker/__init__.py @@ -0,0 +1,2 @@ +# SPDX-License-Identifier: AGPL-3.0-only +# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0 diff --git a/studio/backend/picker/routes/__init__.py b/studio/backend/picker/routes/__init__.py new file mode 100644 index 0000000000..c0e988c8bb --- /dev/null +++ b/studio/backend/picker/routes/__init__.py @@ -0,0 +1,6 @@ +# SPDX-License-Identifier: AGPL-3.0-only +# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0 + +from .templates import router as templates_router + +__all__ = ["templates_router"] diff --git a/studio/backend/picker/routes/templates.py b/studio/backend/picker/routes/templates.py new file mode 100644 index 0000000000..02b8bf7184 --- /dev/null +++ b/studio/backend/picker/routes/templates.py @@ -0,0 +1,45 @@ +# SPDX-License-Identifier: AGPL-3.0-only +# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0 + +from __future__ import annotations + +import asyncio +from typing import Optional + +from fastapi import APIRouter, Body, Depends, Query + +from auth.authentication import get_current_subject +from hub.dependencies import get_hf_token + +from ..schemas import ( + MAX_CHAT_TEMPLATE_BYTES, + ModelTemplateResponse, + ValidateChatTemplateRequest, + ValidateChatTemplateResponse, +) +from ..service import read_default_chat_template, validate_chat_template + +router = APIRouter() + + +@router.post("/validate-chat-template", response_model = ValidateChatTemplateResponse) +async def validate_chat_template_route( + body: ValidateChatTemplateRequest = Body(...), + current_subject: str = Depends(get_current_subject), +) -> ValidateChatTemplateResponse: + return await asyncio.to_thread(validate_chat_template, body.template) + + +@router.get("/chat-template/{model_name:path}", response_model = ModelTemplateResponse) +async def get_default_chat_template_route( + model_name: str, + gguf_variant: Optional[str] = Query(None), + hf_token: Optional[str] = Depends(get_hf_token), + current_subject: str = Depends(get_current_subject), +) -> ModelTemplateResponse: + template = await asyncio.to_thread( + read_default_chat_template, model_name, hf_token, gguf_variant + ) + if template is not None and len(template.encode("utf-8")) > MAX_CHAT_TEMPLATE_BYTES: + template = None + return ModelTemplateResponse(model_name = model_name, chat_template = template) diff --git a/studio/backend/picker/schemas.py b/studio/backend/picker/schemas.py new file mode 100644 index 0000000000..b4f956188f --- /dev/null +++ b/studio/backend/picker/schemas.py @@ -0,0 +1,32 @@ +# SPDX-License-Identifier: AGPL-3.0-only +# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0 + +from typing import Optional + +from pydantic import BaseModel, Field, field_validator + +# Mirror the frontend's 64 KiB chat-template contract (per-model-config.ts) at +# the API boundary so a direct caller cannot make Jinja parse an oversized +# template. MaxBodyMiddleware only caps the whole request body, not this field. +MAX_CHAT_TEMPLATE_BYTES = 65_536 + + +class ValidateChatTemplateRequest(BaseModel): + template: str = Field(default = "") + + @field_validator("template") + @classmethod + def _enforce_template_size(cls, value: str) -> str: + if len(value.encode("utf-8")) > MAX_CHAT_TEMPLATE_BYTES: + raise ValueError(f"Chat template exceeds the {MAX_CHAT_TEMPLATE_BYTES}-byte limit.") + return value + + +class ValidateChatTemplateResponse(BaseModel): + valid: bool + error: Optional[str] = None + + +class ModelTemplateResponse(BaseModel): + model_name: str + chat_template: Optional[str] = None diff --git a/studio/backend/picker/service.py b/studio/backend/picker/service.py new file mode 100644 index 0000000000..13065b2920 --- /dev/null +++ b/studio/backend/picker/service.py @@ -0,0 +1,426 @@ +# SPDX-License-Identifier: AGPL-3.0-only +# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0 + +from __future__ import annotations + +import json +import logging +import os +import re +from pathlib import Path +from typing import Optional + +from hub.services.models.folder_browser import ( + _build_browse_allowlist, + _is_path_inside_allowlist, +) +from hub.utils.gguf import extract_quant_label, iter_hf_cache_snapshots +from utils.models.gguf_metadata import read_gguf_chat_template +from utils.models.model_config import ( + _extract_quant_label, + _is_big_endian_gguf_path, + _is_mmproj, + _is_mtp_drafter, +) +from utils.paths.path_utils import ( + is_local_path, + normalize_path, + resolve_cached_repo_id_case, +) + +from .schemas import MAX_CHAT_TEMPLATE_BYTES, ValidateChatTemplateResponse + +logger = logging.getLogger(__name__) + +_VALID_REPO_ID = re.compile(r"^[A-Za-z0-9._-]+/[A-Za-z0-9._-]+$") + + +def _is_valid_repo_id(repo_id: str) -> bool: + return bool(_VALID_REPO_ID.fullmatch(repo_id)) + + +_TOKENIZER_CONFIG_PATHS = ("tokenizer_config.json", "LLM/tokenizer_config.json") +_JINJA_TEMPLATE_PATHS = ("chat_template.jinja", "LLM/chat_template.jinja") +_PROCESSOR_TEMPLATE_PATHS = ("chat_template.json", "LLM/chat_template.json") + +# Cap sidecar reads so a malformed or hostile metadata file cannot exhaust memory +# before its template is size-checked. The JSON envelope may exceed a bare template +# (it carries other tokenizer metadata); the extracted template is still bounded by +# MAX_CHAT_TEMPLATE_BYTES downstream. +MAX_TEMPLATE_METADATA_BYTES = 4 * 1024 * 1024 + + +def _read_bounded_text(path: Path, limit: int) -> Optional[str]: + """Read at most `limit` bytes of UTF-8 text; None if larger or unreadable.""" + try: + with path.open("rb") as f: + data = f.read(limit + 1) + except OSError: + return None + if len(data) > limit: + return None + try: + return data.decode("utf-8") + except UnicodeError: + return None + + +def _leaf_inside_allowlist(path: Path, allow_roots: Optional[list[Path]]) -> bool: + # Block symlinked children from escaping the validated directory (realpath-checked). + # None = trusted caller (HF cache / remote download). + return allow_roots is None or _is_path_inside_allowlist(path, allow_roots) + + +def validate_chat_template(template: str) -> ValidateChatTemplateResponse: + text = (template or "").strip() + if not text: + return ValidateChatTemplateResponse(valid = True, error = None) + # Import Jinja lazily: optional at runtime (e.g. GGUF-only installs), so a + # missing dependency must not crash API startup. + try: + from jinja2 import TemplateError + from jinja2.ext import Extension + from jinja2.sandbox import ImmutableSandboxedEnvironment + except ImportError: + return ValidateChatTemplateResponse(valid = True, error = None) + + class _GenerationTag(Extension): + # Accept Transformers' {% generation %} assistant-mask tag so a pasted HF + # chat template validates (we only parse it). + tags = {"generation"} + + def parse(self, parser): + next(parser.stream) + return parser.parse_statements(["name:endgeneration"], drop_needle = True) + + try: + env = ImmutableSandboxedEnvironment( + trim_blocks = True, + lstrip_blocks = True, + extensions = ["jinja2.ext.loopcontrols", _GenerationTag], + ) + env.parse(text) + return ValidateChatTemplateResponse(valid = True, error = None) + except TemplateError as exc: + message = getattr(exc, "message", None) or str(exc) + lineno = getattr(exc, "lineno", None) + if lineno: + message = f"Line {lineno}: {message}" + return ValidateChatTemplateResponse(valid = False, error = message) + except Exception as exc: + return ValidateChatTemplateResponse(valid = False, error = str(exc)) + + +def _chat_template_from_tokenizer_config(config: dict) -> Optional[str]: + if not isinstance(config, dict): + return None + raw = config.get("chat_template") + if isinstance(raw, str) and raw.strip(): + return raw + if isinstance(raw, list): + fallback: Optional[str] = None + for entry in raw: + if not isinstance(entry, dict): + continue + template = entry.get("template") + if not isinstance(template, str): + continue + if entry.get("name") == "default": + return template + if fallback is None: + fallback = template + return fallback + return None + + +def _chat_template_from_jinja_file( + dir_path: Path, allow_roots: Optional[list[Path]] = None +) -> Optional[str]: + for rel in _JINJA_TEMPLATE_PATHS: + template_file = dir_path / rel + if not template_file.exists() or not _leaf_inside_allowlist(template_file, allow_roots): + continue + try: + if template_file.stat().st_size > MAX_CHAT_TEMPLATE_BYTES: + continue + template = template_file.read_text(encoding = "utf-8") + except Exception: + continue + if template.strip(): + return template + return None + + +def _chat_template_from_processor_payload(payload: object) -> Optional[str]: + # processor chat_template.json may be the template string itself or a + # {name: template} map, not only a tokenizer_config-shaped object. + if isinstance(payload, str): + return payload if payload.strip() else None + template = _chat_template_from_tokenizer_config(payload) # type: ignore[arg-type] + if template: + return template + if isinstance(payload, dict): + # Named-template map: prefer "default", else the first non-empty entry + # (mirrors the tokenizer-config list fallback). + default = payload.get("default") + if isinstance(default, str) and default.strip(): + return default + for value in payload.values(): + if isinstance(value, str) and value.strip(): + return value + return None + + +def _chat_template_from_processor_json( + dir_path: Path, allow_roots: Optional[list[Path]] = None +) -> Optional[str]: + for rel in _PROCESSOR_TEMPLATE_PATHS: + config_file = dir_path / rel + if not config_file.exists() or not _leaf_inside_allowlist(config_file, allow_roots): + continue + raw = _read_bounded_text(config_file, MAX_TEMPLATE_METADATA_BYTES) + if raw is None: + continue + try: + payload = json.loads(raw) + except Exception: + continue + template = _chat_template_from_processor_payload(payload) + if template: + return template + return None + + +def _chat_template_from_tokenizer_dir( + dir_path: Path, allow_roots: Optional[list[Path]] = None +) -> Optional[str]: + jinja = _chat_template_from_jinja_file(dir_path, allow_roots) + if jinja: + return jinja + for rel in _TOKENIZER_CONFIG_PATHS: + config_file = dir_path / rel + if not config_file.exists() or not _leaf_inside_allowlist(config_file, allow_roots): + continue + raw = _read_bounded_text(config_file, MAX_TEMPLATE_METADATA_BYTES) + if raw is None: + continue + try: + config = json.loads(raw) + except Exception: + continue + template = _chat_template_from_tokenizer_config(config) + if template: + return template + return _chat_template_from_processor_json(dir_path, allow_roots) + + +_GGUF_SCAN_MAX_DEPTH = 2 + + +def _iter_ggufs(dir_path: Path) -> list[Path]: + if dir_path == dir_path.parent: + return [] + root = str(dir_path) + found: list[Path] = [] + for current, dirs, files in os.walk(root, followlinks = False): + rel = os.path.relpath(current, root) + depth = 0 if rel == os.curdir else rel.count(os.sep) + 1 + if depth >= _GGUF_SCAN_MAX_DEPTH: + dirs[:] = [] + for name in files: + if not name.lower().endswith(".gguf") or _is_mmproj(name): + continue + path = Path(current) / name + try: + rel = path.relative_to(dir_path).as_posix() + except ValueError: + rel = name + quant = _extract_quant_label(rel) + if _is_mtp_drafter(rel) or _is_big_endian_gguf_path(rel, quant): + continue + found.append(path) + return found + + +def _variant_matches(relative_path: str, needle: str) -> bool: + quant = _extract_quant_label(relative_path).lower() + if quant == needle: + return True + if extract_quant_label(relative_path).lower() == needle: + return True + prefix = f"{needle}-" + if not quant.startswith(prefix): + return False + suffix = quant[len(prefix) :] + if not suffix.endswith("bpw"): + return False + value = suffix[:-3] + return bool(value) and value.replace(".", "", 1).isdigit() + + +_GGUF_SPLIT_INDEX_RE = re.compile(r"-(\d{3,})-of-\d{3,}$", re.IGNORECASE) + + +def _is_nonfirst_gguf_split(path: Path) -> bool: + match = _GGUF_SPLIT_INDEX_RE.search(path.stem) + return match is not None and int(match.group(1)) != 1 + + +def _find_gguf_in_dir(dir_path: Path, gguf_variant: Optional[str]) -> Optional[Path]: + try: + ggufs = sorted(_iter_ggufs(dir_path)) + except OSError: + return None + if not ggufs: + return None + needle = (gguf_variant or "").strip().lower() + if needle: + for path in ggufs: + try: + relative = path.relative_to(dir_path).as_posix() + except ValueError: + relative = path.name + if _variant_matches(relative, needle): + return path + return None + candidates = [path for path in ggufs if not _is_nonfirst_gguf_split(path)] or ggufs + try: + return max(candidates, key = lambda path: path.stat().st_size) + except OSError: + return candidates[0] + + +def _chat_template_from_dir( + dir_path: Path, + gguf_variant: Optional[str] = None, + allow_roots: Optional[list[Path]] = None, +) -> Optional[str]: + def from_gguf() -> Optional[str]: + gguf = _find_gguf_in_dir(dir_path, gguf_variant) + if gguf is None or not _leaf_inside_allowlist(gguf, allow_roots): + return None + return read_gguf_chat_template(str(gguf)) + + # Sidecar tokenizer files (chat_template.jinja / tokenizer_config.json) are the + # author's maintained template and supersede the GGUF's possibly-stale embedded + # copy. The variant only picks the GGUF fallback, so tokenizer-first precedence + # holds whether or not a variant is given. + return _chat_template_from_tokenizer_dir(dir_path, allow_roots) or from_gguf() + + +def read_default_chat_template( + model_name: str, + hf_token: Optional[str] = None, + gguf_variant: Optional[str] = None, +) -> Optional[str]: + if not isinstance(model_name, str) or not model_name.strip(): + return None + name = model_name.strip() + + if is_local_path(name): + try: + target = Path(normalize_path(name)).expanduser() + allow_roots = _build_browse_allowlist() + if not _is_path_inside_allowlist(target, allow_roots): + logger.debug("Refused chat template read outside allowed folders: %s", name) + return None + if name.lower().endswith(".gguf"): + # Prefer a maintained sidecar next to the file over the GGUF's + # embedded copy (tokenizer-first precedence, as elsewhere). + sidecar = _chat_template_from_tokenizer_dir(target.parent, allow_roots) + if sidecar: + return sidecar + return read_gguf_chat_template(str(target)) + return _chat_template_from_dir(target, gguf_variant, allow_roots) + except Exception as exc: + logger.debug("Could not read local chat template for %s: %s", name, exc) + return None + + if not _is_valid_repo_id(name): + return None + + resolved = resolve_cached_repo_id_case(name) + + try: + # Resolve within each cached revision, newest first. A revision's sidecar + # supersedes its own embedded GGUF copy, but must not override a newer + # revision, so precedence stays per-snapshot rather than global. + for snapshot in iter_hf_cache_snapshots(resolved): + template = _chat_template_from_dir(snapshot, gguf_variant) + if template: + return template + except Exception as exc: + logger.debug("Could not read cached chat template for %s: %s", resolved, exc) + + try: + from huggingface_hub import HfApi, hf_hub_download + + _api = HfApi() + + def _remote_exceeds_cap(rel: str) -> bool: + # Best-effort: skip the download when the remote's advertised size + # exceeds the cap, so a maliciously large sidecar is never fetched. + try: + infos = _api.get_paths_info(resolved, [rel], repo_type = "model", token = hf_token) + except Exception: + return False + for info in infos: + size = getattr(info, "size", None) + if ( + getattr(info, "path", None) == rel + and isinstance(size, int) + and size > MAX_TEMPLATE_METADATA_BYTES + ): + return True + return False + + def _download_text(rel: str) -> Optional[str]: + if _remote_exceeds_cap(rel): + return None + try: + path = hf_hub_download(resolved, rel, token = hf_token) + return _read_bounded_text(Path(path), MAX_TEMPLATE_METADATA_BYTES) + except Exception: + return None + + for rel in _JINJA_TEMPLATE_PATHS: + template = _download_text(rel) + if not template or not template.strip(): + continue + # A raw Jinja sidecar is the whole template, so it must fit the route's + # response cap (the local path skips oversized .jinja too). Download stays + # bounded at MAX_TEMPLATE_METADATA_BYTES so a large JSON embedding a small + # template still extracts below, but an over-cap Jinja is dropped so the + # search falls through to the tokenizer/processor template. + if len(template.encode("utf-8")) > MAX_CHAT_TEMPLATE_BYTES: + continue + return template + + for rel in _TOKENIZER_CONFIG_PATHS: + raw = _download_text(rel) + if not raw: + continue + try: + config = json.loads(raw) + except Exception: + continue + template = _chat_template_from_tokenizer_config(config) + if template: + return template + + for rel in _PROCESSOR_TEMPLATE_PATHS: + raw = _download_text(rel) + if not raw: + continue + try: + payload = json.loads(raw) + except Exception: + continue + template = _chat_template_from_processor_payload(payload) + if template: + return template + + return None + except Exception as exc: + logger.debug("Could not fetch chat template for %s: %s", resolved, exc) + return None diff --git a/studio/backend/routes/chat_history.py b/studio/backend/routes/chat_history.py index 7a27a58a52..24b6dfb36d 100644 --- a/studio/backend/routes/chat_history.py +++ b/studio/backend/routes/chat_history.py @@ -5,7 +5,7 @@ Chat history API routes backed by studio.db. """ -from typing import Any, Literal, Optional +from typing import Annotated, Any, Literal, Optional from fastapi import APIRouter, Depends, HTTPException, Query from pydantic import BaseModel, ConfigDict, Field, ValidationError @@ -19,13 +19,16 @@ from storage.studio_db import ( clear_chat_history, count_chat_threads, count_forks_for_message, + delete_chat_attachment, delete_chat_threads, delete_chat_project, ensure_chat_project_workspace, fork_chat_thread, + get_chat_attachment, get_chat_project, get_chat_thread, get_chat_message, + list_chat_attachments_page, list_chat_projects, list_chat_legacy_imports, list_chat_settings, @@ -279,6 +282,131 @@ async def delete_threads( return {"status": "deleted"} +@router.get("/attachments") +def list_attachments( + limit: Annotated[int, Query(ge = 1, le = 100)] = 50, + offset: Annotated[int, Query(ge = 0)] = 0, + current_subject: str = Depends(get_current_subject), +) -> dict: + """One bounded page of chat uploads for the settings Data tab.""" + attachments, next_offset = list_chat_attachments_page(limit = limit, offset = offset) + return {"attachments": attachments, "nextOffset": next_offset} + + +def _decode_attachment_base64(payload: str) -> bytes: + """Strict base64 decode of a stored payload. + + Normalizes first: strips whitespace, fixes padding, accepts the URL-safe + alphabet. validate=False would silently drop bad characters and serve + corrupted bytes instead of failing, so raise 422 on anything else. + """ + import base64 + + normalized = "".join(payload.split()) + altchars = b"-_" if ("-" in normalized or "_" in normalized) else None + normalized += "=" * (-len(normalized) % 4) + try: + return base64.b64decode(normalized, altchars = altchars, validate = True) + except Exception as exc: # noqa: BLE001 - corrupt stored payload + raise HTTPException(status_code = 422, detail = "Attachment data is corrupt") from exc + + +_AUDIO_FORMAT_MEDIA_TYPES = { + "mp3": "audio/mpeg", + "wav": "audio/wav", + "ogg": "audio/ogg", + "flac": "audio/flac", +} + + +def _safe_image_media_type(media_type: str) -> str: + """Clamp a data-URL media type to something inert to render. + + Imported chats store image parts verbatim, so the embedded type can be + text/html or image/svg+xml; echoing those would execute markup with the + app origin when opened. Anything not a plain raster type downloads as + bytes instead. + """ + lowered = media_type.strip().lower() + if lowered.startswith("image/") and lowered != "image/svg+xml": + return lowered + return "application/octet-stream" + + +@router.get("/attachments/{message_id}/{attachment_id}/file") +def get_attachment_file( + message_id: str, + attachment_id: str, + current_subject: str = Depends(get_current_subject), +): + """Serve one attachment's stored content: image or audio bytes, or + extracted text.""" + import urllib.parse + + from fastapi.responses import Response + + attachment = get_chat_attachment(message_id, attachment_id) + if attachment is None: + raise HTTPException(status_code = 404, detail = "Attachment not found") + + attachment_content_type = attachment.get("contentType") + texts: list[str] = [] + for part in attachment.get("content") or []: + if not isinstance(part, dict): + continue + image = part.get("image") + if isinstance(image, str) and image[:5].lower() == "data:": + header, _, payload = image.partition(",") + media_type = _safe_image_media_type( + header[5:].split(";", 1)[0] or "application/octet-stream" + ) + if "base64" not in header.lower(): + # RFC 2397 non-base64 form stores percent-encoded bytes. + data = urllib.parse.unquote_to_bytes(payload) + return Response(content = data, media_type = media_type) + data = _decode_attachment_base64(payload) + return Response(content = data, media_type = media_type) + # Audio parts: the attachment adapter stores {data, format} with raw + # base64; compare chats store a bare base64 string. + audio = part.get("audio") + if isinstance(audio, dict) or (isinstance(audio, str) and audio): + if isinstance(audio, dict): + payload = audio.get("data") + audio_format = audio.get("format") + else: + payload = audio.rsplit(",", 1)[-1] + audio_format = None + if isinstance(payload, str) and payload: + data = _decode_attachment_base64(payload) + media_type = ( + attachment_content_type + if isinstance(attachment_content_type, str) + and attachment_content_type.startswith("audio/") + else _AUDIO_FORMAT_MEDIA_TYPES.get( + str(audio_format or "").lower(), "application/octet-stream" + ) + ) + return Response(content = data, media_type = media_type) + text = part.get("text") + if isinstance(text, str) and text: + texts.append(text) + if texts: + return Response(content = "\n".join(texts), media_type = "text/plain; charset=utf-8") + raise HTTPException(status_code = 404, detail = "Attachment has no stored content") + + +@router.delete("/attachments/{message_id}/{attachment_id}") +def delete_attachment( + message_id: str, + attachment_id: str, + current_subject: str = Depends(get_current_subject), +) -> dict: + """Remove one attachment from its chat message.""" + if not delete_chat_attachment(message_id, attachment_id): + raise HTTPException(status_code = 404, detail = "Attachment not found") + return {"ok": True} + + @router.get("/projects", response_model = ChatProjectListResponse) async def list_projects( include_archived: bool = Query(False), current_subject: str = Depends(get_current_subject) @@ -409,7 +537,7 @@ async def get_thread_message( @router.put("/threads/{thread_id}/messages/{message_id}", response_model = ChatMessage) -async def save_thread_message( +def save_thread_message( thread_id: str, message_id: str, payload: ChatMessage, @@ -432,7 +560,7 @@ async def save_thread_message( @router.put("/threads/{thread_id}/messages", response_model = ChatMessageListResponse) -async def replace_thread_messages( +def replace_thread_messages( thread_id: str, payload: ChatMessageSyncRequest, current_subject: str = Depends(get_current_subject), diff --git a/studio/backend/routes/inference.py b/studio/backend/routes/inference.py index 3d527bf317..d3e588bb0b 100644 --- a/studio/backend/routes/inference.py +++ b/studio/backend/routes/inference.py @@ -13,7 +13,7 @@ from pathlib import Path from fastapi import APIRouter, Depends, HTTPException, Request, status from fastapi.responses import StreamingResponse, JSONResponse, Response from starlette.requests import ClientDisconnect -from typing import Any, Callable, List, Optional, Union +from typing import Any, Callable, List, Literal, Optional, Union import json import httpx from loggers import get_logger @@ -1778,7 +1778,7 @@ from core.inference.providers import get_base_url from core.inference.external_provider import ExternalProviderClient from core.inference.chat_templates import resolve_effective_chat_template_override from storage import providers_db -from utils.utils import safe_error_detail, log_and_http_error +from utils.utils import is_hf_authentication_error, safe_error_detail, log_and_http_error import io import base64 @@ -3115,13 +3115,16 @@ def _normalise_settings_str(value: Optional[str]) -> Optional[str]: def _should_strip_split_mode(request: LoadRequest, backend_extra: Optional[list[str]]) -> bool: - """Whether an inherited --split-mode should be stripped on reload. + """Whether an inherited --split-mode (and its coupled --tensor-split) should + be stripped on reload. The binary Tensor Parallelism toggle can't carry --split-mode's row/none/ layer modes, so only strip when the toggle overrides it: tensor being turned on, or the inherited mode is tensor (toggle turning it off). Non-tensor modes - survive. Shared by the inheritance strip and the already-loaded stale check - so they agree on what reload would do. + survive. A manual per-GPU ratio is handled by _should_strip_tensor_split, + which strips only --tensor-split so the inherited mode is kept. Shared by the + inheritance strip and the already-loaded stale check so they agree on what + reload would do. """ fields_set = getattr(request, "model_fields_set", set()) return "tensor_parallel" in fields_set and ( @@ -3129,6 +3132,25 @@ def _should_strip_split_mode(request: LoadRequest, backend_extra: Optional[list[ ) +def _should_strip_tensor_split(request: LoadRequest) -> bool: + """Whether an inherited --tensor-split alone should be stripped on reload. + + Manual explicit offload (gpu_layers >= 0) owns the per-GPU split: with a ratio + it emits its own --tensor-split (an inherited one, appended last, would + override it), and with the ratio cleared it wants llama.cpp's default + free-VRAM split. Either way an inherited --tensor-split must go, else the + cleared case silently keeps the stale ratio while status reports None. + Unlike _should_strip_split_mode this leaves --split-mode untouched, so a + user's row/none/layer mode survives a Studio split-ratio edit. When the + Tensor Parallelism toggle IS overriding the mode, _should_strip_split_mode + (called alongside this at every site) strips --split-mode anyway. + """ + return ( + getattr(request, "gpu_memory_mode", "auto") == "manual" + and getattr(request, "gpu_layers", -1) >= 0 + ) + + def _carry_preserved_tensor_intent( *, preserved: bool, same_model: bool, explicit_drop: bool ) -> bool: @@ -3187,12 +3209,44 @@ def _request_matches_loaded_settings( else strip_shadowing_flags( backend_extra, strip_split_mode = _should_strip_split_mode(request, backend_extra), + strip_tensor_split = _should_strip_tensor_split(request), + strip_offload = request.gpu_memory_mode == "manual", ) ) if not _tensor_parallel_matches_loaded( effective_extra, request.tensor_parallel, llama_backend.tensor_parallel ): return False + # The diffusion runner is mode-agnostic (it always reports "auto" and ignores + # the layer/MoE/split knobs), so a standing manual preference in the request + # must not force a needless reload -- only the GPU pick matters. + if not llama_backend.is_diffusion: + if request.gpu_memory_mode != llama_backend.gpu_memory_mode: + return False + # Manual: a layer-count change always reloads; MoE/split only matter with + # an explicit offload (gpu_layers >= 0), so a leftover value under Auto + # must not force one. Mirrors LlamaCppBackend._already_in_target_state. + if request.gpu_memory_mode == "manual" and ( + request.gpu_layers != llama_backend.gpu_layers + or ( + request.gpu_layers >= 0 + and ( + request.n_cpu_moe != llama_backend.n_cpu_moe + or (request.tensor_split or None) != (llama_backend.tensor_split or None) + ) + ) + ): + return False + # A changed GPU pick must reload. The diffusion runner collapses a multi-GPU + # request to its single lowest device (it drives one device only), so the + # backend records just that device; compare the request the same way, or a + # multi-GPU pick that resolves to the same device needlessly reloads. + if llama_backend.is_diffusion: + _req_gpu_ids = [sorted(request.gpu_ids)[0]] if request.gpu_ids else None + else: + _req_gpu_ids = sorted(request.gpu_ids) if request.gpu_ids else None + if _req_gpu_ids != llama_backend.gpu_ids: + return False # Preserved tensor->layer fallback (both report tensor=off, so the check above # matches): if the user now explicitly drops tensor intent, reload so placement # re-selects instead of keeping the all-GPU mask (#6659). The effective check @@ -3235,14 +3289,17 @@ def _request_matches_loaded_settings( # contain any shadow flag, so the reload path strips them rather than # leaving a stale override in effect. (backend_extra computed above.) if request.llama_extra_args is None: - # Mirror the reload's conditional split-mode strip, so a preserved - # non-tensor mode (row/none/layer) isn't seen as stale and doesn't - # trigger a needless reload of a healthy server. + # Mirror the reload's conditional strips, so a preserved non-tensor mode + # (row/none/layer) isn't seen as stale and doesn't trigger a needless + # reload of a healthy server, while an inherited offload/ratio flag that + # the reload *would* strip is correctly seen as stale. if ( backend_extra and strip_shadowing_flags( backend_extra, strip_split_mode = _should_strip_split_mode(request, backend_extra), + strip_tensor_split = _should_strip_tensor_split(request), + strip_offload = request.gpu_memory_mode == "manual", ) != backend_extra ): @@ -3861,6 +3918,46 @@ def _estimate_gguf_required_gb( return None +def _classify_diffusion_gguf(config: ModelConfig) -> Optional[bool]: + """Classify a GGUF as diffusion, normal, or unknown before it is loaded. + + ``None`` is important here: a remote GGUF whose header is not cached can + still be routed to the single-GPU diffusion runner after download. Treating + that case as normal would let Manual mode skip the training guard even + though the runner ignores Manual's llama-server placement controls. + """ + identity = " ".join( + str(getattr(config, attr, "") or "") for attr in ("identifier", "gguf_hf_repo", "gguf_file") + ).lower() + if "diffusion" in identity: + return True + + try: + main = getattr(config, "gguf_file", None) + if not (main and Path(main).is_file()): + repo = getattr(config, "gguf_hf_repo", None) + variant = getattr(config, "gguf_variant", None) + if repo and variant: + from hub.utils.gguf import resolve_local_gguf_path + main = resolve_local_gguf_path(repo, variant) + if not main or not Path(main).is_file(): + return None + + probe = LlamaCppBackend() + probe._read_gguf_metadata(str(main)) + if probe.is_diffusion: + return True + # A successfully decoded architecture proves that this is a normal + # llama-server GGUF. No architecture means the lightweight probe could + # not establish the routing decision, so preserve the unknown state. + if getattr(probe, "_architecture", None): + return False + return None + except Exception as e: + logger.debug("Could not identify diffusion GGUF for training guard: %s", e) + return None + + def _guard_chat_load_against_training( config: ModelConfig, *, @@ -3871,11 +3968,19 @@ def _guard_chat_load_against_training( requested_gpu_ids: Optional[List[int]], llama_extra_args: Optional[list[str]] = None, n_parallel: int = 1, + gpu_memory_mode: Literal["auto", "manual"] = "auto", ) -> None: - """Refuse loading a local chat model that would OOM an active training run. + """Protect active training from automatically placed chat-model loads. + No-op when training is inactive or unknown. `load_in_4bit` must be the - effective quantization (see _effective_load_in_4bit). Raises HTTP 409 when the - model would not fit alongside training.""" + effective quantization (see _effective_load_in_4bit). Manual chat-GGUF + placement is an explicit override: Auto layers delegate fitting to + llama.cpp's ``--fit`` and pinned layers are owned by the user, so neither is + estimated here. Diffusion is still guarded because its mode-agnostic runner + ignores those controls and uses one GPU. An unclassified GGUF is guarded as + potentially diffusion until its local header proves otherwise. Other loads + raise HTTP 409 when they would not fit beside training. + """ from core.training import get_training_backend from routes.training_vram import can_load_chat_during_training @@ -3887,6 +3992,19 @@ def _guard_chat_load_against_training( return is_gguf = bool(getattr(config, "is_gguf", False)) + diffusion_kind = _classify_diffusion_gguf(config) if is_gguf else False + if is_gguf and gpu_memory_mode == "manual" and diffusion_kind is False: + return + + diffusion_gpu = None + if is_gguf and diffusion_kind is not False: + # Use the same token selection as the runner: an explicit pick wins, + # followed by DG_GPU, the first parent-visible token, then GPU 0. + diffusion_gpu = LlamaCppBackend._diffusion_gpu_arg( + requested_gpu_ids, + cpu_only = LlamaCppBackend._effective_gpu_count() == 0, + ) + required_override_gb = ( _estimate_gguf_required_gb( config, @@ -3907,6 +4025,7 @@ def _guard_chat_load_against_training( requested_gpu_ids = requested_gpu_ids, is_gguf = is_gguf, required_override_gb = required_override_gb, + single_device_gpu = diffusion_gpu, ) if ok: return @@ -3934,6 +4053,98 @@ def _guard_chat_load_against_training( raise HTTPException(status_code = 409, detail = detail) +def _resolve_inherited_extra_args( + request, + config: ModelConfig, + model_identifier: str, + extra_llama_args: Optional[list[str]], + effective_chat_template_override: Optional[str] = None, +) -> Optional[list[str]]: + """Effective pass-through extras for a GGUF request that omitted the field: + the previous same-model load's extras, shadow-stripped, so a settings-Apply + reload (which does not round-trip the extras field) keeps them (#5401).""" + if getattr(request, "llama_extra_args", None) is not None: + return extra_llama_args + if not getattr(config, "is_gguf", False): + return extra_llama_args + llama_backend = get_llama_cpp_backend() + if not llama_backend.extra_args: + return extra_llama_args + # Inherit the previous load's extras (the chat-settings Apply path doesn't + # round-trip them; an explicit [] still clears). Gated on (model_identifier, + # hf_variant) to refuse cross-model pickup, and shadowing flags are + # stripped so an inherited override can't win the last-wins CLI + # parse against a freshly-supplied first-class field. + source = llama_backend.extra_args_source + # Compare against the resolved variant, not the request field: callers + # commonly omit gguf_variant for local ``.gguf`` paths and HF auto-pick + # flows. ``config.gguf_variant`` is the variant load_model was actually + # invoked with, so both sides of the comparison key off the same string. + resolved_variant = (config.gguf_variant or "").lower() + request_variant = (request.gguf_variant or "").lower() + stored_variant = (source[1] or "").lower() if source else "" + same_model = bool(source and source[0] and source[0].lower() == model_identifier.lower()) + if request.gguf_variant: + variant_mismatch = request_variant != stored_variant + else: + variant_mismatch = bool(stored_variant and resolved_variant != stored_variant) + same_source = same_model and not variant_mismatch + if not same_source: + logger.info( + "Not inheriting llama_extra_args: stored args came from %s, loading %s", + source, + (model_identifier, resolved_variant), + ) + # Cross-model: clear explicitly so the backend doesn't + # inherit via "no opinion" semantics. + extra_llama_args = [] + else: + # Strip only the groups whose first-class field was set by the caller, so + # an inherited --chat-template-file survives an Apply that omits + # chat_template_override. A bundled family template (e.g. gemma-4) counts as + # a first-class template even when the request omits chat_template_override, + # so strip the inherited --chat-template-file then too -- else the stale arg + # (appended last) shadows the bundled template while Studio reports its caps. + fields_set = getattr(request, "model_fields_set", set()) + stripped = strip_shadowing_flags( + llama_backend.extra_args, + strip_context = "max_seq_length" in fields_set, + strip_cache = "cache_type_kv" in fields_set, + strip_spec = ("speculative_type" in fields_set or "spec_draft_n_max" in fields_set), + strip_template = ( + "chat_template_override" in fields_set + or effective_chat_template_override is not None + ), + strip_split_mode = _should_strip_split_mode(request, llama_backend.extra_args), + # manual + per-GPU ratio emits its own --tensor-split; drop + # an inherited one (appended last would override it) while + # keeping the user's --split-mode row/none/layer choice. + strip_tensor_split = _should_strip_tensor_split(request), + # manual emits its own --fit/--gpu-layers, so an inherited offload flag + # must not last-wins-override it. auto leaves a user's inherited -ngl + # alone. getattr: a validate request reuses this resolver, no offload fields. + strip_offload = getattr(request, "gpu_memory_mode", "auto") == "manual", + ) + try: + extra_llama_args = validate_extra_args(stripped) + except ValueError: + # Shouldn't happen on already-validated args; degrade to + # no-extras rather than 400 if managed flags changed. + logger.warning( + "Stored llama_extra_args failed revalidation; loading without them: %s", + stripped, + ) + extra_llama_args = [] + else: + if extra_llama_args: + logger.info( + "Inheriting llama_extra_args from previous " + "load (same model, shadow-stripped): %s", + extra_llama_args, + ) + return extra_llama_args + + def _model_json_response(model, status_code: int = 200) -> Response: """Serialize a pydantic response once via pydantic-core. @@ -4040,6 +4251,35 @@ async def _load_model_impl(request: LoadRequest, fastapi_request: Request, curre None if request.llama_extra_args is None else extra_llama_args ) + # Manual mode owns the offload flags: strip them from EXPLICIT extras + # too (the inherited path already does), or a last-wins --gpu-layers / + # --fit in extras re-enables GPU offload on a load status reports as + # CPU-only. Manual + per-GPU ratio owns --tensor-split the same way. + if request.gpu_memory_mode == "manual" and extra_llama_args: + _stripped_explicit = strip_shadowing_flags( + extra_llama_args, + strip_context = False, + strip_cache = False, + strip_spec = False, + strip_template = False, + strip_split_mode = False, + strip_tensor_split = _should_strip_tensor_split(request), + strip_offload = True, + ) + if _stripped_explicit != extra_llama_args: + logger.info( + "Manual GPU memory owns the offload flags; stripping them " + "from explicit llama_extra_args: %s -> %s", + extra_llama_args, + _stripped_explicit, + ) + extra_llama_args = _stripped_explicit + + # Keep every downstream consumer on the normalized explicit list. In + # particular, the already-loaded comparator must not compare the raw + # request's managed offload flags against the stripped launch state. + request = request.model_copy(update = {"llama_extra_args": extra_llama_args}) + model_identifier, model_log_label, native_grant_backed = ( _resolve_model_identifier_for_request(request, operation = "load-model") ) @@ -4121,6 +4361,13 @@ async def _load_model_impl(request: LoadRequest, fastapi_request: Request, curre speculative_type = llama_backend.requested_spec_mode, spec_draft_n_max = llama_backend.spec_draft_n_max, tensor_parallel = llama_backend.tensor_parallel, + gpu_memory_mode = llama_backend.gpu_memory_mode, + gpu_layers = llama_backend.gpu_layers, + n_cpu_moe = llama_backend.n_cpu_moe, + tensor_split = llama_backend.tensor_split, + n_layers = llama_backend.n_layers, + n_moe_layers = llama_backend.n_moe_layers, + gpu_ids = llama_backend.gpu_ids, ) else: if ( @@ -4187,12 +4434,41 @@ async def _load_model_impl(request: LoadRequest, fastapi_request: Request, curre # Normalize gpu_ids: empty list means auto-selection, same as None effective_gpu_ids = request.gpu_ids if request.gpu_ids else None - # Reject GGUF + gpu_ids first so the guard can't mask it with a VRAM 409. + # GGUF supports gpu_ids: validate the pick up front (before the training + # guard) so a bad pick is a clean 400, not masked by a VRAM 409. Rejects + # negative / out-of-range / duplicate ids and UUID/MIG parents. XPU hosts + # are rejected outright: the picker's indices are torch-xpu ordinals neither + # applicator speaks (CUDA/HIP masks don't apply, the Vulkan --device pin + # uses ggml's own Vulkan ordinals), so a pick could land on the wrong device. if config.is_gguf and effective_gpu_ids is not None: - raise HTTPException( - status_code = 400, - detail = "gpu_ids is not supported for GGUF models yet.", - ) + from utils.hardware import DeviceType, get_device + from utils.hardware.hardware import resolve_requested_gpu_ids + + if get_device() == DeviceType.XPU: + raise HTTPException( + status_code = 400, + detail = ( + "GPU selection (gpu_ids) is not supported on Intel XPU. " + "Omit gpu_ids to use all devices." + ), + ) + # Same reasoning for a Vulkan-only build: --device pins ggml's own + # Vulkan ordinals, so a physical pick can land on the wrong card on + # masked or non-contiguous hosts. + if LlamaCppBackend._is_vulkan_backend(): + raise HTTPException( + status_code = 400, + detail = ( + "GPU selection (gpu_ids) is not supported with a Vulkan " + "llama.cpp build: physical GPU ids have no defined " + "mapping to Vulkan device ordinals. Omit gpu_ids to use " + "all devices." + ), + ) + try: + resolve_requested_gpu_ids(effective_gpu_ids) + except ValueError as exc: + raise HTTPException(status_code = 400, detail = str(exc)) from exc if not config.is_gguf and _mlx_distributed_launch_detected(): raise HTTPException( status_code = 400, @@ -4222,8 +4498,20 @@ async def _load_model_impl(request: LoadRequest, fastapi_request: Request, curre "architectures)" ) - # Refuse a load that would OOM active training, before the unload step below - # frees the resident model. Off-loop: guard does sync nvidia-smi / HF work. + # Inherit the previous same-model load's pass-through extras when this + # request omits the field (a settings-Apply reload doesn't round-trip + # them); shadow-stripped so an inherited flag can't override a + # first-class field the caller did set (#5401). + extra_llama_args = _resolve_inherited_extra_args( + request, + config, + model_identifier, + extra_llama_args, + effective_chat_template_override, + ) + + # Apply the training coexistence policy before the unload step below + # frees the resident model. Off-loop: the default-mode guard does sync work. await asyncio.to_thread( _guard_chat_load_against_training, config, @@ -4234,6 +4522,7 @@ async def _load_model_impl(request: LoadRequest, fastapi_request: Request, curre requested_gpu_ids = effective_gpu_ids, llama_extra_args = extra_llama_args, n_parallel = getattr(fastapi_request.app.state, "llama_parallel_slots", 1), + gpu_memory_mode = request.gpu_memory_mode, ) # ── GGUF path: load via llama-server ────────────────────── @@ -4245,84 +4534,6 @@ async def _load_model_impl(request: LoadRequest, fastapi_request: Request, curre from core.inference.llama_cpp import gguf_load_in_flight gguf_load_stack.enter_context(gguf_load_in_flight(config.gguf_hf_repo)) - # Inherit llama_extra_args from the previous load when the request - # omits the field (the chat-settings Apply path doesn't round-trip - # them; explicit [] still clears). Gated on (model_identifier, - # hf_variant) to refuse cross-model pickup, and shadowing flags are - # stripped so an inherited override can't win the last-wins CLI - # parse against a freshly-supplied first-class field. - if request.llama_extra_args is None and llama_backend.extra_args: - source = llama_backend.extra_args_source - # Compare against the resolved variant, not the request - # field: callers commonly omit gguf_variant for local - # ``.gguf`` paths and HF auto-pick flows. ``config.gguf_ - # variant`` is the variant load_model was actually - # invoked with (see the HF / local branches below), so - # both sides of the comparison key off the same string. - resolved_variant = (config.gguf_variant or "").lower() - request_variant = (request.gguf_variant or "").lower() - stored_variant = (source[1] or "").lower() if source else "" - same_model = bool( - source and source[0] and source[0].lower() == model_identifier.lower() - ) - if request.gguf_variant: - variant_mismatch = request_variant != stored_variant - else: - variant_mismatch = bool(stored_variant and resolved_variant != stored_variant) - same_source = same_model and not variant_mismatch - if not same_source: - logger.info( - "Not inheriting llama_extra_args: stored args came from %s, loading %s", - source, - (model_identifier, resolved_variant), - ) - # Cross-model: clear explicitly so the backend doesn't - # inherit via "no opinion" semantics. - extra_llama_args = [] - else: - # Strip only the groups whose first-class field was set by - # the caller, so an inherited --chat-template-file survives - # an Apply that omits chat_template_override. A bundled family - # template (e.g. the gemma-4 override) is an effective - # first-class template setting even when the raw request - # omits chat_template_override, so strip the inherited - # --chat-template-file in that case too -- otherwise the stale - # extra arg (appended last) shadows the bundled template while - # Unsloth reports the bundled template's capabilities. - fields_set = getattr(request, "model_fields_set", set()) - stripped = strip_shadowing_flags( - llama_backend.extra_args, - strip_context = "max_seq_length" in fields_set, - strip_cache = "cache_type_kv" in fields_set, - strip_spec = ( - "speculative_type" in fields_set or "spec_draft_n_max" in fields_set - ), - strip_template = ( - "chat_template_override" in fields_set - or effective_chat_template_override is not None - ), - strip_split_mode = _should_strip_split_mode( - request, llama_backend.extra_args - ), - ) - try: - extra_llama_args = validate_extra_args(stripped) - except ValueError: - # Shouldn't happen on already-validated args; degrade to - # no-extras rather than 400 if managed flags changed. - logger.warning( - "Stored llama_extra_args failed revalidation; loading without them: %s", - stripped, - ) - extra_llama_args = [] - else: - if extra_llama_args: - logger.info( - "Inheriting llama_extra_args from previous " - "load (same model, shadow-stripped): %s", - extra_llama_args, - ) - # Block cache writes that would race the download manager. This runs # after pass-through argument inheritance so a carried --no-mmproj # changes the companion requirement exactly as it does for the load. @@ -4370,6 +4581,11 @@ async def _load_model_impl(request: LoadRequest, fastapi_request: Request, curre cache_type_kv = request.cache_type_kv, speculative_type = request.speculative_type, spec_draft_n_max = request.spec_draft_n_max, + gpu_memory_mode = request.gpu_memory_mode, + gpu_layers = request.gpu_layers, + n_cpu_moe = request.n_cpu_moe, + tensor_split = request.tensor_split, + gpu_ids = effective_gpu_ids, n_parallel = _n_parallel, ) if config.gguf_hf_repo: @@ -4494,7 +4710,7 @@ async def _load_model_impl(request: LoadRequest, fastapi_request: Request, curre # Clear any idle-unload reload stash now, not only on the next poll. from core.inference.llama_keepwarm import note_model_loaded - note_model_loaded() + await asyncio.to_thread(note_model_loaded, llama_backend) # A plain load advertises its own identifier; auto-switch overwrites # this with the repo id right after _load_model_impl returns. llama_backend._openai_advertised_id = None @@ -4537,6 +4753,13 @@ async def _load_model_impl(request: LoadRequest, fastapi_request: Request, curre speculative_type = llama_backend.requested_spec_mode, spec_draft_n_max = llama_backend.spec_draft_n_max, tensor_parallel = llama_backend.tensor_parallel, + gpu_memory_mode = llama_backend.gpu_memory_mode, + gpu_layers = llama_backend.gpu_layers, + n_cpu_moe = llama_backend.n_cpu_moe, + tensor_split = llama_backend.tensor_split, + n_layers = llama_backend.n_layers, + n_moe_layers = llama_backend.n_moe_layers, + gpu_ids = llama_backend.gpu_ids, ) # ── Standard path: load via Unsloth/transformers ────────── @@ -4795,7 +5018,9 @@ def _requires_security_review_for_model( @router.post("/validate", response_model = ValidateModelResponse) async def validate_model( - request: ValidateModelRequest, current_subject: str = Depends(get_current_subject) + request: ValidateModelRequest, + fastapi_request: Request = None, + current_subject: str = Depends(get_current_subject), ): """ Lightweight validation endpoint for model identifiers. @@ -4823,15 +5048,39 @@ async def validate_model( detail = f"Invalid model identifier: {model_log_label}", ) - # Refuse early (before the frontend unloads to load this) if it can't fit - # alongside training, using the same settings /load uses so they agree. + # Apply the same training coexistence policy as /load before the frontend + # unloads the current model. effective_gpu_ids = request.gpu_ids if request.gpu_ids else None - # Mirror /load: reject GGUF + gpu_ids before the guard so both return 400. + # Mirror /load: GGUF supports gpu_ids, so validate the pick (a bad one is + # a clean 400) before the guard sizes the model against training VRAM. + # XPU-host picks are rejected like /load (no defined mapping from the + # picker's torch-xpu ordinals to the launcher's device spaces). if config.is_gguf and effective_gpu_ids is not None: - raise HTTPException( - status_code = 400, - detail = "gpu_ids is not supported for GGUF models yet.", - ) + from utils.hardware import DeviceType, get_device + from utils.hardware.hardware import resolve_requested_gpu_ids + + if get_device() == DeviceType.XPU: + raise HTTPException( + status_code = 400, + detail = ( + "GPU selection (gpu_ids) is not supported on Intel XPU. " + "Omit gpu_ids to use all devices." + ), + ) + if LlamaCppBackend._is_vulkan_backend(): + raise HTTPException( + status_code = 400, + detail = ( + "GPU selection (gpu_ids) is not supported with a Vulkan " + "llama.cpp build: physical GPU ids have no defined " + "mapping to Vulkan device ordinals. Omit gpu_ids to use " + "all devices." + ), + ) + try: + resolve_requested_gpu_ids(effective_gpu_ids) + except ValueError as exc: + raise HTTPException(status_code = 400, detail = str(exc)) from exc effective_load_in_4bit = _effective_load_in_4bit(config, request.load_in_4bit) # Both checks cover the [adapter, base] set (matching the scan route and workers): @@ -4895,16 +5144,32 @@ async def validate_model( latest_tier_active_for, config.identifier, request.hf_token ): effective_load_in_4bit = False - # Off-loop: guard does sync nvidia-smi / HF work. - await asyncio.to_thread( - _guard_chat_load_against_training, - config, - model_identifier = model_identifier, - hf_token = request.hf_token, - load_in_4bit = effective_load_in_4bit, - max_seq_length = request.max_seq_length, - requested_gpu_ids = effective_gpu_ids, - ) + # A metadata-only probe reads the GGUF header and allocates no VRAM, so the + # training guard must not refuse it. Real loads omit include_context_length / + # include_chat_template, and /load applies the guard again. + if not (request.include_context_length or request.include_chat_template): + # Match /load's inherited llama.cpp extras and parallel slot count so + # validation cannot pass a smaller estimate than the subsequent load. + effective_extra_args = _resolve_inherited_extra_args( + request, config, model_identifier, None + ) + # Off-loop: guard does sync nvidia-smi / HF work. + await asyncio.to_thread( + _guard_chat_load_against_training, + config, + model_identifier = model_identifier, + hf_token = request.hf_token, + load_in_4bit = effective_load_in_4bit, + max_seq_length = request.max_seq_length, + requested_gpu_ids = effective_gpu_ids, + llama_extra_args = effective_extra_args, + n_parallel = ( + getattr(fastapi_request.app.state, "llama_parallel_slots", 1) + if fastapi_request is not None + else 1 + ), + gpu_memory_mode = request.gpu_memory_mode, + ) # A selected GGUF loads via llama.cpp: auto_map Python and root pickle weights in a # mixed repo are inert for this load, so gating on them is a false positive. Only @@ -4918,10 +5183,21 @@ async def validate_model( # Native context length, read from the local GGUF header when present. # Lets the staged ("Load on selection" off) flow populate the context # slider before the GPU load; None until the file is downloaded. + # Staged header dims (one read): native context, total layer count, and + # MoE expert-layer count -- let the staged flow size the context, GPU- + # layers and manual --n-cpu-moe sliders before the load. context_length: Optional[int] = None - if request.include_context_length and is_gguf: + layer_count: Optional[int] = None + moe_layer_count: Optional[int] = None + chat_template: Optional[str] = None + # Both header probes read the same local GGUF, so resolve it once. + if (request.include_context_length or request.include_chat_template) and is_gguf: from hub.utils.gguf import resolve_local_gguf_path - from utils.models.gguf_metadata import read_gguf_context_length + from picker.schemas import MAX_CHAT_TEMPLATE_BYTES + from utils.models.gguf_metadata import ( + read_gguf_chat_template, + read_gguf_staged_dims, + ) # Best-effort: a header-read failure must never fail validation of an # otherwise-valid model (the outer except turns it into a 400). @@ -4937,9 +5213,26 @@ async def validate_model( model_identifier, request.gguf_variant ) if local_gguf: - context_length = read_gguf_context_length(local_gguf) + if request.include_context_length: + # Header walk reads tokenizer arrays (tens of ms); keep it + # off the event loop. + dims = await asyncio.to_thread(read_gguf_staged_dims, local_gguf) + if dims: + context_length = dims["context_length"] + layer_count = dims["layer_count"] + moe_layer_count = dims["moe_layer_count"] + if request.include_chat_template: + # Read only the leased GGUF's own embedded template (the copy + # llama.cpp loads), never a sibling sidecar: the native grant + # authorizes just this path, so neighbours would be scope escalation. + raw_template = await asyncio.to_thread(read_gguf_chat_template, local_gguf) + if ( + raw_template is not None + and len(raw_template.encode("utf-8")) <= MAX_CHAT_TEMPLATE_BYTES + ): + chat_template = raw_template except Exception as e: - logger.debug("Context-length probe failed for %s: %s", model_log_label, e) + logger.debug("Header probe failed for %s: %s", model_log_label, e) return ValidateModelResponse( valid = True, @@ -4954,6 +5247,9 @@ async def validate_model( requires_trust_remote_code = requires_trust_remote_code, requires_security_review = requires_security_review, context_length = context_length, + layer_count = layer_count, + moe_layer_count = moe_layer_count, + chat_template = chat_template, requires_transformers_upgrade = transformers_upgrade is not None, transformers_upgrade = transformers_upgrade, ) @@ -4966,6 +5262,14 @@ async def validate_model( raise HTTPException(status_code = 400, detail = str(e)) except Exception as e: redacted_msg = redact_native_paths(str(e)) + if is_hf_authentication_error(e): + raise HTTPException( + status_code = 400, + detail = ( + "Hugging Face authentication failed. Check or clear the token " + "in Settings, and confirm access to this gated repository." + ), + ) if _is_unsupported_nvfp4_inference_error(redacted_msg): logger.warning( "NVFP4 inference is not supported yet while validating '%s'", @@ -5593,6 +5897,14 @@ async def get_status(current_subject: str = Depends(get_current_subject)): speculative_type = llama_backend.requested_spec_mode, spec_draft_n_max = llama_backend.spec_draft_n_max, tensor_parallel = llama_backend.tensor_parallel, + gpu_memory_mode = llama_backend.gpu_memory_mode, + gpu_layers = llama_backend.gpu_layers, + n_cpu_moe = llama_backend.n_cpu_moe, + tensor_split = llama_backend.tensor_split, + requested_context_length = llama_backend.requested_n_ctx, + n_layers = llama_backend.n_layers, + n_moe_layers = llama_backend.n_moe_layers, + gpu_ids = llama_backend.gpu_ids, llama_cpp_supports_mtp = _supports_mtp, spec_fallback_reason = llama_backend.spec_fallback_reason, llama_cpp_prebuilt_stale = _stale, diff --git a/studio/backend/routes/models.py b/studio/backend/routes/models.py index bb321695cd..a5ce1a72f0 100644 --- a/studio/backend/routes/models.py +++ b/studio/backend/routes/models.py @@ -60,13 +60,52 @@ def _safe_is_dir(path) -> bool: # Shared with the hub inventory scans; keep the private aliases so existing -# importers (core.inference.local_model_resolver, tests) stay valid. +# importers stay valid. ``_HF_REPO_ID_RE`` is the Hub repo id shape ("owner/name"); +# anything else is treated as a local filesystem path. from utils.hidden_models import ( + _HF_REPO_ID_RE, + _existing_resolved_path, _safe_resolve, is_hidden_model as _is_hidden_model, ) +def hidden_model_matchers() -> tuple[list[str], list[str], list[str]]: + """Substring needles, exact repo ids, and exact resolved paths identifying + infra models (the RAG embedder and the llama.cpp install validation probe) + that pickers hide. Served by the ``/api/hub/hidden-models`` endpoint. A + configured HF-repo embedder is published as its exact lowercased repo id + (mirroring ``utils.hidden_models.is_hidden_model``) and a local-path + embedder as its exact resolved path only: a generic basename like "model" + must not substring-hide unrelated chat models.""" + from core.rag import config as rag_config + + needles = [ + # The validation probe's repo and its exact filename. The filename carries + # .gguf so it won't hide unrelated repos like ``user/stories260K-finetune-GGUF``. + "ggml-org/models", + "stories260k.gguf", + ] + exact_ids: list[str] = [] + exact_paths: list[str] = [] + for model in ( + rag_config.effective_embedding_model(), + rag_config.effective_gguf_repo(), + ): + # Resolve an existing local path before the repo-id regex: a local embedder + # shaped like "models/embedder" is an exact path, not a Hub repo id. + existing_path = _existing_resolved_path(model) + if existing_path: + exact_paths.append(existing_path.lower()) + elif _HF_REPO_ID_RE.match(model): + exact_ids.append(model.lower()) + else: + resolved = _safe_resolve(Path(model).expanduser()) + if resolved: + exact_paths.append(resolved.lower()) + return needles, exact_ids, exact_paths + + backend_path = Path(__file__).parent.parent.parent if str(backend_path) not in sys.path: sys.path.insert(0, str(backend_path)) @@ -91,6 +130,7 @@ try: _pick_best_gguf, _extract_quant_label, _is_big_endian_gguf_path, + _is_mtp_drafter, is_audio_input_type, ) from core.inference import get_inference_backend @@ -123,6 +163,7 @@ except ImportError: _pick_best_gguf, _extract_quant_label, _is_big_endian_gguf_path, + _is_mtp_drafter, is_audio_input_type, ) from core.inference import get_inference_backend @@ -803,7 +844,7 @@ def collect_local_models(models_root: Path) -> List[LocalModelInfo]: models = sorted( deduped.values(), - key = lambda item: (item.updated_at or 0), + key = lambda item: item.updated_at or 0, reverse = True, ) return [m for m in models if not _is_hidden_model(m.id, m.model_id, m.path)] @@ -1750,9 +1791,11 @@ def _get_model_size_bytes(model_name: str, hf_token: Optional[str] = None) -> Op async def get_model_config( model_name: str, hf_token: Optional[str] = Query(None), + header_hf_token: Optional[str] = Depends(get_hf_token), current_subject: str = Depends(get_current_subject), ): """Get configuration for a specific model (wraps load_model_defaults).""" + hf_token = _normalize_hf_token(header_hf_token) or _normalize_hf_token(hf_token) try: if not is_local_path(model_name): resolved = resolve_cached_repo_id_case(model_name) @@ -2471,6 +2514,7 @@ async def get_lora_base_model(lora_path: str, current_subject: str = Depends(get async def check_vision_model( model_name: str, hf_token: Optional[str] = Query(None), + header_hf_token: Optional[str] = Depends(get_hf_token), current_subject: str = Depends(get_current_subject), ): """ @@ -2478,6 +2522,7 @@ async def check_vision_model( This endpoint wraps the backend is_vision_model function. """ + hf_token = _normalize_hf_token(header_hf_token) or _normalize_hf_token(hf_token) try: logger.info(f"Checking if vision model: {model_name}") # Authenticate so a gated/private VLM classifies correctly (else 404 -> non-vision). @@ -2503,6 +2548,7 @@ async def check_vision_model( async def check_embedding_model( model_name: str, hf_token: Optional[str] = Query(None), + header_hf_token: Optional[str] = Depends(get_hf_token), current_subject: str = Depends(get_current_subject), ): """ @@ -2510,6 +2556,7 @@ async def check_embedding_model( This endpoint wraps the backend is_embedding_model function. """ + hf_token = _normalize_hf_token(header_hf_token) or _normalize_hf_token(hf_token) try: logger.info(f"Checking if embedding model: {model_name}") is_embedding = is_embedding_model(model_name, hf_token = hf_token) @@ -2573,12 +2620,6 @@ def _resolve_quant_gguf(repo_id: str, quant: str, is_local: bool) -> tuple[Optio Q8_0 weights). Never raises. """ try: - from utils.models.model_config import ( - _extract_quant_label, - _is_big_endian_gguf_path, - _is_mtp_drafter, - ) - if is_local: roots = [Path(repo_id)] else: @@ -2595,25 +2636,19 @@ def _resolve_quant_gguf(repo_id: str, quant: str, is_local: bool) -> tuple[Optio if snaps.is_dir(): roots.extend(s for s in snaps.iterdir() if s.is_dir()) - want = quant.lower().replace("-", "").replace("_", "") + want = _normalized_quant_label(quant) best_total = 0 best_first: Optional[str] = None for root in roots: matches: list[tuple[str, Path]] = [] total = 0 for f in _iter_gguf_paths(root): - if _is_mmproj_filename(f.name): - continue try: rel = f.relative_to(root).as_posix() except ValueError: rel = f.name - if _is_mtp_drafter(rel): - continue - q = _extract_quant_label(rel) - if _is_big_endian_gguf_path(rel, q): - continue - if q.lower().replace("-", "").replace("_", "") != want: + q = _main_variant_gguf_label(rel) + if q is None or _normalized_quant_label(q) != want: continue try: total += f.stat().st_size @@ -2731,7 +2766,11 @@ async def get_gguf_variants( ], has_vision = response.has_vision, default_variant = response.default_variant, - context_length = _read_native_context_length(repo_id, is_local = local), + # The header walk reads tokenizer arrays on dense models (tens of + # ms per uncached file); keep it off the event loop. + context_length = await asyncio.to_thread( + _read_native_context_length, repo_id, is_local = local + ), ) except HTTPException: raise @@ -3031,6 +3070,22 @@ def _is_main_gguf_filename(name: str) -> bool: return _is_gguf_filename(name) and not _is_mmproj_filename(name) +def _main_variant_gguf_label(rel_path: str) -> Optional[str]: + name = rel_path.rsplit("/", 1)[-1] + if not _is_main_gguf_filename(name): + return None + if _is_mtp_drafter(rel_path): + return None + label = _extract_quant_label(rel_path) + if _is_big_endian_gguf_path(rel_path, label): + return None + return label + + +def _normalized_quant_label(label: str) -> str: + return label.lower().replace("-", "").replace("_", "") + + def _repo_has_mmproj(repo_info) -> bool: """True if the repo ships a GGUF vision adapter (mmproj), so it can take image inputs. Cheap: scans already-listed file names only.""" @@ -3358,6 +3413,170 @@ async def delete_cached_model( ) +def _resolve_cached_model_path(repo_id: str, variant: Optional[str]) -> Path: + """Absolute path of a cached repo (newest snapshot dir) or, with *variant*, + that quant's main GGUF file (first split of a sharded quant). Paths come + from the HF cache scan only, so callers can't probe arbitrary paths.""" + cache_scans = _all_hf_cache_scans() + + matching_repos = [] + for hf_cache in cache_scans: + for repo_info in hf_cache.repos: + if repo_info.repo_type != "model": + continue + if repo_info.repo_id.lower() == repo_id.lower(): + matching_repos.append(repo_info) + if not matching_repos: + raise HTTPException(status_code = 404, detail = "Model not found in cache") + + if variant: + want = _normalized_quant_label(variant) + candidate_revisions = sorted( + (rev for repo_info in matching_repos for rev in repo_info.revisions), + key = lambda rev: getattr(rev, "last_modified", 0) or 0, + reverse = True, + ) + for rev in candidate_revisions: + snapshot = getattr(rev, "snapshot_path", None) + matches = [] + for f in rev.files: + p = Path(f.file_path) + rel = f.file_name + if snapshot: + try: + rel = p.relative_to(snapshot).as_posix() + except ValueError: + pass + label = _main_variant_gguf_label(rel) + if label is None or _normalized_quant_label(label) != want: + continue + if p.exists() or p.is_symlink(): + matches.append((rel, p)) + if matches: + # Path-sorted so a sharded quant deterministically yields its first split. + return sorted(matches, key = lambda m: m[0].lower())[0][1] + raise HTTPException( + status_code = 404, + detail = f"Variant {variant} not found in cache for {repo_id}", + ) + + def repo_size(repo_info) -> int: + gguf_size = _repo_gguf_size_bytes(repo_info) + if gguf_size > 0: + return gguf_size + return sum( + (getattr(f, "size_on_disk", None) or 0) + for rev in repo_info.revisions + for f in rev.files + ) + + def repo_last_modified(repo_info) -> float: + return max( + (getattr(rev, "last_modified", 0) or 0 for rev in repo_info.revisions), + default = 0, + ) + + target_repo = max( + matching_repos, + key = lambda repo_info: (repo_size(repo_info), repo_last_modified(repo_info)), + ) + + # Whole repo: the newest revision's snapshot dir holds the visible files. + revisions = sorted( + (rev for rev in target_repo.revisions if getattr(rev, "snapshot_path", None)), + key = lambda rev: getattr(rev, "last_modified", 0) or 0, + reverse = True, + ) + for rev in revisions: + p = Path(rev.snapshot_path) + if p.exists(): + return p + p = Path(target_repo.repo_path) + if p.exists(): + return p + raise HTTPException(status_code = 404, detail = "Cached model path not found") + + +def _wsl_reveal_in_explorer(path: Path) -> bool: + import subprocess + + from utils.paths.path_utils import _IS_WSL + + if not _IS_WSL: + return False + try: + windows_path = subprocess.run( + ["wslpath", "-w", str(path)], + capture_output = True, + text = True, + check = True, + timeout = 10, + ).stdout.strip() + if not windows_path: + return False + argument = f"/select,{windows_path}" if path.is_file() else windows_path + subprocess.Popen(["explorer.exe", argument]) + return True + except (OSError, subprocess.SubprocessError): + return False + + +def _reveal_in_file_manager(path: Path) -> None: + """Open the OS file manager with *path* selected (best effort per platform).""" + import subprocess + + target = str(path) + if sys.platform == "darwin": + cmd = ["open", "-R", target] if path.is_file() else ["open", target] + subprocess.Popen(cmd) + elif os.name == "nt": + if path.is_file(): + subprocess.Popen(["explorer", f"/select,{target}"]) + else: + os.startfile(target) # noqa: S606 - local user's own file manager + elif not _wsl_reveal_in_explorer(path): + # No cross-desktop "select file" standard on Linux; open the directory. + directory = target if path.is_dir() else str(path.parent) + subprocess.Popen(["xdg-open", directory]) + + +class CachedModelPathResponse(BaseModel): + path: str + is_dir: bool + + +@router.get("/cached-model-path", response_model = CachedModelPathResponse) +async def get_cached_model_path( + repo_id: str = Query(..., description = "HuggingFace repo ID"), + variant: str = Query("", description = "Quantization variant (empty for whole repo)"), + current_subject: str = Depends(get_current_subject), +): + """Absolute on-disk path of a cached repo or one of its GGUF variants.""" + if not _is_valid_repo_id(repo_id): + raise HTTPException(status_code = 400, detail = "Invalid repo_id format") + path = await asyncio.to_thread(_resolve_cached_model_path, repo_id, variant.strip() or None) + return {"path": str(path), "is_dir": path.is_dir()} + + +@router.post("/reveal-cached-model") +async def reveal_cached_model( + repo_id: str = Body(...), + variant: Optional[str] = Body(None), + current_subject: str = Depends(get_current_subject), +): + """Reveal a cached repo (or one GGUF variant's file) in the OS file manager.""" + if not _is_valid_repo_id(repo_id): + raise HTTPException(status_code = 400, detail = "Invalid repo_id format") + variant = (variant or "").strip() or None + path = await asyncio.to_thread(_resolve_cached_model_path, repo_id, variant) + try: + await asyncio.to_thread(_reveal_in_file_manager, path) + except Exception as e: + logger.error(f"Failed to reveal {path}: {e}") + raise HTTPException(status_code = 500, detail = "Failed to open file manager") + return {"status": "ok", "path": str(path)} + + @router.get("/checkpoints", response_model = CheckpointListResponse) async def list_checkpoints( outputs_dir: str = Query( diff --git a/studio/backend/routes/rag.py b/studio/backend/routes/rag.py index e20fea74a3..392a4e0d02 100644 --- a/studio/backend/routes/rag.py +++ b/studio/backend/routes/rag.py @@ -318,6 +318,39 @@ def list_project_documents(project_id: str, subject: str = Depends(get_current_s conn.close() +@router.get("/documents") +def list_all_uploaded_documents(subject: str = Depends(get_current_subject)) -> dict: + """Every uploaded file across chats, projects, and knowledge bases (settings + Data tab).""" + _require_rag() + conn = rag_db.get_connection() + try: + docs = store.list_all_documents(conn) + kb_names = {kb["id"]: kb["name"] for kb in store.list_kbs(conn)} + finally: + conn.close() + + from storage.studio_db import list_chat_projects + + project_names = {p["id"]: p["name"] for p in list_chat_projects(include_archived = True)} + + out = [] + for doc in docs: + view = _doc_view(doc) + stored_path = doc.get("stored_path") + size = None + if stored_path: + try: + size = os.path.getsize(stored_path) + except OSError: + size = None + view["sizeBytes"] = size + view["kbName"] = kb_names.get(doc.get("kb_id")) + view["projectName"] = project_names.get(doc.get("project_id")) + out.append(view) + return {"documents": out} + + @router.delete("/documents/{document_id}") def delete_document(document_id: str, subject: str = Depends(get_current_subject)) -> dict: _require_rag() @@ -424,8 +457,10 @@ _CONTENT_TYPES = { ".txt": "text/plain; charset=utf-8", ".md": "text/markdown; charset=utf-8", ".markdown": "text/markdown; charset=utf-8", - ".html": "text/html; charset=utf-8", - ".htm": "text/html; charset=utf-8", + # Served as plain text, never text/html: an uploaded HTML document rendered + # same-origin would execute its scripts with access to the app's storage. + ".html": "text/plain; charset=utf-8", + ".htm": "text/plain; charset=utf-8", ".docx": "application/vnd.openxmlformats-officedocument.wordprocessingml.document", } diff --git a/studio/backend/routes/settings.py b/studio/backend/routes/settings.py index ab0fd2fd99..17e64df918 100644 --- a/studio/backend/routes/settings.py +++ b/studio/backend/routes/settings.py @@ -36,9 +36,10 @@ from utils.helper_precache_settings import ( ) from utils.coding_agents import CODING_AGENTS, detect_installed_coding_agents from utils.openai_auto_switch_settings import ( - DEFAULT_AUTO_UNLOAD_IDLE_SECONDS, + DEFAULT_AUTO_UNLOAD_KEEP_KV, DEFAULT_OPENAI_AUTO_SWITCH_ENABLED, get_auto_unload_idle_seconds, + get_auto_unload_keep_kv, get_model_overrides, get_openai_auto_switch_enabled, get_stored_auto_unload_idle_seconds, @@ -90,7 +91,9 @@ class HelperPrecacheResponse(BaseModel): class OpenAIAutoSwitchPayload(BaseModel): enabled: bool - auto_unload_idle_seconds: int = Field(default = DEFAULT_AUTO_UNLOAD_IDLE_SECONDS, ge = 0) + # None leaves the stored value untouched (partial updates can't clobber it). + auto_unload_idle_seconds: Optional[int] = Field(default = None, ge = 0) + auto_unload_keep_kv: Optional[bool] = None class OpenAIAutoSwitchResponse(BaseModel): @@ -101,6 +104,7 @@ class OpenAIAutoSwitchResponse(BaseModel): # UNSLOTH_MODEL_IDLE_TTL set and nothing stored, this is true even while enabled # is false, so the UI can show idle-unload as active instead of "needs enable". idle_unload_active: bool = False + auto_unload_keep_kv: bool = DEFAULT_AUTO_UNLOAD_KEEP_KV class ModelOverridePayload(BaseModel): @@ -198,6 +202,7 @@ def get_openai_auto_switch( enabled = get_openai_auto_switch_enabled(), auto_unload_idle_seconds = get_stored_auto_unload_idle_seconds(), idle_unload_active = get_auto_unload_idle_seconds() > 0, + auto_unload_keep_kv = get_auto_unload_keep_kv(), ) @@ -206,8 +211,8 @@ def update_openai_auto_switch( payload: OpenAIAutoSwitchPayload, current_subject: str = Depends(get_current_subject) ) -> OpenAIAutoSwitchResponse: try: - enabled, idle_seconds = set_openai_auto_switch( - payload.enabled, payload.auto_unload_idle_seconds + enabled, idle_seconds, keep_kv = set_openai_auto_switch( + payload.enabled, payload.auto_unload_idle_seconds, payload.auto_unload_keep_kv ) except ValueError as exc: raise log_and_http_error( @@ -217,10 +222,16 @@ def update_openai_auto_switch( event = "settings.update_openai_auto_switch_failed", log = logger, ) from exc + idle_unload_active = get_auto_unload_idle_seconds() > 0 + if not keep_kv or not idle_unload_active: + # Keep-KV off or idle unload disabled: drop already-saved chat context too. + from core.inference.llama_keepwarm import purge_kv_resume + purge_kv_resume() return OpenAIAutoSwitchResponse( enabled = enabled, auto_unload_idle_seconds = idle_seconds, - idle_unload_active = get_auto_unload_idle_seconds() > 0, + idle_unload_active = idle_unload_active, + auto_unload_keep_kv = keep_kv, ) diff --git a/studio/backend/routes/training.py b/studio/backend/routes/training.py index 53b1c4d991..a8a9874b1b 100644 --- a/studio/backend/routes/training.py +++ b/studio/backend/routes/training.py @@ -196,6 +196,7 @@ async def start_training( request.local_eval_datasets, "Local eval dataset" ) resume_output_dir: Optional[str] = None + resume_run: Optional[dict] = None if request.resume_from_checkpoint: try: resume_output_dir = normalize_resume_output_dir(request.resume_from_checkpoint) @@ -208,7 +209,7 @@ async def start_training( if not resume_run or not can_resume_run(resume_run): raise HTTPException( status_code = 400, - detail = "Resume checkpoint must belong to a stopped run with saved trainer state.", + detail = "Resume checkpoint must belong to a stopped or errored run with complete saved trainer state.", ) resume_checkpoint = get_resume_checkpoint_path(resume_output_dir) if not resume_checkpoint: @@ -458,7 +459,10 @@ async def start_training( try: success = backend.start_training( - job_id = job_id, before_spawn = _free_vram_for_training, **training_kwargs + job_id = job_id, + before_spawn = _free_vram_for_training, + resume_source_run_id = resume_run["id"] if resume_run else None, + **training_kwargs, ) except SidecarSwapInProgress as exc: # Expected loss of the race against a sidecar install: a retryable @@ -521,7 +525,10 @@ async def stop_training( status = "idle", message = "No training job is currently running" ) - backend.stop_training(save = body.save) + if not backend.stop_training(save = body.save): + return TrainingStopResponse( + status = "idle", message = "No training job is currently running" + ) return TrainingStopResponse( status = "stopped", @@ -637,9 +644,9 @@ async def get_training_status(current_subject: str = Depends(get_current_subject "loss": getattr(progress, "loss", None), "learning_rate": getattr(progress, "learning_rate", None), } - output_dir = getattr(backend, "_output_dir", None) - if output_dir: - details["output_dir"] = output_dir + # Always present: an explicit null tells the client to drop a cached + # path (stop without save clears the run's output_dir). + details["output_dir"] = getattr(backend, "_output_dir", None) or None # Metric history for chart recovery after SSE reconnection. metric_history = None diff --git a/studio/backend/routes/training_vram.py b/studio/backend/routes/training_vram.py index fb361d3359..fd96fe2175 100644 --- a/studio/backend/routes/training_vram.py +++ b/studio/backend/routes/training_vram.py @@ -197,15 +197,18 @@ def can_load_chat_during_training( requested_gpu_ids: Optional[List[int]], is_gguf: bool = False, required_override_gb: Optional[float] = None, + single_device_gpu: Optional[str] = None, ) -> Tuple[bool, Dict[str, Any]]: """Decide if a NEW chat model can load without OOMing active training (inverse of can_keep_chat_during_training: training is already resident, so size the chat model against the free VRAM that remains). Sizes/places it the same way the loader will: HF auto reuses auto_select_gpu_ids; HF explicit requires an even-share per-GPU floor for device_map="balanced"; GGUF sizes from - required_override_gb over the visible pool. `load_in_4bit` must be effective - (LoRA can flip 4-bit -> 16-bit). Non-CUDA allows the load; default-deny on any - CUDA case it can't size, so a load never OOMs training.""" + required_override_gb over the visible pool. ``single_device_gpu`` is the + exact physical device token selected by a single-device runner. + `load_in_4bit` must be effective (LoRA can flip 4-bit -> 16-bit). Non-CUDA + allows the load; default-deny on any CUDA case it can't size, so a load never + OOMs training.""" try: from utils.hardware import ( DeviceType, @@ -251,26 +254,49 @@ def can_load_chat_during_training( } # Explicit GPUs, or GGUF: size directly and check live free VRAM. + if single_device_gpu is not None: + mode = "single_device" + elif is_gguf: + mode = "gguf" + else: + mode = "explicit" required_gb = required_override_gb if required_gb is None: required_gb, _meta = estimate_required_model_memory_gb(model_name, **est_kwargs) if required_gb is None: - mode = "explicit" if requested_gpu_ids else "gguf" return False, {"mode": mode, "reason": "estimate_unavailable"} free_by_index = _free_vram_by_index(get_visible_gpu_utilization().get("devices", [])) - if requested_gpu_ids: + if single_device_gpu is not None: + token = str(single_device_gpu).strip() + if not token: + # Empty token = a CPU-only single-device runner (e.g. a CPU + # diffusion GGUF): it uses no GPU VRAM, so it never threatens + # active training and can always load. + return True, {"mode": "single_device", "reason": "cpu_only"} + try: + selected_gpu = int(token) + if selected_gpu < 0: + raise ValueError + except (TypeError, ValueError): + # A non-numeric device token (e.g. a CUDA UUID / MIG handle) + # can't be mapped to a free-VRAM index, but the runner still + # drives ONE device. Size against the worst-case visible device + # (min free), never the aggregate pool, so a single-device load + # is never OK'd on capacity it can't use and OOMs training. + free_vals = [min(free_by_index.values())] if free_by_index else [] + else: + free_vals = [free_by_index.get(selected_gpu, 0.0)] + elif requested_gpu_ids: # Invalid ids -> load_model 400s first, so don't block; missing id = 0. try: resolved = resolve_requested_gpu_ids(requested_gpu_ids) except ValueError: - return True, {"mode": "explicit", "reason": "invalid_gpu_ids"} + return True, {"mode": mode, "reason": "invalid_gpu_ids"} free_vals = [free_by_index.get(i, 0.0) for i in resolved] - mode = "explicit" else: # GGUF: llama.cpp picks the GPU(s); any visible GPU is a candidate. free_vals = list(free_by_index.values()) - mode = "gguf" if not free_vals: return False, {"mode": mode, "reason": "no_visible_gpus"} diff --git a/studio/backend/storage/studio_db.py b/studio/backend/storage/studio_db.py index 4e0c711b69..6972e7b7ff 100644 --- a/studio/backend/storage/studio_db.py +++ b/studio/backend/storage/studio_db.py @@ -7,6 +7,7 @@ Like auth/storage.py (module-level functions, raw sqlite3, per-function connections) plus WAL mode and PRAGMA foreign_keys = ON for CASCADE deletes. """ +import hashlib import json import logging import os @@ -100,6 +101,7 @@ _schema_lock = threading.Lock() _schema_ready = False _SQLITE_IN_CHUNK_SIZE = 900 _PROJECT_WORKSPACE_SUBDIRS = ("sandbox",) +_CHAT_ATTACHMENT_INVENTORY_VERSION = 1 def _project_slug(name: str) -> str: @@ -190,13 +192,18 @@ def _ensure_schema(conn: sqlite3.Connection) -> None: error_message TEXT, duration_seconds REAL, loss_sparkline TEXT, - display_name TEXT + display_name TEXT, + resume_blocked INTEGER NOT NULL DEFAULT 0 ) """ ) existing_cols = {row[1] for row in conn.execute("PRAGMA table_info(training_runs)").fetchall()} if "display_name" not in existing_cols: conn.execute("ALTER TABLE training_runs ADD COLUMN display_name TEXT") + if "resume_blocked" not in existing_cols: + conn.execute( + "ALTER TABLE training_runs ADD COLUMN resume_blocked INTEGER NOT NULL DEFAULT 0" + ) conn.execute( """ CREATE TABLE IF NOT EXISTS training_metrics ( @@ -313,6 +320,141 @@ def _ensure_schema(conn: sqlite3.Connection) -> None: ) """ ) + tombstone_schema = """ + CREATE TABLE chat_attachment_tombstones ( + thread_id TEXT NOT NULL REFERENCES chat_threads(id) ON DELETE CASCADE, + message_id TEXT NOT NULL, + attachment_id TEXT NOT NULL, + deleted_at INTEGER NOT NULL, + PRIMARY KEY(thread_id, message_id, attachment_id) + ) WITHOUT ROWID + """ + tombstone_table = conn.execute( + """ + SELECT 1 FROM sqlite_master + WHERE type = 'table' AND name = 'chat_attachment_tombstones' + """ + ).fetchone() + if tombstone_table is None: + conn.execute(tombstone_schema) + else: + tombstone_columns = { + row[1] for row in conn.execute("PRAGMA table_info(chat_attachment_tombstones)") + } + tombstone_fk_targets = { + row[2] for row in conn.execute("PRAGMA foreign_key_list(chat_attachment_tombstones)") + } + if "thread_id" not in tombstone_columns or "chat_threads" not in tombstone_fk_targets: + # The first implementation cascaded through chat_messages, which + # erased deletion knowledge during pruneMissing. Rebuild once, + # retaining every tombstone whose owning thread still exists. + conn.execute("SAVEPOINT migrate_chat_attachment_tombstones") + try: + conn.execute( + "ALTER TABLE chat_attachment_tombstones " + "RENAME TO chat_attachment_tombstones_legacy" + ) + conn.execute(tombstone_schema) + if "thread_id" in tombstone_columns: + conn.execute( + """ + INSERT OR IGNORE INTO chat_attachment_tombstones + (thread_id, message_id, attachment_id, deleted_at) + SELECT legacy.thread_id, legacy.message_id, + legacy.attachment_id, legacy.deleted_at + FROM chat_attachment_tombstones_legacy legacy + JOIN chat_threads thread ON thread.id = legacy.thread_id + """ + ) + else: + conn.execute( + """ + INSERT OR IGNORE INTO chat_attachment_tombstones + (thread_id, message_id, attachment_id, deleted_at) + SELECT message.thread_id, legacy.message_id, + legacy.attachment_id, legacy.deleted_at + FROM chat_attachment_tombstones_legacy legacy + JOIN chat_messages message ON message.id = legacy.message_id + """ + ) + conn.execute("DROP TABLE chat_attachment_tombstones_legacy") + conn.execute("RELEASE SAVEPOINT migrate_chat_attachment_tombstones") + except Exception: + conn.execute("ROLLBACK TO SAVEPOINT migrate_chat_attachment_tombstones") + conn.execute("RELEASE SAVEPOINT migrate_chat_attachment_tombstones") + raise + conn.execute( + """ + CREATE TABLE IF NOT EXISTS chat_attachment_inventory ( + message_id TEXT NOT NULL REFERENCES chat_messages(id) ON DELETE CASCADE, + attachment_id TEXT NOT NULL, + name TEXT NOT NULL, + type TEXT, + content_type TEXT, + size_bytes INTEGER, + PRIMARY KEY(message_id, attachment_id) + ) WITHOUT ROWID + """ + ) + conn.execute( + """ + CREATE TABLE IF NOT EXISTS chat_attachment_inventory_state ( + singleton INTEGER NOT NULL PRIMARY KEY CHECK(singleton = 1), + inventory_version INTEGER NOT NULL DEFAULT 0, + dirty INTEGER NOT NULL DEFAULT 1, + backfilled_at INTEGER NOT NULL + ) + """ + ) + inventory_state_columns = { + row[1] for row in conn.execute("PRAGMA table_info(chat_attachment_inventory_state)") + } + if "inventory_version" not in inventory_state_columns: + conn.execute( + "ALTER TABLE chat_attachment_inventory_state " + "ADD COLUMN inventory_version INTEGER NOT NULL DEFAULT 0" + ) + if "dirty" not in inventory_state_columns: + conn.execute( + "ALTER TABLE chat_attachment_inventory_state " + "ADD COLUMN dirty INTEGER NOT NULL DEFAULT 1" + ) + conn.execute( + """ + CREATE TRIGGER IF NOT EXISTS chat_attachment_inventory_dirty_insert + AFTER INSERT ON chat_messages + BEGIN + INSERT INTO chat_attachment_inventory_state + (singleton, inventory_version, dirty, backfilled_at) + VALUES (1, 0, 1, 0) + ON CONFLICT(singleton) DO UPDATE SET dirty = 1; + END + """ + ) + conn.execute( + """ + CREATE TRIGGER IF NOT EXISTS chat_attachment_inventory_dirty_update + AFTER UPDATE ON chat_messages + BEGIN + INSERT INTO chat_attachment_inventory_state + (singleton, inventory_version, dirty, backfilled_at) + VALUES (1, 0, 1, 0) + ON CONFLICT(singleton) DO UPDATE SET dirty = 1; + END + """ + ) + conn.execute( + """ + CREATE TRIGGER IF NOT EXISTS chat_attachment_inventory_dirty_delete + AFTER DELETE ON chat_messages + BEGIN + INSERT INTO chat_attachment_inventory_state + (singleton, inventory_version, dirty, backfilled_at) + VALUES (1, 0, 1, 0) + ON CONFLICT(singleton) DO UPDATE SET dirty = 1; + END + """ + ) conn.execute( "CREATE INDEX IF NOT EXISTS idx_chat_threads_model_type_created_at ON chat_threads(model_type, created_at)" ) @@ -391,6 +533,21 @@ def _ensure_schema(conn: sqlite3.Connection) -> None: conn.execute( "CREATE INDEX IF NOT EXISTS idx_prompt_lists_created_at ON prompt_lists(created_at)" ) + inventory_state = conn.execute( + """ + SELECT inventory_version, dirty + FROM chat_attachment_inventory_state + WHERE singleton = 1 + """ + ).fetchone() + if ( + inventory_state is None + or inventory_state["inventory_version"] != _CHAT_ATTACHMENT_INVENTORY_VERSION + or inventory_state["dirty"] + ): + _rebuild_chat_attachment_inventory(conn) + _mark_chat_attachment_inventory_clean(conn) + conn.commit() def _prompt_entry_from_row(row: sqlite3.Row) -> dict: @@ -582,16 +739,43 @@ def create_run( config_json: str, started_at: str, total_steps: Optional[int], + *, + output_dir: Optional[str] = None, + cancel_requested: bool = False, + resumed_from_run_id: Optional[str] = None, ) -> None: conn = get_connection() try: conn.execute( """ - INSERT INTO training_runs (id, model_name, dataset_name, config_json, started_at, total_steps) - VALUES (?, ?, ?, ?, ?, ?) + INSERT INTO training_runs ( + id, model_name, dataset_name, config_json, started_at, total_steps, + output_dir, resume_blocked + ) + VALUES (?, ?, ?, ?, ?, ?, ?, ?) """, - (id, model_name, dataset_name, config_json, started_at, total_steps), + ( + id, + model_name, + dataset_name, + config_json, + started_at, + total_steps, + None if cancel_requested else output_dir, + int(cancel_requested), + ), ) + if resumed_from_run_id: + claimed = conn.execute( + """ + UPDATE training_runs SET resume_blocked = 1 + WHERE id = ? AND status IN ('stopped', 'error') + AND output_dir = ? AND resume_blocked = 0 + """, + (resumed_from_run_id, output_dir), + ) + if claimed.rowcount != 1: + raise RuntimeError("Resume source is no longer available") conn.commit() finally: conn.close() @@ -634,6 +818,8 @@ def finish_run( loss_sparkline: Optional[str] = None, output_dir: Optional[str] = None, error_message: Optional[str] = None, + clear_output_dir: bool = False, + resume_blocked: bool = False, ) -> None: conn = get_connection() try: @@ -641,9 +827,16 @@ def finish_run( """ UPDATE training_runs SET status = ?, ended_at = ?, final_step = ?, final_loss = ?, - duration_seconds = ?, loss_sparkline = ?, output_dir = ?, - error_message = ? - WHERE id = ? + duration_seconds = ?, loss_sparkline = ?, + output_dir = CASE + WHEN resume_blocked = 1 OR ? = 1 THEN NULL + WHEN ? IS NOT NULL THEN ? + WHEN ? IN ('error', 'stopped') THEN output_dir + ELSE NULL + END, + error_message = ?, + resume_blocked = CASE WHEN resume_blocked = 1 OR ? = 1 THEN 1 ELSE ? END + WHERE id = ? AND status = 'running' """, ( status, @@ -652,8 +845,13 @@ def finish_run( final_loss, duration_seconds, loss_sparkline, + int(clear_output_dir), output_dir, + output_dir, + status, error_message, + int(clear_output_dir), + int(resume_blocked), id, ), ) @@ -713,6 +911,38 @@ def update_run_display_name(id: str, display_name: Optional[str]) -> None: conn.close() +def update_run_output_dir(id: str, output_dir: Optional[str]) -> None: + conn = get_connection() + try: + conn.execute( + """ + UPDATE training_runs SET output_dir = ? + WHERE id = ? AND status = 'running' AND resume_blocked = 0 + """, + (output_dir, id), + ) + conn.commit() + finally: + conn.close() + + +def mark_run_cancel_requested(id: str) -> bool: + """Clear resume/export state only while the exact run is still active.""" + conn = get_connection() + try: + cursor = conn.execute( + """ + UPDATE training_runs SET output_dir = NULL, resume_blocked = 1 + WHERE id = ? AND status = 'running' + """, + (id,), + ) + conn.commit() + return cursor.rowcount > 0 + finally: + conn.close() + + def list_runs(limit: int = 50, offset: int = 0) -> dict: conn = get_connection() try: @@ -722,15 +952,15 @@ def list_runs(limit: int = 50, offset: int = 0) -> dict: SELECT r.id, r.status, r.model_name, r.dataset_name, r.started_at, r.ended_at, r.total_steps, r.final_step, r.final_loss, r.output_dir, r.duration_seconds, r.error_message, - r.loss_sparkline, r.display_name, r.config_json, + r.loss_sparkline, r.display_name, r.config_json, r.resume_blocked, CASE - WHEN r.status = 'stopped' + WHEN r.status IN ('stopped', 'error') AND r.output_dir IS NOT NULL AND EXISTS ( SELECT 1 FROM training_runs newer WHERE newer.output_dir = r.output_dir - AND newer.status IN ('stopped', 'completed') + AND newer.status IN ('stopped', 'completed', 'error', 'running') AND newer.started_at > r.started_at ) THEN 1 ELSE 0 @@ -765,13 +995,13 @@ def get_run(id: str) -> Optional[dict]: """ SELECT r.*, CASE - WHEN r.status = 'stopped' + WHEN r.status IN ('stopped', 'error') AND r.output_dir IS NOT NULL AND EXISTS ( SELECT 1 FROM training_runs newer WHERE newer.output_dir = r.output_dir - AND newer.status IN ('stopped', 'completed') + AND newer.status IN ('stopped', 'completed', 'error', 'running') AND newer.started_at > r.started_at ) THEN 1 ELSE 0 @@ -806,12 +1036,12 @@ def get_resumable_run_by_output_dir(output_dir: str) -> Optional[dict]: 0 AS resumed_later FROM training_runs r WHERE r.output_dir = ? - AND r.status = 'stopped' + AND r.status IN ('stopped', 'error') AND NOT EXISTS ( SELECT 1 FROM training_runs newer WHERE newer.output_dir = r.output_dir - AND newer.status IN ('stopped', 'completed') + AND newer.status IN ('stopped', 'completed', 'error', 'running') AND newer.started_at > r.started_at ) ORDER BY r.started_at DESC @@ -914,8 +1144,12 @@ def cleanup_orphaned_runs() -> None: conn.execute( """ UPDATE training_runs - SET status = 'error', - error_message = 'Server restarted during training', + SET status = CASE WHEN resume_blocked = 1 THEN 'stopped' ELSE 'error' END, + error_message = CASE + WHEN resume_blocked = 1 THEN NULL + ELSE 'Server restarted during training' + END, + output_dir = CASE WHEN resume_blocked = 1 THEN NULL ELSE output_dir END, ended_at = ? WHERE status = 'running' """, @@ -1219,7 +1453,14 @@ def delete_chat_threads(ids: list[str]) -> None: return conn = get_connection() try: + conn.execute("BEGIN IMMEDIATE") + _ensure_chat_attachment_inventory_current(conn) + conn.executemany( + "DELETE FROM chat_attachment_tombstones WHERE thread_id = ?", + [(id,) for id in ids], + ) conn.executemany("DELETE FROM chat_threads WHERE id = ?", [(id,) for id in ids]) + _mark_chat_attachment_inventory_clean(conn) conn.commit() finally: conn.close() @@ -1228,7 +1469,11 @@ def delete_chat_threads(ids: list[str]) -> None: def clear_chat_history() -> None: conn = get_connection() try: + conn.execute("BEGIN IMMEDIATE") + _ensure_chat_attachment_inventory_current(conn) + conn.execute("DELETE FROM chat_attachment_tombstones") conn.execute("DELETE FROM chat_threads") + _mark_chat_attachment_inventory_clean(conn) conn.commit() finally: conn.close() @@ -1354,6 +1599,7 @@ def delete_chat_project(id: str, delete_files: bool = False) -> Optional[dict]: conn = get_connection() try: conn.execute("BEGIN IMMEDIATE") + _ensure_chat_attachment_inventory_current(conn) row = conn.execute("SELECT * FROM chat_projects WHERE id = ?", (id,)).fetchone() if row is None: conn.rollback() @@ -1361,6 +1607,7 @@ def delete_chat_project(id: str, delete_files: bool = False) -> Optional[dict]: project = _chat_project_from_row(row) conn.execute("DELETE FROM chat_threads WHERE project_id = ?", (id,)) conn.execute("DELETE FROM chat_projects WHERE id = ?", (id,)) + _mark_chat_attachment_inventory_clean(conn) conn.commit() if delete_files: _delete_project_workspace(project) @@ -1483,15 +1730,285 @@ def _recompute_chat_thread_updated_at(conn: sqlite3.Connection, thread_id: str) ) +_CONTENT_PART_ID_PREFIX = "content-part-sha256-" +_URI_SCHEME_RE = re.compile(r"^[A-Za-z][A-Za-z0-9+.-]*:") + + +def _is_locally_stored_blob(value: str) -> bool: + """True for data URIs or bare base64, never external/blob URI references.""" + candidate = value.lstrip() + if not candidate: + return False + if candidate[:5].lower() == "data:": + return True + if candidate.startswith(("//", "\\\\")): + return False + return _URI_SCHEME_RE.match(candidate) is None + + +def _managed_content_part_payload(part: dict) -> Optional[tuple[str, Any]]: + """Return the locally stored blob payload used to identify a content part.""" + image = part.get("image") + if isinstance(image, str) and image[:5].lower() == "data:": + return "image", image + + audio = part.get("audio") + if isinstance(audio, str) and _is_locally_stored_blob(audio): + return "audio", audio + if isinstance(audio, dict): + data = audio.get("data") + if isinstance(data, str) and _is_locally_stored_blob(data): + return "audio", audio + return None + + +def _content_part_id(part: dict) -> Optional[str]: + """Stable managed id derived from blob data, without mutating inference content.""" + payload = _managed_content_part_payload(part) + if payload is None: + return None + canonical = json.dumps( + payload, + ensure_ascii = False, + separators = (",", ":"), + sort_keys = True, + ).encode("utf-8") + return f"{_CONTENT_PART_ID_PREFIX}{hashlib.sha256(canonical).hexdigest()}" + + +def _chat_attachment_tombstones_for_messages( + conn: sqlite3.Connection, thread_id: str, message_ids: list[str] +) -> dict[str, set[str]]: + tombstones = {message_id: set() for message_id in message_ids} + unique_ids = list(dict.fromkeys(message_ids)) + for start in range(0, len(unique_ids), _SQLITE_IN_CHUNK_SIZE): + chunk = unique_ids[start : start + _SQLITE_IN_CHUNK_SIZE] + placeholders = ",".join("?" for _ in chunk) + rows = conn.execute( + f""" + SELECT message_id, attachment_id + FROM chat_attachment_tombstones + WHERE thread_id = ? AND message_id IN ({placeholders}) + """, + (thread_id, *chunk), + ).fetchall() + for row in rows: + tombstones[row["message_id"]].add(row["attachment_id"]) + return tombstones + + +def _reconcile_chat_message_uploads(message: dict, tombstones: set[str]) -> dict: + """Strip uploads previously deleted through the Data tab from a stale write.""" + if not tombstones: + return message + + reconciled = dict(message) + attachments = message.get("attachments") + if isinstance(attachments, list): + reconciled["attachments"] = [ + attachment + for attachment in attachments + if not (isinstance(attachment, dict) and str(attachment.get("id") or "") in tombstones) + ] + + content = message.get("content") + if isinstance(content, list): + reconciled["content"] = [ + part + for part in content + if not (isinstance(part, dict) and (_content_part_id(part) or "") in tombstones) + ] + return reconciled + + +def _chat_attachment_metadata_text(value, fallback: Optional[str] = None) -> Optional[str]: + """Keep untyped legacy/import metadata safe for SQLite binding.""" + if value is None: + return fallback + if isinstance(value, str): + return value or fallback + if isinstance(value, (bool, int, float)): + return str(value) + # Objects and arrays are not useful display metadata and sqlite3 rejects + # binding them directly. + return fallback + + +def _chat_attachment_inventory_entries( + attachments_json: Optional[str], + content_json: Optional[str], + tombstones: Optional[set[str]] = None, +) -> list[dict]: + tombstones = tombstones or set() + attachments = _json_loads(attachments_json, None) + if not isinstance(attachments, list): + attachments = [] + attachments = [ + attachment + for attachment in attachments + if isinstance(attachment, dict) and attachment.get("id") + ] + attachments.extend(_content_part_attachments(content_json)) + + entries: list[dict] = [] + seen: set[str] = set() + for attachment in attachments: + attachment_id = str(attachment["id"]) + if attachment_id in seen or attachment_id in tombstones: + continue + seen.add(attachment_id) + entries.append( + { + "id": attachment_id, + "name": _chat_attachment_metadata_text(attachment.get("name"), "attachment"), + "type": _chat_attachment_metadata_text(attachment.get("type")), + "contentType": _chat_attachment_metadata_text(attachment.get("contentType")), + "sizeBytes": _chat_attachment_size_bytes(attachment), + } + ) + return entries + + +def _replace_chat_attachment_inventory( + conn: sqlite3.Connection, + message_id: str, + attachments_json: Optional[str], + content_json: Optional[str], + tombstones: Optional[set[str]] = None, +) -> None: + conn.execute("DELETE FROM chat_attachment_inventory WHERE message_id = ?", (message_id,)) + entries = _chat_attachment_inventory_entries( + attachments_json, + content_json, + tombstones, + ) + conn.executemany( + """ + INSERT INTO chat_attachment_inventory + (message_id, attachment_id, name, type, content_type, size_bytes) + VALUES (?, ?, ?, ?, ?, ?) + """, + [ + ( + message_id, + entry["id"], + entry["name"], + entry["type"], + entry["contentType"], + entry["sizeBytes"], + ) + for entry in entries + ], + ) + + +def _mark_chat_attachment_inventory_clean(conn: sqlite3.Connection) -> None: + conn.execute( + """ + INSERT INTO chat_attachment_inventory_state + (singleton, inventory_version, dirty, backfilled_at) + VALUES (1, ?, 0, ?) + ON CONFLICT(singleton) DO UPDATE SET + inventory_version = excluded.inventory_version, + dirty = 0, + backfilled_at = excluded.backfilled_at + """, + ( + _CHAT_ATTACHMENT_INVENTORY_VERSION, + int(datetime.now(timezone.utc).timestamp() * 1000), + ), + ) + + +def _rebuild_chat_attachment_inventory(conn: sqlite3.Connection) -> None: + """Rebuild after schema upgrade or a write from an older Studio build.""" + conn.execute("DELETE FROM chat_attachment_inventory") + tombstones: dict[tuple[str, str], set[str]] = {} + for row in conn.execute( + "SELECT thread_id, message_id, attachment_id FROM chat_attachment_tombstones" + ).fetchall(): + tombstones.setdefault((row["thread_id"], row["message_id"]), set()).add( + row["attachment_id"] + ) + rows = conn.execute( + "SELECT id, thread_id, attachments_json, content_json FROM chat_messages" + ).fetchall() + for row in rows: + _replace_chat_attachment_inventory( + conn, + row["id"], + row["attachments_json"], + row["content_json"], + tombstones.get((row["thread_id"], row["id"]), set()), + ) + + +def _ensure_chat_attachment_inventory_current(conn: sqlite3.Connection) -> None: + state = conn.execute( + """ + SELECT inventory_version, dirty + FROM chat_attachment_inventory_state + WHERE singleton = 1 + """ + ).fetchone() + if ( + state is not None + and state["inventory_version"] == _CHAT_ATTACHMENT_INVENTORY_VERSION + and not state["dirty"] + ): + return + + owns_transaction = not conn.in_transaction + if owns_transaction: + conn.execute("BEGIN IMMEDIATE") + try: + state = conn.execute( + """ + SELECT inventory_version, dirty + FROM chat_attachment_inventory_state + WHERE singleton = 1 + """ + ).fetchone() + if ( + state is None + or state["inventory_version"] != _CHAT_ATTACHMENT_INVENTORY_VERSION + or state["dirty"] + ): + _rebuild_chat_attachment_inventory(conn) + _mark_chat_attachment_inventory_clean(conn) + if owns_transaction: + conn.commit() + except Exception: + if owns_transaction: + conn.rollback() + raise + + def upsert_chat_message(message: dict) -> dict: conn = get_connection() try: conn.execute("BEGIN IMMEDIATE") + _ensure_chat_attachment_inventory_current(conn) _raise_if_chat_message_thread_conflicts( conn, message["threadId"], [message["id"]], ) + tombstones = _chat_attachment_tombstones_for_messages( + conn, + message["threadId"], + [message["id"]], + ) + reconciled = _reconcile_chat_message_uploads( + message, + tombstones.get(message["id"], set()), + ) + content_json = json.dumps(reconciled.get("content", [])) + attachments_json = ( + json.dumps(reconciled.get("attachments")) + if reconciled.get("attachments") is not None + else None + ) conn.execute( """ INSERT INTO chat_messages @@ -1507,23 +2024,32 @@ def upsert_chat_message(message: dict) -> dict: WHERE excluded.thread_id = chat_messages.thread_id """, ( - message["id"], - message["threadId"], - message.get("parentId"), - message["role"], - json.dumps(message.get("content", [])), - json.dumps(message.get("attachments")) - if message.get("attachments") is not None + reconciled["id"], + reconciled["threadId"], + reconciled.get("parentId"), + reconciled["role"], + content_json, + attachments_json, + json.dumps(reconciled.get("metadata")) + if reconciled.get("metadata") is not None else None, - json.dumps(message.get("metadata")) - if message.get("metadata") is not None - else None, - int(message["createdAt"]), + int(reconciled["createdAt"]), ), ) - _bump_chat_thread_updated_at(conn, message["threadId"], int(message["createdAt"])) + _replace_chat_attachment_inventory( + conn, + reconciled["id"], + attachments_json, + content_json, + ) + _bump_chat_thread_updated_at( + conn, + reconciled["threadId"], + int(reconciled["createdAt"]), + ) + _mark_chat_attachment_inventory_clean(conn) conn.commit() - return message + return reconciled except Exception: conn.rollback() raise @@ -1539,13 +2065,28 @@ def sync_chat_messages( conn = get_connection() try: conn.execute("BEGIN IMMEDIATE") + _ensure_chat_attachment_inventory_current(conn) _raise_if_chat_message_thread_conflicts( conn, thread_id, [m["id"] for m in messages], ) - if prune_missing: - conn.execute("DELETE FROM chat_messages WHERE thread_id = ?", (thread_id,)) + tombstones = _chat_attachment_tombstones_for_messages( + conn, + thread_id, + [m["id"] for m in messages], + ) + reconciled_messages = [ + _reconcile_chat_message_uploads(m, tombstones.get(m["id"], set())) for m in messages + ] + serialized_messages = [ + ( + m, + json.dumps(m.get("content", [])), + json.dumps(m.get("attachments")) if m.get("attachments") is not None else None, + ) + for m in reconciled_messages + ] conn.executemany( """ INSERT INTO chat_messages @@ -1566,20 +2107,46 @@ def sync_chat_messages( thread_id, m.get("parentId"), m["role"], - json.dumps(m.get("content", [])), - json.dumps(m.get("attachments")) if m.get("attachments") is not None else None, + content_json, + attachments_json, json.dumps(m.get("metadata")) if m.get("metadata") is not None else None, int(m["createdAt"]), ) - for m in messages + for m, content_json, attachments_json in serialized_messages ], ) - if prune_missing: - _recompute_chat_thread_updated_at(conn, thread_id) - elif messages: - _bump_chat_thread_updated_at( - conn, thread_id, max(int(m["createdAt"]) for m in messages) + for m, content_json, attachments_json in serialized_messages: + _replace_chat_attachment_inventory( + conn, + m["id"], + attachments_json, + content_json, ) + if prune_missing: + retained_ids = {m["id"] for m in reconciled_messages} + existing_ids = { + row["id"] + for row in conn.execute( + "SELECT id FROM chat_messages WHERE thread_id = ?", + (thread_id,), + ).fetchall() + } + missing_ids = sorted(existing_ids - retained_ids) + for start in range(0, len(missing_ids), _SQLITE_IN_CHUNK_SIZE): + chunk = missing_ids[start : start + _SQLITE_IN_CHUNK_SIZE] + placeholders = ",".join("?" for _ in chunk) + conn.execute( + f"DELETE FROM chat_messages WHERE thread_id = ? AND id IN ({placeholders})", + (thread_id, *chunk), + ) + _recompute_chat_thread_updated_at(conn, thread_id) + elif reconciled_messages: + _bump_chat_thread_updated_at( + conn, + thread_id, + max(int(m["createdAt"]) for m in reconciled_messages), + ) + _mark_chat_attachment_inventory_clean(conn) conn.commit() return list_chat_messages(thread_id) except ChatMessageConflictError: @@ -1613,6 +2180,7 @@ def fork_chat_thread( conn = get_connection() try: conn.execute("BEGIN IMMEDIATE") + _ensure_chat_attachment_inventory_current(conn) src = conn.execute( "SELECT * FROM chat_threads WHERE id = ?", (source_thread_id,) ).fetchone() @@ -1686,6 +2254,14 @@ def fork_chat_thread( for row in ancestry ], ) + for row in ancestry: + _replace_chat_attachment_inventory( + conn, + id_map[row["id"]], + row["attachments_json"], + row["content_json"], + ) + _mark_chat_attachment_inventory_clean(conn) conn.commit() thread_row = conn.execute( "SELECT * FROM chat_threads WHERE id = ?", (new_thread_id,) @@ -1744,6 +2320,279 @@ def get_chat_message(thread_id: str, message_id: str) -> Optional[dict]: conn.close() +def _blob_part_base64_len(part: dict) -> int: + """Base64 payload length of an image or audio content part, or 0.""" + image = part.get("image") + if isinstance(image, str) and image[:5].lower() == "data:": + return len(image.rsplit(",", 1)[-1]) + audio = part.get("audio") + if isinstance(audio, str) and _is_locally_stored_blob(audio): + return len(audio.rsplit(",", 1)[-1]) + if isinstance(audio, dict): + data = audio.get("data") + if isinstance(data, str) and _is_locally_stored_blob(data): + return len(data) + return 0 + + +def _chat_attachment_size_bytes(attachment: dict) -> Optional[int]: + """Approximate stored size of one attachment's content parts. + + Image and audio parts hold base64 payloads (decoded bytes ~= 3/4 of the + encoded length); text parts count their character length. None when there + is no sizable content (e.g. a stripped/legacy attachment). + """ + total = 0 + found = False + for part in attachment.get("content") or []: + if not isinstance(part, dict): + continue + blob_len = _blob_part_base64_len(part) + if blob_len > 0: + total += (blob_len * 3) // 4 + found = True + continue + text = part.get("text") + if isinstance(text, str) and text: + total += len(text.encode("utf-8", errors = "ignore")) + found = True + return total if found else None + + +def _content_part_attachments(content_json: Optional[str]) -> list[dict]: + """Managed local blobs stored in content_json, with stable payload ids. + + Exact duplicate blobs intentionally share one inventory id. Deleting that + id removes every identical copy, avoiding ambiguous index-based addressing. + """ + content = _json_loads(content_json, None) + if not isinstance(content, list): + return [] + out: list[dict] = [] + seen: set[str] = set() + for part in content: + if not isinstance(part, dict): + continue + attachment_id = _content_part_id(part) + payload = _managed_content_part_payload(part) + if attachment_id is None or payload is None or attachment_id in seen: + continue + seen.add(attachment_id) + kind, value = payload + content_type = None + if kind == "image" and isinstance(value, str): + content_type = value[5:].split(";", 1)[0].split(",", 1)[0] or None + out.append( + { + "id": attachment_id, + "type": kind, + "name": "Chat image" if kind == "image" else "Chat audio", + "contentType": content_type, + "content": [part], + } + ) + return out + + +def list_chat_attachments_page( + limit: int = 50, offset: int = 0 +) -> tuple[list[dict], Optional[int]]: + """One bounded page from the normalized attachment inventory.""" + if not 1 <= limit <= 100: + raise ValueError("limit must be between 1 and 100") + if offset < 0: + raise ValueError("offset must be non-negative") + + conn = get_connection() + try: + _ensure_chat_attachment_inventory_current(conn) + rows = conn.execute( + """ + SELECT i.attachment_id, i.name, i.type, i.content_type, + i.size_bytes, m.id AS message_id, m.thread_id, + m.created_at, t.title AS thread_title, t.pair_id + FROM chat_attachment_inventory i + JOIN chat_messages m ON m.id = i.message_id + LEFT JOIN chat_threads t ON t.id = m.thread_id + ORDER BY m.created_at DESC, m.id ASC, i.attachment_id ASC + LIMIT ? OFFSET ? + """, + (limit + 1, offset), + ).fetchall() + finally: + conn.close() + + has_more = len(rows) > limit + page_rows = rows[:limit] + attachments = [ + { + "id": row["attachment_id"], + "messageId": row["message_id"], + "threadId": row["thread_id"], + "pairId": row["pair_id"], + "threadTitle": row["thread_title"], + "name": row["name"], + "type": row["type"], + "contentType": row["content_type"], + "sizeBytes": row["size_bytes"], + "createdAt": row["created_at"], + } + for row in page_rows + ] + return attachments, offset + limit if has_more else None + + +def list_chat_attachments() -> list[dict]: + """Compatibility helper returning the full normalized inventory.""" + attachments: list[dict] = [] + offset = 0 + while True: + page, next_offset = list_chat_attachments_page(limit = 100, offset = offset) + attachments.extend(page) + if next_offset is None: + return attachments + offset = next_offset + + +def get_chat_attachment(message_id: str, attachment_id: str) -> Optional[dict]: + """One attachment record (full content) from a message, or None.""" + conn = get_connection() + try: + row = conn.execute( + """ + SELECT message.attachments_json, message.content_json, + EXISTS( + SELECT 1 FROM chat_attachment_tombstones tombstone + WHERE tombstone.thread_id = message.thread_id + AND tombstone.message_id = message.id + AND tombstone.attachment_id = ? + ) AS tombstoned + FROM chat_messages message + WHERE message.id = ? + """, + (attachment_id, message_id), + ).fetchone() + finally: + conn.close() + if row is None or row["tombstoned"]: + return None + attachments = _json_loads(row["attachments_json"], None) + if isinstance(attachments, list): + for attachment in attachments: + if isinstance(attachment, dict) and str(attachment.get("id") or "") == attachment_id: + return attachment + if attachment_id.startswith(_CONTENT_PART_ID_PREFIX): + for attachment in _content_part_attachments(row["content_json"]): + if attachment["id"] == attachment_id: + return attachment + return None + + +def _record_chat_attachment_tombstone( + conn: sqlite3.Connection, thread_id: str, message_id: str, attachment_id: str +) -> None: + conn.execute( + """ + INSERT INTO chat_attachment_tombstones + (thread_id, message_id, attachment_id, deleted_at) + VALUES (?, ?, ?, ?) + ON CONFLICT(thread_id, message_id, attachment_id) DO UPDATE SET + deleted_at = excluded.deleted_at + """, + ( + thread_id, + message_id, + attachment_id, + int(datetime.now(timezone.utc).timestamp() * 1000), + ), + ) + + +def delete_chat_attachment(message_id: str, attachment_id: str) -> bool: + """Remove one stored upload from a message. + + The tombstone is retained while the thread exists, so pruning and later + recreating the same message id cannot restore the deleted upload. If an + ordinary attachment id collides with a content-blob id, both are deleted as + one managed item. + """ + conn = get_connection() + try: + conn.execute("BEGIN IMMEDIATE") + _ensure_chat_attachment_inventory_current(conn) + row = conn.execute( + """ + SELECT thread_id, attachments_json, content_json + FROM chat_messages WHERE id = ? + """, + (message_id,), + ).fetchone() + if row is None: + conn.rollback() + return False + + attachments = _json_loads(row["attachments_json"], None) + updated_attachments_json = row["attachments_json"] + deleted_attachment = False + if isinstance(attachments, list): + remaining_attachments = [ + attachment + for attachment in attachments + if not ( + isinstance(attachment, dict) + and str(attachment.get("id") or "") == attachment_id + ) + ] + deleted_attachment = len(remaining_attachments) != len(attachments) + if deleted_attachment: + updated_attachments_json = json.dumps(remaining_attachments) + + content = _json_loads(row["content_json"], None) + updated_content_json = row["content_json"] + deleted_content = False + if attachment_id.startswith(_CONTENT_PART_ID_PREFIX) and isinstance(content, list): + remaining_content = [ + part + for part in content + if not (isinstance(part, dict) and _content_part_id(part) == attachment_id) + ] + deleted_content = len(remaining_content) != len(content) + if deleted_content: + updated_content_json = json.dumps(remaining_content) + + if not deleted_attachment and not deleted_content: + conn.rollback() + return False + conn.execute( + """ + UPDATE chat_messages + SET attachments_json = ?, content_json = ? + WHERE id = ? + """, + (updated_attachments_json, updated_content_json, message_id), + ) + _record_chat_attachment_tombstone( + conn, + row["thread_id"], + message_id, + attachment_id, + ) + _replace_chat_attachment_inventory( + conn, + message_id, + updated_attachments_json, + updated_content_json, + ) + _mark_chat_attachment_inventory_clean(conn) + conn.commit() + return True + except Exception: + conn.rollback() + raise + finally: + conn.close() + + def list_chat_messages_for_threads(thread_ids: list[str]) -> list[dict]: if not thread_ids: return [] diff --git a/studio/backend/tests/test_chat_attachments.py b/studio/backend/tests/test_chat_attachments.py new file mode 100644 index 0000000000..459587ca9e --- /dev/null +++ b/studio/backend/tests/test_chat_attachments.py @@ -0,0 +1,634 @@ +# SPDX-License-Identifier: AGPL-3.0-only +# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0 + +import base64 +import json +import os +import sqlite3 +import sys + +import pytest +from fastapi import HTTPException + +_backend = os.path.join(os.path.dirname(__file__), "..") +sys.path.insert(0, _backend) + +from routes import chat_history +from storage import studio_db +from utils.paths import studio_db_path + +PNG_BYTES = base64.b64decode( + "iVBORw0KGgoAAAANSUhEUgAAAAEAAAABCAYAAAAfFcSJAAAADUlEQVR42mP8z8BQDwAEhQGAhKmMIQAAAABJRU5ErkJggg==" +) +PNG_DATA_URL = "data:image/png;base64," + base64.b64encode(PNG_BYTES).decode("ascii") + + +def _reset_studio_db(tmp_path, monkeypatch): + monkeypatch.setenv("UNSLOTH_STUDIO_HOME", str(tmp_path)) + monkeypatch.setenv("UNSLOTH_STUDIO_PROJECTS_HOME", str(tmp_path / "Projects")) + monkeypatch.setattr(studio_db, "_schema_ready", False) + + +def _thread( + thread_id: str = "thread-1", + title: str = "Test Chat", + pair_id: str | None = None, +) -> dict: + return { + "id": thread_id, + "title": title, + "modelType": "base", + "modelId": "test-model", + "pairId": pair_id, + "archived": False, + "createdAt": 1_700_000_000_000, + } + + +def _message( + message_id: str, + created_at: int = 1_700_000_000_000, + attachments = None, + thread_id: str = "thread-1", +) -> dict: + message = { + "id": message_id, + "threadId": thread_id, + "parentId": None, + "role": "user", + "content": [{"type": "text", "text": "hello"}], + "createdAt": created_at, + } + if attachments is not None: + message["attachments"] = attachments + return message + + +def _image_attachment(attachment_id: str = "att-1", name: str = "photo.png") -> dict: + return { + "id": attachment_id, + "type": "image", + "name": name, + "contentType": "image/png", + "content": [{"type": "image", "image": PNG_DATA_URL}], + "status": {"type": "complete"}, + } + + +def _seed( + tmp_path, + monkeypatch, + attachments, + message_id: str = "msg-1", +): + _reset_studio_db(tmp_path, monkeypatch) + studio_db.upsert_chat_thread(_thread()) + studio_db.upsert_chat_message(_message(message_id, attachments = attachments)) + + +def _set_raw_attachments_json(message_id: str, raw: str) -> None: + conn = sqlite3.connect(studio_db_path()) + try: + conn.execute( + "UPDATE chat_messages SET attachments_json = ? WHERE id = ?", + (raw, message_id), + ) + conn.commit() + finally: + conn.close() + + +def _raw_attachments_json(message_id: str): + conn = sqlite3.connect(studio_db_path()) + try: + row = conn.execute( + "SELECT attachments_json FROM chat_messages WHERE id = ?", + (message_id,), + ).fetchone() + return row[0] if row is not None else None + finally: + conn.close() + + +# --------------------------------------------------------------------------- +# Storage: list_chat_attachments +# --------------------------------------------------------------------------- + + +def test_list_chat_attachments_empty_db(tmp_path, monkeypatch): + _reset_studio_db(tmp_path, monkeypatch) + assert studio_db.list_chat_attachments() == [] + + +def test_list_chat_attachments_round_trip(tmp_path, monkeypatch): + _seed(tmp_path, monkeypatch, [_image_attachment()]) + records = studio_db.list_chat_attachments() + assert len(records) == 1 + record = records[0] + assert record["id"] == "att-1" + assert record["messageId"] == "msg-1" + assert record["threadId"] == "thread-1" + assert record["threadTitle"] == "Test Chat" + assert record["name"] == "photo.png" + assert record["type"] == "image" + assert record["contentType"] == "image/png" + assert record["createdAt"] == 1_700_000_000_000 + # Base64 length estimate is within padding error of the decoded size. + assert abs(record["sizeBytes"] - len(PNG_BYTES)) <= 2 + + +def test_list_chat_attachments_counts_text_utf8(tmp_path, monkeypatch): + text = "héllo wörld é世界" + attachment = { + "id": "att-txt", + "type": "document", + "name": "notes.txt", + "content": [{"type": "text", "text": text}], + } + _seed(tmp_path, monkeypatch, [attachment]) + records = studio_db.list_chat_attachments() + assert records[0]["sizeBytes"] == len(text.encode("utf-8")) + + +def test_list_chat_attachments_no_content_size_is_none(tmp_path, monkeypatch): + attachment = {"id": "att-empty", "name": "ghost.bin", "content": []} + _seed(tmp_path, monkeypatch, [attachment]) + records = studio_db.list_chat_attachments() + assert records[0]["sizeBytes"] is None + assert records[0]["name"] == "ghost.bin" + + +def test_list_chat_attachments_defaults_missing_name(tmp_path, monkeypatch): + attachment = {"id": "att-noname", "content": []} + _seed(tmp_path, monkeypatch, [attachment]) + assert studio_db.list_chat_attachments()[0]["name"] == "attachment" + + +def test_list_chat_attachments_sanitizes_structured_metadata(tmp_path, monkeypatch): + attachment = { + "id": "att-weird", + "name": {"nested": "name"}, + "type": ["image"], + "contentType": {"mime": "image/png"}, + "content": [], + } + _seed(tmp_path, monkeypatch, [attachment]) + record = studio_db.list_chat_attachments()[0] + assert record["name"] == "attachment" + assert record["type"] is None + assert record["contentType"] is None + + +def test_list_chat_attachments_skips_malformed_rows(tmp_path, monkeypatch): + _reset_studio_db(tmp_path, monkeypatch) + studio_db.upsert_chat_thread(_thread()) + for i, raw in enumerate( + [ + "not json at all", + '{"id": "att-obj"}', + "null", + "[]", + '[{"noid": true}, "just a string", 42]', + '[{"id": ""}]', + ] + ): + message_id = f"msg-bad-{i}" + studio_db.upsert_chat_message(_message(message_id)) + _set_raw_attachments_json(message_id, raw) + studio_db.upsert_chat_message(_message("msg-good", attachments = [_image_attachment("att-ok")])) + records = studio_db.list_chat_attachments() + assert [r["id"] for r in records] == ["att-ok"] + + +def test_list_chat_attachments_orders_newest_first(tmp_path, monkeypatch): + _reset_studio_db(tmp_path, monkeypatch) + studio_db.upsert_chat_thread(_thread()) + studio_db.upsert_chat_message( + _message("msg-old", 1_700_000_000_000, [_image_attachment("att-old")]) + ) + studio_db.upsert_chat_message( + _message("msg-new", 1_700_000_100_000, [_image_attachment("att-new")]) + ) + assert [r["id"] for r in studio_db.list_chat_attachments()] == ["att-new", "att-old"] + + +def test_list_chat_attachments_survives_missing_thread_row(tmp_path, monkeypatch): + _reset_studio_db(tmp_path, monkeypatch) + studio_db.upsert_chat_thread(_thread()) + studio_db.upsert_chat_message(_message("msg-1", attachments = [_image_attachment()])) + conn = sqlite3.connect(studio_db_path()) + try: + conn.execute("DELETE FROM chat_threads WHERE id = 'thread-1'") + conn.commit() + finally: + conn.close() + records = studio_db.list_chat_attachments() + assert len(records) == 1 + assert records[0]["threadTitle"] is None + + +def test_list_chat_attachments_includes_compare_pair_id(tmp_path, monkeypatch): + _reset_studio_db(tmp_path, monkeypatch) + studio_db.upsert_chat_thread(_thread(pair_id = "pair-1")) + studio_db.upsert_chat_message(_message("msg-compare", attachments = [_image_attachment()])) + record = studio_db.list_chat_attachments()[0] + assert record["threadId"] == "thread-1" + assert record["pairId"] == "pair-1" + + +def test_list_chat_attachments_gone_after_thread_delete(tmp_path, monkeypatch): + _seed(tmp_path, monkeypatch, [_image_attachment()]) + studio_db.delete_chat_threads(["thread-1"]) + assert studio_db.list_chat_attachments() == [] + + +# --------------------------------------------------------------------------- +# Storage: get_chat_attachment / delete_chat_attachment +# --------------------------------------------------------------------------- + + +def test_get_chat_attachment_found_and_missing(tmp_path, monkeypatch): + _seed(tmp_path, monkeypatch, [_image_attachment()]) + attachment = studio_db.get_chat_attachment("msg-1", "att-1") + assert attachment is not None + assert attachment["content"][0]["image"] == PNG_DATA_URL + assert studio_db.get_chat_attachment("msg-1", "att-missing") is None + assert studio_db.get_chat_attachment("msg-missing", "att-1") is None + + +def test_delete_chat_attachment_keeps_others(tmp_path, monkeypatch): + _seed( + tmp_path, + monkeypatch, + [_image_attachment("att-1"), _image_attachment("att-2", "other.png")], + ) + assert studio_db.delete_chat_attachment("msg-1", "att-1") is True + assert studio_db.get_chat_attachment("msg-1", "att-1") is None + assert studio_db.get_chat_attachment("msg-1", "att-2") is not None + assert [r["id"] for r in studio_db.list_chat_attachments()] == ["att-2"] + + +def test_delete_last_chat_attachment_stores_empty_list(tmp_path, monkeypatch): + _seed(tmp_path, monkeypatch, [_image_attachment()]) + assert studio_db.delete_chat_attachment("msg-1", "att-1") is True + # '[]' rather than NULL: a NULL attachments field reads back as missing + # and triggers the legacy IndexedDB backfill, resurrecting the deleted + # attachment on the next chat load. + assert _raw_attachments_json("msg-1") == "[]" + assert studio_db.list_chat_attachments() == [] + # The message itself must survive with its content intact. + message = studio_db.get_chat_message("thread-1", "msg-1") + assert message is not None + assert message["content"] == [{"type": "text", "text": "hello"}] + assert message["attachments"] == [] + + +def test_delete_chat_attachment_missing_targets(tmp_path, monkeypatch): + _seed(tmp_path, monkeypatch, [_image_attachment()]) + assert studio_db.delete_chat_attachment("msg-missing", "att-1") is False + assert studio_db.delete_chat_attachment("msg-1", "att-missing") is False + _set_raw_attachments_json("msg-1", "not json") + assert studio_db.delete_chat_attachment("msg-1", "att-1") is False + + +# --------------------------------------------------------------------------- +# Routes: /attachments endpoints (real storage, direct calls) +# --------------------------------------------------------------------------- + + +def test_list_attachments_route(tmp_path, monkeypatch): + _seed(tmp_path, monkeypatch, [_image_attachment()]) + result = chat_history.list_attachments(current_subject = "unsloth") + assert [a["id"] for a in result["attachments"]] == ["att-1"] + + +def test_attachment_file_serves_image_bytes(tmp_path, monkeypatch): + _seed(tmp_path, monkeypatch, [_image_attachment()]) + response = chat_history.get_attachment_file("msg-1", "att-1", current_subject = "unsloth") + assert response.body == PNG_BYTES + assert response.media_type == "image/png" + + +def test_attachment_file_tolerates_whitespace_in_base64(tmp_path, monkeypatch): + encoded = base64.b64encode(PNG_BYTES).decode("ascii") + wrapped = "\n".join(encoded[i : i + 8] for i in range(0, len(encoded), 8)) + attachment = _image_attachment() + attachment["content"] = [{"type": "image", "image": "data:image/png;base64," + wrapped}] + _seed(tmp_path, monkeypatch, [attachment]) + response = chat_history.get_attachment_file("msg-1", "att-1", current_subject = "unsloth") + assert response.body == PNG_BYTES + + +def test_attachment_file_corrupt_base64_is_422(tmp_path, monkeypatch): + attachment = _image_attachment() + attachment["content"] = [{"type": "image", "image": "data:image/png;base64,%%%"}] + _seed(tmp_path, monkeypatch, [attachment]) + with pytest.raises(HTTPException) as excinfo: + chat_history.get_attachment_file("msg-1", "att-1", current_subject = "unsloth") + assert excinfo.value.status_code == 422 + + +def test_attachment_file_accepts_urlsafe_base64(tmp_path, monkeypatch): + data = bytes(range(251, 256)) * 3 # encodes to characters remapped by urlsafe + payload = base64.urlsafe_b64encode(data).decode("ascii") + assert "-" in payload or "_" in payload + attachment = _image_attachment() + attachment["content"] = [{"type": "image", "image": "data:image/png;base64," + payload}] + _seed(tmp_path, monkeypatch, [attachment]) + response = chat_history.get_attachment_file("msg-1", "att-1", current_subject = "unsloth") + assert response.body == data + + +def test_attachment_file_accepts_missing_padding(tmp_path, monkeypatch): + payload = base64.b64encode(PNG_BYTES).decode("ascii").rstrip("=") + attachment = _image_attachment() + attachment["content"] = [{"type": "image", "image": "data:image/png;base64," + payload}] + _seed(tmp_path, monkeypatch, [attachment]) + response = chat_history.get_attachment_file("msg-1", "att-1", current_subject = "unsloth") + assert response.body == PNG_BYTES + + +def test_attachment_file_serves_percent_encoded_data_url(tmp_path, monkeypatch): + attachment = _image_attachment() + attachment["content"] = [{"type": "image", "image": "data:text/plain,hello%20world"}] + _seed(tmp_path, monkeypatch, [attachment]) + response = chat_history.get_attachment_file("msg-1", "att-1", current_subject = "unsloth") + assert response.body == b"hello world" + # Non-image data URL types are clamped so markup never renders same-origin. + assert response.media_type == "application/octet-stream" + + +def test_attachment_file_serves_text_parts(tmp_path, monkeypatch): + attachment = { + "id": "att-txt", + "type": "document", + "name": "notes.txt", + "content": [ + {"type": "text", "text": "first"}, + {"type": "text", "text": "second"}, + ], + } + _seed(tmp_path, monkeypatch, [attachment]) + response = chat_history.get_attachment_file("msg-1", "att-txt", current_subject = "unsloth") + assert response.body.decode("utf-8") == "first\nsecond" + assert response.media_type.startswith("text/plain") + + +def test_attachment_file_no_content_is_404(tmp_path, monkeypatch): + _seed(tmp_path, monkeypatch, [{"id": "att-empty", "name": "ghost", "content": []}]) + with pytest.raises(HTTPException) as excinfo: + chat_history.get_attachment_file("msg-1", "att-empty", current_subject = "unsloth") + assert excinfo.value.status_code == 404 + + +def test_attachment_file_missing_message_is_404(tmp_path, monkeypatch): + _reset_studio_db(tmp_path, monkeypatch) + with pytest.raises(HTTPException) as excinfo: + chat_history.get_attachment_file("nope", "att-1", current_subject = "unsloth") + assert excinfo.value.status_code == 404 + + +def test_attachment_file_non_data_url_image_is_404(tmp_path, monkeypatch): + attachment = _image_attachment() + attachment["content"] = [{"type": "image", "image": "https://example.com/a.png"}] + _seed(tmp_path, monkeypatch, [attachment]) + with pytest.raises(HTTPException) as excinfo: + chat_history.get_attachment_file("msg-1", "att-1", current_subject = "unsloth") + assert excinfo.value.status_code == 404 + + +def test_attachment_file_defaults_media_type(tmp_path, monkeypatch): + payload = base64.b64encode(b"raw-bytes").decode("ascii") + attachment = _image_attachment() + attachment["content"] = [{"type": "image", "image": "data:;base64," + payload}] + _seed(tmp_path, monkeypatch, [attachment]) + response = chat_history.get_attachment_file("msg-1", "att-1", current_subject = "unsloth") + assert response.body == b"raw-bytes" + assert response.media_type == "application/octet-stream" + + +def test_attachment_file_svg_media_type(tmp_path, monkeypatch): + svg = b"" + payload = base64.b64encode(svg).decode("ascii") + attachment = _image_attachment() + attachment["content"] = [{"type": "image", "image": "data:image/svg+xml;base64," + payload}] + _seed(tmp_path, monkeypatch, [attachment]) + response = chat_history.get_attachment_file("msg-1", "att-1", current_subject = "unsloth") + assert response.body == svg + # SVG can carry scripts, so it downloads as bytes instead of rendering. + assert response.media_type == "application/octet-stream" + + +def test_delete_attachment_route_then_404(tmp_path, monkeypatch): + _seed(tmp_path, monkeypatch, [_image_attachment()]) + result = chat_history.delete_attachment("msg-1", "att-1", current_subject = "unsloth") + assert result == {"ok": True} + with pytest.raises(HTTPException) as excinfo: + chat_history.delete_attachment("msg-1", "att-1", current_subject = "unsloth") + assert excinfo.value.status_code == 404 + + +# --------------------------------------------------------------------------- +# Audio attachments (adapter {data, format} and compare-chat bare base64) +# --------------------------------------------------------------------------- + +WAV_BYTES = b"RIFF$\x00\x00\x00WAVEfmt \x10\x00\x00\x00\x01\x00\x01\x00" +WAV_B64 = base64.b64encode(WAV_BYTES).decode("ascii") + + +def _audio_attachment(attachment_id: str = "att-audio") -> dict: + return { + "id": attachment_id, + "type": "file", + "name": "clip.wav", + "contentType": "audio/wav", + "content": [{"type": "audio", "audio": {"data": WAV_B64, "format": "wav"}}], + "status": {"type": "complete"}, + } + + +def test_audio_attachment_lists_with_size(tmp_path, monkeypatch): + _seed(tmp_path, monkeypatch, [_audio_attachment()]) + records = studio_db.list_chat_attachments() + assert len(records) == 1 + assert records[0]["id"] == "att-audio" + assert abs(records[0]["sizeBytes"] - len(WAV_BYTES)) <= 2 + + +def test_audio_attachment_file_serves_bytes(tmp_path, monkeypatch): + _seed(tmp_path, monkeypatch, [_audio_attachment()]) + response = chat_history.get_attachment_file("msg-1", "att-audio", current_subject = "unsloth") + assert response.body == WAV_BYTES + assert response.media_type == "audio/wav" + + +def test_audio_attachment_media_type_from_format(tmp_path, monkeypatch): + attachment = _audio_attachment() + attachment["contentType"] = None + attachment["content"] = [{"type": "audio", "audio": {"data": WAV_B64, "format": "mp3"}}] + _seed(tmp_path, monkeypatch, [attachment]) + response = chat_history.get_attachment_file("msg-1", "att-audio", current_subject = "unsloth") + assert response.media_type == "audio/mpeg" + + +def test_audio_attachment_corrupt_payload_is_422(tmp_path, monkeypatch): + attachment = _audio_attachment() + attachment["content"] = [{"type": "audio", "audio": {"data": "%%%", "format": "wav"}}] + _seed(tmp_path, monkeypatch, [attachment]) + with pytest.raises(HTTPException) as excinfo: + chat_history.get_attachment_file("msg-1", "att-audio", current_subject = "unsloth") + assert excinfo.value.status_code == 422 + + +# --------------------------------------------------------------------------- +# Compare-chat uploads stored as message content parts +# --------------------------------------------------------------------------- + + +def _compare_message(message_id: str = "msg-cmp") -> dict: + return { + "id": message_id, + "threadId": "thread-1", + "parentId": None, + "role": "user", + "content": [ + {"type": "image", "image": PNG_DATA_URL}, + {"type": "audio", "audio": WAV_B64}, + {"type": "text", "text": "compare these"}, + ], + "createdAt": 1_700_000_000_000, + } + + +def _seed_compare(tmp_path, monkeypatch): + _reset_studio_db(tmp_path, monkeypatch) + studio_db.upsert_chat_thread(_thread()) + studio_db.upsert_chat_message(_compare_message()) + + +_CONTENT_PART_PREFIX = "content-part-sha256-" + + +def _content_part_id_for(message_id: str, kind: str) -> str: + """Resolve the stable content-hash id for a message's stored blob. + + Content-part ids are SHA-256 hashes of the blob payload, not array + indices, so tests look them up from the listing instead of hardcoding an + index that would shift when an earlier part is deleted. + """ + for record in studio_db.list_chat_attachments(): + if record["messageId"] == message_id and record["type"] == kind: + return record["id"] + raise AssertionError(f"no {kind} content-part upload for {message_id}") + + +def test_content_part_uploads_are_listed(tmp_path, monkeypatch): + _seed_compare(tmp_path, monkeypatch) + records = studio_db.list_chat_attachments() + # Ids are stable content hashes, not array indices. + assert all(r["id"].startswith(_CONTENT_PART_PREFIX) for r in records) + assert {r["type"] for r in records} == {"image", "audio"} + image = next(r for r in records if r["type"] == "image") + assert image["contentType"] == "image/png" + assert abs(image["sizeBytes"] - len(PNG_BYTES)) <= 2 + audio = next(r for r in records if r["type"] == "audio") + assert audio["type"] == "audio" + + +def test_content_part_file_serves_image_bytes(tmp_path, monkeypatch): + _seed_compare(tmp_path, monkeypatch) + image_id = _content_part_id_for("msg-cmp", "image") + response = chat_history.get_attachment_file("msg-cmp", image_id, current_subject = "unsloth") + assert response.body == PNG_BYTES + assert response.media_type == "image/png" + + +def test_content_part_delete_keeps_text(tmp_path, monkeypatch): + _seed_compare(tmp_path, monkeypatch) + image_id = _content_part_id_for("msg-cmp", "image") + assert studio_db.delete_chat_attachment("msg-cmp", image_id) is True + message = studio_db.get_chat_message("thread-1", "msg-cmp") + types = [p["type"] for p in message["content"]] + assert types == ["audio", "text"] + # The surviving audio blob keeps its own stable hash id after the delete. + remaining = studio_db.list_chat_attachments() + assert [r["type"] for r in remaining] == ["audio"] + assert remaining[0]["id"].startswith(_CONTENT_PART_PREFIX) + assert remaining[0]["id"] != image_id + + +def test_content_part_delete_rejects_non_blob(tmp_path, monkeypatch): + _seed_compare(tmp_path, monkeypatch) + # The text part is not a stored upload, so it never gets an id: only the + # image and audio blobs are addressable. + assert len(studio_db.list_chat_attachments()) == 2 + # A well-formed but unknown content-hash id, and malformed ids, all no-op. + assert studio_db.delete_chat_attachment("msg-cmp", _CONTENT_PART_PREFIX + "0" * 64) is False + assert studio_db.delete_chat_attachment("msg-cmp", "content-part-99") is False + assert studio_db.delete_chat_attachment("msg-cmp", "content-part-x") is False + + +def test_text_only_messages_not_listed_as_uploads(tmp_path, monkeypatch): + _reset_studio_db(tmp_path, monkeypatch) + studio_db.upsert_chat_thread(_thread()) + # The word "image" inside text must not create phantom upload rows. + message = _message("msg-txt") + message["content"] = [{"type": "text", "text": 'discussing an "image" and "audio" here'}] + studio_db.upsert_chat_message(message) + assert studio_db.list_chat_attachments() == [] + + +def test_remote_image_urls_are_not_listed_as_uploads(tmp_path, monkeypatch): + _reset_studio_db(tmp_path, monkeypatch) + studio_db.upsert_chat_thread(_thread()) + message = _message("msg-remote") + message["content"] = [ + {"type": "image", "image": "https://example.com/cat.png"}, + {"type": "text", "text": "look at this"}, + ] + studio_db.upsert_chat_message(message) + # No stored bytes: nothing to list, open, or delete. + assert studio_db.list_chat_attachments() == [] + assert studio_db.get_chat_attachment("msg-remote", "content-part-0") is None + assert studio_db.delete_chat_attachment("msg-remote", "content-part-0") is False + stored = studio_db.get_chat_message("thread-1", "msg-remote") + assert [p["type"] for p in stored["content"]] == ["image", "text"] + + +def test_html_data_url_serves_as_octet_stream(tmp_path, monkeypatch): + _reset_studio_db(tmp_path, monkeypatch) + studio_db.upsert_chat_thread(_thread()) + html_b64 = base64.b64encode(b"").decode() + message = _message("msg-html") + message["content"] = [ + {"type": "image", "image": f"data:text/html;base64,{html_b64}"}, + ] + studio_db.upsert_chat_message(message) + attachment_id = _content_part_id_for("msg-html", "image") + response = chat_history.get_attachment_file( + "msg-html", attachment_id, current_subject = "unsloth" + ) + # Never echo a script-capable media type back under the app origin. + assert response.media_type == "application/octet-stream" + assert response.body == b"" + + +def test_svg_data_url_serves_as_octet_stream(tmp_path, monkeypatch): + _reset_studio_db(tmp_path, monkeypatch) + studio_db.upsert_chat_thread(_thread()) + svg_b64 = base64.b64encode(b"").decode() + message = _message("msg-svg") + message["content"] = [ + {"type": "image", "image": f"data:image/svg+xml;base64,{svg_b64}"}, + ] + studio_db.upsert_chat_message(message) + attachment_id = _content_part_id_for("msg-svg", "image") + response = chat_history.get_attachment_file("msg-svg", attachment_id, current_subject = "unsloth") + assert response.media_type == "application/octet-stream" + + +def test_png_data_url_keeps_its_media_type(tmp_path, monkeypatch): + _seed_compare(tmp_path, monkeypatch) + image_id = _content_part_id_for("msg-cmp", "image") + response = chat_history.get_attachment_file("msg-cmp", image_id, current_subject = "unsloth") + assert response.media_type == "image/png" diff --git a/studio/backend/tests/test_chat_load_during_training.py b/studio/backend/tests/test_chat_load_during_training.py index 63dba8579c..a5fd71b6a0 100644 --- a/studio/backend/tests/test_chat_load_during_training.py +++ b/studio/backend/tests/test_chat_load_during_training.py @@ -168,11 +168,14 @@ class TestCanLoadGGUF(_GpuCacheResetMixin, unittest.TestCase): devices, required_override = None, estimate = None, + single_device_gpu = None, + gpu_ids = None, ): with ( patch("utils.hardware.get_device", return_value = DeviceType.CUDA), patch("utils.hardware.estimate_required_model_memory_gb", return_value = (estimate, {})), patch("utils.hardware.get_visible_gpu_utilization", return_value = {"devices": devices}), + patch("utils.hardware.resolve_requested_gpu_ids", return_value = gpu_ids), patch("utils.hardware.auto_select_gpu_ids") as auto_mock, ): ok, info = tv.can_load_chat_during_training( @@ -180,9 +183,10 @@ class TestCanLoadGGUF(_GpuCacheResetMixin, unittest.TestCase): hf_token = None, load_in_4bit = True, max_seq_length = 0, - requested_gpu_ids = None, + requested_gpu_ids = gpu_ids, is_gguf = True, required_override_gb = required_override, + single_device_gpu = single_device_gpu, ) return ok, info, auto_mock @@ -198,6 +202,88 @@ class TestCanLoadGGUF(_GpuCacheResetMixin, unittest.TestCase): ok, _, _ = self._run(devices = _devices((0, 80, 35), (1, 80, 70)), required_override = 20.0) self.assertTrue(ok) + def test_no_per_gpu_floor_for_gguf_with_explicit_gpu_ids(self): + # gpu_ids narrows llama.cpp's candidate pool but does not turn its + # self-placement into HF device_map="balanced". The uneven selected + # pair therefore keeps the aggregate GGUF check without an even-share + # floor on the nearly-full card. + ok, info, _ = self._run( + devices = _devices((0, 80, 35), (1, 80, 70), (2, 80, 0)), + required_override = 20.0, + gpu_ids = [0, 1], + ) + self.assertTrue(ok) + self.assertEqual(info["mode"], "gguf") + + def test_single_device_uses_selected_gpu(self): + # The model needs 27 GB with headroom. GPU 0 has 45 GB free, while an + # unrelated training-heavy GPU 1 has only 10 GB free. + ok, info, _ = self._run( + devices = _devices((0, 80, 35), (1, 80, 70)), + required_override = 20.0, + single_device_gpu = "0", + ) + self.assertTrue(ok) + self.assertEqual(info["usable_gb"], 45.0) + + blocked, blocked_info, _ = self._run( + devices = _devices((0, 80, 35), (1, 80, 70)), + required_override = 20.0, + single_device_gpu = "1", + ) + self.assertFalse(blocked) + self.assertEqual(blocked_info["usable_gb"], 10.0) + + def test_single_device_unresolved_token_sizes_against_worst_device(self): + # A non-numeric device token (a CUDA UUID / MIG handle) can't map to a + # free-VRAM index. The runner still drives ONE device, so size against the + # worst-case visible device (min free), not the aggregate pool: one GPU + # with 80 GB free vs a 20 GB model -> allow. + ok, info, _ = self._run( + devices = _devices((0, 80, 0)), + required_override = 20.0, + single_device_gpu = "GPU-uuid", + ) + self.assertTrue(ok) + self.assertEqual(info["mode"], "single_device") + self.assertNotIn("reason", info) + + def test_single_device_unresolved_token_refuses_when_worst_device_full(self): + # Same UUID fallback, worst-case device nearly full (2 GB for a 20 GB + # model) -> refuse (default-deny), not on an unresolved-token technicality. + ok, info, _ = self._run( + devices = _devices((0, 80, 78)), + required_override = 20.0, + single_device_gpu = "GPU-uuid", + ) + self.assertFalse(ok) + self.assertNotEqual(info.get("reason"), "unresolved_gpu_id") + + def test_single_device_unresolved_token_uses_min_free_not_aggregate(self): + # The single-device runner uses ONE device but we can't tell which from a + # UUID token. Sizing against the aggregate pool would let a 20 GB model + # "fit" 160 GB of pooled free VRAM while landing on a 2 GB card and OOMing + # training. Min-free (2 GB) is the safe worst case -> refuse. + ok, info, _ = self._run( + devices = _devices((0, 80, 78), (1, 80, 0), (2, 80, 0)), + required_override = 20.0, + single_device_gpu = "GPU-uuid", + ) + self.assertFalse(ok) + self.assertEqual(info["mode"], "single_device") + + def test_single_device_cpu_token_allows(self): + # An empty device token = a CPU-only single-device runner (CPU diffusion + # GGUF): it uses no GPU VRAM, so it never threatens training -> allow + # regardless of how full the GPUs are. + ok, info, _ = self._run( + devices = _devices((0, 80, 78)), + required_override = 20.0, + single_device_gpu = "", + ) + self.assertTrue(ok) + self.assertEqual(info["reason"], "cpu_only") + def test_estimate_unavailable_refuses(self): # No override and the estimator can't size it -> default-deny. ok, info, _ = self._run(devices = _devices((0, 80, 0)), required_override = None, estimate = None) @@ -309,6 +395,8 @@ class TestChatLoadGuardRoute(unittest.TestCase): captured = None, training_active, decision, + gpu_memory_mode = "auto", + requested_gpu_ids = None, ): config = config or SimpleNamespace(is_gguf = False, is_lora = False, path = None) with _stub_guard_deps( @@ -320,7 +408,8 @@ class TestChatLoadGuardRoute(unittest.TestCase): hf_token = None, load_in_4bit = True, max_seq_length = 0, - requested_gpu_ids = None, + requested_gpu_ids = requested_gpu_ids, + gpu_memory_mode = gpu_memory_mode, ) def test_noop_when_training_inactive(self): @@ -332,6 +421,141 @@ class TestChatLoadGuardRoute(unittest.TestCase): def test_allows_when_fits(self): self._guard(training_active = True, decision = (True, {"mode": "auto"})) + def test_diffusion_detection_uses_name_before_download(self): + config = SimpleNamespace( + identifier = "unsloth/DiffusionGemma-GGUF", + gguf_hf_repo = "unsloth/DiffusionGemma-GGUF", + gguf_file = None, + ) + self.assertTrue(self.route._classify_diffusion_gguf(config)) + + def test_uncached_gguf_classification_remains_unknown(self): + config = SimpleNamespace( + identifier = "owner/renamed-model", + gguf_hf_repo = "owner/renamed-model", + gguf_variant = "Q4_K_M", + gguf_file = None, + ) + self.assertIsNone(self.route._classify_diffusion_gguf(config)) + + def test_diffusion_detection_reuses_loader_metadata_probe(self): + import tempfile + + seen = [] + + class _Probe: + is_diffusion = False + _architecture = None + + def _read_gguf_metadata(self, path): + seen.append(path) + self.is_diffusion = True + + with tempfile.TemporaryDirectory() as d: + model = Path(d) / "renamed.gguf" + model.write_bytes(b"GGUF") + config = SimpleNamespace(identifier = "local", gguf_file = str(model)) + with patch.object(self.route, "LlamaCppBackend", _Probe): + self.assertTrue(self.route._classify_diffusion_gguf(config)) + self.assertEqual(seen, [str(model)]) + + def test_local_chat_gguf_classification_is_definitive(self): + import tempfile + class _Probe: + is_diffusion = False + _architecture = "llama" + + def _read_gguf_metadata(self, _path): + pass + + with tempfile.TemporaryDirectory() as d: + model = Path(d) / "renamed.gguf" + model.write_bytes(b"GGUF") + config = SimpleNamespace(identifier = "local", gguf_file = str(model)) + with patch.object(self.route, "LlamaCppBackend", _Probe): + self.assertFalse(self.route._classify_diffusion_gguf(config)) + + def test_manual_known_normal_gguf_bypasses_training_estimate(self): + captured = [] + config = SimpleNamespace(is_gguf = True) + with patch.object(self.route, "_classify_diffusion_gguf", return_value = False): + self._guard( + config = config, + captured = captured, + training_active = True, + decision = (False, {"reason": "must not run"}), + gpu_memory_mode = "manual", + ) + self.assertEqual(captured, []) + + def test_manual_unknown_gguf_keeps_single_device_training_guard(self): + captured = [] + config = SimpleNamespace(is_gguf = True) + with ( + patch.object(self.route, "_classify_diffusion_gguf", return_value = None), + patch.object(self.route, "_estimate_gguf_required_gb", return_value = 12.5), + patch.object( + self.route.LlamaCppBackend, + "_diffusion_gpu_arg", + return_value = "2", + ), + ): + self._guard( + config = config, + captured = captured, + training_active = True, + decision = (True, {"mode": "single_device"}), + gpu_memory_mode = "manual", + ) + self.assertEqual(len(captured), 1) + self.assertEqual(captured[0]["single_device_gpu"], "2") + + def test_manual_diffusion_uses_single_device_guard(self): + captured = [] + config = SimpleNamespace(is_gguf = True) + with ( + patch.object(self.route, "_classify_diffusion_gguf", return_value = True), + patch.object(self.route, "_estimate_gguf_required_gb", return_value = 12.5), + ): + self._guard( + config = config, + captured = captured, + training_active = True, + decision = (True, {"mode": "gguf"}), + gpu_memory_mode = "manual", + requested_gpu_ids = [3, 1], + ) + self.assertEqual(len(captured), 1) + self.assertEqual(captured[0]["single_device_gpu"], "1") + self.assertEqual(captured[0]["requested_gpu_ids"], [3, 1]) + + def test_unpinned_diffusion_uses_runner_default_gpu(self): + captured = [] + config = SimpleNamespace(is_gguf = True) + with ( + patch.object(self.route, "_classify_diffusion_gguf", return_value = True), + patch.object(self.route, "_estimate_gguf_required_gb", return_value = 12.5), + patch.object( + self.route.LlamaCppBackend, + "_effective_gpu_count", + return_value = 2, + ), + patch.object( + self.route.LlamaCppBackend, + "_diffusion_gpu_arg", + return_value = "3", + ) as gpu_arg, + ): + self._guard( + config = config, + captured = captured, + training_active = True, + decision = (True, {"mode": "single_device"}), + gpu_memory_mode = "manual", + ) + gpu_arg.assert_called_once_with(None, cpu_only = False) + self.assertEqual(captured[0]["single_device_gpu"], "3") + def test_refuses_with_headroom_number(self): info = {"required_gb": 30.0, "usable_gb": 6.0, "needed_gb": 39.0, "mode": "auto"} with self.assertRaises(HTTPException) as exc: @@ -467,36 +691,189 @@ class TestValidateRefusesDuringTraining(unittest.TestCase): self.assertEqual(captured[0]["load_in_4bit"], False) self.assertEqual(captured[0]["max_seq_length"], 4096) - def test_rejects_gguf_with_gpu_ids_before_guard(self): - # /validate must mirror /load's GGUF + gpu_ids 400, before the VRAM guard. + def test_validate_forwards_manual_gpu_memory_mode_to_guard(self): from models.inference import ValidateModelRequest - request = ValidateModelRequest(model_path = "x.gguf", gpu_ids = [0]) + request = ValidateModelRequest( + model_path = "unsloth/model-GGUF", + gguf_variant = "Q4_K_M", + gpu_memory_mode = "manual", + ) cfg = SimpleNamespace( - identifier = "x.gguf", - display_name = "x", + identifier = "unsloth/model-GGUF", + display_name = "model-GGUF", is_gguf = True, is_lora = False, is_vision = False, path = None, base_model = None, ) - captured = [] + captured = {} with ( patch.object( self.route, "_resolve_model_identifier_for_request", - return_value = ("x.gguf", "x.gguf", False), + return_value = ("unsloth/model-GGUF", "unsloth/model-GGUF", False), ), patch.object(self.route.ModelConfig, "from_identifier", return_value = cfg), patch.object(self.route, "load_inference_config", return_value = {}), - _stub_guard_deps(training_active = True, decision = (True, {}), captured = captured), + patch.object( + self.route, + "_guard_chat_load_against_training", + lambda config, **kw: captured.update(kw), + ), ): - with self.assertRaises(HTTPException) as exc: - asyncio.run(self.route.validate_model(request, current_subject = "u")) - self.assertEqual(exc.exception.status_code, 400) - self.assertIn("gpu_ids is not supported for GGUF", exc.exception.detail) - self.assertEqual(captured, []) # guard never reached + asyncio.run(self.route.validate_model(request, current_subject = "u")) + self.assertEqual(captured.get("gpu_memory_mode"), "manual") + + def test_validate_forwards_inherited_extras_and_parallel_to_guard(self): + # Regression: /load resolves inherited same-model extras and passes the + # real slot count to the guard; validate must do the same, else it sizes + # a smaller estimate (no inherited -c/--model-draft, n_parallel=1) and + # /load then 409s after the frontend has already unloaded. + from models.inference import ValidateModelRequest + + request = ValidateModelRequest(model_path = "unsloth/Qwen3-1.7B", max_seq_length = 4096) + cfg = SimpleNamespace( + identifier = "unsloth/Qwen3-1.7B", + display_name = "Qwen3-1.7B", + is_gguf = False, + is_lora = False, + is_vision = False, + path = None, + base_model = None, + ) + captured = {} + with ( + patch.object( + self.route, + "_resolve_model_identifier_for_request", + return_value = ("unsloth/Qwen3-1.7B", "unsloth/Qwen3-1.7B", False), + ), + patch.object(self.route.ModelConfig, "from_identifier", return_value = cfg), + patch.object(self.route, "load_inference_config", return_value = {}), + patch.object(self.route, "_resolve_inherited_extra_args", return_value = ["-c", "32768"]), + patch.object( + self.route, + "_guard_chat_load_against_training", + lambda config, **kw: captured.update(kw), + ), + ): + asyncio.run(self.route.validate_model(request, current_subject = "u")) + self.assertEqual(captured.get("llama_extra_args"), ["-c", "32768"]) + self.assertIn("n_parallel", captured) + + def test_metadata_probe_skips_training_guard(self): + # A header-only probe (include_context_length) allocates no VRAM, so the + # training guard must not run -- else the staging GPU-layers / MoE sliders + # it feeds are hidden exactly when a during-training user needs them. + from models.inference import ValidateModelRequest + + request = ValidateModelRequest( + model_path = "unsloth/Qwen3-1.7B", + max_seq_length = 4096, + include_context_length = True, + ) + cfg = SimpleNamespace( + identifier = "unsloth/Qwen3-1.7B", + display_name = "Qwen3-1.7B", + is_gguf = False, + is_lora = False, + is_vision = False, + path = None, + base_model = None, + ) + guard_called = [] + with ( + patch.object( + self.route, + "_resolve_model_identifier_for_request", + return_value = ("unsloth/Qwen3-1.7B", "unsloth/Qwen3-1.7B", False), + ), + patch.object(self.route.ModelConfig, "from_identifier", return_value = cfg), + patch.object(self.route, "load_inference_config", return_value = {}), + patch.object( + self.route, + "_guard_chat_load_against_training", + lambda *a, **kw: guard_called.append(True), + ), + ): + asyncio.run(self.route.validate_model(request, current_subject = "u")) + self.assertEqual(guard_called, []) + + def _validate_gguf_template( + self, + *, + template, + canonical_path = "/picked/model.gguf", + ): + # Drive validate_model for a native lease-backed GGUF template probe and + # capture what the embedded-template reader was called with. + from models.inference import ValidateModelRequest + + request = ValidateModelRequest( + model_path = "model.gguf", + gguf_variant = "Q4_K_M", + native_path_lease = "signed-lease", + include_chat_template = True, + ) + cfg = SimpleNamespace( + identifier = canonical_path, + display_name = "model.gguf", + is_gguf = True, + is_lora = False, + is_vision = False, + gguf_file = canonical_path, + path = None, + base_model = None, + ) + import utils.models.gguf_metadata as gguf_meta + + seen = {} + + def _fake_read(path): + seen["path"] = path + return template + + guard_called = [] + with ( + patch.object( + self.route, + "_resolve_model_identifier_for_request", + return_value = (canonical_path, "model.gguf", True), + ), + patch.object(self.route.ModelConfig, "from_identifier", return_value = cfg), + patch.object(self.route, "load_inference_config", return_value = {}), + patch.object(gguf_meta, "read_gguf_chat_template", _fake_read), + patch.object( + self.route, + "_guard_chat_load_against_training", + lambda *a, **kw: guard_called.append(True), + ), + ): + resp = asyncio.run(self.route.validate_model(request, current_subject = "u")) + return resp, seen, guard_called + + def test_include_chat_template_reads_leased_gguf_embedded_template(self): + # The picker chat-template GET has no lease plumbing, so a native picked + # GGUF surfaces its default template through this lease-aware probe: the + # embedded template is read from the granted canonical path and returned. + resp, seen, _ = self._validate_gguf_template(template = "{{ messages }}") + self.assertEqual(resp.chat_template, "{{ messages }}") + # Read strictly the leased file's own embedded template, never a sibling + # sidecar: the grant authorizes just this one path. + self.assertEqual(seen["path"], "/picked/model.gguf") + + def test_include_chat_template_skips_training_guard(self): + # A template-only probe allocates no VRAM, so like include_context_length + # it must not be refused by the training guard. + _, _, guard_called = self._validate_gguf_template(template = "{{ messages }}") + self.assertEqual(guard_called, []) + + def test_include_chat_template_over_cap_is_dropped(self): + from picker.schemas import MAX_CHAT_TEMPLATE_BYTES + resp, _, _ = self._validate_gguf_template(template = "a" * (MAX_CHAT_TEMPLATE_BYTES + 1)) + self.assertIsNone(resp.chat_template) # ── _estimate_gguf_required_gb (sizes the same weights the loader loads) ────── diff --git a/studio/backend/tests/test_cuda_torch_spec.py b/studio/backend/tests/test_cuda_torch_spec.py new file mode 100644 index 0000000000..928cef787e --- /dev/null +++ b/studio/backend/tests/test_cuda_torch_spec.py @@ -0,0 +1,73 @@ +# SPDX-License-Identifier: AGPL-3.0-only +# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0 + +"""Tests for _CUDA_TORCH_PKG_SPEC in install_python_stack.py. + +The CUDA repair path installs the torch trio from an exclusive --index-url (no +PyPI fallback), so these pinned ranges decide which torch the venv gets. The +upper bound is locked to the 2.11.x family to match the base image and rocm7.2 +spec and to keep the companions off a torch-2.12 wheel that would ABI-mismatch. +""" + +from __future__ import annotations + +import sys +from pathlib import Path + +import pytest +from packaging.requirements import Requirement + +# install_python_stack.py lives at repo_root/studio/install_python_stack.py +_INSTALL_SCRIPT = Path(__file__).resolve().parents[2] / "install_python_stack.py" + + +def _load_module(monkeypatch): + """(Re-)import and return install_python_stack (mirrors test_torchao_select).""" + sys.modules.pop("install_python_stack", None) + monkeypatch.syspath_prepend(str(_INSTALL_SCRIPT.parent)) + import install_python_stack + + return install_python_stack + + +def _spec_of(pkg_spec: str): + """Parse 'torch>=2.4,<2.12.0' into a packaging SpecifierSet.""" + return Requirement(pkg_spec).specifier + + +@pytest.mark.parametrize( + "index, allowed, rejected", + [ + # torch: 2.11.x allowed (matches base image); 2.12.x excluded. + (0, ["2.11.0", "2.11.2", "2.10.0", "2.4.0"], ["2.12.0", "2.3.0", "1.13.1"]), + # torchvision: 0.26.x (torch 2.11 companion) allowed; 0.27.x (torch 2.12) out. + (1, ["0.26.0", "0.26.1", "0.19.0"], ["0.27.0", "0.18.0"]), + # torchaudio: same 2.11.x window as torch. + (2, ["2.11.0", "2.10.0", "2.4.0"], ["2.12.0", "2.3.0"]), + ], +) +def test_cuda_spec_bounds(monkeypatch, index, allowed, rejected): + mod = _load_module(monkeypatch) + spec = _spec_of(mod._CUDA_TORCH_PKG_SPEC[index]) + for v in allowed: + assert spec.contains(v, prereleases = True), f"{v} should satisfy {spec}" + for v in rejected: + assert not spec.contains(v, prereleases = True), f"{v} should not satisfy {spec}" + + +def test_cuda_spec_matches_rocm72_upper_bound(monkeypatch): + """CUDA and rocm7.2 target the same torch 2.11.x family, so their upper + bounds must stay in lockstep (bump both together at 2.12.x).""" + mod = _load_module(monkeypatch) + rocm72 = mod._ROCM_TORCH_PKG_SPECS["rocm7.2"] + + def _upper(pkg_spec: str) -> str: + for clause in _spec_of(pkg_spec): + if clause.operator == "<": + return clause.version + raise AssertionError(f"no upper bound in {pkg_spec!r}") + + for cuda_pkg, rocm_pkg in zip(mod._CUDA_TORCH_PKG_SPEC, rocm72, strict = True): + assert _upper(cuda_pkg) == _upper( + rocm_pkg + ), f"CUDA {cuda_pkg!r} upper bound must match rocm7.2 {rocm_pkg!r}" diff --git a/studio/backend/tests/test_export_absolute_paths.py b/studio/backend/tests/test_export_absolute_paths.py index 761ea08e3f..5097f9f53a 100644 --- a/studio/backend/tests/test_export_absolute_paths.py +++ b/studio/backend/tests/test_export_absolute_paths.py @@ -158,6 +158,7 @@ def _install_lightweight_backend_stubs(monkeypatch): utils_model_config._pick_best_gguf = lambda variants: variants[0] if variants else None utils_model_config._extract_quant_label = lambda value: value utils_model_config._is_big_endian_gguf_path = lambda *args, **kwargs: False + utils_model_config._is_mtp_drafter = lambda *args, **kwargs: False utils_model_config.is_audio_input_type = lambda *args, **kwargs: None monkeypatch.setitem( sys.modules, diff --git a/studio/backend/tests/test_gguf_load_cache_reuse.py b/studio/backend/tests/test_gguf_load_cache_reuse.py index 15d91cd324..62596fcc8a 100644 --- a/studio/backend/tests/test_gguf_load_cache_reuse.py +++ b/studio/backend/tests/test_gguf_load_cache_reuse.py @@ -728,9 +728,13 @@ class TestLoadHubDownloadExclusion: source = (Path(__file__).resolve().parent.parent / "routes" / "inference.py").read_text() gguf_branch = source[source.index("if config.is_gguf:") :] + # The gguf_load_in_flight marker must be entered before the hub-download + # guard and the unload so a concurrent load can't race the download + # manager. The llama_extra_args inheritance that used to sit between the + # marker and the guard now runs in _guard_chat_load_against_training, ahead + # of the GGUF branch, so it is no longer a landmark inside this slice. assert ( gguf_branch.index("enter_context(gguf_load_in_flight") - < gguf_branch.index("if request.llama_extra_args is None") < gguf_branch.index("_hub_download_blocks_gguf_load") < gguf_branch.index("unsloth_backend.unload_model") ) diff --git a/studio/backend/tests/test_gguf_metadata.py b/studio/backend/tests/test_gguf_metadata.py index a5be07f8e3..ec0330ce05 100644 --- a/studio/backend/tests/test_gguf_metadata.py +++ b/studio/backend/tests/test_gguf_metadata.py @@ -15,6 +15,7 @@ from utils.models.gguf_metadata import ( pairing_score, read_gguf_context_length, read_gguf_general_metadata, + read_gguf_staged_dims, read_mmproj_audio_capability, ) @@ -153,6 +154,78 @@ def test_context_length_ignores_foreign_arch_key(tmp_path: Path): assert read_gguf_context_length(str(p)) is None +# --- read_gguf_staged_dims (one pass: context + layer + moe counts) ---- + + +def test_staged_dims_none_for_missing_or_non_gguf(tmp_path: Path): + assert read_gguf_staged_dims(str(tmp_path / "nope.gguf")) is None + p = tmp_path / "garbage.gguf" + p.write_bytes(b"not a gguf at all") + assert read_gguf_staged_dims(str(p)) is None + + +def test_staged_dims_moe_with_leading_dense(tmp_path: Path): + # GLM-4.7-Flash shape: context + total layers + MoE layers in one read. + p = _write_synthetic_gguf( + tmp_path / "glm.gguf", + {"general.architecture": "deepseek2"}, + extra_uint32 = { + "deepseek2.context_length": 202752, + "deepseek2.block_count": 47, + "deepseek2.expert_count": 64, + "deepseek2.leading_dense_block_count": 1, + }, + ) + assert read_gguf_staged_dims(str(p)) == { + "context_length": 202752, + "layer_count": 47, + "moe_layer_count": 46, + } + + +def test_staged_dims_dense_model(tmp_path: Path): + # Dense: layer_count present, moe_layer_count 0 (slider hidden). + p = _write_synthetic_gguf( + tmp_path / "dense.gguf", + {"general.architecture": "qwen3"}, + extra_uint32 = {"qwen3.context_length": 40960, "qwen3.block_count": 36}, + ) + assert read_gguf_staged_dims(str(p)) == { + "context_length": 40960, + "layer_count": 36, + "moe_layer_count": 0, + } + + +def test_staged_dims_all_moe_no_leading_dense(tmp_path: Path): + # Experts present, no leading_dense key -> every block is a MoE layer. + p = _write_synthetic_gguf( + tmp_path / "moe.gguf", + {"general.architecture": "qwen35moe"}, + extra_uint32 = {"qwen35moe.block_count": 40, "qwen35moe.expert_count": 256}, + ) + assert read_gguf_staged_dims(str(p)) == { + "context_length": None, + "layer_count": 40, + "moe_layer_count": 40, + } + + +def test_staged_dims_uint64_block_count(tmp_path: Path): + # block_count stored as uint64 (vtype 10) still parses; moe == block_count. + p = _write_synthetic_gguf( + tmp_path / "moe64.gguf", + {"general.architecture": "gpt-oss"}, + extra_uint32 = {"gpt-oss.expert_count": 32}, + extra_uint64 = {"gpt-oss.block_count": 24}, + ) + assert read_gguf_staged_dims(str(p)) == { + "context_length": None, + "layer_count": 24, + "moe_layer_count": 24, + } + + def test_context_length_read_from_uint64(tmp_path: Path): # Some models store context_length as a uint64 (vtype 10). p = _write_synthetic_gguf( diff --git a/studio/backend/tests/test_gpu_memory_mode.py b/studio/backend/tests/test_gpu_memory_mode.py new file mode 100644 index 0000000000..b17274197f --- /dev/null +++ b/studio/backend/tests/test_gpu_memory_mode.py @@ -0,0 +1,879 @@ +# SPDX-License-Identifier: AGPL-3.0-only +# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0 + +"""Backend contract for the GPU Memory mode dropdown. + +The dropdown threads a single ``gpu_memory_mode`` ("auto" | "manual") from the +chat UI through the load request. "manual" lets the user own the offload: with +``gpu_layers < 0`` (Auto, the default) it hands all memory management to +llama.cpp's ``--fit on`` (no CUDA/HIP device masking, no context auto-reduce, no +gpu-layer or tensor-split planning); with ``gpu_layers >= 0`` it pins the layers +and MoE offload itself (``--fit off``). These tests pin: + + * the pydantic request/response/status contract (snake_case key, default + "auto", unknown values rejected), + * the backend ``gpu_memory_mode`` property and its reset on unload, + * the ``_already_in_target_state`` reload-detection branch, and + * that the manual + Auto-layers branch in ``load_model`` empties the probed + GPU set and drops tensor parallelism so the selection below no-ops, while + the explicit-offload branch emits ``--gpu-layers`` / ``--fit off``. +""" + +from __future__ import annotations + +import inspect +import sys +import types as _types +from pathlib import Path + +import pytest + +_BACKEND_DIR = str(Path(__file__).resolve().parent.parent) +if _BACKEND_DIR not in sys.path: + sys.path.insert(0, _BACKEND_DIR) + +# Same external-dep stubs as the other llama_cpp unit tests so importing +# the backend doesn't drag in structlog / httpx / loggers. +_loggers_stub = _types.ModuleType("loggers") +_loggers_stub.get_logger = lambda name: __import__("logging").getLogger(name) +sys.modules.setdefault("loggers", _loggers_stub) + +_structlog_stub = _types.ModuleType("structlog") +_structlog_stub.get_logger = lambda *a, **k: __import__("logging").getLogger("stub") +sys.modules.setdefault("structlog", _structlog_stub) + +# httpx is a real, installed backend dependency: import it so the genuine module +# is in sys.modules. A hand-rolled stub here is inevitably incomplete and, since +# setdefault installs it before real httpx loads, would poison a combined pytest +# run -- routes/inference references httpx.Response (and other attrs) at def time. +import httpx # noqa: F401 + +from core.inference import llama_cpp as llama_cpp_module +from core.inference.llama_cpp import LlamaCppBackend +from models.inference import ( + InferenceStatusResponse, + LoadRequest, + LoadResponse, +) + + +# ── Pydantic contract (snake_case key, default "auto") ─────────────── + + +def test_load_request_defaults_gpu_memory_mode_auto(): + assert LoadRequest(model_path = "owner/repo").gpu_memory_mode == "auto" + + +def test_load_request_round_trips_json_key(): + req = LoadRequest.model_validate({"model_path": "owner/repo", "gpu_memory_mode": "manual"}) + assert req.gpu_memory_mode == "manual" + assert req.model_dump()["gpu_memory_mode"] == "manual" + + +def test_load_request_rejects_unknown_mode(): + with pytest.raises(ValueError): + LoadRequest(model_path = "owner/repo", gpu_memory_mode = "bogus") + + +@pytest.mark.parametrize("model_cls", [LoadResponse, InferenceStatusResponse]) +def test_response_models_emit_gpu_memory_mode(model_cls): + if model_cls is LoadResponse: + default = model_cls( + status = "loaded", + model = "owner/repo", + display_name = "repo", + inference = {}, + ) + manual = model_cls( + status = "loaded", + model = "owner/repo", + display_name = "repo", + inference = {}, + gpu_memory_mode = "manual", + ) + else: + default = model_cls() + manual = model_cls(gpu_memory_mode = "manual") + assert default.model_dump()["gpu_memory_mode"] == "auto" + assert manual.model_dump()["gpu_memory_mode"] == "manual" + + +# ── Backend property + reset ───────────────────────────────────────── + + +class _FakeProcess: + """Stand-in for subprocess.Popen so _kill_process is a no-op.""" + + def terminate(self): + pass + + def wait(self, timeout = None): + return 0 + + def kill(self): + pass + + def poll(self): + return 0 + + +def test_gpu_memory_mode_property_defaults_auto(): + assert LlamaCppBackend().gpu_memory_mode == "auto" + + +def test_gpu_memory_mode_property_reflects_field(): + backend = LlamaCppBackend() + backend._gpu_memory_mode = "manual" + assert backend.gpu_memory_mode == "manual" + + +def test_unload_resets_gpu_memory_mode(): + backend = LlamaCppBackend() + backend._process = _FakeProcess() + backend._gpu_memory_mode = "manual" + backend.unload_model() + assert backend.gpu_memory_mode == "auto" + + +# ── _already_in_target_state reload-detection branch ───────────────── + + +def _loaded_backend(gpu_memory_mode: str) -> LlamaCppBackend: + backend = LlamaCppBackend() + backend._process = _FakeProcess() # is_loaded only checks "is not None" + backend._healthy = True + backend._model_identifier = "owner/repo" + backend._hf_variant = "Q4_K_M" + backend._requested_n_ctx = 8192 + backend._cache_type_kv = None + backend._requested_spec_mode = "auto" + backend._chat_template_override = None + backend._is_vision = False + backend._extra_args = None + backend._gguf_path = None + backend._gpu_memory_mode = gpu_memory_mode + return backend + + +def _target_state(backend: LlamaCppBackend, gpu_memory_mode: str) -> bool: + return backend._already_in_target_state( + gguf_path = None, + model_identifier = "owner/repo", + hf_variant = "Q4_K_M", + n_ctx = 8192, + cache_type_kv = None, + speculative_type = "auto", + chat_template_override = None, + extra_args = None, + is_vision = False, + gpu_memory_mode = gpu_memory_mode, + ) + + +@pytest.mark.parametrize("mode", ["auto", "manual"]) +def test_already_in_target_state_matches_same_mode(mode): + assert _target_state(_loaded_backend(mode), mode) is True + + +@pytest.mark.parametrize("loaded,requested", [("auto", "manual"), ("manual", "auto")]) +def test_already_in_target_state_reloads_on_mode_change(loaded, requested): + # Flipping the dropdown either direction must force a reload so the command + # is rebuilt with/without the Unsloth GPU masking. + assert _target_state(_loaded_backend(loaded), requested) is False + + +def test_already_in_target_state_ignores_mode_for_diffusion(): + # The diffusion runner is mode-agnostic (always "auto"), so a standing manual + # preference must not force a needless reload. + backend = _loaded_backend("auto") + backend._is_diffusion = True + assert _target_state(backend, "manual") is True + + +# ── load_model: manual + Auto layers bypasses Unsloth GPU management ── + + +def _load_model_source() -> str: + return inspect.getsource(llama_cpp_module.LlamaCppBackend.load_model) + + +def test_auto_layers_branch_empties_gpus_and_drops_tensor_parallel(): + # Emptying the probed set makes the selection / TP planning below no-op, so + # gpu_indices stays None and use_fit True (--fit on). + src = _load_model_source() + gate = src.find('if gpu_memory_mode == "manual" and gpu_layers < 0:') + assert gate != -1, "load_model must branch on manual + Auto layers (gpu_layers < 0)" + block = src[gate : gate + 1400] + assert "gpus = []" in block, "Auto-layers branch must empty the probed GPU set" + # --fit aborts under --split-mode tensor, so a raw-extras split-mode is stripped. + assert "strip_split_mode_only(extra_args)" in block + assert "requested_ctx if requested_ctx > 0 else 0" in block + # The branch sits before GPU selection assigns gpu_indices; --fit on is its emission. + assert gate < src.find("gpu_indices, use_fit = None, True") + assert 'cmd.extend(["--fit", "on"])' in src + # TP drops for this path, but at a guard BEFORE the quantized-KV cache-drop, so + # a requested quantized cache survives into the --fit load. + tp_drop = src.find('if tensor_parallel and gpu_memory_mode == "manual" and gpu_layers < 0:') + assert tp_drop != -1, "manual + Auto layers must drop tensor_parallel" + assert "tensor_parallel = False" in src[tp_drop : tp_drop + 400] + cache_drop = src.find("Tensor parallelism requires a non-quantized KV cache") + assert cache_drop != -1 + assert ( + tp_drop < cache_drop + ), "TP must drop before the cache-drop so a quantized KV survives --fit" + + +def test_auto_layers_never_sends_ctx_size_zero(): + # Sending "-c 0" sets fit_params_min_ctx = UINT32_MAX in llama.cpp, pinning + # the full native context and disabling --fit's reduction. So the base cmd + # must never carry -c, "-c 0" is emitted only outside the Auto-layers (--fit) + # case, and a positive context is passed through (which --fit optimizes + # layers around). + src = _load_model_source() + base_start = src.find("cmd = [") + base_end = src.find("\n ]", base_start) + base_block = src[base_start:base_end] + assert '"-c"' not in base_block, "-c must be conditional, not in the base cmd list" + assert 'cmd.extend(["-c", str(effective_ctx)])' in src, "positive ctx must pass -c" + assert 'auto_fit = gpu_memory_mode == "manual" and gpu_layers < 0' in src + zero = src.find('cmd.extend(["-c", "0"])') + assert zero != -1, '"-c 0" emission must exist outside the Auto-layers case' + guard = src.rfind("elif not auto_fit:", 0, zero) + assert guard != -1 and zero - guard < 120, '"-c 0" must sit under the not-auto_fit guard' + + +def test_manual_mode_clears_inherited_main_model_placement_env(): + env = {name: "inherited" for name in LlamaCppBackend._MANUAL_PLACEMENT_ENV_VARS} + env["LLAMA_ARG_N_GPU_LAYERS_DRAFT"] = "7" + env["UNRELATED"] = "kept" + + LlamaCppBackend._clear_manual_placement_env(env) + + assert not (set(env) & set(LlamaCppBackend._MANUAL_PLACEMENT_ENV_VARS)) + assert env["LLAMA_ARG_N_GPU_LAYERS_DRAFT"] == "7" + assert env["UNRELATED"] == "kept" + + +def test_load_model_sanitizes_manual_env_after_building_child_env(): + src = _load_model_source() + env_build = src.find("env = self._llama_server_env_for_binary(binary)") + env_clear = src.find("self._clear_manual_placement_env(env)", env_build) + launch = src.find("subprocess.Popen", env_build) + assert env_build != -1 + assert env_build < env_clear < launch + + +# ── Manual offload (--gpu-layers + --fit off + --n-cpu-moe) ─────────── + + +def test_load_request_accepts_manual(): + req = LoadRequest( + model_path = "owner/repo", + gpu_memory_mode = "manual", + gpu_layers = 20, + n_cpu_moe = 8, + tensor_split = [2, 1], + ) + assert req.gpu_memory_mode == "manual" + assert req.gpu_layers == 20 + assert req.n_cpu_moe == 8 + assert req.tensor_split == [2, 1] + + +def test_load_request_manual_defaults(): + req = LoadRequest(model_path = "owner/repo") + assert req.gpu_layers == -1 + assert req.n_cpu_moe == 0 + assert req.tensor_split is None + + +@pytest.mark.parametrize("bad", [[0, 0], [-1, 2], [float("inf"), 1], [float("nan"), 1]]) +def test_load_request_rejects_degenerate_tensor_split(bad): + # A negative/non-finite/all-zero split is dropped at launch but compared raw + # in the reload dedupe, so it would reload forever -- reject it up front. + with pytest.raises(ValueError): + LoadRequest(model_path = "owner/repo", tensor_split = bad) + + +@pytest.mark.parametrize("good", [[2, 1], [1, 1], [], None]) +def test_load_request_accepts_valid_tensor_split(good): + assert LoadRequest(model_path = "owner/repo", tensor_split = good).tensor_split == good + + +def test_route_normalizes_explicit_extras_before_reload_dedupe(): + route_src = (Path(_BACKEND_DIR) / "routes" / "inference.py").read_text(encoding = "utf-8") + load_impl = route_src[route_src.index("async def _load_model_impl") :] + strip = load_impl.index("_stripped_explicit = strip_shadowing_flags") + normalize = load_impl.index( + 'request = request.model_copy(update = {"llama_extra_args": extra_llama_args})' + ) + dedupe = load_impl.index("and _request_matches_loaded_settings(") + assert strip < normalize < dedupe + + +@pytest.mark.parametrize("model_cls", [LoadResponse, InferenceStatusResponse]) +def test_response_models_emit_manual_fields(model_cls): + if model_cls is LoadResponse: + obj = model_cls( + status = "loaded", + model = "owner/repo", + display_name = "repo", + inference = {}, + gpu_memory_mode = "manual", + gpu_layers = 20, + n_cpu_moe = 8, + tensor_split = [2, 1], + n_layers = 32, + n_moe_layers = 32, + ) + else: + obj = model_cls( + gpu_memory_mode = "manual", + gpu_layers = 20, + n_cpu_moe = 8, + tensor_split = [2, 1], + n_layers = 32, + n_moe_layers = 32, + ) + dumped = obj.model_dump() + assert dumped["gpu_memory_mode"] == "manual" + assert dumped["gpu_layers"] == 20 + assert dumped["n_cpu_moe"] == 8 + assert dumped["tensor_split"] == [2, 1] + assert dumped["n_layers"] == 32 + assert dumped["n_moe_layers"] == 32 + + +def test_manual_properties_default_and_reflect_and_reset(): + backend = LlamaCppBackend() + assert backend.gpu_layers == -1 and backend.n_cpu_moe == 0 + assert backend.tensor_split is None + backend._gpu_layers = 20 + backend._n_cpu_moe = 8 + backend._tensor_split = [2, 1] + assert backend.gpu_layers == 20 and backend.n_cpu_moe == 8 + assert backend.tensor_split == [2, 1] + backend._process = _FakeProcess() + backend.unload_model() + assert backend.gpu_layers == -1 and backend.n_cpu_moe == 0 + assert backend.tensor_split is None + + +def test_n_moe_layers_property(): + # 0 for a dense model (hides the slider); block_count for all-MoE; + # block_count - leading_dense otherwise (GLM-4.7-Flash: 47 - 1 -> 46). + b = LlamaCppBackend() + b._n_layers = 36 + b._n_experts = None + assert b.n_moe_layers == 0 + b._n_experts = 128 + b._leading_dense_block_count = None + assert b.n_moe_layers == 36 + b._n_layers = 47 + b._leading_dense_block_count = 1 + assert b.n_moe_layers == 46 + + +def _target_state_manual( + backend, + *, + gpu_layers, + n_cpu_moe, + tensor_split = None, +): + return backend._already_in_target_state( + gguf_path = None, + model_identifier = "owner/repo", + hf_variant = "Q4_K_M", + n_ctx = 8192, + cache_type_kv = None, + speculative_type = "auto", + chat_template_override = None, + extra_args = None, + is_vision = False, + gpu_memory_mode = "manual", + gpu_layers = gpu_layers, + n_cpu_moe = n_cpu_moe, + tensor_split = tensor_split, + ) + + +def test_manual_reloads_on_gpu_layers_or_n_cpu_moe_or_split_change(): + backend = _loaded_backend("manual") + backend._gpu_layers = 20 + backend._n_cpu_moe = 0 + backend._tensor_split = None + # Same knobs -> no reload. + assert _target_state_manual(backend, gpu_layers = 20, n_cpu_moe = 0) is True + # Changed layer count -> reload. + assert _target_state_manual(backend, gpu_layers = 16, n_cpu_moe = 0) is False + # Changed MoE offload -> reload. + assert _target_state_manual(backend, gpu_layers = 20, n_cpu_moe = 8) is False + # Added a GPU split -> reload. + assert _target_state_manual(backend, gpu_layers = 20, n_cpu_moe = 0, tensor_split = [2, 1]) is False + # Same GPU split -> no reload. + backend._tensor_split = [2, 1] + assert _target_state_manual(backend, gpu_layers = 20, n_cpu_moe = 0, tensor_split = [2, 1]) is True + + +def test_auto_layers_reload_tracks_only_gpu_layers(): + # Under Auto (gpu_layers < 0) the MoE/split knobs don't apply, so a leftover + # request value must not reload -- only a gpu_layers change (Auto -> pinned) does. + backend = _loaded_backend("manual") + backend._gpu_layers = -1 + backend._n_cpu_moe = 0 + backend._tensor_split = None + # Same Auto, leftover MoE/split in the request -> still no reload. + assert _target_state_manual(backend, gpu_layers = -1, n_cpu_moe = 8, tensor_split = [2, 1]) is True + # Auto -> explicit offload reloads. + assert _target_state_manual(backend, gpu_layers = 20, n_cpu_moe = 0) is False + + +def test_manual_offload_emits_gpu_layers_fit_off_and_n_cpu_moe(): + src = _load_model_source() + gate = src.find('elif gpu_memory_mode == "manual":') + assert gate != -1, "load_model must have an explicit-offload manual branch" + block = src[gate : gate + 700] + # Empties the probed set (skips the planner) but keeps the user's TP choice + # (only the Auto-layers branch above drops TP). + assert "gpus = []" in block + assert "tensor_parallel = False" not in block + # The cmd emits the layer count with fit disabled, gated on gpu_layers >= 0. + assert 'if gpu_memory_mode == "manual" and gpu_layers >= 0:' in src + assert 'cmd.extend(["--gpu-layers", str(gpu_layers), "--fit", "off"])' in src + # MoE offload uses --n-cpu-moe via _resolve_cpu_moe_flag (tested behaviorally below). + assert "_resolve_cpu_moe_flag(" in src + assert 'cmd.extend(["--n-cpu-moe", str(moe_flag)])' in src + # A count requested on a dense model is never emitted, so it must also be + # dropped from the recorded state -- else /status and /load report a count + # llama-server never received (same rule as the tensor-split drop below). + moe_emit = src.find('cmd.extend(["--n-cpu-moe", str(moe_flag)])') + assert "elif n_cpu_moe:" in src[moe_emit : moe_emit + 300] + assert "self._n_cpu_moe = 0" in src[moe_emit : moe_emit + 300] + # The offload path forces use_fit False so --fit-ctx is never added under --fit off. + emit = src.find('cmd.extend(["--gpu-layers", str(gpu_layers), "--fit", "off"])') + assert "use_fit = False" in src[src.rfind("\n", 0, emit) - 200 : emit + 80] + + +def test_status_reports_requested_context_length(): + # The hydration path re-seeds a Manual+Auto context pin from the REQUESTED + # n_ctx (0 = Auto); context_length only exposes the resolved value. + assert "requested_context_length" in InferenceStatusResponse.model_fields + s = InferenceStatusResponse(requested_context_length = 8192) + assert s.model_dump()["requested_context_length"] == 8192 + assert InferenceStatusResponse().model_dump()["requested_context_length"] is None + # The /status route must actually wire it from the backend (a declared-but- + # never-populated field would leave hydration silently reverting the pin). + from pathlib import Path as _P + + route_src = (_P(_BACKEND_DIR) / "routes" / "inference.py").read_text(encoding = "utf-8") + assert "requested_context_length = llama_backend.requested_n_ctx" in route_src + + +def test_manual_offload_emits_tensor_split(): + # The offload path emits --tensor-split from the per-GPU shares, only when + # provided, with >1 GPU in use, AND matching that count (a stale ratio on a + # narrowed picker or a mismatched direct-API list must not emit -- llama- + # server aborts on a split/GPU-count mismatch). + src = _load_model_source() + assert "if tensor_split and _split_gpus > 1:" in src + # Emit only on a length match AND a positive sanitized total: a mismatched + # or all-zero split aborts llama-server / assigns nothing, so it's dropped. + # The emitted list is the sanitized one (clamping tested behaviorally below). + assert "_sanitized_split = self._sanitize_tensor_split(tensor_split)" in src + assert "if len(_sanitized_split) == _split_gpus and _split_total > 0:" in src + assert '"--tensor-split"' in src + # Joined as a comma list (e.g. "2,1") within the explicit-offload cmd branch. + gate = src.find('if gpu_memory_mode == "manual" and gpu_layers >= 0:') + nxt = src.find("elif use_fit:", gate) + assert '","' in src[gate:nxt] and "tensor_split" in src[gate:nxt] + # A split with a single effective GPU is never emitted, so it must also be + # dropped from the recorded state -- else /status and /load report a ratio + # llama-server never received and the dedupe baseline preserves it. + assert "elif tensor_split:" in src[gate:nxt] + drop = src.find("elif tensor_split:", gate, nxt) + assert "self._tensor_split = None" in src[drop : drop + 250] + + +def test_sanitize_tensor_split_clamps_negative_and_non_finite(): + # Negative entries would launch a placement different from the ratio the + # UI showed; inf passes a plain > 0 total gate and would emit + # "--tensor-split inf,..." (llama.cpp normalizes shares by the running + # total, so an inf poisons the shares from that entry on). Both clamp to 0. + sanitize = LlamaCppBackend._sanitize_tensor_split + assert sanitize([2, 1]) == [2.0, 1.0] + assert sanitize([-1, 2]) == [0.0, 2.0] + assert sanitize([float("inf"), 1]) == [0.0, 1.0] + assert sanitize([float("nan"), 1]) == [0.0, 1.0] + # All-zero survives sanitization; the call site's total gate drops it. + assert sanitize([0, 0]) == [0.0, 0.0] + # Unreadable input -> []; the call site's length gate drops it. + assert sanitize(["x", 1]) == [] + assert sanitize([10**400, 1]) == [] + + +def test_zero_offload_mask_honors_device_pin_spellings(): + # A user device pin must keep the GPUs visible: llama-server aborts on a + # pin it can't see ('error: invalid device'). The pin can arrive as + # --device or its -dev alias, as the draft forms (parsed even with no + # drafter loaded), or as an inherited LLAMA_ARG_DEVICE env var. + load_src = _load_model_source() + assert "self._zero_offload_keeps_gpu_visible(cmd, env)" in load_src + block = inspect.getsource(LlamaCppBackend._cmd_has_gpu_device_pin) + for flag in ( + '"--device"', + '"-dev"', + '"--spec-draft-device"', + '"-devd"', + '"--device-draft"', + ): + assert flag in block + assert '"LLAMA_ARG_DEVICE"' in block + + +def test_resolve_cpu_moe_flag(): + # Clamp the requested MoE-layer count to the model's MoE layers, then offset + # past leading dense layers (--n-cpu-moe counts from layer 0). + R = LlamaCppBackend._resolve_cpu_moe_flag + assert R(0, 40, 0) is None # nothing requested + assert R(8, 0, 0) is None # dense model (no MoE layers) + assert R(8, 40, 0) == 8 # all-MoE: direct + assert R(100, 40, 0) == 40 # clamp to the MoE layer count + # GLM-4.7-Flash (deepseek2): block_count 47, leading_dense 1, n_moe 46. + assert R(5, 46, 1) == 6 # offset past the 1 dense layer + assert R(46, 46, 1) == 47 # all MoE on CPU == block_count + + +def test_manual_allows_tensor_parallel_via_split_mode(): + # Manual offload keeps the user's TP choice but skips the memory-based planner + # (plan_tp excludes manual, so its empty gpu set can't downgrade TP). The + # --split-mode tensor emission gates on tensor_parallel alone, so manual + # reaches it -- with tp_tensor_split None it's an even split (no + # --tensor-split). --fit off means no fit/tensor abort. + src = _load_model_source() + assert 'plan_tp = tensor_parallel and gpu_memory_mode != "manual"' in src + assert "if plan_tp:" in src + assert "if plan_tp and len(tp_gpus) < 2:" in src + sm = src.find('cmd.extend(["--split-mode", "tensor"])') + assert sm != -1, "TP must emit --split-mode tensor" + guard = src.rfind("if tensor_parallel:", 0, sm) + assert guard != -1 and sm - guard < 200, "split-mode gates on tensor_parallel" + # The tensor-split is only emitted for a planned (non-even) split, which + # manual never produces, so manual stays an even split. + assert "if tp_tensor_split and len(tp_tensor_split) > 1:" in src + + +def test_fit_sets_target_margin(): + # Manual + Auto (auto_fit) tightens the per-device VRAM margin to 512 MiB. + caps = {"supports_fit_target": True} + flags = LlamaCppBackend._ctx_integrity_flags(1, True, True, 0, 0, caps) + assert flags[flags.index("--fit-target") + 1] == "512" + # Not emitted on the legacy auto path (fit on but not auto_fit): -c 0 pins + # native there, so the tighter margin must not ride along. + assert "--fit-target" not in LlamaCppBackend._ctx_integrity_flags(1, True, False, 0, 0, caps) + # Not emitted when fit is off. + assert "--fit-target" not in LlamaCppBackend._ctx_integrity_flags(1, False, False, 0, 0, caps) + # Not emitted when the binary lacks support. + assert "--fit-target" not in LlamaCppBackend._ctx_integrity_flags( + 1, True, True, 0, 0, {"supports_fit_target": False} + ) + + +# ── GPU picker (gpu_ids -> CUDA_VISIBLE_DEVICES) ───────────────────── + + +def test_load_request_accepts_gpu_ids(): + req = LoadRequest(model_path = "owner/repo", gpu_ids = [1, 0]) + assert req.gpu_ids == [1, 0] + assert LoadRequest(model_path = "owner/repo").gpu_ids is None + + +@pytest.mark.parametrize("model_cls", [LoadResponse, InferenceStatusResponse]) +def test_response_models_emit_gpu_ids(model_cls): + if model_cls is LoadResponse: + obj = model_cls(status = "loaded", model = "m", display_name = "m", inference = {}, gpu_ids = [1]) + else: + obj = model_cls(gpu_ids = [1]) + assert obj.model_dump()["gpu_ids"] == [1] + + +def test_gpu_ids_property_default_and_reset(): + backend = LlamaCppBackend() + assert backend.gpu_ids is None + backend._gpu_ids = [0, 1] + assert backend.gpu_ids == [0, 1] + backend._process = _FakeProcess() + backend.unload_model() + assert backend.gpu_ids is None + + +def _target_state_gpu_ids(backend, gpu_ids): + return backend._already_in_target_state( + gguf_path = None, + model_identifier = "owner/repo", + hf_variant = "Q4_K_M", + n_ctx = 8192, + cache_type_kv = None, + speculative_type = "auto", + chat_template_override = None, + extra_args = None, + is_vision = False, + gpu_ids = gpu_ids, + ) + + +def test_gpu_ids_reload_detection_is_order_insensitive(): + backend = _loaded_backend("auto") + backend._gpu_ids = [0, 1] + # Same set, different order -> no reload. + assert _target_state_gpu_ids(backend, [1, 0]) is True + # Different set -> reload. + assert _target_state_gpu_ids(backend, [0]) is False + # Dropping the pick (auto) -> reload. + assert _target_state_gpu_ids(backend, None) is False + + +def test_gpu_ids_reload_detection_collapses_diffusion_to_single_device(): + # The diffusion runner drives only its single lowest device, so the backend + # records [lowest]. A later multi-GPU request that still resolves to that + # same lowest device must dedupe (no needless reload); a request whose lowest + # device moves, or that drops the pick, must reload. + backend = _loaded_backend("auto") + backend._is_diffusion = True + backend._gpu_ids = [1] # loaded on the lowest of an earlier [3, 1] pick + assert _target_state_gpu_ids(backend, [3, 1]) is True + assert _target_state_gpu_ids(backend, [1]) is True + # Lowest device changes (2, not 1) -> reload. + assert _target_state_gpu_ids(backend, [3, 2]) is False + # Dropping the pick (auto) -> reload. + assert _target_state_gpu_ids(backend, None) is False + + +def test_start_diffusion_server_resets_tensor_parallel(): + # A prior tensor-parallel chat load leaves self._tensor_parallel True (load_model + # phase 1 only kills the process, it skips the unload reset). Diffusion is never + # TP, so startup must clear it -- else /status misreports TP and an identical + # diffusion re-Apply reloads against stale tensor-parallel state. + src = inspect.getsource(llama_cpp_module.LlamaCppBackend._start_diffusion_server) + assert "self._tensor_parallel = False" in src + + +def test_route_matches_loaded_settings_collapses_diffusion_gpu_ids(): + # The route-level reload dedupe mirrors the backend: for a loaded diffusion + # model it compares the request against the single recorded device, not the + # full requested list, or a same-device multi-GPU pick reloads needlessly. + route_src = (Path(_BACKEND_DIR) / "routes" / "inference.py").read_text(encoding = "utf-8") + match_impl = route_src[route_src.index("def _request_matches_loaded_settings") :] + guard = match_impl.index("if llama_backend.is_diffusion:") + collapse = match_impl.index("[sorted(request.gpu_ids)[0]] if request.gpu_ids else None") + compare = match_impl.index("if _req_gpu_ids != llama_backend.gpu_ids:") + assert guard < collapse < compare + + +# ── Manual tensor split: child enumeration pinned to the picker's order ────── + + +def _patch_split_pin_env(monkeypatch, *, inherited, reported): + """Point the pin helper at a fake inherited mask and picker report. + ``reported`` None = enumeration unavailable (falls back to ascending).""" + import utils.hardware as hw + + monkeypatch.setattr( + LlamaCppBackend, "_resolve_visible_physical_ids", staticmethod(lambda: inherited) + ) + info = ( + {"available": False} + if reported is None + else { + "available": True, + "index_kind": "physical", + "devices": [{"index": i} for i in reported], + } + ) + monkeypatch.setattr(hw, "get_backend_visible_gpu_info", lambda: info) + + +def test_split_pin_reorders_inherited_numeric_mask(monkeypatch): + # Parent CUDA_VISIBLE_DEVICES=3,1 makes the child enumerate dev0=phys3, but + # nvidia-smi reported the picker's list ascending -- the mask must be + # re-emitted in that order or the per-GPU shares land on the wrong cards. + _patch_split_pin_env(monkeypatch, inherited = [3, 1], reported = [1, 3]) + env = {"CUDA_VISIBLE_DEVICES": "3,1"} + LlamaCppBackend._pin_visible_gpu_order_for_split(env) + assert env["CUDA_DEVICE_ORDER"] == "PCI_BUS_ID" + assert env["CUDA_VISIBLE_DEVICES"] == "1,3" + + +def test_split_pin_keeps_mask_order_when_picker_reported_it(monkeypatch): + # Torch-fallback enumeration (no nvidia-smi) reports devices in inherited + # mask order, so the picker's split list follows the mask -- the pin must + # keep that order, not re-sort it into a mismatch. + _patch_split_pin_env(monkeypatch, inherited = [3, 1], reported = [3, 1]) + env = {"CUDA_VISIBLE_DEVICES": "3,1"} + LlamaCppBackend._pin_visible_gpu_order_for_split(env) + assert env["CUDA_VISIBLE_DEVICES"] == "3,1" + + +def test_split_pin_falls_back_to_ascending_without_report(monkeypatch): + # Enumeration unavailable: ascending physical is the best guess (it matches + # the dominant nvidia-smi report order). + _patch_split_pin_env(monkeypatch, inherited = [3, 1], reported = None) + env = {"CUDA_VISIBLE_DEVICES": "3,1"} + LlamaCppBackend._pin_visible_gpu_order_for_split(env) + assert env["CUDA_VISIBLE_DEVICES"] == "1,3" + + +def test_split_pin_without_mask_only_sets_pci_order(monkeypatch): + # No inherited mask (or a UUID/MIG one resolving to None): enumeration order + # is fully fixed by CUDA_DEVICE_ORDER, so no mask is written. + _patch_split_pin_env(monkeypatch, inherited = None, reported = None) + env = {} + LlamaCppBackend._pin_visible_gpu_order_for_split(env) + assert env == {"CUDA_DEVICE_ORDER": "PCI_BUS_ID"} + + +def test_split_pin_mirrors_hip_mask_on_rocm(monkeypatch): + # ROCm: the pin must land in HIP_VISIBLE_DEVICES too, and an inherited ROCR + # mask is cleared so the mask can't apply twice (ROCR re-indexes, then HIP + # would index into the already-reduced set). + _patch_split_pin_env(monkeypatch, inherited = [3, 1], reported = [1, 3]) + torch_stub = _types.ModuleType("torch") + torch_stub.version = _types.SimpleNamespace(hip = "6.0") + monkeypatch.setitem(sys.modules, "torch", torch_stub) + env = {"CUDA_VISIBLE_DEVICES": "3,1", "ROCR_VISIBLE_DEVICES": "3,1"} + LlamaCppBackend._pin_visible_gpu_order_for_split(env) + assert env["CUDA_VISIBLE_DEVICES"] == "1,3" + assert env["HIP_VISIBLE_DEVICES"] == "1,3" + assert "ROCR_VISIBLE_DEVICES" not in env + + +# ── Diffusion single-device selection ─────────────────────────────────────── + + +def test_diffusion_gpu_arg_uses_lowest_explicit_physical_id(monkeypatch): + monkeypatch.setenv("CUDA_VISIBLE_DEVICES", "3,1") + monkeypatch.setenv("DG_GPU", "7") + assert LlamaCppBackend._diffusion_gpu_arg([3, 1]) == "1" + + +def test_diffusion_gpu_arg_preserves_parent_mask_order(monkeypatch): + monkeypatch.delenv("DG_GPU", raising = False) + monkeypatch.setenv("CUDA_VISIBLE_DEVICES", "3,1") + assert LlamaCppBackend._diffusion_gpu_arg(None) == "3" + + +def test_diffusion_gpu_arg_honors_override_and_cpu_mask(monkeypatch): + monkeypatch.setenv("DG_GPU", "GPU-abc") + assert LlamaCppBackend._diffusion_gpu_arg(None) == "GPU-abc" + assert LlamaCppBackend._diffusion_gpu_arg(None, cpu_only = True) == "" + + +# ── Deliberate zero-offload (manual gpu_layers=0): training-skip flag ───────── + + +def test_zero_offload_flag_false_without_companions(): + # CPU-only by construction: False lets training skip unloading a server that + # holds no VRAM. + cmd = ["llama-server", "-m", "model.gguf", "--gpu-layers", "0", "--fit", "off"] + assert LlamaCppBackend._zero_offload_gpu_flag(cmd, [(0, 8000, 24000)], {}) is False + + +@pytest.mark.parametrize( + "companion", + ["--mmproj", "--model-draft", "-md", "--spec-draft-model", "-hfd"], +) +def test_zero_offload_flag_true_with_companion(companion): + # mmproj / a drafter offload to GPU regardless of --gpu-layers, so the + # server still holds VRAM and training must unload it. Drafter detection + # reuses the extras parser, so pass-through aliases count too. + cmd = ["llama-server", "-m", "model.gguf", "--gpu-layers", "0", companion, "x.gguf"] + assert LlamaCppBackend._zero_offload_gpu_flag(cmd, [(0, 8000, 24000)], {}) is True + + +def test_zero_offload_flag_true_with_inline_companion_forms(): + cmd = ["llama-server", "-m", "model.gguf", "--spec-draft-model=x.gguf"] + assert LlamaCppBackend._zero_offload_gpu_flag(cmd, [(0, 8000, 24000)], {}) is True + cmd = ["llama-server", "-m", "model.gguf", "--mmproj=proj.gguf"] + assert LlamaCppBackend._zero_offload_gpu_flag(cmd, [(0, 8000, 24000)], {}) is True + + +def test_zero_offload_flag_true_with_env_drafter(): + cmd = ["llama-server", "-m", "model.gguf", "--gpu-layers", "0"] + env = {"LLAMA_ARG_SPEC_DRAFT_MODEL": "x.gguf"} + assert LlamaCppBackend._zero_offload_gpu_flag(cmd, [(0, 8000, 24000)], env) is True + + +@pytest.mark.parametrize( + "device_args", + [ + ["--device", "CUDA0"], + ["--device=CUDA0"], + ["-dev", "CUDA0"], + ["--spec-draft-device", "CUDA0"], + ["--device-draft=CUDA0"], + ], +) +def test_zero_offload_flag_true_with_device_pin(device_args): + cmd = ["llama-server", "-m", "model.gguf", "--gpu-layers", "0", *device_args] + assert LlamaCppBackend._zero_offload_gpu_flag(cmd, [(0, 8000, 24000)], {}) is True + + +def test_zero_offload_flag_true_with_env_device_pin(): + cmd = ["llama-server", "-m", "model.gguf", "--gpu-layers", "0"] + env = {"LLAMA_ARG_DEVICE": "CUDA0"} + assert LlamaCppBackend._zero_offload_gpu_flag(cmd, [(0, 8000, 24000)], env) is True + + +@pytest.mark.parametrize( + ("device_args", "env"), + [ + (["--device", "cpu"], {}), + (["--device=none"], {}), + (["--spec-draft-device", "cpu"], {}), + ([], {"LLAMA_ARG_DEVICE": "none"}), + (["--device", "CUDA0", "--device", "cpu"], {}), + ], +) +def test_zero_offload_flag_false_with_cpu_device_pin(device_args, env): + cmd = ["llama-server", "-m", "model.gguf", "--gpu-layers", "0", *device_args] + assert LlamaCppBackend._zero_offload_gpu_flag(cmd, [(0, 8000, 24000)], env) is False + + +def test_zero_offload_flag_true_with_surviving_tensor_mode(): + cmd = ["llama-server", "-m", "model.gguf", "--gpu-layers", "0", "--split-mode", "tensor"] + assert LlamaCppBackend._zero_offload_gpu_flag(cmd, [(0, 8000, 24000)], {}) is True + + +def test_zero_offload_flag_true_for_unmasked_vulkan(monkeypatch): + monkeypatch.setattr(LlamaCppBackend, "_is_vulkan_backend", staticmethod(lambda: True)) + cmd = ["llama-server", "-m", "model.gguf", "--gpu-layers", "0"] + assert LlamaCppBackend._zero_offload_gpu_flag(cmd, [(0, 8000, 24000)], {}) is True + + +def test_zero_offload_flag_none_without_gpus(): + cmd = ["llama-server", "-m", "model.gguf", "--gpu-layers", "0"] + assert LlamaCppBackend._zero_offload_gpu_flag(cmd, [], {}) is None + + +def test_cmd_has_gpu_companion_detection(): + # The env mask for CPU-only zero-offload loads keys off this scan: any + # --mmproj form or a drafter (flag aliases / env) keeps the GPUs visible. + has = LlamaCppBackend._cmd_has_gpu_companion + assert has(["llama-server", "-m", "m.gguf"], {}) is False + assert has(["llama-server", "--mmproj", "p.gguf"], {}) is True + assert has(["llama-server", "--mmproj=p.gguf"], {}) is True + assert has(["llama-server", "-md", "d.gguf"], {}) is True + assert has(["llama-server"], {"LLAMA_ARG_SPEC_DRAFT_MODEL": "d.gguf"}) is True + + +def test_cmd_companion_ignores_cpu_forced_drafter(): + # A CPU-pinned drafter holds no VRAM: the zero-offload mask may hide the GPUs + # and training may leave the server alone. + has = LlamaCppBackend._cmd_has_gpu_companion + cmd = ["llama-server", "-md", "d.gguf", "--spec-draft-ngl", "0"] + assert has(cmd, {}) is False + cmd = ["llama-server", "-md", "d.gguf", "--spec-draft-device", "cpu"] + assert has(cmd, {}) is False + # mmproj still counts even alongside a CPU drafter. + cmd = ["llama-server", "-md", "d.gguf", "--spec-draft-ngl", "0", "--mmproj", "p.gguf"] + assert has(cmd, {}) is True diff --git a/studio/backend/tests/test_gpu_selection.py b/studio/backend/tests/test_gpu_selection.py index 69ad560788..d4f2fbe993 100644 --- a/studio/backend/tests/test_gpu_selection.py +++ b/studio/backend/tests/test_gpu_selection.py @@ -853,7 +853,13 @@ class TestRouteErrors(unittest.TestCase): self.assertIn("only supported on CUDA devices", str(exc_info.exception)) - def test_inference_route_rejects_gpu_ids_for_gguf(self): + def test_inference_route_validates_gpu_ids_for_gguf(self): + # gpu_ids is now SUPPORTED for GGUF (the GPU picker), but still + # validated: a rejected pick surfaces as a clean 400, not the old + # "not supported for GGUF" rejection. Patch the validator so the test + # is deterministic regardless of the host's (or a prior test's) GPU env. + import utils.hardware.hardware as hardware_mod + inference_route = _load_route_module( "inference_route_module_for_gguf_gpu_ids_test", "routes/inference.py", @@ -887,6 +893,11 @@ class TestRouteErrors(unittest.TestCase): ), patch.object(inference_route.asyncio, "to_thread", new = _inline_to_thread), patch.object(inference_route, "_hf_offline_if_dns_dead", nullcontext), + patch.object( + hardware_mod, + "resolve_requested_gpu_ids", + side_effect = ValueError("Invalid gpu_ids [0, 1]: rejected by test"), + ), ): with self.assertRaises(HTTPException) as exc_info: asyncio.run( @@ -901,8 +912,11 @@ class TestRouteErrors(unittest.TestCase): ) ) + # The validator's ValueError becomes a clean 400 (not the removed + # "not supported for GGUF" rejection). self.assertEqual(exc_info.exception.status_code, 400) - self.assertIn("GGUF", exc_info.exception.detail) + self.assertIn("gpu_ids", exc_info.exception.detail.lower()) + self.assertNotIn("not supported", exc_info.exception.detail.lower()) def test_training_route_returns_400_for_invalid_gpu_ids(self): training_route = _load_route_module( diff --git a/studio/backend/tests/test_hf_token_validation.py b/studio/backend/tests/test_hf_token_validation.py new file mode 100644 index 0000000000..31b30fc37d --- /dev/null +++ b/studio/backend/tests/test_hf_token_validation.py @@ -0,0 +1,165 @@ +# SPDX-License-Identifier: AGPL-3.0-only +# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0 + +"""Focused coverage for cached, rate-limited HF token validation.""" + +from __future__ import annotations + +from pathlib import Path +import sys + +import httpx +import pytest + + +_BACKEND_DIR = str(Path(__file__).resolve().parent.parent) +if _BACKEND_DIR not in sys.path: + sys.path.insert(0, _BACKEND_DIR) + +import utils.hf_token_validation as validation + + +@pytest.fixture(autouse = True) +def _reset_validation_state(): + validation.reset_hf_token_validation_state() + yield + validation.reset_hf_token_validation_state() + + +def test_cached_token_does_not_spend_another_attempt(monkeypatch): + calls = [] + + def _check(token): + calls.append(token) + return validation.TokenValidationResult(status = "valid") + + monkeypatch.setattr(validation, "_check_remote", _check) + first = validation.validate_hf_token("hf_valid", rate_key = "user:ip") + second = validation.validate_hf_token("hf_valid", rate_key = "user:ip") + + assert first.status == second.status == "valid" + assert calls == ["hf_valid"] + + +def test_three_uncached_attempts_per_hour(monkeypatch): + monkeypatch.setattr( + validation, + "_check_remote", + lambda _token: validation.TokenValidationResult(status = "invalid"), + ) + + for index in range(3): + result = validation.validate_hf_token(f"hf_bad_{index}", rate_key = "user:ip") + assert result.status == "invalid" + + limited = validation.validate_hf_token("hf_bad_4", rate_key = "user:ip") + assert limited.status == "rate_limited" + assert limited.retry_after_seconds is not None + assert limited.retry_after_seconds > 0 + + other_user = validation.validate_hf_token("hf_other", rate_key = "other:ip") + assert other_user.status == "invalid" + + +def test_window_rolls_forward(monkeypatch): + clock = {"now": 100.0} + monkeypatch.setattr(validation.time, "monotonic", lambda: clock["now"]) + monkeypatch.setattr(validation, "_MAX_ATTEMPTS", 1) + monkeypatch.setattr(validation, "_WINDOW_SECONDS", 10.0) + monkeypatch.setattr( + validation, + "_check_remote", + lambda _token: validation.TokenValidationResult(status = "invalid"), + ) + + assert validation.validate_hf_token("hf_a", rate_key = "user:ip").status == "invalid" + assert validation.validate_hf_token("hf_b", rate_key = "user:ip").status == "rate_limited" + clock["now"] += 11.0 + assert validation.validate_hf_token("hf_b", rate_key = "user:ip").status == "invalid" + + +@pytest.mark.parametrize( + ("status_code", "expected"), + [(200, "valid"), (401, "invalid"), (429, "rate_limited"), (500, "unavailable")], +) +def test_remote_status_classification(monkeypatch, status_code, expected): + response = httpx.Response( + status_code, + request = httpx.Request("GET", "https://huggingface.co/api/whoami-v2"), + headers = {"Retry-After": "42"} if status_code == 429 else None, + ) + + class _Session: + def get(self, url, *, headers, timeout): + assert url == "https://huggingface.co/api/whoami-v2" + assert headers["authorization"] == "Bearer hf_test" + assert timeout == validation._REMOTE_TIMEOUT_SECONDS + return response + + monkeypatch.setattr(validation, "get_session", lambda: _Session()) + result = validation._check_remote("hf_test") + assert result.status == expected + if status_code == 429: + assert result.retry_after_seconds == 42 + + +def test_wrapped_http_401_is_invalid(monkeypatch): + response = httpx.Response( + 401, + request = httpx.Request("GET", "https://huggingface.co/api/whoami-v2"), + ) + + class _Session: + def get(self, _url, **_kwargs): + error = RuntimeError("Invalid user token.") + error.response = response + raise error + + monkeypatch.setattr(validation, "get_session", lambda: _Session()) + assert validation._check_remote("hf_test").status == "invalid" + + +def test_remote_timeout_is_bounded_and_unavailable(monkeypatch): + class _Session: + def get(self, _url, *, headers, timeout): + assert headers["authorization"] == "Bearer hf_test" + assert timeout == validation._REMOTE_TIMEOUT_SECONDS + raise TimeoutError("timed out") + + monkeypatch.setattr(validation, "get_session", lambda: _Session()) + assert validation._check_remote("hf_test").status == "unavailable" + + +def test_raw_token_is_not_retained(monkeypatch): + monkeypatch.setattr( + validation, + "_check_remote", + lambda _token: validation.TokenValidationResult(status = "valid"), + ) + token = "hf_do_not_store_this_value" + validation.validate_hf_token(token, rate_key = "user:ip") + + assert token not in repr(validation._cache) + assert token not in repr(validation._attempts) + + +def test_unexpected_remote_exception_releases_singleflight(monkeypatch): + calls = 0 + monkeypatch.setattr(validation, "_INFLIGHT_WAIT_SECONDS", 0.0) + + def _check(_token): + nonlocal calls + calls += 1 + if calls == 1: + raise RuntimeError("unexpected failure") + return validation.TokenValidationResult(status = "valid") + + monkeypatch.setattr(validation, "_check_remote", _check) + + with pytest.raises(RuntimeError, match = "unexpected failure"): + validation.validate_hf_token("hf_test", rate_key = "user:ip") + + result = validation.validate_hf_token("hf_test", rate_key = "user:ip") + assert result.status == "valid" + assert calls == 2 + assert validation._inflight == {} diff --git a/studio/backend/tests/test_install_resolve_prebuilt.py b/studio/backend/tests/test_install_resolve_prebuilt.py index e97ca47717..3ebad861ad 100644 --- a/studio/backend/tests/test_install_resolve_prebuilt.py +++ b/studio/backend/tests/test_install_resolve_prebuilt.py @@ -445,6 +445,101 @@ def test_route_to_vulkan_prebuilt_cpu_fallback_wins(): assert routed is host +@pytest.mark.parametrize("cpu_flag", ["--cpu-fallback", "--force-cpu"]) +def test_resolve_prebuilt_cpu_fallback_overrides_intel_vulkan(monkeypatch, capsys, cpu_flag): + """Either CPU flag via CLI must suppress Vulkan even on an Intel GPU host: both + drop GPU detection (--force-cpu additionally persists, on the install path).""" + monkeypatch.setattr( + ilp, + "detect_host", + lambda: _host(is_linux = True, is_x86_64 = True, has_intel_gpu = True), + ) + seen = {} + + def _resolver(tag, host, repo, published_release_tag): + seen["host"] = host + seen["repo"] = repo + raise ilp.PrebuiltFallback("no asset") + + monkeypatch.setattr(ilp, "resolve_simple_install_release_plans", _resolver) + monkeypatch.setattr( + sys, + "argv", + [ + "install_llama_prebuilt.py", + "--resolve-prebuilt", + "latest", + cpu_flag, + "--output-format", + "json", + ], + ) + assert ilp.main() == ilp.EXIT_SUCCESS + # The CPU flag must suppress Intel GPU, route to fork (not upstream Vulkan) + assert seen["host"].has_intel_gpu is False + assert seen["repo"] == FORK + + +@pytest.mark.parametrize( + "flags, expect_force, expect_persist", + [ + ([], False, False), + # Automatic/transient last resort (arm64 GPU-build recovery): drops GPU but + # does NOT persist, so a later update heals to a GPU bundle (#6097). + (["--cpu-fallback"], True, False), + # Deliberate CPU-only (UNSLOTH_LLAMA_CPP_BACKEND=cpu): drops GPU AND persists so + # the updater re-asserts it and never revives the Intel iGPU crash (#7213). + (["--force-cpu"], True, True), + (["--cpu-fallback", "--force-cpu"], True, True), + ], +) +def test_cli_cpu_flags_thread_force_and_persist( + monkeypatch, tmp_path, flags, expect_force, expect_persist +): + captured = {} + monkeypatch.setattr(ilp, "install_prebuilt", lambda **kw: captured.update(kw)) + monkeypatch.setattr( + sys, + "argv", + ["install_llama_prebuilt.py", "--install-dir", str(tmp_path / "llama.cpp"), *flags], + ) + assert ilp.main() == ilp.EXIT_SUCCESS + assert captured["force_cpu"] is expect_force + assert captured["persist_force_cpu"] is expect_persist + + +@pytest.mark.parametrize( + "existing, requested, expected", + [ + # A deliberate --force-cpu on top of a naturally-installed CPU bundle (same + # asset, install skipped) must still flip the marker to true (#7213). + (False, True, True), + (None, True, True), + # No spurious writes when already in sync, and a released force syncs down. + (True, True, True), + (False, False, False), + (True, False, False), + ], +) +def test_sync_marker_force_cpu(tmp_path, existing, requested, expected): + marker = {"tag": "b9585", "asset": "llama-b9585-bin-ubuntu-x64.tar.gz"} + if existing is not None: + marker["force_cpu"] = existing + marker_path = tmp_path / "UNSLOTH_PREBUILT_INFO.json" + marker_path.write_text(json.dumps(marker)) + ilp.sync_marker_force_cpu(tmp_path, requested) + written = json.loads(marker_path.read_text()) + assert written["force_cpu"] is expected + # Unrelated fields are preserved. + assert written["asset"] == "llama-b9585-bin-ubuntu-x64.tar.gz" + + +def test_sync_marker_force_cpu_missing_marker_is_noop(tmp_path): + # No marker (or unreadable) must not crash the reuse path. + ilp.sync_marker_force_cpu(tmp_path, True) + assert not (tmp_path / "UNSLOTH_PREBUILT_INFO.json").exists() + + def test_route_to_vulkan_prebuilt_hidden_nvidia_not_rerouted(): # A mixed NVIDIA+Intel host that hid NVIDIA (CUDA_VISIBLE_DEVICES=""/-1): # physical NVIDIA present but not usable. Must NOT auto-route to Vulkan, or diff --git a/studio/backend/tests/test_llama_cpp_mtp_detection.py b/studio/backend/tests/test_llama_cpp_mtp_detection.py index 8fe04c0e39..1d15647967 100644 --- a/studio/backend/tests/test_llama_cpp_mtp_detection.py +++ b/studio/backend/tests/test_llama_cpp_mtp_detection.py @@ -9,6 +9,7 @@ the _already_in_target_state mirror that prevents needless reloads. from __future__ import annotations +import ast import inspect import os import struct @@ -345,10 +346,62 @@ def test_windows_full_offload_flags_use_current_llama_server_args(): stale_checkpoint_flag = "--checkpoint-" + "every-n-tokens" assert '"--cache-ram"' in src assert '"--ctx-checkpoints"' in src - assert '"--no-cache-prompt"' in src + # Prompt caching stays on (in-VRAM prefix reuse); #5692 only needed the host-RAM + # checkpoints (--cache-ram / --ctx-checkpoints) disabled, not prompt reuse. + assert '"--no-cache-prompt"' not in src assert stale_checkpoint_flag not in src +# Backend-wide guard: Unsloth must never inject --no-cache-prompt into a llama-server +# command. It disables in-VRAM prompt-prefix reuse, re-prefilling every repeated prompt +# (#5692 only needed --cache-ram / --ctx-checkpoints off; #7260 dropped the stray flag). +# Detecting it (_is_real) or honouring a user-supplied one (_prompt_cache_off) is fine. +_NO_CACHE_PROMPT_FLAG = "--no-cache-prompt" +_LIST_MUTATORS = frozenset({"append", "extend", "insert"}) + + +def _has_flag_literal(node: ast.AST) -> bool: + return any( + isinstance(n, ast.Constant) and n.value == _NO_CACHE_PROMPT_FLAG for n in ast.walk(node) + ) + + +def _no_cache_prompt_injections(source: str, filename: str) -> list[tuple[str, int]]: + """(file, lineno) for each spot adding --no-cache-prompt to a list.""" + hits: list[tuple[str, int]] = [] + for node in ast.walk(ast.parse(source, filename = filename)): + # cmd.append/extend/insert(... flag ...) or cmd += [... flag ...] + if ( + isinstance(node, ast.Call) + and isinstance(node.func, ast.Attribute) + and node.func.attr in _LIST_MUTATORS + and any(_has_flag_literal(a) for a in node.args) + ) or ( + isinstance(node, ast.AugAssign) + and isinstance(node.op, ast.Add) + and _has_flag_literal(node.value) + ): + hits.append((filename, node.lineno)) + return hits + + +def test_unsloth_never_injects_no_cache_prompt_into_any_command(): + root = Path(_BACKEND_DIR) + files = [p for p in root.rglob("*.py") if "tests" not in p.relative_to(root).parts] + violations: list[tuple[str, int]] = [] + for path in files: + try: + violations += _no_cache_prompt_injections(path.read_text(encoding = "utf-8"), str(path)) + except (OSError, UnicodeDecodeError, SyntaxError): + continue + assert files, "no backend source files were scanned" + assert violations == [], ( + "Unsloth must never add --no-cache-prompt to a llama-server command " + "(it disables prompt-prefix reuse); detecting or honouring a user-supplied " + f"one is fine. Offending sites: {violations}" + ) + + def test_load_model_sets_threads_once(): src = inspect.getsource(LlamaCppBackend.load_model) assert src.count('cmd.extend(["--threads", str(') == 1 @@ -741,6 +794,25 @@ def test_probe_reports_windows_cache_flags_absent_for_older_binary(tmp_path): assert caps["supports_no_cache_prompt"] is False +@_NEEDS_BASH +def test_probe_detects_slot_save_path(tmp_path): + fake = _make_fake_llama_server( + tmp_path / "llama-server", + "--slot-save-path PATH path to save slot kv cache\n--threads N\n", + ) + _clear_caps_cache() + caps = LlamaCppBackend.probe_server_capabilities(str(fake)) + assert caps["supports_slot_save"] is True + + +@_NEEDS_BASH +def test_probe_reports_slot_save_absent_for_older_binary(tmp_path): + fake = _make_fake_llama_server(tmp_path / "llama-server", "--threads N\n") + _clear_caps_cache() + caps = LlamaCppBackend.probe_server_capabilities(str(fake)) + assert caps["supports_slot_save"] is False + + def test_build_ngram_mod_flags_new(): flags = _build_ngram_mod_flags({"ngram_mod_flavor": "new"}) assert flags == [ diff --git a/studio/backend/tests/test_llama_cpp_no_context_shift.py b/studio/backend/tests/test_llama_cpp_no_context_shift.py index f320d29a02..662c918305 100644 --- a/studio/backend/tests/test_llama_cpp_no_context_shift.py +++ b/studio/backend/tests/test_llama_cpp_no_context_shift.py @@ -118,9 +118,17 @@ def test_flag_sits_inside_the_base_cmd_list(): "conditional branch -- otherwise some code paths would still " "run with silent context shift enabled." ) - # Pin that it sits next to -c / --ctx so the grouping makes sense. - assert '"-c"' in block assert '"--flash-attn"' in block + # -c is emitted in the conditional right after the base list, not inside + # it: auto-fit (--fit on with no pinned context) must omit -c entirely, + # because "-c 0" pins the full native context and disables --fit's + # VRAM-based sizing. Pin that it still sits next to the base block so the + # context grouping stays intact. + after = rest[end_rel : end_rel + 1000] + assert '"-c"' in after, ( + "-c must still be emitted in the conditional immediately after the " + "base cmd list (omitted only in auto-fit, where --fit sizes context)." + ) def _iter_lines_with_offset(text: str): diff --git a/studio/backend/tests/test_llama_cpp_props_readback.py b/studio/backend/tests/test_llama_cpp_props_readback.py index 488645ee5a..fe1e67edad 100644 --- a/studio/backend/tests/test_llama_cpp_props_readback.py +++ b/studio/backend/tests/test_llama_cpp_props_readback.py @@ -225,31 +225,46 @@ 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; Unsloth must restore the shared pool so one request can use the full advertised context.""" - flags = LlamaCppBackend._ctx_integrity_flags(4, False, 98304, 98304, _CAPS_ALL) + flags = LlamaCppBackend._ctx_integrity_flags(4, False, False, 98304, 98304, _CAPS_ALL) assert "--kv-unified" in flags def test_kv_unified_skipped_for_single_slot_or_old_build(): assert "--kv-unified" not in LlamaCppBackend._ctx_integrity_flags( - 1, False, 98304, 98304, _CAPS_ALL + 1, False, False, 98304, 98304, _CAPS_ALL ) assert "--kv-unified" not in LlamaCppBackend._ctx_integrity_flags( - 4, False, 98304, 98304, _CAPS_NONE + 4, False, False, 98304, 98304, _CAPS_NONE ) def test_fit_ctx_floors_explicit_request_under_fit(): - flags = LlamaCppBackend._ctx_integrity_flags(1, True, 98304, 98304, _CAPS_ALL) + # An explicit requested ctx floors --fit-ctx at that value on any --fit + # path, including legacy auto (auto_fit False). + flags = LlamaCppBackend._ctx_integrity_flags(1, True, False, 98304, 98304, _CAPS_ALL) assert flags[flags.index("--fit-ctx") + 1] == "98304" -def test_fit_ctx_skipped_without_fit_or_explicit_ctx_or_support(): +def test_fit_ctx_skipped_without_fit_or_support(): + # No --fit on -> no --fit-ctx. assert "--fit-ctx" not in LlamaCppBackend._ctx_integrity_flags( - 1, False, 98304, 98304, _CAPS_ALL + 1, False, False, 98304, 98304, _CAPS_ALL ) - assert "--fit-ctx" not in LlamaCppBackend._ctx_integrity_flags(1, True, 0, 262144, _CAPS_ALL) + # --fit on but the binary doesn't support --fit-ctx. assert "--fit-ctx" not in LlamaCppBackend._ctx_integrity_flags( - 1, True, 98304, 98304, _CAPS_NONE + 1, True, True, 98304, 98304, _CAPS_NONE + ) + + +def test_fit_ctx_floors_auto_request_at_8192_only_under_auto_fit(): + # Manual + Auto (auto_fit) floors the auto window at 8192 so --fit can't + # shrink it to a tiny size. + flags = LlamaCppBackend._ctx_integrity_flags(1, True, True, 0, 262144, _CAPS_ALL) + assert flags[flags.index("--fit-ctx") + 1] == "8192" + # Legacy auto (fit on but not auto_fit) emits -c 0 to pin native, so the + # 8192 floor must NOT ride along and override that pin. + assert "--fit-ctx" not in LlamaCppBackend._ctx_integrity_flags( + 1, True, False, 0, 262144, _CAPS_ALL ) diff --git a/studio/backend/tests/test_llama_cpp_slot_resume.py b/studio/backend/tests/test_llama_cpp_slot_resume.py new file mode 100644 index 0000000000..8b20c952c4 --- /dev/null +++ b/studio/backend/tests/test_llama_cpp_slot_resume.py @@ -0,0 +1,494 @@ +# SPDX-License-Identifier: AGPL-3.0-only +# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0 + +import os +from types import SimpleNamespace + +import core.inference.llama_cpp as llama_cpp +from core.inference.llama_cpp import LlamaCppBackend + + +def _resume_backend(tmp_path, n_slots = 1): + backend = LlamaCppBackend() + backend._healthy = True + # No-op lifecycle methods so the atexit cleanup can kill the fake quietly. + backend._process = SimpleNamespace( + poll = lambda: None, + terminate = lambda: None, + wait = lambda *a, **k: 0, + kill = lambda: None, + pid = 0, + ) + backend._port = 8081 + backend._slot_save_dir = str(tmp_path) + backend._slot_save_binary = ("/bin/llama-server", 1) + (tmp_path / "model.gguf").write_bytes(b"gguf") + backend._gguf_path = str(tmp_path / "model.gguf") + backend._effective_parallel_slots = n_slots + backend._estimate_kv_cache_bytes = lambda *a, **k: 0 + return backend + + +def _fake_disk(monkeypatch, free = 1 << 40): + monkeypatch.setattr(llama_cpp.shutil, "disk_usage", lambda _p: SimpleNamespace(free = free)) + + +class _Resp: + def __init__( + self, + status_code = 200, + body = None, + ): + self.status_code = status_code + self._body = body or {} + + def json(self): + return self._body + + +def test_save_returns_none_when_slot_save_disabled(monkeypatch, tmp_path): + backend = _resume_backend(tmp_path) + backend._slot_save_dir = None + monkeypatch.setattr( + llama_cpp.httpx, + "post", + lambda *a, **k: (_ for _ in ()).throw(AssertionError), + raising = False, + ) + assert backend.save_slots_for_resume() is None + + +def test_save_skipped_when_prompt_cache_disabled(monkeypatch, tmp_path): + backend = _resume_backend(tmp_path) + backend._prompt_cache_disabled = True + monkeypatch.setattr( + llama_cpp.httpx, + "post", + lambda *a, **k: (_ for _ in ()).throw(AssertionError), + raising = False, + ) + assert backend.save_slots_for_resume() is None + + +def test_save_skipped_when_insufficient_free_disk(monkeypatch, tmp_path): + backend = _resume_backend(tmp_path) + backend._estimate_kv_cache_bytes = lambda *a, **k: 1 << 40 + _fake_disk(monkeypatch, free = 1 << 20) + monkeypatch.setattr( + llama_cpp.httpx, + "post", + lambda *a, **k: (_ for _ in ()).throw(AssertionError), + raising = False, + ) + assert backend.save_slots_for_resume() is None + + +def test_save_collects_manifest_across_slots(monkeypatch, tmp_path): + backend = _resume_backend(tmp_path, n_slots = 2) + _fake_disk(monkeypatch) + calls = [] + + def fake_post(url, **kwargs): + calls.append((url, kwargs["params"], kwargs["json"])) + return _Resp(200, {"n_saved": 40, "n_written": 100}) + + monkeypatch.setattr(llama_cpp.httpx, "post", fake_post, raising = False) + manifest = backend.save_slots_for_resume() + assert manifest is not None + assert manifest["dir"] == str(tmp_path) + assert manifest["binary"] == ("/bin/llama-server", 1) + assert manifest["gguf"] == str(tmp_path / "model.gguf") + st = os.stat(manifest["gguf"]) + assert manifest["gguf_stat"] == ((st.st_size, st.st_mtime_ns),) + assert manifest["launch"] == backend._slot_launch_fingerprint() + assert [e["id"] for e in manifest["slots"]] == [0, 1] + assert all(e["n_saved"] == 40 for e in manifest["slots"]) + assert [c[1] for c in calls] == [{"action": "save"}] * 2 + assert "/slots/0" in calls[0][0] and "/slots/1" in calls[1][0] + + +def test_save_unlinks_empty_slot_and_returns_none(monkeypatch, tmp_path): + backend = _resume_backend(tmp_path) + _fake_disk(monkeypatch) + + def fake_post(url, **kwargs): + (tmp_path / kwargs["json"]["filename"]).write_bytes(b"") + return _Resp(200, {"n_saved": 0, "n_written": 0}) + + monkeypatch.setattr(llama_cpp.httpx, "post", fake_post, raising = False) + assert backend.save_slots_for_resume() is None + assert list(tmp_path.glob("resume-*.bin")) == [] # empty-slot file removed + + +def test_save_cap_breach_discards_all_files(monkeypatch, tmp_path): + backend = _resume_backend(tmp_path, n_slots = 2) + _fake_disk(monkeypatch) + monkeypatch.setattr(llama_cpp, "_SLOT_SAVE_MAX_BYTES", 150) + + def fake_post(url, **kwargs): + (tmp_path / kwargs["json"]["filename"]).write_bytes(b"x" * 100) + return _Resp(200, {"n_saved": 40, "n_written": 100}) + + monkeypatch.setattr(llama_cpp.httpx, "post", fake_post, raising = False) + assert backend.save_slots_for_resume() is None # 200 bytes > 150 cap + assert list(tmp_path.glob("resume-*.bin")) == [] + + +def test_save_transport_error_aborts_remaining_slots(monkeypatch, tmp_path): + backend = _resume_backend(tmp_path, n_slots = 3) + _fake_disk(monkeypatch) + calls = [] + + def fake_post(url, **kwargs): + calls.append(url) + raise OSError("connection refused") + + monkeypatch.setattr(llama_cpp.httpx, "post", fake_post, raising = False) + assert backend.save_slots_for_resume() is None + assert len(calls) == 1 # no retries against a dead server + + +def test_save_transport_error_unlinks_partial_file(monkeypatch, tmp_path): + backend = _resume_backend(tmp_path) + _fake_disk(monkeypatch) + + def fake_post(url, **kwargs): + (tmp_path / kwargs["json"]["filename"]).write_bytes(b"partial") + raise OSError("timed out") + + monkeypatch.setattr(llama_cpp.httpx, "post", fake_post, raising = False) + assert backend.save_slots_for_resume() is None + assert list(tmp_path.glob("resume-*.bin")) == [] + + +def test_fingerprint_tracks_lora_sidecar_rewrite(tmp_path): + backend = _resume_backend(tmp_path) + adapter = tmp_path / "adapter.gguf" + adapter.write_bytes(b"v1") + backend._extra_args = ["--lora", str(adapter)] + + before = backend._slot_launch_fingerprint() + adapter.write_bytes(b"v2-different") # re-exported adapter, same path + assert backend._slot_launch_fingerprint() != before + + backend._extra_args = [f"--lora={adapter}"] + assert backend._sidecar_weight_files() == [str(adapter)] + backend._extra_args = ["--lora-scaled", str(adapter), "0.5"] + assert backend._sidecar_weight_files() == [str(adapter)] + backend._extra_args = ["--control-vector", str(adapter), "--threads", "4"] + assert backend._sidecar_weight_files() == [str(adapter)] + + +def test_sidecar_files_parse_csv_and_colon_scale(tmp_path): + backend = _resume_backend(tmp_path) + a, b = tmp_path / "a.gguf", tmp_path / "b.gguf" + + backend._extra_args = ["--lora", f"{a},{b}"] + files = backend._sidecar_weight_files() + assert str(a) in files and str(b) in files + + backend._extra_args = ["--lora-scaled", f"{a}:0.5"] + assert str(a) in backend._sidecar_weight_files() + + backend._extra_args = ["--control-vector-scaled", f"{a}:1.0,{b}:2.0"] + files = backend._sidecar_weight_files() + assert str(a) in files and str(b) in files + + # Windows drive letter must not be mistaken for a scale separator. + backend._extra_args = ["--lora-scaled", "C:\\adapters\\a.gguf:0.75"] + assert "C:\\adapters\\a.gguf" in backend._sidecar_weight_files() + backend._extra_args = ["--lora", "C:\\adapters\\a.gguf"] + assert backend._sidecar_weight_files() == ["C:\\adapters\\a.gguf"] + + +def test_fingerprint_tracks_colon_scaled_adapter_rewrite(tmp_path): + backend = _resume_backend(tmp_path) + adapter = tmp_path / "adapter.gguf" + adapter.write_bytes(b"v1") + backend._extra_args = ["--lora-scaled", f"{adapter}:0.5"] + + before = backend._slot_launch_fingerprint() + adapter.write_bytes(b"v2-different") # re-exported adapter, same path + assert backend._slot_launch_fingerprint() != before + + +def test_fingerprint_tracks_effective_context_length(tmp_path): + backend = _resume_backend(tmp_path) + backend._effective_context_length = 8192 + + before = backend._slot_launch_fingerprint() + backend._effective_context_length = 4096 # auto-fit landed smaller on reload + assert backend._slot_launch_fingerprint() != before + + +def test_gguf_file_identity_covers_split_shards(tmp_path): + backend = _resume_backend(tmp_path) + first = tmp_path / "m-00001-of-00002.gguf" + second = tmp_path / "m-00002-of-00002.gguf" + first.write_bytes(b"a") + second.write_bytes(b"bb") + + before = backend._gguf_file_identity(str(first)) + st1, st2 = os.stat(first), os.stat(second) + assert before == ((st1.st_size, st1.st_mtime_ns), (st2.st_size, st2.st_mtime_ns)) + + second.write_bytes(b"rewritten") # sibling changes, primary untouched + after = backend._gguf_file_identity(str(first)) + assert after is not None and after != before + assert after[0] == before[0] # primary shard unchanged + + second.unlink() + assert backend._gguf_file_identity(str(first)) is None # missing shard + + +def test_save_skipped_when_user_disabled_prompt_cache(monkeypatch, tmp_path): + backend = _resume_backend(tmp_path) + backend._extra_args = ["--no-cache-prompt"] + monkeypatch.setattr( + llama_cpp.httpx, + "post", + lambda *a, **k: (_ for _ in ()).throw(AssertionError), + raising = False, + ) + assert backend.save_slots_for_resume() is None + + +def test_save_skipped_when_env_disables_prompt_cache(monkeypatch, tmp_path): + backend = _resume_backend(tmp_path) + monkeypatch.setenv("LLAMA_ARG_CACHE_PROMPT", "0") + monkeypatch.setattr( + llama_cpp.httpx, + "post", + lambda *a, **k: (_ for _ in ()).throw(AssertionError), + raising = False, + ) + assert backend.save_slots_for_resume() is None + monkeypatch.delenv("LLAMA_ARG_CACHE_PROMPT") + monkeypatch.setenv("LLAMA_ARG_NO_CACHE_PROMPT", "1") # legacy negative form + assert backend.save_slots_for_resume() is None + + +def test_explicit_cache_prompt_flag_overrides_env(monkeypatch, tmp_path): + backend = _resume_backend(tmp_path) + monkeypatch.setenv("LLAMA_ARG_CACHE_PROMPT", "0") + backend._extra_args = ["--cache-prompt"] # CLI wins over env in llama.cpp + _fake_disk(monkeypatch) + monkeypatch.setattr( + llama_cpp.httpx, + "post", + lambda *a, **k: _Resp(200, {"n_saved": 1, "n_written": 1}), + raising = False, + ) + assert backend.save_slots_for_resume() is not None + + +def test_user_cache_prompt_overrides_studio_no_cache_flag(monkeypatch, tmp_path): + # User extras follow Studio's flags, so an explicit --cache-prompt wins. + backend = _resume_backend(tmp_path) + backend._prompt_cache_disabled = True + backend._extra_args = ["--cache-prompt"] + _fake_disk(monkeypatch) + monkeypatch.setattr( + llama_cpp.httpx, + "post", + lambda *a, **k: _Resp(200, {"n_saved": 1, "n_written": 1}), + raising = False, + ) + assert backend.save_slots_for_resume() is not None + # Last flag wins when both appear in extras. + backend._extra_args = ["--cache-prompt", "--no-cache-prompt"] + assert backend.save_slots_for_resume() is None + + +def test_save_stops_writing_once_cap_exceeded(monkeypatch, tmp_path): + backend = _resume_backend(tmp_path, n_slots = 3) + _fake_disk(monkeypatch) + monkeypatch.setattr(llama_cpp, "_SLOT_SAVE_MAX_BYTES", 150) + calls = [] + + def fake_post(url, **kwargs): + calls.append(url) + (tmp_path / kwargs["json"]["filename"]).write_bytes(b"x" * 100) + return _Resp(200, {"n_saved": 1, "n_written": 100}) + + monkeypatch.setattr(llama_cpp.httpx, "post", fake_post, raising = False) + assert backend.save_slots_for_resume() is None + assert len(calls) == 2 # cap blown after slot 1; slot 2 never attempted + assert list(tmp_path.glob("resume-*.bin")) == [] + + +def test_save_aborts_between_slots_when_no_longer_idle(monkeypatch, tmp_path): + backend = _resume_backend(tmp_path, n_slots = 3) + _fake_disk(monkeypatch) + calls = [] + + def fake_post(url, **kwargs): + calls.append(url) + return _Resp(200, {"n_saved": 5, "n_written": 10}) + + monkeypatch.setattr(llama_cpp.httpx, "post", fake_post, raising = False) + aborts = iter([False, True, True]) + manifest = backend.save_slots_for_resume(should_abort = lambda: next(aborts)) + assert len(calls) == 1 # slots 1 and 2 skipped + assert manifest is not None + assert [e["id"] for e in manifest["slots"]] == [0] + + +def test_save_non_200_slot_is_skipped_but_others_kept(monkeypatch, tmp_path): + backend = _resume_backend(tmp_path, n_slots = 2) + _fake_disk(monkeypatch) + + def fake_post(url, **kwargs): + if "/slots/0" in url: + return _Resp(500) + return _Resp(200, {"n_saved": 5, "n_written": 10}) + + monkeypatch.setattr(llama_cpp.httpx, "post", fake_post, raising = False) + manifest = backend.save_slots_for_resume() + assert manifest is not None + assert [e["id"] for e in manifest["slots"]] == [1] + + +def test_restore_posts_each_slot_and_tolerates_failures(monkeypatch, tmp_path): + backend = _resume_backend(tmp_path) + calls = [] + + def fake_post(url, **kwargs): + calls.append((url, kwargs["params"], kwargs["json"])) + return _Resp(500 if "/slots/0" in url else 200, {"n_restored": 5}) + + monkeypatch.setattr(llama_cpp.httpx, "post", fake_post, raising = False) + backend.restore_slots_for_resume( + { + "slots": [ + {"id": 0, "filename": "resume-a-slot0.bin", "n_saved": 5}, + {"id": 1, "filename": "resume-a-slot1.bin", "n_saved": 5}, + ] + } + ) + assert [c[1] for c in calls] == [{"action": "restore"}] * 2 + assert calls[0][2] == {"filename": "resume-a-slot0.bin"} + + +def test_restore_transport_error_stops_early(monkeypatch, tmp_path): + backend = _resume_backend(tmp_path) + calls = [] + + def fake_post(url, **kwargs): + calls.append(url) + raise OSError("connection refused") + + monkeypatch.setattr(llama_cpp.httpx, "post", fake_post, raising = False) + backend.restore_slots_for_resume( + {"slots": [{"id": 0, "filename": "a.bin"}, {"id": 1, "filename": "b.bin"}]} + ) + assert len(calls) == 1 + + +def test_save_deletes_orphan_on_malformed_response(monkeypatch, tmp_path): + # A 200 that writes a file but returns a non-numeric counter must be cleaned + # up like any other save failure, not left orphaned holding chat KV. + backend = _resume_backend(tmp_path) + _fake_disk(monkeypatch) + + def fake_post(url, **kwargs): + (tmp_path / kwargs["json"]["filename"]).write_bytes(b"chat-kv") + return _Resp(200, {"n_saved": "not-an-int"}) + + monkeypatch.setattr(llama_cpp.httpx, "post", fake_post, raising = False) + assert backend.save_slots_for_resume() is None + assert list(tmp_path.glob("resume-*.bin")) == [] + + +def test_save_deletes_orphan_on_non_dict_response(monkeypatch, tmp_path): + backend = _resume_backend(tmp_path) + _fake_disk(monkeypatch) + + def fake_post(url, **kwargs): + (tmp_path / kwargs["json"]["filename"]).write_bytes(b"chat-kv") + return _Resp(200, ["unexpected", "list"]) + + monkeypatch.setattr(llama_cpp.httpx, "post", fake_post, raising = False) + assert backend.save_slots_for_resume() is None + assert list(tmp_path.glob("resume-*.bin")) == [] + + +def test_save_cap_uses_actual_file_size_not_reported_bytes(monkeypatch, tmp_path): + # A binary under-reporting n_written must not slip past the disk cap: the + # cap is enforced against the bytes actually on disk. + backend = _resume_backend(tmp_path) + _fake_disk(monkeypatch) + monkeypatch.setattr(llama_cpp, "_SLOT_SAVE_MAX_BYTES", 150) + + def fake_post(url, **kwargs): + (tmp_path / kwargs["json"]["filename"]).write_bytes(b"x" * 200) + return _Resp(200, {"n_saved": 5, "n_written": 1}) # under-reported + + monkeypatch.setattr(llama_cpp.httpx, "post", fake_post, raising = False) + assert backend.save_slots_for_resume() is None # 200 real bytes > 150 cap + assert list(tmp_path.glob("resume-*.bin")) == [] + + +def test_save_skipped_when_estimate_exceeds_cap(monkeypatch, tmp_path): + # An estimate over the cap skips before writing any slot at all. + backend = _resume_backend(tmp_path) + backend._estimate_kv_cache_bytes = lambda *a, **k: 1 << 40 + monkeypatch.setattr(llama_cpp, "_SLOT_SAVE_MAX_BYTES", 1 << 20) + _fake_disk(monkeypatch) + monkeypatch.setattr( + llama_cpp.httpx, + "post", + lambda *a, **k: (_ for _ in ()).throw(AssertionError), + raising = False, + ) + assert backend.save_slots_for_resume() is None + + +def test_save_skipped_when_model_file_changed_since_load(monkeypatch, tmp_path): + # The GGUF/sidecars were swapped on disk after the server loaded them, so the + # live KV belongs to the old weights: refuse to persist it (no POST at all). + backend = _resume_backend(tmp_path) + backend._slot_loaded_identity = ((("stale", 0),), ()) # != current identity + _fake_disk(monkeypatch) + monkeypatch.setattr( + llama_cpp.httpx, + "post", + lambda *a, **k: (_ for _ in ()).throw(AssertionError), + raising = False, + ) + assert backend.save_slots_for_resume() is None + + +def test_save_proceeds_when_load_identity_matches(monkeypatch, tmp_path): + # Matching load-time snapshot: the save runs normally. + backend = _resume_backend(tmp_path) + backend._slot_loaded_identity = ( + backend._gguf_file_identity(backend._gguf_path), + backend._slot_launch_fingerprint(), + ) + _fake_disk(monkeypatch) + + def fake_post(url, **kwargs): + (tmp_path / kwargs["json"]["filename"]).write_bytes(b"kv") + return _Resp(200, {"n_saved": 5, "n_written": 2}) + + monkeypatch.setattr(llama_cpp.httpx, "post", fake_post, raising = False) + manifest = backend.save_slots_for_resume() + assert manifest is not None + assert [e["id"] for e in manifest["slots"]] == [0] + + +def test_save_skipped_when_estimate_unavailable_and_low_disk(monkeypatch, tmp_path): + # A 0 estimate means metadata was insufficient, not a zero-byte cache: the save + # must demand room for the whole cap, not just 1 GiB, on a low-disk host. + backend = _resume_backend(tmp_path) + backend._estimate_kv_cache_bytes = lambda *a, **k: 0 # metadata unavailable + monkeypatch.setattr(llama_cpp, "_SLOT_SAVE_MAX_BYTES", 8 << 30) # 8 GiB cap + _fake_disk(monkeypatch, free = 2 << 30) # 2 GiB free < 8 + 1 GiB required + monkeypatch.setattr( + llama_cpp.httpx, + "post", + lambda *a, **k: (_ for _ in ()).throw(AssertionError), + raising = False, + ) + assert backend.save_slots_for_resume() is None diff --git a/studio/backend/tests/test_llama_cpp_stall_timeout.py b/studio/backend/tests/test_llama_cpp_stall_timeout.py new file mode 100644 index 0000000000..da36f75e8e --- /dev/null +++ b/studio/backend/tests/test_llama_cpp_stall_timeout.py @@ -0,0 +1,125 @@ +# SPDX-License-Identifier: AGPL-3.0-only +# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0 + +"""Regression test for the post-first-token stall timeout in the cancel-aware read. + +httpcore snapshots ``request.extensions["timeout"]["read"]`` once at body start, so +when ``_iter_text_cancellable`` lowers it after the first token, a one-token-then-silent +server hangs for the full prefill window. The fix re-reads the live extensions timeout +per call; a fake clock and always-silent stream check the read gives up after the live +stall timeout, not the stale prefill one. +""" + +from __future__ import annotations + +import inspect +import sys +import threading +import types as _types +from pathlib import Path + +import pytest + +_BACKEND_DIR = str(Path(__file__).resolve().parent.parent) +if _BACKEND_DIR not in sys.path: + sys.path.insert(0, _BACKEND_DIR) + +# Mirror sibling tests' stubbing so the module imports without fastapi. +_loggers_stub = _types.ModuleType("loggers") +_loggers_stub.get_logger = lambda name: __import__("logging").getLogger(name) +sys.modules.setdefault("loggers", _loggers_stub) +sys.modules.setdefault("structlog", _types.ModuleType("structlog")) + +import httpcore # noqa: E402 + +from core.inference import llama_cpp as llama_cpp_mod # noqa: E402 +from core.inference.llama_cpp import LlamaCppBackend # noqa: E402 + +_PREFILL_TIMEOUT = 1200.0 # what httpcore snapshots from the prefill timeout +_STALL_TIMEOUT = 120.0 # the post-first-token stall timeout the wrapper must honor + + +class _Obj: + pass + + +def _install(response, clock, silent_stream): + """Wire fake client/pool so _install_cancel_aware_read finds the stream; return the wrapped stream.read.""" + inner = _Obj() + inner._network_stream = silent_stream + connection = _Obj() + connection._connection = inner + pool = _Obj() + pool._connections = [connection] + transport = _Obj() + transport._pool = pool + client = _Obj() + client._transport = transport + + cancel_event = threading.Event() # never set: we test the stall path, not cancel + sig = inspect.signature(LlamaCppBackend._install_cancel_aware_read) + if "response" in sig.parameters: + # Fixed signature: wrapper reads the live extensions timeout. + LlamaCppBackend._install_cancel_aware_read(client, cancel_event, response) + else: + # Pre-fix signature: no response, so the stall assertion fails (proves the bug). + LlamaCppBackend._install_cancel_aware_read(client, cancel_event) + return silent_stream.read + + +def test_stall_timeout_honored_after_first_token(monkeypatch): + clock = {"t": 0.0} + monkeypatch.setattr(llama_cpp_mod.time, "monotonic", lambda: clock["t"]) + + # One token then silence: every read times out, advancing fake time by its timeout. + def silent_read(max_bytes, timeout = None): + clock["t"] += timeout if timeout is not None else 0.0 + raise httpcore.ReadTimeout("slice timed out on silence") + + stream = _Obj() + stream.read = silent_read + + # First token seen: the live read timeout is lowered to the stall timeout. + request = _Obj() + request.extensions = {"timeout": {"read": _STALL_TIMEOUT}} + response = _Obj() + response.request = request + + wrapped_read = _install(response, clock, stream) + + # httpcore still passes the stale prefill timeout it snapshotted at body start. + with pytest.raises(httpcore.ReadTimeout): + wrapped_read(65536, timeout = _PREFILL_TIMEOUT) + + # Must give up ~stall timeout after the last token, not the prefill window. + assert clock["t"] <= _STALL_TIMEOUT * 1.5, ( + f"stall timeout not honored: waited {clock['t']}s " + f"(expected ~{_STALL_TIMEOUT}s, not {_PREFILL_TIMEOUT}s)" + ) + assert clock["t"] >= _STALL_TIMEOUT * 0.5 + + +def test_prefill_timeout_used_when_no_live_override(monkeypatch): + """Without a lowered live timeout, the wrapper honors the passed prefill timeout, so the normal first-token wait is unchanged.""" + clock = {"t": 0.0} + monkeypatch.setattr(llama_cpp_mod.time, "monotonic", lambda: clock["t"]) + + def silent_read(max_bytes, timeout = None): + clock["t"] += timeout if timeout is not None else 0.0 + raise httpcore.ReadTimeout("slice timed out on silence") + + stream = _Obj() + stream.read = silent_read + + # No timeout extension: wrapper falls back to httpcore's passed timeout. + request = _Obj() + request.extensions = {} + response = _Obj() + response.request = request + + wrapped_read = _install(response, clock, stream) + + with pytest.raises(httpcore.ReadTimeout): + wrapped_read(65536, timeout = _PREFILL_TIMEOUT) + + assert clock["t"] >= _PREFILL_TIMEOUT * 0.9 diff --git a/studio/backend/tests/test_llama_cpp_update.py b/studio/backend/tests/test_llama_cpp_update.py index 83ea07a066..f12384231f 100644 --- a/studio/backend/tests/test_llama_cpp_update.py +++ b/studio/backend/tests/test_llama_cpp_update.py @@ -83,6 +83,7 @@ def _write_install( repo: str = "unslothai/llama.cpp", asset: str | None = None, release_tag: str | None = None, + force_cpu: bool | None = None, ) -> str: """Create a fake prebuilt install and return the llama-server path.""" bin_dir = dir_ / "build" / "bin" @@ -99,6 +100,8 @@ def _write_install( } if asset is not None: marker["asset"] = asset + if force_cpu is not None: + marker["force_cpu"] = force_cpu (dir_ / MARKER).write_text(json.dumps(marker)) return str(binary) @@ -493,6 +496,47 @@ def test_start_update_preserves_vulkan_via_env(monkeypatch, tmp_path): assert popen_kwargs["env"]["UNSLOTH_FORCE_VULKAN"] == "1" +@pytest.mark.parametrize( + "force_cpu, expect_flag", + [ + # A deliberate CPU install (marker force_cpu=True) re-asserts --force-cpu on + # update so detect_host on a GPU host cannot re-route and revive the crash + # (#7213); --force-cpu also re-persists the flag for the next update. + (True, True), + # A transient fallback (or a legacy marker without the flag) stays free to + # heal to a GPU bundle (#6097). + (False, False), + (None, False), + ], +) +def test_start_update_cpu_fallback_preserved_by_flag(monkeypatch, tmp_path, force_cpu, expect_flag): + asset = "llama-b9493-bin-ubuntu-x64.tar.gz" + install_dir = tmp_path / "llama.cpp" + binary = _write_install(install_dir, "b9493", asset = asset, force_cpu = force_cpu) + monkeypatch.setattr(upd, "_find_binary", lambda: binary) + monkeypatch.setattr(upd, "_installer_script", lambda: tmp_path / "install_llama_prebuilt.py") + monkeypatch.setattr(freshness, "_fetch_latest_release_tag", lambda repo, timeout = 5.0: "b9518") + + captured: dict = {} + + def _on_start(cmd): + captured["cmd"] = cmd + _write_install(install_dir, "b9518", asset = asset, force_cpu = force_cpu) + + _patch_installer_popen(monkeypatch, lines = ["installed\n"], on_start = _on_start) + + assert upd.start_update()["started"] is True + deadline = time.time() + 10 + while time.time() < deadline: + job = upd.get_update_status()["job"] + if job["state"] in ("success", "error"): + break + time.sleep(0.05) + assert job["state"] == "success", job + assert ("--force-cpu" in captured["cmd"]) is expect_flag + assert "--cpu-fallback" not in captured["cmd"] + + def test_start_update_reports_full_release_tag(monkeypatch, tmp_path): install_dir = tmp_path / "llama.cpp" binary = _write_install(install_dir, "b9595") @@ -676,7 +720,7 @@ def test_install_cmd_rocm_marker_forwards_gfx(monkeypatch, tmp_path): assert "--rocm-gfx" in cmd assert cmd[cmd.index("--rocm-gfx") + 1] == "gfx110x" assert "--has-rocm" not in cmd - assert "--cpu-fallback" not in cmd + assert "--force-cpu" not in cmd assert "--simple-policy" not in cmd assert "--published-repo" in cmd and "unslothai/llama.cpp" in cmd @@ -690,17 +734,17 @@ def test_install_cmd_fork_rocm_marker_forwards_has_rocm(monkeypatch, tmp_path): def test_install_cmd_ggml_cpu_marker_has_no_cpu_fallback(monkeypatch, tmp_path): - # Legacy CPU installs recorded a ggml-org marker (new installs use the fork). - # Re-running into the same install-dir/repo reproduces the same CPU bundle; - # --cpu-fallback (which force-drops GPU detection) is reserved for setup.sh's - # arm64 rescue and must not appear here. + # Legacy CPU installs recorded a ggml-org marker (new installs use the fork) with + # no force_cpu field. Re-running into the same install-dir/repo reproduces the same + # CPU bundle; --force-cpu (the persisted-CPU re-assert) must not appear for a marker + # that never recorded a deliberate CPU choice, so it can still heal to GPU (#6097). cmd = _capture_install_cmd( monkeypatch, tmp_path, repo = "ggml-org/llama.cpp", asset = "llama-b9334-bin-ubuntu-x64.tar.gz", ) - assert "--cpu-fallback" not in cmd + assert "--force-cpu" not in cmd assert "--rocm-gfx" not in cmd assert "--has-rocm" not in cmd assert "--simple-policy" not in cmd @@ -714,7 +758,7 @@ def test_install_cmd_cuda_marker_minimal_and_backward_compatible(monkeypatch, tm assert "--simple-policy" not in cmd assert "--rocm-gfx" not in cmd assert "--has-rocm" not in cmd - assert "--cpu-fallback" not in cmd + assert "--force-cpu" not in cmd def test_install_cmd_pins_offered_release_tag(monkeypatch, tmp_path): diff --git a/studio/backend/tests/test_llama_server_args.py b/studio/backend/tests/test_llama_server_args.py index ba52afad1c..fa4ba71791 100644 --- a/studio/backend/tests/test_llama_server_args.py +++ b/studio/backend/tests/test_llama_server_args.py @@ -183,6 +183,8 @@ def test_non_flag_token_passes_through(): "--reranking", # llama-server's own --tools clashes with Unsloth's tool policy. "--tools", + # Slot-state dir: Studio owns it for KV persistence across idle unload. + "--slot-save-path", ], ) def test_denylist_rejects_all_aliases(denied): @@ -224,6 +226,16 @@ def test_denylist_rejects_equals_form(): validate_extra_args(["--port=9000"]) +def test_slot_save_path_is_managed_in_all_forms(): + for args in (["--slot-save-path", "/tmp/x"], ["--slot-save-path=/tmp/x"], ["--slot-save-path"]): + with pytest.raises(ValueError, match = "--slot-save-path"): + validate_extra_args(args) + assert is_managed_flag("--slot-save-path") is True + assert is_managed_flag("--slot-save-path=/tmp/x") is True + # --slots (read-only diagnostics endpoint) stays a user choice. + assert is_managed_flag("--slots") is False + + @pytest.mark.parametrize( "padded", [" --parallel", "--parallel ", "\t--parallel", " -np", "-np \n", "-np\t"], @@ -747,6 +759,34 @@ def test_strip_shadowing_flags_defaults_strip_split_mode_too(): assert strip_shadowing_flags(["--split-mode", "tensor"]) == [] +def test_strip_offload_is_opt_in_and_covers_moe(): + base = dict( + strip_context = False, + strip_cache = False, + strip_spec = False, + strip_template = False, + strip_split_mode = False, + ) + # Default: offload (incl. MoE) flags are NOT stripped. + assert strip_shadowing_flags(["--n-cpu-moe", "8", "--top-k", "20"], **base) == [ + "--n-cpu-moe", + "8", + "--top-k", + "20", + ] + # Opt-in strips layer AND MoE offload flags (value-aware), keeps the rest. + assert strip_shadowing_flags( + ["--n-cpu-moe", "8", "--gpu-layers", "33", "--fit", "off", "--top-k", "20"], + **base, + strip_offload = True, + ) == ["--top-k", "20"] + # Boolean --cpu-moe drops the flag only, not the following value. + assert strip_shadowing_flags(["--cpu-moe", "--seed", "-1"], **base, strip_offload = True) == [ + "--seed", + "-1", + ] + + @pytest.mark.parametrize( "args", [ @@ -796,6 +836,23 @@ def test_strip_split_mode_only_drops_tensor_split_too(): assert strip_split_mode_only(["-sm=tensor", "-ts=3,1"]) == [] +def test_strip_tensor_split_alone_preserves_split_mode(): + # Manual mode emits its own --tensor-split, so an inherited ratio is dropped + # -- but the user's --split-mode row/none/layer choice (which the manual + # ratio toggle can't express) must survive. strip_tensor_split removes only + # the ratio, unlike strip_split_mode which removes the whole group. + out = strip_shadowing_flags( + ["--split-mode", "row", "--tensor-split", "1,1", "--top-k", "20"], + strip_context = False, + strip_cache = False, + strip_spec = False, + strip_template = False, + strip_split_mode = False, + strip_tensor_split = True, + ) + assert out == ["--split-mode", "row", "--top-k", "20"] + + def test_strip_shadowing_flags_keeps_model_draft_without_spec(): out = strip_shadowing_flags( ["--model-draft", "/custom/mtp.gguf"], diff --git a/studio/backend/tests/test_middleware.py b/studio/backend/tests/test_middleware.py index 11aeee6d77..209c6cb90a 100644 --- a/studio/backend/tests/test_middleware.py +++ b/studio/backend/tests/test_middleware.py @@ -14,6 +14,7 @@ import pytest from fastapi import FastAPI, HTTPException, Request from fastapi.responses import Response from fastapi.testclient import TestClient +from starlette.middleware.gzip import GZipMiddleware _BACKEND_ROOT = Path(__file__).resolve().parents[1] @@ -471,6 +472,71 @@ class TestSecurityHeadersMiddleware: assert b"server" in names +class TestFrontendAssets: + def test_hashed_assets_are_compressed_and_cached(self, tmp_path, main_module): + content = b"export const value = 'responsive';\n" * 200 + (tmp_path / "page-abc123.js").write_bytes(content) + app = FastAPI() + assets_app = GZipMiddleware( + main_module.ImmutableStaticFiles(directory = tmp_path), + minimum_size = 1024, + compresslevel = 6, + ) + app.mount("/assets", assets_app, name = "assets") + + response = TestClient(app).get( + "/assets/page-abc123.js", + headers = {"Accept-Encoding": "gzip"}, + ) + + assert response.status_code == 200 + assert response.content == content + assert response.headers["content-encoding"] == "gzip" + assert response.headers["cache-control"] == (main_module._IMMUTABLE_ASSET_CACHE_CONTROL) + assert "accept-encoding" in response.headers["vary"].lower() + + def test_asset_revalidation_keeps_immutable_cache_header(self, tmp_path, main_module): + (tmp_path / "page-abc123.js").write_text("export {};", encoding = "utf-8") + app = FastAPI() + app.mount( + "/assets", + main_module.ImmutableStaticFiles(directory = tmp_path), + name = "assets", + ) + client = TestClient(app) + first = client.get("/assets/page-abc123.js") + + response = client.get( + "/assets/page-abc123.js", + headers = {"If-None-Match": first.headers["etag"]}, + ) + + assert response.status_code == 304 + assert response.headers["cache-control"] == (main_module._IMMUTABLE_ASSET_CACHE_CONTROL) + + def test_range_request_is_not_compressed(self, tmp_path, main_module): + content = b"export const value = 'responsive';\n" * 200 + (tmp_path / "page-abc123.js").write_bytes(content) + app = FastAPI() + assets_app = main_module._AssetGZipMiddleware( + main_module.ImmutableStaticFiles(directory = tmp_path), + minimum_size = 1024, + compresslevel = 6, + ) + app.mount("/assets", assets_app, name = "assets") + + response = TestClient(app).get( + "/assets/page-abc123.js", + headers = {"Accept-Encoding": "gzip", "Range": "bytes=0-99"}, + ) + + assert response.status_code == 206 + assert response.headers.get("content-encoding") != "gzip" + assert response.headers["content-range"] == f"bytes 0-99/{len(content)}" + assert response.content == content[:100] + assert response.headers["cache-control"] == (main_module._IMMUTABLE_ASSET_CACHE_CONTROL) + + # /api/health auth gate diff --git a/studio/backend/tests/test_mlx_stop_checkpoint.py b/studio/backend/tests/test_mlx_stop_checkpoint.py new file mode 100644 index 0000000000..d4a00cc6c8 --- /dev/null +++ b/studio/backend/tests/test_mlx_stop_checkpoint.py @@ -0,0 +1,137 @@ +# SPDX-License-Identifier: AGPL-3.0-only +# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0 + +"""Regression tests for MLX stop-and-save checkpoint handling.""" + +import importlib.util +import json +import sys +import types +from pathlib import Path + +import numpy as np +from safetensors.numpy import save_file + + +_BACKEND = Path(__file__).resolve().parents[1] + + +def _load_worker_module(): + spec = importlib.util.spec_from_file_location( + "training_worker_under_test", + _BACKEND / "core" / "training" / "worker.py", + ) + module = importlib.util.module_from_spec(spec) + assert spec.loader is not None + spec.loader.exec_module(module) + return module + + +worker = _load_worker_module() + + +class _FakeTrainer: + def __init__(self, step: int): + self._global_step = step + self._train_loss_history = [] + self.model = object() + + +def _write_checkpoint(out: Path, step: int) -> Path: + checkpoint = out / f"checkpoint-{step}" + checkpoint.mkdir(parents = True, exist_ok = True) + (checkpoint / "trainer_state.json").write_text( + json.dumps({"global_step": step}), encoding = "utf-8" + ) + save_file({"weight": np.ones(1, dtype = np.float32)}, checkpoint / "adapters.safetensors") + save_file( + {"state": np.ones(1, dtype = np.float32)}, + checkpoint / "optimizer_state.safetensors", + ) + return checkpoint + + +def test_mlx_has_checkpoint_at_step_requires_complete_state(tmp_path): + out = tmp_path / "outputs" / "run_x" + _write_checkpoint(out, 5) + + assert worker._mlx_has_checkpoint_at_step(out, 5) is True + + +def test_write_mlx_stop_checkpoint_returns_true_when_current_step_checkpoint_exists(tmp_path): + out = tmp_path / "outputs" / "run_x" + _write_checkpoint(out, 5) + + assert worker._write_mlx_stop_checkpoint(_FakeTrainer(step = 5), object(), out) is True + + +def test_write_mlx_stop_checkpoint_writes_current_step_when_only_older_checkpoint_exists( + tmp_path, monkeypatch +): + out = tmp_path / "outputs" / "run_x" + _write_checkpoint(out, 5) + + saved_steps: list[int] = [] + + def _save_state(_value, path, name): + save_file({"state": np.ones(1, dtype = np.float32)}, Path(path, name)) + + def _save_trainer_state(state, ckpt_dir, **_kwargs): + Path(ckpt_dir, "trainer_state.json").write_text(json.dumps(state), encoding = "utf-8") + saved_steps.append(int(state["global_step"])) + + fake_utils = types.SimpleNamespace( + save_trainable_adapters = lambda model, path: _save_state( + model, path, "adapters.safetensors" + ), + save_optimizer_state = lambda optimizer, path: _save_state( + optimizer, path, "optimizer_state.safetensors" + ), + save_trainer_state = _save_trainer_state, + ) + monkeypatch.setitem(sys.modules, "unsloth_zoo.mlx.utils", fake_utils) + + assert worker._write_mlx_stop_checkpoint(_FakeTrainer(step = 10), object(), out) is True + assert saved_steps == [10] + assert (out / "checkpoint-10" / "trainer_state.json").is_file() + + +def test_write_mlx_stop_checkpoint_returns_false_without_optimizer(tmp_path): + out = tmp_path / "outputs" / "run_x" + out.mkdir(parents = True) + + assert worker._write_mlx_stop_checkpoint(_FakeTrainer(step = 5), None, out) is False + + +def test_write_mlx_stop_checkpoint_rejects_incomplete_current_checkpoint(tmp_path): + out = tmp_path / "outputs" / "run_x" + ckpt = out / "checkpoint-5" + ckpt.mkdir(parents = True) + (ckpt / "trainer_state.json").write_text('{"global_step": 5}', encoding = "utf-8") + + assert worker._write_mlx_stop_checkpoint(_FakeTrainer(step = 5), None, out) is False + + +def test_write_mlx_stop_checkpoint_ignores_stale_checkpoint_without_optimizer(tmp_path): + # An older checkpoint does not cover the current step, so this still fails. + out = tmp_path / "outputs" / "run_x" + _write_checkpoint(out, 5) + + assert worker._write_mlx_stop_checkpoint(_FakeTrainer(step = 10), None, out) is False + + +def test_write_mlx_stop_checkpoint_returns_false_when_save_fails(tmp_path, monkeypatch): + out = tmp_path / "outputs" / "run_x" + out.mkdir(parents = True) + + def _boom(*_args, **_kwargs): + raise RuntimeError("save failed") + + fake_utils = types.SimpleNamespace( + save_trainable_adapters = _boom, + save_optimizer_state = lambda *_a, **_k: None, + save_trainer_state = lambda *_a, **_k: None, + ) + monkeypatch.setitem(sys.modules, "unsloth_zoo.mlx.utils", fake_utils) + + assert worker._write_mlx_stop_checkpoint(_FakeTrainer(step = 5), object(), out) is False diff --git a/studio/backend/tests/test_model_picker_regression.py b/studio/backend/tests/test_model_picker_regression.py new file mode 100644 index 0000000000..f38a4d0b8d --- /dev/null +++ b/studio/backend/tests/test_model_picker_regression.py @@ -0,0 +1,232 @@ +# SPDX-License-Identifier: AGPL-3.0-only +# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0 + +"""Regression guards for the model-picker per-model-config feature (the set of +bugs that got the predecessor PR reverted). Pure-function / validation checks +only, so they run on CPU in the backend pytest job with no model download. + +Covers, at the backend layer: + - infra-model hiding: the RAG embedder (bge-small-en-v1.5) and the llama.cpp + install-validation probe (ggml-org/models / stories260K) stay hidden, while + normal chat repos are not hidden; + - the HF token is honored from the dedicated header with the query string as a + fallback, never the other way around; + - the chat-template byte caps reject oversized overrides (both the char-count + fast path and the UTF-8 byte path) and the sidecar reader is size-bounded. +""" + +from __future__ import annotations + +import sys +import types + +import pytest + +# Keep this test runnable without the optional structlog dependency (mirrors +# tests/test_cached_gguf_routes.py), since importing routes.models pulls it in. +if "structlog" not in sys.modules: + + class _DummyLogger: + def __getattr__(self, _name): + return lambda *args, **kwargs: None + + sys.modules["structlog"] = types.SimpleNamespace( + BoundLogger = _DummyLogger, + get_logger = lambda *args, **kwargs: _DummyLogger(), + ) + +import routes.models as models_route +from core.rag import config as rag_config +from hub.dependencies import get_hf_token +from models.inference import LoadRequest +from picker.schemas import MAX_CHAT_TEMPLATE_BYTES +from picker.service import _read_bounded_text +from utils.hidden_models import is_hidden_model + + +@pytest.fixture(autouse = True) +def _pin_default_embedder(monkeypatch): + """Pin the effective embedder to Studio's static default so hiding is + deterministic and cannot depend on ambient RAG config / env.""" + default = "unsloth/bge-small-en-v1.5" + monkeypatch.setattr(rag_config, "EMBEDDING_MODEL", default, raising = False) + monkeypatch.setattr(rag_config, "effective_embedding_model", lambda: default) + monkeypatch.setattr(rag_config, "effective_gguf_repo", lambda: default) + monkeypatch.setattr(rag_config, "default_gguf_repo", lambda: default) + + +# --------------------------------------------------------------------------- # +# Infra-model hiding (the "infra models resurfaced in the picker" regression) # +# --------------------------------------------------------------------------- # + + +@pytest.mark.parametrize( + "value", + [ + "ggml-org/models", # the probe repo id + "unsloth/bge-small-en-v1.5", # the RAG embedder repo + "unsloth/bge-small-en-v1.5-GGUF", # its GGUF companion + "/root/.cache/huggingface/hub/x/stories260K.gguf", # probe on disk + "/root/.cache/x/Stories260K.GGUF", # case-insensitive + r"C:\\models\\stories260K.gguf", # windows-style path + "/opt/models/bge-small-en-v1.5", # embedder basename folder + "/opt/models/bge-small-en-v1.5-Q8_0.gguf", # suffixed local weight + ], +) +def test_infra_models_are_hidden(value): + assert is_hidden_model(value) is True + + +@pytest.mark.parametrize( + "value", + [ + "unsloth/gemma-3-270m-it-GGUF", # a normal small chat GGUF + "unsloth/Qwen3-0.6B", # a normal non-GGUF chat model + "user/stories260K-finetune-GGUF", # repo id merely contains "stories260k" + "user/model-chat", # generic repo must not be hidden + "meta-llama/Llama-3.1-8B-Instruct", + ], +) +def test_normal_models_are_not_hidden(value): + assert is_hidden_model(value) is False + + +def test_is_hidden_model_ignores_empty_values(): + assert is_hidden_model(None) is False + assert is_hidden_model("") is False + assert is_hidden_model(None, "", "unsloth/gemma-3-270m-it-GGUF") is False + + +def test_hidden_model_matchers_expose_probe_needles(): + needles, exact_ids, _exact_paths = models_route.hidden_model_matchers() + lowered = [n.lower() for n in needles] + assert "ggml-org/models" in lowered + assert "stories260k.gguf" in lowered + # The configured embedder is exposed as an exact repo id, never as a + # basename needle that would substring-hide unrelated chat models. + assert "bge-small-en-v1.5" not in lowered + assert "unsloth/bge-small-en-v1.5" in exact_ids + + +def test_hidden_model_matchers_custom_repo_publishes_exact_ids(monkeypatch): + monkeypatch.setattr(rag_config, "effective_embedding_model", lambda: "org/model") + monkeypatch.setattr(rag_config, "effective_gguf_repo", lambda: "org/model-GGUF") + needles, exact_ids, exact_paths = models_route.hidden_model_matchers() + assert needles == ["ggml-org/models", "stories260k.gguf"] + assert "org/model" in exact_ids + assert "org/model-gguf" in exact_ids + assert exact_paths == [] + + +def test_hidden_model_matchers_local_owner_name_path_is_exact_path(monkeypatch, tmp_path): + # A local embedder shaped like owner/name that exists on disk must be an + # exact resolved path, not a Hub repo id (mirroring is_hidden_model), so the + # local row stays hidden instead of showing as a chat model. + (tmp_path / "models" / "embedder").mkdir(parents = True) + monkeypatch.chdir(tmp_path) + monkeypatch.setattr(rag_config, "effective_embedding_model", lambda: "models/embedder") + monkeypatch.setattr(rag_config, "effective_gguf_repo", lambda: "ggml-org/models") + _needles, exact_ids, exact_paths = models_route.hidden_model_matchers() + resolved = str((tmp_path / "models" / "embedder").resolve()).lower() + assert resolved in exact_paths + assert "models/embedder" not in exact_ids + + +# --------------------------------------------------------------------------- # +# HF token via header, query string only as a fallback (the token-leak fix) # +# --------------------------------------------------------------------------- # + + +def test_get_hf_token_strips_and_returns(): + assert get_hf_token(" hf_abc ") == "hf_abc" + + +@pytest.mark.parametrize("value", [None, "", " ", "\n\t"]) +def test_get_hf_token_blank_is_none(value): + assert get_hf_token(value) is None + + +@pytest.mark.parametrize( + "value,expected", + [(" hf_x ", "hf_x"), ("", None), (" ", None), (None, None), (1234, None)], +) +def test_normalize_hf_token(value, expected): + assert models_route._normalize_hf_token(value) == expected + + +def test_header_token_wins_over_query(): + header, query = "hf_header", "hf_query" + resolved = models_route._normalize_hf_token(header) or models_route._normalize_hf_token(query) + assert resolved == "hf_header" + + +def test_query_token_is_fallback_when_header_absent(): + resolved = models_route._normalize_hf_token(None) or models_route._normalize_hf_token( + "hf_query" + ) + assert resolved == "hf_query" + + +# --------------------------------------------------------------------------- # +# Chat-template byte caps (the unbounded-template hardening) # +# --------------------------------------------------------------------------- # + + +def _load_request(**overrides): + data = {"model_path": "unsloth/test-model-GGUF", "gguf_variant": "Q4_K_M"} + data.update(overrides) + return LoadRequest.model_validate(data) + + +def test_blank_chat_template_override_normalizes_to_none(): + assert _load_request(chat_template_override = " \n\t").chat_template_override is None + + +def test_nonblank_chat_template_override_preserved_verbatim(): + template = " {{ messages }} " + assert _load_request(chat_template_override = template).chat_template_override == template + + +def test_chat_template_at_byte_limit_is_accepted(): + template = "a" * MAX_CHAT_TEMPLATE_BYTES # exactly the limit, 1 byte/char + assert ( + len(_load_request(chat_template_override = template).chat_template_override) + == MAX_CHAT_TEMPLATE_BYTES + ) + + +def test_chat_template_over_char_limit_is_rejected(): + with pytest.raises(Exception): # pydantic ValidationError wrapping ValueError + _load_request(chat_template_override = "a" * (MAX_CHAT_TEMPLATE_BYTES + 1)) + + +def test_chat_template_over_byte_limit_is_rejected(): + # Char count stays under the limit but UTF-8 bytes exceed it (3 bytes/char), + # so only the byte-count branch can catch this. + multibyte = "€" * (MAX_CHAT_TEMPLATE_BYTES // 2) # euro sign, 3 bytes each + assert len(multibyte) <= MAX_CHAT_TEMPLATE_BYTES + assert len(multibyte.encode("utf-8")) > MAX_CHAT_TEMPLATE_BYTES + with pytest.raises(Exception): + _load_request(chat_template_override = multibyte) + + +def test_read_bounded_text_reads_within_limit(tmp_path): + p = tmp_path / "t.json" + p.write_text("hello", encoding = "utf-8") + assert _read_bounded_text(p, 16) == "hello" + + +def test_read_bounded_text_rejects_over_limit(tmp_path): + p = tmp_path / "big.json" + p.write_bytes(b"x" * 100) + assert _read_bounded_text(p, 50) is None + + +def test_read_bounded_text_at_limit_is_read(tmp_path): + p = tmp_path / "exact.json" + p.write_bytes(b"x" * 50) + assert _read_bounded_text(p, 50) == "x" * 50 + + +def test_read_bounded_text_missing_file_is_none(tmp_path): + assert _read_bounded_text(tmp_path / "nope.json", 50) is None diff --git a/studio/backend/tests/test_model_update_robustness.py b/studio/backend/tests/test_model_update_robustness.py index edf55812e2..7d98766616 100644 --- a/studio/backend/tests/test_model_update_robustness.py +++ b/studio/backend/tests/test_model_update_robustness.py @@ -314,6 +314,7 @@ def test_cached_model_scan_keeps_local_safetensors_repo(monkeypatch, tmp_path): file_name = "model.safetensors", size_on_disk = 100, blob_path = str(repo_path / "blobs" / "modelsha"), + blob_last_modified = 3_000.0, ), ] ) @@ -336,6 +337,51 @@ def test_cached_model_scan_keeps_local_safetensors_repo(monkeypatch, tmp_path): assert rows[0]["repo_id"] == "Org/SafeTensorRepo" assert rows[0]["model_format"] == "safetensors" assert rows[0]["size_bytes"] == 100 + assert rows[0]["last_modified"] == 3_000.0 + + +def test_cached_gguf_scan_keeps_download_timestamp(monkeypatch, tmp_path): + repo_path = tmp_path / "models--Org--GgufRepo" + repo = SimpleNamespace( + repo_id = "Org/GgufRepo", + repo_type = "model", + repo_path = repo_path, + revisions = [ + SimpleNamespace( + files = [ + SimpleNamespace( + file_name = "model-Q4_K_M.gguf", + size_on_disk = 100, + blob_path = None, + blob_last_modified = 5_000.0, + ), + ] + ) + ], + ) + monkeypatch.setattr( + CI, + "all_hf_cache_scans", + lambda: [SimpleNamespace(repos = [repo])], + ) + monkeypatch.setattr( + CI.hf_cache_scan, + "is_gguf_repo_partial", + lambda *args, **kwargs: False, + ) + monkeypatch.setattr( + CI, + "_gguf_variant_state_summary", + lambda _repo_id: (False, 0), + ) + + rows = CI._scan_cached_gguf() + + assert len(rows) == 1 + assert rows[0]["repo_id"] == "Org/GgufRepo" + assert rows[0]["model_format"] == "gguf" + assert rows[0]["size_bytes"] == 100 + assert rows[0]["last_modified"] == 5_000.0 # ── hf_hub_download_with_xet_fallback force_download bypass (X2/F2) ─── @@ -636,3 +682,18 @@ def test_reclaim_replaced_gguf_variant_keeps_no_symlink_current_file(monkeypatch assert snap.exists() is True # the current file must survive assert result["removed_snapshots"] == 0 assert result["deleted_blobs"] == 0 + + +def _mmproj_repo(*file_names: str): + return SimpleNamespace( + revisions = [SimpleNamespace(files = [SimpleNamespace(file_name = n) for n in file_names])] + ) + + +def test_repo_has_mmproj_requires_gguf_projector(): + # A non-GGUF sidecar whose name merely contains "mmproj" must NOT mark the + # repo vision-capable; the runtime's projector detection is GGUF-only. + assert CI._repo_has_mmproj(_mmproj_repo("model-Q4_K_M.gguf", "mmproj_config.json")) is False + assert CI._repo_has_mmproj(_mmproj_repo("model-Q4_K_M.gguf", "README-mmproj.md")) is False + # A real GGUF projector still marks the repo vision-capable. + assert CI._repo_has_mmproj(_mmproj_repo("model-Q4_K_M.gguf", "mmproj-F16.gguf")) is True diff --git a/studio/backend/tests/test_offline_gguf_cache_fallback.py b/studio/backend/tests/test_offline_gguf_cache_fallback.py index 295549c443..35dd3979b1 100644 --- a/studio/backend/tests/test_offline_gguf_cache_fallback.py +++ b/studio/backend/tests/test_offline_gguf_cache_fallback.py @@ -119,6 +119,13 @@ def _build_cache( return snap +def _symlink_or_skip(link: Path, target: Path) -> None: + try: + link.symlink_to(target) + except OSError as exc: + pytest.skip(f"symlinks unavailable: {exc}") + + @pytest.fixture def hf_cache(tmp_path, monkeypatch): """Point ``huggingface_hub.constants.HF_HUB_CACHE`` at a temp dir.""" @@ -1084,7 +1091,7 @@ class TestListLocalGgufVariantsSubdir: target.write_bytes(b"\0" * 20) out = _find_local_gguf_by_variant(str(tmp_path), "Q4_K_M") - assert out == str(target.resolve()) + assert out == str(target.absolute()) def test_find_local_gguf_by_variant_skips_big_endian_only_match(self, tmp_path): from utils.models.model_config import _find_local_gguf_by_variant @@ -1094,6 +1101,57 @@ class TestListLocalGgufVariantsSubdir: assert _find_local_gguf_by_variant(str(tmp_path), "Q4_K_M") is None + def test_find_local_gguf_by_variant_keeps_split_symlink_name(self, tmp_path): + from utils.models.model_config import _find_local_gguf_by_variant + + blobs = tmp_path / "blobs" + blobs.mkdir() + snap = tmp_path / "snapshots" / "rev" / "BF16" + snap.mkdir(parents = True) + (tmp_path / "snapshots" / "rev" / "config.json").write_text("{}") + for i, sha in enumerate(("aa" * 32, "bb" * 32), start = 1): + (blobs / sha).write_bytes(b"\0" * 10) + _symlink_or_skip(snap / f"model-BF16-0000{i}-of-00002.gguf", blobs / sha) + + out = _find_local_gguf_by_variant(str(tmp_path / "snapshots" / "rev"), "BF16") + assert out is not None + assert Path(out).name == "model-BF16-00001-of-00002.gguf" + + def test_detect_gguf_model_keeps_split_symlink_name(self, tmp_path): + from utils.models.model_config import detect_gguf_model + + blobs = tmp_path / "blobs" + blobs.mkdir() + snap = tmp_path / "snapshots" / "rev" + snap.mkdir(parents = True) + for i, (sha, size) in enumerate((("cc" * 32, 10), ("dd" * 32, 20)), start = 1): + (blobs / sha).write_bytes(b"\0" * size) + _symlink_or_skip(snap / f"model-BF16-0000{i}-of-00002.gguf", blobs / sha) + + out = detect_gguf_model(str(snap)) + assert out is not None + assert Path(out).name == "model-BF16-00001-of-00002.gguf" + + def test_lone_split_symlink_uses_colocated_target_shards(self, tmp_path): + from utils.models.model_config import _find_local_gguf_by_variant, detect_gguf_model + + target_dir = tmp_path / "external" / "BF16" + target_dir.mkdir(parents = True) + target = target_dir / "model-BF16-00001-of-00002.gguf" + target.write_bytes(b"\0" * 10) + (target_dir / "model-BF16-00002-of-00002.gguf").write_bytes(b"\0" * 10) + + local = tmp_path / "local" + local.mkdir() + (local / "config.json").write_text("{}") + link = local / target.name + _symlink_or_skip(link, target) + + expected = str(target.absolute()) + assert _find_local_gguf_by_variant(str(local), "BF16") == expected + assert detect_gguf_model(str(local)) == expected + assert detect_gguf_model(str(link)) == expected + def test_model_config_variant_ignores_big_endian_sibling(self, tmp_path): from utils.models.model_config import ModelConfig diff --git a/studio/backend/tests/test_openai_auto_switch.py b/studio/backend/tests/test_openai_auto_switch.py index c4c0ce15c9..1ee9ef36d3 100644 --- a/studio/backend/tests/test_openai_auto_switch.py +++ b/studio/backend/tests/test_openai_auto_switch.py @@ -8,6 +8,7 @@ tests/test_gguf_completion_usage.py. """ import asyncio +import os import pytest @@ -18,6 +19,10 @@ from utils import openai_auto_switch_settings as settings class _FakeBackend: + effective_parallel_slots = 1 + _slot_save_binary = None + _gguf_path = None + def __init__( self, loaded_id = None, @@ -29,6 +34,22 @@ class _FakeBackend: self.hf_variant = hf_variant self._openai_advertised_id = advertised_id + def save_slots_for_resume(self, should_abort = None): + return None + + def restore_slots_for_resume(self, manifest): + return None + + def _slot_launch_fingerprint(self): + return ((), None, None, 1) + + def _gguf_file_identity(self, path): + try: + st = os.stat(path) + except OSError: + return None + return ((st.st_size, st.st_mtime_ns),) + class _LoadRecorder: """Stand-in for the load route: records calls and simulates a load.""" @@ -53,10 +74,15 @@ class _LoadRecorder: from fastapi import HTTPException raise HTTPException(status_code = 503, detail = "load failed") self.backend.model_identifier = request.model_path + self.backend.hf_variant = getattr(request, "gguf_variant", None) + self.backend._gguf_path = request.model_path self.backend.is_loaded = True # Mirror _load_model_impl: a load advertises its own id until the # auto-switch caller overwrites it with the repo id. self.backend._openai_advertised_id = None + from core.inference import llama_keepwarm as kw + + kw.note_model_loaded(self.backend) return None @@ -446,6 +472,75 @@ def test_idle_loop_unloads_after_ttl_and_stashes_for_reload(monkeypatch): assert stash is not None and stash[0] == "unsloth/Idle-GGUF" and stash[1] == "Q4_K_M" +def test_idle_loop_deletes_saved_kv_when_unload_fails(monkeypatch, tmp_path): + import time + from core.inference import llama_keepwarm as kw + + monkeypatch.setattr(settings, "get_auto_unload_idle_seconds", lambda: 0.005) + monkeypatch.setattr(settings, "get_auto_unload_keep_kv", lambda: True) + kw._inflight = 0 + kw._pending = 0 + kw._last_active = time.monotonic() - 3600 + kw._last_unloaded_model = None + kw._kv_resume = None + + saved = tmp_path / "resume-abc-slot0.bin" + backend = _FakeBackend("unsloth/Idle-GGUF") + manifests = [] + + def _save(should_abort = None): + if manifests: + return None + saved.write_bytes(b"kv") + manifest = {"dir": str(tmp_path), "slots": [{"id": 0, "filename": saved.name}]} + manifests.append(manifest) + return manifest + + def _unload(): + raise RuntimeError("cuda teardown failed") + + backend.save_slots_for_resume = _save + backend.unload_model = _unload + monkeypatch.setattr(inference_route, "get_llama_cpp_backend", lambda: backend) + + async def _drive(): + task = asyncio.create_task(kw.idle_unload_loop(poll_seconds = 0.01)) + for _ in range(200): + await asyncio.sleep(0.01) + if manifests and not saved.exists(): + break + task.cancel() + try: + await task + except asyncio.CancelledError: + pass + + asyncio.run(_drive()) + assert manifests and not saved.exists() + assert kw._kv_resume is None + + +def test_disabling_idle_unload_purges_saved_kv(monkeypatch, tmp_path): + # PUT leaves keep-KV on but makes idle unload inactive: saved KV must go too. + import routes.settings as settings_route + from core.inference import llama_keepwarm as kw + + saved = tmp_path / "resume-abc-slot0.bin" + saved.write_bytes(b"kv") + kw._kv_resume = { + "identity": ("m", None, "m"), + "dir": str(tmp_path), + "slots": [{"id": 0, "filename": saved.name}], + } + monkeypatch.setattr(settings_route, "set_openai_auto_switch", lambda *a: (False, 300, True)) + monkeypatch.setattr(settings_route, "get_auto_unload_idle_seconds", lambda: 0) + + payload = settings_route.OpenAIAutoSwitchPayload(enabled = False) + resp = settings_route.update_openai_auto_switch(payload, "tester") + assert resp.idle_unload_active is False and resp.auto_unload_keep_kv is True + assert kw._kv_resume is None and not saved.exists() + + def test_audio_generate_is_tracked_as_inference_path(): # Direct GGUF TTS uses the llama backend and can outlive the idle TTL, so # the keep-warm middleware must count it as in-flight inference. @@ -2912,8 +3007,10 @@ def test_non_gguf_load_clears_reload_stash(): # A non-GGUF (Transformers/Unsloth) load must clear the stash like the GGUF # branch, so it never lingers until the idle poll (or forever, idle-unload off). import inspect + src = inspect.getsource(inference_route._load_model_impl) - assert src.count("note_model_loaded()") >= 2 + assert src.count("note_model_loaded()") >= 1 # non-GGUF branch + assert "to_thread(note_model_loaded, llama_backend)" in src # GGUF branch def test_chat_rejects_malformed_tool_choice_before_switch(monkeypatch): @@ -3121,6 +3218,495 @@ def test_responses_stream_hint_matches_toggle_regardless_of_active_model(monkeyp assert "Model auto-switch" in non_gguf_loaded +# ── idle-unload KV persistence (slot save/restore) ────────────────── + + +def _seed_kv_manifest( + tmp_path, + identity = ("unsloth/A-GGUF", "Q4_K_M", "unsloth/A-GGUF"), + gguf = None, +): + if gguf is None: + gguf_file = tmp_path / "model.gguf" + gguf_file.write_bytes(b"gguf") + gguf = str(gguf_file) + st = os.stat(gguf) + state_file = tmp_path / "resume-abc-slot0.bin" + state_file.write_bytes(b"kv") + return state_file, { + "identity": identity, + "dir": str(tmp_path), + "binary": ("/bin/llama-server", 111), + "gguf": gguf, + "gguf_stat": ((st.st_size, st.st_mtime_ns),), + "launch": ((), None, None, 1), + "slots": [{"id": 0, "filename": state_file.name, "n_saved": 42}], + } + + +def _drive_idle_loop( + kw, + poll_seconds = 0.02, + run_for = 0.2, +): + async def _drive(): + task = asyncio.create_task(kw.idle_unload_loop(poll_seconds = poll_seconds)) + await asyncio.sleep(run_for) + task.cancel() + try: + await task + except asyncio.CancelledError: + pass + + asyncio.run(_drive()) + + +def test_idle_unload_saves_slots_before_unload_and_stashes_manifest(monkeypatch, tmp_path): + import time + from core.inference import llama_keepwarm as kw + + monkeypatch.setattr(settings, "get_auto_unload_idle_seconds", lambda: 0.005) + monkeypatch.setattr(settings, "get_auto_unload_keep_kv", lambda: True) + kw._inflight = 0 + kw._pending = 0 + kw._last_active = time.monotonic() - 3600 + kw._last_unloaded_model = None + kw._kv_resume = None + + events = [] + backend = _FakeBackend("unsloth/Idle-GGUF", hf_variant = "Q4_K_M") + manifest = { + "dir": str(tmp_path), + "binary": ("bin", 1), + "slots": [{"id": 0, "filename": "f.bin", "n_saved": 42}], + } + + def _save(should_abort = None): + events.append("save") + return manifest + + def _unload(): + events.append("unload") + backend.is_loaded = False + + backend.save_slots_for_resume = _save + backend.unload_model = _unload + monkeypatch.setattr(inference_route, "get_llama_cpp_backend", lambda: backend) + + _drive_idle_loop(kw) + # KV must be saved while the server is still alive, then exactly one unload. + assert events == ["save", "unload"] + assert kw.get_last_unloaded_model()[:2] == ("unsloth/Idle-GGUF", "Q4_K_M") + resume = kw.take_kv_resume() + assert resume is not None + assert resume["identity"][:2] == ("unsloth/Idle-GGUF", "Q4_K_M") + assert resume["slots"][0]["filename"] == "f.bin" + + +def test_idle_save_failure_still_unloads_plain(monkeypatch): + import time + from core.inference import llama_keepwarm as kw + + monkeypatch.setattr(settings, "get_auto_unload_idle_seconds", lambda: 0.005) + monkeypatch.setattr(settings, "get_auto_unload_keep_kv", lambda: True) + kw._inflight = 0 + kw._pending = 0 + kw._last_active = time.monotonic() - 3600 + kw._last_unloaded_model = None + kw._kv_resume = None + + unloads = [] + backend = _FakeBackend("unsloth/Idle-GGUF", hf_variant = "Q4_K_M") + + def _save(should_abort = None): + raise RuntimeError("slot save exploded") + + def _unload(): + unloads.append(1) + backend.is_loaded = False + + backend.save_slots_for_resume = _save + backend.unload_model = _unload + monkeypatch.setattr(inference_route, "get_llama_cpp_backend", lambda: backend) + + _drive_idle_loop(kw) + assert unloads == [1] # the save failure must not skip the unload + assert kw.get_last_unloaded_model() is not None + assert kw.take_kv_resume() is None + + +def test_keep_kv_setting_off_skips_save(monkeypatch): + import time + from core.inference import llama_keepwarm as kw + + monkeypatch.setattr(settings, "get_auto_unload_idle_seconds", lambda: 0.005) + monkeypatch.setattr(settings, "get_auto_unload_keep_kv", lambda: False) + kw._inflight = 0 + kw._pending = 0 + kw._last_active = time.monotonic() - 3600 + kw._last_unloaded_model = None + kw._kv_resume = None + + saves, unloads = [], [] + backend = _FakeBackend("unsloth/Idle-GGUF") + + def _unload(): + unloads.append(1) + backend.is_loaded = False + + backend.save_slots_for_resume = lambda *a, **k: saves.append(1) + backend.unload_model = _unload + monkeypatch.setattr(inference_route, "get_llama_cpp_backend", lambda: backend) + + _drive_idle_loop(kw) + assert saves == [] + assert unloads == [1] + assert kw.take_kv_resume() is None + + +def test_keep_kv_disabled_mid_save_discards_manifest(monkeypatch, tmp_path): + import time + from core.inference import llama_keepwarm as kw + + keep = {"on": True} + monkeypatch.setattr(settings, "get_auto_unload_idle_seconds", lambda: 0.005) + monkeypatch.setattr(settings, "get_auto_unload_keep_kv", lambda: keep["on"]) + kw._inflight = 0 + kw._pending = 0 + kw._last_active = time.monotonic() - 3600 + kw._last_unloaded_model = None + kw._kv_resume = None + + unloads = [] + backend = _FakeBackend("unsloth/Idle-GGUF", hf_variant = "Q4_K_M") + state_file = tmp_path / "resume-mid-slot0.bin" + state_file.write_bytes(b"kv") + manifest = { + "dir": str(tmp_path), + "binary": ("bin", 1), + "slots": [{"id": 0, "filename": state_file.name, "n_saved": 1}], + } + + def _save(should_abort = None): + keep["on"] = False # user flips the toggle while the save runs + return manifest + + def _unload(): + unloads.append(1) + backend.is_loaded = False + + backend.save_slots_for_resume = _save + backend.unload_model = _unload + monkeypatch.setattr(inference_route, "get_llama_cpp_backend", lambda: backend) + + _drive_idle_loop(kw) + assert unloads == [1] # still unloads; only the stash is dropped + assert kw.take_kv_resume() is None + assert not state_file.exists() + + +def test_idle_ttl_disabled_mid_save_skips_unload(monkeypatch, tmp_path): + import time + from core.inference import llama_keepwarm as kw + + ttl = {"v": 0.005} + monkeypatch.setattr(settings, "get_auto_unload_idle_seconds", lambda: ttl["v"]) + monkeypatch.setattr(settings, "get_auto_unload_keep_kv", lambda: True) + kw._inflight = 0 + kw._pending = 0 + kw._last_active = time.monotonic() - 3600 + kw._last_unloaded_model = None + kw._kv_resume = None + + unloads = [] + backend = _FakeBackend("unsloth/Idle-GGUF", hf_variant = "Q4_K_M") + state_file = tmp_path / "resume-mid-slot0.bin" + state_file.write_bytes(b"kv") + manifest = { + "dir": str(tmp_path), + "binary": ("bin", 1), + "slots": [{"id": 0, "filename": state_file.name, "n_saved": 1}], + } + + def _save(should_abort = None): + ttl["v"] = 0 # user turns idle unload off while the save runs + return manifest + + backend.save_slots_for_resume = _save + backend.unload_model = lambda: unloads.append(1) + monkeypatch.setattr(inference_route, "get_llama_cpp_backend", lambda: backend) + + _drive_idle_loop(kw) + assert unloads == [] # the unload was cancelled by the setting change + assert kw.take_kv_resume() is None + assert not state_file.exists() + + +def test_alias_reload_restores_slots_and_deletes_files(monkeypatch, tmp_path): + from core.inference import llama_keepwarm as kw + + backend = _FakeBackend(None) # idle-unload emptied the backend + backend._slot_save_binary = ("/bin/llama-server", 111) + restored = [] + backend.restore_slots_for_resume = lambda manifest: restored.append(manifest) + + rec = _LoadRecorder(backend) + _wire(monkeypatch, enabled = True, resolves_to = None, backend = backend, recorder = rec) + monkeypatch.setattr(kw, "_inflight", 0) + state_file, manifest = _seed_kv_manifest(tmp_path) + monkeypatch.setattr(kw, "_last_unloaded_model", (manifest["gguf"], "Q4_K_M")) + monkeypatch.setattr(kw, "_kv_resume", manifest) + + _run_hook("gpt-4o-mini") + assert len(rec.calls) == 1 + assert len(restored) == 1 # same model + binary: restore ran + assert not state_file.exists() # state file deleted after the restore + assert kw._kv_resume is None + + +def test_no_restore_when_different_model_loads(monkeypatch, tmp_path): + from core.inference import llama_keepwarm as kw + + backend = _FakeBackend(None) + backend._slot_save_binary = ("/bin/llama-server", 111) + restored = [] + backend.restore_slots_for_resume = lambda manifest: restored.append(manifest) + rec = _LoadRecorder(backend) + _wire( + monkeypatch, + enabled = True, + resolves_to = ("unsloth/B-GGUF", None, "unsloth/B-GGUF"), + backend = backend, + recorder = rec, + ) + monkeypatch.setattr(kw, "_inflight", 0) + state_file, manifest = _seed_kv_manifest(tmp_path) # manifest is for model A + monkeypatch.setattr(kw, "_kv_resume", manifest) + + _run_hook("unsloth/B-GGUF") + assert len(rec.calls) == 1 + assert restored == [] # different model: never restored + assert not state_file.exists() # but the stale files are gone + assert kw._kv_resume is None + + +def test_restore_skipped_when_binary_changed(monkeypatch, tmp_path): + from core.inference import llama_keepwarm as kw + + state_file, manifest = _seed_kv_manifest(tmp_path) + backend = _FakeBackend("unsloth/A-GGUF", hf_variant = "Q4_K_M") + backend._gguf_path = manifest["gguf"] + backend._slot_save_binary = ("/bin/llama-server", 222) # newer mtime + restored = [] + backend.restore_slots_for_resume = lambda manifest: restored.append(manifest) + + kw.restore_kv_resume(backend, manifest) + assert restored == [] + assert not state_file.exists() + + +def test_restore_skipped_when_launch_config_changed(tmp_path): + from core.inference import llama_keepwarm as kw + + state_file, manifest = _seed_kv_manifest(tmp_path) + backend = _FakeBackend("unsloth/A-GGUF", hf_variant = "Q4_K_M") + backend._gguf_path = manifest["gguf"] + backend._slot_save_binary = ("/bin/llama-server", 111) + backend._slot_launch_fingerprint = lambda: (("--rope-freq-scale", "0.5"), None, None, 1) + restored = [] + backend.restore_slots_for_resume = lambda manifest: restored.append(manifest) + + kw.restore_kv_resume(backend, manifest) + assert restored == [] + assert not state_file.exists() + + +def test_restore_skipped_when_gguf_rewritten_in_place(tmp_path): + from core.inference import llama_keepwarm as kw + + state_file, manifest = _seed_kv_manifest(tmp_path) + with open(manifest["gguf"], "wb") as fh: + fh.write(b"different weights") # same path, new content + backend = _FakeBackend("unsloth/A-GGUF", hf_variant = "Q4_K_M") + backend._gguf_path = manifest["gguf"] + backend._slot_save_binary = ("/bin/llama-server", 111) + restored = [] + backend.restore_slots_for_resume = lambda manifest: restored.append(manifest) + + kw.restore_kv_resume(backend, manifest) + assert restored == [] + assert not state_file.exists() + + +def test_note_model_unloaded_purges_manifest_and_files(tmp_path): + from core.inference import llama_keepwarm as kw + + state_file, manifest = _seed_kv_manifest(tmp_path) + kw._set_last_unloaded(("org/A-GGUF", "Q4_K_M")) + kw._set_kv_resume(manifest) + kw.note_model_unloaded() + assert kw.get_last_unloaded_model() is None + assert kw.take_kv_resume() is None + assert not state_file.exists() + + +def test_note_model_loaded_purges_manifest_and_files(tmp_path): + from core.inference import llama_keepwarm as kw + + state_file, manifest = _seed_kv_manifest(tmp_path) + kw._set_last_unloaded(("org/A-GGUF", "Q4_K_M")) + kw._set_kv_resume(manifest) + kw.note_model_loaded() + assert kw.get_last_unloaded_model() is None + assert kw.take_kv_resume() is None + assert not state_file.exists() + + +def test_new_idle_save_purges_previous_manifest_files(tmp_path): + from core.inference import llama_keepwarm as kw + + old_file, old_manifest = _seed_kv_manifest(tmp_path) + kw._set_kv_resume(old_manifest) + new_file = tmp_path / "resume-def-slot0.bin" + new_file.write_bytes(b"kv2") + kw._set_kv_resume( + { + "identity": ("unsloth/B-GGUF", None, "unsloth/B-GGUF"), + "dir": str(tmp_path), + "binary": ("/bin/llama-server", 111), + "slots": [{"id": 0, "filename": new_file.name, "n_saved": 7}], + } + ) + assert not old_file.exists() # replaced manifest's files purged + assert new_file.exists() + assert kw.take_kv_resume()["slots"][0]["filename"] == new_file.name + + +def test_sweep_slot_save_dir_removes_only_resume_files(monkeypatch, tmp_path): + from core.inference import llama_keepwarm as kw + from utils.paths import storage_roots + + monkeypatch.setattr(storage_roots, "llama_slot_cache_root", lambda: tmp_path) + stale = tmp_path / "resume-old-slot0.bin" + stale.write_bytes(b"kv") + other = tmp_path / "unrelated.txt" + other.write_text("keep") + kw.sweep_slot_save_dir() + assert not stale.exists() + assert other.exists() + + +def test_keep_kv_setting_roundtrip_and_default(monkeypatch): + import storage.studio_db as db + + store = {} + monkeypatch.setattr(db, "upsert_app_settings", lambda m: store.update(m)) + monkeypatch.setattr(settings, "_cached_setting", lambda k, d = None: store.get(k, d)) + + assert settings.get_auto_unload_keep_kv() is True # default when never stored + assert settings.set_openai_auto_switch(True, 60, False)[2] is False + assert store[settings.AUTO_UNLOAD_KEEP_KV_SETTING_KEY] is False + assert settings.get_auto_unload_keep_kv() is False + # None leaves the stored value untouched (older clients can't reset it). + assert settings.set_openai_auto_switch(True, 60, None)[2] is False + assert store[settings.AUTO_UNLOAD_KEEP_KV_SETTING_KEY] is False + with pytest.raises(ValueError, match = "true or false"): + settings.set_openai_auto_switch(True, 60, "garbage") + + +def test_stale_stash_cleanup_waits_for_lifecycle_gate(monkeypatch, tmp_path): + # The loop's stale-stash purge must wait on the gate a mid-reload holds. + import time + from core.inference import llama_keepwarm as kw + + monkeypatch.setattr(settings, "get_auto_unload_idle_seconds", lambda: 3600) + kw._inflight = 0 + kw._pending = 0 + kw._last_active = time.monotonic() + backend = _FakeBackend("unsloth/New-GGUF") + monkeypatch.setattr(inference_route, "get_llama_cpp_backend", lambda: backend) + state_file, manifest = _seed_kv_manifest(tmp_path) + kw._kv_resume = manifest + kw._last_unloaded_model = ("unsloth/A-GGUF", "Q4_K_M") + + assert kw._lifecycle_lock.acquire(blocking = False) # simulate in-flight reload + try: + _drive_idle_loop(kw) + assert kw._kv_resume is manifest # purge deferred while the gate is held + assert state_file.exists() + finally: + kw._lifecycle_lock.release() + _drive_idle_loop(kw) + assert kw._kv_resume is None # gate freed: genuinely stale stash purged + assert not state_file.exists() + + +def test_put_route_disabling_keep_kv_purges_saved_state(monkeypatch, tmp_path): + import routes.settings as settings_route + import storage.studio_db as db + from core.inference import llama_keepwarm as kw + + store = {} + monkeypatch.setattr(db, "upsert_app_settings", lambda m: store.update(m)) + monkeypatch.setattr(settings, "_cached_setting", lambda k, d = None: store.get(k, d)) + state_file, manifest = _seed_kv_manifest(tmp_path) + monkeypatch.setattr(kw, "_kv_resume", manifest) + + payload = settings_route.OpenAIAutoSwitchPayload(enabled = True, auto_unload_keep_kv = False) + resp = settings_route.update_openai_auto_switch(payload, "tester") + assert resp.auto_unload_keep_kv is False + assert kw._kv_resume is None + assert not state_file.exists() + + +def test_keep_kv_only_update_leaves_env_idle_ttl_active(monkeypatch): + # A keep-KV-only update must not materialize the env TTL as a stored value. + import routes.settings as settings_route + import storage.studio_db as db + + store = {} + monkeypatch.setattr(db, "upsert_app_settings", lambda m: store.update(m)) + monkeypatch.setattr(settings, "_cached_setting", lambda k, d = None: store.get(k, d)) + monkeypatch.setenv(settings.MODEL_IDLE_TTL_ENV_VAR, "600") + + assert settings_route.OpenAIAutoSwitchPayload(enabled = False).auto_unload_idle_seconds is None + enabled, idle, keep_kv = settings.set_openai_auto_switch(False, None, False) + assert settings.AUTO_UNLOAD_IDLE_SETTING_KEY not in store # idle untouched + assert settings.get_auto_unload_idle_seconds() == 600 # env TTL still active + assert (enabled, idle, keep_kv) == (False, 600, False) + + +def test_load_impl_notes_loaded_with_backend_off_loop(): + import inspect + src = inspect.getsource(inference_route._load_model_impl) + assert "to_thread(note_model_loaded, llama_backend)" in src + + +def test_restore_matches_gguf_realpath_across_naming(tmp_path): + from core.inference import llama_keepwarm as kw + + blob = tmp_path / "blob.gguf" + blob.write_bytes(b"gguf") + link = tmp_path / "snapshot.gguf" + try: + link.symlink_to(blob) + except OSError: + pytest.skip("symlinks unsupported on this host") + + backend = _FakeBackend("/hf/snapshots/d7f5", hf_variant = None) + backend._gguf_path = str(link) # reload resolved the symlink spelling + backend._slot_save_binary = ("/bin/llama-server", 111) + restored = [] + backend.restore_slots_for_resume = lambda manifest: restored.append(manifest) + state_file, manifest = _seed_kv_manifest( + tmp_path, identity = ("unsloth/A-GGUF", None, "unsloth/A-GGUF"), gguf = str(blob) + ) + + kw.restore_kv_resume(backend, manifest) + assert len(restored) == 1 # names differ, file identical: restore ran + assert not state_file.exists() + + def test_setter_rejects_idle_below_floor(monkeypatch): import storage.studio_db as db diff --git a/studio/backend/tests/test_picker_service.py b/studio/backend/tests/test_picker_service.py new file mode 100644 index 0000000000..be7ea18f03 --- /dev/null +++ b/studio/backend/tests/test_picker_service.py @@ -0,0 +1,266 @@ +# SPDX-License-Identifier: AGPL-3.0-only +# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0 + +import json +from types import SimpleNamespace + +from picker.service import ( + MAX_TEMPLATE_METADATA_BYTES, + _chat_template_from_dir, + _chat_template_from_processor_json, + _chat_template_from_tokenizer_config, + _chat_template_from_tokenizer_dir, + _find_gguf_in_dir, + _iter_ggufs, + read_default_chat_template, + validate_chat_template, +) + + +def test_iter_ggufs_skips_gguf_companions(tmp_path): + mtp_dir = tmp_path / "MTP" + mtp_dir.mkdir() + main = tmp_path / "model-Q8_0.gguf" + main.write_bytes(b"") + (tmp_path / "mmproj-F16.gguf").write_bytes(b"") + (tmp_path / "mtp-model-Q8_0.gguf").write_bytes(b"") + (mtp_dir / "model-Q8_0-MTP.gguf").write_bytes(b"") + (tmp_path / "model-Q8_0-be.gguf").write_bytes(b"") + + assert _iter_ggufs(tmp_path) == [main] + + +def test_find_gguf_in_dir_matches_quant_label(tmp_path): + mtp_dir = tmp_path / "MTP" + mtp_dir.mkdir() + main = tmp_path / "model-Q8_0.gguf" + main.write_bytes(b"") + (mtp_dir / "model-Q8_0-MTP.gguf").write_bytes(b"") + (tmp_path / "model-Q4_K_M.gguf").write_bytes(b"") + + assert _find_gguf_in_dir(tmp_path, "Q8_0") == main + assert _find_gguf_in_dir(tmp_path, "Q4_K") is None + + +def test_find_gguf_in_dir_without_variant_prefers_largest_model(tmp_path): + smaller = tmp_path / "a-model-Q4_K_M.gguf" + larger = tmp_path / "z-model-Q8_0.gguf" + smaller.write_bytes(b"0") + larger.write_bytes(b"00") + + assert _find_gguf_in_dir(tmp_path, None) == larger + + +def test_find_gguf_in_dir_without_variant_prefers_first_split(tmp_path): + first = tmp_path / "model-Q4_K_M-00001-of-00003.gguf" + second = tmp_path / "model-Q4_K_M-00002-of-00003.gguf" + third = tmp_path / "model-Q4_K_M-00003-of-00003.gguf" + first.write_bytes(b"0") + second.write_bytes(b"000") + third.write_bytes(b"00") + + assert _find_gguf_in_dir(tmp_path, None) == first + + first.unlink() + assert _find_gguf_in_dir(tmp_path, None) == second + + +def test_find_gguf_in_dir_matches_bpw_variant_base_label(tmp_path): + target = tmp_path / "model-IQ4_XS-3.53bpw.gguf" + target.write_bytes(b"") + (tmp_path / "model-Q4_K_M.gguf").write_bytes(b"") + + assert _find_gguf_in_dir(tmp_path, "IQ4_XS") == target + assert _find_gguf_in_dir(tmp_path, "IQ4_XS-3.53bpw") == target + assert _find_gguf_in_dir(tmp_path, "Q4_K") is None + + +def test_validate_chat_template_accepts_valid_and_empty(): + assert validate_chat_template("{{ messages[0].content }}").valid is True + assert validate_chat_template("").valid is True + assert validate_chat_template(" ").valid is True + + +def test_validate_chat_template_reports_syntax_error_with_line(): + result = validate_chat_template("{% if %}{% endif %}") + assert result.valid is False + assert result.error is not None + assert result.error.startswith("Line ") + + +def test_chat_template_from_tokenizer_config_reads_string(): + assert _chat_template_from_tokenizer_config({"chat_template": "HELLO"}) == "HELLO" + assert _chat_template_from_tokenizer_config({"chat_template": " "}) is None + assert _chat_template_from_tokenizer_config({}) is None + + +def test_chat_template_from_tokenizer_config_prefers_named_default(): + config = { + "chat_template": [ + {"name": "tool_use", "template": "TOOL"}, + {"name": "default", "template": "DEFAULT"}, + ] + } + assert _chat_template_from_tokenizer_config(config) == "DEFAULT" + + +def test_chat_template_from_tokenizer_config_falls_back_to_first_entry(): + config = { + "chat_template": [ + {"name": "tool_use", "template": "TOOL"}, + {"name": "other", "template": "OTHER"}, + ] + } + assert _chat_template_from_tokenizer_config(config) == "TOOL" + + +def test_chat_template_from_tokenizer_dir_prefers_jinja_file(tmp_path): + (tmp_path / "chat_template.jinja").write_text("FROM_JINJA", encoding = "utf-8") + (tmp_path / "tokenizer_config.json").write_text( + json.dumps({"chat_template": "FROM_CONFIG"}), encoding = "utf-8" + ) + assert _chat_template_from_tokenizer_dir(tmp_path) == "FROM_JINJA" + + +def test_chat_template_from_tokenizer_dir_reads_tokenizer_config(tmp_path): + (tmp_path / "tokenizer_config.json").write_text( + json.dumps({"chat_template": "FROM_CONFIG"}), encoding = "utf-8" + ) + assert _chat_template_from_tokenizer_dir(tmp_path) == "FROM_CONFIG" + + +def test_chat_template_from_dir_without_variant_prefers_tokenizer(tmp_path): + (tmp_path / "tokenizer_config.json").write_text( + json.dumps({"chat_template": "FROM_CONFIG"}), encoding = "utf-8" + ) + assert _chat_template_from_dir(tmp_path) == "FROM_CONFIG" + + +def test_chat_template_from_dir_with_variant_still_prefers_tokenizer(tmp_path, monkeypatch): + (tmp_path / "tokenizer_config.json").write_text( + json.dumps({"chat_template": "FROM_CONFIG"}), encoding = "utf-8" + ) + (tmp_path / "model-Q4_K_M.gguf").write_bytes(b"") + monkeypatch.setattr("picker.service.read_gguf_chat_template", lambda _path: "FROM_GGUF") + # Selecting a variant must not flip precedence to the embedded GGUF template. + assert _chat_template_from_dir(tmp_path, "Q4_K_M") == "FROM_CONFIG" + + +def test_chat_template_from_dir_with_variant_falls_back_to_gguf(tmp_path, monkeypatch): + (tmp_path / "model-Q4_K_M.gguf").write_bytes(b"") + monkeypatch.setattr("picker.service.read_gguf_chat_template", lambda _path: "FROM_GGUF") + # With no tokenizer sidecar, the embedded GGUF template is still the fallback. + assert _chat_template_from_dir(tmp_path, "Q4_K_M") == "FROM_GGUF" + + +def test_chat_template_from_dir_returns_none_when_absent(tmp_path): + assert _chat_template_from_dir(tmp_path) is None + + +def test_read_default_chat_template_direct_gguf_prefers_sidecar(tmp_path, monkeypatch): + gguf = tmp_path / "model-Q4_K_M.gguf" + gguf.write_bytes(b"") + (tmp_path / "tokenizer_config.json").write_text( + json.dumps({"chat_template": "FROM_CONFIG"}), encoding = "utf-8" + ) + monkeypatch.setattr("picker.service._build_browse_allowlist", lambda: [tmp_path]) + monkeypatch.setattr("picker.service.read_gguf_chat_template", lambda _path: "FROM_GGUF") + # A directly selected .gguf must prefer a maintained sidecar over its embedded copy. + assert read_default_chat_template(str(gguf)) == "FROM_CONFIG" + + +def test_read_default_chat_template_direct_gguf_falls_back_to_embedded(tmp_path, monkeypatch): + gguf = tmp_path / "model-Q4_K_M.gguf" + gguf.write_bytes(b"") + monkeypatch.setattr("picker.service._build_browse_allowlist", lambda: [tmp_path]) + monkeypatch.setattr("picker.service.read_gguf_chat_template", lambda _path: "FROM_GGUF") + # With no sidecar next to the file, the embedded GGUF template is the fallback. + assert read_default_chat_template(str(gguf)) == "FROM_GGUF" + + +def test_tokenizer_config_over_size_limit_is_skipped_not_parsed(tmp_path): + # An oversized tokenizer_config.json must be skipped before json.loads so a + # hostile sidecar cannot exhaust memory. + padding = "x" * (MAX_TEMPLATE_METADATA_BYTES + 1024) + (tmp_path / "tokenizer_config.json").write_text( + json.dumps({"chat_template": "HELLO", "_pad": padding}), encoding = "utf-8" + ) + assert _chat_template_from_tokenizer_dir(tmp_path) is None + + +def test_processor_json_over_size_limit_is_skipped_not_parsed(tmp_path): + padding = "x" * (MAX_TEMPLATE_METADATA_BYTES + 1024) + (tmp_path / "chat_template.json").write_text( + json.dumps({"default": "HELLO", "_pad": padding}), encoding = "utf-8" + ) + assert _chat_template_from_processor_json(tmp_path) is None + + +def test_tokenizer_config_at_size_limit_is_still_read(tmp_path): + # A normal-sized config is unaffected by the bound (regression guard). + (tmp_path / "tokenizer_config.json").write_text( + json.dumps({"chat_template": "FROM_CONFIG"}), encoding = "utf-8" + ) + assert _chat_template_from_tokenizer_dir(tmp_path) == "FROM_CONFIG" + + +def test_remote_template_over_size_limit_is_skipped_before_download(monkeypatch): + # An uncached Hub repo whose template exceeds the cap must be skipped via the + # remote size pre-check, never downloaded. + import huggingface_hub + + monkeypatch.setattr("picker.service.resolve_cached_repo_id_case", lambda name: name) + monkeypatch.setattr("picker.service.iter_hf_cache_snapshots", lambda resolved: []) + + def _fail_download(*args, **kwargs): + raise AssertionError("oversized remote template must not be downloaded") + + def _fake_get_paths_info(self, repo_id, paths, **kwargs): + return [SimpleNamespace(path = p, size = MAX_TEMPLATE_METADATA_BYTES + 1) for p in paths] + + monkeypatch.setattr(huggingface_hub, "hf_hub_download", _fail_download) + monkeypatch.setattr(huggingface_hub.HfApi, "get_paths_info", _fake_get_paths_info) + + assert read_default_chat_template("org/oversized-model") is None + + +def test_remote_oversized_jinja_falls_through_to_tokenizer_template(tmp_path, monkeypatch): + # A raw chat_template.jinja between the response cap (MAX_CHAT_TEMPLATE_BYTES) + # and the download bound (MAX_TEMPLATE_METADATA_BYTES) must not be returned: the + # route drops it, so the remote path must skip the oversized Jinja and fall + # through to the smaller tokenizer_config.json. + import huggingface_hub + from picker.schemas import MAX_CHAT_TEMPLATE_BYTES + + big_jinja = tmp_path / "chat_template.jinja" + big_jinja.write_text("{{ x }}" * (MAX_CHAT_TEMPLATE_BYTES // 4), encoding = "utf-8") + assert MAX_CHAT_TEMPLATE_BYTES < big_jinja.stat().st_size < MAX_TEMPLATE_METADATA_BYTES + tokenizer_config = tmp_path / "tokenizer_config.json" + tokenizer_config.write_text(json.dumps({"chat_template": "SMALL_TEMPLATE"}), encoding = "utf-8") + files = { + "chat_template.jinja": big_jinja, + "tokenizer_config.json": tokenizer_config, + } + + monkeypatch.setattr("picker.service.resolve_cached_repo_id_case", lambda name: name) + monkeypatch.setattr("picker.service.iter_hf_cache_snapshots", lambda resolved: []) + + def _fake_download(repo_id, rel, **kwargs): + target = files.get(rel) + if target is None: + raise FileNotFoundError(rel) + return str(target) + + def _fake_get_paths_info(self, repo_id, paths, **kwargs): + return [ + SimpleNamespace( + path = p, + size = files[p].stat().st_size if p in files else 0, + ) + for p in paths + ] + + monkeypatch.setattr(huggingface_hub, "hf_hub_download", _fake_download) + monkeypatch.setattr(huggingface_hub.HfApi, "get_paths_info", _fake_get_paths_info) + + assert read_default_chat_template("org/big-jinja-model") == "SMALL_TEMPLATE" diff --git a/studio/backend/tests/test_rocm_windows_vram_7072.py b/studio/backend/tests/test_rocm_windows_vram_7072.py new file mode 100644 index 0000000000..b4079831b7 --- /dev/null +++ b/studio/backend/tests/test_rocm_windows_vram_7072.py @@ -0,0 +1,361 @@ +# SPDX-License-Identifier: AGPL-3.0-only +# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0 + +"""Regression tests for issue #7072 -- "VRAM Usage in System Tab is wrong". + +Reporter: dual AMD (Radeon PRO W7900 ~48GB + W7500 8GB), Windows 10, ROCm 7.13, +torch 2.11.0+rocm7.13. On Windows without a HIP SDK, amd-smi is permanently +disabled (avoids a UAC/DiskPart prompt) and hipMemGetInfo returns free==total +(used 0). Two symptoms followed: + + * System tab (/api/system -> get_visible_gpu_utilization) showed ~0 VRAM used + on every GPU (torch mem_get_info free==total quirk; ROCm/ROCm#1909). + * get_gpu_utilization()'s Windows fallback SUMMED "GPU Adapter Memory\\Dedicated + Usage" across all adapters into ONE fake device with only GPU 0's total, so + the second GPU never appeared. + +The fix reads the per-adapter (LUID-instanced) Dedicated Usage performance +counter -- Task Manager's source -- for per-GPU used, takes per-GPU total from +torch device properties, and guards the free==total mem_get_info quirk. CI has no +AMD GPU/Windows, so torch, the performance counter, and platform are all mocked. +""" + +from __future__ import annotations + +import subprocess +import sys +import types + +import pytest + +from utils.hardware import hardware as hw + +GB = 1024**3 +MiB = 1024**2 + + +# ----------------------------------------------------------------------------- # +# Fakes +# ----------------------------------------------------------------------------- # +def _fake_torch( + devices, + *, + free_equals_total = False, + used_per_device = None, +): + """Build a fake `torch` module. devices: list of (name, total_bytes).""" + dev = list(devices) + + class _Props: + def __init__(self, name, total): + self.name = name + self.total_memory = total + + def get_device_properties(i): + name, total = dev[i] + return _Props(name, total) + + def mem_get_info(i): + _, total = dev[i] + if free_equals_total: + return (total, total) + used = used_per_device[i] if used_per_device is not None else 0 + return (total - used, total) + + t = types.ModuleType("torch") + t.__version__ = "2.11.0+rocm7.13" + t.version = types.SimpleNamespace(hip = "7.13", cuda = None) + t.cuda = types.SimpleNamespace( + is_available = lambda: len(dev) > 0, + device_count = lambda: len(dev), + current_device = lambda: 0, + get_device_properties = get_device_properties, + mem_get_info = mem_get_info, + memory_allocated = lambda i: 0, + memory_reserved = lambda i: 0, + ) + return t + + +def _adapter_output(adapters): + if not adapters: + return "__NONE__\n" + return "".join(f"{name}|{int(used)}\n" for name, used in adapters) + + +def _subprocess_run(*, adapter_output = "__NONE__\n", util_output = "12.0\n"): + def fake_run(cmd, *a, **k): + joined = " ".join(cmd) if isinstance(cmd, list) else str(cmd) + if "GPU Adapter Memory" in joined and "InstanceName" in joined: + out = adapter_output + elif "engtype_3D" in joined or "GPU Engine" in joined: + out = util_output + else: + out = "-1\n" + return subprocess.CompletedProcess(args = cmd, returncode = 0, stdout = out, stderr = "") + + return fake_run + + +@pytest.fixture +def win_rocm(monkeypatch): + """Configure the hardware module as a Windows ROCm host with 2 visible GPUs.""" + monkeypatch.setattr(hw, "get_device", lambda: hw.DeviceType.CUDA) + monkeypatch.setattr(hw, "IS_ROCM", True) + monkeypatch.setattr(hw.platform, "system", lambda: "Windows") + monkeypatch.setattr(hw.sys, "platform", "win32") + monkeypatch.setattr(hw, "_smi_query", lambda *a, **k: None) # amd-smi disabled + # Visible set via HIP mask so we don't shell out to amd-smi for the count. + monkeypatch.setenv("HIP_VISIBLE_DEVICES", "0,1") + monkeypatch.delenv("CUDA_VISIBLE_DEVICES", raising = False) + monkeypatch.delenv("ROCR_VISIBLE_DEVICES", raising = False) + return monkeypatch + + +REPORTER_ADAPTERS = [ + ("luid_0x00000000_0x0000d1e2_phys_0", 40.0 * GB), # W7900, model loaded + ("luid_0x00000000_0x0000e34a_phys_0", 0.5 * GB), # W7500, idle + ("luid_0x00000000_0x0000f001_phys_0", 3 * MiB), # Basic Render Driver +] +DEVICES = [("AMD Radeon PRO W7900", 48 * GB), ("AMD Radeon PRO W7500", 8 * GB)] + + +# ----------------------------------------------------------------------------- # +# System tab (get_visible_gpu_utilization) -- the reporter's screenshot +# ----------------------------------------------------------------------------- # +def test_system_tab_shows_per_gpu_used(win_rocm, monkeypatch): + monkeypatch.setitem(sys.modules, "torch", _fake_torch(DEVICES, free_equals_total = True)) + monkeypatch.setattr( + hw.subprocess, "run", _subprocess_run(adapter_output = _adapter_output(REPORTER_ADAPTERS)) + ) + + devices = hw.get_visible_gpu_utilization()["devices"] + by_idx = {d["index"]: d for d in devices} + assert len(devices) == 2 + assert by_idx[0]["vram_total_gb"] == 48.0 + assert by_idx[0]["vram_used_gb"] == pytest.approx(40.0, abs = 0.01) # not 0 + assert by_idx[1]["vram_total_gb"] == 8.0 # own total + # The 3 MiB Basic Render Driver counter makes this a hidden-adapter case: only + # the 40 GiB is forced onto the 48 GiB card; the idle card reads Unknown. + assert by_idx[1]["vram_used_gb"] is None + assert by_idx[1]["vram_utilization_pct"] is None + assert all( + d["vram_used_gb"] <= d["vram_total_gb"] for d in devices if d["vram_used_gb"] is not None + ) + + +def test_gpu_utilization_does_not_collapse(win_rocm, monkeypatch): + monkeypatch.setitem(sys.modules, "torch", _fake_torch(DEVICES, free_equals_total = True)) + monkeypatch.setattr( + hw.subprocess, "run", _subprocess_run(adapter_output = _adapter_output(REPORTER_ADAPTERS)) + ) + + result = hw.get_gpu_utilization() + devices = result["devices"] + assert sorted(d["index"] for d in devices) == [0, 1] # both GPUs, no collapse + assert {d["vram_total_gb"] for d in devices} == {48.0, 8.0} + assert result["vram_total_gb"] == 48.0 # legacy primary mirror preserved + + +def test_localized_counter_reports_unknown_not_zero(win_rocm, monkeypatch): + monkeypatch.setitem(sys.modules, "torch", _fake_torch(DEVICES, free_equals_total = True)) + monkeypatch.setattr(hw.subprocess, "run", _subprocess_run(adapter_output = "__NONE__\n")) + + devices = hw.get_visible_gpu_utilization()["devices"] + assert len(devices) == 2 # both still shown with correct totals + assert {d["vram_total_gb"] for d in devices} == {48.0, 8.0} + assert all(d["vram_used_gb"] is None for d in devices) # unknown, not fake 0 + assert all(d["vram_utilization_pct"] is None for d in devices) + + +# ----------------------------------------------------------------------------- # +# mem_get_info free==total guard scoping +# ----------------------------------------------------------------------------- # +def test_mem_get_info_guard_scopes_to_windows_rocm(monkeypatch): + torch_mod = _fake_torch(DEVICES, free_equals_total = True) + monkeypatch.setattr(hw, "get_device", lambda: hw.DeviceType.CUDA) + monkeypatch.setitem(sys.modules, "torch", torch_mod) + + # Windows ROCm -> used unknown (None), total kept. + monkeypatch.setattr(hw, "IS_ROCM", True) + monkeypatch.setattr(hw.sys, "platform", "win32") + win = hw._torch_get_per_device_info([0, 1]) + assert [d["used_gb"] for d in win] == [None, None] + assert [d["total_gb"] for d in win] == [48.0, 8.0] + + # Linux ROCm -> unchanged numeric used. + monkeypatch.setattr(hw.sys, "platform", "linux") + assert [d["used_gb"] for d in hw._torch_get_per_device_info([0, 1])] == [0.0, 0.0] + + # Windows NVIDIA -> guard must not fire. + monkeypatch.setattr(hw, "IS_ROCM", False) + monkeypatch.setattr(hw.sys, "platform", "win32") + assert [d["used_gb"] for d in hw._torch_get_per_device_info([0, 1])] == [0.0, 0.0] + + +# ----------------------------------------------------------------------------- # +# Per-adapter attribution helpers (pure unit) +# ----------------------------------------------------------------------------- # +def test_match_adapter_pairs_and_clamps(): + assert hw._match_adapter_used_to_devices([40 * GB, 0.5 * GB], [48 * GB, 8 * GB]) == [ + 40 * GB, + 0.5 * GB, + ] + assert hw._match_adapter_used_to_devices([100 * GB], [48 * GB]) == [48 * GB] # clamp + assert hw._match_adapter_used_to_devices([40 * GB], [48 * GB, 8 * GB]) == [40 * GB, None] + + +def test_match_adapter_reports_unknown_when_more_active_than_visible(): + # More adapters actively using VRAM than are visible (a GPU outside the mask): + # attribution would fabricate a value, so report unknown for every device. + assert hw._match_adapter_used_to_devices([40 * GB, 0.5 * GB], [8 * GB]) == [None] + + +def test_match_adapter_reports_unknown_when_hidden_high_use_adapter_survives_filter(): + # Idle 8 GiB card (10 MiB noise) beside a hidden 48 GiB card at 40 GiB: the + # 40 GiB can't fit the 8 GiB device, so clamping there would fabricate. Unknown. + assert hw._match_adapter_used_to_devices([40 * GB, 10 * MiB], [8 * GB]) == [None] + # Order of the counters must not matter. + assert hw._match_adapter_used_to_devices([10 * MiB, 40 * GB], [8 * GB]) == [None] + + +def test_match_adapter_reports_unknown_for_placeholder_fallback(): + # Every counter below the 64 MiB floor plus a placeholder: no LUID-to-ordinal + # mapping tells placeholder from idle GPU, so report unknown, not fabricate. + # Single visible 8 GiB card idle (10 MiB) beside a 50 MiB placeholder counter. + assert hw._match_adapter_used_to_devices([50 * MiB, 10 * MiB], [8 * GB]) == [None] + # Order of the counters must not matter. + assert hw._match_adapter_used_to_devices([10 * MiB, 50 * MiB], [8 * GB]) == [None] + # Two idle visible GPUs plus a placeholder: all three counters below the floor. + assert hw._match_adapter_used_to_devices([50 * MiB, 10 * MiB, 5 * MiB], [48 * GB, 8 * GB]) == [ + None, + None, + ] + + +def test_match_adapter_reports_unknown_when_usage_not_capacity_ordered(): + # 8 GiB card at 7 GiB beside a 48 GiB card at 5 GiB: the bigger usage still fits + # the smaller card, so both pairings are feasible -> unknown. + assert hw._match_adapter_used_to_devices([7 * GB, 5 * GB], [8 * GB, 48 * GB]) == [None, None] + # Device order must not matter (same physical situation, ordinals flipped). + assert hw._match_adapter_used_to_devices([7 * GB, 5 * GB], [48 * GB, 8 * GB]) == [None, None] + # Same-capacity cards with unequal usage are equally unattributable. + assert hw._match_adapter_used_to_devices([12 * GB, 8 * GB], [24 * GB, 24 * GB]) == [None, None] + # A single usage that fits both cards can sit on either -> unknown. + assert hw._match_adapter_used_to_devices([5 * GB], [48 * GB, 8 * GB]) == [None, None] + # But a capacity-forced assignment (usage exceeds the smaller card) is kept: + # 40 GiB can only be the 48 GiB card, so it is not fabrication. + assert hw._match_adapter_used_to_devices([40 * GB], [48 * GB, 8 * GB]) == [40 * GB, None] + + +def test_match_adapter_reports_unknown_when_hidden_usage_fits_visible_card(): + # A survivor that merely *fits* a visible card must not be pinned onto it. Two + # cards (48/8 GiB) at 40 GiB / 10 MiB beside a hidden 6 GiB adapter: the 6 GiB + # fits the idle 8 GiB card but isn't forced -> Unknown; only 40 GiB is forced. + assert hw._match_adapter_used_to_devices([40 * GB, 10 * MiB, 6 * GB], [48 * GB, 8 * GB]) == [ + 40 * GB, + None, + ] + # Counter order must not matter. + assert hw._match_adapter_used_to_devices([6 * GB, 40 * GB, 10 * MiB], [48 * GB, 8 * GB]) == [ + 40 * GB, + None, + ] + # A single visible card with a hidden adapter is never attributable: a fitting + # survivor could be the hidden GPU's while the visible card is idle. + assert hw._match_adapter_used_to_devices([6 * GB, 10 * MiB], [8 * GB]) == [None] + + +def test_match_adapter_capacity_forced_matrix(): + """Exhaustive hidden-adapter matrix for the capacity-forced rule. + + A value is emitted only when the supra-threshold counters number exactly the + visible devices AND a device's ranked usage strictly exceeds every smaller + card's capacity. Otherwise (a visible card idle, a merely-fitting usage, or the + smallest card) every device reports unknown. + """ + m = hw._match_adapter_used_to_devices + # -- exactly-n supra-threshold counters, capacity-forced survivors are kept - # + # Both visible cards have a real reading (the 3 MiB is a placeholder): 40 GiB + # forced onto the 48 GiB card, 0.5 GiB not forced -> None. + assert m([40 * GB, 0.5 * GB, 3 * MiB], [48 * GB, 8 * GB]) == [40 * GB, None] + # Three visible cards all active (supra-threshold) + placeholder: 40 > 24 and + # 20 > 8, both forced; the 8 GiB card is not forced -> None. + assert m([40 * GB, 20 * GB, 5 * GB, 3 * MiB], [48 * GB, 24 * GB, 8 * GB]) == [ + 40 * GB, + 20 * GB, + None, + ] + # -- fewer supra-threshold counters than visible cards -> all unknown ------ # + # A visible card is idle, so even a "forced" 40 could be the hidden GPU's. + assert m([40 * GB, 3 * MiB, 3 * MiB], [48 * GB, 8 * GB]) == [None, None] + assert m([40 * GB, 10 * MiB, 10 * MiB], [48 * GB, 8 * GB]) == [None, None] + assert m([40 * GB, 20 * GB, 3 * MiB, 3 * MiB], [48 * GB, 24 * GB, 8 * GB]) == [ + None, + None, + None, + ] + # Middle usage (6 GiB) fits both the 24 and 8 GiB cards, and only two cards are + # active for three visible -> not a bijection -> all unknown. + assert m([40 * GB, 6 * GB, 3 * MiB, 3 * MiB], [48 * GB, 24 * GB, 8 * GB]) == [ + None, + None, + None, + ] + # -- hidden larger than every visible card -> all unknown ----------------- # + assert m([40 * GB, 10 * MiB], [8 * GB]) == [None] + assert m([48 * GB, 3 * MiB, 3 * MiB], [24 * GB, 8 * GB]) == [None, None] + # -- more active adapters than visible cards -> all unknown --------------- # + assert m([40 * GB, 7 * GB, 6 * GB, 3 * MiB], [48 * GB, 8 * GB]) == [None, None] + assert m([40 * GB, 7 * GB, 6 * GB, 3 * MiB, 3 * MiB], [48 * GB, 8 * GB]) == [None, None] + # -- every counter below the noise floor (placeholder fallback) -> unknown - # + assert m([50 * MiB, 10 * MiB], [8 * GB]) == [None] + assert m([50 * MiB, 10 * MiB, 5 * MiB], [48 * GB, 8 * GB]) == [None, None] + # -- equal-capacity cards with a hidden adapter: nothing is forced -------- # + assert m([40 * GB, 40 * GB, 3 * MiB], [48 * GB, 48 * GB]) == [None, None] + assert m([40 * GB, 30 * GB, 3 * MiB], [48 * GB, 48 * GB]) == [None, None] + + +def test_perf_counter_parser_and_sentinel(monkeypatch): + monkeypatch.setattr(hw.platform, "system", lambda: "Windows") + monkeypatch.setattr( + hw.subprocess, "run", _subprocess_run(adapter_output = _adapter_output(REPORTER_ADAPTERS)) + ) + parsed = hw._rocm_windows_perf_counter_vram_by_adapter() + assert parsed is not None and len(parsed) == 3 + assert parsed[0][0].startswith("luid_") + monkeypatch.setattr(hw.subprocess, "run", _subprocess_run(adapter_output = "__NONE__\n")) + assert hw._rocm_windows_perf_counter_vram_by_adapter() is None + + +# ----------------------------------------------------------------------------- # +# Unified-memory (Strix Halo APU) total reconciliation (Codex #7238) +# ----------------------------------------------------------------------------- # +def test_unified_memory_adopts_torch_total_even_when_used_unknown(): + """Windows ROCm unified-memory APU: torch's used is None but its total (the full + GTT pool) is authoritative. The correction must still adopt the larger total; + used stays at amd-smi's figure when torch's is unknown.""" + metrics = {"vram_total_gb": 8.0, "vram_used_gb": 2.0, "vram_utilization_pct": 25.0} + hw._apply_unified_memory_correction(metrics, {"total_gb": 124.0, "used_gb": None, "index": 0}) + assert metrics["vram_total_gb"] == 124.0 # full unified pool, not the 8 GB carve-out + assert metrics["vram_used_gb"] == 2.0 # amd-smi used preserved (torch's was None) + assert metrics["vram_utilization_pct"] == pytest.approx(round(2.0 / 124.0 * 100, 1)) + + +def test_unified_memory_overwrites_used_when_torch_used_known(): + """When torch reports both a larger total and a known used, both are adopted + and utilization is recomputed against the corrected total (unchanged path).""" + metrics = {"vram_total_gb": 8.0, "vram_used_gb": 2.0, "vram_utilization_pct": 25.0} + hw._apply_unified_memory_correction(metrics, {"total_gb": 124.0, "used_gb": 40.0, "index": 0}) + assert metrics["vram_total_gb"] == 124.0 + assert metrics["vram_used_gb"] == 40.0 + assert metrics["vram_utilization_pct"] == pytest.approx(round(40.0 / 124.0 * 100, 1)) + + +def test_unified_memory_no_op_when_torch_total_not_larger(): + """A discrete GPU where torch total does not exceed amd-smi's is left untouched.""" + metrics = {"vram_total_gb": 48.0, "vram_used_gb": 10.0, "vram_utilization_pct": 20.8} + hw._apply_unified_memory_correction(metrics, {"total_gb": 48.0, "used_gb": None, "index": 0}) + assert metrics["vram_total_gb"] == 48.0 + assert metrics["vram_used_gb"] == 10.0 + assert metrics["vram_utilization_pct"] == 20.8 diff --git a/studio/backend/tests/test_setup_llama_cpp_backend.py b/studio/backend/tests/test_setup_llama_cpp_backend.py new file mode 100644 index 0000000000..36928c680c --- /dev/null +++ b/studio/backend/tests/test_setup_llama_cpp_backend.py @@ -0,0 +1,154 @@ +# SPDX-License-Identifier: AGPL-3.0-only +# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0 + +"""setup.sh and setup.ps1 must map UNSLOTH_LLAMA_CPP_BACKEND=cpu to +install_llama_prebuilt.py's --force-cpu so users can force the CPU-only prebuilt +on GPU hosts (#7213). The match is case-insensitive and whitespace-trimmed, an +unrecognized value warns instead of silently falling back, and macOS warns (no +CPU-only bundle). Runs the real block extracted from each script so the tests +track the shipped logic. +""" + +import os +import re +import shutil +import subprocess +from pathlib import Path + +import pytest + +_STUDIO = Path(__file__).resolve().parents[2] +_SETUP_SH = _STUDIO / "setup.sh" +_SETUP_PS1 = _STUDIO / "setup.ps1" +_SKIP_NO_BASH = pytest.mark.skipif(shutil.which("bash") is None, reason = "bash unavailable") +_SKIP_NO_PWSH = pytest.mark.skipif(shutil.which("pwsh") is None, reason = "pwsh unavailable") + + +def _backend_block() -> str: + text = _SETUP_SH.read_text(encoding = "utf-8") + m = re.search(r"_llama_backend=.*?esac", text, re.DOTALL) + assert m, "UNSLOTH_LLAMA_CPP_BACKEND block not found in setup.sh" + return m.group(0) + + +def _run(value: str | None, system: str = "Linux") -> tuple[list[str], str]: + # Pass the value through env (not the script text) so whitespace survives, and + # stub the setup.sh logging helpers the unknown-value branch calls. system sets + # _HOST_SYSTEM so the macOS (Darwin) no-op branch can be exercised. + env = {k: v for k, v in os.environ.items() if k != "UNSLOTH_LLAMA_CPP_BACKEND"} + if value is not None: + env["UNSLOTH_LLAMA_CPP_BACKEND"] = value + harness = ( + f'_PREBUILT_CMD=()\nC_WARN=""\n_HOST_SYSTEM="{system}"\n' + 'step() { printf "STEP: %s\\n" "$*" >&2; }\n' + f"{_backend_block()}\n" + 'printf "%s\\n" "${_PREBUILT_CMD[@]}"' + ) + out = subprocess.run( + ["bash", "-c", harness], capture_output = True, text = True, env = env, check = True + ) + return out.stdout.split(), out.stderr + + +@_SKIP_NO_BASH +@pytest.mark.parametrize("value", ["cpu", "CPU", "Cpu", " cpu ", "CPU\t"]) +def test_backend_cpu_appends_flag(value): + # A deliberate CPU choice persists, so it uses --force-cpu (not the transient + # --cpu-fallback the arm64 GPU-build recovery uses). + args, stderr = _run(value) + assert "--force-cpu" in args + assert "--cpu-fallback" not in args + assert "Ignoring" not in stderr + + +@_SKIP_NO_BASH +@pytest.mark.parametrize("value", ["cpu", "CPU", " cpu "]) +def test_backend_cpu_macos_warns_no_flag(value): + # macOS has no CPU-only bundle (the universal build already runs on CPU), so the + # override warns instead of writing a misleading forced-CPU marker. + args, stderr = _run(value, system = "Darwin") + assert "--force-cpu" not in args + assert "--cpu-fallback" not in args + assert "macOS" in stderr + + +@_SKIP_NO_BASH +@pytest.mark.parametrize("value", [None, "", "auto", "AUTO", " "]) +def test_backend_auto_no_flag_no_warn(value): + args, stderr = _run(value) + assert "--force-cpu" not in args + assert "Ignoring" not in stderr + + +@_SKIP_NO_BASH +@pytest.mark.parametrize("value", ["vulkan", "gpu", "cuda"]) +def test_backend_unknown_warns_and_no_flag(value): + args, stderr = _run(value) + assert "--force-cpu" not in args + assert "Ignoring" in stderr + + +@_SKIP_NO_BASH +def test_arm64_recovery_uses_transient_cpu_fallback(): + # The arm64 Linux GPU-build recovery must stay transient (--cpu-fallback), never + # the persisted --force-cpu, so a later update can still heal to a GPU bundle (#6097). + text = _SETUP_SH.read_text(encoding = "utf-8") + m = re.search(r"_ARM64_CPU_CMD=\((.*?)\)", text, re.DOTALL) + assert m, "arm64 CPU recovery command not found in setup.sh" + block = m.group(1) + assert "--cpu-fallback" in block + assert "--force-cpu" not in block + + +def _ps1_search(pattern: str, flags = 0) -> str: + m = re.search(pattern, _SETUP_PS1.read_text(encoding = "utf-8"), flags) + assert m, f"setup.ps1 block not found: {pattern}" + return m.group(0) + + +def _run_ps1(value: str | None) -> str: + # The override is normalized (assign + warn) at the top of the prebuilt block and + # applied to $prebuiltArgs lower down; compose both real snippets. + normalize = _ps1_search( + r'\$llamaBackend = "\$\(\$env:UNSLOTH_LLAMA_CPP_BACKEND\)".*?Write-Host.*?\n\s*\}', + re.DOTALL, + ) + apply_flag = _ps1_search( + r'if \(\$llamaBackend -eq "cpu"\) \{\s*\$prebuiltArgs \+= "--force-cpu"\s*\}' + ) + env = {k: v for k, v in os.environ.items() if k != "UNSLOTH_LLAMA_CPP_BACKEND"} + if value is not None: + env["UNSLOTH_LLAMA_CPP_BACKEND"] = value + harness = f'$prebuiltArgs = @()\n{normalize}\n{apply_flag}\n"ARGS:" + ($prebuiltArgs -join ",")' + out = subprocess.run( + ["pwsh", "-NoProfile", "-Command", harness], + capture_output = True, + text = True, + env = env, + check = True, + ) + return out.stdout + + +@_SKIP_NO_PWSH +@pytest.mark.parametrize("value", ["cpu", "CPU", "Cpu", " cpu ", "CPU\t"]) +def test_ps1_backend_cpu_appends_flag(value): + out = _run_ps1(value) + assert "--force-cpu" in out + assert "Ignoring" not in out + + +@_SKIP_NO_PWSH +@pytest.mark.parametrize("value", [None, "", "auto", "AUTO", " "]) +def test_ps1_backend_auto_no_flag_no_warn(value): + out = _run_ps1(value) + assert "--force-cpu" not in out + assert "Ignoring" not in out + + +@_SKIP_NO_PWSH +@pytest.mark.parametrize("value", ["vulkan", "gpu", "cuda"]) +def test_ps1_backend_unknown_warns_and_no_flag(value): + out = _run_ps1(value) + assert "--force-cpu" not in out + assert "Ignoring" in out diff --git a/studio/backend/tests/test_tensor_parallel.py b/studio/backend/tests/test_tensor_parallel.py index 06f72d3b9f..00c7aeac69 100644 --- a/studio/backend/tests/test_tensor_parallel.py +++ b/studio/backend/tests/test_tensor_parallel.py @@ -262,9 +262,12 @@ def test_proportional_tensor_split_is_emitted_in_tensor_mode(): src = _load_model_source() assert '"--tensor-split"' in src gate = src.find("if tensor_parallel:") - ts = src.find('"--tensor-split"') + # Find the TP block's emission (after the gate); manual mode emits its own + # --tensor-split earlier in the source from the user's per-GPU shares. + ts = src.find('"--tensor-split"', gate) nxt_else = src.find("self._tensor_parallel = False") assert 0 <= gate < ts < nxt_else, "--tensor-split must be emitted under `if tensor_parallel:`" + assert "tp_tensor_split" in src[gate:nxt_else] def test_mtp_decode_probe_wired_under_tensor_parallel(): diff --git a/studio/backend/tests/test_tp_vision_regression.py b/studio/backend/tests/test_tp_vision_regression.py index fb0989b306..d1372ca415 100644 --- a/studio/backend/tests/test_tp_vision_regression.py +++ b/studio/backend/tests/test_tp_vision_regression.py @@ -126,10 +126,21 @@ _ALLOWED_TP_DROP_GUARDS = { # Capability: --split-mode tensor aborted for this (binary, model) (#6415). # Self-healing -- tried by default, skipped only after a real abort (vs #6416). "tensor_parallel and self._tensor_split_aborts(binary, model_identifier)", - # Capacity: tensor needs >= 2 GPUs clearing the compute-buffer reserve. - "tensor_parallel and len(tp_gpus) < 2", + # Capacity: tensor needs >= 2 GPUs clearing the compute-buffer reserve. Gated + # on plan_tp (not raw tensor_parallel) so manual mode skips this planner (#6414). + "plan_tp and len(tp_gpus) < 2", # Capacity: pooled usable VRAM can't hold weights + MTP reserve -> layer split. "_tp_weight_budget_mib <= _tp_required_mib", + # Manual mode, Auto layers: --fit owns memory and is incompatible with a + # tensor split, so TP is dropped (surfaced via logger.info) before the + # cache-drop, so a quantized KV survives into the --fit load (#6414). + "tensor_parallel and gpu_memory_mode == 'manual' and (gpu_layers < 0)", + # Manual mode, explicit layers: a tensor split still needs >= 2 GPUs in use. + "tensor_parallel and gpu_memory_mode == 'manual' and (gpu_layers >= 0) and (self._effective_gpu_count(sorted(gpu_ids) if gpu_ids else None) < 2)", + # Manual mode, zero layers: nothing to split on the GPU, and a tensor-mode + # launch under the CPU-only GPU mask (no visible devices) aborts the server + # instead of the intended CPU-only load (#6414). + "gpu_memory_mode == 'manual' and gpu_layers == 0", } @@ -364,7 +375,7 @@ def test_compute_buffer_downgrade_preserves_multi_gpu_intent(): full GPU set too, so it is symmetric with the budget/geometry downgrades and doesn't collapse a multi-GPU layer load to one card (reviewer.py P1 on #6659).""" src = inspect.getsource(LlamaCppBackend.load_model) - gate = src.find("tensor_parallel and len(tp_gpus) < 2") + gate = src.find("plan_tp and len(tp_gpus) < 2") assert gate != -1 # Bound to exactly this block: from its gate to the next (budget) downgrade. nxt = src.find("_tp_weight_budget_mib <= _tp_required_mib", gate) diff --git a/studio/backend/tests/test_training_pump_resilience.py b/studio/backend/tests/test_training_pump_resilience.py index d75b205f35..5d2a218482 100644 --- a/studio/backend/tests/test_training_pump_resilience.py +++ b/studio/backend/tests/test_training_pump_resilience.py @@ -310,6 +310,80 @@ def test_pump_finalizes_when_read_keeps_raising_on_dead_worker(monkeypatch): assert b._pump_running is False +def test_interrupted_cancel_clears_in_memory_output_dir(monkeypatch): + # Stop-without-save interrupted before its complete event: /status must not + # keep serving the cleared run's output_dir. + b = TrainingBackend() + finalized: dict = {} + monkeypatch.setattr(b, "_ensure_db_run_created", lambda: None) + monkeypatch.setattr(b, "_finalize_run_in_db", lambda **kw: finalized.update(kw)) + + b._proc = _FakeProc(alive = False) + b._event_queue = _IdleQueue() + b._progress.is_training = True + b._should_stop = True + b._cancel_requested = True + b._output_dir = "/out/x" + + b._pump_loop() + + assert b._output_dir is None + assert finalized.get("status") == "stopped" + assert finalized.get("output_dir") is None + assert finalized.get("clear_output_dir") is True + + +def test_worker_exit_reuses_terminal_stop_save_error(monkeypatch): + b = TrainingBackend() + finalized: dict = {} + monkeypatch.setattr(b, "_ensure_db_run_created", lambda: None) + monkeypatch.setattr(b, "_finalize_run_in_db", lambda **kw: finalized.update(kw)) + + b._proc = _FakeProc(alive = False) + b._event_queue = _IdleQueue() + b._progress.is_training = True + b._should_stop = True + b._cancel_requested = False + b._output_dir = "/out/x" + b.current_job_id = "job-x" + b._terminal_finalize_payload = { + "status": "error", + "error_message": "checkpoint failed", + "output_dir": "/out/x", + "clear_output_dir": False, + "resume_blocked": True, + "expected_job_id": "job-x", + } + + b._pump_loop() + + assert b._output_dir == "/out/x" + assert finalized.get("status") == "error" + assert finalized.get("output_dir") == "/out/x" + assert finalized.get("clear_output_dir") is False + assert finalized.get("resume_blocked") is True + + +def test_dead_worker_crash_preserves_output_dir(monkeypatch): + # A crash (no stop requested) after output_dir was emitted must keep the dir + # in the error finalize: checkpoints under it may still exist. + b = TrainingBackend() + finalized: dict = {} + monkeypatch.setattr(b, "_ensure_db_run_created", lambda: None) + monkeypatch.setattr(b, "_finalize_run_in_db", lambda **kw: finalized.update(kw)) + + b._proc = _FakeProc(alive = False) + b._event_queue = _IdleQueue() + b._progress.is_training = True + b._output_dir = "/out/x" + + b._pump_loop() + + assert finalized.get("status") == "error" + assert finalized.get("output_dir") == "/out/x" + assert finalized.get("clear_output_dir") is False + + def test_start_training_clears_stale_pump_running_flag(): # A prior pump that died abnormally leaves _pump_running True. The next # start_training must clear it during reset so the start-time watchdog can't diff --git a/studio/backend/tests/test_training_resume.py b/studio/backend/tests/test_training_resume.py index 91fdac9961..51425b0428 100644 --- a/studio/backend/tests/test_training_resume.py +++ b/studio/backend/tests/test_training_resume.py @@ -7,6 +7,9 @@ import importlib.util import json from pathlib import Path +import pytest +import torch + _BACKEND = Path(__file__).resolve().parents[1] @@ -25,6 +28,30 @@ def _load_resume_module(): resume = _load_resume_module() +def test_resume_request_accepts_sanitized_null_target_modules(): + from models.training import TrainingStartRequest + request = TrainingStartRequest( + model_name = "unsloth/Qwen3-0.6B", + training_type = "Full Finetuning", + format_type = "alpaca", + target_modules = None, + ) + + assert request.target_modules == [] + + +def _write_checkpoint(out: Path, step: int) -> Path: + checkpoint = out / f"checkpoint-{step}" + checkpoint.mkdir(parents = True, exist_ok = True) + (checkpoint / "trainer_state.json").write_text( + json.dumps({"global_step": step}), encoding = "utf-8" + ) + torch.save({"weight": torch.ones(1)}, checkpoint / "adapter_model.bin") + torch.save({"state": {0: torch.ones(1)}}, checkpoint / "optimizer.pt") + torch.save({"last_epoch": step}, checkpoint / "scheduler.pt") + return checkpoint + + def _stopped_run(**overrides): run = { "status": "stopped", @@ -44,6 +71,36 @@ def test_can_resume_run_allows_checkpointed_non_s3_run(monkeypatch): assert resume.can_resume_run(_stopped_run()) is True +def test_can_resume_run_allows_errored_run_with_checkpoint(monkeypatch): + monkeypatch.setattr(resume, "has_resume_state", lambda _path: True) + + assert resume.can_resume_run(_stopped_run(status = "error")) is True + + +def test_can_resume_run_rejects_errored_run_without_checkpoint(monkeypatch): + monkeypatch.setattr(resume, "has_resume_state", lambda _path: False) + + assert resume.can_resume_run(_stopped_run(status = "error")) is False + + +def test_can_resume_run_allows_errored_run_at_final_step(monkeypatch): + # A save-time crash records final_step == total_steps; resuming re-runs the + # final-save path from the checkpoint. + monkeypatch.setattr(resume, "has_resume_state", lambda _path: True) + + run = _stopped_run(status = "error", final_step = 10, total_steps = 10) + + assert resume.can_resume_run(run) is True + + +def test_can_resume_run_rejects_stopped_run_at_final_step(monkeypatch): + monkeypatch.setattr(resume, "has_resume_state", lambda _path: True) + + run = _stopped_run(final_step = 10, total_steps = 10) + + assert resume.can_resume_run(run) is False + + def test_can_resume_run_rejects_s3_dataset_source(monkeypatch): monkeypatch.setattr(resume, "has_resume_state", lambda _path: True) @@ -91,3 +148,444 @@ def test_list_runs_includes_config_json_for_resume_policy(monkeypatch, tmp_path) result = studio_db.list_runs() assert result["runs"][0]["config_json"] == config_json + + +def test_crashed_run_with_persisted_output_dir_is_resumable(monkeypatch, tmp_path): + from storage import studio_db + + monkeypatch.setenv("UNSLOTH_STUDIO_HOME", str(tmp_path)) + monkeypatch.setattr(studio_db, "_schema_ready", False) + + out = tmp_path / "outputs" / "run_x" + _write_checkpoint(out, 10) + + studio_db.create_run( + id = "run-crash", + model_name = "m", + dataset_name = "d", + config_json = "{}", + started_at = "2026-01-01T00:00:00Z", + total_steps = 20, + ) + studio_db.update_run_output_dir("run-crash", str(out)) + conn = studio_db.get_connection() + conn.execute("UPDATE training_runs SET status = 'error' WHERE id = 'run-crash'") + conn.commit() + conn.close() + + run = studio_db.get_run("run-crash") + assert run["output_dir"] == str(out) + assert resume.can_resume_run(run) is True + + +def test_checkpoint_discovery_skips_malformed_newest(monkeypatch, tmp_path): + monkeypatch.setenv("UNSLOTH_STUDIO_HOME", str(tmp_path)) + out = tmp_path / "outputs" / "run_x" + valid = _write_checkpoint(out, 5) + (_write_checkpoint(out, 8) / "scheduler.pt").unlink() + malformed = out / "checkpoint-10" + malformed.mkdir() + (malformed / "trainer_state.json").write_text(json.dumps({"global_step": 10}), encoding = "utf-8") + (malformed / "adapter_model.bin").write_bytes(b"not a torch archive") + (malformed / "optimizer.pt").write_bytes(b"not a torch archive") + + assert resume.get_resume_checkpoint_path(str(out)) == str(valid) + + +def test_completed_run_keeps_output_dir_and_rejects_stale_cancel(monkeypatch, tmp_path): + from storage import studio_db + + monkeypatch.setenv("UNSLOTH_STUDIO_HOME", str(tmp_path)) + monkeypatch.setattr(studio_db, "_schema_ready", False) + + studio_db.create_run( + id = "r", + model_name = "m", + dataset_name = "d", + config_json = "{}", + started_at = "2026-01-01T00:00:00Z", + total_steps = 10, + ) + studio_db.update_run_output_dir("r", "/out/x") + studio_db.finish_run( + id = "r", + status = "completed", + ended_at = "t", + final_step = 2, + final_loss = None, + duration_seconds = 1, + loss_sparkline = "[]", + output_dir = "/out/x", + error_message = None, + ) + + assert studio_db.get_run("r")["output_dir"] == "/out/x" + assert studio_db.mark_run_cancel_requested("r") is False + assert studio_db.get_run("r")["output_dir"] == "/out/x" + assert studio_db.get_run("r")["resume_blocked"] == 0 + + +def test_finish_run_clears_output_dir_for_stop_without_save(monkeypatch, tmp_path): + from storage import studio_db + + monkeypatch.setenv("UNSLOTH_STUDIO_HOME", str(tmp_path)) + monkeypatch.setattr(studio_db, "_schema_ready", False) + + studio_db.create_run( + id = "r", + model_name = "m", + dataset_name = "d", + config_json = "{}", + started_at = "2026-01-01T00:00:00Z", + total_steps = 10, + ) + studio_db.update_run_output_dir("r", "/out/x") + studio_db.finish_run( + id = "r", + status = "stopped", + ended_at = "t", + final_step = 2, + final_loss = None, + duration_seconds = 1, + loss_sparkline = "[]", + output_dir = None, + error_message = None, + clear_output_dir = True, + ) + + assert studio_db.get_run("r")["output_dir"] is None + conn = studio_db.get_connection() + conn.execute( + "UPDATE training_runs SET status = 'running', output_dir = '/out/x', resume_blocked = 0 WHERE id = 'r'" + ) + conn.commit() + conn.close() + studio_db.mark_run_cancel_requested("r") + studio_db.cleanup_orphaned_runs() + assert studio_db.get_run("r")["status"] == "stopped" + assert studio_db.get_run("r")["output_dir"] is None + + +def test_finish_run_clears_output_dir_on_cancel_error_finalize(monkeypatch, tmp_path): + from storage import studio_db + + monkeypatch.setenv("UNSLOTH_STUDIO_HOME", str(tmp_path)) + monkeypatch.setattr(studio_db, "_schema_ready", False) + + studio_db.create_run( + id = "r", + model_name = "m", + dataset_name = "d", + config_json = "{}", + started_at = "2026-01-01T00:00:00Z", + total_steps = 10, + ) + studio_db.update_run_output_dir("r", "/out/x") + studio_db.finish_run( + id = "r", + status = "stopped", + ended_at = "t", + final_step = 2, + final_loss = None, + duration_seconds = 1, + loss_sparkline = "[]", + output_dir = "/out/x", + error_message = "worker failed during cancel", + clear_output_dir = True, + ) + + assert studio_db.get_run("r")["output_dir"] is None + + +def test_finish_run_preserves_output_dir_for_interrupted_stop_and_save(monkeypatch, tmp_path): + from storage import studio_db + + monkeypatch.setenv("UNSLOTH_STUDIO_HOME", str(tmp_path)) + monkeypatch.setattr(studio_db, "_schema_ready", False) + + studio_db.create_run( + id = "r", + model_name = "m", + dataset_name = "d", + config_json = "{}", + started_at = "2026-01-01T00:00:00Z", + total_steps = 10, + ) + studio_db.update_run_output_dir("r", "/out/x") + studio_db.finish_run( + id = "r", + status = "stopped", + ended_at = "t", + final_step = 2, + final_loss = None, + duration_seconds = 1, + loss_sparkline = "[]", + output_dir = None, + error_message = None, + ) + + assert studio_db.get_run("r")["output_dir"] == "/out/x" + + +def test_resumed_errored_run_is_not_offered_again(monkeypatch, tmp_path): + from storage import studio_db + + monkeypatch.setenv("UNSLOTH_STUDIO_HOME", str(tmp_path)) + monkeypatch.setattr(studio_db, "_schema_ready", False) + + out = tmp_path / "outputs" / "run_x" + _write_checkpoint(out, 10) + + studio_db.create_run( + id = "run-old", + model_name = "m", + dataset_name = "d", + config_json = "{}", + started_at = "2026-01-01T00:00:00Z", + total_steps = 20, + ) + studio_db.update_run_output_dir("run-old", str(out)) + studio_db.finish_run( + id = "run-old", + status = "error", + ended_at = "2026-01-01T00:05:00Z", + final_step = 10, + final_loss = None, + duration_seconds = 1, + loss_sparkline = "[]", + output_dir = None, + error_message = "killed", + ) + studio_db.create_run( + id = "run-new", + model_name = "m", + dataset_name = "d", + config_json = "{}", + started_at = "2026-01-02T00:00:00Z", + total_steps = 20, + output_dir = str(out), + resumed_from_run_id = "run-old", + ) + with pytest.raises(RuntimeError, match = "no longer available"): + studio_db.create_run( + id = "run-duplicate", + model_name = "m", + dataset_name = "d", + config_json = "{}", + started_at = "2026-01-02T00:00:01Z", + total_steps = 20, + output_dir = str(out), + resumed_from_run_id = "run-old", + ) + assert studio_db.get_run("run-duplicate") is None + studio_db.finish_run( + id = "run-new", + status = "error", + ended_at = "2026-01-02T00:05:00Z", + final_step = 15, + final_loss = None, + duration_seconds = 1, + loss_sparkline = "[]", + output_dir = None, + error_message = "killed again", + ) + + old_run = studio_db.get_run("run-old") + new_run = studio_db.get_run("run-new") + assert old_run["resumed_later"] == 1 + assert resume.can_resume_run(old_run) is False + assert new_run["resumed_later"] == 0 + assert resume.can_resume_run(new_run) is True + assert studio_db.get_resumable_run_by_output_dir(str(out))["id"] == "run-new" + + +def test_running_continuation_blocks_older_resume(monkeypatch, tmp_path): + from storage import studio_db + + monkeypatch.setenv("UNSLOTH_STUDIO_HOME", str(tmp_path)) + monkeypatch.setattr(studio_db, "_schema_ready", False) + + out = tmp_path / "outputs" / "run_x" + _write_checkpoint(out, 10) + + studio_db.create_run( + id = "run-old", + model_name = "m", + dataset_name = "d", + config_json = "{}", + started_at = "2026-01-01T00:00:00Z", + total_steps = 20, + ) + studio_db.update_run_output_dir("run-old", str(out)) + studio_db.finish_run( + id = "run-old", + status = "error", + ended_at = "2026-01-01T00:05:00Z", + final_step = 10, + final_loss = None, + duration_seconds = 1, + loss_sparkline = "[]", + output_dir = None, + error_message = "killed", + ) + studio_db.create_run( + id = "run-new", + model_name = "m", + dataset_name = "d", + config_json = "{}", + started_at = "2026-01-02T00:00:00Z", + total_steps = 20, + output_dir = str(out), + resumed_from_run_id = "run-old", + ) + + old_run = studio_db.get_run("run-old") + assert old_run["resumed_later"] == 1 + assert resume.can_resume_run(old_run) is False + assert studio_db.get_resumable_run_by_output_dir(str(out)) is None + + +def test_stop_save_checkpoint_failure_keeps_error_status(monkeypatch, tmp_path): + # A stop-and-save whose checkpoint write failed must finalize as an error so + # history explains the missing resume state (keep_error_status flag). + from core.training.training import TrainingBackend + from storage import studio_db + + monkeypatch.setenv("UNSLOTH_STUDIO_HOME", str(tmp_path)) + monkeypatch.setattr(studio_db, "_schema_ready", False) + + studio_db.create_run( + id = "run-failed-save", + model_name = "m", + dataset_name = "d", + config_json = "{}", + started_at = "2026-01-01T00:00:00Z", + total_steps = 10, + ) + backend = TrainingBackend() + backend.current_job_id = "run-failed-save" + backend._db_run_created = True + backend._should_stop = True + backend._handle_event( + { + "type": "error", + "error": "Failed to save a resumable checkpoint after stop.", + "keep_error_status": True, + } + ) + + run = studio_db.get_run("run-failed-save") + assert run["status"] == "error" + assert "resumable checkpoint" in run["error_message"] + + +def test_can_resume_run_rejects_resume_blocked_run(monkeypatch): + monkeypatch.setattr(resume, "has_resume_state", lambda _path: True) + + assert resume.can_resume_run(_stopped_run(status = "error", resume_blocked = 1)) is False + + +def test_stop_save_checkpoint_failure_with_stale_checkpoint_is_not_resumable(monkeypatch, tmp_path): + # A failed stop-and-save must not offer Resume from an older periodic checkpoint; + # that would roll back past the recorded final step. + from core.training.training import TrainingBackend + from storage import studio_db + + monkeypatch.setenv("UNSLOTH_STUDIO_HOME", str(tmp_path)) + monkeypatch.setattr(studio_db, "_schema_ready", False) + + out = tmp_path / "outputs" / "run_x" + _write_checkpoint(out, 10) + + studio_db.create_run( + id = "run-stale-ckpt", + model_name = "m", + dataset_name = "d", + config_json = "{}", + started_at = "2026-01-01T00:00:00Z", + total_steps = 20, + ) + studio_db.update_run_output_dir("run-stale-ckpt", str(out)) + backend = TrainingBackend() + backend.current_job_id = "run-stale-ckpt" + backend._db_run_created = True + backend._should_stop = True + backend._output_dir = str(out) + backend._handle_event( + { + "type": "error", + "error": "Failed to save a resumable checkpoint after stop.", + "keep_error_status": True, + "resume_blocked": True, + } + ) + + run = studio_db.get_run("run-stale-ckpt") + assert run["status"] == "error" + assert run["resume_blocked"] == 1 + assert run["output_dir"] == str(out) + assert resume.can_resume_run(run) is False + + +def test_user_stop_error_without_checkpoint_ack_is_blocked(monkeypatch, tmp_path): + from core.training.training import TrainingBackend + from storage import studio_db + + monkeypatch.setenv("UNSLOTH_STUDIO_HOME", str(tmp_path)) + monkeypatch.setattr(studio_db, "_schema_ready", False) + + studio_db.create_run( + id = "run-user-stop", + model_name = "m", + dataset_name = "d", + config_json = "{}", + started_at = "2026-01-01T00:00:00Z", + total_steps = 10, + ) + backend = TrainingBackend() + backend.current_job_id = "run-user-stop" + backend._db_run_created = True + backend._should_stop = True + backend._handle_event({"type": "error", "error": "interrupted"}) + + run = studio_db.get_run("run-user-stop") + assert run["status"] == "error" and run["resume_blocked"] == 1 + + +def test_terminal_fallback_keeps_resumable_when_current_checkpoint_landed(monkeypatch, tmp_path): + # Worker died before its terminal event, but a valid current-step checkpoint + # is on disk: the fallback must keep the run resumable, not block it. + from core.training.training import TrainingBackend + + monkeypatch.setenv("UNSLOTH_STUDIO_HOME", str(tmp_path)) + out = tmp_path / "outputs" / "run_ok" + _write_checkpoint(out, 7) + + backend = TrainingBackend() + backend.current_job_id = "run-ok" + backend._should_stop = True + backend._output_dir = str(out) + backend._progress.step = 7 + + kwargs = backend._terminal_finalize_kwargs() + assert kwargs["status"] == "stopped" + assert kwargs["resume_blocked"] is False + + +def test_terminal_fallback_blocks_when_no_current_checkpoint(monkeypatch, tmp_path): + # Same path, but only a stale (older-step) checkpoint exists: must block. + from core.training.training import TrainingBackend + + monkeypatch.setenv("UNSLOTH_STUDIO_HOME", str(tmp_path)) + out = tmp_path / "outputs" / "run_stale" + _write_checkpoint(out, 5) + + backend = TrainingBackend() + backend.current_job_id = "run-stale" + backend._should_stop = True + backend._output_dir = str(out) + backend._progress.step = 7 + + kwargs = backend._terminal_finalize_kwargs() + assert kwargs["status"] == "error" + assert kwargs["resume_blocked"] is True diff --git a/studio/backend/tests/test_training_stop_watchdog.py b/studio/backend/tests/test_training_stop_watchdog.py index 0cd702bce2..cbe2082e82 100644 --- a/studio/backend/tests/test_training_stop_watchdog.py +++ b/studio/backend/tests/test_training_stop_watchdog.py @@ -353,7 +353,7 @@ def test_finalize_after_escalation_clears_state(monkeypatch): # stopped so the UI leaves "Stopping..." and a new run can start. b = TrainingBackend() finstop: list = [] - monkeypatch.setattr(b, "_finish_stopped_run", lambda *a: finstop.append(a)) + monkeypatch.setattr(b, "_finish_stopped_run", lambda *a, **k: finstop.append(a)) b._proc = _FakeProc(alive = True) # wedged: still reports alive b._should_stop = True @@ -365,7 +365,7 @@ def test_finalize_after_escalation_clears_state(monkeypatch): assert b._proc is None, "the wedged handle must be dropped so is_training_active clears" assert b._progress.is_training is False - assert b._progress.status_message == "Training stopped." + assert "valid current-step checkpoint" in b._progress.status_message assert finstop and finstop[0][0] == "job_c", "the captured run must be finalized by id" assert b.is_training_active() is False @@ -375,7 +375,7 @@ def test_finalize_after_escalation_preserves_output_dir(monkeypatch): # must record it even if the watchdog wins the finalize race against the pump. b = TrainingBackend() finstop: list = [] - monkeypatch.setattr(b, "_finish_stopped_run", lambda *a: finstop.append(a)) + monkeypatch.setattr(b, "_finish_stopped_run", lambda *a, **k: finstop.append(a)) b._proc = _FakeProc(alive = True) b._should_stop = True @@ -390,6 +390,28 @@ def test_finalize_after_escalation_preserves_output_dir(monkeypatch): assert finstop[0][1] == "/tmp/outputs/run-123" +def test_finalize_after_escalation_clears_output_dir_on_cancel(monkeypatch): + # Stop-without-saving promises no resume: a cancel that escalates through the + # watchdog clears the persisted output_dir, not a checkpoint path. + b = TrainingBackend() + finstop: list = [] + monkeypatch.setattr(b, "_finish_stopped_run", lambda *a, **k: finstop.append((a, k))) + + b._proc = _FakeProc(alive = True) + b._should_stop = True + b._cancel_requested = True + b.current_job_id = "job_c" + b._db_run_created = True + b._output_dir = "/tmp/outputs/run-123" + + b._finalize_stopped_after_escalation(watched_job_id = "job_c") + + assert finstop and finstop[0][0][0] == "job_c" + assert finstop[0][0][1] is None, "a cancelled run must not record a checkpoint path" + assert finstop[0][1].get("clear_output_dir") is True + assert b._output_dir is None, "/status must stop exposing the cancelled run's dir" + + def test_stop_training_starts_watchdog_only_when_worker_alive(monkeypatch): # No worker -> nothing to escalate; the watchdog must not spawn. b = TrainingBackend() @@ -409,7 +431,7 @@ def test_finalize_after_escalation_no_ops_when_superseded(monkeypatch): # The escalation finalize must then leave the NEW run untouched, not drop its handle. b = TrainingBackend() finstop: list = [] - monkeypatch.setattr(b, "_finish_stopped_run", lambda *a: finstop.append(a)) + monkeypatch.setattr(b, "_finish_stopped_run", lambda *a, **k: finstop.append(a)) old_proc = _FakeProc(alive = False) # force-terminated worker we were watching new_proc = _FakeProc(alive = True) # a new run already took over @@ -430,7 +452,7 @@ def test_finalize_after_escalation_runs_for_its_own_worker(monkeypatch): # finalizes the captured run by id. b = TrainingBackend() finstop: list = [] - monkeypatch.setattr(b, "_finish_stopped_run", lambda *a: finstop.append(a)) + monkeypatch.setattr(b, "_finish_stopped_run", lambda *a, **k: finstop.append(a)) proc = _FakeProc(alive = False) b._proc = proc @@ -451,7 +473,7 @@ def test_finalize_after_escalation_no_ops_on_job_change_during_startup(monkeypat # catch this even though the proc-only guard would not. b = TrainingBackend() finstop: list = [] - monkeypatch.setattr(b, "_finish_stopped_run", lambda *a: finstop.append(a)) + monkeypatch.setattr(b, "_finish_stopped_run", lambda *a, **k: finstop.append(a)) old_proc = _FakeProc(alive = False) # old worker, dead; new _proc not installed yet b._proc = old_proc # still the old handle (== target), so proc guard would pass @@ -509,6 +531,7 @@ def _install_fake_db(monkeypatch): recs["insert_ids"].append(job_id), ) fake_db.update_run_progress = lambda **kw: recs["progress_ids"].append(kw.get("id")) + fake_db.mark_run_cancel_requested = lambda _run_id: True fake_storage.studio_db = fake_db monkeypatch.setitem(sys.modules, "storage", fake_storage) monkeypatch.setitem(sys.modules, "storage.studio_db", fake_db) @@ -518,9 +541,48 @@ def _install_fake_db(monkeypatch): return recs +def test_stop_without_save_creates_missing_row_before_signal(monkeypatch): + recs = _install_fake_db(monkeypatch) + b = TrainingBackend() + b.current_job_id, b._db_config = "job_missing", {"model_name": "m"} + b._stop_queue = queue.Queue() + assert b.stop_training(save = False) is True + assert [run["id"] for run in recs["created"]] == ["job_missing"] + assert b._stop_queue.get_nowait() == {"type": "stop", "save": False} + + b._cancel_requested = b._should_stop = False + sys.modules["storage.studio_db"].mark_run_cancel_requested = lambda _run_id: False + assert b.stop_training(save = False) is False + assert not b._cancel_requested and b._stop_queue.empty() + + new_queue = queue.Queue() + b.current_job_id, b._db_run_created = "job_old", True + b._cancel_requested = b._should_stop = False + + def _supersede(_run_id): + b.current_job_id = "job_new" + b._stop_queue = new_queue + return True + + sys.modules["storage.studio_db"].mark_run_cancel_requested = _supersede + assert b.stop_training(save = False) is False + assert not b._cancel_requested and new_queue.empty() + + def test_finalize_run_in_db_single_winner_under_concurrency(monkeypatch): # The watchdog and pump can both finalize; only one call may reach finish_run. recs = _install_fake_db(monkeypatch) + monkeypatch.setitem(_G, "_DB_FINALIZE_RETRY_S", 0.0) + attempts = 0 + + def flaky_finish(**kw): + nonlocal attempts + attempts += 1 + if attempts < 3: + raise RuntimeError("database is locked") + recs["finished"].append(kw) + + sys.modules["storage.studio_db"].finish_run = flaky_finish b = TrainingBackend() b.current_job_id = "job_x" b._db_run_created = True @@ -539,6 +601,7 @@ def test_finalize_run_in_db_single_winner_under_concurrency(monkeypatch): t.join(timeout = 5) assert len(recs["finished"]) == 1, f"finalize must run once, got {len(recs['finished'])}" + assert attempts == 3 assert b._run_finalized is True @@ -646,7 +709,13 @@ def test_ensure_db_run_created_publishes_only_after_insert(monkeypatch): monkeypatch.setitem(sys.modules, "storage", fake_storage) monkeypatch.setitem(sys.modules, "storage.studio_db", fake_db) - b._ensure_db_run_created() + b._run_intent_lock.acquire() + creator = threading.Thread(target = b._ensure_db_run_created) + creator.start() + time.sleep(0.02) + assert b._db_create_in_progress is False + b._run_intent_lock.release() + creator.join(timeout = 5) assert observed["flag_during_create"] is False, "flag must not be published before insert" assert observed["in_progress_during_create"] is True @@ -718,6 +787,7 @@ def test_escalation_finalizes_watched_run_by_id_end_to_end(monkeypatch): b = TrainingBackend() b.current_job_id = "job_old" b._db_run_created = True + b._should_stop = True b._proc = _FakeProc(alive = False) b._progress.is_training = True b._progress.step = 42 @@ -726,7 +796,8 @@ def test_escalation_finalizes_watched_run_by_id_end_to_end(monkeypatch): b._finalize_stopped_after_escalation(target_proc = b._proc, watched_job_id = "job_old") assert [f["id"] for f in recs["finished"]] == ["job_old"], "must finish the captured run by id" - assert recs["finished"][0]["status"] == "stopped" + assert recs["finished"][0]["status"] == "error" + assert recs["finished"][0]["resume_blocked"] is True assert recs["insert_ids"] == ["job_old"], "buffered metrics must land on the captured run" assert b._metric_buffer == [], "the captured batch must be drained" @@ -737,7 +808,7 @@ def test_escalation_defers_when_row_cannot_be_created_here(monkeypatch): # so the pump's create-then-finalize records the run. Parent state still clears. b = TrainingBackend() called: list = [] - monkeypatch.setattr(b, "_finish_stopped_run", lambda *a: called.append(a)) + monkeypatch.setattr(b, "_finish_stopped_run", lambda *a, **k: called.append(a)) b._proc = _FakeProc(alive = False) b.current_job_id = "job_q" @@ -785,7 +856,7 @@ def test_escalation_does_not_drop_a_new_runs_handle(monkeypatch): new_proc = _FakeProc(alive = True) b._proc = old_proc - def hijack(*a): + def hijack(*a, **k): b._proc = new_proc # a new run takes over during the finalize monkeypatch.setattr(b, "_finish_stopped_run", hijack) diff --git a/studio/backend/tests/test_utils.py b/studio/backend/tests/test_utils.py index 64a3c62156..741f19c67a 100644 --- a/studio/backend/tests/test_utils.py +++ b/studio/backend/tests/test_utils.py @@ -38,7 +38,7 @@ from utils.hardware import ( DeviceType, ) import utils.hardware.hardware as _hw_module -from utils.utils import format_error_message +from utils.utils import format_error_message, is_hf_authentication_error # ========== Helpers ========== @@ -439,6 +439,20 @@ class TestFormatErrorMessage: msg = format_error_message(err, "any/model") assert "invalid" in msg.lower() + def test_hf_authentication_error_follows_wrapped_401(self): + response = type("Response", (), {"status_code": 401})() + auth_error = Exception("request failed") + auth_error.response = response + wrapper = RuntimeError("model validation failed") + wrapper.__cause__ = auth_error + assert is_hf_authentication_error(wrapper) is True + + def test_hf_authentication_error_does_not_treat_429_as_invalid(self): + response = type("Response", (), {"status_code": 429})() + rate_error = Exception("too many requests") + rate_error.response = response + assert is_hf_authentication_error(rate_error) is False + # --- OOM on CUDA --- @needs_torch diff --git a/studio/backend/utils/hardware/hardware.py b/studio/backend/utils/hardware/hardware.py index adc9a54aab..9fef53e65e 100644 --- a/studio/backend/utils/hardware/hardware.py +++ b/studio/backend/utils/hardware/hardware.py @@ -538,21 +538,31 @@ def _torch_get_physical_gpu_count() -> Optional[int]: def _torch_get_per_device_info(device_indices: list[int]) -> list[Dict[str, Any]]: - """Query torch for per-GPU name, total VRAM, and used VRAM.""" + """Query torch for per-GPU name, total VRAM, and used VRAM. + + ``used_gb`` is ``None`` on Windows ROCm when ``hipMemGetInfo`` reports + ``free == total`` (ROCm/ROCm#1909): that 0 means unknown, not empty. + """ mod, _ = _torch_get_device_module() if mod is None: return [] + # free==total is a Windows-ROCm-only quirk. + _win_rocm = sys.platform == "win32" and IS_ROCM devices = [] for ordinal, phys_idx in enumerate(device_indices): try: # torch ordinals are 0-based relative to CUDA_VISIBLE_DEVICES. props = mod.get_device_properties(ordinal) total_bytes = props.total_memory + used_bytes: Optional[int] # Prefer mem_get_info (system-wide) so auto-select sees other consumers. if hasattr(mod, "mem_get_info"): free_bytes, total_bytes = mod.mem_get_info(ordinal) used_bytes = total_bytes - free_bytes + # free==total is the broken-API sentinel, not an idle GPU. + if _win_rocm and free_bytes == total_bytes: + used_bytes = None else: used_bytes = mod.memory_allocated(ordinal) devices.append( @@ -561,7 +571,7 @@ def _torch_get_per_device_info(device_indices: list[int]) -> list[Dict[str, Any] "visible_ordinal": ordinal, "name": props.name, "total_gb": round(total_bytes / (1024**3), 2), - "used_gb": round(used_bytes / (1024**3), 2), + "used_gb": round(used_bytes / (1024**3), 2) if used_bytes is not None else None, } ) except Exception as e: @@ -724,20 +734,30 @@ def _rocm_linux_sysfs_vram_gb() -> tuple[Optional[float], Optional[float]]: return None, None -def _rocm_windows_perf_counter_vram_gb() -> tuple[Optional[float], Optional[float]]: - """Query system-wide dedicated GPU VRAM via Windows Performance Counters. +# ── Windows AMD/ROCm per-adapter VRAM (issue #7072) ────────────────────────── +# amd-smi is disabled and hipMemGetInfo reports free==total, so read used from the +# per-LUID "GPU Adapter Memory" perf counters and take each total from torch, so +# every GPU shows instead of one fake device with GPU 0's total. +# Placeholder adapters (Basic Render Driver / idle iGPU) drop only when they would +# outnumber the real torch devices. +_ROCM_WIN_ADAPTER_MIN_BYTES = 64 * 1024 * 1024 # 64 MiB - Same data source as Task Manager, so cross-process usage is accurate. - Works for any GPU vendor without amd-smi or nvidia-smi. - Returns (used_gb, total_gb) or (None, None) on failure. + +def _rocm_windows_perf_counter_vram_by_adapter() -> Optional[list[tuple[str, float]]]: + """Per-adapter dedicated VRAM usage on Windows via Performance Counters. + + Returns ``[(instance_name, used_bytes)]`` (one per LUID-named adapter), or + ``None`` when the counter is unavailable/localized/empty so callers fall back. """ if platform.system() != "Windows": - return None, None + return None try: + # Emit "|" per sample, or a __NONE__ sentinel. ps = ( "$s=(Get-Counter '\\GPU Adapter Memory(*)\\Dedicated Usage'" " -ErrorAction SilentlyContinue).CounterSamples;" - "if($s){($s|Measure-Object CookedValue -Sum).Sum}else{-1}" + "if($s){$s|ForEach-Object{'{0}|{1}' -f $_.InstanceName,[int64]$_.CookedValue}}" + "else{'__NONE__'}" ) r = subprocess.run( ["powershell", "-NoProfile", "-NonInteractive", "-Command", ps], @@ -746,16 +766,167 @@ def _rocm_windows_perf_counter_vram_gb() -> tuple[Optional[float], Optional[floa timeout = 5, ) if r.returncode != 0 or not r.stdout.strip(): - return None, None - used_bytes = float(r.stdout.strip()) - if used_bytes < 0: - return None, None - import torch as _torch - - total_bytes = _torch.cuda.get_device_properties(0).total_memory - return round(used_bytes / (1024**3), 2), round(total_bytes / (1024**3), 2) + return None + adapters: list[tuple[str, float]] = [] + for line in r.stdout.splitlines(): + line = line.strip() + if not line or line == "__NONE__" or "|" not in line: + continue + instance, _, raw = line.rpartition("|") + try: + used = float(raw.strip()) + except (ValueError, TypeError): + continue + if used < 0: + continue + adapters.append((instance.strip(), used)) + return adapters or None except Exception: - return None, None + return None + + +def _match_adapter_used_to_devices( + adapter_useds: list[float], device_totals: list[float] +) -> list[Optional[float]]: + """Attribute per-adapter used bytes to torch devices by capacity ranking. + + Windows shares no key between LUID counters and torch ordinals, so usages are + ranked against device totals and each is trusted only when capacity *forces* it + (it exceeds every smaller device); an ambiguous ranking reports unknown + (``None``) rather than fabricate a per-index free. + + Extra counters mean a hidden/display adapter, and the noise filter may have + dropped a real reading, so values are emitted only when the supra-threshold + counters number EXACTLY the visible devices AND capacity forces the mapping; + otherwise every device is unknown. Best-effort but correct for the common + loaded-card case (#7072). Returns a list aligned to ``device_totals``. + """ + n = len(device_totals) + if n == 0: + return [] + useds = sorted(adapter_useds, reverse = True) + ranked_positions = sorted(range(n), key = lambda i: -device_totals[i]) + ranked_totals = [device_totals[pos] for pos in ranked_positions] + assigned: list[Optional[float]] + # More counters than devices -> a hidden/display adapter (check before noise filter). + if len(useds) > n: + non_trivial = [u for u in useds if u >= _ROCM_WIN_ADAPTER_MIN_BYTES] + if len(non_trivial) != n: + # Not a clean bijection (a masked GPU is busy or a visible card idle): + # no counter maps to a specific card, so report unknown. + return [None] * n + # Exactly n supra-threshold counters: extras were placeholders, so a + # capacity-ranked bijection is plausible. + useds = non_trivial + ranked_useds = [useds[rank] for rank in range(n)] + # A usage above its ranked capacity is a hidden larger GPU; clamping onto the + # smaller card would fabricate a fully-used reading. + for rank in range(n): + if ranked_useds[rank] > ranked_totals[rank]: + return [None] * n + # Capacity forces the mapping only when the usage exceeds the next-smaller + # capacity; the smallest card and merely-fitting usages stay unknown. + # Keeps 40 GiB over 48/8 GiB -> [40, None]. + assigned = [None] * n + for rank, pos in enumerate(ranked_positions): + if rank + 1 < n and ranked_useds[rank] > ranked_totals[rank + 1]: + assigned[pos] = min(ranked_useds[rank], device_totals[pos]) + return assigned + # No hidden adapters: every counter is a visible card, so ranking is a permutation. + ranked_useds = [useds[rank] if rank < len(useds) else 0.0 for rank in range(n)] + # Ambiguous if a strictly larger usage also fits the next smaller card: the two + # could be swapped without breaking capacity, so ranking can't tell them apart. + for rank in range(n - 1): + upper, lower = ranked_useds[rank], ranked_useds[rank + 1] + if upper > lower and upper <= ranked_totals[rank + 1]: + return [None] * n + assigned = [None] * n + for rank, pos in enumerate(ranked_positions): + if rank < len(useds): + assigned[pos] = min(useds[rank], device_totals[pos]) + return assigned + + +def _rocm_windows_per_device_vram(device_indices: list[int]) -> list[Dict[str, Any]]: + """Per-GPU VRAM on Windows AMD/ROCm: total from torch properties (reliable), + used from the per-adapter Dedicated Usage counter. + + Returns ``{index, visible_ordinal, name, used_gb, total_gb}`` per visible GPU + (``used_gb`` may be ``None`` when the counter is unavailable), or ``[]`` when + torch can't enumerate devices so callers fall through to the torch last resort. + """ + if platform.system() != "Windows": + return [] + mod, _ = _torch_get_device_module() + if mod is None: + return [] + # Totals/names from torch properties (mem_get_info's free==total quirk zeroes used). + dev_meta: list[Dict[str, Any]] = [] + for ordinal, phys_idx in enumerate(device_indices): + try: + props = mod.get_device_properties(ordinal) + dev_meta.append( + { + "index": phys_idx, + "visible_ordinal": ordinal, + "name": props.name, + "total_bytes": int(props.total_memory), + } + ) + except Exception as e: + logger.debug("torch property probe failed for ordinal %d: %s", ordinal, e) + if not dev_meta: + return [] + + adapters = _rocm_windows_perf_counter_vram_by_adapter() + if adapters: + assigned = _match_adapter_used_to_devices( + [used for _, used in adapters], + [d["total_bytes"] for d in dev_meta], + ) + else: + # Counter unavailable: show every GPU with a correct total, used unknown. + assigned = [None] * len(dev_meta) + + devices: list[Dict[str, Any]] = [] + for meta, used_bytes in zip(dev_meta, assigned): + total_gb = round(meta["total_bytes"] / (1024**3), 2) + used_gb = round(used_bytes / (1024**3), 2) if used_bytes is not None else None + devices.append( + { + "index": meta["index"], + "visible_ordinal": meta["visible_ordinal"], + "name": meta["name"], + "used_gb": used_gb, + "total_gb": total_gb, + } + ) + return devices + + +def _rocm_windows_device_payload_entry( + device: DeviceType, dev: Dict[str, Any], gpu_util_pct: Optional[float] +) -> Dict[str, Any]: + """Build a ``get_gpu_utilization`` device entry from a per-device VRAM dict.""" + total_gb = dev["total_gb"] + used_gb = dev["used_gb"] + return { + "available": True, + "backend": _backend_label(device), + "index": dev["index"], + "visible_ordinal": dev["visible_ordinal"], + "name": dev.get("name", "Unknown"), + "gpu_utilization_pct": gpu_util_pct, + "temperature_c": None, + "vram_used_gb": used_gb, + "vram_total_gb": total_gb, + "vram_utilization_pct": round((used_gb / total_gb) * 100, 1) + if total_gb and total_gb > 0 and used_gb is not None + else None, + "power_draw_w": None, + "power_limit_w": None, + "power_utilization_pct": None, + } def _gpu_utilization_payload( @@ -821,30 +992,24 @@ def get_gpu_utilization() -> Dict[str, Any]: index_kind = result.get("index_kind"), ) - # Fallback Windows ROCm + # Fallback Windows ROCm: per-adapter VRAM attribution (issue #7072), so + # every visible GPU is shown instead of a sum collapsed onto one device. if IS_ROCM and platform.system() == "Windows": - _win_used, _win_total = _rocm_windows_perf_counter_vram_gb() - if _win_used is not None and _win_total is not None: - _win_util = _rocm_windows_perf_counter_gpu_util_pct() + _win_ids = _get_parent_visible_gpu_spec().get("numeric_ids") + if not _win_ids: + _win_ids = list(range(_torch_get_physical_gpu_count() or 0)) + _win_devices = _rocm_windows_per_device_vram(_win_ids) + if _win_devices: + # A single visible GPU can own the aggregate 3D-engine utilization; + # across several GPUs the sum isn't per-device, so leave it unset. + _win_util = ( + _rocm_windows_perf_counter_gpu_util_pct() if len(_win_devices) == 1 else None + ) return _gpu_utilization_payload( device, [ - { - "available": True, - "backend": _backend_label(device), - "index": 0, - "visible_ordinal": 0, - "gpu_utilization_pct": _win_util, - "temperature_c": None, - "vram_used_gb": _win_used, - "vram_total_gb": _win_total, - "vram_utilization_pct": round((_win_used / _win_total) * 100, 1) - if _win_total > 0 - else None, - "power_draw_w": None, - "power_limit_w": None, - "power_utilization_pct": None, - } + _rocm_windows_device_payload_entry(device, _wd, _win_util) + for _wd in _win_devices ], ) @@ -901,7 +1066,7 @@ def get_gpu_utilization() -> Dict[str, Any]: "vram_used_gb": _used, "vram_total_gb": _total, "vram_utilization_pct": round((_used / _total) * 100, 1) - if _total > 0 + if _total > 0 and _used is not None else None, "power_draw_w": None, "power_limit_w": None, @@ -995,19 +1160,27 @@ def _apply_unified_memory_correction( endpoints stay in sync on AMD iGPUs with unified memory. """ torch_total_gb = torch_info["total_gb"] + torch_used_gb = torch_info.get("used_gb") smi_total_gb = device_metrics.get("vram_total_gb") or 0.0 + # torch sees the full unified (GTT) pool; amd-smi only the dedicated carve-out. + # Adopt torch's larger total regardless of used: on Windows ROCm torch_used is + # None (free==total sentinel) but its total stays authoritative. Overwrite used + # only when torch's is known, then recompute utilization against whatever remains. if torch_total_gb > smi_total_gb: - torch_used_gb = torch_info["used_gb"] device_metrics["vram_total_gb"] = torch_total_gb - device_metrics["vram_used_gb"] = torch_used_gb + if torch_used_gb is not None: + device_metrics["vram_used_gb"] = torch_used_gb + _used_for_pct = device_metrics.get("vram_used_gb") device_metrics["vram_utilization_pct"] = ( - round((torch_used_gb / torch_total_gb) * 100, 1) if torch_total_gb > 0 else None + round((_used_for_pct / torch_total_gb) * 100, 1) + if torch_total_gb > 0 and _used_for_pct is not None + else None ) logger.debug( - "ROCm unified memory: replaced amd-smi VRAM (%.2f GB) with " - "torch mem_get_info total (%.2f GB) for device %s", - smi_total_gb, + "ROCm unified memory: adopted torch mem_get_info total (%.2f GB) over " + "amd-smi (%.2f GB) for device %s", torch_total_gb, + smi_total_gb, torch_info.get("index"), ) @@ -1067,6 +1240,49 @@ def get_visible_gpu_utilization() -> Dict[str, Any]: _reconcile_rocm_unified_memory(result, numeric_ids) return result + # Windows AMD/ROCm (issue #7072): the System tab's VRAM source. The torch + # fallback below would report used==0 (free==total), so read per-adapter + # Dedicated Usage instead; total from torch properties. + if IS_ROCM and platform.system() == "Windows": + win_numeric_ids = parent_visible_spec.get("numeric_ids") + if win_numeric_ids: + win_ids = win_numeric_ids + win_index_kind = "physical" + else: + win_ids = list(range(_torch_get_physical_gpu_count() or 0)) + win_index_kind = "relative" + win_devices = _rocm_windows_per_device_vram(win_ids) + if win_devices: + devices = [] + for wd in win_devices: + total = wd["total_gb"] + used = wd["used_gb"] + devices.append( + { + "index": wd["index"], + "index_kind": win_index_kind, + "visible_ordinal": wd["visible_ordinal"], + "name": wd.get("name"), + "gpu_utilization_pct": None, + "temperature_c": None, + "vram_used_gb": used, + "vram_total_gb": total, + "vram_utilization_pct": round((used / total) * 100, 1) + if total and total > 0 and used is not None + else None, + "power_draw_w": None, + "power_limit_w": None, + "power_utilization_pct": None, + } + ) + return { + "available": True, + "backend": _backend_label(device), + "parent_visible_gpu_ids": win_numeric_ids or [], + "devices": devices, + "index_kind": win_index_kind, + } + # Torch-based fallback for CUDA (nvidia-smi unavailable, AMD ROCm) and XPU (Intel) if device in (DeviceType.CUDA, DeviceType.XPU): parent_ids = get_parent_visible_gpu_ids() @@ -1094,7 +1310,7 @@ def get_visible_gpu_utilization() -> Dict[str, Any]: "vram_used_gb": used, "vram_total_gb": total, "vram_utilization_pct": round((used / total) * 100, 1) - if total > 0 + if total > 0 and used is not None else None, "power_draw_w": None, "power_limit_w": None, diff --git a/studio/backend/utils/hf_token_validation.py b/studio/backend/utils/hf_token_validation.py new file mode 100644 index 0000000000..7247c6e756 --- /dev/null +++ b/studio/backend/utils/hf_token_validation.py @@ -0,0 +1,208 @@ +# SPDX-License-Identifier: AGPL-3.0-only +# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0 + +"""Cached, rate-limited Hugging Face token validation.""" + +from __future__ import annotations + +import hashlib +import threading +import time +from collections import deque +from dataclasses import dataclass +from typing import Literal + +from huggingface_hub import HfApi +from huggingface_hub.utils import build_hf_headers, get_session + + +TokenValidationStatus = Literal["valid", "invalid", "rate_limited", "unavailable"] + + +@dataclass(frozen = True) +class TokenValidationResult: + status: TokenValidationStatus + retry_after_seconds: int | None = None + + +_WINDOW_SECONDS = 3600.0 +_MAX_ATTEMPTS = 3 +_CACHE_TTL_SECONDS = 3600.0 +_TEMPORARY_CACHE_TTL_SECONDS = 15.0 +_MAX_BUCKETS = 4096 +_MAX_CACHE_ENTRIES = 4096 +_INFLIGHT_WAIT_SECONDS = 30.0 +_REMOTE_TIMEOUT_SECONDS = 10.0 + +_attempts: dict[str, deque[float]] = {} +_cache: dict[str, tuple[float, TokenValidationResult]] = {} +_inflight: dict[str, threading.Event] = {} +_lock = threading.Lock() + + +def _fingerprint(token: str) -> str: + return hashlib.sha256(token.encode("utf-8")).hexdigest() + + +def _prune_attempts(bucket: deque[float], now: float) -> None: + while bucket and now - bucket[0] >= _WINDOW_SECONDS: + bucket.popleft() + + +def _prune_locked(now: float) -> None: + for key in list(_attempts): + bucket = _attempts[key] + _prune_attempts(bucket, now) + if not bucket: + del _attempts[key] + for key, (expires_at, _result) in list(_cache.items()): + if expires_at <= now: + del _cache[key] + + +def _cached_locked(fingerprint: str, now: float) -> TokenValidationResult | None: + cached = _cache.get(fingerprint) + if cached is None: + return None + expires_at, result = cached + if expires_at <= now: + del _cache[fingerprint] + return None + return result + + +def _retry_after(bucket: deque[float], now: float) -> int: + return max(1, int(_WINDOW_SECONDS - (now - bucket[0])) + 1) + + +def _reserve_attempt_locked(rate_key: str, now: float) -> TokenValidationResult | None: + bucket = _attempts.get(rate_key) + if bucket is None: + if len(_attempts) >= _MAX_BUCKETS: + _prune_locked(now) + if len(_attempts) >= _MAX_BUCKETS: + return TokenValidationResult( + status = "rate_limited", + retry_after_seconds = max(1, int(_WINDOW_SECONDS)), + ) + bucket = _attempts[rate_key] = deque() + _prune_attempts(bucket, now) + if len(bucket) >= _MAX_ATTEMPTS: + return TokenValidationResult( + status = "rate_limited", + retry_after_seconds = _retry_after(bucket, now), + ) + bucket.append(now) + return None + + +def _http_status(response: object | None) -> int | None: + status = getattr(response, "status_code", None) + try: + return int(status) if status is not None else None + except (TypeError, ValueError): + return None + + +def _remote_retry_after(response: object | None) -> int | None: + headers = getattr(response, "headers", None) + if not headers: + return None + raw = headers.get("Retry-After") + try: + return max(1, int(float(raw))) if raw is not None else None + except (TypeError, ValueError): + return None + + +def _classify_response(response: object | None) -> TokenValidationResult: + status = _http_status(response) + if status is not None and 200 <= status < 300: + return TokenValidationResult(status = "valid") + if status == 401: + return TokenValidationResult(status = "invalid") + if status == 429: + return TokenValidationResult( + status = "rate_limited", + retry_after_seconds = _remote_retry_after(response), + ) + return TokenValidationResult(status = "unavailable") + + +def _check_remote(token: str) -> TokenValidationResult: + api = HfApi() + try: + # HfApi.whoami has no timeout parameter in the pinned Hub client. + # Use its session and headers against the same whoami endpoint. + response = get_session().get( + f"{api.endpoint}/api/whoami-v2", + headers = build_hf_headers(token = token), + timeout = _REMOTE_TIMEOUT_SECONDS, + ) + except Exception as exc: + # huggingface-hub 0.36.x can wrap a 401 as requests.HTTPError. + return _classify_response(getattr(exc, "response", None)) + return _classify_response(response) + + +def validate_hf_token(token: str, *, rate_key: str) -> TokenValidationResult: + """Validate ``token`` without retaining it, sharing results across callers. + + Cached checks do not consume the caller's three-per-hour network budget. A + single-flight event also prevents simultaneously mounted UI surfaces from + sending duplicate ``whoami`` requests for the same token. + """ + normalized = token.strip() + if not normalized: + return TokenValidationResult(status = "invalid") + token_fingerprint = _fingerprint(normalized) + owner_event: threading.Event | None = None + + try: + while True: + now = time.monotonic() + with _lock: + cached = _cached_locked(token_fingerprint, now) + if cached is not None: + return cached + waiting = _inflight.get(token_fingerprint) + if waiting is None: + limited = _reserve_attempt_locked(rate_key, now) + if limited is not None: + return limited + owner_event = threading.Event() + _inflight[token_fingerprint] = owner_event + break + if not waiting.wait(_INFLIGHT_WAIT_SECONDS): + return TokenValidationResult(status = "unavailable") + + result = _check_remote(normalized) + now = time.monotonic() + ttl = ( + _CACHE_TTL_SECONDS + if result.status in ("valid", "invalid") + else max(_TEMPORARY_CACHE_TTL_SECONDS, float(result.retry_after_seconds or 0)) + ) + with _lock: + if len(_cache) >= _MAX_CACHE_ENTRIES: + _prune_locked(now) + if len(_cache) < _MAX_CACHE_ENTRIES: + _cache[token_fingerprint] = (now + ttl, result) + return result + finally: + if owner_event is not None: + with _lock: + event = _inflight.get(token_fingerprint) + if event is owner_event: + _inflight.pop(token_fingerprint, None) + event.set() + + +def reset_hf_token_validation_state() -> None: + """Clear process state for test isolation.""" + with _lock: + for event in _inflight.values(): + event.set() + _inflight.clear() + _attempts.clear() + _cache.clear() diff --git a/studio/backend/utils/llama_cpp_update.py b/studio/backend/utils/llama_cpp_update.py index 31dbda63ea..67733bde35 100644 --- a/studio/backend/utils/llama_cpp_update.py +++ b/studio/backend/utils/llama_cpp_update.py @@ -479,6 +479,7 @@ def _run_update( asset: Optional[str], script: Path, pin_release_tag: Optional[str] = None, + force_cpu: bool = False, ) -> None: """Worker: put the backend into a maintenance state, run the installer for the latest prebuilt, then refresh caches so the next load uses the new build. @@ -522,6 +523,12 @@ def _run_update( if pin_release_tag: cmd.extend(["--published-release-tag", pin_release_tag]) cmd.extend(_rocm_install_args(asset)) + # Re-assert a deliberate CPU install (--force-cpu) so detect_host on a GPU host + # does not re-route to a GPU/Vulkan bundle and revive the crash (#7213). --force-cpu + # (not --cpu-fallback) also re-persists force_cpu, keeping the choice across future + # updates. A natural fallback (or a legacy marker without the flag) heals to GPU (#6097). + if force_cpu: + cmd.append("--force-cpu") logger.info("llama update: installing", cmd = " ".join(cmd)) # Stream progress lines into job["progress"]. env = dict(os.environ, UNSLOTH_PROGRESS_PERCENT_STEP = "5") @@ -671,6 +678,7 @@ def start_update() -> dict: repo = marker.get("published_repo") or DEFAULT_PUBLISHED_REPO from_tag = marker.get("tag") or marker.get("release_tag") asset = marker.get("asset") + force_cpu = bool(marker.get("force_cpu")) # Install exactly the release the banner offered: the installer's own # "latest" is commit-date ordered and can lag the published_at pick # above, reinstalling the current build in a loop (the #6219 class). @@ -705,6 +713,8 @@ def start_update() -> dict: repo = (res or {}).get("repo") or DEFAULT_PUBLISHED_REPO from_tag = None asset = (res or {}).get("asset") + # Source builds carry no forced-CPU marker, so nothing to preserve here. + force_cpu = False # No pin: source-build detection resolves via --resolve-prebuilt latest, # the same resolver the unpinned apply uses, so the two already agree. pin_release_tag = None @@ -735,7 +745,7 @@ def start_update() -> dict: thread = threading.Thread( target = _run_update, - args = (install_dir, repo, asset, script, pin_release_tag), + args = (install_dir, repo, asset, script, pin_release_tag, force_cpu), name = "llama-cpp-update", daemon = True, ) diff --git a/studio/backend/utils/models/gguf_metadata.py b/studio/backend/utils/models/gguf_metadata.py index c24ec28e1d..749f2c9234 100644 --- a/studio/backend/utils/models/gguf_metadata.py +++ b/studio/backend/utils/models/gguf_metadata.py @@ -50,9 +50,15 @@ _CACHE_MAX_ENTRIES = 4096 # keyed by (file cache key, wanted key). None = key absent / file unreadable. _BOOL_CACHE: Dict[Tuple[_CacheKey, str], Optional[bool]] = {} -# Native training context length (``{arch}.context_length``). None = absent / -# unreadable. Lets the UI show the real context ceiling before a model loads. -_CONTEXT_CACHE: Dict[_CacheKey, Optional[int]] = {} +_STRING_CACHE: Dict[Tuple[_CacheKey, str], Optional[str]] = {} + +# GGUF header dims for the staged/deferred-load UI: context_length, layer_count +# (block_count), and moe_layer_count (block_count minus leading dense layers; 0 +# if not MoE). One cached pass fills all three so the staged sheet can size every +# slider before the model loads. None = unreadable / not a GGUF. The native +# training context length (``{arch}.context_length``) the UI shows before a model +# loads is read from here via read_gguf_context_length. +_DIMS_CACHE: Dict[_CacheKey, Optional[Dict[str, Optional[int]]]] = {} def _cache_key(path: str) -> Optional[_CacheKey]: @@ -142,32 +148,45 @@ def _parse_gguf_header(path: str) -> Optional[Dict[str, str]]: return out -def read_gguf_context_length(path: str) -> Optional[int]: - """Return the GGUF's native training context length (``{arch}.context_length``), - or ``None`` if missing/unreadable/not a GGUF. Cached by (path, mtime, size). - Lets the UI populate the context slider before the model is loaded.""" +def read_gguf_staged_dims(path: str) -> Optional[Dict[str, Optional[int]]]: + """GGUF header dims for the staged-load UI in one cached pass: + ``{"context_length", "layer_count", "moe_layer_count"}``. Each may be None + when absent (moe_layer_count is 0 for a dense model). Returns ``None`` if not + a GGUF / unreadable. Cached by (path, mtime, size). Lets the staged sheet size + the context, GPU-layers and MoE sliders before the model loads.""" key = _cache_key(path) if key is None: return None with _CACHE_LOCK: - if key in _CONTEXT_CACHE: - return _CONTEXT_CACHE[key] - result = _parse_gguf_context_length(path) + if key in _DIMS_CACHE: + return _DIMS_CACHE[key] + result = _parse_gguf_staged_dims(path) with _CACHE_LOCK: - while len(_CONTEXT_CACHE) >= _CACHE_MAX_ENTRIES: + while len(_DIMS_CACHE) >= _CACHE_MAX_ENTRIES: try: - _CONTEXT_CACHE.pop(next(iter(_CONTEXT_CACHE))) + _DIMS_CACHE.pop(next(iter(_DIMS_CACHE))) except StopIteration: break - _CONTEXT_CACHE[key] = result + _DIMS_CACHE[key] = result return result -def _parse_gguf_context_length(path: str) -> Optional[int]: - # The context key is architecture-namespaced (``llama.context_length`` etc.), - # so we learn the key only after reading ``general.architecture``. GGUF writes - # general.* before arch.* keys, matching the loader's own parser. - ctx_key: Optional[str] = None +def read_gguf_context_length(path: str) -> Optional[int]: + """Native training context length (``{arch}.context_length``), or ``None``. + Thin accessor over read_gguf_staged_dims.""" + dims = read_gguf_staged_dims(path) + return dims["context_length"] if dims else None + + +def _parse_gguf_arch_uints(path: str, wanted_suffixes: frozenset[str]) -> Optional[Dict[str, int]]: + """Walk a GGUF header once and return the requested architecture-namespaced + uint (vtype 4/10) keys, e.g. ``{"block_count": 32}``. Keys are + ``{arch}.``; the arch is learned from ``general.architecture`` (GGUF + writes general.* before arch.* keys, matching the loader's own parser). + Returns ``None`` if not a GGUF / unreadable, else a dict (possibly empty or + partial when some keys are absent).""" + arch: Optional[str] = None + found: Dict[str, int] = {} try: with open(path, "rb") as f: head = f.read(24) @@ -204,28 +223,68 @@ def _parse_gguf_context_length(path: str) -> Optional[int]: sbytes = f.read(slen) if len(sbytes) < slen: break - ctx_key = f"{sbytes.decode('utf-8', 'replace')}.context_length" - elif ctx_key is not None and key == ctx_key and vtype in (4, 10): + arch = sbytes.decode("utf-8", "replace") + elif ( + arch is not None + and vtype in (4, 10) + and key.startswith(f"{arch}.") + and key[len(arch) + 1 :] in wanted_suffixes + ): width = 4 if vtype == 4 else 8 n_bytes = f.read(width) if len(n_bytes) < width: break - value = struct.unpack(" 0 else None + found[key[len(arch) + 1 :]] = struct.unpack( + " Optional[Dict[str, Optional[int]]]: + vals = _parse_gguf_arch_uints( + path, + frozenset( + { + "context_length", + "block_count", + "expert_count", + "leading_dense_block_count", + } + ), + ) + if vals is None: + return None + ctx = vals.get("context_length") + block = vals.get("block_count") + # A real context/layer count is positive; treat 0/garbage as absent so the + # UI never builds a slider with max < min. + context_length = ctx if ctx and ctx > 0 else None + layer_count = block if block and block > 0 else None + # MoE layer count = block_count - leading dense layers, only when experts + # exist; else 0 (dense -> slider hidden). Mirrors n_moe_layers in + # core/inference/llama_cpp.py. + if not vals.get("expert_count") or not block: + moe_layer_count: Optional[int] = 0 + else: + moe_layer_count = max(0, block - (vals.get("leading_dense_block_count") or 0)) + return { + "context_length": context_length, + "layer_count": layer_count, + "moe_layer_count": moe_layer_count, + } # Strings (8) and arrays (9) are handled inline. @@ -353,6 +412,83 @@ def _read_gguf_bool(path: str, wanted_key: str) -> Optional[bool]: return result +def _parse_gguf_string(path: str, wanted_key: str) -> Optional[str]: + try: + with open(path, "rb") as f: + head = f.read(24) + if len(head) < 24: + return None + magic, _version, _tcount, kv_count = struct.unpack(" 1 << 20: + break + kbytes = f.read(klen) + if len(kbytes) < klen: + break + key = kbytes.decode("utf-8", "replace") + vt_bytes = f.read(4) + if len(vt_bytes) < 4: + break + vtype = struct.unpack(" 1 << 22: + break + sbytes = f.read(slen) + if len(sbytes) < slen: + break + return sbytes.decode("utf-8", "replace") + if not _skip_gguf_value(f, vtype): + break + except (struct.error, UnicodeDecodeError): + break + except OSError as e: + logger.debug(f"_parse_gguf_string: cannot open {path}: {e}") + return None + except Exception as e: + logger.debug(f"_parse_gguf_string: parse failure on {path}: {e}") + return None + return None + + +def _read_gguf_string(path: str, wanted_key: str) -> Optional[str]: + fkey = _cache_key(path) + if fkey is None: + return None + ckey = (fkey, wanted_key) + with _CACHE_LOCK: + if ckey in _STRING_CACHE: + return _STRING_CACHE[ckey] + result = _parse_gguf_string(path, wanted_key) + with _CACHE_LOCK: + while len(_STRING_CACHE) >= _CACHE_MAX_ENTRIES: + try: + _STRING_CACHE.pop(next(iter(_STRING_CACHE))) + except StopIteration: + break + _STRING_CACHE[ckey] = result + return result + + +def read_gguf_chat_template(path: str) -> Optional[str]: + template = _read_gguf_string(path, "tokenizer.chat_template") + if isinstance(template, str) and template.strip(): + return template + return None + + def read_mmproj_audio_capability(path: str) -> Optional[bool]: """``clip.has_audio_encoder`` from an mmproj GGUF (e.g. Gemma 4's gemma4ua): ``True``/``False`` if present, ``None`` if absent/unreadable. diff --git a/studio/backend/utils/models/model_config.py b/studio/backend/utils/models/model_config.py index dadf103cea..821529083d 100644 --- a/studio/backend/utils/models/model_config.py +++ b/studio/backend/utils/models/model_config.py @@ -1249,6 +1249,77 @@ def _iter_gguf_files(directory: Path, recursive: bool = False): yield f +_GGUF_SPLIT_FILE_RE = re.compile( + r"^(?P.+)-(?P\d{5})-of-(?P\d{5})\.gguf$", + re.IGNORECASE, +) + + +def _colocated_first_split_shard(path: Path) -> tuple[Optional[Path], bool]: + """Return shard 1 and whether every shard is beside *path*.""" + match = _GGUF_SPLIT_FILE_RE.match(path.name) + if match is None: + return None, False + + prefix = match.group("prefix").casefold() + total_text = match.group("total") + total = int(total_text) + if total < 1: + return None, False + + first: Optional[Path] = None + indices: set[int] = set() + try: + siblings = path.parent.iterdir() + for sibling in siblings: + sibling_match = _GGUF_SPLIT_FILE_RE.match(sibling.name) + if ( + sibling_match is None + or sibling_match.group("prefix").casefold() != prefix + or sibling_match.group("total") != total_text + ): + continue + try: + if not sibling.is_file(): + continue + except OSError: + continue + index = int(sibling_match.group("index")) + if not 1 <= index <= total: + continue + indices.add(index) + if index == 1: + first = sibling + except OSError: + return None, False + + return first, first is not None and len(indices) == total + + +def _local_gguf_load_path(path: Path) -> Path: + """Choose a loadable local path while preserving complete symlink sets.""" + if _GGUF_SPLIT_FILE_RE.match(path.name) is None: + return path.absolute() + + first, complete = _colocated_first_split_shard(path) + if complete and first is not None: + return first.absolute() + + try: + is_symlink = path.is_symlink() + except OSError: + is_symlink = False + if is_symlink: + try: + target = path.resolve() + except OSError: + return (first or path).absolute() + target_first, _ = _colocated_first_split_shard(target) + return (target_first or target).absolute() + + return (first or path).absolute() + + def detect_mmproj_file(path: str, search_root: Optional[str] = None) -> Optional[str]: """Find the mmproj GGUF for a model. @@ -1434,7 +1505,7 @@ def detect_gguf_model(path: str) -> Optional[str]: except OSError: is_dir = False # stat() unavailable in the lock window if not is_dir: - return str(p.absolute()) # absolute() keeps symlink names readable + return str(_local_gguf_load_path(p)) # Directory named "*.gguf": fall through to the dir scan below. # Case 2: directory containing .gguf files (skip mmproj / MTP drafter) @@ -1452,7 +1523,7 @@ def detect_gguf_model(path: str) -> Optional[str]: gguf_files.append(f) gguf_files.sort(key = lambda f: f.stat().st_size, reverse = True) if gguf_files: - return str(gguf_files[0].resolve()) + return str(_local_gguf_load_path(gguf_files[0])) return None @@ -1879,7 +1950,7 @@ def _find_local_gguf_by_variant(directory: str, variant: str) -> Optional[str]: For sharded GGUFs (multiple files sharing a quant label), returns the first shard (sorted by name), which is what ``llama-server -m`` expects. - Returns the resolved absolute path, or ``None`` if no match. + Returns the absolute path, or ``None`` if no match. """ p = _resolve_gguf_dir(Path(directory)) if p is None: @@ -1900,7 +1971,7 @@ def _find_local_gguf_by_variant(directory: str, variant: str) -> Optional[str]: matches.append(f) matches.sort() if matches: - return str(matches[0].resolve()) + return str(_local_gguf_load_path(matches[0])) return None diff --git a/studio/backend/utils/openai_auto_switch_settings.py b/studio/backend/utils/openai_auto_switch_settings.py index 462435e5d5..7007440f4c 100644 --- a/studio/backend/utils/openai_auto_switch_settings.py +++ b/studio/backend/utils/openai_auto_switch_settings.py @@ -30,11 +30,13 @@ from typing import Any, Optional OPENAI_AUTO_SWITCH_SETTING_KEY = "openai_api_auto_switch_model" AUTO_UNLOAD_IDLE_SETTING_KEY = "openai_api_auto_unload_idle_seconds" +AUTO_UNLOAD_KEEP_KV_SETTING_KEY = "openai_api_auto_unload_keep_kv" MODEL_OVERRIDES_SETTING_KEY = "openai_api_auto_switch_overrides" MODEL_IDLE_TTL_ENV_VAR = "UNSLOTH_MODEL_IDLE_TTL" DEFAULT_OPENAI_AUTO_SWITCH_ENABLED = False DEFAULT_AUTO_UNLOAD_IDLE_SECONDS = 0 +DEFAULT_AUTO_UNLOAD_KEEP_KV = True MIN_AUTO_UNLOAD_IDLE_SECONDS = 60 _CACHE_TTL_S = 2.0 @@ -158,29 +160,54 @@ def get_auto_unload_idle_seconds() -> int: return env if env is not None else 0 -def set_openai_auto_switch(enabled: Any, idle_seconds: Any) -> tuple[bool, int]: - """Set both auto-switch flags in one transaction so a settings PUT can't leave - one key updated and the other stale. Both values are coerced before any write, - so an invalid value raises without persisting either.""" +def get_auto_unload_keep_kv() -> bool: + """Whether the idle unload persists slot KV to disk for restore on reload.""" + parsed = _coerce_bool(_cached_setting(AUTO_UNLOAD_KEEP_KV_SETTING_KEY, None)) + return parsed if parsed is not None else DEFAULT_AUTO_UNLOAD_KEEP_KV + + +def set_openai_auto_switch( + enabled: Any, + idle_seconds: Any, + keep_kv: Any = None, +) -> tuple[bool, int, bool]: + """One-transaction write; ``None`` leaves a stored value untouched.""" parsed_enabled = _coerce_bool(enabled) if parsed_enabled is None: raise ValueError("OpenAI auto-switch must be true or false.") - parsed_idle = _coerce_int(idle_seconds) - if parsed_idle is None: - raise ValueError("Auto-unload idle seconds must be a non-negative integer.") - if 0 < parsed_idle < MIN_AUTO_UNLOAD_IDLE_SECONDS: - raise ValueError( - f"Auto-unload idle seconds must be 0 (off) or at least " - f"{MIN_AUTO_UNLOAD_IDLE_SECONDS}." - ) + parsed_idle = None + if idle_seconds is not None: + parsed_idle = _coerce_int(idle_seconds) + if parsed_idle is None: + raise ValueError("Auto-unload idle seconds must be a non-negative integer.") + if 0 < parsed_idle < MIN_AUTO_UNLOAD_IDLE_SECONDS: + raise ValueError( + f"Auto-unload idle seconds must be 0 (off) or at least " + f"{MIN_AUTO_UNLOAD_IDLE_SECONDS}." + ) + parsed_keep_kv = None + if keep_kv is not None: + parsed_keep_kv = _coerce_bool(keep_kv) + if parsed_keep_kv is None: + raise ValueError("Keep KV on idle unload must be true or false.") from storage.studio_db import upsert_app_settings - upsert_app_settings( - {OPENAI_AUTO_SWITCH_SETTING_KEY: parsed_enabled, AUTO_UNLOAD_IDLE_SETTING_KEY: parsed_idle} - ) + updates: dict[str, Any] = {OPENAI_AUTO_SWITCH_SETTING_KEY: parsed_enabled} + if parsed_idle is not None: + updates[AUTO_UNLOAD_IDLE_SETTING_KEY] = parsed_idle + if parsed_keep_kv is not None: + updates[AUTO_UNLOAD_KEEP_KV_SETTING_KEY] = parsed_keep_kv + upsert_app_settings(updates) _invalidate(OPENAI_AUTO_SWITCH_SETTING_KEY) - _invalidate(AUTO_UNLOAD_IDLE_SETTING_KEY) - return parsed_enabled, parsed_idle + if parsed_idle is not None: + _invalidate(AUTO_UNLOAD_IDLE_SETTING_KEY) + if parsed_keep_kv is not None: + _invalidate(AUTO_UNLOAD_KEEP_KV_SETTING_KEY) + return ( + parsed_enabled, + parsed_idle if parsed_idle is not None else get_stored_auto_unload_idle_seconds(), + parsed_keep_kv if parsed_keep_kv is not None else get_auto_unload_keep_kv(), + ) def get_model_overrides() -> dict[str, dict]: diff --git a/studio/backend/utils/paths/storage_roots.py b/studio/backend/utils/paths/storage_roots.py index 1faa2b1281..35b8c57e9b 100644 --- a/studio/backend/utils/paths/storage_roots.py +++ b/studio/backend/utils/paths/storage_roots.py @@ -61,6 +61,11 @@ def cache_root() -> Path: return studio_root() / "cache" +def llama_slot_cache_root() -> Path: + """Dir llama-server saves/restores slot KV state in across idle unloads.""" + return cache_root() / "llama-slots" + + def studio_bin_root() -> Path: """Dir for Unsloth-managed executables (the `unsloth` shim, downloaded tools like cloudflared).""" return studio_root() / "bin" diff --git a/studio/backend/utils/utils.py b/studio/backend/utils/utils.py index 3818253ac9..31f5f31bee 100644 --- a/studio/backend/utils/utils.py +++ b/studio/backend/utils/utils.py @@ -123,6 +123,26 @@ def without_hf_auth(): os.environ.pop("HF_HUB_DISABLE_IMPLICIT_TOKEN", None) +def is_hf_authentication_error(error: Exception) -> bool: + """Return whether an exception chain contains a definitive HF auth failure.""" + seen: set[int] = set() + current: BaseException | None = error + while current is not None and id(current) not in seen: + seen.add(id(current)) + response = getattr(current, "response", None) + status = getattr(response, "status_code", None) + try: + if status is not None and int(status) == 401: + return True + except (TypeError, ValueError): + pass + message = str(current).lower() + if "invalid user token" in message or "invalid hf token" in message: + return True + current = current.__cause__ or current.__context__ + return False + + def format_error_message(error: Exception, model_name: str) -> str: """ Format a user-friendly error message for common load issues. diff --git a/studio/frontend/src/app/auth-guards.ts b/studio/frontend/src/app/auth-guards.ts index 6849f380b8..a3523ac580 100644 --- a/studio/frontend/src/app/auth-guards.ts +++ b/studio/frontend/src/app/auth-guards.ts @@ -23,19 +23,47 @@ interface AuthStatus { requires_password_change: boolean; } +const AUTH_STATUS_TTL_MS = 30_000; +let authStatusCheckedAt = 0; +let authStatusRequest: Promise | null = null; + +function hasFreshAuthStatus(): boolean { + return ( + authStatusCheckedAt !== 0 && + Date.now() - authStatusCheckedAt < AUTH_STATUS_TTL_MS + ); +} + async function fetchAuthStatus(): Promise { - try { - const res = await fetch(apiUrl("/api/auth/status")); - if (!res.ok) return { initialized: true, requires_password_change: mustChangePassword() }; - const status = (await res.json()) as AuthStatus; - // Server truth wins; keep localStorage in sync both ways. - if (status.requires_password_change !== mustChangePassword()) { - setMustChangePassword(status.requires_password_change); + if (authStatusRequest) return authStatusRequest; + + const request = (async () => { + try { + const res = await fetch(apiUrl("/api/auth/status")); + if (!res.ok) { + return { + initialized: true, + requires_password_change: mustChangePassword(), + }; + } + const status = (await res.json()) as AuthStatus; + authStatusCheckedAt = Date.now(); + // Server truth wins; keep localStorage in sync both ways. + if (status.requires_password_change !== mustChangePassword()) { + setMustChangePassword(status.requires_password_change); + } + return status; + } catch { + return { + initialized: true, + requires_password_change: mustChangePassword(), + }; } - return status; - } catch { - return { initialized: true, requires_password_change: mustChangePassword() }; - } + })().finally(() => { + authStatusRequest = null; + }); + authStatusRequest = request; + return request; } function authRedirect(to: "/login" | "/change-password"): never { @@ -49,12 +77,17 @@ export async function requireAuth(): Promise { } if (await hasActiveSession()) { - const { requires_password_change } = await fetchAuthStatus(); - if (requires_password_change || mustChangePassword()) { - authRedirect("/change-password"); + // Reconcile periodically so local-only routes cannot outlive a server-side + // password-change requirement, while nearby route switches stay local. + if (mustChangePassword() || !hasFreshAuthStatus()) { + const { requires_password_change } = await fetchAuthStatus(); + if (requires_password_change || mustChangePassword()) { + authRedirect("/change-password"); + } } return; } + const status = await fetchAuthStatus(); if (status.requires_password_change || mustChangePassword()) { authRedirect("/change-password"); diff --git a/studio/frontend/src/app/provider.tsx b/studio/frontend/src/app/provider.tsx index a7e9469cfc..e6c89b9cd7 100644 --- a/studio/frontend/src/app/provider.tsx +++ b/studio/frontend/src/app/provider.tsx @@ -51,9 +51,12 @@ async function showSetupWindow(isCurrent: WindowLayoutGuard): Promise { const { getCurrentWindow, LogicalSize } = await import( "@tauri-apps/api/window" ); + const { invoke } = await import("@tauri-apps/api/core"); if (!isCurrent()) return; const win = getCurrentWindow(); + await invoke("reset_app_window_layout_initialized"); + if (!isCurrent()) return; await win.setResizable(false); if (!isCurrent()) return; await win.setSize(new LogicalSize(SETUP_WINDOW_WIDTH, SETUP_WINDOW_HEIGHT)); @@ -98,18 +101,20 @@ async function applyAppWindowLayout( if (!isCurrent()) return; const win = getCurrentWindow(); - // Decide first-launch vs restore from the on-disk state file BEFORE touching the - // window. Probing the window after restoreStateCurrent is unreliable: on GTK, - // set_size on a hidden window is deferred until show(), so innerSize() reads a - // stale value and a baseline fallback would overwrite the queued restore. On - // macOS the same probe works, hence the inconsistency between prior iterations. - const hasSavedState = await invoke("has_saved_window_state"); + // Setup-window activity may create plugin state before the full app is ever + // shown, so use a dedicated full-app marker to decide whether restoration is + // appropriate. Keep checking plugin state so a missing/corrupt state file + // falls back to a monitor-safe centered layout. + const [hasInitializedAppLayout, hasSavedState] = await Promise.all([ + invoke("has_initialized_app_window_layout"), + invoke("has_saved_window_state"), + ]); if (!isCurrent()) return; await win.setResizable(true); if (!isCurrent()) return; - if (hasSavedState) { + if (hasInitializedAppLayout && hasSavedState) { // Subsequent launch: plugin restores size/position/maximized, with built-in // off-screen protection for positions saved on a now-disconnected display. await restoreStateCurrent( @@ -144,6 +149,9 @@ async function applyAppWindowLayout( }); if (!isCurrent()) return; await enforceMinimumWindowSize(win, LogicalSize, isCurrent); + + if (!isCurrent()) return; + await invoke("mark_app_window_layout_initialized"); } async function showWindowFallback(): Promise { diff --git a/studio/frontend/src/app/routes/__root.tsx b/studio/frontend/src/app/routes/__root.tsx index ba56ce7525..7137fd6f96 100644 --- a/studio/frontend/src/app/routes/__root.tsx +++ b/studio/frontend/src/app/routes/__root.tsx @@ -16,6 +16,7 @@ import { type ChatSearch, } from "@/features/chat"; import { RemoteCodeConsentDialog } from "@/features/security"; +import { HfTokenWarningDialog } from "@/features/hf-auth"; import { TransformersUpgradeDialog } from "@/features/transformers-upgrade"; import { useTrainingUnloadGuard } from "@/features/training"; import { useExportRuntimeLifecycle } from "@/features/export"; @@ -195,9 +196,6 @@ function RootLayout() { chatRuntime.setActiveThreadId(null); chatRuntime.setActiveProjectId(null); chatRuntime.setIncognito(false); - // Detach the staging UI but keep any in-flight download running, like Hub. - if (chatRuntime.pendingSelection) - chatRuntime.abandonStagedModel({ keepDownload: true }); void navigate({ to: "/chat", search: { new: crypto.randomUUID() }, @@ -220,16 +218,13 @@ function RootLayout() { chatRuntime.setActiveProjectId(null); chatRuntime.setActiveThreadId(null); chatRuntime.setIncognito(false); - // Leaving chat must not kill an in-flight download: detach the staging UI - // but keep the transfer running in the manager, like a Hub download. - if (chatRuntime.pendingSelection) - chatRuntime.abandonStagedModel({ keepDownload: true }); }, [isChatRoute]); return ( {!isAuthFlowRoute && } + {hideNavbar ? ( @@ -279,7 +274,7 @@ function RootLayout() { initial={{ opacity: 0 }} animate={{ opacity: 1 }} exit={{ opacity: 0 }} - transition={{ duration: 0.15 }} + transition={{ duration: 0.06 }} className="flex min-h-0 min-w-0 flex-1 basis-0 flex-col overflow-visible" > }> diff --git a/studio/frontend/src/app/routes/data-recipes.tsx b/studio/frontend/src/app/routes/data-recipes.tsx index c35e63da5f..22f87821af 100644 --- a/studio/frontend/src/app/routes/data-recipes.tsx +++ b/studio/frontend/src/app/routes/data-recipes.tsx @@ -1,15 +1,13 @@ // SPDX-License-Identifier: AGPL-3.0-only // Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0 -import { createRoute } from "@tanstack/react-router"; -import { lazy } from "react"; +import { createRoute, lazyRouteComponent } from "@tanstack/react-router"; import { requireAuth } from "../auth-guards"; import { Route as rootRoute } from "./__root"; -const DataRecipesPage = lazy(() => - import("@/features/data-recipes").then((m) => ({ - default: m.DataRecipesPage, - })), +const DataRecipesPage = lazyRouteComponent( + () => import("@/features/data-recipes"), + "DataRecipesPage", ); export const Route = createRoute({ diff --git a/studio/frontend/src/app/routes/export.tsx b/studio/frontend/src/app/routes/export.tsx index 40118c6a92..5a7b586f19 100644 --- a/studio/frontend/src/app/routes/export.tsx +++ b/studio/frontend/src/app/routes/export.tsx @@ -1,15 +1,13 @@ // SPDX-License-Identifier: AGPL-3.0-only // Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0 -import { createRoute } from "@tanstack/react-router"; -import { lazy } from "react"; +import { createRoute, lazyRouteComponent } from "@tanstack/react-router"; import { requireAuth } from "../auth-guards"; import { Route as rootRoute } from "./__root"; -const ExportPage = lazy(() => - import("@/features/export/export-page").then((m) => ({ - default: m.ExportPage, - })), +const ExportPage = lazyRouteComponent( + () => import("@/features/export/export-page"), + "ExportPage", ); export type ExportSearch = { diff --git a/studio/frontend/src/app/routes/hub.tsx b/studio/frontend/src/app/routes/hub.tsx index c623ef9848..2207490e44 100644 --- a/studio/frontend/src/app/routes/hub.tsx +++ b/studio/frontend/src/app/routes/hub.tsx @@ -1,15 +1,13 @@ // SPDX-License-Identifier: AGPL-3.0-only // Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0 -import { createRoute } from "@tanstack/react-router"; -import { lazy } from "react"; +import { createRoute, lazyRouteComponent } from "@tanstack/react-router"; import { requireAuth } from "../auth-guards"; import { Route as rootRoute } from "./__root"; -const ModelsPage = lazy(() => - import("@/features/hub/hub-page").then((m) => ({ - default: m.ModelsPage, - })), +const ModelsPage = lazyRouteComponent( + () => import("@/features/hub/hub-page"), + "ModelsPage", ); export interface ModelsSearch { @@ -31,7 +29,11 @@ export const Route = createRoute({ const model = search.model; if (typeof model === "string" && model.length > 0) next.model = model; const section = search.section; - if (section === "trending" || section === "latest" || section === "finetune") { + if ( + section === "trending" || + section === "latest" || + section === "finetune" + ) { next.section = section; } const kind = search.kind; diff --git a/studio/frontend/src/app/routes/projects.tsx b/studio/frontend/src/app/routes/projects.tsx index c63b1d5838..17f58ef631 100644 --- a/studio/frontend/src/app/routes/projects.tsx +++ b/studio/frontend/src/app/routes/projects.tsx @@ -1,15 +1,13 @@ // SPDX-License-Identifier: AGPL-3.0-only // Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0 -import { createRoute } from "@tanstack/react-router"; -import { lazy } from "react"; +import { createRoute, lazyRouteComponent } from "@tanstack/react-router"; import { requireAuth } from "../auth-guards"; import { Route as rootRoute } from "./__root"; -const ProjectsPage = lazy(() => - import("@/features/chat/projects-page").then((m) => ({ - default: m.ProjectsPage, - })), +const ProjectsPage = lazyRouteComponent( + () => import("@/features/chat/projects-page"), + "ProjectsPage", ); export const Route = createRoute({ diff --git a/studio/frontend/src/app/routes/studio.tsx b/studio/frontend/src/app/routes/studio.tsx index ae7f445e94..798044bf64 100644 --- a/studio/frontend/src/app/routes/studio.tsx +++ b/studio/frontend/src/app/routes/studio.tsx @@ -1,15 +1,13 @@ // SPDX-License-Identifier: AGPL-3.0-only // Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0 -import { createRoute } from "@tanstack/react-router"; -import { lazy } from "react"; +import { createRoute, lazyRouteComponent } from "@tanstack/react-router"; import { requireAuth } from "../auth-guards"; import { Route as rootRoute } from "./__root"; -const StudioPage = lazy(() => - import("@/features/studio/studio-page").then((m) => ({ - default: m.StudioPage, - })), +const StudioPage = lazyRouteComponent( + () => import("@/features/studio/studio-page"), + "StudioPage", ); export const Route = createRoute({ diff --git a/studio/frontend/src/components/app-sidebar.tsx b/studio/frontend/src/components/app-sidebar.tsx index b6e218793a..a13828b06b 100644 --- a/studio/frontend/src/components/app-sidebar.tsx +++ b/studio/frontend/src/components/app-sidebar.tsx @@ -55,6 +55,7 @@ import { Archive03Icon, ArrowRight02Icon, BadgeInfoIcon, + BubbleChatIcon, ChefHatIcon, CloudIcon, CpuIcon, @@ -93,7 +94,12 @@ import { import { Tooltip as TooltipPrimitive } from "radix-ui"; import { HugeiconsIcon } from "@hugeicons/react"; import { ChevronDown, Moon } from "lucide-react"; -import { Link, useNavigate, useRouterState } from "@tanstack/react-router"; +import { + Link, + useNavigate, + useRouter, + useRouterState, +} from "@tanstack/react-router"; import { archiveChatItem, ChatSearchDialog, @@ -103,6 +109,7 @@ import { deleteChatItem, listStoredChatThreads, moveChatItemToProject, + notifyChatHistoryUpdated, renameChatItem, renameChatProject, useChatRuntimeStore, @@ -110,6 +117,7 @@ import { useChatSearchStore, useChatSidebarItems, usePinnedChatsStore, + usePinnedProjectsStore, useChatPreferencesStore, type ProjectRecord, type SidebarItem, @@ -135,7 +143,15 @@ import { } from "@/features/training"; import type { TrainingRunSummary } from "@/features/training"; import { useExportRuntimeStore } from "@/features/export"; -import { useEffect, useMemo, useRef, useState, type ReactNode } from "react"; +import { + Fragment, + useEffect, + useMemo, + useRef, + useState, + type ReactNode, +} from "react"; +import { isDownloadCancelled } from "@/lib/native-files"; import { toast } from "@/lib/toast"; import { ShutdownDialog } from "@/components/shutdown-dialog"; import { translate, useT, type TranslationKey } from "@/i18n"; @@ -193,6 +209,9 @@ const TestTubeOutlineIcon = TestTube01Icon.slice( type ConversationExportFormat = "raw-jsonl" | "csv" | "sharegpt-jsonl"; +// A pinned project shows this many recent chats before "Show more". +const PINNED_PROJECT_CHAT_LIMIT = 4; + const CHAT_EXPORT_OPTIONS: Array<{ label: string; format: ConversationExportFormat; @@ -256,6 +275,10 @@ function createNavigationNonce(): string { return `${Date.now()}-${Math.random().toString(36).slice(2, 10)}`; } +function preloadSilently(request: Promise): void { + void request.catch(() => undefined); +} + function NavItem({ icon, label, @@ -267,6 +290,7 @@ function NavItem({ className, spinner, tooltip, + onIntent, }: { icon: typeof ZapIcon; label: string; @@ -277,6 +301,7 @@ function NavItem({ dataTour?: string; className?: string; spinner?: boolean; + onIntent?: () => void; // Overrides the hover tooltip (defaults to `label`). Used to explain why a // disabled item (e.g. Train/Export on a chat-only host) is greyed out. tooltip?: string; @@ -288,6 +313,8 @@ function NavItem({ tooltip={tooltip ?? label} disabled={disabled} onClick={onClick} + onPointerEnter={disabled ? undefined : onIntent} + onFocus={disabled ? undefined : onIntent} isActive={active} data-tour={dataTour} className="sidebar-nav-btn h-[33px] rounded-full gap-[8.5px] pl-3 pr-2.5 font-medium group-data-[collapsible=icon]:px-2.5 group-data-[collapsible=icon]:!w-[32px] group-data-[collapsible=icon]:mx-auto" @@ -324,6 +351,7 @@ export function AppSidebar() { }); const { togglePinned, isMobile, setOpenMobile } = useSidebar(); const navigate = useNavigate(); + const router = useRouter(); // Web update detection: `webUpdate` is non-null only when the installed // (PyPI) version is behind the latest release, so the card is hidden by @@ -430,14 +458,63 @@ export function AppSidebar() { ), [allChatItems, pinnedIdSet], ); - // Pinned chats, in pin order (most recent first). + const [pinnedOpen, setPinnedOpen] = useState(true); + // "Projects" section: projects the user pinned, in pin order (most recent + // first). The section only appears once at least one project is pinned. + const pinnedProjectIds = usePinnedProjectsStore((s) => s.pinnedIds); + const unpinProject = usePinnedProjectsStore((s) => s.unpin); + const pinnedProjectRecords = useMemo(() => { + const byId = new Map(projects.map((p) => [p.id, p])); + return pinnedProjectIds + .map((id) => byId.get(id)) + .filter((p): p is ProjectRecord => Boolean(p)); + }, [projects, pinnedProjectIds]); + // Pinned chats, in pin order (most recent first). Includes chats that live + // inside a project: pinning promotes a chat into this list, and it is removed + // from the project's nested list below so it never shows twice. const pinnedChatItems = useMemo(() => { const byId = new Map(allChatItems.map((item) => [item.id, item])); return pinnedIds .map((id) => byId.get(id)) .filter((item): item is SidebarItem => Boolean(item)); }, [allChatItems, pinnedIds]); - const [pinnedOpen, setPinnedOpen] = useState(true); + // A pinned project reveals its recent chats (most recent first) nested below. + // Pinned chats are excluded here since they render in the pinned-chats list. + const chatsByProjectId = useMemo(() => { + const map = new Map(); + for (const item of allChatItems) { + if (!item.projectId) continue; + if (pinnedIdSet.has(item.id)) continue; + const list = map.get(item.projectId); + if (list) list.push(item); + else map.set(item.projectId, [item]); + } + for (const list of map.values()) + list.sort((a, b) => b.updatedAt - a.updatedAt); + return map; + }, [allChatItems, pinnedIdSet]); + // Default expanded (not collapsed); the row toggles this. Show-more reveals + // chats past the first PINNED_PROJECT_CHAT_LIMIT. + const [collapsedProjectIds, setCollapsedProjectIds] = useState>( + () => new Set(), + ); + const [expandedChatProjectIds, setExpandedChatProjectIds] = useState< + Set + >(() => new Set()); + const toggleProjectCollapsed = (id: string) => + setCollapsedProjectIds((prev) => { + const next = new Set(prev); + if (next.has(id)) next.delete(id); + else next.add(id); + return next; + }); + const toggleProjectShowAll = (id: string) => + setExpandedChatProjectIds((prev) => { + const next = new Set(prev); + if (next.has(id)) next.delete(id); + else next.add(id); + return next; + }); const storeThreadId = useChatRuntimeStore((s) => s.activeThreadId); const setActiveThreadId = useChatRuntimeStore((s) => s.setActiveThreadId); const anyChatRunning = useChatRuntimeStore((s) => @@ -717,6 +794,8 @@ export function AppSidebar() { const shouldDeleteProjectFiles = target.kind === "project" && deleteProjectFiles; setConfirmingDelete(null); + // Reset so the next project delete never inherits this checkbox. + setDeleteProjectFiles(false); if (target.kind === "chat") { await deleteChatWithCleanup(target.item); return; @@ -726,7 +805,20 @@ export function AppSidebar() { await deleteChatProject(target.project.id, { deleteFiles: shouldDeleteProjectFiles, }); - if (activeProjectId === target.project.id) { + // Refresh chat history so the project's reparented chats don't linger + // as stale top-level rows. + notifyChatHistoryUpdated(); + // activeProjectId is only the ?project= param; on a thread-only URL the + // project is resolved from the thread into the runtime store, so check + // that too or we strand the user on a now-deleted thread. Only redirect + // from a chat route: the runtime store value can be stale elsewhere. + const runtimeProjectId = + useChatRuntimeStore.getState().activeProjectId; + if ( + isChatRoute && + (activeProjectId === target.project.id || + runtimeProjectId === target.project.id) + ) { useChatRuntimeStore.getState().setActiveProjectId(null); navigate({ to: "/chat", search: { new: createNavigationNonce() } }); } @@ -813,8 +905,11 @@ export function AppSidebar() { // pl-3 (12px) over the content's pl-1.5 (6px) = 18px, aligning the // title with the nav items above. variant === "project" ? "pl-[39px]" : "pl-3", + // Pinned chats carry a chat icon, so add the nav-item icon gap. + isPinned && variant !== "project" && "gap-[8.5px]", variant === "project" - ? "group-hover/project-chat-item:pr-6 group-has-[.sidebar-row-action[data-state=open]]/project-chat-item:pr-6" + ? // Room for the hover pin quick-action plus the kebab. + "group-hover/project-chat-item:pr-14 group-has-[.sidebar-row-action[data-state=open]]/project-chat-item:pr-8" : isPinned ? // Pinned rows show an extra unpin button on hover, so reserve more room // (pr-8 when the menu is open keeps the unpin button clear of the title). @@ -874,10 +969,43 @@ export function AppSidebar() { closeMobileIfOpen(); }} > + {isPinned && variant !== "project" && ( + + )} {pendingRename?.id === item.id ? pendingRename.title : item.title} + {variant === "project" && ( + + )} + {variant === "recent" && isPinned && ( + + )} - - - Unpin - - - ) : null} ); } @@ -1218,6 +1326,9 @@ export function AppSidebar() { navigate({ to: "/projects" }); closeMobileIfOpen(); }} + onIntent={() => { + preloadSilently(router.preloadRoute({ to: "/projects" })); + }} className="group/projects-item relative" > + {/* Project options */} + + + + + + openProject(project.id)}> + + Project home + + openNewChat(project.id)}> + + New chat + + { + // Seed the shared draft so the dialog opens + // with the current name, not stale text. + setRenameDraft(project.name); + setRenamingTarget({ + kind: "project", + project, + current: project.name, + }); + }} + > + + Rename project + + unpinProject(project.id)}> + + Unpin project + + + { + // Start each delete with the file toggle off: + // Cancel closes programmatically and skips the + // dialog onOpenChange reset. + setDeleteProjectFiles(false); + setConfirmingDelete({ kind: "project", project }); + }} + > + + Delete project + + + + + {expanded && + visibleChats.map((chat) => + renderChatSidebarItem(chat, "project"), + )} + {expanded && + projectChats.length > PINNED_PROJECT_CHAT_LIMIT && ( + + toggleProjectShowAll(project.id)} + // Force the muted token: .sidebar-nav-btn's own + // color rule outweighs a plain text utility, so + // Show more would otherwise match the chat rows. + className="sidebar-nav-btn h-[30px] rounded-full pl-9 pr-4 font-medium text-nav-fg-muted!" + > + + {showAll ? "Show less" : "Show more"} + + + + )} + + ); + })} + {pinnedChatItems.map((item) => + renderChatSidebarItem(item, "recent"), + )} + + + + + + )} {!isStudioRoute && !showTrainingRecents && ( diff --git a/studio/frontend/src/components/assistant-ui/attachment.tsx b/studio/frontend/src/components/assistant-ui/attachment.tsx index 98ebe5ab5f..b26840cc02 100644 --- a/studio/frontend/src/components/assistant-ui/attachment.tsx +++ b/studio/frontend/src/components/assistant-ui/attachment.tsx @@ -7,6 +7,7 @@ import { TooltipIconButton } from "@/components/assistant-ui/tooltip-icon-button"; import { Dialog, + DialogClose, DialogContent, DialogTitle, DialogTrigger, @@ -27,12 +28,7 @@ import { import { AudioWave01Icon, File02Icon } from "@hugeicons/core-free-icons"; import { HugeiconsIcon } from "@hugeicons/react"; import { PlusIcon, XIcon } from "lucide-react"; -import { - type FC, - type PropsWithChildren, - useEffect, - useState, -} from "react"; +import { type FC, type PropsWithChildren, useEffect, useState } from "react"; import { useShallow } from "zustand/shallow"; const useFileSrc = (file: File | undefined): string | undefined => { @@ -83,7 +79,7 @@ const AttachmentPreview: FC = ({ src }) => { src={src} alt="Preview" className={cn( - "block h-auto max-h-[80vh] w-auto max-w-full object-contain", + "block h-auto max-h-[90dvh] w-auto max-w-[92vw] object-contain", isLoaded ? "aui-attachment-preview-image-loaded" : "aui-attachment-preview-image-loading invisible", @@ -108,12 +104,23 @@ const AttachmentPreviewDialog: FC = ({ children }) => { > {children} - + {/* Chrome-free lightbox: the image floats on the dimmed backdrop with + no dialog panel, and the close button sits in the screen corner. */} + Image Attachment Preview -

- + {/* Clicking the backdrop (anywhere off the image) closes the preview. */} + + @@ -1216,6 +1590,96 @@ function ProjectLanding({
)} + { + if (!open) setConfirmingDelete(null); + }} + > + + + Delete chat + + This permanently deletes "{confirmingDelete?.title}". This cannot + be undone. + + + + Cancel + { + const target = confirmingDelete; + setConfirmingDelete(null); + if (target) void runDelete(target); + }} + > + Delete + + + + + { + if (!open) setRenamingProject(false); + }} + > + + + Rename project + + setProjectNameDraft(e.target.value)} + onKeyDown={(e) => { + if (e.key === "Enter") { + e.preventDefault(); + void commitProjectRename(); + } + }} + autoFocus={true} + maxLength={120} + placeholder="Project name" + aria-label="Project name" + className="focus-visible:border-input focus-visible:ring-0" + /> + + + + + + + { + if (!open) setDeletingProject(false); + }} + > + + + Delete project + + Delete "{projectName}"? Its chats will be permanently deleted. + + + + Cancel + void commitProjectDelete()}> + Delete + + + + ); } @@ -1236,6 +1700,13 @@ export function validateChatSearch(search: Record): ChatSearch }; } +type PendingHubAutoLoad = { + selection: SelectedModelInput; + contextKey: string; + originCheckpoint: string; + originGgufVariant: string | null; +}; + // `search` comes from RootLayout (not useSearch) so ChatPage stays mounted off-route // (keeping an in-flight generation alive), frozen to the last /chat search. `active` // is false off-route: close body-portaled surfaces and stop route-specific listeners @@ -1248,30 +1719,6 @@ export function ChatPage({ const settingsOpen = useChatRuntimeStore((s) => s.settingsPanelOpen); const setSettingsOpen = useChatRuntimeStore((s) => s.setSettingsPanelOpen); - // Deferred-load staging: downloads a staged GGUF (if needed) and reads its - // header context so the sheet can show the context slider before the load. - // autoLoad picks instead load the cached file as soon as the download ends; - // selectModel is defined below, so the load runs through a ref. - const autoLoadStagedRef = useRef< - ((pending: PendingModelSelection) => void) | null - >(null); - const stagedDownload = useStagedModelPreparation({ - onAutoLoad: (pending) => autoLoadStagedRef.current?.(pending), - }); - // Abandon a staged pick: the store action cancels its in-flight download and - // reverts the edited knobs, so nothing lingers after the user walks away. - const abandonStaged = useCallback(() => { - useChatRuntimeStore.getState().abandonStagedModel(); - }, []); - // Detach a staged pick on navigation without cancelling its download: the - // transfer keeps running in the manager and lands in cache, like Hub. - const detachStaged = useCallback(() => { - useChatRuntimeStore.getState().abandonStagedModel({ keepDownload: true }); - }, []); - // Tracks whether the chat page is still mounted, so a staged-load failure that - // resolves after the user left chat doesn't resurrect the abandoned pick. - const mountedRef = useRef(true); - useEffect(() => () => void (mountedRef.current = false), []); const incognito = useChatRuntimeStore((s) => s.incognito); const setIncognito = useChatRuntimeStore((s) => s.setIncognito); const incognitoLabel = incognito @@ -1363,6 +1810,9 @@ export function ChatPage({ const ggufContextLength = useChatRuntimeStore( (state) => state.ggufContextLength, ); + const ggufNativeContextLength = useChatRuntimeStore( + (state) => state.ggufNativeContextLength, + ); const contextUsage = useChatRuntimeStore((state) => state.contextUsage); const modelsFromStore = useChatRuntimeStore((state) => state.models); const lorasFromStore = useChatRuntimeStore((state) => state.loras); @@ -1440,37 +1890,37 @@ export function ChatPage({ refreshRef.current = refresh; selectModelRef.current = selectModel; }, [refresh, selectModel]); - // Load a cached autoLoad pick once its download finishes. The sheet was never - // opened, so on a load failure just drop the orphaned staged knobs. The knobs - // were already seeded on stage, so keepSpeculative only when a config was - // saved -- otherwise the standing speculative preference should win. - autoLoadStagedRef.current = (pending) => { - const remembered = loadRememberedLoadSettings( - rememberedLoadSettingsKey(pending), - ); - void selectModel({ - ...pending, - isDownloaded: true, - forceReload: true, - keepSpeculative: remembered != null, - throwOnError: true, - }).catch(() => { - const store = useChatRuntimeStore.getState(); - // selectModel only clears pendingSelection on success, so a failed - // auto-load leaves our staged pick (and its edited load knobs) behind. - // Abandon it when it is still the active stage; otherwise just revert the - // settings if the stage was already cleared by something else. - if (pendingSelectionMatches(store.pendingSelection, pending)) { - store.abandonStagedModel(); - } else if (!store.pendingSelection) { - store.resetModelSettingsToLoaded(); - } - }); - }; + const rememberedConfigFor = useCallback( + (selection: { + id: string; + ggufVariant?: string | null; + source?: string; + }) => { + if (selection.source === "external") return null; + const resolved = resolveInitialConfig(selection.id, selection.ggufVariant); + return resolved.remembered ? resolved.config : null; + }, + [], + ); const isExternalModel = useMemo( () => isExternalModelId(inferenceParams.checkpoint), [inferenceParams.checkpoint], ); + const { + checkpoint: runtimeCheckpoint, + isGguf: runtimeModelIsGguf, + config: activeModelConfig, + } = useActiveModelConfig(); + const activeModelIsGguf = + runtimeCheckpoint != null && !isExternalModel && runtimeModelIsGguf; + const activeModelIsLora = useMemo(() => { + const checkpoint = inferenceParams.checkpoint; + if (!checkpoint || isExternalModel) return false; + const model = modelsFromStore.find((entry) => entry.id === checkpoint); + if (model) return model.isLora; + const lora = lorasFromStore.find((entry) => entry.id === checkpoint); + return lora?.exportType === "lora"; + }, [inferenceParams.checkpoint, isExternalModel, modelsFromStore, lorasFromStore]); const reasoningEnabled = useChatRuntimeStore((s) => s.reasoningEnabled); const reasoningStyle = useChatRuntimeStore((s) => s.reasoningStyle); const reasoningEffort = useChatRuntimeStore((s) => s.reasoningEffort); @@ -1781,75 +2231,21 @@ export function ChatPage({ closeArtifactSurface(); }, [activeThreadId, closeArtifactSurface, selectedArtifact, view]); - // Abandon a staged (not-yet-loaded) pick when the chat context actually - // changes — switching threads, leaving single view, or starting a new chat / - // project — so a stale Load button can't resurface in a different context. - // New Chat keeps activeThreadId null and only bumps the `new` search nonce, so - // the key includes the route identity, not just the thread. Mirrors the - // incognito reset pattern. (Route exit is handled in __root.tsx, which runs - // after this unmounts.) Clear only on a real change, never on mount: staging - // from the Hub sets pendingSelection then navigates here, and clearing on - // mount would wipe it. Comparing the previous context (rather than a first-run - // flag) is also safe under StrictMode's double-invoke and component remounts. - const chatContextKey = `${view.mode}|${activeThreadId ?? ""}|${search.new ?? ""}|${search.project ?? ""}`; - const chatContextKeyRef = useLatestRef(chatContextKey); - const prevChatContextRef = useRef(null); - useEffect(() => { - const prev = prevChatContextRef.current; - prevChatContextRef.current = chatContextKey; - if (prev === null || prev === chatContextKey) return; - detachStaged(); - }, [chatContextKey, detachStaged]); - const hasActiveModel = Boolean(inferenceParams.checkpoint); - // Load immediately, or — when "Load on selection" is off — stage the pick so - // its load options can be set first. Shared by the main selector, native - // drag-drop/picker, and the dropped-file chip (the Hub stages via the store). + const chatContextKey = `${view.mode}|${activeThreadId ?? ""}|${search.new ?? ""}|${search.project ?? ""}`; + const [pendingHubAutoLoad, setPendingHubAutoLoad] = + useState(null); const stageOrLoad = useCallback( async (selection: SelectedModelInput) => { const store = useChatRuntimeStore.getState(); - // An un-cached HF repo (GGUF variant or a full non-GGUF snapshot) downloads - // through the manager first (global indicator), then auto-loads. Everything - // else -- cached picks, local/native files, LoRA, external -- loads now. const wantManagerDownload = isDownloadableHubRepo(selection) && !selection.isDownloaded; - if ( - (!hasGgufSource(selection) && !wantManagerDownload) || - (store.loadOnSelection && selection.isDownloaded) - ) { - // Detach any staged pick first so its edited knobs (e.g. a custom - // context length) don't leak into this immediate load -- resolveLoad - // reads customContextLength before checking the target is GGUF. Detach - // (not abandon) keeps its download running. - detachStaged(); - // Load-on-selection skips the sheet, so seed the saved knobs here the - // way the sheet's restore effect would; the switch would otherwise reset - // the remembered speculative choice (keepSpeculative below prevents it). - const remembered = hasGgufSource(selection) - ? loadRememberedLoadSettings(rememberedLoadSettingsKey(selection)) - : null; - if (remembered) store.applyRememberedLoadSettings(remembered); - await selectModel( - remembered ? { ...selection, keepSpeculative: true } : selection, - ); - return; - } - // Loads can't queue behind each other, but a download is independent: if - // the pick needs downloading, start it in the manager so it runs alongside - // the load. Nothing to download (already on device) just waits. if (store.modelLoading) { - // Both an uncached non-GGUF snapshot (wantManagerDownload) and an - // uncached remote GGUF quant download through the manager, so either can - // run in the background while another model loads. wantManagerDownload - // excludes GGUF by design, so the GGUF case is checked separately. const wantBackgroundDownload = wantManagerDownload || (selection.source === "hub" && hasGgufSource(selection) && !selection.isDownloaded); - // The model currently loading already downloads as part of its own load - // (the /load flow fetches before setting the checkpoint), so re-picking - // it must not kick off a second transfer against the same cache. const isLoadingThisPick = !!loadingModel && normalizeModelRef(loadingModel.id) === @@ -1860,11 +2256,6 @@ export function ChatPage({ description: "It's downloading as part of the load in progress.", }); } else if (wantBackgroundDownload) { - // Only claim the download started once a job is actually created. A - // transport conflict records state that is only resolvable from the - // Hub download card, so point the user there instead of showing a - // success toast for a transfer that never began; "busy" and "error" - // already surface their own toasts. const outcome = await downloadManager.requestStart({ kind: DOWNLOAD_KIND.MODEL, repoId: selection.id, @@ -1881,6 +2272,11 @@ export function ChatPage({ description: "An earlier partial download used a different transport. Open the Hub tab to resume or restart it.", }); + } else if (outcome === "busy") { + toast.info("Download already in progress", { + description: + "Another download for this model is still running. Reselect it once that finishes to load it.", + }); } } else { toast.info("Another model is already loading", { @@ -1889,23 +2285,128 @@ export function ChatPage({ } return; } - // Detach the prior staged pick (keeping its download) before rebinding, so - // a second pick downloads alongside the first instead of cancelling it. - detachStaged(); - store.stageModel({ - id: selection.id, - isLora: selection.isLora, - ggufVariant: selection.ggufVariant, - isDownloaded: selection.isDownloaded, - expectedBytes: selection.expectedBytes, - nativePathToken: selection.nativePathToken, - isGguf: selection.isGguf, - isHubRepo: wantManagerDownload || undefined, - autoLoad: store.loadOnSelection, + const wantManagerStage = + wantManagerDownload || + (selection.source === "hub" && + hasGgufSource(selection) && + !selection.isDownloaded); + if (wantManagerStage) { + setPendingHubAutoLoad((current) => + current && + current.selection.id === selection.id && + (current.selection.ggufVariant ?? null) === + (selection.ggufVariant ?? null) && + current.contextKey === chatContextKey && + current.originCheckpoint === store.params.checkpoint && + current.originGgufVariant === store.activeGgufVariant + ? current + : { + selection, + contextKey: chatContextKey, + originCheckpoint: store.params.checkpoint, + originGgufVariant: store.activeGgufVariant, + }, + ); + return; + } + setPendingHubAutoLoad(null); + const previousConfig = currentRuntimePerModelConfig({ + includeMaxSeqLength: true, + }); + const hasAppliedConfig = applyModelLoadConfigToRuntime( + selection.config ?? rememberedConfigFor(selection), + ); + await selectModel({ + ...selection, + ...(hasAppliedConfig ? { keepSpeculative: true } : {}), + previousConfig, }); }, - [detachStaged, selectModel, loadingModel], + [selectModel, loadingModel, rememberedConfigFor, chatContextKey], ); + useRepoDownload({ + kind: DOWNLOAD_KIND.MODEL, + repoId: pendingHubAutoLoad?.selection.id ?? "__hub_autoload_idle__", + activeVariant: pendingHubAutoLoad?.selection.ggufVariant ?? null, + onComplete: (variant) => { + const pending = pendingHubAutoLoad; + if ( + !pending || + (pending.selection.ggufVariant ?? null) !== (variant ?? null) + ) { + return; + } + setPendingHubAutoLoad(null); + const store = useChatRuntimeStore.getState(); + if ( + !active || + pending.contextKey !== chatContextKey || + normalizeModelRef(pending.originCheckpoint) !== + normalizeModelRef(store.params.checkpoint) || + pending.originGgufVariant !== store.activeGgufVariant + ) { + return; + } + void stageOrLoad({ ...pending.selection, isDownloaded: true }); + }, + onError: (variant) => { + if ( + pendingHubAutoLoad && + (pendingHubAutoLoad.selection.ggufVariant ?? null) === (variant ?? null) + ) { + setPendingHubAutoLoad(null); + } + }, + onCancelled: (variant) => { + if ( + pendingHubAutoLoad && + (pendingHubAutoLoad.selection.ggufVariant ?? null) === (variant ?? null) + ) { + setPendingHubAutoLoad(null); + } + }, + }); + useEffect(() => { + const pending = pendingHubAutoLoad; + if (!pending) return; + let active = true; + void (async () => { + const outcome = await downloadManager.requestStart({ + kind: DOWNLOAD_KIND.MODEL, + repoId: pending.selection.id, + variant: pending.selection.ggufVariant ?? null, + expectedBytes: pending.selection.expectedBytes ?? 0, + }); + if (!active) return; + if (outcome === "started") { + toast.info("Downloading model", { + description: "It'll load automatically once the download finishes.", + }); + return; + } + if (outcome === "conflict") { + // Keep pendingHubAutoLoad bound so this surface's cleanup does not wipe + // the conflict just recorded by requestStart (which the toast points the + // user to); resolving it from the Hub completes the download and this + // surface's onComplete auto-loads, mirroring the "started" branch. + toast.info("Resume this download from the Hub", { + description: + "An earlier partial download used a different transport. Open the Hub tab to resume or restart it.", + }); + return; + } + if (outcome === "busy") { + toast.info("Download already in progress", { + description: + "Another download for this model is still running. Reselect it once that finishes to load it.", + }); + } + setPendingHubAutoLoad((current) => (current === pending ? null : current)); + })(); + return () => { + active = false; + }; + }, [pendingHubAutoLoad]); const loadNativeModelIntent = useCallback( async (intent: NativeIntent, loadingDescription: string) => { const label = @@ -1913,6 +2414,7 @@ export function ChatPage({ await stageOrLoad({ id: label, nativePathToken: intent.path.token, + nativePathExpiresAtMs: intent.path.expiresAtMs ?? null, isDownloaded: true, loadingDescription, forceReload: true, @@ -1963,28 +2465,20 @@ export function ChatPage({ const handleCheckpointChange = useCallback( ( value: string, - meta?: { - source?: string; - isLora: boolean; - ggufVariant?: string; - isDownloaded?: boolean; - expectedBytes?: number; - isGguf?: boolean; - }, + meta?: ModelSelectorChangeMeta, ) => { const store = useChatRuntimeStore.getState(); const currentCheckpoint = store.params.checkpoint; const currentVariant = store.activeGgufVariant; - if ( - !value || - (value === currentCheckpoint && - (meta?.ggufVariant ?? null) === (currentVariant ?? null)) - ) + if (!value) return; + setPendingHubAutoLoad(null); + const isSameLoadedModel = + value === currentCheckpoint && + (meta?.ggufVariant ?? null) === (currentVariant ?? null); + if (isSameLoadedModel && !meta?.forceReload) { return; + } if (meta?.source === "external" || isExternalModelId(value)) { - // Switching to an external model abandons any staged local pick: cancel - // its download too (setCheckpoint below only clears the pending + knobs). - abandonStaged(); const selectedExternal = parseExternalModelId(value); const selectedProvider = selectedExternal ? externalProvidersForChat.find( @@ -2085,6 +2579,7 @@ export function ChatPage({ ggufMaxContextLength: null, ggufNativeContextLength: null, activeNativePathToken: null, + activeNativePathExpiresAtMs: null, // Clear previous-model counters, else the relaxed external-provider // render gate shows stale stats until the next completion. contextUsage: null, @@ -2156,19 +2651,18 @@ export function ChatPage({ source: meta?.source, isLora: meta?.isLora, ggufVariant: meta?.ggufVariant, - isDownloaded: meta?.isDownloaded, + isDownloaded: meta?.isDownloaded || isSameLoadedModel, expectedBytes: meta?.expectedBytes, isGguf: meta?.isGguf, + config: meta?.config, + nativePathToken: meta?.nativePathToken, + nativePathExpiresAtMs: meta?.nativePathExpiresAtMs, + forceReload: isSameLoadedModel || undefined, }; - // "Load on selection" off: stage the model and open settings so its - // load knobs (tensor parallel, context length…) can be set, then it - // loads once via the sheet's Load button. The currently loaded model - // stays put until the user commits. await stageOrLoad(selection); })(); }, [ - abandonStaged, activeThreadId, externalProvidersForChat, modelsFromStore, @@ -2176,6 +2670,45 @@ export function ChatPage({ view, ], ); + const handleReloadActiveModel = useCallback( + (config: PerModelConfig) => { + const checkpoint = inferenceParams.checkpoint; + if (!checkpoint) return; + const runtime = useChatRuntimeStore.getState(); + const nativeToken = runtime.activeNativePathToken; + const nativeExpiry = runtime.activeNativePathExpiresAtMs; + // A file-picked GGUF is reachable only via its native path token, which + // the desktop host prunes after a TTL. Reusing an expired token makes the + // reload fail with an opaque error, so prompt the user to re-select the + // file instead. + if (nativeToken && nativeExpiry != null && Date.now() >= nativeExpiry) { + toast.error("This local model file's access has expired.", { + description: "Re-select the model file to reload it.", + }); + return; + } + handleCheckpointChange(checkpoint, { + source: "local", + isLora: activeModelIsLora, + ggufVariant: activeGgufVariant ?? undefined, + // Without the native token the reload validates the display label as a + // repo and fails. + nativePathToken: nativeToken ?? undefined, + nativePathExpiresAtMs: nativeExpiry, + isGguf: activeModelIsGguf, + isDownloaded: true, + config, + forceReload: true, + }); + }, + [ + inferenceParams.checkpoint, + activeGgufVariant, + activeModelIsLora, + activeModelIsGguf, + handleCheckpointChange, + ], + ); const handleEject = useCallback(() => { void (async () => { if (await ejectModel()) { @@ -2444,12 +2977,27 @@ export function ChatPage({ const state = useChatRuntimeStore.getState(); const targetLora = pickBestLoraForBase(state.loras, handoff.baseModel); + const selectWithConfig = async ( + selection: Pick, + ) => { + const previousConfig = currentRuntimePerModelConfig({ + includeMaxSeqLength: true, + }); + const hasAppliedConfig = applyModelLoadConfigToRuntime( + rememberedConfigFor(selection), + ); + await selectModelRef.current({ + ...selection, + ...(hasAppliedConfig ? { keepSpeculative: true } : {}), + previousConfig, + }); + }; if (targetLora) { console.info("[chat-handoff] loading lora", { id: targetLora.id, baseModel: targetLora.baseModel, }); - await selectModelRef.current({ id: targetLora.id, isLora: true }); + await selectWithConfig({ id: targetLora.id, isLora: true }); if (canceled) return; useChatRuntimeStore.getState().setActiveThreadId(null); useChatRuntimeStore.getState().setContextUsage(null); @@ -2466,10 +3014,7 @@ export function ChatPage({ console.info("[chat-handoff] no lora match, loading base", { id: handoff.baseModel, }); - await selectModelRef.current({ - id: handoff.baseModel, - isLora: false, - }); + await selectWithConfig({ id: handoff.baseModel, isLora: false }); if (canceled) return; } else { console.warn("[chat-handoff] no lora/base match found", { @@ -2489,7 +3034,7 @@ export function ChatPage({ return () => { canceled = true; }; - }, [active, navigate]); + }, [active, navigate, rememberedConfigFor]); const tourSteps = useMemo( () => @@ -2578,6 +3123,8 @@ export function ChatPage({ externalModels={externalModels} value={inferenceParams.checkpoint} activeGgufVariant={activeGgufVariant} + activeModelConfig={activeModelConfig} + activeGgufContextLength={ggufContextLength} onValueChange={handleCheckpointChange} onEject={handleEject} onFoldersChange={refreshLocalModels} @@ -2631,7 +3178,12 @@ export function ChatPage({ stageOrLoad(selection)} + onLoad={() => + loadNativeModelIntent( + pendingNativeModelIntent, + "Loading selected local GGUF model.", + ) + } /> ) : null} {loadingModel && loadToastDismissed ? ( @@ -2788,13 +3340,22 @@ export function ChatPage({ open={active && settingsOpen} onOpenChange={(open) => { setSettingsOpen(open); - // Closing the sheet abandons a staged (not-yet-loaded) pick: cancel its - // download and revert the staged knobs so nothing lingers as a dirty - // edit (or a background download) on the loaded model. - if (!open) abandonStaged(); }} params={inferenceParams} onParamsChange={setInferenceParams} + modelConfig={ + view.mode !== "compare" && activeModelConfig && !modelLoading ? ( + + ) : null + } isExternalModel={isExternalModel} providerCapabilities={activeProviderCapabilities} activeExternalProvider={activeExternalProvider} @@ -2806,62 +3367,6 @@ export function ChatPage({ ); }} externalProviderType={activeExternalProviderType} - loadingModel={loadingModel} - onReloadModel={() => { - const state = useChatRuntimeStore.getState(); - if (state.params.checkpoint) { - selectModel({ - id: state.params.checkpoint, - ggufVariant: state.activeGgufVariant ?? undefined, - forceReload: true, - isDownloaded: true, - loadingDescription: "Reloading with updated chat template.", - }); - } - }} - onLoadPendingModel={() => { - const pending = useChatRuntimeStore.getState().pendingSelection; - if (!pending) return; - const keyAtLoad = chatContextKey; - // forceReload: the staged model isn't loaded yet, so bypass the - // same-checkpoint dedupe. keepSpeculative: honor the speculative mode - // set on the sidebar. - void selectModel({ - ...pending, - forceReload: true, - keepSpeculative: true, - throwOnError: true, - }).catch(() => { - // Recoverable failure (expired token, gated repo, OOM…): the pick is - // cleared only on success, so it normally stays staged with edited - // knobs intact — nothing to restore. - const store = useChatRuntimeStore.getState(); - // Still staged (this pick, or a newer one queued meanwhile): leave it. - if (store.pendingSelection) return; - // Cleared mid-load (sheet closed / switched chats). Re-stage only if - // the staged-load is still wanted: same chat context, sheet still - // open, page still mounted. - const stillWanted = - mountedRef.current && - store.settingsPanelOpen && - chatContextKeyRef.current === keyAtLoad; - if (stillWanted) { - store.setPendingSelection(pending); - } else { - // Abandoned (closed the sheet / switched chats / left chat): drop - // the orphaned staged knob edits so they don't linger as dirty - // settings over the loaded model. - store.resetModelSettingsToLoaded(); - } - }); - }} - stagedDownloadFraction={stagedDownload.progress?.fraction ?? null} - onCancelStagedDownload={() => - stagedDownload.cancelDownload( - useChatRuntimeStore.getState().pendingSelection?.ggufVariant ?? - null, - ) - } /> diff --git a/studio/frontend/src/features/chat/chat-settings-sheet.tsx b/studio/frontend/src/features/chat/chat-settings-sheet.tsx index cedd298ecf..d4f154882c 100644 --- a/studio/frontend/src/features/chat/chat-settings-sheet.tsx +++ b/studio/frontend/src/features/chat/chat-settings-sheet.tsx @@ -1,19 +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 { - Alert, - AlertDescription, - AlertTitle, -} from "@/components/ui/alert"; import { Button } from "@/components/ui/button"; -import { Checkbox } from "@/components/ui/checkbox"; -import { - clearRememberedLoadSettings, - loadRememberedLoadSettings, - rememberedLoadSettingsKey, - saveRememberedLoadSettings, -} from "@/components/assistant-ui/model-selector/remembered-load-settings"; import { Dialog, DialogContent, @@ -29,7 +17,7 @@ import { DropdownMenuSeparator, DropdownMenuTrigger, } from "@/components/ui/dropdown-menu"; -import { Input } from "@/components/ui/input"; +import { InfoHint } from "@/components/ui/info-hint"; import { InputGroup, InputGroupAddon, @@ -50,26 +38,22 @@ import { SheetTitle, } from "@/components/ui/sheet"; import { Slider } from "@/components/ui/slider"; -import { Spinner } from "@/components/ui/spinner"; import { Switch } from "@/components/ui/switch"; import { Textarea } from "@/components/ui/textarea"; -import { InfoHint } from "@/components/ui/info-hint"; import { Tooltip, TooltipContent } from "@/components/ui/tooltip"; -import { useIsMobile } from "@/hooks/use-mobile"; +import { NumericValueInput, snapToStep } from "@/features/model-picker"; +import { RetrievalSettingsSection } from "@/features/rag"; import { useLlamaUpdateCheck } from "@/hooks/use-llama-update-check"; -import { cn } from "@/lib/utils"; -import { - ArrowTurnBackwardIcon, - Edit03Icon, - LayoutAlignRightIcon, -} from "@hugeicons/core-free-icons"; +import { useIsMobile } from "@/hooks/use-mobile"; import { ChevronDownStandardIcon } from "@/lib/chevron-icons"; +import { toast } from "@/lib/toast"; +import { cn } from "@/lib/utils"; +import { Edit03Icon, LayoutAlignRightIcon } from "@hugeicons/core-free-icons"; import { HugeiconsIcon } from "@hugeicons/react"; import { Braces, ChevronDown, ExternalLink } from "lucide-react"; import { Tooltip as TooltipPrimitive } from "radix-ui"; import { Fragment, type ReactNode } from "react"; import { useCallback, useEffect, useMemo, useRef, useState } from "react"; -import { toast } from "@/lib/toast"; import { OpenAICodeExecSection } from "./components/openai-code-exec-section"; import { PermissionModeDropdown } from "./permission-mode-select"; import { resyncInferenceStatusAfterServerModelChange } from "./hooks/use-chat-model-runtime"; @@ -77,8 +61,8 @@ import { type ExternalProviderConfig, getExternalProviderApiKey, parseExternalModelId, - supportsProviderPromptCaching, supportsProviderPromptCacheTtl, + supportsProviderPromptCaching, } from "./external-providers"; import { BUILTIN_PRESETS, @@ -98,12 +82,7 @@ import { providerSupportsBuiltinCodeExecution, providerSupportsFastMode, } from "./provider-capabilities"; -import { - isPendingGguf, - pendingSelectionMatches, - useChatRuntimeStore, -} from "./stores/chat-runtime-store"; -import { RetrievalSettingsSection } from "@/features/rag/components/retrieval-settings-section"; +import { useChatRuntimeStore } from "./stores/chat-runtime-store"; import type { InferenceParams } from "./types/runtime"; export { defaultInferenceParams, type Preset } from "./presets/preset-policy"; @@ -126,7 +105,7 @@ function getPromptVariablesError(raw: string): string | null { return null; } } catch { - return "Use valid JSON, for example { \"env\": \"staging\" }."; + return 'Use valid JSON, for example { "env": "staging" }.'; } return "Variables must be a JSON object."; } @@ -135,112 +114,7 @@ function hasPromptVariableSyntax(prompt: string): boolean { return PROMPT_VARIABLE_PATTERN.test(prompt); } -/** - * Editable numeric value display, shared by every slider value and the Context - * Length input. An that looks like text (shows `displayValue ?? value`, - * so "Off"/"Max" labels render) until focus, when it swaps to the raw number, - * selects it, and accepts free text. Commits on blur/Enter, reverts on Escape. - * Clamping happens on commit so typing intermediate values isn't fought. - */ -function snapToStep( - value: number, - step: number, - min?: number, - max?: number, -): number { - const lo = min ?? Number.NEGATIVE_INFINITY; - const hi = max ?? Number.POSITIVE_INFINITY; - const clamped = Math.min(Math.max(value, lo), hi); - const stepStr = String(step); - const decimals = stepStr.includes(".") ? stepStr.split(".")[1].length : 0; - const base = Number.isFinite(lo) ? lo : 0; - const snapped = base + Math.round((clamped - base) / step) * step; - const reclamped = Math.min(Math.max(snapped, lo), hi); - return Number(reclamped.toFixed(decimals)); -} - -function NumericValueInput({ - value, - min, - max, - step, - onChange, - displayValue, - className, - ariaLabel, - size: sizeAttr, - disabled = false, -}: { - value: number; - min?: number; - max?: number; - step: number; - onChange: (v: number) => void; - displayValue?: string; - className?: string; - ariaLabel?: string; - size?: number; - disabled?: boolean; -}) { - const [focused, setFocused] = useState(false); - const [draft, setDraft] = useState(""); - const cancelBlurCommitRef = useRef(false); - - const commit = (raw: string) => { - const parsed = Number.parseFloat(raw); - if (!Number.isFinite(parsed)) { - return; - } - const final = snapToStep(parsed, step, min, max); - if (final !== value) { - onChange(final); - } - }; - - const displayed = focused ? draft : (displayValue ?? String(value)); - - return ( - { - cancelBlurCommitRef.current = false; - setDraft(String(value)); - setFocused(true); - // Defer select() so it runs after the value swap above. - const target = e.currentTarget; - requestAnimationFrame(() => target.select()); - }} - onBlur={() => { - if (cancelBlurCommitRef.current) { - cancelBlurCommitRef.current = false; - } else { - commit(draft); - } - setFocused(false); - }} - onChange={(e) => setDraft(e.target.value)} - onKeyDown={(e) => { - if (e.key === "Enter") { - e.currentTarget.blur(); - } else if (e.key === "Escape") { - cancelBlurCommitRef.current = true; - setDraft(String(value)); - e.currentTarget.blur(); - } - }} - className={cn("panel-number-input", className)} - /> - ); -} - -function ParamSlider({ +export function ParamSlider({ label, value, min, @@ -250,6 +124,7 @@ function ParamSlider({ displayValue, info, valueSize, + disabled, }: { label: string; value: number; @@ -260,6 +135,7 @@ function ParamSlider({ displayValue?: string; info?: ReactNode; valueSize?: number; + disabled?: boolean; }) { return (
@@ -279,6 +155,8 @@ function ParamSlider({ displayValue={displayValue} ariaLabel={label} size={valueSize ?? 4} + className="panel-number-input" + disabled={disabled} />
onChange(snapToStep(v, step, min, max))} className="panel-slider" + disabled={disabled} /> ); @@ -377,8 +256,7 @@ function CollapsibleSection({ return (
{labelHref ? ( @@ -450,6 +328,7 @@ interface ChatSettingsPanelProps { onOpenChange?: (open: boolean) => void; params: InferenceParams; onParamsChange: (params: InferenceParams) => void; + modelConfig?: ReactNode; isExternalModel?: boolean; /** * Sampling-param capabilities for the active external provider, or `null` for @@ -464,21 +343,6 @@ interface ChatSettingsPanelProps { * Max Tokens floor in the slider. */ externalProviderType?: string | null; - onReloadModel?: () => void; - /** The in-flight load (id + GGUF variant + native path token), or null when - * idle. Used to show a loading state for the staged pick only — not for an - * unrelated load or a cancel's background unload. */ - loadingModel?: { - id: string; - ggufVariant?: string | null; - nativePathToken?: string | null; - } | null; - /** Loads the staged `pendingSelection` (deferred "Load on selection" flow). */ - onLoadPendingModel?: () => void; - /** Download progress (0–1) for a staged GGUF being fetched, or null when idle. */ - stagedDownloadFraction?: number | null; - /** Cancels the in-flight staged download (paired with abandoning the stage). */ - onCancelStagedDownload?: () => void; } export function ChatSettingsPanel({ @@ -486,16 +350,12 @@ export function ChatSettingsPanel({ onOpenChange, params, onParamsChange, + modelConfig = null, isExternalModel = false, providerCapabilities = null, activeExternalProvider = null, onExternalProviderChange, externalProviderType = null, - onReloadModel, - loadingModel = null, - onLoadPendingModel, - stagedDownloadFraction, - onCancelStagedDownload, }: ChatSettingsPanelProps) { // Local models show every knob; providerCapabilities is only consulted when // isExternalModel. Unknown providers fall back to the OpenAI-compat shape via @@ -510,55 +370,23 @@ export function ChatSettingsPanel({ const showPresencePenalty = !isExternalModel || Boolean(providerCapabilities?.presencePenalty); const isMobile = useIsMobile(); - const pendingSelection = useChatRuntimeStore((s) => s.pendingSelection); - // "Loading" only when the in-flight load IS this staged pick (full id + GGUF - // variant + native token match), not an unrelated load or a cancel's - // background unload. The variant matters: a different quant of the same repo - // staged mid-load must not read as this one loading. - const stagedLoading = - loadingModel != null && - pendingSelectionMatches(pendingSelection, { - id: loadingModel.id, - ggufVariant: loadingModel.ggufVariant, - nativePathToken: loadingModel.nativePathToken, - }); - // Load settings are snapshotted at click time; lock them while loading. - const modelControlsDisabled = stagedLoading; - const abandonStagedModel = useChatRuntimeStore((s) => s.abandonStagedModel); - const resetModelSettingsToLoaded = useChatRuntimeStore( - (s) => s.resetModelSettingsToLoaded, + const isLoadedGguf = useChatRuntimeStore((s) => s.activeGgufVariant) != null; + const currentCheckpoint = params.checkpoint; + const ggufContextLength = useChatRuntimeStore((s) => s.ggufContextLength); + // Direct-file / custom-folder GGUFs load without a variant label but still + // report a GGUF context, so detect them via the context and the checkpoint + // suffix too (mirrors the chat page's activeModelIsGguf). Otherwise Max Tokens + // would fall back to params.maxSeqLength instead of the loaded GGUF context. + const isGguf = + isLoadedGguf || + ggufContextLength != null || + (currentCheckpoint?.toLowerCase().endsWith(".gguf") ?? false); + const ggufMaxContextLength = useChatRuntimeStore( + (s) => s.ggufMaxContextLength, ); - // A staged GGUF pick (deferred load) shows the GGUF load knobs so they can be - // set before the single load. - const pendingIsGguf = isPendingGguf(pendingSelection); - // Short, human-readable name for the staged pick (HF ids carry an org prefix; - // native picks are already a display label). Drives the "staged, not loaded" - // callout so it's obvious the selection hasn't loaded yet. - const stagedLabel = (() => { - const id = pendingSelection?.id ?? ""; - const slash = id.lastIndexOf("/"); - const base = slash >= 0 ? id.slice(slash + 1) : id; - return base || id; - })(); - const isLoadedGguf = - useChatRuntimeStore((s) => s.activeGgufVariant) != null; - // While a pick is staged the sheet configures *that* model, so its GGUF-ness - // (not the currently loaded model's) decides whether the GGUF-only controls - // show. Otherwise a staged non-GGUF Hub repo would inherit the loaded GGUF's - // context/KV/speculative controls. - const isGguf = pendingSelection != null ? pendingIsGguf : isLoadedGguf; - // The Model section (and Load button) shows for any staged pick, even when the - // currently active model is external. - const hasModelContent = - pendingSelection != null || - (!isExternalModel && (isGguf || Boolean(params.checkpoint))); + const customContextLength = useChatRuntimeStore((s) => s.customContextLength); const speculativeType = useChatRuntimeStore((s) => s.speculativeType); - const setSpeculativeType = useChatRuntimeStore((s) => s.setSpeculativeType); - const loadedSpeculativeType = useChatRuntimeStore( - (s) => s.loadedSpeculativeType, - ); const specFallbackReason = useChatRuntimeStore((s) => s.specFallbackReason); - // Only binary fallback states are solved by a newer prebuilt. const mtpUpdatable = specFallbackReason === "binary_no_mtp" || specFallbackReason === "binary_outdated"; @@ -580,43 +408,27 @@ export function ChatSettingsPanel({ `llama.cpp updated to ${result.tag ?? "the latest build"}.${reloadHint}`, ); } else { - toast.error(`llama.cpp update failed: ${result.error ?? "unknown error"}`); + 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( - (s) => s.loadedSpecDraftNMax, - ); - const currentCheckpoint = params.checkpoint; - const ggufContextLength = useChatRuntimeStore((s) => s.ggufContextLength); - const ggufMaxContextLength = useChatRuntimeStore( - (s) => s.ggufMaxContextLength, - ); - const ggufNativeContextLength = useChatRuntimeStore( - (s) => s.ggufNativeContextLength, - ); - const kvCacheDtype = useChatRuntimeStore((s) => s.kvCacheDtype); - const setKvCacheDtype = useChatRuntimeStore((s) => s.setKvCacheDtype); - const applyRememberedLoadSettings = useChatRuntimeStore( - (s) => s.applyRememberedLoadSettings, - ); - const loadedKvCacheDtype = useChatRuntimeStore((s) => s.loadedKvCacheDtype); - const tensorParallel = useChatRuntimeStore((s) => s.tensorParallel); - const setTensorParallel = useChatRuntimeStore((s) => s.setTensorParallel); - const loadedTensorParallel = useChatRuntimeStore( - (s) => s.loadedTensorParallel, - ); - const chatTemplateOverride = useChatRuntimeStore( - (s) => s.chatTemplateOverride, - ); - const loadedChatTemplateOverride = useChatRuntimeStore( - (s) => s.loadedChatTemplateOverride, - ); - const customContextLength = useChatRuntimeStore((s) => s.customContextLength); - const setCustomContextLength = useChatRuntimeStore( - (s) => s.setCustomContextLength, - ); + const loadedEffectiveContext = customContextLength ?? ggufContextLength; + const showSpecFallback = + !isExternalModel && + isGguf && + specFallbackReason != null && + (speculativeType === "auto" || + speculativeType === "mtp" || + speculativeType === "mtp+ngram"); + const showContextVramWarning = + !isExternalModel && + isGguf && + ggufMaxContextLength != null && + loadedEffectiveContext != null && + loadedEffectiveContext > ggufMaxContextLength; + const showLoadedDiagnostics = showSpecFallback || showContextVramWarning; + const hasModelContent = showLoadedDiagnostics; const setActivePresetSource = useChatRuntimeStore( (s) => s.setActivePresetSource, ); @@ -627,49 +439,7 @@ export function ChatSettingsPanel({ const setActivePreset = useChatRuntimeStore((s) => s.setActivePreset); const settingsHydrated = useChatRuntimeStore((s) => s.settingsHydrated); - // A staged (not-yet-loaded) GGUF carries its own header context length on - // pendingSelection, so the slider can use the staged model's real ceiling - // without reading the loaded model's `ggufContextLength`. - const stagedContextLength = pendingSelection?.contextLength ?? null; - // "Remember settings next time" tick for a staged model. Seeds the store from - // the saved per-model settings on stage, so the sheet opens with what was used - // last time; the tick reflects whether a saved entry exists. - const [remember, setRemember] = useState(false); - // Keyed per quant: a different variant of the same repo has its own settings. - const pendingKey = pendingSelection - ? rememberedLoadSettingsKey(pendingSelection) - : null; - useEffect(() => { - if (!pendingKey) return; - const saved = loadRememberedLoadSettings(pendingKey); - setRemember(saved != null); - if (saved) applyRememberedLoadSettings(saved); - }, [pendingKey, applyRememberedLoadSettings]); - // While staging, the sheet reflects the STAGED model, so its header context - // takes precedence over the loaded model's (which may differ or be larger). - const baseContext = pendingIsGguf ? stagedContextLength : ggufContextLength; - const baseNativeContext = pendingIsGguf - ? stagedContextLength - : ggufNativeContextLength; - // Context controls render once we actually have a ceiling: for a staged GGUF, - // once its header metadata arrives (post-download); otherwise post-load. - const showContextControl = pendingIsGguf - ? stagedContextLength != null - : isLoadedGguf; - const stagedDownloading = - stagedDownloadFraction != null && stagedDownloadFraction < 1; - const ctxDisplayValue = customContextLength ?? baseContext ?? ""; - const ctxMaxValue = baseNativeContext ?? baseContext ?? null; - const kvDirty = kvCacheDtype !== loadedKvCacheDtype; - const ctxDirty = customContextLength !== null; - const specDirty = speculativeType !== loadedSpeculativeType; - const specDraftDirty = specDraftNMax !== loadedSpecDraftNMax; - const tpDirty = tensorParallel !== (loadedTensorParallel ?? false); - // A saved chat-template override is a reload-time setting too, so surface - // Apply for a template-only edit (otherwise it could never be applied). - const templateDirty = chatTemplateOverride !== loadedChatTemplateOverride; - const modelSettingsDirty = - kvDirty || ctxDirty || specDirty || specDraftDirty || tpDirty || templateDirty; + const baseContext = ggufContextLength; const [presetNameInput, setPresetNameInput] = useState(activePreset); const [systemPromptEditorOpen, setSystemPromptEditorOpen] = useState(false); const [systemPromptDraft, setSystemPromptDraft] = useState(""); @@ -695,8 +465,7 @@ export function ChatSettingsPanel({ BUILTIN_PRESETS.find((preset) => preset.name === activePreset) ?? null, [activePreset], ); - const hasUnsavedPresetChanges = useMemo( - () => { + const hasUnsavedPresetChanges = useMemo(() => { if (activePresetDefinition == null) { return false; } @@ -704,9 +473,7 @@ export function ChatSettingsPanel({ return activePresetSource === "modified"; } return !isSamePresetConfig(activePresetDefinition.params, params); - }, - [activePresetDefinition, activePresetSource, params], - ); + }, [activePresetDefinition, activePresetSource, params]); const presetSaveState = useMemo( () => getPresetSaveState({ @@ -735,6 +502,14 @@ export function ChatSettingsPanel({ const externalSelection = currentCheckpoint ? parseExternalModelId(currentCheckpoint) : null; + const maxTokensMax = isExternalModel + ? getExternalMaxOutputTokens( + externalProviderType, + externalSelection?.modelId, + ) + : isGguf && baseContext + ? baseContext + : Math.max(64, params.maxSeqLength); const showOpenAICodeExecSection = activeExternalProvider != null && providerSupportsBuiltinCodeExecution( @@ -817,8 +592,7 @@ export function ChatSettingsPanel({ return; } const fallbackPreset = - BUILTIN_PRESETS.find((preset) => preset.name === "Default") ?? - null; + BUILTIN_PRESETS.find((preset) => preset.name === "Default") ?? null; const next = customPresets.filter((preset) => preset.name !== name); setCustomPresets(next); if (activePreset === name) { @@ -930,7 +704,7 @@ export function ChatSettingsPanel({ Run settings - + - )} -
- )} - {(speculativeType === "mtp" || - speculativeType === "mtp+ngram") && ( -
-
- - Draft Tokens - - - Max MTP draft tokens per step - (--spec-draft-n-max). Lower = less wasted - draft decode; higher = bigger speedup when - acceptance stays high. Default: 2 on GPU, - 3 on CPU/Mac. - -
- { - const raw = e.target.value; - if (raw === "") { - setSpecDraftNMax(null); - return; - } - const parsed = Number.parseInt(raw, 10); - if (Number.isFinite(parsed)) { - const clamped = Math.max(1, Math.min(16, parsed)); - setSpecDraftNMax(clamped); - } - }} - data-test-id="spec-draft-n-max-input" - aria-label="Speculative decoding draft tokens" - className="h-7 w-[88px] rounded-full border-border bg-background hover:bg-accent/50 dark:border-transparent dark:bg-white/[0.05] dark:hover:bg-white/[0.1] pl-3 py-0 text-[13px] font-medium text-nav-fg outline-none focus-visible:ring-0" - /> -
- )} - - )} -
-
- - Tensor Parallelism - - - No effect on a single GPU. On multi-GPU setups, improves - tokens/sec during generation when using dense models. MoE - models don't benefit and can be much slower. - -
- -
- - )} - {/* No persistent "enable custom code" toggle: it is consented per model - via the load-time review dialog. */} - {/* Apply/Reset belongs to the model-reload settings above (context - length, KV cache, speculative decoding). Render it here, before - the Chat Template row, so it never reads as attached to Chat - Template (which is edited via its own dialog). When a model is - staged (deferred load), Load/Cancel takes its place: there's - nothing loaded to "apply" against yet. */} - {pendingSelection ? ( -
- {stagedDownloading && ( -

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

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

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

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