"""
Modal inference server — Qwen/Qwen2.5-32B-Instruct on A100-80GB via vLLM.
8B fallback: change MODEL_NAME → "NousResearch/Hermes-3-Llama-3.1-8B" and gpu → "L4".
NOTE: the app's MODAL_MODEL_ID env var must match MODEL_NAME below, or the OpenAI
client requests a model the server does not serve.

Deploy:
    modal deploy modal-setup/inference_app.py

First cold start downloads weights into the huggingface-cache volume (Qwen2.5-32B ≈ 65GB,
~10–15 min). Subsequent cold starts reuse the cached weights (~2 min to load + warm up).

Required Modal secrets (create once, never again):
    modal secret create inference-auth MODAL_API_KEY=<any-random-string>
    modal secret create hf-secret HF_TOKEN=<your-hf-token>
"""

import subprocess

import modal

# ── Image ─────────────────────────────────────────────────────────────────────
# Follows the official modal-examples vllm_inference.py pattern.

vllm_image = (
    modal.Image.from_registry("nvidia/cuda:12.9.0-devel-ubuntu22.04", add_python="3.12")
    .entrypoint([])
    .uv_pip_install("vllm==0.21.0")
    .env({"HF_XET_HIGH_PERFORMANCE": "1"})  # hf-xet is a vllm transitive dep; this enables its high-perf mode
)

# ── Volumes (weight cache + vLLM JIT cache) ───────────────────────────────────

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

# ── Constants ─────────────────────────────────────────────────────────────────

MODEL_NAME = "Qwen/Qwen2.5-32B-Instruct"
VLLM_PORT = 8000
MINUTES = 60

# ── App ───────────────────────────────────────────────────────────────────────

app = modal.App("kobo-analyst-inference")


@app.function(
    image=vllm_image,
    gpu="A100-80GB",
    scaledown_window=5 * MINUTES,
    timeout=10 * MINUTES,
    volumes={
        "/root/.cache/huggingface": hf_cache_vol,
        "/root/.cache/vllm": vllm_cache_vol,
    },
    secrets=[
        modal.Secret.from_name("inference-auth"),   # provides MODAL_API_KEY
        modal.Secret.from_name("hf-secret"),         # provides HF_TOKEN
    ],
)
@modal.concurrent(max_inputs=10)
@modal.web_server(port=VLLM_PORT, startup_timeout=10 * MINUTES)
def serve():
    import os

    cmd = [
        "vllm", "serve", MODEL_NAME,
        "--served-model-name", MODEL_NAME,
        "--host", "0.0.0.0",
        "--port", str(VLLM_PORT),
        "--max-model-len", "32768",
        "--enforce-eager",
        "--enable-auto-tool-choice",
        "--tool-call-parser", "hermes",
        "--api-key", os.environ["MODAL_API_KEY"],
    ]

    subprocess.Popen(cmd)