gpu-fancontrol.py

#!/bin/env python3
"""
Default env vars:

- `CONFIG_FILE`: `/etc/gpu-fancontrol.json`
- `LOG_FILE`: None, defaults to stdout
- `LOG_LEVEL`: `INFO`

Example config:

```json
{
    "fan_device": "/sys/class/hwmon/hwmon2/pwm1",
    "temps_speeds": [
        [0, 63],
        [55, 150],
        [65, 180],
        [75, 255]
    ]
}
```
"""
import json
import logging
import os
import signal
import subprocess
import sys
from pathlib import Path
from time import sleep

log = logging.getLogger(__name__)


def calc_speed(temps_speeds, curr_speed, curr_temp):
    for i in range(len(temps_speeds)):
        desired_temp, desired_speed = temps_speeds[i]
        if i == len(temps_speeds) - 1:
            set_speed = desired_speed
            break

        next_temp, _ = temps_speeds[i + 1]
        margin = min(5, round((desired_temp - next_temp) / 2))

        if curr_temp >= desired_temp or (
            curr_speed >= desired_speed and curr_temp >= (desired_temp - margin)
        ):
            set_speed = desired_speed
            break

    return set_speed


_exit = False


def _main(config):
    fan_device = config["fan_device"]
    temps_speeds = list(sorted(config["temps_speeds"], reverse=True))

    while 1 and not _exit:
        out = (
            subprocess.check_output(
                [
                    "nvidia-smi",
                    "--query-gpu",
                    "temperature.gpu",
                    "--format=csv,noheader,nounits",
                ],
                env={},
                user="nobody",
            )
            .decode("utf8")
            .strip()
        )
        with open(fan_device) as fp:
            curr_speed = int(fp.read())

        temp = float(out)
        set_speed = calc_speed(temps_speeds, curr_speed, temp)

        log.debug("temp %s, speed %s", temp, curr_speed)

        if set_speed != curr_speed:
            log.info("setting speed to %s", set_speed)
            try:
                with open(fan_device, "w") as fp:
                    fp.write(str(set_speed))
            except Exception:
                log.exception("Failed setting speed")

        sleep(1)


def read_config():
    CONFIG_FILE = Path(os.environ.get("CONFIG_FILE", "/etc/gpu-fancontrol.json"))
    if not CONFIG_FILE.exists():
        sys.exit("Config file '%s' doesn't exist", str(CONFIG_FILE))

    try:
        with CONFIG_FILE.open() as fp:
            return json.loads(fp.read())
    except Exception:
        sys.exit("Invalid config file")


def main():
    log_file = os.environ.get("LOG_FILE", None)
    if not log_file:
        log_file = None

    format = "%(asctime)s: %(levelname)s: %(message)s"
    log_level = os.environ.get("LOG_LEVEL", "INFO")
    logging.basicConfig(
        filename=log_file,
        format=format,
        level=logging.getLevelNamesMapping()[log_level],
    )

    config = read_config()

    fan_device = config["fan_device"]

    log.info("starting with fan device '%s'", fan_device)

    def revert_fan_config(*args):
        global _exit
        with open(f"{fan_device}_enable", "w") as fp:
            fp.write("5")

        _exit = True
        if len(args):
            log.info("Got signal, stopping...")

    signal.signal(signal.SIGTERM, revert_fan_config)
    signal.signal(signal.SIGINT, revert_fan_config)

    try:
        with open(f"{fan_device}_enable", "w") as fp:
            fp.write("1")

        _main(config)
    except:
        log.exception("Exception, stopping")

    revert_fan_config()


if __name__ == "__main__":
    main()