263 lines
10 KiB
Python
263 lines
10 KiB
Python
from __future__ import annotations
|
|
|
|
import argparse
|
|
import json
|
|
import os
|
|
import threading
|
|
from collections.abc import Callable
|
|
from pathlib import Path
|
|
from typing import Any
|
|
|
|
|
|
class PerformanceModeError(RuntimeError):
|
|
"""Raised when cpufreq policies cannot be changed safely."""
|
|
|
|
|
|
class PerformanceModeManager:
|
|
RECORD_NAME = "cpufreq-governors.original.json"
|
|
|
|
def __init__(
|
|
self,
|
|
runtime_root: Path,
|
|
*,
|
|
sysfs_root: Path | None = None,
|
|
) -> None:
|
|
configured_root = os.environ.get("MATRIX_CPUFREQ_SYSFS_ROOT")
|
|
self.sysfs_root = Path(
|
|
sysfs_root
|
|
if sysfs_root is not None
|
|
else configured_root or "/sys/devices/system/cpu/cpufreq"
|
|
)
|
|
self.runtime_root = Path(runtime_root)
|
|
self.record_path = self.runtime_root / self.RECORD_NAME
|
|
self._lock = threading.RLock()
|
|
self._restore_governors: dict[str, str] = {}
|
|
self._requested = False
|
|
self._last_error: str | None = None
|
|
self._load_or_capture_baseline()
|
|
|
|
def _policy_paths(self) -> dict[str, Path]:
|
|
if not self.sysfs_root.is_dir():
|
|
return {}
|
|
return {
|
|
path.name: path
|
|
for path in sorted(self.sysfs_root.glob("policy*"), key=lambda item: item.name)
|
|
if path.is_dir()
|
|
}
|
|
|
|
@staticmethod
|
|
def _read_text(path: Path) -> str:
|
|
return path.read_text(encoding="ascii").strip()
|
|
|
|
@staticmethod
|
|
def _write_text(path: Path, value: str) -> None:
|
|
path.write_text(f"{value}\n", encoding="ascii")
|
|
|
|
def _read_governors(self) -> dict[str, str]:
|
|
return {
|
|
name: self._read_text(path / "scaling_governor")
|
|
for name, path in self._policy_paths().items()
|
|
}
|
|
|
|
def _available_governors(self) -> dict[str, set[str]]:
|
|
result: dict[str, set[str]] = {}
|
|
for name, path in self._policy_paths().items():
|
|
available_path = path / "scaling_available_governors"
|
|
try:
|
|
result[name] = set(self._read_text(available_path).split())
|
|
except OSError:
|
|
result[name] = set()
|
|
return result
|
|
|
|
def _write_record(self, governors: dict[str, str]) -> None:
|
|
self.runtime_root.mkdir(parents=True, exist_ok=True)
|
|
temporary = self.record_path.with_suffix(".tmp")
|
|
payload = {
|
|
"schema_version": 1,
|
|
"governors": governors,
|
|
}
|
|
with temporary.open("w", encoding="utf-8", newline="\n") as handle:
|
|
json.dump(payload, handle, ensure_ascii=False, sort_keys=True, separators=(",", ":"))
|
|
handle.write("\n")
|
|
handle.flush()
|
|
os.fsync(handle.fileno())
|
|
os.replace(temporary, self.record_path)
|
|
|
|
def _read_record(self) -> dict[str, str] | None:
|
|
if not self.record_path.is_file():
|
|
return None
|
|
try:
|
|
payload = json.loads(self.record_path.read_text(encoding="utf-8"))
|
|
governors = payload.get("governors")
|
|
if payload.get("schema_version") != 1 or not isinstance(governors, dict):
|
|
raise ValueError("invalid governor record")
|
|
checked = {
|
|
str(name): str(value)
|
|
for name, value in governors.items()
|
|
if str(name).startswith("policy") and str(value).strip()
|
|
}
|
|
if checked != governors:
|
|
raise ValueError("invalid governor record fields")
|
|
return checked
|
|
except (OSError, ValueError, TypeError, json.JSONDecodeError) as exc:
|
|
self._last_error = f"saved governor record is invalid: {exc}"
|
|
return None
|
|
|
|
def _load_or_capture_baseline(self) -> None:
|
|
with self._lock:
|
|
policies = self._policy_paths()
|
|
if not policies:
|
|
return
|
|
saved = self._read_record()
|
|
if saved is not None:
|
|
if set(saved) == set(policies):
|
|
self._restore_governors = saved
|
|
return
|
|
self._last_error = "saved governor policies do not match current cpufreq policies"
|
|
return
|
|
try:
|
|
current = self._read_governors()
|
|
self._write_record(current)
|
|
self._restore_governors = current
|
|
except OSError as exc:
|
|
self._last_error = f"could not capture original governors: {exc}"
|
|
|
|
def _set_governors(self, requested: dict[str, str]) -> None:
|
|
policies = self._policy_paths()
|
|
if set(requested) != set(policies):
|
|
raise PerformanceModeError("cpufreq policies changed while applying governor")
|
|
before = self._read_governors()
|
|
changed: list[str] = []
|
|
try:
|
|
for name, value in requested.items():
|
|
self._write_text(policies[name] / "scaling_governor", value)
|
|
changed.append(name)
|
|
actual = self._read_text(policies[name] / "scaling_governor")
|
|
if actual != value:
|
|
raise PerformanceModeError(
|
|
f"{name} governor is {actual!r} after requesting {value!r}"
|
|
)
|
|
except Exception as exc:
|
|
rollback_errors: list[str] = []
|
|
for name in reversed(changed):
|
|
try:
|
|
self._write_text(policies[name] / "scaling_governor", before[name])
|
|
except OSError as rollback_exc:
|
|
rollback_errors.append(f"{name}: {rollback_exc}")
|
|
suffix = f"; rollback failed for {', '.join(rollback_errors)}" if rollback_errors else ""
|
|
raise PerformanceModeError(f"could not apply cpufreq governor: {exc}{suffix}") from exc
|
|
|
|
def available(self) -> bool:
|
|
policies = self._policy_paths()
|
|
if not policies or set(self._restore_governors) != set(policies):
|
|
return False
|
|
available = self._available_governors()
|
|
return all("performance" in values for values in available.values())
|
|
|
|
def start(self, requested: bool) -> None:
|
|
with self._lock:
|
|
self._requested = requested
|
|
if not requested:
|
|
if self._restore_governors and self._read_governors() != self._restore_governors:
|
|
try:
|
|
self._set_governors(self._restore_governors)
|
|
self._last_error = None
|
|
except PerformanceModeError as exc:
|
|
self._last_error = str(exc)
|
|
return
|
|
try:
|
|
self._apply_requested(True)
|
|
self._last_error = None
|
|
except PerformanceModeError as exc:
|
|
self._last_error = str(exc)
|
|
|
|
def _apply_requested(self, requested: bool) -> None:
|
|
if requested:
|
|
if not self.available():
|
|
raise PerformanceModeError("performance governor is unavailable")
|
|
target = {name: "performance" for name in self._policy_paths()}
|
|
else:
|
|
if not self._restore_governors:
|
|
raise PerformanceModeError("original governors are unavailable")
|
|
target = dict(self._restore_governors)
|
|
self._set_governors(target)
|
|
|
|
def transact(self, requested: bool, persist: Callable[[bool], Any]) -> None:
|
|
with self._lock:
|
|
previous_requested = self._requested
|
|
previous_governors = self._read_governors()
|
|
try:
|
|
self._apply_requested(requested)
|
|
persist(requested)
|
|
except Exception as exc:
|
|
try:
|
|
if previous_governors:
|
|
self._set_governors(previous_governors)
|
|
except Exception as rollback_exc:
|
|
self._last_error = f"{exc}; governor rollback failed: {rollback_exc}"
|
|
raise PerformanceModeError(self._last_error) from exc
|
|
self._requested = previous_requested
|
|
self._last_error = str(exc)
|
|
if isinstance(exc, PerformanceModeError):
|
|
raise
|
|
raise PerformanceModeError(f"performance mode transaction failed: {exc}") from exc
|
|
self._requested = requested
|
|
self._last_error = None
|
|
|
|
def restore(self, *, remove_record: bool = False) -> None:
|
|
with self._lock:
|
|
if self._restore_governors:
|
|
try:
|
|
self._set_governors(self._restore_governors)
|
|
self._last_error = None
|
|
except PerformanceModeError as exc:
|
|
self._last_error = str(exc)
|
|
raise
|
|
if remove_record and self.record_path.exists():
|
|
self.record_path.unlink()
|
|
|
|
def status(self, requested: bool | None = None) -> dict[str, Any]:
|
|
with self._lock:
|
|
current: dict[str, str] = {}
|
|
try:
|
|
current = self._read_governors()
|
|
except OSError as exc:
|
|
self._last_error = f"could not read current governors: {exc}"
|
|
effective_requested = self._requested if requested is None else requested
|
|
error = self._last_error
|
|
if not self._policy_paths() and error is None:
|
|
error = "当前内核未提供 CPU 调频策略；请检查候选内核是否回退"
|
|
elif self._policy_paths() and not self._restore_governors and error is None:
|
|
error = "无法读取 CPU 调频策略的原始值"
|
|
effective = bool(
|
|
effective_requested
|
|
and current
|
|
and all(value == "performance" for value in current.values())
|
|
)
|
|
return {
|
|
"requested": bool(effective_requested),
|
|
"available": self.available(),
|
|
"effective": effective,
|
|
"current_governors": current,
|
|
"restore_governors": dict(self._restore_governors),
|
|
"last_error": error,
|
|
}
|
|
|
|
|
|
def _restore_cli() -> int:
|
|
parser = argparse.ArgumentParser(description="Restore matrix service cpufreq governors")
|
|
parser.add_argument("restore", nargs="?")
|
|
parser.add_argument("--runtime-root", type=Path, required=True)
|
|
parser.add_argument("--sysfs-root", type=Path)
|
|
args = parser.parse_args()
|
|
manager = PerformanceModeManager(args.runtime_root, sysfs_root=args.sysfs_root)
|
|
try:
|
|
manager.restore(remove_record=True)
|
|
except PerformanceModeError as exc:
|
|
parser.error(str(exc))
|
|
return 0
|
|
|
|
|
|
if __name__ == "__main__":
|
|
raise SystemExit(_restore_cli())
|