194 lines
7.8 KiB
Python
Executable File
194 lines
7.8 KiB
Python
Executable File
"""PolicyEngine — applies static rules + learned heuristics + live overrides
|
|
to scheduler decisions.
|
|
|
|
Constructor accepts an optional ``db`` (anything with an ``events``
|
|
dict-like attribute). If ``db`` is None, the heuristics layer falls
|
|
back to the in-process reputation store (``_reputation``) which is
|
|
fed by :meth:`update_reputation`.
|
|
|
|
Learning loop (F-04 remediation):
|
|
:meth:`update_reputation` records per-node success/failure outcomes
|
|
and feeds them back into :meth:`_heuristics` so the scheduler can
|
|
prefer nodes that historically succeed and de-prioritise nodes that
|
|
overheat or fail repeatedly. The reputation store is an
|
|
exponentially-decaying moving average — old observations count for
|
|
less, so a node that recovered from a bad spell isn't penalised
|
|
forever.
|
|
"""
|
|
|
|
from __future__ import annotations
|
|
|
|
import logging
|
|
import threading
|
|
import time
|
|
from typing import Any, Dict, Optional
|
|
|
|
log = logging.getLogger(__name__)
|
|
|
|
|
|
# Decay factor for the moving average. New observation contributes
|
|
# ``_DECAY`` of the new value, old history contributes ``1 - _DECAY``.
|
|
_DECAY: float = 0.3
|
|
|
|
|
|
class PolicyEngine:
|
|
"""Static rules + learned heuristics + live overrides."""
|
|
|
|
def __init__(self, db: Optional[Any] = None) -> None:
|
|
self.db = db
|
|
self.rules: list = []
|
|
self.overrides: Dict[str, Any] = {}
|
|
# Per-node reputation store (populated by update_reputation).
|
|
# shape: {node_name: {"success": float, "failure": float,
|
|
# "thermal_spike": float, "last_seen": float,
|
|
# "samples": int}}
|
|
self._reputation: Dict[str, Dict[str, float]] = {}
|
|
self._lock = threading.Lock()
|
|
|
|
# ------------------------------------------------------------------
|
|
# Rule registration
|
|
# ------------------------------------------------------------------
|
|
def add_rule(self, rule: Any) -> None:
|
|
self.rules.append(rule)
|
|
|
|
def set_override(self, key: str, value: Any) -> None:
|
|
self.overrides[key] = value
|
|
|
|
def clear_override(self, key: str) -> None:
|
|
self.overrides.pop(key, None)
|
|
|
|
# ------------------------------------------------------------------
|
|
# Learning loop (F-04 remediation)
|
|
# ------------------------------------------------------------------
|
|
def update_reputation(
|
|
self,
|
|
node: str,
|
|
success: bool,
|
|
duration: float = 0.0,
|
|
thermal_spike: bool = False,
|
|
) -> Dict[str, float]:
|
|
"""Feed an execution outcome back into the reputation store.
|
|
|
|
Parameters
|
|
----------
|
|
node
|
|
Node identifier (typically ``node["name"]``).
|
|
success
|
|
``True`` if the action completed with rc=0.
|
|
duration
|
|
Wall-clock seconds the action took. Used to track slow
|
|
nodes via the rolling average.
|
|
thermal_spike
|
|
``True`` if the node's temperature rose > 0.15°C during
|
|
the action (overheating risk).
|
|
|
|
Returns the updated reputation dict for ``node`` so callers
|
|
can log / emit it on the bus.
|
|
"""
|
|
if not node:
|
|
return {}
|
|
with self._lock:
|
|
cur = self._reputation.setdefault(node, {
|
|
"success": 0.5, "failure": 0.0, "thermal_spike": 0.0,
|
|
"avg_duration": 0.0, "samples": 0, "last_seen": 0.0,
|
|
})
|
|
# Exponentially-weighted moving average
|
|
cur["success"] = (_DECAY * (1.0 if success else 0.0)
|
|
+ (1 - _DECAY) * cur["success"])
|
|
cur["failure"] = (_DECAY * (0.0 if success else 1.0)
|
|
+ (1 - _DECAY) * cur["failure"])
|
|
cur["thermal_spike"] = (_DECAY * (1.0 if thermal_spike else 0.0)
|
|
+ (1 - _DECAY) * cur["thermal_spike"])
|
|
if duration > 0:
|
|
cur["avg_duration"] = (_DECAY * duration
|
|
+ (1 - _DECAY) * cur["avg_duration"])
|
|
cur["samples"] = cur["samples"] + 1
|
|
cur["last_seen"] = time.time()
|
|
return dict(cur)
|
|
|
|
def get_reputation(self, node: str) -> Dict[str, float]:
|
|
"""Return the current reputation snapshot for ``node`` (or empty)."""
|
|
with self._lock:
|
|
return dict(self._reputation.get(node, {}))
|
|
|
|
def all_reputations(self) -> Dict[str, Dict[str, float]]:
|
|
"""Return a deep-ish snapshot of all known node reputations."""
|
|
with self._lock:
|
|
return {n: dict(v) for n, v in self._reputation.items()}
|
|
|
|
# ------------------------------------------------------------------
|
|
# Apply policy to node selection
|
|
# ------------------------------------------------------------------
|
|
def evaluate(self, action: Dict[str, Any], target: str,
|
|
node: Dict[str, Any]) -> float:
|
|
score_modifier: float = 0.0
|
|
for rule in self.rules:
|
|
try:
|
|
score_modifier += rule.apply(action, target, node)
|
|
except Exception:
|
|
log.debug("policy rule raised", exc_info=True)
|
|
score_modifier += self._heuristics(action, target, node)
|
|
score_modifier += self._overrides(action, target, node)
|
|
return score_modifier
|
|
|
|
# ------------------------------------------------------------------
|
|
# Heuristics (learning layer)
|
|
# ------------------------------------------------------------------
|
|
def _heuristics(self, action: Dict[str, Any], target: str,
|
|
node: Dict[str, Any]) -> float:
|
|
score: float = 0.0
|
|
name = node.get("name") if isinstance(node, dict) else None
|
|
if not name:
|
|
return 0.0
|
|
|
|
# Layer 1: reputation store (always available — even with db=None)
|
|
with self._lock:
|
|
rep = self._reputation.get(name)
|
|
if rep:
|
|
# Reward historical success, penalise failure + thermal spikes
|
|
score += rep["success"] * 20.0
|
|
score -= rep["failure"] * 50.0
|
|
score -= rep["thermal_spike"] * 15.0
|
|
# Penalise slow nodes (only after we have a few samples)
|
|
if rep["samples"] >= 3 and rep["avg_duration"] > 0:
|
|
# Normalise: 60s baseline, 1 point per 60s over baseline
|
|
score -= max(0.0, (rep["avg_duration"] - 60.0) / 60.0) * 5.0
|
|
|
|
# Layer 2: db-backed history (preserved for backward compat)
|
|
if self.db is not None and hasattr(self.db, "events"):
|
|
try:
|
|
history = self.db.events
|
|
except Exception:
|
|
history = None
|
|
if history is not None:
|
|
success_count = 0
|
|
if isinstance(history, dict):
|
|
iterable = (e for sess in history.values() for e in (sess or []))
|
|
else:
|
|
iterable = history or []
|
|
for event in iterable:
|
|
if not isinstance(event, dict):
|
|
continue
|
|
if event.get("type") in ("action_end", "task_update"):
|
|
d = event.get("data", event)
|
|
if (d.get("action") == action.get("name")
|
|
and d.get("state") == "done"
|
|
and d.get("node") == name):
|
|
success_count += 1
|
|
score += success_count * 2
|
|
return score
|
|
|
|
# ------------------------------------------------------------------
|
|
# Live overrides
|
|
# ------------------------------------------------------------------
|
|
def _overrides(self, action: Dict[str, Any], target: str,
|
|
node: Dict[str, Any]) -> float:
|
|
score: float = 0.0
|
|
forced = self.overrides.get("force_node")
|
|
if forced and node.get("name") == forced:
|
|
score += 1000
|
|
avoid = self.overrides.get("avoid_node")
|
|
if avoid and node.get("name") == avoid:
|
|
score -= 1000
|
|
return score
|