TheRock/build_tools/github_actions/generate_jax_manifest.py

#!/usr/bin/env python3
"""
Generate a manifest for JAX external builds.

Writes a JSON manifest containing:
  - jax: git commit + origin repo (+ branch best-effort)
  - therock: repo + commit + branch from GitHub Actions env (best-effort)

Filename format:
  therock-manifest_jax_py<python_version>_<jax_git_ref>.json
"""

import argparse
import json
import os
from pathlib import Path
import sys

from github_actions.manifest_utils import (
    GitSourceInfo,
    git_branch_best_effort,
    git_head,
    log,
    normalize_python_version_for_filename,
    normalize_ref_for_filename,
    resolve_branch,
)


def manifest_filename(*, python_version: str, jax_git_ref: str) -> str:
    py = normalize_python_version_for_filename(python_version)
    ref = normalize_ref_for_filename(jax_git_ref)
    return f"therock-manifest_jax_py{py}_{ref}.json"


def build_sources(*, jax_dir: Path, jax_git_ref: str) -> dict[str, dict[str, str]]:
    jax = git_head(jax_dir, label="jax")

    jax_branch = resolve_branch(
        inferred=git_branch_best_effort(jax_dir),
        provided=jax_git_ref,
    )

    return {
        "jax": GitSourceInfo(
            commit=jax.commit, repo=jax.repo, branch=jax_branch
        ).to_dict(),
    }


def build_manifest(
    *,
    sources: dict[str, dict[str, str]],
    therock_repo: str,
    therock_commit: str,
    therock_branch: str,
) -> dict[str, object]:
    manifest: dict[str, object] = {}
    manifest.update(sources)
    manifest["therock"] = {
        "commit": therock_commit,
        "repo": therock_repo,
        "branch": therock_branch,
    }
    return manifest


def generate_manifest_dict(
    *,
    jax_dir: Path,
    python_version: str,
    jax_git_ref: str,
) -> tuple[str, dict[str, object]]:
    sources = build_sources(jax_dir=jax_dir, jax_git_ref=jax_git_ref)

    server_url = os.environ.get("GITHUB_SERVER_URL")
    repo = os.environ.get("GITHUB_REPOSITORY")
    sha = os.environ.get("GITHUB_SHA")
    ref = os.environ.get("GITHUB_REF")

    therock_repo = "unknown"
    if server_url and repo:
        therock_repo = f"{server_url}/{repo}.git"

    therock_commit = sha or "unknown"

    therock_branch = "unknown"
    if ref:
        if ref.startswith("refs/heads/"):
            therock_branch = ref[len("refs/heads/") :]
        else:
            therock_branch = ref

    name = manifest_filename(
        python_version=python_version,
        jax_git_ref=jax_git_ref,
    )

    manifest = build_manifest(
        sources=sources,
        therock_repo=therock_repo,
        therock_commit=therock_commit,
        therock_branch=therock_branch,
    )
    return name, manifest


def parse_args(argv: list[str]) -> argparse.Namespace:
    ap = argparse.ArgumentParser(description="Generate JAX manifest.")
    ap.add_argument(
        "--manifest-dir",
        type=Path,
        required=True,
        help="Output directory for the manifest JSON.",
    )
    ap.add_argument(
        "--python-version",
        required=True,
        help="Python version for manifest naming (e.g. 3.11 or py3.11).",
    )
    ap.add_argument(
        "--jax-git-ref",
        required=True,
        help=(
            "Git ref used for manifest naming and branch fallback "
            "(e.g. nightly, release/0.4.28, rocm-jaxlib-v0.8.2)."
        ),
    )
    ap.add_argument("--jax-dir", type=Path, required=True)
    return ap.parse_args(argv)


def main(argv: list[str]) -> None:
    args = parse_args(argv)

    manifest_dir = args.manifest_dir.resolve()
    manifest_dir.mkdir(parents=True, exist_ok=True)

    name, manifest = generate_manifest_dict(
        jax_dir=args.jax_dir,
        python_version=args.python_version,
        jax_git_ref=args.jax_git_ref,
    )

    out_path = manifest_dir / name
    out_path.write_text(
        json.dumps(manifest, indent=2, sort_keys=False) + "\n",
        encoding="utf-8",
    )

    if not out_path.exists():
        raise RuntimeError(f"Failed to write manifest: {out_path}")
    if out_path.stat().st_size == 0:
        raise RuntimeError(f"Manifest is empty: {out_path}")

    log(f"[jax-sources-manifest] wrote {out_path}")


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