mirror of
https://github.com/unslothai/unsloth.git
synced 2026-10-21 14:00:32 +00:00
## Summary
Four small detection fixes for hosts that were misreported, without
changing behaviour anywhere else.
## What changes
- **Studio hardware detection:**
- An x86_64 Python on Apple Silicon (Rosetta) now reports chat-only
reason `rosetta` instead of `intel_mac`.
- The check is `sysctl -n hw.optional.arm64`: bounded, cached, and it
never raises. It runs only on Darwin x86_64 and fails open to
`intel_mac`.
- The sidebar tells the user to re-run the installer from a native arm64
terminal, and MLX-only settings stay hidden.
- **audio.cpp installer:**
- The prebuilt CUDA bundles carry sm_75+ kernels only (`CUDA12_ARCHS` /
`CUDA13_ARCHS` in the audio.cpp prebuilt workflow).
- Auto-detection now reads `compute_cap` and installs the CPU build when
every NVIDIA card is older than 7.5.
- An unreadable compute capability keeps CUDA, and an explicit
accelerator still wins.
- The floor applies only to the default publisher: a custom
`UNSLOTH_AUDIO_CPP_REPO` release keeps CUDA, since its build list is
unknown.
- **`unsloth doctor`:**
- A torch that is neither a CUDA nor a ROCm build but can use MPS or XPU
is reported healthy ("multi-GPU CUDA checks skipped") instead of
"PROBLEMS FOUND".
- A CUDA or ROCm build that sees no device still fails, as before.
- **Token preview worker:**
- The child env also sets `ONEAPI_DEVICE_SELECTOR=opencl:cpu`, so
torch-xpu does not open a Level Zero GPU context. Only the child env
changes.
## Hosts affected
- Apple Silicon Macs running Studio under an x86_64 Python.
- GTX 10xx, P40, P100 and V100 hosts installing audio.cpp with
auto-detection.
- Apple Silicon and Intel XPU hosts running `unsloth doctor`.
- Intel XPU hosts using the token preview.
Native arm64 Macs, real Intel Macs, sm_75+ NVIDIA hosts, and CUDA/ROCm
hosts keep today's answers.
## Tests
- `studio/backend/tests/test_chat_only_reason.py`: Rosetta, Intel Mac, a
failing probe, native arm64.
- `tests/studio/install/test_install_audio_cpp_prebuilt_logic.py`:
pre-Turing, mixed, sm_75+, and unreadable or failing compute_cap
queries.
- `unsloth_cli/tests/test_doctor_gpus.py`: MPS, XPU, CUDA/ROCm without a
device, a CPU build, and a raising probe.
- `studio/backend/tests/test_token_preview.py`: the child env hides
CUDA, HIP and SYCL, and the parent env is untouched.
- `studio/frontend/tests/mlx-context-helpers.test.ts`: `rosetta` rules
out MLX.
The new tests fail on the merge base and pass here. The related suites
show no new failures.
## Screenshots
Studio sidebar with `/api/health` reporting `chat_only_reason:
"rosetta"` (injected into the real backend's response on both builds,
since this cannot be produced on a Linux host).
Train tooltip: before it only said "Train"; now it explains the Rosetta
install and the fix.

Video row: before it stayed enabled with no hint; now it is disabled
with the same reason and fix.

Export and Video pages show the backend's own message
(`export_capability()` / `video_capability()` in
`/api/system/hardware`). On a spoofed Rosetta host (Darwin, x86_64
process, `hw.optional.arm64` = 1, CPU torch):
| Page | Before | After |
|---|---|---|
| Export | `no_accelerator`: "Export requires an NVIDIA, AMD, or Intel
GPU, or Apple Silicon (MLX). No supported accelerator was found on this
host. (PyTorch is installed, but Unsloth cannot export on CPU only.)" |
`rosetta`: "Studio is running under Rosetta (x86_64 Python), so export
is disabled. Re-run the installer from a native arm64 terminal to enable
export." |
| Video | `macos_unsupported`: "Video generation requires Apple Silicon.
This Intel Mac has no Metal (MPS) device for the video pipelines to run
on." | `rosetta`: "Studio is running under Rosetta (x86_64 Python),
which has no Metal (MPS) device for the video pipelines. Re-run the
installer from a native arm64 terminal to enable video generation." |
1431 lines
59 KiB
Python
1431 lines
59 KiB
Python
# SPDX-License-Identifier: AGPL-3.0-only
|
|
# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0
|
|
|
|
"""Install a prebuilt ``audiocpp_server`` (audio.cpp) for Studio's audio.cpp engine.
|
|
|
|
audio.cpp vendors its own patched ggml, so unlike whisper.cpp it cannot share the
|
|
llama.cpp runtime: every bundle is self-contained (the server links ggml statically).
|
|
Releases are named ``audio-<tag>-bin-<os>-<arch>-<backend>.{zip,tar.gz}``, the scheme
|
|
upstream uses and the Unsloth fork keeps, so one resolver serves both. Windows CUDA
|
|
bundles need the CUDA runtime DLLs, which upstream publishes as a separate
|
|
``audio-<tag>-cudart-windows-x64-<line>.zip``; that download is skipped when the
|
|
running Python's torch already ships the same CUDA runtime line.
|
|
|
|
The bundle is extracted into a staging directory beside the managed tree, started once
|
|
with ``--help`` so a bundle this host cannot load is refused, and swapped in only then,
|
|
so an interrupted or unusable install never replaces a working tree.
|
|
|
|
The newest unslothai/audio.cpp release is installed unless ``UNSLOTH_AUDIO_CPP_TAG`` (or
|
|
``--published-release-tag``) names one. The release is resolved like whisper.cpp's, through
|
|
prebuilt_core: the download host first, then the GitHub API. Every bundle and CUDA runtime
|
|
archive must match the sha256 in the release's own ``audio-cpp-prebuilt-sha256.json``, which
|
|
is bound to the release tag; an archive the index does not cover is refused. That checksum
|
|
comes from the same release as the archive, so it proves integrity (a corrupt or swapped
|
|
download), not authenticity. A user-picked repo or tag whose release publishes no index
|
|
falls back to GitHub's asset digest.
|
|
|
|
A re-run whose install record already names the newest release (one HEAD to github.com,
|
|
no API call) or the pinned one (no network call) downloads nothing, an intact install newer
|
|
than the release a lookup names is never replaced by it, and a lookup that cannot answer
|
|
keeps a complete install.
|
|
|
|
Usage:
|
|
python studio/install_audio_cpp_prebuilt.py # auto-detect host
|
|
python studio/install_audio_cpp_prebuilt.py --accelerator cpu
|
|
python studio/install_audio_cpp_prebuilt.py --print-asset # resolve only
|
|
python studio/install_audio_cpp_prebuilt.py --published-release-tag v0.9.1-unsloth.1
|
|
"""
|
|
|
|
from __future__ import annotations
|
|
|
|
import argparse
|
|
import datetime
|
|
import fnmatch
|
|
import glob
|
|
import hashlib
|
|
import json
|
|
import os
|
|
import platform
|
|
import re
|
|
import shutil
|
|
import subprocess
|
|
import sys
|
|
import tarfile
|
|
import urllib.error
|
|
import urllib.request
|
|
import zipfile
|
|
from concurrent.futures import ThreadPoolExecutor
|
|
from pathlib import Path
|
|
from typing import Optional, Sequence
|
|
|
|
if __package__:
|
|
from . import prebuilt_core as core
|
|
from .install_sd_cpp_prebuilt import (
|
|
GitHubRateLimited,
|
|
_is_rate_limited,
|
|
_rate_limit_message,
|
|
_safe_extractall,
|
|
)
|
|
else:
|
|
_STUDIO_DIR = os.path.dirname(os.path.abspath(__file__))
|
|
if _STUDIO_DIR not in sys.path:
|
|
sys.path.insert(0, _STUDIO_DIR)
|
|
import prebuilt_core as core # type: ignore[no-redef]
|
|
from install_sd_cpp_prebuilt import ( # type: ignore[no-redef]
|
|
GitHubRateLimited,
|
|
_is_rate_limited,
|
|
_rate_limit_message,
|
|
_safe_extractall,
|
|
)
|
|
|
|
# The Unsloth fork publishes the bundles Studio installs (CI: .github/workflows/unsloth-prebuilt.yml
|
|
# there, like unslothai/whisper.cpp); UNSLOTH_AUDIO_CPP_REPO overrides it. Its tags carry the Unsloth
|
|
# packaging (static eSpeak-ng, lean multi-arch Linux CUDA) and a manifest + sha256 index per release.
|
|
DEFAULT_REPO = "unslothai/audio.cpp"
|
|
|
|
# What prebuilt_core's release lookup and index parsers read through _OPS.
|
|
COMPONENT = "audio.cpp"
|
|
MANIFEST_ASSET_NAME = "audio-cpp-prebuilt-manifest.json"
|
|
SHA256_ASSET_NAME = "audio-cpp-prebuilt-sha256.json"
|
|
|
|
INSTALL_RECORD = "UNSLOTH_AUDIO_CPP_PREBUILT_INFO.json"
|
|
OWNERSHIP_MARKER = ".unsloth-studio-owned"
|
|
SERVER_NAME = "audiocpp_server.exe" if sys.platform == "win32" else "audiocpp_server"
|
|
|
|
ACCELERATORS = ("auto", "cpu", "cuda", "vulkan", "metal")
|
|
# The staged server's --help must exit 0 within this long before it replaces the live tree.
|
|
SMOKE_TIMEOUT_SECONDS = 30.0
|
|
USER_AGENT = "unsloth-audio-cpp-installer"
|
|
|
|
EXIT_OK = 0
|
|
EXIT_FAILED = 1
|
|
EXIT_BUSY = 3
|
|
|
|
KEPT_EXISTING_LINE = "audio.cpp: keeping the existing complete install"
|
|
|
|
|
|
class ReleaseLookupUnavailable(RuntimeError):
|
|
"""No release lookup answered (offline, DNS, GitHub down): says nothing about the tree on disk."""
|
|
|
|
|
|
def log(message: str) -> None:
|
|
# prebuilt_core's download retries report through this.
|
|
print(f"audio.cpp: {message}", file = sys.stderr, flush = True)
|
|
|
|
|
|
# Late-binding seam for prebuilt_core: its downloader and release lookup read USER_AGENT, log and the
|
|
# asset names above from here, and see a test's monkeypatch of any core name set on this module.
|
|
_OPS = core.ModuleOps(globals())
|
|
|
|
_OS_TOKENS = {
|
|
"windows": ("windows", "win"),
|
|
"linux": ("ubuntu", "linux"),
|
|
"darwin": ("macos", "darwin"),
|
|
}
|
|
_ARCH_TOKENS = {
|
|
"x86_64": "x64",
|
|
"amd64": "x64",
|
|
"x64": "x64",
|
|
"arm64": "arm64",
|
|
"aarch64": "arm64",
|
|
}
|
|
# "-cuda12.8" (upstream) or "-cuda12" (unslothai/audio.cpp: one lean bundle per CUDA major line).
|
|
_CUDA_LINE = re.compile(r"-cuda(\d+)(?:\.(\d+))?(?=[-.]|$)")
|
|
|
|
|
|
def _repo() -> str:
|
|
return (os.environ.get("UNSLOTH_AUDIO_CPP_REPO") or "").strip() or DEFAULT_REPO
|
|
|
|
|
|
def _pinned_tag() -> Optional[str]:
|
|
"""``UNSLOTH_AUDIO_CPP_TAG``, the exact release to install; ``None`` (track the newest) when it
|
|
is unset, empty or ``latest``."""
|
|
val = (os.environ.get("UNSLOTH_AUDIO_CPP_TAG") or "").strip()
|
|
return None if not val or val.lower() == "latest" else val
|
|
|
|
|
|
def _user_picked_release() -> bool:
|
|
"""The user chose the repo or tag, so a release without a sha256 index is theirs to trust
|
|
on GitHub's digest."""
|
|
return (
|
|
bool((os.environ.get("UNSLOTH_AUDIO_CPP_REPO") or "").strip()) or _pinned_tag() is not None
|
|
)
|
|
|
|
|
|
# "v0.9.1-unsloth.1", "v0.9.0", "v0.8.2-audio8-perf-hotfix-unsloth.1": a leading version, any label.
|
|
_RELEASE_VERSION = re.compile(r"v?(\d+)\.(\d+)(?:\.(\d+))?(?:-\S*)?")
|
|
_UNSLOTH_SERIAL = re.compile(r"-unsloth\.(\d+)$")
|
|
|
|
|
|
def release_version_key(tag: Optional[str]) -> Optional[tuple[int, int, int, int]]:
|
|
"""``(major, minor, patch, unsloth_serial)`` of a release tag, or ``None``. A label between the
|
|
version and ``-unsloth.N`` is ignored, so a hotfix of an older line still sorts below a newer
|
|
release. Kept in step with ``audio_cpp_freshness.parse_release_version`` (a test compares them)."""
|
|
if not isinstance(tag, str):
|
|
return None
|
|
tag = tag.strip()
|
|
match = _RELEASE_VERSION.fullmatch(tag)
|
|
if not match:
|
|
return None
|
|
serial = _UNSLOTH_SERIAL.search(tag)
|
|
return (
|
|
int(match.group(1)),
|
|
int(match.group(2)),
|
|
int(match.group(3) or 0),
|
|
int(serial.group(1)) if serial else 0,
|
|
)
|
|
|
|
|
|
def _record_is_current(record_tag: Optional[str], latest_tag: Optional[str]) -> bool:
|
|
"""Whether an installed tag needs nothing from ``latest_tag``: the same release, or one that
|
|
parses strictly newer (``/releases/latest`` and the API's newest can name an older release)."""
|
|
if not record_tag or not latest_tag:
|
|
return False
|
|
if record_tag == latest_tag:
|
|
return True
|
|
installed, latest = release_version_key(record_tag), release_version_key(latest_tag)
|
|
return installed is not None and latest is not None and installed > latest
|
|
|
|
|
|
_GLIBC_MIN_OS = re.compile(r"glibc-(\d+)\.(\d+)")
|
|
|
|
|
|
def _glibc_below(min_os: Optional[str]) -> Optional[str]:
|
|
"""The host glibc when it is older than a Linux bundle's manifest ``min_os`` (``glibc-2.35``),
|
|
else ``None``. Checked before the download: without it the start check fails after a full
|
|
download (~1 GB for CUDA) on every setup run. No manifest (a user-picked release) skips it."""
|
|
if not sys.platform.startswith("linux") or not isinstance(min_os, str):
|
|
return None
|
|
floor = _GLIBC_MIN_OS.fullmatch(min_os.strip())
|
|
if not floor:
|
|
return None
|
|
name, version = platform.libc_ver()
|
|
try:
|
|
parsed = tuple(int(x) for x in version.split(".")[:2])
|
|
except ValueError:
|
|
return None
|
|
if name != "glibc" or len(parsed) != 2 or parsed >= (int(floor.group(1)), int(floor.group(2))):
|
|
return None
|
|
return version
|
|
|
|
|
|
def _accelerator_from_env() -> str:
|
|
"""``UNSLOTH_AUDIO_CPP_ACCELERATOR`` (auto|cpu|cuda|vulkan|metal), for hosts whose GPU is not visible
|
|
at install time (a Docker build). Set, it is an explicit request and never downgraded."""
|
|
raw = (os.environ.get("UNSLOTH_AUDIO_CPP_ACCELERATOR") or "").strip().lower()
|
|
if not raw:
|
|
return "auto"
|
|
if raw not in ACCELERATORS:
|
|
print(
|
|
f"audio.cpp: ignoring UNSLOTH_AUDIO_CPP_ACCELERATOR={raw!r} (expected one of {', '.join(ACCELERATORS)})",
|
|
file = sys.stderr,
|
|
flush = True,
|
|
)
|
|
return "auto"
|
|
return raw
|
|
|
|
|
|
def _expected_sha256(release: dict, asset: dict) -> Optional[str]:
|
|
"""The sha256 an asset must hash to: its entry in the release's sha256 index, which any digest
|
|
the manifest embeds must agree with. An asset the index does not cover raises
|
|
``ReleaseIntegrityError``. Only a user-picked release without an index (``checksums`` None)
|
|
falls back to GitHub's digest, if it reports one."""
|
|
checksums = release.get("checksums")
|
|
if checksums is None:
|
|
return core.normalize_sha256_digest(asset.get("digest"))
|
|
return core.expected_sha256_for(
|
|
_OPS, checksums, asset["name"], manifest_sha256 = asset.get("manifest_sha256")
|
|
)
|
|
|
|
|
|
def default_install_dir() -> Path:
|
|
"""``<UNSLOTH_HOME>/audio.cpp``, else ``<STUDIO_HOME>/audio.cpp``, else ``~/.unsloth/audio.cpp``.
|
|
|
|
Kept in step with ``audio_cpp_server.managed_audio_cpp_dir``; separate because this
|
|
script runs before the backend package is importable.
|
|
"""
|
|
master = (os.environ.get("UNSLOTH_HOME") or "").strip()
|
|
if master:
|
|
return Path(master).expanduser().resolve() / "audio.cpp"
|
|
home = (os.environ.get("UNSLOTH_STUDIO_HOME") or os.environ.get("STUDIO_HOME") or "").strip()
|
|
legacy = Path.home() / ".unsloth" / "audio.cpp"
|
|
if not home:
|
|
return legacy
|
|
root = Path(home).expanduser().resolve()
|
|
return legacy if root == (Path.home() / ".unsloth" / "studio").resolve() else root / "audio.cpp"
|
|
|
|
|
|
def nvidia_driver_cuda_version() -> Optional[tuple[int, int]]:
|
|
"""The CUDA version the installed NVIDIA driver supports, from nvidia-smi, or None."""
|
|
exe = shutil.which("nvidia-smi")
|
|
if not exe:
|
|
return None
|
|
try:
|
|
out = subprocess.run([exe], capture_output = True, text = True, timeout = 15).stdout
|
|
except (OSError, subprocess.SubprocessError):
|
|
return None
|
|
# Newer drivers print "CUDA UMD Version: X.Y" in place of "CUDA Version: X.Y" (same rule as the llama installer).
|
|
match = re.search(r"CUDA(?: UMD)? Version:\s*(\d+)\.(\d+)", out or "")
|
|
return (int(match.group(1)), int(match.group(2))) if match else None
|
|
|
|
|
|
def _nvidia_has_prebuilt_arch() -> Optional[bool]:
|
|
"""Any NVIDIA card sm_75+ (audio.cpp CUDA*_ARCHS in unsloth-prebuilt-build.yml start at 75); None if unreadable."""
|
|
exe = shutil.which("nvidia-smi")
|
|
if not exe:
|
|
return None
|
|
try:
|
|
out = subprocess.run(
|
|
[exe, "--query-gpu=compute_cap", "--format=csv,noheader"],
|
|
capture_output = True,
|
|
text = True,
|
|
timeout = 15,
|
|
)
|
|
except (OSError, subprocess.SubprocessError):
|
|
return None
|
|
caps = []
|
|
for line in (out.stdout or "").splitlines() if out.returncode == 0 else ():
|
|
if not line.strip():
|
|
continue
|
|
match = re.fullmatch(r"\s*(\d+)\.(\d+)\s*", line)
|
|
if not match:
|
|
# Old driver ("[N/A]"): unknown, keep the CUDA answer.
|
|
return None
|
|
caps.append((int(match.group(1)), int(match.group(2))))
|
|
return any(cap >= (7, 5) for cap in caps) if caps else None
|
|
|
|
|
|
def _nvidia_gpu_present() -> bool:
|
|
"""An NVIDIA card, also one whose driver reports no CUDA version (wedged, partial)."""
|
|
if shutil.which("nvidia-smi"):
|
|
return True
|
|
try:
|
|
return sys.platform.startswith("linux") and bool(os.listdir("/proc/driver/nvidia/gpus"))
|
|
except OSError:
|
|
return False
|
|
|
|
|
|
def _amd_gpu_present() -> bool:
|
|
"""PCI vendor 0x1002 in Linux DRM sysfs or the Windows display-adapter registry, as install_llama_prebuilt reads it."""
|
|
if sys.platform.startswith("linux"):
|
|
for path in glob.glob("/sys/class/drm/card*/device/vendor"):
|
|
try:
|
|
if Path(path).read_text(encoding = "utf-8").strip().lower() == "0x1002":
|
|
return True
|
|
except (OSError, UnicodeDecodeError):
|
|
continue
|
|
return False
|
|
if sys.platform != "win32":
|
|
return False
|
|
try:
|
|
if __package__:
|
|
from .install_llama_prebuilt import windows_amd_gpu_in_registry
|
|
else:
|
|
from install_llama_prebuilt import windows_amd_gpu_in_registry # type: ignore[no-redef]
|
|
return bool(windows_amd_gpu_in_registry())
|
|
except Exception: # noqa: BLE001 - no registry answer: keep the CPU build
|
|
return False
|
|
|
|
|
|
def detect_accelerator() -> str:
|
|
system = platform.system().lower()
|
|
if system == "darwin":
|
|
return "metal"
|
|
# Pascal / Volta-only hosts: the fork's bundles have no kernels for them. Other publishers set their own archs.
|
|
if nvidia_driver_cuda_version() is not None and (
|
|
_repo().casefold() != DEFAULT_REPO.casefold() or _nvidia_has_prebuilt_arch() is not False
|
|
):
|
|
return "cuda"
|
|
# No HIP bundle is published (#13042); never beside NVIDIA, which Vulkan could enumerate first.
|
|
if _amd_gpu_present() and not _nvidia_gpu_present():
|
|
return "vulkan"
|
|
return "cpu"
|
|
|
|
|
|
def torch_cuda_major() -> Optional[int]:
|
|
"""The CUDA major line of the torch importable here, without importing torch."""
|
|
try:
|
|
import importlib.metadata as md
|
|
version = md.version("torch")
|
|
except Exception: # noqa: BLE001 - no torch
|
|
return None
|
|
match = re.search(r"\+cu(\d+)", version)
|
|
if match:
|
|
digits = match.group(1)
|
|
return int(digits[:-1]) if len(digits) >= 3 else int(digits[:2])
|
|
return None
|
|
|
|
|
|
def resolve_release_asset(
|
|
asset_names: Sequence[str],
|
|
*,
|
|
system: str,
|
|
machine: str,
|
|
accelerator: str,
|
|
driver_cuda: Optional[tuple[int, int]] = None,
|
|
prefer_cuda_major: Optional[int] = None,
|
|
) -> Optional[str]:
|
|
"""Pick the bundle for a host, or None. Pure: the caller passes the release's asset names.
|
|
|
|
CUDA picks the newest line the driver can run, preferring ``prefer_cuda_major`` (the
|
|
torch line, whose runtime DLLs can then be reused). CPU prefers the portable build,
|
|
which avoids host-ISA features a native build may assume.
|
|
"""
|
|
system = system.lower()
|
|
accel = (accelerator or "cpu").lower()
|
|
os_tokens = _OS_TOKENS.get(system)
|
|
arch = _ARCH_TOKENS.get(machine.lower())
|
|
if not os_tokens or not arch:
|
|
return None
|
|
bundles = [
|
|
n
|
|
for n in asset_names
|
|
if "-bin-" in n
|
|
and (n.endswith(".zip") or n.endswith(".tar.gz"))
|
|
and any(f"-bin-{tok}-{arch}-" in n for tok in os_tokens)
|
|
]
|
|
|
|
def backend_of(name: str) -> str:
|
|
stem = name[: -len(".tar.gz")] if name.endswith(".tar.gz") else name[: -len(".zip")]
|
|
return stem.split(f"-{arch}-", 1)[1]
|
|
|
|
if accel == "metal":
|
|
pool = [n for n in bundles if backend_of(n).startswith("metal")]
|
|
return pool[0] if pool else None
|
|
if accel == "cuda":
|
|
candidates = []
|
|
for n in bundles:
|
|
match = _CUDA_LINE.search("-" + backend_of(n))
|
|
if not match:
|
|
continue
|
|
line = (int(match.group(1)), int(match.group(2) or 0))
|
|
if driver_cuda is not None and line > driver_cuda:
|
|
continue
|
|
candidates.append((line, n))
|
|
if not candidates:
|
|
return None
|
|
preferred = [
|
|
c for c in candidates if prefer_cuda_major is not None and c[0][0] == prefer_cuda_major
|
|
]
|
|
# At equal lines a general build beats upstream's Colab one, which carries sm_75 kernels only.
|
|
return max(preferred or candidates, key = lambda c: (c[0], "colab" not in c[1]))[1]
|
|
if accel == "vulkan":
|
|
pool = sorted(
|
|
(n for n in bundles if backend_of(n).startswith("vulkan")),
|
|
key = lambda n: "portable" not in n,
|
|
)
|
|
return pool[0] if pool else None
|
|
pool = sorted(
|
|
(n for n in bundles if backend_of(n).startswith("cpu")), key = lambda n: "portable" not in n
|
|
)
|
|
return pool[0] if pool else None
|
|
|
|
|
|
def cudart_asset_for(asset_names: Sequence[str], bundle: str) -> Optional[str]:
|
|
"""The separate Windows CUDA runtime archive matching a CUDA bundle's line."""
|
|
match = _CUDA_LINE.search(bundle)
|
|
if not match or "-windows-" not in bundle:
|
|
return None
|
|
line = match.group(0)[1:]
|
|
for name in asset_names:
|
|
if "-cudart-windows-" in name and name.endswith(f"-{line}.zip"):
|
|
return name
|
|
return None
|
|
|
|
|
|
def _torch_lib_dir() -> Optional[Path]:
|
|
try:
|
|
import importlib.util
|
|
spec = importlib.util.find_spec("torch")
|
|
except Exception: # noqa: BLE001
|
|
return None
|
|
if spec is None or not spec.submodule_search_locations:
|
|
return None
|
|
return Path(list(spec.submodule_search_locations)[0]) / "lib"
|
|
|
|
|
|
def torch_provides_cuda_runtime(cuda_major: int) -> bool:
|
|
"""Whether this Python's torch ships the CUDA runtime DLLs a line-``cuda_major`` bundle loads."""
|
|
if sys.platform != "win32":
|
|
return False
|
|
lib = _torch_lib_dir()
|
|
if lib is None:
|
|
return False
|
|
needed = (f"cudart64_{cuda_major}*.dll", f"cublas64_{cuda_major}*.dll", "cufft64_*.dll")
|
|
return all(glob.glob(str(lib / pattern)) for pattern in needed)
|
|
|
|
|
|
def _fetch_release(
|
|
repo: str,
|
|
tag: Optional[str],
|
|
token: Optional[str],
|
|
timeout: float = 30.0,
|
|
) -> Optional[dict]:
|
|
headers = {"Accept": "application/vnd.github+json", "User-Agent": USER_AGENT}
|
|
if token:
|
|
headers["Authorization"] = f"Bearer {token}"
|
|
url = f"https://api.github.com/repos/{repo}/releases/" + (f"tags/{tag}" if tag else "latest")
|
|
req = urllib.request.Request(url, headers = headers)
|
|
attempts = core.JSON_FETCH_ATTEMPTS
|
|
for attempt in range(1, attempts + 1):
|
|
try:
|
|
with urllib.request.urlopen(req, timeout = timeout) as resp: # noqa: S310
|
|
return json.loads(resp.read().decode("utf-8"))
|
|
except urllib.error.HTTPError as exc:
|
|
if _is_rate_limited(exc):
|
|
raise GitHubRateLimited(_rate_limit_message()) from exc
|
|
if exc.code == 404:
|
|
return None
|
|
if attempt >= attempts or not core.is_retryable_url_error(exc):
|
|
raise
|
|
core.sleep_backoff(attempt, exc = exc)
|
|
except Exception as exc: # noqa: BLE001 - classified below
|
|
if attempt >= attempts or not core.is_retryable_url_error(exc):
|
|
raise
|
|
core.sleep_backoff(attempt, exc = exc)
|
|
return None # pragma: no cover - the loop returns or raises
|
|
|
|
|
|
# Per process, so the CPU-fallback re-resolve and a full install after the re-run check ask once.
|
|
_LATEST_TAG_MEMO: dict[str, Optional[str]] = {}
|
|
_RELEASE_INDEX_MEMO: dict[tuple[str, Optional[str]], Optional[dict]] = {}
|
|
|
|
|
|
def _latest_tag_via_download_host(repo: str) -> Optional[str]:
|
|
"""The tag ``github.com/<repo>/releases/latest`` redirects to: one HEAD, no api.github.com call.
|
|
``None`` on any failure, so the caller asks the full lookup instead."""
|
|
if repo not in _LATEST_TAG_MEMO:
|
|
try:
|
|
_LATEST_TAG_MEMO[repo] = _OPS._download_host_latest_release_tag(repo)
|
|
except Exception as exc: # noqa: BLE001 - offline, DNS, 5xx: no shortcut
|
|
print(
|
|
f"audio.cpp: {repo} latest release lookup failed ({exc})",
|
|
file = sys.stderr,
|
|
flush = True,
|
|
)
|
|
_LATEST_TAG_MEMO[repo] = None
|
|
return _LATEST_TAG_MEMO[repo]
|
|
|
|
|
|
def _release_index(
|
|
repo: str,
|
|
tag: Optional[str],
|
|
*,
|
|
token: Optional[str] = None,
|
|
allow_unindexed: bool = False,
|
|
) -> Optional[dict]:
|
|
"""``{"tag_name", "assets", "checksums"}`` for ``repo`` at ``tag`` (newest when ``None``), through
|
|
prebuilt_core's release lookup: the manifest's assets with their tag-pinned download URLs, and the
|
|
release's sha256 index. A lookup that cannot answer raises (``PrebuiltFallback``, ``URLError``,
|
|
``OSError``, ``RuntimeError``).
|
|
|
|
``allow_unindexed`` (a user-picked repo or tag) lets a release that publishes no sha256 index
|
|
through on one GitHub API fetch, with ``checksums`` None so GitHub's digest is the expected sha256.
|
|
``None`` when that release does not exist."""
|
|
key = (repo, tag)
|
|
if key in _RELEASE_INDEX_MEMO:
|
|
return _RELEASE_INDEX_MEMO[key]
|
|
try:
|
|
bundle, checksums = _OPS.fetch_release_for_install(repo, published_release_tag = tag)
|
|
except core.ReleaseIntegrityError:
|
|
raise
|
|
except core.PrebuiltFallback:
|
|
if not allow_unindexed:
|
|
raise
|
|
release = _fetch_release(repo, tag, token)
|
|
if release is not None:
|
|
names = {a.get("name") for a in release.get("assets") or [] if isinstance(a, dict)}
|
|
if SHA256_ASSET_NAME in names:
|
|
# An index exists, so the lookup failed for another reason: never skip the index.
|
|
raise
|
|
print(
|
|
f"audio.cpp: {repo} {release.get('tag_name')} publishes no {SHA256_ASSET_NAME}; "
|
|
"using GitHub's asset digests",
|
|
file = sys.stderr,
|
|
flush = True,
|
|
)
|
|
release = {
|
|
"tag_name": release.get("tag_name"),
|
|
"assets": [
|
|
{
|
|
"name": a["name"],
|
|
"browser_download_url": a.get("browser_download_url"),
|
|
"digest": a.get("digest"),
|
|
"manifest_sha256": None,
|
|
"min_os": None,
|
|
}
|
|
for a in release.get("assets") or []
|
|
if isinstance(a, dict) and isinstance(a.get("name"), str)
|
|
],
|
|
"checksums": None,
|
|
"source_commit": None,
|
|
}
|
|
else:
|
|
release = {
|
|
"tag_name": bundle.release_tag,
|
|
"assets": [
|
|
{
|
|
"name": a["asset"],
|
|
"browser_download_url": _OPS.asset_download_url(bundle, a["asset"]),
|
|
"manifest_sha256": a.get("sha256"),
|
|
"min_os": a.get("min_os"),
|
|
}
|
|
for a in bundle.manifest.get("artifacts") or []
|
|
],
|
|
"checksums": checksums,
|
|
"source_commit": bundle.manifest.get("source_commit"),
|
|
}
|
|
_RELEASE_INDEX_MEMO[key] = release
|
|
return release
|
|
|
|
|
|
def _rate_limit_text(exc: BaseException) -> bool:
|
|
# prebuilt_core.fetch_json turns an API 403 / 429 into a RuntimeError with this text.
|
|
return bool(re.search(r"GitHub API returned (403|429)\b", str(exc)))
|
|
|
|
|
|
def resolve(
|
|
accelerator: str,
|
|
token: Optional[str],
|
|
tag: Optional[str] = None,
|
|
) -> tuple[str, Optional[dict], Optional[str]]:
|
|
"""``(repo, release, asset)`` for this host from ``_repo()``: ``tag`` (``--published-release-tag``),
|
|
else ``UNSLOTH_AUDIO_CPP_TAG``, else the newest release. Raises ``ReleaseLookupUnavailable`` (or
|
|
``GitHubRateLimited``) when the lookup could not answer, so a caller can keep what is installed;
|
|
``ReleaseIntegrityError`` (the release answered and cannot be trusted) is never one of those."""
|
|
repo = _repo()
|
|
explicit = (tag or "").strip() or None
|
|
pin = explicit or _pinned_tag()
|
|
want = pin or _latest_tag_via_download_host(repo)
|
|
driver = nvidia_driver_cuda_version() if accelerator == "cuda" else None
|
|
prefer = torch_cuda_major()
|
|
try:
|
|
release = _release_index(
|
|
repo,
|
|
want,
|
|
token = token,
|
|
# Trust follows the user's own repo/tag pick, so the in-app update can install the tag it
|
|
# offered from that repo too; the default fork release always publishes its index.
|
|
allow_unindexed = _user_picked_release(),
|
|
)
|
|
except core.ReleaseIntegrityError:
|
|
raise
|
|
except GitHubRateLimited:
|
|
raise
|
|
except (core.PrebuiltFallback, urllib.error.URLError, OSError, RuntimeError) as exc:
|
|
if _is_rate_limited(exc) or _rate_limit_text(exc):
|
|
raise GitHubRateLimited(_rate_limit_message()) from exc
|
|
print(f"audio.cpp: {repo} release lookup failed ({exc})", file = sys.stderr, flush = True)
|
|
raise ReleaseLookupUnavailable(
|
|
f"no audio.cpp release lookup answered ({repo} {want or 'latest'}: {exc})"
|
|
) from exc
|
|
if not release:
|
|
return repo, None, None
|
|
names = [a["name"] for a in release.get("assets") or []]
|
|
chosen = resolve_release_asset(
|
|
names,
|
|
system = platform.system(),
|
|
machine = platform.machine(),
|
|
accelerator = accelerator,
|
|
driver_cuda = driver,
|
|
prefer_cuda_major = prefer,
|
|
)
|
|
return repo, release, chosen
|
|
|
|
|
|
def resolve_for_request(
|
|
requested: str,
|
|
detected: Optional[str],
|
|
token: Optional[str],
|
|
tag: Optional[str] = None,
|
|
) -> tuple[str, str, Optional[dict], Optional[str]]:
|
|
"""``(accelerator, repo, release, asset)``. An auto-detected GPU the published bundles do not
|
|
cover (an older driver, Linux arm64) falls back to the CPU build, which still runs every model;
|
|
an explicit request is never downgraded."""
|
|
accel = detected if requested == "auto" and detected else requested
|
|
repo, release, chosen = resolve(accel, token, tag)
|
|
if requested == "auto" and accel == "cuda" and chosen:
|
|
major = _cuda_major(chosen)
|
|
if not _linux_cuda_runtime_available(major):
|
|
# The Linux CUDA bundles load cudart / cuBLAS from torch's nvidia-* wheels, as the
|
|
# llama.cpp ones do; a Studio without torch (or with another CUDA line) runs the CPU build.
|
|
print(
|
|
f"audio.cpp: no CUDA {major} runtime (cudart, cuBLAS) for {chosen}; installing the CPU build",
|
|
flush = True,
|
|
)
|
|
accel = "cpu"
|
|
repo, release, chosen = resolve(accel, token, tag)
|
|
if (release is None or not chosen) and requested == "auto" and accel != "cpu":
|
|
print(f"audio.cpp: no {accel} bundle for this host; installing the CPU build", flush = True)
|
|
accel = "cpu"
|
|
repo, release, chosen = resolve(accel, token, tag)
|
|
return accel, repo, release, chosen
|
|
|
|
|
|
def _download(url: str, dest: Path) -> None:
|
|
"""Retried, progress-reporting download (prebuilt_core's), written atomically to ``dest``."""
|
|
core.download_file(_OPS, url, dest)
|
|
|
|
|
|
def _verify(path: Path, expected_sha256: Optional[str]) -> str:
|
|
"""sha256 of ``path``; raises when an expected digest is known and differs."""
|
|
got = core.sha256_file(path)
|
|
if expected_sha256:
|
|
if expected_sha256.lower() != got:
|
|
raise RuntimeError(
|
|
f"sha256 mismatch for {path.name}: expected {expected_sha256.lower()}, got {got}"
|
|
)
|
|
else:
|
|
print(f"audio.cpp: WARNING no published digest for {path.name}", flush = True)
|
|
return got
|
|
|
|
|
|
def _extract(archive: Path, target: Path) -> None:
|
|
if archive.name.endswith(".zip"):
|
|
with zipfile.ZipFile(archive) as zf:
|
|
_safe_extractall(zf, target)
|
|
return
|
|
with tarfile.open(archive, "r:gz") as tf:
|
|
if hasattr(tarfile, "data_filter"):
|
|
# Refuses absolute paths, parent traversal, device files and links leaving the tree.
|
|
tf.extractall(target, filter = "data")
|
|
return
|
|
# Pythons without extraction filters: the same containment rules, checked by hand.
|
|
base = target.resolve()
|
|
for member in tf.getmembers():
|
|
dest = (base / member.name).resolve()
|
|
if dest != base and base not in dest.parents:
|
|
raise RuntimeError(f"unsafe path in archive: {member.name!r}")
|
|
if member.isdev() or member.isfifo():
|
|
raise RuntimeError(f"special file in archive: {member.name!r}")
|
|
if member.issym() or member.islnk():
|
|
link = (
|
|
(dest.parent / member.linkname).resolve()
|
|
if member.issym()
|
|
else (base / member.linkname).resolve()
|
|
)
|
|
if link != base and base not in link.parents:
|
|
raise RuntimeError(f"unsafe link in archive: {member.name!r}")
|
|
tf.extractall(target)
|
|
|
|
|
|
MODEL_SPECS_DIR = "model_specs"
|
|
_COMMIT = re.compile(r"[0-9a-f]{7,40}")
|
|
_REPO_NAME = re.compile(r"[A-Za-z0-9_.-]+/[A-Za-z0-9_.-]+")
|
|
_SPEC_NAME = re.compile(r"[a-z0-9_]+\.json")
|
|
# The specs are optional, so a network that blocks GitHub costs one short try, not retries.
|
|
_SPEC_TIMEOUT = 10
|
|
|
|
|
|
def _spec_tree(repo: str, commit: str) -> object:
|
|
url = f"https://api.github.com/repos/{repo}/git/trees/{commit}:{MODEL_SPECS_DIR}"
|
|
data = core.download_bytes(
|
|
_OPS, url, timeout = _SPEC_TIMEOUT, attempts = 1, headers = _OPS.github_api_headers(url)
|
|
)
|
|
return json.loads(data.decode("utf-8"))
|
|
|
|
|
|
def _spec_blob(repo: str, commit: str, name: str) -> bytes:
|
|
url = f"https://raw.githubusercontent.com/{repo}/{commit}/{MODEL_SPECS_DIR}/{name}"
|
|
return core.download_bytes(_OPS, url, timeout = _SPEC_TIMEOUT, attempts = 1)
|
|
|
|
|
|
def _spec_bytes(repo: str, commit: str, name: str, blob_sha: str) -> bytes:
|
|
data = _spec_blob(repo, commit, name)
|
|
got = hashlib.sha1(b"blob %d\0" % len(data) + data).hexdigest()
|
|
if got != blob_sha:
|
|
raise RuntimeError(f"{name}: git blob sha {got} != {blob_sha}")
|
|
if not isinstance(json.loads(data.decode("utf-8")), dict):
|
|
raise RuntimeError(f"{name} is not a JSON object")
|
|
return data
|
|
|
|
|
|
def _fetch_model_specs(repo: str, commit: Optional[str], dest_dir: Path) -> bool:
|
|
"""Put audio.cpp's ``model_specs/*.json`` at ``commit`` into ``dest_dir/model_specs``, the folder
|
|
Unsloth Studio reads to tell which model families this runtime loads. All or nothing: every file is
|
|
checked against its git blob sha and must parse, else no folder is left. Never raises; True when
|
|
the folder is there afterwards (one already present is kept as is, with no network call).
|
|
|
|
Always by commit (the manifest's ``source_commit``), never the release tag: the fork's release
|
|
tags all point at its packaging commit, whose ``model_specs`` is an older tree."""
|
|
dest_dir = Path(dest_dir)
|
|
dest = dest_dir / MODEL_SPECS_DIR
|
|
if dest.is_dir():
|
|
return True
|
|
if not isinstance(commit, str) or not _COMMIT.fullmatch(commit):
|
|
log("no source commit for this release; model specs not fetched")
|
|
return False
|
|
if not isinstance(repo, str) or not _REPO_NAME.fullmatch(repo):
|
|
log(f"unexpected repo {repo!r}; model specs not fetched")
|
|
return False
|
|
tmp = None
|
|
try:
|
|
tree = _spec_tree(repo, commit)
|
|
if not isinstance(tree, dict) or tree.get("truncated") or not tree.get("tree"):
|
|
raise RuntimeError("empty or truncated model_specs listing")
|
|
entries = {
|
|
e["path"]: e["sha"]
|
|
for e in tree["tree"]
|
|
if isinstance(e, dict)
|
|
and e.get("type") == "blob"
|
|
and isinstance(e.get("path"), str)
|
|
and _SPEC_NAME.fullmatch(e["path"])
|
|
and isinstance(e.get("sha"), str)
|
|
}
|
|
if not entries:
|
|
raise RuntimeError("no spec files in the model_specs listing")
|
|
with ThreadPoolExecutor(max_workers = 8) as pool:
|
|
blobs = dict(
|
|
zip(
|
|
entries,
|
|
pool.map(lambda item: _spec_bytes(repo, commit, *item), entries.items()),
|
|
)
|
|
)
|
|
tmp = core.make_staging_dir(dest_dir, ".model_specs-")
|
|
for name, data in blobs.items():
|
|
(tmp / name).write_bytes(data)
|
|
os.replace(tmp, dest)
|
|
tmp = None
|
|
except Exception as exc: # noqa: BLE001 - offline, 404, rate limit, bad blob: install goes on
|
|
log(f"model specs at {commit[:12]} not fetched ({exc})")
|
|
finally:
|
|
if tmp is not None:
|
|
shutil.rmtree(tmp, ignore_errors = True)
|
|
return dest.is_dir()
|
|
|
|
|
|
def _locate_server(root: Path) -> Optional[Path]:
|
|
for p in sorted(root.rglob(SERVER_NAME), key = lambda p: len(p.parts)):
|
|
if p.is_file():
|
|
return p
|
|
return None
|
|
|
|
|
|
def read_install_record(root: Path) -> dict:
|
|
try:
|
|
with open(root / INSTALL_RECORD, "r", encoding = "utf-8") as f:
|
|
rec = json.load(f)
|
|
return rec if isinstance(rec, dict) else {}
|
|
except (OSError, ValueError):
|
|
return {}
|
|
|
|
|
|
def _backend_of(asset: str) -> str:
|
|
# The Intel macOS bundle keeps the "-metal" name but is built with Metal off (Intel runners
|
|
# have no usable GPU), so it can only run on the CPU backend.
|
|
if "-macos-x64-" in asset:
|
|
return "cpu"
|
|
if "-metal" in asset:
|
|
return "metal"
|
|
if "-cuda" in asset:
|
|
return "cuda"
|
|
if "-vulkan" in asset:
|
|
return "vulkan"
|
|
return "cpu"
|
|
|
|
|
|
# What earlier Studio builds left in the managed dir without an ownership marker: the server's
|
|
# scratch home. A tree holding nothing else is Studio's, not the user's.
|
|
_STUDIO_LEFTOVERS = frozenset({".child_home"})
|
|
|
|
|
|
def _sweep_retired_trees(target: Path) -> None:
|
|
"""Remove ``<target>.old-*`` trees an earlier swap could not delete (Windows keeps a running
|
|
server's files open). Best-effort: one still in use stays for the next run."""
|
|
for old in target.parent.glob(target.name + ".old-*"):
|
|
try:
|
|
owned = old.is_dir() and (old / OWNERSHIP_MARKER).is_file()
|
|
except OSError:
|
|
# Unreadable (an elevated run's tree, #14052): only an elevated run can remove it.
|
|
owned = False
|
|
if owned:
|
|
shutil.rmtree(old, ignore_errors = True)
|
|
# Staging dirs of a killed run; the install lock is held, so none belongs to a live one.
|
|
for staging in target.parent.glob(".audio.cpp-staging-*"):
|
|
if staging.is_dir() and not staging.is_symlink():
|
|
shutil.rmtree(staging, ignore_errors = True)
|
|
|
|
|
|
def _is_unreadable(target: Path) -> bool:
|
|
"""Whether this account cannot look inside ``target``: a tree an elevated run staged under
|
|
mkdtemp's Windows ACL is owned by Administrators and denies the same user unelevated (#14052).
|
|
Renaming it is refused too, so only an elevated shell can clear it."""
|
|
try:
|
|
os.stat(target / OWNERSHIP_MARKER)
|
|
except PermissionError:
|
|
return True
|
|
except OSError:
|
|
pass
|
|
return False
|
|
|
|
|
|
def _server_in_use(server: Optional[Path]) -> bool:
|
|
"""Whether a running audiocpp_server holds this binary. Windows maps a running image
|
|
without write sharing, so opening it for writing fails; elsewhere a replace is safe anyway."""
|
|
if server is None or sys.platform != "win32":
|
|
return False
|
|
try:
|
|
with open(server, "r+b"):
|
|
return False
|
|
except PermissionError:
|
|
return True
|
|
except OSError:
|
|
return False
|
|
|
|
|
|
def _cuda_runtime_satisfied(record: dict) -> bool:
|
|
"""A CUDA install that skipped its cudart archive still has torch's runtime to lean on."""
|
|
if (
|
|
record.get("backend") != "cuda"
|
|
or record.get("cudart_asset")
|
|
or record.get("os") != "windows"
|
|
):
|
|
return True
|
|
match = _CUDA_LINE.search(str(record.get("asset") or ""))
|
|
return bool(match) and torch_provides_cuda_runtime(int(match.group(1)))
|
|
|
|
|
|
def _intact_install(target: Path, record: dict) -> Optional[Path]:
|
|
"""The server of a complete, unaltered Studio install of ``record``, or None."""
|
|
if not record or not (target / OWNERSHIP_MARKER).is_file():
|
|
return None
|
|
rel = record.get("server_relpath")
|
|
want = core.normalize_sha256_digest(record.get("server_sha256"))
|
|
if not isinstance(rel, str) or not rel or not want:
|
|
return None
|
|
server = target / rel
|
|
try:
|
|
if not server.is_file() or core.sha256_file(server) != want:
|
|
return None
|
|
except OSError:
|
|
return None
|
|
# Removed or quarantined phonemizer data: reinstall rather than keep a bundle Kokoro cannot use.
|
|
if record.get("espeak") is True and not (server.parent / "espeak-ng-data.bin").is_file():
|
|
return None
|
|
# A Windows CUDA install that leaned on torch's runtime is broken once torch moves lines.
|
|
return server if _cuda_runtime_satisfied(record) else None
|
|
|
|
|
|
def _record_fits_request(record: dict, requested: str, detected: Optional[str]) -> bool:
|
|
"""The record was installed for the same accelerator request (and, for auto, the same detected
|
|
host), and is not a CPU fallback that torch's CUDA runtime, here now, would replace."""
|
|
if record.get("accelerator_request") != requested:
|
|
return False
|
|
if requested == "auto" and record.get("detected_accelerator") != detected:
|
|
return False
|
|
if (
|
|
requested == "auto"
|
|
and detected == "cuda"
|
|
and record.get("accelerator") == "cpu"
|
|
and _linux_cuda_runtime_available(torch_cuda_major())
|
|
):
|
|
# Installed as the CPU fallback; torch's CUDA runtime is here now, so look up the GPU build.
|
|
return False
|
|
return True
|
|
|
|
|
|
def _current_install_matches(
|
|
target: Path,
|
|
record: dict,
|
|
requested: str,
|
|
detected: Optional[str],
|
|
latest_tag: Optional[str],
|
|
*,
|
|
pinned: bool = False,
|
|
) -> Optional[Path]:
|
|
"""Answer "already matches" before any release lookup: the record names this repo and
|
|
``latest_tag`` (the pin, or the newest release ``/releases/latest`` names), or, unpinned, a
|
|
release newer than it; it fits this accelerator request; and the server on disk is the one it
|
|
installed. ``latest_tag`` None (that HEAD failed) never matches."""
|
|
if not record or not latest_tag or record.get("published_repo") != _repo():
|
|
return None
|
|
tag = record.get("release_tag")
|
|
if not (tag == latest_tag if pinned else _record_is_current(tag, latest_tag)):
|
|
return None
|
|
if not _record_fits_request(record, requested, detected):
|
|
return None
|
|
return _intact_install(target, record)
|
|
|
|
|
|
_WINDOWS_LOADER_FAILURES = {
|
|
0xC0000135: "a DLL it needs was not found",
|
|
0xC000007B: "a DLL it loads is built for another architecture",
|
|
0xC0000139: "a DLL it loads lacks an entry point it needs",
|
|
}
|
|
|
|
|
|
def _scrubbed_environ() -> dict:
|
|
"""``os.environ`` without secrets, through the backend's ``child_env.scrub_env`` so the staged
|
|
server sees what Studio will later launch it with."""
|
|
import importlib.util
|
|
|
|
path = Path(__file__).resolve().parent / "backend" / "utils" / "prebuilt" / "child_env.py"
|
|
try:
|
|
spec = importlib.util.spec_from_file_location("_unsloth_audio_cpp_child_env", path)
|
|
module = importlib.util.module_from_spec(spec)
|
|
spec.loader.exec_module(module)
|
|
return module.scrub_env(os.environ)
|
|
except Exception: # noqa: BLE001 - a partial checkout: drop the obvious secret names
|
|
markers = ("TOKEN", "SECRET", "PASSWORD", "CREDENTIAL", "API_KEY", "PROXY")
|
|
return {k: v for k, v in os.environ.items() if not any(m in k.upper() for m in markers)}
|
|
|
|
|
|
def _loader_env(server: Path, extra_dirs: Sequence[Path] = ()) -> dict:
|
|
"""The environment with the server's own directory first on the loader path, as the backend's
|
|
``child_env`` launches it."""
|
|
env = _scrubbed_environ()
|
|
if sys.platform == "win32":
|
|
var = "PATH"
|
|
elif sys.platform == "darwin":
|
|
var = "DYLD_LIBRARY_PATH"
|
|
else:
|
|
var = "LD_LIBRARY_PATH"
|
|
existing = [p for p in env.get(var, "").split(os.pathsep) if p]
|
|
env[var] = os.pathsep.join([str(server.parent), *(str(d) for d in extra_dirs), *existing])
|
|
return env
|
|
|
|
|
|
def _cuda_runtime_dirs(backend: str) -> list[str]:
|
|
"""The CUDA libraries the backend's ``child_env`` puts after the bundle dir: the Linux CUDA server
|
|
links cuBLAS, cudart, cuFFT and NCCL dynamically and finds them in torch's nvidia-* wheels, and a
|
|
Windows bundle installed without its cudart archive uses torch's DLLs."""
|
|
if backend != "cuda" or sys.platform == "darwin":
|
|
return []
|
|
try:
|
|
if __package__:
|
|
from .install_llama_prebuilt import python_runtime_dirs
|
|
else:
|
|
from install_llama_prebuilt import python_runtime_dirs # type: ignore[no-redef]
|
|
|
|
return list(python_runtime_dirs())
|
|
except Exception: # noqa: BLE001 - no wheel runtime to offer
|
|
return []
|
|
|
|
|
|
def _cuda_major(asset: str) -> Optional[int]:
|
|
match = _CUDA_LINE.search(asset)
|
|
return int(match.group(1)) if match else None
|
|
|
|
|
|
def _ldconfig_libs() -> set[str]:
|
|
try:
|
|
out = subprocess.run(["ldconfig", "-p"], capture_output = True, text = True, timeout = 10).stdout
|
|
except Exception: # noqa: BLE001 - no ldconfig: only the wheel dirs count
|
|
return set()
|
|
return {line.split()[0] for line in out.splitlines()[1:] if line.strip()}
|
|
|
|
|
|
def _linux_cuda_runtime_available(major: Optional[int]) -> bool:
|
|
"""Whether the CUDA libraries a line-``major`` Linux bundle links but does not ship (cudart,
|
|
cuBLAS, cuFFT) are in torch's nvidia-* wheels or on the system loader path. True off Linux,
|
|
where the bundle (or its cudart archive) carries them."""
|
|
if not sys.platform.startswith("linux"):
|
|
return True
|
|
if major is None:
|
|
return False
|
|
dirs = [Path(d) for d in _cuda_runtime_dirs("cuda")]
|
|
system = _ldconfig_libs()
|
|
|
|
def present(pattern: str) -> bool:
|
|
if any(any(d.glob(pattern)) for d in dirs):
|
|
return True
|
|
return any(fnmatch.fnmatch(name, pattern) for name in system)
|
|
|
|
return all(
|
|
present(pattern)
|
|
for pattern in (f"libcudart.so.{major}", f"libcublas.so.{major}", "libcufft.so.*")
|
|
)
|
|
|
|
|
|
def _run_staged(server: Path, arg: str, env: dict, timeout: float) -> subprocess.CompletedProcess:
|
|
return subprocess.run(
|
|
[str(server), arg],
|
|
capture_output = True,
|
|
text = True,
|
|
encoding = "utf-8",
|
|
errors = "replace",
|
|
timeout = timeout,
|
|
env = env,
|
|
cwd = str(server.parent),
|
|
**core.windows_hidden_subprocess_kwargs(),
|
|
)
|
|
|
|
|
|
class NoVulkanDevice(RuntimeError):
|
|
"""The staged Vulkan server loads but enumerates no Vulkan GPU, so every model load would fail."""
|
|
|
|
|
|
# `--list-devices` prints one `Vulkan:<n> "<name>" [GPU]` row per device ggml-vulkan accepts.
|
|
_VULKAN_DEVICE_ROW = re.compile(r"^Vulkan:\d+\s", re.MULTILINE)
|
|
|
|
|
|
def smoke_test_staged_server(
|
|
server: Path,
|
|
*,
|
|
backend: str,
|
|
extra_dirs: Sequence[Path] = (),
|
|
require_vulkan_device: bool = False,
|
|
) -> Optional[str]:
|
|
"""Start the staged server once with ``--help``: a bundle this host cannot load (a newer glibc, a
|
|
missing DLL) raises here, before it replaces a working tree. Returns its ``--version`` text
|
|
(build and enabled backends) when it prints one."""
|
|
env = _loader_env(server, extra_dirs)
|
|
if require_vulkan_device:
|
|
try:
|
|
listed = _run_staged(server, "--list-devices", env, SMOKE_TIMEOUT_SECONDS)
|
|
rows = (listed.stdout or "") if listed.returncode == 0 else ""
|
|
except (OSError, subprocess.SubprocessError):
|
|
rows = ""
|
|
if not _VULKAN_DEVICE_ROW.search(rows):
|
|
raise NoVulkanDevice("the Vulkan build finds no Vulkan GPU on this host")
|
|
try:
|
|
result = _run_staged(server, "--help", env, SMOKE_TIMEOUT_SECONDS)
|
|
except subprocess.TimeoutExpired as exc:
|
|
raise RuntimeError(
|
|
f"staged {SERVER_NAME} --help did not exit within {SMOKE_TIMEOUT_SECONDS:.0f}s; keeping the existing install"
|
|
) from exc
|
|
except OSError as exc:
|
|
raise RuntimeError(
|
|
f"staged {SERVER_NAME} could not be started ({exc}); keeping the existing install"
|
|
) from exc
|
|
if result.returncode != 0:
|
|
output = "\n".join(s.strip() for s in (result.stderr, result.stdout) if s and s.strip())
|
|
if backend == "cuda" and "libcuda.so" in output and nvidia_driver_cuda_version() is None:
|
|
# No NVIDIA driver on this host (a Docker image build): libcuda arrives with the driver at run time.
|
|
print(
|
|
"audio.cpp: no NVIDIA driver here to load the CUDA build; skipping its start check",
|
|
flush = True,
|
|
)
|
|
return None
|
|
hint = (
|
|
_WINDOWS_LOADER_FAILURES.get(result.returncode & 0xFFFFFFFF)
|
|
if sys.platform == "win32"
|
|
else None
|
|
)
|
|
detail = "\n".join(output.splitlines()[-20:]) or hint or "no output"
|
|
raise RuntimeError(
|
|
f"staged {SERVER_NAME} --help exited {result.returncode}; this host cannot run the bundle, "
|
|
f"keeping the existing install:\n{detail}"
|
|
)
|
|
try:
|
|
version = _run_staged(server, "--version", env, 10.0)
|
|
except (OSError, subprocess.SubprocessError):
|
|
return None
|
|
text = (version.stdout or "").strip()
|
|
return text[:2000] if version.returncode == 0 and text else None
|
|
|
|
|
|
def install(
|
|
*,
|
|
install_dir: Optional[Path] = None,
|
|
accelerator: str = "auto",
|
|
token: Optional[str] = None,
|
|
force: bool = False,
|
|
release_tag: Optional[str] = None,
|
|
keep_newer: bool = False,
|
|
) -> Path:
|
|
"""Install (or keep) the bundle for ``accelerator``; ``"auto"`` detects the host here, so the
|
|
CPU fallback for an uncovered GPU applies. ``release_tag`` installs exactly that release (the
|
|
in-app update passes the one it offered); with ``keep_newer`` an intact install of a newer
|
|
release of this repo is kept instead. Serialised on prebuilt_core's install lock."""
|
|
if accelerator not in ACCELERATORS:
|
|
raise RuntimeError(
|
|
f"unknown accelerator {accelerator!r}; expected one of {', '.join(ACCELERATORS)}"
|
|
)
|
|
target = (install_dir or default_install_dir()).resolve()
|
|
if target.exists() and not target.is_dir():
|
|
raise RuntimeError(f"audio.cpp install target is not a directory: {target}")
|
|
if _is_unreadable(target):
|
|
# Matched by setup.ps1, which shows the repair instead of a generic failure.
|
|
repair = (
|
|
f'from an elevated PowerShell with Remove-Item -Recurse -Force "{target}"'
|
|
if sys.platform == "win32"
|
|
else f"with sudo rm -rf '{target}'"
|
|
)
|
|
raise RuntimeError(
|
|
f"{target} cannot be read by this account (left by an elevated run); remove it "
|
|
f"{repair}, then run the install again"
|
|
)
|
|
if target.exists() and not (target / OWNERSHIP_MARKER).is_file():
|
|
entries = {p.name for p in target.iterdir()}
|
|
if entries - _STUDIO_LEFTOVERS:
|
|
raise RuntimeError(
|
|
f"{target} already exists and is not an Unsloth-managed directory; refusing to replace it. "
|
|
"Move it away or pass a different --install-dir."
|
|
)
|
|
with core.install_lock(core.install_lock_path(target)):
|
|
try:
|
|
return _install_locked(
|
|
target, accelerator, token, force, release_tag = release_tag, keep_newer = keep_newer
|
|
)
|
|
except NoVulkanDevice as exc:
|
|
# Forced so the record keeps detected "vulkan": later runs keep this CPU build without a lookup.
|
|
print(f"audio.cpp: {exc}; installing the CPU build", flush = True)
|
|
return _install_locked(
|
|
target,
|
|
accelerator,
|
|
token,
|
|
True,
|
|
vulkan_unusable = True,
|
|
release_tag = release_tag,
|
|
keep_newer = keep_newer,
|
|
)
|
|
|
|
|
|
def _install_locked(
|
|
target: Path,
|
|
requested: str,
|
|
token: Optional[str],
|
|
force: bool,
|
|
vulkan_unusable: bool = False,
|
|
release_tag: Optional[str] = None,
|
|
keep_newer: bool = False,
|
|
) -> Path:
|
|
_sweep_retired_trees(target)
|
|
detected = detect_accelerator() if requested == "auto" else None
|
|
existing = read_install_record(target)
|
|
release_tag = (release_tag or "").strip() or None
|
|
pin = release_tag or _pinned_tag()
|
|
if (
|
|
keep_newer
|
|
and release_tag
|
|
and not force
|
|
and existing.get("published_repo") == _repo()
|
|
and existing.get("release_tag") != release_tag
|
|
and _record_is_current(existing.get("release_tag"), release_tag)
|
|
and _record_fits_request(existing, requested, detected)
|
|
):
|
|
# Checked under the install lock: a setup run that got here first with a newer release wins over a
|
|
# stale offer the in-app update pinned.
|
|
server = _intact_install(target, existing)
|
|
if server is not None:
|
|
print(f"audio.cpp: already matches {existing.get('asset')}", flush = True)
|
|
return server
|
|
if not force:
|
|
# Pinned: no network. Newest: one HEAD, reused as the tag a full install then looks up.
|
|
latest = pin or _latest_tag_via_download_host(_repo())
|
|
server = _current_install_matches(
|
|
target, existing, requested, detected, latest, pinned = pin is not None
|
|
)
|
|
if server is not None:
|
|
print(f"audio.cpp: already matches {existing.get('asset')}", flush = True)
|
|
return server
|
|
token = token or os.environ.get("GH_TOKEN") or os.environ.get("GITHUB_TOKEN")
|
|
try:
|
|
accel, repo, release, chosen = resolve_for_request(
|
|
requested, "cpu" if vulkan_unusable else detected, token, release_tag
|
|
)
|
|
except (ReleaseLookupUnavailable, GitHubRateLimited) as exc:
|
|
# The lookup could not answer, which says nothing about the tree on disk. Keep a complete one
|
|
# unless this run asked for something it is not (--force, another explicit accelerator, or a
|
|
# pinned release: --published-release-tag or UNSLOTH_AUDIO_CPP_TAG, which an installed tree that
|
|
# matched it already answered above without a lookup).
|
|
kept = None if force or pin else _intact_install(target, existing)
|
|
if kept is None or (requested != "auto" and existing.get("accelerator") != requested):
|
|
raise
|
|
# "keeping the existing complete install" is the substring setup.sh and setup.ps1 grep for.
|
|
print(KEPT_EXISTING_LINE, flush = True)
|
|
print(
|
|
f"audio.cpp: {existing.get('release_tag')} stays; release lookup unavailable ({exc})",
|
|
file = sys.stderr,
|
|
flush = True,
|
|
)
|
|
return kept
|
|
if release is None or not chosen:
|
|
raise RuntimeError(
|
|
f"No prebuilt audiocpp_server for {platform.system()}/{platform.machine()} "
|
|
f"(accelerator={accel}) from {repo}. Build it from source: https://github.com/{repo}"
|
|
)
|
|
assets = {a["name"]: a for a in release.get("assets") or []}
|
|
asset = assets[chosen]
|
|
tag = release.get("tag_name")
|
|
expected = _expected_sha256(release, asset)
|
|
if (
|
|
not force
|
|
and expected
|
|
and existing.get("asset") == chosen
|
|
and existing.get("asset_sha256") == expected
|
|
and _intact_install(target, existing) is not None
|
|
):
|
|
print(f"audio.cpp: already matches {chosen}", flush = True)
|
|
return _intact_install(target, existing)
|
|
if (
|
|
not force
|
|
and pin is None
|
|
and existing.get("published_repo") == repo
|
|
and existing.get("release_tag") != tag
|
|
and _record_is_current(existing.get("release_tag"), tag)
|
|
and _record_fits_request(existing, requested, detected)
|
|
and _intact_install(target, existing) is not None
|
|
):
|
|
# The lookup named an older release than the one installed: never downgrade to it.
|
|
print(f"audio.cpp: already matches {existing.get('asset')}", flush = True)
|
|
return _intact_install(target, existing)
|
|
|
|
cudart = cudart_asset_for(list(assets), chosen)
|
|
cuda_match = _CUDA_LINE.search(chosen)
|
|
cuda_major = int(cuda_match.group(1)) if cuda_match else None
|
|
torch_runtime = bool(cudart and cuda_major and torch_provides_cuda_runtime(cuda_major))
|
|
cudart_expected = (
|
|
_expected_sha256(release, assets[cudart]) if cudart and not torch_runtime else None
|
|
)
|
|
old_glibc = _glibc_below(asset.get("min_os"))
|
|
if old_glibc:
|
|
raise RuntimeError(
|
|
f"audio.cpp prebuilt bundles need {asset['min_os'].strip().replace('-', ' ', 1)}+; "
|
|
f"this host has {old_glibc}. Set UNSLOTH_AUDIO_CPP_PATH to a local build instead."
|
|
)
|
|
|
|
if target.exists() and _server_in_use(_locate_server(target)):
|
|
# Checked again before the swap; here it saves the download.
|
|
raise PermissionError(f"audiocpp_server in {target} is running")
|
|
target.parent.mkdir(parents = True, exist_ok = True)
|
|
staging = core.make_staging_dir(target.parent, ".audio.cpp-staging-")
|
|
try:
|
|
archive = staging / chosen
|
|
print(f"audio.cpp: downloading {chosen} from {repo} {tag}", flush = True)
|
|
_download(asset["browser_download_url"], archive)
|
|
digest = _verify(archive, expected)
|
|
tree = staging / "tree"
|
|
tree.mkdir()
|
|
_extract(archive, tree)
|
|
archive.unlink()
|
|
server = _locate_server(tree)
|
|
if server is None:
|
|
raise RuntimeError(f"{chosen} contains no {SERVER_NAME}")
|
|
cudart_used = None
|
|
if cudart and not torch_runtime:
|
|
print(f"audio.cpp: downloading CUDA runtime {cudart}", flush = True)
|
|
rt_archive = staging / cudart
|
|
_download(assets[cudart]["browser_download_url"], rt_archive)
|
|
_verify(rt_archive, cudart_expected)
|
|
with zipfile.ZipFile(rt_archive) as zf:
|
|
_safe_extractall(zf, server.parent)
|
|
rt_archive.unlink()
|
|
cudart_used = cudart
|
|
if sys.platform != "win32":
|
|
server.chmod(server.stat().st_mode | 0o111)
|
|
backend = _backend_of(chosen)
|
|
version = smoke_test_staged_server(
|
|
server,
|
|
backend = backend,
|
|
extra_dirs = _cuda_runtime_dirs(backend),
|
|
require_vulkan_device = requested == "auto" and backend == "vulkan",
|
|
)
|
|
# Beside the server (bin/ on Windows), where Studio reads it; a failed fetch never fails the install.
|
|
source_commit = release.get("source_commit")
|
|
if _fetch_model_specs(repo, source_commit, server.parent):
|
|
specs = len(list((server.parent / MODEL_SPECS_DIR).glob("*.json")))
|
|
print(f"audio.cpp: {specs} model specs beside the server", flush = True)
|
|
(tree / OWNERSHIP_MARKER).touch()
|
|
record = {
|
|
"schema_version": 1,
|
|
"component": "audio.cpp",
|
|
"published_repo": repo,
|
|
"release_tag": tag,
|
|
"source_commit": source_commit,
|
|
"asset": chosen,
|
|
"asset_sha256": digest,
|
|
"cudart_asset": cudart_used,
|
|
"backend": backend,
|
|
# What was asked and seen, so a re-run knows without a lookup whether this still fits.
|
|
"accelerator": accel,
|
|
"accelerator_request": requested,
|
|
"detected_accelerator": detected,
|
|
# Static eSpeak-ng builds ship their phoneme data beside the server; upstream bundles do not.
|
|
"espeak": (server.parent / "espeak-ng-data.bin").is_file(),
|
|
"server_relpath": server.relative_to(tree).as_posix(),
|
|
"server_sha256": core.sha256_file(server),
|
|
"server_version": version,
|
|
"os": platform.system().lower(),
|
|
"arch": _ARCH_TOKENS.get(platform.machine().lower(), platform.machine().lower()),
|
|
"installed_at_utc": datetime.datetime.now(datetime.timezone.utc).isoformat(),
|
|
}
|
|
(tree / INSTALL_RECORD).write_text(json.dumps(record, indent = 2), encoding = "utf-8")
|
|
# A running server would keep its files after the swap and serve the old build; the
|
|
# caller stops it (or reports busy) instead.
|
|
if target.exists() and _server_in_use(_locate_server(target)):
|
|
raise PermissionError(f"audiocpp_server in {target} is running")
|
|
# Swap: the old tree moves aside first so a failed rename can put it back.
|
|
retired = None
|
|
if target.exists():
|
|
retired = target.with_name(target.name + f".old-{os.getpid()}")
|
|
os.replace(target, retired)
|
|
try:
|
|
os.replace(tree, target)
|
|
except OSError:
|
|
if retired is not None:
|
|
os.replace(retired, target)
|
|
raise
|
|
if retired is not None:
|
|
shutil.rmtree(retired, ignore_errors = True)
|
|
finally:
|
|
shutil.rmtree(staging, ignore_errors = True)
|
|
installed = target / record["server_relpath"]
|
|
print(f"audio.cpp: installed {installed}", flush = True)
|
|
return installed
|
|
|
|
|
|
def main(argv: Optional[list[str]] = None) -> int:
|
|
p = argparse.ArgumentParser(description = "Install a prebuilt audiocpp_server (audio.cpp).")
|
|
p.add_argument(
|
|
"--accelerator",
|
|
default = _accelerator_from_env(),
|
|
# setup.sh / setup.ps1 forward UNSLOTH_AUDIO_CPP_ACCELERATOR verbatim ("CUDA", " cpu ").
|
|
type = lambda value: value.strip().lower(),
|
|
choices = list(ACCELERATORS),
|
|
help = "default: UNSLOTH_AUDIO_CPP_ACCELERATOR, else auto (detect; CPU when no bundle covers the GPU)",
|
|
)
|
|
p.add_argument("--install-dir", default = None)
|
|
p.add_argument("--force", action = "store_true", help = "reinstall even when the bundle matches")
|
|
p.add_argument("--print-asset", action = "store_true", help = "resolve and print the asset only")
|
|
p.add_argument(
|
|
"--published-release-tag",
|
|
default = None,
|
|
help = "install exactly this release of the repo (default: UNSLOTH_AUDIO_CPP_TAG, else the newest)",
|
|
)
|
|
p.add_argument(
|
|
"--keep-newer",
|
|
action = "store_true",
|
|
help = "with --published-release-tag: keep an intact install of a newer release of the repo",
|
|
)
|
|
args = p.parse_args(argv)
|
|
token = os.environ.get("GH_TOKEN") or os.environ.get("GITHUB_TOKEN")
|
|
release_tag = (args.published_release_tag or "").strip() or None
|
|
try:
|
|
if args.print_asset:
|
|
detected = detect_accelerator() if args.accelerator == "auto" else None
|
|
_, repo, release, chosen = resolve_for_request(
|
|
args.accelerator, detected, token, release_tag
|
|
)
|
|
print(f"{repo} {release.get('tag_name') if release else '-'} {chosen or '(none)'}")
|
|
return EXIT_OK if chosen else EXIT_FAILED
|
|
install(
|
|
install_dir = Path(args.install_dir).expanduser() if args.install_dir else None,
|
|
accelerator = args.accelerator,
|
|
force = args.force,
|
|
release_tag = release_tag,
|
|
keep_newer = args.keep_newer,
|
|
)
|
|
except GitHubRateLimited as exc:
|
|
print(f"error: {exc}", file = sys.stderr)
|
|
return EXIT_FAILED
|
|
except core.BusyInstallConflict as exc:
|
|
# Another setup (the desktop app, a second shell) holds the install lock.
|
|
print(f"error: another audio.cpp install is running ({exc})", file = sys.stderr)
|
|
return EXIT_BUSY
|
|
except PermissionError as exc:
|
|
# Windows: a running audiocpp_server holds the tree open.
|
|
print(
|
|
f"error: the audio.cpp runtime is in use ({exc}); stop Studio and retry",
|
|
file = sys.stderr,
|
|
)
|
|
return EXIT_BUSY
|
|
except (RuntimeError, OSError) as exc:
|
|
print(f"error: {exc}", file = sys.stderr)
|
|
return EXIT_FAILED
|
|
return EXIT_OK
|
|
|
|
|
|
if __name__ == "__main__":
|
|
raise SystemExit(main())
|