Blame for runtime/memory_manager.py

"""runtime/memory_manager.py — Milestone 4: Memory Manager (PM v0.2 §5)。

职责: 基于 mac16.yaml 的 memory_budget, 动态决定 load/unload/evict。
当前阶段驻留管线的决策模型是确定性的 (AR → unload → Flow), 本模块把
"决策" 从管线代码中抽出来, 并提供真实内存遥测的校验闭环:
  1. 读取 profile 的 memory_budget_breakdown
  2. 决策器: 根据 active/cache/峰值与预算判断某阶段是否可以驻留
  3. 阶段会话: enter → materialize, exit → 释放并校验回落
  4. 测量: 每阶段峰值/耗时/缓存归零检查, 输出 memory report

后续 (M4 扩展) 可加:
  - 动态 eviction: 当 Flow 阶段 chunk 峰值超预算时, 降低 chunk 帧数
  - 双预算档位: preview (更激进) vs create (更保守)
"""
from __future__ import annotations

import gc
import time
from dataclasses import dataclass, field
from pathlib import Path
from typing import Callable, Optional

import mlx.core as mx
import yaml

from runtime.profiler import gb


@dataclass
class StageDecision:
    """一次 load/unload 决策的记录。"""

    stage: str
    action: str  # load | unload | skip
    weight_gb: float
    active_before_gb: float
    active_after_gb: float
    peak_gb: float
    budget_ok: bool
    cache_cleared: bool
    note: str = ""


@dataclass
class MemoryReport:
    stages: list[StageDecision] = field(default_factory=list)
    max_peak_gb: float = 0.0
    profile_name: str = "?"

    def add(self, d: StageDecision) -> None:
        self.stages.append(d)
        self.max_peak_gb = max(self.max_peak_gb, d.peak_gb)

    def to_dict(self) -> dict:
        max_peak_gb = max(
            (stage.peak_gb for stage in self.stages), default=self.max_peak_gb
        )
        return {
            "profile": self.profile_name,
            "max_peak_gb": round(max_peak_gb, 3),
            "stages": [
                {
                    "stage": s.stage,
                    "action": s.action,
                    "weight_gb": round(s.weight_gb, 3),
                    "active_before_gb": round(s.active_before_gb, 3),
                    "active_after_gb": round(s.active_after_gb, 3),
                    "peak_gb": round(s.peak_gb, 3),
                    "budget_ok": s.budget_ok,
                    "cache_cleared": s.cache_cleared,
                }
                for s in self.stages
            ],
        }


class MemoryManager:
    """基于 profile 预算的确定性阶段决策器 + 遥测。"""

    def __init__(
        self,
        profile_path: Optional[Path] = None,
        budget_gb: Optional[float] = None,
    ) -> None:
        self.profile = self._load_profile(profile_path)
        self.budget_gb = budget_gb or float(self.profile.get("memory_budget_gb", 11.0))

    # -------------------------------------------------- 公共 API
    def can_resident(self, weight_gb: float, current_active_gb: float) -> bool:
        """决策: 当前 active + 待加载权重是否在预算内 (留 10% 余量)。"""
        return current_active_gb + weight_gb <= self.budget_gb * 0.9

    def stage_enter(self, stage: str, weight_gb: float, loader: Callable):
        """进入阶段: 检查预算 → 执行加载 → 记录峰值。返回加载的模型。"""
        before = gb(mx.get_active_memory())
        ok = self.can_resident(weight_gb, before)
        if not ok:
            mx.clear_cache()
            before = gb(mx.get_active_memory())
            ok = self.can_resident(weight_gb, before)
        if not ok:
            raise MemoryError(
                f"阶段 {stage} 需要 {weight_gb:.2f}GB，但当前预算 {self.budget_gb:.2f}GB 不足"
            )
        mx.reset_peak_memory()
        obj = loader()
        mx.synchronize()
        peak = gb(mx.get_peak_memory())
        after = gb(mx.get_active_memory())
        return obj, StageDecision(
            stage=stage,
            action="load" if ok else "load_after_evict",
            weight_gb=weight_gb,
            active_before_gb=before,
            active_after_gb=after,
            peak_gb=peak,
            budget_ok=ok,
            cache_cleared=not ok,
        )

    def stage_exit(
        self,
        decision: StageDecision,
        release: Callable[[], None] | None = None,
    ) -> None:
        """退出阶段: 先清掉调用方持有的模型引用, 再执行释放四步曲。"""
        if release is not None:
            release()
        gc.collect()
        mx.synchronize()
        mx.clear_cache()
        cache = gb(mx.get_cache_memory())
        decision.active_after_gb = gb(mx.get_active_memory())
        decision.cache_cleared = cache == 0.0

    def report(self) -> MemoryReport:
        r = MemoryReport(profile_name=self.profile.get("device", "?"))
        return r

    # -------------------------------------------------- helpers
    @staticmethod
    def _load_profile(profile_path: Optional[Path]) -> dict:
        if profile_path is None:
            default = Path(__file__).resolve().parent.parent / "profiles" / "mac16.yaml"
            profile_path = default
        if not profile_path.exists():
            return {"memory_budget_gb": 11.0, "device": "unknown"}
        with open(profile_path) as f:
            return yaml.safe_load(f) or {}


def current_usage_gb() -> float:
    return gb(mx.get_active_memory())


def current_peak_gb() -> float:
    return gb(mx.get_peak_memory())