path: root/app/backends/audiocpp/build.py

"""Checkout lifecycle: clone location, ggml patches, binary build, uninstall."""

import contextlib
import io
import os
import re
import shlex
import shutil
from datetime import datetime
from pathlib import Path
from typing import List, Optional

from backends import common, servers
from backends.common import APP_DIR
from .catalog import _BACKEND_TOKEN_RE, detect_backend, load_server_config
from .constants import AUDIOCPP_DIR_NAME, BACKENDS, PATCH_DIR

def uninstall(*, emit=None, cancel=None) -> int:
    """Remove the audio.cpp backend entirely: stop its server, delete the checkout.

    The checkout (``app/audio.cpp``) holds the built binary, the downloaded
    models, and the server.json, so removing the directory uninstalls the
    backend. A running server this tool started is stopped first
    (best-effort).

    EMIT is accepted for registry symmetry with the other backends but is
    unused here — this uninstall has no subprocess phase, and its prints are
    captured by the task view when run in the TUI. CANCEL is a
    ``threading.Event`` honored between phases only (after the server has
    been stopped, before the checkout is deleted), so a started phase always
    completes and the uninstall never tears halfway. Returns the exit code
    (130 when cancelled before a remaining phase).
    """
    # Only stop when a pid file exists: without one this tool never
    # started the server, so the "not started by this tool" notice would
    # be uninstall-time noise.
    if servers.pid_for("audiocpp") is not None:
        servers.stop("audiocpp")
    if common.cancel_requested(cancel):
        return 130
    checkout = find_local_checkout()
    if checkout is None:
        print("[INFO] No audio.cpp checkout to remove.")
        return 0
    print(f"[INFO] Removing audio.cpp checkout {checkout}...")
    shutil.rmtree(checkout, ignore_errors=True)
    print("[OK] audio.cpp removed.")
    return 0


def update(*, emit=None, cancel=None) -> int:
    """Update the audio.cpp backend: refresh the checkout, rebuild if stale.

    A managed server that is running is stopped first (best-effort): it
    serves the binary whose sources are being replaced. Phases: stop
    server / git update / rebuild — CANCEL is honored between phases only,
    so a started phase always completes. The git update is a fetch plus
    hard reset to origin's HEAD (see ``common.git_update``): everything
    that matters lives untracked in the checkout (models, build trees,
    server.json) and survives, while the vendored-ggml patch edit is
    intentionally wiped — the rebuild re-applies it (the patch step is
    idempotent and fails loudly when upstream re-shaped the file).

    The rebuild target is the backend recorded in server.json, else the
    one detected from existing build directories; when neither names one
    (nothing was ever built) the update stops after the checkout refresh
    — 'Build audio.cpp Server' handles a first build. The rebuild itself
    runs when the sources changed (HEAD moved) or the on-disk binary is
    missing or older than HEAD's commit time — the latter heals an
    interrupted (cancelled or failed) earlier rebuild, which leaves the
    previous binary in place against already-updated sources. An
    up-to-date checkout with a fresh binary costs one fetch. Returns the
    exit code (130 when cancelled before a remaining phase).
    """
    # Only stop when a pid file exists: without one this tool never
    # started the server, so the "not started by this tool" notice would
    # be uninstall-time noise.
    if servers.pid_for("audiocpp") is not None:
        servers.stop("audiocpp")
    if common.cancel_requested(cancel):
        return 130
    checkout = find_local_checkout()
    if checkout is None:
        print("[INFO] No audio.cpp checkout to update.")
        return 0
    head_before = common.git_head(checkout)
    rc = common.git_update(checkout, emit=emit, cancel=cancel)
    if rc != 0:
        print(f"[WARNING] checkout update failed (exit {rc}); update "
              f"manually: git -C {checkout} pull")
        return rc
    head_after = common.git_head(checkout)
    if common.cancel_requested(cancel):
        return 130
    backend = _rebuild_backend(checkout)
    if backend is None:
        print("[INFO] audiocpp_server was never built for a known "
              "backend; skipping the rebuild. 'Build audio.cpp Server' "
              "builds one.")
        return 0
    binary = built_server_binary(checkout, backend)
    if not _rebuild_needed(checkout, binary,
                           moved=head_after not in (None, head_before)):
        print(f"[OK] {checkout} is already at origin's HEAD with an "
              "up-to-date audiocpp_server.")
        return 0
    if head_after in (None, head_before):
        print(f"[INFO] audiocpp_server on disk is older than the "
              f"checked-out sources (earlier build interrupted?); "
              f"rebuilding for {backend}.")
    else:
        print(f"[OK] Updated {checkout} to {head_after[:12]}; rebuilding "
              f"audiocpp_server for {backend}.")
    build_rc = build_audiocpp(checkout, backend, emit=emit, cancel=cancel)
    if build_rc != 0:
        print(f"[WARNING] rebuild exited with code {build_rc}; see the "
              "messages above (the build log under app/logs/ has the "
              "full output). The binary on disk is now older than the "
              "checked-out sources; re-running 'Update Backends' will "
              "retry the rebuild.")
    else:
        print("[OK] rebuild complete.")
    return build_rc


def _rebuild_needed(checkout: Path, binary: Optional[Path],
                    *, moved: bool) -> bool:
    """True when audiocpp_server must be (re)built after an update.

    True when the checkout moved, the binary is missing, its age cannot
    be compared (no commit time), or it predates HEAD's commit — the
    last case is what a cancelled or failed earlier rebuild leaves
    behind (old binary, already-updated sources).
    """
    if moved or binary is None:
        return True
    commit_time = common.git_commit_time(checkout)
    if commit_time is None:
        return True
    try:
        return binary.stat().st_mtime <= commit_time
    except OSError:
        return True


def _rebuild_backend(checkout: Path) -> Optional[str]:
    """The inference backend to rebuild for after an update, or None.

    server.json's recorded backend wins (it is what the managed server
    launches); an existing build directory's token is the fallback for a
    checkout that was built but never configured. None means neither
    names a valid backend — there is no binary to keep fresh.
    """
    server_config = load_server_config(checkout / "server.json") or {}
    recorded = server_config.get("backend")
    if recorded in BACKENDS:
        return recorded
    return detect_backend(checkout)


def find_local_checkout() -> Optional[Path]:
    """Return the managed audio.cpp checkout at ``app/audio.cpp``.

    Returns the path only when it contains a ``model_specs`` directory;
    the checkout is installed there by the setup wizard and nowhere else.
    """
    try:
        resolved = (APP_DIR / AUDIOCPP_DIR_NAME).resolve()
    except OSError:
        return None
    if (resolved / "model_specs").is_dir():
        return resolved
    return None


def find_audiocpp_server_bin(audiocpp_dir: Path) -> Optional[Path]:
    """Return the built audiocpp_server binary, or None when not built.

    Scans ``audiocpp_dir/build/*`` for a build directory containing
    ``bin/audiocpp_server`` (``.exe`` allowed on Windows). When several
    builds exist the first (alphabetical) is returned.
    """
    build_root = audiocpp_dir / "build"
    if not build_root.is_dir():
        return None
    try:
        build_dirs = sorted(build_root.iterdir(),
                            key=lambda p: p.name.lower())
    except OSError:
        return None
    for build_dir in build_dirs:
        if not build_dir.is_dir():
            continue
        for name in ("audiocpp_server", "audiocpp_server.exe"):
            server = build_dir / "bin" / name
            if server.exists():
                return server
    return None


def built_server_binary(audiocpp_dir: Path, backend: str) -> Optional[Path]:
    """Return the built audiocpp_server for BACKEND, or None.

    Like ``find_audiocpp_server_bin`` but limited to build directories whose
    name carries the BACKEND token (``-cuda-``, ``-vulkan-``, ``-hip-``,
    ``-cpu-``; ``-metal-`` counts as ``cpu``). A checkout with builds for
    several backends is asked which one to use without re-offering a build
    for a backend that is already built.
    """
    build_root = audiocpp_dir / "build"
    if not build_root.is_dir():
        return None
    try:
        build_dirs = sorted(build_root.iterdir(),
                            key=lambda p: p.name.lower())
    except OSError:
        return None
    for build_dir in build_dirs:
        if not build_dir.is_dir():
            continue
        match = _BACKEND_TOKEN_RE.search(build_dir.name.lower())
        if not match:
            continue
        token = "cpu" if match.group(1) == "metal" else match.group(1)
        if token != backend:
            continue
        for name in ("audiocpp_server", "audiocpp_server.exe"):
            server = build_dir / "bin" / name
            if server.exists():
                return server
    return None


def find_build_script(audiocpp_dir: Path) -> Optional[Path]:
    """Return the audio.cpp build helper script to run, or None.

    Prefers ``scripts/build_linux.sh``; otherwise the first
    ``scripts/build_*.sh`` it finds. (Windows ``.bat`` scripts are not run
    automatically — build manually there.)
    """
    scripts = audiocpp_dir / "scripts"
    if not scripts.is_dir():
        return None
    preferred = scripts / "build_linux.sh"
    if preferred.exists():
        return preferred
    try:
        candidates = sorted(scripts.glob("build_*.sh"),
                            key=lambda p: p.name.lower())
    except OSError:
        return None
    return candidates[0] if candidates else None


GGML_PATCHES = [
    {
        "file": "ggml-top-k-cuda-iterator.patch",
        "target": "external/ggml/src/ggml-cuda/top-k.cu",
        "marker": r"#\s*include\s*<cuda/iterator>",
        "label": "top-k.cu: add #include <cuda/iterator> (CCCL 3.x build fix)",
    },
]

# No-output watchdog for the build: ninja prints a line per completed
# compile, so 15 minutes of silence means a compiler job wedged (ptxas
# hangs are the known failure mode of a buggy CUDA toolkit). The runner
# kills the build and reports exit 124 (see run_console_subprocess).
BUILD_STALL_TIMEOUT = 900

# Env override for the CUDA architectures passed to build_linux.sh
# (--cuda-arch): a ';'-or-comma separated list of compute capabilities,
# e.g. "86" or "86;89". Set it when GPU detection cannot run.
CUDA_ARCH_ENV = "AUDIOCPP_CUDA_ARCH"

# Compute capability -> common GPUs, shown when detection is impossible
# and in ptxas failure guidance. Terse on purpose; one line.
CUDA_ARCH_GUIDE = ("61 GTX 10xx/P40; 75 RTX 20xx; 80 A100; 86 RTX 30xx "
                   "(3090)/A6000; 89 RTX 40xx (4090)/L40S; 90 H100; "
                   "120/121 RTX 50xx (5090)/B200")


def detect_cuda_arch() -> Optional[str]:
    """The CUDA architecture token to build for, or None when unknown.

    ``AUDIOCPP_CUDA_ARCH`` wins verbatim (validated as a ';'-or-comma
    separated list of compute capabilities like ``86`` or ``86;89``), so a
    user can pin the arch on machines where detection cannot run. Otherwise
    ``nvidia-smi`` reports each GPU's compute capability (works on Linux
    and Windows; it does not exist on macOS, where the CUDA backend is not
    a choice anyway): ``8.6`` becomes ``86``, several distinct GPUs join
    as ``86;89``. audio.cpp's CMake upgrades bare new architectures
    (``120``) to their suffixed forms (``120a``) itself.
    """
    override = os.environ.get(CUDA_ARCH_ENV, "").strip()
    if override:
        parts = [part.strip() for part in
                 override.replace(",", ";").split(";") if part.strip()]
        if parts and all(re.fullmatch(r"\d+(-real|-virtual)?", part)
                         for part in parts):
            return ";".join(parts)
        print(f"[WARNING] {CUDA_ARCH_ENV}={override!r} is not an arch list "
              "(e.g. \"86\" or \"86;89\"); ignoring it")
    proc = common.run_console_subprocess_quiet(
        ["nvidia-smi", "--query-gpu=compute_cap", "--format=csv,noheader"],
        timeout=10)
    if proc is None or proc.returncode != 0:
        return None
    arches: List[str] = []
    for line in proc.stdout.decode("utf-8", errors="replace").splitlines():
        cap = line.strip()
        if not re.fullmatch(r"\d+\.\d+", cap):
            continue
        arch = cap.replace(".", "")
        if arch not in arches:
            arches.append(arch)
    return ";".join(arches) if arches else None


def _cuda_arch_argv(backend: str, emit=None) -> List[str]:
    """The ``--cuda-arch`` flags for a CUDA build, plus a status line.

    EMIT is the in-TUI line sink when building from the task view (the
    line lands in the view, not the real terminal behind curses); without
    it the line prints to the console. Detection failure is not an error:
    the build then uses audio.cpp's portable default arch list, which is
    slower to compile but runs on any GPU.
    """
    if backend != "cuda":
        return []
    arch = detect_cuda_arch()
    say = emit if emit is not None else print
    if arch is None:
        say(f"[INFO] CUDA architecture: portable default list (could not "
            f"detect a GPU; set {CUDA_ARCH_ENV}=<arch> to build only for "
            "this machine's GPU — much faster)")
        return []
    say(f"[INFO] CUDA architecture: {arch} (detected via nvidia-smi; "
        f"override with {CUDA_ARCH_ENV})")
    return ["--cuda-arch", arch]


def _ptxas_failure_hint(log_path: Path) -> str:
    """Guidance appended when the build log shows a ptxas failure.

    ``ptxas fatal`` / ``nvcc error`` lines mean the CUDA toolkit's
    assembler (or nvcc itself) failed — an internal compiler error is a
    toolkit bug, not a broken checkout, and newer ggml template code on
    newer toolkit releases trips it. Building only for the local GPU's
    architecture skips most of the codegen paths ptxas chokes on, so the
    hint points at ``AUDIOCPP_CUDA_ARCH`` (with the detected arch, or the
    GPU table when detection cannot run); a different toolkit version is
    the remaining fix when narrowing the arch is not enough.
    """
    try:
        text = log_path.read_text(encoding="utf-8", errors="ignore")
    except OSError:
        return ""
    if "ptxas fatal" not in text and "nvcc error" not in text:
        return ""
    arch = detect_cuda_arch()
    lines = [
        "  ptxas (the CUDA toolkit's GPU assembler) failed — with an "
        "internal compiler error this is a CUDA toolkit bug, not your "
        "sources.",
        f"  Rebuild for this machine's GPU only: set {CUDA_ARCH_ENV}=<arch> "
        "(semicolon-separated for several GPUs) and re-run the build.",
    ]
    if arch:
        lines.append(f"  Detected arch for this machine: {arch}")
    else:
        lines.append(f"  Arch per GPU: {CUDA_ARCH_GUIDE}")
    lines.append("  If narrowing the arch still fails, a different CUDA "
                 "toolkit version usually does (ptxas bugs are fixed in "
                 "toolkit updates).")
    return "\n".join(lines)


def apply_ggml_patches(audiocpp_dir: Path, *, emit=None, cancel=None) -> int:
    """Apply the shipped ggml build patches to an audio.cpp checkout.

    Idempotent: a patch whose marker already matches its target is skipped
    (it is either already applied, or the fork re-vendored a fixed ggml). A
    patch that no longer applies because the vendored file changed shape is a
    loud, non-interactive failure — the build is aborted so the user
    re-evaluates the patch instead of hitting a known nvcc break minutes
    later. Returns 0 when every patch is applied or already present, 1 on
    drift, 130 when cancelled.
    """
    for patch in GGML_PATCHES:
        if cancel is not None and cancel.is_set():
            return 130
        target = audiocpp_dir / patch["target"]
        if not target.is_file():
            print(f"[INFO] {patch['file']}: target {patch['target']} not "
                  f"present in this checkout; skipping")
            continue
        try:
            text = target.read_text(encoding="utf-8", errors="ignore")
        except OSError as exc:
            print(f"[WARNING] {patch['file']}: could not read {target}: "
                  f"{exc}; skipping")
            continue
        if re.search(patch["marker"], text):
            print(f"[OK] {patch['file']}: fix already present, skipping")
            continue
        patch_path = PATCH_DIR / patch["file"]
        if not patch_path.is_file():
            print(f"[ERROR] {patch['file']}: patch file not found at "
                  f"{patch_path}; cannot apply")
            return 1
        check_argv = ["git", "-C", str(audiocpp_dir), "apply", "--check",
                      "--whitespace=nowarn", str(patch_path)]
        check_rc = common.run_console_subprocess(
            check_argv, emit=emit, cancel=cancel)
        if check_rc == 130 or (cancel is not None and cancel.is_set()):
            return 130
        if check_rc != 0:
            print(f"[ERROR] {patch['file']}: no longer applies to "
                  f"{patch['target']} (git apply --check exit {check_rc}). "
                  f"The audio.cpp fork's vendored ggml changed shape and "
                  f"still lacks the fix. Re-evaluate {patch_path}: "
                  f"regenerate the patch, or drop this entry if the fork "
                  f"now ships the fix.")
            return 1
        apply_argv = ["git", "-C", str(audiocpp_dir), "apply",
                      "--whitespace=nowarn", str(patch_path)]
        rc = common.run_console_subprocess(
            apply_argv, emit=emit, cancel=cancel)
        if rc == 130 or (cancel is not None and cancel.is_set()):
            return 130
        if rc != 0:
            print(f"[ERROR] {patch['file']}: git apply failed (exit {rc})")
            return rc
        print(f"[OK] {patch['file']}: applied ({patch['label']})")
    return 0


def build_audiocpp(audiocpp_dir: Path, backend: str, *,
                   emit=None, cancel=None) -> int:
    """Build audiocpp_server for BACKEND, streaming output.

    With EMIT None the build script runs on the console (inherits the
    terminal); with EMIT given (the in-TUI task view) its output streams line
    by line to EMIT so the view can show progress, and CANCEL aborts it.

    On the EMIT (TUI) path the build output is also tee'd to
    ``app/logs/audiocpp_build_<timestamp>.log`` so it survives the curses
    session; when the build fails (and was not cancelled) a post-TUI notice
    with the copy-pastable command and the log path is queued for the console
    (see ``backends.common.record_post_tui_notice``).

    Returns the build script's exit code (non-zero when the script is
    missing).
    """
    script = find_build_script(audiocpp_dir)
    if script is None:
        message = (f"[ERROR] No build script found in {audiocpp_dir}/scripts; "
                   "build audiocpp_server manually (see the audio.cpp README)")
        print(message)
        if emit is not None:
            common.record_post_tui_notice(message)
        return 1
    argv = ["sh", str(script), "--backend", backend, "--target",
            "audiocpp_server", "--deployment-build"]
    argv += _cuda_arch_argv(backend, emit=emit)
    command = f"cd {audiocpp_dir} && {shlex.join(argv)}"
    if emit is None:
        print(f"[INFO] Building audiocpp_server for {backend} ({command})...")
        patch_rc = apply_ggml_patches(audiocpp_dir, cancel=cancel)
        if patch_rc == 130 or (cancel is not None and cancel.is_set()):
            return 130
        if patch_rc != 0:
            print("[ERROR] ggml build patches could not be applied; "
                  "aborting audiocpp_server build. See the messages above "
                  "and re-evaluate app/backends/patches/.")
            return patch_rc
        return common.run_console_subprocess(argv, cwd=audiocpp_dir)
    return _build_audiocpp_tui(emit, cancel, argv, command, audiocpp_dir)


def _build_audiocpp_tui(emit, cancel, argv: List[str], command: str,
                        audiocpp_dir: Path) -> int:
    """Run the build on the TUI path: tee output to a log file.

    The ggml patch step runs first, inside the same log: every emitted
    line (patch status, build output) is also written (and flushed) to
    ``app/logs/audiocpp_build_<timestamp>.log``. On failure a summary (the
    copy-pastable COMMAND and the log path) is emitted into the TUI,
    written to the log, and queued as a post-TUI console notice. A
    cancelled build (CANCEL set) is not reported as a failure, but its
    partial output stays in the log file.
    """
    log_path = common.LOG_DIR / (
        f"audiocpp_build_{datetime.now():%Y%m%d_%H%M%S}.log")
    log_path.parent.mkdir(parents=True, exist_ok=True)
    log_handle = log_path.open("w", encoding="utf-8")

    def tee(line: str) -> None:
        log_handle.write(line + "\n")
        log_handle.flush()
        emit(line)

    class _TeeWriter(io.TextIOBase):
        """Route print() output from the patch step into the log too."""

        def write(self, s: str) -> int:
            for line in s.splitlines():
                if line:
                    tee(line)
            return len(s)

    try:
        with contextlib.redirect_stdout(_TeeWriter()):
            patch_rc = apply_ggml_patches(audiocpp_dir, emit=tee,
                                          cancel=cancel)
        if patch_rc == 130 or (cancel is not None and cancel.is_set()):
            return 130
        if patch_rc != 0:
            notice = ("[ERROR] ggml build patches could not be applied; "
                      "aborting audiocpp_server build. See the messages "
                      "above and re-evaluate app/backends/patches/.")
            tee(notice)
            common.record_post_tui_notice(notice)
            return patch_rc
        tee(f"[INFO] Building audiocpp_server ({command})...")
        rc = common.run_console_subprocess(
            argv, cwd=audiocpp_dir, emit=tee, cancel=cancel,
            stall_timeout=BUILD_STALL_TIMEOUT)
        if rc != 0 and (cancel is None or not cancel.is_set()):
            if rc == 124:
                head = (f"[ERROR] audio.cpp build stalled — no output for "
                        f"{BUILD_STALL_TIMEOUT // 60} minutes, so it was "
                        "stopped (a wedged compiler job; often a ptxas "
                        "hang from a buggy CUDA toolkit).")
            else:
                head = f"[ERROR] audio.cpp build failed (exit code {rc})."
            notice = (f"{head}\n"
                      f"  Build log: {log_path}\n"
                      f"  Troubleshoot by re-running this command:\n"
                      f"    {command}")
            ptxas_hint = _ptxas_failure_hint(log_path)
            if ptxas_hint:
                notice += "\n" + ptxas_hint
            for line in notice.splitlines():
                tee(line)
            common.record_post_tui_notice(notice)
    finally:
        log_handle.close()
    return rc


def _print_launch_hint(audiocpp_dir: Path, output_path: Path) -> None:
    """Print remediation when audiocpp_server is missing (troubleshooting).

    The hub starts and stops the server itself, so a working install gets
    no manual launch instructions. When no binary was built, though, the
    user needs to know how to build and run it by hand. The commands are
    prefixed with ``cd <checkout> &&`` because the server discovers
    model_specs/<family>.json relative to its working directory.
    """
    if find_audiocpp_server_bin(audiocpp_dir) is not None:
        return
    print("\n[INFO] audiocpp_server binary not found. Build it first, e.g.:")
    script = find_build_script(audiocpp_dir)
    if script is not None:
        print(f"  sh {script} --backend <cuda|vulkan|hip|cpu> "
              "--target audiocpp_server --deployment-build")
    print(f"  then run: cd {audiocpp_dir} && ./build/<platform>-<backend>"
          f"-release/bin/audiocpp_server --config {output_path}")