TheRock/rocm-libraries/projects/rocblas/rtest.py

#!/usr/bin/python3
"""Copyright (C) 2021-2025 Advanced Micro Devices, Inc. All rights reserved.

   Permission is hereby granted, free of charge, to any person obtaining a copy
   of this software and associated documentation files (the "Software"), to deal
   in the Software without restriction, including without limitation the rights
   to use, copy, modify, merge, publish, distribute, sublicense, and/or sell cop-
   ies of the Software, and to permit persons to whom the Software is furnished
   to do so, subject to the following conditions:

   The above copyright notice and this permission notice shall be included in all
   copies or substantial portions of the Software.

   THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR IM-
   PLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, FITNESS
   FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE AUTHORS OR
   COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER LIABILITY, WHETHER
   IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, OUT OF OR IN CONNE-
   CTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE SOFTWARE.
"""

import re
import os
import sys
import subprocess
import shlex
import shutil
import argparse
import pathlib
import platform
from genericpath import exists
from fnmatch import fnmatchcase
from xml.dom import minidom
import multiprocessing
import time

args = {}
OS_info = {}

timeout = False
test_proc = None
stop = 0

test_script = [ '%XML%' ] # [ 'cd %IDIR%', '%XML%' ]

def parse_args():
    """Parse command-line arguments"""
    parser = argparse.ArgumentParser(description="""
    Checks build arguments
    """)
    parser.add_argument(      '--ci_labels', type=str, required=False, default="",
                    help='Semicolon-separated labels that may modify test runs (e.g. "gfx12;TestLevel1Only"). '
                         'If omitted, the GITHUB_PR_LABELS environment variable is used when set (e.g. CI PR labels).')
    parser.add_argument(      '--ci_gfx', type=str, required=False, default="",
                    help='Semi-colon seperated list of gfx targets expected on test runs (optional, e.g. "gfx1030;gfx1201")')
    parser.add_argument('-e', '--emulation', type=str, required=False,
                        help='Emulation test set to run from rtest.xml (e.g. smoke, regression, extended')
    parser.add_argument('-t', '--test', required=False,
                        help='Test set to run from rtest.xml (e.g. osdb)')
    parser.add_argument('-g', '--debug', required=False, default=False,  action='store_true',
                        help='Test Debug build (optional, default: false)')
    parser.add_argument('-o', '--output', type=str, required=False, default="xml",
                        help='Test output file (optional, default: test_detail.xml)')
    parser.add_argument(      '--install_dir', type=str, required=False, default="build",
                        help='Installation directory where build or release folders are (optional, default: build)')
    parser.add_argument(      '--fail_test', default=False, required=False, action='store_true',
                        help='Return as if test failed (optional, default: false)')
    # parser.add_argument('-v', '--verbose', required=False, default = False, action='store_true',
    #                     help='Verbose install (optional, default: False)')
    return parser.parse_args()

def arg_into_list(arg) -> list:
    arg = re.sub(r"['\"]|['\']",'', arg)
    return arg.split(';')

def rocm_executable(exe_name: str) -> str:
  if shutil.which(exe_name):
    return exe_name
  bin_dir = os.environ.get("ROCM_PATH") or os.environ.get("HIP_PATH")
  if bin_dir:
    candidate = pathlib.Path(bin_dir) / "bin" / exe_name
    if candidate.exists():
        return str(candidate)
  return None

def vram_detect():
    global OS_info
    OS_info["VRAM"] = 0
    if os.name == "nt":
        cmd = rocm_executable("hipinfo.exe")
        if cmd is not None: 
            process = subprocess.run([cmd], stdout=subprocess.PIPE)
            for line_in in process.stdout.decode().splitlines():
                if 'totalGlobalMem' in line_in:
                    OS_info["VRAM"] = float(line_in.split()[1])
                    break
    else:
        cmd = rocm_executable("rocminfo")
        if cmd is not None: 
            process = subprocess.run([cmd], stdout=subprocess.PIPE)
            for line_in in process.stdout.decode().splitlines():
                match = re.search(r'.*Size:.*([0-9]+)\(.*\).*KB', line_in, re.IGNORECASE)
                if match:
                    OS_info["VRAM"] = float(match.group(1))/(1024*1024)
                    break

def os_detect():
    global OS_info
    if os.name == "nt":
        OS_info["ID"] = platform.system()
    else:
        inf_file = "/etc/os-release"
        if os.path.exists(inf_file):
            with open(inf_file) as f:
                for line in f:
                    if "=" in line:
                        k,v = line.strip().split("=")
                        OS_info[k] = v.replace('"','')
    OS_info["NUM_PROC"] = os.cpu_count()
    vram_detect()
    print(OS_info)


def create_dir(dir_path):
    if os.path.isabs(dir_path):
        full_path = dir_path
    else:
        full_path = os.path.join( os.getcwd(), dir_path )
    return pathlib.Path(full_path).mkdir(parents=True, exist_ok=True)

def delete_dir(dir_path) :
    if (not os.path.exists(dir_path)):
        return
    if os.name == "nt":
        return run_cmd( "RMDIR" , f"/S /Q {dir_path}")
    else:
        linux_path = pathlib.Path(dir_path).absolute()
        return run_cmd( "rm" , f"-rf {linux_path}")

class TimerProcess(multiprocessing.Process):

    def __init__(self, start, stop, kill_pid):
        multiprocessing.Process.__init__(self)
        self.quit = multiprocessing.Event()
        self.timed_out = multiprocessing.Event()
        self.start_time = start
        self.max_time = stop
        self.kill_pid = kill_pid

    def run(self):
        while not self.quit.is_set():
            #print( f'time_stop {self.start_time} limit {self.max_time}')
            if (self.max_time == 0):
                return
            t = time.monotonic()
            if ( t - self.start_time > self.max_time ):
                print( f'killing {self.kill_pid} t {t}')
                if os.name == "nt":
                    cmd = ['TASKKILL', '/F', '/T', '/PID', str(self.kill_pid)]
                    proc = subprocess.Popen(cmd, stdout=sys.stdout, stderr=sys.stderr)
                else:
                    os.kill(self.kill_pid, signal.SIGKILL)
                self.timed_out.set()
                self.stop()
            pass

    def stop(self):
        self.quit.set()

    def stopped(self):
        return self.timed_out.is_set()


def time_stop(start, pid):
    global timeout, stop
    while (True):
        print( f'time_stop {start} limit {stop}')
        t = time.monotonic()
        if (stop == 0):
            return
        if ( (stop > 0) and (t - start > stop) ):
            print( f'killing {pid} t {t}')
            if os.name == "nt":
                cmd = ['TASKKILL', '/F', '/T', '/PID', str(pid)]
                proc = subprocess.Popen(cmd, stdout=sys.stdout, stderr=sys.stderr)
            else:
                test_proc.kill()
            timeout = True
            stop = 0
        time.sleep(0)

def gfilter_subset(filter, groups) -> str:
    new_filter = ""
    patterns = ["nightly"] if "nightly" in filter else ["quick", "pre_checkin"]
    for i in groups:
        for j in patterns:
            new_filter += f'*{i}*{j}*:'
    return new_filter

def label_modifiers(cmd, labels) -> str:
    original_cmd = cmd
    processed = ["TestTensileOnly", "TestLevel3Only", "TestLevel2Only", "TestLevel1Only"]
    overlap = [v for v in processed if v in labels]
    if len(overlap):
        tok = " --gtest_filter="
        cmd = cmd.split(tok)[0] + tok
    else:
        return cmd

    filter = ""
    if "TestTensileOnly" in overlap:
        filter += gfilter_subset( original_cmd, ["blas3_tensile", "blas2_tensile"] )
    if "TestLevel3Only" in overlap:
        filter += gfilter_subset( original_cmd, ["blas3"] )
    if "TestLevel2Only" in overlap:
        filter += gfilter_subset( original_cmd, ["blas2"] )
    if "TestLevel1Only" in overlap:
        filter += gfilter_subset( original_cmd, ["blas1"] )
    return cmd + f'"{filter}"'

def run_cmd(cmd, test = False, time_limit = 0, path = "" ):
    global args
    global test_proc, timer_thread
    global stop
    if (cmd.startswith('cd ')):
        return os.chdir(cmd[3:])
    if (cmd.startswith('mkdir ')):
        return create_dir(cmd[6:])
    cmdline = f"{cmd}"
    print(cmdline)
    try:
        if not test:
            proc = subprocess.run(cmdline, check=True, stderr=subprocess.STDOUT, shell=True)
            status = proc.returncode
        else:
            sub_env = os.environ.copy()
            if len(path):
                sub_env["PATH"] = path + os.pathsep + sub_env["PATH"]
            error = False
            timeout = False
            test_proc = subprocess.Popen(cmdline, text=True, stdout=subprocess.PIPE, stderr=subprocess.STDOUT, shell=True, env=sub_env)
            if time_limit > 0:
                start = time.monotonic()
                #p = multiprocessing.Process(target=time_stop, args=(start, test_proc.pid))
                p = TimerProcess(start, time_limit, test_proc.pid)
                p.start()
            while True:
                output = test_proc.stdout.readline()
                if output == '' and test_proc.poll() is not None:
                    break
                elif output:
                    outstring = output.strip()
                    print (outstring)
                    # skip output that is optional GTest summary
                    if not re.search(r'\[ FAILED.*\]', outstring):
                        error = error or re.search(r'error|fail', outstring, re.IGNORECASE)
            status = test_proc.poll()
            if time_limit > 0:
                p.stop()
                p.join()
                timeout = p.stopped()
                print(f"timeout {timeout}")
            if error:
                status = 1
            elif timeout:
                status = 2
            else:
                status = test_proc.returncode
    except:
        import traceback
        exc = traceback.format_exc()
        print( "Python Exception: {0}".format(exc) )
        status = 3
    return status

def test_xml():
    # Get the full path of the currently executing script look for installed name first
    exe_path = os.path.dirname(os.path.abspath(sys.argv[0]))
    installed_file = os.path.join( exe_path, 'rocblas_rtest.xml')
    if (os.path.exists(installed_file)):
        xmlPath = installed_file
    else:
        cwd = os.curdir
        xmlPath = os.path.join( cwd, 'rtest.xml')
    return xmlPath

def batch(script, xml):
    global OS_info
    global args
    #
    cwd = os.curdir
    rtest_cwd_path = os.path.abspath( os.path.join( cwd, 'rtest.xml') )
    rtest_path = test_xml()
    if 'rocblas_rtest.xml' in rtest_path:
        test_dir = os.path.dirname( rtest_path )
    elif os.path.isfile(rtest_cwd_path) and os.path.dirname(rtest_cwd_path).endswith( "staging" ):
        # if in a staging directory then test locally
        test_dir = cwd
    else:
        if args.debug: build_type = "debug"
        else: build_type = "release"
        test_dir = f"{args.install_dir}//{build_type}//clients//staging"
    fail = False
    for i in range(len(script)):
        cmdline = script[i]
        xcmd = cmdline.replace('%IDIR%', test_dir)
        cmd = xcmd.replace('%ODIR%', args.output)
        if cmd.startswith('tdir '):
            if pathlib.Path(cmd[5:]).exists():
                return 0 # all further cmds skipped
            else:
                continue
        error = False
        if cmd.startswith('%XML%'):
            # run the matching tests listed in the xml test file
            var_subs = {}
            for var in xml.getElementsByTagName('var'):
                name = var.getAttribute('name')
                val = var.getAttribute('value')
                var_subs[name] = val
            var_subs['EXE_DIR'] = test_dir + os.path.sep
            for test in xml.getElementsByTagName('test'):
                sets = test.getAttribute('sets')
                runset = sets.split(',')
                if args.test in runset:
                    for run in test.getElementsByTagName('run'):
                        name = run.getAttribute('name')
                        vram_limit = run.getAttribute('vram_min')
                        if vram_limit:
                            if OS_info["VRAM"] < float(vram_limit):
                                print( f'***\n*** Skipped: {name} due to VRAM req.\n***')
                                continue
                        if name:
                            print( f'***\n*** Running: {name}\n***')
                        time_limit = run.getAttribute('time_max')
                        if time_limit:
                            timeout = float(time_limit)
                        else:
                            timeout = 0

                        raw_cmd = run.firstChild.data
                        var_cmd = raw_cmd.format_map(var_subs)
                        if args.ci_labels:
                            var_cmd = label_modifiers(var_cmd, arg_into_list(args.ci_labels))
                        error = run_cmd(var_cmd, True, timeout, test_dir)
                        if (error == 2):
                            print( f'***\n*** Timed out when running: {name}\n***')
        else:
            error = run_cmd(cmd)
        fail = fail or error

    if (fail):
        if (cmd == "%XML%"):
            print(f"FAILED xml test suite!")
        else:
            print(f"ERROR running: {cmd}")
        if (os.curdir != cwd):
            os.chdir( cwd )
        return 1
    if (os.curdir != cwd):
        os.chdir( cwd )
    return 0


def run_tests():
    global test_script
    global xmlDoc

    # install
    cwd = os.curdir

    xmlPath = test_xml()
    xmlDoc = minidom.parse( xmlPath )

    scripts = []
    scripts.append( test_script )
    for i in scripts:
        if (batch(i, xmlDoc)):
            #print("Failure in script. ABORTING")
            if (os.curdir != cwd):
                os.chdir( cwd )
            return 1
    if (os.curdir != cwd):
        os.chdir( cwd )
    return 0

def test_xml_reduce():
    ## open test_detail.xml and strip out the larger extraneous attibutes

    test_detail = os.path.join( os.getcwd(), 'test_detail.xml')
    test_detail_reduced = os.path.join( os.getcwd(), 'test_detail_reduced.xml')
    if not os.path.exists(test_detail):
        return

    import xml.etree.ElementTree as ET

    # Load the XML file
    try:
        tree = ET.parse(test_detail)
        root = tree.getroot()

        # Remove specific attributes
        attributes_to_remove = {"value_param", "timestamp", "file", "line"}
        for elem in root.iter():
            for attr in attributes_to_remove:
                if attr in elem.attrib:
                    del elem.attrib[attr]

        # Save the modified XML
        tree.write(test_detail_reduced)
        # copy the test_detail_reduced.xml to test_detail.xml
        os.replace(test_detail_reduced, test_detail)
    except ET.ParseError as e:
        print(f"Error parsing XML file {test_detail}: {e}")


def main():
    global args
    global timer_thread

    os_detect()
    args = parse_args()

    # PR / CI labels (e.g. GitHub Actions): semicolon-separated, same format as --ci_labels.
    # When --ci_labels is omitted, use GITHUB_PR_LABELS from the environment (e.g. set by CI before ctest).
    if not args.ci_labels:
        env_labels = os.environ.get("GITHUB_PR_LABELS", "").strip()
        if env_labels:
            args.ci_labels = env_labels

    if args.emulation:
        args.test = f'emulation_{args.emulation}'

    if not (args.test or args.emulation):
        print('Either -t/--test or -e/--emulation is required.')
        args.fail_test = True
    else:
        status = run_tests()

    if args.fail_test: status = 1

    test_xml_reduce()

    if (status):
        sys.exit(status)

if __name__ == '__main__':
    main()