158 lines
5.2 KiB
Python
158 lines
5.2 KiB
Python
import asyncio
|
|
import re
|
|
import time
|
|
from typing import Annotated, Any
|
|
|
|
from fastapi import APIRouter, Form
|
|
from fastapi.responses import HTMLResponse
|
|
|
|
from app.collect.gpu import shorten
|
|
from app.plugins.base import Plugin
|
|
from app.render import render
|
|
|
|
router = APIRouter(prefix="/api/plugins/lact", tags=["plugins"])
|
|
|
|
LIST_TIMEOUT = 5
|
|
SET_TIMEOUT = 15
|
|
GPU_CACHE_TTL = 60
|
|
|
|
_gpu_cache: tuple[float, list[dict[str, str]]] | None = None
|
|
_set_lock = asyncio.Lock()
|
|
|
|
|
|
async def _run(args: list[str], timeout: float) -> tuple[str, str]:
|
|
try:
|
|
proc = await asyncio.create_subprocess_exec(
|
|
"lact", "cli", *args,
|
|
stdout=asyncio.subprocess.PIPE,
|
|
stderr=asyncio.subprocess.PIPE,
|
|
)
|
|
except FileNotFoundError:
|
|
return "", "lact not found in PATH"
|
|
except OSError as e:
|
|
return "", str(e)[:200]
|
|
try:
|
|
out, err = await asyncio.wait_for(proc.communicate(), timeout)
|
|
except TimeoutError:
|
|
try:
|
|
proc.kill()
|
|
except ProcessLookupError:
|
|
pass
|
|
return "", "lact timed out"
|
|
if proc.returncode != 0:
|
|
return "", (err.decode(errors="replace").strip() or f"lact failed (rc={proc.returncode})")[:200]
|
|
return out.decode(errors="replace"), ""
|
|
|
|
|
|
def _parse_gpus(out: str) -> list[dict[str, str]]:
|
|
gpus: list[dict[str, str]] = []
|
|
for line in out.splitlines():
|
|
m = re.match(r"^\s*(\d+):\s+(.*)$", line)
|
|
if not m:
|
|
continue
|
|
rest = m.group(2)
|
|
name = re.search(r"\(([^)]*)\)", rest)
|
|
gpu_type = re.search(r"\[([^\]]*)\]\s*$", rest)
|
|
gpus.append({
|
|
"id": m.group(1),
|
|
"name": shorten(name.group(1)) if name else rest.strip(),
|
|
"type": gpu_type.group(1) if gpu_type else "",
|
|
})
|
|
return gpus
|
|
|
|
|
|
async def _gpus(force: bool = False) -> tuple[list[dict[str, str]], str]:
|
|
global _gpu_cache
|
|
if not force and _gpu_cache is not None:
|
|
ts, cached = _gpu_cache
|
|
if time.monotonic() - ts < GPU_CACHE_TTL:
|
|
return [dict(g) for g in cached], ""
|
|
out, err = await _run(["list"], LIST_TIMEOUT)
|
|
if err:
|
|
return [], err
|
|
gpus = _parse_gpus(out)
|
|
_gpu_cache = (time.monotonic(), gpus)
|
|
return gpus, ""
|
|
|
|
|
|
async def _gpu_entry(g: dict[str, str], with_profiles: bool) -> dict[str, Any]:
|
|
entry: dict[str, Any] = {**g, "profiles": [], "active": None, "error": ""}
|
|
base = ["--gpu-id", g["id"], "profile"]
|
|
if with_profiles:
|
|
active_p, profiles_p = await asyncio.gather(
|
|
_run([*base, "get"], LIST_TIMEOUT),
|
|
_run([*base, "list"], LIST_TIMEOUT),
|
|
)
|
|
else:
|
|
active_p, profiles_p = await _run([*base, "get"], LIST_TIMEOUT), None
|
|
out, err = active_p
|
|
if err:
|
|
entry["error"] = err
|
|
return entry
|
|
entry["active"] = out.strip() or None
|
|
if profiles_p is not None:
|
|
out2, err2 = profiles_p
|
|
entry["error"] = err2
|
|
if not err2:
|
|
entry["profiles"] = [line.strip() for line in out2.splitlines() if line.strip()]
|
|
return entry
|
|
|
|
|
|
async def _gather(with_profiles: bool, force_gpus: bool = False) -> dict[str, Any]:
|
|
gpus, err = await _gpus(force=force_gpus)
|
|
if err:
|
|
return {"gpus": [], "error": err}
|
|
entries = await asyncio.gather(*[_gpu_entry(g, with_profiles) for g in gpus])
|
|
return {"gpus": list(entries), "error": ""}
|
|
|
|
|
|
async def _state(message: str = "", error: str = "") -> str:
|
|
data = await _gather(with_profiles=False)
|
|
data["message"] = message
|
|
data["error"] = error or data["error"]
|
|
return render("plugins/lact_state.html", **data)
|
|
|
|
|
|
async def _skeleton(message: str = "", error: str = "") -> str:
|
|
data = await _gather(with_profiles=True, force_gpus=True)
|
|
data["message"] = message
|
|
data["error"] = error or data["error"]
|
|
return render("plugins/lact_skeleton.html", **data)
|
|
|
|
|
|
@router.get("/state")
|
|
async def state():
|
|
return HTMLResponse(await _state())
|
|
|
|
|
|
@router.post("/set")
|
|
async def set_profile(gpu_id: Annotated[str, Form()], profile: Annotated[str, Form()]):
|
|
async with _set_lock:
|
|
data = await _gather(with_profiles=True, force_gpus=True)
|
|
if data["error"]:
|
|
return HTMLResponse(await _skeleton(error=data["error"]))
|
|
gpu = next((g for g in data["gpus"] if g["id"] == gpu_id), None)
|
|
if gpu is None:
|
|
return HTMLResponse(await _skeleton(error=f"unknown gpu id: {gpu_id}"))
|
|
if profile not in gpu["profiles"]:
|
|
return HTMLResponse(await _skeleton(error=f"unknown profile: {profile}"))
|
|
if profile == gpu["active"]:
|
|
return HTMLResponse(await _skeleton(message=f"{gpu['name']}: {profile} already active"))
|
|
_out, err = await _run(["--gpu-id", gpu_id, "profile", "set", profile], SET_TIMEOUT)
|
|
if err:
|
|
return HTMLResponse(await _skeleton(error=err))
|
|
return HTMLResponse(await _skeleton(message=f"{gpu['name']}: set profile {profile}"))
|
|
|
|
|
|
@router.post("/reload")
|
|
async def reload(gpu_id: Annotated[str | None, Form()] = None):
|
|
_ = gpu_id
|
|
return HTMLResponse(await _skeleton(message="profiles refreshed"))
|
|
|
|
|
|
plugin = Plugin(
|
|
id="lact",
|
|
title="GPU power profiles",
|
|
description="Active profile per GPU and a profile switcher (lact cli).",
|
|
skeleton_fn=_skeleton,
|
|
)
|