Files
jrich-wakeword-trainer/scripts/wakeword_trainer.py
T

604 lines
22 KiB
Python

"""Standard-library host launcher for the downloadable JRich Wakeword Trainer."""
from __future__ import annotations
import argparse
import hashlib
import json
import os
import re
import shutil
import subprocess
import sys
from pathlib import Path
from typing import Any
RUNNER_VERSION = "0.2.1-jrich-3"
DEFAULT_IMAGE = (
"git.jimandkrista.com/jr-public/jrich-wakeword-training"
"@sha256:cef5b1e769a6160bbde69f67df76f0b865f02c27794a3da096398c18cc9a3c2f"
)
DEFAULT_CACHE = Path.home() / ".cache" / "jrich-wakeword" / "livekit-0.2.1"
DEFAULT_WORK_ROOT = Path.cwd() / "jrich-wakeword-runs"
MIN_FREE_BYTES = 40 * 1024**3
PRESET_PATH = (
Path(__file__).resolve().parents[1]
/ "deploy"
/ "wakeword-training"
/ "presets"
/ "manual-poc-v1.json"
)
_SLUG = re.compile(r"^[a-z0-9]+(?:-[a-z0-9]+)*$")
_MODEL_NAME = re.compile(r"^[A-Za-z0-9][A-Za-z0-9_-]*$")
def _canonical(value: object) -> str:
return json.dumps(value, ensure_ascii=False, sort_keys=True, separators=(",", ":"))
def _sha256(value: bytes) -> str:
return hashlib.sha256(value).hexdigest()
def _deep_merge(base: dict[str, Any], override: dict[str, Any]) -> dict[str, Any]:
result = dict(base)
for key, value in override.items():
if key not in base:
raise ValueError(f"unknown pipeline override: {key}")
if isinstance(base[key], dict) and isinstance(value, dict):
result[key] = _deep_merge(base[key], value)
else:
result[key] = value
return result
def default_pipeline_config(
*,
model_name: str,
family_slug: str,
version: int,
phrases: tuple[str, ...],
near_misses: tuple[str, ...],
) -> dict[str, Any]:
"""Return the measured, versioned manual-POC profile before user overrides."""
preset = json.loads(PRESET_PATH.read_text(encoding="utf-8"))
if not isinstance(preset, dict) or preset.get("preset_version") != "manual-poc-v1":
raise ValueError("the bundled manual POC preset is malformed")
config = preset.get("pipeline_config")
if not isinstance(config, dict):
raise ValueError("the bundled manual POC preset has no pipeline_config")
return {
"model_name": model_name,
"family_slug": family_slug,
"version": version,
"target_phrases": list(phrases),
"custom_negative_phrases": list(near_misses),
**config,
}
def build_setup_command(
*, image_ref: str, gpu_uuid: str, cache_dir: Path, config_path: Path
) -> tuple[str, ...]:
"""Build the explicit networked cache-population command; training stays offline."""
cache_dir.mkdir(parents=True, exist_ok=True)
return (
"docker",
"run",
"--rm",
"--init",
"--gpus",
f"device={gpu_uuid}",
"--read-only",
"--cap-drop",
"ALL",
"--security-opt",
"no-new-privileges",
"--mount",
f"type=bind,src={cache_dir.resolve()},dst=/data",
"--mount",
f"type=bind,src={config_path.resolve()},dst=/config.json,readonly",
"--tmpfs",
"/tmp:rw,noexec,nosuid,size=4g",
image_ref,
"livekit-wakeword",
"setup",
"--config",
"/config.json",
)
def _run(
command: tuple[str, ...] | list[str], *, capture: bool = False
) -> subprocess.CompletedProcess:
return subprocess.run(command, check=False, capture_output=capture, text=True)
def _inspect_image(image_ref: str) -> dict[str, Any] | None:
completed = _run(
["docker", "image", "inspect", "--format", "{{json .}}", image_ref],
capture=True,
)
if completed.returncode != 0:
return None
try:
value = json.loads(completed.stdout)
except json.JSONDecodeError as exc:
raise ValueError("Docker returned malformed image metadata") from exc
if not isinstance(value, dict):
raise ValueError("Docker returned malformed image metadata")
return value
def _image_identity(image_ref: str) -> str:
if image_ref.startswith("sha256:") and len(image_ref) == 71:
return image_ref.lower()
inspected = _inspect_image(image_ref)
if inspected is None:
print(f"Downloading pinned training runtime: {image_ref}", file=sys.stderr)
pulled = _run(["docker", "pull", image_ref])
if pulled.returncode != 0:
raise ValueError(f"training image download failed: {image_ref}")
inspected = _inspect_image(image_ref)
if inspected is None:
raise ValueError(
f"training image is unavailable after download: {image_ref}"
)
if "@sha256:" in image_ref:
expected = image_ref.lower()
repo_digests = inspected.get("RepoDigests", [])
if not isinstance(repo_digests, list) or expected not in {
str(value).lower() for value in repo_digests
}:
raise ValueError("Docker did not retain the requested registry manifest digest")
digest = str(inspected.get("Id", "")).lower()
if not digest.startswith("sha256:") or len(digest) != 71:
raise ValueError("Docker returned an invalid immutable image ID")
return digest
def _default_gpu_uuid() -> str:
completed = _run(
["nvidia-smi", "--query-gpu=uuid", "--format=csv,noheader"], capture=True
)
values = [line.strip() for line in completed.stdout.splitlines() if line.strip()]
if completed.returncode != 0 or not values:
raise ValueError("no NVIDIA GPU UUID was found; verify the driver and nvidia-smi")
if len(values) > 1:
raise ValueError("multiple NVIDIA GPUs found; choose one with --gpu-uuid")
return values[0]
def _prompt(label: str, default: str | None = None) -> str:
suffix = f" [{default}]" if default else ""
value = input(f"{label}{suffix}: ").strip()
return value or (default or "")
def _slug(value: str) -> str:
candidate = re.sub(r"[^a-z0-9]+", "-", value.lower()).strip("-")
if not _SLUG.fullmatch(candidate):
raise ValueError("family slug must use lowercase letters, numbers, and single hyphens")
return candidate
def _interactive_values(args: argparse.Namespace) -> None:
if args.phrase:
return
print("JRich Wakeword Trainer — the defaults below match the measured RTX 3090 POC.")
primary = _prompt("Primary spoken phrase", "Hey Jay Rich")
variants = _prompt("Additional spoken phrases, comma-separated", "Yo Jay Rich")
near_misses = _prompt("Near-miss phrases that must not activate, comma-separated", "Jay Rich")
label = _prompt("Model name", "jrich")
args.phrase = [primary, *[item.strip() for item in variants.split(",") if item.strip()]]
args.near_miss = [item.strip() for item in near_misses.split(",") if item.strip()]
args.model_name = args.model_name or label
args.family_slug = args.family_slug or _slug(label)
def _load_override(path: Path | None) -> dict[str, Any]:
if path is None:
return {}
value = json.loads(path.read_text(encoding="utf-8"))
if not isinstance(value, dict):
raise ValueError("pipeline override must be a JSON object")
return value
def _file_sha256(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 write_local_source_lock(work_root: Path, cache_dir: Path) -> tuple[Path, str]:
"""Write a path-independent, content-addressed inventory of the local cache."""
request_root = work_root / "requests"
request_root.mkdir(parents=True, exist_ok=True)
cache_dir.mkdir(parents=True, exist_ok=True)
files: list[dict[str, object]] = []
for path in sorted(cache_dir.rglob("*")):
if path.is_symlink():
raise ValueError(f"training cache must not contain symlinks: {path}")
if not path.is_file():
continue
files.append(
{
"relative_path": path.relative_to(cache_dir).as_posix(),
"size_bytes": path.stat().st_size,
"sha256": _file_sha256(path),
}
)
value = {
"schema_version": 1,
"run_kind": "local_experiment",
"sources": [
{
"source_id": "livekit-setup-cache-v0.2.1",
"kind": "user-local-cache",
"publication_approved": False,
"files": files,
}
],
}
encoded = (_canonical(value) + "\n").encode("utf-8")
digest = _sha256(encoded)
path = request_root / f"sources-{digest}.lock.json"
if not path.exists():
path.write_bytes(encoded)
path.chmod(0o444)
return path, digest
def resolve_plan(args: argparse.Namespace) -> dict[str, Any]:
phrases = tuple(" ".join(value.split()) for value in (args.phrase or ()) if value.strip())
if not phrases:
raise ValueError("at least one --phrase is required in non-interactive mode")
near_misses = tuple(
" ".join(value.split()) for value in (args.near_miss or ()) if value.strip()
)
model_name = (args.model_name or _slug(phrases[0]).replace("-", "_")).strip()
if not _MODEL_NAME.fullmatch(model_name):
raise ValueError(
"model name must start with a letter or number and contain only letters, "
"numbers, underscores, or hyphens"
)
if args.version < 1:
raise ValueError("version must be a positive integer")
family_slug = _slug(args.family_slug or model_name)
image_digest = _image_identity(args.image_ref)
gpu_uuid = args.gpu_uuid or _default_gpu_uuid()
cache_dir = args.cache_dir.resolve()
work_root = args.work_root.resolve()
pipeline = default_pipeline_config(
model_name=model_name,
family_slug=family_slug,
version=args.version,
phrases=phrases,
near_misses=near_misses,
)
pipeline = _deep_merge(pipeline, _load_override(args.pipeline_config))
identity = {
"model_name": model_name,
"family_slug": family_slug,
"version": args.version,
"target_phrases": list(phrases),
"custom_negative_phrases": list(near_misses),
}
divergent = [name for name, value in identity.items() if pipeline.get(name) != value]
if divergent:
raise ValueError(
"pipeline override cannot change request identity fields: " + ", ".join(divergent)
)
source_lock, source_digest = write_local_source_lock(work_root, cache_dir)
request = {
"schema_version": 1,
"primary_phrase": phrases[0],
"variants": list(phrases[1:]),
"near_misses": list(near_misses),
"evaluation_profile": "local-balanced-v1",
"source_lock_path": str(source_lock),
"source_lock_sha256": source_digest,
"intended_release_id": f"{family_slug}-v{args.version}",
"runner_image_digest": image_digest,
"pipeline_config": pipeline,
"run_kind": "local_experiment",
"derived_from_run_id": args.reuse_from,
}
# Keep the local wizard and future API worker on the exact same fail-closed
# validation boundary. Import shape differs only between the extracted
# archive (sibling module) and repository tests (scripts package).
try:
from wakeword_training import RunKind, WakewordTrainingRequest
except ModuleNotFoundError:
from scripts.wakeword_training import RunKind, WakewordTrainingRequest
WakewordTrainingRequest(
primary_phrase=request["primary_phrase"],
variants=tuple(request["variants"]),
near_misses=tuple(request["near_misses"]),
evaluation_profile=request["evaluation_profile"],
source_lock_path=Path(request["source_lock_path"]),
source_lock_sha256=request["source_lock_sha256"],
intended_release_id=request["intended_release_id"],
runner_image_digest=request["runner_image_digest"],
pipeline_config=request["pipeline_config"],
run_kind=RunKind.LOCAL_EXPERIMENT,
derived_from_run_id=request["derived_from_run_id"],
).canonical_payload()
return {
"schema_version": 1,
"runner_version": RUNNER_VERSION,
"request": request,
"runtime": {
"image_ref": image_digest,
"image_source": args.image_ref,
"gpu_uuid": gpu_uuid,
"cache_dir": str(cache_dir),
"work_root": str(work_root),
"hardware_profile": "shared-conservative"
if args.allow_active_gpu_processes
else "exclusive",
},
}
def _preflight(plan: dict[str, Any]) -> None:
runtime = plan["runtime"]
for executable in ("docker", "nvidia-smi"):
if shutil.which(executable) is None:
raise ValueError(f"required host command is missing: {executable}")
docker = _run(["docker", "info"], capture=True)
if docker.returncode != 0:
raise ValueError("Docker daemon is unavailable to this user")
gpu = _run(
[
"nvidia-smi",
f"--id={runtime['gpu_uuid']}",
"--query-gpu=uuid,name,memory.total,memory.free,driver_version",
"--format=csv,noheader,nounits",
],
capture=True,
)
if gpu.returncode != 0:
raise ValueError("the selected GPU UUID is unavailable")
container_gpu = _run(
[
"docker",
"run",
"--rm",
"--gpus",
f"device={runtime['gpu_uuid']}",
"--network",
"none",
"--read-only",
"--cap-drop",
"ALL",
"--security-opt",
"no-new-privileges",
runtime["image_ref"],
"python",
"-c",
"import torch; assert torch.cuda.is_available(); print(torch.cuda.get_device_name())",
],
capture=True,
)
if container_gpu.returncode != 0:
raise ValueError("NVIDIA Container Toolkit did not expose the selected GPU to the image")
for key in ("cache_dir", "work_root"):
path = Path(runtime[key])
path.mkdir(parents=True, exist_ok=True)
free = shutil.disk_usage(path).free
if free < MIN_FREE_BYTES:
raise ValueError(f"{path} has less than the recommended 40 GiB free")
print(f"GPU: {gpu.stdout.strip()}", file=sys.stderr)
print(f"Container CUDA: {container_gpu.stdout.strip()}", file=sys.stderr)
print(f"Host CPUs: {os.cpu_count() or 'unknown'}", file=sys.stderr)
print("Preflight passed; training will run with network disabled.", file=sys.stderr)
def _write_request(plan: dict[str, Any]) -> Path:
work_root = Path(plan["runtime"]["work_root"])
encoded = (_canonical(plan["request"]) + "\n").encode("utf-8")
path = work_root / "requests" / f"request-{_sha256(encoded)}.json"
path.parent.mkdir(parents=True, exist_ok=True)
path.write_bytes(encoded)
return path
def _setup(plan: dict[str, Any]) -> None:
runtime = plan["runtime"]
config = dict(plan["request"]["pipeline_config"])
config["data_dir"] = "/data"
config["output_dir"] = "/tmp/setup-output"
setup_path = Path(runtime["work_root"]) / "requests" / "setup-config.json"
setup_path.write_text(_canonical(config) + "\n", encoding="utf-8")
command = build_setup_command(
image_ref=runtime["image_ref"],
gpu_uuid=runtime["gpu_uuid"],
cache_dir=Path(runtime["cache_dir"]),
config_path=setup_path,
)
completed = _run(command)
if completed.returncode != 0:
raise ValueError("LiveKit asset setup failed")
def _train(plan: dict[str, Any], request_path: Path, args: argparse.Namespace) -> int:
runtime = plan["runtime"]
command = [
sys.executable,
str(Path(__file__).with_name("run-wakeword-training")),
"--request",
str(request_path),
"--work-root",
runtime["work_root"],
"--enable",
"--maintenance-window-active",
"--image-ref",
runtime["image_ref"],
"--gpu-uuid",
runtime["gpu_uuid"],
"--approved-inputs-dir",
runtime["cache_dir"],
"--timeout-s",
str(args.timeout_s),
"--stream-output",
]
if args.allow_active_gpu_processes:
command.append("--allow-active-gpu-processes")
if args.retry_of:
command.extend(("--retry-of", args.retry_of))
return _run(command).returncode
def _find_run(work_root: Path, run_id: str) -> tuple[Path, dict[str, Any]]:
if len(run_id) != 32 or any(character not in "0123456789abcdef" for character in run_id):
raise ValueError("run ID must be 32 lowercase hexadecimal characters")
matches: list[tuple[Path, dict[str, Any]]] = []
for run_dir in (work_root / "runs").glob(f"*-{run_id}"):
state_path = run_dir / "state.json"
if not state_path.is_file():
continue
state = json.loads(state_path.read_text(encoding="utf-8"))
if isinstance(state, dict) and state.get("run_id") == run_id:
matches.append((run_dir, state))
if len(matches) != 1:
raise ValueError("run ID was not found under the selected work root")
return matches[0]
def run_local_operation(args: argparse.Namespace) -> int | None:
operation = next(
(
(name, run_id)
for name, run_id in (
("status", args.status),
("logs", args.logs),
("cancel", args.cancel),
)
if run_id is not None
),
None,
)
if operation is None:
return None
name, run_id = operation
run_dir, state = _find_run(args.work_root.resolve(), run_id)
if name == "status":
result_path = run_dir / "training-result.json"
result = (
json.loads(result_path.read_text(encoding="utf-8"))
if result_path.is_file()
else None
)
print(json.dumps({"run_dir": str(run_dir), "state": state, "result": result}, indent=2))
return 0
if name == "logs":
for label in ("stdout.log", "stderr.log"):
print(f"== {label} ==")
path = run_dir / label
if path.is_file():
print(path.read_text(encoding="utf-8", errors="replace"), end="")
return 0
if state.get("status") in {
"completed",
"pending_publication",
"failed",
"canceled",
"timed_out",
"abandoned",
}:
raise ValueError("terminal run cannot be canceled")
(run_dir / "cancel-requested").write_text("requested\n", encoding="utf-8")
print(f"Cancellation requested for {run_id}; the runner will stop its own container.")
return 0
def _parser() -> argparse.ArgumentParser:
parser = argparse.ArgumentParser(description="Train a LiveKit Wakeword model in pinned Docker")
parser.add_argument("--non-interactive", action="store_true")
parser.add_argument("--plan", action="store_true", help="print resolved JSON without running")
operation = parser.add_mutually_exclusive_group()
operation.add_argument("--status", metavar="RUN_ID", help="show one local attempt")
operation.add_argument("--logs", metavar="RUN_ID", help="show one local attempt's logs")
operation.add_argument("--cancel", metavar="RUN_ID", help="request safe cancellation")
parser.add_argument("--phrase", action="append", help="spoken phrase; repeat for alternatives")
parser.add_argument("--near-miss", action="append", help="phrase that must not activate")
parser.add_argument("--model-name")
parser.add_argument("--family-slug")
parser.add_argument("--version", type=int, default=1)
parser.add_argument("--pipeline-config", type=Path, help="JSON deep overrides")
parser.add_argument("--image-ref", default=os.getenv("JR_WAKEWORD_IMAGE", DEFAULT_IMAGE))
parser.add_argument("--gpu-uuid")
parser.add_argument("--cache-dir", type=Path, default=DEFAULT_CACHE)
parser.add_argument("--work-root", type=Path, default=DEFAULT_WORK_ROOT)
parser.add_argument("--skip-setup", action="store_true")
parser.add_argument(
"--setup-only",
action="store_true",
help="populate the pinned input cache, print its source-lock digest, and exit",
)
parser.add_argument(
"--reuse-from",
help="terminal local run ID whose sealed generated corpus seeds a new attempt",
)
parser.add_argument(
"--retry-of", help="matching terminal run ID to retry as a new immutable attempt"
)
parser.add_argument("--allow-active-gpu-processes", action="store_true")
parser.add_argument("--timeout-s", type=float, default=6 * 60 * 60)
return parser
def main(argv: list[str] | None = None) -> int:
args = _parser().parse_args(argv)
try:
if getattr(args, "setup_only", False) and args.skip_setup:
raise ValueError("--setup-only and --skip-setup cannot be combined")
operation_result = run_local_operation(args)
if operation_result is not None:
return operation_result
if not args.non_interactive:
_interactive_values(args)
plan = resolve_plan(args)
if args.plan:
print(json.dumps(plan, indent=2, sort_keys=True))
return 0
print(json.dumps(plan, indent=2, sort_keys=True))
if not args.non_interactive and _prompt("Continue with this configuration", "yes") != "yes":
return 0
_preflight(plan)
if not args.skip_setup:
_setup(plan)
refreshed = resolve_plan(args)
if refreshed["request"]["source_lock_sha256"] != plan["request"][
"source_lock_sha256"
]:
plan = refreshed
print("Final content-locked request after setup:", file=sys.stderr)
print(json.dumps(plan, indent=2, sort_keys=True))
if getattr(args, "setup_only", False):
print(
"WAKEWORD_TRAINING_SOURCE_LOCK_SHA256="
+ plan["request"]["source_lock_sha256"]
)
return 0
request_path = _write_request(plan)
return _train(plan, request_path, args)
except (OSError, ValueError, json.JSONDecodeError) as exc:
print(f"error: {exc}", file=sys.stderr)
return 2
if __name__ == "__main__":
raise SystemExit(main())