diff --git a/images/unsloth logo only.png b/images/unsloth logo only.png index 92340eef05..adaafee48d 100644 Binary files a/images/unsloth logo only.png and b/images/unsloth logo only.png differ diff --git a/images/unsloth new logo.png b/images/unsloth new logo.png index 20dac04f38..adaafee48d 100644 Binary files a/images/unsloth new logo.png and b/images/unsloth new logo.png differ diff --git a/install.ps1 b/install.ps1 index 0c36046195..ead4e7368d 100644 --- a/install.ps1 +++ b/install.ps1 @@ -749,7 +749,6 @@ shell.Run cmd, 0, False } else { step "gpu" "none (chat-only / GGUF)" "Yellow" substep "Training and GPU inference require an NVIDIA GPU with drivers installed." "Yellow" - substep "https://www.nvidia.com/Download/index.aspx" "Yellow" } # ── Choose the correct PyTorch index URL based on driver CUDA version ── @@ -777,10 +776,10 @@ shell.Run cmd, 0, False # ── Print CPU-only hint when no GPU detected ── if (-not $SkipTorch -and $TorchIndexUrl -like "*/cpu") { Write-Host "" - Write-Host " NOTE: No NVIDIA GPU detected." -ForegroundColor Yellow - Write-Host " Installing CPU-only PyTorch. If you only need GGUF chat/inference," - Write-Host " re-run with --no-torch for a faster, lighter install:" - Write-Host " .\install.ps1 --no-torch" + substep "No NVIDIA GPU detected." "Yellow" + substep "Installing CPU-only PyTorch. If you only need GGUF chat/inference," "Yellow" + substep "re-run with --no-torch for a faster, lighter install:" "Yellow" + substep ".\install.ps1 --no-torch" "Yellow" Write-Host "" } @@ -820,7 +819,7 @@ shell.Run cmd, 0, False if ($SkipTorch) { # No-torch: install unsloth + unsloth-zoo with --no-deps, then # runtime deps (typer, safetensors, transformers, etc.) with --no-deps. - $baseInstallExit = Invoke-InstallCommand { uv pip install --python $VenvPython --no-deps --reinstall-package unsloth --reinstall-package unsloth-zoo "unsloth>=2026.3.16" unsloth-zoo } + $baseInstallExit = Invoke-InstallCommand { uv pip install --python $VenvPython --no-deps --reinstall-package unsloth --reinstall-package unsloth-zoo "unsloth>=2026.3.18" unsloth-zoo } if ($baseInstallExit -eq 0) { $NoTorchReq = Find-NoTorchRuntimeFile if ($NoTorchReq) { @@ -828,7 +827,7 @@ shell.Run cmd, 0, False } } } else { - $baseInstallExit = Invoke-InstallCommand { uv pip install --python $VenvPython --reinstall-package unsloth --reinstall-package unsloth-zoo "unsloth>=2026.3.16" unsloth-zoo } + $baseInstallExit = Invoke-InstallCommand { uv pip install --python $VenvPython --reinstall-package unsloth --reinstall-package unsloth-zoo "unsloth>=2026.3.18" unsloth-zoo } } if ($baseInstallExit -ne 0) { Write-Host "[ERROR] Failed to install unsloth (exit code $baseInstallExit)" -ForegroundColor Red @@ -858,7 +857,7 @@ shell.Run cmd, 0, False if ($SkipTorch) { # No-torch: install unsloth + unsloth-zoo with --no-deps, then # runtime deps (typer, safetensors, transformers, etc.) with --no-deps. - $baseInstallExit = Invoke-InstallCommand { uv pip install --python $VenvPython --no-deps --upgrade-package unsloth --upgrade-package unsloth-zoo "unsloth>=2026.3.16" unsloth-zoo } + $baseInstallExit = Invoke-InstallCommand { uv pip install --python $VenvPython --no-deps --upgrade-package unsloth --upgrade-package unsloth-zoo "unsloth>=2026.3.18" unsloth-zoo } if ($baseInstallExit -eq 0) { $NoTorchReq = Find-NoTorchRuntimeFile if ($NoTorchReq) { @@ -866,7 +865,7 @@ shell.Run cmd, 0, False } } } elseif ($StudioLocalInstall) { - $baseInstallExit = Invoke-InstallCommand { uv pip install --python $VenvPython --upgrade-package unsloth "unsloth>=2026.3.16" unsloth-zoo } + $baseInstallExit = Invoke-InstallCommand { uv pip install --python $VenvPython --upgrade-package unsloth "unsloth>=2026.3.18" unsloth-zoo } } else { $baseInstallExit = Invoke-InstallCommand { uv pip install --python $VenvPython --upgrade-package unsloth "$PackageName" } } @@ -887,7 +886,7 @@ shell.Run cmd, 0, False # Fallback: GPU detection failed to produce a URL -- let uv resolve torch substep "installing unsloth (this may take a few minutes)..." if ($StudioLocalInstall) { - $baseInstallExit = Invoke-InstallCommand { uv pip install --python $VenvPython unsloth-zoo "unsloth>=2026.3.16" --torch-backend=auto } + $baseInstallExit = Invoke-InstallCommand { uv pip install --python $VenvPython unsloth-zoo "unsloth>=2026.3.18" --torch-backend=auto } if ($baseInstallExit -ne 0) { Write-Host "[ERROR] Failed to install unsloth (exit code $baseInstallExit)" -ForegroundColor Red return diff --git a/install.sh b/install.sh index 9ea80bc161..2d9a368555 100755 --- a/install.sh +++ b/install.sh @@ -949,6 +949,17 @@ if [ -x "$VENV_DIR/bin/python" ]; then substep "${VENV_DIR}" fi +# Default torch constraint -- tightened for Python 3.13+ on arm64 macOS +# (torch <2.6 has no cp313 macOS arm64 wheels) +TORCH_CONSTRAINT="torch>=2.4,<2.11.0" +if [ "$SKIP_TORCH" = false ] && [ "$OS" = "macos" ] && [ "$_ARCH" = "arm64" ]; then + _PY_MINOR=$("$VENV_DIR/bin/python" -c \ + "import sys; print(sys.version_info.minor)" 2>/dev/null || echo "0") + if [ "$_PY_MINOR" -ge 13 ] 2>/dev/null; then + TORCH_CONSTRAINT="torch>=2.6,<2.11.0" + fi +fi + # ── Resolve repo root (for --local installs) ── _REPO_ROOT="$(cd "$(dirname "$0" 2>/dev/null || echo ".")" && pwd)" @@ -1029,7 +1040,7 @@ if [ "$_MIGRATED" = true ]; then # to prevent transitive torch resolution. run_install_cmd "install unsloth (migrated no-torch)" uv pip install --python "$_VENV_PY" --no-deps \ --reinstall-package unsloth --reinstall-package unsloth-zoo \ - "unsloth>=2026.3.16" unsloth-zoo + "unsloth>=2026.3.18" unsloth-zoo _NO_TORCH_RT="$(_find_no_torch_runtime)" if [ -n "$_NO_TORCH_RT" ]; then run_install_cmd "install no-torch runtime deps" uv pip install --python "$_VENV_PY" --no-deps -r "$_NO_TORCH_RT" @@ -1037,7 +1048,7 @@ if [ "$_MIGRATED" = true ]; then else run_install_cmd "install unsloth (migrated)" uv pip install --python "$_VENV_PY" \ --reinstall-package unsloth --reinstall-package unsloth-zoo \ - "unsloth>=2026.3.16" unsloth-zoo + "unsloth>=2026.3.18" unsloth-zoo fi if [ "$STUDIO_LOCAL_INSTALL" = true ]; then substep "overlaying local repo (editable)..." @@ -1049,7 +1060,7 @@ elif [ -n "$TORCH_INDEX_URL" ]; then substep "skipping PyTorch (--no-torch or Intel Mac x86_64)." "$C_WARN" else substep "installing PyTorch ($TORCH_INDEX_URL)..." - run_install_cmd "install PyTorch" uv pip install --python "$_VENV_PY" "torch>=2.4,<2.11.0" torchvision torchaudio \ + run_install_cmd "install PyTorch" uv pip install --python "$_VENV_PY" "$TORCH_CONSTRAINT" torchvision torchaudio \ --index-url "$TORCH_INDEX_URL" fi # Fresh: Step 2 - install unsloth, preserving pre-installed torch @@ -1059,7 +1070,7 @@ elif [ -n "$TORCH_INDEX_URL" ]; then # runtime deps (typer, safetensors, transformers, etc.) with --no-deps. run_install_cmd "install unsloth (no-torch)" uv pip install --python "$_VENV_PY" --no-deps \ --upgrade-package unsloth --upgrade-package unsloth-zoo \ - "unsloth>=2026.3.16" unsloth-zoo + "unsloth>=2026.3.18" unsloth-zoo _NO_TORCH_RT="$(_find_no_torch_runtime)" if [ -n "$_NO_TORCH_RT" ]; then run_install_cmd "install no-torch runtime deps" uv pip install --python "$_VENV_PY" --no-deps -r "$_NO_TORCH_RT" @@ -1070,7 +1081,7 @@ elif [ -n "$TORCH_INDEX_URL" ]; then fi elif [ "$STUDIO_LOCAL_INSTALL" = true ]; then run_install_cmd "install unsloth (local)" uv pip install --python "$_VENV_PY" \ - --upgrade-package unsloth "unsloth>=2026.3.16" unsloth-zoo + --upgrade-package unsloth "unsloth>=2026.3.18" unsloth-zoo substep "overlaying local repo (editable)..." run_install_cmd "overlay local repo" uv pip install --python "$_VENV_PY" -e "$_REPO_ROOT" --no-deps else @@ -1081,7 +1092,7 @@ else # Fallback: GPU detection failed to produce a URL -- let uv resolve torch substep "installing unsloth (this may take a few minutes)..." if [ "$STUDIO_LOCAL_INSTALL" = true ]; then - run_install_cmd "install unsloth (auto torch backend)" uv pip install --python "$_VENV_PY" unsloth-zoo "unsloth>=2026.3.16" --torch-backend=auto + run_install_cmd "install unsloth (auto torch backend)" uv pip install --python "$_VENV_PY" unsloth-zoo "unsloth>=2026.3.18" --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 else diff --git a/install_gemma4_mlx.sh b/install_gemma4_mlx.sh new file mode 100755 index 0000000000..e06339e204 --- /dev/null +++ b/install_gemma4_mlx.sh @@ -0,0 +1,215 @@ +#!/bin/bash +set -e + +# ============================================================ +# Gemma 4 MLX — One-command setup + inference +# +# Usage: +# bash install_gemma4_mlx.sh [--venv-dir DIR] +# +# This script: +# 1. Creates a Python virtual environment +# 2. Installs uv, mlx, mlx-lm, transformers +# 3. Downloads gemma4.py and gemma4_text.py from unsloth repo +# 4. Installs them into mlx-lm's models directory +# ============================================================ + +# ── Output style (inspired by unsloth/install.sh) ───────────── +RULE="" +_rule_i=0 +while [ "$_rule_i" -lt 52 ]; do + RULE="${RULE}─" + _rule_i=$((_rule_i + 1)) +done + +if [ -n "${NO_COLOR:-}" ]; then + C_TITLE= C_DIM= C_OK= C_WARN= C_ERR= C_RST= +elif [ -t 1 ] || [ -n "${FORCE_COLOR:-}" ]; then + _ESC="$(printf '\033')" + C_TITLE="${_ESC}[38;5;117m" + C_DIM="${_ESC}[38;5;245m" + C_OK="${_ESC}[38;5;108m" + C_WARN="${_ESC}[38;5;136m" + C_ERR="${_ESC}[91m" + C_RST="${_ESC}[0m" +else + C_TITLE= C_DIM= C_OK= C_WARN= C_ERR= C_RST= +fi + +step() { printf " ${C_DIM}%-18.18s${C_RST}${3:-$C_OK}%s${C_RST}\n" "$1" "$2"; } +substep() { printf " ${C_DIM}%-18s${2:-$C_DIM}%s${C_RST}\n" "" "$1"; } +fail() { step "error" "$1" "$C_ERR"; exit 1; } + +# ── Parse flags ─────────────────────────────────────────────── +VENV_DIR="" +_next_is_venv=false + +for arg in "$@"; do + if [ "$_next_is_venv" = true ]; then + VENV_DIR="$arg" + _next_is_venv=false + continue + fi + case "$arg" in + --venv-dir) _next_is_venv=true ;; + esac +done + +# Default venv location +if [ -z "$VENV_DIR" ]; then + VENV_DIR="$HOME/.unsloth/unsloth_gemma4_mlx" +fi + +# ── Banner ──────────────────────────────────────────────────── +echo "" +printf " ${C_TITLE}%s${C_RST}\n" "💎 Gemma 4 MLX Installer" +printf " ${C_DIM}%s${C_RST}\n" "$RULE" +echo "" + +# ── Platform check ──────────────────────────────────────────── +if [ "$(uname)" != "Darwin" ]; then + fail "MLX requires macOS with Apple Silicon. Detected: $(uname)" +fi + +_ARCH=$(uname -m) +if [ "$_ARCH" != "arm64" ]; then + step "warning" "Apple Silicon recommended (detected: $_ARCH)" "$C_WARN" +fi + +step "platform" "macOS ($_ARCH)" + +# ── Detect Python ───────────────────────────────────────────── +PYTHON="" +for _candidate in python3.12 python3.11 python3.13 python3; do + if command -v "$_candidate" >/dev/null 2>&1; then + PYTHON="$_candidate" + break + fi +done + +if [ -z "$PYTHON" ]; then + fail "Python 3 not found. Install via: brew install python@3.12" +fi + +_PY_VERSION=$("$PYTHON" -c "import sys; print(f'{sys.version_info.major}.{sys.version_info.minor}.{sys.version_info.micro}')") +step "python" "$PYTHON ($_PY_VERSION)" + +# ── Create virtual environment ──────────────────────────────── +if [ -x "$VENV_DIR/bin/python" ]; then + step "venv" "using existing environment" + substep "$VENV_DIR" +else + step "venv" "creating virtual environment" + substep "$VENV_DIR" + mkdir -p "$(dirname "$VENV_DIR")" + "$PYTHON" -m venv "$VENV_DIR" +fi + +# ── Install uv ─────────────────────────────────────────────── +if ! command -v uv >/dev/null 2>&1; then + step "uv" "installing uv package manager..." + _uv_tmp=$(mktemp) + curl -LsSf "https://astral.sh/uv/install.sh" -o "$_uv_tmp" + sh "$_uv_tmp" /dev/null 2>&1 + rm -f "$_uv_tmp" + if [ -f "$HOME/.local/bin/env" ]; then + . "$HOME/.local/bin/env" + fi + export PATH="$HOME/.local/bin:$PATH" + substep "done" +else + step "uv" "found $(uv --version 2>/dev/null || echo 'uv')" +fi + +_VENV_PY="$VENV_DIR/bin/python" + +# ── Repo config ────────────────────────────────────────────── +BRANCH="fix/ui-fix" +REPO_URL="https://raw.githubusercontent.com/unslothai/unsloth/refs/heads/${BRANCH}" + +# ── Install dependencies ────────────────────────────────────── +step "install" "installing mlx, mlx-lm..." +uv pip install --python "$_VENV_PY" -q mlx mlx-lm 2>/dev/null +substep "done" + +TRANSFORMERS_WHL="transformers-5.5.0-py3-none-any.whl" +TRANSFORMERS_GH="git+https://github.com/huggingface/transformers.git@v5.5-release" + +step "install" "installing transformers>=5.5.0..." +if uv pip install --python "$_VENV_PY" -q "$TRANSFORMERS_GH" 2>/dev/null; then + substep "installed from huggingface/transformers v5.5-release" +elif uv pip install --python "$_VENV_PY" -q "transformers>=5.5.0" 2>/dev/null; then + substep "installed from PyPI" +else + substep "not on PyPI, trying unsloth branch..." + _whl_tmp=$(mktemp -d)/"${TRANSFORMERS_WHL}" + if curl -fsSL "${REPO_URL}/${TRANSFORMERS_WHL}" -o "$_whl_tmp" 2>/dev/null && \ + uv pip install --python "$_VENV_PY" -q "$_whl_tmp" 2>/dev/null; then + substep "installed from branch ${BRANCH}" + elif [ -f "./${TRANSFORMERS_WHL}" ]; then + substep "using local ./${TRANSFORMERS_WHL}" + uv pip install --python "$_VENV_PY" -q "./${TRANSFORMERS_WHL}" + else + rm -f "$_whl_tmp" 2>/dev/null + step "install" "skipping transformers — could not find >=5.5.0" "$C_WARN" + substep "tried: huggingface/transformers v5.5-release, PyPI, branch ${BRANCH}, local ./${TRANSFORMERS_WHL}" + fi + rm -f "$_whl_tmp" 2>/dev/null +fi + +# ── Find mlx-lm models directory ───────────────────────────── +MLX_MODELS=$("$_VENV_PY" -c "import mlx_lm; print(mlx_lm.__path__[0])")/models +step "models dir" "$MLX_MODELS" + +# ── Download and install Gemma 4 model files ────────────────── + +step "download" "installing Gemma 4 model files..." + +_install_model_file() { + _fname="$1" + if curl -fsSL "${REPO_URL}/unsloth/models/${_fname}" -o "${MLX_MODELS}/${_fname}" 2>/dev/null; then + substep "downloaded ${_fname} from branch ${BRANCH}" + elif [ -f "./${_fname}" ]; then + substep "using local ./${_fname}" + cp "./${_fname}" "${MLX_MODELS}/${_fname}" + else + fail "Could not install ${_fname}. Tried: + 1) ${REPO_URL}/unsloth/models/${_fname} + 2) Local file ./${_fname} + + To fix, download the file manually and place it in the current directory, + then re-run this script." + fi +} + +_install_model_file "gemma4.py" +_install_model_file "gemma4_text.py" + +# Verify files were installed correctly +if "$_VENV_PY" -c "from mlx_lm.models.gemma4_text import ProportionalRoPE" 2>/dev/null; then + substep "model files verified" +else + fail "Model files installed but verification failed (ProportionalRoPE import error). + Try manually from: https://github.com/unslothai/unsloth/tree/feature/${BRANCH}" +fi + +# ── Done ────────────────────────────────────────────────────── +echo "" +printf " ${C_TITLE}%s${C_RST}\n" "Gemma 4 MLX installed!" +printf " ${C_DIM}%s${C_RST}\n" "$RULE" +echo "" +step "available models" "unsloth/gemma-4-E2B-it-UD-MLX-4bit (/BF16)" +substep "unsloth/gemma-4-E4B-it-UD-MLX-4bit (/BF16)" +echo "" +step "venv activate" "source ${VENV_DIR}/bin/activate" +echo "" +step "quick start" "python -m mlx_lm chat --model unsloth/gemma-4-E2B-it-UD-MLX-4bit --max-tokens 200" +echo "" +step "python API" "from mlx_lm import load, generate" +substep "model, tokenizer = load('unsloth/gemma-4-E2B-it-UD-MLX-4bit')" +substep "messages = [{'role': 'user', 'content': 'Hello!'}]" +substep "prompt = tokenizer.apply_chat_template(messages, tokenize=False, add_generation_prompt=True)" +substep "print(generate(model, tokenizer, prompt=prompt, max_tokens=200))" +echo "" +printf " ${C_DIM}%s${C_RST}\n" "$RULE" +echo "" diff --git a/studio/backend/core/inference/_html_to_md.py b/studio/backend/core/inference/_html_to_md.py index d96b8168e2..f999120ffb 100644 --- a/studio/backend/core/inference/_html_to_md.py +++ b/studio/backend/core/inference/_html_to_md.py @@ -17,18 +17,26 @@ from html.parser import HTMLParser __all__ = ["html_to_markdown"] -_SKIP_TAGS = frozenset({"script", "style", "head", "noscript", "svg", "math"}) +_SKIP_TAGS = frozenset( + { + "script", + "style", + "head", + "noscript", + "svg", + "math", + "nav", + "footer", + } +) _BLOCK_TAGS = frozenset( { "p", "div", "section", "article", - "header", - "footer", "main", "aside", - "nav", "figure", "figcaption", "details", diff --git a/studio/backend/core/inference/defaults.py b/studio/backend/core/inference/defaults.py index d5e9ca2e97..f3026dddaf 100644 --- a/studio/backend/core/inference/defaults.py +++ b/studio/backend/core/inference/defaults.py @@ -6,6 +6,14 @@ import utils.hardware.hardware as hw DEFAULT_MODELS_GGUF = [ + "unsloth/gemma-4-E2B-it-GGUF", + "unsloth/gemma-4-E4B-it-GGUF", + "unsloth/gemma-4-31B-it-GGUF", + "unsloth/gemma-4-26B-A4B-it-GGUF", + "unsloth/Qwen3.5-4B-GGUF", + "unsloth/Qwen3.5-9B-GGUF", + "unsloth/Qwen3.5-35B-A3B-GGUF", + "unsloth/Qwen3.5-0.8B-GGUF", "unsloth/Llama-3.2-1B-Instruct-GGUF", "unsloth/Llama-3.2-3B-Instruct-GGUF", "unsloth/Llama-3.1-8B-Instruct-GGUF", @@ -15,6 +23,18 @@ DEFAULT_MODELS_GGUF = [ ] DEFAULT_MODELS_STANDARD = [ + "unsloth/gemma-4-E2B-it-GGUF", + "unsloth/gemma-4-E4B-it-GGUF", + "unsloth/gemma-4-31B-it-GGUF", + "unsloth/gemma-4-26B-A4B-it-GGUF", + "unsloth/Qwen3.5-4B-GGUF", + "unsloth/Qwen3.5-9B-GGUF", + "unsloth/Qwen3.5-35B-A3B-GGUF", + "unsloth/Qwen3.5-0.8B-GGUF", + "unsloth/gemma-4-E2B-it", + "unsloth/gemma-4-E4B-it", + "unsloth/gemma-4-31B-it", + "unsloth/gemma-4-26B-A4B-it", "unsloth/Qwen3-4B-Instruct-2507", "unsloth/Meta-Llama-3.1-8B-Instruct-bnb-4bit", "unsloth/Mistral-Nemo-Instruct-2407-bnb-4bit", diff --git a/studio/backend/core/inference/llama_cpp.py b/studio/backend/core/inference/llama_cpp.py index 894dd25cf7..44c700bf3d 100644 --- a/studio/backend/core/inference/llama_cpp.py +++ b/studio/backend/core/inference/llama_cpp.py @@ -27,6 +27,52 @@ import httpx logger = get_logger(__name__) +# ── Pre-compiled patterns for plan-without-action re-prompt ── +# Forward-looking intent signals that indicate the model is +# describing what it *will* do rather than giving a final answer. +_INTENT_SIGNAL = re.compile( + r"(?i)(" + # Direct intent: "I'll ...", "I will ...", "Let me ...", "I am going to ..." + # Handles both straight and curly apostrophes. + # Excludes "I can", "I should", "I want to", "let's" which + # appear frequently in direct answers / explanations. + r"\b(i['\u2019](ll|m going to|m gonna)|i am (going to|gonna)|i will|i shall|let me|allow me)\b" + r"|" + # Step/plan framing: "First ...", "Step 1:", "Here's my plan" + r"\b(?:first\b|step \d+:?|here['\u2019]?s (?:my |the |a )?(?:plan|approach))" + r"|" + # "Now I" / "Next I" patterns + r"\b(?:now i|next i)\b" + r")" +) +_MAX_REPROMPTS = 3 +_REPROMPT_MAX_CHARS = 2000 + +# ── Pre-compiled patterns for GGUF shard detection ─────────── +_SHARD_FULL_RE = re.compile(r"^(.*)-(\d{5})-of-(\d{5})\.gguf$") +_SHARD_RE = re.compile(r"^(.*)-\d{5}-of-\d{5}\.gguf$") + +# Model size extraction (shared with routes/inference.py) +from utils.models import extract_model_size_b as _extract_model_size_b + +# ── Pre-compiled patterns for tool XML stripping ───────────── +_TOOL_CLOSED_PATS = [ + re.compile(r".*?", re.DOTALL), + re.compile(r".*?", re.DOTALL), +] +_TOOL_ALL_PATS = _TOOL_CLOSED_PATS + [ + re.compile(r".*$", re.DOTALL), + re.compile(r".*$", re.DOTALL), +] + +# ── Pre-compiled patterns for tool-call XML parsing ────────── +_TC_JSON_START_RE = re.compile(r"\s*\{") +_TC_FUNC_START_RE = re.compile(r"\s*") +_TC_END_TAG_RE = re.compile(r"") +_TC_FUNC_CLOSE_RE = re.compile(r"\s*\s*$") +_TC_PARAM_START_RE = re.compile(r"\s*") +_TC_PARAM_CLOSE_RE = re.compile(r"\s*\s*$") + class LlamaCppBackend: """ @@ -61,6 +107,15 @@ class LlamaCppBackend: self._n_kv_heads: Optional[int] = None self._n_heads: Optional[int] = None self._embedding_length: Optional[int] = None + # Architecture-aware KV fields (8 new fields for 5-path estimation) + self._kv_key_length: Optional[int] = None + self._kv_value_length: Optional[int] = None + self._sliding_window: Optional[int] = None + self._full_attention_interval: Optional[int] = None + self._kv_lora_rank: Optional[int] = None + self._key_length_mla: Optional[int] = None + self._ssm_inner_size: Optional[int] = None + self._ssm_state_size: Optional[int] = None self._lock = threading.Lock() self._stdout_lines: list[str] = [] self._stdout_thread: Optional[threading.Thread] = None @@ -107,6 +162,11 @@ class LlamaCppBackend: """Return the maximum context currently available on this hardware.""" return self._max_context_length or self._context_length + @property + def native_context_length(self) -> Optional[int]: + """Return the model's native context length from GGUF metadata.""" + return self._context_length + @property def chat_template(self) -> Optional[str]: return self._chat_template @@ -228,14 +288,11 @@ class LlamaCppBackend: @staticmethod def _get_gguf_size_bytes(model_path: str) -> int: """Get total GGUF size in bytes, including split shards.""" - import re - main = Path(model_path) total = main.stat().st_size # Check for split shards (e.g., model-00001-of-00003.gguf) - shard_pat = re.compile(r"^(.*)-(\d{5})-of-(\d{5})\.gguf$") - m = shard_pat.match(main.name) + m = _SHARD_FULL_RE.match(main.name) if m: prefix, _, num_total = m.group(1), m.group(2), m.group(3) sibling_pat = re.compile( @@ -306,11 +363,11 @@ class LlamaCppBackend: """Pick GPU(s) for a model based on estimated VRAM and free memory. ``model_size_bytes`` should include both model weights and estimated - KV cache. The 70% threshold provides headroom for compute buffers, + KV cache. The 90% threshold provides headroom for compute buffers, CUDA context, and other runtime overhead. Returns (gpu_indices, use_fit): - - ([1], False) model fits on 1 GPU at 70% of free + - ([1], False) model fits on 1 GPU at 90% of free - ([1, 2], False) model needs 2 GPUs - (None, True) model too large, let --fit handle it """ @@ -322,8 +379,8 @@ class LlamaCppBackend: # Sort GPUs by free memory descending ranked = sorted(gpus, key = lambda g: g[1], reverse = True) - # Try fitting on 1 GPU (70% of free memory threshold) - if ranked[0][1] * 0.70 >= model_size_mib: + # Try fitting on 1 GPU (90% of free memory threshold) + if ranked[0][1] * 0.90 >= model_size_mib: return [ranked[0][0]], False # Try fitting on N GPUs (accumulate free memory from most-free) @@ -331,7 +388,7 @@ class LlamaCppBackend: selected = [] for idx, free_mib in ranked: selected.append(idx) - cumulative += free_mib * 0.70 + cumulative += free_mib * 0.90 if cumulative >= model_size_mib: return sorted(selected), False @@ -347,10 +404,17 @@ class LlamaCppBackend: def _can_estimate_kv(self) -> bool: """True if we have enough GGUF metadata to estimate KV cache size.""" - return ( - self._n_layers is not None - and self._embedding_length is not None - and (self._n_kv_heads is not None or self._n_heads is not None) + if self._n_layers is None: + return False + # MLA: kv_lora_rank is sufficient (K-only cache) + if self._kv_lora_rank is not None: + return True + # New-style: need both explicit key AND value dimensions + if self._kv_key_length is not None and self._kv_value_length is not None: + return True + # Legacy: need embedding_length + head count + return self._embedding_length is not None and ( + self._n_kv_heads is not None or self._n_heads is not None ) def _estimate_kv_cache_bytes( @@ -358,14 +422,20 @@ class LlamaCppBackend: ) -> int: """Estimate KV cache VRAM for a given context length. + Uses 5-path architecture-aware estimation: + 1. MLA -- compressed KV latent + RoPE, K-only (no separate V) + 2. Hybrid -- only attention layers need KV (Mamba layers don't) + 3. SWA -- sliding-window layers cache min(ctx, window) tokens + 4. GQA -- standard full KV with explicit key/value dimensions + 5. Legacy -- fallback using embed // n_heads + Returns 0 if metadata is insufficient for estimation. """ if not self._can_estimate_kv() or n_ctx <= 0: return 0 n_layers = self._n_layers # type: ignore[assignment] - n_kv_heads = self._n_kv_heads or self._n_heads # type: ignore[assignment] - head_dim = self._embedding_length // self._n_heads if self._n_heads else 128 # type: ignore[operator] + n_kv = self._n_kv_heads or self._n_heads or 1 # type: ignore[assignment] # Bytes per element depends on KV cache quantization bpe = { @@ -380,8 +450,60 @@ class LlamaCppBackend: "iq4_nl": 0.5625, }.get(cache_type_kv or "f16", 2.0) - # K + V caches: 2 * n_kv_heads * head_dim * n_layers * n_ctx * bpe - return int(2 * n_kv_heads * head_dim * n_layers * n_ctx * bpe) + # Path 1: MLA (DeepSeek-V2/V3, GLM-4.7, GLM-5, Kimi-K2.5) + # MLA stores one compressed KV latent per token/layer (shared across heads). + # V is reconstructed from the latent on the fly -- no separate V cache. + # key_length = kv_lora_rank + rope_dim (the full compressed representation). + # MLA GGUFs set head_count_kv=1; default to 1 if absent to avoid + # falling back to n_heads (e.g., 128 for DeepSeek-V3) which would 128x. + if self._kv_lora_rank is not None: + n_kv_mla = self._n_kv_heads or 1 + rope_dim = self._key_length_mla or 64 + key_len = self._kv_key_length or (self._kv_lora_rank + rope_dim) + return int(n_layers * n_ctx * n_kv_mla * key_len * bpe) + + key_len = self._kv_key_length + val_len = self._kv_value_length + + # Path 2: Hybrid Mamba/Attention (Qwen3.5-27B, Qwen3.5-35B-A3B) + # Only 1 in N layers is attention; the rest are Mamba (no KV cache). + if ( + self._ssm_inner_size is not None + and self._full_attention_interval is not None + ): + fai = self._full_attention_interval + n_attn = -(-n_layers // fai) if fai > 0 else n_layers # ceiling division + if key_len is not None and val_len is not None: + return int(n_attn * n_ctx * n_kv * (key_len + val_len) * bpe) + head_dim = self._embedding_length // self._n_heads if self._n_heads else 128 # type: ignore[operator] + return int(n_attn * n_ctx * n_kv * 2 * head_dim * bpe) + + # Path 3: Sliding Window (Gemma-3, gpt-oss) + # SWA layers only cache min(ctx, window) tokens; global layers cache full ctx. + # Most SWA architectures use few global layers (e.g., Gemma-3 uses 1 in 6). + # Without an explicit field, we conservatively assume 1/4 of layers are global + # which is still far more accurate than the legacy formula (which ignores SWA). + if ( + self._sliding_window is not None + and self._sliding_window > 0 + and key_len is not None + and val_len is not None + ): + swa = self._sliding_window + n_global = max(1, n_layers // 4) + n_swa = n_layers - n_global + kv_per_token = n_kv * (key_len + val_len) * bpe + return int( + n_global * n_ctx * kv_per_token + n_swa * min(n_ctx, swa) * kv_per_token + ) + + # Path 4: Standard GQA with explicit key/value dimensions + if key_len is not None and val_len is not None: + return int(n_layers * n_ctx * n_kv * (key_len + val_len) * bpe) + + # Path 5: Legacy fallback (old GGUFs without explicit dimensions) + head_dim = self._embedding_length // self._n_heads if self._n_heads else 128 # type: ignore[operator] + return int(2 * n_kv * head_dim * n_layers * n_ctx * bpe) def _fit_context_to_vram( self, @@ -393,8 +515,8 @@ class LlamaCppBackend: ) -> int: """Return the largest context length that fits in GPU VRAM. - Uses 70% of available VRAM as the budget (matching _select_gpus - threshold -- 30% reserved for compute buffers, CUDA context, + Uses 90% of available VRAM as the budget (matching _select_gpus + threshold -- 10% reserved for compute buffers, CUDA context, scratch space, flash-attn workspace, etc.). If the model weights alone don't fit, returns min_ctx unchanged. """ @@ -406,7 +528,7 @@ class LlamaCppBackend: ) return requested_ctx - budget_bytes = available_mib * 1024 * 1024 * 0.70 + budget_bytes = available_mib * 1024 * 1024 * 0.90 model_footprint = model_size_bytes # Check if requested context already fits @@ -460,8 +582,6 @@ class LlamaCppBackend: Returns (first_shard_filename, total_size_bytes) or None if nothing fits. """ - import re - try: from huggingface_hub import get_paths_info, list_repo_files @@ -477,10 +597,9 @@ class LlamaCppBackend: size_map = {p.path: (p.size or 0) for p in path_infos} # Group files by variant: shards share a prefix before -NNNNN-of-NNNNN - shard_pat = re.compile(r"^(.*)-\d{5}-of-\d{5}\.gguf$") variants: dict[str, list[str]] = {} for f in gguf_files: - m = shard_pat.match(f) + m = _SHARD_RE.match(f) key = m.group(1) if m else f variants.setdefault(key, []).append(f) @@ -585,6 +704,14 @@ class LlamaCppBackend: self._n_kv_heads = None self._n_heads = None self._embedding_length = None + self._kv_key_length = None + self._kv_value_length = None + self._sliding_window = None + self._full_attention_interval = None + self._kv_lora_rank = None + self._key_length_mla = None + self._ssm_inner_size = None + self._ssm_state_size = None try: WANTED = {"general.architecture", "tokenizer.chat_template"} @@ -619,6 +746,15 @@ class LlamaCppBackend: f"{arch}.attention.head_count_kv": "n_kv_heads", f"{arch}.attention.head_count": "n_heads", f"{arch}.embedding_length": "embedding_length", + # Architecture-aware KV cache fields + f"{arch}.attention.key_length": "kv_key_length", + f"{arch}.attention.value_length": "kv_value_length", + f"{arch}.attention.sliding_window": "sliding_window", + f"{arch}.full_attention_interval": "full_attention_interval", + f"{arch}.attention.kv_lora_rank": "kv_lora_rank", + f"{arch}.attention.key_length_mla": "key_length_mla", + f"{arch}.ssm.inner_size": "ssm_inner_size", + f"{arch}.ssm.state_size": "ssm_state_size", } elif key == "tokenizer.chat_template": self._chat_template = val_s @@ -674,7 +810,9 @@ class LlamaCppBackend: # Detect tool calling support from chat template tool_markers = [ "{%- if tools %}", + "{%- if tools -%}", "{% if tools %}", + "{% if tools -%}", '"role" == "tool"', "'role' == 'tool'", 'message.role == "tool"', @@ -714,7 +852,6 @@ class LlamaCppBackend: gguf_extra_shards: list[str] = [] if hf_variant: try: - import re from huggingface_hub import list_repo_files files = list_repo_files(hf_repo, token = hf_token) @@ -729,11 +866,10 @@ class LlamaCppBackend: ) if gguf_files: gguf_filename = gguf_files[0] - shard_pat = re.compile(r"^(.*)-\d{5}-of-(\d{5})\.gguf$") - m = shard_pat.match(gguf_filename) + m = _SHARD_FULL_RE.match(gguf_filename) if m: prefix = m.group(1) - total = m.group(2) + total = m.group(3) sibling_pat = re.compile( r"^" + re.escape(prefix) @@ -790,10 +926,7 @@ class LlamaCppBackend: f"falling back to {fallback_file} ({fallback_size / (1024**3):.1f} GB)" ) gguf_filename = fallback_file - import re as _re - - _shard_pat = _re.compile(r"^(.*)-\d{5}-of-\d{5}\.gguf$") - _m = _shard_pat.match(gguf_filename) + _m = _SHARD_RE.match(gguf_filename) _prefix = _m.group(1) if _m else None if _prefix: gguf_extra_shards = sorted( @@ -1042,7 +1175,7 @@ class LlamaCppBackend: ) kv = self._estimate_kv_cache_bytes(capped, cache_type_kv) total_mib = (model_size + kv) / (1024 * 1024) - if total_mib <= pool_mib * 0.70: + if total_mib <= pool_mib * 0.90: best_cap = max(best_cap, capped) if best_cap > 0: max_available_ctx = best_cap @@ -1071,7 +1204,7 @@ class LlamaCppBackend: capped, cache_type_kv ) total_mib = (model_size + kv) / (1024 * 1024) - if total_mib <= pool_mib * 0.70: + if total_mib <= pool_mib * 0.90: effective_ctx = capped gpu_indices = sorted(idx for idx, _ in subset) use_fit = False @@ -1090,7 +1223,7 @@ class LlamaCppBackend: ) kv = self._estimate_kv_cache_bytes(capped, cache_type_kv) total_mib = (model_size + kv) / (1024 * 1024) - if total_mib <= pool_mib * 0.70: + if total_mib <= pool_mib * 0.90: effective_ctx = capped gpu_indices = sorted(idx for idx, _ in subset) use_fit = False @@ -1196,17 +1329,12 @@ class LlamaCppBackend: # Qwen3.5 models below 9B (0.8B, 2B, 4B) disable thinking by default. # Only 9B and larger enable thinking. if self._supports_reasoning: - import re - thinking_default = True mid = (model_identifier or "").lower() if "qwen3.5" in mid: - # Extract size like "0.8b", "4b", "35b" etc. - size_match = re.search(r"(\d+\.?\d*)\s*b", mid) - if size_match: - size_val = float(size_match.group(1)) - if size_val < 9: - thinking_default = False + size_val = _extract_model_size_b(mid) + if size_val is not None and size_val < 9: + thinking_default = False self._reasoning_default = thinking_default cmd.extend( [ @@ -1422,6 +1550,14 @@ class LlamaCppBackend: self._n_kv_heads = None self._n_heads = None self._embedding_length = None + self._kv_key_length = None + self._kv_value_length = None + self._sliding_window = None + self._full_attention_interval = None + self._kv_lora_rank = None + self._key_length_mla = None + self._ssm_inner_size = None + self._ssm_state_size = None # Clean up temp chat template file if hasattr(self, "_chat_template_file") and self._chat_template_file: try: @@ -1671,13 +1807,11 @@ class LlamaCppBackend: Closing tags (, , ) are all optional since models frequently omit them. """ - import re - tool_calls = [] # Pattern 1: JSON inside tags. # Use balanced-brace extraction that skips braces inside JSON strings. - for m in re.finditer(r"\s*\{", content): + for m in _TC_JSON_START_RE.finditer(content): brace_start = m.end() - 1 # position of the opening { depth, i = 0, brace_start in_string = False @@ -1727,7 +1861,7 @@ class LlamaCppBackend: # boundaries. We avoid using as a boundary because # code parameter values can contain that literal string. # After extracting, we trim a trailing if present. - func_starts = list(re.finditer(r"\s*", content)) + func_starts = list(_TC_FUNC_START_RE.finditer(content)) for idx, fm in enumerate(func_starts): func_name = fm.group(1) body_start = fm.end() @@ -1737,7 +1871,7 @@ class LlamaCppBackend: if idx + 1 < len(func_starts) else len(content) ) - end_tag = re.search(r"", content[body_start:]) + end_tag = _TC_END_TAG_RE.search(content[body_start:]) if end_tag: body_end = body_start + end_tag.start() else: @@ -1745,20 +1879,20 @@ class LlamaCppBackend: body_end = min(body_end, next_func) body = content[body_start:body_end] # Trim trailing if present (it's the real closing tag) - body = re.sub(r"\s*\s*$", "", body) + body = _TC_FUNC_CLOSE_RE.sub("", body) # Step 2: Extract parameters from body. # For single-parameter functions (the common case: code, command, # query), use body end as the only boundary to avoid false matches # on inside code strings. arguments = {} - param_starts = list(re.finditer(r"\s*", body)) + param_starts = list(_TC_PARAM_START_RE.finditer(body)) if len(param_starts) == 1: # Single parameter: value is everything from after the tag # to end of body, trimming any trailing . pm = param_starts[0] val = body[pm.end() :] - val = re.sub(r"\s*\s*$", "", val) + val = _TC_PARAM_CLOSE_RE.sub("", val) arguments[pm.group(1)] = val.strip() else: for pidx, pm in enumerate(param_starts): @@ -1772,7 +1906,7 @@ class LlamaCppBackend: ) val = body[val_start:next_param] # Trim trailing if present - val = re.sub(r"\s*\s*$", "", val) + val = _TC_PARAM_CLOSE_RE.sub("", val) arguments[param_name] = val.strip() tc = { @@ -2145,22 +2279,10 @@ class LlamaCppBackend: _accumulated_predicted_ms = 0.0 _accumulated_predicted_n = 0 - # ── Shared patterns for stripping tool XML from streamed content ── - import re as _re_tool - - _TOOL_CLOSED_PATTERNS = [ - _re_tool.compile(r".*?", _re_tool.DOTALL), - _re_tool.compile(r".*?", _re_tool.DOTALL), - ] - _TOOL_ALL_PATTERNS = _TOOL_CLOSED_PATTERNS + [ - _re_tool.compile(r".*$", _re_tool.DOTALL), - _re_tool.compile(r".*$", _re_tool.DOTALL), - ] - def _strip_tool_markup(text: str, *, final: bool = False) -> str: if not auto_heal_tool_calls: return text - patterns = _TOOL_ALL_PATTERNS if final else _TOOL_CLOSED_PATTERNS + patterns = _TOOL_ALL_PATS if final else _TOOL_CLOSED_PATS for pat in patterns: text = pat.sub("", text) return text.strip() if final else text @@ -2180,7 +2302,19 @@ class LlamaCppBackend: # identical call succeeded). _tool_call_history: list[tuple[str, bool]] = [] # (key, failed) - for iteration in range(max_tool_iterations): + # ── Re-prompt on plan-without-action ───────────────── + # When the model describes what it intends to do (forward-looking + # language) without actually calling a tool, re-prompt once. + # Only triggers on responses that signal intent/planning -- a + # direct answer like "4" or "Hello!" will not match. + # Pattern is compiled once at module level (_INTENT_SIGNAL). + _reprompt_count = 0 + + # Reserve extra iterations for re-prompts so they don't + # consume the caller's tool-call budget. Only add the + # extra slot when tool iterations are actually allowed. + _extra = _MAX_REPROMPTS if max_tool_iterations > 0 else 0 + for iteration in range(max_tool_iterations + _extra): if cancel_event is not None and cancel_event.is_set(): return @@ -2491,6 +2625,57 @@ class LlamaCppBackend: content_accum, ) if not _safety_tc: + # ── Re-prompt on plan-without-action ── + # If the model described what it intends to do + # (forward-looking language) without calling any + # tool, nudge it to act. Only fires once per + # request and only on short responses that + # contain intent signals -- a direct answer + # like "4" or "Hello!" won't trigger this. + # Use content if available, otherwise fall back + # to reasoning text (reasoning-only stalls). + _stripped = content_accum.strip() + if not _stripped: + _stripped = reasoning_accum.strip() + if ( + tools + and _reprompt_count < _MAX_REPROMPTS + and 0 < len(_stripped) < _REPROMPT_MAX_CHARS + and _INTENT_SIGNAL.search(_stripped) + ): + _reprompt_count += 1 + logger.info( + f"Re-prompt {_reprompt_count}/{_MAX_REPROMPTS}: " + f"model responded without calling tools " + f"({len(_stripped)} chars)" + ) + conversation.append( + { + "role": "assistant", + "content": _stripped, + } + ) + conversation.append( + { + "role": "user", + "content": ( + "STOP. Do NOT write code or explain. " + "You MUST call a tool NOW. " + "Call web_search or python immediately." + ), + } + ) + # Accumulate tokens and timing from this iteration + _fu_r = _iter_usage or {} + _accumulated_completion_tokens += _fu_r.get( + "completion_tokens", 0 + ) + _it_r = _iter_timings or {} + _accumulated_predicted_ms += _it_r.get("predicted_ms", 0) + _accumulated_predicted_n += _it_r.get("predicted_n", 0) + yield {"type": "status", "text": ""} + continue + # Content was already streamed. Yield metadata. yield {"type": "status", "text": ""} _fu = _iter_usage or {} @@ -2723,10 +2908,15 @@ class LlamaCppBackend: _error_prefixes ) _tool_call_history.append((_tc_key, _is_error)) + # Strip image sentinel before feeding result to the LLM + # (the full result with sentinel is still yielded via + # tool_end so the frontend can extract image paths). _result_content = result + if "\n__IMAGES__:" in _result_content: + _result_content = _result_content.rsplit("\n__IMAGES__:", 1)[0] if _is_error: _result_content = ( - result + "\n\nThe tool call encountered an issue. " + _result_content + "\n\nThe tool call encountered an issue. " "Please try a different approach or rephrase your request." ) @@ -2740,6 +2930,8 @@ class LlamaCppBackend: tool_msg["tool_call_id"] = tool_call_id conversation.append(tool_msg) + # Clear tool status badge before next generation iteration + yield {"type": "status", "text": ""} # Continue the loop to let model respond with context continue diff --git a/studio/backend/core/inference/orchestrator.py b/studio/backend/core/inference/orchestrator.py index f293a0dcd8..cb5d9da34a 100644 --- a/studio/backend/core/inference/orchestrator.py +++ b/studio/backend/core/inference/orchestrator.py @@ -109,12 +109,13 @@ class InferenceOrchestrator: self._top_models_ready.wait(timeout = 5) top_gguf = self._top_gguf_cache or [] top_hub = self._top_hub_cache or [] - # GGUFs first, then hub models, then static fallbacks. + # Curated static defaults first (editorial picks like new models), + # then HF download-ranked models to backfill. # Send extras so the frontend still has 4 per category # after removing already-downloaded models. result: list[str] = [] seen: set[str] = set() - for m in top_gguf + top_hub + self._static_models: + for m in self._static_models + top_gguf + top_hub: if m not in seen: result.append(m) seen.add(m) diff --git a/studio/backend/core/inference/tools.py b/studio/backend/core/inference/tools.py index 2ac8f76322..b23372b766 100644 --- a/studio/backend/core/inference/tools.py +++ b/studio/backend/core/inference/tools.py @@ -8,20 +8,28 @@ Supports web search (DuckDuckGo), Python code execution, and terminal commands. """ import ast +import http.client import os os.environ["UNSLOTH_IS_PRESENT"] = "1" +import random +import ssl import subprocess import sys import tempfile import threading +import urllib.request from loggers import get_logger logger = get_logger(__name__) _EXEC_TIMEOUT = 300 # 5 minutes + +# Strict raster-image allowlist for sandbox file serving. +# No .svg (XSS risk via embedded scripts), no .html, no .pdf. +_IMAGE_EXTS = frozenset({".png", ".jpg", ".jpeg", ".gif", ".webp", ".bmp"}) _MAX_OUTPUT_CHARS = 8000 # truncate long output _BASH_BLOCKED_WORDS = {"rm", "sudo", "dd", "chmod", "mkfs", "shutdown", "reboot"} @@ -154,8 +162,74 @@ def execute_tool( return f"Unknown tool: {name}" -_MAX_PAGE_CHARS = 16000 # limit fetched page text -_MAX_FETCH_BYTES = _MAX_PAGE_CHARS * 4 + 1 # cap raw download size +_MAX_PAGE_CHARS = 16000 # limit fetched page text (after HTML-to-MD conversion) +# Raw download cap. Must be larger than _MAX_PAGE_CHARS because SSR pages +# embed large sections (CSS, JS, SVGs) that are stripped during +# HTML-to-Markdown conversion. 512 KB is enough to reach article content +# on GitBook / Next.js / Docusaurus pages whose alone can be 200 KB. +_MAX_FETCH_BYTES = 512 * 1024 + +_USER_AGENTS = ( + "Mozilla/5.0 (Windows NT 10.0; Win64; x64) AppleWebKit/537.36 (KHTML, like Gecko) Chrome/131.0.0.0 Safari/537.36", + "Mozilla/5.0 (Macintosh; Intel Mac OS X 10_15_7) AppleWebKit/537.36 (KHTML, like Gecko) Chrome/131.0.0.0 Safari/537.36", + "Mozilla/5.0 (X11; Linux x86_64) AppleWebKit/537.36 (KHTML, like Gecko) Chrome/131.0.0.0 Safari/537.36", + "Mozilla/5.0 (Windows NT 10.0; Win64; x64; rv:133.0) Gecko/20100101 Firefox/133.0", + "Mozilla/5.0 (Macintosh; Intel Mac OS X 10.15; rv:133.0) Gecko/20100101 Firefox/133.0", + "Mozilla/5.0 (Macintosh; Intel Mac OS X 10_15_7) AppleWebKit/605.1.15 (KHTML, like Gecko) Version/18.2 Safari/605.1.15", +) + +_tls_ctx = ssl.create_default_context() + + +class _NoRedirect(urllib.request.HTTPRedirectHandler): + def redirect_request(self, req, fp, code, msg, headers, newurl): + return None + + +class _PinnedHTTPSConnection(http.client.HTTPSConnection): + """HTTPS connection that connects to a pinned IP but uses a different + hostname for SNI and certificate verification. + + The SSRF IP-pinning rewrites URLs to raw IPs. A normal HTTPSConnection + would then send no SNI and verify the cert against the IP, both of which + fail. This subclass splits the two concerns: TCP connects to the pinned + IP (``host`` parameter) while TLS uses ``sni_hostname`` for the + ClientHello and cert check. + """ + + def __init__(self, host: str, *, sni_hostname: str, **kwargs): + super().__init__(host, **kwargs) + self._sni_hostname = sni_hostname + + def connect(self): + # TCP connect to the pinned IP stored in self.host (+ tunnel if + # a proxy is configured via set_tunnel, though we do not use one). + http.client.HTTPConnection.connect(self) + # TLS handshake with the real hostname for SNI + cert verification. + self.sock = self._context.wrap_socket( + self.sock, + server_hostname = self._sni_hostname, + ) + + +class _SNIHTTPSHandler(urllib.request.HTTPSHandler): + """HTTPS handler that sends the correct SNI hostname during TLS handshake. + + The SSRF IP-pinning rewrites URLs to raw IPs, which breaks SNI and cert + verification. This handler returns a ``_PinnedHTTPSConnection`` that + connects to the pinned IP but verifies TLS against the original hostname. + """ + + def __init__(self, hostname: str): + super().__init__(context = _tls_ctx) + self._sni_hostname = hostname + + def https_open(self, req): + return self.do_open(self._sni_connection, req) + + def _sni_connection(self, host, **kwargs): + kwargs["context"] = _tls_ctx + return _PinnedHTTPSConnection(host, sni_hostname = self._sni_hostname, **kwargs) def _validate_and_resolve_host(hostname: str, port: int) -> tuple[bool, str, str]: @@ -215,33 +289,32 @@ def _fetch_page_text( return reason try: - import urllib.request from urllib.error import HTTPError as _HTTPError from urllib.parse import urljoin, urlunparse - # Disable auto-redirect so we can validate each hop for SSRF. - # urllib raises HTTPError for 3xx when the handler returns None, - # so we catch that and extract the Location header manually. - class _NoRedirect(urllib.request.HTTPRedirectHandler): - def redirect_request(self, req, fp, code, msg, headers, newurl): - return None - - opener = urllib.request.build_opener(_NoRedirect) - max_bytes = max_chars * 4 + 1 + max_bytes = _MAX_FETCH_BYTES current_url = url current_host = parsed.hostname + ua = random.choice(_USER_AGENTS) for _hop in range(5): # Pin to the validated IP to prevent DNS rebinding. # Rewrite the URL to use the IP and set the Host header. cp = urlparse(current_url) - ip_netloc = f"{pinned_ip}:{cp.port}" if cp.port else pinned_ip + # Bracket IPv6 addresses so the netloc is valid in a URL. + ip_str = f"[{pinned_ip}]" if ":" in pinned_ip else pinned_ip + ip_netloc = f"{ip_str}:{cp.port}" if cp.port else ip_str pinned_url = urlunparse(cp._replace(netloc = ip_netloc)) + opener = urllib.request.build_opener( + _NoRedirect, + _SNIHTTPSHandler(current_host), + ) + req = urllib.request.Request( pinned_url, headers = { - "User-Agent": "UnslothStudio/1.0", + "User-Agent": ua, "Host": current_host, }, ) @@ -522,6 +595,12 @@ def _check_code_safety(code: str) -> str | None: """ safe, info = _check_signal_escape_patterns(code) if not safe: + # SyntaxError from ast.parse -- let these through so the subprocess + # produces a normal Python traceback instead of a misleading + # "unsafe code detected" message. + if info.get("error"): + return None + reasons = [ item.get("description", "") for item in info.get("signal_tampering", []) ] @@ -565,6 +644,17 @@ def _python_exec( tmp_path = None workdir = _get_workdir(session_id) + # Snapshot image mtimes so we detect both new and overwritten files. + _before: dict[str, int] = {} + if os.path.isdir(workdir): + for _name in os.listdir(workdir): + if os.path.splitext(_name)[1].lower() in _IMAGE_EXTS: + _p = os.path.join(workdir, _name) + if os.path.isfile(_p): + try: + _before[_name] = os.stat(_p).st_mtime_ns + except OSError: + pass try: fd, tmp_path = tempfile.mkstemp( suffix = ".py", prefix = "studio_exec_", dir = workdir @@ -600,7 +690,29 @@ def _python_exec( result = output or "" if proc.returncode != 0: result = f"Exit code {proc.returncode}:\n{result}" - return _truncate(result) if result.strip() else "(no output)" + result = _truncate(result) if result.strip() else "(no output)" + + # Detect new or overwritten image files and append sentinel for frontend + if session_id and os.path.isdir(workdir): + new_images = [] + for _name in os.listdir(workdir): + if os.path.splitext(_name)[1].lower() not in _IMAGE_EXTS: + continue + _p = os.path.join(workdir, _name) + if not os.path.isfile(_p): + continue + try: + _mtime = os.stat(_p).st_mtime_ns + except OSError: + continue + if _name not in _before or _mtime != _before[_name]: + new_images.append(_name) + if new_images: + import json as _json + + result += f"\n__IMAGES__:{_json.dumps(sorted(new_images))}" + + return result except Exception as e: return f"Execution error: {e}" diff --git a/studio/backend/core/training/worker.py b/studio/backend/core/training/worker.py index 699cfe74f7..0454eada89 100644 --- a/studio/backend/core/training/worker.py +++ b/studio/backend/core/training/worker.py @@ -390,11 +390,16 @@ def run_training_process( # Some newer architectures (e.g. NemotronH) have config parsing bugs in # transformers that require trust_remote_code=True as a workaround. # Only auto-enable for unsloth/* prefixed models (trusted source). + # Exclude Gemma 4 since it is a native transformers 5.5 model and + # trust_remote_code=True would bypass the compiler (disabling fused CE). from utils.transformers_version import needs_transformers_5 + _lowered = model_name.lower() + _is_native_t5 = any(x in _lowered for x in ("gemma-4", "gemma4")) if ( needs_transformers_5(model_name) - and model_name.lower().startswith("unsloth/") + and _lowered.startswith("unsloth/") + and not _is_native_t5 and not config.get("trust_remote_code", False) ): config["trust_remote_code"] = True diff --git a/studio/backend/main.py b/studio/backend/main.py index c18f18a743..ad19ee9679 100644 --- a/studio/backend/main.py +++ b/studio/backend/main.py @@ -121,13 +121,13 @@ async def lifespan(app: FastAPI): if storage.ensure_default_admin(): bootstrap_pw = storage.get_bootstrap_password() app.state.bootstrap_password = bootstrap_pw + + bootstrap_path = storage.DB_PATH.parent / ".bootstrap_password" print("\n" + "=" * 60) print("DEFAULT ADMIN ACCOUNT CREATED") - print( - "Sign in with the seeded credentials and change the password immediately:\n" - ) print(f" username: {storage.DEFAULT_ADMIN_USERNAME}") - print(f" password: {bootstrap_pw}\n") + print(f" password saved to: {bootstrap_path}") + print(" Open the Studio UI to sign in and change it.") print("=" * 60 + "\n") else: app.state.bootstrap_password = storage.get_bootstrap_password() diff --git a/studio/backend/models/inference.py b/studio/backend/models/inference.py index 77f70b9bd6..3094df4169 100644 --- a/studio/backend/models/inference.py +++ b/studio/backend/models/inference.py @@ -139,6 +139,10 @@ class LoadResponse(BaseModel): max_context_length: Optional[int] = Field( None, description = "Maximum context length currently available on this hardware" ) + native_context_length: Optional[int] = Field( + None, + description = "Model's native context length from GGUF metadata (not capped by VRAM)", + ) supports_reasoning: bool = Field( False, description = "Whether model supports thinking/reasoning mode (enable_thinking)", @@ -217,6 +221,10 @@ class InferenceStatusResponse(BaseModel): None, description = "Maximum context length currently available for the active model", ) + native_context_length: Optional[int] = Field( + None, + description = "Model's native context length from GGUF metadata (not capped by VRAM)", + ) # ===================================================================== diff --git a/studio/backend/requirements/no-torch-runtime.txt b/studio/backend/requirements/no-torch-runtime.txt index 7133fe8922..3b822ac2a4 100644 --- a/studio/backend/requirements/no-torch-runtime.txt +++ b/studio/backend/requirements/no-torch-runtime.txt @@ -28,6 +28,21 @@ peft>=0.18.0,!=0.11.0 huggingface_hub>=0.34.0 hf_transfer diffusers + +# Transitive deps required because this file is installed with --no-deps. +# Without these, `from transformers import AutoConfig` fails at import time. +regex +typing_extensions +filelock +httpx +httpcore +certifi +idna +anyio +sniffio +h11 + +tokenizers 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.3.0 trl>=0.18.2,!=0.19.0,<=0.24.0 sentence-transformers diff --git a/studio/backend/routes/inference.py b/studio/backend/routes/inference.py index 9bce371775..ced24c1d5f 100644 --- a/studio/backend/routes/inference.py +++ b/studio/backend/routes/inference.py @@ -5,11 +5,12 @@ Inference API routes for model loading and text generation. """ +import os import sys import time import uuid from pathlib import Path -from fastapi import APIRouter, Depends, HTTPException, Request +from fastapi import APIRouter, Depends, HTTPException, Request, status from fastapi.responses import StreamingResponse, JSONResponse from typing import Optional import json @@ -21,6 +22,9 @@ import threading import re as _re +# Model size extraction (shared with core/inference/llama_cpp.py) +from utils.models import extract_model_size_b as _extract_model_size_b + def _friendly_error(exc: Exception) -> str: """Extract a user-friendly message from known llama-server errors.""" @@ -90,6 +94,14 @@ from datetime import date as _date router = APIRouter() +# Appended to tool-use nudge to discourage plan-without-action +_TOOL_ACTION_NUDGE = ( + " IMPORTANT: Always call tools directly -- never write code yourself." + " Never describe what you plan to do -- just call the tool immediately." + " For any code request, call the python tool. For any factual question, call web_search." + " Do NOT output code blocks -- use the python tool instead." +) + # Regex for stripping leaked tool-call XML from assistant messages/stream _TOOL_XML_RE = _re.compile( r".*?|.*?", @@ -163,6 +175,7 @@ async def load_model( inference = inference_config, context_length = llama_backend.context_length, max_context_length = llama_backend.max_context_length, + native_context_length = llama_backend.native_context_length, supports_reasoning = llama_backend.supports_reasoning, reasoning_always_on = llama_backend.reasoning_always_on, chat_template = llama_backend.chat_template, @@ -298,6 +311,7 @@ async def load_model( inference = inference_config, context_length = llama_backend.context_length, max_context_length = llama_backend.max_context_length, + native_context_length = llama_backend.native_context_length, supports_reasoning = llama_backend.supports_reasoning, reasoning_always_on = llama_backend.reasoning_always_on, supports_tools = llama_backend.supports_tools, @@ -637,6 +651,7 @@ async def get_status( supports_tools = llama_backend.supports_tools, context_length = llama_backend.context_length, max_context_length = llama_backend.max_context_length, + native_context_length = llama_backend.native_context_length, ) # Otherwise, report Unsloth backend status @@ -1092,12 +1107,20 @@ async def openai_chat_completions( _date_line = f"The current date is {_date.today().isoformat()}." - _web_tips = ( - "When you search and find a relevant URL in the results, " - "fetch its full content by calling web_search with the url parameter. " - "Do not repeat the same search query. If a search returns " - "no useful results, try rephrasing or fetching a result URL directly." - ) + # Small models (<9B) struggle with multi-step search plans, + # so simplify the web tips to avoid plan-then-stall behavior. + _model_size_b = _extract_model_size_b(model_name) + _is_small_model = _model_size_b is not None and _model_size_b < 9 + + if _is_small_model: + _web_tips = "Do not repeat the same search query." + else: + _web_tips = ( + "When you search and find a relevant URL in the results, " + "fetch its full content by calling web_search with the url parameter. " + "Do not repeat the same search query. If a search returns " + "no useful results, try rephrasing or fetching a result URL directly." + ) _code_tips = ( "Use code execution for math, calculations, data processing, " "or to parse and analyze information from tool results." @@ -1129,6 +1152,7 @@ async def openai_chat_completions( _nudge = "" if _nudge: + _nudge += _TOOL_ACTION_NUDGE # Append nudge to system prompt (preserve user's prompt) if system_prompt: system_prompt = system_prompt.rstrip() + "\n\n" + _nudge @@ -1205,7 +1229,14 @@ async def openai_chat_completions( break if event["type"] == "status": + # Empty status marks an iteration boundary + # in the GGUF tool loop (e.g. after a + # re-prompt). Reset the cumulative cursor + # so the next assistant turn streams cleanly. + if not event["text"]: + prev_text = "" # Emit tool status as a custom SSE event + # (including empty ones to clear UI badges) status_data = json.dumps( { "type": "tool_status", @@ -1652,6 +1683,94 @@ async def openai_chat_completions( raise HTTPException(status_code = 500, detail = str(e)) +# ===================================================================== +# Sandbox file serving (/sandbox/{session_id}/{filename}) +# ===================================================================== + +_SANDBOX_MEDIA_TYPES = { + ".png": "image/png", + ".jpg": "image/jpeg", + ".jpeg": "image/jpeg", + ".gif": "image/gif", + ".webp": "image/webp", + ".bmp": "image/bmp", +} + + +@router.get("/sandbox/{session_id}/{filename}") +async def serve_sandbox_file( + session_id: str, + filename: str, + request: Request, + token: Optional[str] = None, +): + """ + Serve image files created by Python tool execution. + + Accepts auth via Authorization header OR ?token= query param + (needed because cannot send custom headers). + """ + from fastapi.responses import FileResponse + + # ── Authentication (header or query param) ────────────────── + auth_header = request.headers.get("authorization") + if auth_header and auth_header.lower().startswith("bearer "): + jwt_token = auth_header[7:] + elif token: + jwt_token = token + else: + raise HTTPException( + status_code = status.HTTP_401_UNAUTHORIZED, + detail = "Missing authentication token", + ) + from fastapi.security import HTTPAuthorizationCredentials + + creds = HTTPAuthorizationCredentials(scheme = "Bearer", credentials = jwt_token) + await get_current_subject(creds) + + # ── Filename sanitization ─────────────────────────────────── + safe_filename = os.path.basename(filename) + if not safe_filename or safe_filename in (".", ".."): + raise HTTPException(status_code = 404, detail = "Not found") + + # ── Extension allowlist ───────────────────────────────────── + ext = os.path.splitext(safe_filename)[1].lower() + media_type = _SANDBOX_MEDIA_TYPES.get(ext) + if not media_type: + raise HTTPException( + status_code = status.HTTP_403_FORBIDDEN, + detail = "File type not allowed", + ) + + # ── Path containment check ────────────────────────────────── + home = os.path.expanduser("~") + sandbox_root = os.path.realpath(os.path.join(home, "studio_sandbox")) + safe_session = os.path.basename(session_id.replace("..", "")) + if not safe_session: + raise HTTPException(status_code = 404, detail = "Not found") + + file_path = os.path.realpath( + os.path.join(sandbox_root, safe_session, safe_filename) + ) + if not file_path.startswith(sandbox_root + os.sep): + raise HTTPException( + status_code = status.HTTP_403_FORBIDDEN, + detail = "Access denied", + ) + + if not os.path.isfile(file_path): + raise HTTPException(status_code = 404, detail = "Not found") + + return FileResponse( + path = file_path, + media_type = media_type, + headers = { + "Cache-Control": "private, no-store", + "X-Content-Type-Options": "nosniff", + }, + ) + + # ===================================================================== # OpenAI-Compatible Models Listing (/models → /v1/models) # ===================================================================== diff --git a/studio/backend/tests/test_kv_cache_estimation.py b/studio/backend/tests/test_kv_cache_estimation.py new file mode 100644 index 0000000000..2640ded90d --- /dev/null +++ b/studio/backend/tests/test_kv_cache_estimation.py @@ -0,0 +1,929 @@ +# SPDX-License-Identifier: AGPL-3.0-only +# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. + +"""Tests for 5-path architecture-aware KV cache VRAM estimation. + +Covers the GGUF metadata parser, _can_estimate_kv gate, all 5 estimation +paths (MLA, Hybrid Mamba, Sliding Window, Standard GQA, Legacy), KV cache +quantization, edge cases, and lifecycle (init/unload/reparse). + +Requires no GPU, network, or external libraries beyond pytest. +Cross-platform: Linux, macOS, Windows, WSL. +""" + +import io +import struct +import sys +import types as _types +from pathlib import Path + +import pytest + +# --------------------------------------------------------------------------- +# Stub heavy / unavailable external dependencies before importing the +# module under test. Same pattern as test_native_context_length.py. +# --------------------------------------------------------------------------- + +_BACKEND_DIR = str(Path(__file__).resolve().parent.parent) +if _BACKEND_DIR not in sys.path: + sys.path.insert(0, _BACKEND_DIR) + +# loggers +_loggers_stub = _types.ModuleType("loggers") +_loggers_stub.get_logger = lambda name: __import__("logging").getLogger(name) +sys.modules.setdefault("loggers", _loggers_stub) + +# structlog +_structlog_stub = _types.ModuleType("structlog") +sys.modules.setdefault("structlog", _structlog_stub) + +# httpx +_httpx_stub = _types.ModuleType("httpx") +for _exc_name in ( + "ConnectError", + "TimeoutException", + "ReadTimeout", + "ReadError", + "RemoteProtocolError", + "CloseError", +): + setattr(_httpx_stub, _exc_name, type(_exc_name, (Exception,), {})) + + +class _FakeTimeout: + def __init__(self, *a, **kw): + pass + + +_httpx_stub.Timeout = _FakeTimeout +_httpx_stub.Client = type( + "Client", + (), + { + "__init__": lambda self, **kw: None, + "__enter__": lambda self: self, + "__exit__": lambda self, *a: None, + }, +) +sys.modules.setdefault("httpx", _httpx_stub) + +from core.inference.llama_cpp import LlamaCppBackend + +# --------------------------------------------------------------------------- +# Helpers +# --------------------------------------------------------------------------- + + +def _make_gguf_bytes(arch: str, kv_pairs: dict) -> bytes: + """Build a minimal GGUF v3 binary blob with the given KV metadata. + + Only supports UINT32 (type 4), UINT64 (type 10), and STRING (type 8) + values, which is all the metadata parser reads. + """ + buf = io.BytesIO() + # Header: magic, version, tensor_count, kv_count + buf.write(struct.pack(" LlamaCppBackend: + """Create a LlamaCppBackend with parsed GGUF metadata from given fields.""" + kv = {"general.architecture": arch} + for k, v in fields.items(): + kv[f"{arch}.{k}"] = v + import tempfile, os + + data = _make_gguf_bytes(arch, kv) + fd, path = tempfile.mkstemp(suffix = ".gguf") + try: + os.write(fd, data) + os.close(fd) + b = LlamaCppBackend() + b._read_gguf_metadata(path) + return b + finally: + os.unlink(path) + + +# --------------------------------------------------------------------------- +# A. GGUF Parser Tests +# --------------------------------------------------------------------------- + + +class TestGGUFParserNewFields: + """Verify that the 8 new architecture-aware fields are correctly parsed.""" + + @pytest.mark.parametrize( + "field,gguf_key,value", + [ + ("_kv_key_length", "attention.key_length", 128), + ("_kv_value_length", "attention.value_length", 128), + ("_sliding_window", "attention.sliding_window", 1024), + ("_full_attention_interval", "full_attention_interval", 4), + ("_kv_lora_rank", "attention.kv_lora_rank", 512), + ("_key_length_mla", "attention.key_length_mla", 256), + ("_ssm_inner_size", "ssm.inner_size", 6144), + ("_ssm_state_size", "ssm.state_size", 128), + ], + ) + def test_field_parsed(self, field, gguf_key, value): + b = _backend_from_gguf("testarch", {gguf_key: value}) + assert getattr(b, field) == value + + def test_missing_fields_are_none(self): + b = _backend_from_gguf("testarch", {"block_count": 10}) + for attr in [ + "_kv_key_length", + "_kv_value_length", + "_sliding_window", + "_full_attention_interval", + "_kv_lora_rank", + "_key_length_mla", + "_ssm_inner_size", + "_ssm_state_size", + ]: + assert getattr(b, attr) is None + + def test_all_13_fields_parsed_together(self): + fields = { + "context_length": 131072, + "block_count": 62, + "attention.head_count_kv": 16, + "attention.head_count": 32, + "embedding_length": 5376, + "attention.key_length": 128, + "attention.value_length": 128, + "attention.sliding_window": 1024, + "full_attention_interval": 6, + "attention.kv_lora_rank": 512, + "attention.key_length_mla": 256, + "ssm.inner_size": 4096, + "ssm.state_size": 128, + } + b = _backend_from_gguf("testarch", fields) + assert b._context_length == 131072 + assert b._n_layers == 62 + assert b._n_kv_heads == 16 + assert b._n_heads == 32 + assert b._embedding_length == 5376 + assert b._kv_key_length == 128 + assert b._kv_value_length == 128 + assert b._sliding_window == 1024 + assert b._full_attention_interval == 6 + assert b._kv_lora_rank == 512 + assert b._key_length_mla == 256 + assert b._ssm_inner_size == 4096 + assert b._ssm_state_size == 128 + + +class TestGGUFParserReset: + """Verify that fields are properly reset between parses.""" + + def test_reset_between_parses(self): + # First parse with all fields + b = _backend_from_gguf( + "arch1", + { + "block_count": 32, + "attention.key_length": 128, + "attention.kv_lora_rank": 512, + "ssm.inner_size": 4096, + }, + ) + assert b._kv_key_length == 128 + assert b._kv_lora_rank == 512 + assert b._ssm_inner_size == 4096 + + # Second parse without those fields -- they should be None + kv = {"general.architecture": "arch2", "arch2.block_count": 64} + import tempfile, os + + data = _make_gguf_bytes("arch2", kv) + fd, path = tempfile.mkstemp(suffix = ".gguf") + os.write(fd, data) + os.close(fd) + try: + b._read_gguf_metadata(path) + finally: + os.unlink(path) + assert b._kv_key_length is None + assert b._kv_lora_rank is None + assert b._ssm_inner_size is None + assert b._n_layers == 64 + + +# --------------------------------------------------------------------------- +# B. _can_estimate_kv Gate Tests +# --------------------------------------------------------------------------- + + +class TestCanEstimateKV: + """Verify gate logic for all field combinations.""" + + def test_no_layers_returns_false(self): + b = LlamaCppBackend() + b._n_layers = None + b._kv_key_length = 128 + assert not b._can_estimate_kv() + + def test_explicit_both_dims_sufficient(self): + b = LlamaCppBackend() + b._n_layers = 32 + b._kv_key_length = 128 + b._kv_value_length = 128 + assert b._can_estimate_kv() + + def test_key_length_alone_insufficient(self): + """key_length without value_length should NOT be enough.""" + b = LlamaCppBackend() + b._n_layers = 32 + b._kv_key_length = 128 + assert not b._can_estimate_kv() + + def test_kv_lora_rank_sufficient(self): + b = LlamaCppBackend() + b._n_layers = 61 + b._kv_lora_rank = 512 + assert b._can_estimate_kv() + + def test_legacy_embed_plus_heads(self): + b = LlamaCppBackend() + b._n_layers = 28 + b._embedding_length = 1024 + b._n_heads = 16 + assert b._can_estimate_kv() + + def test_legacy_embed_plus_kv_heads(self): + b = LlamaCppBackend() + b._n_layers = 28 + b._embedding_length = 1024 + b._n_kv_heads = 8 + assert b._can_estimate_kv() + + def test_legacy_no_embed_returns_false(self): + b = LlamaCppBackend() + b._n_layers = 28 + b._n_heads = 16 + # No embedding_length, no new-style fields + assert not b._can_estimate_kv() + + def test_fresh_backend_returns_false(self): + b = LlamaCppBackend() + assert not b._can_estimate_kv() + + +# --------------------------------------------------------------------------- +# C. Path 1: MLA Estimation +# --------------------------------------------------------------------------- + + +class TestMLAEstimation: + """MLA: K-only cache using compressed KV latent + RoPE.""" + + def _mla_backend(self, **overrides): + defaults = { + "_n_layers": 61, + "_n_kv_heads": 1, + "_n_heads": 128, + "_embedding_length": 7168, + "_kv_key_length": 576, + "_kv_value_length": 512, + "_kv_lora_rank": 512, + "_key_length_mla": 192, + } + defaults.update(overrides) + b = LlamaCppBackend() + for k, v in defaults.items(): + setattr(b, k, v) + return b + + def test_deepseek_v3_f16(self): + b = self._mla_backend() + # 61 layers * 163840 ctx * 1 head * 576 key_len * 2 bpe + expected = 61 * 163840 * 1 * 576 * 2 + assert b._estimate_kv_cache_bytes(163840, "f16") == expected + + def test_mla_ignores_value_length(self): + """MLA should NOT add value_length -- V is reconstructed from the latent.""" + b = self._mla_backend() + result = b._estimate_kv_cache_bytes(1000, "f16") + # Should be n_layers * ctx * 1 * key_len(576) * 2 + expected = 61 * 1000 * 1 * 576 * 2 + assert result == expected + + def test_mla_fallback_when_no_key_length(self): + """If key_length is missing, fallback to kv_lora_rank + key_length_mla.""" + b = self._mla_backend(_kv_key_length = None) + # _key_length_mla=192 in default, so rope_dim=192 + result = b._estimate_kv_cache_bytes(1000, "f16") + expected = 61 * 1000 * 1 * (512 + 192) * 2 # 704 + assert result == expected + + def test_mla_fallback_no_key_length_mla(self): + """If both key_length and key_length_mla are missing, fallback to +64.""" + b = self._mla_backend(_kv_key_length = None, _key_length_mla = None) + result = b._estimate_kv_cache_bytes(1000, "f16") + expected = 61 * 1000 * 1 * (512 + 64) * 2 # 576 + assert result == expected + + def test_mla_defaults_n_kv_to_1_when_heads_absent(self): + """MLA should use n_kv=1 even if n_kv_heads is None (not n_heads).""" + b = self._mla_backend(_n_kv_heads = None) # n_heads=128 still set + result = b._estimate_kv_cache_bytes(1000, "f16") + # Should use n_kv_mla=1, NOT n_heads=128 + expected = 61 * 1000 * 1 * 576 * 2 + assert result == expected + + def test_mla_q4_quantization(self): + b = self._mla_backend() + result_f16 = b._estimate_kv_cache_bytes(1000, "f16") + result_q4 = b._estimate_kv_cache_bytes(1000, "q4_0") + assert result_q4 < result_f16 + # q4_0 bpe = 0.5625, f16 bpe = 2.0 + assert result_q4 == int(61 * 1000 * 1 * 576 * 0.5625) + + +# --------------------------------------------------------------------------- +# D. Path 2: Hybrid Mamba Estimation +# --------------------------------------------------------------------------- + + +class TestHybridMambaEstimation: + """Hybrid Mamba: only attention layers (1 in N) need KV cache.""" + + def _hybrid_backend(self, **overrides): + defaults = { + "_n_layers": 64, + "_n_kv_heads": 4, + "_n_heads": 24, + "_embedding_length": 5120, + "_kv_key_length": 256, + "_kv_value_length": 256, + "_full_attention_interval": 4, + "_ssm_inner_size": 6144, + "_ssm_state_size": 128, + } + defaults.update(overrides) + b = LlamaCppBackend() + for k, v in defaults.items(): + setattr(b, k, v) + return b + + def test_qwen35_27b(self): + b = self._hybrid_backend() + # n_attn = 64 // 4 = 16 + expected = 16 * 262144 * 4 * (256 + 256) * 2 + assert b._estimate_kv_cache_bytes(262144, "f16") == expected + + def test_qwen35_35b_a3b(self): + b = self._hybrid_backend( + _n_layers = 40, + _n_kv_heads = 2, + _n_heads = 16, + _embedding_length = 2048, + _ssm_inner_size = 4096, + ) + # n_attn = 40 // 4 = 10 + expected = 10 * 262144 * 2 * (256 + 256) * 2 + assert b._estimate_kv_cache_bytes(262144, "f16") == expected + + def test_hybrid_without_explicit_dims(self): + """Fallback to head_dim when key_length/value_length are missing.""" + b = self._hybrid_backend(_kv_key_length = None, _kv_value_length = None) + head_dim = 5120 // 24 # 213 + expected = 16 * 4096 * 4 * 2 * head_dim * 2 + assert b._estimate_kv_cache_bytes(4096, "f16") == expected + + def test_fai_zero_safety(self): + """full_attention_interval=0 should not cause ZeroDivisionError.""" + b = self._hybrid_backend(_full_attention_interval = 0) + result = b._estimate_kv_cache_bytes(4096, "f16") + # fai=0 -> n_attn = n_layers (all layers) + expected = 64 * 4096 * 4 * (256 + 256) * 2 + assert result == expected + + +# --------------------------------------------------------------------------- +# E. Path 3: Sliding Window Estimation +# --------------------------------------------------------------------------- + + +class TestSlidingWindowEstimation: + """SWA: half global (full ctx) + half sliding window.""" + + def _swa_backend(self, **overrides): + defaults = { + "_n_layers": 62, + "_n_kv_heads": 16, + "_n_heads": 32, + "_embedding_length": 5376, + "_kv_key_length": 128, + "_kv_value_length": 128, + "_sliding_window": 1024, + } + defaults.update(overrides) + b = LlamaCppBackend() + for k, v in defaults.items(): + setattr(b, k, v) + return b + + def test_gemma3(self): + b = self._swa_backend() + # 1/4 heuristic: 62 // 4 = 15 global, 47 SWA + n_global = max(1, 62 // 4) # 15 + n_swa = 62 - n_global # 47 + kv_per = 16 * (128 + 128) * 2 + expected = int(n_global * 131072 * kv_per + n_swa * min(131072, 1024) * kv_per) + assert b._estimate_kv_cache_bytes(131072, "f16") == expected + + def test_gpt_oss(self): + b = self._swa_backend( + _n_layers = 24, + _n_kv_heads = 8, + _n_heads = 64, + _embedding_length = 2880, + _kv_key_length = 64, + _kv_value_length = 64, + _sliding_window = 128, + ) + # 1/4 heuristic: 24 // 4 = 6 global, 18 SWA + n_global = max(1, 24 // 4) # 6 + n_swa = 24 - n_global # 18 + kv_per = 8 * (64 + 64) * 2 + expected = int(n_global * 131072 * kv_per + n_swa * min(131072, 128) * kv_per) + assert b._estimate_kv_cache_bytes(131072, "f16") == expected + + def test_ctx_smaller_than_window(self): + """When context < sliding_window, SWA layers use full context anyway.""" + b = self._swa_backend(_sliding_window = 8192) + n_global = max(1, 62 // 4) # 15 + n_swa = 62 - n_global # 47 + kv_per = 16 * (128 + 128) * 2 + ctx = 4096 + expected = int(n_global * ctx * kv_per + n_swa * min(ctx, 8192) * kv_per) + # min(4096, 8192) = 4096, so both pools use full ctx + assert b._estimate_kv_cache_bytes(ctx, "f16") == expected + + def test_odd_layer_count(self): + """Odd layer count: n_global = max(1, n//4), n_swa = n - n_global.""" + b = self._swa_backend(_n_layers = 63) + n_global = max(1, 63 // 4) # 15 + n_swa = 63 - n_global # 48 + kv_per = 16 * (128 + 128) * 2 + expected = int(n_global * 1000 * kv_per + n_swa * min(1000, 1024) * kv_per) + assert b._estimate_kv_cache_bytes(1000, "f16") == expected + + +# --------------------------------------------------------------------------- +# F. Path 4: Standard GQA Estimation +# --------------------------------------------------------------------------- + + +class TestStandardGQAEstimation: + """Standard GQA with explicit key_length/value_length.""" + + def _gqa_backend(self, **overrides): + defaults = { + "_n_layers": 28, + "_n_kv_heads": 8, + "_n_heads": 16, + "_embedding_length": 1024, + "_kv_key_length": 128, + "_kv_value_length": 128, + } + defaults.update(overrides) + b = LlamaCppBackend() + for k, v in defaults.items(): + setattr(b, k, v) + return b + + def test_qwen3_06b(self): + b = self._gqa_backend() + expected = 28 * 40960 * 8 * (128 + 128) * 2 + assert b._estimate_kv_cache_bytes(40960, "f16") == expected + + def test_asymmetric_kv_dims(self): + """key_length != value_length (some architectures have this).""" + b = self._gqa_backend(_kv_key_length = 192, _kv_value_length = 64) + expected = 28 * 4096 * 8 * (192 + 64) * 2 + assert b._estimate_kv_cache_bytes(4096, "f16") == expected + + def test_differs_from_legacy(self): + """GQA path should differ from legacy when key_length != embed//n_heads.""" + b = self._gqa_backend() + head_dim = 1024 // 16 # 64 + gqa_result = b._estimate_kv_cache_bytes(4096, "f16") + # Legacy would use: 2 * 8 * 64 * 28 * 4096 * 2 + legacy_result = int(2 * 8 * head_dim * 28 * 4096 * 2) + # GQA: 28 * 4096 * 8 * (128+128) * 2 -- uses actual key_length=128 + assert gqa_result != legacy_result + assert gqa_result > legacy_result # key_length (128) > head_dim (64) + + +# --------------------------------------------------------------------------- +# G. Path 5: Legacy Fallback Estimation +# --------------------------------------------------------------------------- + + +class TestLegacyEstimation: + """Legacy: embed // n_heads, for old GGUFs without new fields.""" + + def _legacy_backend(self, **overrides): + defaults = { + "_n_layers": 32, + "_n_kv_heads": 8, + "_n_heads": 32, + "_embedding_length": 4096, + } + defaults.update(overrides) + b = LlamaCppBackend() + for k, v in defaults.items(): + setattr(b, k, v) + return b + + def test_basic_legacy(self): + b = self._legacy_backend() + head_dim = 4096 // 32 # 128 + expected = int(2 * 8 * 128 * 32 * 4096 * 2) + assert b._estimate_kv_cache_bytes(4096, "f16") == expected + + def test_legacy_with_only_n_heads(self): + """n_kv_heads is None, falls back to n_heads.""" + b = self._legacy_backend(_n_kv_heads = None) + head_dim = 4096 // 32 + expected = int(2 * 32 * head_dim * 32 * 4096 * 2) + assert b._estimate_kv_cache_bytes(4096, "f16") == expected + + def test_legacy_identical_to_old_formula(self): + """Confirm legacy path produces the same result as the pre-PR formula.""" + b = self._legacy_backend() + n_layers = 32 + n_kv_heads = 8 + head_dim = 4096 // 32 + n_ctx = 8192 + bpe = 2.0 + old_formula = int(2 * n_kv_heads * head_dim * n_layers * n_ctx * bpe) + assert b._estimate_kv_cache_bytes(n_ctx, "f16") == old_formula + + +# --------------------------------------------------------------------------- +# H. Path Priority (selection order) +# --------------------------------------------------------------------------- + + +class TestPathPriority: + """Confirm: MLA > Hybrid Mamba > SWA > GQA > Legacy.""" + + def test_mla_takes_priority_over_all(self): + """If kv_lora_rank is set, MLA path is used even if other fields are present.""" + b = LlamaCppBackend() + b._n_layers = 61 + b._n_kv_heads = 1 + b._n_heads = 128 + b._embedding_length = 7168 + b._kv_key_length = 576 + b._kv_value_length = 512 + b._kv_lora_rank = 512 + b._ssm_inner_size = 4096 # Would trigger Hybrid + b._full_attention_interval = 4 + b._sliding_window = 1024 # Would trigger SWA + + # MLA: 61 * 1000 * 1 * 576 * 2 + expected_mla = int(61 * 1000 * 1 * 576 * 2) + assert b._estimate_kv_cache_bytes(1000, "f16") == expected_mla + + def test_hybrid_over_swa(self): + """Hybrid takes priority over SWA when both fields present.""" + b = LlamaCppBackend() + b._n_layers = 64 + b._n_kv_heads = 4 + b._n_heads = 24 + b._embedding_length = 5120 + b._kv_key_length = 256 + b._kv_value_length = 256 + b._ssm_inner_size = 6144 + b._full_attention_interval = 4 + b._sliding_window = 1024 # Would trigger SWA + + n_attn = 64 // 4 + expected_hybrid = int(n_attn * 1000 * 4 * (256 + 256) * 2) + assert b._estimate_kv_cache_bytes(1000, "f16") == expected_hybrid + + def test_all_paths_produce_different_values(self): + """With carefully chosen params, each path should yield a distinct value.""" + # Use embedding_length=768 so legacy head_dim (768//16=48) differs from + # key_length (256), and MLA key_len (256) != legacy K+V (2*48=96). + params = { + "_n_layers": 40, + "_n_kv_heads": 4, + "_n_heads": 16, + "_embedding_length": 768, + "_kv_key_length": 256, + "_kv_value_length": 256, + } + ctx = 4096 + + # Path 4: Standard GQA + b_gqa = LlamaCppBackend() + for k, v in params.items(): + setattr(b_gqa, k, v) + gqa_val = b_gqa._estimate_kv_cache_bytes(ctx, "f16") + + # Path 1: MLA + b_mla = LlamaCppBackend() + for k, v in params.items(): + setattr(b_mla, k, v) + b_mla._kv_lora_rank = 512 + mla_val = b_mla._estimate_kv_cache_bytes(ctx, "f16") + + # Path 2: Hybrid Mamba + b_hybrid = LlamaCppBackend() + for k, v in params.items(): + setattr(b_hybrid, k, v) + b_hybrid._ssm_inner_size = 4096 + b_hybrid._full_attention_interval = 4 + hybrid_val = b_hybrid._estimate_kv_cache_bytes(ctx, "f16") + + # Path 3: SWA + b_swa = LlamaCppBackend() + for k, v in params.items(): + setattr(b_swa, k, v) + b_swa._sliding_window = 512 + swa_val = b_swa._estimate_kv_cache_bytes(ctx, "f16") + + # Path 5: Legacy (no key_length/value_length) + b_legacy = LlamaCppBackend() + b_legacy._n_layers = 40 + b_legacy._n_kv_heads = 4 + b_legacy._n_heads = 16 + b_legacy._embedding_length = 768 + legacy_val = b_legacy._estimate_kv_cache_bytes(ctx, "f16") + + values = [mla_val, hybrid_val, swa_val, gqa_val, legacy_val] + assert len(set(values)) == 5, f"Expected 5 distinct values, got {values}" + + +# --------------------------------------------------------------------------- +# I. KV Cache Quantization +# --------------------------------------------------------------------------- + + +class TestQuantization: + """Verify all supported cache_type_kv values produce correct scaling.""" + + @pytest.mark.parametrize( + "cache_type,expected_bpe", + [ + ("f32", 4.0), + ("f16", 2.0), + ("bf16", 2.0), + ("q8_0", 34 / 32), + ("q5_1", 0.75), + ("q5_0", 0.6875), + ("q4_1", 0.625), + ("q4_0", 0.5625), + ("iq4_nl", 0.5625), + (None, 2.0), # default is f16 + ("unknown", 2.0), # unknown falls back to f16 + ], + ) + def test_quantization_scaling(self, cache_type, expected_bpe): + b = LlamaCppBackend() + b._n_layers = 10 + b._n_kv_heads = 1 + b._n_heads = 8 + b._embedding_length = 512 + b._kv_key_length = 64 + b._kv_value_length = 64 + result = b._estimate_kv_cache_bytes(1000, cache_type) + expected = int(10 * 1000 * 1 * (64 + 64) * expected_bpe) + assert result == expected + + +# --------------------------------------------------------------------------- +# J. Edge Cases +# --------------------------------------------------------------------------- + + +class TestEdgeCases: + """Boundary conditions and degenerate inputs.""" + + def test_zero_context(self): + b = LlamaCppBackend() + b._n_layers = 32 + b._kv_key_length = 128 + assert b._estimate_kv_cache_bytes(0, "f16") == 0 + + def test_negative_context(self): + b = LlamaCppBackend() + b._n_layers = 32 + b._kv_key_length = 128 + assert b._estimate_kv_cache_bytes(-1, "f16") == 0 + + def test_context_of_one(self): + b = LlamaCppBackend() + b._n_layers = 10 + b._n_kv_heads = 1 + b._kv_key_length = 64 + b._kv_value_length = 64 + result = b._estimate_kv_cache_bytes(1, "f16") + assert result == int(10 * 1 * 1 * (64 + 64) * 2) + + def test_very_large_context(self): + """1M context should not overflow or crash.""" + b = LlamaCppBackend() + b._n_layers = 10 + b._n_kv_heads = 1 + b._kv_key_length = 128 + b._kv_value_length = 128 + result = b._estimate_kv_cache_bytes(1_000_000, "f16") + assert result > 0 + assert isinstance(result, int) + + def test_n_kv_heads_none_falls_to_n_heads(self): + b = LlamaCppBackend() + b._n_layers = 10 + b._n_kv_heads = None + b._n_heads = 8 + b._kv_key_length = 64 + b._kv_value_length = 64 + result = b._estimate_kv_cache_bytes(100, "f16") + expected = int(10 * 100 * 8 * (64 + 64) * 2) + assert result == expected + + def test_both_heads_none_falls_to_one(self): + b = LlamaCppBackend() + b._n_layers = 10 + b._n_kv_heads = None + b._n_heads = None + b._kv_key_length = 64 + b._kv_value_length = 64 + result = b._estimate_kv_cache_bytes(100, "f16") + expected = int(10 * 100 * 1 * (64 + 64) * 2) + assert result == expected + + +# --------------------------------------------------------------------------- +# K. Lifecycle Tests +# --------------------------------------------------------------------------- + + +class TestLifecycle: + """Init, unload, and reparse field management.""" + + def test_init_fields_none(self): + b = LlamaCppBackend() + for attr in [ + "_kv_key_length", + "_kv_value_length", + "_sliding_window", + "_full_attention_interval", + "_kv_lora_rank", + "_key_length_mla", + "_ssm_inner_size", + "_ssm_state_size", + ]: + assert getattr(b, attr) is None + + def test_unload_resets_fields(self): + b = LlamaCppBackend() + b._n_layers = 32 + b._kv_key_length = 128 + b._kv_lora_rank = 512 + b._sliding_window = 1024 + b._ssm_inner_size = 4096 + b._full_attention_interval = 4 + b.unload_model() + for attr in [ + "_kv_key_length", + "_kv_value_length", + "_sliding_window", + "_full_attention_interval", + "_kv_lora_rank", + "_key_length_mla", + "_ssm_inner_size", + "_ssm_state_size", + ]: + assert getattr(b, attr) is None + + def test_end_to_end_synthetic_mla(self): + """Full round-trip: write GGUF -> parse -> estimate.""" + b = _backend_from_gguf( + "deepseek2", + { + "context_length": 163840, + "block_count": 61, + "attention.head_count_kv": 1, + "attention.head_count": 128, + "embedding_length": 7168, + "attention.key_length": 576, + "attention.value_length": 512, + "attention.kv_lora_rank": 512, + "attention.key_length_mla": 192, + }, + ) + assert b._can_estimate_kv() + result = b._estimate_kv_cache_bytes(163840, "f16") + expected = 61 * 163840 * 1 * 576 * 2 + assert result == expected + + def test_end_to_end_synthetic_hybrid(self): + b = _backend_from_gguf( + "qwen35", + { + "context_length": 262144, + "block_count": 64, + "attention.head_count_kv": 4, + "attention.head_count": 24, + "embedding_length": 5120, + "attention.key_length": 256, + "attention.value_length": 256, + "full_attention_interval": 4, + "ssm.inner_size": 6144, + "ssm.state_size": 128, + }, + ) + assert b._can_estimate_kv() + result = b._estimate_kv_cache_bytes(262144, "f16") + n_attn = 64 // 4 + expected = n_attn * 262144 * 4 * (256 + 256) * 2 + assert result == expected + + def test_end_to_end_synthetic_swa(self): + b = _backend_from_gguf( + "gemma3", + { + "context_length": 131072, + "block_count": 62, + "attention.head_count_kv": 16, + "attention.head_count": 32, + "embedding_length": 5376, + "attention.key_length": 128, + "attention.value_length": 128, + "attention.sliding_window": 1024, + }, + ) + assert b._can_estimate_kv() + result = b._estimate_kv_cache_bytes(131072, "f16") + n_global = max(1, 62 // 4) # 15 + n_swa = 62 - n_global # 47 + kv_per = 16 * 256 * 2 + expected = int(n_global * 131072 * kv_per + n_swa * 1024 * kv_per) + assert result == expected + + def test_end_to_end_synthetic_gqa(self): + b = _backend_from_gguf( + "qwen3", + { + "context_length": 40960, + "block_count": 28, + "attention.head_count_kv": 8, + "attention.head_count": 16, + "embedding_length": 1024, + "attention.key_length": 128, + "attention.value_length": 128, + }, + ) + assert b._can_estimate_kv() + result = b._estimate_kv_cache_bytes(40960, "f16") + expected = 28 * 40960 * 8 * 256 * 2 + assert result == expected + + def test_end_to_end_synthetic_legacy(self): + b = _backend_from_gguf( + "llama", + { + "context_length": 4096, + "block_count": 32, + "attention.head_count_kv": 8, + "attention.head_count": 32, + "embedding_length": 4096, + }, + ) + assert b._can_estimate_kv() + result = b._estimate_kv_cache_bytes(4096, "f16") + head_dim = 4096 // 32 + expected = int(2 * 8 * head_dim * 32 * 4096 * 2) + assert result == expected diff --git a/studio/backend/tests/test_native_context_length.py b/studio/backend/tests/test_native_context_length.py new file mode 100644 index 0000000000..7c69e56f89 --- /dev/null +++ b/studio/backend/tests/test_native_context_length.py @@ -0,0 +1,518 @@ +# SPDX-License-Identifier: AGPL-3.0-only +# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0 + +"""Tests for the native_context_length feature (PR #4746). + +Verifies that the new `native_context_length` property on LlamaCppBackend +and the corresponding Pydantic model fields work correctly. The raw GGUF +`_context_length` must never be overwritten by VRAM-capping logic. + +Requires no GPU, network, or external libraries beyond pytest and pydantic. +""" + +import io +import json +import struct +import sys +import types as _types +from pathlib import Path +from unittest.mock import patch + +import pytest + +# --------------------------------------------------------------------------- +# Stub heavy / unavailable external dependencies before importing the +# module under test. Same pattern as test_kv_cache_estimation.py. +# --------------------------------------------------------------------------- + +_BACKEND_DIR = str(Path(__file__).resolve().parent.parent) +if _BACKEND_DIR not in sys.path: + sys.path.insert(0, _BACKEND_DIR) + +# loggers +_loggers_stub = _types.ModuleType("loggers") +_loggers_stub.get_logger = lambda name: __import__("logging").getLogger(name) +sys.modules.setdefault("loggers", _loggers_stub) + +# structlog +_structlog_stub = _types.ModuleType("structlog") +sys.modules.setdefault("structlog", _structlog_stub) + +# httpx -- stub only the names referenced at import / class-definition time +_httpx_stub = _types.ModuleType("httpx") +for _exc_name in ( + "ConnectError", + "TimeoutException", + "ReadTimeout", + "ReadError", + "RemoteProtocolError", + "CloseError", +): + setattr(_httpx_stub, _exc_name, type(_exc_name, (Exception,), {})) + + +class _FakeTimeout: + def __init__(self, *a, **kw): + pass + + +_httpx_stub.Timeout = _FakeTimeout +_httpx_stub.Client = type( + "Client", + (), + { + "__init__": lambda self, **kw: None, + "__enter__": lambda self: self, + "__exit__": lambda self, *a: None, + }, +) +sys.modules.setdefault("httpx", _httpx_stub) + +from core.inference.llama_cpp import LlamaCppBackend +from models.inference import LoadResponse, InferenceStatusResponse + + +# ── Helpers ────────────────────────────────────────────────────────── + + +def _write_kv(buf: io.BytesIO, key: str, value, vtype: int) -> None: + """Append a single GGUF KV pair to *buf*.""" + key_bytes = key.encode("utf-8") + buf.write(struct.pack(" str: + """Create a minimal valid GGUF v3 binary in *tmp_path*.""" + buf = io.BytesIO() + buf.write(struct.pack("= max >= effective holds when VRAM-capped.""" + backend._context_length = 131072 + backend._max_context_length = 65536 + backend._effective_context_length = 32768 + assert backend.native_context_length >= backend.max_context_length + assert backend.max_context_length >= backend.context_length + + def test_all_equal_when_uncapped(self, backend): + """All three equal when no VRAM constraint.""" + backend._context_length = 8192 + # No effective or max set -- properties fall back to _context_length + assert backend.native_context_length == 8192 + assert backend.max_context_length == 8192 + assert backend.context_length == 8192 + + def test_fit_context_does_not_modify(self, backend): + """_fit_context_to_vram() does not touch _context_length.""" + backend._context_length = 131072 + backend._n_layers = 32 + backend._n_kv_heads = 8 + backend._n_heads = 32 + backend._embedding_length = 4096 + original = backend._context_length + + # Simulate a very small VRAM budget that forces capping + result = backend._fit_context_to_vram( + requested_ctx = 131072, + available_mib = 512, # very small + model_size_bytes = 0, + ) + # _fit_context_to_vram returns the capped value, not modifying _context_length + assert backend._context_length == original + assert backend.native_context_length == original + # The returned capped value should be <= requested + assert result <= 131072 + + def test_native_gt_context_when_capped(self, backend): + """native_context_length > context_length after VRAM capping.""" + backend._context_length = 131072 + backend._effective_context_length = 16384 + assert backend.native_context_length > backend.context_length + + +# ===================================================================== +# C. TestPydanticModels -- LoadResponse & InferenceStatusResponse +# ===================================================================== + + +class TestPydanticModels: + """Tests native_context_length field on Pydantic models.""" + + def test_load_response_has_field(self): + """Field exists in LoadResponse.model_fields.""" + assert "native_context_length" in LoadResponse.model_fields + + def test_load_response_defaults_none(self): + """Omitting native_context_length defaults to None.""" + resp = LoadResponse( + status = "loaded", + model = "test", + display_name = "Test", + inference = {}, + ) + assert resp.native_context_length is None + + def test_load_response_accepts_int(self): + """native_context_length=131072 stores correctly.""" + resp = LoadResponse( + status = "loaded", + model = "test", + display_name = "Test", + inference = {}, + native_context_length = 131072, + ) + assert resp.native_context_length == 131072 + + def test_load_response_json_null(self): + """None serializes to JSON null.""" + resp = LoadResponse( + status = "loaded", + model = "test", + display_name = "Test", + inference = {}, + ) + data = json.loads(resp.model_dump_json()) + assert data["native_context_length"] is None + + def test_load_response_json_int(self): + """131072 serializes to JSON number.""" + resp = LoadResponse( + status = "loaded", + model = "test", + display_name = "Test", + inference = {}, + native_context_length = 131072, + ) + data = json.loads(resp.model_dump_json()) + assert data["native_context_length"] == 131072 + + def test_status_response_has_field(self): + """Field exists in InferenceStatusResponse.model_fields.""" + assert "native_context_length" in InferenceStatusResponse.model_fields + + def test_status_response_defaults_none(self): + """Omitting native_context_length defaults to None.""" + resp = InferenceStatusResponse() + assert resp.native_context_length is None + + def test_roundtrip_preserves_value(self): + """model_validate_json(model_dump_json()) round-trips.""" + resp = LoadResponse( + status = "loaded", + model = "test", + display_name = "Test", + inference = {}, + native_context_length = 131072, + ) + roundtripped = LoadResponse.model_validate_json(resp.model_dump_json()) + assert roundtripped.native_context_length == 131072 + + +# ===================================================================== +# D. TestRouteCompleteness -- source-level verification +# ===================================================================== + + +class TestRouteCompleteness: + """All response construction sites in routes/inference.py include native_context_length.""" + + @pytest.fixture(autouse = True) + def _load_source(self): + """Read routes/inference.py source once.""" + routes_path = Path(__file__).resolve().parent.parent / "routes" / "inference.py" + self._source = routes_path.read_text() + + def _find_construction_blocks(self, class_name: str) -> list[str]: + """Extract all code blocks that construct a given response class.""" + blocks = [] + idx = 0 + while True: + start = self._source.find(f"{class_name}(", idx) + if start == -1: + break + # Find matching closing paren (simple depth counter) + depth = 0 + end = start + for i, ch in enumerate(self._source[start:], start): + if ch == "(": + depth += 1 + elif ch == ")": + depth -= 1 + if depth == 0: + end = i + 1 + break + blocks.append(self._source[start:end]) + idx = end + return blocks + + def test_gguf_load_responses_have_field(self): + """Every GGUF LoadResponse (is_gguf = True) includes native_context_length.""" + blocks = self._find_construction_blocks("LoadResponse") + gguf_blocks = [ + b for b in blocks if "is_gguf = True" in b or "is_gguf=True" in b + ] + assert ( + len(gguf_blocks) >= 2 + ), f"Expected at least 2 GGUF LoadResponse blocks, found {len(gguf_blocks)}" + for i, block in enumerate(gguf_blocks): + assert ( + "native_context_length" in block + ), f"GGUF LoadResponse block #{i} missing native_context_length:\n{block[:200]}" + + def test_non_gguf_load_responses_omit_field(self): + """Non-GGUF LoadResponse blocks do not set native_context_length (defaults to None).""" + blocks = self._find_construction_blocks("LoadResponse") + non_gguf = [ + b for b in blocks if "is_gguf = True" not in b and "is_gguf=True" not in b + ] + # Non-GGUF paths should not reference native_context_length + # (Pydantic defaults it to None, so not setting it is correct) + for block in non_gguf: + assert ( + "native_context_length" not in block + ), f"Non-GGUF LoadResponse should not set native_context_length:\n{block[:200]}" + + def test_status_path(self): + """InferenceStatusResponse construction with llama_backend has the field.""" + blocks = self._find_construction_blocks("InferenceStatusResponse") + found = False + for block in blocks: + if "llama_backend" in block and "native_context_length" in block: + found = True + break + assert found, "No InferenceStatusResponse block with llama_backend has native_context_length" + + +# ===================================================================== +# E. TestEdgeCases +# ===================================================================== + + +class TestNativeContextEdgeCases: + """Edge cases for native_context_length.""" + + def test_context_length_zero(self, tmp_path, backend): + """GGUF context_length=0 returns 0, not None.""" + path = make_gguf(tmp_path, "llama", [("context_length", 0, 4)]) + backend._read_gguf_metadata(path) + assert backend.native_context_length == 0 + + def test_context_length_uint32_max(self, tmp_path, backend): + """2^32 - 1 survives without truncation.""" + val = 2**32 - 1 + path = make_gguf(tmp_path, "llama", [("context_length", val, 4)]) + backend._read_gguf_metadata(path) + assert backend.native_context_length == val + + def test_context_length_uint64(self, tmp_path, backend): + """UINT64 type context_length parsed correctly.""" + val = 2**33 # exceeds UINT32 range + path = make_gguf(tmp_path, "llama", [("context_length", val, 10)]) + backend._read_gguf_metadata(path) + assert backend.native_context_length == val + + def test_no_context_length_in_gguf(self, tmp_path, backend): + """GGUF without context_length key yields None.""" + path = make_gguf(tmp_path, "llama", [("block_count", 32, 4)]) + backend._read_gguf_metadata(path) + assert backend.native_context_length is None + + def test_native_equals_context_when_uncapped(self, backend): + """Both equal when no VRAM cap applied.""" + backend._context_length = 8192 + assert backend.native_context_length == backend.context_length + + def test_native_survives_parse_then_cap(self, tmp_path, backend): + """Parse then set effective cap: native unchanged.""" + path = make_gguf( + tmp_path, + "llama", + [ + ("context_length", 131072, 4), + ("block_count", 32, 4), + ("attention.head_count", 32, 4), + ("attention.head_count_kv", 8, 4), + ("embedding_length", 4096, 4), + ], + ) + backend._read_gguf_metadata(path) + assert backend.native_context_length == 131072 + + # Simulate VRAM capping by setting effective and max + backend._effective_context_length = 16384 + backend._max_context_length = 32768 + assert backend.native_context_length == 131072 + + +# ===================================================================== +# F. TestCrossPlatform -- binary I/O and serialization +# ===================================================================== + + +class TestCrossPlatform: + """Binary I/O and serialization correctness across platforms.""" + + def test_le_uint32_context_length(self, tmp_path, backend): + """Little-endian UINT32 parsed correctly.""" + path = make_gguf(tmp_path, "llama", [("context_length", 16384, 4)]) + backend._read_gguf_metadata(path) + assert backend.native_context_length == 16384 + + def test_le_uint64_context_length(self, tmp_path, backend): + """Little-endian UINT64 parsed correctly.""" + path = make_gguf(tmp_path, "llama", [("context_length", 16384, 10)]) + backend._read_gguf_metadata(path) + assert backend.native_context_length == 16384 + + def test_gguf_magic_le_byte_order(self, tmp_path): + """Magic 0x46554747 matches GGUF spec (little-endian 'GGUF').""" + path = tmp_path / "magic_check.gguf" + buf = io.BytesIO() + buf.write(struct.pack(" bool: if model_identifier.lower() in _REVERSE_MODEL_MAPPING: return True - # Check for exact filename match - model_filename = model_identifier.replace("/", "_") + ".yaml" + # For local filesystem paths (e.g. C:\Users\...\model on Windows), + # normalize backslashes so Path().parts splits correctly on POSIX/WSL, + # then try matching the last 1-2 path components against the registry + # (mirrors the logic in load_model_defaults). + _is_local = is_local_path(model_identifier) + _normalized = normalize_path(model_identifier) if _is_local else model_identifier + + if _is_local: + parts = Path(_normalized).parts + for depth in (2, 1): + if len(parts) >= depth: + suffix = "/".join(parts[-depth:]) + if suffix.lower() in _REVERSE_MODEL_MAPPING: + return True + _lookup = Path(_normalized).name + else: + _lookup = model_identifier + + # Check for exact filename match (basename for local paths to avoid + # passing absolute paths into rglob which raises + # "Non-relative patterns are unsupported" on Windows). + model_filename = _lookup.replace("/", "_") + ".yaml" for config_path in defaults_dir.rglob(model_filename): if config_path.is_file(): return True diff --git a/studio/backend/utils/models/__init__.py b/studio/backend/utils/models/__init__.py index 82236d8013..a81682d8b7 100644 --- a/studio/backend/utils/models/__init__.py +++ b/studio/backend/utils/models/__init__.py @@ -19,6 +19,7 @@ from .model_config import ( get_base_model_from_lora, load_model_config, list_gguf_variants, + extract_model_size_b, MODEL_NAME_MAPPING, UI_STATUS_INDICATORS, ) @@ -38,6 +39,7 @@ __all__ = [ "get_base_model_from_lora", "load_model_config", "list_gguf_variants", + "extract_model_size_b", "MODEL_NAME_MAPPING", "UI_STATUS_INDICATORS", "scan_checkpoints", diff --git a/studio/backend/utils/models/model_config.py b/studio/backend/utils/models/model_config.py index 5de3fd2cf9..df1058abf6 100644 --- a/studio/backend/utils/models/model_config.py +++ b/studio/backend/utils/models/model_config.py @@ -31,6 +31,37 @@ import yaml logger = get_logger(__name__) +# ── Model size extraction ──────────────────────────────────── +import re as _re + +_MODEL_SIZE_RE = _re.compile( + r"(?:^|[-_/])(\d+\.?\d*)\s*([bm])(?:$|[-_/])", _re.IGNORECASE +) +# MoE active-parameter pattern: matches "A3B", "A3.5B", etc. +_ACTIVE_SIZE_RE = _re.compile( + r"(?:^|[-_/])a(\d+\.?\d*)\s*([bm])(?:$|[-_/])", _re.IGNORECASE +) + + +def extract_model_size_b(model_id: str) -> float | None: + """Extract model size in billions from a model identifier. + + Prefers MoE active-parameter notation (e.g. ``A3B`` in + ``Qwen3.5-35B-A3B``) over the total parameter count. + Handles both ``B`` (billions) and ``M`` (millions) suffixes. + """ + mid = (model_id or "").lower() + active = _ACTIVE_SIZE_RE.search(mid) + if active: + val = float(active.group(1)) + return val / 1000.0 if active.group(2).lower() == "m" else val + size = _MODEL_SIZE_RE.search(mid) + if not size: + return None + val = float(size.group(1)) + return val / 1000.0 if size.group(2).lower() == "m" else val + + # Model name mapping: maps all equivalent model names to their canonical YAML config file # Format: "canonical_model_name.yaml": [list of all equivalent model names] # Based on the model mapper provided - canonical filename is based on the first model name in the mapper @@ -1420,17 +1451,20 @@ def load_model_defaults(model_name: str) -> Dict[str, Any]: return config # If model_name is a local path (e.g. /home/.../Spark-TTS-0.5B/LLM from - # adapter_config.json), try matching the last 1-2 path components against - # the registry (e.g. "Spark-TTS-0.5B/LLM"). - if model_name not in _REVERSE_MODEL_MAPPING and ( - model_name.startswith("/") or model_name.startswith(".") - ): - parts = Path(model_name).parts + # adapter_config.json, or C:\Users\...\model on Windows), try matching + # the last 1-2 path components against the registry + # (e.g. "Spark-TTS-0.5B/LLM"). + _is_local_path = is_local_path(model_name) + # Normalize Windows backslash paths so Path().parts splits correctly + # on POSIX/WSL hosts (pathlib treats backslashes as literals on Linux). + _normalized = normalize_path(model_name) if _is_local_path else model_name + if model_name.lower() not in _REVERSE_MODEL_MAPPING and _is_local_path: + parts = Path(_normalized).parts for depth in [2, 1]: if len(parts) >= depth: suffix = "/".join(parts[-depth:]) - if suffix in _REVERSE_MODEL_MAPPING: - canonical_file = _REVERSE_MODEL_MAPPING[suffix] + if suffix.lower() in _REVERSE_MODEL_MAPPING: + canonical_file = _REVERSE_MODEL_MAPPING[suffix.lower()] for config_path in defaults_dir.rglob(canonical_file): if config_path.is_file(): with open(config_path, "r", encoding = "utf-8") as f: @@ -1440,8 +1474,12 @@ def load_model_defaults(model_name: str) -> Dict[str, Any]: ) return config - # Try exact model name match (for backward compatibility) - model_filename = model_name.replace("/", "_") + ".yaml" + # Try exact model name match (for backward compatibility). + # For local filesystem paths, use only the directory basename to + # avoid passing absolute paths (e.g. C:\...) into rglob which + # raises "Non-relative patterns are unsupported" on Windows. + _lookup_name = Path(_normalized).name if _is_local_path else model_name + model_filename = _lookup_name.replace("/", "_") + ".yaml" # Search in subfolders and root for config_path in defaults_dir.rglob(model_filename): if config_path.is_file(): diff --git a/studio/backend/utils/transformers_version.py b/studio/backend/utils/transformers_version.py index d8724de723..07e4a5c000 100644 --- a/studio/backend/utils/transformers_version.py +++ b/studio/backend/utils/transformers_version.py @@ -47,6 +47,8 @@ TRANSFORMERS_5_MODEL_SUBSTRINGS: tuple[str, ...] = ( "qwen3.5", # Qwen3.5 family (35B-A3B, etc.) "qwen3-next", # Qwen3-Next and variants "tiny_qwen3_moe", # imdatta0/tiny_qwen3_moe_2.8B_0.7B + "gemma-4", # Gemma-4 (E2B-it, E4B-it, 31B-it, 26B-A4B-it) + "gemma4", # Gemma-4 alternate naming ) # Tokenizer classes that only exist in transformers>=5.x @@ -58,7 +60,7 @@ _TRANSFORMERS_5_TOKENIZER_CLASSES: set[str] = { _tokenizer_class_cache: dict[str, bool] = {} # Versions -TRANSFORMERS_5_VERSION = "5.3.0" +TRANSFORMERS_5_VERSION = "5.5.0" TRANSFORMERS_DEFAULT_VERSION = "4.57.6" # Pre-installed directory for transformers 5.x — created by setup.sh / setup.ps1 @@ -258,7 +260,7 @@ def _purge_modules() -> int: _VENV_T5_PACKAGES = ( f"transformers=={TRANSFORMERS_5_VERSION}", - "huggingface_hub==1.7.1", + "huggingface_hub==1.8.0", "hf_xet==1.4.2", "tiktoken", ) diff --git a/studio/frontend/src/components/assistant-ui/markdown-text.tsx b/studio/frontend/src/components/assistant-ui/markdown-text.tsx index 5e84b9175e..91ef78fcf9 100644 --- a/studio/frontend/src/components/assistant-ui/markdown-text.tsx +++ b/studio/frontend/src/components/assistant-ui/markdown-text.tsx @@ -100,20 +100,27 @@ function getCodeFilename(language: string | null) { function isSvgFence(codeFence: CodeFence): boolean { const lang = codeFence.language?.toLowerCase() ?? ""; if (lang === "svg") return true; - if ((lang === "xml" || lang === "html") && codeFence.source.trimStart().startsWith(" followed by ]|on\w+\s*=|javascript:|]|]|]|]/i; function sanitizeSvg(source: string): string | null { if (UNSAFE_SVG_RE.test(source)) return null; - return source; + // Strip XML declaration () -- not needed for data URI + // rendering and can cause issues with some renderers. + return source.replace(/^\s*<\?xml[^?]*\?>\s*/i, ""); } function SvgPreview({ source }: { source: string }) { @@ -403,7 +410,7 @@ const MarkdownTextImpl = () => { } return ( -
+
void; onEject?: () => void; + onFoldersChange?: () => void; variant?: "outline" | "ghost" | "muted"; size?: "sm" | "default" | "lg"; className?: string; @@ -100,6 +101,7 @@ function ModelSelectorContent({ value, onSelect, onEject, + onFoldersChange, className, dataTour, }: { @@ -108,6 +110,7 @@ function ModelSelectorContent({ value?: string; onSelect: (id: string, meta: ModelSelectorChangeMeta) => void; onEject?: () => void; + onFoldersChange?: () => void; className?: string; dataTour?: string; }) { @@ -124,7 +127,7 @@ function ModelSelectorContent({ )} > {chatOnly ? ( - + ) : ( @@ -133,7 +136,7 @@ function ModelSelectorContent({ - + @@ -171,6 +174,7 @@ export function ModelSelector({ activeGgufVariant, onValueChange, onEject, + onFoldersChange, variant = "outline", size = "default", className, @@ -253,6 +257,7 @@ export function ModelSelector({ value={selected} onSelect={handleSelect} onEject={onEject ? handleEject : undefined} + onFoldersChange={onFoldersChange} className={contentClassName} dataTour={contentDataTour} /> diff --git a/studio/frontend/src/components/assistant-ui/model-selector/pickers.tsx b/studio/frontend/src/components/assistant-ui/model-selector/pickers.tsx index cf8b4cd54e..74ca2542d4 100644 --- a/studio/frontend/src/components/assistant-ui/model-selector/pickers.tsx +++ b/studio/frontend/src/components/assistant-ui/model-selector/pickers.tsx @@ -20,11 +20,15 @@ import { } from "@/components/ui/tooltip"; import { usePlatformStore } from "@/config/env"; import { + type ScanFolderInfo, + addScanFolder, deleteCachedModel, listCachedGguf, listCachedModels, listGgufVariants, listLocalModels, + listScanFolders, + removeScanFolder, } from "@/features/chat/api/chat-api"; import type { CachedGgufRepo, @@ -42,7 +46,7 @@ import { import { cn, formatCompact } from "@/lib/utils"; import type { VramFitStatus } from "@/lib/vram"; import { checkVramFit, estimateLoadingVram } from "@/lib/vram"; -import { Search01Icon } from "@hugeicons/core-free-icons"; +import { Add01Icon, Cancel01Icon, Folder02Icon, Search01Icon } from "@hugeicons/core-free-icons"; import { HugeiconsIcon } from "@hugeicons/react"; import { Trash2Icon } from "lucide-react"; import { @@ -123,23 +127,22 @@ function ModelRow({ className={cn( "flex w-full items-center gap-2 rounded-md px-2.5 py-1.5 text-left text-sm transition-colors hover:bg-accent", selected && "bg-accent/60", - exceeds && "opacity-50", )} > {label} {vramStatus === "exceeds" && ( - OOM + OOM )} {vramStatus === "tight" && ( - TIGHT + TIGHT )} {meta ? ( {meta} @@ -350,7 +353,7 @@ function GgufVariantExpander({ )} > - {v.quant} + {v.quant} {v.downloaded ? ( downloaded @@ -363,12 +366,12 @@ function GgufVariantExpander({ {oom && ( - + OOM )} {tight && ( - + TIGHT )} @@ -415,6 +418,7 @@ let _cachedGgufCache: CachedGgufRepo[] = []; let _cachedModelsCache: CachedModelRepo[] = []; let _lmStudioCache: LocalModelInfo[] = []; let _customFolderCache: LocalModelInfo[] = []; +let _scanFoldersCache: ScanFolderInfo[] = []; /** Sort LM Studio models with unsloth publisher first. */ function sortLmStudio(models: LocalModelInfo[]): LocalModelInfo[] { @@ -434,10 +438,12 @@ export function HubModelPicker({ models, value, onSelect, + onFoldersChange, }: { models: ModelOption[]; value?: string; onSelect: (id: string, meta: ModelSelectorChangeMeta) => void; + onFoldersChange?: () => void; }) { const gpu = useGpuInfo(); const [query, setQuery] = useState(""); @@ -469,6 +475,13 @@ export function HubModelPicker({ const [customFolderModels, setCustomFolderModels] = useState(_customFolderCache); + // Custom scan folders management + const [scanFolders, setScanFolders] = useState(_scanFoldersCache); + const [folderInput, setFolderInput] = useState(""); + const [folderError, setFolderError] = useState(null); + const [showFolderInput, setShowFolderInput] = useState(false); + const [folderLoading, setFolderLoading] = useState(false); + const refreshLocalModelsList = useCallback(() => { listLocalModels() .then((res) => { @@ -484,6 +497,57 @@ export function HubModelPicker({ .catch(() => {}); }, []); + const refreshScanFolders = useCallback(() => { + listScanFolders() + .then((v) => { + _scanFoldersCache = v; + setScanFolders(v); + }) + .catch(() => {}); + }, []); + + const handleAddFolder = useCallback(async () => { + const trimmed = folderInput.trim(); + if (!trimmed || folderLoading) return; + setFolderError(null); + setFolderLoading(true); + try { + const created = await addScanFolder(trimmed); + // Backend returns existing row for duplicates, so deduplicate + const next = _scanFoldersCache.some((f) => f.id === created.id || f.path === created.path) + ? _scanFoldersCache + : [..._scanFoldersCache, created]; + _scanFoldersCache = next; + setScanFolders(next); + setFolderInput(""); + setShowFolderInput(false); + refreshLocalModelsList(); + onFoldersChange?.(); + // Background reconciliation with the server + void refreshScanFolders(); + } catch (e) { + setFolderError(e instanceof Error ? e.message : "Failed to add folder"); + } finally { + setFolderLoading(false); + } + }, [folderInput, folderLoading, refreshScanFolders, refreshLocalModelsList, onFoldersChange]); + + const handleRemoveFolder = useCallback(async (id: number) => { + try { + await removeScanFolder(id); + // Optimistic update so the folder disappears immediately + const next = _scanFoldersCache.filter((f) => f.id !== id); + _scanFoldersCache = next; + setScanFolders(next); + refreshScanFolders(); + refreshLocalModelsList(); + onFoldersChange?.(); + } catch (e) { + toast.error(e instanceof Error ? e.message : "Failed to remove folder"); + refreshScanFolders(); + } + }, [refreshScanFolders, refreshLocalModelsList, onFoldersChange]); + const refreshCachedLists = useCallback(() => { listCachedGguf() .then((v) => { @@ -503,6 +567,7 @@ export function HubModelPicker({ useEffect(() => { // Always refresh LM Studio + custom folder models (not gated by alreadyCached) refreshLocalModelsList(); + refreshScanFolders(); if (alreadyCached) return; let done = 0; @@ -523,7 +588,7 @@ export function HubModelPicker({ }) .catch(() => {}) .finally(check); - }, [alreadyCached]); + }, [alreadyCached, refreshLocalModelsList, refreshScanFolders]); const handleDeleteConfirm = useCallback(async () => { if (!deleteTarget) return; @@ -878,9 +943,95 @@ export function HubModelPicker({ ) : null} - {!showHfSection && customFolderModels.length > 0 ? ( + {!showHfSection ? ( <> - Custom Folders +
+ + Custom Folders + + +
+ + {/* Folder paths */} + {scanFolders.map((f) => ( +
+ + + {f.path} + + +
+ ))} + + {/* Add folder input */} + {showFolderInput && ( +
+
+ + { setFolderInput(e.target.value); setFolderError(null); }} + onKeyDown={(e) => { + if (e.key === "Enter") { e.preventDefault(); handleAddFolder(); } + if (e.key === "Escape") { e.preventDefault(); e.stopPropagation(); setShowFolderInput(false); setFolderInput(""); setFolderError(null); } + }} + placeholder="/path/to/models" + className="h-6 min-w-0 flex-1 rounded border border-border/50 bg-transparent px-1.5 font-mono text-[10px] text-foreground outline-none placeholder:text-muted-foreground/40 focus:border-foreground/20" + disabled={folderLoading} + autoFocus={true} + /> + +
+ {folderError && ( +

{folderError}

+ )} +
+ )} + + {/* Empty state */} + {scanFolders.length === 0 && customFolderModels.length === 0 && !showFolderInput && ( + + )} + + {/* Models from custom folders */} {customFolderModels.map((m) => { const isGguf = isGgufRepo(m.id) || diff --git a/studio/frontend/src/components/assistant-ui/thread.tsx b/studio/frontend/src/components/assistant-ui/thread.tsx index d688822815..4db0eda0ad 100644 --- a/studio/frontend/src/components/assistant-ui/thread.tsx +++ b/studio/frontend/src/components/assistant-ui/thread.tsx @@ -73,7 +73,7 @@ export const Thread: FC<{ hideComposer?: boolean; hideWelcome?: boolean }> = ({ }} > {!hideWelcome && ( thread.isEmpty}> @@ -89,7 +89,13 @@ export const Thread: FC<{ hideComposer?: boolean; hideWelcome?: boolean }> = ({ }} /> - + + {!hideComposer && ( +
+ )} !thread.isEmpty}> {!hideComposer && } @@ -118,7 +124,7 @@ const SUGGESTION_TOOLS: Record> = "How do you fine-tune an audio model with Unsloth?": ["thinking", "search"], "Create a live weather dashboard in HTML using no API key. Show me the code": ["thinking", "code", "search"], "Solve the integral of x·sin(x), and verify it step by step": ["thinking", "code"], - "Draw an SVG of a cute sloth": ["thinking", "code", "search"], + "Draw an SVG of a cute sloth & show the code": ["thinking", "code", "search"], }; const toolIconMap = { @@ -220,14 +226,20 @@ const GeneratingSpinner: FC = () => { const ComposerAnimated: FC = () => { return ( - - - +
+
+ + + +
); }; @@ -262,7 +274,7 @@ const Composer: FC = () => { { const ToolStatusDisplay: FC = () => { const toolStatus = useChatRuntimeStore((s) => s.toolStatus); + const isThreadRunning = useAuiState(({ thread }) => thread.isRunning); const [elapsed, setElapsed] = useState(0); + const [visible, setVisible] = useState(false); useEffect(() => { if (!toolStatus) { setElapsed(0); + if (!isThreadRunning) { + setVisible(false); + } return; } + setElapsed(0); + + // Debounce badge visibility by 300ms when the badge is not + // already on screen. Once visible from a prior tool, consecutive + // tools show immediately so the badge does not flicker. Fast + // tool calls that all complete under 300ms never show the badge. + let showTimer: ReturnType | undefined; + if (!visible) { + showTimer = setTimeout(() => setVisible(true), 300); + } + const interval = setInterval(() => { setElapsed((prev) => prev + 1); }, 1000); - return () => clearInterval(interval); - }, [toolStatus]); + return () => { + clearInterval(interval); + if (showTimer) clearTimeout(showTimer); + }; + }, [toolStatus, isThreadRunning]); - if (!toolStatus) return null; + if (!toolStatus || !visible) return null; const isRunning = toolStatus.startsWith("Running"); const StatusIcon = isRunning ? TerminalIcon : GlobeIcon; return ( @@ -555,10 +586,10 @@ const GeneratingIndicator: FC = () => { const AssistantMessage: FC = () => { return ( -
+
{
-
+
diff --git a/studio/frontend/src/components/assistant-ui/tool-fallback.tsx b/studio/frontend/src/components/assistant-ui/tool-fallback.tsx index c965749f74..82a5b17e04 100644 --- a/studio/frontend/src/components/assistant-ui/tool-fallback.tsx +++ b/studio/frontend/src/components/assistant-ui/tool-fallback.tsx @@ -156,12 +156,12 @@ function ToolFallbackTrigger({ - {label}: {toolName} + {label}: {toolName} {isRunning && ( - {label}: {toolName} + {label}: {toolName} )} diff --git a/studio/frontend/src/components/assistant-ui/tool-ui-python.tsx b/studio/frontend/src/components/assistant-ui/tool-ui-python.tsx index a510ed0d9e..6aa590ae11 100644 --- a/studio/frontend/src/components/assistant-ui/tool-ui-python.tsx +++ b/studio/frontend/src/components/assistant-ui/tool-ui-python.tsx @@ -4,6 +4,7 @@ "use client"; import { copyToClipboard } from "@/lib/copy-to-clipboard"; +import { getAuthToken } from "@/features/auth/session"; import type { ToolCallMessagePartComponent } from "@assistant-ui/react"; import { code as codePlugin } from "@streamdown/code"; import { CheckIcon, CodeIcon, CopyIcon, LoaderIcon } from "lucide-react"; @@ -15,6 +16,12 @@ import { ToolFallbackTrigger, } from "./tool-fallback"; +interface StructuredResult { + text: string; + images: string[]; + sessionId: string; +} + const MAX_DISPLAY = 10_000; const COPY_RESET_MS = 2000; const SHIKI_THEME = ["github-light", "github-dark"] as ["github-light", "github-dark"]; @@ -84,6 +91,16 @@ function HighlightedCode({ code: source, language }: { code: string; language: s ); } +function isStructuredResult(val: unknown): val is StructuredResult { + return ( + typeof val === "object" && + val !== null && + "text" in val && + "images" in val && + "sessionId" in val + ); +} + const PythonToolUIImpl: ToolCallMessagePartComponent = ({ args, result, @@ -92,12 +109,24 @@ const PythonToolUIImpl: ToolCallMessagePartComponent = ({ const code = (args as { code?: string })?.code ?? ""; const firstLine = code.split("\n")[0]?.slice(0, 60) ?? ""; const isRunning = status?.type === "running"; - const output = - typeof result === "string" - ? result - : result - ? JSON.stringify(result, null, 2) - : ""; + + let output: string; + let images: string[] = []; + let sessionId = ""; + + if (isStructuredResult(result)) { + output = result.text; + images = result.images; + sessionId = result.sessionId; + } else if (typeof result === "string") { + output = result; + } else if (result) { + output = JSON.stringify(result, null, 2); + } else { + output = ""; + } + + const authToken = getAuthToken(); return ( @@ -133,6 +162,21 @@ const PythonToolUIImpl: ToolCallMessagePartComponent = ({
) : null} + + {/* Images from Python tool execution */} + {images.length > 0 && sessionId && ( +
+ {images.map((filename) => ( + {filename} + ))} +
+ )}
diff --git a/studio/frontend/src/components/navbar.tsx b/studio/frontend/src/components/navbar.tsx index 63189a0d5d..121c559db8 100644 --- a/studio/frontend/src/components/navbar.tsx +++ b/studio/frontend/src/components/navbar.tsx @@ -372,90 +372,94 @@ export function Navbar() { })} - {/* Right: docs/tour desktop */} -
- - - - - - Learn more - - - - -

- Unsloth Documentation -

-

- Guides on fine-tuning LLMs 2x faster with 70% less memory. - Covers LoRA, QLoRA, data formatting, and deployment. -

- - Visit docs - - -
-
-
- - - - - + {/* Right: docs/tour desktop — one wrapper per control so flex gap is even (HoverCard roots can confuse flex spacing). */} +
+
+ +
+ + {tourId ? ( +
- - - - - - - +
+ ) : null} +
+ + + + + + + + +
+
+ +
{/* Right: mobile */} diff --git a/studio/frontend/src/config/training.ts b/studio/frontend/src/config/training.ts index 42fccde552..913d612838 100644 --- a/studio/frontend/src/config/training.ts +++ b/studio/frontend/src/config/training.ts @@ -141,6 +141,10 @@ export const MODEL_TYPE_TO_HF_TASK: Record = { export const PRIORITY_TRAINING_MODELS: readonly string[] = [ + "unsloth/gemma-4-E2B-it", + "unsloth/gemma-4-E4B-it", + "unsloth/gemma-4-31B-it", + "unsloth/gemma-4-26B-A4B-it", "unsloth/Qwen3.5-2B", "unsloth/Qwen3.5-9B", "unsloth/gpt-oss-20b", diff --git a/studio/frontend/src/features/chat/api/chat-adapter.ts b/studio/frontend/src/features/chat/api/chat-adapter.ts index 2b8a259930..e287daf33a 100644 --- a/studio/frontend/src/features/chat/api/chat-adapter.ts +++ b/studio/frontend/src/features/chat/api/chat-adapter.ts @@ -635,7 +635,23 @@ export function createOpenAIStreamAdapter(): ChatModelAdapter { toolCallParts[toolCallParts.length - 1]?.toolCallId || ""; const idx = toolCallParts.findIndex((p) => p.toolCallId === id); if (idx !== -1) { - toolCallParts[idx] = { ...toolCallParts[idx], result: toolEvent.result as string }; + const rawResult = (toolEvent.result as string) ?? ""; + const imgMarker = "\n__IMAGES__:"; + const imgIdx = rawResult.lastIndexOf(imgMarker); + let parsedResult: string | { text: string; images: string[]; sessionId: string }; + if (imgIdx !== -1) { + const text = rawResult.slice(0, imgIdx); + const sessionId = unstable_threadId || ""; + try { + const images = JSON.parse(rawResult.slice(imgIdx + imgMarker.length)) as string[]; + parsedResult = { text, images, sessionId }; + } catch { + parsedResult = rawResult; + } + } else { + parsedResult = rawResult; + } + toolCallParts[idx] = { ...toolCallParts[idx], result: parsedResult }; } } // Yield cumulative state so tool UI updates (tools first, text after) diff --git a/studio/frontend/src/features/chat/chat-page.tsx b/studio/frontend/src/features/chat/chat-page.tsx index a47a2c6d92..08450c7ec7 100644 --- a/studio/frontend/src/features/chat/chat-page.tsx +++ b/studio/frontend/src/features/chat/chat-page.tsx @@ -162,10 +162,12 @@ const CompareContent = memo(function CompareContent({ pairId, models, loraModels, + onFoldersChange, }: { pairId: string; models: ModelOption[]; loraModels: LoraModelOption[]; + onFoldersChange?: () => void; }): ReactElement { const isLoraCompare = useIsLoraCompare(); @@ -176,6 +178,7 @@ const CompareContent = memo(function CompareContent({ pairId={pairId} models={models} loraModels={loraModels} + onFoldersChange={onFoldersChange} /> ); }); @@ -259,10 +262,12 @@ const GeneralCompareContent = memo(function GeneralCompareContent({ pairId, models, loraModels, + onFoldersChange, }: { pairId: string; models: ModelOption[]; loraModels: LoraModelOption[]; + onFoldersChange?: () => void; }): ReactElement { const handlesRef = useRef>({}); const [model1ThreadId, setModel1ThreadId] = useState(); @@ -327,6 +332,7 @@ const GeneralCompareContent = memo(function GeneralCompareContent({ ggufVariant: meta.ggufVariant, }) } + onFoldersChange={onFoldersChange} variant="ghost" size="sm" className="max-w-[50%]" @@ -359,6 +365,7 @@ const GeneralCompareContent = memo(function GeneralCompareContent({ ggufVariant: meta.ggufVariant, }) } + onFoldersChange={onFoldersChange} variant="ghost" size="sm" className="max-w-[50%]" @@ -846,6 +853,7 @@ export function ChatPage(): ReactElement { activeGgufVariant={activeGgufVariant} onValueChange={handleCheckpointChange} onEject={handleEject} + onFoldersChange={refreshLocalModels} variant="ghost" open={modelSelectorOpen} onOpenChange={handleModelSelectorOpenChange} @@ -858,14 +866,16 @@ export function ChatPage(): ReactElement { label={ loadProgress?.phase === "starting" ? "Starting model…" - : loadingModel.isDownloaded + : loadingModel.isDownloaded || loadingModel.isCachedLora ? "Loading model…" : "Downloading model…" } title={ loadingModel.isDownloaded ? `Loading ${loadingModel.displayName} from cache.` - : `Loading ${loadingModel.displayName}. This may include downloading.` + : loadingModel.isCachedLora + ? `Loading ${loadingModel.displayName} into memory.` + : `Loading ${loadingModel.displayName}. This may include downloading.` } progressPercent={loadProgress?.percent} progressLabel={loadProgress?.label} @@ -911,6 +921,7 @@ export function ChatPage(): ReactElement { pairId={view.pairId} models={models} loraModels={loraModels} + onFoldersChange={refreshLocalModels} /> )}
@@ -934,7 +945,6 @@ export function ChatPage(): ReactElement { }); } }} - onFoldersChange={refreshLocalModels} />
diff --git a/studio/frontend/src/features/chat/chat-settings-sheet.tsx b/studio/frontend/src/features/chat/chat-settings-sheet.tsx index 6e62c7f9c5..e5b0814343 100644 --- a/studio/frontend/src/features/chat/chat-settings-sheet.tsx +++ b/studio/frontend/src/features/chat/chat-settings-sheet.tsx @@ -34,7 +34,6 @@ import { CodeIcon, Delete02Icon, FloppyDiskIcon, - FolderSearchIcon, PencilEdit01Icon, Settings02Icon, SlidersHorizontalIcon, @@ -44,13 +43,7 @@ import { import { HugeiconsIcon } from "@hugeicons/react"; import { AnimatePresence, motion } from "motion/react"; import type { ReactNode } from "react"; -import { useCallback, useEffect, useMemo, useState } from "react"; -import { - type ScanFolderInfo, - addScanFolder, - listScanFolders, - removeScanFolder, -} from "./api/chat-api"; +import { useEffect, useMemo, useState } from "react"; import { useChatRuntimeStore } from "./stores/chat-runtime-store"; import { DEFAULT_INFERENCE_PARAMS, @@ -266,108 +259,6 @@ function CollapsibleSection({ ); } -function ModelFoldersSection({ - onFoldersChange, -}: { onFoldersChange?: () => void }) { - const [folders, setFolders] = useState([]); - const [input, setInput] = useState(""); - const [error, setError] = useState(null); - const [loading, setLoading] = useState(false); - - const refresh = useCallback(() => { - listScanFolders() - .then(setFolders) - .catch(() => {}); - }, []); - - useEffect(() => { - refresh(); - }, [refresh]); - - const handleAdd = async () => { - const trimmed = input.trim(); - if (!trimmed) return; - setError(null); - setLoading(true); - try { - await addScanFolder(trimmed); - setInput(""); - refresh(); - onFoldersChange?.(); - } catch (e) { - setError(e instanceof Error ? e.message : "Failed to add folder"); - } finally { - setLoading(false); - } - }; - - const handleRemove = async (id: number) => { - try { - await removeScanFolder(id); - onFoldersChange?.(); - } catch (e) { - setError(e instanceof Error ? e.message : "Failed to remove folder"); - } finally { - refresh(); - } - }; - - return ( - -
- {folders.length > 0 && ( -
- {folders.map((f) => ( -
- - {f.path} - - -
- ))} -
- )} -
- { - setInput(e.target.value); - setError(null); - }} - onKeyDown={(e) => { - if (e.key === "Enter") handleAdd(); - }} - placeholder="/path/to/models" - className="h-7 flex-1 text-xs font-mono" - disabled={loading} - /> - -
- {error &&

{error}

} -
-
- ); -} - interface ChatSettingsPanelProps { open: boolean; onOpenChange?: (open: boolean) => void; @@ -376,7 +267,6 @@ interface ChatSettingsPanelProps { autoTitle: boolean; onAutoTitleChange: (enabled: boolean) => void; onReloadModel?: () => void; - onFoldersChange?: () => void; } export function ChatSettingsPanel({ @@ -387,7 +277,6 @@ export function ChatSettingsPanel({ autoTitle, onAutoTitleChange, onReloadModel, - onFoldersChange, }: ChatSettingsPanelProps) { const isMobile = useIsMobile(); const isGguf = useChatRuntimeStore((s) => s.activeGgufVariant) != null; @@ -395,6 +284,9 @@ export function ChatSettingsPanel({ 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 loadedKvCacheDtype = useChatRuntimeStore((s) => s.loadedKvCacheDtype); @@ -404,7 +296,7 @@ export function ChatSettingsPanel({ ); const ctxDisplayValue = customContextLength ?? ggufContextLength ?? ""; - const ctxMaxValue = ggufMaxContextLength ?? ggufContextLength ?? null; + const ctxMaxValue = ggufNativeContextLength ?? ggufContextLength ?? null; const kvDirty = kvCacheDtype !== loadedKvCacheDtype; const ctxDirty = customContextLength !== null; const modelSettingsDirty = kvDirty || ctxDirty; @@ -655,6 +547,13 @@ export function ChatSettingsPanel({ ); }} /> + {ggufMaxContextLength != null && + typeof ctxDisplayValue === "number" && + ctxDisplayValue > ggufMaxContextLength && ( +

+ Exceeds estimated VRAM capacity ({ggufMaxContextLength.toLocaleString()} tokens). The model may use system RAM. +

+ )}
@@ -835,8 +734,6 @@ export function ChatSettingsPanel({
- -
(null); const [loadToastDismissed, setLoadToastDismissed] = useState(false); const [loadProgress, setLoadProgress] = useState<{ @@ -246,12 +247,16 @@ export function useChatModelRuntime() { const ggufMaxContextLength = statusRes.is_gguf ? (statusRes.max_context_length ?? null) : null; + const ggufNativeContextLength = statusRes.is_gguf + ? (statusRes.native_context_length ?? null) + : null; useChatRuntimeStore.setState({ supportsReasoning, reasoningAlwaysOn, supportsTools, ggufContextLength: currentGgufContextLength, ggufMaxContextLength, + ggufNativeContextLength, }); // Set reasoning default for Qwen3.5 small models @@ -290,8 +295,11 @@ export function useChatModelRuntime() { setLoadToastDismissedState(false); clearCheckpoint(); if (tid != null) toast.dismiss(tid); + const isCachedOrLocal = model.isDownloaded || model.isCachedLora; toast.info("Stopped loading model", { - description: "The current download may still finish in the background.", + description: isCachedOrLocal + ? undefined + : "The current download may still finish in the background.", }); // Fire-and-forget: tell backend to stop, don't block UI unloadModel({ model_path: model.id }).catch(() => {}); @@ -335,20 +343,24 @@ export function useChatModelRuntime() { : undefined; const previousIsLora = previousModel?.isLora ?? (previousLora ? true : false); + // Covers Unix absolute (/), relative (./ ../), tilde (~/), Windows drive (C:\), UNC (\\server) + const isLocal = /^(\/|\.{1,2}[\\\/]|~[\\\/]|[A-Za-z]:[\\\/]|\\\\)/.test(modelId); + const isCachedLora = isLora && isLocal; const loadingDescription = [ currentCheckpoint ? "Switching models." : null, extraLoadingDescription ?? null, isDownloaded ? "Loading cached model into memory." : null, + !isDownloaded && isCachedLora ? "Loading trained model into memory." : null, ] .filter(Boolean) .join(" "); setModelsError(null); setLoadToastDismissedState(false); - const loadInfo = { id: modelId, displayName, isDownloaded }; + const loadInfo = { id: modelId, displayName, isDownloaded, isCachedLora }; setLoadingModel(loadInfo); useChatRuntimeStore.getState().setModelLoading(true); setLoadProgress( - isDownloaded + isDownloaded || isCachedLora ? { percent: null, label: null, phase: "starting" } : { percent: 0, label: "Preparing download", phase: "downloading" }, ); @@ -425,6 +437,9 @@ export function useChatModelRuntime() { const reportedMaxCtx = loadResponse.is_gguf ? (loadResponse.max_context_length ?? null) : null; + const reportedNativeCtx = loadResponse.is_gguf + ? (loadResponse.native_context_length ?? null) + : null; // A successful reload has applied settings, so clear pending custom // context state and display the backend-reported effective context. const keepCustomCtx = null; @@ -433,6 +448,7 @@ export function useChatModelRuntime() { useChatRuntimeStore.setState({ ggufContextLength: nativeCtx, ggufMaxContextLength, + ggufNativeContextLength: reportedNativeCtx, supportsReasoning: loadResponse.supports_reasoning ?? false, reasoningAlwaysOn, reasoningEnabled: reasoningAlwaysOn ? true : reasoningDefault, @@ -477,15 +493,16 @@ export function useChatModelRuntime() { } } - const toastTitle = isDownloaded ? "Starting model…" : "Downloading model…"; + const isCachedLoad = isDownloaded || isCachedLora; + const toastTitle = isCachedLoad ? "Starting model…" : "Downloading model…"; const toastId = toast( null, { description: renderLoadDescription( toastTitle, loadingDescription, - isDownloaded ? null : 0, - isDownloaded ? null : "Preparing download", + isCachedLoad ? null : 0, + isCachedLoad ? null : "Preparing download", cancelLoading, ), duration: Infinity, @@ -503,7 +520,7 @@ export function useChatModelRuntime() { // Poll download progress for non-cached models (GGUF and non-GGUF) let progressInterval: ReturnType | null = null; - if (!isDownloaded) { + if (!isDownloaded && !isCachedLora) { const expectedBytes = typeof selection !== "string" ? selection.expectedBytes ?? 0 : 0; let hasShownProgress = false; diff --git a/studio/frontend/src/features/chat/runtime-provider.tsx b/studio/frontend/src/features/chat/runtime-provider.tsx index 404271f896..02f792b509 100644 --- a/studio/frontend/src/features/chat/runtime-provider.tsx +++ b/studio/frontend/src/features/chat/runtime-provider.tsx @@ -47,9 +47,9 @@ const DEFAULT_SUGGESTIONS = [ prompt: "Solve the integral of x·sin(x), and verify it step by step", }, { - title: "Draw an SVG of a cute sloth", + title: "Draw an SVG of a cute sloth & show the code", label: "SVG sloth", - prompt: "Draw an SVG of a cute sloth", + prompt: "Draw an SVG of a cute sloth & show the code", }, ]; @@ -569,10 +569,32 @@ function ThreadHistoryProvider({ store.setContextUsage(savedUsage); } + // If any message has a stored parentId, reconstruct the tree + // so retries/regenerations load as branches instead of being + // unrolled into a flat list. For mixed legacy/new threads + // (old messages without parentId + new messages with), infer + // sequential parents for old messages to preserve the chain. + // Fall back to fromArray for fully legacy threads. + const hasParentIds = msgs.some((m) => "parentId" in m); + if (hasParentIds) { + let previousId: string | null = null; + return { + messages: msgs.map((m) => { + const parentId = "parentId" in m + ? (m.parentId ?? null) + : previousId; + previousId = m.id; + return { + parentId, + message: toThreadMessage(m), + }; + }), + }; + } return ExportedMessageRepository.fromArray(msgs.map(toThreadMessage)); }, - async append({ message }: ExportedMessageRepositoryItem) { + async append({ parentId, message }: ExportedMessageRepositoryItem) { const { remoteId } = await aui.threadListItem().initialize(); const content = cloneContent(message.content); const attachments = @@ -586,6 +608,7 @@ function ThreadHistoryProvider({ await db.messages.put({ id: message.id, threadId: remoteId, + parentId: parentId ?? null, role: message.role, content, ...(attachments.length > 0 && { attachments }), diff --git a/studio/frontend/src/features/chat/shared-composer.tsx b/studio/frontend/src/features/chat/shared-composer.tsx index 59b0880add..ac01f77381 100644 --- a/studio/frontend/src/features/chat/shared-composer.tsx +++ b/studio/frontend/src/features/chat/shared-composer.tsx @@ -473,7 +473,7 @@ export function SharedComposer({ onChange={(e) => setText(e.target.value)} onKeyDown={onKeyDown} placeholder="Send to both models..." - className="mb-1 max-h-32 min-h-14 w-full resize-none bg-transparent px-4 pt-2 pb-3 text-sm outline-none placeholder:text-muted-foreground" + className="mb-1 max-h-32 min-h-14 w-full resize-none bg-transparent pl-5 pr-4 pt-2 pb-3 text-sm outline-none placeholder:text-muted-foreground" rows={1} />
diff --git a/studio/frontend/src/features/chat/stores/chat-runtime-store.ts b/studio/frontend/src/features/chat/stores/chat-runtime-store.ts index 8cea234f21..48abaf7580 100644 --- a/studio/frontend/src/features/chat/stores/chat-runtime-store.ts +++ b/studio/frontend/src/features/chat/stores/chat-runtime-store.ts @@ -151,6 +151,7 @@ type ChatRuntimeStore = { activeGgufVariant: string | null; ggufContextLength: number | null; ggufMaxContextLength: number | null; + ggufNativeContextLength: number | null; supportsReasoning: boolean; reasoningAlwaysOn: boolean; reasoningEnabled: boolean; @@ -215,6 +216,7 @@ export const useChatRuntimeStore = create((set) => ({ activeGgufVariant: null, ggufContextLength: null, ggufMaxContextLength: null, + ggufNativeContextLength: null, supportsReasoning: false, reasoningAlwaysOn: false, reasoningEnabled: true, @@ -290,6 +292,7 @@ export const useChatRuntimeStore = create((set) => ({ activeGgufVariant: null, ggufContextLength: null, ggufMaxContextLength: null, + ggufNativeContextLength: null, contextUsage: null, supportsReasoning: false, reasoningEnabled: true, diff --git a/studio/frontend/src/features/chat/types.ts b/studio/frontend/src/features/chat/types.ts index 11e53d76d3..1f370b6ac1 100644 --- a/studio/frontend/src/features/chat/types.ts +++ b/studio/frontend/src/features/chat/types.ts @@ -20,6 +20,7 @@ export interface ThreadRecord { export interface MessageRecord { id: string; threadId: string; + parentId?: string | null; role: import("@assistant-ui/react").ThreadMessage["role"]; content: import("@assistant-ui/react").ThreadMessage["content"]; attachments?: import("@assistant-ui/react").ThreadMessage["attachments"]; diff --git a/studio/frontend/src/features/chat/types/api.ts b/studio/frontend/src/features/chat/types/api.ts index dcc0a980c8..8f0839615f 100644 --- a/studio/frontend/src/features/chat/types/api.ts +++ b/studio/frontend/src/features/chat/types/api.ts @@ -87,6 +87,7 @@ export interface LoadModelResponse { }; context_length?: number | null; max_context_length?: number | null; + native_context_length?: number | null; supports_reasoning?: boolean; reasoning_always_on?: boolean; supports_tools?: boolean; @@ -121,6 +122,7 @@ export interface InferenceStatusResponse { supports_tools?: boolean; context_length?: number | null; max_context_length?: number | null; + native_context_length?: number | null; } export interface AudioGenerationResponse { diff --git a/studio/frontend/src/features/onboarding/components/steps/model-selection-step.tsx b/studio/frontend/src/features/onboarding/components/steps/model-selection-step.tsx index f2a4796c54..1ff23cb0fc 100644 --- a/studio/frontend/src/features/onboarding/components/steps/model-selection-step.tsx +++ b/studio/frontend/src/features/onboarding/components/steps/model-selection-step.tsx @@ -268,12 +268,12 @@ export function ModelSelectionStep() { {id} @@ -287,12 +287,12 @@ export function ModelSelectionStep() { {fitStatus === "exceeds" && ( - + OOM )} {fitStatus === "tight" && ( - + TIGHT )} diff --git a/studio/frontend/src/features/studio/sections/model-section.tsx b/studio/frontend/src/features/studio/sections/model-section.tsx index 755c93c5f0..775073eb64 100644 --- a/studio/frontend/src/features/studio/sections/model-section.tsx +++ b/studio/frontend/src/features/studio/sections/model-section.tsx @@ -489,12 +489,12 @@ export function ModelSection() { {id} @@ -519,12 +519,12 @@ export function ModelSection() { {fitStatus === "exceeds" && ( - + OOM )} {fitStatus === "tight" && ( - + TIGHT )} diff --git a/studio/frontend/src/hooks/use-hf-model-search.ts b/studio/frontend/src/hooks/use-hf-model-search.ts index 32b261956b..e697544d5c 100644 --- a/studio/frontend/src/hooks/use-hf-model-search.ts +++ b/studio/frontend/src/hooks/use-hf-model-search.ts @@ -104,6 +104,11 @@ function makeMapModel(excludeGguf: boolean) { /** Number of unsloth results to pull up-front before yielding general results. */ const UNSLOTH_PREFETCH = 20; +/** When the user searched for a specific publisher, show fewer unsloth results + * before the pinned (original publisher) model. */ +const UNSLOTH_PINNED_PREFETCH = 4; +/** Matches a valid "owner/repo" identifier (exactly two non-empty segments). */ +const PUBLISHER_RE = /^([^/\s]+)\/([^/\s]+)$/; /** * Prime the hf-cache from a listModels result. For public (non-gated, @@ -131,6 +136,7 @@ async function* mergedModelIterator( query: string, task?: PipelineType, accessToken?: string, + pinnedId?: string, ): AsyncGenerator { const common = { additionalFields: ["safetensors", "tags"] as ("safetensors" | "tags")[], @@ -148,6 +154,18 @@ async function* mergedModelIterator( ...common, }); + // Start pinned model lookup immediately so it can run in parallel with + // the Phase 1 unsloth iteration instead of blocking Phase 2. + const pinnedPromise = pinnedId + ? cachedModelInfo({ + name: pinnedId, + additionalFields: ["safetensors", "tags"], + ...(accessToken ? { credentials: { accessToken } } : {}), + }).catch(() => null) + : null; + + const limit = pinnedId ? UNSLOTH_PINNED_PREFETCH : UNSLOTH_PREFETCH; + // Phase 1: pull & yield unsloth models first const seen = new Set(); let count = 0; @@ -159,10 +177,26 @@ async function* mergedModelIterator( } yield model; count++; - if (count >= UNSLOTH_PREFETCH) break; + if (count >= limit) break; } - // Phase 2: yield general results, skipping already-seen unsloth models + // Phase 1b: yield the pinned (original publisher) model before general results + if (pinnedId && !seen.has(pinnedId) && pinnedPromise) { + const pinned = await pinnedPromise; + if (pinned) { + // Record both the raw input and the canonical name returned by HF + // so phase 2 deduplication works even when casing differs + // (e.g. user typed "OpenAI/gpt-oss-20b", HF returns "openai/gpt-oss-20b"). + seen.add(pinnedId); + const canonicalName = (pinned as { name?: string }).name; + if (canonicalName && canonicalName !== pinnedId) { + seen.add(canonicalName); + } + yield pinned; + } + } + + // Phase 2: yield general results, skipping already-seen models for await (const model of generalIter) { const m = model as { name?: string }; if (m.name && seen.has(m.name)) continue; @@ -235,11 +269,24 @@ export function useHfModelSearch( ) { const { task, accessToken, excludeGguf = false, priorityIds } = options ?? {}; + // Parse publisher detection once and share between the iterator factory + // and the secondary sort gate (avoids duplicating the regex + logic). + const { isPublisherQuery, searchQuery, pinnedId, trimmed } = useMemo(() => { + const t = query.trim(); + const m = PUBLISHER_RE.exec(t); + const is = !!m && m[1].toLowerCase() !== "unsloth"; + return { + isPublisherQuery: is, + searchQuery: is ? m![2] : t, + pinnedId: is ? t : undefined, + trimmed: t, + }; + }, [query]); + const createIter = useCallback( () => { - const trimmed = query.trim(); if (!trimmed) { - // No query → show priority models first (with full metadata), then general unsloth listing + // No query: show priority models first (with full metadata), then general unsloth listing if (priorityIds && priorityIds.length > 0) { return priorityThenListingIterator(priorityIds, task, accessToken) as AsyncGenerator; } @@ -250,24 +297,35 @@ export function useHfModelSearch( ...(accessToken ? { credentials: { accessToken } } : {}), }) as AsyncGenerator; } - // Typed query: disable task filter so explicitly searched models still appear even if HF task metadata is wrong/missing. - return mergedModelIterator(trimmed, undefined, accessToken) as AsyncGenerator; + // Typed query: disable task filter so explicitly searched models still + // appear even if HF task metadata is wrong/missing. + // If the query is a valid "owner/repo" identifier (exactly two non-empty, + // slash-free, space-free segments), strip the org prefix so unsloth + // variants surface, then pin the original publisher model after a small + // batch of unsloth results. Queries for unsloth-owned models are left + // as-is so they get the full 20-result prefetch and secondary sort. + return mergedModelIterator(searchQuery, undefined, accessToken, pinnedId) as AsyncGenerator; }, - [query, task, accessToken, priorityIds], + [trimmed, searchQuery, pinnedId, task, accessToken, priorityIds], ); const mapModel = useMemo(() => makeMapModel(excludeGguf), [excludeGguf]); const search = useHfPaginatedSearch(createIter, mapModel); - // Secondary sort guarantee: unsloth models always float to the top + // Secondary sort guarantee: unsloth models always float to the top. + // Skip when the user searched for a specific non-unsloth publisher + // (e.g. "openai/gpt-oss-20b") -- the iterator already handles the + // pinned ordering in that case. const results = useMemo( () => - [...search.results].sort((a, b) => { - const aFirst = a.id.startsWith("unsloth/") ? 0 : 1; - const bFirst = b.id.startsWith("unsloth/") ? 0 : 1; - return aFirst - bFirst; - }), - [search.results], + isPublisherQuery + ? search.results + : [...search.results].sort((a, b) => { + const aFirst = a.id.startsWith("unsloth/") ? 0 : 1; + const bFirst = b.id.startsWith("unsloth/") ? 0 : 1; + return aFirst - bFirst; + }), + [search.results, isPublisherQuery], ); return { ...search, results }; diff --git a/studio/frontend/src/index.css b/studio/frontend/src/index.css index 8dc159a002..a30b4ca287 100644 --- a/studio/frontend/src/index.css +++ b/studio/frontend/src/index.css @@ -12,6 +12,8 @@ @plugin "@toolwind/corner-shape"; @source "../node_modules/streamdown/dist/*.js"; +@custom-variant dark (&:is(.dark *)); + @font-face { font-family: "Hellix"; src: url("/fonts/Hellix-SemiBold.woff2") format("woff2"), @@ -21,8 +23,6 @@ font-display: swap; } -@custom-variant dark (&:is(.dark *)); - :root { /* Animation timing */ --duration-micro: 100ms; @@ -67,7 +67,7 @@ --sidebar-ring: oklch(0.6929 0.1396 166.5513); --destructive-foreground: oklch(1 0 0); --font-sans: "Inter Variable", ui-sans-serif, sans-serif, system-ui; - --font-heading: "Hellix", "Space Grotesk Variable", ui-sans-serif, sans-serif; + --font-heading: "Hellix", "Space Grotesk Variable", var(--font-sans); --font-serif: Source Serif 4, serif; --font-mono: JetBrains Mono, monospace; --shadow-color: hsl(0 0% 0%); @@ -159,7 +159,7 @@ @theme inline { --font-sans: "Inter Variable", ui-sans-serif, sans-serif, system-ui; - --font-heading: "Hellix", "Space Grotesk Variable", ui-sans-serif, sans-serif; + --font-heading: "Hellix", "Space Grotesk Variable", var(--font-sans); --color-sidebar-ring: var(--sidebar-ring); --color-sidebar-border: var(--sidebar-border); --color-sidebar-accent-foreground: var(--sidebar-accent-foreground); @@ -343,10 +343,58 @@ [data-streamdown="code-block"] { gap: 0; padding: 0.5rem; + /* Wide lines must scroll inside the thread column, not widen past the composer (flex min-width:auto). */ + max-width: 100%; + min-width: 0; + overflow-x: auto; } [data-streamdown="code-block-header"] { padding-left: 0.75rem; } + + /* Chat thread: code slightly smaller by default; step up when the thread column is wide. */ + .aui-thread-root [data-streamdown="code-block"] { + font-size: 0.8125rem; + line-height: 1.55; + } + + .aui-thread-root [data-streamdown="code-block-header"] { + font-size: 0.6875rem; + } + + @container (min-width: 36rem) { + .aui-thread-root [data-streamdown="code-block"] { + font-size: 0.875rem; + } + + .aui-thread-root [data-streamdown="code-block-header"] { + font-size: 0.75rem; + } + } + + /* Chat: use the app sans stack for UI + prose. */ + .aui-thread-root { + --font-heading: var(--font-sans); + font-family: var(--font-sans); + } + + /* Keep monospace for code fences and inline code (not KaTeX). */ + .aui-thread-root [data-streamdown="code-block"] pre, + .aui-thread-root [data-streamdown="code-block"] code { + font-family: var(--font-mono), ui-monospace, monospace; + } + + .aui-thread-root :where(p, li, td, th, blockquote, h1, h2, h3, h4, h5, h6) code { + font-family: var(--font-mono), ui-monospace, monospace; + } + + /* Align fenced code blocks with the main chat column even when nested in lists. */ + .aui-thread-root [data-streamdown="list-item"] > [data-streamdown="code-block"], + .aui-thread-root [data-streamdown="list-item"] [data-streamdown="code-block"] { + margin-left: -1.25rem; + width: calc(100% + 1.25rem); + max-width: calc(100% + 1.25rem); + } } /* Minimal scrollbar — thumb only, no track */ diff --git a/studio/install_llama_prebuilt.py b/studio/install_llama_prebuilt.py index 516dc4b6a4..1b02729649 100755 --- a/studio/install_llama_prebuilt.py +++ b/studio/install_llama_prebuilt.py @@ -7,12 +7,14 @@ from __future__ import annotations import argparse +import errno import fnmatch import hashlib import json import os import platform import random +import re import shutil import site import socket @@ -27,7 +29,7 @@ import urllib.parse import urllib.request import zipfile from contextlib import contextmanager -from dataclasses import dataclass +from dataclasses import dataclass, field try: from filelock import FileLock, Timeout as FileLockTimeout @@ -41,12 +43,31 @@ from typing import Any, Iterable, Iterator EXIT_SUCCESS = 0 EXIT_FALLBACK = 2 EXIT_ERROR = 1 +EXIT_BUSY = 3 -APPROVED_PREBUILT_LLAMA_TAG = "b8508" -DEFAULT_LLAMA_TAG = os.environ.get("UNSLOTH_LLAMA_TAG", APPROVED_PREBUILT_LLAMA_TAG) -DEFAULT_PUBLISHED_REPO = os.environ.get( - "UNSLOTH_LLAMA_RELEASE_REPO", "unslothai/llama.cpp" -) + +def env_int(name: str, default: int, *, minimum: int | None = None) -> int: + raw = os.environ.get(name) + if raw is None: + value = default + else: + try: + value = int(str(raw).strip()) + except (TypeError, ValueError): + value = default + if minimum is not None: + value = max(minimum, value) + return value + + +# Prefer "latest" over "master" -- "master" bypasses the prebuilt resolver +# (no matching GitHub release), forces a source build, and causes HTTP 422 +# errors. Only use "master" temporarily when the latest release is missing +# support for a new model architecture. +DEFAULT_LLAMA_TAG = os.environ.get("UNSLOTH_LLAMA_TAG", "latest") +# Force all installs to use mainline llama.cpp from ggml-org. +# Previously: DEFAULT_PUBLISHED_REPO = os.environ.get("UNSLOTH_LLAMA_RELEASE_REPO", "unslothai/llama.cpp") +DEFAULT_PUBLISHED_REPO = "ggml-org/llama.cpp" DEFAULT_PUBLISHED_TAG = os.environ.get("UNSLOTH_LLAMA_RELEASE_TAG") DEFAULT_PUBLISHED_MANIFEST_ASSET = os.environ.get( "UNSLOTH_LLAMA_RELEASE_MANIFEST_ASSET", "llama-prebuilt-manifest.json" @@ -71,6 +92,11 @@ HTTP_FETCH_BASE_DELAY_SECONDS = 0.75 SERVER_PORT_BIND_ATTEMPTS = 3 SERVER_BIND_RETRY_WINDOW_SECONDS = 5.0 TTY_PROGRESS_START_DELAY_SECONDS = 0.5 +DEFAULT_MAX_PREBUILT_RELEASE_FALLBACKS = env_int( + "UNSLOTH_LLAMA_MAX_PREBUILT_RELEASE_FALLBACKS", + 2, + minimum = 1, +) @dataclass @@ -129,10 +155,18 @@ class PublishedReleaseBundle: repo: str release_tag: str upstream_tag: str - assets: dict[str, str] - manifest_asset_name: str - artifacts: list[PublishedLlamaArtifact] - selection_log: list[str] + manifest_sha256: str | None = None + source_repo: str | None = None + source_repo_url: str | None = None + source_ref_kind: str | None = None + requested_source_ref: str | None = None + resolved_source_ref: str | None = None + source_commit: str | None = None + source_commit_short: str | None = None + assets: dict[str, str] = field(default_factory = dict) + manifest_asset_name: str = DEFAULT_PUBLISHED_MANIFEST_ASSET + artifacts: list[PublishedLlamaArtifact] = field(default_factory = list) + selection_log: list[str] = field(default_factory = list) @dataclass @@ -166,16 +200,108 @@ class ApprovedReleaseChecksums: repo: str release_tag: str upstream_tag: str - source_commit: str | None - artifacts: dict[str, ApprovedArtifactHash] + source_repo: str | None = None + source_repo_url: str | None = None + source_ref_kind: str | None = None + requested_source_ref: str | None = None + resolved_source_ref: str | None = None + source_commit: str | None = None + source_commit_short: str | None = None + artifacts: dict[str, ApprovedArtifactHash] = field(default_factory = dict) + + +@dataclass(frozen = True) +class ResolvedPublishedRelease: + bundle: PublishedReleaseBundle + checksums: ApprovedReleaseChecksums + + +@dataclass(frozen = True) +class SourceBuildPlan: + source_url: str + source_ref: str + source_ref_kind: str + compatibility_upstream_tag: str + source_repo: str | None = None + source_repo_url: str | None = None + requested_source_ref: str | None = None + resolved_source_ref: str | None = None + source_commit: str | None = None + + +@dataclass(frozen = True) +class InstallReleasePlan: + requested_tag: str + llama_tag: str + release_tag: str + attempts: list[AssetChoice] + approved_checksums: ApprovedReleaseChecksums class PrebuiltFallback(RuntimeError): pass +class BusyInstallConflict(RuntimeError): + pass + + +class ExistingInstallSatisfied(RuntimeError): + def __init__(self, choice: AssetChoice, used_fallback: bool): + super().__init__(f"existing install already matches candidate {choice.name}") + self.choice = choice + self.used_fallback = used_fallback + + +def _os_error_messages(exc: BaseException) -> list[str]: + messages: list[str] = [] + if isinstance(exc, OSError): + for value in ( + getattr(exc, "strerror", None), + getattr(exc, "filename", None), + getattr(exc, "filename2", None), + ): + if isinstance(value, str) and value: + messages.append(value) + text = str(exc) + if text: + messages.append(text) + return [message.lower() for message in messages if message] + + +def is_busy_lock_error(exc: BaseException) -> bool: + if isinstance(exc, BusyInstallConflict): + return True + if isinstance(exc, OSError): + if exc.errno in { + errno.EACCES, + errno.EBUSY, + errno.EPERM, + errno.ETXTBSY, + }: + return True + if getattr(exc, "winerror", None) in {5, 32, 145}: + return True + for message in _os_error_messages(exc): + if any( + needle in message + for needle in ( + "access is denied", + "being used by another process", + "device or resource busy", + "permission denied", + "text file busy", + "file is in use", + "process cannot access the file", + "cannot create a file when that file already exists", + ) + ): + return True + return False + + def log(message: str) -> None: - print(f"[llama-prebuilt] {message}") + print(f"[llama-prebuilt] {message}", file = sys.stderr) def log_lines(lines: Iterable[str]) -> None: @@ -263,6 +389,10 @@ def source_archive_logical_name(upstream_tag: str) -> str: return f"llama.cpp-source-{upstream_tag}.tar.gz" +def exact_source_archive_logical_name(source_commit: str) -> str: + return f"llama.cpp-source-commit-{source_commit}.tar.gz" + + def sha256_file(path: Path) -> str: digest = hashlib.sha256() with path.open("rb") as handle: @@ -271,6 +401,10 @@ def sha256_file(path: Path) -> str: return digest.hexdigest() +def sha256_bytes(data: bytes) -> str: + return hashlib.sha256(data).hexdigest() + + def normalize_sha256_digest(value: str | None) -> str | None: if not isinstance(value, str) or not value: return None @@ -282,6 +416,183 @@ def normalize_sha256_digest(value: str | None) -> str | None: return lowered +def normalize_source_ref_kind(value: str | None) -> str | None: + if not isinstance(value, str): + return None + normalized = value.strip().lower() + if normalized in {"tag", "branch", "pull", "commit", "custom"}: + return normalized + return None + + +def normalize_source_commit(value: str | None) -> str | None: + if not isinstance(value, str): + return None + normalized = value.strip().lower() + if len(normalized) < 7 or len(normalized) > 40: + return None + if any(ch not in "0123456789abcdef" for ch in normalized): + return None + return normalized + + +def validate_schema_version(payload: dict[str, Any], *, label: str) -> None: + schema_version = payload.get("schema_version") + if schema_version is None: + return + try: + normalized = int(schema_version) + except (TypeError, ValueError) as exc: + raise RuntimeError(f"{label} schema_version was not an integer") from exc + if normalized != 1: + raise RuntimeError(f"{label} schema_version={normalized} is unsupported") + + +def repo_slug_from_source(value: str | None) -> str | None: + if not isinstance(value, str): + return None + normalized = value.strip() + if not normalized: + return None + normalized = normalized.removesuffix(".git") + if normalized.startswith("https://github.com/"): + slug = normalized[len("https://github.com/") :] + elif normalized.startswith("http://github.com/"): + slug = normalized[len("http://github.com/") :] + elif normalized.startswith("git@github.com:"): + slug = normalized[len("git@github.com:") :] + else: + slug = normalized + slug = slug.strip("/") + parts = slug.split("/") + if len(parts) != 2 or not all(parts): + return None + return f"{parts[0]}/{parts[1]}" + + +def source_url_from_repo_slug(repo_slug: str | None) -> str | None: + if not isinstance(repo_slug, str) or not repo_slug: + return None + return f"https://github.com/{repo_slug}" + + +def source_repo_clone_url(repo: str | None, repo_url: str | None) -> str | None: + if isinstance(repo_url, str) and repo_url.strip(): + return repo_url.strip().removesuffix(".git") + return source_url_from_repo_slug(repo_slug_from_source(repo)) + + +def infer_source_ref_kind(ref: str | None) -> str: + if not isinstance(ref, str): + return "tag" + normalized = ref.strip() + lowered = normalized.lower() + if not normalized: + return "tag" + if lowered.startswith("refs/pull/") or lowered.startswith("pull/"): + return "pull" + if ( + lowered.startswith("refs/heads/") + or lowered in {"main", "master", "head"} + or lowered.startswith("origin/") + ): + return "branch" + normalized_commit = normalize_source_commit(normalized) + if normalized_commit is not None: + return "commit" + return "tag" + + +def normalized_ref_aliases(ref: str | None) -> set[str]: + if not isinstance(ref, str): + return set() + normalized = ref.strip() + if not normalized: + return set() + aliases = {normalized} + lowered = normalized.lower() + commit = normalize_source_commit(normalized) + if commit is not None: + aliases.add(commit) + if lowered.startswith("refs/heads/"): + aliases.add(normalized.split("/", 2)[2]) + elif "/" not in normalized and infer_source_ref_kind(normalized) == "branch": + aliases.add(f"refs/heads/{normalized}") + if lowered.startswith("refs/pull/"): + aliases.add(normalized.removeprefix("refs/")) + elif lowered.startswith("pull/"): + aliases.add(f"refs/{normalized}") + return aliases + + +def refs_match(candidate_ref: str | None, requested_ref: str | None) -> bool: + candidate_aliases = normalized_ref_aliases(candidate_ref) + requested_aliases = normalized_ref_aliases(requested_ref) + if not candidate_aliases or not requested_aliases: + return False + if candidate_aliases & requested_aliases: + return True + candidate_commit = normalize_source_commit(candidate_ref) + requested_commit = normalize_source_commit(requested_ref) + if candidate_commit and requested_commit: + return candidate_commit.startswith( + requested_commit + ) or requested_commit.startswith(candidate_commit) + return False + + +def checkout_friendly_ref(ref_kind: str | None, ref: str | None) -> str | None: + """Normalize a source ref to a form that ``git clone --branch`` accepts. + + Fully qualified branch refs like ``refs/heads/main`` are stripped to + ``main``; tag refs like ``refs/tags/b8508`` are stripped to ``b8508``. + Pull refs like ``refs/pull/123/head`` are left as-is since they are + always fetched explicitly rather than cloned with ``--branch``. + """ + if not isinstance(ref, str) or not ref: + return ref + lowered = ref.lower() + if ref_kind == "branch" and lowered.startswith("refs/heads/"): + return ref.split("/", 2)[2] + if ref_kind == "tag" and lowered.startswith("refs/tags/"): + return ref.split("/", 2)[2] + return ref + + +def windows_cuda_upstream_asset_names(llama_tag: str, runtime: str) -> list[str]: + return [ + f"llama-{llama_tag}-bin-win-cuda-{runtime}-x64.zip", + f"cudart-llama-bin-win-cuda-{runtime}-x64.zip", + ] + + +def windows_cuda_asset_aliases( + asset_name: str, + *, + compatibility_tag: str | None = None, +) -> list[str]: + aliases: list[str] = [] + legacy_match = re.fullmatch( + r"llama-(?P[^/]+)-bin-win-cuda-(?P\d+\.\d+)-x64\.zip", + asset_name, + ) + if legacy_match: + runtime = legacy_match.group("runtime") + aliases.append(f"cudart-llama-bin-win-cuda-{runtime}-x64.zip") + if compatibility_tag: + aliases.append(f"llama-{compatibility_tag}-bin-win-cuda-{runtime}-x64.zip") + return aliases + + current_match = re.fullmatch( + r"cudart-llama-bin-win-cuda-(?P\d+\.\d+)-x64\.zip", + asset_name, + ) + if current_match and compatibility_tag: + runtime = current_match.group("runtime") + aliases.append(f"llama-{compatibility_tag}-bin-win-cuda-{runtime}-x64.zip") + return aliases + + def format_byte_count(num_bytes: float) -> str: units = ["B", "KiB", "MiB", "GiB", "TiB"] value = float(num_bytes) @@ -442,13 +753,21 @@ def download_bytes( def fetch_json(url: str) -> Any: - data = download_bytes( - url, - timeout = 30, - headers = github_api_headers(url) - if is_github_api_url(url) - else auth_headers(url), - ) + try: + data = download_bytes( + url, + timeout = 30, + headers = github_api_headers(url) + if is_github_api_url(url) + else auth_headers(url), + ) + except urllib.error.HTTPError as exc: + if exc.code == 403 and is_github_api_url(url): + hint = "" + if not (os.environ.get("GH_TOKEN") or os.environ.get("GITHUB_TOKEN")): + hint = "; set GH_TOKEN or GITHUB_TOKEN to avoid GitHub API rate limits" + raise RuntimeError(f"GitHub API returned 403 for {url}{hint}") from exc + raise if not data: raise RuntimeError(f"downloaded empty JSON payload from {url}") try: @@ -553,6 +872,14 @@ def upstream_source_archive_urls(tag: str) -> list[str]: ] +def commit_source_archive_urls(repo: str, source_commit: str) -> list[str]: + encoded_commit = urllib.parse.quote(source_commit, safe = "") + return [ + f"https://codeload.github.com/{repo}/tar.gz/{encoded_commit}", + f"https://github.com/{repo}/archive/{encoded_commit}.tar.gz", + ] + + def github_release_assets(repo: str, tag: str) -> dict[str, str]: payload = fetch_json( f"https://api.github.com/repos/{repo}/releases/tags/{urllib.parse.quote(tag, safe = '')}" @@ -598,6 +925,14 @@ def latest_upstream_release_tag() -> str: return tag +def normalized_requested_llama_tag(requested_tag: str | None) -> str: + if isinstance(requested_tag, str): + normalized = requested_tag.strip() + if normalized: + return normalized + return "latest" + + def normalize_compute_cap(value: Any) -> str | None: raw = str(value).strip() if not raw: @@ -881,13 +1216,35 @@ def parse_published_release_bundle( # Mixed repos are filtered by an explicit release-side manifest rather than # by release tag or asset filename conventions. - manifest_payload = fetch_json(manifest_url) + manifest_bytes = download_bytes( + manifest_url, + timeout = 30, + headers = auth_headers(manifest_url), + ) + manifest_sha256 = sha256_bytes(manifest_bytes) + try: + manifest_payload = json.loads(manifest_bytes.decode("utf-8")) + except (UnicodeDecodeError, json.JSONDecodeError) as exc: + raise RuntimeError( + f"published manifest {DEFAULT_PUBLISHED_MANIFEST_ASSET} was not valid JSON" + ) from exc if not isinstance(manifest_payload, dict): raise RuntimeError( f"published manifest {DEFAULT_PUBLISHED_MANIFEST_ASSET} was not a JSON object" ) + validate_schema_version( + manifest_payload, + label = f"published manifest {DEFAULT_PUBLISHED_MANIFEST_ASSET} in {repo}@{release_tag}", + ) component = manifest_payload.get("component") upstream_tag = manifest_payload.get("upstream_tag") + source_repo = manifest_payload.get("source_repo") + source_repo_url = manifest_payload.get("source_repo_url") + source_ref_kind = normalize_source_ref_kind(manifest_payload.get("source_ref_kind")) + requested_source_ref = manifest_payload.get("requested_source_ref") + resolved_source_ref = manifest_payload.get("resolved_source_ref") + source_commit = normalize_source_commit(manifest_payload.get("source_commit")) + source_commit_short = manifest_payload.get("source_commit_short") if component != "llama.cpp": return None if not isinstance(upstream_tag, str) or not upstream_tag: @@ -918,10 +1275,32 @@ def parse_published_release_bundle( f"published_release: manifest={DEFAULT_PUBLISHED_MANIFEST_ASSET}", f"published_release: upstream_tag={upstream_tag}", ] + if isinstance(source_repo, str) and source_repo: + selection_log.append(f"published_release: source_repo={source_repo}") + if source_commit: + selection_log.append(f"published_release: source_commit={source_commit}") return PublishedReleaseBundle( repo = repo, release_tag = release_tag, upstream_tag = upstream_tag, + manifest_sha256 = manifest_sha256, + source_repo = source_repo + if isinstance(source_repo, str) and source_repo + else None, + source_repo_url = source_repo_url + if isinstance(source_repo_url, str) and source_repo_url + else None, + source_ref_kind = source_ref_kind, + requested_source_ref = requested_source_ref + if isinstance(requested_source_ref, str) and requested_source_ref + else None, + resolved_source_ref = resolved_source_ref + if isinstance(resolved_source_ref, str) and resolved_source_ref + else None, + source_commit = source_commit, + source_commit_short = source_commit_short + if isinstance(source_commit_short, str) and source_commit_short + else None, assets = assets, manifest_asset_name = DEFAULT_PUBLISHED_MANIFEST_ASSET, artifacts = artifacts, @@ -938,6 +1317,10 @@ def parse_approved_release_checksums( raise RuntimeError( f"published checksum asset {DEFAULT_PUBLISHED_SHA256_ASSET} was not a JSON object" ) + validate_schema_version( + payload, + label = f"published checksum asset {DEFAULT_PUBLISHED_SHA256_ASSET}", + ) if payload.get("component") != "llama.cpp": raise RuntimeError( f"published checksum asset {DEFAULT_PUBLISHED_SHA256_ASSET} did not describe llama.cpp" @@ -987,13 +1370,33 @@ def parse_approved_release_checksums( kind = kind_value if isinstance(kind_value, str) and kind_value else None, ) - source_commit = payload.get("source_commit") + source_commit = normalize_source_commit(payload.get("source_commit")) + source_commit_short = payload.get("source_commit_short") + source_repo = payload.get("source_repo") + source_repo_url = payload.get("source_repo_url") + source_ref_kind = normalize_source_ref_kind(payload.get("source_ref_kind")) + requested_source_ref = payload.get("requested_source_ref") + resolved_source_ref = payload.get("resolved_source_ref") return ApprovedReleaseChecksums( repo = repo, release_tag = release_tag, upstream_tag = upstream_tag, - source_commit = source_commit - if isinstance(source_commit, str) and source_commit + source_repo = source_repo + if isinstance(source_repo, str) and source_repo + else None, + source_repo_url = source_repo_url + if isinstance(source_repo_url, str) and source_repo_url + else None, + source_ref_kind = source_ref_kind, + requested_source_ref = requested_source_ref + if isinstance(requested_source_ref, str) and requested_source_ref + else None, + resolved_source_ref = resolved_source_ref + if isinstance(resolved_source_ref, str) and resolved_source_ref + else None, + source_commit = source_commit, + source_commit_short = source_commit_short + if isinstance(source_commit_short, str) and source_commit_short else None, artifacts = artifacts, ) @@ -1283,17 +1686,163 @@ def pinned_published_release_bundle( return bundle +def validated_checksums_for_bundle( + repo: str, bundle: PublishedReleaseBundle +) -> ApprovedReleaseChecksums: + checksums = load_approved_release_checksums(repo, bundle.release_tag) + manifest_hash = checksums.artifacts.get(bundle.manifest_asset_name) + if manifest_hash is not None and bundle.manifest_sha256 is not None: + if manifest_hash.sha256 != bundle.manifest_sha256: + raise PrebuiltFallback( + "published manifest checksum did not match the approved checksum asset" + ) + # Accept bundles that carry only an exact-commit source archive + # (e.g. llama.cpp-source-commit-.tar.gz) without requiring the + # legacy llama.cpp-source-.tar.gz entry. + if exact_source_archive_hash(checksums) is None: + require_approved_source_hash(checksums, bundle.upstream_tag) + return checksums + + +def published_release_matches_request( + bundle: PublishedReleaseBundle, requested_ref: str +) -> bool: + if requested_ref == "latest": + return True + for candidate in ( + bundle.upstream_tag, + bundle.requested_source_ref, + bundle.resolved_source_ref, + bundle.source_commit, + ): + if refs_match(candidate, requested_ref): + return True + return False + + +def resolve_published_release( + requested_tag: str | None, + published_repo: str, + published_release_tag: str = "", +) -> ResolvedPublishedRelease: + repo = published_repo or DEFAULT_PUBLISHED_REPO + normalized_requested = normalized_requested_llama_tag(requested_tag) + + if published_release_tag: + bundle = pinned_published_release_bundle(repo, published_release_tag) + if not published_release_matches_request(bundle, normalized_requested): + raise PrebuiltFallback( + "published release " + f"{repo}@{published_release_tag} targeted upstream tag {bundle.upstream_tag}, " + f"but requested {normalized_requested}" + ) + return ResolvedPublishedRelease( + bundle = bundle, + checksums = validated_checksums_for_bundle(repo, bundle), + ) + + skipped_invalid = 0 + for bundle in iter_published_release_bundles(repo): + if not published_release_matches_request(bundle, normalized_requested): + continue + try: + checksums = validated_checksums_for_bundle(repo, bundle) + except PrebuiltFallback as exc: + skipped_invalid += 1 + log( + "published release ignored for install resolution: " + f"{repo}@{bundle.release_tag} ({exc})" + ) + continue + return ResolvedPublishedRelease(bundle = bundle, checksums = checksums) + + if normalized_requested == "latest": + if skipped_invalid: + raise PrebuiltFallback( + f"no usable published llama.cpp releases were available in {repo}" + ) + raise PrebuiltFallback( + f"no published llama.cpp releases were available in {repo}" + ) + + raise PrebuiltFallback( + f"no published prebuilt release in {repo} matched upstream tag {normalized_requested}" + ) + + +def iter_resolved_published_releases( + requested_tag: str | None, + published_repo: str, + published_release_tag: str = "", +) -> Iterable[ResolvedPublishedRelease]: + repo = published_repo or DEFAULT_PUBLISHED_REPO + normalized_requested = normalized_requested_llama_tag(requested_tag) + + if published_release_tag: + bundle = pinned_published_release_bundle(repo, published_release_tag) + if not published_release_matches_request(bundle, normalized_requested): + raise PrebuiltFallback( + "published release " + f"{repo}@{published_release_tag} targeted upstream tag {bundle.upstream_tag}, " + f"but requested {normalized_requested}" + ) + yield ResolvedPublishedRelease( + bundle = bundle, + checksums = validated_checksums_for_bundle(repo, bundle), + ) + return + + matched_any = False + skipped_invalid = 0 + yielded_valid = False + for bundle in iter_published_release_bundles(repo): + if not published_release_matches_request(bundle, normalized_requested): + continue + matched_any = True + try: + checksums = validated_checksums_for_bundle(repo, bundle) + except PrebuiltFallback as exc: + skipped_invalid += 1 + log( + "published release ignored for install resolution: " + f"{repo}@{bundle.release_tag} ({exc})" + ) + continue + yielded_valid = True + yield ResolvedPublishedRelease(bundle = bundle, checksums = checksums) + + if yielded_valid: + return + + if matched_any: + if skipped_invalid: + raise PrebuiltFallback( + f"no usable published llama.cpp releases were available in {repo}" + ) + return + + if normalized_requested == "latest": + raise PrebuiltFallback( + f"no published llama.cpp releases were available in {repo}" + ) + + raise PrebuiltFallback( + f"no published prebuilt release in {repo} matched upstream tag {normalized_requested}" + ) + + def resolve_requested_llama_tag( requested_tag: str | None, published_repo: str = "", + published_release_tag: str = "", ) -> str: """Resolve a llama.cpp tag for source-build fallback. Resolution order: 1. Concrete tag (e.g. "b8508") -- returned as-is. - 2. "latest" with published_repo -- query the Unsloth release repo - (e.g. unslothai/llama.cpp) for its latest release tag. This is the - tested/approved version that matches the prebuilt binaries. + 2. "latest" with published_repo -- resolve the latest usable Unsloth + published release bundle and return its upstream_tag. This is the + preferred version that matches the published prebuilt metadata. 3. "latest" without published_repo or if (2) fails -- query the upstream ggml-org/llama.cpp repo. This may return a newer, untested tag. @@ -1301,20 +1850,20 @@ def resolve_requested_llama_tag( upstream tags that have been validated with Unsloth Studio. Using the upstream bleeding-edge tag risks API/ABI incompatibilities. """ - if requested_tag and requested_tag != "latest": - return requested_tag + normalized_requested = normalized_requested_llama_tag(requested_tag) + if normalized_requested != "latest": + return normalized_requested # Prefer the Unsloth release repo tag (tested/approved) over bleeding-edge # upstream. For example, unslothai/llama.cpp may publish b8508 while # ggml-org/llama.cpp latest is b8514. The source-build fallback should # compile the same version the prebuilt path would have installed. if published_repo: try: - payload = fetch_json( - f"https://api.github.com/repos/{published_repo}/releases/latest" - ) - tag = payload.get("tag_name") - if isinstance(tag, str) and tag: - return tag + return resolve_published_release( + "latest", + published_repo, + published_release_tag, + ).bundle.upstream_tag except Exception: pass # Fall back to upstream ggml-org latest release tag @@ -1324,18 +1873,132 @@ def resolve_requested_llama_tag( def resolve_requested_install_tag( requested_tag: str | None, published_release_tag: str = "", + published_repo: str = DEFAULT_PUBLISHED_REPO, ) -> str: - approved_tag = APPROVED_PREBUILT_LLAMA_TAG - normalized_requested = requested_tag or "latest" - if normalized_requested not in {"latest", approved_tag}: - raise PrebuiltFallback( - f"prebuilt installs are pinned to approved release {approved_tag}; requested {normalized_requested}" + return resolve_published_release( + requested_tag, + published_repo, + published_release_tag, + ).bundle.upstream_tag + + +def exact_source_archive_hash( + checksums: ApprovedReleaseChecksums, +) -> ApprovedArtifactHash | None: + if not checksums.source_commit: + return None + return checksums.artifacts.get( + exact_source_archive_logical_name(checksums.source_commit) + ) + + +def source_clone_url_from_checksums(checksums: ApprovedReleaseChecksums) -> str | None: + return source_repo_clone_url(checksums.source_repo, checksums.source_repo_url) + + +def source_build_plan_for_release( + release: ResolvedPublishedRelease, +) -> SourceBuildPlan: + checksums = release.checksums + exact_source = exact_source_archive_hash(checksums) + source_repo = checksums.source_repo or release.bundle.source_repo + source_repo_url = checksums.source_repo_url or release.bundle.source_repo_url + requested_source_ref = ( + checksums.requested_source_ref or release.bundle.requested_source_ref + ) + resolved_source_ref = ( + checksums.resolved_source_ref or release.bundle.resolved_source_ref + ) + source_commit = checksums.source_commit or release.bundle.source_commit + source_ref_kind = checksums.source_ref_kind or release.bundle.source_ref_kind + source_url = source_repo_clone_url(source_repo, source_repo_url) + if exact_source is not None and source_url and source_commit: + return SourceBuildPlan( + source_url = source_url, + source_ref = source_commit, + source_ref_kind = "commit", + compatibility_upstream_tag = release.bundle.upstream_tag, + source_repo = source_repo, + source_repo_url = source_repo_url, + requested_source_ref = requested_source_ref, + resolved_source_ref = resolved_source_ref, + source_commit = source_commit, ) - if published_release_tag and published_release_tag != approved_tag: - raise PrebuiltFallback( - f"prebuilt installs require published release tag {approved_tag}; requested {published_release_tag}" + source_ref = checkout_friendly_ref( + source_ref_kind, resolved_source_ref or requested_source_ref + ) + if ( + source_url + and source_ref + and source_ref_kind in {"tag", "branch", "pull", "commit"} + ): + return SourceBuildPlan( + source_url = source_url, + source_ref = source_ref, + source_ref_kind = source_ref_kind, + compatibility_upstream_tag = release.bundle.upstream_tag, + source_repo = source_repo, + source_repo_url = source_repo_url, + requested_source_ref = requested_source_ref, + resolved_source_ref = resolved_source_ref, + source_commit = source_commit, ) - return approved_tag + return SourceBuildPlan( + source_url = source_url_from_repo_slug(UPSTREAM_REPO) + or "https://github.com/ggml-org/llama.cpp", + source_ref = release.bundle.upstream_tag, + source_ref_kind = "tag", + compatibility_upstream_tag = release.bundle.upstream_tag, + source_repo = source_repo, + source_repo_url = source_repo_url, + requested_source_ref = requested_source_ref, + resolved_source_ref = resolved_source_ref, + source_commit = source_commit, + ) + + +def resolve_source_build_plan( + requested_tag: str | None, + published_repo: str, + published_release_tag: str = "", +) -> SourceBuildPlan: + normalized_requested = normalized_requested_llama_tag(requested_tag) + if normalized_requested != "latest": + try: + release = resolve_published_release( + normalized_requested, + published_repo, + published_release_tag, + ) + return source_build_plan_for_release(release) + except Exception: + pass + inferred_kind = infer_source_ref_kind(normalized_requested) + return SourceBuildPlan( + source_url = "https://github.com/ggml-org/llama.cpp", + source_ref = checkout_friendly_ref(inferred_kind, normalized_requested) + or normalized_requested, + source_ref_kind = inferred_kind, + compatibility_upstream_tag = normalized_requested, + ) + + if published_repo: + try: + release = resolve_published_release( + "latest", + published_repo, + published_release_tag, + ) + return source_build_plan_for_release(release) + except Exception: + pass + latest_tag = latest_upstream_release_tag() + return SourceBuildPlan( + source_url = "https://github.com/ggml-org/llama.cpp", + source_ref = latest_tag, + source_ref_kind = "tag", + compatibility_upstream_tag = latest_tag, + ) def run_capture( @@ -1655,31 +2318,99 @@ def windows_cuda_attempts( attempts: list[AssetChoice] = [] for runtime_line in runtime_order: runtime = runtime_by_line[runtime_line] - upstream_name = f"llama-{llama_tag}-bin-win-cuda-{runtime}-x64.zip" - asset_url = upstream_assets.get(upstream_name) - if not asset_url: + selected_name = None + asset_url = None + for candidate_name in windows_cuda_upstream_asset_names(llama_tag, runtime): + asset_url = upstream_assets.get(candidate_name) + if asset_url: + selected_name = candidate_name + break + if not asset_url or not selected_name: selection_log.append( - f"windows_cuda_selection: skip missing asset {upstream_name}" + "windows_cuda_selection: skip missing assets " + + ",".join(windows_cuda_upstream_asset_names(llama_tag, runtime)) ) continue attempts.append( AssetChoice( repo = UPSTREAM_REPO, tag = llama_tag, - name = upstream_name, + name = selected_name, url = asset_url, source_label = "upstream", install_kind = "windows-cuda", runtime_line = runtime_line, selection_log = list(selection_log) + [ - f"windows_cuda_selection: selected {upstream_name} runtime={runtime}" + f"windows_cuda_selection: selected {selected_name} runtime={runtime}" ], ) ) return attempts +def published_windows_cuda_attempts( + host: HostInfo, + release: PublishedReleaseBundle, + preferred_runtime_line: str | None, + selection_preamble: Iterable[str] = (), +) -> list[AssetChoice]: + selection_log = list(release.selection_log) + list(selection_preamble) + runtime_by_line = {"cuda12": "12.4", "cuda13": "13.1"} + runtime_order = windows_cuda_attempts( + host, + release.upstream_tag, + { + f"llama-{release.upstream_tag}-bin-win-cuda-{runtime}-x64.zip": "published" + for runtime in runtime_by_line.values() + }, + preferred_runtime_line, + selection_log, + ) + published_artifacts = [ + artifact + for artifact in release.artifacts + if artifact.install_kind == "windows-cuda" + ] + artifacts_by_runtime: dict[str, list[PublishedLlamaArtifact]] = {} + for artifact in published_artifacts: + if not artifact.runtime_line: + continue + artifacts_by_runtime.setdefault(artifact.runtime_line, []).append(artifact) + + attempts: list[AssetChoice] = [] + for ordered_attempt in runtime_order: + runtime_line = ordered_attempt.runtime_line + if not runtime_line: + continue + candidates = sorted( + artifacts_by_runtime.get(runtime_line, []), + key = lambda artifact: (artifact.rank, artifact.asset_name), + ) + for artifact in candidates: + asset_url = release.assets.get(artifact.asset_name) + if not asset_url: + continue + attempts.append( + AssetChoice( + repo = release.repo, + tag = release.release_tag, + name = artifact.asset_name, + url = asset_url, + source_label = "published", + install_kind = "windows-cuda", + runtime_line = runtime_line, + selection_log = list(ordered_attempt.selection_log or []) + + [ + "windows_cuda_selection: selected published asset " + f"{artifact.asset_name} for runtime_line={runtime_line}" + ], + ) + ) + break + return attempts + + def resolve_windows_cuda_choices( host: HostInfo, llama_tag: str, upstream_assets: dict[str, str] ) -> list[AssetChoice]: @@ -1695,32 +2426,52 @@ def resolve_windows_cuda_choices( def resolve_linux_cuda_choice( - host: HostInfo, llama_tag: str, published_repo: str, published_release_tag: str + host: HostInfo, release: PublishedReleaseBundle ) -> LinuxCudaSelection: torch_preference = detect_torch_cuda_runtime_preference(host) - skipped_tag_mismatches = 0 - for release in iter_published_release_bundles( - published_repo, published_release_tag - ): - if release.upstream_tag != llama_tag: - skipped_tag_mismatches += 1 - continue - selection = linux_cuda_choice_from_release( - host, - release, - preferred_runtime_line = torch_preference.runtime_line, - selection_preamble = torch_preference.selection_log, - ) - if selection is not None: - return selection - if skipped_tag_mismatches: - log( - "published Linux CUDA selection skipped " - f"{skipped_tag_mismatches} release(s) with upstream_tag != {llama_tag}" - ) + selection = linux_cuda_choice_from_release( + host, + release, + preferred_runtime_line = torch_preference.runtime_line, + selection_preamble = torch_preference.selection_log, + ) + if selection is not None: + return selection raise PrebuiltFallback("no compatible published Linux CUDA bundle was found") +def published_asset_choice_for_kind( + release: PublishedReleaseBundle, + install_kind: str, +) -> AssetChoice | None: + candidates = sorted( + ( + artifact + for artifact in release.artifacts + if artifact.install_kind == install_kind + ), + key = lambda artifact: (artifact.rank, artifact.asset_name), + ) + for artifact in candidates: + asset_url = release.assets.get(artifact.asset_name) + if not asset_url: + continue + return AssetChoice( + repo = release.repo, + tag = release.release_tag, + name = artifact.asset_name, + url = asset_url, + source_label = "published", + install_kind = install_kind, + runtime_line = artifact.runtime_line, + selection_log = list(release.selection_log) + + [ + f"published_selection: selected {artifact.asset_name} install_kind={install_kind}" + ], + ) + return None + + def resolve_upstream_asset_choice(host: HostInfo, llama_tag: str) -> AssetChoice: upstream_assets = github_release_assets(UPSTREAM_REPO, llama_tag) if host.is_linux and host.is_x86_64: @@ -1786,16 +2537,62 @@ def resolve_upstream_asset_choice(host: HostInfo, llama_tag: str) -> AssetChoice ) -def resolve_asset_choice( - host: HostInfo, llama_tag: str, published_repo: str, published_release_tag: str -) -> AssetChoice: +def resolve_asset_choice(host: HostInfo, llama_tag: str) -> AssetChoice: if host.is_linux and host.is_x86_64 and host.has_usable_nvidia: - return resolve_linux_cuda_choice( - host, llama_tag, published_repo, published_release_tag - ).primary + raise PrebuiltFallback( + "Linux CUDA installs require a compatible published bundle; upstream fallback is not available" + ) return resolve_upstream_asset_choice(host, llama_tag) +def resolve_release_asset_choice( + host: HostInfo, + llama_tag: str, + release: PublishedReleaseBundle, + checksums: ApprovedReleaseChecksums, +) -> list[AssetChoice]: + if host.is_windows and host.is_x86_64 and host.has_usable_nvidia: + torch_preference = detect_torch_cuda_runtime_preference(host) + published_attempts = published_windows_cuda_attempts( + host, + release, + torch_preference.runtime_line, + torch_preference.selection_log, + ) + if published_attempts: + try: + return apply_approved_hashes(published_attempts, checksums) + except PrebuiltFallback as exc: + log( + "published Windows CUDA assets ignored for install planning: " + f"{release.repo}@{release.release_tag} ({exc})" + ) + upstream_assets = github_release_assets(UPSTREAM_REPO, llama_tag) + return apply_approved_hashes( + resolve_windows_cuda_choices(host, llama_tag, upstream_assets), + checksums, + ) + + published_choice: AssetChoice | None = None + if host.is_windows and host.is_x86_64: + published_choice = published_asset_choice_for_kind(release, "windows-cpu") + elif host.is_macos and host.is_arm64: + published_choice = published_asset_choice_for_kind(release, "macos-arm64") + elif host.is_macos and host.is_x86_64: + published_choice = published_asset_choice_for_kind(release, "macos-x64") + + if published_choice is not None: + try: + return apply_approved_hashes([published_choice], checksums) + except PrebuiltFallback as exc: + log( + "published platform asset ignored for install planning: " + f"{release.repo}@{release.release_tag} {published_choice.name} ({exc})" + ) + + return apply_approved_hashes([resolve_asset_choice(host, llama_tag)], checksums) + + def extract_archive(archive_path: Path, destination: Path) -> None: def safe_extract_path(base: Path, member_name: str) -> Path: normalized = member_name.replace("\\", "/") @@ -1997,18 +2794,26 @@ def copy_directory_contents(source_dir: Path, destination: Path) -> None: def hydrate_source_tree( - upstream_tag: str, + source_ref: str, install_dir: Path, work_dir: Path, *, + source_repo: str = UPSTREAM_REPO, expected_sha256: str, + source_label: str | None = None, + exact_source: bool = False, ) -> None: - archive_path = work_dir / f"llama.cpp-source-{upstream_tag}.tar.gz" - source_urls = upstream_source_archive_urls(upstream_tag) + archive_path = work_dir / f"llama.cpp-source-{source_ref}.tar.gz" + source_urls = ( + commit_source_archive_urls(source_repo, source_ref) + if exact_source + else upstream_source_archive_urls(source_ref) + ) + label = source_label or f"llama.cpp source tree for {source_ref}" extract_dir = Path(tempfile.mkdtemp(prefix = "source-extract-", dir = work_dir)) try: - log(f"downloading llama.cpp source tree for upstream tag {upstream_tag}") + log(f"downloading {label}") last_exc: Exception | None = None downloaded = False for index, source_url in enumerate(source_urls): @@ -2021,7 +2826,7 @@ def hydrate_source_tree( source_url, archive_path, expected_sha256 = expected_sha256, - label = f"llama.cpp source tree for {upstream_tag}", + label = label, ) downloaded = True break @@ -2054,9 +2859,7 @@ def hydrate_source_tree( except PrebuiltFallback: raise except Exception as exc: - raise PrebuiltFallback( - f"failed to hydrate upstream llama.cpp source tree for {upstream_tag}: {exc}" - ) from exc + raise PrebuiltFallback(f"failed to hydrate {label}: {exc}") from exc finally: remove_tree(extract_dir) @@ -2163,8 +2966,14 @@ def install_lock(lock_path: Path) -> Iterator[None]: while True: try: fd = os.open(str(lock_path), os.O_CREAT | os.O_EXCL | os.O_RDWR) - os.write(fd, f"{os.getpid()}\n".encode()) - os.fsync(fd) + try: + os.write(fd, f"{os.getpid()}\n".encode()) + os.fsync(fd) + except Exception: + os.close(fd) + fd = None + lock_path.unlink(missing_ok = True) + raise break except FileExistsError: # Check if the holder process is still alive @@ -2177,6 +2986,10 @@ def install_lock(lock_path: Path) -> Iterator[None]: if not raw: # File exists but PID not yet written -- another process # just created it. Wait briefly for the write to land. + if time.monotonic() >= deadline: + raise BusyInstallConflict( + f"timed out after {INSTALL_LOCK_TIMEOUT_SECONDS}s waiting for concurrent install lock: {lock_path}" + ) time.sleep(0.1) continue try: @@ -2195,7 +3008,7 @@ def install_lock(lock_path: Path) -> Iterator[None]: lock_path.unlink(missing_ok = True) continue if time.monotonic() >= deadline: - raise RuntimeError( + raise BusyInstallConflict( f"timed out after {INSTALL_LOCK_TIMEOUT_SECONDS}s waiting for concurrent install lock: {lock_path}" ) time.sleep(0.5) @@ -2211,7 +3024,7 @@ def install_lock(lock_path: Path) -> Iterator[None]: with FileLock(lock_path, timeout = INSTALL_LOCK_TIMEOUT_SECONDS): yield except FileLockTimeout as exc: - raise RuntimeError( + raise BusyInstallConflict( f"timed out after {INSTALL_LOCK_TIMEOUT_SECONDS}s waiting for concurrent install lock: {lock_path}" ) from exc @@ -2359,11 +3172,17 @@ def activate_install_tree(staging_dir: Path, install_dir: Path, host: HostInfo) log(f"restoring rollback path {rollback_dir} -> {install_dir}") os.replace(rollback_dir, install_dir) log(f"restored previous install from rollback path {rollback_dir.name}") + if is_busy_lock_error(exc): + raise BusyInstallConflict( + "staged prebuilt validation passed but the existing install could not be replaced " + "because llama.cpp appears to still be in use; restored previous install " + f"({textwrap.shorten(str(exc), width = 200, placeholder = '...')})" + ) from exc raise PrebuiltFallback( "staged prebuilt validation passed but activation failed; restored previous install " f"({textwrap.shorten(str(exc), width = 200, placeholder = '...')})" ) from exc - except PrebuiltFallback: + except (BusyInstallConflict, PrebuiltFallback): raise except Exception as rollback_exc: log(f"rollback after failed activation also failed: {rollback_exc}") @@ -2395,7 +3214,12 @@ def activate_install_tree(staging_dir: Path, install_dir: Path, host: HostInfo) ) from exc else: if rollback_dir: - remove_tree_logged(rollback_dir, "rollback path") + try: + remove_tree_logged(rollback_dir, "rollback path") + except Exception as cleanup_exc: + log( + f"non-fatal: rollback cleanup failed after successful activation: {cleanup_exc}" + ) finally: remove_tree(failed_dir) remove_tree(staging_dir) @@ -3074,10 +3898,42 @@ def apply_approved_hashes( attempts: Iterable[AssetChoice], checksums: ApprovedReleaseChecksums, ) -> list[AssetChoice]: + def approved_hash_for_attempt(attempt: AssetChoice) -> ApprovedArtifactHash | None: + candidate_names = [attempt.name] + if ( + isinstance(attempt.tag, str) + and attempt.tag + and attempt.tag != checksums.upstream_tag + and attempt.name.startswith("llama-") + ): + legacy_prefix = f"llama-{attempt.tag}-" + compatibility_prefix = f"llama-{checksums.upstream_tag}-" + compatibility_name = ( + attempt.name.replace(legacy_prefix, compatibility_prefix, 1) + if attempt.name.startswith(legacy_prefix) + else attempt.name + ) + candidate_names.append(compatibility_name) + candidate_names.extend( + windows_cuda_asset_aliases( + attempt.name, + compatibility_tag = checksums.upstream_tag, + ) + ) + seen_names: set[str] = set() + for candidate_name in candidate_names: + if candidate_name in seen_names: + continue + seen_names.add(candidate_name) + approved = checksums.artifacts.get(candidate_name) + if approved is not None: + return approved + return None + approved_attempts: list[AssetChoice] = [] missing_assets: list[str] = [] for attempt in attempts: - approved = checksums.artifacts.get(attempt.name) + approved = approved_hash_for_attempt(attempt) if approved is None: missing_assets.append(attempt.name) continue @@ -3104,45 +3960,129 @@ def require_approved_source_hash( return approved_source +def preferred_source_archive( + checksums: ApprovedReleaseChecksums, llama_tag: str +) -> tuple[str, str, ApprovedArtifactHash, bool]: + exact_source = exact_source_archive_hash(checksums) + exact_repo = repo_slug_from_source(checksums.source_repo) or repo_slug_from_source( + checksums.source_repo_url + ) + if exact_source is not None and exact_repo and checksums.source_commit: + return ( + exact_repo, + checksums.source_commit, + exact_source, + True, + ) + legacy = require_approved_source_hash(checksums, llama_tag) + return ( + UPSTREAM_REPO, + llama_tag, + legacy, + False, + ) + + +def selected_source_archive_metadata( + checksums: ApprovedReleaseChecksums, + llama_tag: str, +) -> tuple[str, str | None]: + _source_repo, _source_ref, source_archive, _exact_source = preferred_source_archive( + checksums, llama_tag + ) + return source_archive.asset_name, source_archive.sha256 + + def resolve_install_attempts( llama_tag: str, host: HostInfo, published_repo: str, published_release_tag: str, ) -> tuple[str, str, list[AssetChoice], ApprovedReleaseChecksums]: - requested_tag = llama_tag - resolved_tag = resolve_requested_install_tag(llama_tag, published_release_tag) - checksums = load_approved_release_checksums(published_repo, resolved_tag) - require_approved_source_hash(checksums, resolved_tag) - - if host.is_linux and host.is_x86_64 and host.has_usable_nvidia: - linux_cuda_selection = resolve_linux_cuda_choice( - host, resolved_tag, published_repo, published_release_tag - ) - attempts = apply_approved_hashes(linux_cuda_selection.attempts, checksums) - if not attempts: - raise PrebuiltFallback("no compatible Linux CUDA asset was found") - log_lines(linux_cuda_selection.selection_log) - return requested_tag, resolved_tag, attempts, checksums - - if host.is_windows and host.is_x86_64 and host.has_usable_nvidia: - upstream_assets = github_release_assets(UPSTREAM_REPO, resolved_tag) - attempts = apply_approved_hashes( - resolve_windows_cuda_choices(host, resolved_tag, upstream_assets), checksums - ) - if not attempts: - raise PrebuiltFallback("no compatible Windows CUDA asset was found") - if attempts[0].selection_log: - log_lines(attempts[0].selection_log) - return requested_tag, resolved_tag, attempts, checksums - - choice = resolve_asset_choice( - host, resolved_tag, published_repo, published_release_tag + requested_tag, plans = resolve_install_release_plans( + llama_tag, + host, + published_repo, + published_release_tag, ) - approved_attempts = apply_approved_hashes([choice], checksums) - if choice.selection_log: - log_lines(choice.selection_log) - return requested_tag, resolved_tag, approved_attempts, checksums + if not plans: + raise PrebuiltFallback("no prebuilt release plans were available") + plan = plans[0] + return requested_tag, plan.llama_tag, plan.attempts, plan.approved_checksums + + +def resolve_install_release_plans( + llama_tag: str, + host: HostInfo, + published_repo: str, + published_release_tag: str, + *, + max_release_fallbacks: int = DEFAULT_MAX_PREBUILT_RELEASE_FALLBACKS, +) -> tuple[str, list[InstallReleasePlan]]: + requested_tag = normalized_requested_llama_tag(llama_tag) + allow_older_release_fallback = ( + requested_tag == "latest" and not published_release_tag + ) + release_limit = max(1, max_release_fallbacks) + plans: list[InstallReleasePlan] = [] + last_error: PrebuiltFallback | None = None + + for resolved_release in iter_resolved_published_releases( + llama_tag, + published_repo, + published_release_tag, + ): + bundle = resolved_release.bundle + checksums = resolved_release.checksums + resolved_tag = bundle.upstream_tag + try: + if host.is_linux and host.is_x86_64 and host.has_usable_nvidia: + linux_cuda_selection = resolve_linux_cuda_choice(host, bundle) + attempts = apply_approved_hashes( + linux_cuda_selection.attempts, checksums + ) + if not attempts: + raise PrebuiltFallback("no compatible Linux CUDA asset was found") + log_lines(linux_cuda_selection.selection_log) + else: + attempts = resolve_release_asset_choice( + host, + resolved_tag, + bundle, + checksums, + ) + if not attempts: + raise PrebuiltFallback("no compatible prebuilt asset was found") + if attempts[0].selection_log: + log_lines(attempts[0].selection_log) + except PrebuiltFallback as exc: + last_error = exc + if not allow_older_release_fallback: + raise + log( + "published release skipped for install planning: " + f"{bundle.repo}@{bundle.release_tag} upstream_tag={resolved_tag} ({exc})" + ) + continue + + plans.append( + InstallReleasePlan( + requested_tag = requested_tag, + llama_tag = resolved_tag, + release_tag = bundle.release_tag, + attempts = attempts, + approved_checksums = checksums, + ) + ) + + if not allow_older_release_fallback or len(plans) >= release_limit: + break + + if plans: + return requested_tag, plans + if last_error is not None: + raise last_error + raise PrebuiltFallback("no installable published llama.cpp releases were found") def write_prebuilt_metadata( @@ -3150,17 +4090,54 @@ def write_prebuilt_metadata( *, requested_tag: str, llama_tag: str, + release_tag: str, choice: AssetChoice, + approved_checksums: ApprovedReleaseChecksums, prebuilt_fallback_used: bool, ) -> None: + source_asset_name, source_sha256 = selected_source_archive_metadata( + approved_checksums, + llama_tag, + ) + fingerprint_payload = { + "published_repo": approved_checksums.repo, + "release_tag": release_tag, + "upstream_tag": llama_tag, + "asset": choice.name, + "asset_sha256": choice.expected_sha256, + "source": choice.source_label, + "source_asset": source_asset_name, + "source_sha256": source_sha256, + "runtime_line": choice.runtime_line, + "bundle_profile": choice.bundle_profile, + "coverage_class": choice.coverage_class, + } + fingerprint = hashlib.sha256( + json.dumps(fingerprint_payload, sort_keys = True, separators = (",", ":")).encode( + "utf-8" + ) + ).hexdigest() metadata = { "requested_tag": requested_tag, "tag": llama_tag, + "release_tag": release_tag, + "published_repo": approved_checksums.repo, "asset": choice.name, + "asset_sha256": choice.expected_sha256, "source": choice.source_label, + "source_asset": source_asset_name, + "source_sha256": source_sha256, + "source_commit": approved_checksums.source_commit, + "source_commit_short": approved_checksums.source_commit_short, + "source_repo": approved_checksums.source_repo, + "source_repo_url": approved_checksums.source_repo_url, + "source_ref_kind": approved_checksums.source_ref_kind, + "requested_source_ref": approved_checksums.requested_source_ref, + "resolved_source_ref": approved_checksums.resolved_source_ref, "bundle_profile": choice.bundle_profile, "runtime_line": choice.runtime_line, "coverage_class": choice.coverage_class, + "install_fingerprint": fingerprint, "prebuilt_fallback_used": prebuilt_fallback_used, "installed_at_utc": time.strftime("%Y-%m-%dT%H:%M:%SZ", time.gmtime()), } @@ -3169,6 +4146,184 @@ def write_prebuilt_metadata( ) +def expected_install_fingerprint( + *, + llama_tag: str, + release_tag: str, + choice: AssetChoice, + approved_checksums: ApprovedReleaseChecksums, +) -> str | None: + if not choice.expected_sha256: + return None + source_asset_name, source_sha256 = selected_source_archive_metadata( + approved_checksums, + llama_tag, + ) + payload = { + "published_repo": approved_checksums.repo, + "release_tag": release_tag, + "upstream_tag": llama_tag, + "asset": choice.name, + "asset_sha256": choice.expected_sha256, + "source": choice.source_label, + "source_asset": source_asset_name, + "source_sha256": source_sha256, + "runtime_line": choice.runtime_line, + "bundle_profile": choice.bundle_profile, + "coverage_class": choice.coverage_class, + } + return hashlib.sha256( + json.dumps(payload, sort_keys = True, separators = (",", ":")).encode("utf-8") + ).hexdigest() + + +def load_prebuilt_metadata(install_dir: Path) -> dict[str, Any] | None: + metadata_path = install_dir / "UNSLOTH_PREBUILT_INFO.json" + if not metadata_path.is_file(): + return None + try: + payload = json.loads(metadata_path.read_text(encoding = "utf-8")) + except Exception: + return None + if not isinstance(payload, dict): + return None + return payload + + +def runtime_payload_health_groups(choice: AssetChoice) -> list[list[str]]: + if choice.install_kind == "linux-cpu": + return [ + ["libllama.so*"], + ["libggml.so*"], + ["libggml-base.so*"], + ["libggml-cpu-*.so*"], + ["libmtmd.so*"], + ] + if choice.install_kind == "linux-cuda": + return [ + ["libllama.so*"], + ["libggml.so*"], + ["libggml-base.so*"], + ["libggml-cpu-*.so*"], + ["libmtmd.so*"], + ["libggml-cuda.so*"], + ] + if choice.install_kind in {"macos-arm64", "macos-x64"}: + return [ + ["libllama*.dylib"], + ["libggml*.dylib"], + ["libmtmd*.dylib"], + ] + if choice.install_kind == "windows-cpu": + return [["llama.dll"]] + if choice.install_kind == "windows-cuda": + return [["llama.dll"], ["ggml-cuda.dll"]] + return [] + + +def install_runtime_dir(install_dir: Path, host: HostInfo) -> Path: + if host.is_windows: + return install_dir / "build" / "bin" / "Release" + return install_dir / "build" / "bin" + + +def runtime_payload_is_healthy( + install_dir: Path, host: HostInfo, choice: AssetChoice +) -> bool: + runtime_dir = install_runtime_dir(install_dir, host) + if not runtime_dir.exists(): + return False + for pattern_group in runtime_payload_health_groups(choice): + matched = False + for pattern in pattern_group: + if any(runtime_dir.glob(pattern)): + matched = True + break + if not matched: + return False + return True + + +def existing_install_matches_choice( + install_dir: Path, + host: HostInfo, + *, + llama_tag: str, + release_tag: str, + choice: AssetChoice, + approved_checksums: ApprovedReleaseChecksums, +) -> bool: + if not install_dir.exists(): + return False + + metadata = load_prebuilt_metadata(install_dir) + if metadata is None: + return False + + try: + confirm_install_tree(install_dir, host) + except Exception: + return False + + if not runtime_payload_is_healthy(install_dir, host, choice): + return False + + # Verify primary executables still exist (catches partial deletion) + runtime_dir = install_runtime_dir(install_dir, host) + ext = ".exe" if host.is_windows else "" + for binary in ("llama-server", "llama-quantize"): + if not (runtime_dir / f"{binary}{ext}").exists(): + return False + expected_fingerprint = expected_install_fingerprint( + llama_tag = llama_tag, + release_tag = release_tag, + choice = choice, + approved_checksums = approved_checksums, + ) + if not expected_fingerprint: + return False + + recorded_fingerprint = metadata.get("install_fingerprint") + if not isinstance(recorded_fingerprint, str) or not recorded_fingerprint: + return False + + if recorded_fingerprint != expected_fingerprint: + return False + + expected_pairs = { + "release_tag": release_tag, + "published_repo": approved_checksums.repo, + "tag": llama_tag, + "asset": choice.name, + "asset_sha256": choice.expected_sha256, + "source": choice.source_label, + "runtime_line": choice.runtime_line, + "bundle_profile": choice.bundle_profile, + "coverage_class": choice.coverage_class, + } + for key, expected in expected_pairs.items(): + if metadata.get(key) != expected: + return False + return True + + +def existing_install_matches_plan( + install_dir: Path, + host: HostInfo, + plan: InstallReleasePlan, +) -> bool: + if not plan.attempts: + return False + return existing_install_matches_choice( + install_dir, + host, + llama_tag = plan.llama_tag, + release_tag = plan.release_tag, + choice = plan.attempts[0], + approved_checksums = plan.approved_checksums, + ) + + def validate_prebuilt_choice( choice: AssetChoice, host: HostInfo, @@ -3178,23 +4333,32 @@ def validate_prebuilt_choice( *, requested_tag: str, llama_tag: str, + release_tag: str, approved_checksums: ApprovedReleaseChecksums, prebuilt_fallback_used: bool, quantized_path: Path, ) -> tuple[Path, Path]: - source_archive = approved_checksums.artifacts.get( - source_archive_logical_name(llama_tag) + source_repo, source_ref, source_archive, exact_source = preferred_source_archive( + approved_checksums, llama_tag ) - if source_archive is None: - raise PrebuiltFallback( - f"approved checksum asset did not contain source archive {source_archive_logical_name(llama_tag)}" + if exact_source: + log( + f"hydrating exact llama.cpp source for {source_repo}@{source_ref} into {install_dir}" ) - log(f"hydrating upstream llama.cpp source for {llama_tag} into {install_dir}") + else: + log(f"hydrating upstream llama.cpp source for {llama_tag} into {install_dir}") hydrate_source_tree( - llama_tag, + source_ref, install_dir, work_dir, + source_repo = source_repo, expected_sha256 = source_archive.sha256, + source_label = ( + f"llama.cpp source tree for {source_repo}@{source_ref}" + if exact_source + else f"llama.cpp source tree for {llama_tag}" + ), + exact_source = exact_source, ) log(f"overlaying prebuilt bundle {choice.name} into {install_dir}") server_path, quantize_path = install_from_archives( @@ -3206,7 +4370,9 @@ def validate_prebuilt_choice( install_dir, requested_tag = requested_tag, llama_tag = llama_tag, + release_tag = release_tag, choice = choice, + approved_checksums = approved_checksums, prebuilt_fallback_used = prebuilt_fallback_used, ) validate_quantize( @@ -3237,13 +4403,16 @@ def validate_prebuilt_attempts( *, requested_tag: str, llama_tag: str, + release_tag: str, approved_checksums: ApprovedReleaseChecksums, + initial_fallback_used: bool = False, + existing_install_dir: Path | None = None, ) -> tuple[AssetChoice, Path, bool]: attempt_list = list(attempts) if not attempt_list: raise PrebuiltFallback("no prebuilt bundle attempts were available") - tried_fallback = False + tried_fallback = initial_fallback_used for index, attempt in enumerate(attempt_list): if index > 0: tried_fallback = True @@ -3253,6 +4422,20 @@ def validate_prebuilt_attempts( f"runtime_line={attempt.runtime_line} coverage_class={attempt.coverage_class}" ) + if existing_install_dir is not None and existing_install_matches_choice( + existing_install_dir, + host, + llama_tag = llama_tag, + release_tag = release_tag, + choice = attempt, + approved_checksums = approved_checksums, + ): + log( + "existing llama.cpp install already matches fallback candidate " + f"{attempt.name}; skipping reinstall" + ) + raise ExistingInstallSatisfied(attempt, tried_fallback) + staging_dir = create_install_staging_dir(install_dir) quantized_path = work_dir / f"stories260K-q4-{index}.gguf" if quantized_path.exists(): @@ -3266,6 +4449,7 @@ def validate_prebuilt_attempts( probe_path, requested_tag = requested_tag, llama_tag = llama_tag, + release_tag = release_tag, approved_checksums = approved_checksums, prebuilt_fallback_used = tried_fallback, quantized_path = quantized_path, @@ -3307,42 +4491,81 @@ def install_prebuilt( log( f"no existing llama.cpp install detected at {install_dir}; performing fresh prebuilt install" ) - requested_tag, llama_tag, attempts, approved_checksums = ( - resolve_install_attempts( - llama_tag, - host, - published_repo, - published_release_tag, + requested_tag, release_plans = resolve_install_release_plans( + llama_tag, + host, + published_repo, + published_release_tag, + ) + if release_plans and existing_install_matches_plan( + install_dir, host, release_plans[0] + ): + current = release_plans[0] + log( + "existing llama.cpp install already matches selected release " + f"{current.release_tag} upstream_tag={current.llama_tag}; skipping download and install" ) - ) - choice = attempts[0] - log( - f"selected {choice.name} ({choice.source_label}) for {host.system} {host.machine}" - ) + return with tempfile.TemporaryDirectory(prefix = "unsloth-llama-prebuilt-") as tmp: work_dir = Path(tmp) probe_path = work_dir / "stories260K.gguf" download_validation_model( probe_path, validation_model_cache_path(install_dir) ) - choice, selected_staging_dir, _ = validate_prebuilt_attempts( - attempts, - host, - install_dir, - work_dir, - probe_path, - requested_tag = requested_tag, - llama_tag = llama_tag, - approved_checksums = approved_checksums, - ) - activate_install_tree(selected_staging_dir, install_dir, host) - try: - ensure_converter_scripts(install_dir, llama_tag) - except Exception as exc: + release_count = len(release_plans) + for release_index, plan in enumerate(release_plans): + choice = plan.attempts[0] + if existing_install_matches_plan(install_dir, host, plan): + log( + "existing llama.cpp install already matches fallback release " + f"{plan.release_tag} upstream_tag={plan.llama_tag}; skipping reinstall" + ) + return log( - "converter script fetch failed after activation; install remains valid " - f"({textwrap.shorten(str(exc), width = 200, placeholder = '...')})" + "selected " + f"{choice.name} ({choice.source_label}) from published release " + f"{plan.release_tag} for {host.system} {host.machine}" ) + try: + choice, selected_staging_dir, _ = validate_prebuilt_attempts( + plan.attempts, + host, + install_dir, + work_dir, + probe_path, + requested_tag = requested_tag, + llama_tag = plan.llama_tag, + release_tag = plan.release_tag, + approved_checksums = plan.approved_checksums, + initial_fallback_used = release_index > 0, + existing_install_dir = install_dir, + ) + except ExistingInstallSatisfied: + return + except PrebuiltFallback as exc: + if release_index == release_count - 1: + raise + log( + "published release " + f"{plan.release_tag} upstream_tag={plan.llama_tag} failed; " + "trying the next older published prebuilt " + f"({textwrap.shorten(str(exc), width = 200, placeholder = '...')})" + ) + continue + + activate_install_tree(selected_staging_dir, install_dir, host) + try: + ensure_converter_scripts(install_dir, plan.llama_tag) + except Exception as exc: + log( + "converter script fetch failed after activation; install remains valid " + f"({textwrap.shorten(str(exc), width = 200, placeholder = '...')})" + ) + return + except BusyInstallConflict as exc: + log("prebuilt install path is blocked by an in-use llama.cpp install") + log(f"prebuilt busy reason: {exc}") + raise SystemExit(EXIT_BUSY) from exc except PrebuiltFallback as exc: log("prebuilt install path failed; falling back to source build") log(f"prebuilt fallback reason: {exc}") @@ -3359,7 +4582,10 @@ def parse_args() -> argparse.Namespace: parser.add_argument( "--llama-tag", default = DEFAULT_LLAMA_TAG, - help = f"llama.cpp release tag. Prebuilt installs are pinned to the approved tag {APPROVED_PREBUILT_LLAMA_TAG}.", + help = ( + "llama.cpp release tag. Defaults to the latest usable published Unsloth " + "release unless UNSLOTH_LLAMA_TAG overrides it." + ), ) parser.add_argument( "--published-repo", @@ -3369,7 +4595,10 @@ def parse_args() -> argparse.Namespace: parser.add_argument( "--published-release-tag", default = DEFAULT_PUBLISHED_TAG, - help = "Published GitHub release tag to pin. By default, scan releases until a compatible llama.cpp bundle is found.", + help = ( + "Published GitHub release tag to pin. By default, scan releases " + "until a usable published llama.cpp release bundle is found." + ), ) resolve_group = parser.add_mutually_exclusive_group() resolve_group.add_argument( @@ -3382,30 +4611,108 @@ def parse_args() -> argparse.Namespace: "--resolve-install-tag", nargs = "?", const = "latest", - help = "Resolve a llama.cpp tag such as 'latest' to the concrete tag installable on the current host.", + help = ( + "Resolve a llama.cpp tag such as 'latest' to the concrete upstream tag " + "selected by the current published-release policy." + ), + ) + resolve_group.add_argument( + "--resolve-source-build", + nargs = "?", + const = "latest", + help = ("Resolve the source-build fallback plan."), + ) + parser.add_argument( + "--output-format", + choices = ("plain", "json"), + default = "plain", + help = "Resolver output format. Defaults to plain.", ) return parser.parse_args() +def emit_resolver_output(payload: dict[str, Any], *, output_format: str) -> None: + if output_format == "json": + print(json.dumps(payload, sort_keys = True)) + return + if "llama_tag" in payload: + print(payload["llama_tag"]) + return + if { + "source_url", + "source_ref_kind", + "source_ref", + }.issubset(payload): + print( + "\t".join( + ( + str(payload["source_url"]), + str(payload["source_ref_kind"]), + str(payload["source_ref"]), + ) + ) + ) + return + print(json.dumps(payload, sort_keys = True)) + + def main() -> int: args = parse_args() if args.resolve_llama_tag is not None: - # Pass published_repo so the resolver prefers the Unsloth release tag - # (tested/approved) over the upstream ggml-org bleeding-edge tag. - print(resolve_requested_llama_tag(args.resolve_llama_tag, args.published_repo)) + resolved = resolve_requested_llama_tag( + args.resolve_llama_tag, + args.published_repo, + args.published_release_tag or "", + ) + emit_resolver_output( + { + "requested_tag": normalized_requested_llama_tag(args.resolve_llama_tag), + "llama_tag": resolved, + }, + output_format = args.output_format, + ) return EXIT_SUCCESS if args.resolve_install_tag is not None: - print( - resolve_requested_install_tag( - args.resolve_install_tag, args.published_release_tag or "" - ) + resolved = resolve_requested_install_tag( + args.resolve_install_tag, + args.published_release_tag or "", + args.published_repo, + ) + emit_resolver_output( + { + "requested_tag": normalized_requested_llama_tag( + args.resolve_install_tag + ), + "llama_tag": resolved, + }, + output_format = args.output_format, + ) + return EXIT_SUCCESS + + if args.resolve_source_build is not None: + plan = resolve_source_build_plan( + args.resolve_source_build, + args.published_repo, + args.published_release_tag or "", + ) + emit_resolver_output( + { + "requested_tag": normalized_requested_llama_tag( + args.resolve_source_build + ), + "source_url": plan.source_url, + "source_ref_kind": plan.source_ref_kind, + "source_ref": plan.source_ref, + "compatibility_upstream_tag": plan.compatibility_upstream_tag, + }, + output_format = args.output_format, ) return EXIT_SUCCESS if not args.install_dir: raise SystemExit( - "install_llama_prebuilt.py: --install-dir is required unless --resolve-llama-tag or --resolve-install-tag is used" + "install_llama_prebuilt.py: --install-dir is required unless --resolve-llama-tag, --resolve-install-tag, or --resolve-source-build is used" ) install_prebuilt( install_dir = Path(args.install_dir).expanduser().resolve(), @@ -3421,6 +4728,11 @@ if __name__ == "__main__": raise SystemExit(main()) except SystemExit: raise + except BusyInstallConflict as exc: + log( + f"fatal helper busy conflict: {textwrap.shorten(str(exc), width = 400, placeholder = '...')}" + ) + raise SystemExit(EXIT_BUSY) except Exception as exc: message = textwrap.shorten(str(exc), width = 400, placeholder = "...") log(f"fatal helper error: {message}") diff --git a/studio/setup.ps1 b/studio/setup.ps1 index cdba0e6690..d478319098 100644 --- a/studio/setup.ps1 +++ b/studio/setup.ps1 @@ -22,6 +22,19 @@ $ErrorActionPreference = "Stop" $ScriptDir = Split-Path -Parent $MyInvocation.MyCommand.Path $PackageDir = Split-Path -Parent $ScriptDir +# -------------------------------------------------------------------------- +# Maintainer-editable defaults +# Change these in the GitHub-hosted script so users get updated defaults. +# User env vars always override these baked-in values. +# -------------------------------------------------------------------------- +# Prefer "latest" over "master" -- "master" bypasses the prebuilt resolver +# (no matching GitHub release), forces a source build, and causes HTTP 422 +# errors. Only use "master" temporarily when the latest release is missing +# support for a new model architecture. +$DefaultLlamaPrForce = "" +$DefaultLlamaSource = "https://github.com/ggml-org/llama.cpp" +$DefaultLlamaTag = "latest" + # Verbose can be enabled either by CLI flag or by UNSLOTH_VERBOSE=1. $script:UnslothVerbose = ($env:UNSLOTH_VERBOSE -eq '1') foreach ($a in $args) { @@ -62,6 +75,12 @@ function Refresh-Environment { $env:Path = "$machinePath;$userPath" } +# PowerShell 5.1 compatibility helper: avoid relying on New-TemporaryFile. +function New-UnslothTemporaryFile { + $tempPath = [System.IO.Path]::GetTempFileName() + return Get-Item -LiteralPath $tempPath +} + # Find nvcc on PATH, CUDA_PATH, or standard toolkit dirs. # Returns the path to nvcc.exe, or $null if not found. function Find-Nvcc { @@ -490,8 +509,7 @@ if (-not $HasNvidiaSmi) { if (-not $HasNvidiaSmi) { Write-Host "" step "gpu" "none (chat-only / GGUF)" "Yellow" - Write-Host " Training and GPU inference require an NVIDIA GPU with drivers installed." -ForegroundColor Yellow - Write-Host " https://www.nvidia.com/Download/index.aspx" -ForegroundColor Yellow + substep "Training and GPU inference require an NVIDIA GPU with drivers installed." "Yellow" Write-Host "" } else { step "gpu" "NVIDIA GPU detected" @@ -1544,7 +1562,7 @@ if (Test-Path $VenvT5Dir) { Remove-Item -Recurse -Force $VenvT5Dir } New-Item -ItemType Directory -Path $VenvT5Dir -Force | Out-Null $prevEAP_t5 = $ErrorActionPreference $ErrorActionPreference = "Continue" -foreach ($pkg in @("transformers==5.3.0", "huggingface_hub==1.7.1", "hf_xet==1.4.2")) { +foreach ($pkg in @("transformers==5.5.0", "huggingface_hub==1.8.0", "hf_xet==1.4.2")) { if ($script:UnslothVerbose) { Fast-Install --target $VenvT5Dir --no-deps $pkg $t5PkgExit = $LASTEXITCODE @@ -1590,43 +1608,156 @@ if (-not (Test-Path $UnslothHome)) { New-Item -ItemType Directory -Force $Unslot $LlamaCppDir = Join-Path $UnslothHome "llama.cpp" $NeedLlamaSourceBuild = $false $SkipPrebuiltInstall = $false -$RequestedLlamaTag = if ($env:UNSLOTH_LLAMA_TAG) { $env:UNSLOTH_LLAMA_TAG } else { "latest" } -$HelperReleaseRepo = if ($env:UNSLOTH_LLAMA_RELEASE_REPO) { $env:UNSLOTH_LLAMA_RELEASE_REPO } else { "unslothai/llama.cpp" } -$resolveOutput = & python "$PSScriptRoot\install_llama_prebuilt.py" --resolve-install-tag $RequestedLlamaTag --published-repo $HelperReleaseRepo 2>&1 -$resolveExit = $LASTEXITCODE -$ResolvedLlamaTag = if ($resolveOutput) { ($resolveOutput | Select-Object -Last 1).ToString().Trim() } else { "" } -if ($resolveExit -ne 0 -or [string]::IsNullOrWhiteSpace($ResolvedLlamaTag)) { - Write-Host "" - substep "Failed to resolve an installable prebuilt llama.cpp tag via $HelperReleaseRepo" "Yellow" - Write-LlamaFailureLog -Output ($resolveOutput | Out-String) - # Resolve the llama.cpp tag for source-build fallback. Pass --published-repo - # so the resolver prefers Unsloth's tested tag (e.g. b8508) over the upstream - # bleeding-edge tag (e.g. b8514) from ggml-org/llama.cpp. - $fallbackOutput = & python "$PSScriptRoot\install_llama_prebuilt.py" --resolve-llama-tag $RequestedLlamaTag --published-repo $HelperReleaseRepo 2>$null - $fallbackExit = $LASTEXITCODE - $ResolvedLlamaTag = if ($fallbackExit -eq 0 -and $fallbackOutput) { - ($fallbackOutput | Select-Object -Last 1).ToString().Trim() - } elseif ($RequestedLlamaTag -eq "latest") { - # Try Unsloth release repo first, then fall back to ggml-org upstream - $resolvedLatest = $null - try { - $latestRelease = Invoke-RestMethod -Uri "https://api.github.com/repos/$HelperReleaseRepo/releases/latest" -ErrorAction Stop - $resolvedLatest = $latestRelease.tag_name - } catch {} - if (-not $resolvedLatest) { - try { - $latestRelease = Invoke-RestMethod -Uri "https://api.github.com/repos/ggml-org/llama.cpp/releases/latest" -ErrorAction Stop - $resolvedLatest = $latestRelease.tag_name - } catch {} - } - if ($resolvedLatest) { $resolvedLatest } else { $RequestedLlamaTag } - } else { - $RequestedLlamaTag - } +$RequestedLlamaTag = if ($env:UNSLOTH_LLAMA_TAG) { $env:UNSLOTH_LLAMA_TAG } else { $DefaultLlamaTag } +$HelperReleaseRepo = "ggml-org/llama.cpp" +$LlamaPr = if ($env:UNSLOTH_LLAMA_PR) { $env:UNSLOTH_LLAMA_PR.Trim() } else { "" } + +$LlamaPrForce = if ($env:UNSLOTH_LLAMA_PR_FORCE) { $env:UNSLOTH_LLAMA_PR_FORCE.Trim() } else { $DefaultLlamaPrForce } +$LlamaSource = $DefaultLlamaSource +if ($LlamaSource.EndsWith('.git')) { $LlamaSource = $LlamaSource.Substring(0, $LlamaSource.Length - 4) } +$ResolvedSourceUrl = $LlamaSource +$ResolvedSourceRef = $RequestedLlamaTag +$ResolvedSourceRefKind = "tag" + +if ($env:UNSLOTH_LLAMA_FORCE_COMPILE -eq "1") { $NeedLlamaSourceBuild = $true $SkipPrebuiltInstall = $true } +function Invoke-LlamaHelper { + param( + [string[]]$Arguments, + [string]$StderrPath = $null + ) + + $previousErrorActionPreference = $ErrorActionPreference + $previousNativeErrorPreference = $null + $restoreNativeErrorPreference = $false + $ErrorActionPreference = "Continue" + if ($PSVersionTable.PSVersion.Major -ge 7) { + $previousNativeErrorPreference = $PSNativeCommandUseErrorActionPreference + $PSNativeCommandUseErrorActionPreference = $false + $restoreNativeErrorPreference = $true + } + + try { + # Capture all output (stdout + stderr) so that PowerShell does not + # convert stderr lines into visible ErrorRecord objects. Separate + # stdout from stderr afterwards. + $allOutput = & python "$PSScriptRoot\install_llama_prebuilt.py" @Arguments 2>&1 + $exitCode = $LASTEXITCODE + $stdoutLines = @() + $stderrLines = @() + foreach ($line in $allOutput) { + if ($line -is [System.Management.Automation.ErrorRecord]) { + $stderrLines += $line.ToString() + } else { + $stdoutLines += $line + } + } + if ($StderrPath -and $stderrLines.Count -gt 0) { + $stderrLines | Out-File -FilePath $StderrPath -Encoding utf8 + } + return [pscustomobject]@{ + Output = $stdoutLines + ExitCode = $exitCode + } + } finally { + $ErrorActionPreference = $previousErrorActionPreference + if ($restoreNativeErrorPreference) { + $PSNativeCommandUseErrorActionPreference = $previousNativeErrorPreference + } + } +} + +if ($LlamaSource -ne "https://github.com/ggml-org/llama.cpp") { + step "llama.cpp" "custom source: $LlamaSource -- forcing source build" "Yellow" + $NeedLlamaSourceBuild = $true + $SkipPrebuiltInstall = $true +} + +if (-not $LlamaPr -and $LlamaPrForce -and $LlamaPrForce -match '^\d+$' -and [int]$LlamaPrForce -gt 0) { + $LlamaPr = $LlamaPrForce + step "llama.cpp" "baked-in PR_FORCE=$LlamaPrForce" "Yellow" +} + +if ($LlamaPr) { + if ($LlamaPr -notmatch '^\d+$' -or [int]$LlamaPr -le 0) { + Write-Host "[ERROR] UNSLOTH_LLAMA_PR=$LlamaPr is not a valid PR number" -ForegroundColor Red + exit 1 + } + step "llama.cpp" "UNSLOTH_LLAMA_PR=$LlamaPr -- will build from PR head" "Yellow" + $ResolvedLlamaTag = "pr-$LlamaPr" + $ResolvedSourceUrl = $LlamaSource + $ResolvedSourceRef = "pr-$LlamaPr" + $ResolvedSourceRefKind = "pull" + $NeedLlamaSourceBuild = $true + $SkipPrebuiltInstall = $true +} elseif ($SkipPrebuiltInstall) { + # Custom source or other override already forced source build; skip the + # prebuilt release resolution. When building from a custom fork, the fork + # may not carry upstream bNNNN tags. + if ($env:UNSLOTH_LLAMA_FORCE_COMPILE -eq "1") { + $ResolvedLlamaTag = $RequestedLlamaTag + } elseif ($LlamaSource -eq "https://github.com/ggml-org/llama.cpp") { + $resolveTagArgs = @("--resolve-llama-tag", $RequestedLlamaTag, "--published-repo", $HelperReleaseRepo, "--output-format", "json") + if ($env:UNSLOTH_LLAMA_RELEASE_TAG) { $resolveTagArgs += @("--published-release-tag", $env:UNSLOTH_LLAMA_RELEASE_TAG) } + $fallbackResult = Invoke-LlamaHelper -Arguments $resolveTagArgs + $fallbackOutput = $fallbackResult.Output + $fallbackExit = $fallbackResult.ExitCode + $ResolvedLlamaTag = if ($fallbackExit -eq 0 -and $fallbackOutput) { + try { + (($fallbackOutput | Out-String) | ConvertFrom-Json).llama_tag + } catch { + $RequestedLlamaTag + } + } else { + $RequestedLlamaTag + } + } else { + $ResolvedLlamaTag = $RequestedLlamaTag + } +} else { + $resolveInstallArgs = @("--resolve-install-tag", $RequestedLlamaTag, "--published-repo", $HelperReleaseRepo, "--output-format", "json") + if ($env:UNSLOTH_LLAMA_RELEASE_TAG) { $resolveInstallArgs += @("--published-release-tag", $env:UNSLOTH_LLAMA_RELEASE_TAG) } + $resolveErrorLog = New-UnslothTemporaryFile + $resolveResult = Invoke-LlamaHelper -Arguments $resolveInstallArgs -StderrPath $resolveErrorLog + $resolveOutput = $resolveResult.Output + $resolveExit = $resolveResult.ExitCode + $ResolvedLlamaTag = if ($resolveOutput) { + try { + (($resolveOutput | Out-String) | ConvertFrom-Json).llama_tag + } catch { + "" + } + } else { "" } + if ($resolveExit -ne 0 -or [string]::IsNullOrWhiteSpace($ResolvedLlamaTag)) { + Write-Host "" + substep "Failed to resolve a published llama.cpp release via $HelperReleaseRepo" "Yellow" + Write-LlamaFailureLog -Output (Get-Content -Raw $resolveErrorLog) + # Resolve the llama.cpp tag for source-build fallback. Pass --published-repo + # so the resolver prefers the latest usable Unsloth-published upstream tag + # before falling back to the bleeding-edge ggml-org/llama.cpp tag. + $resolveFallbackArgs = @("--resolve-llama-tag", $RequestedLlamaTag, "--published-repo", $HelperReleaseRepo, "--output-format", "json") + if ($env:UNSLOTH_LLAMA_RELEASE_TAG) { $resolveFallbackArgs += @("--published-release-tag", $env:UNSLOTH_LLAMA_RELEASE_TAG) } + $fallbackResult = Invoke-LlamaHelper -Arguments $resolveFallbackArgs + $fallbackOutput = $fallbackResult.Output + $fallbackExit = $fallbackResult.ExitCode + $ResolvedLlamaTag = if ($fallbackExit -eq 0 -and $fallbackOutput) { + try { + (($fallbackOutput | Out-String) | ConvertFrom-Json).llama_tag + } catch { + $RequestedLlamaTag + } + } else { + $RequestedLlamaTag + } + $NeedLlamaSourceBuild = $true + $SkipPrebuiltInstall = $true + } + Remove-Item $resolveErrorLog -Force -ErrorAction SilentlyContinue +} + Write-Host "" substep "Resolved llama.cpp release tag: $ResolvedLlamaTag" @@ -1646,7 +1777,7 @@ if ($env:UNSLOTH_LLAMA_FORCE_COMPILE -eq "1") { $prebuiltArgs = @( "$PSScriptRoot\install_llama_prebuilt.py", "--install-dir", $LlamaCppDir, - "--llama-tag", $ResolvedLlamaTag, + "--llama-tag", $RequestedLlamaTag, "--published-repo", $HelperReleaseRepo ) if ($env:UNSLOTH_LLAMA_RELEASE_TAG) { @@ -1654,21 +1785,46 @@ if ($env:UNSLOTH_LLAMA_FORCE_COMPILE -eq "1") { } $prevEAPPrebuilt = $ErrorActionPreference $ErrorActionPreference = "Continue" - if ($script:UnslothVerbose) { - # Show live output in verbose mode while still capturing for error log - $prebuiltLog = Join-Path $env:TEMP "unsloth-prebuilt-$PID.log" - & python @prebuiltArgs 2>&1 | Tee-Object -FilePath $prebuiltLog | Out-Host - $prebuiltExit = $LASTEXITCODE - $prebuiltOutput = if (Test-Path $prebuiltLog) { Get-Content $prebuiltLog -Raw } else { "" } - Remove-Item $prebuiltLog -ErrorAction SilentlyContinue - } else { - $prebuiltOutput = & python @prebuiltArgs 2>&1 | Out-String - $prebuiltExit = $LASTEXITCODE + $previousNativeErrorPreference = $null + $restoreNativeErrorPreference = $false + if ($PSVersionTable.PSVersion.Major -ge 7) { + $previousNativeErrorPreference = $PSNativeCommandUseErrorActionPreference + $PSNativeCommandUseErrorActionPreference = $false + $restoreNativeErrorPreference = $true + } + try { + if ($script:UnslothVerbose) { + # Show live output in verbose mode while still capturing for error log + $prebuiltLog = Join-Path $env:TEMP "unsloth-prebuilt-$PID.log" + & python @prebuiltArgs 2>&1 | Tee-Object -FilePath $prebuiltLog | Out-Host + $prebuiltExit = $LASTEXITCODE + $prebuiltOutput = if (Test-Path $prebuiltLog) { Get-Content $prebuiltLog -Raw } else { "" } + Remove-Item $prebuiltLog -ErrorAction SilentlyContinue + } else { + $prebuiltOutput = & python @prebuiltArgs 2>&1 | Out-String + $prebuiltExit = $LASTEXITCODE + } + } finally { + if ($restoreNativeErrorPreference) { + $PSNativeCommandUseErrorActionPreference = $previousNativeErrorPreference + } } $ErrorActionPreference = $prevEAPPrebuilt if ($prebuiltExit -eq 0) { - step "llama.cpp" "prebuilt installed and validated" + if ($prebuiltOutput -match "already matches") { + step "llama.cpp" "prebuilt up to date and validated" + } else { + step "llama.cpp" "prebuilt installed and validated" + } + } elseif ($prebuiltExit -eq 3) { + step "llama.cpp" "install blocked by active llama.cpp process" "Yellow" + Write-LlamaFailureLog -Output $prebuiltOutput + if (Test-Path $LlamaCppDir) { + substep "Existing install was restored" "Yellow" + } + substep "Close Studio or other llama.cpp users and retry" "Yellow" + exit 3 } else { step "llama.cpp" "prebuilt install failed (continuing)" "Yellow" Write-LlamaFailureLog -Output $prebuiltOutput @@ -1766,7 +1922,10 @@ if (Test-Path $LlamaServerBin) { if (-not $NeedLlamaSourceBuild) { Write-Host "" step "llama.cpp" "prebuilt (validated)" -} elseif ((Test-Path $LlamaServerBin) -and -not $NeedRebuild) { +} elseif ((Test-Path $LlamaServerBin) -and -not $NeedRebuild -and $RequestedLlamaTag -ne "master") { + # Skip rebuild only for pinned tags (e.g. b8635). When the requested + # tag is "master" (a moving target), always rebuild so the binary picks + # up new model architecture support (e.g. Gemma 4). Write-Host "" step "llama.cpp" "already built" } elseif (-not $HasCmakeForBuild) { @@ -1821,42 +1980,188 @@ if (-not $NeedLlamaSourceBuild) { [Environment]::SetEnvironmentVariable('CudaToolkitDir', "$CudaToolkitRoot\", 'Process') } + if (-not $LlamaPr) { + if ($LlamaSource -eq "https://github.com/ggml-org/llama.cpp") { + $resolveSourceArgs = @("--resolve-source-build", $RequestedLlamaTag, "--published-repo", $HelperReleaseRepo, "--output-format", "json") + if ($env:UNSLOTH_LLAMA_RELEASE_TAG) { $resolveSourceArgs += @("--published-release-tag", $env:UNSLOTH_LLAMA_RELEASE_TAG) } + $sourcePlanResult = Invoke-LlamaHelper -Arguments $resolveSourceArgs + $sourcePlanOutput = $sourcePlanResult.Output + $sourcePlanExit = $sourcePlanResult.ExitCode + if ($sourcePlanExit -eq 0 -and $sourcePlanOutput) { + try { + $sourcePlan = ($sourcePlanOutput | Out-String) | ConvertFrom-Json + $ResolvedSourceUrl = $sourcePlan.source_url + $ResolvedSourceRefKind = $sourcePlan.source_ref_kind + $ResolvedSourceRef = $sourcePlan.source_ref + } catch { + } + } + } + if ([string]::IsNullOrWhiteSpace($ResolvedSourceUrl)) { $ResolvedSourceUrl = $LlamaSource } + if ([string]::IsNullOrWhiteSpace($ResolvedSourceRef)) { $ResolvedSourceRef = $ResolvedLlamaTag } + } + # -- Step A: Clone or pull llama.cpp -- - $UseConcreteRef = ($ResolvedLlamaTag -ne "latest" -and -not [string]::IsNullOrWhiteSpace($ResolvedLlamaTag)) + $UseConcreteRef = ($ResolvedSourceRef -ne "latest" -and -not [string]::IsNullOrWhiteSpace($ResolvedSourceRef)) if (Test-Path (Join-Path $LlamaCppDir ".git")) { - Write-Host " Syncing llama.cpp to $ResolvedLlamaTag..." -ForegroundColor Gray - if ($UseConcreteRef) { - $gitFetchExit = Invoke-SetupCommand -AlwaysQuiet { git -C $LlamaCppDir fetch --depth 1 origin $ResolvedLlamaTag } + Write-Host " Syncing llama.cpp to $ResolvedSourceRef..." -ForegroundColor Gray + # Always sync the remote URL so switching between default/fork sources works + Invoke-SetupCommand -AlwaysQuiet { git -C $LlamaCppDir remote set-url origin "$ResolvedSourceUrl.git" } | Out-Null + if ($LlamaPr) { + $gitFetchExit = Invoke-SetupCommand -AlwaysQuiet { git -C $LlamaCppDir fetch --depth 1 origin "pull/$LlamaPr/head" } + if ($gitFetchExit -ne 0) { + $BuildOk = $false + $FailedStep = "git fetch PR #$LlamaPr" + } else { + $gitCheckoutExit = Invoke-SetupCommand -AlwaysQuiet { git -C $LlamaCppDir checkout -B "pr-$LlamaPr" FETCH_HEAD } + if ($gitCheckoutExit -ne 0) { + $BuildOk = $false + $FailedStep = "git checkout PR #$LlamaPr" + } else { + Invoke-SetupCommand -AlwaysQuiet { git -C $LlamaCppDir clean -fdx } | Out-Null + } + } + } elseif ($ResolvedSourceRefKind -eq "pull") { + $gitFetchExit = Invoke-SetupCommand -AlwaysQuiet { git -C $LlamaCppDir fetch --depth 1 origin $ResolvedSourceRef } + if ($gitFetchExit -ne 0) { + substep "git fetch failed -- using existing source" "Yellow" + } else { + $gitCheckoutExit = Invoke-SetupCommand -AlwaysQuiet { git -C $LlamaCppDir checkout -B unsloth-llama-build FETCH_HEAD } + if ($gitCheckoutExit -ne 0) { + $BuildOk = $false + $FailedStep = "git checkout" + } else { + Invoke-SetupCommand -AlwaysQuiet { git -C $LlamaCppDir clean -fdx } | Out-Null + } + } + } elseif ($ResolvedSourceRefKind -eq "commit") { + $gitFetchExit = Invoke-SetupCommand -AlwaysQuiet { git -C $LlamaCppDir fetch --depth 1 origin $ResolvedSourceRef } + if ($gitFetchExit -ne 0) { + substep "git fetch failed -- using existing source" "Yellow" + } else { + $gitCheckoutExit = Invoke-SetupCommand -AlwaysQuiet { git -C $LlamaCppDir checkout -B unsloth-llama-build FETCH_HEAD } + if ($gitCheckoutExit -ne 0) { + $BuildOk = $false + $FailedStep = "git checkout" + } else { + Invoke-SetupCommand -AlwaysQuiet { git -C $LlamaCppDir clean -fdx } | Out-Null + } + } + } elseif ($UseConcreteRef) { + $gitFetchExit = Invoke-SetupCommand -AlwaysQuiet { git -C $LlamaCppDir fetch --depth 1 origin $ResolvedSourceRef } + if ($gitFetchExit -ne 0) { + substep "git fetch failed -- using existing source" "Yellow" + } else { + $gitCheckoutExit = Invoke-SetupCommand -AlwaysQuiet { git -C $LlamaCppDir checkout -B unsloth-llama-build FETCH_HEAD } + if ($gitCheckoutExit -ne 0) { + $BuildOk = $false + $FailedStep = "git checkout" + } else { + Invoke-SetupCommand -AlwaysQuiet { git -C $LlamaCppDir clean -fdx } | Out-Null + } + } } else { $gitFetchExit = Invoke-SetupCommand -AlwaysQuiet { git -C $LlamaCppDir fetch --depth 1 origin } - } - if ($gitFetchExit -ne 0) { - substep "git fetch failed -- using existing source" "Yellow" - } else { - $gitCheckoutExit = Invoke-SetupCommand -AlwaysQuiet { git -C $LlamaCppDir checkout -B unsloth-llama-build FETCH_HEAD } - if ($gitCheckoutExit -ne 0) { - $BuildOk = $false - $FailedStep = "git checkout" + if ($gitFetchExit -ne 0) { + substep "git fetch failed -- using existing source" "Yellow" } else { - Invoke-SetupCommand -AlwaysQuiet { git -C $LlamaCppDir clean -fdx } | Out-Null + $gitCheckoutExit = Invoke-SetupCommand -AlwaysQuiet { git -C $LlamaCppDir checkout -B unsloth-llama-build FETCH_HEAD } + if ($gitCheckoutExit -ne 0) { + $BuildOk = $false + $FailedStep = "git checkout" + } else { + Invoke-SetupCommand -AlwaysQuiet { git -C $LlamaCppDir clean -fdx } | Out-Null + } } } } else { - Write-Host " Cloning llama.cpp @ $ResolvedLlamaTag..." -ForegroundColor Gray + Write-Host " Cloning llama.cpp @ $ResolvedSourceRef..." -ForegroundColor Gray $buildTmp = "$LlamaCppDir.build.$PID" + $null = New-Item -ItemType Directory -Force -Path (Split-Path $LlamaCppDir -Parent) if (Test-Path $buildTmp) { Remove-Item -Recurse -Force $buildTmp } - $cloneArgs = @("clone", "--depth", "1") - if ($UseConcreteRef) { - $cloneArgs += @("--branch", $ResolvedLlamaTag) - } - $cloneArgs += @("https://github.com/ggml-org/llama.cpp.git", $buildTmp) - $cloneExit = Invoke-SetupCommand -AlwaysQuiet { git @cloneArgs } - if ($cloneExit -ne 0) { - $BuildOk = $false - $FailedStep = "git clone" - if (Test-Path $buildTmp) { Remove-Item -Recurse -Force $buildTmp } + if ($LlamaPr) { + $cloneExit = Invoke-SetupCommand -AlwaysQuiet { git clone --depth 1 "$LlamaSource.git" $buildTmp } + if ($cloneExit -ne 0) { + $BuildOk = $false + $FailedStep = "git clone" + if (Test-Path $buildTmp) { Remove-Item -Recurse -Force $buildTmp } + } + if ($BuildOk) { + $fetchExit = Invoke-SetupCommand -AlwaysQuiet { git -C $buildTmp fetch --depth 1 origin "pull/$LlamaPr/head:pr-$LlamaPr" } + if ($fetchExit -ne 0) { + $BuildOk = $false + $FailedStep = "git fetch PR #$LlamaPr" + if (Test-Path $buildTmp) { Remove-Item -Recurse -Force $buildTmp } + } + } + if ($BuildOk) { + $checkoutExit = Invoke-SetupCommand -AlwaysQuiet { git -C $buildTmp checkout "pr-$LlamaPr" } + if ($checkoutExit -ne 0) { + $BuildOk = $false + $FailedStep = "git checkout PR #$LlamaPr" + if (Test-Path $buildTmp) { Remove-Item -Recurse -Force $buildTmp } + } + } + } elseif ($ResolvedSourceRefKind -eq "pull") { + $cloneExit = Invoke-SetupCommand -AlwaysQuiet { git clone --depth 1 "$ResolvedSourceUrl.git" $buildTmp } + if ($cloneExit -ne 0) { + $BuildOk = $false + $FailedStep = "git clone" + if (Test-Path $buildTmp) { Remove-Item -Recurse -Force $buildTmp } + } + if ($BuildOk) { + $fetchExit = Invoke-SetupCommand -AlwaysQuiet { git -C $buildTmp fetch --depth 1 origin $ResolvedSourceRef } + if ($fetchExit -ne 0) { + $BuildOk = $false + $FailedStep = "git fetch source PR ref" + if (Test-Path $buildTmp) { Remove-Item -Recurse -Force $buildTmp } + } + } + if ($BuildOk) { + $checkoutExit = Invoke-SetupCommand -AlwaysQuiet { git -C $buildTmp checkout -B unsloth-llama-build FETCH_HEAD } + if ($checkoutExit -ne 0) { + $BuildOk = $false + $FailedStep = "git checkout source PR ref" + if (Test-Path $buildTmp) { Remove-Item -Recurse -Force $buildTmp } + } + } + } elseif ($ResolvedSourceRefKind -eq "commit") { + $cloneExit = Invoke-SetupCommand -AlwaysQuiet { git clone --depth 1 "$ResolvedSourceUrl.git" $buildTmp } + if ($cloneExit -ne 0) { + $BuildOk = $false + $FailedStep = "git clone" + if (Test-Path $buildTmp) { Remove-Item -Recurse -Force $buildTmp } + } + if ($BuildOk) { + $fetchExit = Invoke-SetupCommand -AlwaysQuiet { git -C $buildTmp fetch --depth 1 origin $ResolvedSourceRef } + if ($fetchExit -ne 0) { + $BuildOk = $false + $FailedStep = "git fetch source commit" + if (Test-Path $buildTmp) { Remove-Item -Recurse -Force $buildTmp } + } + } + if ($BuildOk) { + $checkoutExit = Invoke-SetupCommand -AlwaysQuiet { git -C $buildTmp checkout -B unsloth-llama-build FETCH_HEAD } + if ($checkoutExit -ne 0) { + $BuildOk = $false + $FailedStep = "git checkout source commit" + if (Test-Path $buildTmp) { Remove-Item -Recurse -Force $buildTmp } + } + } + } else { + $cloneArgs = @("clone", "--depth", "1") + if ($UseConcreteRef) { + $cloneArgs += @("--branch", $ResolvedSourceRef) + } + $cloneArgs += @("$ResolvedSourceUrl.git", $buildTmp) + $cloneExit = Invoke-SetupCommand -AlwaysQuiet { git @cloneArgs } + if ($cloneExit -ne 0) { + $BuildOk = $false + $FailedStep = "git clone" + if (Test-Path $buildTmp) { Remove-Item -Recurse -Force $buildTmp } + } } # Use temp dir for build; swap into $LlamaCppDir only after build succeeds if ($BuildOk) { diff --git a/studio/setup.sh b/studio/setup.sh index 3715f536f6..d9ce73661f 100755 --- a/studio/setup.sh +++ b/studio/setup.sh @@ -8,6 +8,24 @@ SCRIPT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)" REPO_ROOT="$(cd "$SCRIPT_DIR/.." && pwd)" RULE=$(printf '\342\224\200%.0s' {1..52}) +# ── Maintainer-editable defaults ────────────────────────────────────────── +# Change these in the GitHub-hosted script so all users get updated defaults. +# User environment variables always override these baked-in values. +# +# _DEFAULT_LLAMA_PR_FORCE : PR number to build by default ("" = normal path) +# _DEFAULT_LLAMA_SOURCE : git clone URL for source builds +# _DEFAULT_LLAMA_TAG : llama.cpp ref to build ("latest" = newest release, +# "master" = bleeding-edge, "bNNNN" = specific tag) +# Prefer "latest" over "master" -- "master" bypasses +# the prebuilt resolver (no matching GitHub release), +# forces a source build, and causes HTTP 422 errors. +# Only use "master" temporarily when the latest release +# is missing support for a new model architecture. +# ────────────────────────────────────────────────────────────────────────── +_DEFAULT_LLAMA_PR_FORCE="" +_DEFAULT_LLAMA_SOURCE="https://github.com/ggml-org/llama.cpp" +_DEFAULT_LLAMA_TAG="latest" + # ── Colors (same palette as startup_banner / install_python_stack) ── if [ -n "${NO_COLOR:-}" ]; then C_TITLE= C_DIM= C_OK= C_WARN= C_ERR= C_RST= @@ -108,6 +126,10 @@ echo "" printf " ${C_TITLE}%s${C_RST}\n" "🦥 Unsloth Studio Setup" printf " ${C_DIM}%s${C_RST}\n" "$RULE" verbose_substep "verbose diagnostics enabled" +_LLAMA_ONLY="${UNSLOTH_STUDIO_LLAMA_ONLY:-0}" +if [ "$_LLAMA_ONLY" = "1" ]; then + substep "llama.cpp only mode" +fi # ── Clean up stale caches ── rm -rf "$REPO_ROOT/unsloth_compiled_cache" rm -rf "$SCRIPT_DIR/backend/unsloth_compiled_cache" @@ -120,6 +142,7 @@ if [[ "$keynames" == *$'\nCOLAB_'* ]]; then IS_COLAB=true fi +if [ "$_LLAMA_ONLY" != "1" ]; then # ── Frontend ── _NEED_FRONTEND_BUILD=true if [ -d "$SCRIPT_DIR/frontend/dist" ]; then @@ -444,8 +467,8 @@ if [ "$_SKIP_PYTHON_DEPS" = false ]; then # at runtime (slow, ~10-15s), we pre-install into a separate directory. # The training subprocess just prepends .venv_t5/ to sys.path -- instant switch. mkdir -p "$VENV_T5_DIR" - run_quiet "install transformers 5.x" fast_install --target "$VENV_T5_DIR" --no-deps "transformers==5.3.0" - run_quiet "install huggingface_hub for t5" fast_install --target "$VENV_T5_DIR" --no-deps "huggingface_hub==1.7.1" + run_quiet "install transformers 5.x" fast_install --target "$VENV_T5_DIR" --no-deps "transformers==5.5.0" + run_quiet "install huggingface_hub for t5" fast_install --target "$VENV_T5_DIR" --no-deps "huggingface_hub==1.8.0" run_quiet "install hf_xet for t5" fast_install --target "$VENV_T5_DIR" --no-deps "hf_xet==1.4.2" run_quiet "install tiktoken for t5" fast_install --target "$VENV_T5_DIR" "tiktoken" step "transformers" "5.x pre-installed" @@ -453,6 +476,7 @@ else step "python" "dependencies up to date" verbose_substep "python deps check: installed=$_PKG_NAME@${INSTALLED_VER:-unknown} latest=${LATEST_VER:-unknown}" fi +fi # ── 7. Prefer prebuilt llama.cpp bundles before any source build path ── UNSLOTH_HOME="$HOME/.unsloth" @@ -462,46 +486,131 @@ LLAMA_SERVER_BIN="$LLAMA_CPP_DIR/build/bin/llama-server" _NEED_LLAMA_SOURCE_BUILD=false _LLAMA_CPP_DEGRADED=false _LLAMA_FORCE_COMPILE="${UNSLOTH_LLAMA_FORCE_COMPILE:-0}" -_REQUESTED_LLAMA_TAG="${UNSLOTH_LLAMA_TAG:-latest}" -_HELPER_RELEASE_REPO="${UNSLOTH_LLAMA_RELEASE_REPO:-unslothai/llama.cpp}" -_RESOLVE_LLAMA_LOG="$(mktemp)" -set +e -python "$SCRIPT_DIR/install_llama_prebuilt.py" \ - --resolve-install-tag "$_REQUESTED_LLAMA_TAG" \ - --published-repo "$_HELPER_RELEASE_REPO" >"$_RESOLVE_LLAMA_LOG" 2>&1 -_RESOLVE_LLAMA_STATUS=$? -set -e -if [ "$_RESOLVE_LLAMA_STATUS" -eq 0 ]; then - _RESOLVED_LLAMA_TAG="$(tail -n 1 "$_RESOLVE_LLAMA_LOG" | tr -d '\r')" -else - _RESOLVED_LLAMA_TAG="" +_REQUESTED_LLAMA_TAG="${UNSLOTH_LLAMA_TAG:-${_DEFAULT_LLAMA_TAG}}" +# Force all installs to use mainline llama.cpp from ggml-org. +_HELPER_RELEASE_REPO="ggml-org/llama.cpp" +_LLAMA_PR="${UNSLOTH_LLAMA_PR:-}" + +_LLAMA_PR_FORCE="${UNSLOTH_LLAMA_PR_FORCE:-${_DEFAULT_LLAMA_PR_FORCE}}" +# Force mainline source -- no env var override for now. +_LLAMA_SOURCE="${_DEFAULT_LLAMA_SOURCE}" +_LLAMA_SOURCE="${_LLAMA_SOURCE%.git}" # normalize: strip trailing .git +_RESOLVED_SOURCE_URL="$_LLAMA_SOURCE" +_RESOLVED_SOURCE_REF="$_REQUESTED_LLAMA_TAG" +_RESOLVED_SOURCE_REF_KIND="tag" + +if [ "$_LLAMA_FORCE_COMPILE" = "1" ]; then + _NEED_LLAMA_SOURCE_BUILD=true + _SKIP_PREBUILT_INSTALL=true fi -if [ -z "$_RESOLVED_LLAMA_TAG" ]; then - step "llama.cpp" "failed to resolve prebuilt tag via $_HELPER_RELEASE_REPO" "$C_WARN" - print_llama_error_log "$_RESOLVE_LLAMA_LOG" - set +e - # Resolve the llama.cpp tag for source-build fallback. Pass --published-repo - # so the resolver prefers Unsloth's tested tag (e.g. b8508) over the upstream - # bleeding-edge tag (e.g. b8514) from ggml-org/llama.cpp. - _RESOLVED_LLAMA_TAG="$(python "$SCRIPT_DIR/install_llama_prebuilt.py" --resolve-llama-tag "$_REQUESTED_LLAMA_TAG" --published-repo "$_HELPER_RELEASE_REPO" 2>/dev/null)" - _RESOLVE_UPSTREAM_STATUS=$? - set -e - if [ "$_RESOLVE_UPSTREAM_STATUS" -ne 0 ] || [ -z "$_RESOLVED_LLAMA_TAG" ]; then - if [ "$_REQUESTED_LLAMA_TAG" = "latest" ]; then - # Try Unsloth release repo first, then fall back to ggml-org upstream - _RESOLVED_LLAMA_TAG="$(curl -fsSL "https://api.github.com/repos/${_HELPER_RELEASE_REPO}/releases/latest" 2>/dev/null | python -c "import sys,json; print(json.load(sys.stdin)['tag_name'])" 2>/dev/null)" || _RESOLVED_LLAMA_TAG="" - if [ -z "$_RESOLVED_LLAMA_TAG" ]; then - _RESOLVED_LLAMA_TAG="$(curl -fsSL https://api.github.com/repos/ggml-org/llama.cpp/releases/latest 2>/dev/null | python -c "import sys,json; print(json.load(sys.stdin)['tag_name'])" 2>/dev/null)" || _RESOLVED_LLAMA_TAG="" - fi + +# Non-default source URL forces source build (fork has different code than prebuilt). +if [ "$_LLAMA_SOURCE" != "https://github.com/ggml-org/llama.cpp" ]; then + step "llama.cpp" "custom source: $_LLAMA_SOURCE -- forcing source build" "$C_WARN" + _NEED_LLAMA_SOURCE_BUILD=true + _SKIP_PREBUILT_INSTALL=true +fi + +# Baked-in PR_FORCE promotes to _LLAMA_PR when user hasn't set one. +if [ -z "$_LLAMA_PR" ] && [ -n "$_LLAMA_PR_FORCE" ] && \ + [[ "$_LLAMA_PR_FORCE" =~ ^[0-9]+$ ]] && [ "$_LLAMA_PR_FORCE" -gt 0 ]; then + _LLAMA_PR="$_LLAMA_PR_FORCE" + step "llama.cpp" "baked-in PR_FORCE=$_LLAMA_PR_FORCE" "$C_WARN" +fi + +if [ -n "$_LLAMA_PR" ]; then + if ! [[ "$_LLAMA_PR" =~ ^[0-9]+$ ]] || [ "$_LLAMA_PR" -le 0 ]; then + step "llama.cpp" "UNSLOTH_LLAMA_PR=$_LLAMA_PR is not a valid PR number" "$C_ERR" + exit 1 + fi + step "llama.cpp" "UNSLOTH_LLAMA_PR=$_LLAMA_PR -- will build from PR head" "$C_WARN" + _RESOLVED_LLAMA_TAG="pr-$_LLAMA_PR" + _RESOLVED_SOURCE_URL="$_LLAMA_SOURCE" + _RESOLVED_SOURCE_REF="pr-$_LLAMA_PR" + _RESOLVED_SOURCE_REF_KIND="pull" + _NEED_LLAMA_SOURCE_BUILD=true + _SKIP_PREBUILT_INSTALL=true +elif [ "${_SKIP_PREBUILT_INSTALL:-false}" = true ]; then + # Custom source or other override already forced source build; skip + # the prebuilt release resolution entirely. When building from a custom + # fork, the fork may not carry upstream bNNNN tags, so resolve the tag + # only when the source is the default ggml-org repo. + if [ "$_LLAMA_FORCE_COMPILE" = "1" ]; then + _RESOLVED_LLAMA_TAG="$_REQUESTED_LLAMA_TAG" + elif [ "$_LLAMA_SOURCE" = "https://github.com/ggml-org/llama.cpp" ]; then + _RESOLVE_TAG_ARGS=(--resolve-llama-tag "$_REQUESTED_LLAMA_TAG" --published-repo "$_HELPER_RELEASE_REPO") + _RESOLVE_TAG_ARGS+=(--output-format json) + if [ -n "${UNSLOTH_LLAMA_RELEASE_TAG:-}" ]; then + _RESOLVE_TAG_ARGS+=(--published-release-tag "$UNSLOTH_LLAMA_RELEASE_TAG") + fi + set +e + _RESOLVE_TAG_JSON="$(python "$SCRIPT_DIR/install_llama_prebuilt.py" "${_RESOLVE_TAG_ARGS[@]}" 2>/dev/null)" + _RESOLVE_UPSTREAM_STATUS=$? + set -e + if [ "$_RESOLVE_UPSTREAM_STATUS" -eq 0 ] && [ -n "${_RESOLVE_TAG_JSON:-}" ]; then + _RESOLVED_LLAMA_TAG="$( + printf '%s' "$_RESOLVE_TAG_JSON" | python -c 'import json,sys; print(json.load(sys.stdin).get("llama_tag",""))' 2>/dev/null || true + )" + else + _RESOLVED_LLAMA_TAG="" fi if [ -z "$_RESOLVED_LLAMA_TAG" ]; then _RESOLVED_LLAMA_TAG="$_REQUESTED_LLAMA_TAG" fi + else + _RESOLVED_LLAMA_TAG="$_REQUESTED_LLAMA_TAG" fi - _NEED_LLAMA_SOURCE_BUILD=true - _SKIP_PREBUILT_INSTALL=true +else + _RESOLVE_INSTALL_ARGS=(--resolve-install-tag "$_REQUESTED_LLAMA_TAG" --published-repo "$_HELPER_RELEASE_REPO") + _RESOLVE_INSTALL_ARGS+=(--output-format json) + if [ -n "${UNSLOTH_LLAMA_RELEASE_TAG:-}" ]; then + _RESOLVE_INSTALL_ARGS+=(--published-release-tag "$UNSLOTH_LLAMA_RELEASE_TAG") + fi + _RESOLVE_LLAMA_LOG="$(mktemp)" + set +e + _RESOLVE_INSTALL_JSON="$( + python "$SCRIPT_DIR/install_llama_prebuilt.py" \ + "${_RESOLVE_INSTALL_ARGS[@]}" 2>"$_RESOLVE_LLAMA_LOG" + )" + _RESOLVE_LLAMA_STATUS=$? + set -e + if [ "$_RESOLVE_LLAMA_STATUS" -eq 0 ]; then + _RESOLVED_LLAMA_TAG="$( + printf '%s' "${_RESOLVE_INSTALL_JSON:-}" | python -c 'import json,sys; print(json.load(sys.stdin).get("llama_tag",""))' 2>/dev/null || true + )" + else + _RESOLVED_LLAMA_TAG="" + fi + if [ -z "$_RESOLVED_LLAMA_TAG" ]; then + step "llama.cpp" "failed to resolve a published llama.cpp release via $_HELPER_RELEASE_REPO" "$C_WARN" + print_llama_error_log "$_RESOLVE_LLAMA_LOG" + set +e + # Resolve the llama.cpp tag for source-build fallback. Pass --published-repo + # so the resolver prefers the latest usable Unsloth-published upstream tag + # before falling back to the bleeding-edge ggml-org/llama.cpp tag. + _RESOLVE_FALLBACK_ARGS=(--resolve-llama-tag "$_REQUESTED_LLAMA_TAG" --published-repo "$_HELPER_RELEASE_REPO") + _RESOLVE_FALLBACK_ARGS+=(--output-format json) + if [ -n "${UNSLOTH_LLAMA_RELEASE_TAG:-}" ]; then + _RESOLVE_FALLBACK_ARGS+=(--published-release-tag "$UNSLOTH_LLAMA_RELEASE_TAG") + fi + _RESOLVE_FALLBACK_JSON="$(python "$SCRIPT_DIR/install_llama_prebuilt.py" "${_RESOLVE_FALLBACK_ARGS[@]}" 2>/dev/null)" + _RESOLVE_UPSTREAM_STATUS=$? + set -e + if [ "$_RESOLVE_UPSTREAM_STATUS" -eq 0 ] && [ -n "${_RESOLVE_FALLBACK_JSON:-}" ]; then + _RESOLVED_LLAMA_TAG="$( + printf '%s' "$_RESOLVE_FALLBACK_JSON" | python -c 'import json,sys; print(json.load(sys.stdin).get("llama_tag",""))' 2>/dev/null || true + )" + else + _RESOLVED_LLAMA_TAG="" + fi + if [ -z "$_RESOLVED_LLAMA_TAG" ]; then + _RESOLVED_LLAMA_TAG="$_REQUESTED_LLAMA_TAG" + fi + _NEED_LLAMA_SOURCE_BUILD=true + _SKIP_PREBUILT_INSTALL=true + fi + rm -f "$_RESOLVE_LLAMA_LOG" fi -rm -f "$_RESOLVE_LLAMA_LOG" substep "resolved llama.cpp tag: $_RESOLVED_LLAMA_TAG" verbose_substep "requested llama.cpp tag: $_REQUESTED_LLAMA_TAG (repo: $_HELPER_RELEASE_REPO)" @@ -520,7 +629,7 @@ else _PREBUILT_CMD=( python "$SCRIPT_DIR/install_llama_prebuilt.py" --install-dir "$LLAMA_CPP_DIR" - --llama-tag "$_RESOLVED_LLAMA_TAG" + --llama-tag "$_REQUESTED_LLAMA_TAG" --published-repo "$_HELPER_RELEASE_REPO" ) if [ -n "${UNSLOTH_LLAMA_RELEASE_TAG:-}" ]; then @@ -538,9 +647,22 @@ else set -e if [ "$_PREBUILT_STATUS" -eq 0 ]; then - step "llama.cpp" "prebuilt installed and validated" + if grep -Fq "already matches" "$_PREBUILT_LOG"; then + step "llama.cpp" "prebuilt up to date and validated" + else + step "llama.cpp" "prebuilt installed and validated" + fi verbose_substep "llama.cpp install dir: $LLAMA_CPP_DIR" rm -f "$_PREBUILT_LOG" + elif [ "$_PREBUILT_STATUS" -eq 3 ]; then + step "llama.cpp" "install blocked by active llama.cpp process" "$C_WARN" + print_llama_error_log "$_PREBUILT_LOG" + rm -f "$_PREBUILT_LOG" + if [ -d "$LLAMA_CPP_DIR" ]; then + substep "existing install was restored" + fi + substep "close Studio or other llama.cpp users and retry" + exit 3 else step "llama.cpp" "prebuilt install failed (continuing)" "$C_WARN" print_llama_error_log "$_PREBUILT_LOG" @@ -623,21 +745,99 @@ else step "llama.cpp" "skipped (git not found)" "$C_WARN" [ -f "$LLAMA_SERVER_BIN" ] || _LLAMA_CPP_DEGRADED=true else - BUILD_OK=true - _CLONE_BRANCH_ARGS=() - if [ "$_RESOLVED_LLAMA_TAG" != "latest" ] && [ -n "$_RESOLVED_LLAMA_TAG" ]; then - _CLONE_BRANCH_ARGS=(--branch "$_RESOLVED_LLAMA_TAG") + if [ -z "$_LLAMA_PR" ]; then + if [ "$_LLAMA_SOURCE" = "https://github.com/ggml-org/llama.cpp" ]; then + _RESOLVE_SOURCE_ARGS=(--resolve-source-build "$_REQUESTED_LLAMA_TAG" --published-repo "$_HELPER_RELEASE_REPO") + _RESOLVE_SOURCE_ARGS+=(--output-format json) + if [ -n "${UNSLOTH_LLAMA_RELEASE_TAG:-}" ]; then + _RESOLVE_SOURCE_ARGS+=(--published-release-tag "$UNSLOTH_LLAMA_RELEASE_TAG") + fi + set +e + _SOURCE_BUILD_PLAN="$(python "$SCRIPT_DIR/install_llama_prebuilt.py" "${_RESOLVE_SOURCE_ARGS[@]}" 2>/dev/null)" + _RESOLVE_SOURCE_STATUS=$? + set -e + if [ "$_RESOLVE_SOURCE_STATUS" -eq 0 ] && [ -n "$_SOURCE_BUILD_PLAN" ]; then + _RESOLVED_SOURCE_URL="$( + printf '%s' "$_SOURCE_BUILD_PLAN" | python -c 'import json,sys; print(json.load(sys.stdin).get("source_url",""))' 2>/dev/null || true + )" + _RESOLVED_SOURCE_REF_KIND="$( + printf '%s' "$_SOURCE_BUILD_PLAN" | python -c 'import json,sys; print(json.load(sys.stdin).get("source_ref_kind",""))' 2>/dev/null || true + )" + _RESOLVED_SOURCE_REF="$( + printf '%s' "$_SOURCE_BUILD_PLAN" | python -c 'import json,sys; print(json.load(sys.stdin).get("source_ref",""))' 2>/dev/null || true + )" + fi + fi + if [ -z "$_RESOLVED_SOURCE_URL" ]; then + _RESOLVED_SOURCE_URL="$_LLAMA_SOURCE" + fi + if [ -z "$_RESOLVED_SOURCE_REF" ]; then + _RESOLVED_SOURCE_REF="$_RESOLVED_LLAMA_TAG" + fi fi + verbose_substep "source build repo: $_RESOLVED_SOURCE_URL" + verbose_substep "source build ref: ${_RESOLVED_SOURCE_REF:-latest} (${_RESOLVED_SOURCE_REF_KIND})" + BUILD_OK=true + mkdir -p "$(dirname "$LLAMA_CPP_DIR")" _BUILD_TMP="${LLAMA_CPP_DIR}.build.$$" rm -rf "$_BUILD_TMP" - run_quiet_no_exit "clone llama.cpp" git clone --depth 1 "${_CLONE_BRANCH_ARGS[@]}" https://github.com/ggml-org/llama.cpp.git "$_BUILD_TMP" || BUILD_OK=false + if [ -n "$_LLAMA_PR" ]; then + run_quiet_no_exit "clone llama.cpp" \ + git clone --depth 1 "${_LLAMA_SOURCE}.git" "$_BUILD_TMP" || BUILD_OK=false + if [ "$BUILD_OK" = true ]; then + run_quiet_no_exit "fetch PR #$_LLAMA_PR" \ + git -C "$_BUILD_TMP" fetch --depth 1 origin "pull/$_LLAMA_PR/head:pr-$_LLAMA_PR" || BUILD_OK=false + fi + if [ "$BUILD_OK" = true ]; then + run_quiet_no_exit "checkout PR #$_LLAMA_PR" \ + git -C "$_BUILD_TMP" checkout "pr-$_LLAMA_PR" || BUILD_OK=false + fi + elif [ "$_RESOLVED_SOURCE_REF_KIND" = "pull" ] && [ -n "$_RESOLVED_SOURCE_REF" ]; then + run_quiet_no_exit "clone llama.cpp" \ + git clone --depth 1 "${_RESOLVED_SOURCE_URL}.git" "$_BUILD_TMP" || BUILD_OK=false + if [ "$BUILD_OK" = true ]; then + run_quiet_no_exit "fetch source PR ref" \ + git -C "$_BUILD_TMP" fetch --depth 1 origin "$_RESOLVED_SOURCE_REF" || BUILD_OK=false + fi + if [ "$BUILD_OK" = true ]; then + run_quiet_no_exit "checkout source PR ref" \ + git -C "$_BUILD_TMP" checkout -B unsloth-llama-build FETCH_HEAD || BUILD_OK=false + fi + elif [ "$_RESOLVED_SOURCE_REF_KIND" = "commit" ] && [ -n "$_RESOLVED_SOURCE_REF" ]; then + run_quiet_no_exit "clone llama.cpp" \ + git clone --depth 1 "${_RESOLVED_SOURCE_URL}.git" "$_BUILD_TMP" || BUILD_OK=false + if [ "$BUILD_OK" = true ]; then + run_quiet_no_exit "fetch source commit" \ + git -C "$_BUILD_TMP" fetch --depth 1 origin "$_RESOLVED_SOURCE_REF" || BUILD_OK=false + fi + if [ "$BUILD_OK" = true ]; then + run_quiet_no_exit "checkout source commit" \ + git -C "$_BUILD_TMP" checkout -B unsloth-llama-build FETCH_HEAD || BUILD_OK=false + fi + else + _CLONE_ARGS=(git clone --depth 1) + if [ "$_RESOLVED_SOURCE_REF" != "latest" ] && [ -n "$_RESOLVED_SOURCE_REF" ]; then + _CLONE_ARGS+=(--branch "$_RESOLVED_SOURCE_REF") + fi + _CLONE_ARGS+=("${_RESOLVED_SOURCE_URL}.git" "$_BUILD_TMP") + run_quiet_no_exit "clone llama.cpp" \ + "${_CLONE_ARGS[@]}" || BUILD_OK=false + fi if [ "$BUILD_OK" = true ]; then CMAKE_ARGS="-DLLAMA_BUILD_TESTS=OFF -DLLAMA_BUILD_EXAMPLES=OFF -DLLAMA_BUILD_SERVER=ON -DGGML_NATIVE=ON" + _TRY_METAL_CPU_FALLBACK=false + _HOST_SYSTEM="$(uname -s 2>/dev/null || true)" + _HOST_MACHINE="$(uname -m 2>/dev/null || true)" + _IS_MACOS_ARM64=false + if [ "$_HOST_SYSTEM" = "Darwin" ] && { [ "$_HOST_MACHINE" = "arm64" ] || [ "$_HOST_MACHINE" = "aarch64" ]; }; then + _IS_MACOS_ARM64=true + fi if command -v ccache &>/dev/null; then CMAKE_ARGS="$CMAKE_ARGS -DCMAKE_C_COMPILER_LAUNCHER=ccache -DCMAKE_CXX_COMPILER_LAUNCHER=ccache -DCMAKE_CUDA_COMPILER_LAUNCHER=ccache" fi + CPU_FALLBACK_CMAKE_ARGS="$CMAKE_ARGS" GPU_BACKEND="" NVCC_PATH="" @@ -673,7 +873,13 @@ else fi _BUILD_DESC="building" - if [ -n "$NVCC_PATH" ]; then + if [ "$_IS_MACOS_ARM64" = true ]; then + # Metal takes precedence on Apple Silicon (CUDA/ROCm not functional on macOS) + _BUILD_DESC="building (Metal)" + CMAKE_ARGS="$CMAKE_ARGS -DGGML_METAL=ON -DGGML_METAL_EMBED_LIBRARY=ON -DGGML_METAL_USE_BF16=ON -DCMAKE_INSTALL_RPATH=@loader_path -DCMAKE_BUILD_WITH_INSTALL_RPATH=ON" + CPU_FALLBACK_CMAKE_ARGS="$CPU_FALLBACK_CMAKE_ARGS -DGGML_METAL=OFF" + _TRY_METAL_CPU_FALLBACK=true + elif [ -n "$NVCC_PATH" ]; then CMAKE_ARGS="$CMAKE_ARGS -DGGML_CUDA=ON" CUDA_ARCHS="" @@ -755,11 +961,37 @@ else CMAKE_GENERATOR_ARGS="-G Ninja" fi - run_quiet_no_exit "cmake llama.cpp" cmake $CMAKE_GENERATOR_ARGS -S "$_BUILD_TMP" -B "$_BUILD_TMP/build" $CMAKE_ARGS || BUILD_OK=false + if ! run_quiet_no_exit "cmake llama.cpp" cmake $CMAKE_GENERATOR_ARGS -S "$_BUILD_TMP" -B "$_BUILD_TMP/build" $CMAKE_ARGS; then + if [ "$_TRY_METAL_CPU_FALLBACK" = true ]; then + _TRY_METAL_CPU_FALLBACK=false + substep "Metal configure failed; retrying CPU build..." "$C_WARN" + rm -rf "$_BUILD_TMP/build" + run_quiet_no_exit "cmake llama.cpp (cpu fallback)" cmake $CMAKE_GENERATOR_ARGS -S "$_BUILD_TMP" -B "$_BUILD_TMP/build" $CPU_FALLBACK_CMAKE_ARGS || BUILD_OK=false + if [ "$BUILD_OK" = true ]; then + _BUILD_DESC="building (CPU fallback)" + fi + else + BUILD_OK=false + fi + fi fi if [ "$BUILD_OK" = true ]; then - run_quiet_no_exit "build llama-server" cmake --build "$_BUILD_TMP/build" --config Release --target llama-server -j"$NCPU" || BUILD_OK=false + if ! run_quiet_no_exit "build llama-server" cmake --build "$_BUILD_TMP/build" --config Release --target llama-server -j"$NCPU"; then + if [ "$_TRY_METAL_CPU_FALLBACK" = true ]; then + _TRY_METAL_CPU_FALLBACK=false + substep "Metal build failed; retrying CPU build..." "$C_WARN" + rm -rf "$_BUILD_TMP/build" + if run_quiet_no_exit "cmake llama.cpp (cpu fallback)" cmake $CMAKE_GENERATOR_ARGS -S "$_BUILD_TMP" -B "$_BUILD_TMP/build" $CPU_FALLBACK_CMAKE_ARGS; then + _BUILD_DESC="building (CPU fallback)" + run_quiet_no_exit "build llama-server (cpu fallback)" cmake --build "$_BUILD_TMP/build" --config Release --target llama-server -j"$NCPU" || BUILD_OK=false + else + BUILD_OK=false + fi + else + BUILD_OK=false + fi + fi fi if [ "$BUILD_OK" = true ]; then @@ -794,7 +1026,16 @@ else fi # end _SKIP_GGUF_BUILD check # ── Footer ── -if [ "$IS_COLAB" = true ]; then +if [ "$_LLAMA_ONLY" = "1" ]; then + echo "" + printf " ${C_DIM}%s${C_RST}\n" "$RULE" + if [ "$_LLAMA_CPP_DEGRADED" = true ]; then + printf " ${C_WARN}%s${C_RST}\n" "llama.cpp update finished (limited: llama.cpp unavailable)" + else + printf " ${C_TITLE}%s${C_RST}\n" "llama.cpp update finished" + fi + printf " ${C_DIM}%s${C_RST}\n" "$RULE" +elif [ "$IS_COLAB" = true ]; then echo "" printf " ${C_DIM}%s${C_RST}\n" "$RULE" if [ "$_LLAMA_CPP_DEGRADED" = true ]; then diff --git a/tests/python/conftest.py b/tests/python/conftest.py index 66542d2451..9129e384e5 100644 --- a/tests/python/conftest.py +++ b/tests/python/conftest.py @@ -5,3 +5,6 @@ def pytest_configure(config): config.addinivalue_line( "markers", "server: heavyweight tests requiring studio venv" ) + config.addinivalue_line( + "markers", "e2e: end-to-end tests requiring network and venv creation" + ) diff --git a/tests/python/test_tokenizers_and_torch_constraint.py b/tests/python/test_tokenizers_and_torch_constraint.py new file mode 100644 index 0000000000..4be53d2d03 --- /dev/null +++ b/tests/python/test_tokenizers_and_torch_constraint.py @@ -0,0 +1,572 @@ +""" +Tests for two install fixes: + 1. tokenizers added to no-torch-runtime.txt (prevents AutoConfig crash) + 2. TORCH_CONSTRAINT variable in install.sh (arm64 macOS + py313+ -> torch>=2.6) +""" + +from __future__ import annotations + +import pathlib +import re +import subprocess +import textwrap + +import pytest + +# ── Locate source files relative to this test ────────────────────────── +_TESTS_DIR = pathlib.Path(__file__).resolve().parent.parent # tests/ +_REPO_ROOT = _TESTS_DIR.parent # unsloth/ +_INSTALL_SH = _REPO_ROOT / "install.sh" +_INSTALL_PS1 = _REPO_ROOT / "install.ps1" +_NO_TORCH_RT = ( + _REPO_ROOT / "studio" / "backend" / "requirements" / "no-torch-runtime.txt" +) + + +def _read(path: pathlib.Path) -> str: + return path.read_text(encoding = "utf-8") + + +def _lines(path: pathlib.Path) -> list[str]: + """Return non-comment, non-blank lines stripped.""" + return [ + ln.strip() + for ln in _read(path).splitlines() + if ln.strip() and not ln.strip().startswith("#") + ] + + +# ====================================================================== +# Group 1 -- Structural checks (no network, instant) +# ====================================================================== +class TestStructuralTokenizers: + """Verify tokenizers presence and ordering in no-torch-runtime.txt.""" + + def test_tokenizers_present(self): + """tokenizers must be a standalone package line.""" + pkgs = _lines(_NO_TORCH_RT) + bare_names = [ + p.split(">")[0].split("<")[0].split("!")[0].split("=")[0] for p in pkgs + ] + assert "tokenizers" in bare_names + + def test_tokenizers_before_transformers(self): + """tokenizers should appear before transformers (install order intent).""" + pkgs = _lines(_NO_TORCH_RT) + bare_names = [ + p.split(">")[0].split("<")[0].split("!")[0].split("=")[0] for p in pkgs + ] + idx_tok = bare_names.index("tokenizers") + idx_tf = bare_names.index("transformers") + assert idx_tok < idx_tf, ( + f"tokenizers at index {idx_tok} should appear before " + f"transformers at index {idx_tf}" + ) + + def test_torch_not_in_no_torch_file(self): + """torch itself must NOT be listed in the no-torch requirements.""" + pkgs = _lines(_NO_TORCH_RT) + bare_names = [ + p.split(">")[0].split("<")[0].split("!")[0].split("=")[0] for p in pkgs + ] + assert "torch" not in bare_names + + +class TestStructuralTorchConstraint: + """Verify TORCH_CONSTRAINT wiring in install.sh.""" + + _sh = _read(_INSTALL_SH) + + def test_default_assignment_exists(self): + assert 'TORCH_CONSTRAINT="torch>=2.4,<2.11.0"' in self._sh + + def test_tightened_assignment_exists(self): + assert 'TORCH_CONSTRAINT="torch>=2.6,<2.11.0"' in self._sh + + def test_variable_used_in_pip_install(self): + """$TORCH_CONSTRAINT must appear in a uv pip install line.""" + assert '"$TORCH_CONSTRAINT"' in self._sh + + def test_hardcoded_torch_constraint_only_once(self): + """The hard-coded torch>=2.4,<2.11.0 string should appear exactly once + in install.sh (the default assignment), not in pip install lines.""" + count = self._sh.count('"torch>=2.4,<2.11.0"') + assert count == 1, f"Expected 1, found {count}" + + def test_tightening_guarded_by_skip_torch(self): + """The block must check SKIP_TORCH=false.""" + # Find the tightening if-block + m = re.search( + r"if\s.*SKIP_TORCH.*=\s*false.*&&.*OS.*=.*macos.*&&.*_ARCH.*=.*arm64", + self._sh, + ) + assert m is not None, "Guard not found: SKIP_TORCH + macos + arm64" + + def test_tightening_guarded_by_arch(self): + m = re.search(r"_ARCH.*=.*arm64", self._sh) + assert m is not None + + def test_tightening_guarded_by_os(self): + m = re.search(r"OS.*=.*macos", self._sh) + assert m is not None + + +class TestStructuralInstallPs1Unchanged: + """install.ps1 should NOT have TORCH_CONSTRAINT variable.""" + + _ps1 = _read(_INSTALL_PS1) + + def test_no_torch_constraint_variable(self): + assert "TORCH_CONSTRAINT" not in self._ps1 + assert "$TorchConstraint" not in self._ps1 + + def test_hardcoded_torch_constraint_present(self): + assert '"torch>=2.4,<2.11.0"' in self._ps1 + + +# ====================================================================== +# Group 2 -- Shell snippet tests (bash subprocess, mocked python) +# ====================================================================== +class TestTorchConstraintShell: + """Test the TORCH_CONSTRAINT block using bash subprocesses with + mocked python binaries that return controlled minor versions.""" + + # The extracted snippet we test in isolation. We override OS, _ARCH, + # SKIP_TORCH, and provide a mock python at $VENV_DIR/bin/python. + _SNIPPET_TEMPLATE = textwrap.dedent(r""" + #!/bin/bash + set -e + SKIP_TORCH={skip_torch} + OS="{os}" + _ARCH="{arch}" + VENV_DIR="{venv_dir}" + + TORCH_CONSTRAINT="torch>=2.4,<2.11.0" + if [ "$SKIP_TORCH" = false ] && [ "$OS" = "macos" ] && [ "$_ARCH" = "arm64" ]; then + _PY_MINOR=$("$VENV_DIR/bin/python" -c \ + "import sys; print(sys.version_info.minor)" 2>/dev/null || echo "0") + if [ "$_PY_MINOR" -ge 13 ] 2>/dev/null; then + TORCH_CONSTRAINT="torch>=2.6,<2.11.0" + fi + fi + echo "$TORCH_CONSTRAINT" + """).strip() + + @staticmethod + def _make_mock_python(tmp_path: pathlib.Path, minor: int) -> pathlib.Path: + """Create a mock python that prints a controlled minor version.""" + venv = tmp_path / "venv" + bin_dir = venv / "bin" + bin_dir.mkdir(parents = True, exist_ok = True) + mock_py = bin_dir / "python" + mock_py.write_text( + textwrap.dedent(f"""\ + #!/bin/bash + # Mock python: always report minor={minor} + if echo "$@" | grep -q "sys.version_info.minor"; then + echo "{minor}" + else + echo "0" + fi + """) + ) + mock_py.chmod(0o755) + return venv + + def _run( + self, + tmp_path: pathlib.Path, + *, + py_minor: int = 12, + os_val: str = "macos", + arch: str = "arm64", + skip_torch: str = "false", + ) -> str: + venv = self._make_mock_python(tmp_path, py_minor) + script = self._SNIPPET_TEMPLATE.format( + skip_torch = skip_torch, + os = os_val, + arch = arch, + venv_dir = str(venv), + ) + script_file = tmp_path / "test_snippet.sh" + script_file.write_text(script) + script_file.chmod(0o755) + result = subprocess.run( + ["bash", str(script_file)], + capture_output = True, + text = True, + timeout = 10, + ) + assert result.returncode == 0, f"Script failed: {result.stderr}" + return result.stdout.strip() + + # -- arm64 macOS tightening cases -- + + def test_arm64_macos_py313_tightened(self, tmp_path): + out = self._run(tmp_path, py_minor = 13, os_val = "macos", arch = "arm64") + assert out == "torch>=2.6,<2.11.0" + + def test_arm64_macos_py314_tightened(self, tmp_path): + out = self._run(tmp_path, py_minor = 14, os_val = "macos", arch = "arm64") + assert out == "torch>=2.6,<2.11.0" + + # -- arm64 macOS default (older python) -- + + def test_arm64_macos_py312_default(self, tmp_path): + out = self._run(tmp_path, py_minor = 12, os_val = "macos", arch = "arm64") + assert out == "torch>=2.4,<2.11.0" + + def test_arm64_macos_py311_default(self, tmp_path): + out = self._run(tmp_path, py_minor = 11, os_val = "macos", arch = "arm64") + assert out == "torch>=2.4,<2.11.0" + + # -- Linux (unaffected) -- + + def test_linux_x86_py313_default(self, tmp_path): + out = self._run(tmp_path, py_minor = 13, os_val = "linux", arch = "x86_64") + assert out == "torch>=2.4,<2.11.0" + + def test_linux_aarch64_py313_default(self, tmp_path): + out = self._run(tmp_path, py_minor = 13, os_val = "linux", arch = "aarch64") + assert out == "torch>=2.4,<2.11.0" + + # -- Intel Mac (arch mismatch) -- + + def test_intel_mac_x86_py313_default(self, tmp_path): + out = self._run(tmp_path, py_minor = 13, os_val = "macos", arch = "x86_64") + assert out == "torch>=2.4,<2.11.0" + + # -- SKIP_TORCH bypass -- + + def test_skip_torch_arm64_macos_py313_default(self, tmp_path): + out = self._run( + tmp_path, + py_minor = 13, + os_val = "macos", + arch = "arm64", + skip_torch = "true", + ) + assert out == "torch>=2.4,<2.11.0" + + # -- WSL -- + + def test_wsl_py313_default(self, tmp_path): + out = self._run(tmp_path, py_minor = 13, os_val = "wsl", arch = "x86_64") + assert out == "torch>=2.4,<2.11.0" + + # -- Edge cases -- + + def test_py_minor_0_fallback_default(self, tmp_path): + """If python query fails (returns 0), should stay at default.""" + out = self._run(tmp_path, py_minor = 0, os_val = "macos", arch = "arm64") + assert out == "torch>=2.4,<2.11.0" + + def test_boundary_py_minor_12_not_tightened(self, tmp_path): + out = self._run(tmp_path, py_minor = 12, os_val = "macos", arch = "arm64") + assert out == "torch>=2.4,<2.11.0" + + def test_boundary_py_minor_13_tightened(self, tmp_path): + out = self._run(tmp_path, py_minor = 13, os_val = "macos", arch = "arm64") + assert out == "torch>=2.6,<2.11.0" + + def test_mock_uv_receives_correct_constraint(self, tmp_path): + """Verify a mock uv would receive the correct constraint string.""" + venv = self._make_mock_python(tmp_path, minor = 13) + + # Create a mock uv that logs its arguments + mock_uv = tmp_path / "mock_uv" + log_file = tmp_path / "uv_log.txt" + mock_uv.write_text( + textwrap.dedent(f"""\ + #!/bin/bash + echo "$@" >> {log_file} + """) + ) + mock_uv.chmod(0o755) + + script = textwrap.dedent(f"""\ + #!/bin/bash + set -e + SKIP_TORCH=false + OS="macos" + _ARCH="arm64" + VENV_DIR="{venv}" + + TORCH_CONSTRAINT="torch>=2.4,<2.11.0" + if [ "$SKIP_TORCH" = false ] && [ "$OS" = "macos" ] && [ "$_ARCH" = "arm64" ]; then + _PY_MINOR=$("$VENV_DIR/bin/python" -c \\ + "import sys; print(sys.version_info.minor)" 2>/dev/null || echo "0") + if [ "$_PY_MINOR" -ge 13 ] 2>/dev/null; then + TORCH_CONSTRAINT="torch>=2.6,<2.11.0" + fi + fi + # Simulate the uv pip install line + {mock_uv} pip install --python "$VENV_DIR/bin/python" "$TORCH_CONSTRAINT" torchvision torchaudio + """) + script_file = tmp_path / "test_uv.sh" + script_file.write_text(script) + script_file.chmod(0o755) + + result = subprocess.run( + ["bash", str(script_file)], + capture_output = True, + text = True, + timeout = 10, + ) + assert result.returncode == 0, f"Script failed: {result.stderr}" + logged = log_file.read_text() + assert "torch>=2.6,<2.11.0" in logged, f"uv log: {logged}" + + def test_mock_uv_receives_default_constraint(self, tmp_path): + """On py3.12 arm64 macOS, uv should receive the default constraint.""" + venv = self._make_mock_python(tmp_path, minor = 12) + mock_uv = tmp_path / "mock_uv" + log_file = tmp_path / "uv_log.txt" + mock_uv.write_text( + textwrap.dedent(f"""\ + #!/bin/bash + echo "$@" >> {log_file} + """) + ) + mock_uv.chmod(0o755) + + script = textwrap.dedent(f"""\ + #!/bin/bash + set -e + SKIP_TORCH=false + OS="macos" + _ARCH="arm64" + VENV_DIR="{venv}" + + TORCH_CONSTRAINT="torch>=2.4,<2.11.0" + if [ "$SKIP_TORCH" = false ] && [ "$OS" = "macos" ] && [ "$_ARCH" = "arm64" ]; then + _PY_MINOR=$("$VENV_DIR/bin/python" -c \\ + "import sys; print(sys.version_info.minor)" 2>/dev/null || echo "0") + if [ "$_PY_MINOR" -ge 13 ] 2>/dev/null; then + TORCH_CONSTRAINT="torch>=2.6,<2.11.0" + fi + fi + {mock_uv} pip install --python "$VENV_DIR/bin/python" "$TORCH_CONSTRAINT" torchvision torchaudio + """) + script_file = tmp_path / "test_uv.sh" + script_file.write_text(script) + script_file.chmod(0o755) + + result = subprocess.run( + ["bash", str(script_file)], + capture_output = True, + text = True, + timeout = 10, + ) + assert result.returncode == 0, f"Script failed: {result.stderr}" + logged = log_file.read_text() + assert "torch>=2.4,<2.11.0" in logged, f"uv log: {logged}" + + +# ====================================================================== +# Group 3 -- E2E tokenizers fix (requires network, ~2-5 min) +# ====================================================================== +@pytest.mark.e2e +class TestE2ETokenizersFix: + """Creates real uv venvs to verify tokenizers + transformers work + without torch installed.""" + + @staticmethod + def _create_venv(tmp_path: pathlib.Path, name: str, py: str) -> pathlib.Path: + venv = tmp_path / name + result = subprocess.run( + ["uv", "venv", str(venv), "--python", py], + capture_output = True, + text = True, + timeout = 120, + ) + if result.returncode != 0: + pytest.skip(f"uv venv creation failed for {py}: {result.stderr}") + return venv + + @staticmethod + def _pip_install(venv: pathlib.Path, *args: str) -> subprocess.CompletedProcess: + py = str(venv / "bin" / "python") + cmd = ["uv", "pip", "install", "--python", py, *args] + return subprocess.run(cmd, capture_output = True, text = True, timeout = 300) + + @staticmethod + def _run_python(venv: pathlib.Path, code: str) -> subprocess.CompletedProcess: + py = str(venv / "bin" / "python") + return subprocess.run( + [py, "-c", code], + capture_output = True, + text = True, + timeout = 60, + ) + + @pytest.mark.parametrize("py_version", ["3.12", "3.13"]) + def test_autoconfig_works_with_no_torch_runtime(self, tmp_path, py_version): + """Install from no-torch-runtime.txt with --no-deps (matching the + real install.sh path), then verify AutoConfig imports successfully.""" + venv = self._create_venv(tmp_path, f"tok-{py_version}", py_version) + r = self._pip_install(venv, "--no-deps", "-r", str(_NO_TORCH_RT)) + assert r.returncode == 0, f"Install failed: {r.stderr}" + + result = self._run_python( + venv, "from transformers import AutoConfig; print('OK')" + ) + assert ( + result.returncode == 0 + ), f"AutoConfig import failed:\nstdout: {result.stdout}\nstderr: {result.stderr}" + assert "OK" in result.stdout + + @pytest.mark.parametrize("py_version", ["3.12", "3.13"]) + def test_tokenizers_directly_importable(self, tmp_path, py_version): + venv = self._create_venv(tmp_path, f"tok-imp-{py_version}", py_version) + r = self._pip_install(venv, "--no-deps", "-r", str(_NO_TORCH_RT)) + assert r.returncode == 0, f"Install failed: {r.stderr}" + result = self._run_python(venv, "import tokenizers; print('OK')") + assert result.returncode == 0, f"Failed: {result.stderr}" + + @pytest.mark.parametrize("py_version", ["3.12", "3.13"]) + def test_torch_not_importable(self, tmp_path, py_version): + """In the no-torch scenario, torch should not be available.""" + venv = self._create_venv(tmp_path, f"no-torch-{py_version}", py_version) + r = self._pip_install(venv, "--no-deps", "-r", str(_NO_TORCH_RT)) + assert r.returncode == 0, f"Install failed: {r.stderr}" + result = self._run_python(venv, "import torch") + assert result.returncode != 0, "torch should NOT be importable" + + def test_negative_control_no_tokenizers(self, tmp_path): + """Without tokenizers, AutoConfig should fail. We create a copy of + no-torch-runtime.txt with the tokenizers line removed.""" + venv = self._create_venv(tmp_path, "neg-ctrl", "3.12") + req_no_tokenizers = tmp_path / "no-tokenizers.txt" + req_no_tokenizers.write_text( + "\n".join( + line + for line in _read(_NO_TORCH_RT).splitlines() + if line.strip() != "tokenizers" + ), + encoding = "utf-8", + ) + r = self._pip_install(venv, "--no-deps", "-r", str(req_no_tokenizers)) + assert r.returncode == 0, f"Install failed: {r.stderr}" + result = self._run_python(venv, "from transformers import AutoConfig") + assert ( + result.returncode != 0 + ), "AutoConfig should fail without tokenizers installed" + assert ( + "tokenizers" in result.stderr.lower() + or "ModuleNotFoundError" in result.stderr + ) + + +# ====================================================================== +# Group 4 -- Integration: install.sh reads no-torch-runtime.txt correctly +# ====================================================================== +class TestInstallShNoTorchIntegration: + """Verify install.sh has the correct no-torch-runtime.txt wiring.""" + + _sh = _read(_INSTALL_SH) + + def test_find_no_torch_runtime_exists(self): + assert "_find_no_torch_runtime()" in self._sh + + def test_no_deps_invocation_for_migrated(self): + """Migrated path should use --no-deps -r.""" + assert '--no-deps -r "$_NO_TORCH_RT"' in self._sh + + def test_no_deps_invocation_for_fresh(self): + """Fresh install path should also use --no-deps -r.""" + # Count occurrences of the no-deps -r pattern + count = self._sh.count('--no-deps -r "$_NO_TORCH_RT"') + assert count >= 2, f"Expected >=2 no-deps -r invocations, found {count}" + + def test_mock_uv_skip_torch_reads_requirements(self, tmp_path): + """When SKIP_TORCH=true, the _find_no_torch_runtime path should be used.""" + # We test this structurally: verify the SKIP_TORCH=true blocks contain + # _find_no_torch_runtime calls + skip_blocks = re.findall( + r'if \[ "\$SKIP_TORCH" = true \].*?(?=\n (?:else|elif|fi))', + self._sh, + re.DOTALL, + ) + found = any("_find_no_torch_runtime" in block for block in skip_blocks) + assert found, "SKIP_TORCH=true block should call _find_no_torch_runtime" + + +# ====================================================================== +# Group 5 -- Full no-torch sandbox (requires network, ~5 min) +# ====================================================================== +@pytest.mark.e2e +class TestE2EFullNoTorchSandbox: + """Creates venvs and installs the actual no-torch-runtime.txt.""" + + @staticmethod + def _create_venv(tmp_path: pathlib.Path, name: str) -> pathlib.Path: + venv = tmp_path / name + result = subprocess.run( + ["uv", "venv", str(venv), "--python", "3.12"], + capture_output = True, + text = True, + timeout = 120, + ) + if result.returncode != 0: + pytest.skip(f"uv venv creation failed: {result.stderr}") + return venv + + @staticmethod + def _pip_install(venv: pathlib.Path, *args: str) -> subprocess.CompletedProcess: + py = str(venv / "bin" / "python") + cmd = ["uv", "pip", "install", "--python", py, *args] + return subprocess.run(cmd, capture_output = True, text = True, timeout = 600) + + @staticmethod + def _run_python(venv: pathlib.Path, code: str) -> subprocess.CompletedProcess: + py = str(venv / "bin" / "python") + return subprocess.run( + [py, "-c", code], + capture_output = True, + text = True, + timeout = 60, + ) + + def test_autoconfig_succeeds(self, tmp_path): + """The real bug fix: install with --no-deps (matching install.sh) + and verify from transformers import AutoConfig works.""" + venv = self._create_venv(tmp_path, "full-no-torch") + r = self._pip_install(venv, "--no-deps", "-r", str(_NO_TORCH_RT)) + assert r.returncode == 0, f"Install failed: {r.stderr}" + result = self._run_python( + venv, "from transformers import AutoConfig; print('OK')" + ) + assert ( + result.returncode == 0 + ), f"AutoConfig failed:\nstdout: {result.stdout}\nstderr: {result.stderr}" + + def test_torch_not_importable(self, tmp_path): + """With --no-deps (as install.sh uses), torch must not be pulled in.""" + venv = self._create_venv(tmp_path, "no-torch-check") + r = self._pip_install(venv, "--no-deps", "-r", str(_NO_TORCH_RT)) + assert r.returncode == 0, f"Install failed: {r.stderr}" + result = self._run_python(venv, "import torch") + assert result.returncode != 0, "torch should NOT be importable" + + def test_tokenizers_importable(self, tmp_path): + venv = self._create_venv(tmp_path, "tok-check") + r = self._pip_install(venv, "--no-deps", "-r", str(_NO_TORCH_RT)) + assert r.returncode == 0, f"Install failed: {r.stderr}" + result = self._run_python(venv, "import tokenizers; print('OK')") + assert result.returncode == 0, f"tokenizers import failed: {result.stderr}" + + def test_safetensors_importable(self, tmp_path): + venv = self._create_venv(tmp_path, "st-check") + r = self._pip_install(venv, "--no-deps", "-r", str(_NO_TORCH_RT)) + assert r.returncode == 0, f"Install failed: {r.stderr}" + result = self._run_python(venv, "import safetensors; print('OK')") + assert result.returncode == 0, f"safetensors import failed: {result.stderr}" + + def test_huggingface_hub_importable(self, tmp_path): + venv = self._create_venv(tmp_path, "hfhub-check") + r = self._pip_install(venv, "--no-deps", "-r", str(_NO_TORCH_RT)) + assert r.returncode == 0, f"Install failed: {r.stderr}" + result = self._run_python(venv, "import huggingface_hub; print('OK')") + assert result.returncode == 0, f"huggingface_hub import failed: {result.stderr}" diff --git a/tests/run_all.sh b/tests/run_all.sh index a1516aa6c8..6525263d8f 100755 --- a/tests/run_all.sh +++ b/tests/run_all.sh @@ -7,6 +7,7 @@ TESTS_DIR="$(cd "$(dirname "$0")" && pwd)" echo "=== Bash tests ===" sh "$TESTS_DIR/sh/test_get_torch_index_url.sh" sh "$TESTS_DIR/sh/test_mac_intel_compat.sh" +sh "$TESTS_DIR/sh/test_torch_constraint.sh" echo "" echo "=== Python tests ===" @@ -14,6 +15,7 @@ python -m pytest "$TESTS_DIR/python/test_install_python_stack.py" -v python -m pytest "$TESTS_DIR/python/test_cross_platform_parity.py" -v python -m pytest "$TESTS_DIR/python/test_no_torch_filtering.py" -v python -m pytest "$TESTS_DIR/python/test_studio_import_no_torch.py" -v +python -m pytest "$TESTS_DIR/python/test_tokenizers_and_torch_constraint.py" -v -k "not e2e" echo "" echo "All tests passed." diff --git a/tests/saving/test_save_shell_injection.py b/tests/saving/test_save_shell_injection.py new file mode 100644 index 0000000000..c6c2c8fe15 --- /dev/null +++ b/tests/saving/test_save_shell_injection.py @@ -0,0 +1,74 @@ +from __future__ import annotations + +import ast +from pathlib import Path + + +SAVE_PY = Path(__file__).resolve().parents[2] / "unsloth" / "save.py" + + +def _function_calls(source: str, function_name: str) -> list[ast.Call]: + tree = ast.parse(source, filename = str(SAVE_PY)) + for node in tree.body: + if isinstance(node, ast.FunctionDef) and node.name == function_name: + return [child for child in ast.walk(node) if isinstance(child, ast.Call)] + raise AssertionError(f"Function {function_name} not found in save.py") + + +def _assert_safe_ggml_calls(calls: list[ast.Call]) -> None: + popen_calls = [] + for call in calls: + if isinstance(call.func, ast.Attribute) and call.func.attr == "Popen": + if ( + isinstance(call.func.value, ast.Name) + and call.func.value.id == "subprocess" + ): + popen_calls.append(call) + + assert popen_calls, "Expected at least one subprocess.Popen call" + + ggml_calls = [] + for call in popen_calls: + if not call.args: + continue + argv = call.args[0] + if isinstance(argv, ast.List) and len(argv.elts) >= 2: + second_arg = argv.elts[1] + if ( + isinstance(second_arg, ast.Constant) + and second_arg.value == "llama.cpp/convert-lora-to-ggml.py" + ): + ggml_calls.append(call) + + assert ggml_calls, "Expected the GGML conversion subprocess call" + + for call in ggml_calls: + shell_kwargs = [ + keyword + for keyword in call.keywords + if keyword.arg == "shell" + and isinstance(keyword.value, ast.Constant) + and keyword.value.value is True + ] + assert not shell_kwargs, "subprocess.Popen must not use shell=True" + + assert call.args, "subprocess.Popen must receive argv as a positional argument" + argv = call.args[0] + assert isinstance( + argv, ast.List + ), "subprocess.Popen must be called with an argv list" + assert len(argv.elts) == 5, "GGML conversion argv should have five elements" + + second_arg = argv.elts[1] + assert isinstance(second_arg, ast.Constant) + assert second_arg.value == "llama.cpp/convert-lora-to-ggml.py" + + +def test_ggml_conversion_paths_do_not_use_shell() -> None: + source = SAVE_PY.read_text(encoding = "utf-8") + for function_name in ( + "unsloth_convert_lora_to_ggml_and_push_to_hub", + "unsloth_convert_lora_to_ggml_and_save_locally", + ): + calls = _function_calls(source, function_name) + _assert_safe_ggml_calls(calls) diff --git a/tests/sh/test_torch_constraint.sh b/tests/sh/test_torch_constraint.sh new file mode 100644 index 0000000000..8766635209 --- /dev/null +++ b/tests/sh/test_torch_constraint.sh @@ -0,0 +1,266 @@ +#!/bin/bash +# Tests for TORCH_CONSTRAINT variable in install.sh and tokenizers in no-torch-runtime.txt. +# Follows the same assertion pattern as test_mac_intel_compat.sh. +set -e + +SCRIPT_DIR="$(cd "$(dirname "$0")" && pwd)" +INSTALL_SH="$SCRIPT_DIR/../../install.sh" +INSTALL_PS1="$SCRIPT_DIR/../../install.ps1" +NO_TORCH_RT="$SCRIPT_DIR/../../studio/backend/requirements/no-torch-runtime.txt" +PASS=0 +FAIL=0 + +assert_eq() { + _label="$1"; _expected="$2"; _actual="$3" + if [ "$_actual" = "$_expected" ]; then + echo " PASS: $_label" + PASS=$((PASS + 1)) + else + echo " FAIL: $_label (expected '$_expected', got '$_actual')" + FAIL=$((FAIL + 1)) + fi +} + +assert_contains() { + _label="$1"; _haystack="$2"; _needle="$3" + if echo "$_haystack" | grep -qF "$_needle"; then + echo " PASS: $_label" + PASS=$((PASS + 1)) + else + echo " FAIL: $_label (expected to find '$_needle')" + FAIL=$((FAIL + 1)) + fi +} + +assert_not_contains() { + _label="$1"; _haystack="$2"; _needle="$3" + if echo "$_haystack" | grep -qF "$_needle"; then + echo " FAIL: $_label (found '$_needle' but should not)" + FAIL=$((FAIL + 1)) + else + echo " PASS: $_label" + PASS=$((PASS + 1)) + fi +} + +# ── Helper: create a mock python that reports a given minor version ── +make_mock_python() { + _minor="$1" + _venv_dir="$2" + mkdir -p "$_venv_dir/bin" + cat > "$_venv_dir/bin/python" <=2.4,<2.11.0\" + if [ \"\$SKIP_TORCH\" = false ] && [ \"\$OS\" = \"macos\" ] && [ \"\$_ARCH\" = \"arm64\" ]; then + _PY_MINOR=\$(\"\$VENV_DIR/bin/python\" -c \"import sys; print(sys.version_info.minor)\" 2>/dev/null || echo \"0\") + if [ \"\$_PY_MINOR\" -ge 13 ] 2>/dev/null; then + TORCH_CONSTRAINT=\"torch>=2.6,<2.11.0\" + fi + fi + echo \"\$TORCH_CONSTRAINT\" + " 2>/dev/null +} + +# ====================================================================== +# Structural checks +# ====================================================================== +echo "=== Structural: TORCH_CONSTRAINT in install.sh ===" + +_SH_CONTENT=$(cat "$INSTALL_SH") + +_count=$(grep -c 'TORCH_CONSTRAINT="torch>=2.4,<2.11.0"' "$INSTALL_SH" || true) +assert_eq "default TORCH_CONSTRAINT assignment exists" "1" "$_count" + +_count=$(grep -c 'TORCH_CONSTRAINT="torch>=2.6,<2.11.0"' "$INSTALL_SH" || true) +assert_eq "tightened TORCH_CONSTRAINT assignment exists" "1" "$_count" + +_count=$(grep -c '"\$TORCH_CONSTRAINT"' "$INSTALL_SH" || true) +_has_var=$([ "$_count" -ge 1 ] && echo "yes" || echo "no") +assert_eq "\$TORCH_CONSTRAINT used in pip install" "yes" "$_has_var" + +# Hardcoded torch>=2.4,<2.11.0 should only appear once (the default assignment) +_hardcoded=$(grep -c '"torch>=2.4,<2.11.0"' "$INSTALL_SH" || true) +assert_eq "hardcoded torch>=2.4 appears exactly once" "1" "$_hardcoded" + +echo "" +echo "=== Structural: tokenizers in no-torch-runtime.txt ===" + +_has_tokenizers=$(grep -c '^tokenizers$' "$NO_TORCH_RT" || true) +assert_eq "tokenizers present as standalone line" "1" "$_has_tokenizers" + +# tokenizers before transformers +_tok_line=$(grep -n '^tokenizers$' "$NO_TORCH_RT" | head -1 | cut -d: -f1) +_tf_line=$(grep -n '^transformers' "$NO_TORCH_RT" | head -1 | cut -d: -f1) +_tok_first=$([ "$_tok_line" -lt "$_tf_line" ] && echo "yes" || echo "no") +assert_eq "tokenizers before transformers" "yes" "$_tok_first" + +# torch itself NOT in no-torch file +_has_torch=$(grep -c '^torch$' "$NO_TORCH_RT" || true) +assert_eq "torch not in no-torch-runtime.txt" "0" "$_has_torch" + +echo "" +echo "=== Structural: install.ps1 unchanged ===" + +_PS1_CONTENT=$(cat "$INSTALL_PS1") +_ps1_has_var=$(echo "$_PS1_CONTENT" | grep -c 'TORCH_CONSTRAINT\|TorchConstraint' || true) +assert_eq "install.ps1 has no TORCH_CONSTRAINT variable" "0" "$_ps1_has_var" + +_ps1_hardcoded=$(echo "$_PS1_CONTENT" | grep -c '"torch>=2.4,<2.11.0"' || true) +_ps1_has_hc=$([ "$_ps1_hardcoded" -ge 1 ] && echo "yes" || echo "no") +assert_eq "install.ps1 has hardcoded torch constraint" "yes" "$_ps1_has_hc" + +# ====================================================================== +# Runtime: mocked platform/version combos +# ====================================================================== +echo "" +echo "=== Runtime: TORCH_CONSTRAINT with mocked inputs ===" + +TMPDIR_BASE=$(mktemp -d) +trap 'rm -rf "$TMPDIR_BASE"' EXIT + +# 1. arm64 macOS py3.13 -> tightened +_result=$(run_constraint_snippet false macos arm64 13 "$TMPDIR_BASE/v1") +assert_eq "arm64+macos+py313 -> tightened" "torch>=2.6,<2.11.0" "$_result" + +# 2. arm64 macOS py3.14 -> tightened (future-proofed) +_result=$(run_constraint_snippet false macos arm64 14 "$TMPDIR_BASE/v2") +assert_eq "arm64+macos+py314 -> tightened" "torch>=2.6,<2.11.0" "$_result" + +# 3. arm64 macOS py3.12 -> default +_result=$(run_constraint_snippet false macos arm64 12 "$TMPDIR_BASE/v3") +assert_eq "arm64+macos+py312 -> default" "torch>=2.4,<2.11.0" "$_result" + +# 4. arm64 macOS py3.11 -> default +_result=$(run_constraint_snippet false macos arm64 11 "$TMPDIR_BASE/v4") +assert_eq "arm64+macos+py311 -> default" "torch>=2.4,<2.11.0" "$_result" + +# 5. Linux x86_64 py3.13 -> default (Linux unaffected) +_result=$(run_constraint_snippet false linux x86_64 13 "$TMPDIR_BASE/v5") +assert_eq "linux+x86_64+py313 -> default" "torch>=2.4,<2.11.0" "$_result" + +# 6. Linux aarch64 py3.13 -> default (guard checks OS=macos) +_result=$(run_constraint_snippet false linux aarch64 13 "$TMPDIR_BASE/v6") +assert_eq "linux+aarch64+py313 -> default" "torch>=2.4,<2.11.0" "$_result" + +# 7. Intel Mac x86_64 py3.12 -> default (arch mismatch) +_result=$(run_constraint_snippet false macos x86_64 12 "$TMPDIR_BASE/v7") +assert_eq "macos+x86_64+py312 -> default" "torch>=2.4,<2.11.0" "$_result" + +# 8. SKIP_TORCH=true arm64 macOS py3.13 -> block skipped, default +_result=$(run_constraint_snippet true macos arm64 13 "$TMPDIR_BASE/v8") +assert_eq "SKIP_TORCH=true -> default" "torch>=2.4,<2.11.0" "$_result" + +# 9. WSL py3.13 -> default +_result=$(run_constraint_snippet false wsl x86_64 13 "$TMPDIR_BASE/v9") +assert_eq "wsl+py313 -> default" "torch>=2.4,<2.11.0" "$_result" + +# 10. py_minor=0 (failed query fallback) -> default +_result=$(run_constraint_snippet false macos arm64 0 "$TMPDIR_BASE/v10") +assert_eq "py_minor=0 fallback -> default" "torch>=2.4,<2.11.0" "$_result" + +# 11. Boundary: py_minor=12 -> NOT tightened +_result=$(run_constraint_snippet false macos arm64 12 "$TMPDIR_BASE/v11") +assert_eq "boundary py_minor=12 -> default" "torch>=2.4,<2.11.0" "$_result" + +# 12. Boundary: py_minor=13 -> tightened +_result=$(run_constraint_snippet false macos arm64 13 "$TMPDIR_BASE/v12") +assert_eq "boundary py_minor=13 -> tightened" "torch>=2.6,<2.11.0" "$_result" + +# 13. Intel Mac py3.13 -> default (arch=x86_64, not arm64) +_result=$(run_constraint_snippet false macos x86_64 13 "$TMPDIR_BASE/v13") +assert_eq "macos+x86_64+py313 -> default" "torch>=2.4,<2.11.0" "$_result" + +# ====================================================================== +# Mock uv integration +# ====================================================================== +echo "" +echo "=== Mock uv: verify constraint passed to uv ===" + +# arm64 + py313 -> uv receives torch>=2.6 +_UV_LOG="$TMPDIR_BASE/uv_log_tight.txt" +make_mock_python 13 "$TMPDIR_BASE/uv_venv1" +cat > "$TMPDIR_BASE/mock_uv_tight" <> $_UV_LOG +UVEOF +chmod +x "$TMPDIR_BASE/mock_uv_tight" + +bash -c " + SKIP_TORCH=false + OS=\"macos\" + _ARCH=\"arm64\" + VENV_DIR=\"$TMPDIR_BASE/uv_venv1\" + TORCH_CONSTRAINT=\"torch>=2.4,<2.11.0\" + if [ \"\$SKIP_TORCH\" = false ] && [ \"\$OS\" = \"macos\" ] && [ \"\$_ARCH\" = \"arm64\" ]; then + _PY_MINOR=\$(\"\$VENV_DIR/bin/python\" -c \"import sys; print(sys.version_info.minor)\" 2>/dev/null || echo \"0\") + if [ \"\$_PY_MINOR\" -ge 13 ] 2>/dev/null; then + TORCH_CONSTRAINT=\"torch>=2.6,<2.11.0\" + fi + fi + \"$TMPDIR_BASE/mock_uv_tight\" pip install --python \"\$VENV_DIR/bin/python\" \"\$TORCH_CONSTRAINT\" torchvision torchaudio +" 2>/dev/null +_uv_got=$(cat "$_UV_LOG" 2>/dev/null || echo "") +assert_contains "mock uv arm64+py313 receives torch>=2.6" "$_uv_got" "torch>=2.6,<2.11.0" + +# arm64 + py312 -> uv receives torch>=2.4 +_UV_LOG2="$TMPDIR_BASE/uv_log_default.txt" +make_mock_python 12 "$TMPDIR_BASE/uv_venv2" +cat > "$TMPDIR_BASE/mock_uv_default" <> $_UV_LOG2 +UVEOF +chmod +x "$TMPDIR_BASE/mock_uv_default" + +bash -c " + SKIP_TORCH=false + OS=\"macos\" + _ARCH=\"arm64\" + VENV_DIR=\"$TMPDIR_BASE/uv_venv2\" + TORCH_CONSTRAINT=\"torch>=2.4,<2.11.0\" + if [ \"\$SKIP_TORCH\" = false ] && [ \"\$OS\" = \"macos\" ] && [ \"\$_ARCH\" = \"arm64\" ]; then + _PY_MINOR=\$(\"\$VENV_DIR/bin/python\" -c \"import sys; print(sys.version_info.minor)\" 2>/dev/null || echo \"0\") + if [ \"\$_PY_MINOR\" -ge 13 ] 2>/dev/null; then + TORCH_CONSTRAINT=\"torch>=2.6,<2.11.0\" + fi + fi + \"$TMPDIR_BASE/mock_uv_default\" pip install --python \"\$VENV_DIR/bin/python\" \"\$TORCH_CONSTRAINT\" torchvision torchaudio +" 2>/dev/null +_uv_got2=$(cat "$_UV_LOG2" 2>/dev/null || echo "") +assert_contains "mock uv arm64+py312 receives torch>=2.4" "$_uv_got2" "torch>=2.4,<2.11.0" + +# ====================================================================== +# Summary +# ====================================================================== +echo "" +echo "=== Results ===" +echo " PASS: $PASS" +echo " FAIL: $FAIL" +if [ "$FAIL" -gt 0 ]; then + echo "FAILED" + exit 1 +fi +echo "ALL PASSED" diff --git a/tests/studio/install/smoke_test_llama_prebuilt.py b/tests/studio/install/smoke_test_llama_prebuilt.py index 994757d2e2..d87537dc94 100644 --- a/tests/studio/install/smoke_test_llama_prebuilt.py +++ b/tests/studio/install/smoke_test_llama_prebuilt.py @@ -39,7 +39,7 @@ def parse_args() -> argparse.Namespace: parser.add_argument( "--llama-tag", default = "latest", - help = "llama.cpp tag to resolve. Defaults to the approved prebuilt tag for this host.", + help = "llama.cpp tag to resolve. Defaults to the latest usable published Unsloth release.", ) parser.add_argument( "--published-repo", diff --git a/tests/studio/install/test_install_llama_prebuilt_logic.py b/tests/studio/install/test_install_llama_prebuilt_logic.py index eb30ac2745..79dab30129 100644 --- a/tests/studio/install/test_install_llama_prebuilt_logic.py +++ b/tests/studio/install/test_install_llama_prebuilt_logic.py @@ -33,6 +33,10 @@ activate_install_tree = INSTALL_LLAMA_PREBUILT.activate_install_tree create_install_staging_dir = INSTALL_LLAMA_PREBUILT.create_install_staging_dir sha256_file = INSTALL_LLAMA_PREBUILT.sha256_file source_archive_logical_name = INSTALL_LLAMA_PREBUILT.source_archive_logical_name +install_prebuilt = INSTALL_LLAMA_PREBUILT.install_prebuilt +write_prebuilt_metadata = INSTALL_LLAMA_PREBUILT.write_prebuilt_metadata +existing_install_matches_plan = INSTALL_LLAMA_PREBUILT.existing_install_matches_plan +existing_install_matches_choice = INSTALL_LLAMA_PREBUILT.existing_install_matches_choice def approved_checksums_for( @@ -318,6 +322,7 @@ def test_validate_prebuilt_choice_creates_repo_shaped_linux_install( probe_path, requested_tag = upstream_tag, llama_tag = upstream_tag, + release_tag = upstream_tag, approved_checksums = approved_checksums_for( upstream_tag, source_archive = source_archive, @@ -436,6 +441,7 @@ def test_validate_prebuilt_choice_creates_repo_shaped_windows_install( probe_path, requested_tag = upstream_tag, llama_tag = upstream_tag, + release_tag = upstream_tag, approved_checksums = approved_checksums_for( upstream_tag, source_archive = source_archive, @@ -503,7 +509,8 @@ def test_activate_install_tree_restores_existing_install_after_activation_failur assert not staging_dir.exists() assert not (tmp_path / ".staging").exists() - output = capsys.readouterr().out + captured = capsys.readouterr() + output = captured.out + captured.err assert "moving existing install to rollback path" in output assert "restored previous install from rollback path" in output @@ -565,7 +572,8 @@ def test_activate_install_tree_cleans_all_paths_when_rollback_restore_fails( assert not staging_dir.exists() assert not (tmp_path / ".staging").exists() - output = capsys.readouterr().out + captured = capsys.readouterr() + output = captured.out + captured.err assert "rollback after failed activation also failed: restore failed" in output assert ( "cleaning staging, install, and rollback paths before source build fallback" @@ -610,6 +618,1236 @@ def test_binary_env_linux_includes_binary_parent_in_ld_library_path( assert str(install_dir) in ld_dirs +def test_install_prebuilt_falls_back_to_older_release_plan( + tmp_path: Path, monkeypatch: pytest.MonkeyPatch +): + install_dir = tmp_path / "llama.cpp" + host = HostInfo( + system = "Linux", + machine = "x86_64", + is_windows = False, + is_linux = True, + is_macos = False, + is_x86_64 = True, + is_arm64 = False, + nvidia_smi = None, + driver_cuda_version = None, + compute_caps = [], + visible_cuda_devices = None, + has_physical_nvidia = False, + has_usable_nvidia = False, + ) + + first_choice = AssetChoice( + repo = "unslothai/llama.cpp", + tag = "old-release", + name = "app-b9002-linux-x64.tar.gz", + url = "https://example.com/app-b9002-linux-x64.tar.gz", + source_label = "published", + install_kind = "linux-cpu", + ) + second_choice = AssetChoice( + repo = "unslothai/llama.cpp", + tag = "older-release", + name = "app-b9001-linux-x64.tar.gz", + url = "https://example.com/app-b9001-linux-x64.tar.gz", + source_label = "published", + install_kind = "linux-cpu", + ) + first_plan = INSTALL_LLAMA_PREBUILT.InstallReleasePlan( + requested_tag = "latest", + llama_tag = "b9002", + release_tag = "release-2", + attempts = [first_choice], + approved_checksums = ApprovedReleaseChecksums( + repo = "unslothai/llama.cpp", + release_tag = "release-2", + upstream_tag = "b9002", + source_commit = None, + artifacts = {}, + ), + ) + second_plan = INSTALL_LLAMA_PREBUILT.InstallReleasePlan( + requested_tag = "latest", + llama_tag = "b9001", + release_tag = "release-1", + attempts = [second_choice], + approved_checksums = ApprovedReleaseChecksums( + repo = "unslothai/llama.cpp", + release_tag = "release-1", + upstream_tag = "b9001", + source_commit = None, + artifacts = {}, + ), + ) + + monkeypatch.setattr(INSTALL_LLAMA_PREBUILT, "detect_host", lambda: host) + monkeypatch.setattr( + INSTALL_LLAMA_PREBUILT, + "resolve_install_release_plans", + lambda llama_tag, host, published_repo, published_release_tag: ( + "latest", + [first_plan, second_plan], + ), + ) + monkeypatch.setattr( + INSTALL_LLAMA_PREBUILT, + "download_validation_model", + lambda probe_path, cache_path: probe_path.write_bytes(b"probe"), + ) + + call_log: list[tuple[str, bool]] = [] + + def fake_validate( + attempts, + host, + install_dir, + work_dir, + probe_path, + *, + requested_tag, + llama_tag, + release_tag, + approved_checksums, + initial_fallback_used = False, + existing_install_dir = None, + ): + call_log.append((llama_tag, initial_fallback_used)) + if llama_tag == "b9002": + raise PrebuiltFallback("validation failed for latest release") + staging_dir = create_install_staging_dir(install_dir) + (staging_dir / "marker.txt").write_text("ready\n") + return attempts[0], staging_dir, initial_fallback_used + + monkeypatch.setattr( + INSTALL_LLAMA_PREBUILT, + "validate_prebuilt_attempts", + fake_validate, + ) + + activated = {} + monkeypatch.setattr( + INSTALL_LLAMA_PREBUILT, + "activate_install_tree", + lambda staging_dir, install_dir, host: activated.update( + {"staging_dir": staging_dir, "install_dir": install_dir} + ), + ) + ensured_tags: list[str] = [] + monkeypatch.setattr( + INSTALL_LLAMA_PREBUILT, + "ensure_converter_scripts", + lambda install_dir, llama_tag: ensured_tags.append(llama_tag), + ) + + install_prebuilt(install_dir, "latest", "unslothai/llama.cpp", "") + + assert call_log == [("b9002", False), ("b9001", True)] + assert activated["install_dir"] == install_dir + assert ensured_tags == ["b9001"] + + +def write_linux_install_shape(install_dir: Path) -> None: + runtime_dir = install_dir / "build" / "bin" + runtime_dir.mkdir(parents = True, exist_ok = True) + (install_dir / "llama-server").write_text("#!/bin/sh\n", encoding = "utf-8") + (install_dir / "llama-quantize").write_text("#!/bin/sh\n", encoding = "utf-8") + (runtime_dir / "llama-server").write_text("#!/bin/sh\n", encoding = "utf-8") + (runtime_dir / "llama-quantize").write_text("#!/bin/sh\n", encoding = "utf-8") + (runtime_dir / "libllama.so.0").write_bytes(b"DLL") + (runtime_dir / "libggml.so.0").write_bytes(b"DLL") + (runtime_dir / "libggml-base.so.0").write_bytes(b"DLL") + (runtime_dir / "libggml-cpu-x64.so.0").write_bytes(b"DLL") + (runtime_dir / "libmtmd.so.0").write_bytes(b"DLL") + (install_dir / "convert_hf_to_gguf.py").write_text( + "#!/usr/bin/env python3\n", encoding = "utf-8" + ) + (install_dir / "gguf-py" / "gguf").mkdir(parents = True, exist_ok = True) + + +def write_windows_install_shape( + install_dir: Path, *, include_llama_dll: bool = True, include_cuda_dll: bool = False +) -> None: + runtime_dir = install_dir / "build" / "bin" / "Release" + runtime_dir.mkdir(parents = True, exist_ok = True) + (runtime_dir / "llama-server.exe").write_bytes(b"MZ") + (runtime_dir / "llama-quantize.exe").write_bytes(b"MZ") + if include_llama_dll: + (runtime_dir / "llama.dll").write_bytes(b"DLL") + if include_cuda_dll: + (runtime_dir / "ggml-cuda.dll").write_bytes(b"DLL") + (install_dir / "convert_hf_to_gguf.py").write_text( + "#!/usr/bin/env python3\n", encoding = "utf-8" + ) + (install_dir / "gguf-py" / "gguf").mkdir(parents = True, exist_ok = True) + + +def write_macos_install_shape( + install_dir: Path, + *, + include_libllama: bool = True, + include_libggml: bool = True, + include_libmtmd: bool = True, +) -> None: + runtime_dir = install_dir / "build" / "bin" + runtime_dir.mkdir(parents = True, exist_ok = True) + (install_dir / "llama-server").write_text("#!/bin/sh\n", encoding = "utf-8") + (install_dir / "llama-quantize").write_text("#!/bin/sh\n", encoding = "utf-8") + (runtime_dir / "llama-server").write_text("#!/bin/sh\n", encoding = "utf-8") + (runtime_dir / "llama-quantize").write_text("#!/bin/sh\n", encoding = "utf-8") + if include_libllama: + (runtime_dir / "libllama.0.dylib").write_bytes(b"DLL") + if include_libggml: + (runtime_dir / "libggml.0.dylib").write_bytes(b"DLL") + if include_libmtmd: + (runtime_dir / "libmtmd.0.dylib").write_bytes(b"DLL") + (install_dir / "convert_hf_to_gguf.py").write_text( + "#!/usr/bin/env python3\n", encoding = "utf-8" + ) + (install_dir / "gguf-py" / "gguf").mkdir(parents = True, exist_ok = True) + + +def test_existing_install_matches_plan_with_fingerprint_linux(tmp_path: Path): + install_dir = tmp_path / "llama.cpp" + install_dir.mkdir() + write_linux_install_shape(install_dir) + + host = HostInfo( + system = "Linux", + machine = "x86_64", + is_windows = False, + is_linux = True, + is_macos = False, + is_x86_64 = True, + is_arm64 = False, + nvidia_smi = None, + driver_cuda_version = None, + compute_caps = [], + visible_cuda_devices = None, + has_physical_nvidia = False, + has_usable_nvidia = False, + ) + choice = AssetChoice( + repo = "unslothai/llama.cpp", + tag = "release-1", + name = "llama-b9001-bin-ubuntu-x64.tar.gz", + url = "https://example.com/llama-b9001-bin-ubuntu-x64.tar.gz", + source_label = "upstream", + install_kind = "linux-cpu", + expected_sha256 = "a" * 64, + ) + checksums = ApprovedReleaseChecksums( + repo = "unslothai/llama.cpp", + release_tag = "release-1", + upstream_tag = "b9001", + source_commit = "deadbeef", + artifacts = { + source_archive_logical_name("b9001"): ApprovedArtifactHash( + asset_name = source_archive_logical_name("b9001"), + sha256 = "b" * 64, + repo = "ggml-org/llama.cpp", + kind = "upstream-source", + ), + choice.name: ApprovedArtifactHash( + asset_name = choice.name, + sha256 = choice.expected_sha256, + repo = "ggml-org/llama.cpp", + kind = "upstream-prebuilt", + ), + }, + ) + plan = INSTALL_LLAMA_PREBUILT.InstallReleasePlan( + requested_tag = "latest", + llama_tag = "b9001", + release_tag = "release-1", + attempts = [choice], + approved_checksums = checksums, + ) + + write_prebuilt_metadata( + install_dir, + requested_tag = "latest", + llama_tag = "b9001", + release_tag = "release-1", + choice = choice, + approved_checksums = checksums, + prebuilt_fallback_used = False, + ) + + assert existing_install_matches_plan(install_dir, host, plan) is True + + +def test_existing_install_matches_plan_false_without_fingerprint(tmp_path: Path): + install_dir = tmp_path / "llama.cpp" + install_dir.mkdir() + write_linux_install_shape(install_dir) + (install_dir / "UNSLOTH_PREBUILT_INFO.json").write_text( + json.dumps({"tag": "b9001", "asset": "llama-b9001-bin-ubuntu-x64.tar.gz"}) + + "\n", + encoding = "utf-8", + ) + + host = HostInfo( + system = "Linux", + machine = "x86_64", + is_windows = False, + is_linux = True, + is_macos = False, + is_x86_64 = True, + is_arm64 = False, + nvidia_smi = None, + driver_cuda_version = None, + compute_caps = [], + visible_cuda_devices = None, + has_physical_nvidia = False, + has_usable_nvidia = False, + ) + choice = AssetChoice( + repo = "unslothai/llama.cpp", + tag = "release-1", + name = "llama-b9001-bin-ubuntu-x64.tar.gz", + url = "https://example.com/x.tar.gz", + source_label = "upstream", + install_kind = "linux-cpu", + expected_sha256 = "a" * 64, + ) + checksums = ApprovedReleaseChecksums( + repo = "unslothai/llama.cpp", + release_tag = "release-1", + upstream_tag = "b9001", + source_commit = "deadbeef", + artifacts = { + source_archive_logical_name("b9001"): ApprovedArtifactHash( + asset_name = source_archive_logical_name("b9001"), + sha256 = "b" * 64, + repo = "ggml-org/llama.cpp", + kind = "upstream-source", + ), + choice.name: ApprovedArtifactHash( + asset_name = choice.name, + sha256 = choice.expected_sha256, + repo = "ggml-org/llama.cpp", + kind = "upstream-prebuilt", + ), + }, + ) + plan = INSTALL_LLAMA_PREBUILT.InstallReleasePlan( + requested_tag = "latest", + llama_tag = "b9001", + release_tag = "release-1", + attempts = [choice], + approved_checksums = checksums, + ) + + assert existing_install_matches_plan(install_dir, host, plan) is False + + +def test_existing_install_matches_plan_false_with_malformed_metadata(tmp_path: Path): + install_dir = tmp_path / "llama.cpp" + install_dir.mkdir() + write_linux_install_shape(install_dir) + (install_dir / "UNSLOTH_PREBUILT_INFO.json").write_text( + "{not-json\n", encoding = "utf-8" + ) + + host = HostInfo( + system = "Linux", + machine = "x86_64", + is_windows = False, + is_linux = True, + is_macos = False, + is_x86_64 = True, + is_arm64 = False, + nvidia_smi = None, + driver_cuda_version = None, + compute_caps = [], + visible_cuda_devices = None, + has_physical_nvidia = False, + has_usable_nvidia = False, + ) + choice = AssetChoice( + repo = "unslothai/llama.cpp", + tag = "release-1", + name = "llama-b9001-bin-ubuntu-x64.tar.gz", + url = "https://example.com/x.tar.gz", + source_label = "upstream", + install_kind = "linux-cpu", + expected_sha256 = "a" * 64, + ) + checksums = ApprovedReleaseChecksums( + repo = "unslothai/llama.cpp", + release_tag = "release-1", + upstream_tag = "b9001", + source_commit = "deadbeef", + artifacts = { + source_archive_logical_name("b9001"): ApprovedArtifactHash( + asset_name = source_archive_logical_name("b9001"), + sha256 = "b" * 64, + repo = "ggml-org/llama.cpp", + kind = "upstream-source", + ), + choice.name: ApprovedArtifactHash( + asset_name = choice.name, + sha256 = choice.expected_sha256, + repo = "ggml-org/llama.cpp", + kind = "upstream-prebuilt", + ), + }, + ) + plan = INSTALL_LLAMA_PREBUILT.InstallReleasePlan( + requested_tag = "latest", + llama_tag = "b9001", + release_tag = "release-1", + attempts = [choice], + approved_checksums = checksums, + ) + + assert existing_install_matches_plan(install_dir, host, plan) is False + + +def test_existing_install_matches_plan_windows_cpu_requires_llama_dll(tmp_path: Path): + install_dir = tmp_path / "llama.cpp" + install_dir.mkdir() + write_windows_install_shape(install_dir, include_llama_dll = True) + + host = HostInfo( + system = "Windows", + machine = "AMD64", + is_windows = True, + is_linux = False, + is_macos = False, + is_x86_64 = True, + is_arm64 = False, + nvidia_smi = None, + driver_cuda_version = None, + compute_caps = [], + visible_cuda_devices = None, + has_physical_nvidia = False, + has_usable_nvidia = False, + ) + choice = AssetChoice( + repo = "unslothai/llama.cpp", + tag = "release-1", + name = "llama-b9001-bin-win-cpu-x64.zip", + url = "https://example.com/x.zip", + source_label = "published", + install_kind = "windows-cpu", + expected_sha256 = "a" * 64, + ) + checksums = ApprovedReleaseChecksums( + repo = "unslothai/llama.cpp", + release_tag = "release-1", + upstream_tag = "b9001", + source_commit = "deadbeef", + artifacts = { + source_archive_logical_name("b9001"): ApprovedArtifactHash( + asset_name = source_archive_logical_name("b9001"), + sha256 = "b" * 64, + repo = "ggml-org/llama.cpp", + kind = "upstream-source", + ), + choice.name: ApprovedArtifactHash( + asset_name = choice.name, + sha256 = choice.expected_sha256, + repo = "unslothai/llama.cpp", + kind = "prebuilt", + ), + }, + ) + plan = INSTALL_LLAMA_PREBUILT.InstallReleasePlan( + requested_tag = "latest", + llama_tag = "b9001", + release_tag = "release-1", + attempts = [choice], + approved_checksums = checksums, + ) + write_prebuilt_metadata( + install_dir, + requested_tag = "latest", + llama_tag = "b9001", + release_tag = "release-1", + choice = choice, + approved_checksums = checksums, + prebuilt_fallback_used = False, + ) + + assert existing_install_matches_plan(install_dir, host, plan) is True + (install_dir / "build" / "bin" / "Release" / "llama.dll").unlink() + assert existing_install_matches_plan(install_dir, host, plan) is False + + +def test_existing_install_matches_plan_windows_cuda_requires_cuda_dll(tmp_path: Path): + install_dir = tmp_path / "llama.cpp" + install_dir.mkdir() + write_windows_install_shape( + install_dir, include_llama_dll = True, include_cuda_dll = True + ) + + host = HostInfo( + system = "Windows", + machine = "AMD64", + is_windows = True, + is_linux = False, + is_macos = False, + is_x86_64 = True, + is_arm64 = False, + nvidia_smi = None, + driver_cuda_version = (12, 4), + compute_caps = [], + visible_cuda_devices = None, + has_physical_nvidia = False, + has_usable_nvidia = True, + ) + choice = AssetChoice( + repo = "unslothai/llama.cpp", + tag = "release-1", + name = "llama-b9001-bin-win-cuda-12.4-x64.zip", + url = "https://example.com/x.zip", + source_label = "published", + install_kind = "windows-cuda", + runtime_line = "cuda12", + expected_sha256 = "a" * 64, + ) + checksums = ApprovedReleaseChecksums( + repo = "unslothai/llama.cpp", + release_tag = "release-1", + upstream_tag = "b9001", + source_commit = "deadbeef", + artifacts = { + source_archive_logical_name("b9001"): ApprovedArtifactHash( + asset_name = source_archive_logical_name("b9001"), + sha256 = "b" * 64, + repo = "ggml-org/llama.cpp", + kind = "upstream-source", + ), + choice.name: ApprovedArtifactHash( + asset_name = choice.name, + sha256 = choice.expected_sha256, + repo = "unslothai/llama.cpp", + kind = "prebuilt", + ), + }, + ) + plan = INSTALL_LLAMA_PREBUILT.InstallReleasePlan( + requested_tag = "latest", + llama_tag = "b9001", + release_tag = "release-1", + attempts = [choice], + approved_checksums = checksums, + ) + write_prebuilt_metadata( + install_dir, + requested_tag = "latest", + llama_tag = "b9001", + release_tag = "release-1", + choice = choice, + approved_checksums = checksums, + prebuilt_fallback_used = False, + ) + + assert existing_install_matches_plan(install_dir, host, plan) is True + (install_dir / "build" / "bin" / "Release" / "ggml-cuda.dll").unlink() + assert existing_install_matches_plan(install_dir, host, plan) is False + + +def test_existing_install_matches_plan_macos_requires_dylibs(tmp_path: Path): + install_dir = tmp_path / "llama.cpp" + install_dir.mkdir() + write_macos_install_shape(install_dir) + + host = HostInfo( + system = "Darwin", + machine = "arm64", + is_windows = False, + is_linux = False, + is_macos = True, + is_x86_64 = False, + is_arm64 = True, + nvidia_smi = None, + driver_cuda_version = None, + compute_caps = [], + visible_cuda_devices = None, + has_physical_nvidia = False, + has_usable_nvidia = False, + ) + choice = AssetChoice( + repo = "unslothai/llama.cpp", + tag = "release-1", + name = "llama-b9001-bin-macos-arm64.tar.gz", + url = "https://example.com/x.tar.gz", + source_label = "published", + install_kind = "macos-arm64", + expected_sha256 = "a" * 64, + ) + checksums = ApprovedReleaseChecksums( + repo = "unslothai/llama.cpp", + release_tag = "release-1", + upstream_tag = "b9001", + source_commit = "deadbeef", + artifacts = { + source_archive_logical_name("b9001"): ApprovedArtifactHash( + asset_name = source_archive_logical_name("b9001"), + sha256 = "b" * 64, + repo = "ggml-org/llama.cpp", + kind = "upstream-source", + ), + choice.name: ApprovedArtifactHash( + asset_name = choice.name, + sha256 = choice.expected_sha256, + repo = "unslothai/llama.cpp", + kind = "prebuilt", + ), + }, + ) + plan = INSTALL_LLAMA_PREBUILT.InstallReleasePlan( + requested_tag = "latest", + llama_tag = "b9001", + release_tag = "release-1", + attempts = [choice], + approved_checksums = checksums, + ) + write_prebuilt_metadata( + install_dir, + requested_tag = "latest", + llama_tag = "b9001", + release_tag = "release-1", + choice = choice, + approved_checksums = checksums, + prebuilt_fallback_used = False, + ) + + assert existing_install_matches_plan(install_dir, host, plan) is True + (install_dir / "build" / "bin" / "libggml.0.dylib").unlink() + assert existing_install_matches_plan(install_dir, host, plan) is False + + +def test_install_prebuilt_skips_download_when_existing_install_matches( + tmp_path: Path, monkeypatch: pytest.MonkeyPatch +): + install_dir = tmp_path / "llama.cpp" + install_dir.mkdir() + write_linux_install_shape(install_dir) + + host = HostInfo( + system = "Linux", + machine = "x86_64", + is_windows = False, + is_linux = True, + is_macos = False, + is_x86_64 = True, + is_arm64 = False, + nvidia_smi = None, + driver_cuda_version = None, + compute_caps = [], + visible_cuda_devices = None, + has_physical_nvidia = False, + has_usable_nvidia = False, + ) + choice = AssetChoice( + repo = "unslothai/llama.cpp", + tag = "release-1", + name = "llama-b9001-bin-ubuntu-x64.tar.gz", + url = "https://example.com/llama-b9001-bin-ubuntu-x64.tar.gz", + source_label = "upstream", + install_kind = "linux-cpu", + expected_sha256 = "a" * 64, + ) + checksums = ApprovedReleaseChecksums( + repo = "unslothai/llama.cpp", + release_tag = "release-1", + upstream_tag = "b9001", + source_commit = "deadbeef", + artifacts = { + source_archive_logical_name("b9001"): ApprovedArtifactHash( + asset_name = source_archive_logical_name("b9001"), + sha256 = "b" * 64, + repo = "ggml-org/llama.cpp", + kind = "upstream-source", + ), + choice.name: ApprovedArtifactHash( + asset_name = choice.name, + sha256 = choice.expected_sha256, + repo = "ggml-org/llama.cpp", + kind = "upstream-prebuilt", + ), + }, + ) + plan = INSTALL_LLAMA_PREBUILT.InstallReleasePlan( + requested_tag = "latest", + llama_tag = "b9001", + release_tag = "release-1", + attempts = [choice], + approved_checksums = checksums, + ) + + write_prebuilt_metadata( + install_dir, + requested_tag = "latest", + llama_tag = "b9001", + release_tag = "release-1", + choice = choice, + approved_checksums = checksums, + prebuilt_fallback_used = False, + ) + + monkeypatch.setattr(INSTALL_LLAMA_PREBUILT, "detect_host", lambda: host) + monkeypatch.setattr( + INSTALL_LLAMA_PREBUILT, + "resolve_install_release_plans", + lambda llama_tag, host, published_repo, published_release_tag: ( + "latest", + [plan], + ), + ) + monkeypatch.setattr( + INSTALL_LLAMA_PREBUILT, + "download_validation_model", + lambda *args, **kwargs: (_ for _ in ()).throw( + AssertionError( + "matching install should skip before validation model download" + ) + ), + ) + + install_prebuilt(install_dir, "latest", "unslothai/llama.cpp", "") + + +def test_install_prebuilt_does_not_skip_unhealthy_existing_install( + tmp_path: Path, monkeypatch: pytest.MonkeyPatch +): + install_dir = tmp_path / "llama.cpp" + install_dir.mkdir() + write_linux_install_shape(install_dir) + (install_dir / "llama-quantize").unlink() + + host = HostInfo( + system = "Linux", + machine = "x86_64", + is_windows = False, + is_linux = True, + is_macos = False, + is_x86_64 = True, + is_arm64 = False, + nvidia_smi = None, + driver_cuda_version = None, + compute_caps = [], + visible_cuda_devices = None, + has_physical_nvidia = False, + has_usable_nvidia = False, + ) + choice = AssetChoice( + repo = "unslothai/llama.cpp", + tag = "release-1", + name = "llama-b9001-bin-ubuntu-x64.tar.gz", + url = "https://example.com/llama-b9001-bin-ubuntu-x64.tar.gz", + source_label = "upstream", + install_kind = "linux-cpu", + expected_sha256 = "a" * 64, + ) + checksums = ApprovedReleaseChecksums( + repo = "unslothai/llama.cpp", + release_tag = "release-1", + upstream_tag = "b9001", + source_commit = "deadbeef", + artifacts = { + source_archive_logical_name("b9001"): ApprovedArtifactHash( + asset_name = source_archive_logical_name("b9001"), + sha256 = "b" * 64, + repo = "ggml-org/llama.cpp", + kind = "upstream-source", + ), + choice.name: ApprovedArtifactHash( + asset_name = choice.name, + sha256 = choice.expected_sha256, + repo = "ggml-org/llama.cpp", + kind = "upstream-prebuilt", + ), + }, + ) + plan = INSTALL_LLAMA_PREBUILT.InstallReleasePlan( + requested_tag = "latest", + llama_tag = "b9001", + release_tag = "release-1", + attempts = [choice], + approved_checksums = checksums, + ) + + write_prebuilt_metadata( + install_dir, + requested_tag = "latest", + llama_tag = "b9001", + release_tag = "release-1", + choice = choice, + approved_checksums = checksums, + prebuilt_fallback_used = False, + ) + + monkeypatch.setattr(INSTALL_LLAMA_PREBUILT, "detect_host", lambda: host) + monkeypatch.setattr( + INSTALL_LLAMA_PREBUILT, + "resolve_install_release_plans", + lambda llama_tag, host, published_repo, published_release_tag: ( + "latest", + [plan], + ), + ) + monkeypatch.setattr( + INSTALL_LLAMA_PREBUILT, + "download_validation_model", + lambda *args, **kwargs: (_ for _ in ()).throw( + AssertionError("unhealthy install must continue into normal install flow") + ), + ) + + with pytest.raises( + AssertionError, match = "unhealthy install must continue into normal install flow" + ): + install_prebuilt(install_dir, "latest", "unslothai/llama.cpp", "") + + +def test_install_prebuilt_skips_when_older_release_fallback_matches_existing_install( + tmp_path: Path, monkeypatch: pytest.MonkeyPatch +): + install_dir = tmp_path / "llama.cpp" + install_dir.mkdir() + write_linux_install_shape(install_dir) + + host = HostInfo( + system = "Linux", + machine = "x86_64", + is_windows = False, + is_linux = True, + is_macos = False, + is_x86_64 = True, + is_arm64 = False, + nvidia_smi = None, + driver_cuda_version = None, + compute_caps = [], + visible_cuda_devices = None, + has_physical_nvidia = False, + has_usable_nvidia = False, + ) + latest_choice = AssetChoice( + repo = "unslothai/llama.cpp", + tag = "release-2", + name = "llama-b9002-bin-ubuntu-x64.tar.gz", + url = "https://example.com/llama-b9002-bin-ubuntu-x64.tar.gz", + source_label = "upstream", + install_kind = "linux-cpu", + expected_sha256 = "c" * 64, + ) + fallback_choice = AssetChoice( + repo = "unslothai/llama.cpp", + tag = "release-1", + name = "llama-b9001-bin-ubuntu-x64.tar.gz", + url = "https://example.com/llama-b9001-bin-ubuntu-x64.tar.gz", + source_label = "upstream", + install_kind = "linux-cpu", + expected_sha256 = "a" * 64, + ) + latest_checksums = ApprovedReleaseChecksums( + repo = "unslothai/llama.cpp", + release_tag = "release-2", + upstream_tag = "b9002", + source_commit = "beadfeed", + artifacts = { + source_archive_logical_name("b9002"): ApprovedArtifactHash( + asset_name = source_archive_logical_name("b9002"), + sha256 = "d" * 64, + repo = "ggml-org/llama.cpp", + kind = "upstream-source", + ), + latest_choice.name: ApprovedArtifactHash( + asset_name = latest_choice.name, + sha256 = latest_choice.expected_sha256, + repo = "ggml-org/llama.cpp", + kind = "upstream-prebuilt", + ), + }, + ) + fallback_checksums = ApprovedReleaseChecksums( + repo = "unslothai/llama.cpp", + release_tag = "release-1", + upstream_tag = "b9001", + source_commit = "deadbeef", + artifacts = { + source_archive_logical_name("b9001"): ApprovedArtifactHash( + asset_name = source_archive_logical_name("b9001"), + sha256 = "b" * 64, + repo = "ggml-org/llama.cpp", + kind = "upstream-source", + ), + fallback_choice.name: ApprovedArtifactHash( + asset_name = fallback_choice.name, + sha256 = fallback_choice.expected_sha256, + repo = "ggml-org/llama.cpp", + kind = "upstream-prebuilt", + ), + }, + ) + latest_plan = INSTALL_LLAMA_PREBUILT.InstallReleasePlan( + requested_tag = "latest", + llama_tag = "b9002", + release_tag = "release-2", + attempts = [latest_choice], + approved_checksums = latest_checksums, + ) + fallback_plan = INSTALL_LLAMA_PREBUILT.InstallReleasePlan( + requested_tag = "latest", + llama_tag = "b9001", + release_tag = "release-1", + attempts = [fallback_choice], + approved_checksums = fallback_checksums, + ) + + write_prebuilt_metadata( + install_dir, + requested_tag = "latest", + llama_tag = "b9001", + release_tag = "release-1", + choice = fallback_choice, + approved_checksums = fallback_checksums, + prebuilt_fallback_used = True, + ) + + monkeypatch.setattr(INSTALL_LLAMA_PREBUILT, "detect_host", lambda: host) + monkeypatch.setattr( + INSTALL_LLAMA_PREBUILT, + "resolve_install_release_plans", + lambda llama_tag, host, published_repo, published_release_tag: ( + "latest", + [latest_plan, fallback_plan], + ), + ) + monkeypatch.setattr( + INSTALL_LLAMA_PREBUILT, + "download_validation_model", + lambda probe_path, cache_path: probe_path.write_bytes(b"probe"), + ) + + call_log: list[str] = [] + + def fake_validate( + attempts, + host, + install_dir, + work_dir, + probe_path, + *, + requested_tag, + llama_tag, + release_tag, + approved_checksums, + initial_fallback_used = False, + existing_install_dir = None, + ): + call_log.append(llama_tag) + raise PrebuiltFallback("validation failed for latest release") + + monkeypatch.setattr( + INSTALL_LLAMA_PREBUILT, + "validate_prebuilt_attempts", + fake_validate, + ) + monkeypatch.setattr( + INSTALL_LLAMA_PREBUILT, + "activate_install_tree", + lambda *args, **kwargs: (_ for _ in ()).throw( + AssertionError("matching fallback install should not reactivate") + ), + ) + + install_prebuilt(install_dir, "latest", "unslothai/llama.cpp", "") + + assert call_log == ["b9002"] + + +def test_install_prebuilt_skips_same_release_fallback_attempt_when_installed( + tmp_path: Path, monkeypatch: pytest.MonkeyPatch +): + install_dir = tmp_path / "llama.cpp" + install_dir.mkdir() + write_linux_install_shape(install_dir) + + host = HostInfo( + system = "Linux", + machine = "x86_64", + is_windows = False, + is_linux = True, + is_macos = False, + is_x86_64 = True, + is_arm64 = False, + nvidia_smi = None, + driver_cuda_version = None, + compute_caps = [], + visible_cuda_devices = None, + has_physical_nvidia = False, + has_usable_nvidia = False, + ) + first_choice = AssetChoice( + repo = "unslothai/llama.cpp", + tag = "release-1", + name = "llama-b9001-bin-ubuntu-x64-bad.tar.gz", + url = "https://example.com/llama-b9001-bin-ubuntu-x64-bad.tar.gz", + source_label = "published", + install_kind = "linux-cpu", + expected_sha256 = "c" * 64, + ) + fallback_choice = AssetChoice( + repo = "unslothai/llama.cpp", + tag = "release-1", + name = "llama-b9001-bin-ubuntu-x64-good.tar.gz", + url = "https://example.com/llama-b9001-bin-ubuntu-x64-good.tar.gz", + source_label = "upstream", + install_kind = "linux-cpu", + expected_sha256 = "a" * 64, + ) + checksums = ApprovedReleaseChecksums( + repo = "unslothai/llama.cpp", + release_tag = "release-1", + upstream_tag = "b9001", + source_commit = "deadbeef", + artifacts = { + source_archive_logical_name("b9001"): ApprovedArtifactHash( + asset_name = source_archive_logical_name("b9001"), + sha256 = "b" * 64, + repo = "ggml-org/llama.cpp", + kind = "upstream-source", + ), + first_choice.name: ApprovedArtifactHash( + asset_name = first_choice.name, + sha256 = first_choice.expected_sha256, + repo = "unslothai/llama.cpp", + kind = "prebuilt", + ), + fallback_choice.name: ApprovedArtifactHash( + asset_name = fallback_choice.name, + sha256 = fallback_choice.expected_sha256, + repo = "ggml-org/llama.cpp", + kind = "upstream-prebuilt", + ), + }, + ) + plan = INSTALL_LLAMA_PREBUILT.InstallReleasePlan( + requested_tag = "latest", + llama_tag = "b9001", + release_tag = "release-1", + attempts = [first_choice, fallback_choice], + approved_checksums = checksums, + ) + + write_prebuilt_metadata( + install_dir, + requested_tag = "latest", + llama_tag = "b9001", + release_tag = "release-1", + choice = fallback_choice, + approved_checksums = checksums, + prebuilt_fallback_used = True, + ) + assert ( + existing_install_matches_choice( + install_dir, + host, + llama_tag = "b9001", + release_tag = "release-1", + choice = fallback_choice, + approved_checksums = checksums, + ) + is True + ) + + monkeypatch.setattr(INSTALL_LLAMA_PREBUILT, "detect_host", lambda: host) + monkeypatch.setattr( + INSTALL_LLAMA_PREBUILT, + "resolve_install_release_plans", + lambda llama_tag, host, published_repo, published_release_tag: ( + "latest", + [plan], + ), + ) + monkeypatch.setattr( + INSTALL_LLAMA_PREBUILT, + "download_validation_model", + lambda probe_path, cache_path: probe_path.write_bytes(b"probe"), + ) + + attempted_names: list[str] = [] + + def fake_validate_choice( + choice, + host, + staging_dir, + work_dir, + probe_path, + *, + requested_tag, + llama_tag, + release_tag, + approved_checksums, + prebuilt_fallback_used, + quantized_path, + ): + attempted_names.append(choice.name) + if choice.name == first_choice.name: + raise PrebuiltFallback("newest candidate failed") + raise AssertionError("installed fallback candidate should have been skipped") + + monkeypatch.setattr( + INSTALL_LLAMA_PREBUILT, + "validate_prebuilt_choice", + fake_validate_choice, + ) + monkeypatch.setattr( + INSTALL_LLAMA_PREBUILT, + "activate_install_tree", + lambda *args, **kwargs: (_ for _ in ()).throw( + AssertionError("installed fallback candidate should not be activated") + ), + ) + + install_prebuilt(install_dir, "latest", "unslothai/llama.cpp", "") + + assert attempted_names == [first_choice.name] + + +def test_install_prebuilt_same_tag_upstream_failure_uses_older_unsloth_release_plan( + tmp_path: Path, monkeypatch: pytest.MonkeyPatch +): + install_dir = tmp_path / "llama.cpp" + host = HostInfo( + system = "Linux", + machine = "x86_64", + is_windows = False, + is_linux = True, + is_macos = False, + is_x86_64 = True, + is_arm64 = False, + nvidia_smi = None, + driver_cuda_version = None, + compute_caps = [], + visible_cuda_devices = None, + has_physical_nvidia = False, + has_usable_nvidia = False, + ) + + same_tag_upstream_choice = AssetChoice( + repo = "ggml-org/llama.cpp", + tag = "b9002", + name = "llama-b9002-bin-ubuntu-x64.tar.gz", + url = "https://example.com/llama-b9002-bin-ubuntu-x64.tar.gz", + source_label = "upstream", + install_kind = "linux-cpu", + expected_sha256 = "a" * 64, + ) + older_release_choice = AssetChoice( + repo = "unslothai/llama.cpp", + tag = "release-1", + name = "llama-b9001-bin-ubuntu-x64.tar.gz", + url = "https://example.com/llama-b9001-bin-ubuntu-x64.tar.gz", + source_label = "upstream", + install_kind = "linux-cpu", + expected_sha256 = "b" * 64, + ) + latest_plan = INSTALL_LLAMA_PREBUILT.InstallReleasePlan( + requested_tag = "latest", + llama_tag = "b9002", + release_tag = "release-2", + attempts = [same_tag_upstream_choice], + approved_checksums = ApprovedReleaseChecksums( + repo = "unslothai/llama.cpp", + release_tag = "release-2", + upstream_tag = "b9002", + source_commit = None, + artifacts = {}, + ), + ) + older_plan = INSTALL_LLAMA_PREBUILT.InstallReleasePlan( + requested_tag = "latest", + llama_tag = "b9001", + release_tag = "release-1", + attempts = [older_release_choice], + approved_checksums = ApprovedReleaseChecksums( + repo = "unslothai/llama.cpp", + release_tag = "release-1", + upstream_tag = "b9001", + source_commit = None, + artifacts = {}, + ), + ) + + monkeypatch.setattr(INSTALL_LLAMA_PREBUILT, "detect_host", lambda: host) + monkeypatch.setattr( + INSTALL_LLAMA_PREBUILT, + "resolve_install_release_plans", + lambda llama_tag, host, published_repo, published_release_tag: ( + "latest", + [latest_plan, older_plan], + ), + ) + monkeypatch.setattr( + INSTALL_LLAMA_PREBUILT, + "download_validation_model", + lambda probe_path, cache_path: probe_path.write_bytes(b"probe"), + ) + monkeypatch.setattr( + INSTALL_LLAMA_PREBUILT, + "latest_upstream_release_tag", + lambda: (_ for _ in ()).throw( + AssertionError("install fallback should not walk upstream releases") + ), + ) + + attempted = [] + + def fake_validate( + attempts, + host, + install_dir, + work_dir, + probe_path, + *, + requested_tag, + llama_tag, + release_tag, + approved_checksums, + initial_fallback_used = False, + existing_install_dir = None, + ): + attempted.append((llama_tag, release_tag, attempts[0].source_label)) + if llama_tag == "b9002": + raise PrebuiltFallback("same-tag upstream asset failed validation") + staging_dir = create_install_staging_dir(install_dir) + (staging_dir / "marker.txt").write_text("ready\n") + return attempts[0], staging_dir, initial_fallback_used + + monkeypatch.setattr( + INSTALL_LLAMA_PREBUILT, "validate_prebuilt_attempts", fake_validate + ) + + activated = {} + monkeypatch.setattr( + INSTALL_LLAMA_PREBUILT, + "activate_install_tree", + lambda staging_dir, install_dir, host: activated.update( + {"staging_dir": staging_dir, "install_dir": install_dir} + ), + ) + monkeypatch.setattr( + INSTALL_LLAMA_PREBUILT, + "ensure_converter_scripts", + lambda install_dir, llama_tag: None, + ) + + install_prebuilt(install_dir, "latest", "unslothai/llama.cpp", "") + + assert attempted == [ + ("b9002", "release-2", "upstream"), + ("b9001", "release-1", "upstream"), + ] + assert activated["install_dir"] == install_dir + + def io_bytes(data: bytes): return io.BytesIO(data) @@ -628,3 +1866,184 @@ def add_symlink_to_tar(archive: tarfile.TarFile, name: str, target: str) -> None info.type = tarfile.SYMTYPE info.linkname = target archive.addfile(info) + + +def test_existing_install_matches_choice_fails_when_install_tree_incomplete( + tmp_path: Path, +): + """confirm_install_tree guard rejects installs missing critical files.""" + install_dir = tmp_path / "llama.cpp" + install_dir.mkdir() + write_linux_install_shape(install_dir) + + host = HostInfo( + system = "Linux", + machine = "x86_64", + is_windows = False, + is_linux = True, + is_macos = False, + is_x86_64 = True, + is_arm64 = False, + nvidia_smi = None, + driver_cuda_version = None, + compute_caps = [], + visible_cuda_devices = None, + has_physical_nvidia = False, + has_usable_nvidia = False, + ) + choice = AssetChoice( + repo = "unslothai/llama.cpp", + tag = "release-1", + name = "llama-b9001-bin-ubuntu-x64.tar.gz", + url = "https://example.com/llama-b9001-bin-ubuntu-x64.tar.gz", + source_label = "upstream", + install_kind = "linux-cpu", + expected_sha256 = "a" * 64, + ) + checksums = ApprovedReleaseChecksums( + repo = "unslothai/llama.cpp", + release_tag = "release-1", + upstream_tag = "b9001", + source_commit = "deadbeef", + artifacts = { + source_archive_logical_name("b9001"): ApprovedArtifactHash( + asset_name = source_archive_logical_name("b9001"), + sha256 = "b" * 64, + repo = "ggml-org/llama.cpp", + kind = "upstream-source", + ), + choice.name: ApprovedArtifactHash( + asset_name = choice.name, + sha256 = choice.expected_sha256, + repo = "ggml-org/llama.cpp", + kind = "upstream-prebuilt", + ), + }, + ) + write_prebuilt_metadata( + install_dir, + requested_tag = "latest", + llama_tag = "b9001", + release_tag = "release-1", + choice = choice, + approved_checksums = checksums, + prebuilt_fallback_used = False, + ) + + # Full install should match + assert ( + existing_install_matches_choice( + install_dir, + host, + llama_tag = "b9001", + release_tag = "release-1", + choice = choice, + approved_checksums = checksums, + ) + is True + ) + + # Remove convert_hf_to_gguf.py (checked by confirm_install_tree but not + # runtime_payload_is_healthy) and verify the guard catches it + (install_dir / "convert_hf_to_gguf.py").unlink() + assert ( + existing_install_matches_choice( + install_dir, + host, + llama_tag = "b9001", + release_tag = "release-1", + choice = choice, + approved_checksums = checksums, + ) + is False + ) + + +def test_existing_install_matches_choice_fails_when_install_tree_incomplete_macos( + tmp_path: Path, +): + """confirm_install_tree guard rejects macOS arm64 installs missing critical files.""" + install_dir = tmp_path / "llama.cpp" + install_dir.mkdir() + write_macos_install_shape(install_dir) + + host = HostInfo( + system = "Darwin", + machine = "arm64", + is_windows = False, + is_linux = False, + is_macos = True, + is_x86_64 = False, + is_arm64 = True, + nvidia_smi = None, + driver_cuda_version = None, + compute_caps = [], + visible_cuda_devices = None, + has_physical_nvidia = False, + has_usable_nvidia = False, + ) + choice = AssetChoice( + repo = "unslothai/llama.cpp", + tag = "release-1", + name = "llama-b9001-bin-macos-arm64.tar.gz", + url = "https://example.com/llama-b9001-bin-macos-arm64.tar.gz", + source_label = "upstream", + install_kind = "macos-arm64", + expected_sha256 = "a" * 64, + ) + checksums = ApprovedReleaseChecksums( + repo = "unslothai/llama.cpp", + release_tag = "release-1", + upstream_tag = "b9001", + source_commit = "deadbeef", + artifacts = { + source_archive_logical_name("b9001"): ApprovedArtifactHash( + asset_name = source_archive_logical_name("b9001"), + sha256 = "b" * 64, + repo = "ggml-org/llama.cpp", + kind = "upstream-source", + ), + choice.name: ApprovedArtifactHash( + asset_name = choice.name, + sha256 = choice.expected_sha256, + repo = "ggml-org/llama.cpp", + kind = "upstream-prebuilt", + ), + }, + ) + write_prebuilt_metadata( + install_dir, + requested_tag = "latest", + llama_tag = "b9001", + release_tag = "release-1", + choice = choice, + approved_checksums = checksums, + prebuilt_fallback_used = False, + ) + + # Full install should match + assert ( + existing_install_matches_choice( + install_dir, + host, + llama_tag = "b9001", + release_tag = "release-1", + choice = choice, + approved_checksums = checksums, + ) + is True + ) + + # Remove a macOS-specific runtime artifact and verify the guard catches it + (install_dir / "build" / "bin" / "libmtmd.0.dylib").unlink() + assert ( + existing_install_matches_choice( + install_dir, + host, + llama_tag = "b9001", + release_tag = "release-1", + choice = choice, + approved_checksums = checksums, + ) + is False + ) diff --git a/tests/studio/install/test_llama_pr_force_and_source.py b/tests/studio/install/test_llama_pr_force_and_source.py new file mode 100644 index 0000000000..114680f458 --- /dev/null +++ b/tests/studio/install/test_llama_pr_force_and_source.py @@ -0,0 +1,635 @@ +""" +Tests for the current llama.cpp wrapper policy in setup.sh / setup.ps1. + +Tests cover: + - Bash subprocess: PR_FORCE promotion, user-override, zero/empty/invalid ignored + - Bash subprocess: source remains pinned to ggml-org even if env source is set + - Static source checks: mainline repo/source are hardcoded for now + - PowerShell subprocess: PR_FORCE promotion and fixed-source parity + +Run: pytest tests/studio/install/test_llama_pr_force_and_source.py -v +""" + +import os +import shlex +import subprocess +import textwrap +from pathlib import Path + +import pytest + +# --------------------------------------------------------------------------- +# Paths +# --------------------------------------------------------------------------- +PACKAGE_ROOT = Path(__file__).resolve().parents[3] +SETUP_SH = PACKAGE_ROOT / "studio" / "setup.sh" +SETUP_PS1 = PACKAGE_ROOT / "studio" / "setup.ps1" + +BASH = "/bin/bash" +PWSH = "/usr/bin/pwsh" +PWSH_AVAILABLE = os.path.isfile(PWSH) and os.access(PWSH, os.X_OK) +requires_pwsh = pytest.mark.skipif(not PWSH_AVAILABLE, reason = "pwsh not available") + + +# --------------------------------------------------------------------------- +# Helpers +# --------------------------------------------------------------------------- +def run_bash( + script: str, *, timeout: int = 10, env: dict | None = None +) -> subprocess.CompletedProcess: + """Run a bash script fragment and return the CompletedProcess.""" + run_env = os.environ.copy() + if env: + run_env.update(env) + return subprocess.run( + [BASH, "-c", script], + capture_output = True, + text = True, + timeout = timeout, + env = run_env, + ) + + +def run_pwsh( + script: str, *, timeout: int = 10, env: dict | None = None +) -> subprocess.CompletedProcess: + """Run a PowerShell script fragment and return the CompletedProcess.""" + run_env = os.environ.copy() + run_env["NO_COLOR"] = "1" + if env: + run_env.update(env) + return subprocess.run( + [PWSH, "-NoProfile", "-Command", script], + capture_output = True, + text = True, + timeout = timeout, + env = run_env, + ) + + +# --------------------------------------------------------------------------- +# Shared bash stubs +# --------------------------------------------------------------------------- +BASH_STUBS = textwrap.dedent("""\ + step() { echo "step:$1:$2"; } + substep() { :; } + verbose_substep() { :; } + print_llama_error_log() { :; } + C_ERR= C_WARN= C_OK= C_RST= C_TITLE= C_DIM= +""") + +RUN_QUIET_STUB = textwrap.dedent("""\ + run_quiet_no_exit() { local _label="$1"; shift; "$@"; return $?; } +""") + + +def make_mock_git(tmp_path: Path, *, fail_on: str = "") -> tuple[Path, Path]: + """Create a mock git binary that logs calls. Returns (mock_bin, log_file).""" + mock_bin = tmp_path / "mock_bin" + mock_bin.mkdir(exist_ok = True) + log_file = tmp_path / "git_calls.log" + + if fail_on: + script = ( + f'#!/bin/bash\necho "$*" >> {log_file}\n' + f'_args=("$@")\n' + f"_i=0\n" + f'while [ "${{_args[$_i]:-}}" = "-C" ]; do _i=$((_i+2)); done\n' + f'_subcmd="${{_args[$_i]:-}}"\n' + f'if [ "$_subcmd" = "{fail_on}" ]; then exit 1; fi\n' + f"exit 0\n" + ) + else: + script = f'#!/bin/bash\necho "$*" >> {log_file}\nexit 0\n' + + git_bin = mock_bin / "git" + git_bin.write_text(script) + git_bin.chmod(0o755) + return mock_bin, log_file + + +# ========================================================================= +# Bash fragment that exercises PR_FORCE and fixed _LLAMA_SOURCE resolution +# ========================================================================= +def _bash_resolution_fragment( + llama_pr: str = "", + llama_pr_force: str = "", + llama_source: str = "", + default_pr_force: str = "", + default_source: str = "https://github.com/ggml-org/llama.cpp", +) -> str: + """Build the bash fragment that mirrors setup.sh resolution logic.""" + return BASH_STUBS + textwrap.dedent(f"""\ + _LLAMA_PR={shlex.quote(llama_pr) if llama_pr else '""'} + _DEFAULT_LLAMA_PR_FORCE={shlex.quote(default_pr_force) if default_pr_force else '""'} + _DEFAULT_LLAMA_SOURCE={shlex.quote(default_source)} + + _LLAMA_PR_FORCE={shlex.quote(llama_pr_force) if llama_pr_force else '"$_DEFAULT_LLAMA_PR_FORCE"'} + export UNSLOTH_LLAMA_SOURCE={shlex.quote(llama_source) if llama_source else '""'} + _LLAMA_SOURCE="$_DEFAULT_LLAMA_SOURCE" + _LLAMA_SOURCE="${{_LLAMA_SOURCE%.git}}" + + _NEED_LLAMA_SOURCE_BUILD=false + _SKIP_PREBUILT_INSTALL=false + + if [ "$_LLAMA_SOURCE" != "https://github.com/ggml-org/llama.cpp" ]; then + step "llama.cpp" "custom source: $_LLAMA_SOURCE -- forcing source build" + _NEED_LLAMA_SOURCE_BUILD=true + _SKIP_PREBUILT_INSTALL=true + fi + + if [ -z "$_LLAMA_PR" ] && [ -n "$_LLAMA_PR_FORCE" ] && \\ + [[ "$_LLAMA_PR_FORCE" =~ ^[0-9]+$ ]] && [ "$_LLAMA_PR_FORCE" -gt 0 ]; then + _LLAMA_PR="$_LLAMA_PR_FORCE" + step "llama.cpp" "baked-in PR_FORCE=$_LLAMA_PR_FORCE" + fi + + echo "LLAMA_PR=$_LLAMA_PR" + echo "LLAMA_SOURCE=$_LLAMA_SOURCE" + echo "NEED_SOURCE=$_NEED_LLAMA_SOURCE_BUILD" + echo "SKIP_PREBUILT=$_SKIP_PREBUILT_INSTALL" + """) + + +# ========================================================================= +# TEST GROUP A: Bash PR_FORCE promotion (subprocess) +# ========================================================================= +class TestBashPrForcePromotion: + """PR_FORCE promotes to _LLAMA_PR when user hasn't set one.""" + + def test_baked_in_pr_force_promotes(self): + script = _bash_resolution_fragment(default_pr_force = "12345") + r = run_bash(script) + assert r.returncode == 0 + assert "LLAMA_PR=12345" in r.stdout + assert "baked-in PR_FORCE=12345" in r.stdout + + def test_env_pr_force_promotes(self): + script = _bash_resolution_fragment(llama_pr_force = "999") + r = run_bash(script) + assert r.returncode == 0 + assert "LLAMA_PR=999" in r.stdout + + def test_user_pr_overrides_pr_force(self): + """UNSLOTH_LLAMA_PR takes priority over PR_FORCE.""" + script = _bash_resolution_fragment( + llama_pr = "100", + llama_pr_force = "200", + ) + r = run_bash(script) + assert r.returncode == 0 + assert "LLAMA_PR=100" in r.stdout + assert "baked-in PR_FORCE" not in r.stdout + + def test_user_pr_overrides_baked_in(self): + script = _bash_resolution_fragment( + llama_pr = "100", + default_pr_force = "200", + ) + r = run_bash(script) + assert r.returncode == 0 + assert "LLAMA_PR=100" in r.stdout + assert "baked-in PR_FORCE" not in r.stdout + + def test_pr_force_zero_ignored(self): + script = _bash_resolution_fragment(llama_pr_force = "0") + r = run_bash(script) + assert r.returncode == 0 + assert "LLAMA_PR=" in r.stdout + assert "baked-in PR_FORCE" not in r.stdout + + def test_pr_force_empty_ignored(self): + script = _bash_resolution_fragment(default_pr_force = "") + r = run_bash(script) + assert r.returncode == 0 + assert "LLAMA_PR=" in r.stdout + assert "baked-in PR_FORCE" not in r.stdout + + def test_pr_force_alpha_ignored(self): + script = _bash_resolution_fragment(llama_pr_force = "abc") + r = run_bash(script) + assert r.returncode == 0 + assert "LLAMA_PR=" in r.stdout + assert "baked-in PR_FORCE" not in r.stdout + + def test_pr_force_negative_ignored(self): + script = _bash_resolution_fragment(llama_pr_force = "-5") + r = run_bash(script) + assert r.returncode == 0 + assert "LLAMA_PR=" in r.stdout + + def test_pr_force_decimal_ignored(self): + script = _bash_resolution_fragment(llama_pr_force = "12.34") + r = run_bash(script) + assert r.returncode == 0 + assert "LLAMA_PR=" in r.stdout + + +# ========================================================================= +# TEST GROUP B: Bash fixed mainline source (subprocess) +# ========================================================================= +class TestBashFixedMainlineSource: + """Source remains pinned to ggml-org while the temporary policy is active.""" + + def test_default_source_no_force(self): + script = _bash_resolution_fragment() + r = run_bash(script) + assert r.returncode == 0 + assert "NEED_SOURCE=false" in r.stdout + assert "SKIP_PREBUILT=false" in r.stdout + assert "custom source:" not in r.stdout + + def test_env_source_override_is_ignored(self): + script = _bash_resolution_fragment( + llama_source = "https://github.com/unslothai/llama.cpp.git", + ) + r = run_bash(script) + assert r.returncode == 0 + assert "LLAMA_SOURCE=https://github.com/ggml-org/llama.cpp" in r.stdout + assert "NEED_SOURCE=false" in r.stdout + assert "SKIP_PREBUILT=false" in r.stdout + + def test_baked_in_source_stays_mainline(self): + script = _bash_resolution_fragment( + default_source = "https://github.com/ggml-org/llama.cpp", + ) + r = run_bash(script) + assert r.returncode == 0 + assert "LLAMA_SOURCE=https://github.com/ggml-org/llama.cpp" in r.stdout + + +# ========================================================================= +# TEST GROUP C: Bash clone URL parameterization (subprocess with mock git) +# ========================================================================= +class TestBashCloneUrlParameterized: + """Verify git clone uses _LLAMA_SOURCE instead of hardcoded URL.""" + + @staticmethod + def _clone_script( + mock_bin: Path, + build_tmp: str, + llama_pr: str = "", + llama_source: str = "https://github.com/ggml-org/llama.cpp", + resolved_tag: str = "b8508", + ) -> str: + return RUN_QUIET_STUB + textwrap.dedent(f"""\ + export PATH="{mock_bin}:$PATH" + _LLAMA_PR={shlex.quote(llama_pr) if llama_pr else '""'} + _LLAMA_SOURCE={shlex.quote(llama_source)} + _RESOLVED_LLAMA_TAG={shlex.quote(resolved_tag)} + _BUILD_TMP={shlex.quote(build_tmp)} + BUILD_OK=true + + if [ -n "$_LLAMA_PR" ]; then + run_quiet_no_exit "clone llama.cpp" \\ + git clone --depth 1 "${{_LLAMA_SOURCE}}.git" "$_BUILD_TMP" || BUILD_OK=false + else + _CLONE_ARGS=(git clone --depth 1) + if [ "$_RESOLVED_LLAMA_TAG" != "latest" ] && [ -n "$_RESOLVED_LLAMA_TAG" ]; then + _CLONE_ARGS+=(--branch "$_RESOLVED_LLAMA_TAG") + fi + _CLONE_ARGS+=("${{_LLAMA_SOURCE}}.git" "$_BUILD_TMP") + run_quiet_no_exit "clone llama.cpp" \\ + "${{_CLONE_ARGS[@]}}" || BUILD_OK=false + fi + echo "BUILD_OK=$BUILD_OK" + """) + + def test_pr_path_uses_custom_source(self, tmp_path: Path): + mock_bin, log_file = make_mock_git(tmp_path) + build_tmp = str(tmp_path / "build_tmp") + script = self._clone_script( + mock_bin, + build_tmp, + llama_pr = "123", + llama_source = "https://github.com/unslothai/llama.cpp", + ) + r = run_bash(script) + assert r.returncode == 0 + log = log_file.read_text() + assert "unslothai/llama.cpp.git" in log + assert "ggml-org" not in log + + def test_non_pr_path_uses_custom_source(self, tmp_path: Path): + mock_bin, log_file = make_mock_git(tmp_path) + build_tmp = str(tmp_path / "build_tmp") + script = self._clone_script( + mock_bin, + build_tmp, + llama_source = "https://github.com/unslothai/llama.cpp", + ) + r = run_bash(script) + assert r.returncode == 0 + log = log_file.read_text() + assert "unslothai/llama.cpp.git" in log + assert "ggml-org" not in log + + def test_default_source_unchanged(self, tmp_path: Path): + mock_bin, log_file = make_mock_git(tmp_path) + build_tmp = str(tmp_path / "build_tmp") + script = self._clone_script(mock_bin, build_tmp) + r = run_bash(script) + assert r.returncode == 0 + log = log_file.read_text() + assert "ggml-org/llama.cpp.git" in log + + def test_latest_tag_omits_branch_flag(self, tmp_path: Path): + """resolved_tag='latest' should not pass --branch to git clone.""" + mock_bin, log_file = make_mock_git(tmp_path) + build_tmp = str(tmp_path / "build_tmp") + script = self._clone_script( + mock_bin, + build_tmp, + resolved_tag = "latest", + ) + r = run_bash(script) + assert r.returncode == 0 + log = log_file.read_text() + assert "--branch" not in log + assert "ggml-org/llama.cpp.git" in log + + def test_empty_tag_omits_branch_flag(self, tmp_path: Path): + """resolved_tag='' (empty) should not pass --branch to git clone.""" + mock_bin, log_file = make_mock_git(tmp_path) + build_tmp = str(tmp_path / "build_tmp") + script = self._clone_script( + mock_bin, + build_tmp, + resolved_tag = "", + ) + r = run_bash(script) + assert r.returncode == 0 + log = log_file.read_text() + assert "--branch" not in log + assert "ggml-org/llama.cpp.git" in log + + +# ========================================================================= +# TEST GROUP D: Static source patterns -- setup.sh +# ========================================================================= +class TestSourcePatternsSh: + """Verify setup.sh keeps the temporary mainline-only llama.cpp policy.""" + + @pytest.fixture(autouse = True) + def _load_source(self): + self.content = SETUP_SH.read_text() + + def test_has_default_pr_force(self): + assert '_DEFAULT_LLAMA_PR_FORCE=""' in self.content + + def test_has_default_source(self): + assert ( + '_DEFAULT_LLAMA_SOURCE="https://github.com/ggml-org/llama.cpp"' + in self.content + ) + + def test_has_pr_force_env_read(self): + assert "UNSLOTH_LLAMA_PR_FORCE" in self.content + + def test_source_env_override_removed(self): + assert "UNSLOTH_LLAMA_SOURCE:-${_DEFAULT_LLAMA_SOURCE}" not in self.content + assert '_LLAMA_SOURCE="${_DEFAULT_LLAMA_SOURCE}"' in self.content + + def test_release_repo_override_removed(self): + assert "UNSLOTH_LLAMA_RELEASE_REPO:-unslothai/llama.cpp" not in self.content + assert '_HELPER_RELEASE_REPO="ggml-org/llama.cpp"' in self.content + + def test_force_compile_skips_prebuilt_resolution_early(self): + assert 'if [ "$_LLAMA_FORCE_COMPILE" = "1" ]; then' in self.content + assert "_SKIP_PREBUILT_INSTALL=true" in self.content + + def test_force_compile_uses_requested_tag_without_helper(self): + assert 'if [ "$_LLAMA_FORCE_COMPILE" = "1" ]; then' in self.content + assert '_RESOLVED_LLAMA_TAG="$_REQUESTED_LLAMA_TAG"' in self.content + + def test_pr_force_resolution_block(self): + assert '_LLAMA_PR="$_LLAMA_PR_FORCE"' in self.content + + def test_source_trailing_git_strip(self): + assert "${_LLAMA_SOURCE%.git}" in self.content + + def test_clone_urls_parameterized_pr_path(self): + """PR clone path uses ${_LLAMA_SOURCE}.git, not hardcoded URL.""" + pr_clone_idx = self.content.index( + 'if [ -n "$_LLAMA_PR" ]; then\n' + ' run_quiet_no_exit "clone llama.cpp"' + ) + else_idx = self.content.index("else\n", pr_clone_idx) + pr_block = self.content[pr_clone_idx:else_idx] + assert '"${_LLAMA_SOURCE}.git"' in pr_block + assert "ggml-org/llama.cpp.git" not in pr_block + + def test_clone_urls_parameterized_tag_path(self): + """Non-PR clone path uses the resolved source URL, not a hardcoded URL.""" + # Find the non-PR clone line (after _CLONE_ARGS) + idx = self.content.index("_CLONE_ARGS=(git clone --depth 1)") + block = self.content[idx : idx + 400] + assert '"${_RESOLVED_SOURCE_URL}.git"' in block + assert "ggml-org/llama.cpp.git" not in block + + def test_no_hardcoded_clone_urls(self): + """No remaining hardcoded ggml-org clone URLs in clone commands.""" + lines = self.content.splitlines() + for i, line in enumerate(lines, 1): + if "git clone" in line and "ggml-org/llama.cpp.git" in line: + pytest.fail( + f"Line {i} has hardcoded ggml-org clone URL: {line.strip()}" + ) + + +# ========================================================================= +# TEST GROUP E: Static source patterns -- setup.ps1 +# ========================================================================= +class TestSourcePatternsPs1: + """Verify setup.ps1 keeps the temporary mainline-only llama.cpp policy.""" + + @pytest.fixture(autouse = True) + def _load_source(self): + self.content = SETUP_PS1.read_text() + + def test_has_default_pr_force(self): + assert '$DefaultLlamaPrForce = ""' in self.content + + def test_has_default_source(self): + assert ( + '$DefaultLlamaSource = "https://github.com/ggml-org/llama.cpp"' + in self.content + ) + + def test_has_pr_force_env_read(self): + assert "$env:UNSLOTH_LLAMA_PR_FORCE" in self.content + + def test_source_env_override_removed(self): + assert "$LlamaSource = if ($env:UNSLOTH_LLAMA_SOURCE)" not in self.content + assert "$LlamaSource = $DefaultLlamaSource" in self.content + + def test_release_repo_override_removed(self): + assert ( + "$HelperReleaseRepo = if ($env:UNSLOTH_LLAMA_RELEASE_REPO)" + not in self.content + ) + assert '$HelperReleaseRepo = "ggml-org/llama.cpp"' in self.content + + def test_force_compile_skips_prebuilt_resolution_early(self): + assert 'if ($env:UNSLOTH_LLAMA_FORCE_COMPILE -eq "1") {' in self.content + assert "$SkipPrebuiltInstall = $true" in self.content + + def test_force_compile_uses_requested_tag_without_helper(self): + assert 'if ($env:UNSLOTH_LLAMA_FORCE_COMPILE -eq "1") {' in self.content + assert "$ResolvedLlamaTag = $RequestedLlamaTag" in self.content + + def test_pr_force_promotion_block(self): + assert "$LlamaPr = $LlamaPrForce" in self.content + + def test_source_trailing_git_strip(self): + assert ".EndsWith('.git')" in self.content + + def test_clone_urls_parameterized_pr_path(self): + """PR clone path uses $LlamaSource.git, not hardcoded URL.""" + pr_idx = self.content.index( + "if ($LlamaPr) {\n", self.content.index("Cloning llama.cpp") + ) + else_idx = self.content.index("} else {", pr_idx) + pr_block = self.content[pr_idx:else_idx] + assert '"$LlamaSource.git"' in pr_block + assert "ggml-org/llama.cpp.git" not in pr_block + + def test_clone_urls_parameterized_tag_path(self): + """Non-PR clone path uses the resolved source URL, not a hardcoded URL.""" + clone_args_idx = self.content.index('$cloneArgs = @("clone"') + block = self.content[clone_args_idx : clone_args_idx + 400] + assert '"$ResolvedSourceUrl.git"' in block + assert "ggml-org/llama.cpp.git" not in block + + def test_no_hardcoded_clone_urls(self): + """No remaining hardcoded ggml-org clone URLs in clone commands.""" + lines = self.content.splitlines() + for i, line in enumerate(lines, 1): + if "git clone" in line and "ggml-org/llama.cpp.git" in line: + pytest.fail( + f"Line {i} has hardcoded ggml-org clone URL: {line.strip()}" + ) + + +# ========================================================================= +# TEST GROUP F: PowerShell PR_FORCE promotion (subprocess) +# ========================================================================= +@requires_pwsh +class TestPwshPrForcePromotion: + """PR_FORCE promotion and fixed-source logic via pwsh subprocess.""" + + FRAGMENT_TEMPLATE = textwrap.dedent("""\ + function step($a, $b, $c) { Write-Output "step:$a`:$b" } + + $DefaultLlamaPrForce = "%%DEFAULT_PR_FORCE%%" + $DefaultLlamaSource = "%%DEFAULT_SOURCE%%" + + $LlamaPr = if ($env:UNSLOTH_LLAMA_PR) { $env:UNSLOTH_LLAMA_PR.Trim() } else { "" } + $LlamaPrForce = if ($env:UNSLOTH_LLAMA_PR_FORCE) { $env:UNSLOTH_LLAMA_PR_FORCE.Trim() } else { $DefaultLlamaPrForce } + $LlamaSource = $DefaultLlamaSource + if ($LlamaSource.EndsWith('.git')) { $LlamaSource = $LlamaSource.Substring(0, $LlamaSource.Length - 4) } + + $NeedLlamaSourceBuild = $false + $SkipPrebuiltInstall = $false + + if ($LlamaSource -ne "https://github.com/ggml-org/llama.cpp") { + step "llama.cpp" "custom source: $LlamaSource -- forcing source build" "Yellow" + $NeedLlamaSourceBuild = $true + $SkipPrebuiltInstall = $true + } + + if (-not $LlamaPr -and $LlamaPrForce -and $LlamaPrForce -match '^\\d+$' -and [int]$LlamaPrForce -gt 0) { + $LlamaPr = $LlamaPrForce + step "llama.cpp" "baked-in PR_FORCE=$LlamaPrForce" "Yellow" + } + + Write-Output "LLAMA_PR=$LlamaPr" + Write-Output "LLAMA_SOURCE=$LlamaSource" + Write-Output "NEED_SOURCE=$NeedLlamaSourceBuild" + Write-Output "SKIP_PREBUILT=$SkipPrebuiltInstall" + """) + + def _run( + self, + default_pr_force: str = "", + default_source: str = "https://github.com/ggml-org/llama.cpp", + env: dict | None = None, + ) -> subprocess.CompletedProcess: + script = self.FRAGMENT_TEMPLATE.replace( + "%%DEFAULT_PR_FORCE%%", + default_pr_force, + ).replace( + "%%DEFAULT_SOURCE%%", + default_source, + ) + run_env = {} + # Ensure env vars are unset by default + run_env["UNSLOTH_LLAMA_PR"] = "" + run_env["UNSLOTH_LLAMA_PR_FORCE"] = "" + if env: + run_env.update(env) + return run_pwsh(script, env = run_env) + + def test_baked_in_pr_force_promotes(self): + r = self._run(default_pr_force = "12345") + assert r.returncode == 0 + assert "LLAMA_PR=12345" in r.stdout + assert "baked-in PR_FORCE=12345" in r.stdout + + def test_env_pr_force_promotes(self): + r = self._run(env = {"UNSLOTH_LLAMA_PR_FORCE": "999"}) + assert r.returncode == 0 + assert "LLAMA_PR=999" in r.stdout + + def test_user_pr_overrides_pr_force(self): + r = self._run( + env = { + "UNSLOTH_LLAMA_PR": "100", + "UNSLOTH_LLAMA_PR_FORCE": "200", + } + ) + assert r.returncode == 0 + assert "LLAMA_PR=100" in r.stdout + assert "baked-in PR_FORCE" not in r.stdout + + def test_pr_force_zero_ignored(self): + r = self._run(env = {"UNSLOTH_LLAMA_PR_FORCE": "0"}) + assert r.returncode == 0 + assert "LLAMA_PR=" in r.stdout + assert "baked-in PR_FORCE" not in r.stdout + + def test_pr_force_alpha_ignored(self): + r = self._run(env = {"UNSLOTH_LLAMA_PR_FORCE": "abc"}) + assert r.returncode == 0 + assert "baked-in PR_FORCE" not in r.stdout + + def test_env_source_override_is_ignored(self): + r = self._run( + env = { + "UNSLOTH_LLAMA_SOURCE": "https://github.com/unslothai/llama.cpp", + } + ) + assert r.returncode == 0 + assert "LLAMA_SOURCE=https://github.com/ggml-org/llama.cpp" in r.stdout + assert "NEED_SOURCE=False" in r.stdout + assert "SKIP_PREBUILT=False" in r.stdout + + def test_default_source_no_force(self): + r = self._run() + assert r.returncode == 0 + assert "NEED_SOURCE=False" in r.stdout + assert "SKIP_PREBUILT=False" in r.stdout + + def test_trailing_git_override_is_ignored(self): + r = self._run( + env = { + "UNSLOTH_LLAMA_SOURCE": "https://github.com/unslothai/llama.cpp.git", + } + ) + assert r.returncode == 0 + assert "LLAMA_SOURCE=https://github.com/ggml-org/llama.cpp" in r.stdout + + def test_baked_in_source_stays_mainline(self): + r = self._run(default_source = "https://github.com/ggml-org/llama.cpp") + assert r.returncode == 0 + assert "LLAMA_SOURCE=https://github.com/ggml-org/llama.cpp" in r.stdout diff --git a/tests/studio/install/test_pr4562_bugfixes.py b/tests/studio/install/test_pr4562_bugfixes.py index 9b8c6219de..7fa654d845 100644 --- a/tests/studio/install/test_pr4562_bugfixes.py +++ b/tests/studio/install/test_pr4562_bugfixes.py @@ -6,7 +6,7 @@ Tests cover: - Bug 2: Source-build fallback ignores pinned tag (both .sh and .ps1) - Bug 3: Unix fallback deletes install before checking prerequisites - Bug 4: Linux LD_LIBRARY_PATH missing build/bin - - "latest" tag resolution fallback chain (Unsloth -> ggml-org -> raw) + - "latest" tag resolution fallback chain (helper only) - Cross-platform binary_env (Linux, macOS, Windows) - Edge cases: malformed JSON, empty responses, env overrides @@ -14,13 +14,11 @@ Run: pytest tests/studio/install/test_pr4562_bugfixes.py -v """ import importlib.util -import json import os import subprocess import sys import textwrap from pathlib import Path -from unittest.mock import patch import pytest @@ -40,6 +38,10 @@ SPEC.loader.exec_module(MOD) binary_env = MOD.binary_env HostInfo = MOD.HostInfo resolve_requested_llama_tag = MOD.resolve_requested_llama_tag +PublishedReleaseBundle = MOD.PublishedReleaseBundle +ApprovedArtifactHash = MOD.ApprovedArtifactHash +ApprovedReleaseChecksums = MOD.ApprovedReleaseChecksums +source_archive_logical_name = MOD.source_archive_logical_name SETUP_SH = PACKAGE_ROOT / "studio" / "setup.sh" SETUP_PS1 = PACKAGE_ROOT / "studio" / "setup.ps1" @@ -82,6 +84,9 @@ def run_bash(script: str, *, timeout: int = 10, env: dict | None = None) -> str: timeout = timeout, env = run_env, ) + assert ( + result.returncode == 0 + ), f"bash script failed (exit {result.returncode}):\n{result.stderr}" return result.stdout.strip() @@ -240,6 +245,107 @@ class TestResolveRequestedLlamaTag: monkeypatch.setattr(MOD, "latest_upstream_release_tag", lambda: "b5555") assert resolve_requested_llama_tag("") == "b5555" + def test_latest_with_published_repo_uses_latest_valid_published_release( + self, monkeypatch: pytest.MonkeyPatch + ): + invalid = PublishedReleaseBundle( + repo = "unslothai/llama.cpp", + release_tag = "v2.0", + upstream_tag = "b9000", + assets = {}, + manifest_asset_name = "llama-prebuilt-manifest.json", + artifacts = [], + selection_log = [], + ) + valid = PublishedReleaseBundle( + repo = "unslothai/llama.cpp", + release_tag = "v1.0", + upstream_tag = "b8999", + assets = {}, + manifest_asset_name = "llama-prebuilt-manifest.json", + artifacts = [], + selection_log = [], + ) + + monkeypatch.setattr( + MOD, + "iter_published_release_bundles", + lambda repo, published_release_tag = "": iter([invalid, valid]), + ) + + def fake_load(repo, release_tag): + if release_tag == "v2.0": + raise MOD.PrebuiltFallback("checksum asset missing") + return ApprovedReleaseChecksums( + repo = repo, + release_tag = release_tag, + upstream_tag = "b8999", + source_commit = None, + artifacts = { + source_archive_logical_name("b8999"): ApprovedArtifactHash( + asset_name = source_archive_logical_name("b8999"), + sha256 = "a" * 64, + repo = "ggml-org/llama.cpp", + kind = "upstream-source", + ) + }, + ) + + monkeypatch.setattr(MOD, "load_approved_release_checksums", fake_load) + monkeypatch.setattr(MOD, "latest_upstream_release_tag", lambda: "b7777") + + assert resolve_requested_llama_tag("latest", "unslothai/llama.cpp") == "b8999" + + def test_latest_with_published_release_tag_passes_pin_through( + self, monkeypatch: pytest.MonkeyPatch + ): + captured = {} + + def fake_resolve(requested_tag, published_repo, published_release_tag = ""): + captured["requested_tag"] = requested_tag + captured["published_repo"] = published_repo + captured["published_release_tag"] = published_release_tag + return MOD.ResolvedPublishedRelease( + bundle = PublishedReleaseBundle( + repo = published_repo, + release_tag = published_release_tag, + upstream_tag = "b9001", + assets = {}, + manifest_asset_name = "llama-prebuilt-manifest.json", + artifacts = [], + selection_log = [], + ), + checksums = ApprovedReleaseChecksums( + repo = published_repo, + release_tag = published_release_tag, + upstream_tag = "b9001", + artifacts = { + source_archive_logical_name("b9001"): ApprovedArtifactHash( + asset_name = source_archive_logical_name("b9001"), + sha256 = "a" * 64, + repo = "ggml-org/llama.cpp", + kind = "upstream-source", + ) + }, + ), + ) + + monkeypatch.setattr(MOD, "resolve_published_release", fake_resolve) + + assert ( + resolve_requested_llama_tag( + "latest", + "unslothai/llama.cpp", + "llama-prebuilt-main", + ) + == "b9001" + ) + assert captured == { + "requested_tag": "latest", + "published_repo": "unslothai/llama.cpp", + "published_release_tag": "llama-prebuilt-main", + } + # ========================================================================= # TEST GROUP C: setup.sh logic (bash subprocess tests) @@ -429,140 +535,61 @@ class TestSetupShLogic: # TEST GROUP D: "latest" tag resolution (bash subprocess) # ========================================================================= class TestLatestTagResolution: - """Test the fallback chain: Unsloth API -> ggml-org API -> raw.""" + """Test the fallback chain: helper resolver -> raw.""" RESOLVE_TEMPLATE = textwrap.dedent("""\ - export PATH="{mock_bin}:$PATH" _REQUESTED_LLAMA_TAG="{requested_tag}" _RESOLVED_LLAMA_TAG="" - _RESOLVE_UPSTREAM_STATUS=1 - _HELPER_RELEASE_REPO="unslothai/llama.cpp" - if [ "$_RESOLVE_UPSTREAM_STATUS" -ne 0 ] || [ -z "$_RESOLVED_LLAMA_TAG" ]; then - if [ "$_REQUESTED_LLAMA_TAG" = "latest" ]; then - _RESOLVED_LLAMA_TAG="$(curl -fsSL "https://api.github.com/repos/${{_HELPER_RELEASE_REPO}}/releases/latest" 2>/dev/null | python -c "import sys,json; print(json.load(sys.stdin)['tag_name'])" 2>/dev/null)" || _RESOLVED_LLAMA_TAG="" - if [ -z "$_RESOLVED_LLAMA_TAG" ]; then - _RESOLVED_LLAMA_TAG="$(curl -fsSL https://api.github.com/repos/ggml-org/llama.cpp/releases/latest 2>/dev/null | python -c "import sys,json; print(json.load(sys.stdin)['tag_name'])" 2>/dev/null)" || _RESOLVED_LLAMA_TAG="" - fi - fi - if [ -z "$_RESOLVED_LLAMA_TAG" ]; then - _RESOLVED_LLAMA_TAG="$_REQUESTED_LLAMA_TAG" - fi + _RESOLVE_UPSTREAM_STATUS={resolve_status} + if [ "$_RESOLVE_UPSTREAM_STATUS" -eq 0 ] && [ -n "{resolved_tag}" ]; then + _RESOLVED_LLAMA_TAG="{resolved_tag}" + else + _RESOLVED_LLAMA_TAG="$_REQUESTED_LLAMA_TAG" fi echo "$_RESOLVED_LLAMA_TAG" """) - @staticmethod - def _make_curl_mock( - mock_bin: Path, unsloth_response: str | None, ggml_response: str | None - ): - """Create a curl mock that returns different responses per repo.""" - lines = ["#!/bin/bash"] - if unsloth_response is not None: - lines.append( - f'if echo "$*" | grep -q "unslothai/llama.cpp"; then echo \'{unsloth_response}\'; exit 0; fi' - ) - else: - lines.append( - 'if echo "$*" | grep -q "unslothai/llama.cpp"; then exit 1; fi' - ) - if ggml_response is not None: - lines.append( - f'if echo "$*" | grep -q "ggml-org/llama.cpp"; then echo \'{ggml_response}\'; exit 0; fi' - ) - else: - lines.append('if echo "$*" | grep -q "ggml-org/llama.cpp"; then exit 1; fi') - lines.append("exit 1") - curl_path = mock_bin / "curl" - curl_path.write_text("\n".join(lines) + "\n") - curl_path.chmod(0o755) - def _run_resolve( self, tmp_path: Path, requested_tag: str, - unsloth_resp: str | None, - ggml_resp: str | None, + resolved_tag: str, + resolve_status: int, ) -> str: - mock_bin = tmp_path / "mock_bin" - mock_bin.mkdir(exist_ok = True) - self._make_curl_mock(mock_bin, unsloth_resp, ggml_resp) script = self.RESOLVE_TEMPLATE.format( - mock_bin = mock_bin, requested_tag = requested_tag + requested_tag = requested_tag, + resolved_tag = resolved_tag, + resolve_status = resolve_status, ) return run_bash(script) - def test_unsloth_succeeds(self, tmp_path: Path): + def test_helper_resolution_succeeds(self, tmp_path: Path): output = self._run_resolve( tmp_path, "latest", - unsloth_resp = '{"tag_name":"b8508"}', - ggml_resp = '{"tag_name":"b9000"}', + resolved_tag = "b8508", + resolve_status = 0, ) assert output == "b8508" - def test_unsloth_fails_ggml_succeeds(self, tmp_path: Path): + def test_helper_resolution_falls_back_to_raw_requested_tag(self, tmp_path: Path): output = self._run_resolve( tmp_path, "latest", - unsloth_resp = None, - ggml_resp = '{"tag_name":"b9000"}', - ) - assert output == "b9000" - - def test_both_fail_raw_fallback(self, tmp_path: Path): - output = self._run_resolve( - tmp_path, - "latest", - unsloth_resp = None, - ggml_resp = None, + resolved_tag = "", + resolve_status = 1, ) assert output == "latest" - def test_concrete_tag_passes_through(self, tmp_path: Path): + def test_concrete_tag_passes_through_when_helper_fails(self, tmp_path: Path): output = self._run_resolve( tmp_path, "b7777", - unsloth_resp = '{"tag_name":"b8508"}', - ggml_resp = '{"tag_name":"b9000"}', + resolved_tag = "", + resolve_status = 1, ) assert output == "b7777" - def test_unsloth_malformed_json_falls_through(self, tmp_path: Path): - output = self._run_resolve( - tmp_path, - "latest", - unsloth_resp = '{"bad_key":"no_tag"}', - ggml_resp = '{"tag_name":"b9001"}', - ) - assert output == "b9001" - - def test_both_malformed_json_raw_fallback(self, tmp_path: Path): - output = self._run_resolve( - tmp_path, - "latest", - unsloth_resp = '{"bad":"data"}', - ggml_resp = '{"also":"bad"}', - ) - assert output == "latest" - - def test_unsloth_empty_body_falls_through(self, tmp_path: Path): - output = self._run_resolve( - tmp_path, - "latest", - unsloth_resp = "", - ggml_resp = '{"tag_name":"b7000"}', - ) - assert output == "b7000" - - def test_unsloth_empty_tag_name_falls_through(self, tmp_path: Path): - output = self._run_resolve( - tmp_path, - "latest", - unsloth_resp = '{"tag_name":""}', - ggml_resp = '{"tag_name":"b6000"}', - ) - assert output == "b6000" - def test_env_override_unsloth_llama_tag(self): output = run_bash( 'echo "${UNSLOTH_LLAMA_TAG:-latest}"', @@ -593,10 +620,10 @@ class TestSourceCodePatterns: def test_setup_sh_no_rm_before_prereq_check(self): """rm -rf must appear AFTER cmake/git checks, not before.""" content = SETUP_SH.read_text() - # Find the source-build block - idx_else = content.find("# Check prerequisites") - assert idx_else != -1 - block = content[idx_else:] + # Anchor on the source-build cmake check block. + idx_block = content.find("command -v cmake") + assert idx_block != -1 + block = content[idx_block:] # rm -rf should appear after the cmake/git checks idx_cmake = block.find("command -v cmake") idx_git = block.find("command -v git") @@ -605,28 +632,102 @@ class TestSourceCodePatterns: assert idx_rm > idx_git, "rm -rf should come after git check" def test_setup_sh_clone_uses_branch_tag(self): - """git clone in source-build should use --branch via _CLONE_BRANCH_ARGS.""" + """git clone in source-build should use --branch via the clone args array.""" content = SETUP_SH.read_text() - # The clone line should use _CLONE_BRANCH_ARGS (which conditionally includes --branch) + assert "_CLONE_ARGS=(git clone --depth 1)" in content assert ( - "_CLONE_BRANCH_ARGS" in content - ), "Clone should use _CLONE_BRANCH_ARGS array" - assert ( - '--branch "$_RESOLVED_LLAMA_TAG"' in content - ), "_CLONE_BRANCH_ARGS should be set to --branch $_RESOLVED_LLAMA_TAG" + '_CLONE_ARGS+=(--branch "$_RESOLVED_SOURCE_REF")' in content + ), "_CLONE_ARGS should be extended with --branch $_RESOLVED_SOURCE_REF" # Verify the guard: --branch is only used when tag is not "latest" assert ( - '_RESOLVED_LLAMA_TAG" != "latest"' in content + '_RESOLVED_SOURCE_REF" != "latest"' in content ), "Should guard against literal 'latest' tag" - def test_setup_sh_latest_resolution_queries_unsloth_first(self): - """The Unsloth repo should be queried before ggml-org.""" + def test_setup_sh_source_build_uses_helper_resolution(self): + """Shell source fallback should consult the helper for repo/ref planning.""" content = SETUP_SH.read_text() - idx_unsloth = content.find("_HELPER_RELEASE_REPO}/releases/latest") - idx_ggml = content.find("ggml-org/llama.cpp/releases/latest") - assert idx_unsloth != -1, "Unsloth API query not found" - assert idx_ggml != -1, "ggml-org API query not found" - assert idx_unsloth < idx_ggml, "Unsloth should be queried before ggml-org" + assert "--resolve-source-build" in content + assert "--output-format json" in content + assert "_RESOLVED_SOURCE_URL" in content + assert "_RESOLVED_SOURCE_REF_KIND" in content + assert "_RESOLVED_SOURCE_REF" in content + + def test_setup_sh_latest_resolution_uses_helper_only(self): + """Shell fallback should rely on helper output, not raw GitHub API tag_name.""" + content = SETUP_SH.read_text() + assert "--resolve-install-tag" in content + assert "--resolve-llama-tag" in content + assert 'tail -n 1 "$_RESOLVE_LLAMA_LOG"' not in content + assert "json.load" in content + assert "_HELPER_RELEASE_REPO}/releases/latest" not in content + assert "ggml-org/llama.cpp/releases/latest" not in content + + def test_setup_sh_macos_arm64_uses_metal_flags(self): + """Apple Silicon source builds should explicitly enable Metal like upstream.""" + content = SETUP_SH.read_text() + assert "_IS_MACOS_ARM64=true" in content + assert 'if [ "$_IS_MACOS_ARM64" = true ]; then' in content + assert "-DGGML_METAL=ON" in content + assert "-DGGML_METAL_EMBED_LIBRARY=ON" in content + assert "-DGGML_METAL_USE_BF16=ON" in content + assert "-DCMAKE_INSTALL_RPATH=@loader_path" in content + assert "-DCMAKE_BUILD_WITH_INSTALL_RPATH=ON" in content + + def test_setup_sh_macos_metal_configure_has_cpu_fallback(self): + """If Metal configure or build fails, setup should retry with CPU fallback.""" + content = SETUP_SH.read_text() + assert "_TRY_METAL_CPU_FALLBACK=true" in content + assert ( + 'substep "Metal configure failed; retrying CPU build..." "$C_WARN"' + in content + ) + assert ( + 'substep "Metal build failed; retrying CPU build..." "$C_WARN"' in content + ) + assert 'run_quiet_no_exit "cmake llama.cpp (cpu fallback)"' in content + assert "-DGGML_METAL=OFF" in content + # _TRY_METAL_CPU_FALLBACK must be reset to false in both fallback branches + # (1 init + 2 resets = at least 3 occurrences of =false) + assert content.count("_TRY_METAL_CPU_FALLBACK=false") >= 3, ( + "_TRY_METAL_CPU_FALLBACK=false should appear at least 3 times " + "(init + configure fallback + build fallback)" + ) + + def test_macos_arm64_cpu_fallback_args_exclude_rpath(self): + """CPU fallback args must NOT contain Metal-only RPATH flags at runtime.""" + script = ( + '_IS_MACOS_ARM64=true\nNVCC_PATH=""\nGPU_BACKEND=""\n' + + _GPU_BACKEND_FRAGMENT + ) + output = run_bash(script) + fallback_line = next( + line + for line in output.splitlines() + if line.startswith("CPU_FALLBACK_CMAKE_ARGS=") + ) + assert "-DGGML_METAL=OFF" in fallback_line + assert ( + "@loader_path" not in fallback_line + ), "CPU fallback args should not contain RPATH flags" + assert ( + "-DCMAKE_BUILD_WITH_INSTALL_RPATH=ON" not in fallback_line + ), "CPU fallback args should not contain RPATH build flag" + + def test_setup_sh_does_not_enable_metal_for_intel_macos(self): + """Intel macOS should stay on the existing non-Metal path in this patch.""" + content = SETUP_SH.read_text() + assert 'if [ "$_IS_MACOS_ARM64" = true ]; then' in content + assert ( + 'Darwin" ] && { [ "$_HOST_MACHINE" = "arm64" ] || [ "$_HOST_MACHINE" = "aarch64" ]; }' + in content + ) + assert ( + "x86_64" + not in content[ + content.find("-DGGML_METAL=ON") - 200 : content.find("-DGGML_METAL=ON") + + 200 + ] + ) def test_setup_ps1_uses_checkout_b(self): """PS1 should use checkout -B, not checkout --force FETCH_HEAD.""" @@ -637,7 +738,7 @@ class TestSourceCodePatterns: def test_setup_ps1_clone_uses_branch_tag(self): """PS1 clone should use --branch with the resolved tag.""" content = SETUP_PS1.read_text() - assert "--branch" in content and "$ResolvedLlamaTag" in content + assert "--branch" in content and "$ResolvedSourceRef" in content # The old commented-out line should be gone assert "# git clone --depth 1 --branch" not in content @@ -658,14 +759,52 @@ class TestSourceCodePatterns: f"Found 'git pull' in llama.cpp build section at line {i+1}" ) - def test_setup_ps1_latest_resolution_queries_unsloth_first(self): - """PS1 should query Unsloth repo before ggml-org.""" + def test_setup_ps1_latest_resolution_uses_helper_only(self): + """PS1 fallback should rely on helper output, not raw GitHub API tag_name.""" content = SETUP_PS1.read_text() - idx_unsloth = content.find("$HelperReleaseRepo/releases/latest") - idx_ggml = content.find("ggml-org/llama.cpp/releases/latest") - assert idx_unsloth != -1, "Unsloth API query not found in PS1" - assert idx_ggml != -1, "ggml-org API query not found in PS1" - assert idx_unsloth < idx_ggml, "Unsloth should be queried before ggml-org" + assert "--resolve-install-tag" in content + assert "--resolve-llama-tag" in content + assert '--output-format", "json"' in content + assert "ConvertFrom-Json" in content + assert "$HelperReleaseRepo/releases/latest" not in content + assert "ggml-org/llama.cpp/releases/latest" not in content + + def test_setup_ps1_source_build_uses_helper_resolution(self): + """PS1 source fallback should consult the helper for repo/ref planning.""" + content = SETUP_PS1.read_text() + assert "--resolve-source-build" in content + assert '--output-format", "json"' in content + assert "$ResolvedSourceUrl" in content + assert "$ResolvedSourceRefKind" in content + assert "$ResolvedSourceRef" in content + + def test_setup_ps1_prebuilt_install_disables_native_error_abort(self): + """PS1 prebuilt install should not abort setup on helper stderr.""" + content = SETUP_PS1.read_text() + install_idx = content.index("& python @prebuiltArgs 2>&1") + block = content[max(0, install_idx - 800) : install_idx + 800] + assert "$PSNativeCommandUseErrorActionPreference = $false" in block + assert "$restoreNativeErrorPreference = $true" in block + assert ( + "$PSNativeCommandUseErrorActionPreference = $previousNativeErrorPreference" + in block + ) + + def test_setup_ps1_helper_disables_error_action_abort(self): + """Helper resolution should suppress terminating NativeCommandError on PS 5.1.""" + content = SETUP_PS1.read_text() + helper_idx = content.index("function Invoke-LlamaHelper") + block = content[helper_idx : helper_idx + 1200] + assert "$previousErrorActionPreference = $ErrorActionPreference" in block + assert '$ErrorActionPreference = "Continue"' in block + assert "$ErrorActionPreference = $previousErrorActionPreference" in block + + def test_setup_ps1_uses_local_tempfile_helper(self): + """PS1 should not depend on New-TemporaryFile being available.""" + content = SETUP_PS1.read_text() + assert "function New-UnslothTemporaryFile" in content + assert "$resolveErrorLog = New-UnslothTemporaryFile" in content + assert "$resolveErrorLog = New-TemporaryFile" not in content def test_binary_env_linux_has_binary_parent(self): """The Linux branch of binary_env should include binary_path.parent.""" @@ -685,3 +824,274 @@ class TestSourceCodePatterns: found = True break assert found, "binary_path.parent not found in Linux branch of binary_env" + + +# ========================================================================= +# TEST GROUP F: macOS Metal build logic (bash subprocess tests) +# ========================================================================= + +# Minimal bash fragment that mirrors setup.sh's GPU backend decision chain. +# Variables _IS_MACOS_ARM64, NVCC_PATH, GPU_BACKEND are injected by tests. +_GPU_BACKEND_FRAGMENT = textwrap.dedent("""\ + CMAKE_ARGS="-DLLAMA_BUILD_TESTS=OFF" + _TRY_METAL_CPU_FALLBACK=false + CPU_FALLBACK_CMAKE_ARGS="$CMAKE_ARGS" + + _BUILD_DESC="building" + if [ "$_IS_MACOS_ARM64" = true ]; then + _BUILD_DESC="building (Metal)" + CMAKE_ARGS="$CMAKE_ARGS -DGGML_METAL=ON -DGGML_METAL_EMBED_LIBRARY=ON -DGGML_METAL_USE_BF16=ON -DCMAKE_INSTALL_RPATH=@loader_path -DCMAKE_BUILD_WITH_INSTALL_RPATH=ON" + CPU_FALLBACK_CMAKE_ARGS="$CPU_FALLBACK_CMAKE_ARGS -DGGML_METAL=OFF" + _TRY_METAL_CPU_FALLBACK=true + elif [ -n "$NVCC_PATH" ]; then + CMAKE_ARGS="$CMAKE_ARGS -DGGML_CUDA=ON" + _BUILD_DESC="building (CUDA)" + elif [ "$GPU_BACKEND" = "rocm" ]; then + CMAKE_ARGS="$CMAKE_ARGS -DGGML_HIP=ON" + _BUILD_DESC="building (ROCm)" + else + _BUILD_DESC="building (CPU)" + fi + + echo "CMAKE_ARGS=$CMAKE_ARGS" + echo "CPU_FALLBACK_CMAKE_ARGS=$CPU_FALLBACK_CMAKE_ARGS" + echo "BUILD_DESC=$_BUILD_DESC" + echo "TRY_METAL_CPU_FALLBACK=$_TRY_METAL_CPU_FALLBACK" +""") + + +class TestMacOSMetalBuildLogic: + """Behavioral bash subprocess tests for the Metal GPU backend logic.""" + + def test_macos_arm64_cmake_args_contain_metal_flags(self): + """macOS arm64 should enable Metal, not CUDA.""" + script = ( + '_IS_MACOS_ARM64=true\nNVCC_PATH=""\nGPU_BACKEND=""\n' + + _GPU_BACKEND_FRAGMENT + ) + output = run_bash(script) + assert "-DGGML_METAL=ON" in output + assert "-DGGML_CUDA=ON" not in output + assert "BUILD_DESC=building (Metal)" in output + + def test_intel_macos_no_metal_flags(self): + """Intel macOS (not arm64) should not get Metal flags.""" + script = ( + '_IS_MACOS_ARM64=false\nNVCC_PATH=""\nGPU_BACKEND=""\n' + + _GPU_BACKEND_FRAGMENT + ) + output = run_bash(script) + assert "-DGGML_METAL=ON" not in output + assert "BUILD_DESC=building (CPU)" in output + + def test_macos_arm64_metal_precedes_nvcc(self): + """Even with nvcc in PATH, macOS arm64 should use Metal, not CUDA.""" + script = ( + '_IS_MACOS_ARM64=true\nNVCC_PATH="/usr/local/cuda/bin/nvcc"\n' + 'GPU_BACKEND="cuda"\n' + _GPU_BACKEND_FRAGMENT + ) + output = run_bash(script) + assert "-DGGML_METAL=ON" in output + assert "-DGGML_CUDA=ON" not in output + assert "BUILD_DESC=building (Metal)" in output + + def test_metal_cpu_fallback_triggers_on_cmake_failure(self, tmp_path: Path): + """When cmake fails on Metal, the fallback should retry with -DGGML_METAL=OFF.""" + mock_bin = tmp_path / "mock_bin" + mock_bin.mkdir() + calls_file = tmp_path / "cmake_calls.log" + # cmake that logs args and fails on first call (Metal), succeeds on second (CPU fallback) + cmake_script = mock_bin / "cmake" + cmake_script.write_text( + textwrap.dedent(f"""\ + #!/bin/bash + echo "$*" >> "{calls_file}" + COUNTER_FILE="{tmp_path}/cmake_counter" + if [ ! -f "$COUNTER_FILE" ]; then + echo 1 > "$COUNTER_FILE" + exit 1 + fi + exit 0 + """) + ) + cmake_script.chmod(0o755) + + script = textwrap.dedent(f"""\ + export PATH="{mock_bin}:$PATH" + _IS_MACOS_ARM64=true + NVCC_PATH="" + GPU_BACKEND="" + CMAKE_ARGS="-DLLAMA_BUILD_TESTS=OFF" + _TRY_METAL_CPU_FALLBACK=false + CPU_FALLBACK_CMAKE_ARGS="$CMAKE_ARGS" + + _BUILD_DESC="building" + if [ "$_IS_MACOS_ARM64" = true ]; then + _BUILD_DESC="building (Metal)" + CMAKE_ARGS="$CMAKE_ARGS -DGGML_METAL=ON -DGGML_METAL_EMBED_LIBRARY=ON -DGGML_METAL_USE_BF16=ON -DCMAKE_INSTALL_RPATH=@loader_path -DCMAKE_BUILD_WITH_INSTALL_RPATH=ON" + CPU_FALLBACK_CMAKE_ARGS="$CPU_FALLBACK_CMAKE_ARGS -DGGML_METAL=OFF" + _TRY_METAL_CPU_FALLBACK=true + fi + + BUILD_OK=true + _BUILD_TMP="{tmp_path}/build_tmp" + mkdir -p "$_BUILD_TMP" + if ! cmake -S "$_BUILD_TMP" -B "$_BUILD_TMP/build" $CMAKE_ARGS; then + if [ "$_TRY_METAL_CPU_FALLBACK" = true ]; then + _TRY_METAL_CPU_FALLBACK=false + echo "FALLBACK_TRIGGERED" + rm -rf "$_BUILD_TMP/build" + cmake -S "$_BUILD_TMP" -B "$_BUILD_TMP/build" $CPU_FALLBACK_CMAKE_ARGS || BUILD_OK=false + if [ "$BUILD_OK" = true ]; then + _BUILD_DESC="building (CPU fallback)" + fi + else + BUILD_OK=false + fi + fi + + echo "BUILD_OK=$BUILD_OK" + echo "BUILD_DESC=$_BUILD_DESC" + echo "TRY_METAL_CPU_FALLBACK=$_TRY_METAL_CPU_FALLBACK" + """) + output = run_bash(script) + assert "FALLBACK_TRIGGERED" in output + assert "BUILD_OK=true" in output + assert "BUILD_DESC=building (CPU fallback)" in output + assert ( + "TRY_METAL_CPU_FALLBACK=false" in output + ), "Fallback flag should be reset to false after configure fallback" + + # Verify cmake args: first call has Metal ON, second has Metal OFF + calls = calls_file.read_text().splitlines() + assert len(calls) >= 2, f"Expected >= 2 cmake calls, got {len(calls)}" + assert ( + "-DGGML_METAL=ON" in calls[0] + ), f"First cmake call should have Metal ON: {calls[0]}" + assert ( + "-DGGML_METAL=OFF" in calls[1] + ), f"Second cmake call should have Metal OFF: {calls[1]}" + assert ( + "-DGGML_METAL=ON" not in calls[1] + ), f"Second cmake call should NOT have Metal ON: {calls[1]}" + assert ( + "@loader_path" not in calls[1] + ), f"CPU fallback should not have RPATH: {calls[1]}" + assert ( + "-DCMAKE_BUILD_WITH_INSTALL_RPATH=ON" not in calls[1] + ), f"CPU fallback should not have RPATH build flag: {calls[1]}" + + def test_metal_build_failure_retries_cpu_fallback(self, tmp_path: Path): + """When cmake --build fails on Metal, the fallback should re-configure and rebuild with CPU.""" + mock_bin = tmp_path / "mock_bin" + mock_bin.mkdir() + calls_file = tmp_path / "cmake_calls.log" + # cmake mock: configure always succeeds; first --build fails, rest succeed + cmake_script = mock_bin / "cmake" + cmake_script.write_text( + textwrap.dedent(f"""\ + #!/bin/bash + echo "$*" >> "{calls_file}" + if [ "$1" = "--build" ]; then + BUILD_COUNTER_FILE="{tmp_path}/build_counter" + if [ ! -f "$BUILD_COUNTER_FILE" ]; then + echo 1 > "$BUILD_COUNTER_FILE" + exit 1 + fi + fi + exit 0 + """) + ) + cmake_script.chmod(0o755) + + script = textwrap.dedent(f"""\ + export PATH="{mock_bin}:$PATH" + _IS_MACOS_ARM64=true + NVCC_PATH="" + GPU_BACKEND="" + CMAKE_ARGS="-DLLAMA_BUILD_TESTS=OFF" + _TRY_METAL_CPU_FALLBACK=false + CPU_FALLBACK_CMAKE_ARGS="$CMAKE_ARGS" + CMAKE_GENERATOR_ARGS="" + NCPU=2 + + _BUILD_DESC="building" + if [ "$_IS_MACOS_ARM64" = true ]; then + _BUILD_DESC="building (Metal)" + CMAKE_ARGS="$CMAKE_ARGS -DGGML_METAL=ON -DGGML_METAL_EMBED_LIBRARY=ON -DGGML_METAL_USE_BF16=ON -DCMAKE_INSTALL_RPATH=@loader_path -DCMAKE_BUILD_WITH_INSTALL_RPATH=ON" + CPU_FALLBACK_CMAKE_ARGS="$CPU_FALLBACK_CMAKE_ARGS -DGGML_METAL=OFF" + _TRY_METAL_CPU_FALLBACK=true + fi + + BUILD_OK=true + _BUILD_TMP="{tmp_path}/build_tmp" + mkdir -p "$_BUILD_TMP" + + # Configure (succeeds) + if ! cmake $CMAKE_GENERATOR_ARGS -S "$_BUILD_TMP" -B "$_BUILD_TMP/build" $CMAKE_ARGS; then + if [ "$_TRY_METAL_CPU_FALLBACK" = true ]; then + _TRY_METAL_CPU_FALLBACK=false + echo "CONFIGURE_FALLBACK" + rm -rf "$_BUILD_TMP/build" + cmake $CMAKE_GENERATOR_ARGS -S "$_BUILD_TMP" -B "$_BUILD_TMP/build" $CPU_FALLBACK_CMAKE_ARGS || BUILD_OK=false + if [ "$BUILD_OK" = true ]; then + _BUILD_DESC="building (CPU fallback)" + fi + else + BUILD_OK=false + fi + fi + + # Build (first --build fails, triggers fallback) + if [ "$BUILD_OK" = true ]; then + if ! cmake --build "$_BUILD_TMP/build" --config Release --target llama-server -j"$NCPU"; then + if [ "$_TRY_METAL_CPU_FALLBACK" = true ]; then + _TRY_METAL_CPU_FALLBACK=false + echo "BUILD_FALLBACK_TRIGGERED" + rm -rf "$_BUILD_TMP/build" + if cmake $CMAKE_GENERATOR_ARGS -S "$_BUILD_TMP" -B "$_BUILD_TMP/build" $CPU_FALLBACK_CMAKE_ARGS; then + _BUILD_DESC="building (CPU fallback)" + cmake --build "$_BUILD_TMP/build" --config Release --target llama-server -j"$NCPU" || BUILD_OK=false + else + BUILD_OK=false + fi + else + BUILD_OK=false + fi + fi + fi + + echo "BUILD_OK=$BUILD_OK" + echo "BUILD_DESC=$_BUILD_DESC" + echo "TRY_METAL_CPU_FALLBACK=$_TRY_METAL_CPU_FALLBACK" + """) + output = run_bash(script) + assert "CONFIGURE_FALLBACK" not in output, "Configure should have succeeded" + assert "BUILD_FALLBACK_TRIGGERED" in output + assert "BUILD_OK=true" in output + assert "BUILD_DESC=building (CPU fallback)" in output + assert ( + "TRY_METAL_CPU_FALLBACK=false" in output + ), "Fallback flag should be reset to false after build fallback" + + # Verify: configure with Metal ON, build fails, re-configure with Metal OFF, rebuild + calls = calls_file.read_text().splitlines() + assert len(calls) >= 4, f"Expected >= 4 cmake calls, got {len(calls)}: {calls}" + # First call: configure with Metal ON + assert "-DGGML_METAL=ON" in calls[0] + # Second call: build (fails) + assert "--build" in calls[1] + # Third call: re-configure with Metal OFF and no RPATH flags + assert "-DGGML_METAL=OFF" in calls[2] + assert "-DGGML_METAL=ON" not in calls[2] + assert ( + "@loader_path" not in calls[2] + ), f"CPU fallback should not have RPATH: {calls[2]}" + assert ( + "-DCMAKE_BUILD_WITH_INSTALL_RPATH=ON" not in calls[2] + ), f"CPU fallback should not have RPATH build flag: {calls[2]}" + assert ( + "-DLLAMA_BUILD_TESTS=OFF" in calls[2] + ), f"CPU fallback should preserve baseline flags: {calls[2]}" + # Fourth call: rebuild (succeeds) + assert "--build" in calls[3] diff --git a/tests/studio/install/test_selection_logic.py b/tests/studio/install/test_selection_logic.py index 906c978b0d..6a2e367dc8 100644 --- a/tests/studio/install/test_selection_logic.py +++ b/tests/studio/install/test_selection_logic.py @@ -54,6 +54,26 @@ apply_approved_hashes = INSTALL_LLAMA_PREBUILT.apply_approved_hashes linux_cuda_choice_from_release = INSTALL_LLAMA_PREBUILT.linux_cuda_choice_from_release windows_cuda_attempts = INSTALL_LLAMA_PREBUILT.windows_cuda_attempts resolve_upstream_asset_choice = INSTALL_LLAMA_PREBUILT.resolve_upstream_asset_choice +resolve_requested_install_tag = INSTALL_LLAMA_PREBUILT.resolve_requested_install_tag +resolve_install_attempts = INSTALL_LLAMA_PREBUILT.resolve_install_attempts +resolve_install_release_plans = INSTALL_LLAMA_PREBUILT.resolve_install_release_plans +resolve_published_release = INSTALL_LLAMA_PREBUILT.resolve_published_release +resolve_source_build_plan = INSTALL_LLAMA_PREBUILT.resolve_source_build_plan +validated_checksums_for_bundle = INSTALL_LLAMA_PREBUILT.validated_checksums_for_bundle +parse_approved_release_checksums = ( + INSTALL_LLAMA_PREBUILT.parse_approved_release_checksums +) +published_release_matches_request = ( + INSTALL_LLAMA_PREBUILT.published_release_matches_request +) +exact_source_archive_logical_name = ( + INSTALL_LLAMA_PREBUILT.exact_source_archive_logical_name +) +source_archive_logical_name = INSTALL_LLAMA_PREBUILT.source_archive_logical_name +windows_cuda_upstream_asset_names = ( + INSTALL_LLAMA_PREBUILT.windows_cuda_upstream_asset_names +) +env_int = INSTALL_LLAMA_PREBUILT.env_int # --------------------------------------------------------------------------- @@ -104,6 +124,13 @@ def make_release(artifacts, **overrides): repo = "unslothai/llama.cpp", release_tag = "v1.0", upstream_tag = "b8508", + source_repo = None, + source_repo_url = None, + source_ref_kind = None, + requested_source_ref = None, + resolved_source_ref = None, + source_commit = None, + source_commit_short = None, assets = {a.asset_name: f"https://example.com/{a.asset_name}" for a in artifacts}, manifest_asset_name = "llama-prebuilt-manifest.json", artifacts = artifacts, @@ -118,7 +145,13 @@ def make_checksums(asset_names): repo = "unslothai/llama.cpp", release_tag = "v1.0", upstream_tag = "b8508", + source_repo = None, + source_repo_url = None, + source_ref_kind = None, + requested_source_ref = None, + resolved_source_ref = None, source_commit = None, + source_commit_short = None, artifacts = { name: ApprovedArtifactHash( asset_name = name, @@ -131,6 +164,64 @@ def make_checksums(asset_names): ) +def make_checksums_with_source( + asset_names, + *, + release_tag = "v1.0", + upstream_tag = "b8508", + source_repo = None, + source_repo_url = None, + source_ref_kind = None, + requested_source_ref = None, + resolved_source_ref = None, + source_commit = None, +): + artifacts = { + **{ + name: ApprovedArtifactHash( + asset_name = name, + sha256 = "a" * 64, + repo = "unslothai/llama.cpp", + kind = "prebuilt", + ) + for name in asset_names + }, + source_archive_logical_name(upstream_tag): ApprovedArtifactHash( + asset_name = source_archive_logical_name(upstream_tag), + sha256 = "b" * 64, + repo = "ggml-org/llama.cpp", + kind = "upstream-source", + ), + } + normalized_source_commit = ( + source_commit.lower() if isinstance(source_commit, str) else None + ) + if normalized_source_commit: + artifacts[exact_source_archive_logical_name(normalized_source_commit)] = ( + ApprovedArtifactHash( + asset_name = exact_source_archive_logical_name(normalized_source_commit), + sha256 = "c" * 64, + repo = source_repo or "example/custom-llama.cpp", + kind = "exact-source", + ) + ) + return ApprovedReleaseChecksums( + repo = "unslothai/llama.cpp", + release_tag = release_tag, + upstream_tag = upstream_tag, + source_repo = source_repo, + source_repo_url = source_repo_url, + source_ref_kind = source_ref_kind, + requested_source_ref = requested_source_ref, + resolved_source_ref = resolved_source_ref, + source_commit = normalized_source_commit, + source_commit_short = normalized_source_commit[:7] + if normalized_source_commit + else None, + artifacts = artifacts, + ) + + def mock_linux_runtime(monkeypatch, lines): dirs = {line: ["/usr/lib/stub"] for line in lines} monkeypatch.setattr( @@ -395,6 +486,81 @@ class TestApplyApprovedHashes: assert len(result) == 1 assert result[0].name == "a.tar.gz" + def test_upstream_asset_can_match_compatibility_tag_name(self): + choice = AssetChoice( + repo = UPSTREAM_REPO, + tag = "main", + name = "llama-main-bin-macos-arm64.tar.gz", + url = "https://x/llama-main-bin-macos-arm64.tar.gz", + source_label = "upstream", + ) + checksums = ApprovedReleaseChecksums( + repo = "unslothai/llama.cpp", + release_tag = "r1", + upstream_tag = "b9000", + artifacts = { + "llama-b9000-bin-macos-arm64.tar.gz": ApprovedArtifactHash( + asset_name = "llama-b9000-bin-macos-arm64.tar.gz", + sha256 = "a" * 64, + repo = UPSTREAM_REPO, + kind = "macos-arm64-upstream", + ) + }, + ) + + result = apply_approved_hashes([choice], checksums) + assert result[0].expected_sha256 == "a" * 64 + + def test_windows_cuda_legacy_choice_can_match_current_upstream_name(self): + choice = AssetChoice( + repo = UPSTREAM_REPO, + tag = "b9000", + name = "llama-b9000-bin-win-cuda-13.1-x64.zip", + url = "https://x/llama-b9000-bin-win-cuda-13.1-x64.zip", + source_label = "upstream", + ) + checksums = ApprovedReleaseChecksums( + repo = "unslothai/llama.cpp", + release_tag = "r1", + upstream_tag = "b9000", + artifacts = { + "cudart-llama-bin-win-cuda-13.1-x64.zip": ApprovedArtifactHash( + asset_name = "cudart-llama-bin-win-cuda-13.1-x64.zip", + sha256 = "b" * 64, + repo = UPSTREAM_REPO, + kind = "windows-cuda-upstream", + ) + }, + ) + + result = apply_approved_hashes([choice], checksums) + assert result[0].expected_sha256 == "b" * 64 + + def test_windows_cuda_current_choice_can_match_legacy_compatibility_name(self): + choice = AssetChoice( + repo = UPSTREAM_REPO, + tag = "main", + name = "cudart-llama-bin-win-cuda-13.1-x64.zip", + url = "https://x/cudart-llama-bin-win-cuda-13.1-x64.zip", + source_label = "upstream", + ) + checksums = ApprovedReleaseChecksums( + repo = "unslothai/llama.cpp", + release_tag = "r1", + upstream_tag = "b9000", + artifacts = { + "llama-b9000-bin-win-cuda-13.1-x64.zip": ApprovedArtifactHash( + asset_name = "llama-b9000-bin-win-cuda-13.1-x64.zip", + sha256 = "c" * 64, + repo = UPSTREAM_REPO, + kind = "windows-cuda-upstream", + ) + }, + ) + + result = apply_approved_hashes([choice], checksums) + assert result[0].expected_sha256 == "c" * 64 + def test_none_approved(self): c1 = self._choice("missing.tar.gz") checksums = make_checksums(["other.tar.gz"]) @@ -408,7 +574,353 @@ class TestApplyApprovedHashes: # =========================================================================== -# J. linux_cuda_choice_from_release -- core selection +# J. published release resolution +# =========================================================================== + + +class TestPublishedReleaseResolution: + def test_latest_skips_invalid_release_and_uses_next_valid(self, monkeypatch): + invalid = make_release([], release_tag = "v2.0", upstream_tag = "b9000") + valid = make_release([], release_tag = "v1.0", upstream_tag = "b8999") + + monkeypatch.setattr( + INSTALL_LLAMA_PREBUILT, + "iter_published_release_bundles", + lambda repo, published_release_tag = "": iter([invalid, valid]), + ) + + def fake_load(repo, release_tag): + if release_tag == "v2.0": + raise PrebuiltFallback("checksum asset missing") + return make_checksums_with_source( + [], release_tag = "v1.0", upstream_tag = "b8999" + ) + + monkeypatch.setattr( + INSTALL_LLAMA_PREBUILT, + "load_approved_release_checksums", + fake_load, + ) + + resolved = resolve_published_release("latest", "unslothai/llama.cpp") + assert resolved.bundle.release_tag == "v1.0" + assert resolved.bundle.upstream_tag == "b8999" + assert resolved.checksums.release_tag == "v1.0" + + def test_concrete_tag_matches_manifest_upstream_tag(self, monkeypatch): + release = make_release([], release_tag = "release-b8508", upstream_tag = "b8508") + monkeypatch.setattr( + INSTALL_LLAMA_PREBUILT, + "iter_published_release_bundles", + lambda repo, published_release_tag = "": iter([release]), + ) + monkeypatch.setattr( + INSTALL_LLAMA_PREBUILT, + "load_approved_release_checksums", + lambda repo, release_tag: make_checksums_with_source( + [], + release_tag = release_tag, + upstream_tag = "b8508", + ), + ) + + assert ( + resolve_requested_install_tag("b8508", "", "unslothai/llama.cpp") == "b8508" + ) + + def test_concrete_tag_without_matching_release_raises(self, monkeypatch): + release = make_release([], release_tag = "release-b9000", upstream_tag = "b9000") + monkeypatch.setattr( + INSTALL_LLAMA_PREBUILT, + "iter_published_release_bundles", + lambda repo, published_release_tag = "": iter([release]), + ) + + with pytest.raises(PrebuiltFallback, match = "matched upstream tag b8508"): + resolve_requested_install_tag("b8508", "", "unslothai/llama.cpp") + + def test_pinned_release_must_match_requested_upstream_tag(self, monkeypatch): + bundle = make_release( + [], release_tag = "llama-prebuilt-latest", upstream_tag = "b9000" + ) + monkeypatch.setattr( + INSTALL_LLAMA_PREBUILT, + "pinned_published_release_bundle", + lambda repo, release_tag: bundle, + ) + monkeypatch.setattr( + INSTALL_LLAMA_PREBUILT, + "load_approved_release_checksums", + lambda repo, release_tag: make_checksums_with_source( + [], + release_tag = release_tag, + upstream_tag = "b9000", + ), + ) + + with pytest.raises(PrebuiltFallback, match = "but requested b8508"): + resolve_requested_install_tag( + "b8508", + "llama-prebuilt-latest", + "unslothai/llama.cpp", + ) + + def test_request_matches_requested_source_ref(self, monkeypatch): + release = make_release( + [], + release_tag = "release-main", + upstream_tag = "b9000", + requested_source_ref = "main", + resolved_source_ref = "refs/heads/main", + ) + monkeypatch.setattr( + INSTALL_LLAMA_PREBUILT, + "iter_published_release_bundles", + lambda repo, published_release_tag = "": iter([release]), + ) + monkeypatch.setattr( + INSTALL_LLAMA_PREBUILT, + "load_approved_release_checksums", + lambda repo, release_tag: make_checksums_with_source( + [], + release_tag = release_tag, + upstream_tag = "b9000", + requested_source_ref = "main", + resolved_source_ref = "refs/heads/main", + ), + ) + + resolved = resolve_published_release("main", "unslothai/llama.cpp") + assert resolved.bundle.release_tag == "release-main" + + def test_request_matches_source_commit(self, monkeypatch): + commit = "a" * 40 + release = make_release( + [], + release_tag = "release-commit", + upstream_tag = "b9000", + source_commit = commit, + ) + monkeypatch.setattr( + INSTALL_LLAMA_PREBUILT, + "iter_published_release_bundles", + lambda repo, published_release_tag = "": iter([release]), + ) + monkeypatch.setattr( + INSTALL_LLAMA_PREBUILT, + "load_approved_release_checksums", + lambda repo, release_tag: make_checksums_with_source( + [], + release_tag = release_tag, + upstream_tag = "b9000", + source_commit = commit, + ), + ) + + resolved = resolve_published_release(commit, "unslothai/llama.cpp") + assert resolved.bundle.release_tag == "release-commit" + + +class TestSourceBuildPlanResolution: + def test_matches_request_by_non_tag_provenance(self): + bundle = make_release( + [], + requested_source_ref = "main", + resolved_source_ref = "refs/heads/main", + source_commit = "a" * 40, + ) + assert published_release_matches_request(bundle, "main") is True + assert published_release_matches_request(bundle, "refs/heads/main") is True + assert published_release_matches_request(bundle, "a" * 12) is True + assert published_release_matches_request(bundle, "a" * 40) is True + + def test_matches_pull_ref_aliases(self): + bundle = make_release( + [], + requested_source_ref = "refs/pull/123/head", + resolved_source_ref = "pull/123/head", + ) + assert published_release_matches_request(bundle, "refs/pull/123/head") is True + assert published_release_matches_request(bundle, "pull/123/head") is True + + def test_prefers_exact_source_commit_when_available(self, monkeypatch): + commit = "a" * 40 + resolved = INSTALL_LLAMA_PREBUILT.ResolvedPublishedRelease( + bundle = make_release( + [], + release_tag = "release-main", + upstream_tag = "b9000", + source_repo = "example/custom-llama.cpp", + source_repo_url = "https://github.com/example/custom-llama.cpp", + source_ref_kind = "branch", + requested_source_ref = "main", + resolved_source_ref = "refs/heads/main", + source_commit = commit, + ), + checksums = make_checksums_with_source( + [], + release_tag = "release-main", + upstream_tag = "b9000", + source_repo = "example/custom-llama.cpp", + source_repo_url = "https://github.com/example/custom-llama.cpp", + source_ref_kind = "branch", + requested_source_ref = "main", + resolved_source_ref = "refs/heads/main", + source_commit = commit, + ), + ) + monkeypatch.setattr( + INSTALL_LLAMA_PREBUILT, + "resolve_published_release", + lambda requested_tag, published_repo, published_release_tag = "": resolved, + ) + + plan = resolve_source_build_plan("main", "unslothai/llama.cpp") + assert plan.source_url == "https://github.com/example/custom-llama.cpp" + assert plan.source_ref_kind == "commit" + assert plan.source_ref == commit + assert plan.compatibility_upstream_tag == "b9000" + + def test_uses_branch_provenance_without_exact_source_hash(self, monkeypatch): + resolved = INSTALL_LLAMA_PREBUILT.ResolvedPublishedRelease( + bundle = make_release( + [], + release_tag = "release-main", + upstream_tag = "b9000", + source_repo = "example/custom-llama.cpp", + source_repo_url = "https://github.com/example/custom-llama.cpp", + source_ref_kind = "branch", + requested_source_ref = "main", + resolved_source_ref = "main", + ), + checksums = make_checksums_with_source( + [], + release_tag = "release-main", + upstream_tag = "b9000", + source_repo = "example/custom-llama.cpp", + source_repo_url = "https://github.com/example/custom-llama.cpp", + source_ref_kind = "branch", + requested_source_ref = "main", + resolved_source_ref = "main", + source_commit = None, + ), + ) + monkeypatch.setattr( + INSTALL_LLAMA_PREBUILT, + "resolve_published_release", + lambda requested_tag, published_repo, published_release_tag = "": resolved, + ) + + plan = resolve_source_build_plan("main", "unslothai/llama.cpp") + assert plan.source_url == "https://github.com/example/custom-llama.cpp" + assert plan.source_ref_kind == "branch" + assert plan.source_ref == "main" + assert plan.compatibility_upstream_tag == "b9000" + + def test_direct_main_request_without_published_release_uses_branch_kind( + self, monkeypatch + ): + monkeypatch.setattr( + INSTALL_LLAMA_PREBUILT, + "resolve_published_release", + lambda requested_tag, published_repo, published_release_tag = "": ( + _ for _ in () + ).throw(PrebuiltFallback("missing")), + ) + + plan = resolve_source_build_plan("main", "unslothai/llama.cpp") + assert plan.source_url == "https://github.com/ggml-org/llama.cpp" + assert plan.source_ref_kind == "branch" + assert plan.source_ref == "main" + + +class TestParseApprovedReleaseChecksums: + def test_rejects_wrong_component(self): + with pytest.raises(RuntimeError, match = "did not describe llama.cpp"): + parse_approved_release_checksums( + "repo/test", + "r1", + { + "schema_version": 1, + "component": "other", + "release_tag": "r1", + "upstream_tag": "b8508", + "artifacts": {}, + }, + ) + + def test_rejects_mismatched_release_tag(self): + with pytest.raises(RuntimeError, match = "did not match pinned release tag"): + parse_approved_release_checksums( + "repo/test", + "r1", + { + "schema_version": 1, + "component": "llama.cpp", + "release_tag": "r2", + "upstream_tag": "b8508", + "artifacts": {}, + }, + ) + + def test_rejects_bad_sha256(self): + with pytest.raises(RuntimeError, match = "valid sha256"): + parse_approved_release_checksums( + "repo/test", + "r1", + { + "schema_version": 1, + "component": "llama.cpp", + "release_tag": "r1", + "upstream_tag": "b8508", + "artifacts": { + "asset.tar.gz": { + "sha256": "bad-digest", + } + }, + }, + ) + + def test_rejects_unsupported_schema_version(self): + with pytest.raises(RuntimeError, match = "schema_version=2 is unsupported"): + parse_approved_release_checksums( + "repo/test", + "r1", + { + "schema_version": 2, + "component": "llama.cpp", + "release_tag": "r1", + "upstream_tag": "b8508", + "artifacts": {}, + }, + ) + + +class TestValidatedChecksumsForBundle: + def test_rejects_manifest_checksum_mismatch(self, monkeypatch): + bundle = make_release([], release_tag = "r1", upstream_tag = "b8508") + bundle.manifest_sha256 = "a" * 64 + checksums = make_checksums_with_source( + [], release_tag = "r1", upstream_tag = "b8508" + ) + checksums.artifacts[bundle.manifest_asset_name] = ApprovedArtifactHash( + asset_name = bundle.manifest_asset_name, + sha256 = "b" * 64, + repo = "unslothai/llama.cpp", + kind = "published-manifest", + ) + monkeypatch.setattr( + INSTALL_LLAMA_PREBUILT, + "load_approved_release_checksums", + lambda repo, release_tag: checksums, + ) + + with pytest.raises(PrebuiltFallback, match = "manifest checksum"): + validated_checksums_for_bundle("unslothai/llama.cpp", bundle) + + +# =========================================================================== +# K. linux_cuda_choice_from_release -- core selection # =========================================================================== @@ -676,17 +1188,567 @@ class TestLinuxCudaChoiceFromRelease: # =========================================================================== -# K. windows_cuda_attempts +# L. resolve_install_attempts +# =========================================================================== + + +class TestResolveInstallAttempts: + def test_windows_cuda_prefers_published_asset_from_selected_release( + self, monkeypatch + ): + host = make_host(system = "Windows", machine = "AMD64") + host.driver_cuda_version = (12, 4) + mock_windows_runtime(monkeypatch, ["cuda12"]) + asset_name = "llama-b9000-bin-win-cuda-12.4-x64.zip" + release = make_release( + [ + make_artifact( + asset_name, + install_kind = "windows-cuda", + runtime_line = "cuda12", + coverage_class = None, + supported_sms = [], + min_sm = None, + max_sm = None, + bundle_profile = None, + ) + ], + release_tag = "llama-prebuilt-latest", + upstream_tag = "b9000", + assets = {asset_name: f"https://published.example/{asset_name}"}, + ) + checksums = make_checksums_with_source( + [asset_name], + release_tag = release.release_tag, + upstream_tag = "b9000", + ) + + monkeypatch.setattr( + INSTALL_LLAMA_PREBUILT, + "iter_resolved_published_releases", + lambda requested_tag, published_repo, published_release_tag = "": iter( + [ + INSTALL_LLAMA_PREBUILT.ResolvedPublishedRelease( + bundle = release, + checksums = checksums, + ) + ] + ), + ) + monkeypatch.setattr( + INSTALL_LLAMA_PREBUILT, + "github_release_assets", + lambda repo, tag: (_ for _ in ()).throw( + AssertionError( + "published Windows CUDA choice should not query upstream" + ) + ), + ) + + requested_tag, resolved_tag, attempts, approved = resolve_install_attempts( + "latest", + host, + "unslothai/llama.cpp", + "", + ) + + assert requested_tag == "latest" + assert resolved_tag == "b9000" + assert attempts[0].name == asset_name + assert attempts[0].source_label == "published" + assert attempts[0].expected_sha256 == "a" * 64 + assert approved.release_tag == "llama-prebuilt-latest" + + def test_windows_cuda_uses_selected_release_upstream_tag(self, monkeypatch): + host = make_host(system = "Windows", machine = "AMD64") + host.driver_cuda_version = (12, 4) + mock_windows_runtime(monkeypatch, ["cuda12"]) + release = make_release( + [], release_tag = "llama-prebuilt-latest", upstream_tag = "b9000" + ) + checksums = make_checksums_with_source( + ["llama-b9000-bin-win-cuda-12.4-x64.zip"], + release_tag = release.release_tag, + upstream_tag = "b9000", + ) + + monkeypatch.setattr( + INSTALL_LLAMA_PREBUILT, + "iter_resolved_published_releases", + lambda requested_tag, published_repo, published_release_tag = "": iter( + [ + INSTALL_LLAMA_PREBUILT.ResolvedPublishedRelease( + bundle = release, + checksums = checksums, + ) + ] + ), + ) + monkeypatch.setattr( + INSTALL_LLAMA_PREBUILT, + "github_release_assets", + lambda repo, tag: { + f"llama-{tag}-bin-win-cuda-12.4-x64.zip": f"https://example.com/llama-{tag}-bin-win-cuda-12.4-x64.zip" + }, + ) + monkeypatch.setattr( + INSTALL_LLAMA_PREBUILT, + "resolve_windows_cuda_choices", + lambda host, tag, assets: [ + AssetChoice( + repo = UPSTREAM_REPO, + tag = tag, + name = f"llama-{tag}-bin-win-cuda-12.4-x64.zip", + url = assets[f"llama-{tag}-bin-win-cuda-12.4-x64.zip"], + source_label = "upstream", + install_kind = "windows-cuda", + runtime_line = "cuda12", + ) + ], + ) + + requested_tag, resolved_tag, attempts, approved = resolve_install_attempts( + "latest", + host, + "unslothai/llama.cpp", + "", + ) + + assert requested_tag == "latest" + assert resolved_tag == "b9000" + assert attempts[0].name == "llama-b9000-bin-win-cuda-12.4-x64.zip" + assert attempts[0].expected_sha256 == "a" * 64 + assert approved.release_tag == "llama-prebuilt-latest" + + def test_linux_cpu_uses_same_tag_upstream_asset(self, monkeypatch): + host = make_host( + has_usable_nvidia = False, + has_physical_nvidia = False, + nvidia_smi = None, + ) + release = make_release( + [], release_tag = "llama-prebuilt-latest", upstream_tag = "b9000" + ) + checksums = make_checksums_with_source( + ["llama-b9000-bin-ubuntu-x64.tar.gz"], + release_tag = release.release_tag, + upstream_tag = "b9000", + ) + + monkeypatch.setattr( + INSTALL_LLAMA_PREBUILT, + "iter_resolved_published_releases", + lambda requested_tag, published_repo, published_release_tag = "": iter( + [ + INSTALL_LLAMA_PREBUILT.ResolvedPublishedRelease( + bundle = release, + checksums = checksums, + ) + ] + ), + ) + monkeypatch.setattr( + INSTALL_LLAMA_PREBUILT, + "github_release_assets", + lambda repo, tag: { + f"llama-{tag}-bin-ubuntu-x64.tar.gz": f"https://example.com/llama-{tag}-bin-ubuntu-x64.tar.gz" + }, + ) + + _requested_tag, resolved_tag, attempts, _approved = resolve_install_attempts( + "latest", + host, + "unslothai/llama.cpp", + "", + ) + + assert resolved_tag == "b9000" + assert attempts[0].name == "llama-b9000-bin-ubuntu-x64.tar.gz" + assert attempts[0].source_label == "upstream" + assert attempts[0].expected_sha256 == "a" * 64 + + def test_linux_cuda_does_not_fall_back_to_upstream_cpu(self, monkeypatch): + host = make_host(system = "Linux", machine = "x86_64", compute_caps = ["86"]) + release = make_release( + [], release_tag = "llama-prebuilt-latest", upstream_tag = "b9000" + ) + checksums = make_checksums_with_source( + [], + release_tag = release.release_tag, + upstream_tag = "b9000", + ) + + monkeypatch.setattr( + INSTALL_LLAMA_PREBUILT, + "iter_resolved_published_releases", + lambda requested_tag, published_repo, published_release_tag = "": iter( + [ + INSTALL_LLAMA_PREBUILT.ResolvedPublishedRelease( + bundle = release, + checksums = checksums, + ) + ] + ), + ) + mock_linux_runtime(monkeypatch, ["cuda12"]) + + with pytest.raises( + PrebuiltFallback, match = "no compatible published Linux CUDA bundle" + ): + resolve_install_attempts("latest", host, "unslothai/llama.cpp", "") + + def test_windows_cpu_prefers_published_asset(self, monkeypatch): + host = make_host( + system = "Windows", + machine = "AMD64", + has_usable_nvidia = False, + has_physical_nvidia = False, + nvidia_smi = None, + ) + asset_name = "llama-b9000-bin-win-cpu-x64.zip" + release = make_release( + [ + make_artifact( + asset_name, + install_kind = "windows-cpu", + runtime_line = None, + coverage_class = None, + supported_sms = [], + min_sm = None, + max_sm = None, + bundle_profile = None, + ) + ], + release_tag = "llama-prebuilt-latest", + upstream_tag = "b9000", + assets = {asset_name: f"https://published.example/{asset_name}"}, + ) + checksums = make_checksums_with_source( + [asset_name], + release_tag = release.release_tag, + upstream_tag = "b9000", + ) + + monkeypatch.setattr( + INSTALL_LLAMA_PREBUILT, + "iter_resolved_published_releases", + lambda requested_tag, published_repo, published_release_tag = "": iter( + [ + INSTALL_LLAMA_PREBUILT.ResolvedPublishedRelease( + bundle = release, + checksums = checksums, + ) + ] + ), + ) + monkeypatch.setattr( + INSTALL_LLAMA_PREBUILT, + "github_release_assets", + lambda repo, tag: (_ for _ in ()).throw( + AssertionError("published Windows CPU choice should not query upstream") + ), + ) + + _requested_tag, resolved_tag, attempts, _approved = resolve_install_attempts( + "latest", + host, + "unslothai/llama.cpp", + "", + ) + + assert resolved_tag == "b9000" + assert attempts[0].name == asset_name + assert attempts[0].source_label == "published" + + def test_macos_prefers_published_asset(self, monkeypatch): + host = make_host( + system = "Darwin", + machine = "arm64", + nvidia_smi = None, + driver_cuda_version = None, + compute_caps = [], + has_physical_nvidia = False, + has_usable_nvidia = False, + ) + asset_name = "llama-b9000-bin-macos-arm64.tar.gz" + release = make_release( + [ + make_artifact( + asset_name, + install_kind = "macos-arm64", + runtime_line = None, + coverage_class = None, + supported_sms = [], + min_sm = None, + max_sm = None, + bundle_profile = None, + ) + ], + release_tag = "llama-prebuilt-latest", + upstream_tag = "b9000", + assets = {asset_name: f"https://published.example/{asset_name}"}, + ) + checksums = make_checksums_with_source( + [asset_name], + release_tag = release.release_tag, + upstream_tag = "b9000", + ) + + monkeypatch.setattr( + INSTALL_LLAMA_PREBUILT, + "iter_resolved_published_releases", + lambda requested_tag, published_repo, published_release_tag = "": iter( + [ + INSTALL_LLAMA_PREBUILT.ResolvedPublishedRelease( + bundle = release, + checksums = checksums, + ) + ] + ), + ) + monkeypatch.setattr( + INSTALL_LLAMA_PREBUILT, + "github_release_assets", + lambda repo, tag: (_ for _ in ()).throw( + AssertionError("published macOS choice should not query upstream") + ), + ) + + _requested_tag, resolved_tag, attempts, _approved = resolve_install_attempts( + "latest", + host, + "unslothai/llama.cpp", + "", + ) + + assert resolved_tag == "b9000" + assert attempts[0].name == asset_name + assert attempts[0].source_label == "published" + + def test_windows_cpu_missing_checksum_rejects_install(self, monkeypatch): + host = make_host( + system = "Windows", + machine = "AMD64", + has_usable_nvidia = False, + has_physical_nvidia = False, + nvidia_smi = None, + ) + published_name = "llama-b9000-bin-win-cpu-x64.zip" + release = make_release( + [ + make_artifact( + published_name, + install_kind = "windows-cpu", + runtime_line = None, + coverage_class = None, + supported_sms = [], + min_sm = None, + max_sm = None, + bundle_profile = None, + ) + ], + release_tag = "llama-prebuilt-latest", + upstream_tag = "b9000", + assets = {published_name: f"https://published.example/{published_name}"}, + ) + checksums = make_checksums_with_source( + [], + release_tag = release.release_tag, + upstream_tag = "b9000", + ) + + monkeypatch.setattr( + INSTALL_LLAMA_PREBUILT, + "iter_resolved_published_releases", + lambda requested_tag, published_repo, published_release_tag = "": iter( + [ + INSTALL_LLAMA_PREBUILT.ResolvedPublishedRelease( + bundle = release, + checksums = checksums, + ) + ] + ), + ) + monkeypatch.setattr( + INSTALL_LLAMA_PREBUILT, + "github_release_assets", + lambda repo, tag: { + f"llama-{tag}-bin-win-cpu-x64.zip": f"https://upstream.example/llama-{tag}-bin-win-cpu-x64.zip" + }, + ) + + with pytest.raises( + PrebuiltFallback, + match = "approved checksum asset did not contain the selected prebuilt archive", + ): + resolve_install_attempts( + "latest", + host, + "unslothai/llama.cpp", + "", + ) + + +class TestResolveInstallReleasePlans: + def test_latest_collects_multiple_older_release_plans_up_to_limit( + self, monkeypatch + ): + host = make_host( + has_usable_nvidia = False, + has_physical_nvidia = False, + nvidia_smi = None, + ) + releases = [ + INSTALL_LLAMA_PREBUILT.ResolvedPublishedRelease( + bundle = make_release([], release_tag = "r3", upstream_tag = "b9003"), + checksums = make_checksums_with_source( + ["llama-b9003-bin-ubuntu-x64.tar.gz"], + release_tag = "r3", + upstream_tag = "b9003", + ), + ), + INSTALL_LLAMA_PREBUILT.ResolvedPublishedRelease( + bundle = make_release([], release_tag = "r2", upstream_tag = "b9002"), + checksums = make_checksums_with_source( + ["llama-b9002-bin-ubuntu-x64.tar.gz"], + release_tag = "r2", + upstream_tag = "b9002", + ), + ), + INSTALL_LLAMA_PREBUILT.ResolvedPublishedRelease( + bundle = make_release([], release_tag = "r1", upstream_tag = "b9001"), + checksums = make_checksums_with_source( + ["llama-b9001-bin-ubuntu-x64.tar.gz"], + release_tag = "r1", + upstream_tag = "b9001", + ), + ), + ] + + monkeypatch.setattr( + INSTALL_LLAMA_PREBUILT, + "iter_resolved_published_releases", + lambda requested_tag, published_repo, published_release_tag = "": iter( + releases + ), + ) + monkeypatch.setattr( + INSTALL_LLAMA_PREBUILT, + "github_release_assets", + lambda repo, tag: { + f"llama-{tag}-bin-ubuntu-x64.tar.gz": f"https://example.com/llama-{tag}-bin-ubuntu-x64.tar.gz" + }, + ) + + requested_tag, plans = resolve_install_release_plans( + "latest", + host, + "unslothai/llama.cpp", + "", + max_release_fallbacks = 2, + ) + + assert requested_tag == "latest" + assert [plan.release_tag for plan in plans] == ["r3", "r2"] + assert [plan.llama_tag for plan in plans] == ["b9003", "b9002"] + + def test_latest_skips_non_installable_release_and_keeps_searching( + self, monkeypatch + ): + host = make_host( + has_usable_nvidia = False, + has_physical_nvidia = False, + nvidia_smi = None, + ) + releases = [ + INSTALL_LLAMA_PREBUILT.ResolvedPublishedRelease( + bundle = make_release([], release_tag = "r2", upstream_tag = "b9002"), + checksums = make_checksums_with_source( + [], + release_tag = "r2", + upstream_tag = "b9002", + ), + ), + INSTALL_LLAMA_PREBUILT.ResolvedPublishedRelease( + bundle = make_release([], release_tag = "r1", upstream_tag = "b9001"), + checksums = make_checksums_with_source( + ["llama-b9001-bin-ubuntu-x64.tar.gz"], + release_tag = "r1", + upstream_tag = "b9001", + ), + ), + ] + + monkeypatch.setattr( + INSTALL_LLAMA_PREBUILT, + "iter_resolved_published_releases", + lambda requested_tag, published_repo, published_release_tag = "": iter( + releases + ), + ) + monkeypatch.setattr( + INSTALL_LLAMA_PREBUILT, + "github_release_assets", + lambda repo, tag: ( + {} + if tag == "b9002" + else { + f"llama-{tag}-bin-ubuntu-x64.tar.gz": f"https://example.com/llama-{tag}-bin-ubuntu-x64.tar.gz" + } + ), + ) + + _requested_tag, plans = resolve_install_release_plans( + "latest", + host, + "unslothai/llama.cpp", + "", + max_release_fallbacks = 2, + ) + + assert len(plans) == 1 + assert plans[0].release_tag == "r1" + assert plans[0].llama_tag == "b9001" + + def test_malformed_release_fallback_env_uses_default(self, monkeypatch): + monkeypatch.setenv("UNSLOTH_LLAMA_MAX_PREBUILT_RELEASE_FALLBACKS", "not-an-int") + assert ( + env_int("UNSLOTH_LLAMA_MAX_PREBUILT_RELEASE_FALLBACKS", 3, minimum = 1) == 3 + ) + + def test_import_with_malformed_release_fallback_env_does_not_crash( + self, monkeypatch + ): + monkeypatch.setenv("UNSLOTH_LLAMA_MAX_PREBUILT_RELEASE_FALLBACKS", "bad-value") + spec = importlib.util.spec_from_file_location( + "studio_install_llama_prebuilt_env_reload", + MODULE_PATH, + ) + assert spec is not None and spec.loader is not None + module = importlib.util.module_from_spec(spec) + sys.modules[spec.name] = module + try: + spec.loader.exec_module(module) + assert module.DEFAULT_MAX_PREBUILT_RELEASE_FALLBACKS == 2 + finally: + sys.modules.pop(spec.name, None) + + +# =========================================================================== +# N. windows_cuda_attempts # =========================================================================== class TestWindowsCudaAttempts: TAG = "b8508" - def _upstream(self, *runtime_versions): + def _upstream(self, *runtime_versions, current_names: bool = False): assets = {} for rv in runtime_versions: - name = f"llama-{self.TAG}-bin-win-cuda-{rv}-x64.zip" + if current_names: + name = f"cudart-llama-bin-win-cuda-{rv}-x64.zip" + else: + name = f"llama-{self.TAG}-bin-win-cuda-{rv}-x64.zip" assets[name] = f"https://example.com/{name}" return assets @@ -751,9 +1813,18 @@ class TestWindowsCudaAttempts: result = windows_cuda_attempts(host, self.TAG, assets, None) assert len(result) == 2 + def test_current_upstream_names_are_supported(self, monkeypatch): + mock_windows_runtime(monkeypatch, ["cuda13", "cuda12"]) + host = make_host(system = "Windows", machine = "AMD64", driver_cuda_version = (13, 1)) + assets = self._upstream("13.1", "12.4", current_names = True) + result = windows_cuda_attempts(host, self.TAG, assets, None) + assert len(result) == 2 + assert result[0].name == "cudart-llama-bin-win-cuda-13.1-x64.zip" + assert result[1].name == "cudart-llama-bin-win-cuda-12.4-x64.zip" + # =========================================================================== -# L. resolve_upstream_asset_choice -- platform routing +# O. resolve_upstream_asset_choice -- platform routing # =========================================================================== diff --git a/unsloth/chat_templates.py b/unsloth/chat_templates.py index 35eb871529..71f91cc828 100644 --- a/unsloth/chat_templates.py +++ b/unsloth/chat_templates.py @@ -863,6 +863,114 @@ DEFAULT_SYSTEM_MESSAGE["gemma-3n"] = None # No system message in Gemma-3n CHAT_TEMPLATES["gemma3n"] = (gemma3n_template, gemma3n_template_eos_token, False, gemma3n_ollama,) DEFAULT_SYSTEM_MESSAGE["gemma3n"] = None # No system message in Gemma-3n +# =========================================== Gemma-4 +# Gemma-4 uses <|turn>role\n...\n format +gemma4_template = \ +"""{%- if messages[0]['role'] == 'system' -%} + {%- set first_user_prefix = messages[0]['content'] + '\n\n' -%} + {%- set loop_messages = messages[1:] -%} +{%- else -%} + {%- set first_user_prefix = "" -%} + {%- set loop_messages = messages -%} +{%- endif -%} +{%- for message in loop_messages -%} + {%- if (message['role'] == 'user') != (loop.index0 % 2 == 0) -%} + {{ raise_exception("Conversation roles must alternate user/assistant/user/assistant/...") }} + {%- endif -%} + {%- if (message['role'] == 'assistant') -%} + {%- set role = "model" -%} + {%- else -%} + {%- set role = message['role'] -%} + {%- endif -%} + {{ '<|turn>' + role + '\n' + (first_user_prefix if loop.first else "") }} + {%- if message['content'] is string -%} + {{ message['content'] | trim }} + {%- elif message['content'] is iterable -%} + {%- for item in message['content'] -%} + {%- if item['type'] == 'audio' -%} + {{ '<|audio|>' }} + {%- elif item['type'] == 'image' -%} + {{ '<|image|>' }} + {%- elif item['type'] == 'video' -%} + {{ '<|video|>' }} + {%- elif item['type'] == 'text' -%} + {{ item['text'] | trim }} + {%- endif -%} + {%- endfor -%} + {%- else -%} + {{ raise_exception("Invalid content type") }} + {%- endif -%} + {{ '\n' }} +{%- endfor -%} +{%- if add_generation_prompt -%} + {{'<|turn>model\n'}} +{%- endif -%} +""" + +try: + gemma4_ollama = _ollama_template("gemma-4") +except KeyError: + gemma4_ollama = "" +gemma4_template_eos_token = "" +CHAT_TEMPLATES["gemma-4"] = (gemma4_template, gemma4_template_eos_token, False, gemma4_ollama,) +DEFAULT_SYSTEM_MESSAGE["gemma-4"] = None + +CHAT_TEMPLATES["gemma4"] = (gemma4_template, gemma4_template_eos_token, False, gemma4_ollama,) +DEFAULT_SYSTEM_MESSAGE["gemma4"] = None + +# Gemma-4 with empty thought channel (required for larger models like 31B, 26B-A4B) +# Injects <|channel>thought\n at the start of each model response during training +gemma4_thinking_template = \ +"""{%- if messages[0]['role'] == 'system' -%} + {%- set first_user_prefix = messages[0]['content'] + '\n\n' -%} + {%- set loop_messages = messages[1:] -%} +{%- else -%} + {%- set first_user_prefix = "" -%} + {%- set loop_messages = messages -%} +{%- endif -%} +{%- for message in loop_messages -%} + {%- if (message['role'] == 'user') != (loop.index0 % 2 == 0) -%} + {{ raise_exception("Conversation roles must alternate user/assistant/user/assistant/...") }} + {%- endif -%} + {%- if (message['role'] == 'assistant') -%} + {%- set role = "model" -%} + {%- else -%} + {%- set role = message['role'] -%} + {%- endif -%} + {{ '<|turn>' + role + '\n' + (first_user_prefix if loop.first else "") }} + {%- if role == "model" -%} + {{ '<|channel>thought\n' }} + {%- endif -%} + {%- if message['content'] is string -%} + {{ message['content'] | trim }} + {%- elif message['content'] is iterable -%} + {%- for item in message['content'] -%} + {%- if item['type'] == 'audio' -%} + {{ '<|audio|>' }} + {%- elif item['type'] == 'image' -%} + {{ '<|image|>' }} + {%- elif item['type'] == 'video' -%} + {{ '<|video|>' }} + {%- elif item['type'] == 'text' -%} + {{ item['text'] | trim }} + {%- endif -%} + {%- endfor -%} + {%- else -%} + {{ raise_exception("Invalid content type") }} + {%- endif -%} + {{ '\n' }} +{%- endfor -%} +{%- if add_generation_prompt -%} + {{'<|turn>model\n'}} +{%- endif -%} +""" + +CHAT_TEMPLATES["gemma-4-thinking"] = (gemma4_thinking_template, gemma4_template_eos_token, False, gemma4_ollama,) +DEFAULT_SYSTEM_MESSAGE["gemma-4-thinking"] = None + +CHAT_TEMPLATES["gemma4-thinking"] = (gemma4_thinking_template, gemma4_template_eos_token, False, gemma4_ollama,) +DEFAULT_SYSTEM_MESSAGE["gemma4-thinking"] = None + # =========================================== GPT-OSS # Obtained via # print(tokenizer.chat_template.replace("}\n", "####").replace("\n", "\\n").replace("####", "}\n")) diff --git a/unsloth/models/_utils.py b/unsloth/models/_utils.py index d296ac7e74..28526056ba 100644 --- a/unsloth/models/_utils.py +++ b/unsloth/models/_utils.py @@ -12,7 +12,7 @@ # See the License for the specific language governing permissions and # limitations under the License. -__version__ = "2026.3.18" +__version__ = "2026.4.1" __all__ = [ "SUPPORTS_BFLOAT16", @@ -64,7 +64,8 @@ __all__ = [ "patch_compiled_autograd", "process_vision_info", "unsloth_compile_transformers", - "prefer_flex_attn_if_supported", + "determine_attention_implementation", + "_set_attn_impl", "patch_fast_lora", "validate_loftq_config", "RaiseUninitialized", @@ -222,44 +223,76 @@ def apply_unsloth_gradient_checkpointing( return use_gradient_checkpointing -def prefer_flex_attn_if_supported(model_class, config): - if os.environ.get("UNSLOTH_ENABLE_FLEX_ATTENTION", "1") == "0": - return None - try: - from transformers.utils.import_utils import is_torch_flex_attn_available +# Models that don't work with flex_attention: +# GPT-OSS: left padding issues cause incorrect outputs. +# Mllama: BlockMask Q_LEN!=KV_LEN ValueError on decode. +# NemotronH: hybrid Mamba-2 + Transformer, raises NotImplementedError. +# Gemma3N: timm vision wrappers don't support flex_attention. +# ModernBERT: create_block_mask with _compile=True hits CUDA illegal memory +# access on some GPU architectures (B200). Falls back to eager safely. +_FLEX_EXCLUDED_MODELS = ("gpt_oss", "mllama", "nemotron_h", "modernbert") +_EAGER_ONLY_PREFIXES = ("gemma3n",) - if not is_torch_flex_attn_available(): - return None - if model_class is None or not getattr( - model_class, "_supports_flex_attn", False - ): - return None - attention_dropout = getattr(config, "attention_dropout", 0) or 0 - if attention_dropout > 0: - return None - # GPT-OSS, Mllama and Gemma3N use eager/sdpa attention during - # inference since flex attention returns incorrect results or errors out. - # GPT-OSS: left padding issues cause incorrect outputs. - # Mllama: _update_causal_mask uses make_flex_block_causal_mask which - # creates BlockMask with Q_LEN=KV_LEN=total_seq_len, but during - # decode q_len=1, causing ValueError. Needs transformers update. - # Gemma3N: timm vision wrappers (eg Gemma3nVisionConfig) do not - # support flex_attention. - # NemotronH: hybrid Mamba-2 + Transformer model that does not - # support flex_attention (raises NotImplementedError from transformers). - model_type = getattr(config, "model_type", "") if config else "" - if model_type in ("gpt_oss", "mllama", "nemotron_h") or str( - model_type - ).startswith("gemma3n"): - return None - if config is not None: - setattr(config, "_attn_implementation", "flex_attention") - if hasattr(config, "attn_implementation"): - setattr(config, "attn_implementation", "flex_attention") - return "flex_attention" - except Exception: - return None +def _is_flex_excluded(model_type): + return model_type in _FLEX_EXCLUDED_MODELS + + +def _is_eager_only(model_type): + return any(model_type.startswith(p) for p in _EAGER_ONLY_PREFIXES) + + +def _set_attn_impl(config, impl): + """Helper function to set attention implementation on config and return it.""" + if config is not None: + setattr(config, "_attn_implementation", impl) + if hasattr(config, "attn_implementation"): + setattr(config, "attn_implementation", impl) + return impl + + +def determine_attention_implementation(model_class, config): + model_type = getattr(config, "model_type", "").lower() + + # Eager-only models (e.g. gemma3n timm vision towers) + if _is_eager_only(model_type): + _set_attn_impl(config, "eager") + return "eager" + + # Flash Attention 2 + if HAS_FLASH_ATTENTION and model_class is not None: + supports_fa2 = getattr(model_class, "_supports_flash_attn_2", False) or getattr( + model_class, "_supports_flash_attn", False + ) + if supports_fa2: + _set_attn_impl(config, "flash_attention_2") + return "flash_attention_2" + + # Flex Attention + if os.environ.get("UNSLOTH_ENABLE_FLEX_ATTENTION", "1") != "0": + try: + from transformers.utils.import_utils import is_torch_flex_attn_available + + if ( + is_torch_flex_attn_available() + and model_class is not None + and getattr(model_class, "_supports_flex_attn", False) + and not _is_flex_excluded(model_type) + ): + attention_dropout = getattr(config, "attention_dropout", 0) or 0 + if attention_dropout == 0: + _set_attn_impl(config, "flex_attention") + return "flex_attention" + except Exception: + pass + + # SDPA + if model_class is not None and getattr(model_class, "_supports_sdpa", False): + _set_attn_impl(config, "sdpa") + return "sdpa" + + _set_attn_impl(config, "eager") + return "eager" def _run_temporary_patches(phase): @@ -504,6 +537,15 @@ try: except: pass +# Gemma4 It is strongly recommended to train Gemma4 models with the `eager` +try: + from transformers.models.gemma4.modeling_gemma4 import logger as gemma4_logger + + gemma4_logger.addFilter(HideLoggingMessage("strongly recommended")) + del gemma4_logger +except: + pass + # Xet Storage is enabled for this repo, but the 'hf_xet' package is not installed. try: from huggingface_hub.file_download import logger as hub_logger @@ -765,7 +807,16 @@ model_architectures = [ "falcon_h1", ] +# Transformers 5.x uses class-level annotations with @strict, @auto_docstring, +# and interval() in config classes. exec(inspect.getsource(...)) fails because +# those symbols are not in scope. Skip the exec-based config patching for 5.x +# since those configs already use rope_parameters (the v5 replacement for +# rope_scaling). +_skip_config_exec_patch = Version(transformers_version) >= Version("5.0.0") + for model_name in model_architectures: + if _skip_config_exec_patch: + break config_filepath = f"transformers.models.{model_name}.configuration_{model_name}" model_filepath = f"transformers.models.{model_name}.modeling_{model_name}" config_filename = f"{model_name.title().replace('_','')}Config" # qwen3 arch folder is qwen3_moe but config is Qwen3Config. Need to remove underscore(_) for now @@ -799,9 +850,12 @@ for model_name in model_architectures: if Version(transformers_version) <= Version("4.42.4"): config = patch_mistral_nemo_config(config) - exec(config, globals()) - exec(f"import {config_filepath}", globals()) - exec(f"{config_filepath}.{config_filename} = {config_filename}", globals()) + try: + exec(config, globals()) + exec(f"import {config_filepath}", globals()) + exec(f"{config_filepath}.{config_filename} = {config_filename}", globals()) + except Exception: + continue # ============================================= # ============================================= @@ -1885,6 +1939,18 @@ def _unsloth_pre_compute_loss(self, model, inputs, *args, **kwargs): _has_ccm = _mod is not None and hasattr(_mod, "create_causal_mask_mapping") if _has_ccm and _inner.training: inputs["token_type_ids"] = torch.zeros_like(inputs["input_ids"]) + # Gemma4 uses mm_token_type_ids (not token_type_ids) for VLM masking + if "mm_token_type_ids" not in inputs and "input_ids" in inputs: + _inner = model + for _attr in ("base_model", "model", "model"): + _inner = getattr(_inner, _attr, _inner) + if getattr(getattr(_inner, "config", None), "model_type", "") in ("gemma4",): + import sys as _sys + + _mod = _sys.modules.get(type(_inner).__module__) + _has_ccm = _mod is not None and hasattr(_mod, "create_causal_mask_mapping") + if _has_ccm and _inner.training: + inputs["mm_token_type_ids"] = torch.zeros_like(inputs["input_ids"]) outputs = self._old_compute_loss(model, inputs, *args, **kwargs) return outputs diff --git a/unsloth/models/cohere.py b/unsloth/models/cohere.py index 4251f3acd9..294e8d0c7e 100644 --- a/unsloth/models/cohere.py +++ b/unsloth/models/cohere.py @@ -357,6 +357,9 @@ def CohereAttention_fast_forward_inference( # cos, sin = self.rotary_emb(Vn, seq_len = kv_seq_len) # Qn, Kn = inplace_rope_embedding(Qn, Kn, cos, sin, position_ids) cos, sin = self.rotary_emb.get_cached(kv_seq_len, Qn.device.index) + # Transformers 5.x: position_ids may be [batch, full_seq_len]; slice to last + if position_ids.dim() >= 2 and position_ids.shape[-1] > 1: + position_ids = position_ids[:, -1:] cos = cos[position_ids].unsqueeze(1) sin = sin[position_ids].unsqueeze(1) h = self.half_head_dim diff --git a/unsloth/models/falcon_h1.py b/unsloth/models/falcon_h1.py index 6e3b16b21b..659d27de54 100644 --- a/unsloth/models/falcon_h1.py +++ b/unsloth/models/falcon_h1.py @@ -313,6 +313,9 @@ def FalconH1Attention_fast_forward_inference( # or else error self.rotary_emb.extend_rope_embedding(Vn, seq_len + 2) cos, sin = self.rotary_emb.get_cached(kv_seq_len, Qn.device.index) + # Transformers 5.x: position_ids may be [batch, full_seq_len]; slice to last + if position_ids.dim() >= 2 and position_ids.shape[-1] > 1: + position_ids = position_ids[:, -1:] cos = cos[position_ids].unsqueeze(1) sin = sin[position_ids].unsqueeze(1) h = self.half_head_dim diff --git a/unsloth/models/gemma2.py b/unsloth/models/gemma2.py index e59b8d5ebd..720c9a7414 100644 --- a/unsloth/models/gemma2.py +++ b/unsloth/models/gemma2.py @@ -394,6 +394,9 @@ def Gemma2Attention_fast_forward_inference( # cos, sin = self.rotary_emb(Vn, seq_len = kv_seq_len) # Qn, Kn = inplace_rope_embedding(Qn, Kn, cos, sin, position_ids) cos, sin = self.rotary_emb.get_cached(kv_seq_len, Qn.device.index) + # Transformers 5.x: position_ids may be [batch, full_seq_len]; slice to last + if position_ids.dim() >= 2 and position_ids.shape[-1] > 1: + position_ids = position_ids[:, -1:] cos = cos[position_ids].unsqueeze(1) sin = sin[position_ids].unsqueeze(1) h = self.half_head_dim diff --git a/unsloth/models/granite.py b/unsloth/models/granite.py index 79ac41c43f..fea3dc1b36 100644 --- a/unsloth/models/granite.py +++ b/unsloth/models/granite.py @@ -355,6 +355,9 @@ def GraniteAttention_fast_forward_inference( # cos, sin = self.rotary_emb(Vn, seq_len = kv_seq_len) # Qn, Kn = inplace_rope_embedding(Qn, Kn, cos, sin, position_ids) cos, sin = position_embeddings + # Transformers 5.x: position_ids may be [batch, full_seq_len]; slice to last + if position_ids.dim() >= 2 and position_ids.shape[-1] > 1: + position_ids = position_ids[:, -1:] cos, sin = cos[position_ids], sin[position_ids] h = self.half_head_dim diff --git a/unsloth/models/llama.py b/unsloth/models/llama.py index 93d93e26d6..2f61913550 100644 --- a/unsloth/models/llama.py +++ b/unsloth/models/llama.py @@ -496,6 +496,10 @@ def LlamaAttention_fast_forward_inference( # ensure correct shape if position_ids.dim() == 1: position_ids = position_ids[:, None] + # Transformers 5.x generate() accumulates position_ids as [batch, full_seq_len] + # across decode steps. In single-token inference we only need the last position. + if position_ids.shape[-1] > 1: + position_ids = position_ids[:, -1:] position_ids = position_ids.to(Qn.device) if rotary_seq_len is None: @@ -2341,8 +2345,8 @@ class FastLlamaModel: model_function = MODEL_FOR_CAUSAL_LM_MAPPING[model_config.__class__] IS_FALCON_H1 = model_config.model_type.startswith("falcon_h1") - preferred_attn_impl = ( - prefer_flex_attn_if_supported(model_function, model_config) or "eager" + preferred_attn_impl = determine_attention_implementation( + model_function, model_config ) has_rope_scaling = False @@ -2414,14 +2418,24 @@ class FastLlamaModel: raise_handler = RaiseUninitialized() if num_labels is not None: + # Transformers 5.x @strict config classes reject unexpected kwargs + # like num_labels and max_position_embeddings. Set on the config + # object directly and pass config= instead. + model_config.num_labels = num_labels + if max_position_embeddings is not None: + model_config.max_position_embeddings = max_position_embeddings + # Pop config-level attrs that would be rejected by @strict model init + for _cfg_key in ("id2label", "label2id", "rope_scaling"): + _cfg_val = kwargs.pop(_cfg_key, None) + if _cfg_val is not None: + setattr(model_config, _cfg_key, _cfg_val) model = AutoModelForSequenceClassification.from_pretrained( model_name, + config = model_config, device_map = device_map, # torch_dtype = dtype, # transformers changed torch_dtype to dtype - num_labels = num_labels, # quantization_config = bnb_config, token = token, - max_position_embeddings = max_position_embeddings, trust_remote_code = trust_remote_code, attn_implementation = preferred_attn_impl, **kwargs, diff --git a/unsloth/models/loader.py b/unsloth/models/loader.py index b54ceaf842..a811d3fb75 100644 --- a/unsloth/models/loader.py +++ b/unsloth/models/loader.py @@ -78,6 +78,7 @@ SUPPORTS_QWEN3_MOE = transformers_version >= Version("4.50.3") SUPPORTS_FALCON_H1 = transformers_version >= Version("4.53.0") SUPPORTS_GEMMA3N = transformers_version >= Version("4.53.0") SUPPORTS_GPTOSS = transformers_version >= Version("4.55.0") +SUPPORTS_GEMMA4 = transformers_version >= Version("5.5.0") # Transformers v5 meta-device loading corrupts non-persistent buffers (inv_freq). # See _fix_rope_inv_freq() below for details. _NEEDS_ROPE_FIX = transformers_version >= Version("5.0.0") @@ -107,6 +108,8 @@ FORCE_FLOAT32 = [ "gemma3n", "gpt_oss", "qwen3_5", # Qwen3.5 GDN layers produce NaN grad norms in float16 training + "gemma4,", # Add comma bc gemma4 will match gemma4_text + "gemma4_text", ] global DISABLE_COMPILE_MODEL_NAMES @@ -1130,6 +1133,17 @@ class FastModel(FastBaseModel): raise RuntimeError( "Unsloth: Qwen 2.5 only works on transformers >= 4.49.0." + LATEST ) + # Gemma 4 must be before Gemma 3N and Gemma 3 + elif "gemma4" in model_types_all: + if not SUPPORTS_GEMMA4: + raise RuntimeError( + "Unsloth: Gemma 4 requires transformers >= 5.5.0" + LATEST + ) + os.environ["UNSLOTH_DISABLE_STATIC_GENERATION"] = "1" + os.environ["UNSLOTH_HIGH_PRECISION_LAYERNORM"] = "1" + # Disable flex_attention for Gemma-4: flex compile overhead is 2.7x slower + # than SDPA. Our attention patch ensures Q/K/V dtype alignment for SDPA. + os.environ["UNSLOTH_ENABLE_FLEX_ATTENTION"] = "0" # Gemma 3N must be before Gemma 3 elif "gemma3n" in model_types_all: if transformers_version < Version("4.53.0"): @@ -1407,8 +1421,14 @@ class FastModel(FastBaseModel): architectures = [] is_vlm = any(x.endswith("ForConditionalGeneration") for x in architectures) is_vlm = is_vlm or hasattr(model_config, "vision_config") + # If num_labels is set, use AutoModelForSequenceClassification + _num_labels = kwargs.get("num_labels", None) if auto_model is None: - if is_vlm: + if _num_labels is not None: + from transformers import AutoModelForSequenceClassification + + auto_model = AutoModelForSequenceClassification + elif is_vlm: # Check if the model's auto_map supports the VLM auto class. # Some VL models (e.g. Nemotron-VL) only register AutoModelForCausalLM # in their auto_map, not AutoModelForImageTextToText/AutoModelForVision2Seq. diff --git a/unsloth/models/qwen3.py b/unsloth/models/qwen3.py index b93dddb186..3129483be8 100644 --- a/unsloth/models/qwen3.py +++ b/unsloth/models/qwen3.py @@ -302,6 +302,9 @@ def Qwen3Attention_fast_forward_inference( # or else error self.rotary_emb.extend_rope_embedding(Vn, seq_len + 2) cos, sin = self.rotary_emb.get_cached(kv_seq_len, Qn.device.index) + # Transformers 5.x: position_ids may be [batch, full_seq_len]; slice to last + if position_ids.dim() >= 2 and position_ids.shape[-1] > 1: + position_ids = position_ids[:, -1:] cos = cos[position_ids].unsqueeze(1) sin = sin[position_ids].unsqueeze(1) h = self.half_head_dim diff --git a/unsloth/models/rl_replacements.py b/unsloth/models/rl_replacements.py index 9f555416d4..2544afe82e 100755 --- a/unsloth/models/rl_replacements.py +++ b/unsloth/models/rl_replacements.py @@ -542,6 +542,37 @@ def grpo_trainer__generate_and_score_completions(function_name, function): function = patched + # Transformers 5.x: Extend mm_token_type_ids for completion tokens (Qwen3VL M-RoPE). + # TRL handles token_type_ids but not mm_token_type_ids. + _tt_search = ( + 'if "token_type_ids" in forward_kwargs:\n' + ' token_type_ids = forward_kwargs["token_type_ids"]\n' + ' forward_kwargs["token_type_ids"] = torch.cat(\n' + " [token_type_ids, token_type_ids.new_zeros(completion_ids.shape)], dim=1\n" + " )" + ) + _tt_replace = ( + _tt_search + "\n" + ' if "mm_token_type_ids" in forward_kwargs:\n' + ' mm_tti = forward_kwargs["mm_token_type_ids"]\n' + ' forward_kwargs["mm_token_type_ids"] = torch.cat(\n' + " [mm_tti, mm_tti.new_zeros(completion_ids.shape)], dim=1\n" + " )" + ) + function = function.replace(_tt_search, _tt_replace) + + # Save mm_token_type_ids to output dict alongside token_type_ids + _save_search = ( + 'if "token_type_ids" in forward_kwargs:\n' + ' output["token_type_ids"] = forward_kwargs["token_type_ids"]' + ) + _save_replace = ( + _save_search + "\n" + ' if "mm_token_type_ids" in forward_kwargs:\n' + ' output["mm_token_type_ids"] = forward_kwargs["mm_token_type_ids"]' + ) + function = function.replace(_save_search, _save_replace) + return function @@ -714,6 +745,9 @@ def grpo_trainer__get_per_token_logps_and_entropies(function_name, function): kwargs.get("pixel_attention_mask", None), kwargs.get("image_sizes", None), ) + # Transformers 5.x needs token_type_ids/mm_token_type_ids for some vision models + token_type_ids = kwargs.get("token_type_ids", None) + mm_token_type_ids = kwargs.get("mm_token_type_ids", None) unwrapped_model = self.accelerator.unwrap_model( model, keep_fp32_wrapper = False @@ -831,6 +865,10 @@ def grpo_trainer__get_per_token_logps_and_entropies(function_name, function): if logit_scale_divide is None: logit_scale_divide = 0 + # Transformers 5.x needs token_type_ids/mm_token_type_ids for some vision models + token_type_ids_chunks = chunk_optional(token_type_ids, B) + mm_token_type_ids_chunks = chunk_optional(mm_token_type_ids, B) + zipped_inputs = zip( input_ids_chunks, attention_mask_chunks, @@ -838,6 +876,8 @@ def grpo_trainer__get_per_token_logps_and_entropies(function_name, function): image_grid_thw_chunks, pixel_attention_mask_chunks, image_sizes_chunks, + token_type_ids_chunks, + mm_token_type_ids_chunks, ) os.environ["UNSLOTH_RETURN_HIDDEN_STATES"] = "1" @@ -849,7 +889,16 @@ def grpo_trainer__get_per_token_logps_and_entropies(function_name, function): image_grid_thw_chunk, pixel_attention_mask_chunk, image_sizes_chunk, + token_type_ids_chunk, + mm_token_type_ids_chunk, ) in zipped_inputs: + _extra_vision_kwargs = {} + if token_type_ids_chunk is not None: + _extra_vision_kwargs["token_type_ids"] = token_type_ids_chunk + if mm_token_type_ids_chunk is not None: + _extra_vision_kwargs["mm_token_type_ids"] = ( + mm_token_type_ids_chunk + ) with torch.amp.autocast( device_type = "cuda", dtype = self._autocast_dtype ): @@ -861,6 +910,7 @@ def grpo_trainer__get_per_token_logps_and_entropies(function_name, function): image_grid_thw = image_grid_thw_chunk, pixel_attention_mask = pixel_attention_mask_chunk, image_sizes = image_sizes_chunk, + **_extra_vision_kwargs, ).logits completion_input_ids_chunk = input_ids_chunk[ @@ -893,6 +943,7 @@ def grpo_trainer__get_per_token_logps_and_entropies(function_name, function): pixel_attention_mask = pixel_attention_mask_chunk, image_sizes = image_sizes_chunk, logits_to_keep = logits_to_keep + 1, + **_extra_vision_kwargs, ).logits logits_chunk = logits_chunk[:, :-1, :] @@ -993,6 +1044,9 @@ def grpo_trainer_compute_loss(function_name, function): inputs.get("pixel_attention_mask", None), inputs.get("image_sizes", None), ) + # Transformers 5.x needs token_type_ids/mm_token_type_ids for some vision models + token_type_ids = inputs.get("token_type_ids", None) + mm_token_type_ids = inputs.get("mm_token_type_ids", None) num_items_in_batch = inputs.get("num_items_in_batch", None) sampling_per_token_logps = inputs.get("sampling_per_token_logps", None) current_gradient_accumulation_steps = self.current_gradient_accumulation_steps @@ -1136,6 +1190,8 @@ def grpo_trainer_compute_loss(function_name, function): current_gradient_accumulation_steps = current_gradient_accumulation_steps, num_processes = num_processes, sampling_per_token_logps = sampling_per_token_logps, + token_type_ids = token_type_ids, + mm_token_type_ids = mm_token_type_ids, ) else: # to ensure backwards compatibility with trl 0.15.2 and maybe even 0.17 @@ -1154,6 +1210,8 @@ def grpo_trainer_compute_loss(function_name, function): logit_scale_multiply = logit_scale_multiply, logit_scale_divide = logit_scale_divide, attention_mask = attention_mask, + token_type_ids = token_type_ids, + mm_token_type_ids = mm_token_type_ids, ) ) if "train" in self._metrics: diff --git a/unsloth/models/vision.py b/unsloth/models/vision.py index f558aa3f00..5abeb3a81a 100644 --- a/unsloth/models/vision.py +++ b/unsloth/models/vision.py @@ -216,6 +216,10 @@ def unsloth_base_fast_generate( kwargs["pixel_values"] = kwargs["pixel_values"].to(dtype) except: pass + try: + kwargs["pixel_values_videos"] = kwargs["pixel_values_videos"].to(dtype) + except: + pass # Mixed precision autocast if os.environ.get("UNSLOTH_FORCE_FLOAT32", "0") == "1": @@ -597,8 +601,6 @@ class FastBaseModel: custom_datatype = None correct_dtype = None - # Stop SDPA for some archs like Pixtral / Mistral3 - flex_attn_impl = None if auto_config is None: auto_config = AutoConfig.from_pretrained( model_name, @@ -609,7 +611,14 @@ class FastBaseModel: model_class = auto_model._model_mapping[auto_config.__class__] except Exception: model_class = None - flex_attn_impl = prefer_flex_attn_if_supported(model_class, auto_config) + if model_class is None: + # When model_class cannot be resolved (remote-code or unmapped + # configs), preserve the old fallback of sdpa when supported. + attn_impl = _set_attn_impl( + auto_config, "sdpa" if supports_sdpa else "eager" + ) + else: + attn_impl = determine_attention_implementation(model_class, auto_config) # Handle FP8 models: get_model_name has already redirected this to BF16 sibling if the model ships with # FP8 weights. We just need to update it here for sanity. @@ -620,21 +629,15 @@ class FastBaseModel: except Exception: model_class = None - model_type = str(getattr(auto_config, "model_type", "")).lower() - if model_type.startswith("gemma3n"): - # Gemma3N variants initialize timm-based vision towers which do - # not support flex_attention, so default to eager unless overridden. - default_attn_impl = "eager" - else: - default_attn_impl = "flex_attention" if flex_attn_impl else "sdpa" if not ("attn_implementation" in kwargs): - kwargs["attn_implementation"] = default_attn_impl + kwargs["attn_implementation"] = attn_impl if not supports_sdpa and kwargs.get("attn_implementation") == "sdpa": - if os.environ.get("UNSLOTH_ENABLE_FLEX_ATTENTION", "0") == "0": - print( - f"Unsloth: {model_type_arch.title()} does not support SDPA - switching to fast eager." - ) + print( + f"Unsloth: {model_type_arch.title()} does not support SDPA - switching to fast eager." + ) del kwargs["attn_implementation"] + # Re-stamp config so it stays consistent with the actual impl + _set_attn_impl(auto_config, "eager") bnb_config = None user_quantization_config = kwargs.get("quantization_config", None) @@ -788,6 +791,15 @@ class FastBaseModel: if not fast_inference: # Prevent load_in_fp8 from being forwarded into HF internal model loading load_in_fp8 = kwargs.pop("load_in_fp8", None) + # Transformers 5.x @strict config classes reject unexpected kwargs. + # Move config-level attributes onto the config object directly. + _num_labels = kwargs.pop("num_labels", None) + if _num_labels is not None: + model_config.num_labels = _num_labels + for _cfg_key in ("id2label", "label2id", "max_position_embeddings"): + _cfg_val = kwargs.pop(_cfg_key, None) + if _cfg_val is not None: + setattr(model_config, _cfg_key, _cfg_val) model = auto_model.from_pretrained( model_name, config = model_config, @@ -1021,6 +1033,15 @@ class FastBaseModel: f"Unsloth: Warning - VLM processor fallback returned None for model_type={model_type_arch}", file = sys.stderr, ) + # Backwards compat: if processor has no chat_template (e.g. old saves without + # chat_template.jinja) but the inner tokenizer does, copy it to the processor. + if ( + hasattr(tokenizer, "tokenizer") + and getattr(tokenizer, "chat_template", None) is None + and getattr(tokenizer.tokenizer, "chat_template", None) is not None + ): + tokenizer.chat_template = tokenizer.tokenizer.chat_template + if hasattr(tokenizer, "tokenizer"): __tokenizer = tokenizer.tokenizer # Add padding side as well @@ -1277,7 +1298,59 @@ class FastBaseModel: model, use_gradient_checkpointing = use_gradient_checkpointing, ) + # Gemma4 ClippableLinear wraps nn.Linear -- PEFT can't inject LoRA on it directly. + # Monkey-patch PEFT to target the inner .linear child instead. + _clippable_linear_cls = None + try: + from transformers.models.gemma4.modeling_gemma4 import ( + Gemma4ClippableLinear as _clippable_linear_cls, + ) + except ImportError: + pass + if _clippable_linear_cls is not None: + from peft.tuners.lora.model import LoraModel as _LoraModel + + _original_car = _LoraModel._create_and_replace + + def _patched_car( + self, + peft_config, + adapter_name, + target, + target_name, + parent, + current_key = None, + **kwargs, + ): + if isinstance(target, _clippable_linear_cls): + return _original_car( + self, + peft_config, + adapter_name, + target.linear, + "linear", + target, + current_key = current_key, + **kwargs, + ) + return _original_car( + self, + peft_config, + adapter_name, + target, + target_name, + parent, + current_key = current_key, + **kwargs, + ) + + _LoraModel._create_and_replace = _patched_car + model = _get_peft_model(model, lora_config) + + # Restore original PEFT method + if _clippable_linear_cls is not None: + _LoraModel._create_and_replace = _original_car # Apply QAT + LoRA if specified if qat_scheme is not None: print("Unsloth: Applying QAT to mitigate quantization degradation") @@ -1375,7 +1448,7 @@ class FastBaseModel: # after this point, so we intercept gradient_checkpointing_enable # to always force use_reentrant=True for Gemma3N. _model_type = getattr(getattr(model, "config", None), "model_type", "") or "" - if "gemma3n" in _model_type.lower(): + if "gemma3n" in _model_type.lower() or "gemma4" in _model_type.lower(): _original_gc_enable = model.gradient_checkpointing_enable def _gc_enable_reentrant(**kwargs): diff --git a/unsloth/ollama_template_mappers.py b/unsloth/ollama_template_mappers.py index 1bf77461d9..728b08813a 100644 --- a/unsloth/ollama_template_mappers.py +++ b/unsloth/ollama_template_mappers.py @@ -1199,6 +1199,21 @@ TEMPLATE """{{- range $i, $_ := .Messages }} OLLAMA_TEMPLATES["gemma-3n"] = gemma3n_ollama OLLAMA_TEMPLATES["gemma3n"] = gemma3n_ollama +# =========================================== Gemma-4 +gemma4_ollama = ''' +FROM {__FILE_LOCATION__} +TEMPLATE """{{- range $i, $_ := .Messages }} +{{- $last := eq (len (slice $.Messages $i)) 1 }} +<|turn>{{ .Role }} +{{ .Content }}{{ if not $last }} +{{ end }} +{{- end }} +<|turn>model +""" +''' +OLLAMA_TEMPLATES["gemma-4"] = gemma4_ollama +OLLAMA_TEMPLATES["gemma4"] = gemma4_ollama + # =========================================== GPT-OSS # Ollama from https://ollama.com/library/gpt-oss:latest/blobs/fa6710a93d78 @@ -1961,6 +1976,16 @@ OLLAMA_TEMPLATE_TO_MODEL_MAPPER = { "google/medgemma-27b-text-it", "unsloth/medgemma-27b-text-it-bnb-4bit", ), + "gemma4": ( + "unsloth/gemma-4-E2B-it", + "unsloth/gemma-4-E2B", + "unsloth/gemma-4-E4B-it", + "unsloth/gemma-4-E4B", + "unsloth/gemma-4-31B-it", + "unsloth/gemma-4-31B", + "unsloth/gemma-4-26B-A4B-it", + "unsloth/gemma-4-26B-A4B", + ), "gemma3n": ( "unsloth/gemma-3n-E4B-it-unsloth-bnb-4bit", "unsloth/gemma-3n-E4B-it", diff --git a/unsloth/save.py b/unsloth/save.py index 1759d86fb1..3c318fab02 100644 --- a/unsloth/save.py +++ b/unsloth/save.py @@ -2533,12 +2533,15 @@ def unsloth_convert_lora_to_ggml_and_push_to_hub( ) print(f"The output file will be {output_file}") - command = f"python3 llama.cpp/convert-lora-to-ggml.py {lora_directory_push} {output_file} llama" - try: with subprocess.Popen( - command, - shell = True, + [ + sys.executable, + "llama.cpp/convert-lora-to-ggml.py", + lora_directory_push, + output_file, + "llama", + ], stdout = subprocess.PIPE, stderr = subprocess.PIPE, bufsize = 1, @@ -2550,7 +2553,7 @@ def unsloth_convert_lora_to_ggml_and_push_to_hub( print(line, end = "", flush = True) sp.wait() if sp.returncode != 0: - raise subprocess.CalledProcessError(sp.returncode, command) + raise subprocess.CalledProcessError(sp.returncode, sp.args) except subprocess.CalledProcessError as e: print(f"Error: Conversion failed with return code {e.returncode}") return @@ -2612,12 +2615,15 @@ def unsloth_convert_lora_to_ggml_and_save_locally( ) print(f"The output file will be {output_file}") - command = f"python3 llama.cpp/convert-lora-to-ggml.py {save_directory} {output_file} llama" - try: with subprocess.Popen( - command, - shell = True, + [ + sys.executable, + "llama.cpp/convert-lora-to-ggml.py", + save_directory, + output_file, + "llama", + ], stdout = subprocess.PIPE, stderr = subprocess.PIPE, bufsize = 1, @@ -2629,7 +2635,7 @@ def unsloth_convert_lora_to_ggml_and_save_locally( print(line, end = "", flush = True) sp.wait() if sp.returncode != 0: - raise subprocess.CalledProcessError(sp.returncode, command) + raise subprocess.CalledProcessError(sp.returncode, sp.args) except subprocess.CalledProcessError as e: print(f"Error: Conversion failed with return code {e.returncode}") return @@ -2777,19 +2783,79 @@ def unsloth_generic_save( elif save_method == "merged_4bit_forced": save_method = "merged_4bit" - merge_and_overwrite_lora( - get_model_name, - model = model, - tokenizer = tokenizer, - save_directory = save_directory, - push_to_hub = push_to_hub, - private = private, - token = token, - save_method = save_method, - output_dtype = None, - low_disk_space_usage = True, - use_temp_file = False, - ) + # Full-finetuned models (no LoRA) cannot use merge_and_overwrite_lora + # since there are no adapters to merge. Fall back to save_pretrained. + # This mirrors the non-PeftModel handling in save_pretrained_torchao + # and the GGUF save path. + _is_peft = isinstance(model, PeftModel) + if not _is_peft: + if not is_main_process: + return + + # Honor merged_16bit by casting to the target dtype if needed + _save_kwargs = dict( + safe_serialization = safe_serialization, + max_shard_size = max_shard_size, + variant = variant, + ) + if "16bit" in save_method: + _target_dtype = ( + torch.bfloat16 if torch.cuda.is_bf16_supported() else torch.float16 + ) + _save_kwargs["state_dict"] = { + k: v.to(dtype = _target_dtype) if v.is_floating_point() else v + for k, v in model.state_dict().items() + } + + if push_to_hub: + print(f"Unsloth: Pushing full fine-tuned model to '{save_directory}' ...") + model.push_to_hub( + repo_id = save_directory, + token = token, + private = private, + commit_message = commit_message, + create_pr = create_pr, + revision = revision, + commit_description = commit_description, + tags = tags, + **_save_kwargs, + ) + if tokenizer is not None: + old_padding_side = tokenizer.padding_side + tokenizer.padding_side = "left" + tokenizer.push_to_hub( + save_directory, + token = token, + private = private, + commit_message = commit_message, + create_pr = create_pr, + revision = revision, + ) + tokenizer.padding_side = old_padding_side + else: + print(f"Unsloth: Saving full fine-tuned model to '{save_directory}' ...") + model.save_pretrained(save_directory, **_save_kwargs) + if tokenizer is not None: + old_padding_side = tokenizer.padding_side + tokenizer.padding_side = "left" + tokenizer.save_pretrained(save_directory) + tokenizer.padding_side = old_padding_side + + print(f"Unsloth: Model saved successfully to '{save_directory}'") + else: + merge_and_overwrite_lora( + get_model_name, + model = model, + tokenizer = tokenizer, + save_directory = save_directory, + push_to_hub = push_to_hub, + private = private, + token = token, + save_method = save_method, + output_dtype = None, + low_disk_space_usage = True, + use_temp_file = False, + ) if push_to_hub and datasets: try: