modal_app.py

"""Modal inference service scaffold for Snap2Sim.

The endpoints intentionally keep placeholder inference as the default until the
Nemotron runtime passes a GPU smoke test. The deployment helpers here pin the
target model assets and make that smoke test explicit instead of silently
claiming multimodal GGUF support before image input has been proven.
"""

from __future__ import annotations

import os
from pathlib import Path
from typing import Any
import secrets as token_secrets

import modal
from fastapi import Header, HTTPException

from snap2sim.model_io import coerce_analysis_response, parse_analysis_response
from snap2sim.prompts import build_vision_messages
from snap2sim.schema import EXAMPLE_ANALYSIS, select_render_mode, validate_analysis


DEFAULT_MODEL_REPO = "unsloth/NVIDIA-Nemotron-3-Nano-Omni-30B-A3B-Reasoning-GGUF"
DEFAULT_GGUF_QUANT = "UD-Q4_K_M"
DEFAULT_MMPROJ_FILE = "mmproj-F16.gguf"
DEFAULT_RUNTIME_MODE = "placeholder"
CACHE_DIR = "/cache"
HF_CACHE_DIR = f"{CACHE_DIR}/huggingface"
MODEL_ASSET_DIR = f"{CACHE_DIR}/models"
MAX_IMAGE_BASE64_CHARS = 12 * 1024 * 1024
MAX_IMAGE_BYTES = 9 * 1024 * 1024
MAX_IMAGE_PIXELS = 12_000_000


model_cache = modal.Volume.from_name("snap2sim-hf-cache", create_if_missing=True)
api_auth_secret = modal.Secret.from_name("snap2sim-api-auth")

image = (
    modal.Image.debian_slim(python_version="3.11")
    .pip_install(
        "fastapi[standard]",
        "huggingface_hub[hf_xet]",
        "pillow",
        "requests",
    )
    .env(
        {
            "HF_HUB_CACHE": HF_CACHE_DIR,
            "HF_XET_HIGH_PERFORMANCE": "1",
        }
    )
    .add_local_python_source("snap2sim")
)

llamacpp_image = (
    modal.Image.from_registry("nvidia/cuda:12.9.0-devel-ubuntu22.04", add_python="3.11")
    .entrypoint([])
    .apt_install("build-essential", "cmake", "curl", "git", "libcurl4-openssl-dev")
    .pip_install("fastapi[standard]", "huggingface_hub[hf_xet]", "pillow")
    .env(
        {
            "HF_HUB_CACHE": HF_CACHE_DIR,
            "HF_XET_HIGH_PERFORMANCE": "1",
            "LIBRARY_PATH": "/usr/local/cuda/lib64/stubs:/usr/local/cuda/targets/x86_64-linux/lib/stubs",
        }
    )
    .run_commands(
        "if [ -f /usr/local/cuda/lib64/stubs/libcuda.so ] && [ ! -f /usr/local/cuda/lib64/stubs/libcuda.so.1 ]; then ln -s /usr/local/cuda/lib64/stubs/libcuda.so /usr/local/cuda/lib64/stubs/libcuda.so.1; fi",
        "if [ -f /usr/local/cuda/targets/x86_64-linux/lib/stubs/libcuda.so ] && [ ! -f /usr/local/cuda/targets/x86_64-linux/lib/stubs/libcuda.so.1 ]; then ln -s /usr/local/cuda/targets/x86_64-linux/lib/stubs/libcuda.so /usr/local/cuda/targets/x86_64-linux/lib/stubs/libcuda.so.1; fi",
        "git clone --depth 1 https://github.com/ggml-org/llama.cpp.git /opt/llama.cpp",
        "cmake -S /opt/llama.cpp -B /opt/llama.cpp/build -DGGML_CUDA=ON -DCMAKE_BUILD_TYPE=Release -DCMAKE_EXE_LINKER_FLAGS='-L/usr/local/cuda/lib64/stubs -L/usr/local/cuda/targets/x86_64-linux/lib/stubs -Wl,-rpath-link,/usr/local/cuda/lib64/stubs -Wl,-rpath-link,/usr/local/cuda/targets/x86_64-linux/lib/stubs'",
        "cmake --build /opt/llama.cpp/build --target llama-mtmd-cli llama-cli -j",
    )
    .add_local_python_source("snap2sim")
)

app = modal.App("snap2sim-inside-the-machine")


@app.function(image=image, timeout=120)
def check_remote_imports() -> dict[str, Any]:
    """Lightweight Modal check that local project modules are packaged."""
    import snap2sim.model_io
    import snap2sim.schema

    return {
        "ok": True,
        "modules": [
            snap2sim.model_io.__name__,
            snap2sim.schema.__name__,
        ],
    }


def runtime_config() -> dict[str, str]:
    return {
        "model_repo": os.getenv("SNAP2SIM_MODEL_REPO", DEFAULT_MODEL_REPO),
        "gguf_quant": os.getenv("SNAP2SIM_GGUF_QUANT", DEFAULT_GGUF_QUANT),
        "mmproj_file": os.getenv("SNAP2SIM_MMPROJ_FILE", DEFAULT_MMPROJ_FILE),
        "runtime_mode": os.getenv("SNAP2SIM_RUNTIME_MODE", DEFAULT_RUNTIME_MODE),
    }


def require_authorization(authorization: str) -> None:
    expected_token = os.getenv("SNAP2SIM_API_TOKEN", "")
    if not expected_token:
        raise HTTPException(status_code=503, detail="API authentication is not configured.")
    scheme, separator, provided_token = authorization.partition(" ")
    if (
        not separator
        or scheme.lower() != "bearer"
        or not token_secrets.compare_digest(provided_token, expected_token)
    ):
        raise HTTPException(status_code=401, detail="Unauthorized")


def asset_patterns(config: dict[str, str]) -> list[str]:
    """Return the repo file patterns needed by the llama.cpp runtime."""
    return [
        f"*{config['gguf_quant']}*.gguf",
        config["mmproj_file"],
        "README.md",
    ]


def cached_asset_paths(config: dict[str, str]) -> tuple[Path, Path] | None:
    """Return cached model paths when both required runtime files exist."""
    local_dir = Path(MODEL_ASSET_DIR) / config["model_repo"]
    model_matches = sorted(local_dir.glob(f"*{config['gguf_quant']}*.gguf"))
    mmproj_path = local_dir / config["mmproj_file"]
    if model_matches and mmproj_path.exists():
        return model_matches[0], mmproj_path
    return None


@app.function(
    image=image,
    volumes={CACHE_DIR: model_cache},
    timeout=60 * 60,
)
def download_runtime_assets() -> dict[str, Any]:
    """Cache the selected GGUF quant and projector on the Modal Volume."""
    from huggingface_hub import snapshot_download

    config = runtime_config()
    patterns = asset_patterns(config)
    path = snapshot_download(
        repo_id=config["model_repo"],
        local_dir=f"{MODEL_ASSET_DIR}/{config['model_repo']}",
        allow_patterns=patterns,
    )
    model_cache.commit()
    return {
        "repo": config["model_repo"],
        "quant": config["gguf_quant"],
        "mmproj": config["mmproj_file"],
        "path": path,
        "patterns": patterns,
    }


def ensure_runtime_assets() -> tuple[Path, Path]:
    """Download configured model assets if needed and return local paths."""
    from huggingface_hub import snapshot_download

    config = runtime_config()
    cached_paths = cached_asset_paths(config)
    if cached_paths:
        return cached_paths

    local_dir = Path(MODEL_ASSET_DIR) / config["model_repo"]
    snapshot_download(
        repo_id=config["model_repo"],
        local_dir=local_dir,
        allow_patterns=asset_patterns(config),
    )
    cached_paths = cached_asset_paths(config)
    if not cached_paths:
        raise FileNotFoundError(f"No GGUF file matched {config['gguf_quant']} in {local_dir}")
    return cached_paths


@app.function(
    image=llamacpp_image,
    gpu=os.getenv("SNAP2SIM_SMOKE_GPU", "L40S"),
    volumes={CACHE_DIR: model_cache},
    timeout=60 * 60,
)
def smoke_test_llamacpp_image() -> dict[str, Any]:
    """Run one image prompt through llama.cpp's multimodal CLI on a GPU."""
    import base64
    import subprocess
    import tempfile
    import time

    from PIL import Image, ImageDraw

    model_path, mmproj_path = ensure_runtime_assets()
    temp_file = tempfile.NamedTemporaryFile(prefix="snap2sim-smoke-", suffix=".jpg", delete=False)
    temp_file.close()
    test_image = Path(temp_file.name)
    img = Image.new("RGB", (512, 384), "#d8d0bd")
    draw = ImageDraw.Draw(img)
    draw.rectangle((82, 96, 430, 288), outline="#2c3138", width=8)
    draw.ellipse((178, 112, 334, 268), outline="#b06c23", width=14)
    draw.line((256, 112, 256, 268), fill="#2c3138", width=6)
    draw.line((178, 190, 334, 190), fill="#2c3138", width=6)
    img.save(test_image, format="JPEG", quality=92)

    prompt = """Answer with only this compact JSON shape. Do not include markdown.
{
  "component": "short component name",
  "confidence": 0.7,
  "summary": "one sentence about the visible test image mechanism",
  "trigger": "manual alignment",
  "motion_sequence": ["first motion", "second motion"],
  "parts": [
    {
      "id": "ring",
      "name": "outer ring",
      "role": "frames the mechanism",
      "geometry": {"shape": "cylinder", "size": [1, 0.1, 1], "position": [0, 0, 0]},
      "motion": {"type": "rotate", "axis": [0, 1, 0], "speed": 0.2}
    },
    {
      "id": "crossbar",
      "name": "crossbar",
      "role": "shows alignment",
      "geometry": {"shape": "rod", "size": [0.05, 0.05, 1], "position": [0, 0.05, 0]},
      "motion": {"type": "static"}
    }
  ]
}"""
    cmd = [
        "/opt/llama.cpp/build/bin/llama-mtmd-cli",
        "-m",
        str(model_path),
        "--mmproj",
        str(mmproj_path),
        "--image",
        str(test_image),
        "-p",
        prompt,
        "-n",
        "1024",
        "-c",
        "4096",
        "--temp",
        "0.2",
    ]
    start = time.monotonic()
    proc = subprocess.run(cmd, capture_output=True, text=True, timeout=45 * 60)
    elapsed_seconds = round(time.monotonic() - start, 2)
    stdout = proc.stdout.strip()
    stderr = proc.stderr.strip()
    combined_output = "\n".join(part for part in [stdout, stderr] if part).strip()
    parsed_component = ""
    valid_json = False
    parse_error = ""
    try:
        parsed_component = parse_analysis_response(stdout)["component"]
        valid_json = True
    except Exception as exc:
        parse_error = str(exc)

    image_supported = (
        proc.returncode == 0
        and bool(stdout)
        and "image input is not supported" not in combined_output.lower()
        and "failed to load projector" not in combined_output.lower()
    )
    image_base64_prefix = base64.b64encode(test_image.read_bytes()).decode("ascii")[:80]
    test_image.unlink(missing_ok=True)
    return {
        "ok": image_supported and valid_json,
        "image_supported": image_supported,
        "valid_json": valid_json,
        "parsed_component": parsed_component,
        "parse_error": parse_error,
        "returncode": proc.returncode,
        "elapsed_seconds": elapsed_seconds,
        "model_path": str(model_path),
        "mmproj_path": str(mmproj_path),
        "image_base64_prefix": image_base64_prefix,
        "stdout_tail": stdout[-4000:],
        "stderr_tail": stderr[-4000:],
    }


def run_llamacpp_prompt(
    prompt: str,
    image_path: Path | None = None,
    system_prompt: str | None = None,
    max_tokens: int = 3072,
    timeout_seconds: int = 300,
    ctx_size: int = 8192,
) -> str:
    """Run one prompt through the llama.cpp multimodal CLI."""
    import subprocess

    model_path, mmproj_path = ensure_runtime_assets()

    def build_cmd(user_prompt: str, system: str | None) -> list[str]:
        cmd = [
            "/opt/llama.cpp/build/bin/llama-mtmd-cli",
            "-m",
            str(model_path),
            "--mmproj",
            str(mmproj_path),
        ]
        if system:
            cmd.extend(["-sys", system])
        if image_path is not None:
            cmd.extend(["--image", str(image_path)])
        cmd.extend(
            [
                "-p",
                user_prompt,
                "-n",
                str(max_tokens),
                "-c",
                str(ctx_size),
                "--temp",
                "0.2",
            ]
        )
        return cmd

    def run_cmd(cmd: list[str]) -> subprocess.CompletedProcess[str]:
        try:
            return subprocess.run(cmd, capture_output=True, text=True, timeout=timeout_seconds)
        except subprocess.TimeoutExpired as exc:
            partial_output = "\n".join(
                part.decode("utf-8", errors="replace") if isinstance(part, bytes) else part
                for part in [exc.stdout, exc.stderr]
                if part
            ).strip()
            raise TimeoutError(
                f"llama.cpp timed out after {timeout_seconds}s: {partial_output[-2000:]}"
            ) from exc

    proc = run_cmd(build_cmd(prompt, system_prompt))
    output = "\n".join(part for part in [proc.stdout, proc.stderr] if part).strip()
    if proc.returncode != 0 and system_prompt and _system_prompt_flag_unsupported(output):
        combined_prompt = f"{system_prompt.strip()}\n\nUser request:\n{prompt}"
        proc = run_cmd(build_cmd(combined_prompt, None))
        output = "\n".join(part for part in [proc.stdout, proc.stderr] if part).strip()
    if proc.returncode != 0:
        raise RuntimeError(f"llama.cpp exited with {proc.returncode}: {output[-2000:]}")
    return output


def _system_prompt_flag_unsupported(output: str) -> bool:
    lowered = output.lower()
    return any(
        marker in lowered
        for marker in [
            "unknown argument: -sys",
            "unknown argument '-sys'",
            "unknown option: -sys",
            "unrecognized option '-sys'",
            "invalid option -- 'sys'",
            "error: unknown argument",
        ]
    )


def write_payload_image(payload: dict[str, Any]) -> Path:
    """Decode an image_base64 payload into a temporary RGB JPEG."""
    import base64
    import binascii
    import tempfile
    import warnings
    from io import BytesIO

    from PIL import Image, UnidentifiedImageError

    Image.MAX_IMAGE_PIXELS = MAX_IMAGE_PIXELS
    image_base64 = payload.get("image_base64")
    if not isinstance(image_base64, str) or not image_base64:
        raise ValueError("Request payload must include image_base64.")
    if "," in image_base64 and image_base64.lstrip().startswith("data:"):
        image_base64 = image_base64.split(",", 1)[1]
    if len(image_base64) > MAX_IMAGE_BASE64_CHARS:
        raise ValueError("Image upload is too large.")

    try:
        raw = base64.b64decode(image_base64, validate=True)
    except (binascii.Error, ValueError) as exc:
        raise ValueError("Image payload is not valid base64.") from exc
    if len(raw) > MAX_IMAGE_BYTES:
        raise ValueError("Image upload is too large.")

    try:
        with warnings.catch_warnings():
            warnings.simplefilter("error", Image.DecompressionBombWarning)
            image = Image.open(BytesIO(raw))
            image.load()
    except Image.DecompressionBombWarning as exc:
        raise ValueError("Image dimensions are too large.") from exc
    except Image.DecompressionBombError as exc:
        raise ValueError("Image dimensions are too large.") from exc
    except (UnidentifiedImageError, OSError, ValueError) as exc:
        raise ValueError("Upload a valid image file.") from exc
    if image.width * image.height > MAX_IMAGE_PIXELS:
        raise ValueError("Image dimensions are too large.")

    temp_file = tempfile.NamedTemporaryFile(prefix="snap2sim-request-", suffix=".jpg", delete=False)
    temp_file.close()
    path = Path(temp_file.name)
    image = image.convert("RGB")
    image.save(path, format="JPEG", quality=92)
    return path


def analyze_image_llamacpp_payload(payload: dict[str, Any]) -> dict[str, Any]:
    image_path = write_payload_image(payload)
    try:
        system_prompt, user_prompt = build_vision_messages()
        response = run_llamacpp_prompt(
            user_prompt,
            image_path=image_path,
            system_prompt=system_prompt,
            max_tokens=4096,
            timeout_seconds=300,
            ctx_size=8192,
        )
        try:
            return parse_analysis_response(response)
        except Exception:
            return coerce_analysis_response(response)
    finally:
        image_path.unlink(missing_ok=True)


@app.local_entrypoint()
def run_runtime_preflight() -> None:
    """Cache assets, then run the Modal GPU llama.cpp image smoke test."""
    print(download_runtime_assets.remote())
    print(smoke_test_llamacpp_image.remote())


@app.local_entrypoint()
def run_smoke_test() -> None:
    """Run and print the Modal GPU llama.cpp image smoke test result."""
    import json

    print(json.dumps(smoke_test_llamacpp_image.remote(), indent=2))


@app.local_entrypoint()
def run_analysis_endpoint_check() -> None:
    """Run the experimental image-analysis endpoint logic with a test image."""
    import base64
    import json
    from io import BytesIO

    from PIL import Image, ImageDraw

    img = Image.new("RGB", (512, 384), "#d8d0bd")
    draw = ImageDraw.Draw(img)
    draw.rectangle((82, 96, 430, 288), outline="#2c3138", width=8)
    draw.ellipse((178, 112, 334, 268), outline="#b06c23", width=14)
    draw.line((256, 112, 256, 268), fill="#2c3138", width=6)
    draw.line((178, 190, 334, 190), fill="#2c3138", width=6)
    buffer = BytesIO()
    img.save(buffer, format="JPEG", quality=92)
    result = analyze_image_llamacpp_task.remote(
        {"image_base64": base64.b64encode(buffer.getvalue()).decode("ascii")}
    )
    print(json.dumps(result, indent=2))


@app.local_entrypoint()
def run_analysis_raw_check() -> None:
    """Print raw llama.cpp analysis output diagnostics for prompt tuning."""
    import base64
    import json
    from io import BytesIO

    from PIL import Image, ImageDraw

    img = Image.new("RGB", (512, 384), "#d8d0bd")
    draw = ImageDraw.Draw(img)
    draw.rectangle((82, 96, 430, 288), outline="#2c3138", width=8)
    draw.ellipse((178, 112, 334, 268), outline="#b06c23", width=14)
    draw.line((256, 112, 256, 268), fill="#2c3138", width=6)
    draw.line((178, 190, 334, 190), fill="#2c3138", width=6)
    buffer = BytesIO()
    img.save(buffer, format="JPEG", quality=92)
    payload = {"image_base64": base64.b64encode(buffer.getvalue()).decode("ascii")}
    result = analyze_image_llamacpp_raw_task.remote(payload)
    print(json.dumps(result, indent=2))


@app.function(image=image, timeout=120, secrets=[api_auth_secret])
@modal.fastapi_endpoint(method="GET")
def runtime_probe(authorization: str = Header(default="")) -> dict[str, Any]:
    """Expose the currently selected runtime path for deployment diagnostics."""
    require_authorization(authorization)
    config = runtime_config()
    return {
        "runtime_mode": config["runtime_mode"],
        "model_repo": config["model_repo"],
        "gguf_quant": config["gguf_quant"],
        "mmproj_file": config["mmproj_file"],
        "status": (
            "placeholder endpoint active; llama.cpp endpoint verified separately"
            if config["runtime_mode"] == "placeholder"
            else "runtime selected"
        ),
        "verified_endpoint": "analyze_image_llamacpp",
        "recommended_generate_endpoint": "generate_scene",
    }


@app.function(image=image, volumes={CACHE_DIR: model_cache}, timeout=600, secrets=[api_auth_secret])
@modal.fastapi_endpoint(method="POST")
def analyze_image(payload: dict[str, Any], authorization: str = Header(default="")) -> dict[str, Any]:
    require_authorization(authorization)
    if runtime_config()["runtime_mode"] != "placeholder":
        raise NotImplementedError(
            "Use the analyze_image_llamacpp endpoint for the verified Nemotron "
            "llama.cpp runtime path."
        )
    return validate_analysis(dict(EXAMPLE_ANALYSIS))


@app.function(image=llamacpp_image, gpu=os.getenv("SNAP2SIM_RUNTIME_GPU", "L40S"), volumes={CACHE_DIR: model_cache}, timeout=60 * 60)
def analyze_image_llamacpp_task(payload: dict[str, Any]) -> dict[str, Any]:
    """Remote-callable task for testing the llama.cpp image analysis path."""
    return analyze_image_llamacpp_payload(payload)


@app.function(image=llamacpp_image, gpu=os.getenv("SNAP2SIM_RUNTIME_GPU", "L40S"), volumes={CACHE_DIR: model_cache}, timeout=60 * 60)
def analyze_image_llamacpp_raw_task(payload: dict[str, Any]) -> dict[str, Any]:
    """Remote-callable diagnostic task for prompt and budget tuning."""
    import time

    image_path = write_payload_image(payload)
    start = time.monotonic()
    try:
        system_prompt, user_prompt = build_vision_messages()
        response = run_llamacpp_prompt(
            user_prompt,
            image_path=image_path,
            system_prompt=system_prompt,
            max_tokens=4096,
            timeout_seconds=300,
            ctx_size=8192,
        )
    finally:
        image_path.unlink(missing_ok=True)

    elapsed_seconds = round(time.monotonic() - start, 2)
    parse_error = ""
    parsed: dict[str, Any] | None = None
    try:
        parsed = parse_analysis_response(response)
    except Exception as exc:
        parse_error = str(exc)
    coerced = coerce_analysis_response(response)
    return {
        "elapsed_seconds": elapsed_seconds,
        "parse_ok": parsed is not None,
        "parse_error": parse_error,
        "parsed_component": parsed["component"] if parsed else "",
        "coerced_component": coerced["component"],
        "coerced_render_mode": select_render_mode(coerced),
        "coerced_confidence": coerced.get("confidence"),
        "stdout_tail": response[-4000:],
    }


@app.function(image=llamacpp_image, gpu=os.getenv("SNAP2SIM_RUNTIME_GPU", "L40S"), volumes={CACHE_DIR: model_cache}, timeout=60 * 60, secrets=[api_auth_secret])
@modal.fastapi_endpoint(method="POST")
def analyze_image_llamacpp(payload: dict[str, Any], authorization: str = Header(default="")) -> dict[str, Any]:
    """Experimental GPU endpoint for llama.cpp multimodal image analysis."""
    require_authorization(authorization)
    return analyze_image_llamacpp_payload(payload)


@app.function(image=image, volumes={CACHE_DIR: model_cache}, timeout=600, secrets=[api_auth_secret])
@modal.fastapi_endpoint(method="POST")
def generate_scene(payload: dict[str, Any], authorization: str = Header(default="")) -> dict[str, Any]:
    require_authorization(authorization)
    analysis = validate_analysis(payload.get("analysis") or EXAMPLE_ANALYSIS)
    if runtime_config()["runtime_mode"] != "placeholder":
        raise NotImplementedError(
            "Scene generation is deterministic in the browser from validated JSON."
        )
    render_mode = select_render_mode(analysis)
    renderer = "three" if render_mode == "three" else "photo"
    return {"renderer": renderer, "render_mode": render_mode, "analysis": analysis}


@app.function(image=llamacpp_image, gpu=os.getenv("SNAP2SIM_RUNTIME_GPU", "L40S"), volumes={CACHE_DIR: model_cache}, timeout=60 * 60, secrets=[api_auth_secret])
@modal.fastapi_endpoint(method="POST")
def generate_scene_llamacpp(payload: dict[str, Any], authorization: str = Header(default="")) -> dict[str, Any]:
    """Compatibility endpoint; scene rendering is deterministic browser-side."""
    require_authorization(authorization)
    analysis = validate_analysis(payload.get("analysis") or EXAMPLE_ANALYSIS)
    render_mode = select_render_mode(analysis)
    renderer = "three" if render_mode == "three" else "photo"
    return {"renderer": renderer, "render_mode": render_mode, "analysis": analysis}