[ci] Balance component test batches by estimated build time (#19980)

Co-authored-by: Jonathan Swoboda <154711427+swoboda1337@users.noreply.github.com>

script/determine-jobs.py

@@ -99,10 +99,6 @@ from split_components_for_ci import create_intelligent_batches
 # For large PRs (>= 65 files), use split for better parallelization
 CLANG_TIDY_SPLIT_THRESHOLD = 65
 
-# Component test batch size (weighted)
-# Isolated components count as 10x, groupable components count as 1x
-COMPONENT_TEST_BATCH_SIZE = 40
-
 # 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
@@ -1534,7 +1530,6 @@ def main() -> None:
         batches, _ = create_intelligent_batches(
             components=changed_components_with_tests,
             tests_dir=tests_dir,
-            batch_size=COMPONENT_TEST_BATCH_SIZE,
             directly_changed=batch_directly_changed,
         )
         # Convert batches to CI matrix entries: the component list plus which

script/split_components_for_ci.py

@@ -13,7 +13,9 @@ from __future__ import annotations
 
 import argparse
 from collections import defaultdict
+from dataclasses import dataclass, field
 import json
+import math
 from pathlib import Path
 import sys
 
@@ -28,13 +30,34 @@ from script.analyze_component_buses import (
     create_grouping_signature,
     merge_compatible_bus_groups,
 )
-from script.helpers import get_component_test_files, split_conflicting_groups
+from script.helpers import (
+    get_component_test_files,
+    parse_test_filename,
+    split_conflicting_groups,
+)
 
-# Weighting for batch creation
-# Isolated components can't be grouped/merged, so they count as 10x
-# Groupable components can be merged into single builds, so they count as 1x
-ISOLATED_WEIGHT = 10
-GROUPABLE_WEIGHT = 1
+# Estimated CI seconds for one build per test platform, read from batch job
+# logs. Absolute accuracy does not matter, only the ratios between platforms.
+# A full run takes over an hour, so recording real timings is not practical.
+PLATFORM_BUILD_SECONDS = {
+    "host": 10,
+    "esp8266-ard": 30,
+    "esp32-idf": 45,
+    "rtl87xx-ard": 35,
+    "ln882x-ard": 50,
+    "rp2040-ard": 50,
+    "rp2350-ard": 50,
+    "bk72xx-ard": 60,
+}
+NRF52_BUILD_SECONDS = 45
+# Other ESP32 chips and Arduino on ESP32 share less of the ccache
+DEFAULT_BUILD_SECONDS = 90
+# Each extra component merged into a grouped build makes it larger
+GROUPED_COMPONENT_SECONDS = 5
+# Estimated build seconds per CI runner; sets how many runners are used.
+# Approximate: the runner count charges each grouped build once, but a large
+# group spreads over several runners, which each pay for that build.
+TARGET_BATCH_SECONDS = 600
 
 # Platform used for batching (platform-agnostic batching)
 # Batches are split across CI runners and each runner tests all platforms
@@ -61,10 +84,115 @@ def has_test_files(component_name: str, tests_dir: Path) -> bool:
     )
 
 
+def build_seconds(platform: str) -> int:
+    """Return the estimated CI seconds for one build on a test platform."""
+    if (seconds := PLATFORM_BUILD_SECONDS.get(platform)) is not None:
+        return seconds
+    if platform.startswith("nrf52"):
+        return NRF52_BUILD_SECONDS
+    return DEFAULT_BUILD_SECONDS
+
+
+@dataclass
+class _BatchItem:
+    """A component and the builds it adds to the batch it lands in."""
+
+    component: str
+    # Builds that always run on their own (isolated tests, variants on a
+    # platform without a base test)
+    own_seconds: int
+    # Per (signature, platform) shared build: the seconds of this component's
+    # files on that platform when no other batch member shares the build
+    grouped_builds: dict[tuple[str, str], int] = field(default_factory=dict)
+
+    def standalone_seconds(self) -> int:
+        return self.own_seconds + sum(self.grouped_builds.values())
+
+
+def _grouped_build_seconds(platform: str, solo_seconds: list[int]) -> int:
+    """Return the seconds of one (signature, platform) build in a batch.
+
+    test_build_components only groups when two or more members share the
+    build; a grouped member then skips its variants on that platform.
+    """
+    if len(solo_seconds) <= 1:
+        return sum(solo_seconds)
+    return build_seconds(platform) + GROUPED_COMPONENT_SECONDS * (len(solo_seconds) - 1)
+
+
+def _make_item(
+    tests_dir: Path, component: str, signature: str, is_isolated: bool
+) -> _BatchItem:
+    files_by_platform: dict[str, list[str]] = defaultdict(list)
+    for test_file in (tests_dir / component).glob("test[.-]*.yaml"):
+        test_name, platform = parse_test_filename(test_file)
+        files_by_platform[platform].append(test_name)
+    own_seconds = 0
+    grouped: dict[tuple[str, str], int] = {}
+    for platform, test_names in files_by_platform.items():
+        seconds = build_seconds(platform) * len(test_names)
+        if is_isolated or "test" not in test_names:
+            own_seconds += seconds
+        else:
+            grouped[(signature, platform)] = seconds
+    return _BatchItem(component, own_seconds, grouped)
+
+
+class _Batch:
+    """A CI runner's components and its estimated build seconds."""
+
+    def __init__(self) -> None:
+        self.components: list[str] = []
+        self.seconds = 0
+        self.grouped_builds: dict[tuple[str, str], list[int]] = defaultdict(list)
+
+    def added_seconds(self, item: _BatchItem) -> int:
+        """Return the seconds item would add; joining an existing build is cheap."""
+        added = item.own_seconds
+        for (signature, platform), solo in item.grouped_builds.items():
+            members = self.grouped_builds.get((signature, platform), [])
+            added += _grouped_build_seconds(
+                platform, [*members, solo]
+            ) - _grouped_build_seconds(platform, members)
+        return added
+
+    def add(self, item: _BatchItem) -> None:
+        self.seconds += self.added_seconds(item)
+        for build, solo in item.grouped_builds.items():
+            self.grouped_builds[build].append(solo)
+        self.components.append(item.component)
+
+
+def balance_batches(items: list[_BatchItem], target_seconds: int) -> list[list[str]]:
+    """Spread items over enough runners to stay near target_seconds each.
+
+    The runner count comes from the total estimate with every grouped build
+    counted once. Heaviest items go first, each to the runner that ends up
+    lightest, so grouped components follow the builds they can join.
+    """
+    if not items:
+        return []
+    grouped_builds: dict[tuple[str, str], list[int]] = defaultdict(list)
+    for item in items:
+        for build, solo in item.grouped_builds.items():
+            grouped_builds[build].append(solo)
+    total = sum(item.own_seconds for item in items) + sum(
+        _grouped_build_seconds(platform, solos)
+        for (_, platform), solos in grouped_builds.items()
+    )
+    count = min(len(items), max(1, math.ceil(total / target_seconds)))
+    batches = [_Batch() for _ in range(count)]
+    for item in sorted(items, key=lambda i: (-i.standalone_seconds(), i.component)):
+        min(batches, key=lambda b: (b.seconds + b.added_seconds(item), b.seconds)).add(
+            item
+        )
+    return [batch.components for batch in batches if batch.components]
+
+
 def create_intelligent_batches(
     components: list[str],
     tests_dir: Path,
-    batch_size: int = 40,
+    target_seconds: int = TARGET_BATCH_SECONDS,
     directly_changed: set[str] | None = None,
 ) -> tuple[list[list[str]], dict[tuple[str, str], list[str]]]:
     """Create batches optimized for component grouping.
@@ -76,7 +204,7 @@ def create_intelligent_batches(
     Args:
         components: List of component names to batch
         tests_dir: Path to tests/components directory
-        batch_size: Target size for each batch
+        target_seconds: Estimated build seconds per batch
         directly_changed: Set of directly changed components (for logging only)
 
     Returns:
@@ -157,51 +285,17 @@ def create_intelligent_batches(
     # actually be split into two at build time -- throwing off CI distribution.
     signature_groups = split_conflicting_groups(signature_groups)
 
-    # Create batches by keeping signature groups together
-    # Components with the same signature stay in the same batches
-    batches = []
-
-    # Sort signature groups to prioritize groupable components
-    # 1. Put "isolated_*" signatures last (can't be grouped with others)
-    # 2. Sort groupable signatures by size (largest first)
-    # 3. "no_buses" components CAN be grouped together
-    def sort_key(item):
-        (_platform, signature), components = item
-        is_isolated = signature.startswith(ISOLATED_SIGNATURE_PREFIX)
-        # Put "isolated_*" last (1), groupable first (0)
-        # Within each category, sort by size (largest first)
-        return (is_isolated, -len(components))
-
-    sorted_groups = sorted(signature_groups.items(), key=sort_key)
-
-    # Strategy: Create batches using weighted sizes
-    # - Isolated components count as 10x (since they can't be grouped/merged)
-    # - Groupable components count as 1x (can be merged into single builds)
-    # - This distributes isolated components across more runners
-    # - Ensures each runner has a good mix of groupable vs isolated components
-
-    current_batch = []
-    current_weight = 0
-
-    for (_platform, signature), group_components in sorted_groups:
-        is_isolated = signature.startswith(ISOLATED_SIGNATURE_PREFIX)
-        weight_per_component = ISOLATED_WEIGHT if is_isolated else GROUPABLE_WEIGHT
-
-        for component in group_components:
-            # Check if adding this component would exceed the batch size
-            if current_weight + weight_per_component > batch_size and current_batch:
-                # Start a new batch
-                batches.append(current_batch)
-                current_batch = []
-                current_weight = 0
-
-            # Add component to current batch
-            current_batch.append(component)
-            current_weight += weight_per_component
-
-    # Don't forget the last batch
-    if current_batch:
-        batches.append(current_batch)
+    items = [
+        _make_item(
+            tests_dir,
+            component,
+            signature,
+            signature.startswith(ISOLATED_SIGNATURE_PREFIX),
+        )
+        for (_platform, signature), group_components in sorted(signature_groups.items())
+        for component in group_components
+    ]
+    batches = balance_batches(items, target_seconds)
 
     return batches, signature_groups
 
@@ -218,11 +312,11 @@ def main() -> int:
         help="JSON array of component names",
     )
     parser.add_argument(
-        "--batch-size",
-        "-b",
+        "--target-seconds",
+        "-t",
         type=int,
-        default=40,
-        help="Target batch size (default: 40, weighted)",
+        default=TARGET_BATCH_SECONDS,
+        help=f"Estimated build seconds per batch (default: {TARGET_BATCH_SECONDS})",
     )
     parser.add_argument(
         "--tests-dir",
@@ -268,7 +362,7 @@ def main() -> int:
     batches, signature_groups = create_intelligent_batches(
         components=components,
         tests_dir=args.tests_dir,
-        batch_size=args.batch_size,
+        target_seconds=args.target_seconds,
         directly_changed=directly_changed,
     )
 
@@ -378,15 +472,15 @@ def main() -> int:
             file=sys.stderr,
         )
 
-    print(f"  - Groupable (weight=1): {groupable_count}", file=sys.stderr)
-    print(f"  - Isolated (weight=10): {isolated_count}", file=sys.stderr)
+    print(f"  - Groupable: {groupable_count}", file=sys.stderr)
+    print(f"  - Isolated: {isolated_count}", file=sys.stderr)
     if actual_components < len(components):
         print(
             f"Components skipped (no test files): {len(components) - actual_components}",
             file=sys.stderr,
         )
     print(f"Number of batches: {len(batches)}", file=sys.stderr)
-    print(f"Batch size target (weighted): {args.batch_size}", file=sys.stderr)
+    print(f"Target build seconds per batch: {args.target_seconds}", file=sys.stderr)
     if len(batches) > 0:
         print(
             f"Average components per batch: {actual_components / len(batches):.1f}",

tests/script/test_determine_jobs.py

@@ -2594,7 +2594,7 @@ def test_detect_platform_hint_from_filename_case_insensitive(
     assert result == expected_platform
 
 
-def test_component_batching_beta_branch_40_per_batch(
+def test_component_batching_beta_branch_groups_evenly(
     tmp_path: Path,
     mock_determine_integration_tests: Mock,
     mock_should_run_clang_tidy: Mock,
@@ -2604,11 +2604,10 @@ def test_component_batching_beta_branch_40_per_batch(
     mock_determine_cpp_unit_tests: Mock,
     capsys: pytest.CaptureFixture[str],
 ) -> None:
-    """Test that beta/release branches create batches with 40 actual components each.
+    """Test that beta/release branches group every component and split evenly.
 
     For beta/release branches, all components should be groupable (not isolated),
-    and each batch should contain 40 actual components with weight 1 each.
-    This matches the original behavior before consolidation.
+    so they share one grouped build per runner and spread evenly.
     """
     # Create 120 test components with test files
     component_names = [f"comp_{i:03d}" for i in range(120)]
@@ -2672,15 +2671,9 @@ def test_component_batching_beta_branch_40_per_batch(
     assert "component_test_batches" in output
     batches = output["component_test_batches"]
 
-    # Should have 3 batches (120 components / 40 per batch = 3)
-    assert len(batches) == 3, f"Expected 3 batches, got {len(batches)}"
-
-    # Each batch should have approximately 40 components (all weight=1, groupable)
-    for i, batch in enumerate(batches):
-        batch_components = batch["components"].split()
-        assert len(batch_components) == 40, (
-            f"Batch {i} should have 40 components, got {len(batch_components)}"
-        )
+    # One grouped esp32-idf build plus 5 s per extra component is 640 s,
+    # so two runners of 60 components each
+    assert [len(batch["components"].split()) for batch in batches] == [60, 60]
 
     # Verify all 120 components are in batches
     all_components = []

tests/script/test_split_components_for_ci.py

@@ -0,0 +1,157 @@
+"""Unit tests for script/split_components_for_ci.py module."""
+
+from collections.abc import Generator
+from pathlib import Path
+import sys
+from unittest.mock import patch
+
+import pytest
+
+script_dir = str((Path(__file__).parent / ".." / ".." / "script").resolve())
+sys.path.insert(0, script_dir)
+
+import split_components_for_ci  # noqa: E402
+
+import script.helpers  # noqa: E402
+
+
+@pytest.fixture
+def tests_dir(tmp_path: Path) -> Generator[Path, None, None]:
+    """Point the helpers at an empty tests/components tree."""
+    path = tmp_path / "tests" / "components"
+    path.mkdir(parents=True)
+    with patch.object(script.helpers, "root_path", str(tmp_path)):
+        yield path
+
+
+def _add_component(tests_dir: Path, name: str, files: list[str]) -> None:
+    comp_dir = tests_dir / name
+    comp_dir.mkdir()
+    for file in files:
+        (comp_dir / file).write_text(f"# {name}\n")
+
+
+def test_build_seconds() -> None:
+    """Known platforms use the table; nrf52 boards share one estimate."""
+    build_seconds = split_components_for_ci.build_seconds
+    assert build_seconds("host") == 10
+    assert build_seconds("esp32-idf") == 45
+    assert (
+        build_seconds("nrf52-xiao-ble") == split_components_for_ci.NRF52_BUILD_SECONDS
+    )
+    assert (
+        build_seconds("esp32-s3-idf") == split_components_for_ci.DEFAULT_BUILD_SECONDS
+    )
+
+
+def test_make_item_isolated_and_grouped(tests_dir: Path) -> None:
+    """Isolated components own every build; groupable ones join shared builds."""
+    _add_component(
+        tests_dir,
+        "comp",
+        [
+            "test.esp32-idf.yaml",
+            "test.host.yaml",
+            "test-extra.esp32-idf.yaml",
+            "test-only.rp2040-ard.yaml",
+            "validate.esp32-idf.yaml",
+        ],
+    )
+
+    isolated = split_components_for_ci._make_item(
+        tests_dir, "comp", "isolated_comp", True
+    )
+    assert isolated.own_seconds == 45 + 10 + 45 + 50
+    assert isolated.grouped_builds == {}
+
+    # The rp2040 variant has no base test to group with, so it always runs
+    grouped = split_components_for_ci._make_item(tests_dir, "comp", "i2c", False)
+    assert grouped.own_seconds == 50
+    assert grouped.grouped_builds == {("i2c", "esp32-idf"): 90, ("i2c", "host"): 10}
+
+
+def test_grouped_component_joins_existing_build() -> None:
+    """A second member turns a lone build into a group that skips variants."""
+    item = split_components_for_ci._BatchItem
+    batch = split_components_for_ci._Batch()
+    batch.add(item("a", 0, {("i2c", "esp32-idf"): 90}))
+    assert batch.seconds == 90
+
+    other = item("b", 0, {("i2c", "esp32-idf"): 45, ("i2c", "host"): 10})
+    grouped = 45 + split_components_for_ci.GROUPED_COMPONENT_SECONDS
+    assert batch.added_seconds(other) == grouped - 90 + 10
+    batch.add(other)
+    assert batch.seconds == grouped + 10
+
+
+def test_isolated_components_balance_across_runners(tests_dir: Path) -> None:
+    """Heavy isolated components spread out and light ones fill the gaps."""
+    platforms = ["esp32-idf", "esp8266-ard", "rp2040-ard", "bk72xx-ard"]
+    for name in ("big_a", "big_b"):
+        _add_component(
+            tests_dir,
+            name,
+            [f"test.{p}.yaml" for p in platforms]
+            + [f"test-variant.{p}.yaml" for p in platforms],
+        )
+    for name in ("small_a", "small_b"):
+        _add_component(tests_dir, name, ["test.esp32-idf.yaml"])
+
+    batches, _ = split_components_for_ci.create_intelligent_batches(
+        components=["big_a", "big_b", "small_a", "small_b"],
+        tests_dir=tests_dir,
+        target_seconds=500,
+        directly_changed={"big_a", "big_b", "small_a", "small_b"},
+    )
+
+    assert sorted(sorted(batch) for batch in batches) == [
+        ["big_a", "small_a"],
+        ["big_b", "small_b"],
+    ]
+
+
+def test_groupable_variants_skipped_when_grouped(tests_dir: Path) -> None:
+    """Grouped members skip their variants, so the group fits one runner."""
+    names = [f"comp_{i:02d}" for i in range(12)]
+    for name in names:
+        _add_component(
+            tests_dir,
+            name,
+            ["test.esp32-idf.yaml", "test-a.esp32-idf.yaml", "test-b.esp32-idf.yaml"],
+        )
+
+    batches, _ = split_components_for_ci.create_intelligent_batches(
+        components=names, tests_dir=tests_dir, target_seconds=400
+    )
+
+    assert batches == [names]
+
+
+def test_groupable_components_split_evenly(tests_dir: Path) -> None:
+    """A group too large for one runner spreads evenly."""
+    platforms = ["esp32-idf", "esp8266-ard", "rp2040-ard", "bk72xx-ard"]
+    names = [f"comp_{i:02d}" for i in range(12)]
+    for name in names:
+        _add_component(tests_dir, name, [f"test.{p}.yaml" for p in platforms])
+
+    batches, _ = split_components_for_ci.create_intelligent_batches(
+        components=names, tests_dir=tests_dir, target_seconds=100
+    )
+
+    assert sorted(len(batch) for batch in batches) == [2, 2, 2, 3, 3]
+    assert sorted(c for batch in batches for c in batch) == names
+
+
+def test_make_item_reads_tests_dir(tmp_path: Path) -> None:
+    """Costs come from the given tests_dir, not the repository tree."""
+    other = tmp_path / "other"
+    other.mkdir()
+    _add_component(other, "comp", ["test.esp32-idf.yaml"])
+
+    item = split_components_for_ci._make_item(other, "comp", "isolated_comp", True)
+    assert item.own_seconds == 45
+
+
+def test_balance_batches_empty() -> None:
+    """No components means no runners."""
+    assert split_components_for_ci.balance_batches([], 600) == []