[core] Prefetch PlatformIO packages in parallel (#18769)

esphome/espidf/framework.py

@@ -2,7 +2,6 @@
 
 from collections.abc import Callable
 from ctypes.util import find_library
-from functools import partial
 import json
 import logging
 import os
@@ -24,16 +23,17 @@ from esphome.framework_helpers import (
     create_venv,
     download_and_extract,
     download_from_mirrors,
-    download_with_resume,
     failure_reason,
     get_python_env_executable_path,
     get_system_python_path,
+    resume_fetch_job,
     rmdir,
     run_batch_downloads,
     run_command,
     run_command_ok,
     str_to_lst_of_str,
     tool_version_runs,
+    warn_prefetch_failures,
 )
 from esphome.helpers import write_file_if_changed
 
@@ -686,18 +686,6 @@ def _patch_tools_json_demote_unused_tools(framework_path: Path) -> None:
     )
 
 
-def _download_tool(
-    dist_path: Path, entry: dict, tracker: Callable[[int], None]
-) -> None:
-    download_with_resume(
-        entry["url"],
-        dist_path / entry["dest"],
-        sha256=entry["sha256"],
-        size=entry["size"],
-        progress=tracker,
-    )
-
-
 def _prefetch_idf_tool_archives(
     framework_path: Path,
     targets_str: str,
@@ -775,15 +763,17 @@ def _prefetch_idf_tool_archives(
                 (
                     entry["name"],
                     entry["size"],
-                    partial(_download_tool, dist_path, entry),
+                    resume_fetch_job(
+                        entry["url"],
+                        dist_path / entry["dest"],
+                        sha256=entry["sha256"],
+                        size=entry["size"],
+                    ),
                 )
                 for entry in entries
             ],
         )
-        for name, e in failures:
-            # failure_reason: a message-less exception must not log blank
-            _LOGGER.warning("Could not prefetch %s: %s", name, failure_reason(e))
-            _LOGGER.debug("Prefetch failure detail", exc_info=e)
+        warn_prefetch_failures(failures)
         if len(failures) == len(entries):
             # A systematic fault, not one flaky mirror: the resume
             # workaround (#17703) is off for this whole install

esphome/framework_helpers.py

@@ -701,7 +701,7 @@ def _write_download_meta(
         _LOGGER.debug("Could not update download metadata %s: %s", meta, e)
 
 
-def _content_length(resp: "requests.Response") -> int:
+def content_length(resp: "requests.Response") -> int:
     """Return the response's Content-Length, or 0 when absent or malformed.
 
     0 means "unknown", which downstream disables the progress bar and the
@@ -744,7 +744,7 @@ def _stream_response_to_file(
     """
     f.seek(offset)
     f.truncate(offset)
-    total_size = size or offset + _content_length(resp)
+    total_size = size or offset + content_length(resp)
     downloaded = offset
     own_bar: ProgressBar | None = None
     if progress is None:
@@ -909,6 +909,19 @@ def _part_path(dest: Path) -> Path:
     return dest.with_name(dest.name + ".part")
 
 
+def discard_partial_download(dest: Path) -> None:
+    """Remove ``dest`` and the resume sidecars of an abandoned download."""
+    part = _part_path(dest)
+    for stale in (dest, part, part.with_name(part.name + ".meta")):
+        try:
+            stale.unlink()
+        except FileNotFoundError:
+            continue
+        except OSError as err:
+            # The caller's cache is never pruned; leave a trace
+            _LOGGER.debug("Could not remove %s: %s", stale, err)
+
+
 def _cancellable_sleep(
     delay: float, progress: Callable[[int], None] | None, done: int
 ) -> None:
@@ -922,6 +935,31 @@ def _cancellable_sleep(
         time.sleep(min(0.5, remaining))
 
 
+def resume_fetch_job(
+    url: str, dest: PathType, **kwargs
+) -> Callable[[Callable[[int], None]], None]:
+    """A ``run_batch_downloads`` job callable wrapping ``download_with_resume``.
+
+    Forwards the runner's positional tracker as the ``progress`` keyword.
+    """
+
+    def fetch(tracker: Callable[[int], None]) -> None:
+        download_with_resume(url, dest, progress=tracker, **kwargs)
+
+    return fetch
+
+
+def warn_prefetch_failures(
+    failures: list[tuple[str, BaseException]],
+    message: str = "Could not prefetch %s: %s",
+) -> None:
+    """Warn per failed batch-prefetch job; the caller's installer retries them."""
+    for name, err in failures:
+        # failure_reason: a message-less exception must not log blank
+        _LOGGER.warning(message, name, failure_reason(err))
+        _LOGGER.debug("Prefetch failure detail", exc_info=err)
+
+
 def download_with_resume(
     url: str,
     dest: PathType,
@@ -1022,7 +1060,7 @@ def download_with_resume(
                         streamed = True
                         if offset == 0:
                             validator = _response_validator(resp)
-                            expected_total = _content_length(resp)
+                            expected_total = content_length(resp)
                             # Recorded so a later run can prove an If-Range
                             # resume of this part file safe.
                             _write_download_meta(meta, url, validator, expected_total)

esphome/platformio/library.py

@@ -35,6 +35,7 @@ from esphome.framework_helpers import (
     failure_reason,
     rmdir,
     run_batch_downloads,
+    warn_prefetch_failures,
 )
 
 _LOGGER = logging.getLogger(__name__)
@@ -977,14 +978,10 @@ def _prefetch_wave(
                 for c in components
             ],
         )
-        for name, err in failures:
-            # The sequential call below retries and raises the real error
-            _LOGGER.warning(
-                "Prefetch of %s failed (retrying sequentially): %s",
-                name,
-                failure_reason(err),
-            )
-            _LOGGER.debug("Prefetch failure detail", exc_info=err)
+        # The sequential call below retries and raises the real error
+        warn_prefetch_failures(
+            failures, "Prefetch of %s failed (retrying sequentially): %s"
+        )
     except Exception as err:  # noqa: BLE001  # pylint: disable=broad-exception-caught
         # Same policy as the ESP-IDF twin: the prefetch must never become a
         # new way for the build to fail

esphome/platformio/toolchain.py

@@ -96,7 +96,7 @@ def _clean_platformio_python_env(config: "ProjectConfig", core_dir: Path) -> Non
         rmtree(penv)
 
 
-def _current_python_minor() -> str:
+def current_python_minor() -> str:
     """Return the running interpreter's ``major.minor`` (e.g. ``3.13``)."""
     return f"{sys.version_info.major}.{sys.version_info.minor}"
 
@@ -161,7 +161,7 @@ def heal_platformio_python_env() -> None:
 
 def _check_platformio_python_stamp(config: "ProjectConfig") -> None:
     """Compare the stamp to the running interpreter; wipe and restamp on mismatch."""
-    current = _current_python_minor()
+    current = current_python_minor()
     stamp_dir = _pio_stamp_dir(config)
     # Host the stamp/lock even before PlatformIO's first run creates the dir.
     stamp_dir.mkdir(parents=True, exist_ok=True)
@@ -289,6 +289,12 @@ def copy_ccache_script() -> None:
     )
 
 
+def default_libdeps_dir() -> str:
+    """The PLATFORMIO_LIBDEPS_DIR value a pio run defaults to; the package
+    prefetch must resolve installed libraries against the same dir."""
+    return str(CORE.relative_piolibdeps_path().absolute())
+
+
 def run_platformio_cli(*args, **kwargs) -> str | int:
     # Re-provision the PlatformIO cache if the interpreter's major.minor changed
     # since it was last built; a stale platform otherwise rejects the new Python
@@ -296,9 +302,7 @@ def run_platformio_cli(*args, **kwargs) -> str | int:
     heal_platformio_python_env()
     os.environ["PLATFORMIO_FORCE_COLOR"] = "true"
     os.environ["PLATFORMIO_BUILD_DIR"] = str(CORE.relative_pioenvs_path().absolute())
-    os.environ.setdefault(
-        "PLATFORMIO_LIBDEPS_DIR", str(CORE.relative_piolibdeps_path().absolute())
-    )
+    os.environ.setdefault("PLATFORMIO_LIBDEPS_DIR", default_libdeps_dir())
     # Suppress Python syntax warnings from third-party scripts during compilation
     os.environ.setdefault("PYTHONWARNINGS", "ignore::SyntaxWarning")
     # Increase uv retry count to handle transient network errors (default is 3)
@@ -346,6 +350,9 @@ def run_platformio_cli_run(config, verbose, *args, **kwargs) -> str | int:
 
 
 def run_compile(config, verbose):
+    from esphome.platformio.prefetch import prefetch_platformio_packages
+
+    prefetch_platformio_packages()
     args = []
     if CONF_COMPILE_PROCESS_LIMIT in config[CONF_ESPHOME]:
         args += [f"-j{config[CONF_ESPHOME][CONF_COMPILE_PROCESS_LIMIT]}"]

tests/unit_tests/test_espidf_framework.py

@@ -911,7 +911,7 @@ def test_prefetch_leaves_unverifiable_entries_to_the_installer(
             "esphome.espidf.framework.run_command",
             return_value=(True, json.dumps(entries), ""),
         ),
-        patch("esphome.espidf.framework.download_with_resume") as download,
+        patch("esphome.framework_helpers.download_with_resume") as download,
         patch("esphome.espidf.framework.get_system_python_path", return_value="python"),
         patch("esphome.framework_helpers._BatchDownloadProgress") as progress_cls,
     ):
@@ -934,7 +934,7 @@ def test_prefetch_all_entries_unverifiable_is_a_noop(tmp_path: Path) -> None:
             "esphome.espidf.framework.run_command",
             return_value=(True, json.dumps(entries), ""),
         ),
-        patch("esphome.espidf.framework.download_with_resume") as download,
+        patch("esphome.framework_helpers.download_with_resume") as download,
         patch("esphome.espidf.framework.get_system_python_path", return_value="python"),
     ):
         _prefetch_idf_tool_archives(tmp_path, "esp32", ["required"], None)
@@ -952,7 +952,7 @@ def test_prefetch_dedupes_entries_by_dest(tmp_path: Path) -> None:
             "esphome.espidf.framework.run_command",
             return_value=(True, json.dumps(entries), ""),
         ),
-        patch("esphome.espidf.framework.download_with_resume") as download,
+        patch("esphome.framework_helpers.download_with_resume") as download,
         patch("esphome.espidf.framework.get_system_python_path", return_value="python"),
         patch("esphome.framework_helpers._BatchDownloadProgress"),
     ):
@@ -967,7 +967,7 @@ def test_prefetch_downloads_each_archive_with_resume(tmp_path: Path) -> None:
             "esphome.espidf.framework.run_command",
             return_value=(True, _PREFETCH_JSON, ""),
         ),
-        patch("esphome.espidf.framework.download_with_resume") as download,
+        patch("esphome.framework_helpers.download_with_resume") as download,
         patch("esphome.espidf.framework.get_system_python_path", return_value="python"),
         patch("esphome.framework_helpers._BatchDownloadProgress") as progress_cls,
     ):
@@ -1011,7 +1011,7 @@ def test_prefetch_downloads_archives_concurrently(tmp_path: Path) -> None:
             "esphome.espidf.framework.run_command",
             return_value=(True, json.dumps(entries), ""),
         ),
-        patch("esphome.espidf.framework.download_with_resume") as download,
+        patch("esphome.framework_helpers.download_with_resume") as download,
         patch("esphome.espidf.framework.get_system_python_path", return_value="python"),
         patch(
             "esphome.framework_helpers.ThreadPoolExecutor", wraps=ThreadPoolExecutor
@@ -1032,7 +1032,7 @@ def test_prefetch_skips_already_downloaded_archives(tmp_path: Path) -> None:
             "esphome.espidf.framework.run_command",
             return_value=(True, _PREFETCH_JSON, ""),
         ),
-        patch("esphome.espidf.framework.download_with_resume") as download,
+        patch("esphome.framework_helpers.download_with_resume") as download,
         patch("esphome.espidf.framework.get_system_python_path", return_value="python"),
     ):
         _prefetch_idf_tool_archives(tmp_path, "esp32", ["required"], None)
@@ -1065,7 +1065,7 @@ def test_prefetch_failures_never_raise(
     with (
         patch("esphome.espidf.framework.run_command", return_value=run_result),
         patch(
-            "esphome.espidf.framework.download_with_resume",
+            "esphome.framework_helpers.download_with_resume",
             side_effect=download_error,
         ),
         patch("esphome.espidf.framework.get_system_python_path", return_value="python"),
@@ -1087,7 +1087,7 @@ def test_prefetch_total_failure_logs_error(
             return_value=(True, _PREFETCH_JSON, ""),
         ),
         patch(
-            "esphome.espidf.framework.download_with_resume",
+            "esphome.framework_helpers.download_with_resume",
             side_effect=OSError("proxy refuses everything"),
         ),
         patch("esphome.espidf.framework.get_system_python_path", return_value="python"),
@@ -1112,7 +1112,7 @@ def test_prefetch_one_failed_archive_does_not_stop_the_rest(
             return_value=(True, _PREFETCH_JSON, ""),
         ),
         patch(
-            "esphome.espidf.framework.download_with_resume",
+            "esphome.framework_helpers.download_with_resume",
             side_effect=_fail_cmake_download,
         ) as download,
         patch("esphome.espidf.framework.get_system_python_path", return_value="python"),
@@ -1133,7 +1133,7 @@ def test_prefetch_finishes_progress_bar_and_cancels_queue(tmp_path: Path) -> Non
             "esphome.espidf.framework.run_command",
             return_value=(True, _PREFETCH_JSON, ""),
         ),
-        patch("esphome.espidf.framework.download_with_resume"),
+        patch("esphome.framework_helpers.download_with_resume"),
         patch("esphome.espidf.framework.get_system_python_path", return_value="python"),
         patch("esphome.framework_helpers._BatchDownloadProgress") as progress_cls,
         patch("esphome.framework_helpers.ThreadPoolExecutor") as pool_cls,

tests/unit_tests/test_framework_helpers.py

@@ -2280,6 +2280,32 @@ class TestGetProjectCxxCompileFlags:
             assert get_project_cxx_compile_flags() == []
 
 
+def test_resume_fetch_job_threads_tracker(tmp_path: Path) -> None:
+    """The batch runner passes the tracker positionally; the shared adapter
+    must deliver it as download_with_resume's progress keyword."""
+    from esphome.framework_helpers import resume_fetch_job
+
+    with patch("esphome.framework_helpers.download_with_resume") as mock_download:
+        fetch = resume_fetch_job("https://x/a.zip", tmp_path / "a", sha256="ff", size=9)
+        tracker = lambda done: None  # noqa: E731
+        fetch(tracker)
+    mock_download.assert_called_once_with(
+        "https://x/a.zip", tmp_path / "a", progress=tracker, sha256="ff", size=9
+    )
+
+
+def test_warn_prefetch_failures_names_each_failure(
+    caplog: pytest.LogCaptureFixture,
+) -> None:
+    """The shared failure loop warns per job with the failure reason."""
+    from esphome.framework_helpers import warn_prefetch_failures
+
+    warn_prefetch_failures([("toolchain-x@1", OSError("down"))])
+    assert "Could not prefetch toolchain-x@1: down" in caplog.text
+    warn_prefetch_failures([("lib", OSError("gone"))], "Prefetch of %s failed: %s")
+    assert "Prefetch of lib failed: gone" in caplog.text
+
+
 @pytest.mark.parametrize(
     ("platform", "input_path", "expected"),
     [
@@ -2312,3 +2338,18 @@ def test_strip_win_long_path_prefix(
     r"""``\\?\`` and ``\\?\UNC\`` prefixes are stripped only on win32."""
     with patch("esphome.framework_helpers.sys.platform", platform):
         assert framework_helpers.strip_win_long_path_prefix(input_path) == expected
+
+
+def test_discard_partial_download_logs_undeletable(
+    tmp_path: Path, caplog: pytest.LogCaptureFixture
+) -> None:
+    """An unremovable staging file leaves a debug trace; the caller's
+    cache is never pruned, so silence would hide unbounded growth."""
+    dest = tmp_path / "archive"
+    dest.write_bytes(b"stale")
+    with (
+        patch.object(Path, "unlink", side_effect=OSError("busy")),
+        caplog.at_level(logging.DEBUG),
+    ):
+        framework_helpers.discard_partial_download(dest)
+    assert "Could not remove" in caplog.text

tests/unit_tests/test_platformio_toolchain.py

@@ -932,8 +932,13 @@ def test_run_compile(setup_core: Path, mock_run_platformio_cli_run: Mock) -> Non
     config = {CONF_ESPHOME: {CONF_COMPILE_PROCESS_LIMIT: 4}}
     mock_run_platformio_cli_run.return_value = 0
 
-    toolchain.run_compile(config, verbose=True)
+    with patch(
+        "esphome.platformio.prefetch.prefetch_platformio_packages"
+    ) as mock_prefetch:
+        toolchain.run_compile(config, verbose=True)
 
+    # The only wiring of the prefetch into a build lives here
+    mock_prefetch.assert_called_once_with()
     mock_run_platformio_cli_run.assert_called_once_with(config, True, "-j4")
 
 
@@ -947,7 +952,8 @@ def test_run_compile_without_process_limit(
     config = {CONF_ESPHOME: {}}
     mock_run_platformio_cli_run.return_value = 0
 
-    toolchain.run_compile(config, verbose=False)
+    with patch("esphome.platformio.prefetch.prefetch_platformio_packages"):
+        toolchain.run_compile(config, verbose=False)
 
     mock_run_platformio_cli_run.assert_called_once_with(config, False)
 
@@ -1677,8 +1683,8 @@ def pio_core_dir(tmp_path: Path) -> Path:
 
 
 def test_current_python_minor_matches_running_interpreter() -> None:
-    """_current_python_minor returns major.minor of the running interpreter."""
-    assert toolchain._current_python_minor() == _CURRENT_MINOR
+    """current_python_minor returns major.minor of the running interpreter."""
+    assert toolchain.current_python_minor() == _CURRENT_MINOR
 
 
 def test_pio_stamp_round_trip(tmp_path: Path) -> None: