run_tests.py

#!/usr/bin/env python
# Script to run tests and benchmarks.
import argparse
import glob
import os
import subprocess
import sys
from typing import Sequence


def parse_arguments():
  parser = argparse.ArgumentParser(description="Select tests to run")
  parser.add_argument(
      "--gpu-integration-tests",
      help="Run GPU integration tests - requires a GPU with CUDA installation.",
      dest="gpu_integration_tests",
      default=False,
      action=argparse.BooleanOptionalAction,
  )
  parser.add_argument(
      "--iterators-tests",
      help="Run Iterators tests.",
      dest="iterators_tests",
      default=False,
      action=argparse.BooleanOptionalAction,
  )
  return parser.parse_args()


def _convert_path_to_module(test_script: str) -> str:
  """Convert the path of the test script to its module name."""
  test_script = test_script.replace(os.sep, ".")
  test_script = test_script.strip(".")
  if test_script.endswith(".py"):
    return test_script[:-3]
  return test_script


def _configure_env():
  env = os.environ
  build_dir = env["IREE_LLVM_SANDBOX_BUILD_DIR"]
  # TODO: just grab from .env.
  env["PYTHONPATH"] = (
      os.path.join(build_dir, "tools/sandbox/python_packages") +
      ((":" + env["PYTHONPATH"]) if "PYTHONPATH" in env else ""))
  env["MLIR_RUNNER_UTILS_LIB"] = os.path.join(build_dir,
                                              "lib/libmlir_runner_utils.so")
  env["MLIR_C_RUNNER_UTILS_LIB"] = os.path.join(
      build_dir, "lib/libmlir_c_runner_utils.so")
  env["MLIR_RUNNER_EXTRA_LIBS"] = os.path.join(
      build_dir, "lib/libmlir_async_runtime_copy.so")
  return env


def _run_test(test_script: str, test_args: Sequence[str] = []) -> bool:
  """Run the provided test script an return failure or success.
  A test succeeds if:
  - it does not time out
  - it returns zero
  - it does not print FAILURE
  """
  print(f"- running {test_script}: ", end="")
  module = _convert_path_to_module(test_script)
  env = _configure_env()
  proc = subprocess.Popen(["python", "-m", module] + test_args,
                          stdout=subprocess.PIPE,
                          stderr=subprocess.PIPE,
                          env=env)
  try:
    outs, errs = proc.communicate(timeout=20)
  except subprocess.TimeoutExpired:
    proc.kill()
    print("\033[31m" + "FAILED" + "\033[m")
    print("  -> test execution timed out")
    return False
  if proc.returncode != 0:
    print("\033[31m" + "FAILED" + "\033[m")
    print(f"  -> test returned code {proc.returncode}")
    print(errs.decode("utf-8"))
    return False
  # Check the output for numerical failures.
  outs = outs.decode("utf-8")
  errs = errs.decode("utf-8")
  for line in outs.splitlines() + errs.splitlines():
    if line.count("FAILURE") != 0:
      print("\033[31m" + "FAILED" + "\033[m")
      print(f"  -> test failure: {line}")
      return False
  print("\033[32m" + "SUCCESS" + "\033[m")
  return True


def run_small_search():
  return _run_test(
      "./python/examples/tuning/test_nevergrad_small_matmul.py",
      ['--search-budget 500', '--n_iters 3000', '--num-parallel-tasks 10'])


def main(args):
  results = []
  for f in glob.glob("./python/**/*test.py", recursive=True):
    results.append(_run_test(f))
  # Tun a small search.
  results.append(
      _run_test("./python/examples/tuning/test_nevergrad_small_matmul.py", [
          '--search-budget', '10', '--n_iters', '10', '--num-parallel-tasks',
          '4'
      ]))
  errors = results.count(False)
  if errors:
    print(f"-> {errors} tests failed!")

  # Additionally run the lit tests.
  print(f"- running lit tests:")
  lit_args = ["lit", "-v"]
  if not args.gpu_integration_tests:
    lit_args.append("--filter-out=Integration/Dialect/VectorExt/GPU")
  test_dirs = ["test"]
  if args.iterators_tests:
    test_dirs += [
        "experimental/iterators/unittests", "experimental/iterators/test"
    ]
  returncode = subprocess.call(lit_args + test_dirs, env=_configure_env())
  if returncode != 0:
    print(f"-> lit tests failed!")
  if returncode != 0 or errors:
    exit(1)


if __name__ == '__main__':
  main(parse_arguments())