PaddleX/zhch/ppstructurev3_multi_gpu_multiprocess_official.py

# zhch/ppstructurev3_multi_gpu_multiprocess_official.py
import json
import time
import os
import glob
import traceback
import argparse
import sys
from pathlib import Path
from typing import List, Dict, Any, Tuple
from multiprocessing import Manager, Process, Queue
from queue import Empty
import cv2
import numpy as np
from paddlex import create_pipeline
from paddlex.utils.device import constr_device, parse_device
from tqdm import tqdm
import paddle
from cuda_utils import detect_available_gpus, monitor_gpu_memory

from dotenv import load_dotenv
load_dotenv(override=True)

def worker(pipeline_name_or_config_path: str, 
          device: str, 
          task_queue: Queue, 
          result_queue: Queue,
          batch_size: int, 
          output_dir: str,
          worker_id: int):
    """
    工作进程函数 - 基于官方parallel_inference.md实现
    
    Args:
        pipeline_name_or_config_path: Pipeline名称或配置路径
        device: 设备字符串
        task_queue: 任务队列
        result_queue: 结果队列
        batch_size: 批处理大小
        output_dir: 输出目录
        worker_id: 工作进程ID
    """
    try:
        # 创建pipeline实例
        from dotenv import load_dotenv
        load_dotenv(override=True)
        print(f"CUDA_VISIBLE_DEVICES: {os.environ.get('CUDA_VISIBLE_DEVICES')}")
        import paddle
        paddle.set_device(device)
        pipeline = create_pipeline(pipeline_name_or_config_path, device=device)
        print(f"Worker {worker_id} initialized with device {device}")
    except Exception as e:
        print(f"Worker {worker_id} ({device}) initialization failed: {e}", file=sys.stderr)
        traceback.print_exc()
        # 发送错误信息到结果队列
        result_queue.put([{
            "error": f"Worker initialization failed: {str(e)}",
            "worker_id": worker_id,
            "device": device,
            "success": False
        }])
        return
        
    try:
        should_end = False
        batch = []
        processed_count = 0
        
        while not should_end:
            try:
                input_path = task_queue.get_nowait()
            except Empty:
                should_end = True
            else:
                batch.append(input_path)
            
            if batch and (len(batch) == batch_size or should_end):
                try:
                    start_time = time.time()
                    
                    # 使用pipeline预测
                    results = pipeline.predict(
                        batch,
                        use_doc_orientation_classify=True,
                        use_doc_unwarping=False,
                        use_seal_recognition=True,
                        use_chart_recognition=True,
                        use_table_recognition=True,
                        use_formula_recognition=True,
                    )
                    
                    batch_processing_time = time.time() - start_time
                    batch_results = []
                    
                    for result in results:
                        try:
                            input_path = Path(result.input_path)
                            
                            # 保存结果
                            if result.get("page_index") is not None:
                                output_filename = f"{input_path.stem}_{result['page_index']}"
                            else:
                                output_filename = f"{input_path.stem}"
                            
                            # 保存JSON和Markdown
                            json_output_path = str(Path(output_dir, f"{output_filename}.json"))
                            md_output_path = str(Path(output_dir, f"{output_filename}.md"))
                            
                            result.save_to_json(json_output_path)
                            result.save_to_markdown(md_output_path)
                            
                            # 记录处理结果
                            batch_results.append({
                                "image_path": input_path.name,
                                "processing_time": batch_processing_time / len(batch),  # 平均时间
                                "success": True,
                                "device": device,
                                "worker_id": worker_id,
                                "output_json": json_output_path,
                                "output_md": md_output_path
                            })
                            
                            processed_count += 1
                            
                        except Exception as e:
                            batch_results.append({
                                "image_path": Path(result.input_path).name if hasattr(result, 'input_path') else "unknown",
                                "processing_time": 0,
                                "success": False,
                                "device": device,
                                "worker_id": worker_id,
                                "error": str(e)
                            })
                    
                    # 将结果放入结果队列
                    result_queue.put(batch_results)
                    
                    print(f"Worker {worker_id} ({device}) processed batch of {len(batch)} files. Total: {processed_count}")
                    
                except Exception as e:
                    # 批处理失败
                    error_results = []
                    for img_path in batch:
                        error_results.append({
                            "image_path": Path(img_path).name,
                            "processing_time": 0,
                            "success": False,
                            "device": device,
                            "worker_id": worker_id,
                            "error": str(e)
                        })
                    result_queue.put(error_results)
                    
                    print(f"Error processing batch {batch} on {device}: {e}", file=sys.stderr)
                
                batch.clear()
    
    except Exception as e:
        print(f"Worker {worker_id} ({device}) initialization failed: {e}", file=sys.stderr)
        traceback.print_exc()
    finally:
        print(f"Worker {worker_id} ({device}) finished")

def parallel_process_with_official_approach(image_paths: List[str],
                                          pipeline_name: str = "PP-StructureV3",
                                          device_str: str = "gpu:0,1",
                                          instances_per_device: int = 1,
                                          batch_size: int = 1,
                                          output_dir: str = "./output") -> List[Dict[str, Any]]:
    """
    使用官方推荐的方法进行多GPU多进程并行处理
    
    Args:
        image_paths: 图像路径列表
        pipeline_name: Pipeline名称
        device_str: 设备字符串，如"gpu:0,1,2,3"
        instances_per_device: 每个设备的实例数
        batch_size: 批处理大小
        output_dir: 输出目录
        
    Returns:
        处理结果列表
    """
    # 创建输出目录
    output_path = Path(output_dir)
    output_path.mkdir(parents=True, exist_ok=True)
    
    # 解析设备
    try:
        device_type, device_ids = parse_device(device_str)
        if device_ids is None or len(device_ids) < 1:
            print("No valid devices specified.", file=sys.stderr)
            return []
        
        print(f"Parsed devices: {device_type}:{device_ids}")
        
    except Exception as e:
        print(f"Failed to parse device string '{device_str}': {e}", file=sys.stderr)
        return []
    
    # 验证批处理大小
    if batch_size <= 0:
        print("Batch size must be greater than 0.", file=sys.stderr)
        return []
    
    total_instances = len(device_ids) * instances_per_device
    print(f"Configuration:")
    print(f"  Devices: {device_ids}")
    print(f"  Instances per device: {instances_per_device}")
    print(f"  Total instances: {total_instances}")
    print(f"  Batch size: {batch_size}")
    print(f"  Total images: {len(image_paths)}")
    
    # 在主进程中初始化paddle，防止子进程CUDA初始化冲突
    try:
        import paddle
        # 只在主进程中设置一个默认设备
        paddle.set_device("cpu")  # 主进程使用CPU
    except Exception as e:
        print(f"Warning: Failed to initialize paddle in main process: {e}")

    # 使用Manager创建队列
    with Manager() as manager:
        task_queue = manager.Queue()
        result_queue = manager.Queue()
        
        # 将任务放入队列
        for img_path in image_paths:
            task_queue.put(str(img_path))
        
        print(f"Added {len(image_paths)} tasks to queue")
        
        # 创建并启动工作进程
        processes = []
        worker_id = 0
        
        for device_id in device_ids:
            for instance_idx in range(instances_per_device):
                device = constr_device(device_type, [device_id])
                
                p = Process(
                    target=worker,
                    args=(
                        pipeline_name,
                        device,
                        task_queue,
                        result_queue,
                        batch_size,
                        str(output_path),
                        worker_id,
                    ),
                    name=f"Worker-{worker_id}-{device}"
                )
                p.start()
                processes.append(p)
                worker_id += 1
        
        print(f"Started {len(processes)} worker processes")
        
        # 收集结果
        all_results = []
        completed_images = 0
        total_images = len(image_paths)
        
        with tqdm(total=total_images, desc="Processing images", unit="img") as pbar:
            # 等待所有结果
            active_workers = len(processes)
            
            while completed_images < total_images and active_workers > 0:
                try:
                    # 设置较短的超时时间，定期检查进程状态
                    batch_results = result_queue.get(timeout=5.0)
                    
                    all_results.extend(batch_results)
                    batch_size_actual = len(batch_results)
                    completed_images += batch_size_actual
                    
                    pbar.update(batch_size_actual)
                    
                    # 更新进度条信息
                    success_count = sum(1 for r in batch_results if r.get('success', False))
                    total_success = sum(1 for r in all_results if r.get('success', False))
                    
                    # 按设备统计
                    device_stats = {}
                    for r in all_results:
                        device = r.get('device', 'unknown')
                        if device not in device_stats:
                            device_stats[device] = {'success': 0, 'total': 0}
                        device_stats[device]['total'] += 1
                        if r.get('success', False):
                            device_stats[device]['success'] += 1
                    
                    device_info = ', '.join([f"{k}:{v['success']}/{v['total']}" 
                                           for k, v in device_stats.items()])
                    
                    pbar.set_postfix({
                        'batch_success': f"{success_count}/{batch_size_actual}",
                        'total_success': f"{total_success}/{completed_images}",
                        'devices': device_info
                    })
                    
                except Exception as e:
                    # 检查是否还有活跃的进程
                    active_workers = sum(1 for p in processes if p.is_alive())
                    if active_workers == 0:
                        print("All workers have finished")
                        break
                    
                    # 超时或其他错误，继续等待
                    continue
        
        # 等待所有进程结束
        print("Waiting for all processes to finish...")
        for p in processes:
            p.join(timeout=10.0)
            if p.is_alive():
                print(f"Force terminating process: {p.name}")
                p.terminate()
                p.join(timeout=5.0)
    
    return all_results

def main():
    """主函数"""
    parser = argparse.ArgumentParser(description="PaddleX PP-StructureV3 Multi-GPU Parallel Processing")
    
    # 必需参数
    parser.add_argument("--input_dir", type=str, required=True, 
                       help="Input directory containing images")
    parser.add_argument("--output_dir", type=str, default="./output", 
                       help="Output directory")
    
    # Pipeline配置
    parser.add_argument("--pipeline", type=str, default="PP-StructureV3",
                       help="Pipeline name or config path")
    parser.add_argument("--device", type=str, default="gpu:0,1",
                       help="Devices for parallel inference (e.g., 'gpu:0,1,2,3')")
    
    # 并行配置
    parser.add_argument("--instances_per_device", type=int, default=1,
                       help="Number of pipeline instances per device")
    parser.add_argument("--batch_size", type=int, default=1,
                       help="Inference batch size for each pipeline instance")
    
    # 输入文件配置
    parser.add_argument("--input_glob_pattern", type=str, default="*",
                       help="Pattern to find input files")
    
    # 测试模式
    parser.add_argument("--test_mode", action="store_true",
                       help="Test mode: only process first 20 images")
    
    args = parser.parse_args()
    
    # 验证输入目录
    input_dir = Path(args.input_dir)
    if not input_dir.exists():
        print(f"Input directory does not exist: {input_dir}", file=sys.stderr)
        return 2
    if not input_dir.is_dir():
        print(f"{input_dir} is not a directory", file=sys.stderr)
        return 2
    
    # 验证输出目录
    output_dir = Path(args.output_dir)
    if output_dir.exists() and not output_dir.is_dir():
        print(f"{output_dir} is not a directory", file=sys.stderr)
        return 2
    
    print("="*70)
    print("PaddleX PP-StructureV3 Multi-GPU Parallel Processing")
    print("="*70)
    print(f"Input directory: {input_dir}")
    print(f"Output directory: {output_dir}")
    print(f"Pipeline: {args.pipeline}")
    print(f"Device: {args.device}")
    print(f"Instances per device: {args.instances_per_device}")
    print(f"Batch size: {args.batch_size}")
    print(f"Input pattern: {args.input_glob_pattern}")
    
    # 查找图像文件
    image_extensions = ['*.jpg', '*.jpeg', '*.png', '*.bmp', '*.tiff', '*.pdf']
    image_files = []
    
    for ext in image_extensions:
        pattern = args.input_glob_pattern if args.input_glob_pattern != "*" else ext
        image_files.extend(input_dir.glob(pattern))
    
    # 如果没有找到文件，尝试使用用户指定的模式
    if not image_files and args.input_glob_pattern != "*":
        image_files = list(input_dir.glob(args.input_glob_pattern))
    
    if not image_files:
        print(f"No image files found in {input_dir} with pattern {args.input_glob_pattern}")
        return 1
    
    # 转换为字符串路径
    image_paths = [str(f) for f in image_files]
    
    print(f"Found {len(image_paths)} image files")
    
    # 测试模式
    if args.test_mode:
        image_paths = image_paths[:20]
        print(f"Test mode: processing only {len(image_paths)} images")
    
    # 开始处理
    start_time = time.time()
    
    try:
        results = parallel_process_with_official_approach(
            image_paths=image_paths,
            pipeline_name=args.pipeline,
            device_str=args.device,
            instances_per_device=args.instances_per_device,
            batch_size=args.batch_size,
            output_dir=args.output_dir
        )
        
        total_time = time.time() - start_time
        
        # 统计信息
        success_count = sum(1 for r in results if r.get('success', False))
        error_count = len(results) - success_count
        total_processing_time = sum(r.get('processing_time', 0) for r in results if r.get('success', False))
        avg_processing_time = total_processing_time / success_count if success_count > 0 else 0
        
        # 按设备统计
        device_stats = {}
        worker_stats = {}
        
        for r in results:
            device = r.get('device', 'unknown')
            worker_id = r.get('worker_id', 'unknown')
            
            # 设备统计
            if device not in device_stats:
                device_stats[device] = {'success': 0, 'total': 0, 'total_time': 0}
            device_stats[device]['total'] += 1
            if r.get('success', False):
                device_stats[device]['success'] += 1
                device_stats[device]['total_time'] += r.get('processing_time', 0)
            
            # Worker统计
            if worker_id not in worker_stats:
                worker_stats[worker_id] = {'success': 0, 'total': 0, 'device': device}
            worker_stats[worker_id]['total'] += 1
            if r.get('success', False):
                worker_stats[worker_id]['success'] += 1
        
        # 保存详细结果
        detailed_results = {
            "configuration": {
                "pipeline": args.pipeline,
                "device": args.device,
                "instances_per_device": args.instances_per_device,
                "batch_size": args.batch_size,
                "input_glob_pattern": args.input_glob_pattern,
                "test_mode": args.test_mode
            },
            "statistics": {
                "total_files": len(image_paths),
                "success_count": success_count,
                "error_count": error_count,
                "success_rate": success_count / len(image_paths) if image_paths else 0,
                "total_time": total_time,
                "avg_processing_time": avg_processing_time,
                "throughput": len(image_paths) / total_time if total_time > 0 else 0,
                "device_stats": device_stats,
                "worker_stats": worker_stats
            },
            "results": results
        }
        
        # 保存结果文件
        result_file = output_dir / "processing_results.json"
        with open(result_file, 'w', encoding='utf-8') as f:
            json.dump(detailed_results, f, ensure_ascii=False, indent=2)
        
        # 打印统计信息
        print("\n" + "="*70)
        print("Processing completed!")
        print("="*70)
        print(f"Total files: {len(image_paths)}")
        print(f"Successfully processed: {success_count}")
        print(f"Failed: {error_count}")
        print(f"Success rate: {success_count / len(image_paths) * 100:.2f}%")
        print(f"Total time: {total_time:.2f} seconds")
        print(f"Average processing time: {avg_processing_time:.2f} seconds/image")
        print(f"Throughput: {len(image_paths) / total_time:.2f} images/second")
        
        # 设备统计
        print(f"\nDevice Statistics:")
        for device, stats in device_stats.items():
            if stats['total'] > 0:
                success_rate = stats['success'] / stats['total'] * 100
                avg_time = stats['total_time'] / stats['success'] if stats['success'] > 0 else 0
                print(f"  {device}: {stats['success']}/{stats['total']} "
                      f"({success_rate:.1f}%), avg {avg_time:.2f}s/image")
        
        # Worker统计
        print(f"\nWorker Statistics:")
        for worker_id, stats in worker_stats.items():
            if stats['total'] > 0:
                success_rate = stats['success'] / stats['total'] * 100
                print(f"  Worker {worker_id} ({stats['device']}): {stats['success']}/{stats['total']} "
                      f"({success_rate:.1f}%)")
        
        print(f"\nDetailed results saved to: {result_file}")
        print("All done!")
        
        return 0
        
    except Exception as e:
        print(f"Processing failed: {e}", file=sys.stderr)
        traceback.print_exc()
        return 1

if __name__ == "__main__":
    print(f"🚀 启动OCR程序...")
    print(f"CUDA 版本: {paddle.device.cuda.get_device_name()}")
    print(f"CUDA_VISIBLE_DEVICES: {os.environ.get('CUDA_VISIBLE_DEVICES')}")
    available_gpus = detect_available_gpus()
    monitor_gpu_memory(available_gpus)

    if len(sys.argv) == 1:
        # 如果没有命令行参数，使用默认配置运行
        print("No command line arguments provided. Running with default configuration...")
        
        # 默认配置
        default_config = {
            "input_dir": "../../OmniDocBench/OpenDataLab___OmniDocBench/images",
            "output_dir": "./OmniDocBench_Results_Official",
            "pipeline": "PP-StructureV3",
            "device": "gpu:0",
            "instances_per_device": 1,
            "batch_size": 4,
            # "test_mode": False
        }
        
        # 构造参数
        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())