TheRock/build_tools/github_actions/upload_jax_manifest.py

#!/usr/bin/env python3
"""
Upload the generated JAX manifest JSON to S3.

Upload layout:
  s3://{bucket}/{external_repo}{run_id}-{platform}/manifests/{amdgpu_family}/{manifest_name}
"""

import argparse
from pathlib import Path
import platform
import sys

sys.path.insert(0, str(Path(__file__).resolve().parent.parent))
from _therock_utils.workflow_outputs import WorkflowOutputRoot
from _therock_utils.storage_location import StorageLocation
from _therock_utils.storage_backend import create_storage_backend

from github_actions.manifest_utils import (
    normalize_python_version_for_filename,
    normalize_ref_for_filename,
)


PLATFORM = platform.system().lower()


def _log(*args: object) -> None:
    print(*args)
    sys.stdout.flush()


def _make_output_root(
    run_id: str, bucket_override: str | None = None
) -> WorkflowOutputRoot:
    if bucket_override:
        return WorkflowOutputRoot(
            bucket=bucket_override,
            external_repo="",
            run_id=run_id,
            platform=PLATFORM,
        )
    return WorkflowOutputRoot.from_workflow_run(run_id=run_id, platform=PLATFORM)


def parse_args(argv: list[str]) -> argparse.Namespace:
    parser = argparse.ArgumentParser(description="Upload a JAX manifest JSON to S3.")
    parser.add_argument(
        "--dist-dir",
        type=Path,
        required=True,
        help="Wheel dist dir (contains manifests/).",
    )
    parser.add_argument(
        "--run-id",
        type=str,
        required=True,
        help="Workflow run ID (e.g. 21440027240).",
    )
    parser.add_argument(
        "--amdgpu-family",
        type=str,
        required=True,
        help="AMDGPU family (e.g. gfx94X-dcgpu).",
    )
    parser.add_argument(
        "--python-version",
        type=str,
        required=True,
        help="Python version (e.g. 3.12 or py3.12).",
    )
    parser.add_argument(
        "--jax-git-ref",
        type=str,
        required=True,
        help=(
            "JAX git ref used in manifest naming "
            '(e.g. "nightly", "release/0.4.28", "rocm-jaxlib-v0.8.2").'
        ),
    )
    parser.add_argument(
        "--bucket",
        type=str,
        default=None,
        help="Override S3 bucket (default: auto-select from workflow run).",
    )
    parser.add_argument(
        "--output-dir",
        type=Path,
        default=None,
        help="Output to local directory instead of S3 (for testing).",
    )
    parser.add_argument(
        "--dry-run",
        action="store_true",
        help="Print what would be uploaded without actually uploading.",
    )
    return parser.parse_args(argv)


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

    py = normalize_python_version_for_filename(args.python_version)
    track = normalize_ref_for_filename(args.jax_git_ref)

    manifest_name = f"therock-manifest_jax_py{py}_{track}.json"
    manifest_path = (args.dist_dir / "manifests" / manifest_name).resolve()

    _log(f"Manifest expected at: {manifest_path}")
    if not manifest_path.is_file():
        raise FileNotFoundError(f"Manifest not found: {manifest_path}")

    output_root = _make_output_root(args.run_id, bucket_override=args.bucket)
    manifest_dir_loc = output_root.manifest_dir(args.amdgpu_family)
    dest = StorageLocation(
        manifest_dir_loc.bucket,
        f"{manifest_dir_loc.relative_path}/{manifest_name}",
    )

    backend = create_storage_backend(staging_dir=args.output_dir, dry_run=args.dry_run)
    backend.upload_file(manifest_path, dest)


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