from __future__ import annotations

import modal

APP_NAME = "jackailocal-gemma-4-12b-it-vllm"
MODEL_NAME = "google/gemma-4-12B-it"
VLLM_PORT = 8000
MINUTES = 60

app = modal.App(APP_NAME)

hf_cache_vol = modal.Volume.from_name("jackailocal-gemma-hf-cache", create_if_missing=True)
vllm_cache_vol = modal.Volume.from_name("jackailocal-gemma-vllm-cache", create_if_missing=True)

image = (
    modal.Image.from_registry(
        "nvidia/cuda:12.9.0-devel-ubuntu22.04",
        add_python="3.12",
    )
    .entrypoint([])
    .uv_pip_install(
        "vllm",
        "huggingface-hub",
        "fastapi[standard]",
    )
    .env(
        {
            "HF_XET_HIGH_PERFORMANCE": "1",
            "VLLM_LOG_STATS_INTERVAL": "10",
        }
    )
)


@app.function(
    image=image,
    gpu="A100-80GB",
    timeout=10 * MINUTES,
    scaledown_window=10 * MINUTES,
    secrets=[modal.Secret.from_name("gemma-secrets")],
    volumes={
        "/root/.cache/huggingface": hf_cache_vol,
        "/root/.cache/vllm": vllm_cache_vol,
    },
)
@modal.concurrent(max_inputs=32)
@modal.web_server(port=VLLM_PORT, startup_timeout=10 * MINUTES)
def serve() -> None:
    import os
    import subprocess

    api_key = os.environ["VLLM_API_KEY"]

    cmd = [
        "vllm",
        "serve",
        MODEL_NAME,
        "--served-model-name",
        MODEL_NAME,
        "--host",
        "0.0.0.0",
        "--port",
        str(VLLM_PORT),
        "--api-key",
        api_key,
        "--max-model-len",
        "8192",
        "--gpu-memory-utilization",
        "0.90",
        "--uvicorn-log-level",
        "info",
    ]

    subprocess.Popen(cmd).wait()