install_rocm_from_artifacts.py

#!/usr/bin/env python
"""install_rocm_from_artifacts.py

This script helps CI workflows, developers and testing suites easily install
TheRock to their environment using artifacts. It installs TheRock to an output
directory from one of these sources:

  - GitHub CI workflow run
  - Release tag
  - An existing installation of TheRock

Usage:
python build_tools/install_rocm_from_artifacts.py
    (--artifact-group ARTIFACT_GROUP | --amdgpu_family AMDGPU_FAMILY)
    [--output-dir OUTPUT_DIR]
    (--run-id RUN_ID | --release RELEASE | --latest-release | --input-dir INPUT_DIR)
    [--dry-run]
    [--run-github-repo RUN_GITHUB_REPO]
    [--aqlprofile | --no-aqlprofile]
    [--blas | --no-blas]
    [--debug-tools | --no-debug-tools]
    [--fft | --no-fft]
    [--hipdnn | --no-hipdnn]
    [--hipdnn-samples | --no-hipdnn-samples]
    [--miopen | --no-miopen]
    [--miopen-plugin | --no-miopen-plugin]
    [--fusilli-plugin | --no-fusilli-plugin]
    [--hipblaslt-plugin | --no-hipblaslt-plugin]
    [--prim | --no-prim]
    [--rand | --no-rand]
    [--rccl | --no-rccl]
    [--rocdecode | --no-rocdecode]
    [--rocjpeg | --no-rocjpeg]
    [--rocprofiler-compute | --no-rocprofiler-compute]
    [--rocprofiler-sdk | --no-rocprofiler-sdk ]
    [--rocprofiler-systems | --no-rocprofiler-systems]
    [--rocrtst | --no-rocrtst]
    [--rocwmma | --no-rocwmma]
    [--libhipcxx | --no-libhipcxx]
    [--tests | --no-tests]
    [--base-only]

Examples:
- Downloads and unpacks the gfx94X S3 artifacts from GitHub CI workflow run 14474448215
  (from https://github.com/ROCm/TheRock/actions/runs/14474448215) to the
  default output directory `therock-build`:
    ```
    python build_tools/install_rocm_from_artifacts.py \
        --run-id 14474448215 \
        --amdgpu-family gfx94X-dcgpu \
        --tests
    ```
- Downloads and unpacks the version `6.4.0rc20250416` gfx110X artifacts from
  release tag `nightly-tarball` to the specified output directory `build`:
    ```
    python build_tools/install_rocm_from_artifacts.py \
        --release 6.4.0rc20250416 \
        --amdgpu-family gfx110X-all \
        --output-dir build
    ```
- Downloads and unpacks the version `6.4.0.dev0+8f6cdfc0d95845f4ca5a46de59d58894972a29a9`
  gfx120X artifacts from release tag `dev-tarball` to the default output directory `therock-build`:
    ```
    python build_tools/install_rocm_from_artifacts.py \
        --release 6.4.0.dev0+8f6cdfc0d95845f4ca5a46de59d58894972a29a9 \
        --amdgpu-family gfx120X-all
    ```
- Downloads and unpacks the gfx94X S3 artifacts from GitHub CI workflow run 19644138192
  (from https://github.com/ROCm/rocm-libraries/actions/runs/19644138192) in the `ROCm/rocm-libraries` repository to the
  default output directory `therock-build`:
    ```
    python build_tools/install_rocm_from_artifacts.py \
        --run-id 19644138192 \
        --amdgpu-family gfx94X-dcgpu \
        --tests \
        --run-github-repo ROCm/rocm-libraries
    ```
- Downloads and unpacks the latest nightly release for gfx110X:
    ```
    python build_tools/install_rocm_from_artifacts.py \
        --latest-release \
        --amdgpu-family gfx110X-all
    ```
- Shows what would be downloaded without actually downloading (works with any mode):
    ```
    python build_tools/install_rocm_from_artifacts.py \
        --latest-release \
        --amdgpu-family gfx110X-all \
        --dry-run

    python build_tools/install_rocm_from_artifacts.py \
        --release 7.11.0a20260119 \
        --amdgpu-family gfx110X-all \
        --dry-run
    ```
You can select your AMD GPU family from therock_amdgpu_targets.cmake.

By default for CI workflow retrieval, all artifacts (excluding test artifacts)
will be downloaded. For specific artifacts, pass in the flag such as `--rand`
(RAND artifacts) For test artifacts, pass in the flag `--tests` (test artifacts).
For base artifacts only, pass in the flag `--base-only`

Note that the ARTIFACT_GROUP controls which sub-directory of the run contains
the artifacts. If not specified, it defaults to the AMDGPU_FAMILY, which was
the historic interpretation.

Note: the script will overwrite the output directory argument. If no argument
is passed, it will overwrite the default "therock-build" directory.
"""

import argparse
import boto3
from botocore import UNSIGNED
from botocore.config import Config
from datetime import datetime
from fetch_artifacts import main as fetch_artifacts_main
from pathlib import Path
import platform
import re
import shutil
import subprocess
import sys
import tarfile
from typing import Optional

PLATFORM = platform.system().lower()
s3_client = boto3.client(
    "s3",
    verify=False,
    config=Config(max_pool_connections=100, signature_version=UNSIGNED),
)
# S3 bucket names for TheRock releases.
# NOTE: These buckets will be restricted to CloudFront-only access in the future.
# When that happens, direct S3 API calls (list_objects, download_fileobj) will fail
# and this script will need to be updated to use CloudFront URLs instead.
NIGHTLY_BUCKET_NAME = "therock-nightly-tarball"
DEV_BUCKET_NAME = "therock-dev-tarball"


def parse_nightly_version(version: str) -> Optional[datetime]:
    """
    Parse nightly version like '7.11.0a20251124' to extract date.
    Returns datetime for sorting, None if not parseable.
    """
    match = re.search(r"(\d+)\.(\d+)\.(\d+)(a|rc)(\d{4})(\d{2})(\d{2})", version)
    if match:
        year, month, day = int(match.group(5)), int(match.group(6)), int(match.group(7))
        return datetime(year, month, day)
    return None


def extract_version_from_asset_name(
    asset_name: str, artifact_group: str, platform_str: str
) -> Optional[str]:
    """
    Extract version string from asset name.
    E.g., 'therock-dist-linux-gfx110X-all-7.11.0a20251124.tar.gz' -> '7.11.0a20251124'
    """
    prefix = f"therock-dist-{platform_str}-{artifact_group}-"
    suffix = ".tar.gz"
    if asset_name.startswith(prefix) and asset_name.endswith(suffix):
        return asset_name[len(prefix) : -len(suffix)]
    return None


def list_available_nightly_gpu_families(platform_str: str = PLATFORM) -> set[str]:
    """
    Query S3 to find all GPU families that have nightly releases.
    Useful for error messages when an invalid GPU family is specified.
    """
    prefix = f"therock-dist-{platform_str}-"

    paginator = s3_client.get_paginator("list_objects_v2")
    families: set[str] = set()

    for page in paginator.paginate(Bucket=NIGHTLY_BUCKET_NAME, Prefix=prefix):
        for obj in page.get("Contents", []):
            # Extract family from: therock-dist-linux-{family}-{version}.tar.gz
            match = re.match(rf"{prefix}([\w-]+)-", obj["Key"])
            if match:
                families.add(match.group(1))

    return families


def _fetch_and_sort_nightly_releases(
    artifact_group: str,
    platform_str: str = PLATFORM,
) -> list[dict]:
    """
    Fetch and sort nightly releases from S3 bucket for a given artifact group.

    Returns:
        List of dicts with keys: version, asset_name, last_modified, size, parsed_date
        Sorted by recency (newest first).
    """
    prefix = f"therock-dist-{platform_str}-{artifact_group}-"

    paginator = s3_client.get_paginator("list_objects_v2")
    releases: list[dict] = []

    for page in paginator.paginate(Bucket=NIGHTLY_BUCKET_NAME, Prefix=prefix):
        for obj in page.get("Contents", []):
            key = obj["Key"]
            if not key.endswith(".tar.gz"):
                continue
            version = extract_version_from_asset_name(key, artifact_group, platform_str)
            if version:
                releases.append(
                    {
                        "version": version,
                        "asset_name": key,
                        "last_modified": obj["LastModified"],
                        "size": obj["Size"],
                        "parsed_date": parse_nightly_version(version),
                    }
                )

    # Sort by parsed date (newest first), falling back to last_modified
    releases.sort(
        key=lambda x: (
            x["parsed_date"] if x["parsed_date"] else datetime.min,
            x["last_modified"],
        ),
        reverse=True,
    )

    return releases


def discover_latest_release(
    artifact_group: str,
    platform_str: str = PLATFORM,
) -> Optional[tuple[str, str]]:
    """
    Query S3 bucket to find the latest nightly release for given artifact group.

    Returns:
        Tuple of (version_string, full_asset_name) or None if not found.
    """
    releases = _fetch_and_sort_nightly_releases(artifact_group, platform_str)
    if not releases:
        return None
    return (releases[0]["version"], releases[0]["asset_name"])


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


def _untar_files(output_dir: Path, destination: Path):
    """
    Retrieves all tar files in the output_dir, then extracts all files to the output_dir
    """
    log(f"Extracting {destination.name} to {str(output_dir)}")
    with tarfile.open(destination) as extracted_tar_file:
        extracted_tar_file.extractall(output_dir)
    destination.unlink()


def _create_output_directory(output_dir: Path):
    """
    If the output directory already exists, delete it and its contents.
    Then, create the output directory.
    """
    log(f"Creating output directory '{output_dir.resolve()}'")
    if output_dir.is_dir():
        log(
            f"Directory '{output_dir}' already exists, removing existing directory and files"
        )
        shutil.rmtree(output_dir)
    output_dir.mkdir(parents=True)
    log(f"Created output directory '{output_dir.resolve()}'")


def _retrieve_s3_release_assets(
    release_bucket, artifact_group, release_version, output_dir
):
    """
    Makes an API call to retrieve the release's assets, then retrieves the asset matching the amdgpu family
    """
    asset_name = f"therock-dist-{PLATFORM}-{artifact_group}-{release_version}.tar.gz"
    destination = output_dir / asset_name

    with open(destination, "wb") as f:
        s3_client.download_fileobj(release_bucket, asset_name, f)

    # After downloading the asset, untar-ing the file
    _untar_files(output_dir, destination)


def retrieve_artifacts_by_run_id(args):
    """
    If the user requested TheRock artifacts by CI (run ID), this function will retrieve those assets
    """
    run_id = args.run_id
    log(f"Retrieving artifacts for run ID {run_id}")
    argv = [
        "--run-id",
        run_id,
        "--artifact-group",
        args.artifact_group,
        "--output-dir",
        str(args.output_dir),
        "--flatten",
    ]
    if args.amdgpu_targets:
        argv.extend(["--amdgpu-targets", args.amdgpu_targets])
    if args.dry_run:
        argv.append("--dry-run")
    if args.run_github_repo:
        argv.extend(["--run-github-repo", args.run_github_repo])

    # These artifacts are the "base" requirements for running tests.
    base_artifact_patterns = [
        "core-hipinfo_run",
        "core-runtime_run",
        "core-runtime_lib",
        "sysdeps_lib",
        "base_run",
        "base_lib",
        "amd-llvm_run",
        "amd-llvm_lib",
        "core-amdsmi_run",
        "core-amdsmi_lib",
        "core-hip_lib",
        "core-hip_dev",
        "core-ocl_lib",
        "core-ocl_dev",
        "rocprofiler-sdk_lib",
        "host-suite-sparse_lib",
    ]

    if args.base_only:
        argv.extend(base_artifact_patterns)
    elif any(
        [
            args.aqlprofile,
            args.blas,
            args.debug_tools,
            args.fft,
            args.hipdnn,
            args.hipdnn_samples,
            args.miopen,
            args.miopen_plugin,
            args.fusilli_plugin,
            args.iree_compiler,
            args.hipblaslt_plugin,
            args.prim,
            args.rand,
            args.rccl,
            args.rocdecode,
            args.rocjpeg,
            args.rocprofiler_compute,
            args.rocprofiler_sdk,
            args.rocprofiler_systems,
            args.rocrtst,
            args.rocwmma,
            args.libhipcxx,
        ]
    ):
        argv.extend(base_artifact_patterns)

        extra_artifacts = []
        if args.aqlprofile:
            extra_artifacts.append("aqlprofile-tests")
        if args.blas:
            extra_artifacts.append("blas")
        if args.debug_tools:
            # Add extra artifacts so we generate _lib and _test artifact
            # entries later.
            extra_artifacts.append("amd-dbgapi")
            extra_artifacts.append("rocgdb")
            extra_artifacts.append("rocr-debug-agent")
            extra_artifacts.append("rocr-debug-agent-tests")

            # Add the rest of the artifacts not handled automatically (non-lib
            # and non-test).
            argv.append("rocgdb_run")

            # Libraries rocgdb depends on.
            extra_artifacts.append("gmp")
            extra_artifacts.append("mpfr")
            extra_artifacts.append("expat")
            extra_artifacts.append("ncurses")
        if args.fft:
            extra_artifacts.append("fft")
            extra_artifacts.append("fftw3")
        if args.hipdnn:
            extra_artifacts.append("hipdnn")
        if args.hipdnn_samples:
            extra_artifacts.append("hipdnn-samples")
        if args.miopen:
            extra_artifacts.append("miopen")
            # We need bin/MIOpenDriver executable for tests.
            argv.append("miopen_run")
            # Also need these for runtime kernel compilation (rocrand includes).
            argv.append("rand_dev")
        if args.miopen_plugin:
            extra_artifacts.append("miopen-plugin")
        if args.fusilli_plugin:
            extra_artifacts.append("fusilli-plugin")
        if args.iree_compiler:
            extra_artifacts.append("iree-compiler")
        if args.rocdecode:
            extra_artifacts.append("sysdeps-amd-mesa")
            extra_artifacts.append("rocdecode")
            argv.append("rocdecode_dev")
            argv.append("rocdecode_test")
            argv.append("base_dev")
            argv.append("amd-llvm_dev")
        if args.rocjpeg:
            extra_artifacts.append("sysdeps-amd-mesa")
            extra_artifacts.append("rocjpeg")
            argv.append("rocjpeg_dev")
            argv.append("rocjpeg_test")
            argv.append("base_dev")
            argv.append("amd-llvm_dev")
        if args.hipblaslt_plugin:
            extra_artifacts.append("hipblaslt-plugin")
        if args.prim:
            extra_artifacts.append("prim")
        if args.rand:
            extra_artifacts.append("rand")
        if args.rccl:
            extra_artifacts.append("rccl")
        if args.rocprofiler_sdk:
            extra_artifacts.append("rocprofiler-sdk")
            # Contains rocprofiler-sdk-rocpd
            argv.append("rocprofiler-sdk_run")
        if args.rocprofiler_compute:
            extra_artifacts.append("rocprofiler-compute")
        if args.rocprofiler_systems:
            extra_artifacts.append("rocprofiler-systems")
            # Contains executables (rocprof-sys-run, rocprof-sys-instrument, etc.)
            argv.append("rocprofiler-systems_run")
        if args.rocrtst:
            extra_artifacts.append("rocrtst")
            # rocrtst depends on sysdeps-hwloc (which depends on sysdeps-libpciaccess)
            extra_artifacts.append("sysdeps-hwloc")
            extra_artifacts.append("sysdeps-libpciaccess")
        if args.rocwmma:
            extra_artifacts.append("rocwmma")
            argv.append("rocwmma_dev")
        if args.libhipcxx:
            extra_artifacts.append("libhipcxx")
            argv.append("amd-llvm_dev")
            argv.append("amd-llvm_lib")
            argv.append("base_dev_generic")

        extra_artifact_patterns = [f"{a}_lib" for a in extra_artifacts]
        if args.tests:
            extra_artifact_patterns.extend([f"{a}_test" for a in extra_artifacts])

        argv.extend(extra_artifact_patterns)
    else:
        # No include (or exclude) patterns, so all artifacts will be fetched.
        pass

    log(f"\nCalling fetch_artifacts_main with args:\n  {' '.join(argv)}\n")
    fetch_artifacts_main(argv)

    log(f"Retrieved artifacts for run ID {run_id}")


def retrieve_artifacts_by_release(args):
    """
    If the user requested TheRock artifacts by release version, this function will retrieve those assets
    """
    output_dir = args.output_dir
    artifact_group = args.artifact_group
    # Determine if version is nightly-tarball or dev-tarball
    nightly_regex_expression = (
        "(\\d+\\.)?(\\d+\\.)?(\\*|\\d+)(a|rc)(\\d{4})(\\d{2})(\\d{2})"
    )
    dev_regex_expression = "(\\d+\\.)?(\\d+\\.)?(\\*|\\d+).dev0+"
    nightly_release = re.search(nightly_regex_expression, args.release) != None
    dev_release = re.search(dev_regex_expression, args.release) != None
    if not nightly_release and not dev_release:
        log("This script requires a nightly-tarball or dev-tarball version.")
        log("Please retrieve the correct release version from:")
        log(
            "\t - https://therock-nightly-tarball.s3.amazonaws.com/ (nightly-tarball examples: 6.4.0rc20250416, 7.10.0a20251024)"
        )
        log(
            "\t - https://therock-dev-tarball.s3.amazonaws.com/ (dev-tarball example: 6.4.0.dev0+8f6cdfc0d95845f4ca5a46de59d58894972a29a9)"
        )
        log("Exiting...")
        return

    release_bucket = NIGHTLY_BUCKET_NAME if nightly_release else DEV_BUCKET_NAME
    release_version = args.release

    log(f"Retrieving artifacts from release bucket {release_bucket}")

    if args.dry_run:
        asset_name = (
            f"therock-dist-{PLATFORM}-{artifact_group}-{release_version}.tar.gz"
        )
        log(f"[DRY RUN] Would download: {asset_name} (version {release_version})")
        return

    _retrieve_s3_release_assets(
        release_bucket, artifact_group, release_version, output_dir
    )


def retrieve_artifacts_by_input_dir(args):
    input_dir = args.input_dir
    output_dir = args.output_dir
    log(f"Retrieving artifacts from input dir {input_dir}")

    if args.dry_run:
        log(f"[DRY RUN] Would rsync from {input_dir} to {output_dir}")
        return

    # Check to make sure rsync exists
    if not shutil.which("rsync"):
        log("Error: rsync command not found.")
        if platform.system() == "Windows":
            log("Please install rsync via MSYS2 or WSL to your Windows system")
        return

    cmd = [
        "rsync",
        "-azP",  # archive, compress and progress indicator
        input_dir,
        output_dir,
    ]
    try:
        subprocess.run(cmd, check=True)
        log(f"Retrieved artifacts from input dir {input_dir} to {output_dir}")
    except Exception as ex:
        # rsync is not available
        log(f"Error when running [{cmd}]")
        log(str(ex))


def retrieve_artifacts_by_latest_release(args):
    """
    Find and retrieve the latest nightly release from S3.
    """
    log(f"Finding latest nightly release for {args.artifact_group}...")

    result = discover_latest_release(artifact_group=args.artifact_group)

    if result is None:
        log(f"ERROR: No nightly release found for '{args.artifact_group}'")
        log("")
        log("Available GPU families in the nightly bucket:")
        available = list_available_nightly_gpu_families()
        for family in sorted(available):
            log(f"  - {family}")
        sys.exit(1)

    version, asset_name = result
    log(f"Found latest release: {version}")

    if args.dry_run:
        log(f"[DRY RUN] Would download: {asset_name} (version {version})")
        return

    # Reuse existing download logic
    _retrieve_s3_release_assets(
        release_bucket=NIGHTLY_BUCKET_NAME,
        artifact_group=args.artifact_group,
        release_version=version,
        output_dir=args.output_dir,
    )


def run(args):
    log("### Installing TheRock using artifacts ###")

    # Skip directory creation for dry-run
    if not args.dry_run:
        _create_output_directory(args.output_dir)

    if args.run_id:
        retrieve_artifacts_by_run_id(args)
    elif args.release:
        retrieve_artifacts_by_release(args)
    elif args.latest_release:
        retrieve_artifacts_by_latest_release(args)

    if args.input_dir:
        retrieve_artifacts_by_input_dir(args)


def main(argv):
    parser = argparse.ArgumentParser(prog="provision")
    parser.add_argument(
        "--output-dir",
        type=Path,
        default="./therock-build",
        help="Path of the output directory for TheRock",
    )

    artifact_group_parser = parser.add_mutually_exclusive_group(required=True)
    artifact_group_parser.add_argument(
        "--artifact-group",
        dest="artifact_group",
        type=str,
        help="Explicit artifact group to install",
    )
    artifact_group_parser.add_argument(
        "--amdgpu-family",
        dest="artifact_group",
        type=str,
        help="AMD GPU family to install (please refer to this: https://github.com/ROCm/TheRock/blob/59c324a759e8ccdfe5a56e0ebe72a13ffbc04c1f/cmake/therock_amdgpu_targets.cmake#L44-L81 for family choices)",
    )

    # This mutually exclusive group will ensure that only one argument is present
    group = parser.add_mutually_exclusive_group(required=True)
    group.add_argument("--run-id", type=str, help="GitHub run ID of TheRock to install")

    group.add_argument(
        "--release",
        type=str,
        help="Release version of TheRock to install, from the nightly-tarball (X.Y.ZrcYYYYMMDD) or dev-tarball (X.Y.Z.dev0+{hash})",
    )

    group.add_argument(
        "--latest-release",
        action="store_true",
        help="Install the latest nightly release (built daily from main branch)",
    )

    parser.add_argument(
        "--dry-run",
        action="store_true",
        help="Show what would be downloaded/copied without actually doing it",
    )

    artifacts_group = parser.add_argument_group("artifacts_group")
    artifacts_group.add_argument(
        "--aqlprofile",
        default=False,
        help="Include 'aqlprofile' artifacts",
        action=argparse.BooleanOptionalAction,
    )

    artifacts_group.add_argument(
        "--blas",
        default=False,
        help="Include 'blas' artifacts",
        action=argparse.BooleanOptionalAction,
    )

    artifacts_group.add_argument(
        "--debug-tools",
        default=False,
        help="Include ROCm debugging tools (amd-dbgapi, rocgdb and rocr_debug_agent) artifacts",
        action=argparse.BooleanOptionalAction,
    )

    artifacts_group.add_argument(
        "--fft",
        default=False,
        help="Include 'fft' artifacts",
        action=argparse.BooleanOptionalAction,
    )

    artifacts_group.add_argument(
        "--hipdnn",
        default=False,
        help="Include 'hipdnn' artifacts",
        action=argparse.BooleanOptionalAction,
    )

    artifacts_group.add_argument(
        "--hipdnn-samples",
        default=False,
        help="Include 'hipdnn-samples' artifacts",
        action=argparse.BooleanOptionalAction,
    )

    artifacts_group.add_argument(
        "--miopen",
        default=False,
        help="Include 'miopen' artifacts",
        action=argparse.BooleanOptionalAction,
    )

    artifacts_group.add_argument(
        "--miopen-plugin",
        default=False,
        help="Include 'miopen-plugin' artifacts",
        action=argparse.BooleanOptionalAction,
    )

    artifacts_group.add_argument(
        "--fusilli-plugin",
        default=False,
        help="Include 'fusilli-plugin' artifacts",
        action=argparse.BooleanOptionalAction,
    )

    artifacts_group.add_argument(
        "--iree-compiler",
        default=False,
        help="Include 'iree-compiler' artifacts",
        action=argparse.BooleanOptionalAction,
    )

    artifacts_group.add_argument(
        "--rocdecode",
        default=False,
        help="Include 'rocdecode' artifacts",
        action=argparse.BooleanOptionalAction,
    )

    artifacts_group.add_argument(
        "--rocjpeg",
        default=False,
        help="Include 'rocjpeg' artifacts",
        action=argparse.BooleanOptionalAction,
    )

    artifacts_group.add_argument(
        "--hipblaslt-plugin",
        default=False,
        help="Include 'hipblaslt-plugin' artifacts",
        action=argparse.BooleanOptionalAction,
    )

    artifacts_group.add_argument(
        "--prim",
        default=False,
        help="Include 'prim' artifacts",
        action=argparse.BooleanOptionalAction,
    )

    artifacts_group.add_argument(
        "--rand",
        default=False,
        help="Include 'rand' artifacts",
        action=argparse.BooleanOptionalAction,
    )

    artifacts_group.add_argument(
        "--rccl",
        default=False,
        help="Include 'rccl' artifacts",
        action=argparse.BooleanOptionalAction,
    )

    artifacts_group.add_argument(
        "--rocprofiler-compute",
        default=False,
        help="Include 'rocprofiler-compute' artifacts",
        action=argparse.BooleanOptionalAction,
    )

    artifacts_group.add_argument(
        "--rocprofiler-sdk",
        default=False,
        help="Include 'rocprofiler-sdk' artifacts",
        action=argparse.BooleanOptionalAction,
    )

    artifacts_group.add_argument(
        "--rocprofiler-systems",
        default=False,
        help="Include 'rocprofiler-systems' artifacts",
        action=argparse.BooleanOptionalAction,
    )

    artifacts_group.add_argument(
        "--rocrtst",
        default=False,
        help="Include 'rocrtst' artifacts",
        action=argparse.BooleanOptionalAction,
    )

    artifacts_group.add_argument(
        "--rocwmma",
        default=False,
        help="Include 'rocwmma' artifacts",
        action=argparse.BooleanOptionalAction,
    )

    artifacts_group.add_argument(
        "--libhipcxx",
        default=False,
        help="Include 'libhipcxx' artifacts",
        action=argparse.BooleanOptionalAction,
    )

    artifacts_group.add_argument(
        "--tests",
        default=False,
        help="Include all test artifacts for enabled libraries",
        action=argparse.BooleanOptionalAction,
    )

    artifacts_group.add_argument(
        "--base-only", help="Include only base artifacts", action="store_true"
    )

    group.add_argument(
        "--input-dir",
        type=str,
        help="Pass in an existing directory of TheRock to provision and test",
    )

    parser.add_argument(
        "--amdgpu-targets",
        type=str,
        default="",
        help="Comma-separated individual GPU targets for fetching split artifacts (e.g. 'gfx942')",
    )

    parser.add_argument(
        "--run-github-repo",
        type=str,
        help="GitHub repository for --run-id in 'owner/repo' format (e.g. 'ROCm/TheRock'). Defaults to GITHUB_REPOSITORY env var or 'ROCm/TheRock'",
    )

    args = parser.parse_args(argv)

    if not args.artifact_group:
        raise argparse.ArgumentTypeError(
            "Either --amdgpu-family or --artifact-group must be specified"
        )

    run(args)


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