char-fork/llama.cpp

git clone https://git.t4t.associates/char-fork/llama.cpp

kurquharsnapdragon: python SDK setup (Windows) (#27903)d7bd3bfca

char/aliceai-t5-moe
brancheschar/aliceai-t5-moeall refs →
8.9 KiB233 linesraw
1#!/usr/bin/env python3
2#
3# Install Windows on Snapdragon SDKs for llama.cpp.
4#
5
6import sys
7import os
8import argparse
9import shutil
10import logging
11import json
12import hashlib
13import tarfile
14import tempfile
15from pathlib import Path
16from urllib.error import HTTPError, URLError
17from urllib.request import Request, urlopen
18
19from sdk import SDK_CONFIGS, get_hexagon_tools_dir, is_valid_sdk
20
21
22logger = logging.getLogger("setup_sdk")
23
24DEFAULT_SDK_BASE_DIR = r"C:\Qualcomm"
25
26
27def get_sdk_releases(config):
28    request = Request(
29        f"https://api.github.com/repos/{config['repo']}/releases?per_page=100",
30        headers={"Accept": "application/vnd.github+json", "User-Agent": "llama.cpp"},
31    )
32    try:
33        with urlopen(request, timeout=30) as response:
34            releases = json.load(response)
35    except (HTTPError, URLError, TimeoutError) as err:
36        raise RuntimeError(f"Cannot query {config['name']} releases: {err}") from err
37
38    result = []
39    for release in releases:
40        if release["draft"] or release["prerelease"]:
41            continue
42        version = release["tag_name"].removeprefix("v")
43        archive_name = f"{config['archive_prefix']}{version}-arm64-wos.tar.xz"
44        for asset in release["assets"]:
45            if asset["name"] != archive_name:
46                continue
47            result.append({
48                "version": version,
49                "name": asset["name"],
50                "url": asset["browser_download_url"],
51                "sha256": (asset.get("digest") or "").removeprefix("sha256:"),
52            })
53    return result
54
55
56def list_sdk_releases():
57    for config in SDK_CONFIGS:
58        logger.info("%s:", config["name"])
59        releases = get_sdk_releases(config)
60        if not releases:
61            logger.info("  no Windows on Snapdragon releases found")
62            continue
63        for release in releases:
64            logger.info("  %s: %s", release["version"], release["name"])
65
66
67def get_sdk_release(config, version):
68    version = version or config["default_version"]
69    version = version.removeprefix("v")
70    for release in get_sdk_releases(config):
71        if release["version"] == version:
72            if not release["sha256"]:
73                raise RuntimeError(f"{config['name']} {version} does not provide a SHA-256 digest")
74            return release
75    raise RuntimeError(
76        f"No Windows on Snapdragon release for {config['name']} {version}. "
77        "Run scripts/snapdragon/setup-sdk.py --list-sdk-releases to see available versions."
78    )
79
80
81def sha256sum(path):
82    digest = hashlib.sha256()
83    with open(path, "rb") as file:
84        for chunk in iter(lambda: file.read(1024 * 1024), b""):
85            digest.update(chunk)
86    return digest.hexdigest()
87
88
89def download_sdk(release, archive):
90    while True:
91        if archive.exists() and sha256sum(archive) == release["sha256"]:
92            logger.info("Using existing archive %s", archive)
93            return
94
95        offset = archive.stat().st_size if archive.exists() else 0
96        headers = {"User-Agent": "llama.cpp"}
97        if offset:
98            headers["Range"] = f"bytes={offset}-"
99            logger.info("Resuming download of %s at %d MiB", release["name"], offset // (1024 * 1024))
100        else:
101            logger.info("Downloading %s", release["name"])
102
103        try:
104            with urlopen(Request(release["url"], headers=headers), timeout=30) as response:
105                mode = "ab" if offset and response.status == 206 else "wb"
106                with open(archive, mode) as file:
107                    shutil.copyfileobj(response, file)
108        except HTTPError as err:
109            if err.code != 416:
110                raise RuntimeError(f"Cannot download {release['name']}: {err}") from err
111            archive.unlink(missing_ok=True)
112            continue
113        except (URLError, TimeoutError) as err:
114            raise RuntimeError(f"Cannot download {release['name']}: {err}") from err
115
116        if sha256sum(archive) == release["sha256"]:
117            return
118        raise RuntimeError(f"SHA-256 mismatch for {archive}. Re-run the command to resume the download.")
119
120
121def extract_sdk(config, archive, target_dir):
122    if not hasattr(tarfile, "data_filter"):
123        raise RuntimeError("SDK extraction requires Python 3.10.12 or later")
124
125    with tempfile.TemporaryDirectory(prefix=f".{target_dir.name}.tmp-", dir=target_dir.parent) as staging_path:
126        staging_dir = Path(staging_path)
127        with tarfile.open(archive, "r:xz") as tar:
128            tar.extractall(staging_dir, filter=tarfile.data_filter)
129
130        candidates = [staging_dir] + [path for path in staging_dir.iterdir() if path.is_dir()]
131        extracted_dirs = [path for path in candidates if is_valid_sdk(config, path)]
132        if len(extracted_dirs) != 1:
133            raise RuntimeError(f"{config['name']} archive does not contain the expected files")
134        extracted_dir = extracted_dirs[0]
135
136        backup_dir = None
137        if target_dir.exists():
138            backup_dir = target_dir.parent / f".{target_dir.name}.backup"
139            if backup_dir.exists():
140                raise RuntimeError(f"Cannot replace {target_dir}: backup directory {backup_dir} already exists")
141            target_dir.replace(backup_dir)
142        try:
143            extracted_dir.replace(target_dir)
144        except Exception:
145            if backup_dir:
146                backup_dir.replace(target_dir)
147            raise
148        if backup_dir:
149            shutil.rmtree(backup_dir)
150
151
152def install_sdk(config, version, base_dir, force):
153    version = (version or config["default_version"]).removeprefix("v")
154    target_dir = base_dir / config["parent_dir"] / version
155    if is_valid_sdk(config, target_dir) and not force:
156        logger.info("Using existing %s at %s", config["name"], target_dir)
157        return target_dir
158
159    release = get_sdk_release(config, version)
160    target_dir.parent.mkdir(parents=True, exist_ok=True)
161    archive = target_dir.parent / release["name"]
162    download_sdk(release, archive)
163    logger.info("Extracting %s to %s", config["name"], target_dir)
164    extract_sdk(config, archive, target_dir)
165    archive.unlink(missing_ok=True)
166    return target_dir
167
168
169def set_user_environment(values):
170    if os.name != "nt":
171        raise RuntimeError("SDK setup must run on Windows")
172
173    import winreg
174
175    with winreg.CreateKey(winreg.HKEY_CURRENT_USER, "Environment") as key:
176        for name, value in values.items():
177            winreg.SetValueEx(key, name, 0, winreg.REG_SZ, str(value))
178            os.environ[name] = str(value)
179
180    import ctypes
181
182    result = ctypes.c_ulong()
183    ctypes.windll.user32.SendMessageTimeoutW(0xffff, 0x001a, 0, "Environment", 0x0002, 5000, ctypes.byref(result))
184
185
186def setup_sdks(args):
187    base_dir = Path(args.sdk_base_dir).expanduser().resolve()
188    hexagon_config, opencl_config = SDK_CONFIGS
189    environment = {}
190
191    if args.hexagon is not None:
192        hexagon_dir = install_sdk(hexagon_config, args.hexagon, base_dir, args.force)
193        environment["HEXAGON_SDK_ROOT"] = hexagon_dir
194        environment["HEXAGON_TOOLS_ROOT"] = get_hexagon_tools_dir(hexagon_dir)
195    if args.opencl is not None:
196        opencl_dir = install_sdk(opencl_config, args.opencl, base_dir, args.force)
197        environment["OPENCL_SDK_ROOT"] = opencl_dir
198
199    set_user_environment(environment)
200    logger.info("SDK environment variables were updated. Start a new terminal before building.")
201
202
203def main():
204    logging.basicConfig(level=logging.INFO, format="%(message)s")
205    parser = argparse.ArgumentParser(description="Install Windows on Snapdragon SDKs for llama.cpp.")
206    parser.add_argument("--list-sdk-releases", action="store_true", help="List available Windows on Snapdragon SDK releases")
207    parser.add_argument("--sdk-base-dir", default=DEFAULT_SDK_BASE_DIR, help=r"SDK installation directory (default: C:\Qualcomm)")
208    parser.add_argument("--hexagon", nargs="?", const=SDK_CONFIGS[0]["default_version"], metavar="VERSION", help="Install the Hexagon SDK, optionally selecting a version")
209    parser.add_argument("--opencl", nargs="?", const=SDK_CONFIGS[1]["default_version"], metavar="VERSION", help="Install the OpenCL SDK, optionally selecting a version")
210    parser.add_argument("--force", action="store_true", help="Reinstall selected SDKs even when they already exist")
211    args = parser.parse_args()
212
213    if args.list_sdk_releases:
214        if args.sdk_base_dir != DEFAULT_SDK_BASE_DIR or args.hexagon is not None or args.opencl is not None or args.force:
215            parser.error("Installation options cannot be combined with --list-sdk-releases")
216        list_sdk_releases()
217        return
218    if args.hexagon is None and args.opencl is None:
219        parser.error("Select at least one SDK with --hexagon or --opencl")
220    if os.name != "nt":
221        parser.error("SDK setup must run on Windows")
222    setup_sdks(args)
223
224
225if __name__ == "__main__":
226    try:
227        main()
228    except KeyboardInterrupt:
229        logger.info("\nInterrupted by user.")
230        sys.exit(130)
231    except RuntimeError as err:
232        logger.error("Error: %s", err)
233        sys.exit(1)