#!/usr/bin/env python3 # SPDX-License-Identifier: AGPL-3.0-only # Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0 """Cross platform llama.cpp prebuilt installer for Unsloth Studio""" from __future__ import annotations import argparse import atexit import errno import fnmatch import hashlib import json import os import platform import random import re import shutil import site import socket import struct import subprocess import sys import tarfile import tempfile import textwrap import time import urllib.error import urllib.parse import urllib.request import zipfile from contextlib import contextmanager from dataclasses import dataclass, field, replace as dataclasses_replace try: from filelock import FileLock, Timeout as FileLockTimeout except ImportError: FileLock = None FileLockTimeout = None from pathlib import Path from typing import Any, Iterable, Iterator EXIT_SUCCESS = 0 EXIT_FALLBACK = 2 EXIT_ERROR = 1 EXIT_BUSY = 3 # DiskPart-prompt suppression. RunAsInvoker does NOT stop amd-smi's runtime # elevation (its manifest is asInvoker), so this is just harmless belt-and- # suspenders for manifest-elevating tools. The real guard is _amd_smi_allowed(): # we don't spawn amd-smi on Windows w/o a HIP SDK (or opt-in). if platform.system() == "Windows": os.environ.setdefault("__COMPAT_LAYER", "RunAsInvoker") def _path_inside_venv(path: str) -> bool: """True if ``path`` is inside the active venv (sys.prefix). The venv hipInfo.exe (AMD wheel, put on PATH by the bnb fix) is NOT a HIP SDK (_amd_smi_allowed).""" try: # realpath (not abspath): resolve symlinks/8.3 names so an aliased venv matches. _root = os.path.normcase(os.path.realpath(sys.prefix)) # Guard a root-dir prefix (C:\ or /): commonpath would match every path on # it. A venv is never at root, so treat that as outside. if os.path.dirname(_root) == _root: return False return os.path.normcase(os.path.commonpath([os.path.realpath(path), _root])) == _root except (ValueError, OSError): # Different drive / unresolvable -> treat as outside the venv. return False def _external_hipinfo_on_path() -> bool: """True if a hipinfo OUTSIDE the venv is on PATH. shutil.which returns only the first hit, so the venv hipInfo could shadow a real HIP SDK's; scan every PATH entry and skip the venv copy.""" for _dir in os.environ.get("PATH", "").split(os.pathsep): _dir = _dir.strip('"') # PATH entries can be quoted on Windows if not _dir: continue _candidate = os.path.join(_dir, "hipinfo.exe") if os.path.isfile(_candidate) and not _path_inside_venv(_candidate): return True return False def _amd_smi_allowed() -> bool: """Whether it is safe to spawn amd-smi here. On Windows w/o a working HIP runtime, amd-smi elevates a child and pops a UAC/DiskPart prompt RunAsInvoker can't suppress. Only call it on Windows when a HIP SDK is detectable (hipinfo present) or UNSLOTH_ENABLE_AMD_SMI=1; Linux/macOS always allowed. When skipped, the gfx arch still arrives via the forwarded --rocm-gfx, so prebuilt selection is unaffected. """ if platform.system() != "Windows": return True flag = os.environ.get("UNSLOTH_ENABLE_AMD_SMI", "").strip().lower() if flag in ("1", "true", "yes", "on"): return True if flag in ("0", "false", "no", "off"): return False # A real HIP SDK lets amd-smi run un-elevated; hipinfo-on-PATH is the proxy. # Ignore the venv hipInfo.exe (AMD wheel via bnb fix): not a HIP SDK, doesn't # stop amd-smi's DiskPart UAC. if _external_hipinfo_on_path(): return True for _var in ("HIP_PATH", "HIP_PATH_57", "ROCM_PATH"): _root = os.environ.get(_var) if not _root: continue _candidate = os.path.join(_root, "bin", "hipinfo.exe") if os.path.isfile(_candidate) and not _path_inside_venv(_candidate): return True return False def windows_hidden_subprocess_kwargs() -> dict[str, object]: """Return Windows-only subprocess kwargs that suppress console windows.""" if sys.platform != "win32": return {} kwargs: dict[str, object] = {} create_no_window = getattr(subprocess, "CREATE_NO_WINDOW", 0) if create_no_window: kwargs["creationflags"] = create_no_window startupinfo_factory = getattr(subprocess, "STARTUPINFO", None) startf_use_showwindow = getattr(subprocess, "STARTF_USESHOWWINDOW", 0) sw_hide = getattr(subprocess, "SW_HIDE", 0) if startupinfo_factory is not None and startf_use_showwindow: startupinfo = startupinfo_factory() startupinfo.dwFlags |= startf_use_showwindow startupinfo.wShowWindow = sw_hide kwargs["startupinfo"] = startupinfo return kwargs def env_int( name: str, default: int, *, minimum: int | None = None, ) -> int: raw = os.environ.get(name) if raw is None: value = default else: try: value = int(str(raw).strip()) except (TypeError, ValueError): value = default if minimum is not None: value = max(minimum, value) return value # Prefer "latest" over "master" -- "master" bypasses the prebuilt resolver # (no matching GitHub release), forces a source build, and causes HTTP 422 # errors. Only use "master" temporarily when the latest release is missing # support for a new model architecture. DEFAULT_LLAMA_TAG = os.environ.get("UNSLOTH_LLAMA_TAG", "latest") # Default published repo for prebuilt release resolution. Linux uses # Unsloth prebuilts; setup.sh/setup.ps1 pass --published-repo explicitly # for macOS/Windows to override with ggml-org/llama.cpp when needed. DEFAULT_PUBLISHED_REPO = "unslothai/llama.cpp" DEFAULT_PUBLISHED_TAG = os.environ.get("UNSLOTH_LLAMA_RELEASE_TAG") DEFAULT_PUBLISHED_MANIFEST_ASSET = os.environ.get( "UNSLOTH_LLAMA_RELEASE_MANIFEST_ASSET", "llama-prebuilt-manifest.json" ) DEFAULT_PUBLISHED_SHA256_ASSET = os.environ.get( "UNSLOTH_LLAMA_RELEASE_SHA256_ASSET", "llama-prebuilt-sha256.json" ) UPSTREAM_REPO = "ggml-org/llama.cpp" UPSTREAM_RELEASES_API = f"https://api.github.com/repos/{UPSTREAM_REPO}/releases/latest" TEST_MODEL_URL = "https://huggingface.co/ggml-org/models/resolve/main/tinyllamas/stories260K.gguf" TEST_MODEL_SHA256 = "270cba1bd5109f42d03350f60406024560464db173c0e387d91f0426d3bd256d" VALIDATION_MODEL_CACHE_DIRNAME = ".cache" VALIDATION_MODEL_CACHE_FILENAME = "stories260K.gguf" # Master switch for the staged runtime smoke test (llama-quantize + llama-server) # in validate_prebuilt_choice. Disabled for now: the llama-server GPU forward pass # JIT-compiles CUDA kernels on first load and stalls every install and update by # minutes on Blackwell (sm_100). The check and the source-build fallback it triggers # are kept intact -- set this to True to re-enable them. _RUN_STAGED_PREBUILT_VALIDATION = False INSTALL_LOCK_TIMEOUT_SECONDS = 300 INSTALL_STAGING_ROOT_NAME = ".staging" GITHUB_AUTH_HOSTS = {"api.github.com", "github.com"} HF_AUTH_HOSTS = {"huggingface.co", "www.huggingface.co"} RETRYABLE_HTTP_STATUS = {408, 429, 500, 502, 503, 504} HTTP_FETCH_ATTEMPTS = 4 HTTP_FETCH_BASE_DELAY_SECONDS = 0.75 JSON_FETCH_ATTEMPTS = 3 DEFAULT_GITHUB_RELEASE_SCAN_MAX_PAGES = env_int( "UNSLOTH_LLAMA_GITHUB_RELEASE_SCAN_MAX_PAGES", 5, minimum = 1, ) SERVER_PORT_BIND_ATTEMPTS = 3 SERVER_BIND_RETRY_WINDOW_SECONDS = 5.0 TTY_PROGRESS_START_DELAY_SECONDS = 0.5 DEFAULT_MAX_PREBUILT_RELEASE_FALLBACKS = env_int( "UNSLOTH_LLAMA_MAX_PREBUILT_RELEASE_FALLBACKS", 2, minimum = 1, ) # Deeper macOS-only walk-back: upstream can ship a run of prebuilts built for a # newer macOS than the host, only caught at validate time, so an older host must # skip the whole run. Free on new hosts (first plan validates, extras unused). DEFAULT_MAX_MACOS_RELEASE_FALLBACKS = env_int( "UNSLOTH_LLAMA_MAX_MACOS_RELEASE_FALLBACKS", 16, minimum = 1, ) # Last upstream macOS release before ggml-org moved macOS builds to Tahoe. _PINNED_MACOS_FALLBACK_TAG = "b9415" _PINNED_MACOS_LATEST_FLOOR = (26, 0) FORCE_COMPILE_DEFAULT_REF = os.environ.get("UNSLOTH_LLAMA_FORCE_COMPILE_REF", "master") # Lowest CUDA major we ship prebuilts for, and the highest major we probe for # installed runtime libraries. Detection and runtime-line derivation are # generated per major so a new toolkit (cuda14, ...) needs no code change while # llama.cpp keeps the cudart64_.dll / libcudart.so. naming. _MIN_CUDA_MAJOR = 12 _MAX_PROBE_CUDA_MAJOR = 19 # Blackwell floor is sm_100: data-center parts (B100/B200 sm_100, B300/GB300 # sm_103) sit below consumer Blackwell (RTX 50 sm_120); the family needs toolkit # >= 12.8, except sm_103/sm_121 which need 12.9. (120 here wrongly excluded the # sm_100/103 data-center hosts.) _BLACKWELL_MIN_SM = 100 _BLACKWELL_MIN_TOOLKIT = (12, 8) # SMs that need a newer toolkit than the family floor (CUDA 12.9 added native # sm_103/sm_121 targets; 12.8 covers sm_100/101/120). _BLACKWELL_SM_MIN_TOOLKIT = {103: (12, 9), 121: (12, 9)} def _cuda_runtime_lines_for_major(major: int) -> list[str]: """Runtime lines a driver of this CUDA major can use, newest major first down to the minimum we ship. A driver runs its own major and any older one (backward compatibility).""" return [f"cuda{m}" for m in range(major, _MIN_CUDA_MAJOR - 1, -1)] @dataclass class HostInfo: system: str machine: str is_windows: bool is_linux: bool is_macos: bool is_x86_64: bool is_arm64: bool nvidia_smi: str | None driver_cuda_version: tuple[int, int] | None compute_caps: list[str] visible_cuda_devices: str | None has_physical_nvidia: bool has_usable_nvidia: bool has_rocm: bool = False rocm_gfx_target: str | None = None # (major, minor) from platform.mac_ver(); None off macOS or if unparseable. # Skips a macos prebuilt whose minimum-OS exceeds this host. macos_version: tuple[int, int] | None = None @dataclass class AssetChoice: repo: str tag: str name: str url: str source_label: str # Paired runtime archive (Windows CUDA cudart bundle). When set, # install_from_archives also downloads it and overlays its DLLs on # top of the main install. See unslothai/unsloth#5106. runtime_name: str | None = None runtime_url: str | None = None runtime_sha256: str | None = None is_ready_bundle: bool = False install_kind: str = "" bundle_profile: str | None = None runtime_line: str | None = None coverage_class: str | None = None supported_sms: list[str] | None = None min_sm: int | None = None max_sm: int | None = None selection_log: list[str] | None = None expected_sha256: str | None = None @dataclass(frozen = True) class PublishedLlamaArtifact: asset_name: str install_kind: str runtime_line: str | None coverage_class: str | None supported_sms: list[str] min_sm: int | None max_sm: int | None bundle_profile: str | None rank: int # ROCm bundles only: the umbrella gfx target (e.g. "gfx110X") and the # concrete gfx archs it covers (e.g. ["gfx1100", "gfx1101", ...]). gfx_target: str | None = None mapped_targets: list[str] = field(default_factory = list) @dataclass class PublishedReleaseBundle: repo: str release_tag: str upstream_tag: str manifest_sha256: str | None = None source_repo: str | None = None source_repo_url: str | None = None source_ref_kind: str | None = None requested_source_ref: str | None = None resolved_source_ref: str | None = None source_commit: str | None = None source_commit_short: str | None = None assets: dict[str, str] = field(default_factory = dict) manifest_asset_name: str = DEFAULT_PUBLISHED_MANIFEST_ASSET artifacts: list[PublishedLlamaArtifact] = field(default_factory = list) selection_log: list[str] = field(default_factory = list) @dataclass class LinuxCudaSelection: attempts: list[AssetChoice] selection_log: list[str] @property def primary(self) -> AssetChoice: if not self.attempts: raise RuntimeError("linux CUDA selection unexpectedly had no attempts") return self.attempts[0] @dataclass class CudaRuntimePreference: runtime_line: str | None selection_log: list[str] @dataclass(frozen = True) class ApprovedArtifactHash: asset_name: str sha256: str repo: str | None kind: str | None @dataclass class ApprovedReleaseChecksums: repo: str release_tag: str upstream_tag: str source_repo: str | None = None source_repo_url: str | None = None source_ref_kind: str | None = None requested_source_ref: str | None = None resolved_source_ref: str | None = None source_commit: str | None = None source_commit_short: str | None = None artifacts: dict[str, ApprovedArtifactHash] = field(default_factory = dict) @dataclass(frozen = True) class ResolvedPublishedRelease: bundle: PublishedReleaseBundle checksums: ApprovedReleaseChecksums @dataclass(frozen = True) class SourceBuildPlan: source_url: str source_ref: str source_ref_kind: str compatibility_upstream_tag: str source_repo: str | None = None source_repo_url: str | None = None requested_source_ref: str | None = None resolved_source_ref: str | None = None source_commit: str | None = None @dataclass(frozen = True) class InstallReleasePlan: requested_tag: str llama_tag: str release_tag: str attempts: list[AssetChoice] approved_checksums: ApprovedReleaseChecksums class PrebuiltFallback(RuntimeError): pass class BusyInstallConflict(RuntimeError): pass class ExistingInstallSatisfied(RuntimeError): def __init__(self, choice: AssetChoice, used_fallback: bool): super().__init__(f"existing install already matches candidate {choice.name}") self.choice = choice self.used_fallback = used_fallback def _os_error_messages(exc: BaseException) -> list[str]: messages: list[str] = [] if isinstance(exc, OSError): for value in ( getattr(exc, "strerror", None), getattr(exc, "filename", None), getattr(exc, "filename2", None), ): if isinstance(value, str) and value: messages.append(value) text = str(exc) if text: messages.append(text) return [message.lower() for message in messages if message] def is_busy_lock_error(exc: BaseException) -> bool: if isinstance(exc, BusyInstallConflict): return True if isinstance(exc, OSError): if exc.errno in { errno.EACCES, errno.EBUSY, errno.EPERM, errno.ETXTBSY, }: return True if getattr(exc, "winerror", None) in {5, 32, 145}: return True for message in _os_error_messages(exc): if any( needle in message for needle in ( "access is denied", "being used by another process", "device or resource busy", "permission denied", "text file busy", "file is in use", "process cannot access the file", "cannot create a file when that file already exists", ) ): return True return False # Status logs default to stderr so resolver modes keep stdout machine-readable # (setup.sh json.load()s the whole stdout). main() flips this for the install # path, where PowerShell otherwise renders stderr as NativeCommandError noise. _LOG_TO_STDOUT = False def log(message: str) -> None: print(f"[llama-prebuilt] {message}", file = sys.stdout if _LOG_TO_STDOUT else sys.stderr) def log_lines(lines: Iterable[str]) -> None: for line in lines: log(line) def parsed_hostname(url: str | None) -> str | None: if not url: return None try: hostname = urllib.parse.urlparse(url).hostname except Exception: return None if not hostname: return None return hostname.lower() def should_send_github_auth(url: str | None) -> bool: return parsed_hostname(url) in GITHUB_AUTH_HOSTS def should_send_hf_auth(url: str | None) -> bool: return parsed_hostname(url) in HF_AUTH_HOSTS def auth_headers(url: str | None = None) -> dict[str, str]: headers = { "User-Agent": "unsloth-studio-llama-prebuilt", } token = os.environ.get("GH_TOKEN") or os.environ.get("GITHUB_TOKEN") if token and should_send_github_auth(url): headers["Authorization"] = f"Bearer {token}" return headers # Anonymous huggingface.co fetches share a per-IP rate limit that CI # fleets exhaust (HTTP 429), sinking the prebuilt path into a source # build. Authenticate when a token is available. hf_token = os.environ.get("HF_TOKEN") or os.environ.get("HUGGING_FACE_HUB_TOKEN") if hf_token and should_send_hf_auth(url): headers["Authorization"] = f"Bearer {hf_token}" return headers class _CrossHostAuthStrippingRedirectHandler(urllib.request.HTTPRedirectHandler): """Drop Authorization when a redirect leaves the original host. huggingface.co redirects file downloads to CDN hosts whose signed URLs can reject foreign Authorization headers; urllib forwards headers to redirect targets by default (requests/huggingface_hub strip them). """ def redirect_request(self, req, fp, code, msg, headers, newurl): new_request = super().redirect_request(req, fp, code, msg, headers, newurl) if new_request is not None and parsed_hostname(newurl) != parsed_hostname(req.full_url): new_request.headers.pop("Authorization", None) new_request.unredirected_hdrs.pop("Authorization", None) return new_request _URL_OPENER = urllib.request.build_opener(_CrossHostAuthStrippingRedirectHandler()) def github_api_headers(url: str | None = None) -> dict[str, str]: return { "Accept": "application/vnd.github+json", **auth_headers(url), } def is_github_api_url(url: str | None) -> bool: return parsed_hostname(url) == "api.github.com" def is_retryable_url_error(exc: Exception) -> bool: if isinstance(exc, urllib.error.HTTPError): # GitHub returns 403 (not the standard 429) when the API rate # limit is hit. Anonymous calls share a 60-req/hour bucket per # runner IP, which CI fleets can exhaust trivially. Treat 403 # against api.github.com as retryable so we get one or two # backoff cycles before the source-build fallback fires; honour # Retry-After / X-RateLimit-Reset in sleep_backoff for accurate # waits. Real 403s on other hosts (private artefact downloads, # auth failures) stay non-retryable. if exc.code == 403: return is_github_api_url(getattr(exc, "url", None)) return exc.code in RETRYABLE_HTTP_STATUS if isinstance(exc, urllib.error.URLError): return True if isinstance(exc, TimeoutError): return True if isinstance(exc, socket.timeout): return True return False _RATE_LIMIT_WAIT_CAP_SECONDS = 60.0 def _http_error_retry_delay(exc: Exception) -> float | None: """Extract a recommended wait from rate-limit headers on a 403/429. Returns None when no header is present or the indicated wait is longer than _RATE_LIMIT_WAIT_CAP_SECONDS (in which case the caller should not block on it -- the source-build fallback is faster). """ if not isinstance(exc, urllib.error.HTTPError): return None headers = getattr(exc, "headers", None) if headers is None: return None retry_after = headers.get("Retry-After") if retry_after and retry_after.strip().isdigit(): wait = float(retry_after.strip()) return wait if wait <= _RATE_LIMIT_WAIT_CAP_SECONDS else None rate_reset = headers.get("X-RateLimit-Reset") if rate_reset and rate_reset.strip().isdigit(): wait = float(rate_reset.strip()) - time.time() if 0.0 < wait <= _RATE_LIMIT_WAIT_CAP_SECONDS: return wait + 1.0 # +1s of slack so the bucket is fresh return None def sleep_backoff( attempt: int, *, base_delay: float = HTTP_FETCH_BASE_DELAY_SECONDS, exc: Exception | None = None, ) -> None: delay = base_delay * (2 ** max(attempt - 1, 0)) header_delay = _http_error_retry_delay(exc) if exc is not None else None if header_delay is not None: delay = max(delay, header_delay) delay += random.uniform(0.0, 0.2) time.sleep(delay) def atomic_write_bytes(destination: Path, data: bytes) -> None: destination.parent.mkdir(parents = True, exist_ok = True) with tempfile.NamedTemporaryFile( prefix = destination.name + ".tmp-", dir = destination.parent, delete = False, ) as handle: tmp_path = Path(handle.name) handle.write(data) handle.flush() os.fsync(handle.fileno()) os.replace(tmp_path, destination) def atomic_replace_from_tempfile(tmp_path: Path, destination: Path) -> None: destination.parent.mkdir(parents = True, exist_ok = True) os.replace(tmp_path, destination) def source_archive_logical_name(upstream_tag: str) -> str: return f"llama.cpp-source-{upstream_tag}.tar.gz" def exact_source_archive_logical_name(source_commit: str) -> str: return f"llama.cpp-source-commit-{source_commit}.tar.gz" def sha256_file(path: Path) -> str: digest = hashlib.sha256() with path.open("rb") as handle: for chunk in iter(lambda: handle.read(1024 * 1024), b""): digest.update(chunk) return digest.hexdigest() def sha256_bytes(data: bytes) -> str: return hashlib.sha256(data).hexdigest() def normalize_sha256_digest(value: str | None) -> str | None: if not isinstance(value, str) or not value: return None lowered = value.lower() if lowered.startswith("sha256:"): lowered = lowered.split(":", 1)[1] if len(lowered) != 64 or any(ch not in "0123456789abcdef" for ch in lowered): return None return lowered def normalize_source_ref_kind(value: str | None) -> str | None: if not isinstance(value, str): return None normalized = value.strip().lower() if normalized in {"tag", "branch", "pull", "commit", "custom"}: return normalized return None def normalize_source_commit(value: str | None) -> str | None: if not isinstance(value, str): return None normalized = value.strip().lower() if len(normalized) < 7 or len(normalized) > 40: return None if any(ch not in "0123456789abcdef" for ch in normalized): return None return normalized def validate_schema_version(payload: dict[str, Any], *, label: str) -> None: schema_version = payload.get("schema_version") if schema_version is None: return try: normalized = int(schema_version) except (TypeError, ValueError) as exc: raise RuntimeError(f"{label} schema_version was not an integer") from exc if normalized != 1: raise RuntimeError(f"{label} schema_version={normalized} is unsupported") def repo_slug_from_source(value: str | None) -> str | None: if not isinstance(value, str): return None normalized = value.strip() if not normalized: return None normalized = normalized.removesuffix(".git") if normalized.startswith("https://github.com/"): slug = normalized[len("https://github.com/") :] elif normalized.startswith("http://github.com/"): slug = normalized[len("http://github.com/") :] elif normalized.startswith("git@github.com:"): slug = normalized[len("git@github.com:") :] else: slug = normalized slug = slug.strip("/") parts = slug.split("/") if len(parts) != 2 or not all(parts): return None return f"{parts[0]}/{parts[1]}" def source_url_from_repo_slug(repo_slug: str | None) -> str | None: if not isinstance(repo_slug, str) or not repo_slug: return None return f"https://github.com/{repo_slug}" def source_repo_clone_url(repo: str | None, repo_url: str | None) -> str | None: if isinstance(repo_url, str) and repo_url.strip(): return repo_url.strip().removesuffix(".git") return source_url_from_repo_slug(repo_slug_from_source(repo)) def infer_source_ref_kind(ref: str | None) -> str: if not isinstance(ref, str): return "tag" normalized = ref.strip() lowered = normalized.lower() if not normalized: return "tag" if lowered.startswith("refs/pull/") or lowered.startswith("pull/"): return "pull" if ( lowered.startswith("refs/heads/") or lowered in {"main", "master", "head"} or lowered.startswith("origin/") ): return "branch" normalized_commit = normalize_source_commit(normalized) if normalized_commit is not None: return "commit" return "tag" def normalized_ref_aliases(ref: str | None) -> set[str]: if not isinstance(ref, str): return set() normalized = ref.strip() if not normalized: return set() aliases = {normalized} lowered = normalized.lower() commit = normalize_source_commit(normalized) if commit is not None: aliases.add(commit) if lowered.startswith("refs/heads/"): aliases.add(normalized.split("/", 2)[2]) elif "/" not in normalized and infer_source_ref_kind(normalized) == "branch": aliases.add(f"refs/heads/{normalized}") if lowered.startswith("refs/pull/"): aliases.add(normalized.removeprefix("refs/")) elif lowered.startswith("pull/"): aliases.add(f"refs/{normalized}") return aliases def refs_match(candidate_ref: str | None, requested_ref: str | None) -> bool: candidate_aliases = normalized_ref_aliases(candidate_ref) requested_aliases = normalized_ref_aliases(requested_ref) if not candidate_aliases or not requested_aliases: return False if candidate_aliases & requested_aliases: return True candidate_commit = normalize_source_commit(candidate_ref) requested_commit = normalize_source_commit(requested_ref) if candidate_commit and requested_commit: return candidate_commit.startswith(requested_commit) or requested_commit.startswith( candidate_commit ) return False def checkout_friendly_ref(ref_kind: str | None, ref: str | None) -> str | None: """Normalize a source ref to a form that ``git clone --branch`` accepts. Fully qualified branch refs like ``refs/heads/main`` are stripped to ``main``; tag refs like ``refs/tags/b8508`` are stripped to ``b8508``. Pull refs like ``refs/pull/123/head`` are left as-is since they are always fetched explicitly rather than cloned with ``--branch``. """ if not isinstance(ref, str) or not ref: return ref lowered = ref.lower() if ref_kind == "branch" and lowered.startswith("refs/heads/"): return ref.split("/", 2)[2] if ref_kind == "tag" and lowered.startswith("refs/tags/"): return ref.split("/", 2)[2] return ref def windows_cuda_upstream_asset_names(llama_tag: str, runtime: str) -> list[str]: return [ f"llama-{llama_tag}-bin-win-cuda-{runtime}-x64.zip", f"cudart-llama-bin-win-cuda-{runtime}-x64.zip", ] def windows_cuda_asset_aliases( asset_name: str, *, compatibility_tag: str | None = None ) -> list[str]: aliases: list[str] = [] legacy_match = re.fullmatch( r"llama-(?P[^/]+)-bin-win-cuda-(?P\d+\.\d+)-x64\.zip", asset_name, ) if legacy_match: runtime = legacy_match.group("runtime") aliases.append(f"cudart-llama-bin-win-cuda-{runtime}-x64.zip") if compatibility_tag: aliases.append(f"llama-{compatibility_tag}-bin-win-cuda-{runtime}-x64.zip") return aliases current_match = re.fullmatch( r"cudart-llama-bin-win-cuda-(?P\d+\.\d+)-x64\.zip", asset_name, ) if current_match and compatibility_tag: runtime = current_match.group("runtime") aliases.append(f"llama-{compatibility_tag}-bin-win-cuda-{runtime}-x64.zip") return aliases def _published_windows_cuda_runtime( upstream_assets: dict[str, str], major: int, driver: tuple[int, int] | None ) -> str | None: """Highest cuda-. published upstream that `driver` can run by default CUDA compatibility, i.e. (major, minor) <= driver. None if nothing qualifies. Gating on the driver (not just the major) keeps a 13.3 build off a driver that only advertises 13.1, where it would otherwise rely on the unguaranteed minor-version-compatibility path.""" if driver is None: return None best: int | None = None for name in upstream_assets: m = re.search(r"-bin-win-cuda-(\d+)\.(\d+)-x64\.zip$", name) if m and int(m.group(1)) == major: minor = int(m.group(2)) if (major, minor) <= driver and (best is None or minor > best): best = minor return f"{major}.{best}" if best is not None else None def format_byte_count(num_bytes: float) -> str: units = ["B", "KiB", "MiB", "GiB", "TiB"] value = float(num_bytes) for unit in units: if abs(value) < 1024.0 or unit == units[-1]: if unit == "B": return f"{int(value)} {unit}" return f"{value:.1f} {unit}" value /= 1024.0 return f"{num_bytes:.1f} B" def _progress_percent_step() -> int: """Non-tty milestone granularity. The in-app updater sets UNSLOTH_PROGRESS_PERCENT_STEP=5 to stream finer progress lines.""" try: step = int(os.environ.get("UNSLOTH_PROGRESS_PERCENT_STEP", "25")) except ValueError: return 25 return min(max(step, 1), 50) class DownloadProgress: def __init__(self, label: str, total_bytes: int | None) -> None: self.label = label self.total_bytes = total_bytes if total_bytes and total_bytes > 0 else None self.start_time = time.monotonic() self.last_emit = 0.0 term_ok = os.environ.get("TERM", "").lower() != "dumb" self.stream = ( sys.stderr if sys.stderr.isatty() else sys.stdout if sys.stdout.isatty() else sys.stderr ) self.is_tty = term_ok and self.stream.isatty() self.completed = False self.milestone_step = _progress_percent_step() self.last_milestone_percent = -1 self.last_milestone_bytes = 0 self.has_rendered_tty_progress = False def _render( self, downloaded_bytes: int, *, final: bool = False, ) -> str: elapsed = max(time.monotonic() - self.start_time, 1e-6) speed = downloaded_bytes / elapsed speed_text = f"{format_byte_count(speed)}/s" if self.total_bytes is not None: percent = min(100.0, (downloaded_bytes / self.total_bytes) * 100.0) return ( f"{self.label}: {percent:5.1f}% " f"({format_byte_count(downloaded_bytes)}/{format_byte_count(self.total_bytes)}) " f"at {speed_text}" ) if final: return f"{self.label}: {format_byte_count(downloaded_bytes)} downloaded at {speed_text}" return f"{self.label}: {format_byte_count(downloaded_bytes)} downloaded at {speed_text}" def update(self, downloaded_bytes: int) -> None: now = time.monotonic() if self.is_tty: elapsed = now - self.start_time if not self.has_rendered_tty_progress: if self.total_bytes is not None and downloaded_bytes >= self.total_bytes: return if elapsed < TTY_PROGRESS_START_DELAY_SECONDS: return min_interval = 0.2 if ( self.has_rendered_tty_progress and not self.completed and (now - self.last_emit) < min_interval ): return self.last_emit = now line = self._render(downloaded_bytes) self.stream.write("\r\033[K" + line) self.stream.flush() self.has_rendered_tty_progress = True return should_emit = False if self.total_bytes is not None: percent = int((downloaded_bytes * 100) / max(self.total_bytes, 1)) step = self.milestone_step milestone_percent = min((percent // step) * step, 100) if milestone_percent > self.last_milestone_percent and milestone_percent < 100: self.last_milestone_percent = milestone_percent should_emit = True else: byte_step = 25 * 1024 * 1024 if ( downloaded_bytes - self.last_milestone_bytes >= byte_step and (now - self.last_emit) >= 5.0 ): self.last_milestone_bytes = downloaded_bytes should_emit = True if not should_emit: return self.last_emit = now self.stream.write(self._render(downloaded_bytes) + "\n") self.stream.flush() def finish(self, downloaded_bytes: int) -> None: self.completed = True line = self._render(downloaded_bytes, final = True) if self.is_tty: if not self.has_rendered_tty_progress: return self.stream.write("\r\033[K") else: self.stream.write(line + "\n") self.stream.flush() def download_label_from_url(url: str) -> str: name = Path(urllib.parse.urlparse(url).path).name return name or url def download_bytes( url: str, *, timeout: int = 120, attempts: int = HTTP_FETCH_ATTEMPTS, headers: dict[str, str] | None = None, progress_label: str | None = None, ) -> bytes: last_exc: Exception | None = None for attempt in range(1, attempts + 1): try: request = urllib.request.Request(url, headers = headers or auth_headers(url)) with _URL_OPENER.open(request, timeout = timeout) as response: total_bytes: int | None = None content_length = response.headers.get("Content-Length") if content_length and content_length.isdigit(): total_bytes = int(content_length) progress = DownloadProgress(progress_label, total_bytes) if progress_label else None data = bytearray() while True: chunk = response.read(1024 * 1024) if not chunk: break data.extend(chunk) if progress is not None: progress.update(len(data)) if progress is not None: progress.finish(len(data)) return bytes(data) except Exception as exc: last_exc = exc if attempt >= attempts or not is_retryable_url_error(exc): raise log(f"fetch failed ({attempt}/{attempts}) for {url}: {exc}; retrying") sleep_backoff(attempt, exc = exc) assert last_exc is not None raise last_exc def fetch_json(url: str) -> Any: attempts = JSON_FETCH_ATTEMPTS if is_github_api_url(url) else 1 last_decode_exc: Exception | None = None for attempt in range(1, attempts + 1): try: data = download_bytes( url, timeout = 30, headers = github_api_headers(url) if is_github_api_url(url) else auth_headers(url), ) except urllib.error.HTTPError as exc: if exc.code == 403 and is_github_api_url(url): hint = "" if not (os.environ.get("GH_TOKEN") or os.environ.get("GITHUB_TOKEN")): hint = "; set GH_TOKEN or GITHUB_TOKEN to avoid GitHub API rate limits" raise RuntimeError(f"GitHub API returned 403 for {url}{hint}") from exc raise if not data: last_decode_exc = RuntimeError(f"downloaded empty JSON payload from {url}") else: try: payload = json.loads(data.decode("utf-8")) except (UnicodeDecodeError, json.JSONDecodeError) as exc: last_decode_exc = RuntimeError(f"downloaded invalid JSON from {url}: {exc}") else: if not isinstance(payload, dict) and not isinstance(payload, list): raise RuntimeError( f"downloaded unexpected JSON type from {url}: {type(payload).__name__}" ) return payload if attempt >= attempts: assert last_decode_exc is not None raise last_decode_exc log(f"json fetch failed ({attempt}/{attempts}) for {url}; retrying") sleep_backoff(attempt) assert last_decode_exc is not None raise last_decode_exc def download_file(url: str, destination: Path) -> None: destination.parent.mkdir(parents = True, exist_ok = True) last_exc: Exception | None = None for attempt in range(1, HTTP_FETCH_ATTEMPTS + 1): tmp_path: Path | None = None try: request = urllib.request.Request(url, headers = auth_headers(url)) with tempfile.NamedTemporaryFile( prefix = destination.name + ".tmp-", dir = destination.parent, delete = False, ) as handle: tmp_path = Path(handle.name) with _URL_OPENER.open(request, timeout = 120) as response: total_bytes: int | None = None content_length = response.headers.get("Content-Length") if content_length and content_length.isdigit(): total_bytes = int(content_length) progress = DownloadProgress(f"Downloading {destination.name}", total_bytes) downloaded_bytes = 0 while True: chunk = response.read(1024 * 1024) if not chunk: break handle.write(chunk) downloaded_bytes += len(chunk) progress.update(downloaded_bytes) progress.finish(downloaded_bytes) handle.flush() os.fsync(handle.fileno()) if not tmp_path.exists() or tmp_path.stat().st_size == 0: raise RuntimeError(f"downloaded empty file from {url}") atomic_replace_from_tempfile(tmp_path, destination) return except Exception as exc: last_exc = exc if tmp_path is not None: try: tmp_path.unlink(missing_ok = True) except Exception: pass if attempt >= HTTP_FETCH_ATTEMPTS or not is_retryable_url_error(exc): raise log(f"download failed ({attempt}/{HTTP_FETCH_ATTEMPTS}) for {url}: {exc}; retrying") sleep_backoff(attempt, exc = exc) assert last_exc is not None raise last_exc def download_file_verified( url: str, destination: Path, *, expected_sha256: str | None, label: str ) -> None: normalized_expected = normalize_sha256_digest(expected_sha256) if not normalized_expected: download_file(url, destination) log(f"downloaded {label} without a published sha256; relying on install validation") return for attempt in range(1, 3): download_file(url, destination) actual_sha256 = sha256_file(destination) if actual_sha256 == normalized_expected: log(f"verified {label} sha256={actual_sha256}") return log( f"{label} checksum mismatch on attempt {attempt}/2: " f"expected={normalized_expected} actual={actual_sha256}" ) destination.unlink(missing_ok = True) if attempt == 2: raise PrebuiltFallback( f"{label} checksum mismatch after retry: expected={normalized_expected} actual={actual_sha256}" ) log(f"retrying {label} download after checksum mismatch") def upstream_source_archive_urls(tag: str) -> list[str]: encoded_tag = urllib.parse.quote(tag, safe = "") return [ f"https://codeload.github.com/{UPSTREAM_REPO}/tar.gz/refs/tags/{encoded_tag}", f"https://github.com/{UPSTREAM_REPO}/archive/refs/tags/{encoded_tag}.tar.gz", ] def commit_source_archive_urls(repo: str, source_commit: str) -> list[str]: encoded_commit = urllib.parse.quote(source_commit, safe = "") return [ f"https://codeload.github.com/{repo}/tar.gz/{encoded_commit}", f"https://github.com/{repo}/archive/{encoded_commit}.tar.gz", ] def release_asset_download_url( repo: str | None, release_tag: str | None, asset_name: str | None ) -> str | None: """Direct download URL for a release asset, or None if any part is missing. A mix build's merged commit is never pushed, so its source tree is only reachable as this asset (codeload would 404 on the merge commit).""" if not repo or not release_tag or not asset_name: return None return ( f"https://github.com/{repo}/releases/download/" f"{urllib.parse.quote(release_tag, safe = '')}/{urllib.parse.quote(asset_name, safe = '')}" ) def github_release_assets(repo: str, tag: str) -> dict[str, str]: payload = fetch_json( f"https://api.github.com/repos/{repo}/releases/tags/{urllib.parse.quote(tag, safe = '')}" ) if not isinstance(payload, dict): raise RuntimeError(f"unexpected release payload for {repo}@{tag}") return release_asset_map(payload) def github_release(repo: str, tag: str) -> dict[str, Any]: payload = fetch_json( f"https://api.github.com/repos/{repo}/releases/tags/{urllib.parse.quote(tag, safe = '')}" ) if not isinstance(payload, dict): raise RuntimeError(f"unexpected release payload for {repo}@{tag}") return payload def github_releases( repo: str, *, per_page: int = 100, max_pages: int = 0, ) -> list[dict[str, Any]]: releases: list[dict[str, Any]] = [] page = 1 while True: payload = fetch_json( f"https://api.github.com/repos/{repo}/releases?per_page={per_page}&page={page}" ) if not isinstance(payload, list): raise RuntimeError(f"unexpected releases payload for {repo}") page_items = [item for item in payload if isinstance(item, dict)] releases.extend(page_items) if len(payload) < per_page: break page += 1 if max_pages > 0 and page > max_pages: break return releases def latest_upstream_release_tag() -> str: payload = fetch_json(UPSTREAM_RELEASES_API) tag = payload.get("tag_name") if not isinstance(tag, str) or not tag: raise RuntimeError(f"latest release tag was missing from {UPSTREAM_RELEASES_API}") return tag def is_release_tag_like(value: str | None) -> bool: return isinstance(value, str) and bool(re.fullmatch(r"b\d+", value.strip())) def release_time_sort_key(release: dict[str, Any]) -> tuple[str, int]: published_at = release.get("published_at") created_at = release.get("created_at") release_id = release.get("id") timestamp = ( published_at if isinstance(published_at, str) and published_at else created_at if isinstance(created_at, str) and created_at else "" ) try: normalized_id = int(release_id) except (TypeError, ValueError): normalized_id = 0 return (timestamp, normalized_id) def iter_release_payloads_by_time( repo: str, published_release_tag: str = "", requested_tag: str = "", ) -> Iterable[dict[str, Any]]: if published_release_tag: yield github_release(repo, published_release_tag) return if requested_tag and requested_tag != "latest" and is_release_tag_like(requested_tag): try: yield github_release(repo, requested_tag) return except urllib.error.HTTPError as exc: if exc.code == 404: log(f"release tag {requested_tag} not found in {repo}; scanning recent releases") else: raise except Exception: raise releases = [ release for release in github_releases(repo, max_pages = DEFAULT_GITHUB_RELEASE_SCAN_MAX_PAGES) if isinstance(release, dict) and not release.get("draft") and not release.get("prerelease") ] releases.sort(key = release_time_sort_key, reverse = True) for release in releases: yield release def direct_release_matches_request(*, release_tag: str, llama_tag: str, requested_tag: str) -> bool: if requested_tag == "latest": return True for candidate in (release_tag, llama_tag): if refs_match(candidate, requested_tag): return True return False def synthetic_checksums_for_release( repo: str, release_tag: str, upstream_tag: str ) -> ApprovedReleaseChecksums: return ApprovedReleaseChecksums( repo = repo, release_tag = release_tag, upstream_tag = upstream_tag, artifacts = {}, ) def parse_direct_linux_release_bundle( repo: str, release: dict[str, Any] ) -> PublishedReleaseBundle | None: release_tag = release.get("tag_name") if not isinstance(release_tag, str) or not release_tag: return None assets = release_asset_map(release) artifacts: list[PublishedLlamaArtifact] = [] inferred_labels: list[str] = [] linux_asset_re = re.compile( r"^app-(?P