From 563083b4a97759171ef90b2db64d0d3520a52794 Mon Sep 17 00:00:00 2001 From: "J. Nick Koston" Date: Thu, 27 Aug 2026 17:30:29 -0500 Subject: [PATCH] Extract native toolchain package archives in parallel --- esphome/arduino8266/framework.py | 7 +- esphome/framework_helpers.py | 66 ++++++--- esphome/platformio/registry.py | 96 ++++++++++++- .../unit_tests/test_arduino8266_framework.py | 37 ++--- tests/unit_tests/test_framework_helpers.py | 51 +++++++ tests/unit_tests/test_platformio_registry.py | 134 +++++++++++++++++- 6 files changed, 340 insertions(+), 51 deletions(-) diff --git a/esphome/arduino8266/framework.py b/esphome/arduino8266/framework.py index dae0ed4b39..139f9f8dbe 100644 --- a/esphome/arduino8266/framework.py +++ b/esphome/arduino8266/framework.py @@ -23,7 +23,7 @@ from esphome.build_helpers.pch import ccache_pch_env from esphome.build_helpers.tools_cache import ARDUINO8266_TOOLS_CACHE, tools_cache_path from esphome.core import EsphomeError, Version from esphome.framework_helpers import str_to_lst_of_str -from esphome.platformio.registry import install_package, prefetch_packages +from esphome.platformio.registry import install_packages, prefetch_packages FRAMEWORK_PACKAGE = "framework-arduinoespressif8266" TOOLCHAIN_PACKAGE = "toolchain-xtensa" @@ -121,10 +121,9 @@ def check_and_install(framework_version: Version) -> InstalledPaths: ("bin", "xtensa-lx106-elf"), ), ) - # Fetch both archives at once; the installs below verify and extract + # Fetch both archives at once; the install verifies and extracts them prefetch_packages([spec[:4] for spec in specs], downloads_dir) - for name, version, dest, mirrors, expect in specs: - install_package(name, version, dest, mirrors, downloads_dir, expect=expect) + install_packages(specs, downloads_dir) return InstalledPaths( framework=framework_path, toolchain=toolchain_path, ninja=ninja_path ) diff --git a/esphome/framework_helpers.py b/esphome/framework_helpers.py index 82bc0d3727..843ce8b1cb 100644 --- a/esphome/framework_helpers.py +++ b/esphome/framework_helpers.py @@ -287,10 +287,25 @@ def _detect_archive_root(names: Iterable[str]) -> str | None: return root if has_descendant else None +def _resolve_progress( + progress: Callable[[float], None] | None, + progress_header: str | None, + has_work: bool, +) -> Callable[[float], None] | None: + """Fraction reporter for an extractor: the caller's callback wins over a + private ``progress_header`` bar.""" + if progress is not None: + return progress + if progress_header and has_work: + return ProgressBar(progress_header).update + return None + + def _tar_extract_all( data: io.BufferedIOBase, extract_dir: PathType = ".", progress_header: str | None = None, + progress: Callable[[float], None] | None = None, ): """ Extract a TAR archive to the specified directory. @@ -305,6 +320,8 @@ def _tar_extract_all( data: File-like object containing the TAR archive extract_dir: Directory to extract contents to progress_header: If set, show a progress bar with this header + progress: If set, receives fractions in [0, 1] ending at 1.0 and + overrides ``progress_header`` """ import tarfile @@ -363,21 +380,20 @@ def _tar_extract_all( safe_members.append(member) total = len(safe_members) - progress = ( - ProgressBar(progress_header) if progress_header and total > 0 else None - ) + report = _resolve_progress(progress, progress_header, total > 0) for i, member in enumerate(safe_members, 1): tar_ref.extract(member, abs_dest) - if progress is not None: - progress.update(i / total) - if progress is not None: - progress.update(1) + if report is not None: + report(i / total) + if report is not None: + report(1) def _zip_extract_all( data: io.BufferedIOBase, extract_dir: PathType = ".", progress_header: str | None = None, + progress: Callable[[float], None] | None = None, ): """ Extract a ZIP archive to the specified directory. @@ -386,6 +402,8 @@ def _zip_extract_all( data: File-like object containing the ZIP archive extract_dir: Directory to extract contents to progress_header: If set, show a progress bar with this header + progress: If set, receives fractions in [0, 1] ending at 1.0 and + overrides ``progress_header`` """ import zipfile @@ -402,9 +420,7 @@ def _zip_extract_all( strip_prefix = f"{strip_root}/" if strip_root is not None else None total = len(all_members) - progress = ( - ProgressBar(progress_header) if progress_header and total > 0 else None - ) + report = _resolve_progress(progress, progress_header, total > 0) for i, member in enumerate(all_members, 1): # 1. Normalize name @@ -437,10 +453,10 @@ def _zip_extract_all( # 6. Extract zip_ref.extract(member, extract_dir) - if progress is not None: - progress.update(i / total) - if progress is not None: - progress.update(1) + if report is not None: + report(i / total) + if report is not None: + report(1) def _rename_with_retry( @@ -471,6 +487,7 @@ def _7z_extract_all( data: io.BufferedIOBase, extract_dir: PathType = ".", progress_header: str | None = None, + progress: Callable[[float], None] | None = None, ): """ Extract a 7z archive to the specified directory. @@ -485,6 +502,8 @@ def _7z_extract_all( data: File-like object containing the 7z archive (must be seekable) extract_dir: Directory to extract contents to progress_header: If set, show a progress bar with this header + progress: If set, called with 1.0 on completion and overrides + ``progress_header`` """ import py7zr @@ -523,19 +542,15 @@ def _7z_extract_all( continue safe_targets.append(raw) - progress = ( - ProgressBar(progress_header) - if progress_header and safe_targets - else None - ) + report = _resolve_progress(progress, progress_header, bool(safe_targets)) if len(safe_targets) == len(all_names): z.extractall(path=staging) else: z.extract(path=staging, targets=safe_targets) - if progress is not None: - progress.update(1) + if report is not None: + report(1) src_root = staging / strip_root if strip_root else staging for item in src_root.iterdir(): @@ -566,6 +581,7 @@ def archive_extract_all( archive: PathType | io.RawIOBase | IO[bytes], extract_dir: PathType = ".", progress_header: str | None = None, + progress: Callable[[float], None] | None = None, ): """ Extract an archive file to the specified directory. @@ -574,6 +590,8 @@ def archive_extract_all( archive: Path to archive file or file-like object extract_dir: Directory to extract contents to progress_header: If set, show a progress bar with this header + progress: If set, receives fractions in [0, 1] ending at 1.0 and + overrides ``progress_header`` Raises: TypeError: If archive is not a valid type @@ -604,7 +622,9 @@ def archive_extract_all( break if matched_fct is None: raise ValueError("Unsupported archive format") - matched_fct(archive_ref, extract_dir, progress_header=progress_header) + matched_fct( + archive_ref, extract_dir, progress_header=progress_header, progress=progress + ) def _open_ranged( @@ -774,7 +794,7 @@ def run_batch_downloads( jobs: list[tuple[str, int, Callable[[Callable[[int], None]], None]]], max_workers: int = BATCH_DOWNLOAD_WORKERS, ) -> list[tuple[str, BaseException]]: - """Run ``(name, size, fetch)`` download jobs concurrently under one bar. + """Run ``(name, size, fetch)`` jobs concurrently under one bar. Each ``fetch(tracker)`` reports absolute byte counts; the bar total is the sum of the sizes. Failures are returned after the bar is done so diff --git a/esphome/platformio/registry.py b/esphome/platformio/registry.py index 9538a28ff4..7cf639514c 100644 --- a/esphome/platformio/registry.py +++ b/esphome/platformio/registry.py @@ -19,7 +19,9 @@ from esphome.framework_helpers import ( download_with_resume, rmdir, run_batch_downloads, + warn_prefetch_failures, ) +from esphome.helpers import get_usable_cpu_count from esphome.net_retry import fetch_with_retry, http_request _LOGGER = logging.getLogger(__name__) @@ -160,6 +162,14 @@ def _check_layout(name: str, dest: Path, expect: Collection[str]) -> None: ) +# (name, version, dest, mirrors, expect) as accepted by install_packages +PackageSpec = tuple[str, str, Path, list[str], Collection[str]] + + +def _archive_path(downloads_dir: Path, name: str, version: str) -> Path: + return downloads_dir / f"{name}-{version}" + + class _PendingArchive(NamedTuple): name: str version: str @@ -259,6 +269,7 @@ def install_package( mirrors: list[str], downloads_dir: Path, expect: Collection[str], + extract_progress: Callable[[float], None] | None = None, ) -> None: """Download, verify, and extract one package if not already installed. @@ -266,6 +277,9 @@ def install_package( publishes; a mirror override (URL templates with ``{VERSION}``/``{SYSTEM}`` substitution) is trusted as configured. ``downloads_dir`` holds the archive between runs so an interrupted download resumes. + + ``extract_progress`` receives extraction fractions in [0, 1] instead of + the private per-file bars (see ``install_packages``). """ if not expect: # Layout validation before marker.touch() is the only guard against @@ -288,8 +302,10 @@ def install_package( rmdir(dest, msg=f"Clean up incomplete {name} install") # Persistent location so an interrupted download resumes across runs. downloads_dir.mkdir(parents=True, exist_ok=True) - archive = downloads_dir / f"{name}-{version}" - _LOGGER.info("Downloading %s %s ...", name, version) + archive = _archive_path(downloads_dir, name, version) + # The batch header already names each package + log = _LOGGER.debug if extract_progress is not None else _LOGGER.info + log("Downloading %s %s ...", name, version) if mirrors: _LOGGER.warning( "Downloading %s from a mirror override; checksum verification " @@ -301,11 +317,81 @@ def install_package( ) else: url, sha256, size = registry_download(name, version) - download_with_resume(url, archive, sha256=sha256, size=size) - _LOGGER.info("Extracting %s ...", name) - archive_extract_all(archive, dest, progress_header="Extracting") + # Batched: no private bar, no bytes (the shared bar must never + # run backwards), but the zero tick keeps cancellation observable + download_progress = ( + None + if extract_progress is None + else lambda _done: extract_progress(0.0) + ) + download_with_resume( + url, archive, sha256=sha256, size=size, progress=download_progress + ) + log("Extracting %s ...", name) + archive_extract_all( + archive, dest, progress_header="Extracting", progress=extract_progress + ) # Validate the layout before recording success, so an unexpected # package is never cached as a working install. _check_layout(name, dest, expect) marker.touch() archive.unlink(missing_ok=True) + + +def install_packages(specs: Collection[PackageSpec], downloads_dir: Path) -> None: + """Install several packages, extracting verified archives in parallel. + + Prefetched archives extract concurrently under one shared bar; the rest + (marker hits, mirror overrides, missing archives) take the sequential + ``install_package`` path. The first failure is re-raised. + """ + pending: list[tuple[PackageSpec, int]] = [] + rest: list[PackageSpec] = [] + for spec in specs: + name, version, dest, mirrors, _expect = spec + if _already_installed(dest) or mirrors: + rest.append(spec) + continue + try: + # An archive at its final name already passed sha256/size + # verification + size = _archive_path(downloads_dir, name, version).stat().st_size + except OSError: + rest.append(spec) + continue + pending.append((spec, size)) + if len(pending) < 2: + rest = list(specs) + pending = [] + for name, version, dest, mirrors, expect in rest: + install_package(name, version, dest, mirrors, downloads_dir, expect=expect) + if not pending: + return + workers = min(get_usable_cpu_count(), len(pending)) + _LOGGER.info( + "Extracting %d package archive(s) with %d worker(s): %s", + len(pending), + workers, + ", ".join(spec[0] for spec, _ in pending), + ) + + def _install(spec: PackageSpec, size: int, tracker: Callable[[int], None]) -> None: + name, version, dest, mirrors, expect = spec + install_package( + name, + version, + dest, + mirrors, + downloads_dir, + expect=expect, + extract_progress=lambda frac: tracker(int(frac * size)), + ) + + failures = run_batch_downloads( + "Extracting packages", + [(spec[0], size, partial(_install, spec, size)) for spec, size in pending], + max_workers=workers, + ) + if failures: + warn_prefetch_failures(failures[1:], "Could not install %s: %s") + raise failures[0][1] diff --git a/tests/unit_tests/test_arduino8266_framework.py b/tests/unit_tests/test_arduino8266_framework.py index bd0a620e10..f983fbc1f1 100644 --- a/tests/unit_tests/test_arduino8266_framework.py +++ b/tests/unit_tests/test_arduino8266_framework.py @@ -62,7 +62,7 @@ def test_tools_path_default_and_prefix(tmp_path: Path) -> None: def test_check_and_install_returns_paths(tmp_path: Path) -> None: with ( patch.dict(os.environ, {"ESPHOME_ARDUINO8266_PREFIX": str(tmp_path)}), - patch.object(framework, "install_package") as mock_install, + patch.object(framework, "install_packages") as mock_install, patch.object(framework, "prefetch_packages") as mock_prefetch, patch.object(framework, "find_ninja", return_value=tmp_path / "ninja"), ): @@ -70,26 +70,27 @@ def test_check_and_install_returns_paths(tmp_path: Path) -> None: assert paths.framework == tmp_path / "frameworks" / "3.30102.0" assert paths.toolchain == tmp_path / "toolchains" / framework.TOOLCHAIN_VERSION assert paths.ninja == tmp_path / "ninja" - assert mock_install.call_count == 2 # Full argument pinning: a copy-paste swap between the two near-identical - # calls (mirrors, destination) must not stay green - fw_call, tc_call = mock_install.call_args_list - assert fw_call.args == ( - framework.FRAMEWORK_PACKAGE, - "3.30102.0", - tmp_path / "frameworks" / "3.30102.0", - framework.ESPHOME_ARDUINO8266_FRAMEWORK_MIRRORS, + # specs (mirrors, destination) must not stay green + assert mock_install.call_args.args == ( + ( + ( + framework.FRAMEWORK_PACKAGE, + "3.30102.0", + tmp_path / "frameworks" / "3.30102.0", + framework.ESPHOME_ARDUINO8266_FRAMEWORK_MIRRORS, + ("cores/esp8266", "tools/sdk", "libraries"), + ), + ( + framework.TOOLCHAIN_PACKAGE, + framework.TOOLCHAIN_VERSION, + tmp_path / "toolchains" / framework.TOOLCHAIN_VERSION, + framework.ESPHOME_ARDUINO8266_TOOLCHAIN_MIRRORS, + ("bin", "xtensa-lx106-elf"), + ), + ), tmp_path / "downloads", ) - assert fw_call.kwargs["expect"] == ("cores/esp8266", "tools/sdk", "libraries") - assert tc_call.args == ( - framework.TOOLCHAIN_PACKAGE, - framework.TOOLCHAIN_VERSION, - tmp_path / "toolchains" / framework.TOOLCHAIN_VERSION, - framework.ESPHOME_ARDUINO8266_TOOLCHAIN_MIRRORS, - tmp_path / "downloads", - ) - assert tc_call.kwargs["expect"] == ("bin", "xtensa-lx106-elf") # The prefetch sees the same package specs as the installs assert mock_prefetch.call_args.args == ( [ diff --git a/tests/unit_tests/test_framework_helpers.py b/tests/unit_tests/test_framework_helpers.py index fcc5572f51..81c1948605 100644 --- a/tests/unit_tests/test_framework_helpers.py +++ b/tests/unit_tests/test_framework_helpers.py @@ -523,6 +523,17 @@ class TestArchiveExtractAll: archive_extract_all(archive, dest) assert (dest / "file.txt").read_text() == "hi" + def test_progress_callback_passed_through(self, tmp_path: Path) -> None: + """The progress kwarg reaches the dispatched extractor.""" + archive = tmp_path / "test.tar.gz" + archive.write_bytes(_gzip_tar_bytes({"file.txt": b"hello"})) + dest = tmp_path / "out" + dest.mkdir() + fractions: list[float] = [] + archive_extract_all(archive, dest, progress=fractions.append) + assert fractions[-1] == 1 + assert (dest / "file.txt").read_bytes() == b"hello" + def test_invalid_type_raises_type_error(self) -> None: with pytest.raises(TypeError, match="archive must be"): archive_extract_all(42, ".") # type: ignore[arg-type] @@ -1951,6 +1962,19 @@ class TestTarExtractAllBranches: mock_pb.assert_called_once_with("Extracting") mock_pb.return_value.update.assert_called() + def test_progress_callback_replaces_bar(self, tmp_path: Path) -> None: + """A progress callback wins over progress_header and ends at 1.0.""" + buf = _make_tar([_reg("a.txt"), _reg("b.txt")], {"a.txt": b"x", "b.txt": b"y"}) + fractions: list[float] = [] + with patch("esphome.framework_helpers.ProgressBar") as mock_pb: + _tar_extract_all( + buf, tmp_path, progress_header="Extracting", progress=fractions.append + ) + mock_pb.assert_not_called() + assert fractions == sorted(fractions) + assert fractions[-1] == 1 + assert (tmp_path / "a.txt").is_file() + # --------------------------------------------------------------------------- # _zip_extract_all — additional branch coverage @@ -1980,6 +2004,19 @@ class TestZipExtractAllBranches: mock_pb.assert_called_once_with("Unzipping") mock_pb.return_value.update.assert_called() + def test_progress_callback_replaces_bar(self, tmp_path: Path) -> None: + """A progress callback wins over progress_header and ends at 1.0.""" + buf = _make_zip([("a.txt", "aaa"), ("b.txt", "bbb")]) + fractions: list[float] = [] + with patch("esphome.framework_helpers.ProgressBar") as mock_pb: + _zip_extract_all( + buf, tmp_path, progress_header="Unzipping", progress=fractions.append + ) + mock_pb.assert_not_called() + assert fractions == sorted(fractions) + assert fractions[-1] == 1 + assert (tmp_path / "a.txt").is_file() + # --------------------------------------------------------------------------- # _rename_with_retry @@ -2137,6 +2174,20 @@ class TestSevenZipExtractAll: mock_pb.assert_called_once_with("Unpacking 7z") mock_pb.return_value.update.assert_called() + def test_progress_callback_replaces_bar(self, tmp_path: Path) -> None: + """A progress callback wins over progress_header; 7z reports 1.0 once.""" + buf = self._make_7z({"file.txt": b"x"}) + out = tmp_path / "out" + out.mkdir() + fractions: list[float] = [] + with patch("esphome.framework_helpers.ProgressBar") as mock_pb: + _7z_extract_all( + buf, out, progress_header="Unpacking 7z", progress=fractions.append + ) + mock_pb.assert_not_called() + assert fractions == [1] + assert (out / "file.txt").is_file() + def test_absolute_path_in_names_skipped(self, tmp_path: Path) -> None: """Names that resolve as absolute are silently skipped.""" import py7zr diff --git a/tests/unit_tests/test_platformio_registry.py b/tests/unit_tests/test_platformio_registry.py index 6ba8691c4e..081ecf6b1f 100644 --- a/tests/unit_tests/test_platformio_registry.py +++ b/tests/unit_tests/test_platformio_registry.py @@ -4,6 +4,7 @@ from __future__ import annotations from contextlib import contextmanager import json +import logging import os from pathlib import Path from unittest.mock import MagicMock, patch @@ -307,7 +308,11 @@ def test_install_package_downloads_via_registry(tmp_path: Path) -> None: "pkg", "1.0.0", dest, [], tmp_path / "dl", expect=("payload",) ) assert mock_download.call_args[0][0] == "http://x/pkg.tar.gz" - assert mock_download.call_args[1] == {"sha256": "abc123", "size": 42} + assert mock_download.call_args[1] == { + "sha256": "abc123", + "size": 42, + "progress": None, + } def test_install_package_validates_expected_layout(tmp_path: Path) -> None: @@ -723,3 +728,130 @@ def test_prefetch_packages_unexpected_failure_warns( tmp_path / "dl", ) assert "TypeError" in caplog.text + + +def _spec(name: str, version: str, dest: Path, mirrors=None, expect=("payload",)): + return (name, version, dest, mirrors or [], expect) + + +def test_install_packages_extracts_verified_archives_in_parallel( + tmp_path: Path, +) -> None: + """Two prefetched archives install concurrently under one shared bar.""" + dl = tmp_path / "dl" + dl.mkdir() + (dl / "a-1.0").write_bytes(b"x" * 10) + (dl / "b-2.0").write_bytes(b"y" * 20) + with patch.object(registry, "install_package") as mock_install: + registry.install_packages( + [_spec("a", "1.0", tmp_path / "a"), _spec("b", "2.0", tmp_path / "b")], dl + ) + assert mock_install.call_count == 2 + calls = sorted(mock_install.call_args_list, key=lambda c: c[0][0]) + for c, (name, version) in zip(calls, [("a", "1.0"), ("b", "2.0")], strict=True): + assert c[0][:3] == (name, version, tmp_path / name) + assert c[1]["expect"] == ("payload",) + assert callable(c[1]["extract_progress"]) + # Driving the tracker exercises the fraction-to-bytes scaling + c[1]["extract_progress"](0.5) + c[1]["extract_progress"](1.0) + + +def test_install_packages_single_archive_stays_sequential(tmp_path: Path) -> None: + """One verified archive has nothing to parallelize; original order kept.""" + dl = tmp_path / "dl" + dl.mkdir() + (dl / "a-1.0").write_bytes(b"x") + specs = [_spec("a", "1.0", tmp_path / "a"), _spec("b", "2.0", tmp_path / "b")] + with patch.object(registry, "install_package") as mock_install: + registry.install_packages(specs, dl) + assert [c[0][0] for c in mock_install.call_args_list] == ["a", "b"] + for c in mock_install.call_args_list: + assert "extract_progress" not in c[1] + + +def test_install_packages_mirror_and_marker_stay_sequential(tmp_path: Path) -> None: + """Mirror overrides and marker hits never enter the parallel batch.""" + dl = tmp_path / "dl" + dl.mkdir() + for name, ver in (("a", "1.0"), ("b", "2.0"), ("c", "3.0"), ("d", "4.0")): + (dl / f"{name}-{ver}").write_bytes(b"x") + marked = tmp_path / "c" + marked.mkdir() + (marked / ".esphome_extracted").touch() + specs = [ + _spec("a", "1.0", tmp_path / "a"), + _spec("b", "2.0", tmp_path / "b", mirrors=["http://m"]), + _spec("c", "3.0", marked), + _spec("d", "4.0", tmp_path / "d"), + ] + with patch.object(registry, "install_package") as mock_install: + registry.install_packages(specs, dl) + sequential = [ + c for c in mock_install.call_args_list if "extract_progress" not in c[1] + ] + batched = [c for c in mock_install.call_args_list if "extract_progress" in c[1]] + assert sorted(c[0][0] for c in sequential) == ["b", "c"] + assert sorted(c[0][0] for c in batched) == ["a", "d"] + + +def test_install_packages_first_failure_reraised( + tmp_path: Path, caplog: pytest.LogCaptureFixture +) -> None: + """Installs are mandatory: the first failure propagates, extras are logged.""" + dl = tmp_path / "dl" + dl.mkdir() + (dl / "a-1.0").write_bytes(b"x") + (dl / "b-2.0").write_bytes(b"y") + boom = EsphomeError("bad layout") + + def _fail(name: str, *_a, **_kw) -> None: + raise boom if name == "a" else EsphomeError("also bad") + + with ( + patch.object(registry, "install_package", side_effect=_fail), + pytest.raises(EsphomeError), + ): + registry.install_packages( + [_spec("a", "1.0", tmp_path / "a"), _spec("b", "2.0", tmp_path / "b")], dl + ) + assert "Could not install" in caplog.text + + +def test_install_package_extract_progress_suppresses_bars( + tmp_path: Path, caplog: pytest.LogCaptureFixture +) -> None: + """A batched install routes extraction fractions to the caller and keeps + both private bars and per-package INFO lines off the shared bar.""" + dest = tmp_path / "pkg" + fractions: list[float] = [] + with ( + caplog.at_level(logging.INFO), + patch.object(registry, "download_with_resume") as mock_download, + patch.object(registry, "archive_extract_all") as mock_extract, + patch.object( + registry, + "registry_download", + return_value=("http://x/pkg.tar.gz", "abc123", 42), + ), + ): + mock_extract.side_effect = lambda *_a, **_kw: (dest / "payload").mkdir( + parents=True + ) + registry.install_package( + "pkg", + "1.0.0", + dest, + [], + tmp_path / "dl", + expect=("payload",), + extract_progress=fractions.append, + ) + assert mock_extract.call_args[1]["progress"] == fractions.append + # The download tracker reports zero bytes, keeping the shared bar honest + download_progress = mock_download.call_args[1]["progress"] + assert callable(download_progress) + download_progress(42) + assert fractions == [0.0] + assert "Downloading pkg" not in caplog.text + assert "Extracting pkg" not in caplog.text