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

"""NVML fan control with JSON curve, modes, and graceful shutdown."""

from __future__ import annotations

import json
import logging
import signal
import threading
import time
from pathlib import Path
from typing import Any

import pynvml

log = logging.getLogger(__name__)

MIN_CURVE_POINTS = 3
MAX_CURVE_POINTS = 7
MIN_FAN_SPEED = 30
MAX_FAN_SPEED = 100


class FanControlError(Exception):
    pass


def _clamp_speed(speed: int) -> int:
    if speed == 0:
        return 0
    return max(MIN_FAN_SPEED, min(MAX_FAN_SPEED, speed))


def parse_curve(raw: dict[str, Any]) -> list[tuple[int, int]]:
    if not raw:
        raise FanControlError("Curve is empty")

    points: list[tuple[int, int]] = []
    for temp_str, speed in raw.items():
        try:
            temp = int(temp_str)
            speed_int = int(speed)
        except (TypeError, ValueError) as exc:
            raise FanControlError(f"Invalid curve point {temp_str!r}: {speed!r}") from exc

        if not 0 <= temp <= 120:
            raise FanControlError(f"Temperature {temp} out of range 0-120")
        if speed_int != 0 and not MIN_FAN_SPEED <= speed_int <= MAX_FAN_SPEED:
            raise FanControlError(
                f"Fan speed {speed_int}% invalid (use 0 or {MIN_FAN_SPEED}-{MAX_FAN_SPEED})"
            )
        points.append((temp, speed_int))

    points.sort(key=lambda p: p[0])
    temps = [p[0] for p in points]
    if len(set(temps)) != len(temps):
        raise FanControlError("Temperatures must be unique")
    if len(points) < MIN_CURVE_POINTS or len(points) > MAX_CURVE_POINTS:
        raise FanControlError(
            f"Curve must have {MIN_CURVE_POINTS}-{MAX_CURVE_POINTS} points"
        )
    return points


def curve_to_dict(points: list[tuple[int, int]]) -> dict[str, int]:
    return {str(temp): speed for temp, speed in points}


def interpolate_speed(temp: int, curve: list[tuple[int, int]]) -> int:
    if temp <= curve[0][0]:
        return curve[0][1]
    if temp >= curve[-1][0]:
        return curve[-1][1]

    for i in range(len(curve) - 1):
        t1, s1 = curve[i]
        t2, s2 = curve[i + 1]
        if t1 <= temp <= t2:
            if t2 == t1:
                return s2
            ratio = (temp - t1) / (t2 - t1)
            return int(s1 + ratio * (s2 - s1))
    return curve[-1][1]


class FanController:
    def __init__(
        self,
        curve_path: Path,
        gpu_index: int = 0,
        poll_interval: float = 2.0,
    ) -> None:
        self.curve_path = curve_path
        self.gpu_index = gpu_index
        self.poll_interval = poll_interval
        self._lock = threading.Lock()
        self._running = False
        self._mode = "curve"
        self._manual_speed = 100
        self._curve: list[tuple[int, int]] = []
        self._handle: Any = None
        self._fan_count = 0
        self._gpu_name = ""
        self._last_metrics: dict[str, Any] = {}
        self._auto_active = False
        self.dry_run = False

    @property
    def mode(self) -> str:
        with self._lock:
            return self._mode

    @property
    def manual_speed(self) -> int:
        with self._lock:
            return self._manual_speed

    def load_curve_file(self) -> list[tuple[int, int]]:
        if not self.curve_path.exists():
            raise FanControlError(f"Curve file not found: {self.curve_path}")
        raw = json.loads(self.curve_path.read_text(encoding="utf-8"))
        return parse_curve(raw)

    def save_curve_file(self, points: list[tuple[int, int]]) -> None:
        validated = parse_curve(curve_to_dict(points))
        self.curve_path.parent.mkdir(parents=True, exist_ok=True)
        self.curve_path.write_text(
            json.dumps(curve_to_dict(validated), indent=2) + "\n",
            encoding="utf-8",
        )
        with self._lock:
            self._curve = validated
        log.info("Saved curve: %s", validated)

    def reload_curve(self) -> None:
        curve = self.load_curve_file()
        with self._lock:
            self._curve = curve
        log.info("Reloaded curve: %s", curve)

    def set_mode(self, mode: str, manual_speed: int | None = None) -> None:
        if mode not in ("auto", "curve", "manual"):
            raise FanControlError(f"Unknown mode: {mode}")
        with self._lock:
            self._mode = mode
            if manual_speed is not None:
                self._manual_speed = _clamp_speed(manual_speed)
        log.info("Mode set to %s (manual_speed=%s)", mode, manual_speed)

    def get_curve(self) -> list[tuple[int, int]]:
        with self._lock:
            return list(self._curve)

    def get_metrics(self) -> dict[str, Any]:
        with self._lock:
            return dict(self._last_metrics)

    def init_nvml(self) -> None:
        pynvml.nvmlInit()
        count = pynvml.nvmlDeviceGetCount()
        if self.gpu_index >= count:
            raise FanControlError(
                f"GPU index {self.gpu_index} out of range (found {count} GPU(s))"
            )

        handle = pynvml.nvmlDeviceGetHandleByIndex(self.gpu_index)
        name = pynvml.nvmlDeviceGetName(handle)
        fan_count = pynvml.nvmlDeviceGetNumFans(handle)

        self._handle = handle
        self._fan_count = fan_count
        self._gpu_name = name if isinstance(name, str) else name.decode()
        self.reload_curve()
        log.info("GPU %d: %s (%d fan(s))", self.gpu_index, self._gpu_name, fan_count)

    def _read_metrics(self) -> dict[str, Any]:
        handle = self._handle
        assert handle is not None

        temp = pynvml.nvmlDeviceGetTemperature(handle, pynvml.NVML_TEMPERATURE_GPU)
        util = pynvml.nvmlDeviceGetUtilizationRates(handle)
        power_mw = pynvml.nvmlDeviceGetPowerUsage(handle)

        mem_info = pynvml.nvmlDeviceGetMemoryInfo(handle)
        memory_used_mb = mem_info.used // (1024**2)
        memory_total_mb = mem_info.total // (1024**2)

        power_limit_w: float | None = None
        clock_graphics_mhz: int | None = None
        clock_memory_mhz: int | None = None
        try:
            power_limit_w = round(
                pynvml.nvmlDeviceGetEnforcedPowerLimit(handle) / 1000, 1
            )
        except pynvml.NVMLError:
            pass
        try:
            clock_graphics_mhz = pynvml.nvmlDeviceGetClockInfo(
                handle, pynvml.NVML_CLOCK_GRAPHICS
            )
        except pynvml.NVMLError:
            pass
        try:
            clock_memory_mhz = pynvml.nvmlDeviceGetClockInfo(
                handle, pynvml.NVML_CLOCK_MEM
            )
        except pynvml.NVMLError:
            pass

        fan_speeds: list[int] = []
        for fan_idx in range(self._fan_count):
            try:
                fan_speeds.append(pynvml.nvmlDeviceGetFanSpeed_v2(handle, fan_idx))
            except pynvml.NVMLError:
                fan_speeds.append(-1)

        with self._lock:
            mode = self._mode
            manual_speed = self._manual_speed
            curve = list(self._curve)

        if mode == "curve":
            target = interpolate_speed(temp, curve)
        elif mode == "manual":
            target = manual_speed
        else:
            target = None

        metrics = {
            "gpu_index": self.gpu_index,
            "gpu_name": self._gpu_name,
            "temperature_c": temp,
            "fan_speeds_pct": fan_speeds,
            "target_speed_pct": target,
            "power_w": round(power_mw / 1000, 1),
            "power_limit_w": power_limit_w,
            "utilization_pct": util.gpu,
            "memory_utilization_pct": util.memory,
            "memory_used_mb": memory_used_mb,
            "memory_total_mb": memory_total_mb,
            "clock_graphics_mhz": clock_graphics_mhz,
            "clock_memory_mhz": clock_memory_mhz,
            "mode": mode,
            "manual_speed_pct": manual_speed,
            "curve": curve_to_dict(curve),
        }
        with self._lock:
            self._last_metrics = metrics
        return metrics

    def _set_manual_policy(self) -> None:
        handle = self._handle
        assert handle is not None
        for fan_idx in range(self._fan_count):
            pynvml.nvmlDeviceSetFanControlPolicy(
                handle, fan_idx, pynvml.NVML_FAN_POLICY_MANUAL
            )

    def _restore_auto_policy(self) -> None:
        if self._handle is None:
            return
        log.info("Restoring automatic fan control...")
        for fan_idx in range(self._fan_count):
            try:
                pynvml.nvmlDeviceSetFanControlPolicy(
                    self._handle,
                    fan_idx,
                    pynvml.NVML_FAN_POLICY_TEMPERATURE_CONTINOUS_SW,
                )
                log.info("Fan %d: restored to auto", fan_idx)
            except pynvml.NVMLError as exc:
                log.error("Fan %d: could not restore auto: %s", fan_idx, exc)

    def _apply_fan_speed(self, speed: int) -> None:
        handle = self._handle
        assert handle is not None
        speed = _clamp_speed(speed)
        self._set_manual_policy()
        for fan_idx in range(self._fan_count):
            pynvml.nvmlDeviceSetFanSpeed_v2(handle, fan_idx, speed)

    def update_once(self) -> dict[str, Any]:
        metrics = self._read_metrics()
        mode = metrics["mode"]

        if mode == "auto":
            if not self._auto_active and not self.dry_run:
                self._restore_auto_policy()
                self._auto_active = True
            elif self.dry_run:
                self._auto_active = True
            return metrics

        self._auto_active = False
        target = metrics["target_speed_pct"]
        if target is None:
            return metrics

        if self.dry_run:
            log.info(
                "[dry-run] GPU %d: %d°C -> %d%% (fans: %s)",
                self.gpu_index,
                metrics["temperature_c"],
                target,
                metrics["fan_speeds_pct"],
            )
            return metrics

        try:
            self._apply_fan_speed(target)
            log.info(
                "GPU %d: %d°C -> %d%% (fans: %s)",
                self.gpu_index,
                metrics["temperature_c"],
                target,
                metrics["fan_speeds_pct"],
            )
        except pynvml.NVMLError as exc:
            log.error("Failed to set fan speed: %s", exc)
            metrics["error"] = str(exc)
        return metrics

    def run_loop(self) -> None:
        self._running = True
        log.info("Fan control loop started (interval=%ss)", self.poll_interval)
        while self._running:
            try:
                self.update_once()
            except Exception:
                log.exception("Fan control loop error")
            time.sleep(self.poll_interval)

    def stop(self) -> None:
        self._running = False

    def shutdown(self) -> None:
        self.stop()
        if not self.dry_run:
            self._restore_auto_policy()
        try:
            pynvml.nvmlShutdown()
        except pynvml.NVMLError:
            pass
        log.info("Fan controller shut down")

    def install_signal_handlers(self) -> None:
        def handler(signum: int, _frame: Any) -> None:
            log.info("Received signal %s", signum)
            if signum == signal.SIGHUP:
                try:
                    self.reload_curve()
                except FanControlError as exc:
                    log.error("Curve reload failed: %s", exc)
                return
            self.shutdown()

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