[ci] Refit component batch estimates to a full run (#20401)

script/split_components_for_ci.py

@@ -13,15 +13,20 @@ from __future__ import annotations
 
 import argparse
 from collections import defaultdict
+from collections.abc import Set as AbstractSet
+import contextlib
 from dataclasses import dataclass, field
+import functools
 import json
 import math
 from pathlib import Path
 import sys
+from typing import Any, NamedTuple
 
 # Add esphome to path
 sys.path.insert(0, str(Path(__file__).parent.parent))
 
+from esphome import yaml_util
 from script.analyze_component_buses import (
     ISOLATED_COMPONENTS,
     ISOLATED_SIGNATURE_PREFIX,
@@ -31,29 +36,52 @@ from script.analyze_component_buses import (
     merge_compatible_bus_groups,
 )
 from script.helpers import (
+    _extract_components_from_yaml,
     get_component_test_files,
     parse_test_filename,
     split_conflicting_groups,
 )
 
-# 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.
+# Estimated CI seconds for one build per test platform, fitted to the batch
+# job logs of a full run (every component, PR #20389). 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,
+    "rtl87xx-ard": 40,
+    "esp32-c3-idf": 40,
+    "esp32-s2-idf": 40,
+    "rp2350-ard": 40,
+    "rp2040-ard": 45,
+    "esp32-idf": 55,
+    "esp32-p4-idf": 55,
+    "ln882x-ard": 55,
+    "esp32-s3-idf": 60,
+    "bk72xx-ard": 70,
+    "esp32-c6-idf": 70,
 }
-NRF52_BUILD_SECONDS = 45
-# Other ESP32 chips and Arduino on ESP32 share less of the ccache
-DEFAULT_BUILD_SECONDS = 90
+NRF52_BUILD_SECONDS = 50
+# Arduino on any ESP32 chip builds the IDF plus the Arduino core on top
+ESP32_ARDUINO_BUILD_SECONDS = 90
+# The remaining ESP32 chips on ESP-IDF
+DEFAULT_BUILD_SECONDS = 50
 # Each extra component merged into a grouped build makes it larger
-GROUPED_COMPONENT_SECONDS = 5
+GROUPED_COMPONENT_SECONDS = 3
+# Components that pull in a large library the build compiles from source
+# (TensorFlow Lite Micro, FastLED, LVGL, the HTTP server), added to every
+# build whose test config enables them. Fitted the same way as the platforms.
+HEAVY_COMPONENT_SECONDS: dict[str, int] = {
+    "micro_wake_word": 200,
+    "fastled_clockless": 90,
+    "fastled_spi": 90,
+    "wled": 60,
+    "lvgl": 50,
+    "web_server": 50,
+    "online_image": 45,
+    "prometheus": 35,
+    "bluetooth_proxy": 30,
+}
 # 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.
@@ -90,9 +118,62 @@ def build_seconds(platform: str) -> int:
         return seconds
     if platform.startswith("nrf52"):
         return NRF52_BUILD_SECONDS
+    if platform.startswith("esp32") and platform.endswith("-ard"):
+        return ESP32_ARDUINO_BUILD_SECONDS
     return DEFAULT_BUILD_SECONDS
 
 
+def _config_components(data: Any) -> set[str]:
+    """Return the components a loaded test config and its ``packages:`` use."""
+    if not isinstance(data, dict):
+        return set()
+    found = _extract_components_from_yaml(data)
+    if isinstance(packages := data.get("packages"), dict):
+        for package in packages.values():
+            if isinstance(package, yaml_util.IncludeFile):
+                found |= _file_components(
+                    (package.parent_file.parent / package.file).resolve()
+                )
+            else:
+                found |= _config_components(package)
+    return found
+
+
+@functools.cache
+def _file_components(path: Path) -> frozenset[str]:
+    """Return the components a config file uses; shared packages load once."""
+    return frozenset(
+        _config_components(yaml_util.load_yaml(path, track_document_range=False))
+    )
+
+
+def heavy_components(component: str, test_file: Path) -> frozenset[str]:
+    """Return the heavy components a test compiles.
+
+    A test counts as heavy when it belongs to a heavy component or when its
+    config enables one, such as a display test that draws with lvgl.
+    """
+    heavy = {component} & HEAVY_COMPONENT_SECONDS.keys()
+    # A config that does not load is left to the build to report
+    with contextlib.suppress(Exception):
+        heavy |= _file_components(test_file.resolve()) & HEAVY_COMPONENT_SECONDS.keys()
+    return frozenset(heavy)
+
+
+def heavy_seconds(heavy: AbstractSet[str]) -> int:
+    """Return the library seconds a set of heavy components adds to a build."""
+    return sum(HEAVY_COMPONENT_SECONDS[name] for name in heavy)
+
+
+class _GroupedShare(NamedTuple):
+    """A component's part in one (signature, platform) shared build."""
+
+    # Seconds of its files on that platform when no other member shares it
+    solo: int
+    # Heavy components its base test brings into the shared build
+    heavy: frozenset[str]
+
+
 @dataclass
 class _BatchItem:
     """A component and the builds it adds to the batch it lands in."""
@@ -101,40 +182,52 @@ class _BatchItem:
     # 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)
+    grouped_builds: dict[tuple[str, str], _GroupedShare] = field(default_factory=dict)
 
     def standalone_seconds(self) -> int:
-        return self.own_seconds + sum(self.grouped_builds.values())
+        return self.own_seconds + sum(
+            share.solo for share in self.grouped_builds.values()
+        )
 
 
-def _grouped_build_seconds(platform: str, solo_seconds: list[int]) -> int:
+def _grouped_build_seconds(platform: str, members: list[_GroupedShare]) -> int:
     """Return the seconds of one (signature, platform) build in a batch.
 
+    A library two members both enable compiles once, so the union is charged.
     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)
+    if len(members) <= 1:
+        return sum(share.solo for share in members)
+    return (
+        build_seconds(platform)
+        + GROUPED_COMPONENT_SECONDS * (len(members) - 1)
+        + heavy_seconds(set().union(*(share.heavy for share in members)))
+    )
 
 
 def _make_item(
     tests_dir: Path, component: str, signature: str, is_isolated: bool
 ) -> _BatchItem:
-    files_by_platform: dict[str, list[str]] = defaultdict(list)
+    # Per platform: the seconds of its test files, and the base test's heavy
+    # components; a platform without a base test cannot join a shared build
+    seconds_by_platform: dict[str, int] = defaultdict(int)
+    base_heavy: dict[str, frozenset[str]] = {}
     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)
+        heavy = heavy_components(component, test_file)
+        seconds_by_platform[platform] += build_seconds(platform) + heavy_seconds(heavy)
+        if test_name == "test":
+            base_heavy[platform] = heavy
     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:
+    grouped: dict[tuple[str, str], _GroupedShare] = {}
+    for platform, seconds in seconds_by_platform.items():
+        if is_isolated or platform not in base_heavy:
             own_seconds += seconds
         else:
-            grouped[(signature, platform)] = seconds
+            grouped[(signature, platform)] = _GroupedShare(
+                seconds, base_heavy[platform]
+            )
     return _BatchItem(component, own_seconds, grouped)
 
 
@@ -144,22 +237,24 @@ class _Batch:
     def __init__(self) -> None:
         self.components: list[str] = []
         self.seconds = 0
-        self.grouped_builds: dict[tuple[str, str], list[int]] = defaultdict(list)
+        self.grouped_builds: dict[tuple[str, str], list[_GroupedShare]] = 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():
+        for (signature, platform), share in item.grouped_builds.items():
             members = self.grouped_builds.get((signature, platform), [])
             added += _grouped_build_seconds(
-                platform, [*members, solo]
+                platform, [*members, share]
             ) - _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)
+        for build, share in item.grouped_builds.items():
+            self.grouped_builds[build].append(share)
         self.components.append(item.component)
 
 
@@ -172,14 +267,11 @@ def balance_batches(items: list[_BatchItem], target_seconds: int) -> list[list[s
     """
     if not items:
         return []
-    grouped_builds: dict[tuple[str, str], list[int]] = defaultdict(list)
+    # One batch holding every item counts each grouped build once
+    whole = _Batch()
     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()
-    )
+        whole.add(item)
+    total = whole.seconds
     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)):

tests/script/test_determine_jobs.py

@@ -2613,8 +2613,8 @@ def test_component_batching_beta_branch_groups_evenly(
     For beta/release branches, all components should be groupable (not isolated),
     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)]
+    # Create 200 test components with test files
+    component_names = [f"comp_{i:03d}" for i in range(200)]
     tests_dir = tmp_path / "tests" / "components"
 
     for comp in component_names:
@@ -2675,15 +2675,15 @@ def test_component_batching_beta_branch_groups_evenly(
     assert "component_test_batches" in output
     batches = output["component_test_batches"]
 
-    # 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]
+    # One grouped esp32-idf build plus 3 s per extra component is 652 s,
+    # so two runners of 100 components each
+    assert [len(batch["components"].split()) for batch in batches] == [100, 100]
 
-    # Verify all 120 components are in batches
+    # Verify all 200 components are in batches
     all_components = []
     for batch in batches:
         all_components.extend(batch["components"].split())
-    assert len(all_components) == 120
+    assert len(all_components) == 200
     assert set(all_components) == set(component_names)
 
 

tests/script/test_split_components_for_ci.py

@@ -35,13 +35,46 @@ 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("esp32-idf") == 55
     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
+        build_seconds("esp32-c5-idf") == split_components_for_ci.DEFAULT_BUILD_SECONDS
     )
+    assert (
+        build_seconds("esp32-c3-ard")
+        == split_components_for_ci.ESP32_ARDUINO_BUILD_SECONDS
+    )
+
+
+def test_heavy_components(tests_dir: Path) -> None:
+    """A heavy component, or a test that enables one, costs its library time."""
+    heavy_components = split_components_for_ci.heavy_components
+    heavy_seconds = split_components_for_ci.heavy_seconds
+    table = split_components_for_ci.HEAVY_COMPONENT_SECONDS
+    _add_component(tests_dir, "dht", ["test.esp32-idf.yaml"])
+    _add_component(tests_dir, "micro_wake_word", ["test.esp32-idf.yaml"])
+    assert heavy_components("dht", tests_dir / "dht/test.esp32-idf.yaml") == set()
+    assert heavy_components(
+        "micro_wake_word", tests_dir / "micro_wake_word/test.esp32-idf.yaml"
+    ) == {"micro_wake_word"}
+    assert heavy_seconds(frozenset({"micro_wake_word"})) == table["micro_wake_word"]
+
+    display = tests_dir / "mipi_spi"
+    display.mkdir()
+    (display / "common.yaml").write_text("web_server:\n  port: 80\n")
+    (display / "test-lvgl.esp32-s3-idf.yaml").write_text(
+        "packages:\n"
+        "  mipi_spi: !include common.yaml\n"
+        "lvgl:\n  pages: []\n"
+        "light:\n  - platform: fastled_spi\n    num_leds: 1\n"
+    )
+    assert heavy_components("mipi_spi", display / "test-lvgl.esp32-s3-idf.yaml") == {
+        "lvgl",
+        "web_server",
+        "fastled_spi",
+    }
 
 
 def test_make_item_isolated_and_grouped(tests_dir: Path) -> None:
@@ -61,29 +94,93 @@ def test_make_item_isolated_and_grouped(tests_dir: Path) -> None:
     isolated = split_components_for_ci._make_item(
         tests_dir, "comp", "isolated_comp", True
     )
-    assert isolated.own_seconds == 45 + 10 + 45 + 50
+    assert isolated.own_seconds == 55 + 10 + 55 + 45
     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}
+    assert grouped.own_seconds == 45
+    assert grouped.grouped_builds == {
+        ("i2c", "esp32-idf"): (110, frozenset()),
+        ("i2c", "host"): (10, frozenset()),
+    }
+
+
+def test_make_item_heavy_component(tests_dir: Path) -> None:
+    """A heavy component's library cost follows it into every build."""
+    _add_component(tests_dir, "lvgl", ["test.esp32-idf.yaml", "test-a.host.yaml"])
+    extra = split_components_for_ci.HEAVY_COMPONENT_SECONDS["lvgl"]
+
+    item = split_components_for_ci._make_item(tests_dir, "lvgl", "spi", False)
+    assert item.own_seconds == 10 + extra
+    assert item.grouped_builds == {
+        ("spi", "esp32-idf"): (55 + extra, frozenset({"lvgl"}))
+    }
+
+
+def test_heavy_cost_is_per_platform(tests_dir: Path) -> None:
+    """A base test that enables a heavy component only charges its own platform."""
+    comp = tests_dir / "web_server_idf"
+    comp.mkdir()
+    (comp / "test.esp32-idf.yaml").write_text("web_server:\n  port: 80\n")
+    (comp / "test.host.yaml").write_text("logger:\n")
+    extra = split_components_for_ci.HEAVY_COMPONENT_SECONDS["web_server"]
+
+    item = split_components_for_ci._make_item(
+        tests_dir, "web_server_idf", "none", False
+    )
+    assert item.grouped_builds == {
+        ("none", "esp32-idf"): (55 + extra, frozenset({"web_server"})),
+        ("none", "host"): (10, frozenset()),
+    }
 
 
 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
+    share = split_components_for_ci._GroupedShare
     batch = split_components_for_ci._Batch()
-    batch.add(item("a", 0, {("i2c", "esp32-idf"): 90}))
-    assert batch.seconds == 90
+    batch.add(item("a", 0, {("i2c", "esp32-idf"): share(110, frozenset())}))
+    assert batch.seconds == 110
 
-    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
+    other = item(
+        "b",
+        0,
+        {
+            ("i2c", "esp32-idf"): share(55, frozenset()),
+            ("i2c", "host"): share(10, frozenset()),
+        },
+    )
+    grouped = 55 + split_components_for_ci.GROUPED_COMPONENT_SECONDS
+    assert batch.added_seconds(other) == grouped - 110 + 10
     batch.add(other)
     assert batch.seconds == grouped + 10
 
 
+def test_heavy_component_charges_the_shared_build() -> None:
+    """A heavy member makes the whole grouped build slower, once per library."""
+    item = split_components_for_ci._BatchItem
+    share = split_components_for_ci._GroupedShare
+    batch = split_components_for_ci._Batch()
+    extra = split_components_for_ci.HEAVY_COMPONENT_SECONDS["lvgl"]
+    per_member = split_components_for_ci.GROUPED_COMPONENT_SECONDS
+    lvgl = frozenset({"lvgl"})
+    batch.add(item("a", 0, {("spi", "esp32-idf"): share(55, frozenset())}))
+
+    heavy = item("lvgl", 0, {("spi", "esp32-idf"): share(55 + extra, lvgl)})
+    grouped = 55 + per_member + extra
+    assert batch.added_seconds(heavy) == grouped - 55
+    batch.add(heavy)
+    assert batch.seconds == grouped
+
+    light = item("b", 0, {("spi", "esp32-idf"): share(55, frozenset())})
+    assert batch.added_seconds(light) == per_member
+
+    # A second member that also draws with lvgl compiles the library once
+    also_heavy = item("c", 0, {("spi", "esp32-idf"): share(55 + extra, lvgl)})
+    assert batch.added_seconds(also_heavy) == per_member
+
+
 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"]
@@ -138,7 +235,7 @@ def test_groupable_components_split_evenly(tests_dir: Path) -> None:
         components=names, tests_dir=tests_dir, target_seconds=100
     )
 
-    assert sorted(len(batch) for batch in batches) == [2, 2, 2, 3, 3]
+    assert sorted(len(batch) for batch in batches) == [3, 3, 3, 3]
     assert sorted(c for batch in batches for c in batch) == names
 
 
@@ -149,7 +246,7 @@ def test_make_item_reads_tests_dir(tmp_path: Path) -> None:
     _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
+    assert item.own_seconds == 55
 
 
 def test_balance_batches_empty() -> None: