tinymind-native-colab-handoff/bundle/model/auto_upgrade.py

"""Auto-upgrade governor for TinyMind.

Accepts self-improvement candidates only when capability is preserved or
improved and resource use does not regress without a real power gain.
"""

from __future__ import annotations

from dataclasses import asdict, dataclass


@dataclass(frozen=True)
class AutoUpgradeCandidate:
    name: str
    quality: float
    purity: float
    power: float
    vram_mb: float
    latency_ms: float

    def efficiency(self) -> float:
        resource = max(self.vram_mb / 1024.0, 0.001) + max(self.latency_ms / 1000.0, 0.001)
        return (self.quality + self.purity + self.power) / resource


class AutoUpgradeGovernor:
    policy = {
        "self_upgrade_enabled": True,
        "rollback_on_power_regression": True,
        "rollback_on_quality_or_purity_regression": True,
        "resource_pressure_minimized": True,
        "external_world_best_claim_blocked_without_official_results": True,
    }

    def __init__(
        self,
        max_quality_drop: float = 0.0,
        max_purity_drop: float = 0.0,
        max_power_drop: float = 0.0,
        min_power_gain_for_resource_increase: float = 0.03,
    ):
        self.max_quality_drop = float(max_quality_drop)
        self.max_purity_drop = float(max_purity_drop)
        self.max_power_drop = float(max_power_drop)
        self.min_power_gain_for_resource_increase = float(min_power_gain_for_resource_increase)

    def _block_reason(self, base: AutoUpgradeCandidate, candidate: AutoUpgradeCandidate) -> str | None:
        if base.quality - candidate.quality > self.max_quality_drop:
            return f"quality_regression:{base.quality - candidate.quality:.6f}"
        if base.purity - candidate.purity > self.max_purity_drop:
            return f"purity_regression:{base.purity - candidate.purity:.6f}"
        if base.power - candidate.power > self.max_power_drop:
            return f"power_regression:{base.power - candidate.power:.6f}"
        resource_worse = candidate.vram_mb > base.vram_mb and candidate.latency_ms > base.latency_ms
        power_gain = candidate.power - base.power
        if resource_worse and power_gain < self.min_power_gain_for_resource_increase:
            return "resource_regression_without_power_gain"
        return None

    def decide(self, base: AutoUpgradeCandidate, candidates: list[AutoUpgradeCandidate]) -> dict:
        accepted = []
        blocked = []
        for candidate in candidates:
            reason = self._block_reason(base, candidate)
            if reason:
                blocked.append(f"{candidate.name}:{reason}")
            else:
                accepted.append(candidate)

        selected = base
        if accepted:
            selected = max([base] + accepted, key=lambda item: item.efficiency())
        accepted_flag = selected.name != base.name

        return {
            "schema_version": "tinymind-auto-upgrade-governor-v1",
            "policy": dict(self.policy),
            "base": asdict(base),
            "candidates": [asdict(candidate) for candidate in candidates],
            "selected": selected.name,
            "accepted": accepted_flag,
            "blocked": blocked,
            "selected_efficiency": selected.efficiency(),
            "resource_delta": {
                "vram_mb": selected.vram_mb - base.vram_mb,
                "latency_ms": selected.latency_ms - base.latency_ms,
            },
            "capability_delta": {
                "quality": selected.quality - base.quality,
                "purity": selected.purity - base.purity,
                "power": selected.power - base.power,
            },
            "world_best_claim_allowed": False,
        }