launch.py

#!/usr/bin/env python3
"""Download and serve the public, verified Gemma 4 12B TT-native all-BFP8 P150 release.

Requires Python 3.11+, Docker, and an exclusively caller-owned P150 to serve.
No Hugging Face login, Python ML packages, or original model weights are needed.
Downloads are anonymous; serving is offline inside the ID-pinned container.
"""

import argparse
import hashlib
import json
import os
from pathlib import Path
import re
import stat
import sys
import subprocess
import tempfile
from urllib.error import URLError
from urllib.parse import quote
from urllib.request import Request, urlopen


MODEL_REPO = "Lottolabs/gemma-4-12B-it-TT-BFP8-P150"
RELEASE_ID = "p3c-20260930"
IMAGE_TAG = "lottolabs/gemma4-12b-tt-p150:p3c"
CHECKPOINT_SUBDIR = "gemma-4-12B-it"
DRAFTER_SUBDIR = "gemma-4-12B-it-assistant"
HF_BASE = "https://huggingface.co"
SHA256 = re.compile(r"[0-9a-f]{64}")
CHUNK_SIZE = 1024 * 1024


def request(url):
    # urllib does not load HF tokens, netrc credentials, or ML dependencies.
    return urlopen(Request(url, headers={"User-Agent": "gemma4-12b-native-launch/1"}), timeout=120)


def fetch_json(url):
    with request(url) as response:
        raw = response.read(16 * CHUNK_SIZE + 1)
    if len(raw) > 16 * CHUNK_SIZE:
        raise ValueError("Remote JSON exceeds the 16 MiB metadata limit")
    value = json.loads(raw)
    if not isinstance(value, dict):
        raise ValueError("Expected a JSON object")
    return value, raw


def safe_relative(value):
    if (not isinstance(value, str) or not value
            or any(part in ("", ".", "..") for part in value.split("/"))
            or "\\" in value or ":" in value
            or any(ord(char) < 32 or ord(char) == 127 for char in value)):
        raise ValueError(f"Unsafe release path: {value!r}")
    return value


def file_records(value):
    if not isinstance(value, dict) or not value:
        raise ValueError("Manifest must contain a nonempty files object")
    for relative, record in value.items():
        safe_relative(relative)
        if (not isinstance(record, dict)
                or not isinstance(record.get("sha256"), str)
                or SHA256.fullmatch(record["sha256"]) is None
                or type(record.get("bytes")) is not int or record["bytes"] < 0):
            raise ValueError(f"Invalid file checksum/size: {relative}")
        parts = relative.split("/")
        if any("/".join(parts[:index]) in value for index in range(1, len(parts))):
            raise ValueError(f"File/directory collision: {relative}")
    return value


def validate_release(release):
    if (type(release.get("schema_version")) is not int or release["schema_version"] != 1
            or release.get("model_repo") != MODEL_REPO
            or release.get("release_id") != RELEASE_ID
            or release.get("checkpoint_subdir") != CHECKPOINT_SUBDIR
            or release.get("drafter_subdir") != DRAFTER_SUBDIR):
        raise ValueError("Unsupported runtime release identity or schema")
    image = release.get("image")
    image_id = release.get("image_id")
    image_archive = release.get("image_archive")
    if (image != IMAGE_TAG or not isinstance(image_id, str)
            or re.fullmatch(r"sha256:[0-9a-f]{64}", image_id) is None
            or image_archive != "runtime-image.tar.gz"):
        raise ValueError("Release must identify the checksum-pinned bundled runtime image")
    files = file_records(release.get("files"))
    target_manifest = f"{CHECKPOINT_SUBDIR}/native_manifest.json"
    drafter_manifest = f"{DRAFTER_SUBDIR}/drafter_manifest.json"
    required = {"serve_native.py", "native_checkpoint.py", target_manifest, f"{CHECKPOINT_SUBDIR}/equivalence.json",
                drafter_manifest, f"{DRAFTER_SUBDIR}/spec_equivalence.json", image_archive}
    if not required.issubset(files):
        raise ValueError("Release is missing the server, native manifests, or equivalence proofs")
    if "runtime-release.json" in files or "launch.py" in files:
        raise ValueError("Release inventory must exclude runtime-release.json and launch.py")
    if (release.get("checkpoint_manifest_sha256") != files[target_manifest]["sha256"]
            or release.get("drafter_manifest_sha256") != files[drafter_manifest]["sha256"]):
        raise ValueError("Release manifest checksums do not match its file inventory")
    return files


def ensure_image(image, expected_id, archive):
    inspect = subprocess.run(
        ["docker", "image", "inspect", image, "--format", "{{.Id}}"],
        text=True, capture_output=True)
    if inspect.returncode != 0 or inspect.stdout.strip() != expected_id:
        print(f"Loading {archive.name} into Docker", file=sys.stderr, flush=True)
        loaded = subprocess.run(["docker", "load", "--input", str(archive)], text=True)
        if loaded.returncode != 0:
            raise ValueError("Docker could not load the verified runtime image archive")
        inspect = subprocess.run(
            ["docker", "image", "inspect", image, "--format", "{{.Id}}"],
            text=True, capture_output=True)
    if inspect.returncode != 0 or inspect.stdout.strip() != expected_id:
        raise ValueError("Loaded runtime image identity does not match runtime-release.json")


def cache_path(root, relative, directory=False):
    """Do not follow pre-existing links or special files within a snapshot."""
    current = root
    parts = safe_relative(relative).split("/")
    for index, part in enumerate(parts):
        current = current / part
        if current.is_symlink():
            raise ValueError(f"Refusing a symlink in the release cache: {current}")
        if current.exists():
            mode = current.stat().st_mode
            expected = stat.S_ISDIR if directory or index < len(parts) - 1 else stat.S_ISREG
            if not expected(mode):
                raise ValueError(f"Unexpected cache entry type: {current}")
    return current


def verified(path, record):
    if not path.is_file() or path.stat().st_size != record["bytes"]:
        return False
    with path.open("rb") as source:
        return hashlib.file_digest(source, "sha256").hexdigest() == record["sha256"]


def atomic_write(destination, chunks, record=None):
    destination.parent.mkdir(parents=True, exist_ok=True)
    temporary = None
    try:
        with tempfile.NamedTemporaryFile(prefix=".download-", dir=destination.parent, delete=False) as target:
            temporary = Path(target.name)
            digest = hashlib.sha256()
            size = 0
            for chunk in chunks:
                size += len(chunk)
                if record is not None and size > record["bytes"]:
                    raise ValueError(f"Download exceeds expected size: {destination.name}")
                target.write(chunk)
                digest.update(chunk)
            if record is not None and (size != record["bytes"] or digest.hexdigest() != record["sha256"]):
                raise ValueError(f"Download failed SHA256/size verification: {destination}")
            target.flush()
            os.fsync(target.fileno())
            # The container must be able to read the read-only checkpoint mounts.
            os.fchmod(target.fileno(), 0o644)
        os.replace(temporary, destination)
    finally:
        if temporary is not None:
            temporary.unlink(missing_ok=True)


def download(snapshot, base_url, relative, record):
    destination = cache_path(snapshot, relative)
    if verified(destination, record):
        return
    print(f"Downloading {relative} ({record['bytes']:,} bytes)", file=sys.stderr, flush=True)
    with request(base_url + quote(relative, safe="/")) as response:
        atomic_write(destination, iter(lambda: response.read(CHUNK_SIZE), b""), record)


def covered(snapshot, files, subdir, manifest_name, excluded):
    """The published inventory must be exactly the manifest's file set plus the manifest and its proof."""
    manifest = json.loads((snapshot / subdir / manifest_name).read_bytes())
    if not isinstance(manifest, dict):
        raise ValueError(f"{subdir}/{manifest_name} must be a JSON object")
    listed = file_records(manifest.get("files"))
    for relative, record in listed.items():
        published = files.get(f"{subdir}/{relative}")
        if published is None or any(published[key] != record[key] for key in ("sha256", "bytes")):
            raise ValueError(f"Release inventory does not cover {subdir}: {relative}")
    extra = {name for name in files if name.startswith(subdir + "/")} - {f"{subdir}/{name}" for name in (*listed, *excluded)}
    if extra:
        raise ValueError(f"Release inventory has files outside the {subdir} manifest: {sorted(extra)[:3]}")


def independent_kernel_cache(snapshot, kernel_cache):
    if (kernel_cache == snapshot or kernel_cache.is_relative_to(snapshot)
            or snapshot.is_relative_to(kernel_cache)):
        raise ValueError("Kernel cache must be independent of the downloaded release directory")
    if ":" in str(snapshot) or ":" in str(kernel_cache):
        raise ValueError("Docker bind-mount paths cannot contain ':'")
    kernel_cache.mkdir(parents=True, exist_ok=True)
    with tempfile.TemporaryFile(dir=kernel_cache) as probe:
        probe.write(b"writable")
        probe.flush()


def parser():
    cache_home = Path(os.environ.get("XDG_CACHE_HOME") or Path.home() / ".cache")
    result = argparse.ArgumentParser(description=__doc__, formatter_class=argparse.RawDescriptionHelpFormatter,
        epilog="""Examples:
  python3 launch.py --download-only --cache-root /large-disk/gemma4-12b
  python3 launch.py --cache-root /large-disk/gemma4-12b --print-command
  python3 launch.py --cache-root /large-disk/gemma4-12b --device-ownership-confirmed
  python3 launch.py --speculation off --port 8001 --device-ownership-confirmed

--print-command downloads/verifies the release but does not run Docker.
--download-only downloads/verifies without requiring Docker or device ownership.
Serving never stops containers, resets the device, or takes ownership for you.
Stop your other accelerator workloads yourself before confirming ownership.
The API binds to 127.0.0.1 only. Keep the printed immutable revision to reproduce.
""")
    result.add_argument("--cache-root", type=Path, default=cache_home / "gemma4-12b-tt-native",
                        help="Download/cache root (default: $XDG_CACHE_HOME/gemma4-12b-tt-native or ~/.cache/gemma4-12b-tt-native)")
    result.add_argument("--kernel-cache", type=Path,
                        help="Independent writable TT kernel cache (default: CACHE_ROOT/kernel-cache)")
    result.add_argument("--port", type=int, default=8000, help="Loopback API port (default: 8000)")
    result.add_argument("--speculation", choices=("on", "off"), default="on",
                        help="Greedy speculative decoding with the proven assistant drafter, K=5 (default: on)")
    result.add_argument("--max-model-len", type=int, default=262144,
                        help="Context limit, 1024–262144 (default: 262144)")
    result.add_argument("--name", default="gemma4-12b-native-api",
                        help="Docker container name (default: gemma4-12b-native-api)")
    result.add_argument("--revision", default="main", help="Public HF branch, tag, or commit (default: main)")
    result.add_argument("--device-ownership-confirmed", action="store_true",
                        help="Confirm you exclusively own the P150 and have stopped your other workloads")
    mode = result.add_mutually_exclusive_group()
    mode.add_argument("--print-command", action="store_true", help="Download, verify, and print the Docker command without running it")
    mode.add_argument("--download-only", action="store_true", help="Download and verify the native package without serving")
    return result


def main():
    cli = parser()
    args = cli.parse_args()
    if sys.version_info < (3, 11):
        cli.error("Python 3.11 or newer is required")
    if not 1 <= args.port <= 65535 or not 1024 <= args.max_model_len <= 262144:
        cli.error("Port must be 1–65535 and max-model-len must be 1024–262144")
    if re.fullmatch(r"[A-Za-z0-9][A-Za-z0-9_.-]*", args.name) is None:
        cli.error("Invalid Docker container name")
    if not args.revision.strip():
        cli.error("Revision must not be empty")
    if not (args.print_command or args.download_only or args.device_ownership_confirmed):
        cli.error("Confirm exclusive P150 ownership with --device-ownership-confirmed, or use --print-command/--download-only")
    try:
        info, _ = fetch_json(f"{HF_BASE}/api/models/{MODEL_REPO}/revision/{quote(args.revision, safe='')}")
        commit = info.get("sha")
        if not isinstance(commit, str) or re.fullmatch(r"[0-9a-f]{40}", commit) is None:
            raise ValueError("Hugging Face did not resolve the revision to an immutable commit")
        base_url = f"{HF_BASE}/{MODEL_REPO}/resolve/{commit}/"
        release, release_raw = fetch_json(base_url + "runtime-release.json")
        files = validate_release(release)
        root = args.cache_root.expanduser().resolve()
        root.mkdir(parents=True, exist_ok=True)
        snapshot = cache_path(root, f"snapshots/{commit}", directory=True)
        snapshot.mkdir(parents=True, exist_ok=True)
        kernel_cache = (args.kernel_cache.expanduser() if args.kernel_cache else root / "kernel-cache").resolve()
        independent_kernel_cache(snapshot, kernel_cache)
        print(f"Revision: {commit}\nImage: {release['image']} ({release['image_id']})"
              f"\nCheckpoint: {snapshot / CHECKPOINT_SUBDIR}\nDrafter: {snapshot / DRAFTER_SUBDIR}"
              f"\nKernel cache: {kernel_cache}", file=sys.stderr, flush=True)
        # Check both native inventories before downloading the multi-gigabyte tensors.
        for subdir, manifest_name, proof in ((CHECKPOINT_SUBDIR, "native_manifest.json", "equivalence.json"),
                                             (DRAFTER_SUBDIR, "drafter_manifest.json", "spec_equivalence.json")):
            relative = f"{subdir}/{manifest_name}"
            download(snapshot, base_url, relative, files[relative])
            covered(snapshot, files, subdir, manifest_name, (manifest_name, proof))
        for relative, record in files.items():
            download(snapshot, base_url, relative, record)
        atomic_write(cache_path(snapshot, "runtime-release.json"), (release_raw,))
        print(f"All {len(files)} release files verified.", file=sys.stderr, flush=True)
        if args.download_only:
            return 0
        ensure_image(release["image"], release["image_id"], snapshot / release["image_archive"])
        command = [sys.executable, "-B", str(snapshot / "serve_native.py"),
                   "--cache-root", str(kernel_cache), "--image", release["image"],
                   "--port", str(args.port), "--name", args.name,
                   "--max-model-len", str(args.max_model_len)]
        if args.speculation == "off":
            command.append("--no-speculation")
        if args.print_command:
            command.append("--print-command")
        if args.device_ownership_confirmed:
            command.append("--device-ownership-confirmed")
        os.execv(sys.executable, command)
    except (OSError, URLError, ValueError) as error:
        cli.exit(1, f"launch.py: {error}\n")
    return 0


if __name__ == "__main__":
    raise SystemExit(main())