|
|
@@ -0,0 +1,380 @@
|
|
|
+import json
|
|
|
+import time
|
|
|
+import os
|
|
|
+import argparse
|
|
|
+import sys
|
|
|
+import subprocess
|
|
|
+import tempfile
|
|
|
+from pathlib import Path
|
|
|
+from typing import List, Dict, Any, Tuple
|
|
|
+from concurrent.futures import ProcessPoolExecutor, as_completed
|
|
|
+import threading
|
|
|
+from queue import Queue
|
|
|
+from tqdm import tqdm
|
|
|
+
|
|
|
+def split_files(file_list: List[str], num_splits: int) -> List[List[str]]:
|
|
|
+ """
|
|
|
+ 将文件列表分割成指定数量的子列表
|
|
|
+
|
|
|
+ Args:
|
|
|
+ file_list: 文件路径列表
|
|
|
+ num_splits: 分割数量
|
|
|
+
|
|
|
+ Returns:
|
|
|
+ 分割后的文件列表
|
|
|
+ """
|
|
|
+ if num_splits <= 0:
|
|
|
+ return [file_list]
|
|
|
+
|
|
|
+ chunk_size = len(file_list) // num_splits
|
|
|
+ remainder = len(file_list) % num_splits
|
|
|
+
|
|
|
+ chunks = []
|
|
|
+ start = 0
|
|
|
+
|
|
|
+ for i in range(num_splits):
|
|
|
+ # 前remainder个chunk多分配一个文件
|
|
|
+ current_chunk_size = chunk_size + (1 if i < remainder else 0)
|
|
|
+ if current_chunk_size > 0:
|
|
|
+ chunks.append(file_list[start:start + current_chunk_size])
|
|
|
+ start += current_chunk_size
|
|
|
+
|
|
|
+ return [chunk for chunk in chunks if chunk] # 过滤空列表
|
|
|
+
|
|
|
+def create_temp_file_list(file_chunk: List[str]) -> str:
|
|
|
+ """
|
|
|
+ 创建临时文件列表文件
|
|
|
+
|
|
|
+ Args:
|
|
|
+ file_chunk: 文件路径列表
|
|
|
+
|
|
|
+ Returns:
|
|
|
+ 临时文件路径
|
|
|
+ """
|
|
|
+ with tempfile.NamedTemporaryFile(mode='w', suffix='.txt', delete=False) as f:
|
|
|
+ for file_path in file_chunk:
|
|
|
+ f.write(f"{file_path}\n")
|
|
|
+ return f.name
|
|
|
+
|
|
|
+def run_single_process(args: Tuple[List[str], Dict[str, Any], int]) -> Dict[str, Any]:
|
|
|
+ """
|
|
|
+ 运行单个ppstructurev3_single_process.py进程
|
|
|
+
|
|
|
+ Args:
|
|
|
+ args: (file_chunk, config, process_id)
|
|
|
+
|
|
|
+ Returns:
|
|
|
+ 处理结果
|
|
|
+ """
|
|
|
+ file_chunk, config, process_id = args
|
|
|
+
|
|
|
+ if not file_chunk:
|
|
|
+ return {"process_id": process_id, "success": False, "error": "Empty file chunk"}
|
|
|
+
|
|
|
+ # 创建临时文件列表
|
|
|
+ temp_file_list = create_temp_file_list(file_chunk)
|
|
|
+
|
|
|
+ try:
|
|
|
+ # 创建进程专用的输出目录
|
|
|
+ process_output_dir = Path(config["output_dir"]) / f"process_{process_id}"
|
|
|
+ process_output_dir.mkdir(parents=True, exist_ok=True)
|
|
|
+
|
|
|
+ # 构建命令行参数
|
|
|
+ cmd = [
|
|
|
+ sys.executable,
|
|
|
+ config["single_process_script"],
|
|
|
+ "--input_file_list", temp_file_list, # 需要修改single_process脚本支持文件列表
|
|
|
+ "--output_dir", str(process_output_dir),
|
|
|
+ "--pipeline", config["pipeline"],
|
|
|
+ "--device", config["device"],
|
|
|
+ "--batch_size", str(config["batch_size"]),
|
|
|
+ ]
|
|
|
+
|
|
|
+ # 添加可选参数
|
|
|
+ if config.get("test_mode", False):
|
|
|
+ cmd.append("--test_mode")
|
|
|
+
|
|
|
+ print(f"Process {process_id} starting with {len(file_chunk)} files on device {config['device']}")
|
|
|
+
|
|
|
+ # 执行子进程
|
|
|
+ start_time = time.time()
|
|
|
+ result = subprocess.run(
|
|
|
+ cmd,
|
|
|
+ capture_output=True,
|
|
|
+ text=True,
|
|
|
+ timeout=config.get("timeout", 3600) # 1小时超时
|
|
|
+ )
|
|
|
+
|
|
|
+ processing_time = time.time() - start_time
|
|
|
+
|
|
|
+ if result.returncode == 0:
|
|
|
+ print(f"Process {process_id} completed successfully in {processing_time:.2f}s")
|
|
|
+
|
|
|
+ # 读取结果文件
|
|
|
+ result_files = list(process_output_dir.glob("*.json"))
|
|
|
+
|
|
|
+ return {
|
|
|
+ "process_id": process_id,
|
|
|
+ "success": True,
|
|
|
+ "processing_time": processing_time,
|
|
|
+ "file_count": len(file_chunk),
|
|
|
+ "device": config["device"],
|
|
|
+ "output_dir": str(process_output_dir),
|
|
|
+ "result_files": [str(f) for f in result_files],
|
|
|
+ "stdout": result.stdout,
|
|
|
+ "stderr": result.stderr
|
|
|
+ }
|
|
|
+ else:
|
|
|
+ print(f"Process {process_id} failed with return code {result.returncode}")
|
|
|
+ return {
|
|
|
+ "process_id": process_id,
|
|
|
+ "success": False,
|
|
|
+ "error": f"Process failed with return code {result.returncode}",
|
|
|
+ "stdout": result.stdout,
|
|
|
+ "stderr": result.stderr
|
|
|
+ }
|
|
|
+
|
|
|
+ except subprocess.TimeoutExpired:
|
|
|
+ print(f"Process {process_id} timed out")
|
|
|
+ return {
|
|
|
+ "process_id": process_id,
|
|
|
+ "success": False,
|
|
|
+ "error": "Process timeout"
|
|
|
+ }
|
|
|
+ except Exception as e:
|
|
|
+ print(f"Process {process_id} error: {e}")
|
|
|
+ return {
|
|
|
+ "process_id": process_id,
|
|
|
+ "success": False,
|
|
|
+ "error": str(e)
|
|
|
+ }
|
|
|
+ finally:
|
|
|
+ # 清理临时文件
|
|
|
+ try:
|
|
|
+ os.unlink(temp_file_list)
|
|
|
+ except:
|
|
|
+ pass
|
|
|
+
|
|
|
+def monitor_progress(total_files: int, completed_queue: Queue):
|
|
|
+ """
|
|
|
+ 监控处理进度
|
|
|
+ """
|
|
|
+ with tqdm(total=total_files, desc="Total Progress", unit="files") as pbar:
|
|
|
+ completed_count = 0
|
|
|
+ while completed_count < total_files:
|
|
|
+ try:
|
|
|
+ batch_count = completed_queue.get(timeout=1)
|
|
|
+ completed_count += batch_count
|
|
|
+ pbar.update(batch_count)
|
|
|
+ except:
|
|
|
+ continue
|
|
|
+
|
|
|
+def main():
|
|
|
+ """主函数"""
|
|
|
+ parser = argparse.ArgumentParser(description="PaddleX PP-StructureV3 Multi-Process Scheduler")
|
|
|
+
|
|
|
+ # 输入输出参数
|
|
|
+ parser.add_argument("--input_dir", type=str, required=True, help="Input directory")
|
|
|
+ parser.add_argument("--output_dir", type=str, required=True, help="Output directory")
|
|
|
+ parser.add_argument("--single_process_script", type=str,
|
|
|
+ default="./ppstructurev3_single_process.py",
|
|
|
+ help="Path to single process script")
|
|
|
+
|
|
|
+ # 并行参数
|
|
|
+ parser.add_argument("--num_processes", type=int, default=4, help="Number of parallel processes")
|
|
|
+ parser.add_argument("--devices", type=str, default="gpu:0,gpu:1,gpu:2,gpu:3",
|
|
|
+ help="Device list (comma separated)")
|
|
|
+
|
|
|
+ # Pipeline参数
|
|
|
+ parser.add_argument("--pipeline", type=str, default="PP-StructureV3", help="Pipeline name")
|
|
|
+ parser.add_argument("--batch_size", type=int, default=4, help="Batch size per process")
|
|
|
+ parser.add_argument("--timeout", type=int, default=3600, help="Process timeout in seconds")
|
|
|
+
|
|
|
+ # 其他参数
|
|
|
+ parser.add_argument("--test_mode", action="store_true", help="Test mode")
|
|
|
+ parser.add_argument("--max_files", type=int, default=None, help="Maximum files to process")
|
|
|
+
|
|
|
+ args = parser.parse_args()
|
|
|
+
|
|
|
+ try:
|
|
|
+ # 获取图像文件列表
|
|
|
+ input_dir = Path(args.input_dir).resolve()
|
|
|
+ output_dir = Path(args.output_dir).resolve()
|
|
|
+
|
|
|
+ print(f"Input dir: {input_dir}")
|
|
|
+ print(f"Output dir: {output_dir}")
|
|
|
+
|
|
|
+ if not input_dir.exists():
|
|
|
+ print(f"Input directory does not exist: {input_dir}")
|
|
|
+ return 1
|
|
|
+
|
|
|
+ # 查找图像文件
|
|
|
+ image_extensions = ['.jpg', '.jpeg', '.png', '.bmp', '.tiff', '.tif']
|
|
|
+ image_files = []
|
|
|
+ for ext in image_extensions:
|
|
|
+ image_files.extend(list(input_dir.glob(f"*{ext}")))
|
|
|
+ image_files.extend(list(input_dir.glob(f"*{ext.upper()}")))
|
|
|
+
|
|
|
+ if not image_files:
|
|
|
+ print(f"No image files found in {input_dir}")
|
|
|
+ return 1
|
|
|
+
|
|
|
+ # 去重并排序
|
|
|
+ image_files = sorted(list(set(str(f) for f in image_files)))
|
|
|
+
|
|
|
+ # 限制文件数量
|
|
|
+ if args.max_files:
|
|
|
+ image_files = image_files[:args.max_files]
|
|
|
+
|
|
|
+ if args.test_mode:
|
|
|
+ image_files = image_files[:20]
|
|
|
+ print(f"Test mode: processing only {len(image_files)} images")
|
|
|
+
|
|
|
+ print(f"Found {len(image_files)} image files")
|
|
|
+
|
|
|
+ # 解析设备列表
|
|
|
+ devices = [d.strip() for d in args.devices.split(',')]
|
|
|
+ if len(devices) < args.num_processes:
|
|
|
+ # 如果设备数少于进程数，循环使用设备
|
|
|
+ devices = devices * ((args.num_processes // len(devices)) + 1)
|
|
|
+ devices = devices[:args.num_processes]
|
|
|
+
|
|
|
+ print(f"Using {args.num_processes} processes with devices: {devices}")
|
|
|
+
|
|
|
+ # 分割文件列表
|
|
|
+ file_chunks = split_files(image_files, args.num_processes)
|
|
|
+ print(f"Split into {len(file_chunks)} chunks: {[len(chunk) for chunk in file_chunks]}")
|
|
|
+
|
|
|
+ # 创建输出目录
|
|
|
+ output_dir.mkdir(parents=True, exist_ok=True)
|
|
|
+
|
|
|
+ # 准备进程参数
|
|
|
+ process_configs = []
|
|
|
+ for i, (chunk, device) in enumerate(zip(file_chunks, devices)):
|
|
|
+ config = {
|
|
|
+ "single_process_script": str(Path(args.single_process_script).resolve()),
|
|
|
+ "output_dir": str(output_dir),
|
|
|
+ "pipeline": args.pipeline,
|
|
|
+ "device": device,
|
|
|
+ "batch_size": args.batch_size,
|
|
|
+ "timeout": args.timeout,
|
|
|
+ "test_mode": args.test_mode
|
|
|
+ }
|
|
|
+ process_configs.append((chunk, config, i))
|
|
|
+
|
|
|
+ # 启动进度监控
|
|
|
+ completed_queue = Queue()
|
|
|
+ progress_thread = threading.Thread(
|
|
|
+ target=monitor_progress,
|
|
|
+ args=(len(image_files), completed_queue)
|
|
|
+ )
|
|
|
+ progress_thread.daemon = True
|
|
|
+ progress_thread.start()
|
|
|
+
|
|
|
+ # 执行并行处理
|
|
|
+ start_time = time.time()
|
|
|
+ results = []
|
|
|
+
|
|
|
+ with ProcessPoolExecutor(max_workers=args.num_processes) as executor:
|
|
|
+ # 提交所有任务
|
|
|
+ future_to_process = {
|
|
|
+ executor.submit(run_single_process, config): i
|
|
|
+ for i, config in enumerate(process_configs)
|
|
|
+ }
|
|
|
+
|
|
|
+ # 收集结果
|
|
|
+ for future in as_completed(future_to_process):
|
|
|
+ process_id = future_to_process[future]
|
|
|
+ try:
|
|
|
+ result = future.result()
|
|
|
+ results.append(result)
|
|
|
+
|
|
|
+ # 更新进度
|
|
|
+ if result.get("success", False):
|
|
|
+ completed_queue.put(result.get("file_count", 0))
|
|
|
+
|
|
|
+ print(f"Process {process_id} finished: {result.get('success', False)}")
|
|
|
+
|
|
|
+ except Exception as e:
|
|
|
+ print(f"Process {process_id} generated an exception: {e}")
|
|
|
+ results.append({
|
|
|
+ "process_id": process_id,
|
|
|
+ "success": False,
|
|
|
+ "error": str(e)
|
|
|
+ })
|
|
|
+
|
|
|
+ total_time = time.time() - start_time
|
|
|
+
|
|
|
+ # 统计结果
|
|
|
+ successful_processes = sum(1 for r in results if r.get('success', False))
|
|
|
+ total_processed_files = sum(r.get('file_count', 0) for r in results if r.get('success', False))
|
|
|
+
|
|
|
+ print(f"\n" + "="*60)
|
|
|
+ print(f"🎉 Parallel processing completed!")
|
|
|
+ print(f"📊 Statistics:")
|
|
|
+ print(f" Total processes: {len(results)}")
|
|
|
+ print(f" Successful processes: {successful_processes}")
|
|
|
+ print(f" Total files processed: {total_processed_files}/{len(image_files)}")
|
|
|
+ print(f" Success rate: {total_processed_files/len(image_files)*100:.2f}%")
|
|
|
+ print(f"⏱️ Performance:")
|
|
|
+ print(f" Total time: {total_time:.2f} seconds")
|
|
|
+ print(f" Throughput: {total_processed_files/total_time:.2f} files/second")
|
|
|
+ print(f" Avg time per file: {total_time/total_processed_files:.2f} seconds")
|
|
|
+
|
|
|
+ # 保存调度结果
|
|
|
+ scheduler_stats = {
|
|
|
+ "total_files": len(image_files),
|
|
|
+ "total_processes": len(results),
|
|
|
+ "successful_processes": successful_processes,
|
|
|
+ "total_processed_files": total_processed_files,
|
|
|
+ "success_rate": total_processed_files / len(image_files) if len(image_files) > 0 else 0,
|
|
|
+ "total_time": total_time,
|
|
|
+ "throughput": total_processed_files / total_time if total_time > 0 else 0,
|
|
|
+ "avg_time_per_file": total_time / total_processed_files if total_processed_files > 0 else 0,
|
|
|
+ "num_processes": args.num_processes,
|
|
|
+ "devices": devices,
|
|
|
+ "batch_size": args.batch_size,
|
|
|
+ "pipeline": args.pipeline,
|
|
|
+ "timestamp": time.strftime("%Y-%m-%d %H:%M:%S")
|
|
|
+ }
|
|
|
+
|
|
|
+ final_results = {
|
|
|
+ "scheduler_stats": scheduler_stats,
|
|
|
+ "process_results": results
|
|
|
+ }
|
|
|
+
|
|
|
+ # 保存结果
|
|
|
+ output_file = output_dir / f"scheduler_results_{args.num_processes}procs.json"
|
|
|
+ with open(output_file, 'w', encoding='utf-8') as f:
|
|
|
+ json.dump(final_results, f, ensure_ascii=False, indent=2)
|
|
|
+
|
|
|
+ print(f"💾 Scheduler results saved to: {output_file}")
|
|
|
+
|
|
|
+ return 0 if successful_processes == len(results) else 1
|
|
|
+
|
|
|
+ except Exception as e:
|
|
|
+ print(f"❌ Scheduler failed: {e}")
|
|
|
+ import traceback
|
|
|
+ traceback.print_exc()
|
|
|
+ return 1
|
|
|
+
|
|
|
+if __name__ == "__main__":
|
|
|
+ print(f"🚀 启动多进程调度程序...")
|
|
|
+
|
|
|
+ if len(sys.argv) == 1:
|
|
|
+ # 默认配置
|
|
|
+ default_config = {
|
|
|
+ "input_dir": "../../OmniDocBench/OpenDataLab___OmniDocBench/images",
|
|
|
+ "output_dir": "./OmniDocBench_Results_Scheduler",
|
|
|
+ "num_processes": 4,
|
|
|
+ "devices": "gpu:0,gpu:1,gpu:2,gpu:3",
|
|
|
+ "batch_size": 2,
|
|
|
+ }
|
|
|
+
|
|
|
+ sys.argv = [sys.argv[0]]
|
|
|
+ for key, value in default_config.items():
|
|
|
+ sys.argv.extend([f"--{key}", str(value)])
|
|
|
+
|
|
|
+ sys.argv.append("--test_mode")
|
|
|
+
|
|
|
+ sys.exit(main())
|
