diff --git a/studio/install_llama_prebuilt.py b/studio/install_llama_prebuilt.py index cf9b2cf45f..cca8886777 100644 --- a/studio/install_llama_prebuilt.py +++ b/studio/install_llama_prebuilt.py @@ -2356,10 +2356,9 @@ def pinned_published_release_bundle( return bundle -def validated_checksums_for_bundle( - repo: str, bundle: PublishedReleaseBundle +def _validate_checksums_against_bundle( + repo: str, bundle: PublishedReleaseBundle, checksums: ApprovedReleaseChecksums ) -> ApprovedReleaseChecksums: - checksums = load_approved_release_checksums(repo, bundle.release_tag) manifest_hash = checksums.artifacts.get(bundle.manifest_asset_name) if manifest_hash is not None and bundle.manifest_sha256 is not None: if manifest_hash.sha256 != bundle.manifest_sha256: @@ -2383,6 +2382,129 @@ def validated_checksums_for_bundle( return checksums +def validated_checksums_for_bundle( + repo: str, bundle: PublishedReleaseBundle +) -> ApprovedReleaseChecksums: + checksums = load_approved_release_checksums(repo, bundle.release_tag) + return _validate_checksums_against_bundle(repo, bundle, checksums) + + +def _download_host_resolve_enabled() -> bool: + """Escape hatch to force the legacy GitHub API path instead of the + download-host fast path (which avoids the api.github.com rate limit).""" + return os.environ.get( + "UNSLOTH_LLAMA_DISABLE_DOWNLOAD_HOST_RESOLVE", "" + ).strip().lower() not in {"1", "true", "yes", "on"} + + +def _release_asset_download_url(repo: str, tag: str, asset_name: str) -> str: + """Tag-pinned asset URL on the release-assets CDN (not api.github.com, so no + rate limit).""" + return ( + f"https://github.com/{urllib.parse.quote(repo, safe = '/')}/releases/download/" + f"{urllib.parse.quote(tag, safe = '')}/" + f"{urllib.parse.quote(asset_name, safe = '')}" + ) + + +def _download_host_latest_release_tag(repo: str) -> str | None: + """Authoritative latest tag from GitHub's /releases/latest redirect target + (github.com, no api.github.com rate limit); the fast path pins URLs to it rather + than the checksum asset's self-reported release_tag. /releases/latest resolves by + created_at/make_latest, which can lag the published_at newest the freshness + detection uses. None on 404 so the caller falls back to the API.""" + url = f"https://github.com/{urllib.parse.quote(repo, safe = '/')}/releases/latest" + request = urllib.request.Request( + url, + method = "HEAD", + headers = {"User-Agent": "unsloth-studio-llama-prebuilt"}, + ) + try: + with _URL_OPENER.open(request, timeout = 30) as response: + final_url = response.geturl() + except urllib.error.HTTPError as exc: + if exc.code == 404: + return None + raise + marker = "/releases/tag/" + index = final_url.find(marker) + if index == -1: + return None + tag = urllib.parse.unquote(final_url[index + len(marker) :]).strip("/") + return tag or None + + +def _fetch_download_host_json(url: str) -> Any: + # Public CDN asset: plain unauthenticated GET, not the rate-limited API. + data = download_bytes( + url, + timeout = 30, + headers = {"User-Agent": "unsloth-studio-llama-prebuilt"}, + ) + return json.loads(data.decode("utf-8")) + + +def _download_host_resolved_release(repo: str) -> ResolvedPublishedRelease | None: + """Resolve the latest fork release from the download host with zero + api.github.com calls, reusing the API path's parsing and validation. The latest + tag is the authoritative /releases/latest redirect tag, and the checksum asset's + self-reported release_tag is cross-checked against it. Returns None (caller falls + back to the API) on a missing JSON asset or a tag mismatch.""" + release_tag = _download_host_latest_release_tag(repo) + if not release_tag: + return None + sha_url = _release_asset_download_url(repo, release_tag, DEFAULT_PUBLISHED_SHA256_ASSET) + try: + sha_payload = _fetch_download_host_json(sha_url) + except urllib.error.HTTPError as exc: + if exc.code == 404: + return None + raise + if not isinstance(sha_payload, dict): + return None + # Cross-check the asset's self-reported release_tag against the authoritative + # redirect tag: parse_approved_release_checksums raises on a mismatch. + checksums = parse_approved_release_checksums(repo, release_tag, sha_payload) + # Synthesize the API release payload with tag-pinned CDN URLs for every named + # asset; parse_published_release_bundle then reads the manifest, still no API. + asset_names = set(checksums.artifacts) | { + DEFAULT_PUBLISHED_MANIFEST_ASSET, + DEFAULT_PUBLISHED_SHA256_ASSET, + } + synthetic_release: dict[str, Any] = { + "tag_name": release_tag, + "draft": False, + "prerelease": False, + "assets": [ + { + "name": name, + "browser_download_url": _release_asset_download_url(repo, release_tag, name), + } + for name in sorted(asset_names) + ], + } + try: + bundle = parse_published_release_bundle(repo, synthetic_release) + except urllib.error.HTTPError as exc: + # In-progress release: the checksum asset can land before the manifest; + # treat a manifest 404 like the sha256 404 above and fall back to the API. + if exc.code == 404: + return None + raise + if bundle is None: + return None + # A manifest artifact can be keyed in the checksum JSON under an upstream-tag + # alias, so add a tag-pinned URL for any manifest artifact missing above (the + # API path gets these from the real asset list); sha256 is still verified. + for artifact in bundle.artifacts: + bundle.assets.setdefault( + artifact.asset_name, + _release_asset_download_url(repo, release_tag, artifact.asset_name), + ) + _validate_checksums_against_bundle(repo, bundle, checksums) + return ResolvedPublishedRelease(bundle = bundle, checksums = checksums) + + def published_release_matches_request(bundle: PublishedReleaseBundle, requested_ref: str) -> bool: if requested_ref == "latest": return True @@ -2449,6 +2571,8 @@ def iter_resolved_published_releases( requested_tag: str | None, published_repo: str, published_release_tag: str = "", + *, + allow_download_host_fast_path: bool = True, ) -> Iterable[ResolvedPublishedRelease]: repo = published_repo or DEFAULT_PUBLISHED_REPO normalized_requested = normalized_requested_llama_tag(requested_tag) @@ -2467,6 +2591,29 @@ def iter_resolved_published_releases( ) return + # Fast path: resolve the fork's latest release from the download host (no + # api.github.com rate limit). It surfaces only the single latest release, so the + # caller disables it when the multi-release walk-back is needed (macOS skipping + # too-new prebuilts); a broken latest then drops to source build, not an older + # release. Any rejection/network error is non-fatal and falls through to the API. + if ( + allow_download_host_fast_path + and repo == DEFAULT_PUBLISHED_REPO + and normalized_requested == "latest" + and _download_host_resolve_enabled() + ): + try: + resolved = _download_host_resolved_release(repo) + except PrebuiltFallback as exc: + log(f"download-host latest release rejected for {repo} ({exc}); trying GitHub API") + resolved = None + except Exception as exc: + log(f"download-host latest resolve unavailable for {repo} ({exc}); trying GitHub API") + resolved = None + if resolved is not None: + yield resolved + return + matched_any = False skipped_invalid = 0 yielded_valid = False @@ -6236,6 +6383,9 @@ def _fork_manifest_release_plans( llama_tag, published_repo, published_release_tag, + # macOS relies on the multi-release walk-back to skip too-new prebuilts, + # which the single-latest download-host path cannot provide. + allow_download_host_fast_path = not host.is_macos, ): bundle = resolved_release.bundle checksums = resolved_release.checksums diff --git a/tests/studio/install/test_download_host_resolve.py b/tests/studio/install/test_download_host_resolve.py new file mode 100644 index 0000000000..5108c84062 --- /dev/null +++ b/tests/studio/install/test_download_host_resolve.py @@ -0,0 +1,284 @@ +"""Routing of iter_resolved_published_releases between the download-host fast +path and the GitHub API enumeration; all I/O monkeypatched.""" + +import importlib.util +import json +import sys +import urllib.error +from pathlib import Path + +import pytest + + +PACKAGE_ROOT = Path(__file__).resolve().parents[3] +MODULE_PATH = PACKAGE_ROOT / "studio" / "install_llama_prebuilt.py" +SPEC = importlib.util.spec_from_file_location("studio_install_llama_prebuilt_dlhost", MODULE_PATH) +assert SPEC is not None and SPEC.loader is not None +ILP = importlib.util.module_from_spec(SPEC) +sys.modules[SPEC.name] = ILP +SPEC.loader.exec_module(ILP) + +FORK_REPO = ILP.DEFAULT_PUBLISHED_REPO +PrebuiltFallback = ILP.PrebuiltFallback + + +def _api_raises(*_a, **_k): + raise AssertionError("GitHub API enumeration was used") + + +def _fast_path_raises(_repo): + raise AssertionError("download-host fast path was used") + + +def _empty_api(*_a, **_k): + return iter(()) + + +def _resolve(tag = "latest", **kw): + return list(ILP.iter_resolved_published_releases(tag, FORK_REPO, "", **kw)) + + +def test_fast_path_yields_latest_without_api(monkeypatch): + sentinel = object() + monkeypatch.delenv("UNSLOTH_LLAMA_DISABLE_DOWNLOAD_HOST_RESOLVE", raising = False) + monkeypatch.setattr(ILP, "_download_host_resolved_release", lambda _repo: sentinel) + monkeypatch.setattr(ILP, "iter_published_release_bundles", _api_raises) + assert _resolve() == [sentinel] + + +def test_fast_path_disabled_by_caller_uses_api(monkeypatch): + # macOS passes allow_download_host_fast_path = False to keep the walk-back. + monkeypatch.delenv("UNSLOTH_LLAMA_DISABLE_DOWNLOAD_HOST_RESOLVE", raising = False) + monkeypatch.setattr(ILP, "_download_host_resolved_release", _fast_path_raises) + monkeypatch.setattr(ILP, "iter_published_release_bundles", _empty_api) + with pytest.raises(PrebuiltFallback): + _resolve(allow_download_host_fast_path = False) + + +def test_fast_path_disabled_by_env_uses_api(monkeypatch): + monkeypatch.setenv("UNSLOTH_LLAMA_DISABLE_DOWNLOAD_HOST_RESOLVE", "1") + monkeypatch.setattr(ILP, "_download_host_resolved_release", _fast_path_raises) + monkeypatch.setattr(ILP, "iter_published_release_bundles", _empty_api) + with pytest.raises(PrebuiltFallback): + _resolve() + + +def test_fast_path_skipped_for_non_latest_request(monkeypatch): + monkeypatch.delenv("UNSLOTH_LLAMA_DISABLE_DOWNLOAD_HOST_RESOLVE", raising = False) + monkeypatch.setattr(ILP, "_download_host_resolved_release", _fast_path_raises) + monkeypatch.setattr(ILP, "iter_published_release_bundles", _empty_api) + with pytest.raises(PrebuiltFallback): + _resolve(tag = "b9964") + + +def test_fast_path_none_falls_back_to_api(monkeypatch): + monkeypatch.delenv("UNSLOTH_LLAMA_DISABLE_DOWNLOAD_HOST_RESOLVE", raising = False) + monkeypatch.setattr(ILP, "_download_host_resolved_release", lambda _repo: None) + used = {"api": False} + + def _api(*_a, **_k): + used["api"] = True + return iter(()) + + monkeypatch.setattr(ILP, "iter_published_release_bundles", _api) + with pytest.raises(PrebuiltFallback): + _resolve() + assert used["api"] + + +def test_fast_path_rejected_checksum_falls_back_to_api(monkeypatch): + monkeypatch.delenv("UNSLOTH_LLAMA_DISABLE_DOWNLOAD_HOST_RESOLVE", raising = False) + + def _reject(_repo): + raise PrebuiltFallback("checksum mismatch") + + monkeypatch.setattr(ILP, "_download_host_resolved_release", _reject) + used = {"api": False} + + def _api(*_a, **_k): + used["api"] = True + return iter(()) + + monkeypatch.setattr(ILP, "iter_published_release_bundles", _api) + with pytest.raises(PrebuiltFallback): + _resolve() + assert used["api"] + + +# --- _download_host_resolved_release body (only download_bytes stubbed) -------- + +RELEASE_TAG = "b9964-mix-53618c5" +UPSTREAM_TAG = "b9964" +BINARY_ASSET = "app-b9964-mix-53618c5-windows-x64-rocm-gfx1151.zip" +MANIFEST_ASSET = ILP.DEFAULT_PUBLISHED_MANIFEST_ASSET +SHA256_ASSET = ILP.DEFAULT_PUBLISHED_SHA256_ASSET + + +def _manifest_bytes(): + return json.dumps( + { + "schema_version": 1, + "component": "llama.cpp", + "upstream_tag": UPSTREAM_TAG, + "source_repo": "unslothai/llama.cpp", + "source_repo_url": "https://github.com/unslothai/llama.cpp", + "artifacts": [ + { + "asset_name": BINARY_ASSET, + "install_kind": "windows-rocm-app", + "supported_sms": [], + "gfx_target": "gfx1151", + "mapped_targets": ["gfx1151"], + "rank": 10, + } + ], + } + ).encode("utf-8") + + +def _sha_payload(*, manifest_sha256 = None): + # BINARY_ASSET is deliberately absent: real releases key its hash under an + # upstream-tag alias, so the manifest name must still get a tag-pinned URL. + # The source archive entry keeps _validate_checksums_against_bundle happy. + artifacts = {"llama.cpp-source-b9964.tar.gz": {"sha256": "a" * 64}} + if manifest_sha256 is not None: + artifacts[MANIFEST_ASSET] = {"sha256": manifest_sha256} + return { + "schema_version": 1, + "component": "llama.cpp", + "release_tag": RELEASE_TAG, + "upstream_tag": UPSTREAM_TAG, + "artifacts": artifacts, + } + + +def _stub_downloads( + monkeypatch, + sha_payload, + manifest_bytes, + *, + latest_tag = RELEASE_TAG, +): + def _no_api(*_a, **_k): + raise AssertionError("GitHub API was used") + + monkeypatch.setattr(ILP, "github_release", _no_api) + monkeypatch.setattr(ILP, "fetch_json", _no_api) + # The authoritative latest tag comes from the /releases/latest redirect + # (github.com, no api.github.com); stub it so no real request is made. + monkeypatch.setattr(ILP, "_download_host_latest_release_tag", lambda _repo: latest_tag) + + def _download_bytes(url, *_a, **_k): + if SHA256_ASSET in url: + return json.dumps(sha_payload).encode("utf-8") + if MANIFEST_ASSET in url: + if isinstance(manifest_bytes, Exception): + raise manifest_bytes + return manifest_bytes + raise AssertionError(f"unexpected download: {url}") + + monkeypatch.setattr(ILP, "download_bytes", _download_bytes) + + +def test_resolved_release_adds_tag_pinned_url_for_manifest_only_asset(monkeypatch): + _stub_downloads(monkeypatch, _sha_payload(), _manifest_bytes()) + resolved = ILP._download_host_resolved_release(FORK_REPO) + assert resolved is not None + assert resolved.bundle.release_tag == RELEASE_TAG + # Binary is named only in the manifest, yet the fast path must expose a + # tag-pinned CDN URL for it (the API path gets it from the real asset list). + assert resolved.bundle.assets[BINARY_ASSET] == ( + f"https://github.com/{FORK_REPO}/releases/download/{RELEASE_TAG}/{BINARY_ASSET}" + ) + + +def test_resolved_release_rejects_manifest_checksum_mismatch(monkeypatch): + # A wrong manifest hash in the checksum payload must fail closed, so the + # router falls back to the API rather than trusting the fast path. + _stub_downloads(monkeypatch, _sha_payload(manifest_sha256 = "b" * 64), _manifest_bytes()) + with pytest.raises(PrebuiltFallback, match = "manifest checksum"): + ILP._download_host_resolved_release(FORK_REPO) + + +def test_resolved_release_rejects_release_tag_mismatch(monkeypatch): + # The checksum asset self-reports RELEASE_TAG, but the authoritative + # /releases/latest redirect resolves a different tag: the fast path must not + # pin to the stale self-reported tag (it raises, so the router falls back). + _stub_downloads(monkeypatch, _sha_payload(), _manifest_bytes(), latest_tag = "b9999-mix-other") + with pytest.raises(RuntimeError, match = "did not match pinned release tag"): + ILP._download_host_resolved_release(FORK_REPO) + + +def test_resolved_release_manifest_404_falls_back(monkeypatch): + # An in-progress release can serve the checksum asset before the manifest; a + # manifest 404 returns None so the router falls back to the API. + not_found = urllib.error.HTTPError( + f"https://github.com/{FORK_REPO}/releases/download/{RELEASE_TAG}/{MANIFEST_ASSET}", + 404, + "Not Found", + {}, # type: ignore[arg-type] + None, + ) + _stub_downloads(monkeypatch, _sha_payload(), not_found) + assert ILP._download_host_resolved_release(FORK_REPO) is None + + +# --- _download_host_latest_release_tag (redirect resolution) ------------------- + + +class _FakeResponse: + def __init__(self, url): + self._url = url + + def __enter__(self): + return self + + def __exit__(self, *_a): + return False + + def geturl(self): + return self._url + + +class _FakeOpener: + def __init__( + self, + *, + url = None, + exc = None, + ): + self._url = url + self._exc = exc + + def open( + self, + _request, + timeout = None, + ): + if self._exc is not None: + raise self._exc + return _FakeResponse(self._url) + + +def test_latest_release_tag_parses_redirect(monkeypatch): + final = f"https://github.com/{FORK_REPO}/releases/tag/{RELEASE_TAG}" + monkeypatch.setattr(ILP, "_URL_OPENER", _FakeOpener(url = final)) + assert ILP._download_host_latest_release_tag(FORK_REPO) == RELEASE_TAG + + +def test_latest_release_tag_none_on_404(monkeypatch): + not_found = urllib.error.HTTPError( + f"https://github.com/{FORK_REPO}/releases/latest", + 404, + "Not Found", + {}, + None, # type: ignore[arg-type] + ) + monkeypatch.setattr(ILP, "_URL_OPENER", _FakeOpener(exc = not_found)) + assert ILP._download_host_latest_release_tag(FORK_REPO) is None + + +def test_latest_release_tag_none_when_not_a_tag_url(monkeypatch): + # No /releases/tag/ segment (e.g. redirected somewhere unexpected) -> None. + monkeypatch.setattr(ILP, "_URL_OPENER", _FakeOpener(url = f"https://github.com/{FORK_REPO}")) + assert ILP._download_host_latest_release_tag(FORK_REPO) is None diff --git a/tests/studio/install/test_selection_logic.py b/tests/studio/install/test_selection_logic.py index 7d3120adb8..7575dea581 100644 --- a/tests/studio/install/test_selection_logic.py +++ b/tests/studio/install/test_selection_logic.py @@ -69,6 +69,14 @@ pinned_macos_release_tag = INSTALL_LLAMA_PREBUILT.pinned_macos_release_tag resolve_simple_install_release_plans = INSTALL_LLAMA_PREBUILT.resolve_simple_install_release_plans +@pytest.fixture(autouse = True) +def _disable_download_host_fast_path(monkeypatch): + # This module exercises the GitHub API enumeration and asset selection against + # mocked releases; keep the download-host fast path (real CDN) out of the way. + # test_download_host_resolve.py covers the fast path itself. + monkeypatch.setenv("UNSLOTH_LLAMA_DISABLE_DOWNLOAD_HOST_RESOLVE", "1") + + def load_studio_run_module(monkeypatch): logger = types.SimpleNamespace( debug = lambda *a, **k: None, @@ -1629,7 +1637,7 @@ class TestResolveInstallAttempts: monkeypatch.setattr( INSTALL_LLAMA_PREBUILT, "iter_resolved_published_releases", - lambda requested_tag, published_repo, published_release_tag = "": iter( + lambda requested_tag, published_repo, published_release_tag = "", **_kwargs: iter( [ INSTALL_LLAMA_PREBUILT.ResolvedPublishedRelease( bundle = release, @@ -1674,7 +1682,7 @@ class TestResolveInstallAttempts: monkeypatch.setattr( INSTALL_LLAMA_PREBUILT, "iter_resolved_published_releases", - lambda requested_tag, published_repo, published_release_tag = "": iter( + lambda requested_tag, published_repo, published_release_tag = "", **_kwargs: iter( [ INSTALL_LLAMA_PREBUILT.ResolvedPublishedRelease( bundle = release, @@ -1739,7 +1747,7 @@ class TestResolveInstallAttempts: monkeypatch.setattr( INSTALL_LLAMA_PREBUILT, "iter_resolved_published_releases", - lambda requested_tag, published_repo, published_release_tag = "": iter( + lambda requested_tag, published_repo, published_release_tag = "", **_kwargs: iter( [ INSTALL_LLAMA_PREBUILT.ResolvedPublishedRelease( bundle = release, @@ -1771,7 +1779,7 @@ class TestResolveInstallAttempts: monkeypatch.setattr( INSTALL_LLAMA_PREBUILT, "iter_resolved_published_releases", - lambda requested_tag, published_repo, published_release_tag = "": iter( + lambda requested_tag, published_repo, published_release_tag = "", **_kwargs: iter( [ INSTALL_LLAMA_PREBUILT.ResolvedPublishedRelease( bundle = release, @@ -1820,7 +1828,7 @@ class TestResolveInstallAttempts: monkeypatch.setattr( INSTALL_LLAMA_PREBUILT, "iter_resolved_published_releases", - lambda requested_tag, published_repo, published_release_tag = "": iter( + lambda requested_tag, published_repo, published_release_tag = "", **_kwargs: iter( [ INSTALL_LLAMA_PREBUILT.ResolvedPublishedRelease( bundle = release, @@ -1911,7 +1919,7 @@ class TestResolveInstallAttempts: monkeypatch.setattr( INSTALL_LLAMA_PREBUILT, "iter_resolved_published_releases", - lambda requested_tag, published_repo, published_release_tag = "": iter( + lambda requested_tag, published_repo, published_release_tag = "", **_kwargs: iter( [ INSTALL_LLAMA_PREBUILT.ResolvedPublishedRelease( bundle = release, @@ -1980,7 +1988,7 @@ class TestResolveInstallAttempts: monkeypatch.setattr( INSTALL_LLAMA_PREBUILT, "iter_resolved_published_releases", - lambda requested_tag, published_repo, published_release_tag = "": iter( + lambda requested_tag, published_repo, published_release_tag = "", **_kwargs: iter( [ INSTALL_LLAMA_PREBUILT.ResolvedPublishedRelease( bundle = release, @@ -2030,7 +2038,7 @@ class TestResolveInstallAttempts: monkeypatch.setattr( INSTALL_LLAMA_PREBUILT, "iter_resolved_published_releases", - lambda requested_tag, published_repo, published_release_tag = "": iter( + lambda requested_tag, published_repo, published_release_tag = "", **_kwargs: iter( [ INSTALL_LLAMA_PREBUILT.ResolvedPublishedRelease( bundle = release, @@ -2093,7 +2101,7 @@ class TestResolveInstallAttempts: monkeypatch.setattr( INSTALL_LLAMA_PREBUILT, "iter_resolved_published_releases", - lambda requested_tag, published_repo, published_release_tag = "": iter( + lambda requested_tag, published_repo, published_release_tag = "", **_kwargs: iter( [ INSTALL_LLAMA_PREBUILT.ResolvedPublishedRelease( bundle = release, @@ -2156,7 +2164,9 @@ class TestResolveInstallReleasePlans: monkeypatch.setattr( INSTALL_LLAMA_PREBUILT, "iter_resolved_published_releases", - lambda requested_tag, published_repo, published_release_tag = "": iter(releases), + lambda requested_tag, published_repo, published_release_tag = "", **_kwargs: iter( + releases + ), ) requested_tag, plans = _fork_manifest_release_plans( @@ -2190,7 +2200,9 @@ class TestResolveInstallReleasePlans: monkeypatch.setattr( INSTALL_LLAMA_PREBUILT, "iter_resolved_published_releases", - lambda requested_tag, published_repo, published_release_tag = "": iter(releases), + lambda requested_tag, published_repo, published_release_tag = "", **_kwargs: iter( + releases + ), ) _requested_tag, plans = _fork_manifest_release_plans(