diff --git a/install.ps1 b/install.ps1 index 515f700fea..fa0d160a33 100644 --- a/install.ps1 +++ b/install.ps1 @@ -1574,7 +1574,7 @@ shell.Run cmd, 0, False # popping a UAC/DiskPart prompt RunAsInvoker can't suppress (manifest is # asInvoker). So only probe when a HIP SDK is present (hipinfo found -> # un-elevated) or the user opts in; else fall through to WMI name inference - # (enough to pick ROCm wheels + lemonade llama.cpp). + # (enough to pick ROCm wheels + the ROCm llama.cpp prebuilt). # An explicit opt-out (UNSLOTH_ENABLE_AMD_SMI=0/false/no/off) wins over the # HIP-SDK heuristic: a HIP SDK binary with a broken runtime can still pop the # prompt, so $HipSdkInstalled must NOT silently re-enable it. @@ -1625,7 +1625,7 @@ shell.Run cmd, 0, False # ── Arch resolution: env-var override → name inference ────────────── # Runs even when the hipinfo/amd-smi probe could NOT confirm a runtime # ($HasROCm false): the gfx arch inferred from the WMI GPU name lets the - # studio setup forward --rocm-gfx and pull a GPU-accelerated (lemonade) + # studio setup forward --rocm-gfx and pull a GPU-accelerated ROCm # llama.cpp, which bundles its own ROCm runtime. PyTorch's ROCm wheels # still require a confirmed HIP SDK -- they stay gated on $HasROCm below. if (-not $ROCmGfxArch) { @@ -1636,7 +1636,7 @@ shell.Run cmd, 0, False substep "gfx arch from UNSLOTH_ROCM_GFX_ARCH env override: $ROCmGfxArch" "Cyan" } # 2. Best-effort name → arch lookup from marketing name (amd-smi / WMI). - # Targets only arches the lemonade-sdk ROCm prebuilts cover + # Targets only arches the ROCm prebuilts cover # (gfx120X/110X/1151/1150/103X); unknown names fall back to CPU. elseif ($ROCmGpuLabel) { $nameArchTable = @( @@ -1647,9 +1647,9 @@ shell.Run cmd, 0, False @{ P = "RX 7900|RX 7800|RX 7700(?!S)|PRO W7900|PRO W7800|PRO W7700"; A = "gfx1100" } # RDNA 3 desktop/workstation (Navi 31) @{ P = "RX 7600|RX 7700S|RX 7650|PRO W7600|PRO W7500|PRO V710"; A = "gfx1102" } # RDNA 3 (Navi 33) @{ P = "780M|760M|740M|Phoenix|Hawk Point|Z1 Extreme|Z2 Extreme"; A = "gfx1103" } # RDNA 3 iGPU (Phoenix / Hawk Point) - @{ P = "RX 6900|RX 6800|RX 6750|RX 6700|PRO W6800|PRO W6900"; A = "gfx1030" } # RDNA 2 (Navi 21) -- lemonade gfx103X - @{ P = "RX 6650|RX 6600|PRO W6600|PRO W6650"; A = "gfx1032" } # RDNA 2 (Navi 23) -- lemonade gfx103X - @{ P = "RX 6500|RX 6400|RX 6300|PRO W6400|PRO W6500"; A = "gfx1034" } # RDNA 2 (Navi 24) -- lemonade gfx103X + @{ P = "RX 6900|RX 6800|RX 6750|RX 6700|PRO W6800|PRO W6900"; A = "gfx1030" } # RDNA 2 (Navi 21) -- gfx103X family + @{ P = "RX 6650|RX 6600|PRO W6600|PRO W6650"; A = "gfx1032" } # RDNA 2 (Navi 23) -- gfx103X family + @{ P = "RX 6500|RX 6400|RX 6300|PRO W6400|PRO W6500"; A = "gfx1034" } # RDNA 2 (Navi 24) -- gfx103X family ) foreach ($row in $nameArchTable) { if ($ROCmGpuLabel -match $row.P) { @@ -1921,7 +1921,7 @@ shell.Run cmd, 0, False if ($SkipTorch) { # No-torch: install unsloth + unsloth-zoo with --no-deps, then # runtime deps (typer, safetensors, transformers, etc.) with --no-deps. - $baseInstallExit = Invoke-InstallCommand { uv pip install --python $VenvPython --no-deps --reinstall-package unsloth --reinstall-package unsloth-zoo "unsloth>=2026.6.4" unsloth-zoo } + $baseInstallExit = Invoke-InstallCommand { uv pip install --python $VenvPython --no-deps --reinstall-package unsloth --reinstall-package unsloth-zoo "unsloth>=2026.6.5" unsloth-zoo } if ($baseInstallExit -eq 0) { # Resolve pydantic WITH deps so pip pins pydantic-core # to the matching version (no-torch-runtime.txt below @@ -1935,7 +1935,7 @@ shell.Run cmd, 0, False } } } else { - $baseInstallExit = Invoke-InstallCommand { uv pip install --python $VenvPython --reinstall-package unsloth --reinstall-package unsloth-zoo "unsloth>=2026.6.4" unsloth-zoo } + $baseInstallExit = Invoke-InstallCommand { uv pip install --python $VenvPython --reinstall-package unsloth --reinstall-package unsloth-zoo "unsloth>=2026.6.5" unsloth-zoo } } if ($baseInstallExit -ne 0) { Write-Host "[ERROR] Failed to install unsloth (exit code $baseInstallExit)" -ForegroundColor Red @@ -1982,7 +1982,7 @@ shell.Run cmd, 0, False if ($SkipTorch) { # No-torch: install unsloth + unsloth-zoo with --no-deps, then # runtime deps (typer, safetensors, transformers, etc.) with --no-deps. - $baseInstallExit = Invoke-InstallCommand { uv pip install --python $VenvPython --no-deps --upgrade-package unsloth --upgrade-package unsloth-zoo "unsloth>=2026.6.4" unsloth-zoo } + $baseInstallExit = Invoke-InstallCommand { uv pip install --python $VenvPython --no-deps --upgrade-package unsloth --upgrade-package unsloth-zoo "unsloth>=2026.6.5" unsloth-zoo } if ($baseInstallExit -eq 0) { # Same pydantic-with-deps trick as the migrated branch. $baseInstallExit = Invoke-InstallCommand { uv pip install --python $VenvPython pydantic } @@ -1994,7 +1994,7 @@ shell.Run cmd, 0, False } } } elseif ($StudioLocalInstall) { - $baseInstallExit = Invoke-InstallCommand { uv pip install --python $VenvPython --upgrade-package unsloth "unsloth>=2026.6.4" unsloth-zoo } + $baseInstallExit = Invoke-InstallCommand { uv pip install --python $VenvPython --upgrade-package unsloth "unsloth>=2026.6.5" unsloth-zoo } } else { $baseInstallExit = Invoke-InstallCommand { uv pip install --python $VenvPython --upgrade-package unsloth -- "$PackageName" } } @@ -2022,7 +2022,7 @@ shell.Run cmd, 0, False Write-TauriLog "STEP" "Installing unsloth" substep "installing unsloth (this may take a few minutes)..." if ($StudioLocalInstall) { - $baseInstallExit = Invoke-InstallCommand { uv pip install --python $VenvPython unsloth-zoo "unsloth>=2026.6.4" --torch-backend=auto } + $baseInstallExit = Invoke-InstallCommand { uv pip install --python $VenvPython unsloth-zoo "unsloth>=2026.6.5" --torch-backend=auto } if ($baseInstallExit -ne 0) { Write-Host "[ERROR] Failed to install unsloth (exit code $baseInstallExit)" -ForegroundColor Red return (Exit-InstallFailure "Failed to install unsloth (exit code $baseInstallExit)" $baseInstallExit) diff --git a/install.sh b/install.sh index d13689150a..a9a74ac4e6 100755 --- a/install.sh +++ b/install.sh @@ -2422,7 +2422,7 @@ if [ "$_MIGRATED" = true ]; then # to prevent transitive torch resolution. run_install_cmd "install unsloth (migrated no-torch)" uv pip install --python "$_VENV_PY" --no-deps \ --reinstall-package unsloth --reinstall-package unsloth-zoo \ - "unsloth>=2026.6.4" unsloth-zoo + "unsloth>=2026.6.5" unsloth-zoo # Resolve pydantic WITH deps so pip pins pydantic-core to the # matching version (no-torch-runtime.txt below is --no-deps). # All transitive deps are torch-free. @@ -2435,7 +2435,7 @@ if [ "$_MIGRATED" = true ]; then else run_install_cmd "install unsloth (migrated)" uv pip install --python "$_VENV_PY" \ --reinstall-package unsloth --reinstall-package unsloth-zoo \ - "unsloth>=2026.6.4" unsloth-zoo + "unsloth>=2026.6.5" unsloth-zoo fi if [ "$STUDIO_LOCAL_INSTALL" = true ]; then substep "overlaying local repo (editable)..." @@ -2639,7 +2639,7 @@ elif [ -n "$TORCH_INDEX_URL" ]; then # runtime deps (typer, safetensors, transformers, etc.) with --no-deps. run_install_cmd "install unsloth (no-torch)" uv pip install --python "$_VENV_PY" --no-deps \ --upgrade-package unsloth --upgrade-package unsloth-zoo \ - "unsloth>=2026.6.4" unsloth-zoo + "unsloth>=2026.6.5" unsloth-zoo # Same pydantic-with-deps trick as the migrated branch. run_install_cmd "install pydantic (with deps for compatible core)" \ uv pip install --python "$_VENV_PY" pydantic @@ -2657,7 +2657,7 @@ elif [ -n "$TORCH_INDEX_URL" ]; then fi elif [ "$STUDIO_LOCAL_INSTALL" = true ]; then run_install_cmd "install unsloth (local)" uv pip install --python "$_VENV_PY" \ - --upgrade-package unsloth "unsloth>=2026.6.4" unsloth-zoo + --upgrade-package unsloth "unsloth>=2026.6.5" unsloth-zoo substep "overlaying local repo (editable)..." run_install_cmd "overlay local repo" uv pip install --python "$_VENV_PY" -e "$_REPO_ROOT" --no-deps substep "overlaying unsloth-zoo from git main..." @@ -2689,7 +2689,7 @@ else tauri_log "STEP" "Installing Unsloth" substep "installing unsloth (this may take a few minutes)..." if [ "$STUDIO_LOCAL_INSTALL" = true ]; then - run_install_cmd "install unsloth (auto torch backend)" uv pip install --python "$_VENV_PY" unsloth-zoo "unsloth>=2026.6.4" --torch-backend=auto + run_install_cmd "install unsloth (auto torch backend)" uv pip install --python "$_VENV_PY" unsloth-zoo "unsloth>=2026.6.5" --torch-backend=auto substep "overlaying local repo (editable)..." run_install_cmd "overlay local repo" uv pip install --python "$_VENV_PY" -e "$_REPO_ROOT" --no-deps substep "overlaying unsloth-zoo from git main..." diff --git a/studio/backend/cloudflare_tunnel.py b/studio/backend/cloudflare_tunnel.py index 7c088b6e1e..e5dba69452 100644 --- a/studio/backend/cloudflare_tunnel.py +++ b/studio/backend/cloudflare_tunnel.py @@ -24,12 +24,20 @@ from pathlib import Path from typing import Optional, Tuple # cloudflared logs the quick-tunnel URL; match only the URL so we do not depend -# on the surrounding wording, which Cloudflare may change. -_URL_RE = re.compile(r"https://[A-Za-z0-9-]+\.trycloudflare\.com") +# on the surrounding wording, which Cloudflare may change. The negative lookahead +# drops cloudflared's own API host, which appears in failure lines such as +# failed to request quick Tunnel: Post "https://api.trycloudflare.com/tunnel" +# and must never be mistaken for a usable tunnel URL. +_URL_RE = re.compile(r"https://(?!api\.)[A-Za-z0-9-]+\.trycloudflare\.com") + +# cloudflared logs this once per edge connection it establishes. Until at least +# one appears the quick-tunnel URL returns Cloudflare error 1033 (HTTP 530), so +# we wait for it before advertising the URL. +_REGISTERED_MARKER = "Registered tunnel connection" _RELEASE_BASE = "https://github.com/cloudflare/cloudflared/releases/latest/download" -_URL_TIMEOUT = 15.0 # seconds to wait for the public URL before giving up +_READY_TIMEOUT = 15.0 # seconds to wait for the URL + a registered edge connection _DOWNLOAD_TIMEOUT = 60 # urlopen timeout for the one-time binary download @@ -180,13 +188,24 @@ class CloudflareTunnel: upstream stays local-only. """ - def __init__(self, port: int, binary: str): + def __init__( + self, + port: int, + binary: str, + protocol: Optional[str] = None, + ): self.port = port self.binary = binary + # None lets cloudflared pick its default (quic, with its own http2 + # fallback); set to "http2" to force it when quic is blocked. + self.protocol = protocol self._proc: Optional[subprocess.Popen] = None self._lock = threading.Lock() + self._stopped = False self._url_event = threading.Event() + self._ready_event = threading.Event() self.url: Optional[str] = None + self.ready = False self.error: Optional[str] = None def start(self) -> None: @@ -197,25 +216,33 @@ class CloudflareTunnel: f"http://localhost:{self.port}", "--no-autoupdate", ] - proc = subprocess.Popen( - cmd, - stdout = subprocess.PIPE, - stderr = subprocess.STDOUT, - stdin = subprocess.DEVNULL, - text = True, - errors = "replace", - bufsize = 1, - **_windows_hidden_kwargs(), - ) + if self.protocol: + cmd += ["--protocol", self.protocol] with self._lock: + # A stop() that landed before us (e.g. a shutdown in the caller's + # register->start window) marks the tunnel stopped; spawning now would + # orphan a process nobody owns, so refuse. + if self._stopped: + return + proc = subprocess.Popen( + cmd, + stdout = subprocess.PIPE, + stderr = subprocess.STDOUT, + stdin = subprocess.DEVNULL, + text = True, + errors = "replace", + bufsize = 1, + **_windows_hidden_kwargs(), + ) self._proc = proc threading.Thread( target = self._reader, args = (proc,), name = "cloudflared-reader", daemon = True ).start() def _reader(self, proc: subprocess.Popen) -> None: - # Drain cloudflared's output, capture the first trycloudflare URL, and - # keep draining so it never blocks on a full pipe. + # Drain cloudflared's output: capture the first trycloudflare URL and the + # first edge-connection registration, and keep draining so it never + # blocks on a full pipe. try: if proc.stdout is not None: for line in proc.stdout: @@ -224,20 +251,35 @@ class CloudflareTunnel: if match: self.url = match.group(0) self._url_event.set() + if not self.ready and _REGISTERED_MARKER in line: + self.ready = True + self._ready_event.set() except Exception: pass finally: + # stdout closed -> cloudflared has exited. Record why, and unblock any + # waiters at once instead of letting them wait out the full timeout. if self.url is None: self.error = "cloudflared exited before emitting a tunnel URL" - self._url_event.set() + elif not self.ready: + self.error = "cloudflared exited before the tunnel connection registered" + self._url_event.set() + self._ready_event.set() - def wait_for_url(self, timeout: float = _URL_TIMEOUT) -> Optional[str]: - self._url_event.wait(timeout) - return self.url + def wait_for_ready(self, timeout: float = _READY_TIMEOUT) -> Optional[str]: + """Block until the tunnel is actually serving -- the URL has been minted + *and* at least one edge connection has registered -- or until timeout. + + Returns the URL only when ready, so callers never advertise a URL that + would return Cloudflare error 1033 (HTTP 530).""" + self._ready_event.wait(timeout) + return self.url if self.ready else None def stop(self) -> None: """Terminate the tunnel. Idempotent and safe to call from a signal handler.""" with self._lock: + # Mark stopped so a start() racing behind us refuses to spawn. + self._stopped = True proc, self._proc = self._proc, None if proc is None: return @@ -260,43 +302,74 @@ class CloudflareTunnel: # enough; the lock guards the start/stop/shutdown races. _active_tunnel: Optional[CloudflareTunnel] = None _active_lock = threading.Lock() +# Latched by stop_studio_tunnel so a shutdown landing *between* a start's retry +# attempts aborts the loop instead of starting a tunnel nobody will ever stop. +_shutdown_requested = False -def start_studio_tunnel(port: int, timeout: float = _URL_TIMEOUT) -> Optional[str]: - """Start a quick tunnel and return its public URL, or None (best-effort). +def start_studio_tunnel(port: int, timeout: float = _READY_TIMEOUT) -> Optional[str]: + """Start a quick tunnel and return its public URL once it is actually + serving, or None (best-effort). - On any failure (no binary, no URL within timeout, early crash) the tunnel is - stopped and None is returned, so the caller prints a hint and continues. + Waits for cloudflared to both mint the URL and register an edge connection + before returning, so the caller never advertises a URL that yields Cloudflare + error 1033 (HTTP 530). If a URL is minted but no connection registers within + the window (e.g. quic is blocked on this network), retries once forcing the + http2 protocol. On any failure the tunnel is stopped and None is returned. """ - global _active_tunnel + global _active_tunnel, _shutdown_requested binary = ensure_cloudflared() if not binary: return None - tunnel = CloudflareTunnel(port, binary) - # Register before start/wait so a shutdown during the URL wait can stop it. with _active_lock: - prior, _active_tunnel = _active_tunnel, tunnel - if prior is not None: - prior.stop() - try: - tunnel.start() - url = tunnel.wait_for_url(timeout) - except Exception: - url = None - if url: - return url - # No URL (or crash): drop it unless a concurrent shutdown already replaced it. - with _active_lock: - if _active_tunnel is tunnel: - _active_tunnel = None - tunnel.stop() + _shutdown_requested = False # fresh session + # Default protocol first (quic, with cloudflared's own http2 fallback); if a + # URL appears but no connection registers, quic is likely blocked -> retry + # once forcing http2. + for protocol in (None, "http2"): + # Create + register under the lock, and bail if a stop already landed + # (e.g. between this and the previous attempt) so we never start a tunnel + # after shutdown has run. + with _active_lock: + if _shutdown_requested: + _active_tunnel = None + return None + tunnel = CloudflareTunnel(port, binary, protocol = protocol) + prior, _active_tunnel = _active_tunnel, tunnel + if prior is not None: + prior.stop() + try: + tunnel.start() + url = tunnel.wait_for_ready(timeout) + except Exception: + url = None + if url: + return url + saw_url = tunnel.url is not None + # Not ready: drop it, but only if we are still the active tunnel. + with _active_lock: + was_active = _active_tunnel is tunnel + if was_active: + _active_tunnel = None + tunnel.stop() + # A concurrent shutdown or start took over while we waited; retrying would + # spawn a tunnel nobody owns (orphaned after shutdown), so bail instead. + if not was_active: + return None + # No URL at all is an API/network failure, not a protocol one; forcing + # http2 will not help, so do not burn another window on it. + if not saw_url: + return None return None def stop_studio_tunnel() -> None: """Terminate the active tunnel, if any. Idempotent.""" - global _active_tunnel + global _active_tunnel, _shutdown_requested with _active_lock: + # Latch so an in-flight start_studio_tunnel won't start a fresh tunnel + # (e.g. its http2 retry) after we have already torn down. + _shutdown_requested = True tunnel, _active_tunnel = _active_tunnel, None if tunnel is not None: tunnel.stop() diff --git a/studio/backend/core/inference/llama_cpp.py b/studio/backend/core/inference/llama_cpp.py index 75da00451e..16f6376aab 100644 --- a/studio/backend/core/inference/llama_cpp.py +++ b/studio/backend/core/inference/llama_cpp.py @@ -122,15 +122,11 @@ _INTENT_SIGNAL = re.compile( ) _MAX_REPROMPTS = 1 -# Without max_tokens, llama-server defaults n_predict = n_ctx (up to 262144 for -# Qwen3.5), causing many-minute zombie decodes when cancel fails. -# t_max_predict_ms is a wall-clock backstop but per the llama.cpp README only -# fires after a newline, so we keep a token cap as the front-line limiter. -# The cap is the effective context length when known, else this floor. 4096 was -# too low: Qwen3 / gpt-oss reasoning traces and max_tokens-omitting OpenAI-API -# callers (langchain, llama-index, curl) got truncated mid-sentence. +# Default max_tokens to the effective context when known. The floor is high +# enough for reasoning-heavy GGUFs and max_tokens-omitting API clients. _DEFAULT_MAX_TOKENS_FLOOR = 32768 -_DEFAULT_T_MAX_PREDICT_MS = 600_000 # 10 min +_DEFAULT_FIRST_TOKEN_TIMEOUT_S = 1200.0 # 20 min +_DEFAULT_STREAM_STALL_TIMEOUT_S = 120.0 # 2 min _REPROMPT_MAX_CHARS = 2000 _FORCED_REPEAT_PLAN_SIGNAL = re.compile( r"\b(?:i\s+will|i'll|let\s+me|going\s+to|need\s+to|call|use|run|search|fetch|render)\b", @@ -1075,7 +1071,7 @@ class LlamaCppBackend: # ── Binary discovery ────────────────────────────────────────── @staticmethod - def _find_llama_server_binary() -> Optional[str]: + def _find_llama_server_binary(*, include_denied: bool = False) -> Optional[str]: """ Locate the llama-server binary. @@ -1092,28 +1088,70 @@ class LlamaCppBackend: """ binary_name = "llama-server.exe" if sys.platform == "win32" else "llama-server" + def _file_status(p: Path) -> str: + # "file", "absent", or "denied" (exists but stays access-denied + # across a short retry: Windows AV/ACL or an install replace in + # flight). is_file() raises PermissionError (WinError 5) instead of + # returning False for the locked case, so never treat it as missing. + for _ in range(5): + try: + return "file" if p.is_file() else "absent" + except PermissionError: + time.sleep(0.2) + except OSError: + return "absent" + return "denied" + + def _is_file(p: Path) -> bool: + return _file_status(p) == "file" + + def _layout_candidates(d: Path) -> list: + # build layouts probed under a llama.cpp dir, highest priority first + cands = [d / binary_name, d / "build" / "bin" / binary_name] + if sys.platform == "win32": + cands.append(d / "build" / "bin" / "Release" / binary_name) + return cands + + def _unavailable(p: object) -> None: + # a pinned or managed binary that exists but is access-denied: report + # it instead of silently downgrading to a lower-priority llama-server + logger.warning( + f"llama-server at {p} exists but is access-denied (antivirus or " + "an in-flight install); not falling back to another binary, " + "retry once it is released" + ) + return None + + def _scan_pinned(paths: list): + # first existing candidate wins -> (path, None); a present-but-denied + # one -> (None, denied_path) so the caller reports it rather than + # skipping to a lower-priority location. include_denied returns the + # locked path instead: diffusion asset lookup only needs its dir. + for p in paths: + st = _file_status(p) + if st == "file": + return str(p), None + if st == "denied": + return (str(p), None) if include_denied else (None, p) + return None, None + # 1. Env var: direct path to binary env_path = os.environ.get("LLAMA_SERVER_PATH") - if env_path and Path(env_path).is_file(): - return env_path + if env_path: + hit, locked = _scan_pinned([Path(env_path)]) + if locked is not None: + return _unavailable(locked) + if hit: + return hit # 1b. UNSLOTH_LLAMA_CPP_PATH: custom llama.cpp install dir custom_llama_cpp = os.environ.get("UNSLOTH_LLAMA_CPP_PATH") if custom_llama_cpp: - custom_dir = Path(custom_llama_cpp) - # Root dir (make builds) - root_bin = custom_dir / binary_name - if root_bin.is_file(): - return str(root_bin) - # build/bin/ (cmake on Linux) - cmake_bin = custom_dir / "build" / "bin" / binary_name - if cmake_bin.is_file(): - return str(cmake_bin) - # build/bin/Release/ (cmake on Windows) - if sys.platform == "win32": - win_bin = custom_dir / "build" / "bin" / "Release" / binary_name - if win_bin.is_file(): - return str(win_bin) + hit, locked = _scan_pinned(_layout_candidates(Path(custom_llama_cpp))) + if locked is not None: + return _unavailable(locked) + if hit: + return hit # 2-4. Match installer layout: env-mode -> $STUDIO_HOME/llama.cpp; # default/HOME-redirect -> ~/.unsloth/llama.cpp (sibling of studio). @@ -1145,31 +1183,18 @@ class LlamaCppBackend: _seen_roots.add(k) _unique_roots.append(r) for unsloth_home in _unique_roots: - home_root = unsloth_home / binary_name - if home_root.is_file(): - return str(home_root) - home_linux = unsloth_home / "build" / "bin" / binary_name - if home_linux.is_file(): - return str(home_linux) - if sys.platform == "win32": - home_win = unsloth_home / "build" / "bin" / "Release" / binary_name - if home_win.is_file(): - return str(home_win) + hit, locked = _scan_pinned(_layout_candidates(unsloth_home)) + if locked is not None: + return _unavailable(locked) + if hit: + return hit - # 5-6. Legacy: in-tree build (older setup.sh / setup.ps1) + # 5-6. Legacy: in-tree build (older setup.sh / setup.ps1). A fallback, + # so a denied candidate here just continues (no no-fallback halt). project_root = Path(__file__).resolve().parents[4] - # Root dir (make builds) - root_path = project_root / "llama.cpp" / binary_name - if root_path.is_file(): - return str(root_path) - # build/bin/ (cmake builds) - build_path = project_root / "llama.cpp" / "build" / "bin" / binary_name - if build_path.is_file(): - return str(build_path) - if sys.platform == "win32": - win_path = project_root / "llama.cpp" / "build" / "bin" / "Release" / binary_name - if win_path.is_file(): - return str(win_path) + for p in _layout_candidates(project_root / "llama.cpp"): + if _is_file(p): + return str(p) # 7. System PATH system_path = shutil.which("llama-server") @@ -1178,7 +1203,7 @@ class LlamaCppBackend: # 8. Legacy: extracted to bin/ bin_path = project_root / "bin" / binary_name - if bin_path.is_file(): + if _is_file(bin_path): return str(bin_path) return None @@ -2442,7 +2467,9 @@ class LlamaCppBackend: visual_bin = os.environ.get("DG_VISUAL_BIN") if not visual_bin: name = "llama-diffusion-gemma-visual-server" + (".exe" if os.name == "nt" else "") - base = self._find_llama_server_binary() + # include_denied: a transiently locked llama-server still pins the + # install dir so the adjacent visual-server can be found + base = self._find_llama_server_binary(include_denied = True) if base: base_dir = Path(base).parent for cand in ( @@ -3507,6 +3534,15 @@ class LlamaCppBackend: ) if not binary: + # distinguish a transiently locked binary (antivirus / in-flight + # install) from a missing one so the user retries, not reinstalls + locked = self._find_llama_server_binary(include_denied = True) + if locked: + raise RuntimeError( + f"llama-server at {locked} is temporarily unavailable " + "(access-denied; antivirus or an in-flight install). " + "Retry the load once it is released." + ) raise RuntimeError( "llama-server binary not found. " "Run setup.sh to build it, install llama.cpp, " @@ -5308,28 +5344,84 @@ class LlamaCppBackend: @staticmethod def _iter_text_cancellable( - response: "httpx.Response", cancel_event: Optional[threading.Event] = None + response: "httpx.Response", + cancel_event: Optional[threading.Event] = None, + stall_timeout_s: float = _DEFAULT_STREAM_STALL_TIMEOUT_S, + first_token_deadline: Optional[float] = None, + post_first_chunk_read_timeout_s: Optional[float] = _DEFAULT_STREAM_STALL_TIMEOUT_S, ) -> Generator[str, None, None]: - """Iterate an httpx streaming response with cancel support. - - Checks cancel_event between chunks and on ReadTimeout; the - _stream_with_retry watcher also closes the response on cancel. - """ + """Iterate a stream while polling cancel and stall timeouts.""" text_iter = response.iter_text() + if first_token_deadline is None: + first_token_deadline = time.monotonic() + _DEFAULT_FIRST_TOKEN_TIMEOUT_S + last_chunk_at: Optional[float] = None while True: if cancel_event is not None and cancel_event.is_set(): response.close() return try: + if last_chunk_at is None: + remaining_s = first_token_deadline - time.monotonic() + if remaining_s <= 0: + raise httpx.ReadTimeout("The model did not produce a first token in time.") + LlamaCppBackend._set_stream_read_timeout(response, remaining_s) chunk = next(text_iter) + if chunk: + if last_chunk_at is None and post_first_chunk_read_timeout_s is not None: + LlamaCppBackend._set_stream_read_timeout( + response, + post_first_chunk_read_timeout_s, + ) + last_chunk_at = time.monotonic() yield chunk except StopIteration: return except httpx.ReadTimeout: - # No data within the timeout window -- loop back and re-check - # cancel_event. + now = time.monotonic() + if last_chunk_at is None: + if now >= first_token_deadline: + raise + elif now - last_chunk_at >= stall_timeout_s: + raise httpx.ReadTimeout("The model stopped producing tokens mid-response.") continue + @staticmethod + def _set_stream_read_timeout(response: "httpx.Response", read_timeout_s: float) -> None: + """Lower only post-header stream reads; keep prefill timeout long.""" + try: + timeout_ext = response.request.extensions.get("timeout") + if isinstance(timeout_ext, dict): + timeout_ext["read"] = read_timeout_s + except Exception: + logger.debug("Could not lower response read timeout", exc_info = True) + + @staticmethod + def _shutdown_active_httpx_sockets(client: "httpx.Client") -> None: + """Best-effort interrupt for a sync httpx request blocked before headers.""" + try: + pool = getattr(getattr(client, "_transport", None), "_pool", None) + connections = list(getattr(pool, "_connections", []) or []) + for connection in connections: + inner = getattr(connection, "_connection", None) + stream = getattr(inner, "_network_stream", None) + sock = getattr(stream, "_sock", None) + if sock is None: + continue + try: + sock.shutdown(socket.SHUT_RDWR) + except OSError: + pass + try: + sock.close() + except OSError: + pass + except Exception: + logger.debug("Could not shutdown active httpx socket", exc_info = True) + try: + client.close() + except Exception: + logger.debug("Could not close httpx client", exc_info = True) + @staticmethod @contextlib.contextmanager def _stream_with_retry( @@ -5338,38 +5430,28 @@ class LlamaCppBackend: payload: dict, cancel_event: Optional[threading.Event] = None, headers: Optional[dict] = None, + first_token_deadline: Optional[float] = None, ): - """Open an httpx streaming POST with cancel support. - - Sends once with a long read timeout (120 s) so prefill finishes without - a retry storm (the old 0.5 s timeout caused duplicate POSTs every half - second). A watcher thread cancels by closing the response. httpx can't - interrupt a blocked read before the response exists, so cancel during - the header wait (1-5 s prefill) is deferred until headers arrive. - """ + """Open one streaming POST and let cancel interrupt prefill or reads.""" if cancel_event is not None and cancel_event.is_set(): raise GeneratorExit - # Background watcher: close the response if cancel is requested. - # Only effective after response headers arrive (httpx limitation). _cancel_closed = threading.Event() _response_ref: list = [None] def _cancel_watcher(): while not _cancel_closed.is_set(): if cancel_event.wait(timeout = 0.3): - # Cancel requested. Poll until the response object exists - # so we can close it, or until the main thread finishes - # (_cancel_closed set in finally). while not _cancel_closed.is_set(): r = _response_ref[0] - if r is not None: - try: + try: + if r is not None: r.close() - return - except Exception as e: - logger.debug(f"Error closing response in cancel watcher: {e}") - # Response not created yet -- wait briefly and retry + else: + LlamaCppBackend._shutdown_active_httpx_sockets(client) + return + except Exception as e: + logger.debug(f"Error closing request in cancel watcher: {e}") _cancel_closed.wait(timeout = 0.1) return @@ -5379,12 +5461,12 @@ class LlamaCppBackend: watcher.start() try: - # Long read timeout so prefill can finish without a retry storm. - # Cancel during prefill and streaming is handled by the watcher - # thread closing the response, unblocking any httpx read. + if first_token_deadline is None: + first_token_deadline = time.monotonic() + _DEFAULT_FIRST_TOKEN_TIMEOUT_S + prefill_read_timeout = max(0.1, first_token_deadline - time.monotonic()) prefill_timeout = httpx.Timeout( connect = 30, - read = 120.0, + read = prefill_read_timeout, write = 10, pool = 10, ) @@ -5400,7 +5482,7 @@ class LlamaCppBackend: raise GeneratorExit yield response return - except (httpx.ReadError, httpx.RemoteProtocolError, httpx.CloseError): + except (httpx.RequestError, RuntimeError): # Response was closed by the cancel watcher if cancel_event is not None and cancel_event.is_set(): raise GeneratorExit @@ -5455,14 +5537,12 @@ class LlamaCppBackend: ) if _reasoning_kw is not None: payload["chat_template_kwargs"] = _reasoning_kw - # Cap to the effective context length when known, else the floor. - # The wall-clock backstop below stops a stuck model regardless. + # Default cap to the model context when known. payload["max_tokens"] = ( max_tokens if max_tokens is not None else (self._effective_context_length or _DEFAULT_MAX_TOKENS_FLOOR) ) - payload["t_max_predict_ms"] = _DEFAULT_T_MAX_PREDICT_MS if stop: payload["stop"] = stop if seed is not None: @@ -5478,20 +5558,20 @@ class LlamaCppBackend: _metadata_finish_reason = None try: - # _stream_with_retry uses a 120 s read timeout so prefill can - # finish. Cancel during streaming is handled by the watcher - # thread (closes the response on cancel_event). + # Prefill can use the long first-token timeout; body reads are lowered after headers. stream_timeout = httpx.Timeout(connect = 10, read = 0.5, write = 10, pool = 10) _auth_headers = {"Authorization": f"Bearer {self._api_key}"} if self._api_key else None with httpx.Client( timeout = stream_timeout, limits = httpx.Limits(max_keepalive_connections = 0) ) as client: + first_token_deadline = time.monotonic() + _DEFAULT_FIRST_TOKEN_TIMEOUT_S with self._stream_with_retry( client, url, payload, cancel_event, headers = _auth_headers, + first_token_deadline = first_token_deadline, ) as response: if response.status_code != 200: error_body = response.read().decode() @@ -5502,7 +5582,11 @@ class LlamaCppBackend: buffer = "" has_content_tokens = False reasoning_text = "" - for raw_chunk in self._iter_text_cancellable(response, cancel_event): + for raw_chunk in self._iter_text_cancellable( + response, + cancel_event, + first_token_deadline = first_token_deadline, + ): buffer += raw_chunk while "\n" in buffer: line, buffer = buffer.split("\n", 1) @@ -5732,7 +5816,6 @@ class LlamaCppBackend: if max_tokens is not None else (self._effective_context_length or _DEFAULT_MAX_TOKENS_FLOOR) ) - payload["t_max_predict_ms"] = _DEFAULT_T_MAX_PREDICT_MS if stop: payload["stop"] = stop if seed is not None: @@ -5778,12 +5861,14 @@ class LlamaCppBackend: timeout = stream_timeout, limits = httpx.Limits(max_keepalive_connections = 0), ) as client: + first_token_deadline = time.monotonic() + _DEFAULT_FIRST_TOKEN_TIMEOUT_S with self._stream_with_retry( client, url, payload, cancel_event, headers = _auth_headers, + first_token_deadline = first_token_deadline, ) as response: if response.status_code != 200: error_body = response.read().decode() @@ -5795,6 +5880,7 @@ class LlamaCppBackend: for raw_chunk in self._iter_text_cancellable( response, cancel_event, + first_token_deadline = first_token_deadline, ): raw_buf += raw_chunk while "\n" in raw_buf: @@ -6444,7 +6530,6 @@ class LlamaCppBackend: if max_tokens is not None else (self._effective_context_length or _DEFAULT_MAX_TOKENS_FLOOR) ) - stream_payload["t_max_predict_ms"] = _DEFAULT_T_MAX_PREDICT_MS if stop: stream_payload["stop"] = stop if seed is not None: @@ -6467,12 +6552,14 @@ class LlamaCppBackend: with httpx.Client( timeout = stream_timeout, limits = httpx.Limits(max_keepalive_connections = 0) ) as client: + first_token_deadline = time.monotonic() + _DEFAULT_FIRST_TOKEN_TIMEOUT_S with self._stream_with_retry( client, url, stream_payload, cancel_event, headers = _auth_headers, + first_token_deadline = first_token_deadline, ) as response: if response.status_code != 200: error_body = response.read().decode() @@ -6481,7 +6568,11 @@ class LlamaCppBackend: ) buffer = "" - for raw_chunk in self._iter_text_cancellable(response, cancel_event): + for raw_chunk in self._iter_text_cancellable( + response, + cancel_event, + first_token_deadline = first_token_deadline, + ): buffer += raw_chunk while "\n" in buffer: line, buffer = buffer.split("\n", 1) diff --git a/studio/backend/main.py b/studio/backend/main.py index 2b34b11fef..064e261753 100644 --- a/studio/backend/main.py +++ b/studio/backend/main.py @@ -193,6 +193,11 @@ if _STUDIO_ROOT_RESOLVED != _LEGACY_STUDIO_ROOT: if not os.environ.get("UNSLOTH_LLAMA_CPP_PATH"): os.environ["UNSLOTH_LLAMA_CPP_PATH"] = str(_STUDIO_ROOT_RESOLVED / "llama.cpp") +# The studio bundles unsloth_zoo; declare unsloth present (as `import unsloth` +# does) so its lazy submodule imports (export, hardware, mlx) and the +# DiffusionGemma runner never trip the install guard on a clean install. +os.environ.setdefault("UNSLOTH_IS_PRESENT", "1") + import hashlib import mimetypes import re as _re diff --git a/studio/backend/routes/inference.py b/studio/backend/routes/inference.py index fa708d0e39..e9353cd803 100644 --- a/studio/backend/routes/inference.py +++ b/studio/backend/routes/inference.py @@ -121,6 +121,19 @@ def _template_raise_message(error_text: str, chat_template: Optional[str]) -> Op def _friendly_error(exc: Exception) -> str: """Extract a user-friendly message from known llama-server errors.""" + if isinstance(exc, httpx.ReadTimeout): + if "stopped producing tokens" in str(exc).lower(): + return ( + "The model stopped producing tokens before the response " + "completed. Try stopping and retrying, or reduce max tokens." + ) + return ( + "The model is still processing the prompt but did not produce a " + "first token within 20 minutes. Try reducing context length, " + "using more GPU offload, or loading a smaller model." + ) + if isinstance(exc, httpx.TimeoutException): + return "Timed out communicating with the model server. Try again shortly." # httpx transport failures from the async pass-through helpers. Any # RequestError subclass (ConnectError, ReadError, RemoteProtocolError, # WriteError, PoolTimeout, ...) means the llama-server subprocess is @@ -224,7 +237,11 @@ def _openai_stream_error_chunk(exc) -> dict: (a code-less error hides it).""" _cls = _classify_llama_generation_error(exc) if _cls: - return openai_error_body(_friendly_error(exc), status = 400, code = "context_length_exceeded") + return openai_error_body( + _friendly_error(exc), + status = 400, + code = "context_length_exceeded", + ) if _cls is False: return openai_error_body(_friendly_error(exc), status = 400) return openai_error_body(_friendly_error(exc), status = 500) @@ -415,17 +432,15 @@ def _apply_overflow_truncation(body: dict, err_text: str) -> bool: return True -def _anthropic_stream_error_event(exc): - """Anthropic in-band SSE ``error`` event for a mid-stream failure, or ``None`` - to fall through to a normal message_delta finish. Returns an event only for a - classifiable upstream client error (context overflow / 4xx) so a streaming - over-context request surfaces a real error instead of a silent empty - end_turn message.""" - if _classify_llama_generation_error(exc) is None: +def _anthropic_stream_error_event(exc, *, force: bool = False): + """Return an Anthropic in-band stream error event when one is useful.""" + _cls = _classify_llama_generation_error(exc) + if _cls is None and not force: return None + status = 400 if _cls is not None else 500 return build_anthropic_sse_event( "error", - anthropic_error_body(_friendly_error(exc), status = 400), + anthropic_error_body(_friendly_error(exc), status = status), ) @@ -588,8 +603,9 @@ try: from core.inference import get_inference_backend from core.inference.llama_cpp import ( LlamaCppBackend, + _DEFAULT_FIRST_TOKEN_TIMEOUT_S, _DEFAULT_MAX_TOKENS_FLOOR, - _DEFAULT_T_MAX_PREDICT_MS, + _DEFAULT_STREAM_STALL_TIMEOUT_S, _canonicalize_spec_mode, _extra_args_set_spec_type, _hf_offline_if_dns_dead, @@ -621,8 +637,9 @@ except ImportError: from core.inference import get_inference_backend from core.inference.llama_cpp import ( LlamaCppBackend, + _DEFAULT_FIRST_TOKEN_TIMEOUT_S, _DEFAULT_MAX_TOKENS_FLOOR, - _DEFAULT_T_MAX_PREDICT_MS, + _DEFAULT_STREAM_STALL_TIMEOUT_S, _canonicalize_spec_mode, _extra_args_set_spec_type, _hf_offline_if_dns_dead, @@ -648,6 +665,152 @@ except ImportError: verify_native_path_lease, ) + +def _llama_non_streaming_generation_timeout() -> httpx.Timeout: + return httpx.Timeout( + connect = _DEFAULT_FIRST_TOKEN_TIMEOUT_S, + read = _DEFAULT_FIRST_TOKEN_TIMEOUT_S, + write = _DEFAULT_FIRST_TOKEN_TIMEOUT_S, + pool = _DEFAULT_FIRST_TOKEN_TIMEOUT_S, + ) + + +def _llama_streaming_generation_timeout() -> httpx.Timeout: + return httpx.Timeout( + connect = _DEFAULT_FIRST_TOKEN_TIMEOUT_S, + read = _DEFAULT_FIRST_TOKEN_TIMEOUT_S, + write = _DEFAULT_FIRST_TOKEN_TIMEOUT_S, + pool = _DEFAULT_FIRST_TOKEN_TIMEOUT_S, + ) + + +def _set_stream_response_read_timeout( + response: httpx.Response, read_timeout_s: float = _DEFAULT_STREAM_STALL_TIMEOUT_S +) -> None: + try: + timeout_ext = response.request.extensions.get("timeout") + if isinstance(timeout_ext, dict): + timeout_ext["read"] = read_timeout_s + except Exception: + pass + + +async def _preheader_cancelled(cancel_event = None, request: Optional[Request] = None) -> bool: + if cancel_event is not None and cancel_event.is_set(): + return True + if request is not None and await request.is_disconnected(): + if cancel_event is not None: + cancel_event.set() + return True + return False + + +async def _wait_preheader_cancel(cancel_event = None, request: Optional[Request] = None) -> None: + while not await _preheader_cancelled(cancel_event, request): + await asyncio.sleep(0.05) + + +async def _send_stream_with_preheader_cancel( + client: httpx.AsyncClient, + req: httpx.Request, + cancel_event = None, + request: Optional[Request] = None, +) -> Optional[httpx.Response]: + if cancel_event is None and request is None: + return await client.send(req, stream = True) + if await _preheader_cancelled(cancel_event, request): + return None + + send_task = asyncio.create_task(client.send(req, stream = True)) + cancel_task = asyncio.create_task(_wait_preheader_cancel(cancel_event, request)) + + async def _stop_send_task() -> None: + try: + await client.aclose() + except Exception: + pass + send_task.cancel() + try: + await send_task + except (asyncio.CancelledError, Exception): + pass + + try: + done, _pending = await asyncio.wait( + {send_task, cancel_task}, + return_when = asyncio.FIRST_COMPLETED, + ) + if send_task in done: + return await send_task + + await _stop_send_task() + return None + except asyncio.CancelledError: + if cancel_event is not None: + cancel_event.set() + await _stop_send_task() + raise + finally: + cancel_task.cancel() + try: + await cancel_task + except (asyncio.CancelledError, Exception): + pass + + +async def _aiter_llama_stream_items( + async_iter, + *, + cancel_event = None, + request: Optional[Request] = None, + first_token_deadline: Optional[float] = None, + response: Optional[httpx.Response] = None, + post_first_item_read_timeout_s: Optional[float] = _DEFAULT_STREAM_STALL_TIMEOUT_S, +): + if first_token_deadline is None: + first_token_deadline = time.monotonic() + _DEFAULT_FIRST_TOKEN_TIMEOUT_S + last_item_at: Optional[float] = None + while True: + if cancel_event is not None and cancel_event.is_set(): + return + if request is not None and await request.is_disconnected(): + if cancel_event is not None: + cancel_event.set() + return + waiting_first_item = last_item_at is None + try: + if waiting_first_item: + remaining_s = first_token_deadline - time.monotonic() + if remaining_s <= 0: + raise httpx.ReadTimeout("The model did not produce a first token in time.") + if response is not None: + _set_stream_response_read_timeout(response, remaining_s) + item = await asyncio.wait_for(async_iter.__anext__(), timeout = remaining_s) + else: + item = await async_iter.__anext__() + except asyncio.TimeoutError as exc: + if waiting_first_item: + raise httpx.ReadTimeout("The model did not produce a first token in time.") from exc + raise + except StopAsyncIteration: + return + except httpx.ReadTimeout: + now = time.monotonic() + if last_item_at is None: + if now >= first_token_deadline: + raise + continue + raise httpx.ReadTimeout("The model stopped producing tokens mid-response.") + if ( + last_item_at is None + and response is not None + and post_first_item_read_timeout_s is not None + ): + _set_stream_response_read_timeout(response, post_first_item_read_timeout_s) + last_item_at = time.monotonic() + yield item + + from models.inference import ( LoadRequest, UnloadRequest, @@ -4935,6 +5098,8 @@ async def openai_completions(request: Request, current_subject: str = Depends(ge ) body = await request.json() + if body.get("max_tokens") is None: + body["max_tokens"] = llama_backend.context_length or _DEFAULT_MAX_TOKENS_FLOOR target_url = f"{llama_backend.base_url}/v1/completions" is_stream = body.get("stream", False) @@ -4952,15 +5117,27 @@ async def openai_completions(request: Request, current_subject: str = Depends(ge # honor stream_options.include_usage per event, while keeping SSE # framing and token bytes intact. _include_usage = bool((body.get("stream_options") or {}).get("include_usage")) - client = httpx.AsyncClient(timeout = 600) + client = httpx.AsyncClient(timeout = _llama_streaming_generation_timeout()) resp = None bytes_iter = None try: req = client.build_request("POST", target_url, json = body) - resp = await client.send(req, stream = True) + first_token_deadline = time.monotonic() + _DEFAULT_FIRST_TOKEN_TIMEOUT_S + resp = await _send_stream_with_preheader_cancel(client, req, request = request) + if resp is None: + return + if resp.status_code != 200: + err_bytes = await resp.aread() + err_text = err_bytes.decode("utf-8", errors = "replace") + raise RuntimeError(f"llama-server returned {resp.status_code}: {err_text}") bytes_iter = resp.aiter_bytes() buffer = b"" - async for chunk in bytes_iter: + async for chunk in _aiter_llama_stream_items( + bytes_iter, + request = request, + first_token_deadline = first_token_deadline, + response = resp, + ): buffer += chunk while b"\n\n" in buffer: event, buffer = buffer.split(b"\n\n", 1) @@ -4976,6 +5153,9 @@ async def openai_completions(request: Request, current_subject: str = Depends(ge yield out + b"\n\n" except Exception as e: logger.error("openai_completions stream error: %s", e) + error_chunk = _openai_stream_error_chunk(e) + yield f"data: {json.dumps(error_chunk)}\n\n".encode("utf-8") + return finally: if bytes_iter is not None: try: @@ -4995,7 +5175,11 @@ async def openai_completions(request: Request, current_subject: str = Depends(ge return StreamingResponse(_stream(), media_type = "text/event-stream") else: async with httpx.AsyncClient() as client: - resp = await client.post(target_url, json = body, timeout = 600) + resp = await client.post( + target_url, + json = body, + timeout = _llama_non_streaming_generation_timeout(), + ) if resp.status_code != 200: raise _openai_passthrough_error(resp.status_code, resp.text) @@ -5033,7 +5217,7 @@ async def openai_embeddings(request: Request, current_subject: str = Depends(get target_url = f"{llama_backend.base_url}/v1/embeddings" async with httpx.AsyncClient() as client: - resp = await client.post(target_url, json = body, timeout = 600) + resp = await client.post(target_url, json = body, timeout = _DEFAULT_FIRST_TOKEN_TIMEOUT_S) return Response( content = resp.content, status_code = resp.status_code, @@ -5775,6 +5959,28 @@ async def _responses_stream( ) return [item for _, item in sorted(indexed_items, key = lambda pair: pair[0])] + def _failed_response_payload(exc: Exception, status_code: int) -> dict: + return { + "type": "response.failed", + "response": { + "id": resp_id, + "object": "response", + "created_at": created_at, + "status": "failed", + "model": payload.model, + "output": _snapshot_output(), + "usage": { + "input_tokens": input_tokens, + "output_tokens": output_tokens, + "total_tokens": input_tokens + output_tokens, + }, + "error": { + "code": status_code, + "message": _friendly_error(exc), + }, + }, + } + # ── Preamble events ── yield _sse( "response.created", @@ -5798,13 +6004,16 @@ async def _responses_stream( # `async with`, explicit aclose of lines_iter BEFORE resp / client so # the innermost httpcore byte stream is finalised in this task (not via # the asyncgen GC in a sibling task). - client = httpx.AsyncClient(timeout = 600) + client = httpx.AsyncClient(timeout = _llama_streaming_generation_timeout()) resp = None lines_iter = None try: req = client.build_request("POST", target_url, json = body) + first_token_deadline = time.monotonic() + _DEFAULT_FIRST_TOKEN_TIMEOUT_S try: - resp = await client.send(req, stream = True) + resp = await _send_stream_with_preheader_cancel(client, req, request = request) + if resp is None: + return except httpx.RequestError as e: logger.error("responses stream: upstream unreachable: %s", e) yield _sse( @@ -5853,9 +6062,12 @@ async def _responses_stream( return lines_iter = resp.aiter_lines() - async for raw_line in lines_iter: - if await request.is_disconnected(): - break + async for raw_line in _aiter_llama_stream_items( + lines_iter, + request = request, + first_token_deadline = first_token_deadline, + response = resp, + ): if not raw_line: continue if not raw_line.startswith("data: "): @@ -5974,6 +6186,12 @@ async def _responses_stream( output_tokens = usage.get("completion_tokens", output_tokens) except Exception as e: logger.error("responses stream error: %s", e) + status_code = 400 if _classify_llama_generation_error(e) is not None else 500 + yield _sse( + "response.failed", + _failed_response_payload(e, status_code), + ) + return finally: if lines_iter is not None: try: @@ -7083,7 +7301,6 @@ def _build_passthrough_payload( body["max_tokens"] = ( max_tokens if max_tokens is not None else (backend_ctx or _DEFAULT_MAX_TOKENS_FLOOR) ) - body["t_max_predict_ms"] = _DEFAULT_T_MAX_PREDICT_MS # Normalize stop the same way the non-passthrough path does (the passthrough # was previously the one path that forwarded an empty stop string verbatim). _stop = _normalize_stop_sequences(stop) @@ -7193,7 +7410,7 @@ async def _anthropic_passthrough_stream( # `try: ... except Exception: pass` so nested anyio cleanup noise can't # bubble out. client = httpx.AsyncClient( - timeout = 600, + timeout = _llama_streaming_generation_timeout(), limits = httpx.Limits(max_keepalive_connections = 0), ) resp = None @@ -7201,7 +7418,12 @@ async def _anthropic_passthrough_stream( cancel_watcher = None try: req = client.build_request("POST", target_url, json = body) - resp = await client.send(req, stream = True) + first_token_deadline = time.monotonic() + _DEFAULT_FIRST_TOKEN_TIMEOUT_S + resp = await _send_stream_with_preheader_cancel( + client, req, cancel_event, request = request + ) + if resp is None: + return # Upstream client error (e.g. over-context 400) arrives before any # SSE. The 200 stream headers are already flushed, so surface it as @@ -7230,12 +7452,13 @@ async def _anthropic_passthrough_stream( # The watcher closes `resp` on cancel, raising in aiter_lines. cancel_watcher = asyncio.create_task(_await_cancel_then_close(cancel_event, resp)) lines_iter = resp.aiter_lines() - async for raw_line in lines_iter: - if cancel_event.is_set(): - break - if await request.is_disconnected(): - cancel_event.set() - break + async for raw_line in _aiter_llama_stream_items( + lines_iter, + cancel_event = cancel_event, + request = request, + first_token_deadline = first_token_deadline, + response = resp, + ): if not raw_line or not raw_line.startswith("data: "): continue data_str = raw_line[6:] @@ -7249,11 +7472,26 @@ async def _anthropic_passthrough_stream( _drop_parallel_tool_call_deltas(chunk) for line in emitter.feed_chunk(chunk): yield line - except (httpx.RemoteProtocolError, httpx.ReadError, httpx.CloseError): + except (httpx.RemoteProtocolError, httpx.ReadError, httpx.CloseError) as e: if not cancel_event.is_set(): - raise + logger.error("anthropic_messages passthrough stream error: %s", e) + event = _anthropic_stream_error_event( + e, + force = True, + ) + if event is not None: + yield event + return except Exception as e: - logger.error("anthropic_messages passthrough stream error: %s", e) + if not cancel_event.is_set(): + logger.error("anthropic_messages passthrough stream error: %s", e) + event = _anthropic_stream_error_event( + e, + force = True, + ) + if event is not None: + yield event + return finally: if cancel_watcher is not None: cancel_watcher.cancel() @@ -7327,7 +7565,11 @@ async def _anthropic_passthrough_non_streaming( ) async with httpx.AsyncClient() as client: - resp = await client.post(target_url, json = body, timeout = 600) + resp = await client.post( + target_url, + json = body, + timeout = _llama_non_streaming_generation_timeout(), + ) if resp.status_code != 200: raise HTTPException( @@ -7681,15 +7923,13 @@ async def _openai_passthrough_stream( _tracker = _TrackedCancel(cancel_event, *_cancel_keys) _tracker.__enter__() - # Outer guard: asyncio.CancelledError at `await client.send(...)` is a - # BaseException that bypasses `except httpx.RequestError`; without this the - # tracker leaks. The generator's finally only runs once iteration starts. + # Keep tracker cleanup paired if pre-header dispatch is cancelled. try: # Dispatch BEFORE returning StreamingResponse so transport errors and # non-200 upstream statuses surface as real HTTP errors -- OpenAI SDKs # rely on status codes to raise APIError/BadRequestError. client = httpx.AsyncClient( - timeout = 600, + timeout = _llama_streaming_generation_timeout(), limits = httpx.Limits(max_keepalive_connections = 0), ) resp = None @@ -7699,7 +7939,10 @@ async def _openai_passthrough_stream( while True: try: req = client.build_request("POST", target_url, json = body) - resp = await client.send(req, stream = True) + first_token_deadline = time.monotonic() + _DEFAULT_FIRST_TOKEN_TIMEOUT_S + resp = await _send_stream_with_preheader_cancel( + client, req, cancel_event, request = request + ) except httpx.RequestError as e: # llama-server subprocess crashed / starting / unreachable. logger.error("openai passthrough stream: upstream unreachable: %s", e) @@ -7716,6 +7959,21 @@ async def _openai_passthrough_stream( status_code = 502, detail = _friendly_error(e), ) + if resp is None: + try: + await client.aclose() + except Exception: + pass + _tracker.__exit__(None, None, None) + return StreamingResponse( + iter(()), + media_type = "text/event-stream", + headers = { + "Cache-Control": "no-cache", + "Connection": "keep-alive", + "X-Accel-Buffering": "no", + }, + ) if resp.status_code == 200: break @@ -7759,12 +8017,13 @@ async def _openai_passthrough_stream( cancel_watcher = asyncio.create_task(_await_cancel_then_close(cancel_event, resp)) try: lines_iter = resp.aiter_lines() - async for raw_line in lines_iter: - if cancel_event.is_set(): - break - if await request.is_disconnected(): - cancel_event.set() - break + async for raw_line in _aiter_llama_stream_items( + lines_iter, + cancel_event = cancel_event, + request = request, + first_token_deadline = first_token_deadline, + response = resp, + ): if not raw_line: continue if not raw_line.startswith("data: "): @@ -7843,7 +8102,11 @@ async def _openai_passthrough_non_streaming(llama_backend, payload, model_name): while True: try: async with httpx.AsyncClient() as client: - resp = await client.post(target_url, json = body, timeout = 600) + resp = await client.post( + target_url, + json = body, + timeout = _llama_non_streaming_generation_timeout(), + ) except httpx.RequestError as e: # llama-server subprocess crashed / starting / unreachable. Surface the # same friendly message the sync chat path emits so operators don't see diff --git a/studio/backend/run.py b/studio/backend/run.py index a883154c4a..84991a71bb 100644 --- a/studio/backend/run.py +++ b/studio/backend/run.py @@ -545,6 +545,11 @@ if _STUDIO_ROOT_RESOLVED != _LEGACY_STUDIO_ROOT: if not os.environ.get("UNSLOTH_LLAMA_CPP_PATH"): os.environ["UNSLOTH_LLAMA_CPP_PATH"] = str(_STUDIO_ROOT_RESOLVED / "llama.cpp") +# The studio bundles unsloth_zoo; declare unsloth present (as `import unsloth` +# does) so its lazy submodule imports (export, hardware, mlx) and the +# DiffusionGemma runner never trip the install guard on a clean install. +os.environ.setdefault("UNSLOTH_IS_PRESENT", "1") + def _write_pid_file(): """Write the current process PID to the studio PID file.""" @@ -1027,9 +1032,13 @@ def run_server( _cloudflare_enabled = cloudflare and host == "0.0.0.0" and not api_only and not _IS_COLAB if _cloudflare_enabled: try: # best-effort: any failure must not block startup - from cloudflare_tunnel import start_studio_tunnel + from cloudflare_tunnel import start_studio_tunnel, stop_studio_tunnel + _cloudflare_url = start_studio_tunnel(port) app.state.cloudflare_url = _cloudflare_url + # Backstop: tear the tunnel down even on an abnormal exit that bypasses + # _graceful_shutdown (e.g. an exception after startup -> sys.exit). Idempotent. + atexit.register(stop_studio_tunnel) except Exception as e: logger.debug("Cloudflare tunnel skipped: %s", e) diff --git a/studio/backend/tests/test_cloudflare_tunnel.py b/studio/backend/tests/test_cloudflare_tunnel.py index 873547631d..8042240b64 100644 --- a/studio/backend/tests/test_cloudflare_tunnel.py +++ b/studio/backend/tests/test_cloudflare_tunnel.py @@ -52,6 +52,26 @@ def test_url_regex_no_match_on_unrelated(): assert ct._URL_RE.search("INF connecting to https://api.cloudflare.com/v4") is None +def test_url_regex_ignores_api_endpoint(): + # cloudflared's failure line names its own API host; it must never be taken + # as the tunnel URL (it returns a 404 and is not a quick tunnel). + line = ( + 'failed to request quick Tunnel: Post "https://api.trycloudflare.com/tunnel": ' + "context deadline exceeded" + ) + assert ct._URL_RE.search(line) is None + + +def test_url_regex_skips_api_host_but_matches_real_url(): + blob = ( + 'ERR failed to request quick Tunnel: Post "https://api.trycloudflare.com/tunnel"\n' + "INF | https://brave-mountain-river-clouds.trycloudflare.com |\n" + ) + m = ct._URL_RE.search(blob) + assert m is not None + assert m.group(0) == "https://brave-mountain-river-clouds.trycloudflare.com" + + # ── asset mapping ──────────────────────────────────────────────────── @@ -295,9 +315,88 @@ def test_stop_terminates_process(): t.stop() -def test_wait_for_url_times_out_without_blocking(): +def test_start_after_stop_does_not_spawn(monkeypatch): + # If stop() lands before start() (a concurrent shutdown in the caller's + # register->start window), start() must NOT spawn a cloudflared process -- + # nobody would own it and it would be orphaned. t = ct.CloudflareTunnel(8080, "/bin/cloudflared") - assert t.wait_for_url(timeout = 0.05) is None + spawned = [] + + class _FakeProc: + stdout = None + + def poll(self): + return 0 + + monkeypatch.setattr(ct.subprocess, "Popen", lambda *a, **k: (spawned.append(a), _FakeProc())[1]) + t.stop() # proc is None -> no-op terminate, but marks the tunnel stopped + t.start() # must short-circuit before Popen + assert spawned == [] + assert t._proc is None + + +def test_wait_for_ready_times_out_without_blocking(): + t = ct.CloudflareTunnel(8080, "/bin/cloudflared") + assert t.wait_for_ready(timeout = 0.05) is None + + +def _fake_proc(text): + return types.SimpleNamespace(stdout = io.StringIO(text)) + + +def test_reader_captures_url_and_registration(): + t = ct.CloudflareTunnel(8080, "/bin/cloudflared") + t._reader( + _fake_proc( + "INF Requesting new quick Tunnel on trycloudflare.com...\n" + "INF | https://words-here-abc.trycloudflare.com |\n" + "INF Registered tunnel connection connIndex=0 protocol=http2\n" + ) + ) + assert t.url == "https://words-here-abc.trycloudflare.com" + assert t.ready is True + assert t.wait_for_ready(0) == t.url + assert t.error is None # a fully-registered tunnel records no error + + +def test_reader_url_without_registration_is_not_ready(): + # A URL but no "Registered tunnel connection" (e.g. quic control stream + # fails) must not be advertised -- it returns Cloudflare error 1033. + t = ct.CloudflareTunnel(8080, "/bin/cloudflared") + t._reader( + _fake_proc( + "INF | https://words-here-abc.trycloudflare.com |\n" + 'ERR failed to serve tunnel connection error="control stream failure"\n' + ) + ) + assert t.url == "https://words-here-abc.trycloudflare.com" + assert t.ready is False + assert t.wait_for_ready(0) is None + assert t.error == "cloudflared exited before the tunnel connection registered" + + +def test_reader_handles_none_stdout(): + # Popen.stdout can be None; _reader must not crash and must leave the tunnel + # un-ready so wait_for_ready returns None. + t = ct.CloudflareTunnel(8080, "/bin/cloudflared") + t._reader(types.SimpleNamespace(stdout = None)) + assert t.url is None + assert t.ready is False + assert t.wait_for_ready(0) is None + assert t.error == "cloudflared exited before emitting a tunnel URL" + + +def test_reader_ignores_api_endpoint_failure_line(): + t = ct.CloudflareTunnel(8080, "/bin/cloudflared") + t._reader( + _fake_proc( + "ERR failed to request quick Tunnel: Post " + '"https://api.trycloudflare.com/tunnel": context deadline exceeded\n' + ) + ) + assert t.url is None + assert t.wait_for_ready(0) is None + assert t.error == "cloudflared exited before emitting a tunnel URL" def test_start_studio_tunnel_no_binary(monkeypatch): @@ -306,18 +405,23 @@ def test_start_studio_tunnel_no_binary(monkeypatch): def test_start_studio_tunnel_registers_before_wait(monkeypatch): - # The tunnel must be visible to stop_studio_tunnel() during the URL wait, - # else a shutdown in that window orphans cloudflared. + # The tunnel must be visible to stop_studio_tunnel() during the readiness + # wait, else a shutdown in that window orphans cloudflared. seen = {} class _Stub: - def __init__(self, port, binary): + def __init__( + self, + port, + binary, + protocol = None, + ): self.url = None def start(self): pass - def wait_for_url(self, timeout): + def wait_for_ready(self, timeout): seen["active_during_wait"] = ct._active_tunnel is self self.url = "https://x.trycloudflare.com" return self.url @@ -338,13 +442,18 @@ def test_start_studio_tunnel_clears_and_stops_on_no_url(monkeypatch): seen = {} class _Stub: - def __init__(self, port, binary): + def __init__( + self, + port, + binary, + protocol = None, + ): self.url = None def start(self): pass - def wait_for_url(self, timeout): + def wait_for_ready(self, timeout): return None def stop(self): @@ -359,13 +468,18 @@ def test_start_studio_tunnel_clears_and_stops_on_no_url(monkeypatch): def test_start_studio_tunnel_returns_url(monkeypatch): class _StubTunnel: - def __init__(self, port, binary): + def __init__( + self, + port, + binary, + protocol = None, + ): self.url = None def start(self): self.url = "https://stub-xyz.trycloudflare.com" - def wait_for_url(self, timeout): + def wait_for_ready(self, timeout): return self.url def stop(self): @@ -379,6 +493,169 @@ def test_start_studio_tunnel_returns_url(monkeypatch): ct.stop_studio_tunnel() +def test_start_studio_tunnel_falls_back_to_http2(monkeypatch): + # First attempt mints a URL but never registers (quic blocked); the http2 + # retry registers and wins. + attempts = [] + + class _Stub: + def __init__( + self, + port, + binary, + protocol = None, + ): + self.protocol = protocol + self.url = None + attempts.append(protocol) + + def start(self): + self.url = "https://words.trycloudflare.com" # URL always minted + + def wait_for_ready(self, timeout): + return self.url if self.protocol == "http2" else None + + def stop(self): + pass + + monkeypatch.setattr(ct, "ensure_cloudflared", lambda: "/bin/cloudflared") + monkeypatch.setattr(ct, "CloudflareTunnel", _Stub) + try: + assert ct.start_studio_tunnel(8080) == "https://words.trycloudflare.com" + assert attempts == [None, "http2"] # default first, then forced http2 + finally: + ct.stop_studio_tunnel() + + +def test_start_studio_tunnel_no_retry_when_shutdown_between_attempts(monkeypatch): + # A stop() landing in the gap AFTER the failed first attempt is cleaned up but + # BEFORE the http2 retry registers must abort the loop -- not start a second + # tunnel that nobody will ever stop (Codex review). Simulated by having the + # first attempt's stop() (called during cleanup) trigger the shutdown. + attempts = [] + + class _Stub: + def __init__( + self, + port, + binary, + protocol = None, + ): + self.url = None + attempts.append(protocol) + + def start(self): + self.url = "https://words.trycloudflare.com" # URL minted, never ready + + def wait_for_ready(self, timeout): + return None + + def stop(self): + ct.stop_studio_tunnel() # a concurrent shutdown lands in the gap + + monkeypatch.setattr(ct, "ensure_cloudflared", lambda: "/bin/cloudflared") + monkeypatch.setattr(ct, "CloudflareTunnel", _Stub) + assert ct.start_studio_tunnel(8080) is None + assert attempts == [None] # http2 retry aborted after shutdown + assert ct._active_tunnel is None + + +def test_start_studio_tunnel_no_http2_retry_when_no_url(monkeypatch): + # No URL at all is an API/network failure; the http2 fallback would not help, + # so it must be skipped (don't burn a second timeout window). + attempts = [] + + class _Stub: + def __init__( + self, + port, + binary, + protocol = None, + ): + self.url = None + attempts.append(protocol) + + def start(self): + pass # never mints a URL + + def wait_for_ready(self, timeout): + return None + + def stop(self): + pass + + monkeypatch.setattr(ct, "ensure_cloudflared", lambda: "/bin/cloudflared") + monkeypatch.setattr(ct, "CloudflareTunnel", _Stub) + assert ct.start_studio_tunnel(8080) is None + assert attempts == [None] + + +def test_start_studio_tunnel_both_protocols_fail_registration(monkeypatch): + # Both quic and http2 mint a URL but neither registers -> both attempts are + # exhausted and None is returned (no dead URL advertised). + attempts = [] + + class _Stub: + def __init__( + self, + port, + binary, + protocol = None, + ): + self.url = None + attempts.append(protocol) + + def start(self): + self.url = "https://words.trycloudflare.com" # URL minted, never ready + + def wait_for_ready(self, timeout): + return None + + def stop(self): + pass + + monkeypatch.setattr(ct, "ensure_cloudflared", lambda: "/bin/cloudflared") + monkeypatch.setattr(ct, "CloudflareTunnel", _Stub) + assert ct.start_studio_tunnel(8080) is None + assert attempts == [None, "http2"] + assert ct._active_tunnel is None + + +def test_start_studio_tunnel_aborts_retry_on_concurrent_shutdown(monkeypatch): + # If a concurrent stop_studio_tunnel() clears _active_tunnel while we wait, + # the retry loop must NOT start a second (http2) tunnel: shutdown is already + # done, so nothing would ever stop it and it would be orphaned. + attempts = [] + + class _Stub: + def __init__( + self, + port, + binary, + protocol = None, + ): + self.url = None + attempts.append(protocol) + + def start(self): + self.url = "https://words.trycloudflare.com" # URL minted (saw_url True) + + def wait_for_ready(self, timeout): + # Simulate stop_studio_tunnel() landing during the wait. + with ct._active_lock: + ct._active_tunnel = None + return None # never registered + + def stop(self): + pass + + monkeypatch.setattr(ct, "ensure_cloudflared", lambda: "/bin/cloudflared") + monkeypatch.setattr(ct, "CloudflareTunnel", _Stub) + assert ct.start_studio_tunnel(8080) is None + assert attempts == [None] # no http2 retry -> no orphaned second tunnel + assert ct._active_tunnel is None + + # ── run.py source-level pins (AST / source, no heavy import) ───────── @@ -419,6 +696,13 @@ def test_argparse_cloudflare_default_true(): assert _argparse_default(_RUN_PY.read_text(), "--cloudflare") is True +def test_run_server_registers_tunnel_atexit_backstop(): + # An abnormal exit (exception after startup -> sys.exit) bypasses + # _graceful_shutdown; an atexit backstop must still stop the tunnel. + src = _RUN_PY.read_text() + assert "atexit.register(stop_studio_tunnel)" in src + + def test_run_server_gates_tunnel_on_wildcard(): # Guard against accidentally widening the trigger beyond 0.0.0.0. source = _RUN_PY.read_text() diff --git a/studio/backend/tests/test_gguf_route_cursor_reset.py b/studio/backend/tests/test_gguf_route_cursor_reset.py index fdbf31f9de..dc397b2d08 100644 --- a/studio/backend/tests/test_gguf_route_cursor_reset.py +++ b/studio/backend/tests/test_gguf_route_cursor_reset.py @@ -59,11 +59,16 @@ def _make_backend(monkeypatch, streams: list[list[str]], payloads: list[dict]): payload, _cancel_event, headers = None, + first_token_deadline = None, ): payloads.append(copy.deepcopy(payload)) yield type("FakeResponse", (), {"status_code": 200, "chunks": streams.pop(0)})() - def fake_iter_text_cancellable(response, _cancel_event): + def fake_iter_text_cancellable( + response, + _cancel_event, + first_token_deadline = None, + ): yield from response.chunks monkeypatch.setattr(backend, "_stream_with_retry", fake_stream_with_retry) diff --git a/studio/backend/tests/test_lemonade_llamacpp_rocm_bins_mock.py b/studio/backend/tests/test_lemonade_llamacpp_rocm_bins_mock.py deleted file mode 100644 index 7c6c514b8f..0000000000 --- a/studio/backend/tests/test_lemonade_llamacpp_rocm_bins_mock.py +++ /dev/null @@ -1,523 +0,0 @@ -# SPDX-License-Identifier: AGPL-3.0-only -# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0 - -"""Validates that the installer correctly resolves lemonade ROCm prebuilt assets. - -Uses a faked HostInfo so no AMD GPU is needed. Network calls to the lemonade -GitHub API are stubbed out so the suite runs without internet access and is -not subject to rate limits. -""" - -from __future__ import annotations - -import importlib -import sys -from pathlib import Path -from unittest.mock import patch - -import pytest - -_studio = Path(__file__).resolve().parent.parent.parent -if str(_studio) not in sys.path: - sys.path.insert(0, str(_studio)) - -_mod = importlib.import_module("install_llama_prebuilt") -HostInfo = _mod.HostInfo -resolve_lemonade_rocm_choice = getattr(_mod, "resolve_lemonade_rocm_choice", None) -_LEMONADE_GFX_FAMILIES = getattr(_mod, "_LEMONADE_GFX_FAMILIES", None) - -if resolve_lemonade_rocm_choice is None or _LEMONADE_GFX_FAMILIES is None: - pytest.skip("PR symbols not present - check branch", allow_module_level = True) - - -@pytest.fixture(autouse = True) -def _clear_lemonade_release_cache(): - """Prevent cross-test pollution of the lemonade release lru_cache and - selection-log dedup set when tests vary the fetch_json mock return value.""" - _cache = getattr(_mod, "_fetch_lemonade_release_cached", None) - _logged: set | None = getattr(_mod, "_lemonade_selection_logged", None) - if _cache is not None and hasattr(_cache, "cache_clear"): - _cache.cache_clear() - if _logged is not None: - _logged.clear() - yield - if _cache is not None and hasattr(_cache, "cache_clear"): - _cache.cache_clear() - if _logged is not None: - _logged.clear() - - -_STUB_TAG = "b1262" -_STUB_OS_PREFIXES = ("ubuntu", "windows") -_STUB_FAMILIES = ("gfx1151", "gfx1150", "gfx120X", "gfx110X", "gfx103X") - - -def _stub_lemonade_release() -> dict: - """Minimal lemonade release payload covering all supported GPU/OS combinations.""" - assets = [ - { - "name": f"llama-{_STUB_TAG}-{prefix}-rocm-{family}-x64.zip", - "browser_download_url": ( - f"https://github.com/lemonade-sdk/llamacpp-rocm/releases/download/" - f"{_STUB_TAG}/llama-{_STUB_TAG}-{prefix}-rocm-{family}-x64.zip" - ), - } - for prefix in _STUB_OS_PREFIXES - for family in _STUB_FAMILIES - ] - return {"tag_name": _STUB_TAG, "assets": assets} - - -def _make_rocm_host(gfx_target: str, *, windows: bool = False) -> HostInfo: - return HostInfo( - system = "Windows" if windows else "Linux", - machine = "amd64" if windows else "x86_64", - is_windows = windows, - is_linux = not windows, - 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, - has_rocm = True, - rocm_gfx_target = gfx_target, - ) - - -def _lookup_family(gfx: str) -> str | None: - for prefix, family in _LEMONADE_GFX_FAMILIES: - if gfx.startswith(prefix): - return family - return None - - -# --------------------------------------------------------------------------- -# GPU family mapping -# --------------------------------------------------------------------------- - - -@pytest.mark.parametrize( - "gfx,expected_family", - [ - ("gfx1151", "gfx1151"), - ("gfx1150", "gfx1150"), - ("gfx1201", "gfx120X"), - ("gfx1200", "gfx120X"), - ("gfx1100", "gfx110X"), - ("gfx1030", "gfx103X"), - ], -) -def test_gpu_family_mapping(gfx, expected_family): - assert _lookup_family(gfx) == expected_family - - -def test_unknown_gpu_not_in_families(): - assert _lookup_family("gfx999") is None - - -# --------------------------------------------------------------------------- -# Asset resolution - hits real lemonade GitHub API -# --------------------------------------------------------------------------- - - -@pytest.mark.parametrize( - "gfx,os_prefix,windows", - [ - ("gfx1151", "ubuntu", False), - ("gfx1150", "ubuntu", False), - ("gfx1201", "ubuntu", False), - ("gfx1100", "ubuntu", False), - ("gfx1030", "ubuntu", False), - ("gfx1151", "windows", True), - ("gfx1100", "windows", True), - ], -) -def test_asset_resolves_for_known_gpu(gfx, os_prefix, windows): - host = _make_rocm_host(gfx, windows = windows) - with patch.object(_mod, "fetch_json", return_value = _stub_lemonade_release()): - result = resolve_lemonade_rocm_choice(host, os_prefix, "default", llama_tag = "latest") - assert result is not None, f"Installer will NOT fetch lemonade binary for {gfx} ({os_prefix})" - assert _lookup_family(gfx) in result.name - assert result.url.startswith("https://github.com/lemonade-sdk/llamacpp-rocm") - - -def test_unknown_gpu_falls_through_to_upstream(): - host = _make_rocm_host("gfx999") - result = resolve_lemonade_rocm_choice(host, "ubuntu", "default", llama_tag = "latest") - assert result is None - - -# --------------------------------------------------------------------------- -# The Linux attempt builder must plan a lemonade ROCm attempt for AMD-only hosts. -# This is the path setup.sh actually invokes (fork hosts now select from the -# manifest), so the lemonade integration is useless if it isn't wired in here. -# --------------------------------------------------------------------------- - -_linux_published_attempts = getattr(_mod, "_linux_published_attempts", None) -direct_upstream_release_plan = getattr(_mod, "direct_upstream_release_plan", None) - -PublishedLlamaArtifact = _mod.PublishedLlamaArtifact -PublishedReleaseBundle = _mod.PublishedReleaseBundle - - -def _rocm_bundle(gfx_family: str, mapped_targets: list[str]) -> "PublishedReleaseBundle": - """A fork manifest bundle exposing a per-gfx linux-rocm artifact, so - published_rocm_choice_for_host can match the host before the lemonade - fallback is appended.""" - asset_name = f"app-b9457-linux-x64-rocm-{gfx_family}.tar.gz" - artifact = PublishedLlamaArtifact( - asset_name = asset_name, - install_kind = "linux-rocm", - runtime_line = None, - coverage_class = None, - supported_sms = [], - min_sm = None, - max_sm = None, - bundle_profile = None, - rank = 1000, - gfx_target = gfx_family, - mapped_targets = mapped_targets, - ) - return PublishedReleaseBundle( - repo = "unslothai/llama.cpp", - release_tag = "v1.0", - upstream_tag = "b9457", - assets = {asset_name: f"https://example.invalid/{asset_name}"}, - artifacts = [artifact], - ) - - -@pytest.mark.skipif( - _linux_published_attempts is None, - reason = "Linux attempt builder not present on this branch", -) -def test_linux_attempts_include_fork_rocm_and_lemonade_for_rocm_host(): - host = _make_rocm_host("gfx1151") - bundle = _rocm_bundle("gfx1151", ["gfx1151"]) - with patch.object(_mod, "fetch_json", return_value = _stub_lemonade_release()): - attempts = _linux_published_attempts(host, bundle, "latest") - kinds = [a.install_kind for a in attempts] - assert "linux-rocm" in kinds, f"builder did not include any linux-rocm attempt; got {kinds}" - sources = {a.source_label for a in attempts if a.install_kind == "linux-rocm"} - # The fork's own per-gfx bundle is preferred, with the lemonade prebuilt as - # the fallback -- both must be present for a covered ROCm host. - assert "published" in sources, f"fork ROCm bundle missing; got {sources}" - assert "lemonade" in sources, f"lemonade ROCm fallback missing; got {sources}" - lemonade_attempt = next(a for a in attempts if a.source_label == "lemonade") - assert "gfx1151" in lemonade_attempt.name - - -@pytest.mark.skipif( - direct_upstream_release_plan is None, - reason = "direct release planners not present on this branch", -) -def test_direct_upstream_plan_includes_lemonade_for_windows_hip_host(): - host = _make_rocm_host("gfx1151", windows = True) - release = { - "tag_name": "b9022", - "name": "b9022", - "assets": [], - } - with patch.object(_mod, "fetch_json", return_value = _stub_lemonade_release()): - plan = direct_upstream_release_plan(release, host, "ggml-org/llama.cpp", "latest") - assert plan is not None, "Windows ROCm host should plan a lemonade HIP attempt" - kinds = [a.install_kind for a in plan.attempts] - assert "windows-hip" in kinds, f"planner did not include a lemonade HIP attempt; got {kinds}" - - -@pytest.mark.skipif( - direct_upstream_release_plan is None, - reason = "direct release planners not present on this branch", -) -def test_windows_hip_falls_back_to_upstream_when_lemonade_unavailable(): - """If lemonade returns None (e.g. gfx999 or transient API failure), the planner - must still include the upstream HIP asset rather than silently downgrading to CPU.""" - host = _make_rocm_host("gfx999", windows = True) - hip_asset = "llama-b9022-bin-win-hip-radeon-x64.zip" - release = { - "tag_name": "b9022", - "name": "b9022", - "assets": [ - { - "name": hip_asset, - "browser_download_url": f"https://example.invalid/{hip_asset}", - }, - ], - } - plan = direct_upstream_release_plan(release, host, "ggml-org/llama.cpp", "latest") - assert plan is not None - kinds = [a.install_kind for a in plan.attempts] - assert "windows-hip" in kinds, f"upstream HIP asset not included as fallback; got {kinds}" - hip_attempt = next(a for a in plan.attempts if a.install_kind == "windows-hip") - assert hip_attempt.source_label == "upstream" - - -# ── Follow-up: pinned-tag URL helper, URL trust pinning, opt-out env, autouse cache clear ── - - -def test_lemonade_release_api_url_pinned_tag(): - """A pinned llama_tag must produce the /releases/tags/ URL.""" - assert _mod._lemonade_release_api_for("b1262").endswith("/releases/tags/b1262") - assert _mod._lemonade_release_api_for("latest").endswith("/releases/latest") - assert _mod._lemonade_release_api_for("").endswith("/releases/latest") - - -def test_lemonade_release_api_url_encodes_tag(): - """Unexpected slashes / hashes in the tag must be URL-encoded so the URL - cannot be reshaped (defence in depth -- tags should already be sanitised - upstream).""" - url = _mod._lemonade_release_api_for("b1260/../latest") - assert "/releases/tags/b1260%2F..%2Flatest" in url - assert "//latest" not in url.split("/releases/tags/", 1)[1] - - -def test_lemonade_resolver_skipped_by_opt_out_env(monkeypatch): - """UNSLOTH_DISABLE_LEMONADE_ROCM=1 must short-circuit the resolver.""" - monkeypatch.setenv("UNSLOTH_DISABLE_LEMONADE_ROCM", "1") - host = _make_rocm_host("gfx1151") - res = resolve_lemonade_rocm_choice(host, "ubuntu", "linux-rocm", llama_tag = "latest") - assert res is None - - -def test_lemonade_resolver_rejects_non_github_url(monkeypatch): - """If the GitHub API response somehow contained an off-host download URL, - the resolver must refuse to use it (lemonade assets are not in the - approved-hash manifest).""" - bad_release = { - "tag_name": _STUB_TAG, - "assets": [ - { - "name": f"llama-{_STUB_TAG}-ubuntu-rocm-gfx1151-x64.zip", - "browser_download_url": "https://attacker.invalid/llama.zip", - }, - ], - } - host = _make_rocm_host("gfx1151") - with patch.object(_mod, "fetch_json", return_value = bad_release): - res = resolve_lemonade_rocm_choice(host, "ubuntu", "linux-rocm", llama_tag = "latest") - assert res is None - - -def test_lemonade_resolver_rejects_http_scheme(): - assert not _mod._is_trusted_github_release_url( - "http://github.com/lemonade-sdk/llamacpp-rocm/releases/download/x/y.zip", - "lemonade-sdk/llamacpp-rocm", - ) - - -def test_lemonade_resolver_accepts_github_cdn(): - # Real GitHub release CDN URLs carry the /github-production-release-asset- prefix. - assert _mod._is_trusted_github_release_url( - "https://objects.githubusercontent.com/github-production-release-asset-abc123/456/789?token=x", - "lemonade-sdk/llamacpp-rocm", - ) - - -def test_lemonade_resolver_rejects_arbitrary_cdn_path(): - # A CDN URL without the release-asset path prefix must be rejected. - assert not _mod._is_trusted_github_release_url( - "https://objects.githubusercontent.com/abc/def", - "lemonade-sdk/llamacpp-rocm", - ) - - -def test_lemonade_resolver_accepts_release_path(): - url = "https://github.com/lemonade-sdk/llamacpp-rocm/releases/download/b1262/llama-b1262-ubuntu-rocm-gfx1151-x64.zip" - assert _mod._is_trusted_github_release_url(url, "lemonade-sdk/llamacpp-rocm") - - -def test_lemonade_resolver_rejects_wrong_repo(): - """A github.com release URL for a different repo must be rejected.""" - assert not _mod._is_trusted_github_release_url( - "https://github.com/attacker/llamacpp-rocm/releases/download/x/y.zip", - "lemonade-sdk/llamacpp-rocm", - ) - - -def test_lemonade_resolver_rejects_empty_browser_download_url(): - """An asset entry with an empty browser_download_url must fall through.""" - release = { - "tag_name": _STUB_TAG, - "assets": [ - { - "name": f"llama-{_STUB_TAG}-ubuntu-rocm-gfx1151-x64.zip", - "browser_download_url": "", - }, - ], - } - host = _make_rocm_host("gfx1151") - with patch.object(_mod, "fetch_json", return_value = release): - res = resolve_lemonade_rocm_choice(host, "ubuntu", "linux-rocm", llama_tag = "latest") - assert res is None - - -def test_lemonade_runtime_patterns_include_hip_runtime(): - """linux-rocm overlay must use a broad lib glob to catch all bundled .so files. - - Lemonade ZIPs carry transitive deps (libamd_comgr, libLLVM, libclang-cpp, - ...) whose names change across ROCm releases. A broad ``lib*.so*`` glob - avoids having to enumerate every transitive dependency by name. - """ - from install_llama_prebuilt import runtime_patterns_for_choice, AssetChoice - - choice = AssetChoice( - repo = "lemonade-sdk/llamacpp-rocm", - tag = "b1262", - name = "llama-b1262-ubuntu-rocm-gfx1151-x64.zip", - url = "https://github.com/lemonade-sdk/llamacpp-rocm/releases/download/b1262/x.zip", - source_label = "lemonade", - install_kind = "linux-rocm", - ) - pats = runtime_patterns_for_choice(choice) - # The broad glob must be present so every .so in the lemonade bundle - # (including transitive deps added in future ROCm releases) gets overlaid. - assert "lib*.so*" in pats, f"'lib*.so*' missing from linux-rocm patterns: {pats}" - - -_pick_rocm_gfx_target = getattr(_mod, "_pick_rocm_gfx_target", None) - - -@pytest.mark.skipif( - _pick_rocm_gfx_target is None, - reason = "_pick_rocm_gfx_target not present on this branch", -) -def test_pick_rocm_gfx_target_honors_cuda_visible_devices(monkeypatch): - """AMD HIP honours CUDA_VISIBLE_DEVICES identically to HIP_VISIBLE_DEVICES; - on a gfx1151 + gfx1100 mixed host, CUDA_VISIBLE_DEVICES=1 must select gfx1100.""" - # Two GPUs; rocminfo reports each token twice (as in the real tool output). - probe_out = "gfx1151\ngfx1151\ngfx1100\ngfx1100" - monkeypatch.delenv("HIP_VISIBLE_DEVICES", raising = False) - monkeypatch.delenv("ROCR_VISIBLE_DEVICES", raising = False) - monkeypatch.setenv("CUDA_VISIBLE_DEVICES", "1") - assert _pick_rocm_gfx_target(probe_out) == "gfx1100" - - -@pytest.mark.skipif( - _pick_rocm_gfx_target is None, - reason = "_pick_rocm_gfx_target not present on this branch", -) -def test_pick_rocm_gfx_target_cuda_visible_devices_minus_one_returns_none(monkeypatch): - """CUDA_VISIBLE_DEVICES=-1 means no GPU visible; resolver must return None.""" - probe_out = "gfx1151\ngfx1100" - monkeypatch.delenv("HIP_VISIBLE_DEVICES", raising = False) - monkeypatch.delenv("ROCR_VISIBLE_DEVICES", raising = False) - monkeypatch.setenv("CUDA_VISIBLE_DEVICES", "-1") - assert _pick_rocm_gfx_target(probe_out) is None - - -@pytest.mark.skipif( - _pick_rocm_gfx_target is None, - reason = "_pick_rocm_gfx_target not present on this branch", -) -def test_pick_rocm_gfx_target_same_arch_multi_gpu(monkeypatch): - """Regression: [gfx1100, gfx1100, gfx1151] with HIP_VISIBLE_DEVICES=2 must - return gfx1151, not fall back to GPU 0 due to dict.fromkeys collapsing the - two gfx1100 entries into one and making index 2 out of range.""" - # Simulate rocminfo output for 3 GPUs (2x gfx1100 dGPU + 1x gfx1151 APU). - # Each GPU gets its own Agent section with a few token mentions. - probe_out = ( - "***\nAgent 1\n***\n gfx1100 some info\n gfx1100\n" - "***\nAgent 2\n***\n gfx1100 some info\n gfx1100\n" - "***\nAgent 3\n***\n gfx1151 some info\n gfx1151\n" - ) - monkeypatch.delenv("ROCR_VISIBLE_DEVICES", raising = False) - monkeypatch.delenv("CUDA_VISIBLE_DEVICES", raising = False) - monkeypatch.setenv("HIP_VISIBLE_DEVICES", "2") - assert _pick_rocm_gfx_target(probe_out) == "gfx1151" - - -# --------------------------------------------------------------------------- -# Fork release scan: Windows ROCm resolves lemonade by the requested tag -# --------------------------------------------------------------------------- - -_resolve_release_asset_choice = getattr(_mod, "resolve_release_asset_choice", None) -_ApprovedReleaseChecksums = getattr(_mod, "ApprovedReleaseChecksums", None) - - -@pytest.mark.skipif( - _resolve_release_asset_choice is None or _ApprovedReleaseChecksums is None, - reason = "fork release planner not present on this branch", -) -def test_fork_scan_windows_rocm_resolves_lemonade_by_requested_tag(): - """The fork release scan pins llama_tag to per-release upstream tags - (b9457, ...) that lemonade's own tag series never contains, so the - lemonade lookup must use the requested tag ("latest") instead. Pinning - lemonade to the per-release tag 404s on every scanned release and a - Windows ROCm host ends in a rate-limited fatal instead of the lemonade - prebuilt.""" - host = _make_rocm_host("gfx1151", windows = True) - # No windows-rocm artifact in the bundle, matching current fork releases. - bundle = _rocm_bundle("gfx1151", ["gfx1151"]) - checksums = _ApprovedReleaseChecksums( - repo = "unslothai/llama.cpp", - release_tag = "v1.0", - upstream_tag = "b9457", - artifacts = {}, - ) - seen_urls: list[str] = [] - - def _fake_fetch(api_url, *args, **kwargs): - seen_urls.append(api_url) - if "lemonade-sdk" in api_url: - if api_url.endswith("/releases/latest"): - return _stub_lemonade_release() - raise RuntimeError(f"unexpected pinned lemonade fetch: {api_url}") - # ggml-org asset listing for the upstream HIP/CPU filename fallbacks. - return {"tag_name": "b9457", "assets": []} - - with patch.object(_mod, "fetch_json", side_effect = _fake_fetch): - attempts = _resolve_release_asset_choice( - host, - "b9457", # concrete per-release upstream tag from the scan loop - bundle, - checksums, - requested_tag = "latest", - ) - - lemonade = [a for a in attempts if a.source_label == "lemonade"] - assert lemonade, f"lemonade attempt missing for Windows ROCm host; got {attempts}" - assert "gfx1151" in lemonade[0].name - assert any( - u.endswith("/releases/latest") for u in seen_urls - ), f"lemonade was never resolved via /releases/latest; fetches: {seen_urls}" - assert not any( - "lemonade-sdk" in u and "/releases/tags/" in u for u in seen_urls - ), f"lemonade lookup was pinned to the fork release tag: {seen_urls}" - - -@pytest.mark.skipif( - direct_upstream_release_plan is None, - reason = "direct release planners not present on this branch", -) -def test_direct_upstream_plan_includes_lemonade_for_linux_rocm_host(): - """A Linux ROCm host on the ggml-org direct path (e.g. a --published-repo - override) must plan lemonade before the CPU tarball, mirroring the Windows - branch. The lemonade planning previously lived in the removed - --simple-policy dispatcher, so without this leg such hosts silently - install the CPU build.""" - host = _make_rocm_host("gfx1151") - release = { - "tag_name": "b9022", - "name": "b9022", - "assets": [ - { - "name": "llama-b9022-bin-ubuntu-x64.tar.gz", - "browser_download_url": ( - "https://github.com/ggml-org/llama.cpp/releases/download/" - "b9022/llama-b9022-bin-ubuntu-x64.tar.gz" - ), - } - ], - } - with patch.object(_mod, "fetch_json", return_value = _stub_lemonade_release()): - plan = direct_upstream_release_plan(release, host, "ggml-org/llama.cpp", "latest") - assert plan is not None, "Linux ROCm host should produce a direct plan" - kinds = [a.install_kind for a in plan.attempts] - sources = [a.source_label for a in plan.attempts] - assert "linux-rocm" in kinds, f"lemonade ROCm attempt missing; got {kinds}" - assert sources[0] == "lemonade", f"lemonade must be the first attempt; got {sources}" - assert "gfx1151" in plan.attempts[0].name diff --git a/studio/backend/tests/test_llama_cpp_tool_loop.py b/studio/backend/tests/test_llama_cpp_tool_loop.py index 3c121c281d..aacb029ff6 100644 --- a/studio/backend/tests/test_llama_cpp_tool_loop.py +++ b/studio/backend/tests/test_llama_cpp_tool_loop.py @@ -52,11 +52,16 @@ def _make_backend(monkeypatch, streams: list[list[str]], payloads: list[dict]): payload, _cancel_event, headers = None, + first_token_deadline = None, ): payloads.append(copy.deepcopy(payload)) yield type("FakeResponse", (), {"status_code": 200, "chunks": streams.pop(0)})() - def fake_iter_text_cancellable(response, _cancel_event): + def fake_iter_text_cancellable( + response, + _cancel_event, + first_token_deadline = None, + ): yield from response.chunks monkeypatch.setattr(backend, "_stream_with_retry", fake_stream_with_retry) diff --git a/studio/backend/tests/test_llama_cpp_update.py b/studio/backend/tests/test_llama_cpp_update.py index 3b8f511a7c..e207306c91 100644 --- a/studio/backend/tests/test_llama_cpp_update.py +++ b/studio/backend/tests/test_llama_cpp_update.py @@ -420,8 +420,8 @@ def test_start_update_installer_failure_reports_error(monkeypatch, tmp_path): # --- installer-argument construction (mirrors the post-#5963 setup scripts) --- -def test_rocm_install_args_lemonade_gfx(): - # Lemonade HIP app bundle: gfx family lives in the asset name. +def test_rocm_install_args_gfx_family(): + # Per-gfx ROCm bundle: gfx family lives in the asset name. assert upd._rocm_install_args("app-b9585-linux-x64-rocm-gfx110X.tar.gz") == [ "--rocm-gfx", "gfx110x", diff --git a/studio/backend/tests/test_llama_route_timeouts.py b/studio/backend/tests/test_llama_route_timeouts.py new file mode 100644 index 0000000000..5aee6198ba --- /dev/null +++ b/studio/backend/tests/test_llama_route_timeouts.py @@ -0,0 +1,87 @@ +# SPDX-License-Identifier: AGPL-3.0-only +# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. + +import asyncio +import os +import sys +import time +from types import SimpleNamespace + +_backend = os.path.join(os.path.dirname(__file__), "..") +sys.path.insert(0, _backend) + +import routes.inference as inf_mod # noqa: E402 + + +def test_non_streaming_generation_timeout_has_read_deadline(): + timeout = inf_mod._llama_non_streaming_generation_timeout() + assert timeout.read == inf_mod._DEFAULT_FIRST_TOKEN_TIMEOUT_S + + +def test_stream_first_item_deadline_after_headers(): + async def _run(): + class _Never: + async def __anext__(self): + await asyncio.Future() + + started = time.monotonic() + try: + async for _ in inf_mod._aiter_llama_stream_items( + _Never(), + first_token_deadline = started + 0.02, + ): + pass + except inf_mod.httpx.ReadTimeout: + pass + else: + raise AssertionError("first item deadline did not fire") + assert time.monotonic() - started < 0.5 + + asyncio.run(_run()) + + +def test_preheader_send_cleanup_on_disconnect_and_cancel(): + async def _run(cancel_parent): + state = SimpleNamespace(disconnected = False, closed = False, cancelled = False) + started = asyncio.Event() + + class _Client: + async def send( + self, + req, + stream = False, + ): + started.set() + try: + await asyncio.Future() + except asyncio.CancelledError: + state.cancelled = True + raise + + async def aclose(self): + state.closed = True + + class _Request: + async def is_disconnected(self): + return state.disconnected + + task = asyncio.create_task( + inf_mod._send_stream_with_preheader_cancel(_Client(), object(), request = _Request()) + ) + await started.wait() + if cancel_parent: + task.cancel() + try: + await task + except asyncio.CancelledError: + pass + else: + raise AssertionError("helper cancellation did not propagate") + else: + state.disconnected = True + assert await task is None + assert state.closed + assert state.cancelled + + asyncio.run(_run(False)) + asyncio.run(_run(True)) diff --git a/studio/backend/tests/test_openai_tool_passthrough.py b/studio/backend/tests/test_openai_tool_passthrough.py index baa3e50b95..b54f4c130d 100644 --- a/studio/backend/tests/test_openai_tool_passthrough.py +++ b/studio/backend/tests/test_openai_tool_passthrough.py @@ -1,13 +1,7 @@ # SPDX-License-Identifier: AGPL-3.0-only # Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. -"""Tests for the OpenAI /v1/chat/completions client-side tool pass-through. - -Covers ChatMessage tool/assistant roles, ChatCompletionRequest tool fields and -extra="allow", anthropic_tool_choice_to_openai, _build_passthrough_payload -tool_choice propagation, and _friendly_error's httpx-to-"Lost connection" -mapping. No server or GPU required. -""" +"""Tests for the OpenAI /v1/chat/completions client-side tool pass-through.""" import os import sys @@ -914,12 +908,6 @@ class TestOpenAICompatibilityHelpers: class TestFriendlyErrorHttpx: - """When llama-server is down, httpx RequestError strings lack the - "Lost connection to llama-server" substring the sync path keys off, so the - old substring-only `_friendly_error` returned a useless generic message. - These tests pin the new isinstance-based mapping. - """ - def _req(self): return httpx.Request("POST", "http://127.0.0.1:65535/v1/chat/completions") @@ -937,7 +925,7 @@ class TestFriendlyErrorHttpx: def test_read_timeout_mapped(self): exc = httpx.ReadTimeout("timed out", request = self._req()) - assert "Lost connection" in _friendly_error(exc) + assert "first token within 20 minutes" in _friendly_error(exc) def test_non_httpx_unchanged(self): # Non-httpx exceptions still fall through to the substring heuristics diff --git a/studio/backend/utils/llama_cpp_update.py b/studio/backend/utils/llama_cpp_update.py index 6eb34ffe34..8b90d36bc5 100644 --- a/studio/backend/utils/llama_cpp_update.py +++ b/studio/backend/utils/llama_cpp_update.py @@ -310,8 +310,9 @@ def get_update_status(*, force_refresh: bool = False) -> dict: def _rocm_install_args(asset: Optional[str]) -> list[str]: """Forward --rocm-gfx/--has-rocm from the marker asset, mirroring setup.sh. - The installer probe can miss the gfx arch on amd-smi-only hosts; lemonade - bundles carry the family in the name (rocm-gfx110X), fork bundles only rocm/hip.""" + The installer probe can miss the gfx arch on amd-smi-only hosts; per-gfx + ROCm bundles carry the family in the name (rocm-gfx110X), version-tagged + bundles only rocm/hip.""" if not asset: return [] low = asset.lower() diff --git a/studio/backend/utils/models/model_config.py b/studio/backend/utils/models/model_config.py index 3381ae0391..7245343eeb 100644 --- a/studio/backend/utils/models/model_config.py +++ b/studio/backend/utils/models/model_config.py @@ -1343,6 +1343,7 @@ def _extract_quant_label(filename: str) -> str: "model-UD-IQ1_S.gguf" → "UD-IQ1_S" "model-UD-TQ1_0.gguf" → "UD-TQ1_0" "MXFP4_MOE/model-MXFP4_MOE-0001.gguf"→ "MXFP4_MOE" + "Qwen3.6-IQ4_XS-3.53bpw.gguf" → "IQ4_XS-3.53bpw" """ import re @@ -1358,6 +1359,10 @@ def _extract_quant_label(filename: str) -> str: r"|Q[0-9]+_[0-9]+" # Standard: Q8_0, Q5_1 r"|Q[0-9]+_K" # Short K-quant: Q6_K r"|BF16|F16|F32)" # Full precision + # Optional bits-per-weight modifier so repos that ship multiple + # files at the same base quant (e.g. byteshape's IQ4_XS at 3.53, + # 3.97, 4.19 bpw) don't collapse into a single merged variant. + r"(-[0-9]+(?:\.[0-9]+)?bpw)?" ) match = re.search(quant_re, stem, re.IGNORECASE) # Subdir layouts like ``BF16/foo.gguf`` keep the quant in the directory, @@ -1372,7 +1377,8 @@ def _extract_quant_label(filename: str) -> str: break if match: prefix = match.group(1) or "" - return f"{prefix}{match.group(2)}" + bpw = match.group(3) or "" + return f"{prefix}{match.group(2)}{bpw}" # Fallback: last hyphen-separated segment return stem.split("-")[-1] @@ -2355,10 +2361,13 @@ class ModelConfig: # Does the HF repo contain GGUF files? gguf_filename = detect_gguf_model_remote(identifier, hf_token = hf_token) if gguf_filename: - # Preflight: verify llama-server binary exists before a multi-GB download + # Preflight: verify llama-server binary exists before a multi-GB + # download. include_denied: a transiently locked binary still + # exists (the lock clears long before the download finishes; the + # load itself reports a still-locked binary distinctly). from core.inference.llama_cpp import LlamaCppBackend - if not LlamaCppBackend._find_llama_server_binary(): + if not LlamaCppBackend._find_llama_server_binary(include_denied = True): raise RuntimeError( "llama-server binary not found — cannot load GGUF models. " "Run setup.sh to build it, or set LLAMA_SERVER_PATH." diff --git a/studio/frontend/src/app/provider.tsx b/studio/frontend/src/app/provider.tsx index f806a0f405..39e04843a9 100644 --- a/studio/frontend/src/app/provider.tsx +++ b/studio/frontend/src/app/provider.tsx @@ -259,7 +259,7 @@ function TauriWrapper({ children }: { children: ReactNode }) { <> {children} -
+
{displayTitle} Unsloth
- + {/* settings cog (replaces the up/down chevron) */} + useSettingsDialogStore.getState().openDialog()} > - + {t("shell.navigation.settings")} ⌘, @@ -1239,7 +1245,7 @@ export function AppSidebar() { ref={anchorRef as React.Ref} onSelect={(e) => { e.preventDefault(); toggleTheme(); }} > - {isDark ? : } + {isDark ? : } {isDark ? t("shell.navigation.lightMode") diff --git a/studio/frontend/src/components/llama-update-banner.tsx b/studio/frontend/src/components/llama-update-banner.tsx index 88826e57e5..8f8bc5ca7f 100644 --- a/studio/frontend/src/components/llama-update-banner.tsx +++ b/studio/frontend/src/components/llama-update-banner.tsx @@ -6,6 +6,7 @@ import { useLlamaUpdateCheck } from "@/hooks/use-llama-update-check"; import { useShowLlamaUpdateBanner } from "@/hooks/use-llama-update-pref"; import { toast } from "@/lib/toast"; import { cn } from "@/lib/utils"; +import { Download } from "lucide-react"; import { AnimatePresence, motion } from "motion/react"; import { type ReactElement, useEffect, useRef, useState } from "react"; @@ -132,12 +133,12 @@ export function LlamaUpdateBanner({ transition={{ duration: 0.35, ease: EASE_OUT_QUART }} className={cn( positioned - ? "fixed bottom-4 right-4 z-[9998] w-[calc(100vw-2rem)] max-w-[340px]" + ? "fixed bottom-4 right-4 z-[9998] w-[calc(100vw-2rem)] max-w-[400px]" : "pointer-events-auto w-full", )} data-testid="llama-update-banner" > -
+
{applying ? null : ( +
+
)}
diff --git a/studio/frontend/src/components/web/update-banner.tsx b/studio/frontend/src/components/web/update-banner.tsx index 4aa3943d4f..57fcb98e95 100644 --- a/studio/frontend/src/components/web/update-banner.tsx +++ b/studio/frontend/src/components/web/update-banner.tsx @@ -7,6 +7,7 @@ import { useWebUpdateCheck } from "@/hooks/use-web-update-check"; import { isTauri } from "@/lib/api-base"; import { copyToClipboard } from "@/lib/copy-to-clipboard"; import { cn } from "@/lib/utils"; +import { Download } from "lucide-react"; import { AnimatePresence, motion } from "motion/react"; import { type ReactElement, useEffect, useRef, useState } from "react"; @@ -78,12 +79,12 @@ export function WebUpdateBanner({ transition={{ duration: 0.35, ease: EASE_OUT_QUART }} className={cn( positioned - ? "fixed bottom-4 right-4 z-[9999] w-[calc(100vw-2rem)] max-w-[340px]" + ? "fixed bottom-4 right-4 z-[9999] w-[calc(100vw-2rem)] max-w-[400px]" : "pointer-events-auto w-full", )} data-testid="web-update-banner" > -
+
-
-

- New Unsloth version -

-
-

+

+
-
- + Release notes + +
diff --git a/studio/install_llama_prebuilt.py b/studio/install_llama_prebuilt.py index 5473491bde..46aa14b041 100644 --- a/studio/install_llama_prebuilt.py +++ b/studio/install_llama_prebuilt.py @@ -9,7 +9,6 @@ from __future__ import annotations import argparse import errno import fnmatch -import functools import hashlib import json import os @@ -140,35 +139,6 @@ DEFAULT_PUBLISHED_SHA256_ASSET = os.environ.get( UPSTREAM_REPO = "ggml-org/llama.cpp" UPSTREAM_RELEASES_API = f"https://api.github.com/repos/{UPSTREAM_REPO}/releases/latest" -LEMONADE_ROCM_REPO = "lemonade-sdk/llamacpp-rocm" -LEMONADE_ROCM_RELEASES_API = f"https://api.github.com/repos/{LEMONADE_ROCM_REPO}/releases/latest" - - -def _lemonade_release_api_for(llama_tag: str) -> str: - """Return the GitHub API URL for the lemonade release that matches a - requested llama.cpp tag. - - When llama_tag is unset or "latest", point at /releases/latest. When the - caller has pinned a specific tag (e.g. "b1260"), point at the same tag in - lemonade. Lemonade tracks `ggml-org/llama.cpp` build tags (e.g. "b1260") - but is NOT guaranteed to publish every upstream build -- lemonade may be - several builds behind ggml-org. Pinning to a specific tag that lemonade - skipped will produce a 404 and the caller falls through to the upstream - tarball; that is intentional so pinned installs stay reproducible. - Do NOT pass a `unslothai/llama.cpp` fork tag -- the fork uses its own - namespace and will always 404 against lemonade. - - The tag is URL-encoded with `safe=""` so an unexpected slash / hash / query - character cannot reshape the URL. - """ - normalized = (llama_tag or "").strip() - if not normalized or normalized.lower() == "latest": - return LEMONADE_ROCM_RELEASES_API - return ( - f"https://api.github.com/repos/{LEMONADE_ROCM_REPO}/releases/tags/" - f"{urllib.parse.quote(normalized, safe = '')}" - ) - TEST_MODEL_URL = "https://huggingface.co/ggml-org/models/resolve/main/tinyllamas/stories260K.gguf" TEST_MODEL_SHA256 = "270cba1bd5109f42d03350f60406024560464db173c0e387d91f0426d3bd256d" @@ -1395,22 +1365,10 @@ def direct_linux_release_plan( ) if selection is not None: attempts.extend(selection.attempts) - if host.has_rocm and not host.has_usable_nvidia: - # Per-GPU lemonade prebuilts ship the ROCm runtime libs alongside - # llama.cpp, so they install cleanly even on hosts (e.g. gfx1151 - # Strix Halo) the upstream combined-ROCm tarball doesn't cover. - # "ubuntu" is lemonade's asset naming convention only -- the binary - # is a manylinux-style glibc build that runs on Arch, Fedora, - # openSUSE, etc. with a recent-enough glibc. Do NOT append the CPU - # asset for ROCm-only hosts: if lemonade fails validation we want - # validate_prebuilt_attempts to raise PrebuiltFallback so the caller - # triggers the HIP source build, not silently install a CPU binary. - lemonade_choice = resolve_lemonade_rocm_choice( - host, "ubuntu", "linux-rocm", llama_tag = requested_tag - ) - if lemonade_choice is not None: - attempts.append(lemonade_choice) - elif not host.has_usable_nvidia: + elif not host.has_rocm: + # A ROCm-only host gets no CPU asset: leaving attempts empty lets the + # raise below trigger a HIP source build instead of shipping a CPU + # binary on a GPU host (this ggml-org path has no per-gfx ROCm asset). cpu_choice = published_asset_choice_for_kind(bundle, "linux-cpu") if cpu_choice is not None: attempts.append(cpu_choice) @@ -1489,11 +1447,6 @@ def direct_upstream_release_plan( if pinned is not None: attempts.insert(0, pinned) elif host.has_rocm: - lemonade_choice = resolve_lemonade_rocm_choice( - host, "windows", "windows-hip", llama_tag = requested_tag - ) - if lemonade_choice is not None: - attempts.append(lemonade_choice) hip_asset = f"llama-{release_tag}-bin-win-hip-radeon-x64.zip" hip_url = assets.get(hip_asset) if hip_url: @@ -1566,15 +1519,10 @@ def direct_upstream_release_plan( install_kind = "macos-x64", ) ) - elif host.is_linux and host.is_x86_64 and not host.has_usable_nvidia: - if host.has_rocm: - # Lemonade first, mirroring the Windows ROCm branch above, so a - # ROCm host routed to ggml-org does not silently get the CPU build. - lemonade_choice = resolve_lemonade_rocm_choice( - host, "ubuntu", "linux-rocm", llama_tag = requested_tag - ) - if lemonade_choice is not None: - attempts.append(lemonade_choice) + elif host.is_linux and host.is_x86_64 and not host.has_usable_nvidia and not host.has_rocm: + # ROCm hosts are excluded: this ggml-org path ships no per-gfx ROCm + # asset, so they fall through to the empty-attempts raise (HIP source + # build) rather than silently getting a CPU binary on a GPU host. asset_name = f"llama-{release_tag}-bin-ubuntu-x64.tar.gz" asset_url = assets.get(asset_name) if asset_url: @@ -3115,7 +3063,7 @@ def _apply_host_overrides( A forwarded gfx (--rocm-gfx or UNSLOTH_ROCM_GFX_ARCH) is authoritative and implies ROCm: the installer's own hipinfo/amd-smi probe can miss the arch on amd-smi-only hosts or when setup inferred it from the GPU name, leaving - rocm_gfx_target None and no lemonade prebuilt selected. force_cpu is the + rocm_gfx_target None and no per-gfx ROCm prebuilt selected. force_cpu is the opposite explicit signal (arm64 Linux GPU host whose source build failed): drop GPU attributes so the CPU prebuilt for this OS/arch is selected.""" if force_cpu: @@ -3473,8 +3421,7 @@ def _pinned_windows_cuda_fallback( once upstream ships a driver-runnable build again. The b9360 binary reuses the current release's source tree and convert scripts - and is recorded via binary_release_tag, the same binary/source split used for - the lemonade prebuilt.""" + and is recorded via binary_release_tag.""" if not (host.is_windows and host.is_x86_64 and host.has_usable_nvidia): return None driver = host.driver_cuda_version @@ -3850,46 +3797,30 @@ def _detect_host_rocm_version() -> tuple[int, int] | None: return None -# Map detected gfx IDs to lemonade-sdk asset family suffixes. -# More-specific prefixes must come before shorter ones (e.g. gfx1151 before gfx110). -_LEMONADE_GFX_FAMILIES: list[tuple[str, str]] = [ - ("gfx1151", "gfx1151"), - ("gfx1150", "gfx1150"), - ("gfx120", "gfx120X"), - ("gfx110", "gfx110X"), - ("gfx103", "gfx103X"), -] - - -def _lemonade_gfx_family(gfx_id: str) -> str | None: - gfx_id = gfx_id.lower().strip() - for prefix, family in _LEMONADE_GFX_FAMILIES: - if gfx_id.startswith(prefix): - return family - return None - - def published_rocm_choice_for_host( release: PublishedReleaseBundle, host: HostInfo, install_kind: str ) -> AssetChoice | None: """Select the published ROCm bundle whose gfx target covers the host GPU. - The manifest's gfx_target uses the same umbrella family labels that - _lemonade_gfx_family produces (gfx110X, gfx120X, ...), so the host's detected - gfx is matched either to that family or to the bundle's concrete - mapped_targets list. Returns None when no published bundle covers the GPU, so - the caller can fall back (lemonade / upstream HIP).""" + The manifest's gfx_target uses umbrella family labels (gfx110X, gfx120X, + ...). A host's detected gfx is matched against the bundle's concrete + mapped_targets list, or against the family label itself. Returns None when no + published bundle covers the GPU, so the caller falls back to a HIP source + build.""" if not host.rocm_gfx_target: return None gfx = host.rocm_gfx_target.lower().strip() for artifact in release.artifacts: if artifact.install_kind != install_kind: continue - # Match on the concrete built-arch list, not the family prefix: an - # in-generation-but-unbuilt arch (e.g. gfx1033 in the gfx103 prefix) must - # NOT be served the family bundle. None makes the caller fall back to a - # source build for that GPU. - if gfx not in {target.lower() for target in artifact.mapped_targets}: + # Match the concrete built-arch list so an in-generation-but-unbuilt arch + # (e.g. gfx1033 in the gfx103 family) is NOT served the family bundle and + # falls back to a source build. Also accept the family label itself: the + # llama.cpp update path re-derives --rocm-gfx from the family-named marker + # asset, so an update forwards the family token (gfx110X), not a concrete + # arch. + mapped = {target.lower() for target in artifact.mapped_targets} + if gfx not in mapped and gfx != (artifact.gfx_target or "").lower(): continue asset_url = release.assets.get(artifact.asset_name) if not asset_url: @@ -3910,184 +3841,7 @@ def published_rocm_choice_for_host( return None -def _is_trusted_github_release_url(url: str, expected_repo: str) -> bool: - """Validate a release asset URL points at GitHub's expected hosts. - - Accepts: - https://github.com/{expected_repo}/releases/download/... - https://objects.githubusercontent.com/... (GitHub's release CDN) - Anything else (including http://, raw.githubusercontent.com, gist, etc.) - is rejected so a malicious API response cannot redirect downloads to an - attacker-chosen host. - """ - if not isinstance(url, str) or not url: - return False - try: - parsed = urllib.parse.urlparse(url) - except Exception: - return False - if parsed.scheme != "https": - return False - host = (parsed.netloc or "").lower() - if host == "objects.githubusercontent.com": - # GitHub's release CDN. Restrict to release-asset paths so a tampered - # API response pointing at an arbitrary CDN object is still rejected. - # Real release asset URLs carry the "/github-production-release-asset-" - # prefix; gist / raw / avatar CDN paths do not. - return parsed.path.startswith("/github-production-release-asset-") - if host == "github.com": - return parsed.path.startswith(f"/{expected_repo}/releases/download/") - return False - - -# (gfx_target, asset_name) pairs already logged. resolve_lemonade_rocm_choice() -# runs twice per install (direct planner + resolve_upstream_asset_choice), so -# this stops its selection banner and hash-manifest NOTE printing twice. -_lemonade_selection_logged: "set[tuple[str, str]]" = set() - - -@functools.lru_cache(maxsize = 8) -def _fetch_lemonade_release_cached(api_url: str, llama_tag: str) -> "dict | None": - """Cached wrapper around fetch_json for lemonade release lookups. - - resolve_lemonade_rocm_choice() is called twice per install (once from the - direct planner, once from resolve_upstream_asset_choice) with identical - arguments. Without memoisation, each install hits api.github.com twice, - doubling the rate-limit failure surface on busy CI runners. Cache is - process-scoped; tests that need to vary fetch_json's return value across - invocations should call cache_clear(). - """ - try: - return fetch_json(api_url) - except Exception as exc: - normalized = (llama_tag or "").strip().lower() - if normalized and normalized != "latest": - log( - f"Could not fetch {LEMONADE_ROCM_REPO} release for " - f"llama_tag={llama_tag!r} ({exc}); skipping lemonade prebuilt" - ) - else: - log(f"Could not fetch {LEMONADE_ROCM_REPO} latest release: {exc}") - return None - - -def resolve_lemonade_rocm_choice( - host: HostInfo, - os_prefix: str, - install_kind: str, - llama_tag: str = "latest", -) -> "AssetChoice | None": - """Return an AssetChoice from lemonade-sdk/llamacpp-rocm for the detected GPU, or None. - - os_prefix: lemonade's asset filename label, NOT a host-distro filter. - Pass "ubuntu" for any Linux host (Arch, Fedora, openSUSE, - Debian, ...) -- lemonade only publishes one Linux variant - and it is a manylinux-style glibc build that runs on any - distro with a recent-enough glibc. Pass "windows" for - Windows hosts. - install_kind: "linux-rocm" or "windows-hip" - llama_tag: the requested upstream llama.cpp tag ("latest" or a pinned - release like "b1260"). When pinned, the resolver fetches - the matching lemonade release. When the pinned tag is not - published by lemonade we skip silently (and the caller - falls through to upstream) rather than drift to whatever - lemonade ships as latest. - """ - if not host.rocm_gfx_target: - return None - # Opt-out for users who want the upstream HIP build path only -- lemonade - # binaries are downloaded without entries in the approved-hash manifest, so - # the integrity gate is functional validation only. - if os.environ.get("UNSLOTH_DISABLE_LEMONADE_ROCM", "").strip().lower() in ( - "1", - "true", - "yes", - ): - log("UNSLOTH_DISABLE_LEMONADE_ROCM is set; skipping lemonade-sdk prebuilt") - return None - gfx_family = _lemonade_gfx_family(host.rocm_gfx_target) - if gfx_family is None: - log( - f"AMD GPU {host.rocm_gfx_target!r} is not covered by lemonade-sdk ROCm prebuilts; " - "skipping lemonade prebuilt" - ) - return None - api_url = _lemonade_release_api_for(llama_tag) - release = _fetch_lemonade_release_cached(api_url, llama_tag) - if release is None: - return None - release_tag = release.get("tag_name") if isinstance(release, dict) else None - if not isinstance(release_tag, str) or not release_tag: - log(f"Unexpected {LEMONADE_ROCM_REPO} release payload; skipping lemonade prebuilt") - return None - assets = release_asset_map(release) - asset_name = f"llama-{release_tag}-{os_prefix}-rocm-{gfx_family}-x64.zip" - if asset_name not in assets: - log( - f"{LEMONADE_ROCM_REPO}@{release_tag} has no asset {asset_name!r}; " - "skipping lemonade prebuilt" - ) - return None - asset_url = assets[asset_name] - if not asset_url: - # release_asset_map defaults to "" when an asset row is missing - # browser_download_url; skip cleanly instead of letting - # download_file("") raise a less obvious error downstream. - log( - f"{LEMONADE_ROCM_REPO}@{release_tag} asset {asset_name!r} has no " - "browser_download_url; skipping lemonade prebuilt" - ) - return None - # Defence in depth: lemonade browser_download_url should be on github.com - # or githubusercontent.com. A compromised GitHub API response that - # redirects to an attacker-chosen host would otherwise be honoured - # silently (lemonade assets are not in the approved-hash manifest). - if not _is_trusted_github_release_url(asset_url, LEMONADE_ROCM_REPO): - log( - f"{LEMONADE_ROCM_REPO}@{release_tag} asset {asset_name!r} points " - f"to an unexpected host ({asset_url!r}); refusing to download " - "lemonade prebuilt" - ) - return None - # Note: lemonade tags Linux assets with "ubuntu" but the binary is a - # generic glibc build that runs on any distro (Arch, Fedora, ...), so - # this attempt is selected for all Linux ROCm hosts, not just Ubuntu. - # Log once per (gfx_target, asset); see _lemonade_selection_logged. - log_key = (host.rocm_gfx_target, asset_name) - if log_key not in _lemonade_selection_logged: - _lemonade_selection_logged.add(log_key) - log( - f"AMD GPU {host.rocm_gfx_target!r} ({gfx_family}) -- " - f"trying lemonade-sdk ROCm prebuilt {asset_name} " - f"(works on any glibc Linux, not just Ubuntu)" - ) - log( - f"NOTE: lemonade-sdk/llamacpp-rocm releases are not covered by the " - f"Unsloth approved-hash manifest; download integrity relies on " - f"functional validation (llama-bench / llama-server smoke tests) " - f"after extraction. Set UNSLOTH_DISABLE_LEMONADE_ROCM=1 to skip " - f"lemonade and fall back to the upstream HIP build path." - ) - return AssetChoice( - repo = LEMONADE_ROCM_REPO, - tag = release_tag, - name = asset_name, - url = asset_url, - source_label = "lemonade", - install_kind = install_kind, - ) - - -def resolve_upstream_asset_choice( - host: HostInfo, - llama_tag: str, - lemonade_tag: "str | None" = None, -) -> AssetChoice: - # lemonade_tag: tag for the lemonade lookup only. The release scan pins - # llama_tag to per-release upstream tags (b9518, ...) that lemonade's own - # tag series (b1292, ...) never contains, so pinning lemonade to them 404s - # on every scanned release. Scan callers pass the original request - # (normally "latest") here; upstream asset names keep the pinned tag. +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: # AMD ROCm: try upstream ROCm prebuilt first, then fall back to source build. @@ -4095,15 +3849,7 @@ def resolve_upstream_asset_choice( # the exact GPU target via rocminfo, which is more reliable for consumer # GPUs (e.g. gfx1151) that may not be in the prebuilt. if host.has_rocm and not host.has_usable_nvidia: - # Try lemonade-sdk per-GPU prebuilt first: these are built against - # specific gfx targets and bundle all required ROCm runtime libs. - lemonade_choice = resolve_lemonade_rocm_choice( - host, "ubuntu", "linux-rocm", llama_tag = lemonade_tag or llama_tag - ) - if lemonade_choice is not None: - return lemonade_choice - - # Fall back to upstream combined ROCm tarball. + # Upstream combined ROCm tarball. # Scan upstream assets for any rocm- prebuilt. When the # host ROCm runtime version is known, pick the newest candidate # whose major.minor is <= host version -- otherwise a ROCm 6.4 @@ -4181,14 +3927,8 @@ def resolve_upstream_asset_choice( return attempts[0] raise PrebuiltFallback("no compatible Windows CUDA asset was found") - # AMD ROCm on Windows: try lemonade per-GPU prebuilt first, then upstream HIP + # AMD ROCm on Windows: try upstream HIP prebuilt, then fall back to CPU if host.has_rocm: - lemonade_choice = resolve_lemonade_rocm_choice( - host, "windows", "windows-hip", llama_tag = lemonade_tag or llama_tag - ) - if lemonade_choice is not None: - return lemonade_choice - hip_name = f"llama-{llama_tag}-bin-win-hip-radeon-x64.zip" if hip_name in upstream_assets: log(f"AMD ROCm detected on Windows -- trying upstream HIP prebuilt {hip_name}") @@ -4243,16 +3983,12 @@ def resolve_upstream_asset_choice( raise PrebuiltFallback(f"no prebuilt policy exists for {host.system} {host.machine}") -def resolve_asset_choice( - host: HostInfo, - llama_tag: str, - lemonade_tag: "str | None" = None, -) -> 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: raise PrebuiltFallback( "Linux CUDA installs require a compatible published bundle; upstream fallback is not available" ) - return resolve_upstream_asset_choice(host, llama_tag, lemonade_tag = lemonade_tag) + return resolve_upstream_asset_choice(host, llama_tag) def resolve_release_asset_choice( @@ -4260,7 +3996,6 @@ def resolve_release_asset_choice( llama_tag: str, release: PublishedReleaseBundle, checksums: ApprovedReleaseChecksums, - requested_tag: "str | None" = None, ) -> list[AssetChoice]: if host.is_windows and host.is_x86_64 and host.has_usable_nvidia: torch_preference = detect_torch_cuda_runtime_preference(host) @@ -4317,7 +4052,7 @@ def resolve_release_asset_choice( ) return apply_approved_hashes( - [resolve_asset_choice(host, llama_tag, lemonade_tag = requested_tag)], + [resolve_asset_choice(host, llama_tag)], checksums, ) @@ -6108,15 +5843,6 @@ def apply_approved_hashes( approved_attempts: list[AssetChoice] = [] missing_assets: list[str] = [] for attempt in attempts: - # External prebuilts (e.g. lemonade-sdk) are not listed in the - # approved-hash manifest; they are explicitly documented as relying - # on functional validation only (llama-bench / smoke tests). - # Passing them through here lets the caller include both a lemonade - # attempt and a hash-approved upstream fallback in the same list - # without apply_approved_hashes discarding the lemonade entry. - if attempt.source_label == "lemonade": - approved_attempts.append(attempt) - continue approved = approved_hash_for_attempt(attempt) if approved is None: missing_assets.append(attempt.name) @@ -6204,13 +5930,11 @@ def resolve_install_attempts( return requested_tag, plan.llama_tag, plan.attempts, plan.approved_checksums -def _linux_published_attempts( - host: HostInfo, bundle: PublishedReleaseBundle, requested_tag: str -) -> list[AssetChoice]: +def _linux_published_attempts(host: HostInfo, bundle: PublishedReleaseBundle) -> list[AssetChoice]: """Build the install attempts for a fork Linux host from a manifest-described - bundle: CUDA (with a CPU fallback), per-gfx ROCm (with a lemonade fallback), - or CPU. Same selection the upstream filename path used, just sourced from the - manifest instead of reconstructed from asset names.""" + bundle: CUDA (with a CPU fallback), per-gfx ROCm, or CPU. Same selection the + upstream filename path used, just sourced from the manifest instead of + reconstructed from asset names.""" attempts: list[AssetChoice] = [] if host.has_usable_nvidia: # Prefer the cudart major Studio loads at runtime (torch's bundled @@ -6226,20 +5950,14 @@ def _linux_published_attempts( if selection is not None: attempts.extend(selection.attempts) if host.has_rocm and not host.has_usable_nvidia: - # Prefer the fork's own per-gfx ROCm bundle (hash-approved, ships the - # full ROCm runtime) and fall back to the external lemonade prebuilt. - # Do NOT append the CPU asset for ROCm-only hosts: if lemonade fails - # validation we want validate_prebuilt_attempts to raise PrebuiltFallback - # so the caller triggers the HIP source build, not silently install a - # CPU-only binary. + # Use the fork's own per-gfx ROCm bundle (hash-approved, ships the full + # ROCm runtime). Do NOT append the CPU asset for ROCm-only hosts: if no + # bundle covers the GPU we want validate_prebuilt_attempts to raise + # PrebuiltFallback so the caller triggers the HIP source build, not + # silently install a CPU-only binary. published_rocm = published_rocm_choice_for_host(bundle, host, "linux-rocm") if published_rocm is not None: attempts.append(published_rocm) - lemonade_choice = resolve_lemonade_rocm_choice( - host, "ubuntu", "linux-rocm", llama_tag = requested_tag - ) - if lemonade_choice is not None: - attempts.append(lemonade_choice) else: cpu_choice = published_asset_choice_for_kind(bundle, "linux-cpu") if cpu_choice is not None: @@ -6279,7 +5997,7 @@ def _fork_manifest_release_plans( resolved_tag = bundle.upstream_tag try: if host.is_linux: - linux_attempts = _linux_published_attempts(host, bundle, requested_tag) + linux_attempts = _linux_published_attempts(host, bundle) if not linux_attempts: raise PrebuiltFallback("no compatible Linux prebuilt asset was found") attempts = apply_approved_hashes(linux_attempts, checksums) @@ -6293,7 +6011,6 @@ def _fork_manifest_release_plans( resolved_tag, bundle, checksums, - requested_tag = requested_tag, ) if not attempts: raise PrebuiltFallback("no compatible prebuilt asset was found") @@ -6364,10 +6081,10 @@ def write_prebuilt_metadata( "asset": choice.name, "asset_sha256": choice.expected_sha256, "source": choice.source_label, - # Binary-side repo/tag for non-upstream sources (e.g. lemonade). - # published_repo/release_tag always refer to the unsloth source tree; - # these capture where the actual binaries came from so the install - # summary can show both (e.g. "unslothai/llama.cpp@b9334 + lemonade@b1280"). + # Binary-side repo/tag for non-fork sources (e.g. the ggml-org upstream + # CPU/HIP prebuilts). published_repo/release_tag always refer to the + # unsloth source tree; these capture where the actual binaries came from + # so the install summary can show both. "binary_repo": choice.repo, "binary_release_tag": choice.tag, "source_asset": source_asset_name, @@ -6663,7 +6380,7 @@ def validate_prebuilt_choice( approved_checksums = approved_checksums, prebuilt_fallback_used = prebuilt_fallback_used, ) - # Hashless external prebuilts (e.g. lemonade) are not in the approved-sha256 + # Hashless external prebuilts are not in the approved-sha256 # manifest and rely on the functional smoke test as their only integrity gate, # so they are always validated. For an approved bundle the sha256 manifest # already proves integrity, so its runtime smoke test -- a cold CUDA-JIT pass @@ -6773,6 +6490,31 @@ def validate_prebuilt_attempts( raise PrebuiltFallback("no prebuilt bundle passed validation") +def diffusion_visual_server_backfill_needed( + install_dir: Path, host: HostInfo, choice: AssetChoice +) -> bool: + """True when an existing install matches the tag but lacks the DiffusionGemma + visual-server the chosen bundle ships. An install made before the visual-server + entered the copy allowlist matches on tag yet is missing the binary, so the + tag-match skip never backfills it (DiffusionGemma then fails with "runner not + found"). Gated to the fork ("published") bundles that actually carry it, so + upstream installs -- which never ship it -- can't thrash on repeated updates. + Once a re-extract lands the binary this returns False, so it self-limits.""" + if choice.source_label != "published": + return False + name = "llama-diffusion-gemma-visual-server" + (".exe" if host.is_windows else "") + if name not in runtime_patterns_for_choice(choice): + return False + for cand in ( + install_dir / name, + install_dir / "build" / "bin" / name, + install_dir / "build" / "bin" / "Release" / name, + ): + if cand.is_file(): + return False + return True + + def install_prebuilt( install_dir: Path, llama_tag: str, @@ -6811,11 +6553,17 @@ def install_prebuilt( ) 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" - ) - return + if diffusion_visual_server_backfill_needed(install_dir, host, current.attempts[0]): + log( + f"existing install matches {current.release_tag} but is missing the " + "DiffusionGemma visual-server; re-extracting the bundle to backfill it" + ) + else: + log( + "existing llama.cpp install already matches selected release " + f"{current.release_tag} upstream_tag={current.llama_tag}; skipping download and install" + ) + return with tempfile.TemporaryDirectory(prefix = "unsloth-llama-prebuilt-") as tmp: work_dir = Path(tmp) probe_path = work_dir / "stories260K.gguf" @@ -6823,12 +6571,19 @@ def install_prebuilt( release_count = len(release_plans) for release_index, plan in enumerate(release_plans): choice = plan.attempts[0] + backfill = diffusion_visual_server_backfill_needed(install_dir, host, choice) 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 + if backfill: + log( + f"existing install matches fallback {plan.release_tag} but is missing " + "the DiffusionGemma visual-server; re-extracting to backfill it" + ) + else: + log( + "existing llama.cpp install already matches fallback release " + f"{plan.release_tag} upstream_tag={plan.llama_tag}; skipping reinstall" + ) + return log( "selected " f"{choice.name} ({choice.source_label}) from published release " @@ -6846,7 +6601,9 @@ def install_prebuilt( release_tag = plan.release_tag, approved_checksums = plan.approved_checksums, initial_fallback_used = release_index > 0, - existing_install_dir = install_dir, + # a backfill must reinstall, so do not let the inner + # existing-install match short-circuit the re-extract + existing_install_dir = None if backfill else install_dir, ) except ExistingInstallSatisfied: return @@ -6931,7 +6688,7 @@ def parse_args() -> argparse.Namespace: default = os.environ.get("UNSLOTH_ROCM_GFX_ARCH"), help = ( "Forward the AMD gfx target (e.g. gfx1151) that setup.ps1/setup.sh " - "resolved, so the lemonade HIP prebuilt is selected even when the " + "resolved, so the per-gfx ROCm prebuilt is selected even when the " "installer's own hipinfo/amd-smi probe cannot report it. Implies " "--has-rocm. Defaults to the UNSLOTH_ROCM_GFX_ARCH environment variable." ), diff --git a/studio/install_python_stack.py b/studio/install_python_stack.py index e540aac305..d404041ca7 100644 --- a/studio/install_python_stack.py +++ b/studio/install_python_stack.py @@ -433,7 +433,7 @@ def _detect_windows_gfx_arch() -> str | None: # GPU marketing-name → gfx arch table, mirroring setup.ps1's $nameArchTable. -# Most-specific first; first match wins. Covers only arches the lemonade-sdk +# Most-specific first; first match wins. Covers only arches the ROCm # prebuilts / AMD Windows torch indexes support; unknown names return None # (callers then fall back cleanly to CPU). _WIN_GPU_NAME_ARCH_TABLE: "list[tuple[str, str]]" = [ diff --git a/studio/setup.ps1 b/studio/setup.ps1 index 542846ac70..acc7e92875 100644 --- a/studio/setup.ps1 +++ b/studio/setup.ps1 @@ -850,7 +850,7 @@ if (-not $HasNvidiaSmi) { # popping a UAC/DiskPart prompt RunAsInvoker can't suppress (its manifest is # asInvoker; even 'amd-smi version' hangs). So only probe when a HIP SDK is present # (hipinfo found -> un-elevated) or the user opts in; else fall through to WMI name - # inference (enough to pick ROCm wheels + lemonade llama.cpp). + # inference (enough to pick ROCm wheels + the ROCm llama.cpp prebuilt). # An explicit opt-out (UNSLOTH_ENABLE_AMD_SMI=0/false/no/off) wins over the HIP-SDK # heuristic: a HIP SDK binary with a broken runtime can still pop the prompt, so # $HipSdkInstalled must NOT silently re-enable it. @@ -923,7 +923,7 @@ if (-not $HasNvidiaSmi) { } # ── Arch resolution: env-var override → name inference ────────────────── # Runs after all probes, even when none confirmed a ROCm runtime ($HasROCm false): - # the Adrenalin driver alone runs the lemonade-sdk llama.cpp prebuilt (bundles its + # the Adrenalin driver alone runs the per-gfx ROCm llama.cpp prebuilt (bundles its # own runtime), and all it needs is the gfx arch, inferable from the WMI GPU name. # Resolving it here lets setup.ps1 forward --rocm-gfx so a GPU llama.cpp is pulled # instead of CPU. (PyTorch ROCm wheels still require a HIP SDK -- gated on $HasROCm @@ -936,7 +936,7 @@ if (-not $HasNvidiaSmi) { substep "gfx arch from UNSLOTH_ROCM_GFX_ARCH env override: $script:ROCmGfxArch" "Cyan" } # 2. Best-effort name → arch lookup (amd-smi / WMI). Most-specific first, - # first match wins. Covers only arches the lemonade-sdk prebuilts support + # first match wins. Covers only arches the ROCm prebuilts support # (gfx120X/110X/1151/1150/103X); unknown names fall back cleanly to CPU. elseif ($ROCmGpuLabel) { $nameArchTable = @( @@ -947,9 +947,9 @@ if (-not $HasNvidiaSmi) { @{ P = "RX 7900|RX 7800|RX 7700(?!S)|PRO W7900|PRO W7800|PRO W7700"; A = "gfx1100" } # RDNA 3 desktop / workstation (Navi 31) @{ P = "RX 7600|RX 7700S|RX 7650|PRO W7600|PRO W7500|PRO V710"; A = "gfx1102" } # RDNA 3 (Navi 33) @{ P = "780M|760M|740M|Phoenix|Hawk Point|Z1 Extreme|Z2 Extreme"; A = "gfx1103" } # RDNA 3 iGPU (Phoenix / Hawk Point) - @{ P = "RX 6900|RX 6800|RX 6750|RX 6700|PRO W6800|PRO W6900"; A = "gfx1030" } # RDNA 2 (Navi 21) -- lemonade gfx103X - @{ P = "RX 6650|RX 6600|PRO W6600|PRO W6650"; A = "gfx1032" } # RDNA 2 (Navi 23) -- lemonade gfx103X - @{ P = "RX 6500|RX 6400|RX 6300|PRO W6400|PRO W6500"; A = "gfx1034" } # RDNA 2 (Navi 24) -- lemonade gfx103X + @{ P = "RX 6900|RX 6800|RX 6750|RX 6700|PRO W6800|PRO W6900"; A = "gfx1030" } # RDNA 2 (Navi 21) -- gfx103X family + @{ P = "RX 6650|RX 6600|PRO W6600|PRO W6650"; A = "gfx1032" } # RDNA 2 (Navi 23) -- gfx103X family + @{ P = "RX 6500|RX 6400|RX 6300|PRO W6400|PRO W6500"; A = "gfx1034" } # RDNA 2 (Navi 24) -- gfx103X family ) foreach ($row in $nameArchTable) { if ($ROCmGpuLabel -match $row.P) { @@ -2644,7 +2644,10 @@ $SkipPrebuiltInstall = $false $RequestedLlamaTag = if ($env:UNSLOTH_LLAMA_TAG) { $env:UNSLOTH_LLAMA_TAG } else { $DefaultLlamaTag } # GPU Windows (CUDA / ROCm) installs the fork's app-* prebuilts; CPU-only stays # on ggml-org (the fork ships no windows-cpu bundle). Mirrors setup.sh's routing. -$HelperReleaseRepo = if ($HasNvidiaSmi -or $HasROCm) { "unslothai/llama.cpp" } else { "ggml-org/llama.cpp" } +# A resolved gfx arch counts as a GPU host even when $HasROCm is false (Adrenalin +# driver only, no HIP runtime): the fork's per-gfx bundle ships its own runtime, +# so route there instead of ggml-org / a CPU build. +$HelperReleaseRepo = if ($HasNvidiaSmi -or $HasROCm -or $script:ROCmGfxArch) { "unslothai/llama.cpp" } else { "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 } @@ -2754,7 +2757,7 @@ if ($env:UNSLOTH_LLAMA_FORCE_COMPILE -eq "1") { # or the upstream windows-hip fallback, so accept either and never # treat a valid ROCm install as mismatched. A name-inferred gfx # arch (Adrenalin-only, no confirmed runtime) still counts as - # ROCm-capable -- the lemonade prebuilt bundles its own runtime, + # ROCm-capable -- the ROCm prebuilt bundles its own runtime, # mirroring the --rocm-gfx forward below. NOTE: this block is # currently inert -- write_prebuilt_metadata does not persist an # install_kind key, so $existingKind is always null. If that changes, @@ -2785,7 +2788,7 @@ if ($env:UNSLOTH_LLAMA_FORCE_COMPILE -eq "1") { if ($HasROCm) { $prebuiltArgs += "--has-rocm" } - # Forward the resolved gfx arch so the lemonade HIP prebuilt is picked even + # Forward the resolved gfx arch so the per-gfx ROCm prebuilt is picked even # when the installer's probe can't confirm the runtime (amd-smi-only / # Adrenalin-only, name-inferred arch). --rocm-gfx is authoritative and # implies ROCm in install_llama_prebuilt.py, so the GPU prebuilt is selected @@ -2971,10 +2974,10 @@ if (-not $NeedLlamaSourceBuild) { } elseif ($HasROCm -or $script:ROCmGfxArch) { # AMD GPU present but in the CPU-only source-build fallback: a HIP source # build needs the full HIP SDK + ROCm clang toolchain. AMD GPU acceleration - # comes from the lemonade prebuilt (bundles the runtime, no SDK) -- reaching + # comes from the per-gfx ROCm prebuilt (bundles the runtime, no SDK) -- reaching # here means it couldn't be installed. Warn loudly, don't ship a slow CPU build. $_amdArch = if ($script:ROCmGfxArch) { $script:ROCmGfxArch } else { "ROCm" } - substep "[WARN] AMD GPU ($_amdArch) detected, but the GPU-accelerated lemonade" "Yellow" + substep "[WARN] AMD GPU ($_amdArch) detected, but the GPU-accelerated ROCm" "Yellow" substep " llama.cpp prebuilt could not be installed -- falling back to a CPU build." "Yellow" substep " The prebuilt is the AMD GPU path (no HIP SDK required). To restore GPU" "Yellow" substep " acceleration: re-run the installer (check your network / proxy), or set" "Yellow" diff --git a/studio/setup.sh b/studio/setup.sh index e56fab36ef..613d311708 100755 --- a/studio/setup.sh +++ b/studio/setup.sh @@ -340,9 +340,9 @@ binary_tag = str(payload.get("binary_release_tag") or "").strip() if not repo or not release_tag: raise SystemExit(0) -# For non-upstream sources (e.g. lemonade) the published_repo/release_tag -# refer to the unsloth source tree while the actual binaries came from a -# different repo. Show both so the log is unambiguous. +# For non-fork sources (e.g. ggml-org upstream prebuilts) the published_repo/ +# release_tag refer to the unsloth source tree while the actual binaries came +# from a different repo. Show both so the log is unambiguous. if source and source != "upstream" and binary_repo and binary_tag and binary_repo != repo: message = f"installed release: {repo}@{release_tag} + {source}@{binary_tag}" else: @@ -1004,6 +1004,19 @@ else fi done fi +# UNSLOTH_ROCM_GFX_ARCH may be set on a host where no probe fired, so the override +# nested in the AMD-detected branch above never ran and _setup_gfx is still empty. +# Honour it here so the routing guard below and the --rocm-gfx forwarding both see +# it (install_llama_prebuilt.py reads the same env var as the --rocm-gfx default). +if [ "$_setup_nvidia_usable" != true ] && [ -z "${_setup_gfx:-}" ] && [ -n "${UNSLOTH_ROCM_GFX_ARCH:-}" ]; then + _setup_gfx="${UNSLOTH_ROCM_GFX_ARCH}" +fi +# A resolved/forwarded gfx arch (UNSLOTH_ROCM_GFX_ARCH) means an AMD GPU even when +# no ROCm tooling is on PATH; route it to the fork so the per-gfx prebuilt is +# picked instead of ggml-org / a source build. +if [ "$_LINUX_HAS_GPU" = false ] && [ -n "${_setup_gfx:-}" ]; then + _LINUX_HAS_GPU=true +fi if [ "$_HOST_SYSTEM" = "Linux" ] \ && [ "$_HOST_MACHINE" = "x86_64" ] \ @@ -1086,7 +1099,7 @@ else if [ -n "${UNSLOTH_LLAMA_RELEASE_TAG:-}" ]; then _PREBUILT_CMD+=(--published-release-tag "$UNSLOTH_LLAMA_RELEASE_TAG") fi - # Forward the gfx arch resolved above so the lemonade HIP prebuilt is picked + # Forward the gfx arch resolved above so the per-gfx ROCm prebuilt is picked # even when the installer's own probe cannot report it (amd-smi-only hosts, # name-inferred arch). Implies --has-rocm on the installer side. if [ -n "${_setup_gfx:-}" ]; then diff --git a/tests/studio/install/test_install_llama_prebuilt_logic.py b/tests/studio/install/test_install_llama_prebuilt_logic.py index 313cfa9225..9ccfe15a88 100644 --- a/tests/studio/install/test_install_llama_prebuilt_logic.py +++ b/tests/studio/install/test_install_llama_prebuilt_logic.py @@ -2854,7 +2854,7 @@ def test_validate_prebuilt_choice_approved_validation_skipped_when_flag_off(tmp_ def test_validate_prebuilt_choice_hashless_build_always_validated(tmp_path, monkeypatch): - # A hashless external build (e.g. lemonade) has no approved sha256, so the + # A hashless external build has no approved sha256, so the # functional smoke test is its only integrity gate and must run even while the # flag is off -- otherwise a corrupted/replaced archive could be activated. calls = _run_validate_prebuilt_choice(monkeypatch, tmp_path, expected_sha256 = None) diff --git a/tests/studio/install/test_llama_pr_force_and_source.py b/tests/studio/install/test_llama_pr_force_and_source.py index 6c780057ad..8643806127 100644 --- a/tests/studio/install/test_llama_pr_force_and_source.py +++ b/tests/studio/install/test_llama_pr_force_and_source.py @@ -467,7 +467,7 @@ class TestSourcePatternsPs1: # (GPU -> fork, CPU -> ggml-org), mirroring setup.sh. assert "$HelperReleaseRepo = if ($env:UNSLOTH_LLAMA_RELEASE_REPO)" not in self.content assert ( - "$HelperReleaseRepo = if ($HasNvidiaSmi -or $HasROCm) " + "$HelperReleaseRepo = if ($HasNvidiaSmi -or $HasROCm -or $script:ROCmGfxArch) " '{ "unslothai/llama.cpp" } else { "ggml-org/llama.cpp" }' in self.content ) diff --git a/tests/studio/install/test_rocm_support.py b/tests/studio/install/test_rocm_support.py index 1a7de1af6d..a3c9555d0b 100644 --- a/tests/studio/install/test_rocm_support.py +++ b/tests/studio/install/test_rocm_support.py @@ -3130,7 +3130,7 @@ class TestHipSdkInstalledButDeviceInaccessible: # TEST: --rocm-gfx forwarding -- setup.sh/setup.ps1 hand their resolved gfx arch -# to install_llama_prebuilt.py so the lemonade HIP prebuilt is selected even when +# to install_llama_prebuilt.py so the per-gfx ROCm prebuilt is selected even when # the installer's own hipinfo/amd-smi probe cannot report it. _SETUP_SH_PATH = PACKAGE_ROOT / "studio" / "setup.sh" @@ -3219,6 +3219,176 @@ class TestRocmGfxForwarding: assert "--rocm-gfx" in source assert "$script:ROCmGfxArch" in source + def test_setup_sh_routes_inferred_gfx_to_fork(self): + # A forwarded/inferred gfx arch must route to the fork even without ROCm + # tooling on PATH, so the per-gfx prebuilt is picked over ggml-org. Pin + # the routing guard specifically -- a bare "${_setup_gfx:-}" check also + # appears in the unrelated --rocm-gfx forwarding block. + source = _SETUP_SH_PATH.read_text(encoding = "utf-8") + assert '[ "$_LINUX_HAS_GPU" = false ] && [ -n "${_setup_gfx:-}" ]' in source + + def test_setup_ps1_routes_inferred_gfx_to_fork(self): + # Same on Windows: a resolved $script:ROCmGfxArch counts as a fork/GPU + # install even when $HasROCm is false (Adrenalin-only, no HIP runtime). + source = _SETUP_PS1_PATH.read_text(encoding = "utf-8") + assert "$HasNvidiaSmi -or $HasROCm -or $script:ROCmGfxArch" in source + + # The two assertions above pin the guard *text*. The tests below *execute* + # the real routing block from setup.sh / setup.ps1 and assert the resolved + # release repo, so a refactor that keeps the literal but breaks (or drops) + # the inferred-gfx -> fork decision is still caught. All inputs are faked -- + # no GPU, no ROCm tooling on PATH, no network. + + @staticmethod + def _resolve_setup_sh_repo( + host_machine, + nvidia_usable, + setup_gfx, + rocm_gfx_arch_env = "", + ): + """Run setup.sh's release-repo routing block under bash and return the + resolved _HELPER_RELEASE_REPO. PATH is emptied so the rocminfo/amd-smi/ + hipconfig/hipinfo `command -v` probes all miss (no ROCm tooling). + rocm_gfx_arch_env populates UNSLOTH_ROCM_GFX_ARCH for the env-forwarded + path that fires when no probe set _setup_gfx.""" + import shutil + + bash = shutil.which("bash") + if bash is None: + pytest.skip("bash not available") + source = _SETUP_SH_PATH.read_text(encoding = "utf-8") + start = source.index("\n_LINUX_HAS_GPU=false\n") + 1 + end = source.index("\nunset _GPU_TOOL", start) + len("\nunset _GPU_TOOL") + block = source[start:end] + assert "_HELPER_RELEASE_REPO" in block, "setup.sh routing anchors not found" + env = { + "PATH": "", # no rocminfo/amd-smi/hipconfig/hipinfo discoverable + "ROUTING_BLOCK": block, + "_HOST_SYSTEM": "Linux", + "_HOST_MACHINE": host_machine, + "_setup_nvidia_usable": "true" if nvidia_usable else "false", + "_setup_gfx": setup_gfx, + "UNSLOTH_ROCM_GFX_ARCH": rocm_gfx_arch_env, + } + result = subprocess.run( + [bash, "-c", 'eval "$ROUTING_BLOCK"; printf "%s" "$_HELPER_RELEASE_REPO"'], + capture_output = True, + text = True, + timeout = 30, + env = env, + ) + assert result.returncode == 0, result.stderr + return result.stdout.strip() + + def test_setup_sh_inferred_gfx_resolves_to_fork(self): + # No usable NVIDIA, no ROCm tooling on PATH, only a name-inferred gfx + # arch -> the host must still be treated as a GPU host and routed to the + # fork's per-gfx prebuilt, not ggml-org / a source build. Linux x64 and + # arm64 both go through the same fork branch. + assert self._resolve_setup_sh_repo("x86_64", False, "gfx1100") == "unslothai/llama.cpp" + assert self._resolve_setup_sh_repo("aarch64", False, "gfx1100") == "unslothai/llama.cpp" + + def test_setup_sh_env_forwarded_gfx_resolves_to_fork(self): + # UNSLOTH_ROCM_GFX_ARCH set on a host where no probe fired (_setup_gfx + # empty, no usable NVIDIA, no ROCm tooling): setup.sh adopts the env arch + # and routes to the fork, same as the name-inference path. + repo = self._resolve_setup_sh_repo("x86_64", False, "", rocm_gfx_arch_env = "gfx1100") + assert repo == "unslothai/llama.cpp" + + def test_setup_sh_cpu_host_still_resolves_to_ggml(self): + # Guard against over-correcting the fix: a real CPU host (no usable GPU, + # no inferred gfx, no env override) must keep routing to ggml-org for the + # CPU prebuilt. + assert self._resolve_setup_sh_repo("x86_64", False, "") == "ggml-org/llama.cpp" + + @staticmethod + def _resolve_setup_ps1_repo(has_nvidia, has_rocm, gfx_arch): + """Run setup.ps1's $HelperReleaseRepo selection under pwsh and return the + resolved repo.""" + import shutil + + pwsh = shutil.which("pwsh") + if pwsh is None: + pytest.skip("pwsh not available") + source = _SETUP_PS1_PATH.read_text(encoding = "utf-8") + line = next( + ( + ln + for ln in source.splitlines() + if ln.strip().startswith("$HelperReleaseRepo = if (") + ), + None, + ) + assert line is not None, "$HelperReleaseRepo selection not found in setup.ps1" + harness = ( + f"$HasNvidiaSmi = ${'true' if has_nvidia else 'false'}\n" + f"$HasROCm = ${'true' if has_rocm else 'false'}\n" + f"$script:ROCmGfxArch = '{gfx_arch}'\n" + f"{line}\n" + "Write-Output $HelperReleaseRepo" + ) + result = subprocess.run( + [pwsh, "-NoProfile", "-Command", harness], + capture_output = True, + text = True, + timeout = 60, + ) + assert result.returncode == 0, result.stderr + return result.stdout.strip() + + def test_setup_ps1_inferred_gfx_resolves_to_fork(self): + # Adrenalin-only Windows host: $HasROCm is false (no HIP runtime) but a + # gfx arch was inferred -> route to the fork's windows-rocm bundle. + assert self._resolve_setup_ps1_repo(False, False, "gfx1100") == "unslothai/llama.cpp" + + def test_setup_ps1_cpu_host_still_resolves_to_ggml(self): + # No NVIDIA, no ROCm, no inferred gfx -> CPU host stays on ggml-org. + assert self._resolve_setup_ps1_repo(False, False, "") == "ggml-org/llama.cpp" + + +# TEST: _pick_rocm_gfx_target -- visible-device selection from rocminfo output. +# Honours CUDA_VISIBLE_DEVICES/HIP_VISIBLE_DEVICES so a mixed-arch host installs +# the prebuilt for the GPU actually selected, not GPU 0. + +_pick_rocm_gfx_target = prebuilt_mod._pick_rocm_gfx_target + + +def test_pick_rocm_gfx_target_honors_cuda_visible_devices(monkeypatch): + """AMD HIP honours CUDA_VISIBLE_DEVICES identically to HIP_VISIBLE_DEVICES; + on a gfx1151 + gfx1100 mixed host, CUDA_VISIBLE_DEVICES=1 must select gfx1100.""" + # Two GPUs; rocminfo reports each token twice (as in the real tool output). + probe_out = "gfx1151\ngfx1151\ngfx1100\ngfx1100" + monkeypatch.delenv("HIP_VISIBLE_DEVICES", raising = False) + monkeypatch.delenv("ROCR_VISIBLE_DEVICES", raising = False) + monkeypatch.setenv("CUDA_VISIBLE_DEVICES", "1") + assert _pick_rocm_gfx_target(probe_out) == "gfx1100" + + +def test_pick_rocm_gfx_target_cuda_visible_devices_minus_one_returns_none(monkeypatch): + """CUDA_VISIBLE_DEVICES=-1 means no GPU visible; resolver must return None.""" + probe_out = "gfx1151\ngfx1100" + monkeypatch.delenv("HIP_VISIBLE_DEVICES", raising = False) + monkeypatch.delenv("ROCR_VISIBLE_DEVICES", raising = False) + monkeypatch.setenv("CUDA_VISIBLE_DEVICES", "-1") + assert _pick_rocm_gfx_target(probe_out) is None + + +def test_pick_rocm_gfx_target_same_arch_multi_gpu(monkeypatch): + """Regression: [gfx1100, gfx1100, gfx1151] with HIP_VISIBLE_DEVICES=2 must + return gfx1151, not fall back to GPU 0 due to dict.fromkeys collapsing the + two gfx1100 entries into one and making index 2 out of range.""" + # Simulate rocminfo output for 3 GPUs (2x gfx1100 dGPU + 1x gfx1151 APU). + # Each GPU gets its own Agent section with a few token mentions. + probe_out = ( + "***\nAgent 1\n***\n gfx1100 some info\n gfx1100\n" + "***\nAgent 2\n***\n gfx1100 some info\n gfx1100\n" + "***\nAgent 3\n***\n gfx1151 some info\n gfx1151\n" + ) + monkeypatch.delenv("ROCR_VISIBLE_DEVICES", raising = False) + monkeypatch.delenv("CUDA_VISIBLE_DEVICES", raising = False) + monkeypatch.setenv("HIP_VISIBLE_DEVICES", "2") + assert _pick_rocm_gfx_target(probe_out) == "gfx1151" + # TEST: WSL ROCDXG fixes -- drop-in persistence + system-HIP-before-bundle diff --git a/tests/studio/install/test_selection_logic.py b/tests/studio/install/test_selection_logic.py index a4fc09a183..4e8beeb888 100644 --- a/tests/studio/install/test_selection_logic.py +++ b/tests/studio/install/test_selection_logic.py @@ -2963,6 +2963,29 @@ class TestPublishedRocmGfxSelection: is None ), unbuilt + def test_family_token_matches_family_bundle(self): + # The llama.cpp update path re-derives --rocm-gfx from the family-named + # marker asset, so it forwards a family token (gfx110X, lowercased to + # gfx110x by _normalize_forwarded_gfx), not a concrete arch. That must + # still select the family bundle instead of falling to a source build. + release = self._release("linux-rocm", "app-b9457-linux-x64-rocm") + for token in ("gfx110X", "gfx110x"): + choice = INSTALL_LLAMA_PREBUILT.published_rocm_choice_for_host( + release, self._host(token), "linux-rocm" + ) + assert choice is not None, token + assert choice.name == "app-b9457-linux-x64-rocm-gfx110X.tar.gz", token + + def test_windows_family_token_matches_family_bundle(self): + # The Windows update path forwards the same family token (gfx120X) for a + # windows-rocm bundle, so the family-label match must cover it too. + release = self._release("windows-rocm", "app-b9457-windows-x64-rocm") + choice = INSTALL_LLAMA_PREBUILT.published_rocm_choice_for_host( + release, self._host("gfx120x"), "windows-rocm" + ) + assert choice is not None + assert choice.name == "app-b9457-windows-x64-rocm-gfx120X.zip" + class TestPublishedMacosForkSelection: """macOS now routes to the fork (setup.sh), which ships diff --git a/tests/studio/test_llama_cpp_wall_clock_cap.py b/tests/studio/test_llama_cpp_wall_clock_cap.py index 671abea823..f173cfbc55 100644 --- a/tests/studio/test_llama_cpp_wall_clock_cap.py +++ b/tests/studio/test_llama_cpp_wall_clock_cap.py @@ -1,22 +1,4 @@ -""" -Tests for the llama-server wall-clock cap (t_max_predict_ms). - -The UI always sends max_tokens = context_length, so gating -t_max_predict_ms on `max_tokens is None` makes the safety net dead -code. The fix applies the wall-clock cap unconditionally on all three -streaming payload sites and raises the default to 10 minutes so slow -CPU / macOS / Windows installs are not cut off mid-generation. - -Verifies: - - t_max_predict_ms is assigned unconditionally at the three - payload-builder sites (not inside an `if max_tokens is None` else - branch). - - _DEFAULT_T_MAX_PREDICT_MS is at least 10 minutes (previously - 120_000). - - The default max_tokens path still applies _DEFAULT_MAX_TOKENS. - - The three payload variable names (payload x2, stream_payload x1) - each get both `max_tokens` and `t_max_predict_ms`. -""" +"""Timeout policy checks for Studio's local llama-server path.""" from __future__ import annotations @@ -36,88 +18,25 @@ SRC = SOURCE_PATH.read_text() TREE = ast.parse(SRC) -def _is_subscript_assign(stmt: ast.stmt, target_name: str, key: str) -> bool: - if not isinstance(stmt, ast.Assign) or len(stmt.targets) != 1: - return False - t = stmt.targets[0] - if not isinstance(t, ast.Subscript): - return False - if not (isinstance(t.value, ast.Name) and t.value.id == target_name): - return False - slc = t.slice - return isinstance(slc, ast.Constant) and slc.value == key - - -def _collect_assignments(tree, target_name, key): - """Return list of (node, stack_of_enclosing_ifs) for each match.""" - hits = [] - - def visit(node, stack): - if _is_subscript_assign(node, target_name, key): - hits.append((node, stack)) - for child in ast.iter_child_nodes(node): - if isinstance(child, ast.If): - for sub in child.body: - visit(sub, stack + [(child, "body")]) - for sub in child.orelse: - visit(sub, stack + [(child, "orelse")]) - else: - visit(child, stack) - - visit(tree, []) - return hits - - -def test_default_t_max_predict_ms_is_at_least_ten_minutes(): +def _module_constant(name: str): for node in TREE.body: if isinstance(node, ast.Assign) and len(node.targets) == 1: t = node.targets[0] - if isinstance(t, ast.Name) and t.id == "_DEFAULT_T_MAX_PREDICT_MS": + if isinstance(t, ast.Name) and t.id == name: value = node.value assert isinstance(value, ast.Constant) - assert value.value >= 600_000, ( - f"_DEFAULT_T_MAX_PREDICT_MS must be >= 10 minutes " - f"(600_000 ms) to avoid cutting off slow-CPU generations; " - f"got {value.value}" - ) - return - raise AssertionError("_DEFAULT_T_MAX_PREDICT_MS constant missing") + return value.value + raise AssertionError(f"{name} constant missing") -def test_t_max_predict_ms_set_unconditionally_at_three_sites(): - hits_payload = _collect_assignments(TREE, "payload", "t_max_predict_ms") - hits_stream = _collect_assignments(TREE, "stream_payload", "t_max_predict_ms") - total = len(hits_payload) + len(hits_stream) - assert total == 3, ( - f"expected 3 total t_max_predict_ms assignments " - f"(payload x2 + stream_payload x1), got {total}" - ) - for node, stack in hits_payload + hits_stream: - for parent_if, branch in stack: - # The assignment must not be gated by a test that checks - # `max_tokens is None` (which would make it dead code for - # the UI path where max_tokens is always set). - test_src = ast.unparse(parent_if.test) - assert "max_tokens" not in test_src, ( - f"t_max_predict_ms at line {node.lineno} is nested under " - f"`if {test_src}:` -- it must be applied unconditionally so " - f"the wall-clock cap is not dead code for callers that set " - f"max_tokens" - ) +def test_first_token_timeout_is_at_least_twenty_minutes(): + value = _module_constant("_DEFAULT_FIRST_TOKEN_TIMEOUT_S") + assert value >= 1200.0 + + +def test_studio_chat_payloads_do_not_set_wall_clock_generation_cap(): + assert "t_max_predict_ms" not in SRC def test_max_tokens_default_cap_still_applied(): - # _DEFAULT_MAX_TOKENS must still kick in when caller passes None. - # We check the conditional expression `max_tokens if max_tokens is not - # None else _DEFAULT_MAX_TOKENS` appears at each site. - matches = 0 - for node in ast.walk(TREE): - if not isinstance(node, ast.IfExp): - continue - src = ast.unparse(node) - if "max_tokens" in src and "_DEFAULT_MAX_TOKENS" in src: - matches += 1 - assert matches >= 3, ( - f"expected >=3 `max_tokens if max_tokens is not None else " - f"_DEFAULT_MAX_TOKENS` expressions; got {matches}" - ) + assert SRC.count("_DEFAULT_MAX_TOKENS_FLOOR") >= 3 diff --git a/unsloth/models/_utils.py b/unsloth/models/_utils.py index 977ac74689..0dd46a036b 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.6.5" +__version__ = "2026.6.6" __all__ = [ "SUPPORTS_BFLOAT16",