ceos/unsloth
1
0
Fork 0
Code Issues Pull requests Projects Releases Packages Wiki Activity Actions

Merge branch 'main' into pip

Browse source
This commit is contained in:
Daniel Han 2026-06-12 11:20:56 -07:00
commit a370521879
30 changed files with 1466 additions and 1262 deletions
Download patch file Download diff file Expand all files Collapse all files

22
install.ps1
View file

@ -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)

10
install.sh
View file

@ -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..."

159
studio/backend/cloudflare_tunnel.py
View file

@ -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()

275
studio/backend/core/inference/llama_cpp.py
View file

@ -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)

5
studio/backend/main.py
View file

@ -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

355
studio/backend/routes/inference.py
View file

@ -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

11
studio/backend/run.py
View file

@ -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)

304
studio/backend/tests/test_cloudflare_tunnel.py
View file

@ -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()

7
studio/backend/tests/test_gguf_route_cursor_reset.py
View file

@ -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)

523
studio/backend/tests/test_lemonade_llamacpp_rocm_bins_mock.py
View file

@ -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/<tag> 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

7
studio/backend/tests/test_llama_cpp_tool_loop.py
View file

@ -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)

4
studio/backend/tests/test_llama_cpp_update.py
View file

@ -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",

87
studio/backend/tests/test_llama_route_timeouts.py Normal file
View file

@ -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))

16
studio/backend/tests/test_openai_tool_passthrough.py
View file

@ -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

5
studio/backend/utils/llama_cpp_update.py
View file

@ -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()

15
studio/backend/utils/models/model_config.py
View file

@ -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."

2
studio/frontend/src/app/provider.tsx
View file

@ -259,7 +259,7 @@ function TauriWrapper({ children }: { children: ReactNode }) {
<>
{children}
<DownloadManagerPanel />
<div className="pointer-events-none fixed bottom-4 right-4 z-[9998] flex w-[calc(100vw-2rem)] max-w-[340px] flex-col items-stretch gap-2">
<div className="pointer-events-none fixed bottom-4 right-4 z-[9998] flex w-[calc(100vw-2rem)] max-w-[400px] flex-col items-stretch gap-2">
<WebUpdateBanner
positioned={false}
enabled={!WEB_UPDATE_HIDDEN_ROUTES.has(pathname)}

16
studio/frontend/src/components/app-sidebar.tsx
View file

@ -66,7 +66,8 @@ import {
PowerIcon,
PencilEdit02Icon,
LayoutAlignLeftIcon,
Settings02Icon,
Setting07Icon,
Sun03Icon,
TestTube01Icon,
ZapIcon,
} from "@hugeicons/core-free-icons";
@ -82,7 +83,7 @@ import {
} from "@/components/ui/tooltip";
import { Tooltip as TooltipPrimitive } from "radix-ui";
import { HugeiconsIcon } from "@hugeicons/react";
import { ChevronDown, ChevronsUpDown, MoreHorizontalIcon, Moon, Sun } from "lucide-react";
import { ChevronDown, MoreHorizontalIcon, Moon } from "lucide-react";
import { Link, useNavigate, useRouterState } from "@tanstack/react-router";
import {
archiveChatItem,
@ -1209,7 +1210,12 @@ export function AppSidebar() {
<span className="truncate font-heading text-[13.5px] tracking-[0.025em] dark:tracking-[0.04em] font-semibold text-nav-fg">{displayTitle}</span>
<span className="truncate text-[11.5px] tracking-nav text-muted-foreground">Unsloth</span>
</div>
<ChevronsUpDown strokeWidth={1.25} className="ml-auto size-4 text-muted-foreground group-data-[collapsible=icon]:hidden" />
{/* settings cog (replaces the up/down chevron) */}
<HugeiconsIcon
icon={Setting07Icon}
strokeWidth={1.5}
className="ml-auto !size-[18px] text-muted-foreground group-data-[collapsible=icon]:hidden"
/>
</SidebarMenuButton>
</DropdownMenuTrigger>
<DropdownMenuContent
@ -1222,7 +1228,7 @@ export function AppSidebar() {
<DropdownMenuItem
onSelect={() => useSettingsDialogStore.getState().openDialog()}
>
<HugeiconsIcon icon={Settings02Icon} strokeWidth={1.75} className="size-icon" />
<HugeiconsIcon icon={Setting07Icon} strokeWidth={1.75} className="size-icon" />
<span>{t("shell.navigation.settings")}</span>
<DropdownMenuShortcut>⌘,</DropdownMenuShortcut>
</DropdownMenuItem>
@ -1239,7 +1245,7 @@ export function AppSidebar() {
ref={anchorRef as React.Ref<HTMLDivElement>}
onSelect={(e) => { e.preventDefault(); toggleTheme(); }}
>
{isDark ? <Sun strokeWidth={1.75} className="size-icon" /> : <Moon strokeWidth={1.75} className="size-icon" />}
{isDark ? <HugeiconsIcon icon={Sun03Icon} strokeWidth={1.75} className="size-icon" /> : <Moon strokeWidth={1.75} className="size-icon" />}
<span>
{isDark
? t("shell.navigation.lightMode")

53
studio/frontend/src/components/llama-update-banner.tsx
View file

@ -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"
>
<div className="relative overflow-hidden rounded-[24px] bg-white px-4 pb-[22px] pl-6 pt-5 shadow-[0_2px_8px_-2px_rgba(0,0,0,0.16)] dark:bg-card dark:shadow-[0_8px_28px_-6px_rgba(0,0,0,0.28)]">
<div className="relative overflow-hidden rounded-[24px] bg-white px-5 pb-4 pt-5 shadow-[0_2px_8px_-2px_rgba(0,0,0,0.16)] dark:bg-card dark:shadow-[0_8px_28px_-6px_rgba(0,0,0,0.28)]">
{applying ? null : (
<button
type="button"
@ -163,16 +164,23 @@ export function LlamaUpdateBanner({
</button>
)}
<div className="min-w-0 pr-6">
<p className="font-heading text-base font-medium text-foreground">
{applying ? "Updating llama.cpp..." : "New llama.cpp version"}
</p>
<p className="mt-0.5 text-xs text-muted-foreground">
{status?.installed_tag ?? "unknown"} &rarr;{" "}
<span className="font-medium text-foreground">
{status?.latest_tag ?? ""}
</span>
</p>
<div className="flex min-w-0 items-start gap-4 pr-6">
<Download
aria-hidden="true"
className="mt-1 size-5 shrink-0 text-foreground"
strokeWidth={1.75}
/>
<div className="min-w-0">
<p className="font-heading text-base font-medium text-foreground">
{applying ? "Updating llama.cpp..." : "New llama.cpp version"}
</p>
<p className="mt-0.5 text-xs text-muted-foreground">
{status?.installed_tag ?? "unknown"} &rarr;{" "}
<span className="font-medium text-foreground">
{status?.latest_tag ?? ""}
</span>
</p>
</div>
</div>
{applying ? (
@ -195,24 +203,25 @@ export function LlamaUpdateBanner({
/>
</div>
) : (
<div className="mt-3 flex items-center gap-2">
<Button
size="sm"
className="h-auto rounded-full px-3.5 py-2 text-[13px]"
onClick={handleUpdate}
data-testid="llama-update-button"
>
Update
</Button>
<div className="mt-4 flex flex-wrap items-center justify-between gap-y-2">
<Button
size="sm"
variant="ghost"
className="h-auto rounded-full px-2.5 py-2 text-[13px] text-muted-foreground hover:text-foreground"
className="-ml-2 h-auto rounded-full px-2.5 py-2 text-[13px] font-medium text-foreground"
onClick={snooze}
data-testid="llama-update-snooze-button"
>
Remind me later
</Button>
<Button
size="sm"
// ml offsets the pill's filled edge so visual gaps stay equal
className="ml-2.5 h-auto rounded-full px-3.5 py-2 text-[13px]"
onClick={handleUpdate}
data-testid="llama-update-button"
>
Update
</Button>
</div>
)}
</div>

59
studio/frontend/src/components/web/update-banner.tsx
View file

@ -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"
>
<div className="relative overflow-hidden rounded-[24px] bg-white px-4 pb-[22px] pl-6 pt-5 shadow-[0_2px_8px_-2px_rgba(0,0,0,0.16)] dark:bg-card dark:shadow-[0_8px_28px_-6px_rgba(0,0,0,0.28)]">
<div className="relative overflow-hidden rounded-[24px] bg-white px-5 pb-4 pt-5 shadow-[0_2px_8px_-2px_rgba(0,0,0,0.16)] dark:bg-card dark:shadow-[0_8px_28px_-6px_rgba(0,0,0,0.28)]">
<button
type="button"
onClick={dismiss}
@ -107,47 +108,53 @@ export function WebUpdateBanner({
</svg>
</button>
<div className="min-w-0 pr-6">
<p className="font-heading text-base font-medium text-foreground">
New Unsloth version
</p>
<div className="mt-0.5 flex items-center justify-between gap-3">
<p className="text-xs text-muted-foreground">
<div className="flex min-w-0 items-start gap-4 pr-6">
<Download
aria-hidden="true"
className="mt-1 size-5 shrink-0 text-foreground"
strokeWidth={1.75}
/>
<div className="min-w-0">
<p className="font-heading text-base font-medium text-foreground">
New Unsloth version
</p>
<p className="mt-0.5 text-xs text-muted-foreground">
{status.currentVersion} &rarr;{" "}
<span className="font-medium text-foreground">
{status.latestVersion}
</span>
</p>
<a
href={RELEASE_NOTES_URL}
target="_blank"
rel="noopener noreferrer"
className="shrink-0 text-xs font-medium text-muted-foreground transition-colors hover:text-foreground"
data-testid="web-update-release-notes-link"
>
Release notes
</a>
</div>
</div>
<div className="mt-3 flex items-center gap-2">
<Button
size="sm"
className="h-auto rounded-full px-3.5 py-2 text-[13px]"
onClick={handleCopyCommand}
data-testid="web-update-copy-button"
<div className="mt-4 flex flex-wrap items-center justify-between gap-y-2">
<a
href={RELEASE_NOTES_URL}
target="_blank"
rel="noopener noreferrer"
className="-ml-2 whitespace-nowrap rounded-full px-2.5 py-2 text-[13px] font-medium text-foreground transition-colors hover:bg-muted"
data-testid="web-update-release-notes-link"
>
{copied ? "Copied" : "Copy command"}
</Button>
Release notes
</a>
<Button
size="sm"
variant="ghost"
className="h-auto rounded-full px-2.5 py-2 text-[13px] text-muted-foreground hover:text-foreground"
className="h-auto rounded-full px-2.5 py-2 text-[13px] font-medium text-foreground"
onClick={snooze}
data-testid="web-update-snooze-button"
>
Remind me later
</Button>
<Button
size="sm"
// ml offsets the pill's filled edge so visual gaps stay equal
className="ml-2.5 h-auto rounded-full px-3.5 py-2 text-[13px]"
onClick={handleCopyCommand}
data-testid="web-update-copy-button"
>
{copied ? "Copied" : "Copy command"}
</Button>
</div>
</div>
</motion.div>

435
studio/install_llama_prebuilt.py
View file

@ -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-<version> 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."
),

2
studio/install_python_stack.py
View file

@ -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]]" = [

25
studio/setup.ps1
View file

@ -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"

21
studio/setup.sh
View file

@ -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

2
tests/studio/install/test_install_llama_prebuilt_logic.py
View file

@ -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)

2
tests/studio/install/test_llama_pr_force_and_source.py
View file

@ -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
)

172
tests/studio/install/test_rocm_support.py
View file

@ -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

23
tests/studio/install/test_selection_logic.py
View file

@ -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

107
tests/studio/test_llama_cpp_wall_clock_cap.py
View file

@ -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

2
unsloth/models/_utils.py
View file

@ -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",