diff --git a/studio/backend/tests/test_lemonade_llamacpp_rocm_bins_mock.py b/studio/backend/tests/test_lemonade_llamacpp_rocm_bins_mock.py new file mode 100644 index 0000000000..5d2d672890 --- /dev/null +++ b/studio/backend/tests/test_lemonade_llamacpp_rocm_bins_mock.py @@ -0,0 +1,430 @@ +# 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 when + future tests vary the fetch_json mock return value.""" + _cache = getattr(_mod, "_fetch_lemonade_release_cached", None) + if _cache is not None and hasattr(_cache, "cache_clear"): + _cache.cache_clear() + yield + if _cache is not None and hasattr(_cache, "cache_clear"): + _cache.cache_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 + + +# --------------------------------------------------------------------------- +# Simple-policy dispatcher must plan a lemonade ROCm attempt for AMD-only hosts. +# This is the path setup.sh actually invokes (via --simple-policy), so the +# lemonade integration is useless if it isn't wired in here. +# --------------------------------------------------------------------------- + +direct_linux_release_plan = getattr(_mod, "direct_linux_release_plan", None) +direct_upstream_release_plan = getattr(_mod, "direct_upstream_release_plan", None) + + +def _stub_unsloth_release(release_tag: str = "b9022") -> dict: + # Minimal payload that parse_direct_linux_release_bundle accepts. It + # requires at least one `app-{label}-linux-x64*.tar.gz` asset for the + # bundle to be recognised; we ship a bare CPU one so the planner has a + # baseline non-ROCm attempt to fall through to. + asset_name = f"app-{release_tag}-linux-x64.tar.gz" + return { + "tag_name": release_tag, + "name": release_tag, + "assets": [ + { + "name": asset_name, + "browser_download_url": f"https://example.invalid/{asset_name}", + }, + ], + } + + +@pytest.mark.skipif( + direct_linux_release_plan is None, + reason = "simple-policy dispatcher not present on this branch", +) +def test_simple_policy_plans_lemonade_for_rocm_host(): + host = _make_rocm_host("gfx1151") + with patch.object(_mod, "fetch_json", return_value = _stub_lemonade_release()): + plan = direct_linux_release_plan( + _stub_unsloth_release(), + host, + "unslothai/llama.cpp", + "latest", + ) + assert plan is not None, "ROCm host should not be skipped by simple-policy planner" + kinds = [a.install_kind for a in plan.attempts] + assert ( + "linux-rocm" in kinds + ), f"simple-policy planner did not include a lemonade ROCm attempt; got {kinds}" + rocm_attempt = next(a for a in plan.attempts if a.install_kind == "linux-rocm") + assert rocm_attempt.source_label == "lemonade" + assert "gfx1151" in rocm_attempt.name + + +@pytest.mark.skipif( + direct_upstream_release_plan is None, + reason = "simple-policy dispatcher not present on this branch", +) +def test_simple_policy_plans_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"simple-policy planner did not include a lemonade HIP attempt; got {kinds}" + + +@pytest.mark.skipif( + direct_upstream_release_plan is None, + reason = "simple-policy dispatcher not present on this branch", +) +def test_simple_policy_windows_hip_falls_back_to_upstream_when_lemonade_unavailable(): + """If lemonade returns None (e.g. gfx999 or transient API failure), the planner + must still include the upstream HIP asset rather than silently downgrading to CPU.""" + host = _make_rocm_host("gfx999", windows = True) + hip_asset = "llama-b9022-bin-win-hip-radeon-x64.zip" + release = { + "tag_name": "b9022", + "name": "b9022", + "assets": [ + { + "name": hip_asset, + "browser_download_url": f"https://example.invalid/{hip_asset}", + }, + ], + } + plan = direct_upstream_release_plan(release, host, "ggml-org/llama.cpp", "latest") + assert plan is not None + kinds = [a.install_kind for a in plan.attempts] + assert ( + "windows-hip" in kinds + ), f"upstream HIP asset not included as fallback; got {kinds}" + hip_attempt = next(a for a in plan.attempts if a.install_kind == "windows-hip") + assert hip_attempt.source_label == "upstream" + + +# ── Follow-up: pinned-tag URL helper, URL trust pinning, opt-out env, autouse cache clear ── + + +def test_lemonade_release_api_url_pinned_tag(): + """A pinned llama_tag must produce the /releases/tags/ URL.""" + assert _mod._lemonade_release_api_for("b1262").endswith("/releases/tags/b1262") + assert _mod._lemonade_release_api_for("latest").endswith("/releases/latest") + assert _mod._lemonade_release_api_for("").endswith("/releases/latest") + + +def test_lemonade_release_api_url_encodes_tag(): + """Unexpected slashes / hashes in the tag must be URL-encoded so the URL + cannot be reshaped (defence in depth -- tags should already be sanitised + upstream).""" + url = _mod._lemonade_release_api_for("b1260/../latest") + assert "/releases/tags/b1260%2F..%2Flatest" in url + assert "//latest" not in url.split("/releases/tags/", 1)[1] + + +def test_lemonade_resolver_skipped_by_opt_out_env(monkeypatch): + """UNSLOTH_DISABLE_LEMONADE_ROCM=1 must short-circuit the resolver.""" + monkeypatch.setenv("UNSLOTH_DISABLE_LEMONADE_ROCM", "1") + host = _make_rocm_host("gfx1151") + res = resolve_lemonade_rocm_choice(host, "ubuntu", "linux-rocm", llama_tag = "latest") + assert res is None + + +def test_lemonade_resolver_rejects_non_github_url(monkeypatch): + """If the GitHub API response somehow contained an off-host download URL, + the resolver must refuse to use it (lemonade assets are not in the + approved-hash manifest).""" + bad_release = { + "tag_name": _STUB_TAG, + "assets": [ + { + "name": f"llama-{_STUB_TAG}-ubuntu-rocm-gfx1151-x64.zip", + "browser_download_url": "https://attacker.invalid/llama.zip", + }, + ], + } + host = _make_rocm_host("gfx1151") + with patch.object(_mod, "fetch_json", return_value = bad_release): + res = resolve_lemonade_rocm_choice( + host, "ubuntu", "linux-rocm", llama_tag = "latest" + ) + assert res is None + + +def test_lemonade_resolver_rejects_http_scheme(): + assert not _mod._is_trusted_github_release_url( + "http://github.com/lemonade-sdk/llamacpp-rocm/releases/download/x/y.zip", + "lemonade-sdk/llamacpp-rocm", + ) + + +def test_lemonade_resolver_accepts_github_cdn(): + # Real GitHub release CDN URLs carry the /github-production-release-asset- prefix. + assert _mod._is_trusted_github_release_url( + "https://objects.githubusercontent.com/github-production-release-asset-abc123/456/789?token=x", + "lemonade-sdk/llamacpp-rocm", + ) + + +def test_lemonade_resolver_rejects_arbitrary_cdn_path(): + # A CDN URL without the release-asset path prefix must be rejected. + assert not _mod._is_trusted_github_release_url( + "https://objects.githubusercontent.com/abc/def", + "lemonade-sdk/llamacpp-rocm", + ) + + +def test_lemonade_resolver_accepts_release_path(): + url = "https://github.com/lemonade-sdk/llamacpp-rocm/releases/download/b1262/llama-b1262-ubuntu-rocm-gfx1151-x64.zip" + assert _mod._is_trusted_github_release_url(url, "lemonade-sdk/llamacpp-rocm") + + +def test_lemonade_resolver_rejects_wrong_repo(): + """A github.com release URL for a different repo must be rejected.""" + assert not _mod._is_trusted_github_release_url( + "https://github.com/attacker/llamacpp-rocm/releases/download/x/y.zip", + "lemonade-sdk/llamacpp-rocm", + ) + + +def test_lemonade_resolver_rejects_empty_browser_download_url(): + """An asset entry with an empty browser_download_url must fall through.""" + release = { + "tag_name": _STUB_TAG, + "assets": [ + { + "name": f"llama-{_STUB_TAG}-ubuntu-rocm-gfx1151-x64.zip", + "browser_download_url": "", + }, + ], + } + host = _make_rocm_host("gfx1151") + with patch.object(_mod, "fetch_json", return_value = release): + res = resolve_lemonade_rocm_choice( + host, "ubuntu", "linux-rocm", llama_tag = "latest" + ) + assert res is None + + +def test_lemonade_runtime_patterns_include_hip_runtime(): + """linux-rocm overlay must use a broad lib glob to catch all bundled .so files. + + Lemonade ZIPs carry transitive deps (libamd_comgr, libLLVM, libclang-cpp, + ...) whose names change across ROCm releases. A broad ``lib*.so*`` glob + avoids having to enumerate every transitive dependency by name. + """ + from install_llama_prebuilt import runtime_patterns_for_choice, AssetChoice + + choice = AssetChoice( + repo = "lemonade-sdk/llamacpp-rocm", + tag = "b1262", + name = "llama-b1262-ubuntu-rocm-gfx1151-x64.zip", + url = "https://github.com/lemonade-sdk/llamacpp-rocm/releases/download/b1262/x.zip", + source_label = "lemonade", + install_kind = "linux-rocm", + ) + pats = runtime_patterns_for_choice(choice) + # The broad glob must be present so every .so in the lemonade bundle + # (including transitive deps added in future ROCm releases) gets overlaid. + assert "lib*.so*" in pats, f"'lib*.so*' missing from linux-rocm patterns: {pats}" + + +_pick_rocm_gfx_target = getattr(_mod, "_pick_rocm_gfx_target", None) + + +@pytest.mark.skipif( + _pick_rocm_gfx_target is None, + reason = "_pick_rocm_gfx_target not present on this branch", +) +def test_pick_rocm_gfx_target_honors_cuda_visible_devices(monkeypatch): + """AMD HIP honours CUDA_VISIBLE_DEVICES identically to HIP_VISIBLE_DEVICES; + on a gfx1151 + gfx1100 mixed host, CUDA_VISIBLE_DEVICES=1 must select gfx1100.""" + # Two GPUs; rocminfo reports each token twice (as in the real tool output). + probe_out = "gfx1151\ngfx1151\ngfx1100\ngfx1100" + monkeypatch.delenv("HIP_VISIBLE_DEVICES", raising = False) + monkeypatch.delenv("ROCR_VISIBLE_DEVICES", raising = False) + monkeypatch.setenv("CUDA_VISIBLE_DEVICES", "1") + assert _pick_rocm_gfx_target(probe_out) == "gfx1100" + + +@pytest.mark.skipif( + _pick_rocm_gfx_target is None, + reason = "_pick_rocm_gfx_target not present on this branch", +) +def test_pick_rocm_gfx_target_cuda_visible_devices_minus_one_returns_none(monkeypatch): + """CUDA_VISIBLE_DEVICES=-1 means no GPU visible; resolver must return None.""" + probe_out = "gfx1151\ngfx1100" + monkeypatch.delenv("HIP_VISIBLE_DEVICES", raising = False) + monkeypatch.delenv("ROCR_VISIBLE_DEVICES", raising = False) + monkeypatch.setenv("CUDA_VISIBLE_DEVICES", "-1") + assert _pick_rocm_gfx_target(probe_out) is None + + +@pytest.mark.skipif( + _pick_rocm_gfx_target is None, + reason = "_pick_rocm_gfx_target not present on this branch", +) +def test_pick_rocm_gfx_target_same_arch_multi_gpu(monkeypatch): + """Regression: [gfx1100, gfx1100, gfx1151] with HIP_VISIBLE_DEVICES=2 must + return gfx1151, not fall back to GPU 0 due to dict.fromkeys collapsing the + two gfx1100 entries into one and making index 2 out of range.""" + # Simulate rocminfo output for 3 GPUs (2x gfx1100 dGPU + 1x gfx1151 APU). + # Each GPU gets its own Agent section with a few token mentions. + probe_out = ( + "***\nAgent 1\n***\n gfx1100 some info\n gfx1100\n" + "***\nAgent 2\n***\n gfx1100 some info\n gfx1100\n" + "***\nAgent 3\n***\n gfx1151 some info\n gfx1151\n" + ) + monkeypatch.delenv("ROCR_VISIBLE_DEVICES", raising = False) + monkeypatch.delenv("CUDA_VISIBLE_DEVICES", raising = False) + monkeypatch.setenv("HIP_VISIBLE_DEVICES", "2") + assert _pick_rocm_gfx_target(probe_out) == "gfx1151" diff --git a/studio/install_llama_prebuilt.py b/studio/install_llama_prebuilt.py index 31fb841414..9443b13fd7 100644 --- a/studio/install_llama_prebuilt.py +++ b/studio/install_llama_prebuilt.py @@ -9,6 +9,7 @@ from __future__ import annotations import argparse import errno import fnmatch +import functools import hashlib import json import os @@ -100,6 +101,39 @@ 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" ) @@ -196,6 +230,7 @@ class HostInfo: has_physical_nvidia: bool has_usable_nvidia: bool has_rocm: bool = False + rocm_gfx_target: str | None = None @dataclass @@ -1268,9 +1303,26 @@ def direct_linux_release_plan( selection = linux_cuda_choice_from_release(host, bundle) if selection is not None: attempts.extend(selection.attempts) - cpu_choice = published_asset_choice_for_kind(bundle, "linux-cpu") - if cpu_choice is not None: - attempts.append(cpu_choice) + 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) that the upstream combined-ROCm tarball doesn't cover. + # The "ubuntu" label is lemonade's asset naming convention only -- + # the binary is a manylinux-style glibc build that runs on Arch, + # Fedora, openSUSE, etc. as long as the host glibc is recent enough. + # 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. + 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: + attempts.append(cpu_choice) if not attempts: raise PrebuiltFallback("no compatible Linux prebuilt asset was found") approved_checksums = synthetic_checksums_for_release( @@ -1337,6 +1389,11 @@ def direct_upstream_release_plan( ) ) 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: @@ -2618,6 +2675,72 @@ def run_capture( return result +def _pick_rocm_gfx_target(out: str) -> str | None: + """Choose the gfx target rocminfo / hipinfo report for the active GPU. + + A bare first-match picked the wrong device on mixed APU + dGPU hosts + (e.g. Strix Halo gfx1151 + discrete RX 7900 gfx1100). Respect + HIP_VISIBLE_DEVICES / ROCR_VISIBLE_DEVICES / CUDA_VISIBLE_DEVICES so the + asset matches what HIP actually runs on. Falls back to the first GPU when + no env var is set. + + rocminfo / hipinfo print the same gfx token multiple times per GPU (Name, + ISA, marketing-name). We first try to split the output on per-GPU section + headers (rocminfo: "Agent N" blocks, hipinfo: "device#N" entries) and take + exactly one gfx token per section. This gives the correct per-GPU list even + on same-arch multi-GPU hosts (e.g. two RX 7900 XTX cards) where global + dict.fromkeys dedup would collapse both cards to a single entry and make + HIP_VISIBLE_DEVICES=1 point out of range. + + Falls back to insertion-order dedup when the output has no recognisable + section markers (flat gfx-string inputs, unit-test stubs, etc.). + + Empty / "-1" env values mean no AMD GPU is visible to HIP: return None. + """ + # Try to build a per-GPU token list by splitting on section boundaries. + # rocminfo sections are introduced by "Agent N" lines (optionally between + # rows of asterisks). hipinfo sections start with "device#N". + _sections = re.split( + r"(?mi)^\s*\*+\s*$\s*agent\s+\d+\s*$|\bdevice\s*#\s*\d+\b", + out, + ) + if len(_sections) > 1: + # Section-based: one gfx token per GPU section preserves physical order. + _tokens: list[str] = [] + for _sec in _sections[1:]: + _m = re.search(r"gfx[1-9][0-9a-z]{2,3}", _sec.lower()) + if _m: + _tokens.append(_m.group(0)) + else: + # Fallback: insertion-order dedup (handles flat strings / unknown formats). + _raw = re.findall(r"gfx[1-9][0-9a-z]{2,3}", out.lower()) + _tokens = list(dict.fromkeys(_raw)) + + if not _tokens: + return None + + _vis_raw = None + # AMD's HIP runtime honours all three env vars with identical semantics. + for _env in ("HIP_VISIBLE_DEVICES", "ROCR_VISIBLE_DEVICES", "CUDA_VISIBLE_DEVICES"): + _val = os.environ.get(_env) + if _val is not None: + _vis_raw = _val + break + if _vis_raw is not None: + _vis = _vis_raw.strip() + # Empty or "-1" means "no AMD GPU visible" (matches the rest of Studio). + if _vis == "" or _vis == "-1": + return None + _first = _vis.split(",")[0].strip() + try: + _idx = int(_first) + if 0 <= _idx < len(_tokens): + return _tokens[_idx] + except ValueError: + pass + return _tokens[0] + + def detect_host() -> HostInfo: system = platform.system() machine = platform.machine().lower() @@ -2721,6 +2844,7 @@ def detect_host() -> HostInfo: return bool(re.search(r"(?im)^gpu\s*[:\[]\s*\d", stdout)) has_rocm = False + rocm_gfx_target: str | None = None if is_linux: for _cmd, _check in ( # rocminfo: look for a real gfx GPU id (3-4 chars, nonzero first digit). @@ -2743,6 +2867,7 @@ def detect_host() -> HostInfo: if _result.returncode == 0 and _result.stdout.strip(): if _check(_result.stdout): has_rocm = True + rocm_gfx_target = _pick_rocm_gfx_target(_result.stdout) break elif is_windows: # Windows: prefer active probes that validate GPU presence. @@ -2778,6 +2903,8 @@ def detect_host() -> HostInfo: if _result.returncode == 0 and _result.stdout.strip(): if _check(_result.stdout): has_rocm = True + # hipinfo reports "gcnArchName: gfx1100" -- extract if present + rocm_gfx_target = _pick_rocm_gfx_target(_result.stdout) break # Note: amdhip64.dll presence alone is NOT treated as GPU evidence # since the HIP SDK can be installed without an AMD GPU. @@ -2797,6 +2924,7 @@ def detect_host() -> HostInfo: has_physical_nvidia = has_physical_nvidia, has_usable_nvidia = has_usable_nvidia, has_rocm = has_rocm, + rocm_gfx_target = rocm_gfx_target, ) @@ -3295,6 +3423,185 @@ 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 _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 + + +@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( + 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) -> AssetChoice: upstream_assets = github_release_assets(UPSTREAM_REPO, llama_tag) if host.is_linux and host.is_x86_64: @@ -3303,6 +3610,15 @@ def resolve_upstream_asset_choice(host: HostInfo, llama_tag: str) -> AssetChoice # 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 = llama_tag + ) + if lemonade_choice is not None: + return lemonade_choice + + # Fall back to upstream combined ROCm tarball. # Scan upstream assets for any rocm- prebuilt. When the # host ROCm runtime version is known, pick the newest candidate # whose major.minor is <= host version -- otherwise a ROCm 6.4 @@ -3382,8 +3698,14 @@ def resolve_upstream_asset_choice(host: HostInfo, llama_tag: str) -> AssetChoice return attempts[0] raise PrebuiltFallback("no compatible Windows CUDA asset was found") - # AMD ROCm on Windows: try HIP prebuilt + # AMD ROCm on Windows: try lemonade per-GPU prebuilt first, then upstream HIP if host.has_rocm: + lemonade_choice = resolve_lemonade_rocm_choice( + host, "windows", "windows-hip", llama_tag = 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( @@ -3924,6 +4246,22 @@ def runtime_patterns_for_choice(choice: AssetChoice) -> list[str]: ) +def runtime_subdirs_for_choice(choice: AssetChoice) -> list[str]: + """Subdirectory names within the archive root that must be copied into + the overlay directory alongside the flat shared libraries. + + hipBLASLt and rocBLAS expect their Tensile kernel catalog trees + (hipblaslt/library// and rocblas/library//) to sit next to + their shared libraries at runtime. These trees are multi-level and + cannot be handled by copy_globs (filename-only matching, flat copy).""" + if choice.source_label == "lemonade" and choice.install_kind in { + "linux-rocm", + "windows-hip", + }: + return ["hipblaslt", "rocblas"] + return [] + + def metadata_patterns_for_choice(choice: AssetChoice) -> list[str]: patterns = ["BUILD_INFO.txt", "THIRD_PARTY_LICENSES.txt"] if choice.install_kind.startswith("windows"): @@ -4248,6 +4586,10 @@ def install_from_archives( copy_globs( source_dir, overlay_dir, runtime_patterns_for_choice(choice), required = True ) + for _subdir in runtime_subdirs_for_choice(choice): + _src_subdir = source_dir / _subdir + if _src_subdir.is_dir(): + shutil.copytree(_src_subdir, overlay_dir / _subdir, dirs_exist_ok = True) if runtime_extract_dir is not None: # The runtime archive only contributes the CUDA DLLs. # Restrict the overlay to the cudart bundle's known @@ -4996,6 +5338,15 @@ 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) @@ -5196,6 +5547,12 @@ 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_repo": choice.repo, + "binary_release_tag": choice.tag, "source_asset": source_asset_name, "source_sha256": source_sha256, "source_commit": approved_checksums.source_commit, @@ -5274,7 +5631,7 @@ def runtime_payload_health_groups(choice: AssetChoice) -> list[list[str]]: ["libllama.so*"], ["libggml.so*"], ["libggml-base.so*"], - ["libggml-cpu-*.so*"], + ["libggml-cpu*.so*"], ["libmtmd.so*"], ] if choice.install_kind == "linux-cuda": @@ -5283,7 +5640,7 @@ def runtime_payload_health_groups(choice: AssetChoice) -> list[list[str]]: ["libllama.so*"], ["libggml.so*"], ["libggml-base.so*"], - ["libggml-cpu-*.so*"], + ["libggml-cpu*.so*"], ["libmtmd.so*"], ["libggml-cuda.so*"], ] @@ -5299,7 +5656,7 @@ def runtime_payload_health_groups(choice: AssetChoice) -> list[list[str]]: ["libllama.so*"], ["libggml.so*"], ["libggml-base.so*"], - ["libggml-cpu-*.so*"], + ["libggml-cpu*.so*"], ["libmtmd.so*"], ["libggml-hip.so*"], ] diff --git a/studio/setup.sh b/studio/setup.sh index f178a5e5d7..3b263a08da 100755 --- a/studio/setup.sh +++ b/studio/setup.sh @@ -181,12 +181,21 @@ if not isinstance(payload, dict): repo = str(payload.get("published_repo") or "").strip() release_tag = str(payload.get("release_tag") or "").strip() llama_tag = str(payload.get("tag") or "").strip() +source = str(payload.get("source") or "").strip() +binary_repo = str(payload.get("binary_repo") or "").strip() +binary_tag = str(payload.get("binary_release_tag") or "").strip() if not repo or not release_tag: raise SystemExit(0) -message = f"installed release: {repo}@{release_tag}" -if llama_tag and llama_tag != release_tag: - message += f" (tag {llama_tag})" +# 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. +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: + message = f"installed release: {repo}@{release_tag}" + if llama_tag and llama_tag != release_tag: + message += f" (tag {llama_tag})" print(message) PY }