#!/usr/bin/env python3 # -*- coding: utf-8 -*- """ AIM医疗问诊工作流批处理系统 使用多线程并行处理数据集中的所有病例样本 """ import argparse import json import logging import os import sys import time import threading import glob from concurrent.futures import ThreadPoolExecutor, as_completed from datetime import datetime from typing import Dict, Any, List, Optional # 导入本地模块 from workflow import MedicalWorkflow from config import LLM_CONFIG class BatchProcessor: """批处理管理器，负责协调多线程执行和状态管理""" def __init__(self, num_threads: int = 20): self.num_threads = num_threads self.lock = threading.Lock() # 线程安全锁 self.processed_count = 0 # 已处理样本数 self.success_count = 0 # 成功处理数 self.failed_count = 0 # 失败处理数 self.skipped_count = 0 # 跳过的样本数 self.results = [] # 结果列表 self.failed_samples = [] # 失败样本列表 self.start_time = None # 开始时间 def update_progress(self, success: bool, result: Dict[str, Any] = None, error: Exception = None, sample_index: int = None): """线程安全地更新处理进度""" with self.lock: self.processed_count += 1 if success: self.success_count += 1 if result: self.results.append(result) else: self.failed_count += 1 if error and sample_index is not None: self.failed_samples.append({ 'sample_index': sample_index, 'error': str(error), 'timestamp': datetime.now().isoformat() }) def update_skipped(self, sample_index: int): """线程安全地更新跳过样本计数""" with self.lock: self.skipped_count += 1 logging.info(f"样本 {sample_index} 已完成，跳过处理") def get_progress_stats(self) -> Dict[str, Any]: """获取当前进度统计""" with self.lock: elapsed_time = time.time() - self.start_time if self.start_time else 0 return { 'processed': self.processed_count, 'success': self.success_count, 'failed': self.failed_count, 'skipped': self.skipped_count, 'success_rate': self.success_count / max(self.processed_count, 1), 'elapsed_time': elapsed_time, 'samples_per_minute': self.processed_count / max(elapsed_time / 60, 0.01) } def setup_logging(log_level: str = "INFO") -> None: """设置日志记录配置""" logging.basicConfig( level=getattr(logging, log_level.upper()), format='%(asctime)s - %(levelname)s - %(message)s', handlers=[ logging.StreamHandler(), logging.FileHandler(f'batch_processing_{datetime.now().strftime("%Y%m%d_%H%M%S")}.log') ] ) def parse_arguments() -> argparse.Namespace: """解析命令行参数""" parser = argparse.ArgumentParser( description="AIM医疗问诊工作流批处理系统", formatter_class=argparse.ArgumentDefaultsHelpFormatter ) # 数据和输出配置 parser.add_argument( '--dataset-path', type=str, default='dataset/bbb.json', help='数据集JSON文件路径' ) parser.add_argument( '--log-dir', type=str, default='results/results09010-score_driven', help='日志文件保存目录' ) parser.add_argument( '--output-dir', type=str, default='batch_results', help='批处理结果保存目录' ) # 执行参数 parser.add_argument( '--num-threads', type=int, default=45, help='并行处理线程数' ) parser.add_argument( '--max-steps', type=int, default=30, help='每个工作流的最大执行步数' ) parser.add_argument( '--start-index', type=int, default=0, help='开始处理的样本索引' ) parser.add_argument( '--end-index', type=int, default=5000, help='结束处理的样本索引（不包含）' ) parser.add_argument( '--sample-limit', type=int, default=None, help='限制处理的样本数量（用于测试）' ) # 模型配置 available_models = list(LLM_CONFIG.keys()) parser.add_argument( '--model-type', type=str, choices=available_models, default='openai-mirror/gpt-oss-20b', help=f'使用的语言模型类型，可选: {", ".join(available_models)}' ) parser.add_argument( '--list-models', action='store_true', help='显示所有可用的模型配置并退出' ) parser.add_argument( '--model-config', type=str, default=None, help='模型配置JSON字符串（可选，覆盖默认配置）' ) parser.add_argument( '--controller-mode', type=str, choices=['normal', 'sequence', 'score_driven'], default='score_driven', help='任务控制器模式：normal为智能模式（需要LLM推理），sequence为顺序模式（直接选择第一个任务），score_driven为分数驱动模式（选择当前任务组中分数最低的任务）' ) # 调试和日志 parser.add_argument( '--log-level', type=str, choices=['DEBUG', 'INFO', 'WARNING', 'ERROR'], default='INFO', help='日志记录级别' ) parser.add_argument( '--progress-interval', type=int, default=10, help='进度报告间隔（秒）' ) parser.add_argument( '--dry-run', action='store_true', help='试运行模式，只验证配置不执行处理' ) return parser.parse_args() def is_case_completed(log_dir: str, case_index: int) -> bool: """ 检查指定case是否已经完成工作流 如果存在不完整的文件则删除，确保每个case在目录中只出现一次 Args: log_dir: 日志目录 case_index: case序号 Returns: bool: 如果case已完成返回True，否则返回False """ # 构建文件路径模式：workflow_*_case_{case_index:04d}.jsonl pattern = os.path.join(log_dir, f"workflow_*_case_{case_index:04d}.jsonl") matching_files = glob.glob(pattern) if not matching_files: return False # 应该只有一个匹配的文件 if len(matching_files) > 1: logging.warning(f"发现多个匹配文件 case {case_index}: {matching_files}") # 检查每个匹配的文件 for log_file in matching_files: try: with open(log_file, 'r', encoding='utf-8') as f: # 读取最后一行 lines = f.readlines() if not lines: # 文件为空，删除 os.remove(log_file) logging.info(f"删除空文件: {log_file}") continue last_line = lines[-1].strip() if not last_line: # 最后一行为空，删除 os.remove(log_file) logging.info(f"删除最后一行为空的文件: {log_file}") continue # 解析最后一行的JSON try: last_entry = json.loads(last_line) if last_entry.get("event_type") == "workflow_complete": # 找到完整的文件 logging.info(f"发现已完成的case {case_index}: {log_file}") return True else: # 文件不完整，删除 os.remove(log_file) logging.info(f"删除不完整的文件: {log_file}") continue except json.JSONDecodeError: # JSON解析失败，删除文件 os.remove(log_file) logging.info(f"删除JSON格式错误的文件: {log_file}") continue except Exception as e: logging.warning(f"检查文件 {log_file} 时出错: {e}") # 出现异常也删除文件，避免后续问题 try: os.remove(log_file) logging.info(f"删除异常文件: {log_file}") except: pass continue # 所有匹配的文件都被删除或没有完整的文件 return False def load_dataset(dataset_path: str, start_index: int = 0, end_index: Optional[int] = None, sample_limit: Optional[int] = None) -> List[Dict[str, Any]]: """加载和验证数据集""" logging.info(f"正在加载数据集: {dataset_path}") if not os.path.exists(dataset_path): raise FileNotFoundError(f"数据集文件不存在: {dataset_path}") try: with open(dataset_path, 'r', encoding='utf-8') as f: full_dataset = json.load(f) except json.JSONDecodeError as e: raise ValueError(f"数据集JSON格式错误: {e}") except Exception as e: raise Exception(f"加载数据集失败: {e}") if not isinstance(full_dataset, list): raise ValueError("数据集应该是包含病例的JSON数组") total_samples = len(full_dataset) logging.info(f"数据集总样本数: {total_samples}") # 确定处理范围 if end_index is None: end_index = total_samples end_index = min(end_index, total_samples) start_index = max(0, start_index) if sample_limit: end_index = min(start_index + sample_limit, end_index) if start_index >= end_index: raise ValueError(f"无效的索引范围: start_index={start_index}, end_index={end_index}") # 提取指定范围的数据 dataset = full_dataset[start_index:end_index] logging.info(f"将处理样本范围: [{start_index}, {end_index}), 共 {len(dataset)} 个样本") # 验证数据格式 for i, sample in enumerate(dataset[:5]): # 只验证前5个样本 if not isinstance(sample, dict): raise ValueError(f"样本 {start_index + i} 格式错误，应为字典类型") required_keys = ['病案介绍'] for key in required_keys: if key not in sample: logging.warning(f"样本 {start_index + i} 缺少必需字段: {key}") return dataset def process_single_sample(sample_data: Dict[str, Any], sample_index: int, args: argparse.Namespace, processor: BatchProcessor) -> Dict[str, Any]: """处理单个样本的工作函数""" thread_id = threading.current_thread().ident start_time = time.time() try: # 使用 LLM_CONFIG 作为基础配置 # BaseAgent 会根据 model_type 自动选择正确的模型配置 llm_config = LLM_CONFIG.copy() # 如果用户提供了额外的模型配置，则合并到对应的模型配置中 if args.model_config: try: user_config = json.loads(args.model_config) # 更新选定模型的配置 if args.model_type in llm_config: llm_config[args.model_type]["params"].update(user_config.get("params", {})) else: logging.warning(f"样本 {sample_index}: 模型类型 {args.model_type} 不存在，忽略用户配置") except json.JSONDecodeError: logging.warning(f"样本 {sample_index}: 模型配置JSON格式错误，使用默认配置") # 创建工作流实例 workflow = MedicalWorkflow( case_data=sample_data, model_type=args.model_type, llm_config=llm_config, max_steps=args.max_steps, log_dir=args.log_dir, case_index=sample_index, controller_mode=args.controller_mode ) # 执行工作流 logging.debug(f"线程 {thread_id}: 开始处理样本 {sample_index}") log_file_path = workflow.run() execution_time = time.time() - start_time # 获取执行结果 workflow_status = workflow.get_current_status() medical_summary = workflow.get_medical_summary() # 构建结果 result = { 'sample_index': sample_index, 'thread_id': thread_id, 'execution_time': execution_time, 'log_file_path': log_file_path, 'workflow_status': workflow_status, 'medical_summary': medical_summary, 'processed_at': datetime.now().isoformat() } # 更新进度 processor.update_progress(success=True, result=result) logging.info(f"样本 {sample_index} 处理完成 (耗时: {execution_time:.2f}s, " f"步数: {workflow_status['current_step']}, " f"成功: {workflow_status['workflow_success']})") return result except Exception as e: execution_time = time.time() - start_time error_msg = f"样本 {sample_index} 处理失败: {str(e)}" logging.error(error_msg) processor.update_progress(success=False, error=e, sample_index=sample_index) # 返回错误结果 return { 'sample_index': sample_index, 'thread_id': thread_id, 'execution_time': execution_time, 'error': str(e), 'processed_at': datetime.now().isoformat(), 'success': False } def print_progress_report(processor: BatchProcessor, total_samples: int): """打印进度报告""" stats = processor.get_progress_stats() print(f"\n=== 处理进度报告 ===") print(f"已处理: {stats['processed']}/{total_samples} ({stats['processed']/total_samples:.1%})") print(f"成功: {stats['success']} | 失败: {stats['failed']} | 跳过: {stats['skipped']} | 成功率: {stats['success_rate']:.1%}") print(f"耗时: {stats['elapsed_time']:.1f}s | 处理速度: {stats['samples_per_minute']:.1f} 样本/分钟") remaining_samples = total_samples - stats['processed'] - stats['skipped'] print(f"预计剩余时间: {remaining_samples / max(stats['samples_per_minute'] / 60, 0.01):.1f}s") print("=" * 50) def run_workflow_batch(dataset: List[Dict[str, Any]], args: argparse.Namespace) -> Dict[str, Any]: """执行批量工作流处理""" total_samples = len(dataset) logging.info(f"开始批量处理 {total_samples} 个样本，使用 {args.num_threads} 个线程") # 创建批处理管理器 processor = BatchProcessor(num_threads=args.num_threads) processor.start_time = time.time() # 创建输出目录 os.makedirs(args.output_dir, exist_ok=True) os.makedirs(args.log_dir, exist_ok=True) # 启动进度监控线程 def progress_monitor(): while processor.processed_count + processor.skipped_count < total_samples: time.sleep(args.progress_interval) if processor.processed_count + processor.skipped_count < total_samples: print_progress_report(processor, total_samples) progress_thread = threading.Thread(target=progress_monitor, daemon=True) progress_thread.start() try: # 使用线程池执行批处理 with ThreadPoolExecutor(max_workers=args.num_threads) as executor: # 提交所有任务 future_to_index = {} for i, sample_data in enumerate(dataset): sample_index = args.start_index + i # 检查case是否已经完成 if is_case_completed(args.log_dir, sample_index): processor.update_skipped(sample_index) continue future = executor.submit( process_single_sample, sample_data, sample_index, args, processor ) future_to_index[future] = sample_index # 等待所有任务完成 for future in as_completed(future_to_index): sample_index = future_to_index[future] try: _ = future.result() # 结果已经在process_single_sample中处理 except Exception as e: logging.error(f"线程执行异常 (样本 {sample_index}): {e}") except KeyboardInterrupt: logging.warning("收到中断信号，正在停止处理...") executor.shutdown(wait=False) raise # 最终进度报告 total_time = time.time() - processor.start_time stats = processor.get_progress_stats() print_progress_report(processor, total_samples) # 构建最终结果摘要 summary = { 'total_samples': total_samples, 'processed_samples': processor.processed_count, 'successful_samples': processor.success_count, 'failed_samples': processor.failed_count, 'skipped_samples': processor.skipped_count, 'success_rate': stats['success_rate'], 'total_execution_time': total_time, 'average_time_per_sample': total_time / max(processor.processed_count, 1), 'samples_per_minute': stats['samples_per_minute'], 'failed_sample_details': processor.failed_samples, 'processing_config': { 'num_threads': args.num_threads, 'model_type': args.model_type, 'max_steps': args.max_steps, 'dataset_range': f"[{args.start_index}, {args.start_index + len(dataset)})" } } return { 'summary': summary, 'results': processor.results } def generate_summary_report(batch_results: Dict[str, Any], output_path: str) -> None: """生成详细的执行摘要报告""" summary = batch_results['summary'] results = batch_results['results'] timestamp = datetime.now().strftime('%Y%m%d_%H%M%S') # 生成JSON格式的详细报告 detailed_report = { 'batch_execution_summary': summary, 'sample_results': results, 'generated_at': datetime.now().isoformat(), 'report_version': '1.0' } report_file = os.path.join(output_path, f'batch_report_{timestamp}.json') try: with open(report_file, 'w', encoding='utf-8') as f: json.dump(detailed_report, f, ensure_ascii=False, indent=2) logging.info(f"详细报告已保存: {report_file}") # 生成人类可读的摘要 summary_file = os.path.join(output_path, f'batch_summary_{timestamp}.txt') with open(summary_file, 'w', encoding='utf-8') as f: f.write("AIM医疗问诊工作流批处理执行摘要\n") f.write("=" * 50 + "\n\n") f.write(f"执行时间: {datetime.now().strftime('%Y-%m-%d %H:%M:%S')}\n") f.write(f"总样本数: {summary['total_samples']}\n") f.write(f"处理样本数: {summary['processed_samples']}\n") f.write(f"成功样本数: {summary['successful_samples']}\n") f.write(f"失败样本数: {summary['failed_samples']}\n") f.write(f"跳过样本数: {summary['skipped_samples']}\n") f.write(f"成功率: {summary['success_rate']:.2%}\n") f.write(f"总执行时间: {summary['total_execution_time']:.2f} 秒\n") f.write(f"平均处理时间: {summary['average_time_per_sample']:.2f} 秒/样本\n") f.write(f"处理速度: {summary['samples_per_minute']:.2f} 样本/分钟\n\n") f.write("处理配置:\n") for key, value in summary['processing_config'].items(): f.write(f" {key}: {value}\n") if summary['failed_samples'] > 0: f.write(f"\n失败样本详情:\n") for failed in summary['failed_sample_details']: f.write(f" 样本 {failed['sample_index']}: {failed['error']}\n") logging.info(f"摘要报告已保存: {summary_file}") except Exception as e: logging.error(f"生成报告失败: {e}") def main(): """主入口函数""" # 解析参数 args = parse_arguments() # 处理 --list-models 参数 if args.list_models: print("可用的语言模型配置:") print("=" * 50) for model_name, config in LLM_CONFIG.items(): print(f"模型名称: {model_name}") print(f" 类别: {config['class']}") print(f" 模型ID: {config['params']['id']}") print(f" API端点: {config['params']['base_url']}") print("-" * 30) return 0 # 设置日志 setup_logging(args.log_level) logging.info("=" * 60) logging.info("AIM医疗问诊工作流批处理系统启动") logging.info("=" * 60) try: # 验证参数 if args.num_threads <= 0: raise ValueError("线程数必须大于0") if args.max_steps <= 0: raise ValueError("最大步数必须大于0") # 验证模型类型 if args.model_type not in LLM_CONFIG: available_models = ', '.join(LLM_CONFIG.keys()) raise ValueError(f"不支持的模型类型: {args.model_type}，可用模型: {available_models}") logging.info(f"使用模型: {args.model_type} ({LLM_CONFIG[args.model_type]['class']})") # 试运行模式 if args.dry_run: logging.info("试运行模式：验证配置...") dataset = load_dataset( args.dataset_path, args.start_index, args.end_index, min(args.sample_limit or 5, 5) # 试运行只验证前5个样本 ) logging.info(f"配置验证成功，将处理 {len(dataset)} 个样本") return 0 # 加载数据集 dataset = load_dataset( args.dataset_path, args.start_index, args.end_index, args.sample_limit ) if len(dataset) == 0: logging.warning("没有样本需要处理") return 0 # 执行批处理 logging.info("开始批量处理...") batch_results = run_workflow_batch(dataset, args) # 生成报告 generate_summary_report(batch_results, args.output_dir) # 输出最终统计 summary = batch_results['summary'] logging.info("=" * 60) logging.info("批处理执行完成!") logging.info(f"成功率: {summary['success_rate']:.2%} ({summary['successful_samples']}/{summary['total_samples']})") logging.info(f"总耗时: {summary['total_execution_time']:.2f} 秒") logging.info(f"处理速度: {summary['samples_per_minute']:.2f} 样本/分钟") logging.info("=" * 60) return 0 if summary['success_rate'] > 0.8 else 1 except KeyboardInterrupt: logging.warning("程序被用户中断") return 1 except Exception as e: logging.error(f"程序执行失败: {e}") return 1 if __name__ == "__main__": exit_code = main() sys.exit(exit_code)