TheRock/build_tools/build_tarballs.py

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

"""Fetch multi-arch build artifacts and package them into per-family tarballs.

For each GPU family in --dist-amdgpu-families, this script:
1. Fetches artifacts (generic + family-specific) using artifact_manager.py
2. Flattens them into a single install-prefix-like layout
3. Compresses the result into a tarball

When KPACK_SPLIT_ARTIFACTS is enabled in the build manifest, device-specific
files are split by individual GPU target and don't conflict across families.
In that case, this script also produces a combined multi-arch tarball
containing all targets in a single install prefix.

A shared download cache avoids re-downloading generic (host) artifacts
when processing multiple families.

By default, generated tarballs exclude test artifacts and fftw3. Pass
``--include-test-tarballs`` to also generate full tarballs, named with a
``-tests`` suffix, that include test artifacts.

Tarball naming follows the existing release convention:
    therock-dist-{platform}-{family}-{version}.tar.gz
    therock-dist-{platform}-multiarch-{version}.tar.gz  (KPACK split only)

Example
-------
    python build_tools/build_tarballs.py \\
        --run-id=24104028483 \\
        --dist-amdgpu-families="gfx94X-dcgpu;gfx110X-all" \\
        --platform=linux \\
        --package-version="7.13.0.dev0+abc123" \\
        --output-dir=/tmp/tarballs

Manual testing
--------------
Find a recent multi-arch CI run at
https://github.com/ROCm/TheRock/actions/workflows/multi_arch_ci.yml
and use its run ID. Use ``--platform`` to select which platform's
artifacts to fetch (defaults to the current system).

Expected output: one .tar.gz per family in ``--output-dir``, named
``therock-dist-{platform}-{family}-{version}.tar.gz``. If
KPACK_SPLIT_ARTIFACTS is enabled in the build, also a
``therock-dist-{platform}-multiarch-{version}.tar.gz``.

Each tarball should contain a standard install prefix layout
(``bin/``, ``lib/``, ``include/``, ``share/``, etc.) with GPU-specific
files (e.g. ``lib/hipblaslt/library/*.co``) only for the target family.
"""

import argparse
from concurrent.futures import ProcessPoolExecutor, as_completed
import json
import shlex
import subprocess
import sys
from pathlib import Path

DEFAULT_EXCLUDED_ARTIFACTS: list[str] = ["fftw3"]
DEFAULT_EXCLUDED_COMPONENTS: list[str] = ["test"]


def log(msg: str) -> None:
    print(msg, flush=True)


def run_command(args: list[str | Path], cwd: Path | None = None) -> None:
    args = [str(arg) for arg in args]
    log(f"++ Exec{f' [{cwd}]' if cwd else ''}$ {shlex.join(args)}")
    subprocess.check_call(args, cwd=str(cwd) if cwd else None, stdin=subprocess.DEVNULL)


def fetch_and_flatten(
    *,
    run_id: str,
    amdgpu_families: list[str],
    platform: str,
    output_dir: Path,
    download_cache_dir: Path,
    run_github_repo: str | None = None,
    exclude_components: list[str] | None = None,
    exclude_artifacts: list[str] | None = None,
) -> None:
    """Fetch artifacts for one or more families and flatten into output_dir."""
    families_str = ";".join(amdgpu_families)
    log(f"\n{'='*60}")
    log(f"Fetching artifacts for {families_str}")
    if exclude_components:
        log(f"Excluding components: {', '.join(exclude_components)}")
    if exclude_artifacts:
        log(f"Excluding artifacts: {', '.join(exclude_artifacts)}")
    log(f"{'='*60}")

    cmd = [
        sys.executable,
        "build_tools/artifact_manager.py",
        "fetch",
        f"--run-id={run_id}",
        "--stage=all",
        f"--amdgpu-families={families_str}",
        "--expand-family-to-targets",
        f"--platform={platform}",
        f"--output-dir={output_dir}",
        "--flatten",
        f"--download-cache-dir={download_cache_dir}",
    ]
    if exclude_components:
        cmd.append(f"--exclude-components={','.join(exclude_components)}")
    if exclude_artifacts:
        cmd.append(f"--exclude-artifacts={','.join(exclude_artifacts)}")
    if run_github_repo:
        cmd.append(f"--run-github-repo={run_github_repo}")
    run_command(cmd)


def is_kpack_split(flatten_dir: Path) -> bool:
    """Check if KPACK_SPLIT_ARTIFACTS is enabled from the build manifest."""
    manifest_path = flatten_dir / "share" / "therock" / "therock_manifest.json"
    if not manifest_path.exists():
        return False
    manifest = json.loads(manifest_path.read_text())
    return manifest.get("flags", {}).get("KPACK_SPLIT_ARTIFACTS", False)


def compress_tarball(*, source_dir: Path, tarball_path: Path) -> None:
    """Compress a directory into a .tar.gz tarball.

    Uses subprocess ``tar cfz`` rather than Python's ``tarfile`` module
    (tarfile was significantly slower and produced larger output with default
    settings — its ``compresslevel`` parameter may help but was not tuned).

    Uses gzip to match the existing release tarball format. Switching to
    zstd (``tar cf - . | zstd``) would be faster with better compression,
    but requires downstream consumers to support ``.tar.zst``.
    """
    log(f"\nCompressing {source_dir} -> {tarball_path}")
    tarball_path.parent.mkdir(parents=True, exist_ok=True)
    run_command(["tar", "cfz", str(tarball_path), "."], cwd=source_dir)
    size_mb = tarball_path.stat().st_size / (1024 * 1024)
    log(f"  Created {tarball_path.name} ({size_mb:.1f} MB)")


def main(argv: list[str] | None = None) -> None:
    parser = argparse.ArgumentParser(
        description="Fetch multi-arch artifacts and package into per-family tarballs"
    )
    parser.add_argument("--run-id", required=True, help="Workflow run ID to fetch from")
    parser.add_argument(
        "--run-github-repo",
        type=str,
        default=None,
        help="GitHub repository for --run-id in 'owner/repo' format. "
        "Defaults to GITHUB_REPOSITORY env var or 'ROCm/TheRock'",
    )
    parser.add_argument(
        "--dist-amdgpu-families",
        required=True,
        help="Semicolon-separated GPU families (e.g. 'gfx94X-dcgpu;gfx110X-all')",
    )
    parser.add_argument(
        "--platform",
        default="linux",
        choices=["linux", "windows"],
        help="Platform to fetch artifacts for",
    )
    parser.add_argument(
        "--package-version",
        required=True,
        help="ROCm package version string for tarball naming",
    )
    parser.add_argument(
        "--output-dir",
        type=Path,
        required=True,
        help="Output directory for tarballs",
    )
    parser.add_argument(
        "--include-test-tarballs",
        action="store_true",
        help="Also produce -tests tarballs that include test artifacts",
    )
    args = parser.parse_args(argv)
    # Normalize empty string to None (workflow inputs default to "")
    args.run_github_repo = args.run_github_repo or None

    families = [f.strip() for f in args.dist_amdgpu_families.split(";") if f.strip()]
    if not families:
        raise ValueError("No GPU families specified")

    work_dir = args.output_dir / ".work"
    download_cache_dir = work_dir / "download-cache"
    download_cache_dir.mkdir(parents=True, exist_ok=True)

    log(f"Building tarballs for {len(families)} families: {', '.join(families)}")
    log(f"  Platform: {args.platform}")
    log(f"  Version: {args.package_version}")
    log(f"  Output: {args.output_dir}")
    log(f"  Include test tarballs: {args.include_test_tarballs}")

    # Phase 1: Fetch and flatten sequentially.
    # Sequential so the shared download cache avoids re-downloading generic
    # (host) artifacts for each family.
    family_dirs = []
    compress_tasks = []
    for family in families:
        flatten_dir = work_dir / family
        fetch_and_flatten(
            run_id=args.run_id,
            amdgpu_families=[family],
            platform=args.platform,
            output_dir=flatten_dir,
            download_cache_dir=download_cache_dir,
            run_github_repo=args.run_github_repo,
            exclude_components=DEFAULT_EXCLUDED_COMPONENTS,
            exclude_artifacts=DEFAULT_EXCLUDED_ARTIFACTS,
        )
        family_dirs.append(flatten_dir)
        tarball_name = (
            f"therock-dist-{args.platform}-{family}-{args.package_version}.tar.gz"
        )
        compress_tasks.append((flatten_dir, args.output_dir / tarball_name))
        if args.include_test_tarballs:
            tests_dir = work_dir / "tests" / family
            fetch_and_flatten(
                run_id=args.run_id,
                amdgpu_families=[family],
                platform=args.platform,
                output_dir=tests_dir,
                download_cache_dir=download_cache_dir,
                run_github_repo=args.run_github_repo,
            )
            tests_tarball_name = (
                f"therock-dist-{args.platform}-{family}-tests-"
                f"{args.package_version}.tar.gz"
            )
            compress_tasks.append((tests_dir, args.output_dir / tests_tarball_name))

    # Phase 1.5: If KPACK_SPLIT_ARTIFACTS is enabled, fetch all families
    # into a single combined directory. With KPACK split, device-specific
    # files are per individual GPU target and don't conflict, so all
    # families can coexist in a single install prefix.
    kpack_split = is_kpack_split(family_dirs[0])
    if kpack_split:
        log("::: KPACK_SPLIT_ARTIFACTS detected — building multi-arch tarball")
        multiarch_dir = work_dir / "multiarch"
        fetch_and_flatten(
            run_id=args.run_id,
            amdgpu_families=families,
            platform=args.platform,
            output_dir=multiarch_dir,
            download_cache_dir=download_cache_dir,
            run_github_repo=args.run_github_repo,
            exclude_components=DEFAULT_EXCLUDED_COMPONENTS,
            exclude_artifacts=DEFAULT_EXCLUDED_ARTIFACTS,
        )
        tarball_name = (
            f"therock-dist-{args.platform}-multiarch-{args.package_version}.tar.gz"
        )
        compress_tasks.append((multiarch_dir, args.output_dir / tarball_name))
        if args.include_test_tarballs:
            tests_multiarch_dir = work_dir / "tests" / "multiarch"
            fetch_and_flatten(
                run_id=args.run_id,
                amdgpu_families=families,
                platform=args.platform,
                output_dir=tests_multiarch_dir,
                download_cache_dir=download_cache_dir,
                run_github_repo=args.run_github_repo,
            )
            tests_tarball_name = (
                f"therock-dist-{args.platform}-multiarch-tests-"
                f"{args.package_version}.tar.gz"
            )
            compress_tasks.append(
                (tests_multiarch_dir, args.output_dir / tests_tarball_name)
            )

    # Phase 2: Compress all tarballs in parallel.
    # Each tar cfz is single-threaded, so running N families concurrently
    # on a multi-core runner scales well with minimal per-job slowdown.
    # TODO: Add --compress-workers flag to cap concurrency on smaller runners.
    log(f"\nCompressing {len(compress_tasks)} tarballs in parallel...")
    with ProcessPoolExecutor(max_workers=len(compress_tasks)) as executor:
        futures = {
            executor.submit(compress_tarball, source_dir=src, tarball_path=dst): dst
            for src, dst in compress_tasks
        }
        for future in as_completed(futures):
            future.result()  # Raises on failure

    log(f"\nDone. Tarballs in {args.output_dir}:")
    for tb in sorted(args.output_dir.glob("*.tar.gz")):
        size_mb = tb.stat().st_size / (1024 * 1024)
        log(f"  {tb.name} ({size_mb:.1f} MB)")


if __name__ == "__main__":
    main()