ubuntu-bare-metal/stacks/gpu-fan/app.py

#!/usr/bin/env python3
"""GPU fan control web UI + NVML daemon (single process)."""

from __future__ import annotations

import logging
import os
import signal
import sys
import threading
from pathlib import Path
from typing import Any

import uvicorn
from fastapi import Depends, FastAPI, HTTPException, Request
from fastapi.responses import FileResponse
from pydantic import BaseModel, Field

from fan_controller import (
    FanControlError,
    FanController,
    MAX_CURVE_POINTS,
    MIN_CURVE_POINTS,
    MIN_FAN_SPEED,
    MAX_FAN_SPEED,
    curve_to_dict,
    parse_curve,
)

STACK_DIR = Path(__file__).resolve().parent

# Unified control-plane env (see stacks/control-plane/env_loader.py)
for _cp in ("/opt/control-plane", "/repo/stacks/control-plane", str(STACK_DIR.parent / "control-plane")):
    if _cp not in sys.path and Path(_cp).exists():
        sys.path.insert(0, _cp)
from env_loader import load_control_plane_env  # noqa: E402

STATIC_DIR = STACK_DIR / "static"

logging.basicConfig(
    level=logging.INFO,
    format="%(asctime)s %(levelname)s %(message)s",
    handlers=[logging.StreamHandler(sys.stdout)],
)
log = logging.getLogger(__name__)


def _is_localhost_bind(host: str) -> bool:
    return host.strip().lower() in ("127.0.0.1", "localhost", "::1")


def _resolve_host_port(env: dict[str, str]) -> tuple[str, int]:
    host = env.get("GPU_FAN_API_HOST") or env.get("GPU_FAN_HOST", "0.0.0.0")
    port_raw = env.get("GPU_FAN_API_PORT") or env.get("GPU_FAN_PORT", "8090")
    return host, int(port_raw)


ENV = load_control_plane_env(STACK_DIR)
HOST, PORT = _resolve_host_port(ENV)
API_KEY = ENV.get("API_KEY", "")
CURVE_PATH = Path(ENV.get("CURVE_PATH", "/etc/gpu-fan/curve.json"))
POLL_INTERVAL = float(ENV.get("POLL_INTERVAL", "2.0"))
GPU_INDEX = int(ENV.get("GPU_INDEX", "0"))
DRY_RUN = ENV.get("DRY_RUN", "").lower() in ("1", "true", "yes")

controller = FanController(
    curve_path=CURVE_PATH,
    gpu_index=GPU_INDEX,
    poll_interval=POLL_INTERVAL,
)
controller.dry_run = DRY_RUN
app = FastAPI(title="GPU Fan Control", version="1.0.0")


class CurvePoint(BaseModel):
    temp: int = Field(ge=0, le=120)
    speed: int = Field(ge=0, le=100)


class CurveUpdate(BaseModel):
    points: list[CurvePoint] = Field(min_length=MIN_CURVE_POINTS, max_length=MAX_CURVE_POINTS)


class ModeUpdate(BaseModel):
    mode: str
    speed: int | None = Field(default=None, ge=MIN_FAN_SPEED, le=MAX_FAN_SPEED)


def require_auth(request: Request) -> None:
    if not API_KEY:
        return
    key = request.headers.get("X-API-Key", "")
    if key != API_KEY:
        raise HTTPException(status_code=401, detail="Invalid or missing API key")


@app.get("/")
def index() -> FileResponse:
    return FileResponse(STATIC_DIR / "index.html")


@app.get("/api/status")
def api_status(_: None = Depends(require_auth)) -> dict[str, Any]:
    return controller.get_metrics()


@app.get("/api/curve")
def api_get_curve(_: None = Depends(require_auth)) -> dict[str, Any]:
    points = controller.get_curve()
    return {"points": [{"temp": t, "speed": s} for t, s in points]}


@app.put("/api/curve")
def api_put_curve(body: CurveUpdate, _: None = Depends(require_auth)) -> dict[str, Any]:
    points = [(p.temp, p.speed) for p in body.points]
    try:
        parse_curve(curve_to_dict(points))
        controller.save_curve_file(points)
        controller.set_mode("curve")
    except FanControlError as exc:
        raise HTTPException(status_code=400, detail=str(exc)) from exc
    return {"ok": True, "curve": curve_to_dict(controller.get_curve())}


@app.post("/api/mode")
def api_set_mode(body: ModeUpdate, _: None = Depends(require_auth)) -> dict[str, Any]:
    try:
        controller.set_mode(body.mode, body.speed)
        if body.mode != "auto":
            controller.update_once()
    except FanControlError as exc:
        raise HTTPException(status_code=400, detail=str(exc)) from exc
    return {"ok": True, "mode": controller.mode, "manual_speed": controller.manual_speed}


@app.post("/api/reload")
def api_reload(_: None = Depends(require_auth)) -> dict[str, Any]:
    try:
        controller.reload_curve()
    except FanControlError as exc:
        raise HTTPException(status_code=400, detail=str(exc)) from exc
    return {"ok": True, "curve": curve_to_dict(controller.get_curve())}


def run_daemon_thread() -> None:
    controller.run_loop()


def main() -> None:
    if os.geteuid() != 0 and not DRY_RUN:
        log.error("GPU fan control requires root (NVML write access). Run with sudo.")
        sys.exit(1)

    if not _is_localhost_bind(HOST) and not API_KEY:
        log.error(
            "API_KEY is required when GPU_FAN_HOST=%s (LAN/public bind). "
            "Set API_KEY in .env",
            HOST,
        )
        sys.exit(1)

    controller.init_nvml()

    def shutdown_handler(signum: int, _frame: object) -> None:
        log.info("Received signal %s", signum)
        if signum == signal.SIGHUP:
            try:
                controller.reload_curve()
            except FanControlError as exc:
                log.error("Curve reload failed: %s", exc)
            return
        controller.shutdown()
        sys.exit(0)

    signal.signal(signal.SIGTERM, shutdown_handler)
    signal.signal(signal.SIGINT, shutdown_handler)
    signal.signal(signal.SIGHUP, shutdown_handler)

    thread = threading.Thread(target=run_daemon_thread, daemon=True)
    thread.start()

    if _is_localhost_bind(HOST):
        log.info("Web UI at http://127.0.0.1:%d", PORT)
    else:
        log.info("Web UI listening on 0.0.0.0:%d (LAN — API key required)", PORT)
    uvicorn.run(app, host=HOST, port=PORT, log_level="info")


if __name__ == "__main__":
    if os.environ.get("GPU_FAN_LEGACY_UI", "").lower() in ("1", "true", "yes"):
        main()
    else:
        from fan_daemon import main as daemon_main

        daemon_main()