feat: 实现大规模并行预测功能 (v2.0.0)

新增功能:
- 新增统一批量预测工具 utils/batch_predictor.py
  * 支持单进程/多进程并行模式
  * 灵活的 GPU 配置和显存自动计算
  * 自动临时文件管理和断点续传
  * 完整的 CLI 参数支持(Click 框架)

- 新增 Shell 脚本集合 scripts/
  * run_parallel_predict.sh - 并行预测脚本
  * run_single_predict.sh - 单进程预测脚本
  * merge_results.sh - 结果合并脚本

性能优化:
- 解决 CUDA + multiprocessing fork 死锁问题
  * 使用 spawn 模式替代 fork
  * 文件描述符级别的输出重定向

- 优化预测性能
  * XGBoost OpenMP 多线程(利用所有 CPU 核心)
  * 预加载模型减少重复加载
  * 大批量处理降低函数调用开销
  * 实际加速比:2-3x(12进程 vs 单进程)

- 优化输出显示
  * 抑制模型加载时的权重信息
  * 只显示进度条和关键统计
  * 临时文件自动保存到专门目录

文档更新:
- README.md 新增"大规模并行预测"章节
- README.md 新增"性能优化说明"章节
- 添加详细的使用示例和参数说明
- 更新项目结构和版本信息

技术细节:
- 每个模型实例约占用 2.5GB GPU 显存
- 显存计算公式:建议进程数 = GPU显存(GB) / 2.5
- GPU 瓶颈占比:MolE 表示生成 94%
- 非 GIL 问题:计算密集任务在 C/CUDA 层

Breaking Changes:
- 废弃旧的独立预测脚本,统一使用新工具

相关 Issue: 解决 #并行预测卡死问题
测试平台: Linux, 256 CPU cores, NVIDIA RTX 5090 32GB
This commit is contained in:
2025-10-18 20:53:39 +08:00
parent 4745ce3884
commit a8fea027ac
8 changed files with 1202 additions and 51 deletions

117
scripts/run_parallel_predict.sh Executable file
View File

@@ -0,0 +1,117 @@
#!/bin/bash
#
# 并行预测脚本 - 将大型数据集分成 N 份并行处理
#
# 用法:
# bash scripts/run_parallel_predict.sh [输入文件] [输出文件] [进程数] [GPU设备]
#
# 示例:
# bash scripts/run_parallel_predict.sh Data/fragment/Frags-Enamine-18M.csv output.csv 4 cuda:0
#
# 说明:
# - 进程数建议 = GPU显存(GB) / 2.5
# - 32GB 显存建议使用 4-12 个进程
# - 所有进程将使用同一个 GPU串行使用
#
# 作者: AI Assistant
# 日期: 2025-10-17
set -e
# ============================================================================
# 参数设置
# ============================================================================
# 默认参数
INPUT_FILE="${1:-Data/fragment/Frags-Enamine-18M.csv}"
OUTPUT_FILE="${2:-Data/fragment/Frags-Enamine-18M_predicted.csv}"
N_PROCESSES="${3:-4}"
GPU_DEVICE="${4:-cuda:0}"
# SMILES 和 ID 列名
SMILES_COLUMN="smiles"
ID_COLUMN="chem_id"
# ============================================================================
# 打印配置信息
# ============================================================================
echo "============================================================"
echo "🚀 并行预测 Enamine 抗菌活性"
echo "============================================================"
echo ""
echo "配置:"
echo " 输入文件: $INPUT_FILE"
echo " 输出文件: $OUTPUT_FILE"
echo " 并行进程数: $N_PROCESSES"
echo " GPU 设备: $GPU_DEVICE"
echo " SMILES 列: $SMILES_COLUMN"
echo " ID 列: $ID_COLUMN"
echo ""
echo "说明:"
echo " - 使用批量预测工具: utils/batch_predictor.py"
echo " - 每个模型实例约占用 2.5GB GPU 显存"
echo " - 建议进程数 = GPU显存(GB) / 2.5"
echo " - 32GB 显存 → 建议 4-12 个进程"
echo ""
echo "============================================================"
# ============================================================================
# 检查文件
# ============================================================================
if [ ! -f "$INPUT_FILE" ]; then
echo "❌ 错误: 输入文件不存在: $INPUT_FILE"
exit 1
fi
if [ ! -f "utils/batch_predictor.py" ]; then
echo "❌ 错误: 预测脚本不存在: utils/batch_predictor.py"
exit 1
fi
# ============================================================================
# 运行预测
# ============================================================================
echo ""
echo "🔄 开始预测..."
echo ""
# 使用 nohup 在后台运行(可选)
if [ "$5" == "--background" ]; then
LOG_FILE="${OUTPUT_FILE%.csv}.log"
echo "后台运行模式,日志保存到: $LOG_FILE"
nohup pixi run python utils/batch_predictor.py \
--input "$INPUT_FILE" \
--output "$OUTPUT_FILE" \
--smiles-column "$SMILES_COLUMN" \
--id-column "$ID_COLUMN" \
--device "$GPU_DEVICE" \
--n-processes "$N_PROCESSES" \
--batch-size 1000 \
--verbose \
> "$LOG_FILE" 2>&1 &
echo "✓ 进程已在后台启动"
echo " 进程 ID: $!"
echo " 查看日志: tail -f $LOG_FILE"
else
# 前台运行
pixi run python utils/batch_predictor.py \
--input "$INPUT_FILE" \
--output "$OUTPUT_FILE" \
--smiles-column "$SMILES_COLUMN" \
--id-column "$ID_COLUMN" \
--device "$GPU_DEVICE" \
--n-processes "$N_PROCESSES" \
--batch-size 1000 \
--verbose
fi
echo ""
echo "============================================================"
echo "✅ 完成"
echo "============================================================"