TheRock/rocm-libraries/projects/hipthreads/test/utils/run.py

#!/usr/bin/env python
# ===----------------------------------------------------------------------===##
#
# Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
# See https://llvm.org/LICENSE.txt for license information.
# SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
#
# ===----------------------------------------------------------------------===##

"""run.py is a utility for running a program.

It can perform code signing, forward arguments to the program, and return the
program's error code.
"""

import argparse
import os
import platform
import subprocess


def main():
    parser = argparse.ArgumentParser()
    parser.add_argument("--execdir", type=str, required=True)
    parser.add_argument("--codesign_identity", type=str, required=False, default=None)
    parser.add_argument("--env", type=str, nargs="*", required=False, default=[])
    parser.add_argument(
        "--prepend_env", type=str, nargs="*", required=False, default=[]
    )
    parser.add_argument("command", nargs=argparse.ONE_OR_MORE)
    args = parser.parse_args()
    commandLine = args.command

    # HACK:
    # If an argument is a file that ends in `.tmp.exe`, assume it is the name
    # of an executable generated by a test file. We call these test-executables
    # below. This allows us to do custom processing like codesigning test-executables.
    # It's also possible for there to be no such executable, for example in the case
    # of a .sh.cpp test.
    isTestExe = lambda exe: exe.endswith(".tmp.exe") and os.path.exists(exe)

    # Do any necessary codesigning of test-executables found in the command line.
    if args.codesign_identity:
        for exe in filter(isTestExe, commandLine):
            codesign = ["codesign", "-f", "-s", args.codesign_identity, exe]
            subprocess.check_call(codesign, env={})

    # Extract environment variables into a dictionary
    env = {k: v for (k, v) in map(lambda s: s.split("=", 1), args.env)}

    # Set environment variables where we prepend the given value to the
    # existing environment variable.
    for (k, v) in map(lambda s: s.split("=", 1), args.prepend_env):
        if k in os.environ:
            v = v + os.pathsep + os.environ[k]
        env[k] = v

    if platform.system() == "Windows":
        # Pass some extra variables through on Windows:
        # COMSPEC is needed for running subprocesses via std::system().
        if "COMSPEC" in os.environ:
            env["COMSPEC"] = os.environ.get("COMSPEC")
        # TEMP is needed for placing temp files in a sensible directory.
        if "TEMP" in os.environ:
            env["TEMP"] = os.environ.get("TEMP")

    # Run the command line with the given environment in the execution directory.
    #
    # A hipThreads test can hang (e.g. a hip::wthread whose persistent scheduler
    # never reaches the state join() waits for). Without a timeout, one hung test
    # would block the whole lit run until the CI job's global timeout, cancelling
    # every later test. So we run the child in its own process group and kill the
    # entire group if it exceeds HIPTHREADS_TEST_TIMEOUT seconds (default 30, 0 to
    # disable). Using killpg to kill the whole process group (not just the
    # immediate child like the timeout command) is required in case a test
    # launches a sub-process. A timed-out test returns non-zero, which lit records
    # as a normal FAIL/XFAIL instead of a hang.
    try:
        timeout = float(os.environ.get("HIPTHREADS_TEST_TIMEOUT", "30"))
    except ValueError:
        timeout = 30.0

    if not timeout or platform.system() == "Windows":
        # No timeout requested, or process-group semantics below are POSIX-only.
        return subprocess.call(commandLine, cwd=args.execdir, env=env, shell=False)

    proc = subprocess.Popen(
        commandLine, cwd=args.execdir, env=env, shell=False, start_new_session=True
    )
    try:
        return proc.wait(timeout=timeout)
    except subprocess.TimeoutExpired:
        import signal

        # Graduated shutdown - send SIGTERM to request graceful shutdown and give
        # it 5 seconds to gracefully shut down before forcefully killing it with
        # SIGKILL.
        pgid = os.getpgid(proc.pid)
        os.killpg(pgid, signal.SIGTERM)
        try:
            proc.wait(timeout=5)
        except subprocess.TimeoutExpired:
            os.killpg(pgid, signal.SIGKILL)
            proc.wait()
        print(
            f"run.py: command timed out after {timeout:g}s and was killed: "
            f"{' '.join(commandLine)}"
        )
        return 124


if __name__ == "__main__":
    exit(main())