Extract prefetched ESP-IDF tool archives in parallel

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),

esphome/espidf/install_tool_archives.py

@@ -0,0 +1,106 @@
+"""Extract prefetched ESP-IDF tool archives in parallel.
+
+Run via ``python <this file> <idf_framework_root> <targets-csv> <workers>
+<tool-spec>...``. PYTHONPATH must include ``<idf_framework_root>/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
+``<IDF_TOOLS_PATH>/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()

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

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()