[micro_wake_word] Download models in parallel (#17701)

esphome/components/micro_wake_word/__init__.py

@@ -2,6 +2,7 @@ import hashlib
 import json
 import logging
 from pathlib import Path
+import re
 from urllib.parse import urljoin
 
 from esphome import automation, external_files, git
@@ -9,6 +10,7 @@ from esphome.automation import register_action, register_condition
 from esphome.bundle import add_bundle_file
 import esphome.codegen as cg
 from esphome.components import esp32, microphone, ota, psram
+from esphome.components.http_request import validate_url
 import esphome.config_validation as cv
 from esphome.const import (
     CONF_FILE,
@@ -209,33 +211,13 @@ def _validate_manifest_version(manifest_data):
         raise cv.Invalid("Invalid manifest file, missing 'version' key")
 
 
-def _process_http_source(config):
-    url = config[CONF_URL]
-    path = _compute_local_file_path(config)
-
-    json_path = path / "manifest.json"
-
-    json_contents = external_files.download_content(url, json_path)
-
-    manifest_data = json.loads(json_contents)
-    if not isinstance(manifest_data, dict):
-        raise cv.Invalid("Manifest file must contain a JSON object")
-
-    model = manifest_data[CONF_MODEL]
-    model_url = urljoin(url, model)
-
-    model_path = path / model
-
-    external_files.download_content(str(model_url), model_path)
-
-    return config
-
-
-HTTP_SCHEMA = cv.All(
+HTTP_SCHEMA = cv.Schema(
     {
-        cv.Required(CONF_URL): cv.url,
-    },
-    _process_http_source,
+        # validate_url only accepts http(s); the shorthand validator relies
+        # on this branch rejecting git shorthands ("github://...") so they
+        # fall through to the git branch.
+        cv.Required(CONF_URL): validate_url,
+    }
 )
 
 
@@ -280,6 +262,13 @@ LOCAL_SCHEMA = cv.All(
 )
 
 
+# Bare model names in the official model repository ("okay_nabu"). Must not
+# overlap with local paths, http(s) urls, or git shorthands
+# ("github://user/repo/file.json@ref"), which the shorthand validator tries
+# next; anything containing "/", ":" or "@" is not a model name.
+_MODEL_NAME_RE = re.compile(r"[A-Za-z0-9_.-]+")
+
+
 def _validate_source_model_name(value):
     if not isinstance(value, str):
         raise cv.Invalid("Model name must be a string")
@@ -287,6 +276,9 @@ def _validate_source_model_name(value):
     if value.endswith(".json"):
         raise cv.Invalid("Model name must not end with .json")
 
+    if not _MODEL_NAME_RE.fullmatch(value):
+        raise cv.Invalid("Model name may only contain letters, numbers, . _ -")
+
     return MODEL_SOURCE_SCHEMA(
         {
             CONF_TYPE: TYPE_HTTP,
@@ -376,6 +368,58 @@ def _maybe_empty_vad_schema(value):
     return VAD_MODEL_SCHEMA(value)
 
 
+def _download_http_models(config: ConfigType) -> ConfigType:
+    """Download every http-sourced manifest and model file in two concurrent
+    batches (all manifests, then all model files).
+
+    The model file's URL only becomes known once its manifest has been
+    fetched and parsed, so the two stages cannot be merged into one batch.
+    """
+    model_parameters = [*config[CONF_MODELS]]
+    if vad := config.get(CONF_VAD):
+        model_parameters.append(vad)
+    # Keyed by cache path so a URL referenced twice is fetched and parsed once
+    http_models: dict[Path, str] = {
+        _compute_local_file_path(model_config): model_config[CONF_URL]
+        for parameters in model_parameters
+        if (model_config := parameters.get(CONF_MODEL)) is not None
+        and model_config.get(CONF_TYPE) == TYPE_HTTP
+    }
+    if not http_models:
+        return config
+
+    external_files.download_content_many(
+        ((url, path / "manifest.json") for path, url in http_models.items()),
+        description="wake word manifest(s)",
+    )
+
+    model_files: list[tuple[str, Path]] = []
+    errors: list[cv.Invalid] = []
+    for path, url in http_models.items():
+        try:
+            manifest_data = json.loads((path / "manifest.json").read_bytes())
+        except (OSError, ValueError) as e:
+            errors.append(cv.Invalid(f"Invalid manifest file at {url}: {e}"))
+            continue
+        if not isinstance(manifest_data, dict):
+            errors.append(
+                cv.Invalid(f"Manifest file at {url} must contain a JSON object")
+            )
+            continue
+        model = manifest_data.get(CONF_MODEL)
+        if not isinstance(model, str):
+            errors.append(
+                cv.Invalid(f"Manifest file at {url} is missing the 'model' key")
+            )
+            continue
+        model_files.append((urljoin(url, model), path / model))
+    if errors:
+        raise cv.MultipleInvalid(errors)
+
+    external_files.download_content_many(model_files, description="wake word model(s)")
+    return config
+
+
 CONFIG_SCHEMA = cv.All(
     cv.Schema(
         {
@@ -409,6 +453,7 @@ CONFIG_SCHEMA = cv.All(
         }
     ).extend(cv.COMPONENT_SCHEMA),
     cv.only_on_esp32,
+    _download_http_models,
 )
 
 

esphome/external_files.py

@@ -165,11 +165,8 @@ def download_content(url: str, path: Path, timeout: int = NETWORK_TIMEOUT) -> by
         _LOGGER.debug("Remote file has not changed %s", url)
         return path.read_bytes()
 
-    _LOGGER.debug(
-        "Remote file has changed, downloading from %s to %s",
-        url,
-        path,
-    )
+    _LOGGER.info("Downloading %s", url)
+    _LOGGER.debug("Saving to %s", path)
 
     try:
         req = requests.get(
@@ -210,9 +207,13 @@ def download_content_many(
     items: Iterable[tuple[str, Path]],
     timeout: int = NETWORK_TIMEOUT,
     max_workers: int = DEFAULT_DOWNLOAD_WORKERS,
+    description: str = "remote file(s)",
 ) -> None:
     """Run `download_content` for each (url, path) pair concurrently.
 
+    `description` names the kind of files in the progress log line, e.g.
+    "wake word manifest(s)".
+
     Wall time drops from `sum(latency)` to roughly `max(latency)` for cached
     files where the HEAD round-trip dominates. All workers run to
     completion before this returns; every `cv.Invalid` raised by a worker
@@ -230,6 +231,7 @@ def download_content_many(
     seen: dict[Path, str] = {path: url for url, path in items}
     if not seen:
         return
+    _LOGGER.info("Checking %d %s for updates", len(seen), description)
     if len(seen) == 1:
         path, url = next(iter(seen.items()))
         download_content(url, path, timeout)

tests/unit_tests/components/micro_wake_word/test_init.py

@@ -0,0 +1,169 @@
+"""Tests for the micro_wake_word model source validation and downloads."""
+
+import json
+from pathlib import Path
+from unittest.mock import MagicMock, patch
+
+import pytest
+
+from esphome.components import micro_wake_word as mww
+import esphome.config_validation as cv
+from esphome.const import (
+    CONF_FILE,
+    CONF_MODEL,
+    CONF_PATH,
+    CONF_REF,
+    CONF_TYPE,
+    CONF_URL,
+)
+
+
+@pytest.fixture
+def mock_download_content_many() -> MagicMock:
+    """Patch the concurrent download helper so no network is involved."""
+    with patch(
+        "esphome.components.micro_wake_word.external_files.download_content_many"
+    ) as m:
+        yield m
+
+
+def test_shorthand_model_name_resolves_without_network(
+    mock_download_content_many: MagicMock,
+) -> None:
+    config = mww._validate_source_shorthand("okay_nabu")
+    assert config[CONF_TYPE] == mww.TYPE_HTTP
+    assert config[CONF_URL] == (
+        "https://github.com/esphome/micro-wake-word-models/raw/main/models/v2/okay_nabu.json"
+    )
+    mock_download_content_many.assert_not_called()
+
+
+def test_shorthand_git_with_ref_not_captured_as_model_name(
+    setup_core: Path, tmp_path: Path
+) -> None:
+    repo_dir = tmp_path / "repo"
+    repo_dir.mkdir()
+    (repo_dir / "model.json").write_text("{}")
+    with patch(
+        "esphome.components.micro_wake_word.git.clone_or_update",
+        return_value=(repo_dir, None),
+    ):
+        config = mww._validate_source_shorthand("github://user/repo/model.json@main")
+    assert config[CONF_TYPE] == "git"
+    assert config[CONF_URL] == "https://github.com/user/repo.git"
+    assert config[CONF_FILE] == "model.json"
+    assert config[CONF_REF] == "main"
+
+
+def test_shorthand_local_path_not_captured_as_model_name(
+    setup_core: Path, tmp_path: Path
+) -> None:
+    manifest = tmp_path / "model.json"
+    manifest.write_text("{}")
+    config = mww.MODEL_SOURCE_SCHEMA(str(manifest))
+    assert config[CONF_TYPE] == "local"
+    assert Path(config[CONF_PATH]) == manifest
+
+
+@pytest.mark.parametrize(
+    "value", ["some/path/file", "name@ref", "bad:name", "okay_nabu\n", "héllo"]
+)
+def test_model_name_rejects_non_identifiers(value: str) -> None:
+    with pytest.raises(cv.Invalid):
+        mww._validate_source_model_name(value)
+
+
+def _http_model(name: str) -> dict:
+    return {
+        CONF_MODEL: {
+            CONF_TYPE: mww.TYPE_HTTP,
+            CONF_URL: f"https://example.com/models/{name}.json",
+        }
+    }
+
+
+def _write_manifest(model_config: dict, contents: str) -> Path:
+    path = mww._compute_local_file_path(model_config[CONF_MODEL])
+    path.mkdir(parents=True, exist_ok=True)
+    manifest = path / "manifest.json"
+    manifest.write_text(contents)
+    return path
+
+
+def test_download_http_models_batches_manifests_then_models(
+    setup_core: Path, mock_download_content_many: MagicMock
+) -> None:
+    names = ("okay_nabu", "hey_mycroft", "vad")
+    models = {name: _http_model(name) for name in names}
+    paths = {
+        name: _write_manifest(models[name], json.dumps({"model": f"{name}.tflite"}))
+        for name in names
+    }
+    config = {
+        mww.CONF_MODELS: [
+            models["okay_nabu"],
+            models["hey_mycroft"],
+            # non-http sources must be ignored
+            {CONF_MODEL: {CONF_TYPE: "local", CONF_PATH: "x"}},
+        ],
+        mww.CONF_VAD: models["vad"],
+    }
+
+    assert mww._download_http_models(config) is config
+
+    assert mock_download_content_many.call_count == 2
+    manifest_items = list(mock_download_content_many.call_args_list[0].args[0])
+    assert manifest_items == [
+        (f"https://example.com/models/{name}.json", paths[name] / "manifest.json")
+        for name in names
+    ]
+    model_items = list(mock_download_content_many.call_args_list[1].args[0])
+    assert model_items == [
+        (f"https://example.com/models/{name}.tflite", paths[name] / f"{name}.tflite")
+        for name in names
+    ]
+
+
+def test_download_http_models_no_http_sources_skips_download(
+    mock_download_content_many: MagicMock,
+) -> None:
+    config = {mww.CONF_MODELS: [{CONF_MODEL: {CONF_TYPE: "local", CONF_PATH: "x"}}]}
+    assert mww._download_http_models(config) is config
+    mock_download_content_many.assert_not_called()
+
+
+@pytest.mark.parametrize(
+    ("contents", "message"),
+    [
+        ("not json", "Invalid manifest file"),
+        ("[1, 2]", "must contain a JSON object"),
+        ("{}", "missing the 'model' key"),
+    ],
+)
+def test_download_http_models_bad_manifest_raises(
+    setup_core: Path,
+    mock_download_content_many: MagicMock,
+    contents: str,
+    message: str,
+) -> None:
+    model = _http_model("okay_nabu")
+    config = {mww.CONF_MODELS: [model]}
+    _write_manifest(model, contents)
+
+    with pytest.raises(cv.Invalid, match=message):
+        mww._download_http_models(config)
+    # manifests were still fetched in one batch; the model batch never ran
+    assert mock_download_content_many.call_count == 1
+
+
+def test_download_http_models_collects_all_manifest_errors(
+    setup_core: Path, mock_download_content_many: MagicMock
+) -> None:
+    models = {name: _http_model(name) for name in ("one", "two")}
+    config = {mww.CONF_MODELS: list(models.values())}
+    _write_manifest(models["one"], "not json")
+    _write_manifest(models["two"], "[1]")
+
+    with pytest.raises(cv.MultipleInvalid) as excinfo:
+        mww._download_http_models(config)
+    assert len(excinfo.value.errors) == 2