searchi/tools/speech-runtime/build.py

"""Build pinned isolated speech runtime; auto distributions require GPU SDKs."""
import argparse
import hashlib
import json
import os
from pathlib import Path
import shutil
import subprocess
import sys
import importlib.util

UPSTREAM = "d1be6fde11ac6e0407606b4e42fe72d34add8037"


def shared_helper(name):
    path = Path(__file__).resolve().parent.parent / "embeddinggemma2" / name
    spec = importlib.util.spec_from_file_location("speech_shared_" + path.stem, path)
    module = importlib.util.module_from_spec(spec)
    sys.path.insert(0, str(path.parent))
    try:
        spec.loader.exec_module(module)
    finally:
        sys.path.pop(0)
    return module


def build_inputs():
    own = Path(__file__).resolve().parent
    paths = {name: own / name for name in ["CMakeLists.txt", "worker.cpp", "build.py", "cpu-only-loader.patch", "windows-model-path.patch"]}
    paths.update({"embeddinggemma2/" + name: own.parent / "embeddinggemma2" / name
                  for name in ["build_runtime.py", "provision_vulkan.py", "runtime_provenance.py", "macos_runtime.py"]})
    return {name: sha(path) for name, path in paths.items()}


def sha(path):
    result = hashlib.sha256()
    with path.open("rb") as stream:
        for chunk in iter(lambda: stream.read(65536), b""):
            result.update(chunk)
    return result.hexdigest()


def validate_source(source):
    revision = subprocess.check_output(["git", "-C", str(source), "rev-parse", "HEAD"], text=True).strip()
    dirty = subprocess.check_output(["git", "-C", str(source), "status", "--porcelain"], text=True).strip()
    if revision != UPSTREAM or dirty:
        raise ValueError("Speech build requires clean pinned whisper.cpp source")


def configuration(source, output, mode, cmake):
    return [cmake, "-S", str(Path(__file__).resolve().parent), "-B", str(output / "build"),
            "-DWHISPER_SOURCE=" + str(source), "-DCMAKE_BUILD_TYPE=Release",
            "-DSEARCHI_INTERNAL_CPU=" + ("ON" if mode == "cpu" else "OFF"),
            "-DWHISPER_BUILD_IS_DEV=OFF"]


def main():
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument("--source", type=Path, required=True)
    parser.add_argument("--output", type=Path, required=True)
    parser.add_argument("--device", choices=["auto", "cpu"], default="auto")
    parser.add_argument("--cmake", default="cmake")
    args = parser.parse_args()
    source, output = args.source.resolve(), args.output.resolve()
    validate_source(source)
    own = Path(__file__).resolve().parent
    source_hashes = build_inputs()
    # A fresh owned destination prevents stale CPU/DLL artifacts masquerading as auto.
    output.mkdir(parents=True, exist_ok=False)
    loader = None
    if args.device == "auto" and sys.platform == "win32":
        sdk, loader = shared_helper("provision_vulkan.py").provision()
        os.environ["VULKAN_SDK"] = str(sdk)
    patched = output / "whisper-source"
    subprocess.run(["git", "clone", "--quiet", "--no-hardlinks", str(source), str(patched)], check=True)
    subprocess.run(["git", "-C", str(patched), "checkout", "--quiet", "--detach", UPSTREAM], check=True)
    for patch in ["cpu-only-loader.patch", "windows-model-path.patch"]:
        subprocess.run(["git", "-C", str(patched), "apply", "--check", str(own / patch)], check=True)
        subprocess.run(["git", "-C", str(patched), "apply", str(own / patch)], check=True)
    for name, command in [("configure", configuration(patched, output, args.device, args.cmake)),
                          ("build", [args.cmake, "--build", str(output / "build"), "--config", "Release", "--target", "searchi-speech-worker", "--parallel", "1"])]:
        with (output / (name + ".log")).open("wb") as log:
            result = subprocess.run(command, stdout=log, stderr=subprocess.STDOUT, check=False)
        if result.returncode:
            raise RuntimeError(f"Speech {name} failed; inspect {output / (name + '.log')}")
    validate_source(source)
    if source_hashes != build_inputs():
        raise ValueError("Speech build sources changed during compilation")
    build = output / "build"
    if args.device == "auto" and sys.platform == "win32":
        shared_helper("build_runtime.py").bundle_windows_crt(build)
    suffix = ".exe" if sys.platform == "win32" else ""
    candidates = list(build.rglob("searchi-speech-worker" + suffix))
    if len(candidates) != 1:
        raise ValueError("Speech build executable is missing or ambiguous")
    stage = output / "runtime"
    stage.mkdir()
    files = [candidates[0]]
    if args.device == "auto":
        files.extend(path for path in build.rglob("*") if path.is_file()
                     and (path.suffix.lower() in [".dll", ".dylib"] or ".so" in path.name))
        if loader is not None:
            files.extend([loader / "vulkan-1.dll", loader / "VULKAN-LICENSE.txt"])
            files.extend(build.rglob("MSVC-REDISTRIBUTION.txt"))
        names = [path.name.lower() for path in files]
        required = "metal" if sys.platform == "darwin" else "vulkan"
        if not any(required in name for name in names) or not any("ggml-cpu" in name for name in names):
            raise ValueError("Auto speech runtime requires both acceleration and CPU backends")
    records = {}
    for path in files:
        digest = sha(path)
        if path.name in records and records[path.name]["sha256"] != digest:
            raise ValueError("Ambiguous speech dependency basename")
        shutil.copy2(path, stage / path.name)
        records[path.name] = {"path": path.name, "bytes": path.stat().st_size, "sha256": digest}
    manifest = {"schema": 1, "upstream_commit": UPSTREAM, "device": args.device,
                "distribution": args.device == "auto", "source_hashes": source_hashes,
                "files": list(records.values()), "native_os_gpu_checks": "unexecuted"}
    (stage / "runtime-manifest.json").write_text(json.dumps(manifest, indent=2) + "\n", encoding="utf-8")
    print(json.dumps({"runtime": str(stage), "mode": args.device, "files": len(records)}))


if __name__ == "__main__":
    main()