[TRTLLM-4866] [test] Support waiving unit tests by waives.txt (NVIDIA#8359)

Signed-off-by: Xiwen Yu <13230610+VALLIS-NERIA@users.noreply.github.com>

Author: VALLIS-NERIA

tests/integration/defs/test_list_parser.py

@@ -83,7 +83,7 @@ def applyMarkerIfPrefix(marker, prefix, reason):
     return lines
 
 
-def parse_test_list_lines(test_list, lines, test_prefix):
+def parse_test_list_lines(test_list, lines, test_prefix, convert_unittest=True):
     """Parses the lines of a test list. Test names returned contain all values within square brackets. Does not process
     each test id value.
 
@@ -146,19 +146,21 @@ def parse_test_line(enumerated_line):
                                  test_list, lineno, reason_raw))
                 break
 
-        # extract full:XXX/ prefix
-        full_prefix = ""
-        match = re.match(r'(full:.*?/)(.+)', test_name)
-        if match:
-            full_prefix = match.group(1)
-            test_name = match.group(2)
+        if convert_unittest:
+            # extract full:XXX/ prefix
+            full_prefix = ""
+            match = re.match(r'(full:.*?/)(.+)', test_name)
+            if match:
+                full_prefix = match.group(1)
+                test_name = match.group(2)
 
-        # convert unittest to actual test name
-        if test_name.startswith("unittest/"):
-            test_name = f"test_unittests.py::test_unittests_v2[{test_name}]"
+            # convert unittest to actual test name
+            if test_name.startswith("unittest/"):
+                test_name = f"test_unittests.py::test_unittests_v2[{test_name}]"
+
+            # combine back
+            test_name = full_prefix + test_name
 
-        # combine back
-        test_name = full_prefix + test_name
         test_name = parse_test_name(test_name)
 
         return (test_name, marker, reason, timeout)
@@ -585,8 +587,6 @@ def parse_and_validate_test_list(
     check_for_corrections,
 ):
     test_prefix = config.getoption("--test-prefix")
-    apply_test_list_correction = config.getoption(
-        "--apply-test-list-correction")
     test_names, test_name_to_marker_dict = parse_test_list(
         test_list, test_prefix)
 
@@ -597,6 +597,8 @@ def parse_and_validate_test_list(
             TestCorrectionMode.EXACT_MATCH,
         )
 
+        apply_test_list_correction = config.getoption(
+            "--apply-test-list-correction")
         if apply_test_list_correction and corrections:
             apply_test_list_corrections(test_list, corrections, items,
                                         test_prefix)

tests/integration/defs/test_unittests.py

@@ -76,6 +76,8 @@ def test_unittests_v2(llm_root, llm_venv, case: str, output_dir, request):
     else:
         test_prefix = "unittest"
 
+    waives_file = request.config.getoption("--waives-file")
+
     num_workers = 1
 
     # This dataframe is not manually edited. Infra team will regularly generate this dataframe based on test execution results.
@@ -132,9 +134,13 @@ def test_unittests_v2(llm_root, llm_venv, case: str, output_dir, request):
     if dry_run:
         command += ['--collect-only']
 
+    if waives_file:
+        waives_file = os.path.abspath(waives_file)
+        command += [f"--waives-file={waives_file}"]
+
     command += arg_list
 
-    print(f"Running unit test:'{command}'")
+    print(f"Running unit test:\"python {' '.join(command)}\"")
 
     def run_command(cmd, num_workers=1):
         try:

tests/unittest/conftest.py

@@ -16,6 +16,7 @@
 import os
 import sys
 import traceback
+from functools import partial
 from typing import Any
 
 import _pytest.outcomes
@@ -24,6 +25,9 @@
 import tqdm
 from mpi4py.futures import MPIPoolExecutor
 
+sys.path.append(os.path.dirname(os.path.dirname(os.path.abspath(__file__))))
+from integration.defs import test_list_parser
+
 
 def pytest_configure(config):
     # avoid thread leak of tqdm's TMonitor
@@ -87,11 +91,84 @@ def pytest_addoption(parser):
         default=False,
         help="Run Ray-marked tests (by default they are skipped).",
     )
+    parser.addoption(
+        "--waives-file",
+        "-S",
+        action="store",
+        default=None,
+        help=
+        "Specify a file containing a list of waives, one per line. After filtering collected tests, Pytest will "
+        "apply the waive state specified by this file to the set of tests to be run.",
+    )
+
+
+def apply_waives_ut(waives_file, items: list[pytest.Item], config):
+    """Apply waives based on the waive state specified by the given waives_file."""
+
+    # Corrections don't make sense for the waives file as it specifies global negative
+    # filters that may or may not be applicable to the current platform (i.e., the test names
+    # being waived may not be generated on the current platform).
+    try:
+        parse_test_list_lines_bak = test_list_parser.parse_test_list_lines
+        test_list_parser.parse_test_list_lines = partial(
+            test_list_parser.parse_test_list_lines, convert_unittest=False)
+        ret = test_list_parser.parse_and_validate_test_list(
+            waives_file,
+            config,
+            items,
+            check_for_corrections=False,
+        )
+    finally:
+        test_list_parser.parse_test_list_lines = parse_test_list_lines_bak
+    if not ret:
+        return
+    _, test_name_to_marker_dict = ret
+
+    filtered_dict = {}
+    for waiver in test_name_to_marker_dict.keys():
+        if "unittest/" not in waiver:
+            continue
+        elif "unittest/unittest/" in waiver:
+            filtered_dict[waiver.replace(
+                "unittest/unittest/",
+                "unittest/")] = test_name_to_marker_dict[waiver]
+        else:
+            filtered_dict[waiver] = test_name_to_marker_dict[waiver]
+
+    # Fuzzy match is supported in the following order:
+    # 1. exact match
+    # 2. remove parameterization part in square brackets and try again (function level)
+    # 3. remove the last part after '::' and try again, until no '::' left (file level)
+    # Note: directory level match is not supported.
+    def match_waiver(id: str):
+        if id in filtered_dict:
+            return filtered_dict[id]
+        if id.endswith("]"):
+            id = id.split("[")[0]
+            if id in filtered_dict:
+                return filtered_dict[id]
+        while "::" in id:
+            id = id.rsplit("::", 1)[0]
+            if id in filtered_dict:
+                return filtered_dict[id]
+
+        return None
+
+    # For each item in the list, apply waives if a waive entry exists
+    for item in items:
+        waiver = match_waiver(item.nodeid)
+        if waiver:
+            marker, reason, _ = waiver
+            if marker:
+                mark_func = getattr(pytest.mark, marker.lower())
+                mark = mark_func(reason=reason)
+                item.add_marker(mark)
 
 
 @pytest.hookimpl(tryfirst=True, hookwrapper=True)
 def pytest_collection_modifyitems(session, config, items):
     test_prefix = config.getoption("--test-prefix")
+    waives_file = config.getoption("--waives-file")
 
     yield
 
@@ -102,6 +179,8 @@ def pytest_collection_modifyitems(session, config, items):
         for item in items:
             item._nodeid = f"{test_prefix}/{item._nodeid}"
 
+    if waives_file:
+        apply_waives_ut(waives_file, items, config)
     # Ray tests are disabled by default
     run_ray = config.getoption("--run-ray") or os.environ.get(
         "TLLM_RUN_RAY_TESTS") == "1"