Refactor pytest config + add default dump dir option

Differential Revision: D65481575

Pull Request resolved: pytorch#6637

Author: Erik-Lundell

backends/arm/test/common.py

@@ -11,6 +11,10 @@
 import subprocess
 import sys
 import tempfile
+from datetime import datetime
+from enum import auto, Enum
+from pathlib import Path
+from typing import Any
 
 import pytest
 
@@ -19,27 +23,46 @@
 from executorch.backends.arm.arm_backend import ArmCompileSpecBuilder
 from executorch.exir.backend.compile_spec_schema import CompileSpec
 
-_enabled_options: list[str] = []
+
+class arm_test_options(Enum):
+    quantize_io = auto()
+    corstone300 = auto()
+    dump_path = auto()
+    date_format = auto()
+
+
+_test_options: dict[arm_test_options, Any] = {}
 
 # ==== Pytest hooks ====
 
 
 def pytest_addoption(parser):
     parser.addoption("--arm_quantize_io", action="store_true")
     parser.addoption("--arm_run_corstone300", action="store_true")
+    parser.addoption("--default_dump_path", default=None)
+    parser.addoption("--date_format", default="%d-%b-%H:%M:%S")
 
 
 def pytest_configure(config):
     if config.option.arm_quantize_io:
         load_libquantized_ops_aot_lib()
-        _enabled_options.append("quantize_io")
+        _test_options[arm_test_options.quantize_io] = True
     if config.option.arm_run_corstone300:
         corstone300_exists = shutil.which("FVP_Corstone_SSE-300_Ethos-U55")
         if not corstone300_exists:
             raise RuntimeError(
                 "Tests are run with --arm_run_corstone300 but corstone300 FVP is not installed."
             )
-        _enabled_options.append("corstone300")
+        _test_options[arm_test_options.corstone300] = True
+    if config.option.default_dump_path:
+        dump_path = Path(config.option.default_dump_path).expanduser()
+        if dump_path.exists() and os.path.isdir(dump_path):
+            _test_options[arm_test_options.dump_path] = dump_path
+        else:
+            raise RuntimeError(
+                f"Supplied argument 'default_dump_path={dump_path}' that does not exist or is not a directory."
+            )
+    _test_options[arm_test_options.date_format] = config.option.date_format
     logging.basicConfig(level=logging.INFO, stream=sys.stdout)
 
 
@@ -54,6 +77,18 @@ def pytest_collection_modifyitems(config, items):
                 item.add_marker(skip_if_aot_lib_not_loaded)
 
 
+def pytest_sessionstart(session):
+    pass
+
+
+def pytest_sessionfinish(session, exitstatus):
+    if get_option(arm_test_options.dump_path):
+        _clean_dir(
+            get_option(arm_test_options.dump_path),
+            f"ArmTester_{get_option(arm_test_options.date_format)}.log",
+        )
+
+
 # ==== End of Pytest hooks =====
 
 
@@ -76,7 +111,9 @@ def load_libquantized_ops_aot_lib():
         torch.ops.load_library(library_path)
 
 
-def is_option_enabled(option: str, fail_if_not_enabled: bool = False) -> bool:
+def is_option_enabled(
+    option: str | arm_test_options, fail_if_not_enabled: bool = False
+) -> bool:
     """
     Returns whether an option is successfully enabled, i.e. if the flag was
     given to pytest and the necessary requirements are available.
@@ -87,7 +124,10 @@ def is_option_enabled(option: str, fail_if_not_enabled: bool = False) -> bool:
     The optional parameter 'fail_if_not_enabled' makes the function raise
       a RuntimeError instead of returning False.
     """
-    if option.lower() in _enabled_options:
+    if isinstance(option, str):
+        option = arm_test_options[option.lower()]
+
+    if option in _test_options and _test_options[option]:
         return True
     else:
         if fail_if_not_enabled:
@@ -96,6 +136,12 @@ def is_option_enabled(option: str, fail_if_not_enabled: bool = False) -> bool:
             return False
 
 
+def get_option(option: arm_test_options) -> Any | None:
+    if option in _test_options:
+        return _test_options[option]
+    return None
+
+
 def maybe_get_tosa_collate_path() -> str | None:
     """
     Checks the environment variable TOSA_TESTCASES_BASE_PATH and returns the
@@ -219,3 +265,32 @@ def get_u85_compile_spec_unbuilt(
         .dump_intermediate_artifacts_to(artifact_path)
     )
     return compile_spec
+
+
+def current_time_formated() -> str:
+    """Return current time as a formated string"""
+    return datetime.now().strftime(get_option(arm_test_options.date_format))
+
+
+def _clean_dir(dir: Path, filter: str, num_save=10):
+    sorted_files: list[tuple[datetime, Path]] = []
+    for file in dir.iterdir():
+        try:
+            creation_time = datetime.strptime(file.name, filter)
+            insert_index = -1
+            for i, to_compare in enumerate(sorted_files):
+                compare_time = to_compare[0]
+                if creation_time < compare_time:
+                    insert_index = i
+                    break
+            if insert_index == -1 and len(sorted_files) < num_save:
+                sorted_files.append((creation_time, file))
+            else:
+                sorted_files.insert(insert_index, (creation_time, file))
+        except ValueError:
+            continue
+
+    if len(sorted_files) > num_save:
+        for remove in sorted_files[0 : len(sorted_files) - num_save]:
+            file = remove[1]
+            file.unlink()

backends/arm/test/tester/arm_tester.py

@@ -22,6 +22,11 @@
     ArmQuantizer,
     get_symmetric_quantization_config,
 )
+from executorch.backends.arm.test.common import (
+    arm_test_options,
+    current_time_formated,
+    get_option,
+)
 
 from executorch.backends.arm.test.runner_utils import (
     _get_input_quantization_params,
@@ -575,6 +580,9 @@ def _get_tosa_operator_distribution(
 
 
 def _dump_str(to_print: str, path_to_dump: Optional[str] = None):
+    default_dump_path = get_option(arm_test_options.dump_path)
+    if not path_to_dump and default_dump_path:
+        path_to_dump = default_dump_path / f"ArmTester_{current_time_formated()}.log"
     if path_to_dump:
         with open(path_to_dump, "a") as fp:
             fp.write(to_print)