396 lines
17 KiB
Python
396 lines
17 KiB
Python
"""Pinned-image input and provenance contract for Wakeword Training (#163).
|
|
|
|
This module deliberately does not implement wakeword generation, augmentation,
|
|
feature extraction, training, export, or evaluation. It validates the
|
|
commercial source lock, materializes only its declared inputs, and leaves the
|
|
actual ML stages to the released LiveKit CLI inside the isolated image.
|
|
"""
|
|
|
|
from __future__ import annotations
|
|
|
|
import hashlib
|
|
import json
|
|
import os
|
|
import re
|
|
import shutil
|
|
from dataclasses import dataclass
|
|
from enum import StrEnum
|
|
from pathlib import Path
|
|
from typing import Any
|
|
|
|
LIVEKIT_WAKEWORD_VERSION = "0.2.1"
|
|
LIVEKIT_WAKEWORD_REVISION = "1ec7f680df30ff4ca0ebae6b5983441e94b10980"
|
|
LIVEKIT_WAKEWORD_SDIST_SHA256 = "cf2d9cf4867812c06788f64c15e49abd909d9d6291f0a13f1c3f9cb649fa6127"
|
|
|
|
_COMMERCIAL_PURPOSES = frozenset({"wakeword_training", "classifier_distribution"})
|
|
_BLOCKED_LICENSE_MARKERS = ("unknown", "nc", "non-commercial", "research-only")
|
|
_BLOCKED_SOURCE_MARKERS = ("acav", "unknown_rir_mirror")
|
|
_SYNTHETIC_SOURCE_KINDS = frozenset({"piper_output", "voxcpm_output"})
|
|
_SAFE_TOKEN = re.compile(r"^[a-z0-9][a-z0-9._-]{0,127}$")
|
|
_ALLOWED_PURPOSES = frozenset(
|
|
{
|
|
"positive",
|
|
"adversarial_negative",
|
|
"near_miss",
|
|
"general_negative_speech",
|
|
"background",
|
|
"rir",
|
|
}
|
|
)
|
|
_ALLOWED_PARTITIONS = frozenset({"train", "validation", "calibration", "release_test"})
|
|
_REQUIRED_PURPOSE_PARTITIONS = frozenset(
|
|
{
|
|
("positive", "train"),
|
|
("positive", "release_test"),
|
|
("general_negative_speech", "train"),
|
|
("general_negative_speech", "validation"),
|
|
("background", "train"),
|
|
("rir", "train"),
|
|
}
|
|
)
|
|
|
|
|
|
class SourceLockError(ValueError):
|
|
"""A source lock cannot safely enter the requested training lane."""
|
|
|
|
|
|
class RunClassification(StrEnum):
|
|
COMMERCIAL_REVIEW_REQUIRED = "commercial_review_required"
|
|
LOCAL_ONLY = "local_only"
|
|
QUARANTINED = "quarantined"
|
|
|
|
@property
|
|
def commercial_promotion_eligible(self) -> bool:
|
|
return self is RunClassification.COMMERCIAL_REVIEW_REQUIRED
|
|
|
|
|
|
def classify_run(run_kind: str) -> RunClassification:
|
|
if run_kind == "commercial":
|
|
return RunClassification.COMMERCIAL_REVIEW_REQUIRED
|
|
if run_kind == "local_experiment":
|
|
return RunClassification.LOCAL_ONLY
|
|
if run_kind == "quarantined_smoke":
|
|
return RunClassification.QUARANTINED
|
|
raise SourceLockError("run_kind must be commercial, local_experiment, or quarantined_smoke")
|
|
|
|
|
|
def _sha256_file(path: Path) -> str:
|
|
digest = hashlib.sha256()
|
|
with path.open("rb") as file:
|
|
for block in iter(lambda: file.read(1024 * 1024), b""):
|
|
digest.update(block)
|
|
return digest.hexdigest()
|
|
|
|
|
|
def _require_mapping(value: object, label: str) -> dict[str, Any]:
|
|
if not isinstance(value, dict):
|
|
raise SourceLockError(f"{label} must be an object")
|
|
return value
|
|
|
|
|
|
def _require_string(value: object, label: str) -> str:
|
|
if not isinstance(value, str) or not value.strip():
|
|
raise SourceLockError(f"{label} must be a non-empty string")
|
|
return value.strip()
|
|
|
|
|
|
def _require_sha256(value: object, label: str) -> str:
|
|
digest = _require_string(value, label).removeprefix("sha256:").lower()
|
|
if len(digest) != 64 or any(character not in "0123456789abcdef" for character in digest):
|
|
raise SourceLockError(f"{label} must be a SHA-256 digest")
|
|
return digest
|
|
|
|
|
|
def _safe_input_path(root: Path, declared_path: str) -> Path:
|
|
candidate = (root / declared_path).resolve(strict=True)
|
|
if not candidate.is_file() or os.path.commonpath((root, candidate)) != str(root):
|
|
raise SourceLockError("source lock file path must identify a regular file below input_root")
|
|
return candidate
|
|
|
|
|
|
@dataclass(frozen=True)
|
|
class MaterializedInputs:
|
|
provenance_manifest_path: Path
|
|
recorded_manifest_path: Path
|
|
|
|
|
|
class InputMaterializer:
|
|
"""Verify and copy a source lock's declared inputs into one sealed run.
|
|
|
|
The input root is read-only when this code runs in the training image. The
|
|
run directory is the sole persistence target, so no upstream ``setup``
|
|
download can become an undeclared training input.
|
|
"""
|
|
|
|
def __init__(self, lock_path: Path, input_root: Path, run_dir: Path, *, run_kind: str) -> None:
|
|
self.lock_path = lock_path.resolve(strict=True)
|
|
self.input_root = input_root.resolve(strict=True)
|
|
self.run_dir = run_dir.resolve()
|
|
self.run_kind = run_kind
|
|
|
|
def materialize(self) -> MaterializedInputs:
|
|
classification = classify_run(self.run_kind)
|
|
lock = self._read_lock()
|
|
if lock.get("run_kind") != self.run_kind:
|
|
raise SourceLockError("source lock run_kind does not match the requested run")
|
|
sources = lock.get("sources")
|
|
if not isinstance(sources, list) or not sources:
|
|
raise SourceLockError("source lock must contain at least one source")
|
|
|
|
materialized_root = self.run_dir / "materialized"
|
|
materialized_root.mkdir(parents=True, exist_ok=False)
|
|
source_manifest: list[dict[str, object]] = []
|
|
supplied_purposes: set[tuple[str, str]] = set()
|
|
seen_source_ids: set[str] = set()
|
|
materialized_destinations: set[Path] = set()
|
|
for index, raw_source in enumerate(sources):
|
|
source = _require_mapping(raw_source, f"sources[{index}]")
|
|
source_id, purpose, partition, file_records = self._validate_source(
|
|
source, index, commercial=classification.commercial_promotion_eligible
|
|
)
|
|
if source_id in seen_source_ids:
|
|
raise SourceLockError(f"duplicate source_id: {source_id}")
|
|
seen_source_ids.add(source_id)
|
|
supplied_purposes.add((purpose, partition))
|
|
target_root = materialized_root / partition / purpose / source_id
|
|
target_root.mkdir(parents=True)
|
|
materialized_files: list[dict[str, str]] = []
|
|
for record in file_records:
|
|
source_path = _safe_input_path(self.input_root, record["path"])
|
|
if _sha256_file(source_path) != record["sha256"]:
|
|
raise SourceLockError(f"source hash mismatch: {record['path']}")
|
|
target = target_root / source_path.name
|
|
if target in materialized_destinations:
|
|
raise SourceLockError(f"duplicate materialized destination: {target.name}")
|
|
materialized_destinations.add(target)
|
|
shutil.copyfile(source_path, target)
|
|
target.chmod(0o444)
|
|
materialized_files.append(
|
|
{
|
|
"source_path": record["path"],
|
|
"materialized_path": target.relative_to(self.run_dir).as_posix(),
|
|
"sha256": record["sha256"],
|
|
}
|
|
)
|
|
source_manifest.append(
|
|
{
|
|
"source_id": source_id,
|
|
"purpose": purpose,
|
|
"partition": partition,
|
|
"files": materialized_files,
|
|
}
|
|
)
|
|
|
|
if classification.commercial_promotion_eligible:
|
|
missing = _REQUIRED_PURPOSE_PARTITIONS - supplied_purposes
|
|
if missing:
|
|
formatted = ", ".join(
|
|
f"{purpose}:{partition}" for purpose, partition in sorted(missing)
|
|
)
|
|
raise SourceLockError(
|
|
f"commercial source lock is missing required inputs: {formatted}"
|
|
)
|
|
|
|
provenance_path = self.run_dir / "sources.lock.json"
|
|
if self.lock_path != provenance_path:
|
|
shutil.copyfile(self.lock_path, provenance_path)
|
|
provenance_path.chmod(0o444)
|
|
attribution_dir = self.run_dir / "attribution"
|
|
attribution_dir.mkdir(exist_ok=True)
|
|
self._materialize_attribution(lock["sources"], attribution_dir)
|
|
recorded_manifest_path = self.run_dir / "recorded-inputs.json"
|
|
recorded_manifest_path.write_text(
|
|
json.dumps(
|
|
{
|
|
"schema_version": 1,
|
|
"classification": classification.value,
|
|
"sources": source_manifest,
|
|
},
|
|
sort_keys=True,
|
|
separators=(",", ":"),
|
|
)
|
|
+ "\n",
|
|
encoding="utf-8",
|
|
)
|
|
recorded_manifest_path.chmod(0o444)
|
|
return MaterializedInputs(
|
|
provenance_manifest_path=provenance_path,
|
|
recorded_manifest_path=recorded_manifest_path,
|
|
)
|
|
|
|
def _read_lock(self) -> dict[str, Any]:
|
|
try:
|
|
lock = json.loads(self.lock_path.read_text(encoding="utf-8"))
|
|
except (OSError, json.JSONDecodeError) as exc:
|
|
raise SourceLockError("source lock must be readable JSON") from exc
|
|
lock = _require_mapping(lock, "source lock")
|
|
if lock.get("schema_version") != 1:
|
|
raise SourceLockError("source lock schema_version must be 1")
|
|
return lock
|
|
|
|
def _validate_source(
|
|
self, source: dict[str, Any], index: int, *, commercial: bool
|
|
) -> tuple[str, str, str, list[dict[str, str]]]:
|
|
source_id = _require_string(source.get("source_id"), f"sources[{index}].source_id")
|
|
purpose = _require_string(source.get("purpose"), f"sources[{index}].purpose")
|
|
partition = _require_string(source.get("partition"), f"sources[{index}].partition")
|
|
source_kind = _require_string(source.get("source_kind"), f"sources[{index}].source_kind")
|
|
if not _SAFE_TOKEN.fullmatch(source_id):
|
|
raise SourceLockError(f"sources[{index}].source_id must be a safe token")
|
|
if purpose not in _ALLOWED_PURPOSES:
|
|
raise SourceLockError(f"sources[{index}].purpose is not an allowed purpose")
|
|
if partition not in _ALLOWED_PARTITIONS:
|
|
raise SourceLockError(f"sources[{index}].partition is not an allowed partition")
|
|
publisher = _require_mapping(
|
|
source.get("original_publisher"), f"sources[{index}].original_publisher"
|
|
)
|
|
_require_string(publisher.get("name"), f"sources[{index}].original_publisher.name")
|
|
_require_string(publisher.get("url"), f"sources[{index}].original_publisher.url")
|
|
_require_string(source.get("version"), f"sources[{index}].version")
|
|
_require_string(source.get("revision"), f"sources[{index}].revision")
|
|
_require_string(source.get("attribution"), f"sources[{index}].attribution")
|
|
license_data = _require_mapping(source.get("license"), f"sources[{index}].license")
|
|
license_id = _require_string(
|
|
license_data.get("spdx_id"), f"sources[{index}].license.spdx_id"
|
|
)
|
|
_require_string(license_data.get("terms_url"), f"sources[{index}].license.terms_url")
|
|
captured_text = _require_string(
|
|
license_data.get("captured_text_path"),
|
|
f"sources[{index}].license.captured_text_path",
|
|
)
|
|
captured_hash = _require_sha256(
|
|
license_data.get("captured_text_sha256"),
|
|
f"sources[{index}].license.captured_text_sha256",
|
|
)
|
|
license_path = _safe_input_path(self.input_root, captured_text)
|
|
if _sha256_file(license_path) != captured_hash:
|
|
raise SourceLockError(f"captured license hash mismatch: {captured_text}")
|
|
files = source.get("files")
|
|
if not isinstance(files, list) or not files:
|
|
raise SourceLockError(f"sources[{index}].files must be a non-empty list")
|
|
records: list[dict[str, str]] = []
|
|
for file_index, raw_file in enumerate(files):
|
|
file_data = _require_mapping(raw_file, f"sources[{index}].files[{file_index}]")
|
|
path = _require_string(
|
|
file_data.get("path"), f"sources[{index}].files[{file_index}].path"
|
|
)
|
|
digest = _require_sha256(
|
|
file_data.get("sha256"), f"sources[{index}].files[{file_index}].sha256"
|
|
)
|
|
size = file_data.get("size_bytes")
|
|
if not isinstance(size, int) or size < 1:
|
|
raise SourceLockError(
|
|
f"sources[{index}].files[{file_index}].size_bytes must be positive"
|
|
)
|
|
if _safe_input_path(self.input_root, path).stat().st_size != size:
|
|
raise SourceLockError(f"source size mismatch: {path}")
|
|
records.append({"path": path, "sha256": digest})
|
|
if commercial:
|
|
self._validate_commercial_source(source_id, source_kind, license_id, source, index)
|
|
return source_id, purpose, partition, records
|
|
|
|
@staticmethod
|
|
def _validate_commercial_source(
|
|
source_id: str, source_kind: str, license_id: str, source: dict[str, Any], index: int
|
|
) -> None:
|
|
normalized_id = source_id.lower()
|
|
normalized_license = license_id.lower()
|
|
if any(marker in normalized_id for marker in _BLOCKED_SOURCE_MARKERS) or any(
|
|
marker in normalized_license for marker in _BLOCKED_LICENSE_MARKERS
|
|
):
|
|
raise SourceLockError(f"blocked source for commercial mode: {source_id}")
|
|
commercial_data = _require_mapping(source.get("commercial"), f"sources[{index}].commercial")
|
|
if commercial_data.get("disposition") != "approved":
|
|
raise SourceLockError(f"commercial source is not approved: {source_id}")
|
|
allowed_purposes = commercial_data.get("allowed_purposes")
|
|
if not isinstance(allowed_purposes, list) or not _COMMERCIAL_PURPOSES.issubset(
|
|
set(allowed_purposes)
|
|
):
|
|
raise SourceLockError(f"commercial source has incomplete allowed_purposes: {source_id}")
|
|
_require_string(commercial_data.get("approver"), f"sources[{index}].commercial.approver")
|
|
_require_string(
|
|
commercial_data.get("approved_at"), f"sources[{index}].commercial.approved_at"
|
|
)
|
|
if (
|
|
source_kind in _SYNTHETIC_SOURCE_KINDS
|
|
and commercial_data.get("tts_output_approved") is not True
|
|
):
|
|
raise SourceLockError(f"unapproved synthetic source for commercial mode: {source_id}")
|
|
|
|
def _materialize_attribution(self, sources: object, attribution_dir: Path) -> None:
|
|
assert isinstance(sources, list)
|
|
for raw_source in sources:
|
|
source = _require_mapping(raw_source, "source")
|
|
source_id = _require_string(source.get("source_id"), "source.source_id")
|
|
license_data = _require_mapping(source.get("license"), "source.license")
|
|
license_path = _safe_input_path(
|
|
self.input_root,
|
|
_require_string(
|
|
license_data.get("captured_text_path"), "source.license.captured_text_path"
|
|
),
|
|
)
|
|
target = attribution_dir / f"{source_id}-license.txt"
|
|
shutil.copyfile(license_path, target)
|
|
target.chmod(0o444)
|
|
|
|
|
|
def build_stage_command(run_kind: str) -> tuple[str, ...]:
|
|
"""Return the only command an isolated training container may execute."""
|
|
|
|
classification = classify_run(run_kind)
|
|
if classification is RunClassification.COMMERCIAL_REVIEW_REQUIRED:
|
|
return (
|
|
"python",
|
|
"/opt/jr-wakeword/run-commercial-pipeline.py",
|
|
"/run/resolved-config.yaml",
|
|
)
|
|
if classification is RunClassification.LOCAL_ONLY:
|
|
return (
|
|
"python",
|
|
"/opt/jr-wakeword/run-local-experiment.py",
|
|
"/run/resolved-config.yaml",
|
|
)
|
|
return (
|
|
"python",
|
|
"/opt/jr-wakeword/run-quarantined-smoke.py",
|
|
"/run/resolved-config.yaml",
|
|
)
|
|
|
|
|
|
def collect_image_provenance(
|
|
*,
|
|
dependency_lock_path: Path,
|
|
frontend_asset_paths: list[Path],
|
|
installed_distributions: dict[str, str],
|
|
image_digest: str,
|
|
python_version: str,
|
|
cuda_version: str,
|
|
os_packages: list[str],
|
|
) -> dict[str, object]:
|
|
"""Return manifest-safe image evidence; callers persist it in the run bundle."""
|
|
|
|
if not image_digest.startswith("sha256:"):
|
|
raise ValueError("image_digest must be sha256-pinned")
|
|
lock_path = dependency_lock_path.resolve(strict=True)
|
|
assets = []
|
|
for path in frontend_asset_paths:
|
|
resolved = path.resolve(strict=True)
|
|
assets.append({"path": resolved.name, "sha256": _sha256_file(resolved)})
|
|
return {
|
|
"schema_version": 1,
|
|
"image_digest": image_digest,
|
|
"python_version": python_version,
|
|
"cuda_version": cuda_version,
|
|
"os_packages": sorted(os_packages),
|
|
"livekit_wakeword": {
|
|
"distribution": "livekit-wakeword",
|
|
"version": LIVEKIT_WAKEWORD_VERSION,
|
|
"source_revision": LIVEKIT_WAKEWORD_REVISION,
|
|
"source_distribution_sha256": LIVEKIT_WAKEWORD_SDIST_SHA256,
|
|
"installed_version": installed_distributions.get("livekit-wakeword"),
|
|
},
|
|
"installed_distributions": dict(sorted(installed_distributions.items())),
|
|
"dependency_lock": {"path": lock_path.name, "sha256": _sha256_file(lock_path)},
|
|
"frontend_assets": assets,
|
|
}
|