searchi/tools/embeddinggemma2/linux_loader_smoke.py

"""Prove Linux auto CPU fallback with the Vulkan loader physically absent.

Run as root on x86_64 Linux with util-linux unshare/mount/umount and chroot.
This tool never escalates privileges. It launches itself inside a private mount
namespace and verifies that namespace before changing mount propagation. The
sealed chroot includes approved runtime files, standard C/C++ ABI libraries and
a read-only Google model, without Vulkan loaders, ICDs, drivers or GPU devices.
This tests missing-loader fallback, not GPU execution or distro compatibility.
"""
import argparse
import hashlib
import json
import math
import os
from pathlib import Path
import platform
import shutil
import subprocess
import sys
import tempfile

ABI = ("libc.so.6", "libm.so.6", "libdl.so.2", "libpthread.so.0",
       "librt.so.1", "libstdc++.so.6", "libgcc_s.so.1")


def require(condition, message):
    if not condition:
        raise RuntimeError(message)


def checksum(path):
    with Path(path).open("rb") as stream:
        return hashlib.file_digest(stream, "sha256").hexdigest()


def namespace(pid="self"):
    return Path(f"/proc/{pid}/ns/mnt").stat().st_ino


def isolated_check(args, commands):
    # Check the still-live outer process as well as its recorded namespace.
    # A direct child invocation in the outer namespace fails before any mount.
    outer = namespace(args.outer_pid)
    require(outer == args.outer_namespace and namespace() != outer,
            "Private mount namespace was not established; refusing mount changes")
    subprocess.run([commands["mount"], "--make-rprivate", "/"], check=True)
    from stage_runtime import verify_runtime
    worker = args.worker.resolve()
    dependencies = verify_runtime(worker, True)
    manifest = json.loads(Path(__file__).with_name("model-manifest.json").read_text())
    artifact = manifest["files"][0]
    model = args.models.resolve() / artifact["path"]
    require(model.is_file() and model.stat().st_size == artifact["bytes"]
            and checksum(model) == artifact["sha256"], "Google model verification failed")
    base = Path(tempfile.gettempdir()).resolve()
    root = Path(tempfile.mkdtemp(prefix="searchi-loader-chroot-", dir=base)).resolve()
    require(root.parent == base and root.stat().st_uid == 0,
            "Invalid owned chroot directory")
    mounted = []
    try:
        app = root / "app"
        app.mkdir()
        for source in [worker, *dependencies]:
            destination = app / source.name
            shutil.copy2(source, destination)
            require(checksum(source) == checksum(destination), "Copied runtime checksum mismatch")
        abi_dir = root / "lib/x86_64-linux-gnu"
        abi_dir.mkdir(parents=True)
        for name in ABI:
            shutil.copy2(Path("/lib/x86_64-linux-gnu") / name, abi_dir / name)
        (root / "lib64").mkdir()
        shutil.copy2("/lib64/ld-linux-x86-64.so.2", root / "lib64/ld-linux-x86-64.so.2")
        (root / "models").mkdir()
        mounted_model = root / "models/model.gguf"
        mounted_model.touch()
        proc = root / "proc"
        proc.mkdir()
        subprocess.run([commands["mount"], "--bind", str(model), str(mounted_model)], check=True)
        mounted.append(mounted_model)
        subprocess.run([commands["mount"], "-o", "remount,bind,ro", str(mounted_model)], check=True)
        subprocess.run([commands["mount"], "-t", "proc", "-o", "ro,nosuid,nodev,noexec", "proc", str(proc)], check=True)
        mounted.append(proc)
        require(not (root / "dev").exists() and not (root / "etc").exists(),
                "Unexpected devices or driver configuration in chroot")
        require({path.name for path in abi_dir.iterdir()} == set(ABI), "Unexpected system library")
        require({path.name for path in app.iterdir()} == {worker.name, *(p.name for p in dependencies)},
                "Unexpected runtime file")
        # The worker runs with no inherited loader/ggml variables and no CPU
        # override: the absence of libvulkan must cause its actual auto fallback.
        environment = {"PATH": "/usr/bin:/bin", "LANG": "C.UTF-8"}
        process = subprocess.Popen([commands["chroot"], str(root), "/app/searchi-gemma-worker",
            "--model", "/models/model.gguf", "--threads", "1", "--device", "auto"],
            stdin=subprocess.PIPE, stdout=subprocess.PIPE, stderr=subprocess.PIPE,
            text=True, encoding="utf-8", env=environment)
        try:
            stdout, stderr = process.communicate(json.dumps({"op": "code_query",
                "text": "retry failed requests with exponential backoff"}) + "\n", timeout=180)
        except BaseException:
            process.kill()
            process.wait(timeout=10)
            raise
        require(process.returncode == 0, f"Worker failed: {stderr[-4096:]}")
        lines = [json.loads(line) for line in stdout.splitlines()]
        require(len(lines) == 2, "Unexpected protocol output")
        hello, result = lines
        require(hello.get("ready") and hello.get("device") == "cpu"
                and not hello.get("vulkan_module_loaded") and result.get("ok")
                and result.get("device") == "cpu", "Expected auto fallback to CPU without Vulkan loading")
        vectors = result.get("vectors", [])
        require(len(vectors) == 1 and len(vectors[0]) == 768
                and all(isinstance(v, (float, int)) and math.isfinite(v) for v in vectors[0]),
                "Invalid real Google vector")
        norm = math.sqrt(sum(value * value for value in vectors[0]))
        require(abs(norm - 1) < .001, "Google vector is not normalized")
        evidence = {"platform": "linux", "device_requested": "auto", "forced_cpu": False,
            "gpu_execution": False, "vulkan_loader_present": False,
            "seal": {"system_abi": list(ABI), "proc_read_only": True, "model_read_only": True,
                "private_mount_namespace_verified": True, "gpu_devices": False,
                "icd_manifests": False, "environment": environment},
            "worker_sha256": checksum(worker), "model_sha256": checksum(model),
            "hello": hello, "dimensions": len(vectors[0]), "norm": norm,
            "limitations": ["Actual GPU and driver recovery untested", "Distribution ABI compatibility untested"]}
        args.output.parent.mkdir(parents=True, exist_ok=True)
        args.output.write_text(json.dumps(evidence, indent=2) + "\n", encoding="utf-8")
        print(json.dumps({"auto_without_loader": "cpu", "cpu_module": hello["cpu_backend_module"],
                          "vulkan_loaded": False}))
    finally:
        failures = []
        for path in reversed(mounted):
            result = subprocess.run([commands["umount"], str(path)], check=False)
            if result.returncode:
                failures.append(str(path))
        # Never recursively remove a directory containing a live bind mount.
        require(not failures, f"Mount cleanup failed; preserved owned directory {root}")
        require(root.parent == base and root.name.startswith("searchi-loader-chroot-"),
                "Refusing cleanup outside owned temporary directory")
        shutil.rmtree(root)


def main():
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument("--worker", type=Path, required=True)
    parser.add_argument("--models", type=Path, required=True)
    parser.add_argument("--output", type=Path, required=True)
    parser.add_argument("--outer-pid", type=int, help=argparse.SUPPRESS)
    parser.add_argument("--outer-namespace", type=int, help=argparse.SUPPRESS)
    args = parser.parse_args()
    require(sys.platform.startswith("linux") and platform.machine().lower() in {"amd64", "x86_64"},
            "Run this check on x86_64 Linux")
    require(os.geteuid() == 0, "Run explicitly as root; this tool does not escalate privileges")
    commands = {name: shutil.which(name) for name in ("unshare", "mount", "umount", "chroot")}
    require(all(commands.values()), "Install util-linux mount tools and chroot")
    if args.outer_pid is not None or args.outer_namespace is not None:
        require(args.outer_pid is not None and args.outer_namespace is not None,
                "Incomplete namespace handoff")
        isolated_check(args, commands)
    else:
        subprocess.run([commands["unshare"], "--mount", "--fork", sys.executable, str(Path(__file__).resolve()),
            "--worker", str(args.worker.resolve()), "--models", str(args.models.resolve()),
            "--output", str(args.output.resolve()), "--outer-pid", str(os.getpid()),
            "--outer-namespace", str(namespace())], check=True)


if __name__ == "__main__":
    main()