TheRock/rocm-systems/projects/rocprofiler-systems/tests/validate-rocpd.py

#!/usr/bin/env python3

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

import argparse
import os
import sys
import sqlite3
from pathlib import Path


class validation_rule:
    """Class to represent a validation rule as defined in JSON file"""

    def __init__(
        self,
        description,
        query,
        expected_result,
        comparison,
        error_message,
        requires=None,
        expected_result_max=None,
        gpu_category_to_skip=None,
    ):
        self.description = description
        self.query = query
        self.expected_result = expected_result
        self.comparison = comparison
        self.error_message = error_message
        self.requires = requires
        self.expected_result_max = expected_result_max
        self.gpu_category_to_skip = gpu_category_to_skip or []

    def __repr__(self):
        return f"validation_rule(description={self.description}, query={self.query})"

    def validate_query(self, result):
        """
        Validate the actual result against expected using the specified comparison
        defined in validation_queries in rules definition.
        NOTE: see default_rules.json
        """
        if self.comparison == "equals":
            return result == self.expected_result
        elif self.comparison == "greater_than":
            return result > self.expected_result
        elif self.comparison == "less_than":
            return result < self.expected_result
        elif self.comparison == "greater_than_or_equal":
            return result >= self.expected_result
        elif self.comparison == "less_than_or_equal":
            return result <= self.expected_result
        elif self.comparison == "not_equals":
            return result != self.expected_result
        elif self.comparison == "between_inclusive":
            if self.expected_result_max is None:
                raise ValueError(
                    "between_inclusive requires expected_result (min) and "
                    "expected_result_max (max) in the rule JSON"
                )
            return self.expected_result <= result <= self.expected_result_max
        else:
            raise ValueError(f"Unknown comparison operator: {self.comparison}")

    def expected_summary(self):
        """Human-readable expected value for failure messages."""
        if self.comparison == "between_inclusive":
            return (
                f"{self.comparison} [{self.expected_result}, "
                f"{self.expected_result_max}]"
            )
        return f"{self.comparison} {self.expected_result}"


class required_table:
    """Class to represent a required table as defined in JSON rules file"""

    def __init__(
        self, name, name_prefix, required_columns, min_rows=1, validation_queries=None
    ):
        if name is None and name_prefix is None:
            raise ValueError("Either 'name' or 'name_prefix' must be specified")
        if name is not None and name_prefix is not None:
            raise ValueError("Cannot specify both 'name' and 'name_prefix'")

        self.name = name
        self.name_prefix = name_prefix
        self.required_columns = required_columns
        self.min_rows = min_rows
        self.validation_queries = validation_queries or []

    def __repr__(self):
        identifier = (
            f"name={self.name}" if self.name else f"name_prefix={self.name_prefix}"
        )
        return f"required_table({identifier}, required_columns={self.required_columns})"

    def get_table_identifier(self):
        """Returns the table identifier (name or prefix) for display purposes"""
        return self.name if self.name else f"{self.name_prefix}*"


def print_help():
    """Print out the help message"""
    print(f"""
    ROCPD Database Validation Tool

    DESCRIPTION:
        This tool validates ROCm Profiler Database (ROCPD) files against a set of predefined rules.
        It checks for required tables, columns, minimum row counts, and executes custom validation queries.

    USAGE:
        {os.path.basename(__file__)} --database <path_to_database> [OPTIONS]

    REQUIRED ARGUMENTS:
        -db, --database PATH        Path to the ROCPD database file (.db) to validate

    OPTIONAL ARGUMENTS:
        -r, --validation_rules PATH [PATH ...]  One or more JSON rules files (default: default_rules.json)
        --gpu-category-to-skip CAT [CAT ...]    Skip validation queries tagged with these GPU
                                                categories in the rules JSON via
                                                "gpu_category_to_skip" (e.g. apu). Default: none.
        -h, --help                  Show this help message and exit

    EXAMPLES:
        # Validate database with default rules
        {os.path.basename(__file__)} --database my_profile.db

        # Validate database with custom rules file
        {os.path.basename(__file__)} --database my_profile.db -r custom_rules.json

        # Validate database with multiple rules files
        {os.path.basename(__file__)} --database my_profile.db -r validation_rules.json amd_smi_rules.json

        # Skip queries tagged for APUs (e.g. when shared HBM means no PAGE_MIGRATE events)
        {os.path.basename(__file__)} --database my_profile.db --gpu-category-to-skip apu

    VALIDATION FEATURES:
        - Checks for presence of required tables
        - Verifies required columns exist in each table
        - Ensures minimum row count requirements are met
        - Executes custom SQL validation queries
        - Supports various comparison operators (equals, greater_than, less_than,
          between_inclusive, etc.)

    EXIT CODES:
        0  - All validations passed successfully
        64 - Invalid command line arguments (EX_USAGE)
        65 - Validation failures detected (EX_DATAERR)
        1  - General error (database connection, file not found, etc.)
    """)


def _query_skip_categories(
    query_gpu_category_to_skip: list[str], gpu_category_to_skip: set[str]
) -> list[str]:
    """Return GPU categories that validation query should be skipped for."""
    if not query_gpu_category_to_skip or not gpu_category_to_skip:
        return []
    return sorted(set(query_gpu_category_to_skip) & gpu_category_to_skip)


def validate_table(
    cursor, rule, tables, available_metrics=None, gpu_category_to_skip=None
) -> bool:
    """
    Validates a database table against a set of rules.
    This function checks if a table specified by `rule` exists in the provided `tables` list,
    verifies that all required columns are present, ensures the table meets a minimum row count,
    and executes custom validation queries defined in the rule.

    Args:
        cursor: Database cursor used to execute SQL queries.
        rule: An object containing validation rules for the table.
        bool: True if the table passes all validation checks, False otherwise.

    Returns:
        bool: True if table is found in the database and if all validation queries pass,
              False if any validation fails or matching table not found in database.
    """

    if gpu_category_to_skip is None:
        gpu_category_to_skip = set()

    matching_tables = []

    if rule.name:
        for table in tables:
            if table["name"] == rule.name:
                matching_tables.append(table)
                break
    elif rule.name_prefix:
        for table in tables:
            if table["name"].startswith(rule.name_prefix):
                matching_tables.append(table)

    if not matching_tables:
        if rule.name:
            print(f"❌ ERROR: Required table '{rule.name}' not found in database")
        elif rule.name_prefix:
            print(
                f"❌ ERROR: No tables found with prefix '{rule.name_prefix}' in database"
            )
        return False

    all_tables_passed = True

    for matching_table in matching_tables:
        table_name = matching_table["name"]

        try:
            cursor.execute(f"PRAGMA table_info({table_name})")
            columns = cursor.fetchall()
            column_names = [col["name"] for col in columns]

            missing_columns = [
                col for col in rule.required_columns if col not in column_names
            ]
            if missing_columns:
                print(
                    f"❌ ERROR: Table '{table_name}' missing required columns: {missing_columns}"
                )
                all_tables_passed = False
                continue
            else:
                print(
                    f"✅ All required columns present in '{table_name}': {rule.required_columns}"
                )

            cursor.execute(f"SELECT COUNT(*) as count FROM {table_name}")
            row_count = cursor.fetchone()["count"]

            if row_count < rule.min_rows:
                print(
                    f"❌ ERROR: Table '{table_name}' has {row_count} rows, minimum required: {rule.min_rows}"
                )
                all_tables_passed = False
                continue
            else:
                print(
                    f"✅ Row count check passed for '{table_name}': {row_count} rows (minimum: {rule.min_rows})"
                )

            all_queries_passed = True
            for validation_query in rule.validation_queries:
                # Check if metric is available (based on union across all GPUs for now)
                if (
                    validation_query.requires
                    and available_metrics is not None
                    and validation_query.requires not in available_metrics
                ):
                    print(
                        f"⏭️  Skipping '{validation_query.description}' on '{table_name}' "
                        f"(requires '{validation_query.requires}', not available)"
                    )
                    continue

                matched_categories = _query_skip_categories(
                    validation_query.gpu_category_to_skip, gpu_category_to_skip
                )
                if matched_categories:
                    print(
                        f"⏭️  Skipping '{validation_query.description}' on '{table_name}' "
                        f"(gpu_category_to_skip: matched {matched_categories})"
                    )
                    continue

                try:
                    query = validation_query.query.replace("{table_name}", table_name)
                    cursor.execute(query)
                    result = cursor.fetchone()

                    if result and "count" in result.keys():
                        actual_result = result["count"]
                    else:
                        actual_result = result[0] if result else None

                    if not validation_query.validate_query(actual_result):
                        print(
                            f"❌ ERROR: {validation_query.error_message} (Table: '{table_name}')"
                        )
                        print(
                            f"   Expected: {validation_query.expected_summary()}, "
                            f"Got: {actual_result}"
                        )
                        all_queries_passed = False
                    else:
                        print(
                            f"✅ Validation query passed for '{table_name}': {validation_query.description}"
                        )

                except sqlite3.Error as e:
                    print(
                        f"❌ ERROR: Failed to execute validation query on '{table_name}': {e}"
                    )
                    print(f"Query: {validation_query.query}")
                    all_queries_passed = False

            if not all_queries_passed:
                all_tables_passed = False

        except sqlite3.Error as e:
            print(f"❌ ERROR: Failed to validate table '{table_name}': {e}")
            all_tables_passed = False

    return all_tables_passed


def validate_rocpd(
    cursor, rules, tables, available_metrics=None, gpu_category_to_skip=None
) -> bool:
    """
    Validation of a ROCPD database by applying a set of validation rules to specified tables.
    It iterates through each rule, validates the corresponding table, and provides feedback on the validation status.

    Args:
        cursor: Database cursor object for executing SQL queries
        rules: List of validation rule objects containing validation criteria for a specific table
        tables: Collection of table definitions or table objects to validate against

    Returns:
        bool: True if all validation checks pass for all tables,
              False if any validation fails.
    """

    print("Starting ROCPD database validation...")
    db_valid = True

    for rule in rules:
        print(f"\nValidating table: {rule.get_table_identifier()}")
        table_valid = validate_table(
            cursor, rule, tables, available_metrics, gpu_category_to_skip
        )
        db_valid = db_valid and table_valid

    if db_valid:
        print("\n✅ All validation checks passed!")
    else:
        print("\n❌ Some validation checks failed!")

    return db_valid


def load_validation_rules(validation_rules) -> list:
    """
    Load validation rules from a JSON file and convert them to validation objects.

    Args:
        rules_file: Path to the JSON rules file containing validation configuration.

    Returns:
        list: A list of required_table objects.
              Returns empty list if any file doesn't exist or on error.
    """
    import json

    all_rules = []

    for rules_file in validation_rules:
        try:
            rules_path = Path(rules_file)
            if not rules_path.exists():
                print(
                    f"Warning: Rules file '{rules_file}' not found, using default rules"
                )
                return []

            with open(rules_path, "r") as f:
                rules_data = json.load(f)
                rules = []

                for table_data in rules_data["required_tables"]:
                    validation_queries = []
                    for vq in table_data.get("validation_queries", []):
                        validation_query_obj = validation_rule(
                            description=vq["description"],
                            query=vq["query"],
                            expected_result=vq["expected_result"],
                            comparison=vq.get("comparison", "equals"),
                            error_message=vq["error_message"],
                            requires=vq.get("requires", None),
                            expected_result_max=vq.get("expected_result_max"),
                            gpu_category_to_skip=vq.get("gpu_category_to_skip", []),
                        )
                        validation_queries.append(validation_query_obj)

                    required_table_obj = required_table(
                        name=table_data.get("name", None),
                        name_prefix=table_data.get("name_prefix", None),
                        required_columns=table_data["required_columns"],
                        min_rows=table_data.get("min_rows", 1),
                        validation_queries=validation_queries,
                    )
                    rules.append(required_table_obj)
                    print(f"Loaded required table rule: {required_table_obj}")

                all_rules.extend(rules)

        except Exception as e:
            print(f"Error loading rules file: {e}")
            return []

    if not all_rules:
        print("Warning: No validation rules loaded from any file")
    else:
        print(f"Total rules loaded: {len(all_rules)}")

    return all_rules


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

    parser.add_argument(
        "-db", "--database", type=Path, help="Database file to validate", default=None
    )

    parser.add_argument(
        "-r",
        "--validation-rules",
        type=Path,
        nargs="+",
        help="Rules against which to validate database",
        default=[
            Path(
                f"{os.path.dirname(os.path.abspath(__file__))}/rocpd-validation-rules/default-rules.json"
            )
        ],
    )

    parser.add_argument(
        "--gpu-category-to-skip",
        nargs="*",
        default=[],
        help=(
            "GPU categories to skip tagged validation queries for "
            "(e.g. apu instinct). Passed from pytest tests after detect_gpu()."
        ),
    )

    parser.add_argument(
        "-h", "--help", action="store_true", help="Prints out the help message"
    )

    args = parser.parse_args()

    if args.help:
        print_help()
        sys.exit(os.EX_OK)

    if not args.database:
        print("Database file not provided!")
        print_help()

        sys.exit(os.EX_USAGE)

    # Auto-detect available GPU metrics via amd-smi
    available_metrics = None
    try:
        sys.path.insert(0, os.path.dirname(os.path.abspath(__file__)))
        from check_amd_smi_metrics import get_available_metrics

        gpus = get_available_metrics()
        available_metrics = set()
        from check_amd_smi_metrics import collect_metric_names

        print("\n--- Platform GPU Metric Availability ---")
        for gpu in gpus:
            gpu_metrics = collect_metric_names(gpu)
            available_metrics |= gpu_metrics
            print(f"GPU {gpu.gpu_id}:")
            print(
                f"  Activity:    gfx={gpu.gfx_activity}  umc={gpu.umc_activity}  mm={gpu.mm_activity}"
            )
            print(
                f"  Temperature: hotspot={gpu.hotspot_temperature}  edge={gpu.edge_temperature}"
            )
            print(f"  Power:       socket={gpu.current_socket_power}")
            print(
                f"  VCN/JPEG:    vcn_activity={gpu.vcn_activity}  vcn_busy={gpu.vcn_busy}  jpeg_activity={gpu.jpeg_activity}  jpeg_busy={gpu.jpeg_busy}"
            )
            print(
                f"  Other:       mem_usage={gpu.mem_usage}  xgmi={gpu.xgmi}  pcie={gpu.pcie}"
            )
        print(
            f"Detected available metrics (union): {', '.join(sorted(available_metrics))}"
        )
        print("---\n")
    except Exception as e:
        print(f"Warning: Could not detect GPU metrics ({e}), running all queries")

    gpu_category_to_skip = set(args.gpu_category_to_skip)
    if gpu_category_to_skip:
        categories = ", ".join(sorted(gpu_category_to_skip))
        print(f"GPU category to skip: {categories}")

    print(f"Validating ROCPD. Database file: {args.database}")

    db_path = args.database
    validation_rules_files = args.validation_rules
    rules = load_validation_rules(validation_rules_files)

    if not rules:
        print("❌ No validation rules loaded. Exiting.")
        sys.exit(1)

    try:
        if not Path(db_path).exists():
            print(f"❌ Error: Database file '{db_path}' not found")
            sys.exit(1)

        conn = sqlite3.connect(db_path)
        conn.row_factory = sqlite3.Row
        cursor = conn.cursor()

        print(f"✅ Successfully connected to database: {db_path}")

        cursor.execute("SELECT name FROM sqlite_master WHERE type IN ('table', 'view');")
        tables = cursor.fetchall()

        validation_result = validate_rocpd(
            cursor, rules, tables, available_metrics, gpu_category_to_skip
        )

        conn.close()

        if validation_result:
            print(f"✅ {db_path} validated")
        else:
            print(f"❌ Failure validating {db_path}")

        sys.exit(os.EX_OK if validation_result else os.EX_DATAERR)

    except sqlite3.Error as e:
        print(f"SQLite error: {e}")
        sys.exit(1)
    except Exception as e:
        print(f"Error: {e}")
        sys.exit(1)