fetch_gpu_perf.py

import argparse
import re
from dataclasses import dataclass
from html.parser import HTMLParser
from pathlib import Path
from urllib.request import Request, urlopen


DEFAULT_URL = "https://www.autodl.com/docs/gpu_perf/"
DEFAULT_OUTPUT_DIR = Path("data/gpu_perf")


@dataclass(frozen=True)
class GpuPerfRecord:
    gpu_name: str
    gpu_id: str
    content: str


def gpu_id(name: str) -> str:
    normalized = re.sub(r"[^\w]+", "_", name.strip(), flags=re.UNICODE)
    return normalized.strip("_")


class GpuPerfParser(HTMLParser):
    def __init__(self) -> None:
        super().__init__(convert_charrefs=True)
        self.labels: list[str] = []
        self.code_blocks: list[str] = []
        self._in_tabbed_labels = False
        self._tabbed_labels_depth = 0
        self._in_label = False
        self._label_parts: list[str] = []
        self._in_code = False
        self._code_parts: list[str] = []

    def handle_starttag(self, tag: str, attrs: list[tuple[str, str | None]]) -> None:
        attrs_dict = dict(attrs)
        classes = set((attrs_dict.get("class") or "").split())

        if tag == "div" and "tabbed-labels" in classes:
            self._in_tabbed_labels = True
            self._tabbed_labels_depth = 1
        elif self._in_tabbed_labels and tag == "div":
            self._tabbed_labels_depth += 1

        if self._in_tabbed_labels and tag == "label":
            self._in_label = True
            self._label_parts = []

        if tag == "code":
            self._in_code = True
            self._code_parts = []

    def handle_data(self, data: str) -> None:
        if self._in_label:
            self._label_parts.append(data)
        if self._in_code:
            self._code_parts.append(data)

    def handle_endtag(self, tag: str) -> None:
        if self._in_label and tag == "label":
            label = " ".join("".join(self._label_parts).split())
            if label:
                self.labels.append(label)
            self._in_label = False

        if self._in_code and tag == "code":
            content = "".join(self._code_parts).strip("\n") + "\n"
            if content.strip():
                self.code_blocks.append(content)
            self._in_code = False

        if self._in_tabbed_labels and tag == "div":
            self._tabbed_labels_depth -= 1
            if self._tabbed_labels_depth == 0:
                self._in_tabbed_labels = False


def extract_gpu_perf(html: str) -> list[GpuPerfRecord]:
    parser = GpuPerfParser()
    parser.feed(html)

    if len(parser.labels) != len(parser.code_blocks):
        raise ValueError(
            f"GPU labels/code block count mismatch: "
            f"{len(parser.labels)} labels, {len(parser.code_blocks)} code blocks"
        )

    return [
        GpuPerfRecord(gpu_name=name, gpu_id=gpu_id(name), content=content)
        for name, content in zip(parser.labels, parser.code_blocks)
    ]


def fetch_html(url: str) -> str:
    request = Request(url, headers={"User-Agent": "Mozilla/5.0"})
    with urlopen(request, timeout=30) as response:
        charset = response.headers.get_content_charset() or "utf-8"
        return response.read().decode(charset)


def write_gpu_perf(records: list[GpuPerfRecord], output_dir: Path) -> list[Path]:
    output_dir.mkdir(parents=True, exist_ok=True)

    written: list[Path] = []
    for record in records:
        path = output_dir / f"{record.gpu_id}.txt"
        path.write_text(record.content, encoding="utf-8")
        written.append(path)
    return written


def fetch_and_write(url: str = DEFAULT_URL, output_dir: Path = DEFAULT_OUTPUT_DIR) -> list[Path]:
    html = fetch_html(url)
    records = extract_gpu_perf(html)
    return write_gpu_perf(records, output_dir)


def main() -> None:
    parser = argparse.ArgumentParser(
        description="Fetch AutoDL GPU performance logs into per-GPU text files."
    )
    parser.add_argument("--url", default=DEFAULT_URL, help="AutoDL GPU performance page URL.")
    parser.add_argument(
        "--output-dir",
        default=str(DEFAULT_OUTPUT_DIR),
        help="Directory for generated GPU_ID.txt files.",
    )
    args = parser.parse_args()

    written = fetch_and_write(args.url, Path(args.output_dir))
    for path in written:
        print(path)


if __name__ == "__main__":
    main()