257 lines
13 KiB
Python
257 lines
13 KiB
Python
import argparse
|
||
import io
|
||
import shutil
|
||
import sys
|
||
import tempfile
|
||
import traceback
|
||
|
||
from engine.harness import calibrate_cprofile_overhead, collect_environment, load_source
|
||
from engine.schema import SCHEMA_VERSION, AnalysisResult, Calibration, FlameNode, WallTime
|
||
from engine.structure import profile_and_measure
|
||
|
||
|
||
# Result 文件大小防护（防用户脚本异常把磁盘撑爆）
|
||
# - error.message 上限 1KB:1KB 一般足够覆盖正常异常说明
|
||
# - error.traceback 上限 8KB:100 帧左右 traceback 的体量,UI 可读范围
|
||
# - functions 列表上限 5000:UI 表格默认只看前 200,再多只是给后续解析加负载
|
||
_MAX_ERROR_MESSAGE_BYTES = 1024
|
||
_MAX_TRACEBACK_BYTES = 8 * 1024
|
||
_MAX_FUNCTIONS = 5000
|
||
|
||
|
||
def _truncate_text(text: str, max_bytes: int) -> str:
|
||
"""把文本按 UTF-8 字节截断到 max_bytes 以内,补一个明确截断标记。
|
||
|
||
直接 [:max_chars] 在多字节字符中间切会破坏 UTF-8,UI 解析挂。encode → 截
|
||
→ 走 errors='ignore' 让不完整尾部自然掉落 → decode,保证结果仍是合法 UTF-8。
|
||
"""
|
||
if not text:
|
||
return text
|
||
encoded = text.encode("utf-8", errors="replace")
|
||
if len(encoded) <= max_bytes:
|
||
return text
|
||
truncated = encoded[:max_bytes].decode("utf-8", errors="ignore")
|
||
# 留 64 字节给标记(总长仍 ≤ max_bytes)
|
||
suffix = f"\n\n[... 截断:共 {len(encoded)} 字节,仅保留前 {max_bytes} 字节 ...]"
|
||
# 如果 suffix 太长导致总超,再砍一次
|
||
return (truncated + suffix)[:max_bytes]
|
||
|
||
|
||
def _cap_functions(functions: list, flame: FlameNode | None) -> tuple:
|
||
"""截断 functions 列表到 _MAX_FUNCTIONS,并把超出部分累加到 flame 的 root。
|
||
|
||
为什么不只截 functions 不动 flame?
|
||
- flame 是 UI 看到的"时间分布",丢了 functions 里那些小函数会让 flame
|
||
总和 < 用户报的总耗时,UI 显示出现 100% 但 sum 不到 80% 的诡异情况。
|
||
- 把被丢函数的 tottime 加到 root.value 上,UI 那边仍按 flame 内部占比显示,
|
||
多出来的部分归到 root 节点（最大最宽的 tile）,用户视觉上只是「其他」变
|
||
大了一点 —— 总和不撒谎。
|
||
"""
|
||
if len(functions) <= _MAX_FUNCTIONS:
|
||
return functions, flame
|
||
kept = functions[:_MAX_FUNCTIONS]
|
||
dropped_tottime = sum(f.tottime for f in functions[_MAX_FUNCTIONS:])
|
||
if flame is None:
|
||
flame = FlameNode(name="root", value=dropped_tottime, children=[])
|
||
else:
|
||
flame.value += dropped_tottime
|
||
return kept, flame
|
||
|
||
|
||
def _compute_wall_time(instrumented_sec: float, calibration: Calibration) -> float:
|
||
"""用校准 ratio 折算 instrumented 时间为估计的干净耗时。
|
||
|
||
思路:cProfile 是插桩式的,wall-time 会被膨胀 ratio 倍(instrumented / clean);
|
||
用同一段 tight-loop 测出 ratio 后,把用户的 instrumented 时间除回去,就是
|
||
「去掉 cProfile 开销后的估计耗时」。
|
||
|
||
边界:calibration.ratio 是 0 / 负数 / NaN 时直接信 instrumented 时间 —— 跟
|
||
calibrate_cprofile_overhead 在「校准耗时过短」时回退 ratio=1.0 的兜底策略
|
||
保持一致,UI 显示真实测得的 instrumented 时间而不是除以无效值。
|
||
"""
|
||
ratio = calibration.ratio
|
||
# `ratio == ratio` 排除 NaN;`ratio > 0` 排除 0 / 负数
|
||
if ratio == ratio and ratio > 0:
|
||
return instrumented_sec / ratio
|
||
return instrumented_sec
|
||
|
||
|
||
def _cap_result(result: AnalysisResult) -> AnalysisResult:
|
||
"""给 result.json 上一道防线,防止异常情况把磁盘/IPC 通道撑爆。
|
||
|
||
触发场景:
|
||
- 用户代码 raise 一个带 50MB 文件内容的 ValueError → error.message 撑爆
|
||
- 用户代码触发深递归异常,Python 自动 dump 几千层 traceback → 撑爆
|
||
- 用户脚本 import 一个庞大的库(numpy/pandas)scope=all 把整棵树都归因 →
|
||
38MB+ 的 functions 列表
|
||
|
||
截断而非拒绝:UI 仍能看到「有错」「发生了什么」,只是具体细节截断;functions
|
||
截断后仍按 tottime 排序,前 5000 一定是最值得用户关注的热点。
|
||
"""
|
||
if result.error:
|
||
msg = result.error.get("message")
|
||
if isinstance(msg, str) and msg:
|
||
result.error["message"] = _truncate_text(msg, _MAX_ERROR_MESSAGE_BYTES)
|
||
tb = result.error.get("traceback")
|
||
if isinstance(tb, str) and tb:
|
||
result.error["traceback"] = _truncate_text(tb, _MAX_TRACEBACK_BYTES)
|
||
if result.functions:
|
||
result.functions, result.flame = _cap_functions(result.functions, result.flame)
|
||
return result
|
||
|
||
|
||
def _progress(phase, pct):
|
||
# 前导 \n：tqdm / 进度条类库会写"working..."到 stderr 但不换行,
|
||
# 直接接 PROGRESS 会拼成 "working...PROGRESS running 10",主进程正则匹配挂、
|
||
# 进度事件掉一档且把工程协议行混进 stderrTail 让用户看到 \r 进度条残片。
|
||
print(f"\nPROGRESS {phase} {pct}", file=sys.stderr, flush=True)
|
||
|
||
|
||
def _emit(result: AnalysisResult, out_path=None):
|
||
"""发出结果。
|
||
|
||
out_path 给定时结果写入独立文件，stdout/stderr 就完整留给用户代码 ——
|
||
否则用户代码里一句 sys.stdout.write("done") 或 print(..., end='') 就会和
|
||
JSON 挤在同一行，主进程按"最后一行"取 JSON 时解析失败，
|
||
用户看到的却是"结果解析失败"。不给 out_path 时退回 stdout（直接跑 CLI / 老测试）。
|
||
"""
|
||
payload = result.to_json()
|
||
if out_path:
|
||
with io.open(out_path, "w", encoding="utf-8", newline="\n") as f:
|
||
f.write(payload)
|
||
else:
|
||
print(payload, flush=True)
|
||
|
||
|
||
def _base(env, config, status, error=None, wall_time=None, functions=None, flame=None, calibration=None):
|
||
return AnalysisResult(
|
||
schemaVersion=SCHEMA_VERSION,
|
||
environment=env,
|
||
config=config,
|
||
status=status,
|
||
error=error,
|
||
wallTime=WallTime(seconds=wall_time, unit="s") if wall_time is not None else None,
|
||
functions=functions or [],
|
||
flame=flame,
|
||
calibration=calibration,
|
||
)
|
||
|
||
|
||
def main(argv=None):
|
||
ap = argparse.ArgumentParser()
|
||
ap.add_argument("--script", required=True)
|
||
ap.add_argument("--out", default=None,
|
||
help="结果 JSON 的输出文件；不给则打到 stdout")
|
||
ap.add_argument("--scope", choices=["user", "all"], default="user",
|
||
help="剖析范围：user（默认）只归因用户脚本里的函数；"
|
||
"all 包含所有非 cProfile 内部帧（标准库 + 第三方包 + 用户代码），"
|
||
"让耗时可以下钻到 import 的包里。")
|
||
ap.add_argument("--workdir", default=None,
|
||
help="校准文件落盘的临时目录；不给则 runner 自己 mkdtemp。"
|
||
"profiler-service 一般传过来(它创建的 pyprof-xxx),"
|
||
"直跑 CLI 不传也行。")
|
||
ap.add_argument("--hide-internal", dest="hide_internal",
|
||
action=argparse.BooleanOptionalAction, default=True,
|
||
help="过滤掉测试代码(tests/ / test_*.py / _pyrof_calib 等)+ "
|
||
"tottime=0 的量化噪声帧。默认 on,用户原话:「软件内部的测试部分"
|
||
"默认百分百过滤掉」。--no-hide-internal 关掉这条过滤,看完整数据。")
|
||
args = ap.parse_args(argv)
|
||
|
||
# 统一走这个 emit：结果去 --out 指定的文件，stdout 留给用户代码。
|
||
# 每条 emit 之前过 _cap_result 给 result.json 上一道防线 —— error.traceback /
|
||
# error.message / functions 列表都设硬上限，防止用户异常把磁盘 / IPC 撑爆。
|
||
def emit(result):
|
||
_emit(_cap_result(result), args.out)
|
||
|
||
env = collect_environment()
|
||
config: dict = {}
|
||
|
||
# 版本预检：主进程的 validateInterpreter 已经挡过一层，这里是纵深防御——
|
||
# 用户可能绕过 UI 直接调引擎，或 PATH 上的 python 在探测后被换掉。
|
||
# 低于 3.9 时后面的 list[X] 泛型注解会直接 SyntaxError，拿不到结构化错误。
|
||
if sys.version_info < (3, 9):
|
||
got = ".".join(str(p) for p in sys.version_info[:3])
|
||
emit(_base(env, config, "runtime_error",
|
||
error={"type": "PythonTooOld",
|
||
"message": f"引擎需要 Python 3.9 或更高版本，当前为 {got}"}))
|
||
return
|
||
|
||
# 一次 IO 拿到 src：语法预检和后续 profile 都用这份文本
|
||
# 之前没 try/except：临时目录被删 / 路径权限不够 → 进程崩，无 JSON 输出 → UI 显示 "引擎无输出"
|
||
# UnicodeDecodeError 不是 OSError 的子类：二进制文件 / 非 UTF-8 脚本会绕过这层保护再次崩溃。
|
||
try:
|
||
src = load_source(args.script)
|
||
except (OSError, UnicodeDecodeError) as e:
|
||
emit(_base(env, config, "runtime_error",
|
||
error={"type": type(e).__name__, "message": f"无法读取脚本: {e}"}))
|
||
return
|
||
|
||
# Syntax precheck (用文本而非再读一次)
|
||
# compile() 还会在源码含 null 字节时抛 ValueError，也一并按 syntax_error 归类
|
||
# 把 code object 留起来 —— profile_and_measure 还需要再 exec 一次,
|
||
# 直接复用这份省一次 read + compile。
|
||
try:
|
||
code = compile(src, args.script, "exec")
|
||
except (SyntaxError, ValueError) as e:
|
||
emit(_base(env, config, "syntax_error", error={"type": type(e).__name__, "message": str(e)}))
|
||
return
|
||
|
||
# v4 单跑架构：校准 + 用户代码 under cProfile 一次拿到归因 + instrumented 时间。
|
||
# 跟 v3 的「先裸跑测 wall_time 再 cProfile 跑」相比，
|
||
# 用户脚本的副作用(plot / file write / etc)只发生一次,语义更直观。
|
||
# workdir 优先用 caller 传的(profiler-service 已经 mkdir pyprof-xxx 落 _pyrof_calib.py
|
||
# 和用户脚本);caller 不传则自己 mkdtemp —— 这种情况跑完自己清掉。
|
||
workdir = args.workdir
|
||
workdir_owned = False
|
||
if not workdir:
|
||
workdir = tempfile.mkdtemp(prefix="pyprof-")
|
||
workdir_owned = True
|
||
|
||
try:
|
||
try:
|
||
_progress("calibrating", 5)
|
||
# 校准脚本落在 caller 提供的 workdir 里 —— profiler-service 已经 mkdir
|
||
# pyrof-xxx 装用户脚本,把 _pyrof_calib.py 放一起能让 rmtree(workdir) 一次
|
||
# 清干净。runner 自己 mkdtemp 时(直跑 CLI 场景)workdir_owned=True,
|
||
# calibration 拿到的是同一个目录,临时文件一起被 finally 兜底删。
|
||
calibration = calibrate_cprofile_overhead(workdir)
|
||
|
||
_progress("running", 30)
|
||
st, instrumented = profile_and_measure(
|
||
args.script, scope=args.scope, hide_internal=args.hide_internal,
|
||
code=code, src=src,
|
||
)
|
||
wall_time = _compute_wall_time(instrumented, calibration)
|
||
|
||
_progress("done", 100)
|
||
emit(_base(env, config, "ok",
|
||
wall_time=wall_time,
|
||
functions=st.functions,
|
||
flame=st.flame,
|
||
calibration=calibration))
|
||
return
|
||
except SystemExit as e:
|
||
# sys.exit() 抛 SystemExit（BaseException 的子类）—— 不被下面的 except BaseException 捕获，
|
||
# 子进程直接退出无 JSON，UI 就会显示 "引擎无输出 (exit N)"。转成 runtime_error。
|
||
# e.code 默认 None（无参 sys.exit()），文案要单独处理，否则会出现"sys.exit(None)"
|
||
code = e.code if isinstance(e.code, int) else (0 if e.code is None else f"非整数 {e.code!r}")
|
||
emit(_base(env, config, "runtime_error",
|
||
error={"type": "SystemExit", "message": f"用户代码调用 sys.exit({code})"}))
|
||
except BaseException as e:
|
||
# KeyboardInterrupt / GeneratorExit / 自定义 BaseException 子类 以及
|
||
# 普通 Exception 一并兜底 —— 之前再写一条 `except Exception` 是 dead code
|
||
# （BaseException 是 Exception 的父类,先匹配的赢）。父进程端有 cancel +
|
||
# 进程树 kill 在管 Ctrl+C,这里只负责兜底写 JSON。
|
||
emit(_base(env, config, "runtime_error",
|
||
error={"type": type(e).__name__, "message": str(e),
|
||
"traceback": traceback.format_exc()}))
|
||
finally:
|
||
# runner 自己 mkdtemp 的目录要自己清 —— profiler-service 已经传 --workdir 时
|
||
# 整个 workdir 由 caller 管(主进程清理临时目录的逻辑在 before-quit),这里只兜
|
||
# 「直接 python -m engine.runner」的直跑场景。
|
||
if workdir_owned:
|
||
shutil.rmtree(workdir, ignore_errors=True)
|
||
|
||
|
||
if __name__ == "__main__":
|
||
main() |