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
import getpass
import os
from pathlib import Path
import shutil
import tempfile
from typing import Optional
import re


@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:
            In build mode, this is the root of the build directory.
            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
        - mpiexec: Path to MPI launcher executable
        - is_installed: Whether this is an installed configuration
    """

    rocprofsys_build_dir: Path
    rocprofsys_instrument: Path
    rocprofsys_run: Path
    rocprofsys_sample: Path
    rocprofsys_causal: Path
    rocprofsys_avail: Path
    rocm_path: 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
    mpiexec: Path
    is_installed: bool = False
    rocm_version: Optional[tuple[int, int, int]] = None

    def get_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",
            ]
            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())]

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

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

        return ":".join(paths)

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

        When is_installed is True, searches in the following order:
        1. rocprofsys_build_dir/name (build directory layout)
        2. rocprofsys_examples_dir/name/name (build directory layout)
        3. PATH lookup

        When is_installed is False, searches in the following order:
        1. rocprofsys_examples_dir/name
        2. rocprofsys_bin_dir/name
        3. PATH lookup

        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

            # PATH lookup via shutil.which
            exe = shutil.which(name)
            if exe:
                return Path(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"
                f"  - PATH"
            )

        else:
            # Build directory mode
            exe = self.rocprofsys_examples_dir / name
            if exe.exists() and exe.is_file():
                return exe

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

            # rccl tests lie in their own directory
            exe = self.rocprofsys_examples_dir / "examples" / "rccl" / 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

            # PATH lookup via shutil.which
            exe = shutil.which(name)
            if exe:
                return Path(exe)

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

    def get_fundamental_environment(self) -> dict[str, str]:
        """Get fundamental environment variables inherited from parent process."""
        return {
            "PATH": os.environ.get("PATH", ""),
            "HOME": os.environ.get("HOME", ""),
            "USER": os.environ.get("USER", ""),
            "SHELL": os.environ.get("SHELL", ""),
            "TERM": os.environ.get("TERM", ""),
            "LANG": os.environ.get("LANG", ""),
        }

    def get_base_environment(self) -> dict[str, str]:
        """Get base environment variables for test execution."""
        return {
            "ROCPROFSYS_CI": "ON",
            "ROCPROFSYS_CONFIG_FILE": "",
            "ROCPROFSYS_TRACE": "ON",
            "ROCPROFSYS_PROFILE": "ON",
            "ROCPROFSYS_USE_SAMPLING": "ON",
            "ROCPROFSYS_USE_PROCESS_SAMPLING": "ON",
            "ROCPROFSYS_TIME_OUTPUT": "OFF",
            "ROCPROFSYS_FILE_OUTPUT": "ON",
            "ROCPROFSYS_USE_PID": "OFF",
            "ROCPROFSYS_VERBOSE": "1",
            "ROCPROFSYS_SAMPLING_FREQ": "300",
            "ROCPROFSYS_SAMPLING_DELAY": "0.05",
            "OMP_PROC_BIND": "spread",
            "OMP_PLACES": "threads",
            "OMP_NUM_THREADS": "2",
            "LD_LIBRARY_PATH": self.get_library_path(),
        }

    def get_base_binary_environment(self) -> dict[str, str]:
        """Get base environment variables for rocprof-sys binary test execution."""
        return {
            "ROCPROFSYS_TRACE": "ON",
            "ROCPROFSYS_PROFILE": "ON",
            "ROCPROFSYS_USE_SAMPLING": "ON",
            "ROCPROFSYS_TIME_OUTPUT": "OFF",
            "LD_LIBRARY_PATH": self.get_library_path(),
            "ROCPROFSYS_CI": "ON",
            "ROCPROFSYS_CI_TIMEOUT": "300",
            "ROCPROFSYS_CONFIG_FILE": "",
        }


def _find_rocm_path() -> Optional[Path]:
    """Find ROCm installation path."""
    for candidate in [
        os.environ.get("ROCM_PATH"),
        "/opt/rocm",
        "/usr/local/rocm",
    ]:
        if candidate and Path(candidate).exists():
            return Path(candidate).resolve()
    return None


def _get_rocm_version() -> 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()
    if not rocm_path:
        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_mpiexec() -> Optional[Path]:
    """Find MPI launcher executable."""
    for candidate in ["mpiexec", "mpirun"]:
        path = shutil.which(candidate)
        if path:
            return Path(path)
    return None


def _find_executable(name: str, search_paths: list[Path]) -> Optional[Path]:
    """Find an executable in search paths or via PATH."""
    for search_dir in search_paths:
        exe = search_dir / name
        if exe.exists() and exe.is_file():
            return exe.resolve()

    # Fallback to PATH
    path_exe = shutil.which(name)
    if path_exe:
        return Path(path_exe)

    return None


def discover_install_config(
    install_dir: Optional[Path] = None,
    output_dir: Optional[Path] = None,
) -> RocprofsysConfig:
    """Discover rocprofiler-systems installation configuration.

    Creates configuration for testing against installed binaries.

    Args:
        install_dir: Installation prefix (e.g., /opt/rocm or /usr/local)

    Returns:
        RocprofsysConfig configured for installed binaries

    Raises:
        FileNotFoundError: If installation cannot be found
    """

    if install_dir is None:
        env_install = os.environ.get("ROCPROFSYS_INSTALL_DIR")
        if env_install:
            install_dir = Path(env_install).resolve()
        else:
            for candidate in [
                _find_rocm_path(),
                Path("/usr/local"),
                Path("/usr"),
                Path(
                    "/opt/rocprofiler-systems"
                ),  # Standard install location from README.md
            ]:
                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
    try:
        username = getpass.getuser()
    except Exception:
        username = str(os.getuid())

    if output_dir is None:
        output_dir = Path(tempfile.gettempdir()) / username / "rocprof-sys-pytest-output"
    else:
        output_dir = Path(output_dir)

    output_dir.mkdir(parents=True, exist_ok=True)

    rocm_path = _find_rocm_path()
    mpiexec = _find_mpiexec()

    search_paths = [bin_dir]
    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)

    # If any of the executables are not found, raise an error
    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 executables not found: {', '.join(missing)}. "
            f"Searched in: {search_paths}"
        )

    return RocprofsysConfig(
        rocprofsys_build_dir=install_dir,
        rocprofsys_instrument=rocprof_instrument,
        rocprofsys_run=rocprof_run,
        rocprofsys_sample=rocprof_sample,
        rocprofsys_causal=rocprof_causal,
        rocprofsys_avail=rocprof_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,
        mpiexec=mpiexec,
        rocm_version=_get_rocm_version(),
        is_installed=True,
    )


def discover_build_config(
    build_dir: Optional[Path] = None,
    output_dir: Optional[Path] = None,
) -> RocprofsysConfig:
    """Discover rocprofiler-systems build configuration.

    Attempts to find the build directory and source directory automatically
    if not provided, checking common locations and environment variables.

    If no build directory is found but an installation is available,
    falls back to discover_install_config().

    Args:
        build_dir: Explicit build directory path

    Returns:
        RocprofsysConfig with discovered paths

    Raises:
        FileNotFoundError: If neither build directory nor installation found
    """

    # Explicit install directory check
    if os.environ.get("ROCPROFSYS_INSTALL_DIR"):
        return discover_install_config(output_dir=output_dir)

    # 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()

    # All files should be in the build directory
    if build_dir is None:
        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 build_dir is None or 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()
    mpiexec = _find_mpiexec()

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

    search_paths = [bin_dir]
    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)

    # If any of the executables are not found, raise an error
    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 executables not found: {', '.join(missing)}. "
            f"Searched in: {search_paths}"
        )

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

    if output_dir is None:
        output_dir = build_dir / "rocprof-sys-pytest-output"
    else:
        output_dir = Path(output_dir)

    return RocprofsysConfig(
        rocprofsys_build_dir=build_dir,
        rocprofsys_instrument=rocprof_instrument,
        rocprofsys_run=rocprof_run,
        rocprofsys_sample=rocprof_sample,
        rocprofsys_causal=rocprof_causal,
        rocprofsys_avail=rocprof_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=share_path / "tests",
        rocpd_validation_rules=share_path / "tests" / "rocpd-validation-rules",
        test_output_dir=output_dir,
        mpiexec=mpiexec,
        rocm_version=_get_rocm_version(),
        is_installed=False,
    )