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

#!/usr/bin/env python3
"""GPU fan control NVML daemon + API only (no web UI)."""

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 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

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


def _resolve_host_port(env: dict[str, str]) -> tuple[str, int]:
    """Agent API bind — prefer GPU_FAN_API_*; never fall back to legacy LAN :8090."""
    if env.get("GPU_FAN_API_HOST") or env.get("GPU_FAN_API_PORT"):
        host = env.get("GPU_FAN_API_HOST", "127.0.0.1")
        port_raw = env.get("GPU_FAN_API_PORT", "18090")
    else:
        host = "127.0.0.1"
        port_raw = "18090"
    return host, int(port_raw)


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


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 Agent", 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("/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_API_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()

    log.info("GPU fan agent API at http://%s:%d", HOST, PORT)
    log.info("Web UI: use Server UI at :8091 (GPU Fan tab)")
    uvicorn.run(app, host=HOST, port=PORT, log_level="info")


if __name__ == "__main__":
    main()