From 61df3aaef00bdb54dde443b358a6001bc9b5d33f Mon Sep 17 00:00:00 2001 From: Leo Borcherding Date: Sat, 30 May 2026 00:46:03 -0500 Subject: [PATCH] feature/use lemonade-sdk llamacpp-rocm binaries (#5303) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit * feat(studio): use lemonade-sdk/llamacpp-rocm per-GPU prebuilts for ROCm hosts For AMD GPUs that rocminfo/hipinfo reports a recognised gfx target (gfx103X / gfx110X / gfx1150 / gfx1151 / gfx120X), resolve_lemonade_rocm_choice() now fetches the latest lemonade-sdk/llamacpp-rocm release and returns the matching per-architecture zip, bundling all required ROCm runtime libs. This runs before the existing upstream ggml-org combined-ROCm tarball fallback on Linux and before the upstream HIP zip on Windows, so both platforms benefit from the more targeted build when available. Changes: - Add LEMONADE_ROCM_REPO / LEMONADE_ROCM_RELEASES_API constants - Add HostInfo.rocm_gfx_target populated from rocminfo (Linux) / hipinfo (Windows) - Add _LEMONADE_GFX_FAMILIES prefix map and _lemonade_gfx_family() helper - Add resolve_lemonade_rocm_choice() that fetches latest lemonade release and constructs the llama-{tag}-{os}-rocm-{gfxFamily}-x64.zip asset URL - Wire into resolve_upstream_asset_choice() for both Linux (ubuntu) and Windows paths * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * Honor pinned llama.cpp tag in lemonade ROCm resolver The resolver always fetched lemonade-sdk/llamacpp-rocm's /releases/latest, ignoring the upstream llama.cpp tag the caller had pinned. On a reproducible install where the user requested 'b1260' that meant we would silently pick up whatever lemonade had published as latest at install time, with no way to roll back to the matching tag. Lemonade tags llama.cpp upstream tags 1:1, so: - When llama_tag is unset or 'latest', keep hitting /releases/latest. - When llama_tag is pinned (e.g. 'b1260'), hit /releases/tags/b1260. - When the pinned tag is not published by lemonade (404), skip silently and let the caller fall through to the upstream tarball -- this keeps pinned installs reproducible instead of drifting. resolve_lemonade_rocm_choice now takes llama_tag (default 'latest' for backward compatibility) and both call sites in resolve_upstream_asset_choice forward the upstream llama_tag to it. Note: this PR still has open integration concerns flagged in review -- the simple-policy planner and approved-checksum manifest don't yet route or accept lemonade assets. Those are larger changes and not in scope for this commit; addressing the pinned-tag drift independently because it is small, localized, and self-contained. * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * Add mock test for lemonade ROCm prebuilt asset resolution Validates GPU family mapping and that resolve_lemonade_rocm_choice returns real lemonade release URLs for all supported gfx targets on both Linux and Windows, without requiring AMD hardware. * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * Wire lemonade ROCm prebuilts into the simple-policy install path setup.sh invokes install_llama_prebuilt.py with --simple-policy, which dispatches through resolve_simple_install_release_plans -> direct_linux_release_plan (or direct_upstream_release_plan on Windows). Those planners only handled CUDA + CPU attempts, so ROCm-only hosts (e.g. gfx1151 Strix Halo) had no compatible prebuilt asset and silently fell through to source build, even though resolve_lemonade_rocm_choice already knew how to fetch a per-GPU lemonade-sdk binary. Add a lemonade ROCm/HIP attempt to both simple-policy planners for ROCm-only hosts, and add regression tests that drive the dispatchers end-to-end so this can't be skipped silently again. * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * Document that lemonade ROCm prebuilts work on any glibc Linux The lemonade-sdk asset filename uses "ubuntu" as a label, but the binary is a manylinux-style glibc build with no Ubuntu-specific dependencies. It runs on Arch, Fedora, openSUSE, Debian, etc. as long as the host glibc is recent enough. No behavior change -- the dispatch already runs for any Linux ROCm host. This commit only clarifies the comment, docstring, and log message so users on non-Ubuntu distros (e.g. Strix Halo on Arch) don't mistake the asset name for distro gating. * Pattern matching fix for libggml-cpu*.so* * fix(lemonade): pass resolved tag to lemonade resolver; add upstream HIP fallback; stub API in tests - direct_linux_release_plan: pass bundle.upstream_tag (not requested_tag) to resolve_lemonade_rocm_choice so a "latest" request doesn't mix a newer lemonade binary with an older planned unsloth release (Codex P2) - direct_upstream_release_plan: same fix on the Windows path (release_tag instead of requested_tag); also add the upstream HIP asset (llama--bin-win-hip-radeon-x64.zip) as a fallback between lemonade and CPU so unsupported GPUs or transient lemonade failures don't silently downgrade to CPU when an upstream ROCm prebuilt exists (Codex P2) - test file: stub fetch_json with a synthetic lemonade release payload so the suite is hermetic and not subject to GitHub API rate limits (Codex P1); add test_simple_policy_windows_hip_falls_back_to_upstream_when_lemonade_unavailable to cover the new HIP fallback path * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * fix(lemonade): revert bad tag fix; keep upstream HIP fallback + hermetic tests The previous commit wrongly passed bundle.upstream_tag / release_tag to resolve_lemonade_rocm_choice. Lemonade uses its own versioning (b1262, b1264, …) completely independent of unslothai's tags (b9186, …), so passing a resolved unslothai tag caused a 404 and silently skipped the lemonade binary entirely. Revert both call sites to requested_tag. Keep the two valid fixes from the prior commit: - Upstream HIP fallback (llama--bin-win-hip-radeon-x64.zip) between lemonade and CPU in direct_upstream_release_plan, so unsupported GPUs or transient lemonade failures don't silently downgrade to CPU (Codex P2) - Stub fetch_json in tests so the suite is hermetic (Codex P1) * fix(studio/rocm): respect HIP_VISIBLE_DEVICES when picking lemonade gfx target The rocminfo / hipinfo regex took the first gfx match in the agent listing. On mixed APU + dGPU hosts (e.g. Strix Halo gfx1151 + discrete RX 7900 gfx1100) this picked whichever GPU appeared first in the tool's stdout, not the one HIP actually runs on. The downloaded lemonade asset could then be a binary for a different arch than the active device. Extracted a module-level _pick_rocm_gfx_target() helper that: - collects every gfx token in order via re.findall (skips gfx000 / generic ISAs) - if HIP_VISIBLE_DEVICES or ROCR_VISIBLE_DEVICES is set, parses the first comma-separated entry as an integer index into that list - falls back to the first GPU for non-integer (UUID-style) or out-of-range values, matching the previous default behaviour Both Linux (rocminfo) and Windows (hipinfo) branches use the helper. Existing 18 lemonade tests pass; no behavioural change for single-GPU hosts. * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * fix(studio/rocm): dedup rocminfo gfx tokens and honor disabled visibility Follow-up to 8793aef0. rocminfo / hipinfo emit each gfx target multiple times per GPU (Name, ISA triple, marketing name), so the prior re.findall indexing returned the wrong device when HIP_VISIBLE_DEVICES picked GPU 1 on a mixed-arch host -- the helper picked the second occurrence of GPU 0 instead. Collapse to unique tokens (insertion-ordered) before indexing. Also handle HIP_VISIBLE_DEVICES / ROCR_VISIBLE_DEVICES values of '' and '-1' as "no AMD visible" (matches the rest of Studio's visibility code), returning None so the planner does not pick a Lemonade asset for a hidden GPU. * fix(studio/rocm): memoise lemonade release lookup + HIP_PATH hipinfo fallback Robustness pass on top of 25d4ab63: 1. resolve_lemonade_rocm_choice() is called twice per install (planner + resolve_upstream_asset_choice), so every install previously hit api.github.com twice with identical args -- doubling the 403/rate-limit failure surface on busy CI runners. Extract the fetch into a functools.lru_cache(maxsize=8) helper keyed on (api_url, llama_tag). The cached helper also owns the error-path logging the resolver was doing inline. Tests that need to vary fetch_json output across invocations should call _fetch_lemonade_release_cached.cache_clear(). 2. Windows detect_host probe for hipinfo / amd-smi only used shutil.which(), so HIP SDK installs that set HIP_PATH but do not put %HIP_PATH%\bin on system PATH classified the host as non-ROCm. The PowerShell installer and studio/install_python_stack.py already resolve HIP_PATH\bin\hipinfo.exe as a fallback; mirror that here so the install planner agrees with the rest of Studio on what counts as a ROCm host. Tests: 18 passed (shipped); 52 extra sim cases pass (added 2 for the lru_cache deduplication path). * fix(studio/rocm): lemonade URL trust pinning, runtime overlay covers HIP libs, opt-out env Robustness pass driven by 5 parallel reviewers of head c08b15e6: 1. URL trust pinning. AssetChoice.url comes from the GitHub API response's browser_download_url field. Lemonade attempts are not in the approved-hash manifest, so a compromised API response could redirect the download to an attacker-chosen host without the integrity gate catching it. New _is_trusted_github_release_url() validates https + github.com/ release path OR objects.githubusercontent.com (GitHub's CDN). Resolver refuses to download otherwise. 2. UNSLOTH_DISABLE_LEMONADE_ROCM opt-out. Users who prefer the upstream HIP build path can set this env var to skip lemonade outright (e.g. for air-gapped installs or stricter trust requirements). The install log already prints a NOTE explaining that lemonade lacks approved-hash coverage so users know the trust model. 3. linux-rocm runtime overlay patterns extended to cover lemonade's bundled HIP/ROCm runtime libs (libamdhip64.so*, libhsa-runtime64.so*, libhipblas*, librocblas*, librocsolver*, librocsparse*, librocrand*, libMIOpen*, libmagma*). The upstream tarball does not ship these (links against system /opt/rocm), so the new glob entries are no-op for upstream and load-bearing for lemonade. Without this, install_from_archives would drop the bundled runtime libs from the lemonade ZIP and llama-server's RPATH would fail to load amdhip64 at first inference. 4. _lemonade_release_api_for now URL-encodes llama_tag with quote(safe=""). Defence in depth: a tag containing /, ?, #, or whitespace cannot reshape the request URL. Tags come from internal resolution today, but this removes the risk if a future caller passes user-controlled input. 5. Empty browser_download_url skipped explicitly in the resolver. The release_asset_map helper defaults missing URLs to "". Previously this would have been passed to download_file("") which raises a less obvious error than the new clean log + return None. 6. Docstring on _lemonade_release_api_for clarified: lemonade tags match ggml-org/llama.cpp tags 1:1, NOT unslothai/llama.cpp fork tags. The earlier P1 review report misread this and re-asked for the tag-drift fix that the author intentionally reverted in f256cea950. 7. Autouse pytest fixture in test_lemonade_llamacpp_rocm_bins_mock.py clears _fetch_lemonade_release_cached between tests. Today's tests all mock the same payload so no pollution surfaces, but the lru_cache becomes a footgun the moment any future test parametrises return values. Tests: 28 passed (was 18, added 10 covering URL pinning, opt-out env, pinned-tag helper, URL encoding, empty URL, runtime patterns). * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * fix(studio/rocm): complete lemonade runtime overlay + honor CUDA_VISIBLE_DEVICES runtime_patterns_for_choice("linux-rocm") was missing libamd_comgr.so*, librocm_kpack.so*, and librocm_sysdeps_*.so*, which are direct NEEDED entries of libamdhip64.so.7 in every lemonade bundle. install_from_archives did not copy them, so the preflight ldd walk failed on llama-server and every lemonade attempt fell back to source build. Confirmed on gfx1151 b1272 by h34v3nzc0dex; manually staging the full bundle with the three missing patterns restored the install and preserved bench parity. Also extend _pick_rocm_gfx_target to check CUDA_VISIBLE_DEVICES after HIP_VISIBLE_DEVICES and ROCR_VISIBLE_DEVICES: AMD's HIP runtime honours all three with identical semantics, so mixed-arch hosts where users set only CUDA_VISIBLE_DEVICES were still picking the first gfx token instead of the user-selected device. Tests: 30 passed (was 28); added 2 for CUDA_VISIBLE_DEVICES (multi-GPU pick + -1 opt-out) and extended the runtime-patterns test to assert the three newly added lib globs are present. * fix: tighten CDN trust check and fix multi-GPU arch selection - _is_trusted_github_release_url: require /github-production-release-asset- path prefix so only real GitHub release CDN URLs are accepted - _pick_rocm_gfx_target: parse rocminfo Agent N section boundaries to build a per-physical-GPU arch list; same arch across multiple GPUs no longer collapses to a single token, so HIP_VISIBLE_DEVICES indexing works correctly - tests: update CDN test to use realistic path prefix, add rejection test for arbitrary CDN path, add regression test for same-arch multi-GPU case * fix: use broad lib*.so* glob for linux-rocm runtime overlay Replace the explicit lib allowlist in runtime_patterns_for_choice with lib*.so* for the linux-rocm path. The lemonade ROCm ZIPs carry a full HIP/ROCm runtime including transitive deps like libLLVM.so.23.0git and libclang-cpp.so.23.0git (pulled in by libamd_comgr.so.3). These names change across ROCm releases and were not in the allowlist, so preflight would see them as unresolved NEEDED entries and fall back to a source build. The broad glob catches everything in the bundle now and in future releases without needing to enumerate each library. * fix: show lemonade binary tag in install summary log line Store binary_repo and binary_release_tag in UNSLOTH_PREBUILT_INFO.json so the setup.sh summary can distinguish the source tree (unslothai/llama.cpp) from the actual binary origin (lemonade-sdk/llamacpp-rocm). Before: 'installed release: unslothai/llama.cpp@b9334' After: 'installed release: unslothai/llama.cpp@b9334 + lemonade@b1280' Upstream installs are unchanged (binary_repo == published_repo). * fix(merge-compat): align Windows ROCm guard and helper name with strix branch - elif host.has_rocm instead of if...not to match fix/rocm-strix-halo-unified-memory - _resolve_exe instead of _resolve_amd_exe with identical body/docstring Eliminates the two conflict hunks that would otherwise block a clean bot merge of feature/lemonade-rocm-prebuilts on top of fix/rocm-strix-halo-unified-memory. * fix: three PR review corrections for lemonade ROCm prebuilt integration - _lemonade_release_api_for: fix docstring claiming lemonade tags match ggml-org 1:1 -- lemonade may be several builds behind ggml-org (noted by oobabooga, confirmed: lemonade b1281 vs ggml-org b9370) - direct_linux_release_plan: move cpu_choice into else-branch so ROCm-only hosts never get a CPU fallback in the attempts list -- a failed lemonade binary now raises PrebuiltFallback and triggers the HIP source build instead of silently installing a CPU-only binary - apply_approved_hashes: pass lemonade attempts through without requiring a manifest entry; lemonade is explicitly documented as relying on functional validation only, so rejecting it here caused PrebuiltFallback on the non-simple-policy path before the upstream ROCm/HIP fallback could be considered * fix: copy hipblaslt/rocblas library subdirs from lemonade ROCm archives copy_globs matches filenames only and copies flat, so it cannot preserve the hipblaslt/library// and rocblas/library// Tensile kernel catalog trees that lemonade ROCm zips ship alongside the .so files. Add runtime_subdirs_for_choice() and call shutil.copytree for each named subdir after the copy_globs pass in install_from_archives. --------- Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com> Co-authored-by: Daniel Han Co-authored-by: Jaeic Lee --- .../test_lemonade_llamacpp_rocm_bins_mock.py | 430 ++++++++++++++++++ studio/install_llama_prebuilt.py | 371 ++++++++++++++- studio/setup.sh | 15 +- 3 files changed, 806 insertions(+), 10 deletions(-) create mode 100644 studio/backend/tests/test_lemonade_llamacpp_rocm_bins_mock.py 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 }