modal_app.py

import os

import modal

app = modal.App("rootstock-admin")
app.image = modal.Image.debian_slim().pip_install("fastapi[standard]")

env = os.environ.get("MODAL_ENVIRONMENT", "dev")
vol = modal.Volume.from_name(f"rootstock-admin-{env}", create_if_missing=True)


@app.function(volumes={"/data": vol})
@modal.fastapi_endpoint(method="POST", requires_proxy_auth=True)
def manifest(manifest: dict):
    import json

    vol.reload()
    name = manifest.get("cluster")
    with open(f"/data/{name}.json", "w") as f:
        json.dump(manifest, f)

    vol.commit()
    return manifest


@app.function(volumes={"/data": vol})
@modal.fastapi_endpoint(method="POST", requires_proxy_auth=True)
def usage(payload: dict):
    """Ingest a cluster's usage rollups: {"cluster": str, "rows": [...]}.

    Rows are the aggregated rollup rows produced by `rootstock usage report`
    (keyed by month/cluster/env/checkpoint/device/client, with sessions,
    n_calculations, duration_s, unique_users). Rows are stored per month and
    a push only replaces the months it contains, so history survives even if
    a cluster's spool loses old months (e.g. scratch purge policies).
    """
    import json
    import os
    import re
    from collections import defaultdict

    from fastapi import HTTPException

    cluster = payload.get("cluster")
    rows = payload.get("rows")
    if not isinstance(cluster, str) or not re.fullmatch(r"[A-Za-z0-9_-]+", cluster):
        raise HTTPException(422, "invalid or missing 'cluster'")
    if not isinstance(rows, list):
        raise HTTPException(422, "'rows' must be a list of rollup rows")

    by_month = defaultdict(list)
    for row in rows:
        if not isinstance(row, dict) or not re.fullmatch(
            r"\d{4}-\d{2}", str(row.get("month"))
        ):
            raise HTTPException(422, "every row needs a 'month' of the form YYYY-MM")
        # Only derived counts ever leave a cluster; drop user hashes if a
        # client ever sends them by mistake.
        row.pop("users", None)
        by_month[row["month"]].append(row)

    vol.reload()
    os.makedirs(f"/data/usage/{cluster}", exist_ok=True)
    for month, month_rows in by_month.items():
        with open(f"/data/usage/{cluster}/{month}.json", "w") as f:
            json.dump({"cluster": cluster, "month": month, "rows": month_rows}, f)

    vol.commit()
    return {"cluster": cluster, "months": sorted(by_month)}


@app.function(volumes={"/data": vol}, min_containers=1)
@modal.asgi_app()
def dashboard():
    from fastapi import FastAPI
    from fastapi.middleware.cors import CORSMiddleware

    web_app = FastAPI()
    web_app.add_middleware(
        CORSMiddleware,
        allow_origins=["*"],
        allow_methods=["GET"],
        allow_headers=["*"],
    )

    @web_app.get("/")
    def get_manifests():
        import json
        from pathlib import Path

        vol.reload()
        manifests = [json.loads(p.read_text()) for p in Path("/data").glob("*.json")]
        return {"manifests": manifests}

    @web_app.get("/usage")
    def get_usage():
        import json
        from pathlib import Path

        vol.reload()
        rollups = [
            json.loads(p.read_text())
            for p in sorted(Path("/data/usage").glob("*/*.json"))
        ]
        return {"usage": rollups}

    return web_app