parse_gpu_perf.py

import argparse
import ast
import json
import re
from pathlib import Path
from typing import Any


DEFAULT_INPUT_DIR = Path("data/gpu_perf")
DEFAULT_OUTPUT = Path("data/gpu_perf_records.jsonl")

NAMESPACE_RE = re.compile(r"^Namespace\((?P<body>.*)\)$")
ITERATION_RE = re.compile(
    r"^Iteration (?P<iteration>\d+), "
    r"(?P<images_per_second>\d+(?:\.\d+)?) images/s in "
    r"(?P<seconds>\d+(?:\.\d+)?)s\.?$"
)


def parse_namespace(line: str) -> dict[str, Any]:
    match = NAMESPACE_RE.match(line.strip())
    if not match:
        raise ValueError(f"not a Namespace line: {line!r}")

    values: dict[str, Any] = {}
    for part in re.split(r",\s*", match.group("body")):
        if not part:
            continue
        key, raw_value = part.split("=", 1)
        try:
            value = ast.literal_eval(raw_value)
        except (SyntaxError, ValueError):
            value = raw_value
        values[key] = value
    return values


def parse_gpu_perf_text(
    text: str,
    gpu_id: str | None = None,
    source_file: str | None = None,
) -> list[dict[str, Any]]:
    hardware_model: str | None = None
    framework: str | None = None
    benchmark: str | None = None
    context: dict[str, Any] = {}
    records: list[dict[str, Any]] = []

    for raw_line in text.splitlines():
        line = raw_line.strip()
        if not line:
            continue

        if line.startswith("GPU 具体型号为："):
            hardware_model = line.split("：", 1)[1].strip()
            continue

        if line.startswith("使用"):
            framework = line.removeprefix("使用").strip()
            continue

        if line.startswith(">>>"):
            benchmark = line.removeprefix(">>>").strip()
            context = {}
            continue

        if NAMESPACE_RE.match(line):
            context = parse_namespace(line)
            continue

        iteration_match = ITERATION_RE.match(line)
        if iteration_match:
            if benchmark is None:
                raise ValueError(f"iteration without benchmark context: {line!r}")
            if not context:
                raise ValueError(f"iteration without Namespace context: {line!r}")

            record = {
                "gpu_id": gpu_id,
                "source_file": source_file,
                "hardware_model": hardware_model,
                "framework": framework,
                "benchmark": benchmark,
                "model": context.get("model"),
                "precision": context.get("precision"),
                "train": context.get("train"),
                "device": context.get("device"),
                "iteration": int(iteration_match.group("iteration")),
                "images_per_second": float(iteration_match.group("images_per_second")),
                "seconds": float(iteration_match.group("seconds")),
            }
            records.append(record)

    return records


def parse_gpu_perf_file(path: Path) -> list[dict[str, Any]]:
    return parse_gpu_perf_text(
        path.read_text(encoding="utf-8"),
        gpu_id=path.stem,
        source_file=str(path),
    )


def parse_gpu_perf_dir(input_dir: Path) -> list[dict[str, Any]]:
    records: list[dict[str, Any]] = []
    for path in sorted(input_dir.glob("*.txt")):
        records.extend(parse_gpu_perf_file(path))
    return records


def write_jsonl(records: list[dict[str, Any]], output_path: Path) -> None:
    output_path.parent.mkdir(parents=True, exist_ok=True)
    with output_path.open("w", encoding="utf-8") as f:
        for record in records:
            f.write(json.dumps(record, ensure_ascii=False, sort_keys=True) + "\n")


def main() -> None:
    parser = argparse.ArgumentParser(
        description="Parse AutoDL GPU performance text logs into JSONL iteration records."
    )
    parser.add_argument(
        "--input-dir",
        default=str(DEFAULT_INPUT_DIR),
        help="Directory containing GPU_ID.txt performance logs.",
    )
    parser.add_argument(
        "--output",
        default=str(DEFAULT_OUTPUT),
        help="JSONL output path. Use '-' to print to stdout.",
    )
    args = parser.parse_args()

    records = parse_gpu_perf_dir(Path(args.input_dir))
    if args.output == "-":
        for record in records:
            print(json.dumps(record, ensure_ascii=False, sort_keys=True))
    else:
        write_jsonl(records, Path(args.output))
        print(f"Wrote {len(records)} records to {args.output}")


if __name__ == "__main__":
    main()