175 lines
6.5 KiB
Python
175 lines
6.5 KiB
Python
"""AUTOPIN：专家热度持久化 + 启动预热钉死。
|
||
|
||
两部分：
|
||
- **热度计数与持久化**：搭已有的物化点零同步计频——block host 路径的 flat、
|
||
acquire_gpu 的 LFU piggyback(全命中)与 miss 回退 flat；dual decode/verify 无 Python
|
||
物化点，由 C++ g_real 的 LFU freq 在落盘时按增量差分合并。节流写
|
||
{EXPERT_DIR}/pool_usage.json(同目录临时文件 + rename 防半写；读不到/损坏静默从零)。
|
||
- **启动预热钉死**：build_streaming_model 末尾按历史热度把每层 top-N 热专家
|
||
(N = floor(cap_for × AUTOPIN_BUDGET_FRAC))经 blob 并行读预填进常驻池并 pin 住
|
||
(不参与任何驱逐)，消灭冷启动慢热。
|
||
|
||
覆盖率口径：prefill/host 路径与 GPU remap miss 回退全计；非 dual 的 decode 全命中
|
||
快路径搭 LFU piggyback 计(inds 已物化，零额外同步)；dual decode/verify 靠 C++ freq
|
||
兜底(仅 lfu 策略计频)。EVICT_POLICY≠lfu 时两条 GPU 快路径均不计——绝不为计数给
|
||
decode 热路径新增 GPU→host 同步。
|
||
"""
|
||
import atexit
|
||
import json
|
||
import os
|
||
import tempfile
|
||
from collections import Counter
|
||
from typing import Dict, List
|
||
|
||
from mlx_streaming import config
|
||
|
||
_USAGE_NAME = "pool_usage.json"
|
||
|
||
|
||
def usage_path() -> str:
|
||
"""热度文件路径：AUTOPIN_USAGE_FILE 覆盖，缺省 {EXPERT_DIR}/pool_usage.json。"""
|
||
p = config.autopin_usage_file()
|
||
return p or os.path.join(config.expert_dir(), _USAGE_NAME)
|
||
|
||
|
||
def load_usage(path: str) -> "Dict[int, Counter]":
|
||
"""读热度文件；不存在/损坏/结构非法一律静默返回空(从零开始)。"""
|
||
try:
|
||
with open(path) as f:
|
||
raw = json.load(f)
|
||
return {int(l): Counter({int(e): int(c) for e, c in d.items()})
|
||
for l, d in raw.items()}
|
||
except Exception:
|
||
return {}
|
||
|
||
|
||
def save_usage(path: str, counts: "Dict[int, Counter]") -> None:
|
||
"""原子写：先写同目录临时文件再 rename，防半写损坏。失败静默(不阻断推理)。"""
|
||
try:
|
||
d = os.path.dirname(path) or "."
|
||
fd, tmp = tempfile.mkstemp(dir=d, prefix=".pool_usage.", suffix=".tmp")
|
||
with os.fdopen(fd, "w") as f:
|
||
json.dump({str(l): {str(e): c for e, c in cnt.items()}
|
||
for l, cnt in counts.items()}, f)
|
||
os.replace(tmp, path)
|
||
except Exception:
|
||
pass
|
||
|
||
|
||
class HeatCounter:
|
||
"""进程内热度累计器：note() 增量计频，节流 + atexit 落盘。
|
||
|
||
前向边界用「MoE 层号回绕」近似(与 VirtualPool.begin_forward 同判据)：MoE 块按层号
|
||
递增被调，出现 layer <= 上次即新前向。路径在创建时快照，避免进程退出时 env 已变
|
||
导致 atexit 写错位置。
|
||
"""
|
||
|
||
def __init__(self, path: str):
|
||
self.path = path
|
||
self.counts: "Dict[int, Counter]" = {}
|
||
self._last_layer = -1
|
||
self._forwards_since_save = 0
|
||
# C++ 累计频次的已合并基线(real_freq_dump 是进程累计值，按差分合并防重复计)
|
||
self._cpp_baseline: "Dict[int, Dict[int, int]]" = {}
|
||
|
||
def note(self, layer: int, ids: List[int]) -> None:
|
||
layer = int(layer)
|
||
if layer <= self._last_layer:
|
||
self._forwards_since_save += 1
|
||
self._last_layer = layer
|
||
c = self.counts.get(layer)
|
||
if c is None:
|
||
c = self.counts[layer] = Counter()
|
||
for e in ids:
|
||
c[int(e)] += 1
|
||
if self._forwards_since_save >= config.autopin_save_every():
|
||
self.flush()
|
||
|
||
def tick(self) -> None:
|
||
"""显式前向边界(dual decode 无 Python 计频点，由 begin_forward 驱动周期落盘)。"""
|
||
self._forwards_since_save += 1
|
||
if self._forwards_since_save >= config.autopin_save_every():
|
||
self.flush()
|
||
|
||
def merge_cpp_freq(self) -> None:
|
||
"""把 dual decode 的 C++ LFU 累计频次按增量差分合并进来。"""
|
||
try:
|
||
import mlx_streaming.native_moe_ext as _N
|
||
flat = _N.real_freq_dump()
|
||
except Exception:
|
||
return
|
||
cur: "Dict[int, Dict[int, int]]" = {}
|
||
for i in range(0, len(flat), 3):
|
||
l, e, c = int(flat[i]), int(flat[i + 1]), int(flat[i + 2])
|
||
cur.setdefault(l, {})[e] = c
|
||
for l, em in cur.items():
|
||
base = self._cpp_baseline.get(l, {})
|
||
cnt = self.counts.setdefault(l, Counter())
|
||
for e, c in em.items():
|
||
b = base.get(e, 0)
|
||
cnt[e] += c - b if c >= b else c # real_reset 后 c<b：以 c 为新基线
|
||
self._cpp_baseline = cur
|
||
|
||
def flush(self) -> None:
|
||
self.merge_cpp_freq()
|
||
save_usage(self.path, self.counts)
|
||
self._forwards_since_save = 0
|
||
|
||
|
||
_counter: "HeatCounter | None" = None
|
||
|
||
|
||
def _get_counter() -> HeatCounter:
|
||
global _counter
|
||
if _counter is None:
|
||
_counter = HeatCounter(usage_path())
|
||
atexit.register(_counter.flush)
|
||
return _counter
|
||
|
||
|
||
def note(layer: int, ids: List[int]) -> None:
|
||
"""路由热度计数入口(热路径调用)：AUTOPIN=0 立即返回，零行为。"""
|
||
if not config.autopin():
|
||
return
|
||
_get_counter().note(layer, ids)
|
||
|
||
|
||
def tick() -> None:
|
||
"""前向边界入口(virtual_pool begin_forward 调用)：AUTOPIN=0 立即返回。"""
|
||
if not config.autopin():
|
||
return
|
||
_get_counter().tick()
|
||
|
||
|
||
def counter() -> "HeatCounter | None":
|
||
"""测试/诊断用：当前计数器(未触发过计数则为 None)。"""
|
||
return _counter
|
||
|
||
|
||
def warm_start_pin(store) -> "dict":
|
||
"""启动预热钉死：按 usage 文件每层取 top-N 热专家，经 store.pin_batch 预填进池并 pin 住。
|
||
|
||
N = floor(该层 cap_for × AUTOPIN_BUDGET_FRAC)。pin 加载由 pin_batch 收口：有 blob_loader
|
||
走并行 pread(生产唯一数据源)，无则退回逐专家 per-expert 文件(兼容旧部署)。
|
||
usage 缺失/为空 → 返回零摘要(后台计数仍在累计，供下次启动用)；单层失败不阻断整体。
|
||
"""
|
||
summary = {"layers": 0, "pinned": 0, "skipped_layers": 0, "seconds": 0.0}
|
||
usage = load_usage(usage_path())
|
||
if not usage:
|
||
return summary
|
||
import time
|
||
t0 = time.perf_counter()
|
||
frac = config.autopin_budget_frac()
|
||
for layer in sorted(usage):
|
||
n = int(store.cap_for(layer) * frac)
|
||
if n <= 0:
|
||
continue
|
||
top = [e for e, _ in usage[layer].most_common(n)]
|
||
try:
|
||
summary["pinned"] += store.pin_batch(layer, top)
|
||
summary["layers"] += 1
|
||
except Exception:
|
||
summary["skipped_layers"] += 1
|
||
summary["seconds"] = round(time.perf_counter() - t0, 2)
|
||
return summary
|