JOJO
faef998a9a
- 改用 psutil 实时 RSS 取代历史最大 ru_maxrss - 新增后台每秒采样线程，检测到 >100MB/s 增长立即记录 - WebTerminal 初始化时自动启动采样 - 增加子智能体恢复加载 conversation.json、httpx 流式/非流式响应、 工具返回结果、模型输出组装等埋点
310 lines
9.7 KiB
Python
310 lines
9.7 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
|
||
|
||
|
||
_peak_rss_mb: Optional[float] = None
|
||
_last_sample_rss_mb: Optional[float] = None
|
||
_sample_thread: Optional[threading.Thread] = None
|
||
|
||
|
||
def _current_rss_mb() -> Optional[float]:
|
||
"""当前进程 RSS（MB），优先使用 psutil（实时），失败回退到 resource。"""
|
||
global _peak_rss_mb
|
||
try:
|
||
import psutil
|
||
proc = psutil.Process()
|
||
rss = proc.memory_info().rss / (1024 * 1024)
|
||
if _peak_rss_mb is None or rss > _peak_rss_mb:
|
||
_peak_rss_mb = rss
|
||
return rss
|
||
except Exception:
|
||
pass
|
||
try:
|
||
import resource
|
||
usage = resource.getrusage(resource.RUSAGE_SELF)
|
||
# macOS: ru_maxrss 单位是 bytes；Linux: 单位是 KB
|
||
if sys.platform == "darwin":
|
||
rss = usage.ru_maxrss / (1024 * 1024)
|
||
else:
|
||
rss = usage.ru_maxrss / 1024.0
|
||
if _peak_rss_mb is None or rss > _peak_rss_mb:
|
||
_peak_rss_mb = rss
|
||
return rss
|
||
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
|
||
|
||
|
||
def _periodic_sampling(
|
||
interval_seconds: float = 1.0,
|
||
spike_threshold_mb: float = 100.0,
|
||
max_records: int = 3,
|
||
) -> None:
|
||
"""后台线程：每秒采样一次内存，发现快速增长或峰值时记录。"""
|
||
global _last_sample_rss_mb, _peak_rss_mb
|
||
records_since_spike = 0
|
||
while ENABLED:
|
||
try:
|
||
info = get_memory_info()
|
||
rss = info.get("rss_mb")
|
||
vms = info.get("vms_mb")
|
||
if rss is not None:
|
||
spike = False
|
||
if _last_sample_rss_mb is not None:
|
||
delta = rss - _last_sample_rss_mb
|
||
if delta >= spike_threshold_mb:
|
||
spike = True
|
||
records_since_spike = max_records
|
||
_last_sample_rss_mb = rss
|
||
if records_since_spike > 0:
|
||
records_since_spike -= 1
|
||
log_memory_event(
|
||
"memory_periodic_spike",
|
||
rss_mb=rss,
|
||
vms_mb=vms,
|
||
peak_rss_mb=_peak_rss_mb,
|
||
delta_from_last_mb=(delta if spike else None),
|
||
)
|
||
else:
|
||
# 每 5 秒记录一次常规采样，避免日志过多
|
||
if int(time.time()) % 5 == 0:
|
||
log_memory_event(
|
||
"memory_periodic_sample",
|
||
rss_mb=rss,
|
||
vms_mb=vms,
|
||
peak_rss_mb=_peak_rss_mb,
|
||
)
|
||
except Exception:
|
||
pass
|
||
time.sleep(interval_seconds)
|
||
|
||
|
||
def start_periodic_sampling(
|
||
interval_seconds: float = 1.0,
|
||
spike_threshold_mb: float = 100.0,
|
||
) -> None:
|
||
"""启动后台内存采样线程。幂等：重复调用不会启动多个线程。"""
|
||
global _sample_thread
|
||
if not ENABLED:
|
||
return
|
||
if _sample_thread is not None and _sample_thread.is_alive():
|
||
return
|
||
_sample_thread = threading.Thread(
|
||
target=_periodic_sampling,
|
||
args=(interval_seconds, spike_threshold_mb),
|
||
name="memory-debug-sampler",
|
||
daemon=True,
|
||
)
|
||
_sample_thread.start()
|