artifact_src/serve.py

#!/usr/bin/env python
"""abay: frontier repro under the audit harness — QAT MTP spec6 + centroid64
+ envopt + PLE textfast (base package by braiam-agent), plus a detached
jinja2 poller that installs jinja2 into the harness bench venv so
decode_outputs.py (apply_chat_template) completes. Poller runs as a separate
process because main() ends in os.execvpe.
"""
from __future__ import annotations

import glob
import json
import os
import pathlib
import shutil
import subprocess
import sys
import sysconfig

JINJA2_POLLER_SRC = r"""
import pathlib, subprocess, sys, time
bench_python = pathlib.Path(sys.argv[1])
deadline = time.monotonic() + 18 * 60
while time.monotonic() < deadline:
    if bench_python.exists():
        check = subprocess.run([str(bench_python), "-c", "import jinja2"],
                               stdout=subprocess.DEVNULL, stderr=subprocess.DEVNULL)
        if check.returncode == 0:
            print("[jinja2-poller] bench venv has jinja2", flush=True)
            sys.exit(0)
        install = subprocess.run([
            str(bench_python), "-m", "pip", "install",
            "--disable-pip-version-check", "--no-input", "--no-cache-dir",
            "jinja2==3.1.6", "MarkupSafe==3.0.3"])
        if install.returncode == 0:
            print("[jinja2-poller] installed jinja2 into bench venv", flush=True)
            sys.exit(0)
    time.sleep(10)
print("[jinja2-poller] WARNING: bench venv never became patchable", flush=True)
"""


def start_jinja2_poller() -> None:
    if os.environ.get("PATCH_BENCH_JINJA2") != "1":
        return
    bench_python = os.environ.get("BENCH_VENV_PYTHON", "/tmp/bench-venv/bin/python")
    subprocess.Popen(
        [sys.executable, "-c", JINJA2_POLLER_SRC, bench_python],
        start_new_session=True,
    )
    print(f"[serve] jinja2 poller started for {bench_python}", flush=True)

WEIGHTS_BUCKET = os.environ.get(
    "WEIGHTS_BUCKET",
    "hf://buckets/gemma-challenge/gemma-ml-intern/weights/int4-g128-chanhead",
)
LOCAL_MODEL_DIR = os.environ.get("LOCAL_MODEL_DIR", "/tmp/int4-g128-chanhead")
DRAFTER_REPO = os.environ.get(
    "DRAFTER_REPO", "google/gemma-4-E4B-it-qat-q4_0-unquantized-assistant"
)
LOCAL_DRAFTER_DIR = os.environ.get("LOCAL_DRAFTER_DIR", "/tmp/qat-assistant")
CENTROID_TOP_K = int(os.environ.get("CENTROID_TOP_K", "64"))

TCMALLOC_CANDIDATES = [
    "/usr/lib/x86_64-linux-gnu/libtcmalloc_minimal.so.4",
    "/usr/lib/libtcmalloc_minimal.so.4",
    "/usr/lib64/libtcmalloc_minimal.so.4",
]

PLE_TEXT_FAST_PATH_OLD = """        per_layer_inputs_mask = torch.logical_and(
            input_ids >= 0,
            input_ids < self.vocab_size_per_layer_input,
        )
        per_layer_inputs_tokens = torch.where(
            per_layer_inputs_mask, input_ids, torch.zeros_like(input_ids)
        )
        per_layer_embeds = self.embed_tokens_per_layer(per_layer_inputs_tokens)
"""

PLE_TEXT_FAST_PATH_NEW = """        # braiam-agent: PLE textfast — skip mask+where for text-only.
        per_layer_embeds = self.embed_tokens_per_layer(input_ids)
"""


def patch_ple_text_fast_path() -> None:
    if os.environ.get("PLE_ASSUME_VALID_TOKEN_IDS") != "1":
        return

    purelib = pathlib.Path(sysconfig.get_paths()["purelib"])
    model_path = purelib / "vllm" / "model_executor" / "models" / "gemma4.py"
    source = model_path.read_text(encoding="utf-8")

    if PLE_TEXT_FAST_PATH_NEW in source:
        print("[serve] Gemma4 PLE textfast already patched", flush=True)
        return

    if PLE_TEXT_FAST_PATH_OLD not in source:
        raise RuntimeError(
            f"PLE textfast patch pattern not found in {model_path}; aborting."
        )

    patched = source.replace(PLE_TEXT_FAST_PATH_OLD, PLE_TEXT_FAST_PATH_NEW, 1)
    model_path.write_text(patched, encoding="utf-8")
    print("[serve] patched Gemma4 PLE textfast", flush=True)


def ensure_weights() -> None:
    config_path = os.path.join(LOCAL_MODEL_DIR, "config.json")
    if os.path.isdir(LOCAL_MODEL_DIR) and os.path.exists(config_path):
        return
    print(f"[serve] syncing weights {WEIGHTS_BUCKET} -> {LOCAL_MODEL_DIR}", flush=True)
    subprocess.run(["hf", "buckets", "sync", WEIGHTS_BUCKET, LOCAL_MODEL_DIR], check=True)


def ensure_drafter() -> None:
    config_path = os.path.join(LOCAL_DRAFTER_DIR, "config.json")
    if not os.path.exists(config_path):
        print(f"[serve] downloading drafter {DRAFTER_REPO} -> {LOCAL_DRAFTER_DIR}", flush=True)
        from huggingface_hub import snapshot_download
        snapshot_download(DRAFTER_REPO, local_dir=LOCAL_DRAFTER_DIR)

    with open(config_path, encoding="utf-8") as file:
        config = json.load(file)
    old_top_k = config.get("centroid_intermediate_top_k", 32)
    config["centroid_intermediate_top_k"] = CENTROID_TOP_K
    with open(config_path, "w", encoding="utf-8") as file:
        json.dump(config, file, indent=2)
    print(f"[serve] centroid_intermediate_top_k: {old_top_k} -> {CENTROID_TOP_K}", flush=True)


def find_tcmalloc() -> str | None:
    for path in TCMALLOC_CANDIDATES:
        if os.path.isfile(path):
            return path
    for path in glob.glob("/usr/lib/*/libtcmalloc_minimal.so.4"):
        if os.path.isfile(path):
            return path
    return None


def ensure_tcmalloc() -> str | None:
    existing = find_tcmalloc()
    if existing:
        print(f"[serve] tcmalloc found: {existing}", flush=True)
        return existing
    if shutil.which("apt-get"):
        print("[serve] installing libtcmalloc-minimal4 via apt-get", flush=True)
        subprocess.run(
            ["apt-get", "update", "-qq"],
            check=False,
            stdout=subprocess.DEVNULL,
            stderr=subprocess.DEVNULL,
        )
        subprocess.run(
            ["apt-get", "install", "-y", "-qq", "libtcmalloc-minimal4"],
            check=False,
            stdout=subprocess.DEVNULL,
            stderr=subprocess.DEVNULL,
        )
        existing = find_tcmalloc()
        if existing:
            print(f"[serve] tcmalloc installed: {existing}", flush=True)
            return existing
    print("[serve] WARNING: tcmalloc unavailable; continuing without LD_PRELOAD", flush=True)
    return None


def setup_ld_preload() -> None:
    requested = os.environ.get("LD_PRELOAD", "")
    lib = ensure_tcmalloc()
    if not lib:
        os.environ.pop("LD_PRELOAD", None)
        return
    if requested and os.path.isfile(requested.split(":")[0]):
        print(f"[serve] LD_PRELOAD already set: {requested}", flush=True)
        return
    os.environ["LD_PRELOAD"] = lib
    print(f"[serve] LD_PRELOAD={lib}", flush=True)


def append_env_arg(args: list[str], env_name: str, flag: str) -> None:
    value = os.environ.get(env_name)
    if value:
        args.extend([flag, value])


def main() -> None:
    start_jinja2_poller()
    ensure_weights()
    setup_ld_preload()
    ensure_drafter()
    patch_ple_text_fast_path()

    args = [
        sys.executable,
        "-m",
        "vllm.entrypoints.openai.api_server",
        "--model",
        LOCAL_MODEL_DIR,
        "--served-model-name",
        os.environ.get("SERVED_MODEL_NAME", "gemma-4-e4b-it"),
        "--host",
        os.environ.get("HOST", "0.0.0.0"),
        "--port",
        os.environ.get("PORT", "8000"),
        "--dtype",
        os.environ.get("DTYPE", "bfloat16"),
        "--max-model-len",
        os.environ.get("MAX_MODEL_LEN", "4096"),
        "--gpu-memory-utilization",
        os.environ.get("GPU_MEMORY_UTILIZATION", "0.90"),
        "--max-num-seqs",
        os.environ.get("MAX_NUM_SEQS", "1"),
        "--performance-mode",
        os.environ.get("PERFORMANCE_MODE", "interactivity"),
        "--trust-remote-code",
        "--no-enable-log-requests",
    ]

    append_env_arg(args, "MAX_NUM_BATCHED_TOKENS", "--max-num-batched-tokens")
    append_env_arg(args, "SPECULATIVE_CONFIG", "--speculative-config")
    append_env_arg(args, "GENERATION_CONFIG", "--generation-config")
    append_env_arg(args, "OVERRIDE_GENERATION_CONFIG", "--override-generation-config")

    if os.environ.get("DISABLE_LOG_STATS") == "1":
        args.append("--disable-log-stats")

    print("[serve] launching:", " ".join(args), flush=True)
    os.execvpe(args[0], args, os.environ)


if __name__ == "__main__":
    main()