#!/usr/bin/env python3
"""
power-governor.py — Dynamic clock and power governor for CMP 50HX GPUs.

Directly queries llama-server slot state for 100% accurate active/idle tracking:
- When generating (is_processing: True): Instantly sets full 1900 MHz boost clocks.
- When idle (is_processing: False for > 15s): Downclocks CMP cards to 600 MHz (~32W/card).
"""

import time
import subprocess
import urllib.request
import json
import sys

CMP_GPUS = "1,2"
IDLE_TIMEOUT_SEC = 15
LOW_POWER_CLOCK = 300

def run_cmd(*args):
    try:
        subprocess.run(args, stdout=subprocess.DEVNULL, stderr=subprocess.DEVNULL, check=False)
    except Exception:
        pass

def is_server_processing():
    try:
        req = urllib.request.Request("http://127.0.0.1:8080/slots", headers={"User-Agent": "power-governor"})
        with urllib.request.urlopen(req, timeout=1.0) as r:
            slots = json.loads(r.read().decode())
            return any(bool(s.get("is_processing")) for s in slots)
    except Exception:
        return False

def main():
    print("[power-governor] Initializing GPU persistence and 150W power cap...", flush=True)
    run_cmd("nvidia-smi", "-pm", "1")
    run_cmd("nvidia-smi", f"-i={CMP_GPUS}", "-pl", "150")
    run_cmd("nvidia-smi", f"-i={CMP_GPUS}", "-rgc")

    is_low_power = False
    last_active = time.time()

    print("[power-governor] Monitoring llama-server processing state...", flush=True)
    while True:
        try:
            active = is_server_processing()
            now = time.time()

            if active:
                last_active = now
                if is_low_power:
                    run_cmd("nvidia-smi", f"-i={CMP_GPUS}", "-rgc")
                    is_low_power = False
                    print(f"[{time.strftime('%X')}] Inference Active -> Boost clocks engaged (1900 MHz)", flush=True)
            else:
                if not is_low_power and (now - last_active) > IDLE_TIMEOUT_SEC:
                    run_cmd("nvidia-smi", f"-i={CMP_GPUS}", f"-lgc={LOW_POWER_CLOCK},{LOW_POWER_CLOCK}")
                    is_low_power = True
                    print(f"[{time.strftime('%X')}] Inference Idle -> Low-power state engaged (600 MHz, ~32W/card)", flush=True)

            time.sleep(0.5 if is_low_power else 1.0)
        except Exception as e:
            time.sleep(3.0)

if __name__ == "__main__":
    main()