TheRock/external-builds/pytorch/checkout_from_manifest.py

#!/usr/bin/env python3
# Copyright Advanced Micro Devices, Inc.
# SPDX-License-Identifier: MIT

"""Check out pytorch ecosystem repositories from a manifest file.

Reads a manifest JSON (produced by ``generate_pytorch_source_manifest.py``)
and calls the individual ``pytorch_*_repo.py checkout`` scripts with the
exact commit SHAs and repo URLs from the manifest.

This provides a single command to reproduce the same source tree that CI
builds use, given only a manifest file.

Usage::

    # Check out all projects in the manifest under ./checkouts/:
    python checkout_from_manifest.py \\
        --manifest manifests/therock-manifest_torch_linux_release-2.10.json \\
        --checkout-root ./checkouts

    # Check out only pytorch (skip audio, vision, etc.):
    python checkout_from_manifest.py \\
        --manifest manifest.json \\
        --checkout-root ./checkouts \\
        --projects pytorch
"""

import argparse
import json
import subprocess
import sys
import urllib.request
from pathlib import Path

THIS_DIR = Path(__file__).resolve().parent
SCHEMA_VERSION = 1

# Maps manifest project names to their checkout scripts.
CHECKOUT_SCRIPTS: dict[str, str] = {
    "pytorch": "pytorch_torch_repo.py",
    "pytorch_audio": "pytorch_audio_repo.py",
    "pytorch_vision": "pytorch_vision_repo.py",
    "triton": "pytorch_triton_repo.py",
    "apex": "pytorch_apex_repo.py",
}


def _split_words(value: str) -> list[str]:
    return value.replace(";", " ").split() if value else []


def log(*args, **kwargs) -> None:
    print(*args, **kwargs)
    sys.stdout.flush()


def checkout_project(
    *,
    name: str,
    source_info: dict[str, str],
    checkout_root: Path,
    no_hipify: bool,
) -> None:
    """Check out a single project using its checkout script."""
    script = THIS_DIR / CHECKOUT_SCRIPTS[name]
    checkout_dir = checkout_root / name

    commit = source_info["commit"]
    repo = source_info["repo"]

    log(f"  {name}: {repo} @ {commit[:12]}")

    cmd: list[str] = [
        sys.executable,
        str(script),
        "checkout",
        "--gitrepo-origin",
        repo,
        "--repo-hashtag",
        commit,
        "--checkout-dir",
        str(checkout_dir),
    ]

    if no_hipify:
        cmd.append("--no-hipify")

    log(f"  Running: {' '.join(cmd)}")
    subprocess.check_call(cmd)


def download_manifest(*, manifest_url: str, output_path: Path) -> Path:
    """Download a manifest URL to output_path."""
    output_path.parent.mkdir(parents=True, exist_ok=True)
    log(f"Downloading manifest: {manifest_url} -> {output_path}")
    urllib.request.urlretrieve(manifest_url, output_path)
    if not output_path.is_file() or output_path.stat().st_size == 0:
        raise RuntimeError(f"Failed to download manifest: {manifest_url}")
    return output_path


def load_manifest(manifest_path: Path) -> dict[str, object]:
    manifest = json.loads(manifest_path.read_text(encoding="utf-8"))
    if not isinstance(manifest, dict):
        raise ValueError("Manifest root must be a JSON object")
    schema_version = manifest.get("schema_version")
    if schema_version is not None and schema_version != SCHEMA_VERSION:
        raise ValueError(
            f"Unsupported manifest schema_version {schema_version!r}; "
            f"expected {SCHEMA_VERSION}"
        )
    return manifest


def require_project_source_info(
    manifest: dict[str, object], project: str
) -> dict[str, str]:
    source_info = manifest.get(project)
    if not isinstance(source_info, dict):
        raise ValueError(f"Manifest entry {project!r} must be a JSON object")

    for field in ("repo", "commit"):
        value = source_info.get(field)
        if not isinstance(value, str) or not value:
            raise ValueError(
                f"Manifest entry {project!r} is missing required string "
                f"field {field!r}"
            )

    return {
        str(field): value
        for field, value in source_info.items()
        if isinstance(value, str)
    }


def validate_manifest(
    manifest: dict[str, object], *, expected_pytorch_git_ref: str
) -> None:
    """Validate manifest contents against workflow expectations."""
    if not expected_pytorch_git_ref:
        return
    pytorch = manifest.get("pytorch")
    if not isinstance(pytorch, dict):
        raise ValueError("Manifest is missing pytorch entry")
    actual_ref = pytorch.get("branch")
    if actual_ref != expected_pytorch_git_ref:
        raise ValueError(
            f"Manifest PyTorch ref {actual_ref!r} does not match "
            f"{expected_pytorch_git_ref!r}"
        )


def main(argv: list[str]) -> None:
    parser = argparse.ArgumentParser(
        description="Check out pytorch repos from a manifest file"
    )
    manifest_input = parser.add_mutually_exclusive_group(required=True)
    manifest_input.add_argument(
        "--manifest",
        type=Path,
        help="Path to manifest JSON file",
    )
    manifest_input.add_argument(
        "--manifest-url",
        help="URL to a manifest JSON file to download before checkout",
    )
    parser.add_argument(
        "--checkout-root",
        type=Path,
        required=True,
        help="Root directory for checkouts (each project gets a subdirectory)",
    )
    parser.add_argument(
        "--projects",
        default="",
        help=(
            "Semicolon- or space-separated list of projects to check out "
            "(default: all projects in the manifest)."
        ),
    )
    parser.add_argument(
        "--expected-pytorch-git-ref",
        default="",
        help="Validate manifest pytorch.branch against this ref before checkout",
    )
    parser.add_argument(
        "--no-hipify",
        action="store_true",
        default=False,
        help="Skip HIPIFY for all checkouts (e.g. for test-only runs)",
    )
    args = parser.parse_args(argv)

    checkout_root = args.checkout_root.resolve()
    manifest_path = args.manifest
    if args.manifest_url:
        manifest_path = checkout_root / "pytorch_manifest.json"
        download_manifest(manifest_url=args.manifest_url, output_path=manifest_path)

    manifest = load_manifest(manifest_path)
    validate_manifest(manifest, expected_pytorch_git_ref=args.expected_pytorch_git_ref)

    # Determine which projects to check out.
    available = [name for name in CHECKOUT_SCRIPTS if name in manifest]
    if args.projects:
        projects = _split_words(args.projects)
        unknown = set(projects) - set(available)
        if unknown:
            parser.error(
                f"Projects not in manifest: {unknown}. " f"Available: {available}"
            )
    else:
        projects = available
    if not projects:
        supported = ", ".join(CHECKOUT_SCRIPTS)
        raise ValueError(
            "Manifest contains no supported PyTorch projects. "
            f"Supported projects: {supported}"
        )

    source_infos = {
        name: require_project_source_info(manifest, name) for name in projects
    }

    log(f"Manifest: {manifest_path}")
    log(f"Checkout root: {checkout_root}")
    log(f"Projects: {projects}")
    log("")

    checkout_root.mkdir(parents=True, exist_ok=True)

    # Keep pytorch first so a full checkout tree is easy to inspect.
    if "pytorch" in projects:
        projects = ["pytorch"] + [p for p in projects if p != "pytorch"]

    for name in projects:
        checkout_project(
            name=name,
            source_info=source_infos[name],
            checkout_root=checkout_root,
            no_hipify=args.no_hipify,
        )
        log("")

    log("All checkouts complete.")


if __name__ == "__main__":
    main(sys.argv[1:])