ocr-training/scripts/benchmark_worker.py

#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""
Isolated Benchmark Worker for PaddleOCR Training.
Executes a single (batch_size, num_workers) trial in an isolated subprocess.
Exits cleanly on success (code 0), OOM (code 137), NaN loss (code 138), or error (code 1).
"""

import os
import sys
import time
import json
import argparse
import traceback
import threading

# 1. NumPy 2.x compatibility monkeypatch
import numpy as np
if not hasattr(np, 'sctypes'):
    np.sctypes = {
        'int': [np.int8, np.int16, np.int32, np.int64],
        'uint': [np.uint8, np.uint16, np.uint32, np.uint64],
        'float': [np.float16, np.float32, np.float64],
        'complex': [np.complex64, np.complex128],
        'others': [bool, object, bytes, str]
    }
if not hasattr(np, 'bool'): np.bool = bool
if not hasattr(np, 'int'): np.int = int
if not hasattr(np, 'float'): np.float = float
if not hasattr(np, 'typeDict'): np.typeDict = {}

# 2. Add PaddleOCR path
current_dir = os.path.dirname(os.path.abspath(__file__))
# Check possible PaddleOCR locations
paddleocr_dirs = [
    os.path.join(current_dir, "..", "PaddleOCR"),
    os.path.join(current_dir, "..", "..", "ocr-studio", "PaddleOCR"),
    os.path.join(current_dir, "..", "..", "PaddleOCR"),
    "/teamspace/studios/this_studio/Cham-OCR/ocr-training/PaddleOCR",
]
for p in paddleocr_dirs:
    if os.path.isdir(p):
        sys.path.insert(0, os.path.abspath(p))
        break

import yaml
import paddle
import paddle.distributed as dist

from ppocr.data import build_dataloader, set_signal_handlers
from ppocr.modeling.architectures import build_model, apply_to_static
from ppocr.losses import build_loss
from ppocr.optimizer import build_optimizer
from ppocr.postprocess import build_post_process
from ppocr.utils.logging import get_logger

try:
    import psutil
except ImportError:
    psutil = None

def to_float32(preds):
    if isinstance(preds, dict):
        for k in preds:
            if isinstance(preds[k], dict):
                for sub_k in preds[k]:
                    preds[k][sub_k] = preds[k][sub_k].astype(paddle.float32)
            else:
                preds[k] = preds[k].astype(paddle.float32)
    elif isinstance(preds, list):
        preds = [pred.astype(paddle.float32) for pred in preds]
    else:
        preds = preds.astype(paddle.float32)
    return preds


class GPUMonitor(threading.Thread):
    """Monitors GPU utilization and VRAM in a background thread."""
    def __init__(self, interval=0.3):
        super().__init__(daemon=True)
        self.interval = interval
        self.running = True
        self.gpu_utils = []
        self.vram_used_mb = []
        self.total_vram_mb = 0.0
        self.cpu_percents = []

    def run(self):
        while self.running:
            try:
                # Query nvidia-smi
                import subprocess
                res = subprocess.run(
                    ["nvidia-smi", "--query-gpu=utilization.gpu,memory.used,memory.total", "--format=csv,noheader,nounits"],
                    capture_output=True, text=True, timeout=1.5
                )
                if res.returncode == 0 and res.stdout.strip():
                    parts = [float(x.strip()) for x in res.stdout.strip().split("\n")[0].split(",")]
                    self.gpu_utils.append(parts[0])
                    self.vram_used_mb.append(parts[1])
                    self.total_vram_mb = parts[2]
            except Exception:
                pass

            if psutil:
                try:
                    self.cpu_percents.append(psutil.cpu_percent(interval=None))
                except Exception:
                    pass

            time.sleep(self.interval)

    def stop(self):
        self.running = False


def run_benchmark(config_path, batch_size, num_workers, warmup_steps=100, measure_steps=300, output_json="trial_result.json"):
    result = {
        "batch_size": batch_size,
        "num_workers": num_workers,
        "warmup_steps": warmup_steps,
        "measure_steps": measure_steps,
        "status": "UNKNOWN",
        "ips": 0.0,
        "avg_reader_cost": 0.0,
        "avg_batch_cost": 0.0,
        "reader_ratio": 0.0,
        "peak_vram_mb": 0.0,
        "total_vram_mb": 0.0,
        "vram_pct": 0.0,
        "avg_gpu_util": 0.0,
        "avg_cpu_pct": 0.0,
        "error": None
    }

    monitor = GPUMonitor(interval=0.3)
    monitor.start()

    try:
        with open(config_path, 'r', encoding='utf-8') as f:
            config = yaml.safe_load(f)

        # Override configurations
        config['Global']['distributed'] = False
        config['Global']['use_gpu'] = True
        config['Global']['cal_metric_during_train'] = False
        config['Global']['epoch_num'] = 100
        config['Train']['loader']['batch_size_per_card'] = batch_size
        config['Train']['loader']['num_workers'] = num_workers
        config['Train']['loader']['drop_last'] = True
        config['Eval'] = None

        device = 'gpu'
        paddle.set_device('gpu:0')
        logger = get_logger(name="benchmark_worker")

        # Set signal handlers & build DataLoader
        set_signal_handlers()
        train_dataloader = build_dataloader(config, 'Train', device, logger, seed=42)
        if len(train_dataloader) == 0:
            raise RuntimeError("DataLoader is empty! Check dataset paths in config.")

        global_config = config['Global']
        post_process_class = build_post_process(config['PostProcess'], global_config)

        # MultiHead channel adjustments
        if hasattr(post_process_class, 'character'):
            char_num = len(getattr(post_process_class, 'character'))
            if config['Architecture']['Head']['name'] == 'MultiHead':
                if config['PostProcess']['name'] == 'NRTRLabelDecode':
                    char_num = char_num - 3
                out_channels_list = {'CTCLabelDecode': char_num}
                if len(config['Loss']['loss_config_list']) > 1 and list(config['Loss']['loss_config_list'][1].keys())[0] == 'NRTRLoss':
                    out_channels_list['NRTRLabelDecode'] = char_num + 3
                config['Architecture']['Head']['out_channels_list'] = out_channels_list
            else:
                config['Architecture']['Head']['out_channels'] = char_num

        model = build_model(config['Architecture'])
        model = apply_to_static(model, config, logger)
        model.train()
        loss_class = build_loss(config['Loss'])
        optimizer, lr_scheduler = build_optimizer(
            config['Optimizer'],
            epochs=config['Global']['epoch_num'],
            step_each_epoch=len(train_dataloader),
            model=model
        )

        use_amp = config['Global'].get('use_amp', False)
        amp_level = config['Global'].get('amp_level', 'O1')
        amp_dtype = config['Global'].get('amp_dtype', 'float16')
        scaler = paddle.amp.GradScaler() if use_amp else None

        extra_input_models = [
            "SRN", "NRTR", "SAR", "SEED", "SVTR", "SVTR_LCNet", "SPIN", "VisionLAN",
            "RobustScanner", "RFL", 'DRRG', 'SATRN', 'SVTR_HGNet', "ParseQ", "CPPD"
        ]
        extra_input = config['Architecture'].get('algorithm', '') in extra_input_models or config['Architecture']['Head'].get('name', '') == 'MultiHead'

        total_steps = warmup_steps + measure_steps
        step_idx = 0

        total_reader_time = 0.0
        total_batch_time = 0.0
        total_samples = 0

        data_iter = iter(train_dataloader)
        print(f"[Worker] Starting trial: Batch={batch_size}, Workers={num_workers}, Warmup={warmup_steps}, Measure={measure_steps}, ExtraInput={extra_input}", flush=True)

        reader_start = time.perf_counter()
        while step_idx < total_steps:
            try:
                batch = next(data_iter)
            except StopIteration:
                data_iter = iter(train_dataloader)
                batch = next(data_iter)

            t_read = time.perf_counter() - reader_start
            images = batch[0]
            batch_samples = len(images)

            # Forward + Backward
            if scaler:
                with paddle.amp.auto_cast(level=amp_level, dtype=amp_dtype):
                    if extra_input:
                        preds = model(images, data=batch[1:])
                    else:
                        preds = model(images)
                preds = to_float32(preds)
                loss = loss_class(preds, batch)
                avg_loss = loss['loss']
                loss_val = float(avg_loss)
                if np.isnan(loss_val) or np.isinf(loss_val):
                    raise ValueError(f"Loss exploded: {loss_val}")
                scaled_avg_loss = scaler.scale(avg_loss)
                scaled_avg_loss.backward()
                scaler.minimize(optimizer, scaled_avg_loss)
            else:
                if extra_input:
                    preds = model(images, data=batch[1:])
                else:
                    preds = model(images)
                loss = loss_class(preds, batch)
                avg_loss = loss['loss']
                loss_val = float(avg_loss)
                if np.isnan(loss_val) or np.isinf(loss_val):
                    raise ValueError(f"Loss exploded: {loss_val}")
                avg_loss.backward()
                optimizer.step()

            optimizer.clear_grad()
            if not isinstance(lr_scheduler, float):
                lr_scheduler.step()

            # Synchronize CUDA to get exact batch execution time
            try:
                paddle.device.cuda.synchronize()
            except Exception:
                pass
            t_batch = time.perf_counter() - reader_start

            if step_idx >= warmup_steps:
                total_reader_time += t_read
                total_batch_time += t_batch
                total_samples += batch_samples

            step_idx += 1
            if step_idx % 50 == 0:
                cur_ips = batch_samples / t_batch if t_batch > 0 else 0
                phase = "WARMUP" if step_idx <= warmup_steps else "MEASURE"
                print(f"[Worker] Step {step_idx}/{total_steps} [{phase}] - IPS: {cur_ips:.1f} - BatchTime: {t_batch:.4f}s - Reader: {t_read:.4f}s", flush=True)

            reader_start = time.perf_counter()

        # Stop monitoring
        monitor.stop()

        avg_reader_cost = total_reader_time / measure_steps if measure_steps > 0 else 0
        avg_batch_cost = total_batch_time / measure_steps if measure_steps > 0 else 0
        ips = total_samples / total_batch_time if total_batch_time > 0 else 0

        # Memory stats
        paddle_peak_bytes = paddle.device.cuda.max_memory_allocated()
        paddle_peak_mb = paddle_peak_bytes / (1024 * 1024)
        smi_peak_mb = max(monitor.vram_used_mb) if monitor.vram_used_mb else paddle_peak_mb
        peak_vram_mb = max(paddle_peak_mb, smi_peak_mb)
        total_vram_mb = monitor.total_vram_mb if monitor.total_vram_mb > 0 else 40960.0
        vram_pct = (peak_vram_mb / total_vram_mb) * 100 if total_vram_mb > 0 else 0

        avg_gpu = sum(monitor.gpu_utils) / len(monitor.gpu_utils) if monitor.gpu_utils else 0.0
        avg_cpu = sum(monitor.cpu_percents) / len(monitor.cpu_percents) if monitor.cpu_percents else 0.0

        result.update({
            "status": "SUCCESS",
            "ips": round(ips, 2),
            "avg_reader_cost": round(avg_reader_cost, 4),
            "avg_batch_cost": round(avg_batch_cost, 4),
            "reader_ratio": round(avg_reader_cost / avg_batch_cost, 3) if avg_batch_cost > 0 else 0,
            "peak_vram_mb": round(peak_vram_mb, 1),
            "total_vram_mb": round(total_vram_mb, 1),
            "vram_pct": round(vram_pct, 1),
            "avg_gpu_util": round(avg_gpu, 1),
            "avg_cpu_pct": round(avg_cpu, 1)
        })
        print(f"[Worker SUCCESS] Batch={batch_size} Workers={num_workers} => IPS={ips:.1f} samples/s, VRAM={vram_pct:.1f}%, ReaderRatio={result['reader_ratio']}", flush=True)
        exit_code = 0

    except Exception as e:
        monitor.stop()
        err_str = str(e).lower()
        full_tb = traceback.format_exc()
        print(f"[Worker ERROR] {e}\n{full_tb}", file=sys.stderr, flush=True)

        if "out of memory" in err_str or "bad_alloc" in err_str or "cudabadalloc" in err_str or "memoryerror" in err_str:
            result["status"] = "OOM"
            result["error"] = "CUDA Out Of Memory"
            exit_code = 137
        elif "loss exploded" in err_str or "nan" in err_str or "inf" in err_str:
            result["status"] = "NAN_LOSS"
            result["error"] = str(e)
            exit_code = 138
        else:
            result["status"] = "CRASH"
            result["error"] = str(e)
            exit_code = 1

    with open(output_json, 'w', encoding='utf-8') as f:
        json.dump(result, f, indent=2, ensure_ascii=False)

    sys.exit(exit_code)


if __name__ == "__main__":
    parser = argparse.ArgumentParser(description="Isolated PaddleOCR Benchmark Worker")
    parser.add_argument("-c", "--config_path", type=str, required=True, help="Path to YAML training config")
    parser.add_argument("-b", "--batch_size", type=int, required=True, help="Batch size per card")
    parser.add_argument("-w", "--num_workers", type=int, required=True, help="Number of DataLoader workers")
    parser.add_argument("--warmup_steps", type=int, default=100, help="Warmup steps before measuring")
    parser.add_argument("--measure_steps", type=int, default=300, help="Measured steps")
    parser.add_argument("-o", "--output_json", type=str, default="trial_result.json", help="Path to write trial JSON")
    args = parser.parse_args()

    run_benchmark(
        config_path=args.config_path,
        batch_size=args.batch_size,
        num_workers=args.num_workers,
        warmup_steps=args.warmup_steps,
        measure_steps=args.measure_steps,
        output_json=args.output_json
    )