serving/launch.py

#!/usr/bin/env python3
# SPDX-License-Identifier: Apache-2.0
# Copyright 2026 DevelopingDad
"""Print or run the optional pinned Gemma4 Fast vLLM recipe."""
import argparse
import hashlib
import json
import os
from pathlib import Path
import shlex
import subprocess
import sys

ROOT = Path(__file__).resolve().parent.parent
IMAGE = "vllm/vllm-openai@sha256:3f7dd5b777d34d1724456ce71f87385dca288c3bb23029ab27dee358f5d2b971"
SITE = "/usr/local/lib/python3.12/dist-packages"
MODEL_ID = "DevelopingDad/Gemma4-12B-it-Fast"
ASSISTANT_REVISION = "46d4c6f13f0ac0ad827b915669b8df9b81c64c51"

def command(a):
    model = a.model_dir.resolve()
    draft = a.assistant_dir.resolve()
    cache = a.cache_dir.resolve()
    serving = ROOT / "serving"
    argv = ["docker", "run", "--name", a.name, "--user", str(a.uid) + ":" + str(a.gid),
            "--gpus", "all", "--network", "host", "--ipc", "host",
            "--memory", "96g", "--memory-swap", "96g", "--cap-drop", "ALL",
            "--security-opt", "no-new-privileges", "--workdir", "/cache",
            "--restart", "no"]
    for name, value in {
        "HF_HOME": "/cache/huggingface", "HF_HUB_OFFLINE": "1",
        "TRANSFORMERS_OFFLINE": "1", "XDG_CACHE_HOME": "/cache",
        "VLLM_CACHE_ROOT": "/cache/vllm",
        "TORCHINDUCTOR_CACHE_DIR": "/cache/inductor", "TRITON_CACHE_DIR": "/cache/triton",
        "FLASHINFER_WORKSPACE_BASE": "/cache/flashinfer",
        "PYTORCH_CUDA_ALLOC_CONF": "expandable_segments:True",
        "VLLM_EXPERIMENTAL_DRAFT_HEAD_FP8": "1",
        "VLLM_DRAFT_HEAD_FP8_RECEIPT": "/cache/draft-head-fp8-receipt.json",
    }.items():
        argv.extend(["-e", name + "=" + value])
    for source, dest, readonly in [
        (model, "/model", True), (draft, "/draft-model", True),
        (cache, "/cache", False), (serving, "/profile", True),
    ]:
        argv.extend(["--mount", "type=bind,src=" + str(source) + ",dst=" + dest + (",readonly" if readonly else "")])
    manifest = json.loads((serving / "overlay-manifest.json").read_text())
    for f in manifest["files"]:
        install = f["install_destination"]
        if install:
            source = ROOT / f["path"]
            argv.extend(["--mount", "type=bind,src=" + str(source) + ",dst=" + SITE + "/" + install + ",readonly"])
    argv.extend(["--entrypoint", "python3", IMAGE, "/profile/receipt_bootstrap.py",
        "/model", "--chat-template", "/profile/native-compat-chat-template.jinja",
        "--chat-template-content-format", "openai", "--served-model-name", MODEL_ID,
        "--host", "127.0.0.1", "--port", str(a.port), "--dtype", "bfloat16",
        "--tensor-parallel-size", "1", "--max-model-len", "65536",
        "--max-num-seqs", "8", "--max-num-batched-tokens", "4096",
        "--enable-chunked-prefill", "--kv-cache-dtype", "fp8",
        "--kv-cache-memory-bytes", "64424509440", "--gpu-memory-utilization", "0.65",
        "--reasoning-parser", "gemma4", "--enable-auto-tool-choice",
        "--tool-call-parser", "gemma4",
        "--default-chat-template-kwargs", json.dumps({"enable_thinking": False}),
        "--limit-mm-per-prompt", json.dumps({"image": 4, "audio": 1}),
        "--speculative-config", json.dumps({"method": "mtp", "model": "/draft-model",
            "num_speculative_tokens": 8, "quantization": "fp8_per_tensor"}),
        "--no-enable-prefix-caching" if a.disable_prefix_caching else "--enable-prefix-caching",
        "--long-prefill-token-threshold", "1024", "--long-prefill-token-threshold-adaptive"])
    return argv

def preflight(a):
    if not (a.model_dir / "model.safetensors").is_file():
        raise RuntimeError("model.safetensors is missing; download the model repository first")
    if not (a.assistant_dir / "model.safetensors").is_file():
        raise RuntimeError("assistant model is missing; download the pinned Google assistant first")
    artifacts = json.loads((ROOT / "serving/artifact-manifest.json").read_text())
    if artifacts["assistant_revision"] != ASSISTANT_REVISION:
        raise RuntimeError("assistant revision fence changed")
    for directory, files in [(a.model_dir, artifacts["target_files"]),
                             (a.assistant_dir, artifacts["assistant_files"])]:
        for f in files:
            path = directory / f["name"]
            if not path.is_file() or path.stat().st_size != f["bytes"]:
                raise RuntimeError("artifact missing or size mismatch: " + f["name"])
            digest = hashlib.sha256()
            with path.open("rb") as handle:
                for chunk in iter(lambda: handle.read(16 * 1024 * 1024), b""):
                    digest.update(chunk)
            if digest.hexdigest() != f["sha256"]:
                raise RuntimeError("artifact checksum mismatch: " + f["name"])
    for f in artifacts.get("serving_files", []):
        path = ROOT / f["name"]
        if not path.is_file() or hashlib.sha256(path.read_bytes()).hexdigest() != f["sha256"]:
            raise RuntimeError("serving asset checksum mismatch: " + f["name"])
    manifest = json.loads((ROOT / "serving/overlay-manifest.json").read_text())
    for f in manifest["files"]:
        if hashlib.sha256((ROOT / f["path"]).read_bytes()).hexdigest() != f["published_sha256"]:
            raise RuntimeError("overlay checksum mismatch: " + f["path"])
    # Check the actual base-image source bytes without allocating the GPU.
    check = {
        f["install_destination"]: f["upstream_original_sha256"]
        for f in manifest["files"] if f["upstream_original_sha256"]
    }
    code = "from pathlib import Path;import hashlib,json;expected=json.loads(" + repr(json.dumps(check)) + ");base=Path(" + repr(SITE) + ");assert all(hashlib.sha256((base/p).read_bytes()).hexdigest()==h for p,h in expected.items()),'Pinned runtime source mismatch'"
    subprocess.run(["docker", "run", "--rm", "--network", "none",
                    "--entrypoint", "python3", IMAGE, "-c", code], check=True)
    a.cache_dir.mkdir(parents=True, exist_ok=True)

def main():
    p = argparse.ArgumentParser(description=__doc__)
    p.add_argument("--model-dir", type=Path, default=ROOT)
    p.add_argument("--assistant-dir", type=Path, default=ROOT / "assistant")
    p.add_argument("--cache-dir", type=Path, default=ROOT / "runtime-cache")
    p.add_argument("--name", default="gemma4-fast-server")
    p.add_argument("--port", type=int, default=8000)
    p.add_argument("--uid", type=int, default=os.getuid())
    p.add_argument("--gid", type=int, default=os.getgid())
    p.add_argument("--disable-prefix-caching", action="store_true",
                   help="Disable prefix reuse for cold-performance comparisons")
    p.add_argument("--run", action="store_true",
                   help="Execute the printed plan; default prints only")
    a = p.parse_args()
    if not 1 <= a.port <= 65535:
        p.error("port must be between 1 and 65535")
    if a.run:
        preflight(a)
    argv = command(a)
    print(shlex.join(argv), flush=True)
    if a.run:
        return subprocess.run(argv).returncode
    return 0

if __name__ == "__main__":
    sys.exit(main())