ceos/unsloth
1
0
Fork 0
Code Issues Pull requests Projects Releases Packages Wiki Activity Actions
unsloth/studio/backend/tests/test_lemonade_llamacpp_rocm_bins_mock.py
Leo Borcherding 84b42c9283
 fix: deduplicate lemonade ROCm prebuilt selection log (#6021) 
* fix: deduplicate lemonade ROCm prebuilt selection log

resolve_lemonade_rocm_choice() is called twice per install (direct
planner + resolve_upstream_asset_choice). The API fetch is already
memoised via _fetch_lemonade_release_cached but the selection log
lines were still emitted on both calls, printing the 'trying
lemonade-sdk ROCm prebuilt' banner and hash-manifest NOTE twice.

Add _lemonade_selection_logged set keyed on (gfx_target, asset_name)
and guard the two log() calls behind a membership check so they print
exactly once per process regardless of call count.

Also extend the _clear_lemonade_release_cache test fixture to clear
the new set between tests to prevent cross-test state bleed.

Fixes #6020

* fix: write log() output to stdout to avoid PowerShell NativeCommandError

On Windows, PowerShell treats any stderr output from a native process as
an error record and prefixes it with 'python.exe :' and sets the
ErrorId to NativeCommandError. Since log() wrote to sys.stderr, every
[llama-prebuilt] status line triggered this, making normal progress
output look like errors in the installer console.

Switch log() to sys.stdout. The download progress bar (DownloadProgress)
retains its stderr/tty logic unchanged -- that path is for interactive
terminal rendering, not status logging.

* fix: remove redundant 'or ""' in lemonade log_key

host.rocm_gfx_target is already guaranteed truthy by the early
return at the top of resolve_lemonade_rocm_choice. The fallback
was dead code.

* Keep resolver stdout machine-readable, route install logs to stdout

log() sending everything to stdout breaks the resolver modes: setup.sh
json.load()s the whole stdout, so one helper log line (network retry,
release-tag scan) corrupts the parse and silently drops back to building
"latest". Default log() to stderr and flip to stdout only on the install
path, where PowerShell otherwise renders stderr as NativeCommandError
noise. Also tighten the lemonade dedup comments.

---------

Co-authored-by: Daniel Han <danielhanchen@gmail.com>
2026-06-11 20:14:22 -07:00

523 lines
21 KiB
Python
Raw Permalink Blame History

# 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
Reference in a new issue View git blame Copy permalink