[ci] Balance integration test buckets by recorded durations (#18895)

script/determine-jobs.py

@@ -53,8 +53,10 @@ from collections import Counter
 from enum import StrEnum
 from functools import cache
 import json
+import math
 import os
 from pathlib import Path
+import statistics
 import sys
 from typing import Any
 
@@ -67,7 +69,9 @@ from clang_tidy_hash import (
 from helpers import (
     CPP_FILE_EXTENSIONS,
     ESPHOME_TESTS_COMPONENTS_PATH,
+    INTEGRATION_TESTS_PATH,
     PYTHON_FILE_EXTENSIONS,
+    all_integration_test_files,
     base_python_changed,
     changed_files,
     core_changed,
@@ -83,6 +87,8 @@ from helpers import (
     get_target_branch,
     git_ls_files,
     is_validate_only_file,
+    load_integration_durations,
+    lpt_partition,
     root_path,
 )
 from split_components_for_ci import create_intelligent_batches
@@ -96,10 +102,13 @@ CLANG_TIDY_SPLIT_THRESHOLD = 65
 # Isolated components count as 10x, groupable components count as 1x
 COMPONENT_TEST_BATCH_SIZE = 40
 
-# Integration test bucketing: when more than the threshold tests are scheduled,
-# fan out across this many parallel jobs. Below the threshold, a single job runs.
+# Above the threshold, fan out across up to this many jobs, balanced by the
+# recorded per-file durations. The target is serial junit-time weight per
+# bucket, not wall time (calibrated with the conftest compile cap); it
+# sizes the bucket count for small subsets.
 INTEGRATION_TESTS_SPLIT_THRESHOLD = 10
-INTEGRATION_TESTS_SPLIT_BUCKETS = 3
+INTEGRATION_TESTS_SPLIT_BUCKETS = 5
+INTEGRATION_TESTS_TARGET_BUCKET_WEIGHT = 360.0
 
 # platformio and aioesphomeapi (requirements.txt), the pytest stack
 # (requirements_test.txt) and the fixture every session compiles; a change
@@ -113,27 +122,13 @@ INTEGRATION_TESTS_TRIGGER_FILES = frozenset(
 )
 
 
-def _split_list(items: list[str], n: int) -> list[list[str]]:
-    """Split a list into n roughly-equal contiguous parts (matches script/clang-tidy)."""
-    k, m = divmod(len(items), n)
-    return [items[i * k + min(i, m) : (i + 1) * k + min(i + 1, m)] for i in range(n)]
-
-
-def _all_integration_test_files() -> list[str]:
-    """Return all integration test file paths, sorted, relative to repo root."""
-    return sorted(
-        str(p.relative_to(root_path))
-        for p in (Path(root_path) / "tests" / "integration").glob("test_*.py")
-    )
-
-
 def _compute_integration_test_buckets(
     integration_run_all: bool,
     integration_test_files: list[str],
 ) -> tuple[bool, list[dict[str, Any]]]:
     """Compute (run_integration, buckets) from the determine_integration_tests result.
 
-    Pure function for unit testing — no I/O beyond `_all_integration_test_files`
+    Pure function for unit testing — no I/O beyond `all_integration_test_files`
     when `integration_run_all` is set.
 
     `buckets` is a list of `{name, tests}` dicts where `tests` is a JSON-friendly
@@ -141,7 +136,7 @@ def _compute_integration_test_buckets(
     shell word-splitting / glob hazards.
     """
     if integration_run_all:
-        files = _all_integration_test_files()
+        files = all_integration_test_files()
     else:
         files = sorted(integration_test_files)
 
@@ -152,12 +147,23 @@ def _compute_integration_test_buckets(
         return False, []
 
     if len(files) > INTEGRATION_TESTS_SPLIT_THRESHOLD:
-        parts = [
-            part for part in _split_list(files, INTEGRATION_TESTS_SPLIT_BUCKETS) if part
-        ]
+        durations = load_integration_durations()
+        # Unrecorded files weigh the recording's median; with no recording a
+        # file weighs a whole bucket, which keeps the full fan-out
+        default = (
+            statistics.median(durations.values())
+            if durations
+            else INTEGRATION_TESTS_TARGET_BUCKET_WEIGHT
+        )
+        weights = {f: durations.get(f, default) for f in files}
+        count = min(
+            INTEGRATION_TESTS_SPLIT_BUCKETS,
+            math.ceil(sum(weights.values()) / INTEGRATION_TESTS_TARGET_BUCKET_WEIGHT),
+        )
+        # count <= SPLIT_BUCKETS < threshold < len(files): no group is empty
+        parts = [sorted(part) for part in lpt_partition(files, weights, count)]
         buckets = [
-            {"name": f"{i + 1}/{len(parts)}", "tests": part}
-            for i, part in enumerate(parts)
+            {"name": f"{i + 1}/{count}", "tests": part} for i, part in enumerate(parts)
         ]
     else:
         buckets = [{"name": "1/1", "tests": files}]
@@ -264,9 +270,9 @@ def determine_integration_tests(branch: str | None = None) -> tuple[bool, list[s
     # If infrastructure Python files changed (conftest, utils, etc.), run all tests
     # Excludes test files (test_*.py), fixtures, and non-Python files (README.md)
     if any(
-        f.startswith("tests/integration/")
+        f.startswith(INTEGRATION_TESTS_PATH)
         and f.endswith(".py")
-        and not f.startswith("tests/integration/test_")
+        and not f.startswith(f"{INTEGRATION_TESTS_PATH}test_")
         and "/fixtures/" not in f
         for f in files
     ):
@@ -277,9 +283,9 @@ def determine_integration_tests(branch: str | None = None) -> tuple[bool, list[s
     fixture_to_test_files = get_fixture_to_test_files()
 
     for f in files:
-        if f.startswith("tests/integration/test_") and f.endswith(".py"):
+        if f.startswith(f"{INTEGRATION_TESTS_PATH}test_") and f.endswith(".py"):
             test_files.add(f)
-        elif f.startswith("tests/integration/fixtures/"):
+        elif f.startswith(f"{INTEGRATION_TESTS_PATH}fixtures/"):
             if f.endswith(".yaml"):
                 # Fixture YAML changed - add corresponding test file(s)
                 test_files.update(fixture_to_test_files.get(Path(f).stem, ()))
@@ -1415,6 +1421,7 @@ def main() -> None:
     output: dict[str, Any] = {
         "core_ci": run_core_ci,
         "integration_tests": run_integration,
+        "integration_run_all": integration_run_all,
         "integration_test_buckets": integration_test_buckets,
         "clang_tidy": run_clang_tidy,
         "clang_tidy_mode": clang_tidy_mode,

script/helpers.py

@@ -43,6 +43,53 @@ ESPHOME_TESTS_COMPONENTS_PATH = "tests/components/"
 # Tuple of component and test paths for efficient startswith checks
 COMPONENT_AND_TESTS_PATHS = (ESPHOME_COMPONENTS_PATH, ESPHOME_TESTS_COMPONENTS_PATH)
 
+# Integration tests path prefix
+INTEGRATION_TESTS_PATH = "tests/integration/"
+
+# Per-file integration test durations from CI junit output; shared by the
+# reader (determine-jobs) and writer (update_integration_test_durations)
+INTEGRATION_TEST_DURATIONS_FILE = "tests/integration/integration_test_durations.json"
+
+
+def all_integration_test_files() -> list[str]:
+    """Return all integration test file paths, sorted, relative to repo root."""
+    return sorted(
+        p.relative_to(root_path).as_posix()
+        for p in (Path(root_path) / "tests" / "integration").glob("test_*.py")
+    )
+
+
+def load_integration_durations() -> dict[str, float]:
+    """Return recorded per-file pytest durations in seconds; empty when unavailable."""
+    try:
+        raw = json.loads(
+            (Path(root_path) / INTEGRATION_TEST_DURATIONS_FILE).read_text()
+        )
+        if not isinstance(raw, dict):
+            print(
+                f"integration durations unavailable: expected an object, "
+                f"got {type(raw).__name__}",
+                file=sys.stderr,
+            )
+            return {}
+    except (OSError, ValueError) as err:
+        # The file ships in the repo; degrade to unweighted bucketing, loudly
+        print(f"integration durations unavailable: {err}", file=sys.stderr)
+        return {}
+    durations = {
+        key: seconds
+        for key, value in raw.items()
+        if isinstance(value, (int, float)) and (seconds := float(value)) > 0
+    }
+    if len(durations) != len(raw):
+        # One bad entry must not discard the whole recording
+        print(
+            f"dropped {len(raw) - len(durations)} invalid duration entries",
+            file=sys.stderr,
+        )
+    return durations
+
+
 # Base bus components - these ARE the bus implementations and should not
 # be flagged as needing migration since they are the platform/base components
 BASE_BUS_COMPONENTS = {
@@ -1545,3 +1592,21 @@ def get_cpp_changed_components(files: list[str]) -> list[str]:
         if file.startswith(ESPHOME_COMPONENTS_PATH):
             affected.update(find_children_of_component(components_graph, component))
     return sorted(c for c in affected if has_cpp_unit_tests(c, tests_dir))
+
+
+def lpt_partition(
+    items: list[str], weights: dict[str, float], count: int
+) -> list[list[str]]:
+    """Partition items into `count` weight-balanced groups (LPT greedy).
+
+    Heaviest item first into the lightest group. Ties keep input order, so
+    pass pre-sorted items for deterministic output. script/clang-tidy's
+    split_list is the unweighted contiguous sibling.
+    """
+    groups: list[list[str]] = [[] for _ in range(count)]
+    group_weights = [0.0] * count
+    for item in sorted(items, key=lambda i: -weights[i]):
+        lightest = min(range(count), key=group_weights.__getitem__)
+        groups[lightest].append(item)
+        group_weights[lightest] += weights[item]
+    return groups

script/update_integration_test_durations.py

@@ -0,0 +1,119 @@
+#!/usr/bin/env python3
+"""Merge CI junit output into tests/integration/integration_test_durations.json.
+
+The integration-tests CI job uploads one junit XML artifact per bucket on
+full matrix dev runs. Download a run's artifacts and merge the per file
+durations into the recording used by script/determine-jobs.py:
+
+    gh run download <run-id> --repo esphome/esphome -p "junit-integration-*" -D /tmp/junit
+    script/update_integration_test_durations.py /tmp/junit
+
+Missing files keep their previous recording and deleted files drop out; a
+run covering under 90% of the test files aborts unless --allow-partial.
+"""
+
+from __future__ import annotations
+
+import argparse
+from collections import defaultdict
+import json
+from pathlib import Path
+import sys
+import xml.etree.ElementTree as ET
+
+from helpers import (
+    INTEGRATION_TEST_DURATIONS_FILE,
+    INTEGRATION_TESTS_PATH,
+    all_integration_test_files,
+    load_integration_durations,
+    root_path,
+)
+
+DURATIONS_FILE = Path(root_path) / INTEGRATION_TEST_DURATIONS_FILE
+MIN_COVERAGE = 0.9
+# Exit code for the expected "run covers too few files" refusal, so the
+# refresh workflow can move on to the next candidate run
+EXIT_LOW_COVERAGE = 3
+
+
+def collect_durations(junit_dir: Path, known_files: set[str]) -> dict[str, float]:
+    """Sum junit testcase times per integration test file, in seconds."""
+    durations: defaultdict[str, float] = defaultdict(float)
+    unmatched = 0
+    xml_files = sorted(junit_dir.rglob("*.xml"))
+    if not xml_files:
+        raise SystemExit(f"no junit XML files found under {junit_dir}")
+    for xml_file in xml_files:
+        for testcase in ET.parse(xml_file).getroot().iter("testcase"):
+            # Skipped/errored testcases carry time="0"; recording them would
+            # overwrite a good previous duration
+            if any(
+                testcase.find(tag) is not None
+                for tag in ("skipped", "error", "failure")
+            ):
+                continue
+            # classname is the dotted module plus any test class, e.g.
+            # tests.integration.test_x or tests.integration.test_x.TestFoo
+            parts = testcase.get("classname", "").split(".")
+            if parts[:2] != ["tests", "integration"] or len(parts) < 3:
+                unmatched += 1
+                continue
+            path = f"{INTEGRATION_TESTS_PATH}{parts[2]}.py"
+            if path not in known_files:
+                print(f"skipping unknown test module {path}", file=sys.stderr)
+                continue
+            durations[path] += float(testcase.get("time", "0"))
+    if unmatched:
+        # A junit naming change would otherwise shrink the recording silently
+        raise SystemExit(
+            f"{unmatched} testcases with unexpected classnames; the junit layout changed"
+        )
+    # An all-skipped file totals 0.0; let the merge keep its previous entry
+    return {k: v for k, v in durations.items() if v > 0}
+
+
+def main() -> int:
+    parser = argparse.ArgumentParser(description=__doc__)
+    parser.add_argument(
+        "junit_dir", type=Path, help="directory containing downloaded junit XML files"
+    )
+    parser.add_argument(
+        "--allow-partial",
+        action="store_true",
+        help="merge a run covering under 90%% of the test files",
+    )
+    args = parser.parse_args()
+
+    on_disk = set(all_integration_test_files())
+    if not on_disk:
+        raise SystemExit("no integration test files found; wrong checkout root?")
+    collected = collect_durations(args.junit_dir, on_disk)
+    coverage = len(collected.keys() & on_disk) / len(on_disk)
+    if coverage < MIN_COVERAGE and not args.allow_partial:
+        print(
+            f"artifacts cover only {coverage:.0%} of {len(on_disk)} test files; "
+            "use a full matrix run or pass --allow-partial to merge anyway",
+            file=sys.stderr,
+        )
+        return EXIT_LOW_COVERAGE
+
+    # Validated load: a bad previous entry cannot survive the round trip, and
+    # an unreadable file aborts rather than being overwritten
+    previous = load_integration_durations()
+    if DURATIONS_FILE.is_file() and not previous:
+        raise SystemExit(f"{DURATIONS_FILE} is unreadable; refusing to overwrite it")
+    # New recordings win, absent files keep theirs, deleted files drop out
+    merged = {
+        path: collected.get(path, previous.get(path))
+        for path in sorted(on_disk)
+        if path in collected or path in previous
+    }
+    DURATIONS_FILE.write_text(
+        json.dumps({k: round(v, 2) for k, v in merged.items()}, indent=2) + "\n"
+    )
+    print(f"wrote {len(merged)} entries to {DURATIONS_FILE} ({coverage:.0%} fresh)")
+    return 0
+
+
+if __name__ == "__main__":
+    sys.exit(main())