PlantAI-rag-new/agent/legacy-rag/scripts/run_v16_aqua_batch.py

#!/usr/bin/env python3
"""V16.aqua.1.0 batch extraction with resume support.

Input manifest: JSON array of {paper_id, path, kind: "md"|"pdf"}.
Each paper runs in its own subprocess (run_v16_aqua_one_paper.py); stdout/stderr
are captured to <output_root>/<paper_id>/run.log. Progress is persisted to
<output_root>/batch_progress.json after every completed future; rows whose
status is "success" are skipped on startup, so re-running the same command
resumes an interrupted batch. On completion, BATCH_RUN_SUMMARY.json is written.
Failed rows keep their status and are retried on the next run.
"""
from __future__ import annotations

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

PROJECT_ROOT = Path(__file__).resolve().parents[1]  # agent/legacy-rag（one_paper 的 sys.path 根）
REPO_ROOT = Path(__file__).resolve().parents[3]  # 仓库根(PlantAI-rag-new)：manifest 相对路径以此解析


def _row_id(row: dict) -> str:
    return str(row.get("paper_id") or Path(str(row.get("path", ""))).stem)


def _safe_dir(paper_id: str) -> str:
    """Filesystem-safe directory name for a paper_id (DOIs contain '/')."""
    return re.sub(r"[^A-Za-z0-9._-]+", "_", str(paper_id))


def build_run_command(
    row: dict,
    root: Path,
    output_dir: Path,
    python: Path,
    *,
    use_pc_data: bool = False,
    disable_llm: bool = False,
) -> list[str]:
    kind = str(row.get("kind") or "md").lower()
    source = Path(str(row["path"]))
    if not source.is_absolute():
        source = REPO_ROOT / source  # manifest 路径为仓库根相对路径
    if kind == "pdf":
        flag, value = "--pdf", source
    else:
        flag, value = "--md", source
    command = [
        str(python),
        str(root / "scripts/run_v16_aqua_one_paper.py"),
        flag,
        str(value),
        "--output-dir",
        str(output_dir),
    ]
    pmid = str(row.get("pmid") or "").strip()
    doi = str(row.get("doi") or "").strip()
    paper_id = str(row.get("paper_id") or "").strip()
    if pmid:
        command.extend(["--pmid", pmid])
    if doi:
        command.extend(["--doi", doi])
    if paper_id:
        command.extend(["--paper-id", paper_id])
    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:
    paper_id = _row_id(row)
    out = output_root / _safe_dir(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(),
    }
    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=4)
    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"))
    # Resolve to absolute at the caller's cwd, then to the caller's cwd so the
    # subprocess (cwd=project_root) writes into the same tree the progress file
    # lives in. Relative output roots would otherwise land under project_root.
    args.output_root = args.output_root.resolve()
    args.manifest = args.manifest.resolve()
    args.output_root.mkdir(parents=True, exist_ok=True)

    # Resume: load previous progress, skip success rows.
    progress_path = args.output_root / "batch_progress.json"
    prior = {}
    if progress_path.exists():
        try:
            prior = {_row_id(item): item for item in json.loads(progress_path.read_text(encoding="utf-8"))}
        except Exception:
            prior = {}
    todo = [row for row in rows if prior.get(_row_id(row), {}).get("status") != "success"]
    done_results = [prior[_row_id(row)] for row in rows if prior.get(_row_id(row), {}).get("status") == "success"]
    print(
        f"manifest={len(rows)}  already_success={len(done_results)}  to_run={len(todo)}",
        flush=True,
    )
    if not todo:
        print("nothing to run; writing summary and exiting", flush=True)
    else:
        results: list[dict] = []
        started = time.time()
        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 todo
            }
            for future in concurrent.futures.as_completed(future_map):
                row = future_map[future]
                try:
                    result = future.result()
                    result["status"] = "success" if result["success"] else "failed"
                except Exception as exc:
                    result = {
                        **row,
                        "paper_id": _row_id(row),
                        "return_code": -1,
                        "wall_seconds": 0,
                        "success": False,
                        "status": "failed",
                        "error": repr(exc),
                    }
                results.append(result)
                all_results = done_results + results
                (progress_path).write_text(
                    json.dumps(all_results, ensure_ascii=False, indent=2),
                    encoding="utf-8",
                )
                print(
                    f"[{len(all_results)}/{len(rows)}] {result['paper_id']} "
                    f"status={result['status']} wall={result['wall_seconds']}s",
                    flush=True,
                )
        done_results.extend(results)

    summary = {
        "started_epoch": time.time(),
        "workers": args.workers,
        "total": len(done_results),
        "success": sum(bool(item.get("status") == "success") for item in done_results),
        "failed": sum(bool(item.get("status") != "success") for item in done_results),
        "results": done_results,
    }
    (args.output_root / "BATCH_RUN_SUMMARY.json").write_text(
        json.dumps(summary, ensure_ascii=False, indent=2),
        encoding="utf-8",
    )
    print(
        f"BATCH DONE success={summary['success']} failed={summary['failed']}",
        flush=True,
    )
    return 0 if summary["failed"] == 0 else 1


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