listen_py/main.sh

#!/bin/bash

# 性能测试脚本，用于顺序执行多种性能分析工具
# 颜色定义
RED='\033[0;31m'
GREEN='\033[0;32m'
YELLOW='\033[1;33m'
BLUE='\033[0;34m'
NC='\033[0m' # No Color

# 默认参数
RESULT_DIR="results"
TEST_SCRIPT=""  # 待测试的脚本
ARGS_SCRIPT="" # 待测试脚本的参数
PSUTIL_SCRIPT="monitor_util.py"  # 监测工具脚本
INTEGRATED_SCRIPT="integrate_results.py"  # 融合工具脚本
PLOT_SCRIPT="plot_offline.py"  # 画图工具脚本
CPROFILE_OUTPUT="$RESULT_DIR/cprofile_output.prof"  # cProfile 结果文件
MEMORY_PROFILE_OUTPUT="$RESULT_DIR/mprofile_output.dat"  # memory_profiler输出文件
#MEMORY_PLOT_OUTPUT="$RESULT_DIR/mprofile_memory_plot.png"  # memory_profiler 绘制的曲线图
FLAMEGRAPH_OUTPUT="$RESULT_DIR/pyspy_output.svg"  # py-spy 输出文件
PSUTIL_OUTPUT="$RESULT_DIR/psutil_output.csv"  # psutil 输出文件
INTEGRATED_REPORT="$RESULT_DIR/performance_analysis_report.json"  # 融合结果文件
INTEGRATED_PLOT="$RESULT_DIR/performance_analysis_report.png"  # 融合结果图

# 日志函数
log_info() {
    echo -e "${BLUE}[INFO]${NC} $1"
}

log_success() {
    echo -e "${GREEN}[SUCCESS]${NC} $1"
}

log_warning() {
    echo -e "${YELLOW}[WARNING]${NC} $1"
}

log_error() {
    echo -e "${RED}[ERROR]${NC} $1"
}

# 显示标题
show_header() {
    echo "========================================="
    echo "Python脚本性能分析工具链"
    echo "========================================="
    sleep 2
}

# 显示步骤标题
show_step() {
    echo ""
    echo "-----------------------------------------"
    echo "$1"
    echo "-----------------------------------------"
}

# 检查脚本文件是否存在
check_script_exists() {
    if [ -f "config.json" ]; then
        # 读取脚本路径
        TEST_SCRIPT=$(python -c "
import json
try:
    with open('config.json', 'r') as f:
        config = json.load(f)
        script_path = config.get('test_script_dir', '')
        print(script_path)
except Exception as e:
    print('')
" 2>/dev/null)

        if [ -n "$TEST_SCRIPT" ]; then
            log_info "从 config.json 读取测试脚本: $TEST_SCRIPT"
        else
            log_warning "config.json 中未找到有效的 test_script_dir"
            return 1
        fi
        # 确认脚本存在
        if [ ! -f "$TEST_SCRIPT" ]; then
        log_error "测试脚本 $TEST_SCRIPT 不存在"
        return 1
        fi
        # 读取脚本参数
        ARGS_SCRIPT=$(python -c "
import json
try:
    with open('config.json', 'r') as f:
        config = json.load(f)
        args = config.get('args', {})
        cmd_args = ' '.join('--{} {}'.format(k, v) for k, v in args.items())
        print(cmd_args)
except Exception as e:
    print('')
" 2>/dev/null)
        if [ -n "$ARGS_SCRIPT" ]; then
            log_info "从 config.json 读取脚本参数: $ARGS_SCRIPT"
        else
            log_warning "config.json 中未配置有效的 args"
        fi

    else
        log_warning "未找到 config.json 配置文件"
        return 1
    fi
    sleep 2
    return 0
}
# 创建结果保存目录
mk_dir(){
    if [ ! -d "$RESULT_DIR" ]; then
        mkdir -p "$RESULT_DIR"
        log_info "创建结果保存目录: $RESULT_DIR"
    else
        log_info "使用现有结果保存目录: $RESULT_DIR"
    fi
}

# 1. 使用cProfile进行性能分析
run_cprofile() {
    show_step "步骤 1: 使用 cProfile 进行性能分析"
    log_info "准备对 $TEST_SCRIPT 进行性能分析..."

    # 将之前的结果删除
    rm -rf "$CPROFILE_OUTPUT"
    log_info "键入指令: python -m cProfile -o $CPROFILE_OUTPUT $TEST_SCRIPT $ARGS_SCRIPT"
    if python -m cProfile -o "$CPROFILE_OUTPUT" "$TEST_SCRIPT" $ARGS_SCRIPT; then
        log_success "cProfile 分析完成，结果保存至 $CPROFILE_OUTPUT"
        return 0
    else
        log_error "cProfile 分析失败"
        return 1
    fi
}

# 3. 使用memory_profiler分析内存使用情况
run_memory_profiler() {
    show_step "步骤 3: 使用 memory_profiler 分析内存使用情况"
    # 将之前的结果删除
    rm -rf "$MEMORY_PROFILE_OUTPUT"
    # 运行内存分析
    if command -v mprof &> /dev/null; then
        log_info "键入指令: mprof run --output $MEMORY_PROFILE_OUTPUT $TEST_SCRIPT $ARGS_SCRIPT"
        if mprof run --output "$MEMORY_PROFILE_OUTPUT" "$TEST_SCRIPT" $ARGS_SCRIPT; then
            log_success "memory_profiler 分析完成，结果保存至 $MEMORY_PROFILE_OUTPUT"
            # 绘制内存占用曲线（弃用）
#            log_info "绘制内存占用曲线"
#            if MPLBACKEND=Agg mprof plot --flame  "$MEMORY_PROFILE_OUTPUT" -o "$MEMORY_PLOT_OUTPUT"; then
#                log_success "绘制完成，结果保存至 $MEMORY_PLOT_OUTPUT"
#            else
#                log_error "绘制失败"
#            fi
        else
            log_error "memory_profiler 分析失败"
        fi
    else
        log_warning "无法安装 memory_profiler，跳过内存分析步骤"
    fi
}

# 2. 使用py-spy记录性能数据
run_py_spy() {
    show_step "步骤 2: 使用 py-spy record 进行性能分析"
    # 删除之前的结果
    rm -rf "$FLAMEGRAPH_OUTPUT"
    # 根据操作系统类型选择合适的py-spy参数，需要统计子线程和C拓展
    if [[ "$OSTYPE" == "msys" ]] || [[ "$OSTYPE" == "win32" ]]; then
        # Windows环境下使用--nonblocking选项
        log_info "Windows环境下使用--nonblocking选项"
        PYSPY_CMD="py-spy record -o $FLAMEGRAPH_OUTPUT --duration 35 --native --subprocesses --nonblocking -- python $TEST_SCRIPT $ARGS_SCRIPT"
    elif [[ "$OSTYPE" == "linux-gnu"* ]]; then
        # Linux环境下使用sudo运行
        log_info "Linux环境下使用sudo运行py-spy"
        PYSPY_CMD="sudo env \"PATH=\$PATH\" py-spy record -o $FLAMEGRAPH_OUTPUT --duration 35 --native --subprocesses -- python $TEST_SCRIPT $ARGS_SCRIPT"
    else
        # 其他Unix系统使用默认选项
        PYSPY_CMD="py-spy record -o $FLAMEGRAPH_OUTPUT --duration 35 --native --subprocesses -- python $TEST_SCRIPT $ARGS_SCRIPT"
    fi
    
    log_info "执行命令: $PYSPY_CMD"
    eval "$PYSPY_CMD"
    PYSPY_PID=$!
    if wait $PYSPY_PID; then
        log_success "py-spy record 分析完成，火焰图保存至 $FLAMEGRAPH_OUTPUT"
        return 0
    else
        log_error "py-spy record 分析失败"
        if [[ "$OSTYPE" == "linux-gnu"* ]]; then
            log_warning "提示: 可能需要管理员权限，请使用 'sudo ./$0' 运行此脚本"
        fi
        return 1
    fi
}
# 4. 使用psutil获取进程信息
run_psutil() {
  show_step "步骤 4: 使用 psutil 获取进程信息"

  # 将之前的结果文件删除
  rm -rf "$PSUTIL_OUTPUT"
  # 在后台运行test_script.py并获取其PID
  log_info "在后台启动 $TEST_SCRIPT $ARGS_SCRIP ..."
  python "$TEST_SCRIPT" $ARGS_SCRIPT &
  TEST_PID=$!
  
  # 等待一段时间确保进程已经完全启动
  sleep 2
  
  # 验证进程是否仍在运行
  if ! kill -0 $TEST_PID 2>/dev/null; then
    log_error "测试脚本启动失败"
    return 1
  fi
  
  log_info "测试脚本 PID: $TEST_PID"
  
  # 使用monitor.py监控该进程
  log_info "开始监控进程资源使用情况..."
  if [[ "$OSTYPE" == "linux-gnu"* ]]; then
      # Linux环境下使用sudo运行monitor.py
      sudo env "PATH=$PATH" python "$PSUTIL_SCRIPT" --pid $TEST_PID --output "$PSUTIL_OUTPUT" &
  else
      python "$PSUTIL_SCRIPT" --pid $TEST_PID --output "$PSUTIL_OUTPUT" &
  fi
  MONITOR_PID=$!
  
  # 等待test_script.py执行完成
  wait $TEST_PID
  sleep 4
  log_success "测试脚本执行完成"
  # 终止monitor.py
  kill $MONITOR_PID 2>/dev/null || true
  
  log_success "资源监控完成，结果保存至 $PSUTIL_OUTPUT"
}
# 0.检查环境完备性
check_env(){
    # 检查是否安装了py-spy
    if ! command -v py-spy &> /dev/null; then
        log_warning "未找到 py-spy ，尝试安装..."
        if pip3 install py-spy -i https://pypi.tuna.tsinghua.edu.cn/simple; then
            log_success "py-spy 安装成功"
        else
            log_error "py-spy 安装失败"
            return 1
        fi
    fi

    # 检查是否安装了memory_profiler
    if ! command -v mprof &> /dev/null; then
        log_warning "未找到 memory_profiler (mprof)，尝试安装..."
        if pip3 install memory-profiler -i https://pypi.tuna.tsinghua.edu.cn/simple; then
            log_success "memory_profiler 安装成功"
            echo "适配安装matplotlib: "
            conda install matplotlib
        else
            log_error "memory_profiler 安装失败"
            return 1
        fi
    fi
    # 检查是否安装了psutil
    if ! python -c "import psutil" &> /dev/null; then
      log_warning "未找到 psutil ，尝试安装..."
      if pip3 install psutil -i https://pypi.tuna.tsinghua.edu.cn/simple; then
        log_success "psutil 安装成功"
      else
        log_error "psutil 安装失败"
        return 1
      fi
    fi
    # 检查是否安装nvidia-ml-py
    if ! python -c "import pynvml" &> /dev/null; then
      log_warning "未找到 nvidia_ml_py ，尝试安装..."
      if pip3 install nvidia-ml-py -i https://pypi.tuna.tsinghua.edu.cn/simple; then
        log_success "nvidia_ml_py 安装成功"
      else
        log_error "nvidia_ml_py 安装失败"
        return 1
      fi
    fi
}
# 自动打开SVG文件
open_svg_file() {
    if [ -f "$FLAMEGRAPH_OUTPUT" ]; then
        log_info "正在尝试在浏览器中打开火焰图..."
        
        # 检测是否在WSL环境中
        if grep -qE "(Microsoft|WSL)" /proc/version 2>/dev/null; then
            # WSL环境 - 使用PowerShell打开
            log_info "检测到WSL环境，使用PowerShell打开文件"
            if powershell.exe -Command "Start-Process '$(wslpath -w "$(pwd)/$FLAMEGRAPH_OUTPUT")'" 2>/dev/null; then
                log_success "SVG文件已在Windows默认应用中打开"
                return 0
            elif cmd.exe /c "start $(wslpath -w "$(pwd)/$FLAMEGRAPH_OUTPUT")" 2>/dev/null; then
                log_success "SVG文件已在Windows默认应用中打开"
                return 0
            else
                log_warning "无法在Windows主机上打开文件，请手动打开 $FLAMEGRAPH_OUTPUT"
            fi
        elif command -v xdg-open >/dev/null 2>&1; then
            # 标准Linux环境
            if xdg-open "$FLAMEGRAPH_OUTPUT"; then
                log_success "SVG文件已在默认应用中打开"
                return 0
            else
                log_warning "无法使用xdg-open打开文件"
            fi
        elif command -v gnome-open >/dev/null 2>&1; then
            if gnome-open "$FLAMEGRAPH_OUTPUT"; then
                log_success "SVG文件已在默认应用中打开"
                return 0
            else
                log_warning "无法使用gnome-open打开文件"
            fi
        else
            log_warning "找不到合适的打开命令，请手动在浏览器中打开 $FLAMEGRAPH_OUTPUT"
        fi
    fi
    return 1
}

# 自动打开prof文件
open_prof_file() {
  # 检查snakeviz是否存在，不存在就尝试自动安装
  if [ -f "$CPROFILE_OUTPUT" ]; then
    log_info "自动打开cProfile分析结果"
    if ! command -v snakeviz &> /dev/null; then
      log_warning "未找到snakeviz，尝试安装..."
      if pip3 install snakeviz -i https://pypi.tuna.tsinghua.edu.cn/simple; then
        log_success "snakeviz 安装成功"
      else
        log_error "snakeviz 安装失败"
        return 1
      fi
    fi

    if command -v snakeviz &> /dev/null; then
      log_info "正在尝试使用snakeviz打开prof文件..."
      if snakeviz "$CPROFILE_OUTPUT"; then
        log_success "prof文件浏览完成"
        return 0
      else
        log_warning "无法使用snakeviz打开prof文件"
      fi
    else
      log_warning "snakeviz未安装，请手动打开 $CPROFILE_OUTPUT"
    fi
  else
    log_error "$CPROFILE_OUTPUT 文件不存在"
  fi
   return 1
}

# 显示结果摘要
show_summary() {
    show_step "性能分析结果摘要"
    echo "1. cProfile 输出: $CPROFILE_OUTPUT (可用于 pstats 分析)"
    echo "2. py-spy 火焰图: $FLAMEGRAPH_OUTPUT (可在浏览器中打开查看)"
    echo "3. memory_profiler 输出: $MEMORY_PROFILE_OUTPUT (内存使用情况)"
    echo "4. psutil 输出: $PSUTIL_OUTPUT (硬件资源使用情况)"
    log_success "性能分析工具链执行完毕！"
}
# 整合所有分析结果
integrate_all_results(){
  show_step "步骤5 整合所有分析结果"

  # 检查python脚本是否存在
  if [ ! -f "$INTEGRATED_SCRIPT" ]; then
     log_warning "整合脚本 $INTEGRATED_SCRIPT 不存在"
     return 1
  else
    echo "结果整合中..."
    # 构建命令参数
    CMD="python $INTEGRATED_SCRIPT"
    if [ -f "$CPROFILE_OUTPUT" ]; then
        CMD="$CMD --cprofile $CPROFILE_OUTPUT"
    fi
    if [ -f "$MEMORY_PROFILE_OUTPUT" ]; then
        CMD="$CMD --memory $MEMORY_PROFILE_OUTPUT"
    fi
    if [ -f "$FLAMEGRAPH_OUTPUT" ]; then
        CMD="$CMD --flamegraph $FLAMEGRAPH_OUTPUT"
    fi
    if [ -f "$PSUTIL_OUTPUT" ]; then
        CMD="$CMD --psutil $PSUTIL_OUTPUT"
    fi
    CMD="$CMD --output $INTEGRATED_REPORT"
     # 执行整合脚本
    log_info "执行整合命令: $CMD"
    if eval "$CMD"; then
        log_success "所有分析结果已整合至 $INTEGRATED_REPORT"
    else
        log_error "整合过程出错"
        # 检查pandas是否存在
        if ! python -c "import pandas" &> /dev/null; then
           log_warning "未找到pandas，尝试安装..."
           if pip3 install pandas -i https://pypi.tuna.tsinghua.edu.cn/simple; then
              log_success "pandas 安装成功"
           else
              log_error "pandas 安装失败"
              return 1
           fi
           # 尝试再次执行命令
           if eval "$CMD"; then
              log_success "所有分析结果已整合至 $INTEGRATED_REPORT"
           else
             log_error "无法执行整合命令"
             return 1
           fi
        else
           return 1
        fi
    fi
    # 绘制融合结果图
    if [ -f "$PLOT_SCRIPT" ]; then
      log_info "绘制融合结果图"
      if python "$PLOT_SCRIPT" --input "$INTEGRATED_REPORT" --output "$INTEGRATED_PLOT"; then
        log_success "融合结果图已保存至 $INTEGRATED_PLOT"
      else
        log_error "无法绘制融合结果图"
        return 1
      fi
    else
       log_warning "未找到绘制脚本 $PLOT_SCRIPT"
       return 1
    fi
  fi
  return 0
}



# 主函数
main() {
    # 显示标题
    show_heade
    # 检查环境
    if ! check_env;then
       log_error "请检查环境是否满足要求."
    fi
    # 检查脚本是否存在
    if ! check_script_exists; then
       log_error "请检查测试脚本是否存在."
       exit 1
    fi
    # 创建结果保存目录
    mk_dir
    # 执行各项分析
    run_cprofile  # cProfile 算法性能分析，统计主线程的函数调用次数和运行时间
    run_py_spy  # py-spy 火焰图
    run_memory_profiler  # 内存使用情况分析
    run_psutil  # 使用psutil监控进程资源使用情况
    show_summary  # 显示结果摘要
    integrate_all_results  # 整合所有结果
    # 自动打开结果文件
    open_svg_file
    open_prof_file
}

# 执行主函数
main "$@"