234 lines
7.0 KiB
Python
234 lines
7.0 KiB
Python
"""内存调试工具：在关键路径记录内存变化，帮助定位暴涨点。
|
||
|
||
使用方式：
|
||
from modules.memory_debug import mem_snapshot, log_memory_event, get_memory_info
|
||
|
||
# 简单记录
|
||
log_memory_event("my_event", foo="bar")
|
||
|
||
# 上下文：记录进入/退出内存差
|
||
with mem_snapshot("deep_compress", conversation_id=conv_id):
|
||
await run_deep_compression(...)
|
||
|
||
输出：
|
||
~/.astrion/astrion/<mode>/logs/memory_debug.log（JSONL）
|
||
"""
|
||
from __future__ import annotations
|
||
|
||
import json
|
||
import os
|
||
import sys
|
||
import threading
|
||
import time
|
||
import tracemalloc
|
||
from contextlib import asynccontextmanager, contextmanager
|
||
from pathlib import Path
|
||
from typing import Any, AsyncIterator, Dict, Optional
|
||
|
||
from utils.log_rotation import append_line
|
||
|
||
# 开关：默认关闭，避免生产环境大量写日志。
|
||
# 用户需要排查时设置环境变量 ASTRION_MEMORY_DEBUG=1 开启。
|
||
ENABLED = str(os.environ.get("ASTRION_MEMORY_DEBUG", "1")).strip().lower() in {
|
||
"1", "true", "yes", "on",
|
||
}
|
||
|
||
_log_lock = threading.Lock()
|
||
_log_path: Optional[Path] = None
|
||
|
||
|
||
def _resolve_log_path() -> Path:
|
||
"""解析内存调试日志路径，复用项目运行态 logs 目录。"""
|
||
global _log_path
|
||
if _log_path is not None:
|
||
return _log_path
|
||
try:
|
||
from config import LOGS_DIR
|
||
base = Path(LOGS_DIR)
|
||
except Exception:
|
||
base = Path.home() / ".astrion" / "astrion" / "host" / "logs"
|
||
_log_path = base / "memory_debug.log"
|
||
return _log_path
|
||
|
||
|
||
def _current_rss_mb() -> Optional[float]:
|
||
"""当前进程 RSS（MB），失败返回 None。"""
|
||
try:
|
||
import resource
|
||
usage = resource.getrusage(resource.RUSAGE_SELF)
|
||
# macOS: ru_maxrss 单位是 bytes；Linux: 单位是 KB
|
||
if sys.platform == "darwin":
|
||
return usage.ru_maxrss / (1024 * 1024)
|
||
return usage.ru_maxrss / 1024.0
|
||
except Exception:
|
||
pass
|
||
try:
|
||
import psutil
|
||
proc = psutil.Process()
|
||
return proc.memory_info().rss / (1024 * 1024)
|
||
except Exception:
|
||
return None
|
||
|
||
|
||
def _current_vms_mb() -> Optional[float]:
|
||
"""当前进程 VMS（MB），失败返回 None。"""
|
||
try:
|
||
import psutil
|
||
proc = psutil.Process()
|
||
return proc.memory_info().vms / (1024 * 1024)
|
||
except Exception:
|
||
return None
|
||
|
||
|
||
def get_memory_info() -> Dict[str, Optional[float]]:
|
||
return {
|
||
"rss_mb": _current_rss_mb(),
|
||
"vms_mb": _current_vms_mb(),
|
||
}
|
||
|
||
|
||
def _estimate_text_size(text: Any) -> int:
|
||
"""估算文本/列表/字典在 JSON 序列化后的大小（字符数）。"""
|
||
try:
|
||
return len(json.dumps(text, ensure_ascii=False))
|
||
except Exception:
|
||
return len(str(text))
|
||
|
||
|
||
def estimate_messages_size(messages: Any) -> Dict[str, Any]:
|
||
"""估算消息列表的规模和体积。"""
|
||
if not messages:
|
||
return {"count": 0, "total_chars": 0, "avg_chars": 0}
|
||
if not isinstance(messages, list):
|
||
return {"count": 0, "total_chars": _estimate_text_size(messages), "avg_chars": 0}
|
||
total = 0
|
||
for m in messages:
|
||
total += _estimate_text_size(m)
|
||
return {
|
||
"count": len(messages),
|
||
"total_chars": total,
|
||
"avg_chars": total // max(len(messages), 1),
|
||
}
|
||
|
||
|
||
def log_memory_event(event: str, **kwargs: Any) -> None:
|
||
"""写一条内存调试日志。"""
|
||
if not ENABLED:
|
||
return
|
||
try:
|
||
info = get_memory_info()
|
||
payload: Dict[str, Any] = {
|
||
"t": time.time(),
|
||
"event": event,
|
||
"rss_mb": info.get("rss_mb"),
|
||
"vms_mb": info.get("vms_mb"),
|
||
}
|
||
# 过滤掉不可序列化的对象
|
||
for k, v in kwargs.items():
|
||
try:
|
||
json.dumps({k: v}, ensure_ascii=False)
|
||
payload[k] = v
|
||
except Exception:
|
||
payload[k] = str(v)[:500]
|
||
append_line(_resolve_log_path(), json.dumps(payload, ensure_ascii=False))
|
||
except Exception:
|
||
pass
|
||
|
||
|
||
@contextmanager
|
||
def mem_snapshot(event: str, **kwargs: Any):
|
||
"""同步上下文管理器：记录进入和退出时的内存变化。"""
|
||
if not ENABLED:
|
||
yield
|
||
return
|
||
|
||
start_info = get_memory_info()
|
||
start_t = time.time()
|
||
log_memory_event(f"{event}_enter", duration_ms=0, **kwargs)
|
||
try:
|
||
yield
|
||
finally:
|
||
end_info = get_memory_info()
|
||
delta_rss = None
|
||
delta_vms = None
|
||
if start_info.get("rss_mb") is not None and end_info.get("rss_mb") is not None:
|
||
delta_rss = round(end_info["rss_mb"] - start_info["rss_mb"], 2)
|
||
if start_info.get("vms_mb") is not None and end_info.get("vms_mb") is not None:
|
||
delta_vms = round(end_info["vms_mb"] - start_info["vms_mb"], 2)
|
||
log_memory_event(
|
||
f"{event}_exit",
|
||
duration_ms=round((time.time() - start_t) * 1000, 2),
|
||
delta_rss_mb=delta_rss,
|
||
delta_vms_mb=delta_vms,
|
||
start_rss_mb=start_info.get("rss_mb"),
|
||
end_rss_mb=end_info.get("rss_mb"),
|
||
**kwargs,
|
||
)
|
||
|
||
|
||
@asynccontextmanager
|
||
async def async_mem_snapshot(event: str, **kwargs: Any) -> AsyncIterator[None]:
|
||
"""异步上下文管理器：记录进入和退出时的内存变化。"""
|
||
if not ENABLED:
|
||
yield
|
||
return
|
||
|
||
start_info = get_memory_info()
|
||
start_t = time.time()
|
||
log_memory_event(f"{event}_enter", duration_ms=0, **kwargs)
|
||
try:
|
||
yield
|
||
finally:
|
||
end_info = get_memory_info()
|
||
delta_rss = None
|
||
delta_vms = None
|
||
if start_info.get("rss_mb") is not None and end_info.get("rss_mb") is not None:
|
||
delta_rss = round(end_info["rss_mb"] - start_info["rss_mb"], 2)
|
||
if start_info.get("vms_mb") is not None and end_info.get("vms_mb") is not None:
|
||
delta_vms = round(end_info["vms_mb"] - start_info["vms_mb"], 2)
|
||
log_memory_event(
|
||
f"{event}_exit",
|
||
duration_ms=round((time.time() - start_t) * 1000, 2),
|
||
delta_rss_mb=delta_rss,
|
||
delta_vms_mb=delta_vms,
|
||
start_rss_mb=start_info.get("rss_mb"),
|
||
end_rss_mb=end_info.get("rss_mb"),
|
||
**kwargs,
|
||
)
|
||
|
||
|
||
class TopAllocator:
|
||
"""基于 tracemalloc 的 Top 分配点采样（开销较大，按需显式调用）。"""
|
||
|
||
def __init__(self, top_n: int = 10):
|
||
self.top_n = top_n
|
||
|
||
def start(self) -> None:
|
||
try:
|
||
tracemalloc.start()
|
||
except Exception:
|
||
pass
|
||
|
||
def snapshot(self, label: str = "") -> None:
|
||
try:
|
||
snap = tracemalloc.take_snapshot()
|
||
top = snap.statistics("lineno")[: self.top_n]
|
||
lines = []
|
||
for stat in top:
|
||
lines.append(
|
||
f"{stat.size / (1024 * 1024):.2f}MB {stat.count} {stat.traceback.format()[-1]}"
|
||
)
|
||
log_memory_event(
|
||
"tracemalloc_top",
|
||
label=label,
|
||
top=lines,
|
||
)
|
||
except Exception:
|
||
pass
|
||
|
||
def stop(self) -> None:
|
||
try:
|
||
tracemalloc.stop()
|
||
except Exception:
|
||
pass
|