From c1c9408035a88b65eb4b638c180d81aafadf1a07 Mon Sep 17 00:00:00 2001 From: "J. Nick Koston" Date: Thu, 27 Aug 2026 17:46:39 -0500 Subject: [PATCH] Extract prefetched ESP-IDF tool archives in parallel --- esphome/espidf/framework.py | 40 +++- esphome/espidf/install_tool_archives.py | 106 ++++++++++ .../fixtures/idf_tools_stub/idf_tools.py | 15 +- tests/unit_tests/test_espidf_framework.py | 195 +++++++++++++++++- 4 files changed, 351 insertions(+), 5 deletions(-) create mode 100644 esphome/espidf/install_tool_archives.py diff --git a/esphome/espidf/framework.py b/esphome/espidf/framework.py index e71a333f68..410d71750a 100644 --- a/esphome/espidf/framework.py +++ b/esphome/espidf/framework.py @@ -36,7 +36,7 @@ from esphome.framework_helpers import ( tool_version_runs, warn_batch_failures, ) -from esphome.helpers import write_file_if_changed +from esphome.helpers import get_usable_cpu_count, write_file_if_changed _LOGGER = logging.getLogger(__name__) @@ -292,6 +292,7 @@ def _run_idf_tools_script( msg: str, args: list[str] | None = None, env: dict[str, str] | None = None, + stream_output: bool = False, ) -> tuple[bool, str | None, str | None]: """Run one of the sibling idf_tools-backed helper scripts. @@ -309,6 +310,7 @@ def _run_idf_tools_script( msg=msg, env=(env or os.environ) | {"PYTHONPATH": str(Path(idf_framework_root) / "tools")}, + stream_output=stream_output, ) @@ -789,6 +791,41 @@ def _prefetch_idf_tool_archives( _LOGGER.debug("Prefetch failure detail", exc_info=True) +def _preinstall_idf_tool_archives( + framework_path: Path, + targets_str: str, + tools: list[str], + env: dict[str, str] | None, +) -> None: + """Extract the prefetched tool archives in parallel before the installer. + + ``idf_tools.py install`` unpacks one archive at a time on a single core; + ``install_tool_archives.py`` drives idf_tools' own ``IDFTool.install()`` + with one worker per usable core over the archives the prefetch verified. + The sequential installer still runs afterwards as the authority, skipping + the tools installed here and redoing anything this pass failed on, so + this is strictly best-effort. + """ + try: + success, _stdout, _stderr = _run_idf_tools_script( + framework_path, + "install_tool_archives.py", + "ESP-IDF tool archive extraction", + args=[targets_str, str(get_usable_cpu_count()), *tools], + env=env, + stream_output=True, + ) + if not success: + # Detail already streamed to the terminal by the script + _LOGGER.warning( + "ESP-IDF tool pre-extraction failed; the installer will " + "extract sequentially" + ) + except Exception as e: # noqa: BLE001 # pylint: disable=broad-exception-caught + _LOGGER.warning("ESP-IDF tool pre-extraction failed: %s", failure_reason(e)) + _LOGGER.debug("Pre-extraction failure detail", exc_info=True) + + def _check_esphome_idf_framework_install( version: str, targets: list[str], @@ -940,6 +977,7 @@ def _check_esphome_idf_framework_install( _LOGGER.info("Installing ESP-IDF %s framework ...", version) targets_str = ",".join(targets) _prefetch_idf_tool_archives(framework_path, targets_str, tools, env) + _preinstall_idf_tool_archives(framework_path, targets_str, tools, env) cmd = [ get_system_python_path(), str(idf_tools_path), diff --git a/esphome/espidf/install_tool_archives.py b/esphome/espidf/install_tool_archives.py new file mode 100644 index 0000000000..709e23b8c3 --- /dev/null +++ b/esphome/espidf/install_tool_archives.py @@ -0,0 +1,106 @@ +"""Extract prefetched ESP-IDF tool archives in parallel. + +Run via ``python +...``. PYTHONPATH must include ``/tools`` so +``idf_tools`` is importable, and IDF_TOOLS_PATH must be set. + +``idf_tools.py install`` unpacks one archive at a time; this extracts every +tool whose verified archive the prefetch already placed in +``/dist``, several at once, using idf_tools' own +``IDFTool.install()`` so unpacking, container-dir stripping, and the binary +check match the sequential installer exactly. That installer still runs +afterwards as the authority: it skips the tools installed here and redoes +anything this pass failed on, so per-tool failures only warn on stderr. +""" + +# pylint: disable=import-error # idf_tools is on PYTHONPATH at runtime only + +from concurrent.futures import ThreadPoolExecutor +import os +from pathlib import Path +import sys + +from idf_tools import ( + CURRENT_PLATFORM, + TOOLS_FILE, + IDFEnv, + ToolBinaryError, + add_and_check_targets, + expand_tools_arg, + g, + load_tools_info, +) + + +def collect_pending() -> list[tuple[object, str, str]]: + """The (tool, name, version) jobs whose verified archive is on disk.""" + g.idf_path = sys.argv[1] + g.idf_tools_path = os.environ.get("IDF_TOOLS_PATH") + g.tools_json = str(Path(g.idf_path) / TOOLS_FILE) + + targets = add_and_check_targets(IDFEnv.get_idf_env(), sys.argv[2]) + tools_info = load_tools_info() + dist_path = Path(g.idf_tools_path) / "dist" + pending: list[tuple[object, str, str]] = [] + seen: set[tuple[str, str]] = set() + + for name in expand_tools_arg(sys.argv[4:], tools_info, targets): + if "@" in name: + name, version = name.split("@", 1) + else: + version = None + tool = tools_info.get(name) + if tool is None or not tool.compatible_with_platform(): + continue + version = version or tool.get_recommended_version() + if version is None: + continue + try: + tool.find_installed_versions() + except ToolBinaryError as e: + # Repairing a broken installed binary is the installer's job + print(f"leaving broken {name} to the installer: {e}", file=sys.stderr) + continue + if version in tool.versions_installed or version not in tool.versions: + continue + download = tool.versions[version].get_download_for_platform(CURRENT_PLATFORM) + if download is None: + continue + # An archive at its final name was sha256-verified by the prefetch + archive = dist_path / (download.rename_dist or Path(download.url).name) + if not archive.is_file() or (name, version) in seen: + continue + seen.add((name, version)) + pending.append((tool, name, version)) + return pending + + +def install_one(job: tuple[object, str, str]) -> None: + tool, name, version = job + try: + tool.install(version) + # check_binary_valid exits via SystemExit; the installer redoes failures + except (Exception, SystemExit) as e: # noqa: BLE001 # pylint: disable=broad-exception-caught + print( + f"pre-extracting {name}@{version} failed, leaving it to the installer: {e}", + file=sys.stderr, + ) + + +def main() -> None: + pending = collect_pending() + if len(pending) < 2: + # Nothing to parallelize; the installer keeps its normal output + return + workers = min(int(sys.argv[3]), len(pending)) + print( + f"Extracting {len(pending)} ESP-IDF tool archive(s) with " + f"{workers} worker(s): " + + ", ".join(f"{name}@{version}" for _, name, version in pending), + flush=True, + ) + with ThreadPoolExecutor(max_workers=workers) as ex: + list(ex.map(install_one, pending)) + + +main() diff --git a/tests/unit_tests/fixtures/idf_tools_stub/idf_tools.py b/tests/unit_tests/fixtures/idf_tools_stub/idf_tools.py index aeff24fb57..c24aa21013 100644 --- a/tests/unit_tests/fixtures/idf_tools_stub/idf_tools.py +++ b/tests/unit_tests/fixtures/idf_tools_stub/idf_tools.py @@ -1,7 +1,8 @@ -"""Minimal idf_tools stand-in for get_tool_downloads.py tests.""" +"""Minimal idf_tools stand-in for the espidf helper-script tests.""" from collections.abc import Iterable import os +import pathlib CURRENT_PLATFORM = "linux-amd64" TOOLS_FILE = "tools/tools.json" @@ -54,6 +55,7 @@ class _Tool: installed: Iterable[str] = (), broken: bool = False, ) -> None: + self.name = "" # filled in from the _TOOLS key below self.versions = versions self._recommended = recommended self.versions_installed = list(installed) @@ -69,6 +71,14 @@ class _Tool: if self._broken: raise ToolBinaryError("broken binary") + def install(self, version: str) -> None: + # Real idf_tools' check_binary_valid failure path exits the process + if os.environ.get("TEST_FAIL_INSTALL") == self.name: + raise SystemExit(1) + dest = pathlib.Path(g.idf_tools_path) / "tools" / self.name / version + dest.mkdir(exist_ok=True, parents=True) + (dest / ".installed").write_text("ok", encoding="utf-8") + _TOOLS = { "cmake": _Tool( @@ -97,6 +107,9 @@ _TOOLS = { "no-download-tool": _Tool({"4.0": _Version(None)}, "4.0"), } +for _name, _tool in _TOOLS.items(): + _tool.name = _name + def load_tools_info() -> dict[str, _Tool]: return _TOOLS diff --git a/tests/unit_tests/test_espidf_framework.py b/tests/unit_tests/test_espidf_framework.py index c130ebec35..e2bc29aff4 100644 --- a/tests/unit_tests/test_espidf_framework.py +++ b/tests/unit_tests/test_espidf_framework.py @@ -37,6 +37,7 @@ from esphome.espidf.framework import ( _patch_tools_json_demote_unused_tools, _patch_tools_json_for_linux_arm64, _prefetch_idf_tool_archives, + _preinstall_idf_tool_archives, _read_stamp, _stamp_covers, _windows_long_paths_enabled, @@ -403,6 +404,7 @@ def espidf_mocks(setup_core: Path): patch("esphome.espidf.framework._patch_tools_json_for_linux_arm64"), patch("esphome.espidf.framework._patch_tools_json_demote_unused_tools"), patch("esphome.espidf.framework._prefetch_idf_tool_archives"), + patch("esphome.espidf.framework._preinstall_idf_tool_archives"), patch("esphome.espidf.framework._write_stamp"), patch("esphome.espidf.framework._check_stamp", return_value=True), patch("esphome.espidf.framework._stamp_covers", return_value=True), @@ -1170,21 +1172,25 @@ def test_prefetch_passes_targets_and_tools_to_script(tmp_path: Path) -> None: def test_framework_install_prefetches_before_installer( espidf_mocks: SimpleNamespace, ) -> None: - """The prefetch runs before idf_tools.py install so the installer finds - the archives already in dist/.""" + """The prefetch downloads and the pre-extraction both run before + idf_tools.py install so the installer finds the tools in place.""" calls: list[str] = [] with ( patch( "esphome.espidf.framework._prefetch_idf_tool_archives", side_effect=lambda *a, **k: calls.append("prefetch"), ), + patch( + "esphome.espidf.framework._preinstall_idf_tool_archives", + side_effect=lambda *a, **k: calls.append("preinstall"), + ), ): espidf_mocks.run_ok.side_effect = lambda *a, **k: ( calls.append("install") or True ) check_esp_idf_install(_IDF_VERSION, force=True) - assert calls.index("prefetch") < calls.index("install") + assert calls.index("prefetch") < calls.index("preinstall") < calls.index("install") # --------------------------------------------------------------------------- @@ -1981,3 +1987,186 @@ def test_check_windows_path_length_long_path_warns( assert "long path support" in message # The install is global now; the remedy is the prefix env, not moving the project. assert "ESPHOME_ESP_IDF_PREFIX" in message + + +# --------------------------------------------------------------------------- +# _preinstall_idf_tool_archives +# --------------------------------------------------------------------------- + + +def test_preinstall_streams_script_with_workers(tmp_path: Path) -> None: + """The pre-extraction streams install_tool_archives.py with the worker + count and the same targets/tools the installer will get.""" + with ( + patch( + "esphome.espidf.framework._run_idf_tools_script", + return_value=(True, None, None), + ) as run_script, + patch("esphome.espidf.framework.get_usable_cpu_count", return_value=3), + ): + _preinstall_idf_tool_archives( + tmp_path, "esp32,esp32c3", ["required", "cmake"], {"IDF_TOOLS_PATH": "x"} + ) + run_script.assert_called_once_with( + tmp_path, + "install_tool_archives.py", + "ESP-IDF tool archive extraction", + args=["esp32,esp32c3", "3", "required", "cmake"], + env={"IDF_TOOLS_PATH": "x"}, + stream_output=True, + ) + + +def test_preinstall_script_failure_only_warns( + tmp_path: Path, caplog: pytest.LogCaptureFixture +) -> None: + """A failed pre-extraction leaves the install to the sequential path.""" + with ( + patch( + "esphome.espidf.framework._run_idf_tools_script", + return_value=(False, None, None), + ), + patch("esphome.espidf.framework.get_usable_cpu_count", return_value=1), + ): + _preinstall_idf_tool_archives(tmp_path, "esp32", ["required"], None) + assert "pre-extraction failed" in caplog.text + + +def test_preinstall_exception_only_warns( + tmp_path: Path, caplog: pytest.LogCaptureFixture +) -> None: + """An unexpected error must not become a new way for the install to fail.""" + with ( + patch( + "esphome.espidf.framework._run_idf_tools_script", + side_effect=TypeError("bad call"), + ), + patch("esphome.espidf.framework.get_usable_cpu_count", return_value=1), + ): + _preinstall_idf_tool_archives(tmp_path, "esp32", ["required"], None) + assert "pre-extraction failed" in caplog.text + + +# --------------------------------------------------------------------------- +# install_tool_archives.py (against the stub idf_tools module in fixtures/) +# --------------------------------------------------------------------------- + + +def _run_install_script( + tmp_path: Path, *args: str, env_extra: dict[str, str] | None = None +) -> subprocess.CompletedProcess[str]: + """Run the real install_tool_archives.py against the stub idf_tools.""" + script = ( + Path(__file__).parents[2] / "esphome" / "espidf" / "install_tool_archives.py" + ) + env = os.environ | { + "PYTHONPATH": str(_IDF_TOOLS_STUB_DIR), + "IDF_TOOLS_PATH": str(tmp_path / "tp"), + } + if env_extra: + env |= env_extra + return subprocess.run( + [sys.executable, str(script), str(tmp_path / "fw"), *args], + capture_output=True, + text=True, + env=env, + check=False, + ) + + +def _make_dist(tmp_path: Path, *names: str) -> None: + dist = tmp_path / "tp" / "dist" + dist.mkdir(parents=True, exist_ok=True) + for name in names: + (dist / name).write_bytes(b"x") + + +def test_install_tool_archives_extracts_pending_in_parallel(tmp_path: Path) -> None: + """Tools with a prefetched archive install concurrently; installed tools, + broken tools, and tools without an archive stay with the installer.""" + _make_dist(tmp_path, "cmake.tar.gz", "ninja-v1.zip", "x.tar.gz", "y.tar.gz") + result = _run_install_script(tmp_path, "esp32", "8", "required") + assert result.returncode == 0, result.stderr + tools = tmp_path / "tp" / "tools" + assert (tools / "cmake" / "3.30.2" / ".installed").is_file() + assert (tools / "ninja" / "1.12.1" / ".installed").is_file() + assert not (tools / "installed-tool").exists() + assert not (tools / "broken-tool").exists() + # The worker count clamps to the pending count + assert ( + "Extracting 2 ESP-IDF tool archive(s) with 2 worker(s): " + "cmake@3.30.2, ninja@1.12.1" in result.stdout + ) + assert "leaving broken broken-tool to the installer" in result.stderr + + +def test_install_tool_archives_single_pending_stays_sequential( + tmp_path: Path, +) -> None: + """One pending archive has nothing to parallelize; the installer keeps + its normal output.""" + _make_dist(tmp_path, "cmake.tar.gz") + result = _run_install_script(tmp_path, "esp32", "4", "required") + assert result.returncode == 0, result.stderr + assert not (tmp_path / "tp" / "tools").exists() + assert "Extracting" not in result.stdout + + +def test_install_tool_archives_failed_install_left_to_installer( + tmp_path: Path, +) -> None: + """A per-tool failure (SystemExit from the binary check) warns and moves + on; the other tools still install.""" + _make_dist(tmp_path, "cmake.tar.gz", "ninja-v1.zip") + result = _run_install_script( + tmp_path, "esp32", "4", "required", env_extra={"TEST_FAIL_INSTALL": "ninja"} + ) + assert result.returncode == 0, result.stderr + tools = tmp_path / "tp" / "tools" + assert (tools / "cmake" / "3.30.2" / ".installed").is_file() + assert not (tools / "ninja").exists() + assert "pre-extracting ninja@1.12.1 failed" in result.stderr + + +def _run_install_inprocess( + tmp_path: Path, + monkeypatch: pytest.MonkeyPatch, + *args: str, +) -> None: + """Execute install_tool_archives.py in-process against the stub idf_tools. + + Unlike the subprocess variant this runs under coverage, exercising the + script's own lines. + """ + spec = importlib.util.spec_from_file_location( + "idf_tools", _IDF_TOOLS_STUB_DIR / "idf_tools.py" + ) + stub = importlib.util.module_from_spec(spec) + spec.loader.exec_module(stub) + monkeypatch.setitem(sys.modules, "idf_tools", stub) + monkeypatch.setenv("IDF_TOOLS_PATH", str(tmp_path / "tp")) + script = ( + Path(__file__).parents[2] / "esphome" / "espidf" / "install_tool_archives.py" + ) + monkeypatch.setattr(sys, "argv", [str(script), str(tmp_path / "fw"), *args]) + runpy.run_path(str(script)) + + +def test_install_tool_archives_inprocess_dedupes_specs( + tmp_path: Path, + monkeypatch: pytest.MonkeyPatch, + capsys: pytest.CaptureFixture[str], +) -> None: + """In-process full flow: duplicate tool@version specs collapse to one + job and both tools install.""" + _make_dist(tmp_path, "cmake.tar.gz", "ninja-v1.zip") + _run_install_inprocess( + tmp_path, monkeypatch, "esp32", "8", "cmake", "ninja", "cmake@3.30.2" + ) + out = capsys.readouterr().out + assert ( + "Extracting 2 ESP-IDF tool archive(s) with 2 worker(s): " + "cmake@3.30.2, ninja@1.12.1" in out + ) + assert (tmp_path / "tp" / "tools" / "cmake" / "3.30.2" / ".installed").is_file() + assert (tmp_path / "tp" / "tools" / "ninja" / "1.12.1" / ".installed").is_file()