TheRock/rocm-systems/projects/rocprofiler-systems/tests/pytest/rocprofsys/config.py

# Copyright (c) Advanced Micro Devices, Inc.
# SPDX-License-Identifier: MIT

from __future__ import annotations
from dataclasses import dataclass, field
import os
from pathlib import Path
import shutil
import subprocess
import tempfile
from typing import Optional
import re

from .cache import persistent_cache, persistent_cached_property, resolve_username
from .capabilities import SystemCapabilities


@dataclass
class RocprofsysConfig:
    """Configuration for rocprofiler-systems test execution

    Contains necessary paths to configure tests for build or for install modes.

        Attributes:
        - rocprofsys_build_dir: Path to either the build or install directory
        - rocprofsys_instrument: Path to rocprof-sys-instrument executable
        - rocprofsys_run: Path to rocprof-sys-run executable
        - rocprofsys_sample: Path to rocprof-sys-sample executable
        - rocprofsys_causal: Path to rocprof-sys-causal executable
        - rocprofsys_avail: Path to rocprof-sys-avail executable
        - rocm_path: Path to ROCm installation directory
        - rocprofsys_lib_dir: Path to rocprofsys library directory
        - rocprofsys_bin_dir: Path to rocprofsys binary directory
        - rocprofsys_examples_dir: Path
            In build mode, this is the root of the build directory (most examples lie here).
            In install mode, this is the examples/ directory.
        - rocprofsys_tests_dir: Path to rocprofsys tests directory
        - test_output_dir: Path to test output directory
        - rocpd_validation_rules: Path to rocprofiler-systems rocpd validation rules directory
        - rocm_version: Tuple of (major, minor, patch) of the installed ROCm version
        - is_installed: Whether this is an installed configuration
        - rocprofsys_python: Path to rocprof-sys-python executable
        - rocprofsys_site_packages: Path to rocprofsys site-packages directory
        - _python_versions_hint: List of python versions available (hint for pytest_generate_tests)
        - _python_root_dirs_hint: List of python root directories available (hint for pytest_generate_tests)
        - Capabilities: SystemCapabilities instance
    """

    rocprofsys_build_dir: Path
    rocprofsys_instrument: Path
    rocprofsys_run: Path
    rocprofsys_sample: Path
    rocprofsys_causal: Path
    rocprofsys_avail: Path
    rocm_path: Optional[Path]
    rocprofsys_lib_dir: Path
    rocprofsys_bin_dir: Path
    rocprofsys_examples_dir: Path
    rocprofsys_tests_dir: Path
    rocpd_validation_rules: Path
    test_output_dir: Path
    rocprofsys_python: Optional[Path] = None
    rocprofsys_site_packages: Optional[Path] = None
    is_installed: bool = False
    rocm_version: Optional[tuple[int, int, int]] = None
    _python_versions_hint: Optional[list[str]] = field(default=None, repr=False)
    _python_root_dirs_hint: Optional[list[Path]] = field(default=None, repr=False)
    _capabilities: Optional[SystemCapabilities] = field(
        default=None, init=False, repr=False
    )

    @property
    def capabilities(self) -> SystemCapabilities:
        """Lazy-initialized system capabilities"""
        if self._capabilities is None:
            self._capabilities = SystemCapabilities.from_config(self)
        return self._capabilities

    @persistent_cached_property
    def llvm_lib_paths(self) -> list[Path]:
        """Get list of found ROCm LLVM lib paths.

        Returns:
            List of existing LLVM lib paths found, empty list if none found.
        """
        found_paths = []
        if self.rocm_path:
            # Match discover_llvm_libdir_for_ompt() logic
            candidates = [
                self.rocm_path / "llvm" / "lib",
                self.rocm_path / "lib" / "llvm" / "lib",
            ]
            # Determine the per-target runtime subdir used when
            # LLVM_ENABLE_PER_TARGET_RUNTIME_DIR=ON
            clang_candidates = [
                self.rocm_path / "bin" / "amdclang",
                self.rocm_path / "llvm" / "bin" / "amdclang",
                self.rocm_path / "lib" / "llvm" / "bin" / "amdclang",
            ]
            clang_path = next((c for c in clang_candidates if c.exists()), None)
            if clang_path is None:
                which_clang = shutil.which("amdclang")
                if which_clang:
                    clang_path = Path(which_clang)
            if clang_path is not None:
                try:
                    host_triple = subprocess.run(
                        [str(clang_path), "--print-target-triple"],
                        capture_output=True,
                        text=True,
                        timeout=10,
                        check=True,
                    ).stdout.strip()
                except (subprocess.SubprocessError, OSError) as exc:
                    raise RuntimeError(
                        f"'{clang_path} --print-target-triple' failed; "
                        "this suggests a broken toolchain."
                    ) from exc
                if host_triple:
                    candidates.append(
                        self.rocm_path / "lib" / "llvm" / "lib" / host_triple
                    )
            for candidate in candidates:
                if candidate.exists():
                    found_paths.append(candidate)
        return found_paths

    def get_library_path(self) -> str:
        """Get LD_LIBRARY_PATH including rocprofiler-systems libraries.

        Returns:
            LD_LIBRARY_PATH string with rocprofiler-systems libraries
        """
        paths = [str(self.rocprofsys_lib_dir.resolve())]

        # Where libraries for the examples live
        if self.is_installed:
            examples_lib = self.rocprofsys_examples_dir / "lib"
            if examples_lib.is_dir():
                paths.append(str(examples_lib.resolve()))

        existing = os.environ.get("LD_LIBRARY_PATH", "")
        if existing:
            paths.append(existing)

        # Add ROCm LLVM lib as fallback
        for llvm_path in self.llvm_lib_paths:
            paths.append(str(llvm_path))

        return ":".join(paths)

    def get_preload_path(self) -> Optional[str]:
        """Get LD_PRELOAD for sanitizer builds, or None when not applicable.

        When $ASAN_LIBRARY is set (exported by the sanitizer CI) it is prepended
        to the shell's LD_PRELOAD so the asan runtime stays first in the link
        order once rocprof-sys-run appends librocprof-sys-dl.so.

        Returns:
            LD_PRELOAD string, or None when not building with the sanitizer.
        """
        asan_library = os.environ.get("ASAN_LIBRARY")
        if not asan_library:
            return None
        existing = os.environ.get("LD_PRELOAD", "")
        return f"{asan_library}:{existing}" if existing else asan_library

    def get_target_executable(self, name: str) -> Path:
        """Get path to a test target executable.

        The executable is always resolved from the build/install tree; the
        system PATH is never consulted, so the suite cannot silently pick up an
        unrelated copy (e.g. one under /opt/rocm/bin).

        When is_installed is True, searches in the following order:
        1. rocprofsys_examples_dir/name
        2. rocprofsys_bin_dir/name

        When is_installed is False, searches in the following order:
        1. rocprofsys_examples_dir/examples/python/name (python scripts)
        2. rocprofsys_examples_dir/name
        3. rocprofsys_examples_dir/examples/code-coverage/name (code-coverage.py only)
        4. rocprofsys_bin_dir/name

        Args:
            name: Name of the target executable

        Returns:
            Path to the executable

        Raises:
            FileNotFoundError: If the executable is not found
        """

        if self.is_installed:
            # examples directory layout
            exe = self.rocprofsys_examples_dir / name
            if exe.exists() and exe.is_file():
                return exe

            # binary directory
            exe = self.rocprofsys_bin_dir / name
            if exe.exists() and exe.is_file():
                return exe

            raise FileNotFoundError(
                f"Target executable '{name}' not found. Searched in:\n"
                f"  - {self.rocprofsys_examples_dir}/{name}\n"
                f"  - {self.rocprofsys_bin_dir}/{name}\n"
            )

        else:
            # Python check
            exe = self.rocprofsys_examples_dir / "examples" / "python" / name
            if exe.exists() and exe.is_file():
                return exe

            # examples directory layout
            exe = self.rocprofsys_examples_dir / name
            if exe.exists() and exe.is_file():
                return exe

            # code-coverage.py lies in the code-coverage directory
            if name == "code-coverage.py":
                exe = self.rocprofsys_examples_dir / "examples" / "code-coverage" / name
                if exe.exists() and exe.is_file():
                    return exe

            # binary directory
            exe = self.rocprofsys_bin_dir / name
            if exe.exists() and exe.is_file():
                return exe

            raise FileNotFoundError(
                f"Target executable '{name}' not found. Searched in:\n"
                f"  - {self.rocprofsys_examples_dir}/examples/python/{name}\n"
                f"  - {self.rocprofsys_examples_dir}/{name}\n"
                f"  - {self.rocprofsys_examples_dir}/examples/code-coverage/{name}\n"
                f"  - {self.rocprofsys_bin_dir}/{name}\n"
            )


def _find_rocm_path(optional: bool = False) -> Optional[Path]:
    """Find ROCm installation path.

    Args:
        optional: If True, return None instead of raising when ROCm is not found.
    """
    for candidate in [
        os.environ.get("ROCM_PATH"),
        "/opt/rocm",
        "/usr/local/rocm",
    ]:
        if candidate and Path(candidate).exists():
            return Path(candidate).resolve()
    if optional:
        return None
    raise FileNotFoundError(
        "Could not find ROCm installation. Set ROCM_PATH environment variable."
    )


def _get_rocm_version(rocm_optional: bool = False) -> Optional[tuple[int, int, int]]:
    """Get the installed ROCm version as a tuple (major, minor, patch).

    Returns:
        Tuple of (major, minor, patch) or None if ROCm not found or version undetectable.
    """
    rocm_path = _find_rocm_path(optional=rocm_optional)
    if rocm_path is None:
        return None

    # Check .info/version file
    version_file = rocm_path / ".info" / "version"
    if not version_file.exists():
        # Try alternative location
        version_file = rocm_path / "share" / "rocm" / "version"

    if version_file.exists():
        try:
            version_str = version_file.read_text().strip()
            match = re.match(r"(\d+)\.(\d+)\.(\d+)", version_str)
            if match:
                return (int(match.group(1)), int(match.group(2)), int(match.group(3)))
        except (OSError, ValueError):
            pass

    return None


def _find_executable(name: str, search_paths: list[Path]) -> Optional[Path]:
    """Find a rocprof-sys executable within the given search paths.

    Only the explicit search paths are consulted -- these are the build/install
    ``bin`` directory, which already honors the ROCPROFSYS_BUILD_DIR and
    ROCPROFSYS_INSTALL_DIR overrides. PATH is intentionally NOT searched: the
    suite must exercise the rocprof-sys binaries from the build/install under
    test, never an unrelated copy that happens to be on PATH (e.g.
    ``/opt/rocm/bin``), which would silently test the wrong build.
    """
    for search_dir in search_paths:
        exe = search_dir / name
        if exe.exists() and exe.is_file():
            return exe.resolve()

    return None


def _find_rocprofsys_core_executables(
    search_paths: list[Path],
) -> dict[str, Optional[Path]]:
    """Return a dictionary of rocprofiler-systems executables and their paths.

    Throws a FileNotFoundError if any of the executables are not found.
    """
    rocprof_instrument = _find_executable("rocprof-sys-instrument", search_paths)
    rocprof_sample = _find_executable("rocprof-sys-sample", search_paths)
    rocprof_run = _find_executable("rocprof-sys-run", search_paths)
    rocprof_causal = _find_executable("rocprof-sys-causal", search_paths)
    rocprof_avail = _find_executable("rocprof-sys-avail", search_paths)

    required_executables = {
        "rocprof-sys-instrument": rocprof_instrument,
        "rocprof-sys-sample": rocprof_sample,
        "rocprof-sys-run": rocprof_run,
        "rocprof-sys-causal": rocprof_causal,
        "rocprof-sys-avail": rocprof_avail,
    }

    missing = [name for name, path in required_executables.items() if path is None]
    if missing:
        raise FileNotFoundError(
            f"Required rocprof-sys executables not found: {', '.join(missing)}. "
            f"Searched in: {search_paths}. Build the project first, or point the "
            f"suite at the correct location with ROCPROFSYS_BUILD_DIR (build tree) "
            f"or ROCPROFSYS_INSTALL_DIR (install prefix)."
        )

    return required_executables


def _find_rocprofsys_python(
    search_paths: list[Path], rocprofsys_build_dir: Path
) -> tuple[Optional[Path], Optional[Path]]:
    """Return the Python executable and rocprofsys site-package path.

    Both must be found, otherwise (None, None) is returned.
    """
    rocprof_python = _find_executable("rocprof-sys-python", search_paths)
    if not rocprof_python:
        return None, None

    # It is either in the agnostic path or in the versioned path
    agnostic_path = rocprofsys_build_dir / "lib" / "python" / "site-packages"
    if (agnostic_path / "rocprofsys").is_dir():
        return rocprof_python, agnostic_path

    # Only one versioned path will exist
    # Otherwise, it will be in the agnostic path
    for child in (rocprofsys_build_dir / "lib").iterdir():
        if child.is_dir() and re.match(r"python\d+\.\d+", child.name):
            site_packages = child / "site-packages"
            if (site_packages / "rocprofsys").is_dir():
                return rocprof_python, site_packages

    return None, None


def _merge_python_root_dirs(
    explicit_dirs: Optional[list[Path]],
) -> Optional[list[Path]]:
    """Merge explicit --python-root-dirs with ROCPROFSYS_PYTHON_HINTS env var.

    ROCPROFSYS_PYTHON_HINTS entries may point to bin/ directories or their parents;
    both are normalized to the parent directory for pythonX.Y lookup.
    """
    result: list[Path] = list(explicit_dirs or [])
    env_hints = os.environ.get("ROCPROFSYS_PYTHON_HINTS", "")
    if env_hints:
        for hint in env_hints.split(";"):
            hint = hint.strip()
            if hint:
                hint_path = Path(hint)
                parent = hint_path.parent if hint_path.name == "bin" else hint_path
                if parent not in result:
                    result.append(parent)
    return result or None


@persistent_cache("config.discover_install_config")
def discover_install_config(
    python_versions: Optional[list[str]] = None,
    python_root_dirs: Optional[list[Path]] = None,
    rocm_optional: bool = False,
) -> RocprofsysConfig:
    """Discover rocprofiler-systems installation configuration.

    Creates configuration for testing against installed binaries. The install
    prefix is resolved from ROCPROFSYS_INSTALL_DIR or common locations.

    Returns:
        RocprofsysConfig configured for installed binaries

    Raises:
        FileNotFoundError: If build/installation dirs and executables are not found
    """

    install_dir: Optional[Path] = None
    env_install = os.environ.get("ROCPROFSYS_INSTALL_DIR")
    if env_install:
        install_dir = Path(env_install).resolve()
    else:
        _rocm_candidate = _find_rocm_path(optional=True)
        _install_candidates = [
            Path("/usr/local"),
            Path("/usr"),
            Path("/opt/rocprofiler-systems"),  # Standard install location from README.md
        ]
        if _rocm_candidate is not None:
            _install_candidates.insert(0, _rocm_candidate)
        for candidate in _install_candidates:
            if (
                candidate
                and (candidate / "share" / "rocprofiler-systems" / "tests").is_dir()
                and (candidate / "share" / "rocprofiler-systems" / "examples").is_dir()
            ):
                install_dir = candidate
                break

    if install_dir is None:
        raise FileNotFoundError(
            "Could not find a suitable rocprofiler-systems installation. Set ROCPROFSYS_INSTALL_DIR "
            "environment variable."
            "A suitable installation is one that has the following directory: share/rocprofiler-systems/examples "
            "and share/rocprofiler-systems/tests"
        )

    install_dir = install_dir.resolve()

    # Determine directory layout
    bin_dir = install_dir / "bin"
    lib_dir = install_dir / "lib"

    # For lib64 systems
    if not lib_dir.exists() and (install_dir / "lib64").exists():
        lib_dir = install_dir / "lib64"

    examples_dir = install_dir / "share" / "rocprofiler-systems" / "examples"
    tests_dir = install_dir / "share" / "rocprofiler-systems" / "tests"
    rocpd_validation_rules = tests_dir / "rocpd-validation-rules"

    # Create a temporary directory for test outputs
    username = resolve_username()
    output_dir = Path(tempfile.gettempdir()) / username / "rocprof-sys-pytest-output"
    output_dir.mkdir(parents=True, exist_ok=True)

    rocm_path = _find_rocm_path(optional=rocm_optional)

    search_paths = [bin_dir]
    sys_execs = _find_rocprofsys_core_executables(search_paths)
    rocprofsys_python, rocprofsys_site_packages = _find_rocprofsys_python(
        search_paths, install_dir
    )

    return RocprofsysConfig(
        rocprofsys_build_dir=install_dir,
        rocprofsys_instrument=sys_execs["rocprof-sys-instrument"],
        rocprofsys_run=sys_execs["rocprof-sys-run"],
        rocprofsys_sample=sys_execs["rocprof-sys-sample"],
        rocprofsys_causal=sys_execs["rocprof-sys-causal"],
        rocprofsys_avail=sys_execs["rocprof-sys-avail"],
        rocm_path=rocm_path,
        rocprofsys_lib_dir=lib_dir,
        rocprofsys_bin_dir=bin_dir,
        rocprofsys_examples_dir=examples_dir,
        rocprofsys_tests_dir=tests_dir,
        rocpd_validation_rules=rocpd_validation_rules,
        test_output_dir=output_dir,
        rocm_version=_get_rocm_version(rocm_optional=rocm_optional),
        is_installed=True,
        rocprofsys_python=rocprofsys_python,
        rocprofsys_site_packages=rocprofsys_site_packages,
        _python_versions_hint=python_versions,
        _python_root_dirs_hint=_merge_python_root_dirs(python_root_dirs),
    )


@persistent_cache("config.discover_build_config")
def discover_build_config(
    python_versions: Optional[list[str]] = None,
    python_root_dirs: Optional[list[Path]] = None,
    rocm_optional: bool = False,
) -> RocprofsysConfig:
    """Discover rocprofiler-systems build configuration.

    Finds the build directory automatically (from ROCPROFSYS_BUILD_DIR or by
    walking up from this file). If no build directory is found but an
    installation is available, falls back to discover_install_config().

    Returns:
        RocprofsysConfig with discovered paths

    Raises:
        FileNotFoundError: If build/installation dirs and executables are not found
    """

    # Explicit install directory check
    if os.environ.get("ROCPROFSYS_INSTALL_DIR"):
        return discover_install_config(
            python_versions=python_versions,
            python_root_dirs=python_root_dirs,
            rocm_optional=rocm_optional,
        )

    # When running from pyz package (extracted to /tmp), fall back to install config
    # The pyz extracts to paths like /tmp/rocprofsys-tests-*/tests/rocprofsys/config.py
    current_file = Path(__file__).resolve()
    if str(current_file).startswith(tempfile.gettempdir()):
        return discover_install_config(
            python_versions=python_versions,
            python_root_dirs=python_root_dirs,
            rocm_optional=rocm_optional,
        )

    # All files should be in the build directory
    env_build = os.environ.get("ROCPROFSYS_BUILD_DIR")
    if env_build:
        build_dir = Path(env_build).resolve()
    else:
        build_dir = Path(__file__).resolve().parent.parent.parent.parent.parent.parent

    if not build_dir.exists():
        raise FileNotFoundError(
            "Could not find build directory or installation. Set one of:\n"
            "  - ROCPROFSYS_BUILD_DIR: Path to build directory\n"
            "  - ROCPROFSYS_INSTALL_DIR: Path to installation prefix"
        )

    rocm_path = _find_rocm_path(optional=rocm_optional)

    bin_dir = build_dir / "bin"
    lib_dir = build_dir / "lib"

    search_paths = [bin_dir]
    sys_execs = _find_rocprofsys_core_executables(search_paths)
    rocprofsys_python, rocprofsys_site_packages = _find_rocprofsys_python(
        search_paths, build_dir
    )

    share_path = build_dir / "share" / "rocprofiler-systems"

    output_dir = build_dir / "rocprof-sys-pytest-output"

    tests_dir = share_path / "tests"

    return RocprofsysConfig(
        rocprofsys_build_dir=build_dir,
        rocprofsys_instrument=sys_execs["rocprof-sys-instrument"],
        rocprofsys_run=sys_execs["rocprof-sys-run"],
        rocprofsys_sample=sys_execs["rocprof-sys-sample"],
        rocprofsys_causal=sys_execs["rocprof-sys-causal"],
        rocprofsys_avail=sys_execs["rocprof-sys-avail"],
        rocm_path=rocm_path,
        rocprofsys_lib_dir=lib_dir,
        rocprofsys_bin_dir=bin_dir,
        rocprofsys_examples_dir=build_dir,  # Example binaries are (almost always) in root of build directory
        rocprofsys_tests_dir=tests_dir,
        rocpd_validation_rules=tests_dir / "rocpd-validation-rules",
        test_output_dir=output_dir,
        rocm_version=_get_rocm_version(rocm_optional=rocm_optional),
        is_installed=False,
        rocprofsys_python=rocprofsys_python,
        rocprofsys_site_packages=rocprofsys_site_packages,
        _python_versions_hint=python_versions,
        _python_root_dirs_hint=_merge_python_root_dirs(python_root_dirs),
    )