chore: import upstream snapshot with attribution

python/sglang/multimodal_gen/test/runner/pytest_runner.py

@@ -0,0 +1,364 @@
+from __future__ import annotations
+
+import subprocess
+import sys
+import xml.etree.ElementTree as ET
+from pathlib import Path
+from typing import Sequence
+
+import tabulate
+
+
+def collect_test_items(
+    files: Sequence[str], filter_expr: str | None = None
+) -> list[str]:
+    """Collect pytest node IDs from the given files or node selectors."""
+    cmd = [sys.executable, "-m", "pytest", "--collect-only", "-q"]
+    if filter_expr:
+        cmd.extend(["-k", filter_expr])
+    cmd.extend(files)
+
+    filter_note = f" with filter: {filter_expr}" if filter_expr else ""
+    print(f"Collecting tests from {len(files)} item(s){filter_note}")
+    result = subprocess.run(cmd, capture_output=True, text=True)
+
+    if result.returncode not in (0, 5):
+        error_msg = (
+            f"pytest --collect-only failed with exit code {result.returncode}\n"
+            f"Command: {' '.join(cmd)}\n"
+        )
+        if result.stderr:
+            error_msg += f"stderr:\n{result.stderr}\n"
+        if result.stdout:
+            error_msg += f"stdout:\n{result.stdout}\n"
+        raise RuntimeError(error_msg)
+
+    if result.returncode == 5:
+        print(
+            "No tests were collected (exit code 5). This may be expected with filters."
+        )
+
+    test_items = []
+    for line in result.stdout.strip().split("\n"):
+        line = line.strip()
+        if line and "::" in line and not line.startswith(("=", "-", " ")):
+            test_id = line.split()[0] if " " in line else line
+            if "::" in test_id:
+                test_items.append(test_id)
+
+    print(f"Collected {len(test_items)} test items")
+    return test_items
+
+
+def parse_junit_xml_for_executed_cases(xml_path: str) -> list[str]:
+    if not Path(xml_path).exists():
+        return []
+
+    executed_cases = []
+    tree = ET.parse(xml_path)
+    root = tree.getroot()
+
+    for testcase in root.iter("testcase"):
+        if testcase.find("skipped") is not None:
+            continue
+
+        name = testcase.get("name", "")
+        if "[" in name and "]" in name:
+            case_id = name[name.index("[") + 1 : name.index("]")]
+            executed_cases.append(case_id)
+
+    return executed_cases
+
+
+def parse_junit_xml_for_case_results(xml_path: str) -> dict[str, str]:
+    if not Path(xml_path).exists():
+        return {}
+
+    case_results = {}
+    tree = ET.parse(xml_path)
+    root = tree.getroot()
+
+    for testcase in root.iter("testcase"):
+        if testcase.find("skipped") is not None:
+            continue
+
+        name = testcase.get("name", "")
+        if "[" not in name or "]" not in name:
+            continue
+
+        case_id = name[name.index("[") + 1 : name.index("]")]
+        if testcase.find("failure") is not None:
+            case_results[case_id] = "fail"
+        elif testcase.find("error") is not None:
+            case_results[case_id] = "error"
+        else:
+            case_results[case_id] = "pass"
+
+    return case_results
+
+
+def _run_pytest_attempt(cmd: list[str]) -> tuple[int, str]:
+    process = subprocess.Popen(
+        cmd,
+        stdout=subprocess.PIPE,
+        stderr=subprocess.STDOUT,
+        bufsize=0,
+    )
+
+    output_bytes = bytearray()
+    while True:
+        chunk = process.stdout.read(4096)
+        if not chunk:
+            break
+        sys.stdout.buffer.write(chunk)
+        sys.stdout.buffer.flush()
+        output_bytes.extend(chunk)
+
+    process.wait()
+    return process.returncode, output_bytes.decode("utf-8", errors="replace")
+
+
+def _extract_collection_line(full_output: str) -> str | None:
+    for line in full_output.splitlines():
+        stripped = line.strip()
+        if stripped.startswith("collected "):
+            return stripped
+    return None
+
+
+def _extract_short_test_summary(full_output: str) -> list[str]:
+    summary_lines = []
+    in_summary = False
+    for line in full_output.splitlines():
+        stripped = line.strip()
+        if "short test summary info" in stripped:
+            in_summary = True
+            continue
+        if not in_summary:
+            continue
+        if stripped.startswith("="):
+            break
+        if not stripped or stripped.startswith("!"):
+            continue
+        summary_lines.append(stripped)
+    return summary_lines
+
+
+def _extract_failure_tail(full_output: str, max_lines: int = 20) -> list[str]:
+    summary_lines = _extract_short_test_summary(full_output)
+    if summary_lines:
+        return summary_lines
+
+    lines = [line.rstrip() for line in full_output.splitlines() if line.strip()]
+    return lines[-max_lines:]
+
+
+def _summary_has_retryable_failure(summary_lines: list[str]) -> bool:
+    for line in summary_lines:
+        lowered = line.lower()
+        if (
+            "[performance]" in line
+            or "SafetensorError" in line
+            or "FileNotFoundError" in line
+            or "TimeoutError" in line
+            or "out of memory" in lowered
+            or "oom killer" in lowered
+        ):
+            return True
+    return False
+
+
+def _is_consistency_failure(full_output: str) -> bool:
+    summary_lines = _extract_short_test_summary(full_output)
+    for line in summary_lines:
+        if "Consistency check failed for" in line or "GT not found for" in line:
+            return True
+
+    return (
+        "Consistency check failed for " in full_output
+        or "GT not found for " in full_output
+        or "--- MISSING GROUND TRUTH DETECTED ---" in full_output
+    )
+
+
+def _is_retryable_failure(full_output: str) -> bool:
+    if _is_consistency_failure(full_output):
+        return False
+
+    summary_lines = _extract_short_test_summary(full_output)
+    is_perf_assertion = (
+        "multimodal_gen/test/server/test_server_utils.py" in full_output
+        and "AssertionError" in full_output
+    )
+    is_aggregated_retryable_failure = _summary_has_retryable_failure(summary_lines)
+
+    is_flaky_ci_assertion = (
+        "SafetensorError" in full_output
+        or "FileNotFoundError" in full_output
+        or "TimeoutError" in full_output
+    )
+
+    is_oom_error = (
+        "out of memory" in full_output.lower() or "oom killer" in full_output.lower()
+    )
+
+    return (
+        is_perf_assertion
+        or is_aggregated_retryable_failure
+        or is_flaky_ci_assertion
+        or is_oom_error
+    )
+
+
+def _print_attempt_tail_summary(
+    attempt_reports: list[dict], assigned_count: int
+) -> None:
+    if len(attempt_reports) == 1 and attempt_reports[0]["returncode"] in (0, 5):
+        return
+
+    rows = []
+    for report in attempt_reports:
+        if report["returncode"] in (0, 5):
+            result = "success"
+        elif report["retryable"]:
+            result = "retryable failure"
+        else:
+            result = "failure"
+        rows.append(
+            [
+                report["attempt"],
+                report["mode"],
+                result,
+                report["collection_line"] or "-",
+            ]
+        )
+
+    print("\n" + "=" * 32 + " Pytest Tail Summary " + "=" * 32, flush=True)
+    print(f"Assigned {assigned_count} test item(s)", flush=True)
+    print(
+        tabulate.tabulate(
+            rows,
+            headers=["Attempt", "Mode", "Result", "Collection"],
+            tablefmt="psql",
+        ),
+        flush=True,
+    )
+
+    for report in attempt_reports:
+        if not report["failure_tail"]:
+            continue
+        print(f"\nAttempt {report['attempt']} failure summary:", flush=True)
+        for line in report["failure_tail"]:
+            print(f"  {line}", flush=True)
+
+    print("=" * 84, flush=True)
+
+
+def run_pytest(
+    files: Sequence[str],
+    filter_expr: str | None = None,
+    junit_xml_path: str | None = None,
+    exitfirst: bool = False,
+) -> tuple[int, list[str], dict[str, str]]:
+    if not files:
+        print("No files to run.")
+        return (0, [], {})
+
+    all_executed_cases: set[str] = set()
+    all_case_results: dict[str, str] = {}
+
+    base_cmd = [
+        sys.executable,
+        "-m",
+        "pytest",
+        "-s",
+        "-v",
+        "--tb=short",
+        "--no-header",
+    ]
+    if exitfirst:
+        base_cmd.append("-x")
+    if junit_xml_path:
+        base_cmd.extend(["--junit-xml", junit_xml_path])
+    if filter_expr:
+        base_cmd.extend(["-k", filter_expr])
+
+    max_retries = 6
+    attempt_reports = []
+    for i in range(max_retries + 1):
+        is_retry = i > 0
+        cmd = list(base_cmd)
+        if is_retry:
+            cmd.append("--last-failed")
+        cmd.extend(files)
+
+        mode = "retry failed items" if is_retry else "initial pass"
+        print(
+            f"Starting pytest attempt {i + 1}/{max_retries + 1}: {mode} "
+            f"for {len(files)} assigned item(s)"
+        )
+
+        returncode, full_output = _run_pytest_attempt(cmd)
+        retryable = returncode not in (0, 5) and _is_retryable_failure(full_output)
+        attempt_reports.append(
+            {
+                "attempt": i + 1,
+                "mode": mode,
+                "returncode": returncode,
+                "retryable": retryable,
+                "collection_line": _extract_collection_line(full_output),
+                "failure_tail": (
+                    _extract_failure_tail(full_output)
+                    if returncode not in (0, 5)
+                    else []
+                ),
+            }
+        )
+
+        if junit_xml_path:
+            all_executed_cases.update(
+                parse_junit_xml_for_executed_cases(junit_xml_path)
+            )
+            all_case_results.update(parse_junit_xml_for_case_results(junit_xml_path))
+
+        if returncode == 0:
+            if is_retry:
+                print(f"Recovered retryable failures on attempt {i + 1}.")
+            _print_attempt_tail_summary(attempt_reports, len(files))
+            return (0, list(all_executed_cases), all_case_results)
+        if returncode == 5:
+            print(
+                "No tests collected (exit code 5). This is expected when filters "
+                "deselect all tests in a partition. Treating as success."
+            )
+            _print_attempt_tail_summary(attempt_reports, len(files))
+            return (0, list(all_executed_cases), all_case_results)
+
+        if not retryable:
+            _print_attempt_tail_summary(attempt_reports, len(files))
+            return (returncode, list(all_executed_cases), all_case_results)
+
+        if i == max_retries:
+            print(f"Max retry exceeded ({max_retries})")
+            _print_attempt_tail_summary(attempt_reports, len(files))
+            return (returncode, list(all_executed_cases), all_case_results)
+
+        print(
+            f"Retryable failure detected on attempt {i + 1}. "
+            "Retrying only previously failed items."
+        )
+
+    _print_attempt_tail_summary(attempt_reports, len(files))
+    return (
+        attempt_reports[-1]["returncode"],
+        list(all_executed_cases),
+        all_case_results,
+    )
+
+
+def partition_items_by_index(
+    items: Sequence[str], partition_id: int, total_partitions: int
+) -> list[str]:
+    return [
+        item for i, item in enumerate(items) if i % total_partitions == partition_id
+    ]