PlantAI-rag/scripts/run_v16_batch10_20260625.py

#!/usr/bin/env python3
from __future__ import annotations

import argparse
import concurrent.futures
import json
import os
import subprocess
import sys
import time
from pathlib import Path


PROJECT_ROOT = Path(__file__).resolve().parents[1]


def build_run_command(
    row: dict,
    root: Path,
    output_dir: Path,
    python: Path,
    *,
    use_pc_data: bool = False,
    disable_llm: bool = False,
) -> list[str]:
    command = [
        str(python),
        str(root / "scripts/run_v16_1_2_one_paper.py"),
        "--pdf",
        str(Path(row["destination_path"])),
        "--output-dir",
        str(output_dir),
    ]
    pmid = str(row.get("pmid") or "").strip()
    doi = str(row.get("doi") or "").strip()
    if pmid:
        command.extend(["--pmid", pmid])
    if doi:
        command.extend(["--doi", doi])
    if use_pc_data:
        command.append("--use-pc-data")
    if disable_llm:
        command.append("--disable-llm")
    return command


def run_one(
    row: dict,
    root: Path,
    output_root: Path,
    python: Path,
    *,
    use_pc_data: bool = False,
    disable_llm: bool = False,
) -> dict:
    pdf = Path(row["destination_path"])
    paper_id = str(row.get("paper_id") or pdf.stem)
    out = output_root / paper_id
    out.mkdir(parents=True, exist_ok=True)
    command = [
        *build_run_command(
            row,
            root,
            out,
            python,
            use_pc_data=use_pc_data,
            disable_llm=disable_llm,
        )
    ]
    log_path = out / "run.log"
    started = time.time()
    with log_path.open("w", encoding="utf-8") as log:
        process = subprocess.run(
            command,
            cwd=root,
            env=os.environ.copy(),
            stdout=log,
            stderr=subprocess.STDOUT,
            text=True,
        )
    ended = time.time()
    result_path = out / "result.json"
    result = {
        **row,
        "paper_id": paper_id,
        "output_dir": str(out),
        "result_json": str(result_path),
        "return_code": process.returncode,
        "wall_seconds": round(ended - started, 3),
        "success": process.returncode == 0 and result_path.exists(),
    }
    if result["success"]:
        command_text = " ".join(
            f"'{arg}'" if " " in arg else arg for arg in command
        )
        (out / "SUCCESSFUL_COMMAND.md").write_text(
            "# Successful V16.1.2 command\n\n"
            "Environment: current process environment/private `.env`.\n\n"
            f"```bash\n{command_text}\n```\n",
            encoding="utf-8",
        )
    return result


def main() -> int:
    parser = argparse.ArgumentParser()
    parser.add_argument("--manifest", required=True, type=Path)
    parser.add_argument("--output-root", required=True, type=Path)
    parser.add_argument("--workers", type=int, default=3)
    parser.add_argument("--use-pc-data", action="store_true")
    parser.add_argument("--disable-llm", action="store_true")
    parser.add_argument(
        "--project-root",
        type=Path,
        default=Path(os.environ.get("PLANTAI_PROJECT_ROOT", PROJECT_ROOT)),
    )
    parser.add_argument(
        "--python",
        type=Path,
        default=Path(os.environ.get("PLANTAI_PYTHON", sys.executable)),
    )
    args = parser.parse_args()
    if args.workers < 1:
        raise SystemExit("--workers must be >= 1")
    rows = json.loads(args.manifest.read_text(encoding="utf-8"))
    args.output_root.mkdir(parents=True, exist_ok=True)
    started = time.time()
    results: list[dict] = []
    with concurrent.futures.ThreadPoolExecutor(max_workers=args.workers) as pool:
        future_map = {
            pool.submit(
                run_one,
                row,
                args.project_root,
                args.output_root,
                args.python,
                use_pc_data=args.use_pc_data,
                disable_llm=args.disable_llm,
            ): row
            for row in rows
        }
        for future in concurrent.futures.as_completed(future_map):
            row = future_map[future]
            try:
                result = future.result()
            except Exception as exc:
                result = {
                    **row,
                    "paper_id": Path(row["file_name"]).stem,
                    "return_code": -1,
                    "wall_seconds": 0,
                    "success": False,
                    "error": repr(exc),
                }
            results.append(result)
            results.sort(key=lambda item: int(item["index"]))
            (args.output_root / "batch_progress.json").write_text(
                json.dumps(results, ensure_ascii=False, indent=2),
                encoding="utf-8",
            )
            print(
                f"[{len(results)}/{len(rows)}] "
                f"{result['paper_id']} success={result['success']} "
                f"wall={result['wall_seconds']}s",
                flush=True,
            )
    summary = {
        "started_epoch": started,
        "finished_epoch": time.time(),
        "wall_seconds": round(time.time() - started, 3),
        "workers": args.workers,
        "total": len(results),
        "success": sum(bool(item["success"]) for item in results),
        "failed": sum(not bool(item["success"]) for item in results),
        "results": results,
    }
    (args.output_root / "BATCH_RUN_SUMMARY.json").write_text(
        json.dumps(summary, ensure_ascii=False, indent=2),
        encoding="utf-8",
    )
    return 0 if summary["failed"] == 0 else 1


if __name__ == "__main__":
    raise SystemExit(main())