rocm-systems/cmake/Templates/omnitrace-install.py.in

#!/usr/bin/env python3

import os
import re
import sys
import stat
import argparse
import tempfile
import subprocess as sp
from urllib import request
from urllib.error import HTTPError

omnitrace_version = "@OMNITRACE_VERSION@"
_rocm_path = os.environ.get("ROCM_PATH", "/opt/rocm")
_rocm_version = None


def get_rocm_version(rocm_hint):
    global _rocm_path
    global _rocm_version

    if rocm_hint is not None and rocm_hint is not True:
        if rocm_hint.replace(".", "0").isnumeric():
            _rocm_version = rocm_hint
        else:
            _rocm_path = rocm_hint

    def _parse_version(_v):
        return re.split(r"[\\.-]", _v) if _v is not None else None

    _version = _parse_version(_rocm_version)
    for fname in [
        "version",
        "version-dev",
        "version-hip-libraries",
        "version-hiprt",
        "version-hiprt-devel",
        "version-hip-sdk",
        "version-libs",
        "version-utils",
    ]:
        if _version is not None and len(_version) > 0:
            break
        _fname = os.path.join(_rocm_path, ".info", fname)
        if os.path.exists(_fname):
            with open(_fname, "r") as f:
                _version = _parse_version(f.readlines()[0].strip("\n"))

    if _version is not None and len(_version) > 0:
        _major = int(_version[0])
        _minor = int(_version[1]) if len(_version) >= 2 else 0
        _rocm_version = f"{_major}.{_minor}"
        return "-ROCm-{}".format((10000 * _major) + (100 * _minor))
    return None


def get_os_info(os_distrib, os_version):
    _os_info = {}
    with open("/etc/os-release", "r") as f:
        for line in [_v.strip() for _v in f.readlines()]:
            if "=" not in line:
                continue
            _key, _data = line.split("=", 1)
            _os_info[_key] = _data.strip('"')

    def _parse_version(_v):
        _version = re.split(r"[\\.-]", _v)
        return (
            "{}.{}".format(_version[0], _version[1])
            if len(_version) > 1
            else "{}".format(_version[0])
        )

    if os_distrib is None or os_distrib == "auto":
        if "ubuntu" in _os_info["ID"]:
            os_distrib = "ubuntu"
        elif "opensuse" in _os_info["ID"]:
            os_distrib = "opensuse"
        elif "rhel" in _os_info["ID"]:
            os_distrib = "rhel"
        elif "centos" in _os_info["ID"]:
            os_distrib = "rhel"
        elif "rockylinux" in _os_info["ID"]:
            os_distrib = "rhel"
        elif "debian" in _os_info["ID"]:
            os_distrib = "ubuntu"
            if "debian" in _os_info["ID"] and os_version is None:
                _debian_version = float(_parse_version(_os_info["VERSION_ID"]))
                if _debian_version >= 11.0:
                    os_version = "20.04"
                else:
                    os_version = "18.04"
        elif "fedora" in _os_info["ID"]:
            os_distrib = "rhel"
            # fedora has different versioning system so fallback to 8.7
            if os_version is None:
                os_version = "8.7"
        else:
            # if we don't have an exact match, check ID_LIKE
            if "ID_LIKE" not in _os_info.keys():
                _os_info["ID_LIKE"] = _os_info["ID"]

            if "debian" in _os_info["ID_LIKE"]:
                os_distrib = "ubuntu"
                if os_version is None:
                    # fallback on 18.04 if ID is not ubuntu but debian-like
                    os_version = "18.04"
            elif "suse" in _os_info["ID_LIKE"]:
                os_distrib = "opensuse"
                # fallback on 15.3 if ID is not opensuse but suse-like
                if os_version is None:
                    os_version = "15.3"
            elif "rhel" in _os_info["ID_LIKE"] or "centos" in _os_info["ID_LIKE"]:
                os_distrib = "rhel"
                if os_version is None:
                    os_version = "8.7"
            else:
                raise RuntimeError(
                    "Unknown ID_LIKE value in /etc/os-release: {}".format(
                        _os_info["ID_LIKE"]
                    )
                )
    elif os_distrib == "centos":
        os_distrib = "rhel"
        # uses same versioning system
    elif os_distrib == "fedora":
        os_distrib = "rhel"
        if os_version is None:
            # fedora has different versioning system so fallback to 8.7
            os_version = "8.7"

    if os_version is None:
        os_version = _parse_version(_os_info["VERSION_ID"])

    return (os_distrib, os_version)


def print_log(*args, **kwargs):
    sys.stdout.flush()
    sys.stderr.flush()
    sys.stderr.write("### ")
    sys.stderr.write(*args, **kwargs)
    sys.stderr.write("\n")
    sys.stderr.flush()


def run(*args, **kwargs):
    print_log("Executing: {}\n".format(" ".join(*args)))
    sp.run(*args, **kwargs, check=True)
    sys.stderr.write("\n")


if __name__ == "__main__":
    parser = argparse.ArgumentParser()

    parser.add_argument(
        "--version",
        help="Print omnitrace version which will be installed",
        action="store_true",
    )
    parser.add_argument(
        "-p", "--prefix", help="Installation prefix", type=str, default="/opt/omnitrace"
    )
    parser.add_argument(
        "-i",
        "--interactive",
        help="Prompt to accept the license and include/exclude subdirectory",
        action="store_true",
    )
    parser.add_argument(
        "-D",
        "--download-path",
        help="Download directory (default: temporary directory)",
        type=str,
        default=None,
    )
    parser.add_argument(
        "-d",
        "--os-distrib",
        help="Target OS distribution",
        type=str,
        default=None,
        choices=("auto", "ubuntu", "opensuse", "rhel", "centos", "fedora"),
    )
    parser.add_argument(
        "-v", "--os-version", help="Target OS version", type=str, default=None
    )
    parser.add_argument(
        "-k",
        "--keep-download",
        help="Do not delete downloaded file as installation",
        action="store_true",
    )
    parser.add_argument(
        "--rocm",
        help="Install omnitrace with ROCm support. Accepts either a ROCm version (e.g. '5.4') or the root path to the ROCm install containing .info/version* file(s) (e.g. /opt/rocm if /opt/rocm/.info/version exists). If no argument is provided, the ROCm version will attempted to be deduced from $ENV{ROCM_PATH}/.info/version",
        nargs="?",
        default=None,
        const=True,
        metavar="VERSION or ROCM_PATH with .info/version file(s)",
    )
    # right now, only valid set of extensions are: papi + ompt + python3
    # in the future, this might change, e.g. MPI variants
    parser.add_argument(
        "-e",
        "--extensions",
        help="Omnitrace extensions, e.g. PAPI, OMPT, and Python3",
        nargs="*",
        default=("papi", "ompt", "python3"),
        choices=("papi", "ompt", "python3"),
    )

    args = parser.parse_args()

    if args.version:
        print(f"omnitrace {omnitrace_version}")
        sys.exit(0)

    os_distrib, os_version = get_os_info(args.os_distrib, args.os_version)
    rocm_version = get_rocm_version(args.rocm) if args.rocm is not None else ""
    extensions = ""
    if "papi" in args.extensions:
        extensions += "-PAPI"
    if "ompt" in args.extensions:
        extensions += "-OMPT"
    if "python3" in args.extensions:
        extensions += "-Python3"

    if rocm_version is None:
        raise RuntimeError(
            f"Error! ROCm version could not be determined from {_rocm_path}/.info/version*. Please provide a ROCm version or the root path to the ROCm install containing the .info directory, e.g. '--rocm 5.4' or '--rocm /path/to/rocm/install'"
        )

    script = f"omnitrace-{omnitrace_version}-{os_distrib}-{os_version}{rocm_version}{extensions}.sh"
    url = f"https://github.com/AMDResearch/omnitrace/releases/download/v{omnitrace_version}/{script}"
    download_dir = (
        tempfile.mkdtemp(prefix="omnitrace-install-")
        if args.download_path is None
        else args.download_path
    )
    install_script = os.path.join(download_dir, script)

    try:
        if not os.path.exists(download_dir):
            print_log(f"Creating download directory: {download_dir} ...")
            os.makedirs(download_dir)

        print_log(f"Downloading {url} ...")

        try:
            response = request.urlretrieve(url, install_script)
        except HTTPError as e:
            print_log(f"")
            print_log(f"Error: {e}")
            print_log(f"")
            print_log(f"Error: Installer script download from {url} failed!")
            if args.rocm is not None:
                print_log(
                    f"There may not be a pre-built installer for ROCm version {_rocm_version}"
                )
            sys.exit(-1)

        if os.path.exists(install_script):
            print_log(f"Download completed: {install_script}")
        else:
            raise RuntimeError(f"Download completed but {install_script} does not exist")

        os.chmod(install_script, stat.S_IRWXU)

        if not os.path.exists(args.prefix):
            print_log(f"Creating directory: {args.prefix} ...")
            os.makedirs(args.prefix)

        install_args = (
            ["--exclude-subdir", "--skip-license"] if not args.interactive else []
        )

        print_log(f"Installing omnitrace to {args.prefix} ...")

        run([install_script, f"--prefix={args.prefix}"] + install_args)

        print_log(
            f"omnitrace v{omnitrace_version} installation to {args.prefix} succeeded!"
        )

    finally:
        if not args.keep_download:
            print_log(f"Removing install script {install_script} ...")
            os.remove(install_script)
            # remove the directory if it is a temporary directory
            if args.download_path is None:
                print_log(f"Removing temporary directory {download_dir} ...")
                os.rmdir(download_dir)