serve.py

#!/usr/bin/env python
import os, subprocess, sys

# The optimized weights location
WEIGHTS_BUCKET = os.environ.get(
    "WEIGHTS_BUCKET", "hf://buckets/gemma-challenge/gemma-foffee/weights/int4-lmhead")
LOCAL = os.environ.get("LOCAL_MODEL_DIR", "/tmp/int4-qat-lmhead-v1")

if not os.path.isdir(LOCAL) or not os.path.exists(os.path.join(LOCAL, "config.json")):
    print(f"[serve] syncing weights {WEIGHTS_BUCKET} -> {LOCAL}", flush=True)
    subprocess.run(["hf", "buckets", "sync", WEIGHTS_BUCKET, LOCAL], check=True)

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

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