github/esphome
2
0
Fork 0
mirror of https://github.com/esphome/esphome.git synced 2026-09-11 23:37:34 +00:00
Code Issues Packages Projects Releases Wiki Activity

Extract native toolchain package archives in parallel

Browse Source
This commit is contained in:
J. Nick Koston
2026-08-27 17:30:29 -05:00
parent 1f77384389
commit 563083b4a9
6 changed files with 340 additions and 51 deletions
Download Patch File Download Diff File Expand all files Collapse all files
esphome/arduino8266/framework.py
+3 -4
View File
@@ -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
)
esphome/framework_helpers.py
+43 -23
View File
@@ -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
esphome/platformio/registry.py
+91 -5
View File
@@ -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]
tests/unit_tests/test_arduino8266_framework.py
+19 -18
View File
@@ -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 == (
[
tests/unit_tests/test_framework_helpers.py
+51
View File
@@ -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
tests/unit_tests/test_platformio_registry.py
+133 -1
View File
@@ -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