Files

2136 lines
84 KiB
Python

"""Host-only, sealed Wakeword Training attempt runner.
This module deliberately has no dependency on FastAPI, the gateway queue, Docker,
LiveKit, a GPU, or a dataset. It is the durable-on-disk operator boundary around
the pinned LiveKit CLI. A deployment supplies a process executor; CI uses
``FakeTrainingExecutor`` to exercise the same lifecycle and handoff contract.
"""
from __future__ import annotations
import asyncio
import contextlib
import fcntl
import hashlib
import json
import os
import re
import shutil
import subprocess
import sys
import tarfile
import tempfile
import time
import uuid
from collections.abc import Mapping
from dataclasses import asdict, dataclass, field
from datetime import UTC, datetime
from enum import StrEnum
from pathlib import Path, PurePosixPath
from typing import Any, Protocol
_ARTIFACT_ID = re.compile(r"^[A-Za-z0-9][A-Za-z0-9_-]*$")
_FAMILY_SLUG = re.compile(r"^[a-z0-9]+(?:-[a-z0-9]+)*$")
_PACKAGE_ROOT = Path(__file__).resolve().parents[1]
_PIPELINE_FIELDS = frozenset(
{
"augmentation",
"batch_n_per_class",
"custom_negative_phrases",
"family_slug",
"label_smoothing",
"learning_rate",
"length_scales",
"max_negative_weight",
"max_speakers",
"model",
"model_name",
"n_background_samples",
"n_background_samples_val",
"n_samples",
"n_samples_val",
"noise_scale_ws",
"noise_scales",
"slerp_weights",
"steps",
"target_fp_per_hour",
"target_phrases",
"threshold",
"tts_backend",
"tts_batch_size",
"version",
"weight_decay",
}
)
_NESTED_PIPELINE_FIELDS = {
"augmentation": frozenset(
{"background_paths", "batch_size", "clip_duration", "rir_paths", "rounds"}
),
"model": frozenset({"model_size", "model_type"}),
}
_BATCH_CLASS_FIELDS = frozenset(
{"positive", "adversarial_negative", "ACAV100M_sample", "background_noise"}
)
class RunKind(StrEnum):
"""Whether an attempt may ever be considered for commercial publication."""
QUARANTINED_SMOKE = "quarantined_smoke"
LOCAL_EXPERIMENT = "local_experiment"
COMMERCIAL = "commercial"
class RunStatus(StrEnum):
CREATED = "created"
PREFLIGHT = "preflight"
ADMITTED = "admitted"
RUNNING = "running"
SUCCEEDED = "succeeded"
COMPLETED = "completed"
PENDING_PUBLICATION = "pending_publication"
FAILED = "failed"
CANCELED = "canceled"
TIMED_OUT = "timed_out"
ABANDONED = "abandoned"
TERMINAL_STATUSES = frozenset(
{
RunStatus.PENDING_PUBLICATION,
RunStatus.COMPLETED,
RunStatus.FAILED,
RunStatus.CANCELED,
RunStatus.TIMED_OUT,
RunStatus.ABANDONED,
}
)
_TRANSITIONS: dict[RunStatus, frozenset[RunStatus]] = {
RunStatus.CREATED: frozenset({RunStatus.PREFLIGHT, RunStatus.FAILED, RunStatus.ABANDONED}),
RunStatus.PREFLIGHT: frozenset({RunStatus.ADMITTED, RunStatus.FAILED, RunStatus.ABANDONED}),
RunStatus.ADMITTED: frozenset(
{
RunStatus.RUNNING,
RunStatus.FAILED,
RunStatus.CANCELED,
RunStatus.TIMED_OUT,
RunStatus.ABANDONED,
}
),
RunStatus.RUNNING: frozenset(
{
RunStatus.SUCCEEDED,
RunStatus.FAILED,
RunStatus.CANCELED,
RunStatus.TIMED_OUT,
RunStatus.ABANDONED,
}
),
RunStatus.SUCCEEDED: frozenset(
{RunStatus.COMPLETED, RunStatus.PENDING_PUBLICATION, RunStatus.FAILED}
),
RunStatus.COMPLETED: frozenset(),
RunStatus.PENDING_PUBLICATION: frozenset(),
RunStatus.FAILED: frozenset(),
RunStatus.CANCELED: frozenset(),
RunStatus.TIMED_OUT: frozenset(),
RunStatus.ABANDONED: frozenset(),
}
class AdmissionRefusedError(RuntimeError):
"""The operator's exclusive-maintenance admission rule rejected an attempt."""
class AdmissionDeferredError(AdmissionRefusedError):
"""A healthy attempt should remain queued until GPU admission is available."""
class InvalidTransitionError(RuntimeError):
"""An attempted lifecycle transition would rewrite immutable evidence."""
class StageOutputHashMismatchError(ValueError):
"""An executor-declared output digest did not match the sealed bytes."""
def _utc_now() -> str:
return datetime.now(UTC).isoformat().replace("+00:00", "Z")
def _canonical_json(value: object) -> str:
return json.dumps(value, ensure_ascii=False, sort_keys=True, separators=(",", ":"))
def _sha256_bytes(value: bytes) -> str:
return hashlib.sha256(value).hexdigest()
def _normalize_phrase(value: str, field_name: str) -> str:
normalized = " ".join(value.split())
if not normalized:
raise ValueError(f"{field_name} must not be empty")
if not any(character.isascii() and character.isalpha() for character in normalized):
raise ValueError(f"{field_name} must contain English letters")
return normalized
def _validate_sha256(value: str, field_name: str) -> str:
candidate = value.removeprefix("sha256:").lower()
if len(candidate) != 64 or any(character not in "0123456789abcdef" for character in candidate):
raise ValueError(f"{field_name} must be a SHA-256 digest")
return candidate
def _validate_image_digest(value: str) -> str:
if not value.startswith("sha256:"):
raise ValueError("runner_image_digest must use the sha256:<hex> form")
_validate_sha256(value, "runner_image_digest")
return value.lower()
def _validate_artifact_id(value: object, field_name: str) -> str:
if not isinstance(value, str) or not _ARTIFACT_ID.fullmatch(value):
raise ValueError(
f"{field_name} must start with a letter or number and contain only "
"letters, numbers, underscores, or hyphens"
)
return value
def _unexpected_fields(value: Mapping[str, object], allowed: frozenset[str], prefix: str) -> None:
unknown = sorted(set(value) - allowed)
if unknown:
raise ValueError(f"unknown pipeline configuration field: {prefix}{unknown[0]}")
def _validate_pipeline_config(value: dict[str, object]) -> None:
"""Reject fields the pinned LiveKit 0.2.1 contract would silently ignore."""
_unexpected_fields(value, _PIPELINE_FIELDS, "")
for required in ("model_name", "target_phrases"):
if required not in value:
raise ValueError(f"pipeline_config.{required} is required")
numeric_lists = ("length_scales", "noise_scale_ws", "noise_scales", "slerp_weights")
positive_integers = (
"n_background_samples",
"n_background_samples_val",
"n_samples",
"n_samples_val",
"steps",
"tts_batch_size",
"version",
)
numeric_fields = (
"label_smoothing",
"learning_rate",
"max_negative_weight",
"target_fp_per_hour",
"weight_decay",
)
target_phrases = value.get("target_phrases")
if (
not isinstance(target_phrases, list)
or not target_phrases
or any(not isinstance(entry, str) or not entry.strip() for entry in target_phrases)
):
raise ValueError("pipeline_config.target_phrases must be a non-empty string array")
negative_phrases = value.get("custom_negative_phrases")
if negative_phrases is not None and (
not isinstance(negative_phrases, list)
or any(not isinstance(entry, str) or not entry.strip() for entry in negative_phrases)
):
raise ValueError("pipeline_config.custom_negative_phrases must be a string array")
for name in numeric_lists:
item = value.get(name)
if item is not None and (
not isinstance(item, list)
or not item
or any(
not isinstance(entry, (int, float)) or isinstance(entry, bool) for entry in item
)
):
raise ValueError(f"pipeline_config.{name} must be a non-empty number array")
for name in positive_integers:
item = value.get(name)
if item is not None and (
not isinstance(item, int) or isinstance(item, bool) or item < 1
):
raise ValueError(f"pipeline_config.{name} must be a positive integer")
for name in numeric_fields:
item = value.get(name)
if item is not None and (
not isinstance(item, (int, float)) or isinstance(item, bool) or item < 0
):
raise ValueError(f"pipeline_config.{name} must be a non-negative number")
for name, allowed in _NESTED_PIPELINE_FIELDS.items():
nested = value.get(name)
if nested is None:
continue
if not isinstance(nested, dict):
raise ValueError(f"pipeline_config.{name} must be an object")
_unexpected_fields(nested, allowed, f"{name}.")
augmentation = value.get("augmentation")
if isinstance(augmentation, dict):
for name in ("batch_size", "rounds"):
item = augmentation.get(name)
if item is not None and (
not isinstance(item, int) or isinstance(item, bool) or item < 1
):
raise ValueError(f"pipeline_config.augmentation.{name} must be positive")
clip_duration = augmentation.get("clip_duration")
if clip_duration is not None and (
not isinstance(clip_duration, (int, float))
or isinstance(clip_duration, bool)
or clip_duration <= 0
):
raise ValueError("pipeline_config.augmentation.clip_duration must be positive")
for name in ("background_paths", "rir_paths"):
item = augmentation.get(name)
if item is not None and (
not isinstance(item, list)
or any(
not isinstance(path, str)
or not path.startswith("/inputs/")
or ".." in PurePosixPath(path).parts
for path in item
)
):
raise ValueError(
f"pipeline_config.augmentation.{name} must contain /inputs paths"
)
model = value.get("model")
if isinstance(model, dict):
if "model_type" in model and model["model_type"] not in {
"dnn",
"rnn",
"conv_attention",
}:
raise ValueError("pipeline_config.model.model_type is invalid")
if "model_size" in model and model["model_size"] not in {
"tiny",
"small",
"medium",
"large",
}:
raise ValueError("pipeline_config.model.model_size is invalid")
threshold = value.get("threshold")
if threshold is not None and not (
threshold == "auto"
or (
isinstance(threshold, (int, float))
and not isinstance(threshold, bool)
and 0 <= threshold <= 1
)
):
raise ValueError("pipeline_config.threshold must be between 0 and 1 or 'auto'")
family_slug = value.get("family_slug")
if family_slug is not None and (
not isinstance(family_slug, str) or not _FAMILY_SLUG.fullmatch(family_slug)
):
raise ValueError("pipeline_config.family_slug must be a lowercase hyphenated slug")
if value.get("tts_backend", "piper_vits") != "piper_vits":
raise ValueError("pipeline_config.tts_backend must be piper_vits in runner version 1")
max_speakers = value.get("max_speakers")
if max_speakers is not None and (
not isinstance(max_speakers, int) or isinstance(max_speakers, bool) or max_speakers < 1
):
raise ValueError("pipeline_config.max_speakers must be null or a positive integer")
batch = value.get("batch_n_per_class")
if batch is not None and (
not isinstance(batch, dict)
or any(
not isinstance(key, str)
or not isinstance(count, int)
or isinstance(count, bool)
or count < 1
for key, count in batch.items()
)
):
raise ValueError("pipeline_config.batch_n_per_class must map names to positive integers")
if isinstance(batch, dict):
_unexpected_fields(batch, _BATCH_CLASS_FIELDS, "batch_n_per_class.")
def _json_compatible(value: object) -> object:
"""Return a JSON-only copy and reject host-path-like configuration values."""
try:
canonical = json.loads(_canonical_json(value))
except (TypeError, ValueError) as exc:
raise ValueError("pipeline_config must be JSON-compatible") from exc
if not isinstance(canonical, dict):
raise ValueError("pipeline_config must be a JSON object")
if "data_dir" in canonical or "output_dir" in canonical:
raise ValueError("pipeline_config cannot choose data_dir or output_dir")
if "model_name" in canonical:
_validate_artifact_id(canonical["model_name"], "pipeline_config.model_name")
_validate_pipeline_config(canonical)
return canonical
def terminal_exit_code(status: RunStatus) -> int:
"""Map a durable terminal outcome to a truthful shell/API-worker exit code."""
return 0 if status in {RunStatus.COMPLETED, RunStatus.PENDING_PUBLICATION} else 1
@dataclass(frozen=True)
class WakewordTrainingRequest:
"""The complete semantic input to one offline training attempt.
``source_lock_path`` is deliberately excluded from the hash and persisted
request. It is an operator-local location, while the locked bytes and their
digest are the reproducible input.
"""
primary_phrase: str
variants: tuple[str, ...]
near_misses: tuple[str, ...]
evaluation_profile: str
source_lock_path: Path
source_lock_sha256: str
intended_release_id: str
runner_image_digest: str
pipeline_config: Mapping[str, Any]
run_kind: RunKind = RunKind.COMMERCIAL
derived_from_run_id: str | None = None
def canonical_payload(self) -> dict[str, object]:
primary_phrase = _normalize_phrase(self.primary_phrase, "primary_phrase")
variants = tuple(_normalize_phrase(variant, "variant") for variant in self.variants)
near_misses = tuple(_normalize_phrase(phrase, "near_miss") for phrase in self.near_misses)
if not self.evaluation_profile.strip():
raise ValueError("evaluation_profile must not be empty")
intended_release_id = _validate_artifact_id(
self.intended_release_id.strip(), "intended_release_id"
)
source_lock_sha256 = _validate_sha256(self.source_lock_sha256, "source_lock_sha256")
if self.derived_from_run_id is not None and (
len(self.derived_from_run_id) != 32
or any(character not in "0123456789abcdef" for character in self.derived_from_run_id)
):
raise ValueError("derived_from_run_id must be an opaque 32-character run ID")
pipeline_config = _json_compatible(self.pipeline_config)
if self.run_kind is RunKind.LOCAL_EXPERIMENT:
expected_phrases = [primary_phrase, *variants]
expected_near_misses = list(near_misses)
if pipeline_config.get("target_phrases") != expected_phrases:
raise ValueError("pipeline_config.target_phrases must match request phrases")
if pipeline_config.get("custom_negative_phrases") != expected_near_misses:
raise ValueError(
"pipeline_config.custom_negative_phrases must match request near_misses"
)
return {
"schema_version": 1,
"primary_phrase": primary_phrase,
# Ordered lists are intentionally not sorted: pronunciation and
# near-miss order are a declared part of a repeatable experiment.
"variants": list(variants),
"near_misses": list(near_misses),
"evaluation_profile": self.evaluation_profile.strip(),
"source_lock_sha256": source_lock_sha256,
"intended_release_id": intended_release_id,
"runner_image_digest": _validate_image_digest(self.runner_image_digest),
"pipeline_config": pipeline_config,
"run_kind": RunKind(self.run_kind).value,
"derived_from_run_id": self.derived_from_run_id,
}
def request_sha256(self) -> str:
return _sha256_bytes(_canonical_json(self.canonical_payload()).encode("utf-8"))
def verified_source_lock(self) -> bytes:
path = self.source_lock_path.resolve(strict=True)
if not path.is_file():
raise ValueError("source_lock_path must identify a regular file")
value = path.read_bytes()
if _sha256_bytes(value) != _validate_sha256(self.source_lock_sha256, "source_lock_sha256"):
raise ValueError("source lock digest does not match source_lock_path")
# It is evidence, not just an arbitrary blob. Preserve unknown fields
# for provenance extensions, but require a JSON object at this seam.
try:
parsed = json.loads(value)
except json.JSONDecodeError as exc:
raise ValueError("source lock must be valid JSON") from exc
if not isinstance(parsed, dict):
raise ValueError("source lock must be a JSON object")
return value
@classmethod
def from_file(cls, path: Path) -> WakewordTrainingRequest:
"""Load the versioned request file used by local and future API entrances."""
value = json.loads(path.read_text(encoding="utf-8"))
if not isinstance(value, dict):
raise ValueError("training request must be a JSON object")
schema_version = value.pop("schema_version", 1)
if schema_version != 1:
raise ValueError("unsupported training request schema_version")
value["source_lock_path"] = Path(value["source_lock_path"])
value["variants"] = tuple(value.get("variants", ()))
value["near_misses"] = tuple(value.get("near_misses", ()))
value["run_kind"] = RunKind(value.get("run_kind", RunKind.COMMERCIAL.value))
return cls(**value)
@dataclass
class RunState:
run_id: str
request_sha256: str
status: RunStatus
stage: str
attempt: int
created_at: str
updated_at: str
heartbeat_at: str
retry_of: str | None = None
terminal_reason: str | None = None
container_id: str | None = None
exit_code: int | None = None
lock_path: str | None = None
wrapper_pid: int = field(default_factory=os.getpid)
@classmethod
def from_dict(cls, value: Mapping[str, Any]) -> RunState:
return cls(
run_id=str(value["run_id"]),
request_sha256=str(value["request_sha256"]),
status=RunStatus(value["status"]),
stage=str(value["stage"]),
attempt=int(value["attempt"]),
created_at=str(value["created_at"]),
updated_at=str(value["updated_at"]),
heartbeat_at=str(value["heartbeat_at"]),
retry_of=value.get("retry_of"),
terminal_reason=value.get("terminal_reason"),
container_id=value.get("container_id"),
exit_code=value.get("exit_code"),
lock_path=value.get("lock_path"),
wrapper_pid=int(value.get("wrapper_pid", 0)),
)
@property
def terminal(self) -> bool:
return self.status in TERMINAL_STATUSES
@dataclass(frozen=True)
class RunResult:
run_dir: Path
state: RunState
reused: bool = False
class MaintenanceWindow(Protocol):
def is_active(self) -> bool: ...
@dataclass(frozen=True)
class StaticMaintenanceWindow:
active: bool
def is_active(self) -> bool:
return self.active
class GpuProcessInspector(Protocol):
def has_unapproved_processes(self) -> bool: ...
@dataclass(frozen=True)
class StaticGpuProcessInspector:
active_unapproved_processes: bool = False
def has_unapproved_processes(self) -> bool:
return self.active_unapproved_processes
class NvidiaSmiProcessInspector:
"""Fail-closed host preflight; it observes processes and never terminates them."""
def __init__(self, gpu_uuid: str) -> None:
self.gpu_uuid = gpu_uuid
def has_unapproved_processes(self) -> bool:
command = (
"nvidia-smi",
f"--id={self.gpu_uuid}",
"--query-compute-apps=pid",
"--format=csv,noheader",
)
try:
completed = subprocess.run(command, capture_output=True, check=False, text=True)
except OSError:
return True
if completed.returncode != 0:
return True
return bool(completed.stdout.strip())
class GpuEvidenceSampler(Protocol):
"""Read one host-side, UUID-addressed GPU observation without mutating it."""
gpu_uuid: str
def sample(self) -> Mapping[str, Any]: ...
class NvidiaSmiGpuEvidenceSampler:
"""Fail-closed ``nvidia-smi`` sampler for a specifically configured GPU UUID."""
_FIELDS = (
"uuid",
"memory.used",
"memory.total",
"utilization.gpu",
"temperature.gpu",
"power.draw",
"driver_version",
)
def __init__(self, gpu_uuid: str) -> None:
if not gpu_uuid.startswith("GPU-"):
raise ValueError("gpu_uuid must be a stable NVIDIA UUID, never a device ordinal")
self.gpu_uuid = gpu_uuid
def sample(self) -> Mapping[str, Any]:
command = (
"nvidia-smi",
f"--id={self.gpu_uuid}",
f"--query-gpu={','.join(self._FIELDS)}",
"--format=csv,noheader,nounits",
)
try:
completed = subprocess.run(command, capture_output=True, check=False, text=True)
except OSError as exc:
raise AdmissionRefusedError(
"nvidia-smi is unavailable; GPU identity is unverified"
) from exc
if completed.returncode != 0:
raise AdmissionRefusedError("configured GPU UUID could not be verified")
values = [item.strip() for item in completed.stdout.strip().split(",")]
if len(values) != len(self._FIELDS) or values[0] != self.gpu_uuid:
raise AdmissionRefusedError("nvidia-smi returned an unexpected GPU UUID")
try:
return {
"gpu_uuid": values[0],
"vram_used_mib": float(values[1]),
"vram_total_mib": float(values[2]),
"utilization_pct": float(values[3]),
"temperature_c": float(values[4]),
"power_w": float(values[5]),
"driver_version": values[6],
}
except ValueError as exc:
raise AdmissionRefusedError("nvidia-smi returned malformed GPU evidence") from exc
@dataclass
class StaticGpuEvidenceSampler:
"""Deterministic CPU-only resource fixture; it never contacts a GPU."""
gpu_uuid: str
samples: tuple[Mapping[str, Any], ...]
_index: int = field(default=0, init=False)
def sample(self) -> Mapping[str, Any]:
if not self.samples:
raise AdmissionRefusedError("GPU evidence fixture has no samples")
value = dict(self.samples[min(self._index, len(self.samples) - 1)])
self._index += 1
value["gpu_uuid"] = self.gpu_uuid
return value
class AdmissionLock(Protocol):
path: Path
def try_acquire(self) -> bool: ...
def release(self) -> None: ...
def held_by_anyone(self) -> bool: ...
class FileAdmissionLock:
"""A non-queueing host ``flock`` around the selected maintenance GPU."""
def __init__(self, path: Path) -> None:
self.path = path
self._file: Any | None = None
def try_acquire(self) -> bool:
self.path.parent.mkdir(parents=True, exist_ok=True)
self._file = self.path.open("a+", encoding="utf-8")
try:
fcntl.flock(self._file.fileno(), fcntl.LOCK_EX | fcntl.LOCK_NB)
except BlockingIOError:
self._file.close()
self._file = None
return False
return True
def release(self) -> None:
if self._file is not None:
fcntl.flock(self._file.fileno(), fcntl.LOCK_UN)
self._file.close()
self._file = None
def held_by_anyone(self) -> bool:
self.path.parent.mkdir(parents=True, exist_ok=True)
with self.path.open("a+", encoding="utf-8") as candidate:
try:
fcntl.flock(candidate.fileno(), fcntl.LOCK_EX | fcntl.LOCK_NB)
except BlockingIOError:
return True
fcntl.flock(candidate.fileno(), fcntl.LOCK_UN)
return False
def gpu_admission_lock_path(gpu_uuid: str) -> Path:
"""Return one host-global, path-safe lock for a physical GPU identity."""
identity = gpu_uuid.strip()
if not identity:
raise ValueError("gpu_uuid must not be empty")
token = hashlib.sha256(identity.encode("utf-8")).hexdigest()[:24]
return Path(tempfile.gettempdir()) / f"jrich-wakeword-training-gpu-{token}.lock"
class ContainerInspector(Protocol):
def is_running(self, container_id: str | None) -> bool | None:
"""False is an affirmative stopped check; None is deliberately unknown."""
class UnknownContainerInspector:
def is_running(self, container_id: str | None) -> bool | None:
return None
class DockerContainerInspector:
"""Inspect only an explicitly named runner container.
``False`` is returned only when Docker proves that the container is stopped
or absent. Any daemon, permission, or output problem remains ``None`` so a
recovery loop cannot turn uncertainty into a second GPU execution.
"""
def is_running(self, container_id: str | None) -> bool | None:
if not container_id or not re.fullmatch(r"wakeword-train-[a-f0-9]{32}", container_id):
return None
try:
result = subprocess.run(
("docker", "inspect", "--format", "{{.State.Running}}", container_id),
check=False,
capture_output=True,
text=True,
timeout=10,
)
except (OSError, subprocess.TimeoutExpired):
return None
if result.returncode == 0:
value = result.stdout.strip().lower()
return True if value == "true" else False if value == "false" else None
if "no such object" in result.stderr.lower():
return False
return None
@dataclass(frozen=True)
class ExecutionContext:
run_id: str
run_dir: Path
resolved_config_path: Path
output_dir: Path
runner_image_digest: str
run_kind: RunKind
container_name: str
@dataclass(frozen=True)
class ExecutionResult:
exit_code: int = 0
container_id: str | None = None
stdout: str = ""
stderr: str = ""
resource_summary: Mapping[str, Any] = field(default_factory=dict)
output_hashes: Mapping[str, str] = field(default_factory=dict)
stage_evidence: tuple[Mapping[str, Any], ...] = ()
class TrainingExecutor(Protocol):
async def execute(self, context: ExecutionContext) -> ExecutionResult: ...
async def cancel(self, context: ExecutionContext, *, force: bool) -> None: ...
class FakeTrainingExecutor:
"""Deterministic complete-contract executor used by CI and local dry runs."""
def __init__(self, *, wait_for_release: bool = False, exit_code: int = 0) -> None:
self.wait_for_release = wait_for_release
self.exit_code = exit_code
self.started = asyncio.Event()
self.release = asyncio.Event()
self.calls = 0
self.cancel_calls: list[bool] = []
async def execute(self, context: ExecutionContext) -> ExecutionResult:
self.calls += 1
self.started.set()
if self.wait_for_release:
await self.release.wait()
if self.exit_code == 0:
context.output_dir.mkdir(parents=True, exist_ok=True)
# These fixture bytes are a runner-contract stand-in only. ONNX
# graph/load verification is a pinned-runner/catalog concern.
(context.output_dir / "classifier.onnx").write_bytes(b"fake-onnx-classifier-v1")
(context.output_dir / "manifest.json").write_text(
_canonical_json({"schema_version": 1, "runner": "fake"}), encoding="utf-8"
)
(context.output_dir / "evaluation.json").write_text(
_canonical_json({"schema_version": 1, "result": "quarantined"}), encoding="utf-8"
)
return ExecutionResult(
exit_code=self.exit_code,
container_id="fake-container",
stdout="fake executor completed\n",
resource_summary={"executor": "fake", "output_files": 3},
)
async def cancel(self, context: ExecutionContext, *, force: bool) -> None:
self.cancel_calls.append(force)
self.release.set()
class DockerTrainingExecutor:
"""Small process adapter for the operator-owned, digest-pinned GPU image.
The gateway never imports this class. Docker is resolved only when an
operator invokes it, not while importing/running the fake contract suite.
"""
def __init__(
self,
*,
image_ref: str,
gpu_uuid: str,
approved_inputs_dir: Path,
stream_output: bool = False,
) -> None:
if "@sha256:" not in image_ref and not image_ref.startswith("sha256:"):
raise ValueError("image_ref must be a digest-pinned reference or immutable image ID")
if not gpu_uuid.strip():
raise ValueError("gpu_uuid must not be empty")
if not approved_inputs_dir.resolve().is_dir():
raise ValueError("approved_inputs_dir must be an existing directory")
self.image_ref = image_ref
self.gpu_uuid = gpu_uuid
self.approved_inputs_dir = approved_inputs_dir.resolve()
self.stream_output = stream_output
self._processes: dict[str, asyncio.subprocess.Process] = {}
async def execute(self, context: ExecutionContext) -> ExecutionResult:
digest = self.image_ref.rsplit("@", maxsplit=1)[-1]
if digest != context.runner_image_digest:
raise ValueError("configured image_ref does not match the sealed runner image digest")
# The image wrapper controls the commercial stage sequence and rejects
# an upstream setup/download path. The executor deliberately exposes no
# caller override that could bypass the sealed image contract.
from scripts.wakeword_training_image import build_stage_command
stage_command = build_stage_command(context.run_kind.value)
container_name = context.container_name
command = (
"docker",
"run",
"--rm",
"--init",
"--name",
container_name,
"--gpus",
f"device={self.gpu_uuid}",
"--network",
"none",
"--read-only",
"--cap-drop",
"ALL",
"--security-opt",
"no-new-privileges",
"--env",
"PYTORCH_CUDA_ALLOC_CONF=expandable_segments:True",
"--mount",
f"type=bind,src={self.approved_inputs_dir},dst=/inputs,readonly",
"--mount",
f"type=bind,src={context.run_dir.resolve()},dst=/run",
"--tmpfs",
"/tmp:rw,noexec,nosuid,size=4g",
self.image_ref,
"/usr/bin/time",
"-v",
"-o",
"/run/pipeline-time-v.txt",
*stage_command,
)
process = await asyncio.create_subprocess_exec(
*command,
stdout=asyncio.subprocess.PIPE,
stderr=asyncio.subprocess.PIPE,
)
self._processes[context.run_id] = process
try:
if self.stream_output:
stdout, stderr = await asyncio.gather(
self._read_stream(process.stdout), self._read_stream(process.stderr)
)
await process.wait()
else:
stdout, stderr = await process.communicate()
except BaseException:
if process.returncode is None:
await self.cancel(context, force=False)
try:
await asyncio.wait_for(process.wait(), timeout=15)
except TimeoutError:
await self.cancel(context, force=True)
await process.wait()
raise
finally:
self._processes.pop(context.run_id, None)
resource_summary = {"executor": "docker", "container_name": container_name}
time_report = context.run_dir / "pipeline-time-v.txt"
if time_report.is_file():
resource_summary.update(self._parse_time_report(time_report))
return ExecutionResult(
exit_code=process.returncode or 0,
container_id=container_name,
stdout=stdout.decode("utf-8", errors="replace"),
stderr=stderr.decode("utf-8", errors="replace"),
resource_summary=resource_summary,
stage_evidence=self._load_image_stage_evidence(context.run_dir),
)
@staticmethod
async def _read_stream(stream: asyncio.StreamReader | None) -> bytes:
if stream is None:
return b""
chunks: list[bytes] = []
while chunk := await stream.read(64 * 1024):
chunks.append(chunk)
print(chunk.decode("utf-8", errors="replace"), end="", file=sys.stderr, flush=True)
return b"".join(chunks)
async def cancel(self, context: ExecutionContext, *, force: bool) -> None:
process = self._processes.get(context.run_id)
if process is None:
return
# We address only the wrapper-owned deterministic name; no unrelated
# GPU process is inspected, stopped, or killed here.
signal = "kill" if force else "stop"
command = ("docker", signal, context.container_name)
with contextlib.suppress(OSError):
cancel_process = await asyncio.create_subprocess_exec(
*command,
stdout=asyncio.subprocess.DEVNULL,
stderr=asyncio.subprocess.DEVNULL,
)
await cancel_process.wait()
@staticmethod
def _parse_time_report(path: Path) -> dict[str, Any]:
values: dict[str, Any] = {}
for line in path.read_text(encoding="utf-8", errors="replace").splitlines():
key, separator, value = line.partition(":")
if not separator:
continue
if key.strip() == "Maximum resident set size (kbytes)":
with contextlib.suppress(ValueError):
values["max_rss_kib"] = int(value.strip())
elif key.strip() == "Elapsed (wall clock) time (h:mm:ss or m:ss)":
values["elapsed_time_reported"] = value.strip()
return values
@staticmethod
def _load_image_stage_evidence(run_dir: Path) -> tuple[Mapping[str, Any], ...]:
path = run_dir / "image-stage-events.ndjson"
if not path.is_file():
return ()
try:
records = tuple(
json.loads(line) for line in path.read_text(encoding="utf-8").splitlines() if line
)
except json.JSONDecodeError:
return ()
return tuple(record for record in records if isinstance(record, dict))
@dataclass(frozen=True)
class RunnerSettings:
work_root: Path
enabled: bool = False
maintenance_window: MaintenanceWindow = field(
default_factory=lambda: StaticMaintenanceWindow(active=False)
)
process_inspector: GpuProcessInspector = field(default_factory=StaticGpuProcessInspector)
gpu_uuid: str | None = None
gpu_sampler: GpuEvidenceSampler | None = None
timeout_s: float = 6 * 60 * 60
cancel_grace_s: float = 30.0
heartbeat_interval_s: float = 5.0
redaction_patterns: tuple[str, ...] = ()
defer_when_gpu_busy: bool = False
admission_lock_path: Path | None = None
class _ResourceRecorder:
"""Persist one-second host samples while an admitted executor owns the GPU lock."""
def __init__(self, run_dir: Path, sampler: GpuEvidenceSampler | None) -> None:
self.run_dir = run_dir
self.path = run_dir / "resource-samples.ndjson"
self.sampler = sampler
self.started = time.monotonic()
self.started_at = _utc_now()
self.ended_at: str | None = None
self.initial_disk_bytes = WakewordTrainingRunner._disk_bytes(run_dir, default=0)
self.samples: list[dict[str, Any]] = []
self._stop = asyncio.Event()
self._task: asyncio.Task[None] | None = None
async def start(self) -> None:
self.path.touch(exist_ok=True)
if self.sampler is None:
return
self._task = asyncio.create_task(self._collect())
# An immediate observation makes short, successful smoke runs evidenced;
# subsequent observations are exactly one second apart.
await asyncio.sleep(0)
if self._task.done():
await self._task
async def stop(self) -> None:
self._stop.set()
try:
if self._task is not None:
await self._task
finally:
self.ended_at = _utc_now()
async def _collect(self) -> None:
while True:
raw = await asyncio.to_thread(self.sampler.sample) # type: ignore[union-attr]
if raw.get("gpu_uuid") != self.sampler.gpu_uuid: # type: ignore[union-attr]
raise AdmissionRefusedError("resource sampler returned an unexpected GPU UUID")
sample = {
"timestamp": _utc_now(),
"sample_interval_s": 1,
"stage": self._active_stage(),
**dict(raw),
}
self.samples.append(sample)
with self.path.open("a", encoding="utf-8") as file:
file.write(_canonical_json(sample) + "\n")
file.flush()
os.fsync(file.fileno())
try:
await asyncio.wait_for(self._stop.wait(), timeout=1.0)
return
except TimeoutError:
continue
def _active_stage(self) -> str:
path = self.run_dir / "active-stage"
try:
value = path.read_text(encoding="utf-8").strip()
except OSError:
return "pipeline"
return (
value if value and all(char.isalnum() or char in "_-" for char in value) else "pipeline"
)
def summary(self) -> dict[str, Any]:
numeric = {
field: [
float(sample[field])
for sample in self.samples
if isinstance(sample.get(field), (int, float))
and not isinstance(sample.get(field), bool)
]
for field in ("vram_used_mib", "utilization_pct", "temperature_c", "power_w")
}
return {
"sample_count": len(self.samples),
"sample_interval_s": 1,
"peak_vram_mib": max(numeric["vram_used_mib"], default=0),
"mean_vram_mib": (
sum(numeric["vram_used_mib"]) / len(numeric["vram_used_mib"])
if numeric["vram_used_mib"]
else 0
),
"peak_utilization_pct": max(numeric["utilization_pct"], default=0),
"peak_temperature_c": max(numeric["temperature_c"], default=0),
"peak_power_w": max(numeric["power_w"], default=0),
"elapsed_s": time.monotonic() - self.started,
}
class _FileLock:
"""Short-lived lock for request-hash lookup/create races across wrappers."""
def __init__(self, path: Path) -> None:
self.path = path
self.file: Any | None = None
def __enter__(self) -> _FileLock:
self.path.parent.mkdir(parents=True, exist_ok=True)
self.file = self.path.open("a+", encoding="utf-8")
fcntl.flock(self.file.fileno(), fcntl.LOCK_EX)
return self
def __exit__(self, exc_type: object, exc: object, traceback: object) -> None:
assert self.file is not None
fcntl.flock(self.file.fileno(), fcntl.LOCK_UN)
self.file.close()
class WakewordTrainingRunner:
"""Own immutable attempts and turn a completed executor result into evidence."""
def __init__(
self,
*,
settings: RunnerSettings,
executor: TrainingExecutor,
admission_lock: AdmissionLock | None = None,
container_inspector: ContainerInspector | None = None,
) -> None:
self.settings = settings
self.executor = executor
self.admission_lock = admission_lock or FileAdmissionLock(
settings.admission_lock_path
or (
gpu_admission_lock_path(settings.gpu_uuid)
if settings.gpu_uuid
else settings.work_root / "locks" / "wakeword-training.lock"
)
)
self.container_inspector = container_inspector or UnknownContainerInspector()
async def run(
self,
request: WakewordTrainingRequest,
*,
retry_of: str | None = None,
cancel_event: asyncio.Event | None = None,
run_id: str | None = None,
) -> RunResult:
if not self.settings.enabled:
raise AdmissionRefusedError("Wakeword Training runner is disabled")
canonical_request = request.canonical_payload()
request_sha256 = request.request_sha256()
source_lock = request.verified_source_lock()
run_dir, state, reused = self._create_or_reuse(
request,
canonical_request,
request_sha256,
source_lock,
retry_of=retry_of,
requested_run_id=run_id,
)
if reused:
return RunResult(run_dir=run_dir, state=state, reused=True)
self._transition(run_dir, state, RunStatus.PREFLIGHT, stage="preflight")
if not self.settings.maintenance_window.is_active():
self._update_admission_evidence(run_dir, maintenance_window_active=False)
self._transition(
run_dir,
state,
RunStatus.FAILED,
stage="admission",
terminal_reason="maintenance_window_inactive",
)
raise AdmissionRefusedError("GPU maintenance window is inactive")
self._update_admission_evidence(run_dir, maintenance_window_active=True)
if self.settings.gpu_uuid is not None:
if self.settings.gpu_sampler is None:
self._transition(
run_dir,
state,
RunStatus.FAILED,
stage="admission",
terminal_reason="gpu_resource_sampler_missing",
)
raise AdmissionRefusedError("configured GPU UUID requires a resource sampler")
if self.settings.gpu_sampler.gpu_uuid != self.settings.gpu_uuid:
self._transition(
run_dir,
state,
RunStatus.FAILED,
stage="admission",
terminal_reason="gpu_uuid_mismatch",
)
raise AdmissionRefusedError(
"resource sampler does not identify the configured GPU UUID"
)
try:
preflight_sample = self.settings.gpu_sampler.sample()
except Exception:
self._update_admission_evidence(run_dir, gpu_identity_check="failed")
self._transition(
run_dir,
state,
RunStatus.FAILED,
stage="admission",
terminal_reason="gpu_preflight_failed",
)
raise
if preflight_sample.get("gpu_uuid") != self.settings.gpu_uuid:
self._transition(
run_dir,
state,
RunStatus.FAILED,
stage="admission",
terminal_reason="gpu_uuid_mismatch",
)
raise AdmissionRefusedError("preflight did not verify the configured GPU UUID")
self._write_json_atomic(
run_dir / "gpu-preflight.json",
{"timestamp": _utc_now(), "sample": dict(preflight_sample)},
)
self._update_admission_evidence(
run_dir,
configured_gpu_uuid=self.settings.gpu_uuid,
observed_gpu_uuid=preflight_sample["gpu_uuid"],
)
if not self.admission_lock.try_acquire():
self._update_admission_evidence(run_dir, gpu_lock_acquired=False)
if self.settings.defer_when_gpu_busy:
# No executor has started. Remove only this unadmitted attempt
# shell so the caller-owned ID can be safely tried after requeue.
shutil.rmtree(run_dir)
raise AdmissionDeferredError("GPU training admission lock is held")
self._transition(
run_dir,
state,
RunStatus.FAILED,
stage="admission",
terminal_reason="gpu_admission_lock_unavailable",
)
raise AdmissionRefusedError("GPU training admission lock is held; work was not queued")
self._update_admission_evidence(
run_dir, gpu_lock_acquired=True, gpu_lock_path=str(self.admission_lock.path)
)
if self.settings.process_inspector.has_unapproved_processes():
self._update_admission_evidence(run_dir, process_check="unapproved_process_present")
self._transition(
run_dir,
state,
RunStatus.FAILED,
stage="admission",
terminal_reason="unapproved_gpu_process_present",
)
self.admission_lock.release()
raise AdmissionRefusedError(
"unapproved GPU process is active; no process was terminated"
)
self._update_admission_evidence(run_dir, process_check="no_unapproved_processes")
recorder: _ResourceRecorder | None = None
try:
self._transition(
run_dir,
state,
RunStatus.ADMITTED,
stage="admitted",
lock_path=str(self.admission_lock.path),
)
self._transition(run_dir, state, RunStatus.RUNNING, stage="pipeline")
context = ExecutionContext(
run_id=state.run_id,
run_dir=run_dir,
resolved_config_path=run_dir / "resolved-config.yaml",
output_dir=run_dir / "output",
runner_image_digest=request.runner_image_digest,
run_kind=request.run_kind,
container_name=f"wakeword-train-{state.run_id}",
)
state.container_id = context.container_name
self._heartbeat(run_dir, state)
recorder = _ResourceRecorder(run_dir, self.settings.gpu_sampler)
try:
await recorder.start()
result, outcome = await self._await_executor(context, run_dir, state, cancel_event)
finally:
await recorder.stop()
self._write_execution_evidence(run_dir, result, outcome, recorder)
if outcome is not None:
self._transition(
run_dir, state, outcome, stage="pipeline", terminal_reason=outcome.value
)
return RunResult(run_dir=run_dir, state=state)
# Recovery uses a separate wrapper process. Never let an older
# in-memory writer resurrect a run recovery has terminalized on disk.
persisted = self._read_state(run_dir)
if persisted.terminal:
return RunResult(run_dir=run_dir, state=persisted)
assert result is not None
state.container_id = result.container_id
state.exit_code = result.exit_code
self._append_logs(run_dir, result, self.settings.redaction_patterns)
if result.exit_code != 0:
self._transition(
run_dir,
state,
RunStatus.FAILED,
stage="pipeline",
terminal_reason="executor_nonzero_exit",
)
return RunResult(run_dir=run_dir, state=state)
self._transition(run_dir, state, RunStatus.SUCCEEDED, stage="validation")
try:
self._verify_declared_output_hashes(run_dir, result.output_hashes)
if request.run_kind is RunKind.LOCAL_EXPERIMENT:
self._seal_local_result(run_dir, state, canonical_request)
else:
self._seal_candidate(run_dir, state, canonical_request)
except StageOutputHashMismatchError:
self._transition(
run_dir,
state,
RunStatus.FAILED,
stage="validation",
terminal_reason="stage_output_hash_mismatch",
)
except (OSError, ValueError) as exc:
self._transition(
run_dir,
state,
RunStatus.FAILED,
stage="validation",
terminal_reason=f"candidate_validation_failed:{type(exc).__name__}",
)
else:
self._transition(
run_dir,
state,
RunStatus.COMPLETED
if request.run_kind is RunKind.LOCAL_EXPERIMENT
else RunStatus.PENDING_PUBLICATION,
stage="sealed",
)
return RunResult(run_dir=run_dir, state=state)
except asyncio.CancelledError:
# A supervisor SIGTERM/cancel still gets a durable terminal record;
# propagate cancellation afterwards so supervisors retain their normal
# shutdown semantics.
if not state.terminal:
if recorder is not None:
self._write_execution_evidence(run_dir, None, RunStatus.CANCELED, recorder)
self._transition(
run_dir,
state,
RunStatus.CANCELED,
stage="pipeline",
terminal_reason="wrapper_cancelled",
)
raise
except Exception as exc:
if not state.terminal:
if recorder is not None:
self._write_execution_evidence(run_dir, None, RunStatus.FAILED, recorder)
self._transition(
run_dir,
state,
RunStatus.FAILED,
stage="pipeline",
terminal_reason=f"executor_error:{type(exc).__name__}",
)
return RunResult(run_dir=run_dir, state=state)
finally:
self.admission_lock.release()
def recover_abandoned(self, *, stale_after_s: float) -> list[RunResult]:
"""Mark only affirmatively stopped, stale attempts abandoned.
An unknown container state or a held lock is intentionally a no-op. A
later phase-two worker may attach a durable lease/reaper, but this POC
must never double-run merely because a timestamp is old.
"""
now = datetime.now(UTC).timestamp()
abandoned: list[RunResult] = []
for run_dir, state in self._all_runs():
if state.terminal:
continue
try:
heartbeat = datetime.fromisoformat(
state.heartbeat_at.replace("Z", "+00:00")
).timestamp()
except ValueError:
continue
if now - heartbeat <= stale_after_s or self.admission_lock.held_by_anyone():
continue
if self.container_inspector.is_running(state.container_id) is not False:
continue
self._transition(
run_dir,
state,
RunStatus.ABANDONED,
stage="recovery",
terminal_reason="stale_heartbeat_container_confirmed_stopped",
)
abandoned.append(RunResult(run_dir=run_dir, state=state))
return abandoned
async def _await_executor(
self,
context: ExecutionContext,
run_dir: Path,
state: RunState,
cancel_event: asyncio.Event | None,
) -> tuple[ExecutionResult | None, RunStatus | None]:
task = asyncio.create_task(self.executor.execute(context))
started = time.monotonic()
try:
while not task.done():
await asyncio.wait({task}, timeout=max(0.001, self.settings.heartbeat_interval_s))
# A recovery wrapper may have terminalized this state while the
# original process was still unwinding. Do not overwrite it.
if self._read_state(run_dir).terminal:
return None, None
self._heartbeat(run_dir, state)
if task.done():
break
cancel_requested = (run_dir / "cancel-requested").is_file()
if (cancel_event is not None and cancel_event.is_set()) or cancel_requested:
await self._stop_executor(task, context)
return None, RunStatus.CANCELED
timed_out = (
self.settings.timeout_s > 0
and time.monotonic() - started >= self.settings.timeout_s
)
if timed_out:
await self._stop_executor(task, context)
return None, RunStatus.TIMED_OUT
return task.result(), None
finally:
if not task.done():
task.cancel()
with contextlib.suppress(asyncio.CancelledError, Exception):
await task
async def _stop_executor(
self, task: asyncio.Task[ExecutionResult], context: ExecutionContext
) -> None:
await self.executor.cancel(context, force=False)
try:
await asyncio.wait_for(
asyncio.shield(task), timeout=max(0.001, self.settings.cancel_grace_s)
)
return
except TimeoutError:
pass
await self.executor.cancel(context, force=True)
task.cancel()
with contextlib.suppress(asyncio.CancelledError, Exception):
await task
def _create_or_reuse(
self,
request: WakewordTrainingRequest,
canonical_request: Mapping[str, object],
request_sha256: str,
source_lock: bytes,
*,
retry_of: str | None,
requested_run_id: str | None,
) -> tuple[Path, RunState, bool]:
if requested_run_id is not None and not re.fullmatch(
r"[a-f0-9]{32}", requested_run_id
):
raise ValueError("run_id must be 32 lowercase hexadecimal characters")
lock = self.settings.work_root / "locks" / "request-identity.lock"
with _FileLock(lock):
if requested_run_id is not None:
existing = next(
(
(run_dir, state)
for run_dir, state in self._all_runs()
if state.run_id == requested_run_id
),
None,
)
if existing is not None:
if existing[1].request_sha256 != request_sha256:
raise ValueError("run_id already identifies a different request")
return existing[0], existing[1], True
matching = [
(run_dir, state)
for run_dir, state in self._all_runs()
if state.request_sha256 == request_sha256
]
if requested_run_id is None:
active = next(
((run_dir, state) for run_dir, state in matching if not state.terminal), None
)
if active is not None:
return active[0], active[1], True
if matching and retry_of is None:
raise ValueError(
"a terminal attempt already exists; retry_of is required for a new attempt"
)
predecessor: RunState | None = None
if retry_of is not None:
predecessor_item = next(
((run_dir, state) for run_dir, state in matching if state.run_id == retry_of),
None,
)
if predecessor_item is None or not predecessor_item[1].terminal:
raise ValueError(
"retry_of must identify a terminal attempt with the same request hash"
)
predecessor = predecessor_item[1]
reuse_source = self._resolve_reuse_source(request, canonical_request)
run_id = requested_run_id or uuid.uuid4().hex
stamp = datetime.now(UTC).strftime("%Y%m%dT%H%M%SZ")
run_dir = self.settings.work_root / "runs" / f"{stamp}-{run_id}"
run_dir.mkdir(parents=True, exist_ok=False)
(run_dir / "output").mkdir()
now = _utc_now()
state = RunState(
run_id=run_id,
request_sha256=request_sha256,
status=RunStatus.CREATED,
stage="created",
attempt=(predecessor.attempt + 1) if predecessor else 1,
created_at=now,
updated_at=now,
heartbeat_at=now,
retry_of=predecessor.run_id if predecessor else None,
)
self._write_state(run_dir, state)
if reuse_source is not None:
try:
self._copy_reusable_corpus(
reuse_source[0], run_dir, reuse_source[1], canonical_request
)
except (OSError, ValueError) as exc:
self._transition(
run_dir,
state,
RunStatus.FAILED,
stage="reuse_preparation",
terminal_reason=f"reuse_preparation_failed:{type(exc).__name__}",
)
raise
self._write_json_atomic(
run_dir / "request.json",
{"request_sha256": request_sha256, "request": dict(canonical_request)},
)
self._write_source_lock(run_dir / "sources.lock.json", source_lock)
resolved_config = dict(canonical_request["pipeline_config"])
resolved_config["data_dir"] = "/inputs"
resolved_config["output_dir"] = "/run/output"
# JSON is valid YAML 1.2, so this is deterministic and does not add
# PyYAML to normal gateway/runner test environments.
self._write_text_atomic(
run_dir / "resolved-config.yaml", _canonical_json(resolved_config) + "\n"
)
self._write_json_atomic(
run_dir / "runner-provenance.json",
{
"schema_version": 1,
"runner": "jr-wakeword-training-host",
"runner_image_digest": canonical_request["runner_image_digest"],
"run_kind": canonical_request["run_kind"],
"wrapper_pid": os.getpid(),
"created_at": _utc_now(),
},
)
(run_dir / "stdout.log").touch()
(run_dir / "stderr.log").touch()
return run_dir, state, False
def _resolve_reuse_source(
self,
request: WakewordTrainingRequest,
canonical_request: Mapping[str, object],
) -> tuple[Path, RunState] | None:
derived_from = request.derived_from_run_id
if derived_from is None:
return None
if request.run_kind is not RunKind.LOCAL_EXPERIMENT:
raise ValueError("only a local experiment may reuse a prior local corpus")
prior = next(
(
(run_dir, state)
for run_dir, state in self._all_runs()
if state.run_id == derived_from
),
None,
)
reusable_terminal = {
RunStatus.COMPLETED,
RunStatus.FAILED,
RunStatus.CANCELED,
RunStatus.TIMED_OUT,
RunStatus.ABANDONED,
}
if prior is None or prior[1].status not in reusable_terminal:
raise ValueError("derived_from_run_id must identify a terminal local experiment")
prior_envelope = json.loads((prior[0] / "request.json").read_text(encoding="utf-8"))
prior_request = prior_envelope.get("request")
if (
not isinstance(prior_request, dict)
or prior_request.get("run_kind") != "local_experiment"
):
raise ValueError("derived run is not a local experiment")
immutable_fields = (
"primary_phrase",
"variants",
"near_misses",
"source_lock_sha256",
"runner_image_digest",
)
if any(
prior_request.get(name) != canonical_request.get(name) for name in immutable_fields
):
raise ValueError("corpus reuse requires matching phrases, sources, and runner image")
prior_config = prior_request.get("pipeline_config")
new_config = canonical_request.get("pipeline_config")
if not isinstance(prior_config, dict) or not isinstance(new_config, dict):
raise ValueError("corpus reuse requires valid pipeline configurations")
if prior_config.get("model_name") != new_config.get("model_name"):
raise ValueError("corpus reuse requires the same model_name")
count_fields = (
"n_samples",
"n_samples_val",
"n_background_samples",
"n_background_samples_val",
)
for name in count_fields:
before = prior_config.get(name, 0)
after = new_config.get(name, 0)
if not isinstance(before, int) or not isinstance(after, int) or after < before:
raise ValueError("corpus reuse cannot reduce generated sample counts")
return prior
def _copy_reusable_corpus(
self,
prior_dir: Path,
run_dir: Path,
prior_state: RunState,
canonical_request: Mapping[str, object],
) -> None:
config = canonical_request["pipeline_config"]
assert isinstance(config, dict)
model_name = str(config["model_name"])
prior_model = prior_dir / "output" / model_name
copied: list[str] = []
copied_files: list[dict[str, object]] = []
sealed_hashes = self._reusable_output_hashes(prior_dir, prior_state)
reusable_names = (
"positive_train",
"positive_test",
"negative_train",
"negative_test",
"background_train",
"background_test",
)
reusable_prefix = prior_model.relative_to(prior_dir).as_posix() + "/"
expected_reusable = {
relative
for relative in sealed_hashes
if relative.startswith(reusable_prefix)
and relative[len(reusable_prefix) :].split("/", maxsplit=1)[0]
in reusable_names
and not re.search(r"_r\d+\.wav$", relative)
}
if not expected_reusable:
raise ValueError("derived local experiment has no reusable generated corpus")
copied_directories: set[str] = set()
for relative in sorted(expected_reusable):
path = prior_dir / relative
expected = sealed_hashes[relative]
if not path.is_file() or path.is_symlink():
raise ValueError(f"reusable corpus integrity file is missing: {relative}")
actual = _sha256_bytes(path.read_bytes())
if actual != expected:
raise ValueError(f"reusable corpus integrity check failed: {relative}")
target = run_dir / relative
target.parent.mkdir(parents=True, exist_ok=True)
shutil.copy2(path, target)
copied_directories.add(target.parent.relative_to(run_dir).as_posix())
copied_files.append(
{"path": relative, "sha256": actual, "size_bytes": path.stat().st_size}
)
copied = sorted(copied_directories)
self._write_json_atomic(
run_dir / "reuse-evidence.json",
{
"schema_version": 1,
"mode": "corpus_reuse_retrain_from_scratch",
"derived_from_run_id": prior_state.run_id,
"derived_from_request_sha256": prior_state.request_sha256,
"copied_directories": copied,
"copied_files": copied_files,
"warm_started": False,
"exact_resume": False,
},
)
@staticmethod
def _reusable_output_hashes(prior_dir: Path, prior_state: RunState) -> dict[str, str]:
evidence = prior_dir / "image-stage-events.ndjson"
if not evidence.is_file():
raise ValueError("terminal local experiment has no completed-stage evidence")
records = [
json.loads(line)
for line in evidence.read_text(encoding="utf-8").splitlines()
if line.strip()
]
generated = next(
(
record
for record in reversed(records)
if record.get("name") == "generate" and record.get("status") == "succeeded"
),
None,
)
if generated is None or not isinstance(generated.get("output_hashes"), dict):
raise ValueError("terminal local experiment has no sealed generation stage")
hashes = {
str(relative): _validate_sha256(str(digest), "stage output hash")
for relative, digest in generated["output_hashes"].items()
if str(relative).startswith("output/")
}
if prior_state.status is RunStatus.COMPLETED:
checksum_path = prior_dir / "checksums.sha256"
if not checksum_path.is_file():
raise ValueError("completed local experiment has no sealed checksums")
final_hashes = {
relative: _validate_sha256(digest, "sealed output hash")
for line in checksum_path.read_text(encoding="utf-8").splitlines()
for digest, separator, relative in [line.partition(" ")]
if separator and relative.startswith("output/")
}
if any(final_hashes.get(relative) != digest for relative, digest in hashes.items()):
raise ValueError("completed generation evidence does not match final checksums")
return hashes
def _all_runs(self) -> list[tuple[Path, RunState]]:
runs_root = self.settings.work_root / "runs"
if not runs_root.exists():
return []
result: list[tuple[Path, RunState]] = []
for run_dir in sorted(path for path in runs_root.iterdir() if path.is_dir()):
state_path = run_dir / "state.json"
if not state_path.is_file():
continue
try:
loaded = json.loads(state_path.read_text(encoding="utf-8"))
result.append((run_dir, RunState.from_dict(loaded)))
except (OSError, ValueError, KeyError, json.JSONDecodeError):
# A malformed half-created directory is unsafe to reuse; it is
# left for an operator rather than being treated as a duplicate.
continue
return result
@staticmethod
def _read_state(run_dir: Path) -> RunState:
value = json.loads((run_dir / "state.json").read_text(encoding="utf-8"))
return RunState.from_dict(value)
def _transition(
self,
run_dir: Path,
state: RunState,
next_status: RunStatus,
*,
stage: str,
terminal_reason: str | None = None,
lock_path: str | None = None,
) -> None:
if next_status not in _TRANSITIONS[state.status]:
raise InvalidTransitionError(f"cannot transition {state.status} to {next_status}")
state.status = next_status
state.stage = stage
state.updated_at = _utc_now()
state.heartbeat_at = state.updated_at
if terminal_reason is not None:
state.terminal_reason = terminal_reason
if lock_path is not None:
state.lock_path = lock_path
self._write_state(run_dir, state)
def _heartbeat(self, run_dir: Path, state: RunState) -> None:
state.heartbeat_at = _utc_now()
state.updated_at = state.heartbeat_at
self._write_state(run_dir, state)
def _write_execution_evidence(
self,
run_dir: Path,
result: ExecutionResult | None,
outcome: RunStatus | None,
recorder: _ResourceRecorder,
) -> None:
"""Persist a sealed, reviewable record even when the pipeline is terminally bad."""
resource_summary = dict(result.resource_summary) if result is not None else {}
resource_summary.update(recorder.summary())
disk_after = self._disk_bytes(run_dir)
resource_summary.update(
{
"disk_before_bytes": recorder.initial_disk_bytes,
"disk_after_bytes": disk_after,
"disk_growth_bytes": max(0, disk_after - recorder.initial_disk_bytes),
}
)
self._write_json_atomic(run_dir / "resource-summary.json", resource_summary)
input_hashes = self._existing_hashes(
run_dir, ("request.json", "resolved-config.yaml", "sources.lock.json")
)
preflight_outputs = self._existing_hashes(run_dir, ("gpu-preflight.json", "state.json"))
if not preflight_outputs:
preflight_outputs = {"state.json": _sha256_bytes((run_dir / "state.json").read_bytes())}
pipeline_outputs = self._tree_hashes(run_dir / "output")
if not pipeline_outputs:
pipeline_outputs = {"output": _sha256_bytes(b"")}
pipeline_status = (
outcome.value
if outcome is not None
else "failed"
if result is not None and result.exit_code != 0
else "succeeded"
)
pipeline_stage = {
"name": "pipeline",
"started_at": recorder.started_at,
"ended_at": recorder.ended_at or _utc_now(),
"status": pipeline_status,
"exit_status": result.exit_code if result is not None else None,
"input_hashes": input_hashes,
"output_hashes": pipeline_outputs,
"elapsed_s": resource_summary["elapsed_s"],
"max_rss_kib": resource_summary.get("max_rss_kib", 0),
"disk_growth_bytes": resource_summary["disk_growth_bytes"],
}
executed_stages = (
list(result.stage_evidence)
if result is not None and result.stage_evidence
else [pipeline_stage]
)
stages = {
"schema_version": 1,
"stages": [
{
"name": "preflight",
"started_at": recorder.started_at,
"ended_at": recorder.started_at,
"status": "succeeded",
"exit_status": 0,
"input_hashes": input_hashes,
"output_hashes": preflight_outputs,
"elapsed_s": 0,
"max_rss_kib": 0,
"disk_growth_bytes": 0,
},
*executed_stages,
],
}
self._write_json_atomic(run_dir / "stage-evidence.json", stages)
@staticmethod
def _update_admission_evidence(run_dir: Path, **values: object) -> None:
path = run_dir / "admission-evidence.json"
try:
current = json.loads(path.read_text(encoding="utf-8"))
except (OSError, json.JSONDecodeError):
current = {"schema_version": 1, "updated_at": _utc_now()}
if not isinstance(current, dict):
raise ValueError("admission evidence is malformed")
current.update(values)
current["updated_at"] = _utc_now()
WakewordTrainingRunner._write_json_atomic(path, current)
@staticmethod
def _disk_bytes(path: Path, *, default: int = 0) -> int:
if not path.exists():
return default
return sum(candidate.stat().st_size for candidate in path.rglob("*") if candidate.is_file())
@staticmethod
def _existing_hashes(run_dir: Path, names: tuple[str, ...]) -> dict[str, str]:
return {
name: _sha256_bytes((run_dir / name).read_bytes())
for name in names
if (run_dir / name).is_file()
}
@staticmethod
def _tree_hashes(root: Path) -> dict[str, str]:
if not root.is_dir():
return {}
return {
path.relative_to(root.parent).as_posix(): _sha256_bytes(path.read_bytes())
for path in sorted(candidate for candidate in root.rglob("*") if candidate.is_file())
}
@staticmethod
def _verify_declared_output_hashes(run_dir: Path, declared: Mapping[str, str]) -> None:
for relative, expected in declared.items():
if not isinstance(relative, str) or not isinstance(expected, str):
raise StageOutputHashMismatchError("executor output hash is malformed")
try:
normalized_expected = _validate_sha256(expected, "output hash")
except ValueError as exc:
raise StageOutputHashMismatchError("executor output hash is malformed") from exc
candidate = (run_dir / relative).resolve()
if (
candidate == run_dir.resolve()
or run_dir.resolve() not in candidate.parents
or not candidate.is_file()
or _sha256_bytes(candidate.read_bytes()) != normalized_expected
):
raise StageOutputHashMismatchError(
"executor output hash does not match sealed output"
)
def _seal_candidate(
self, run_dir: Path, state: RunState, canonical_request: Mapping[str, object]
) -> None:
output_dir = run_dir / "output"
classifiers = sorted(output_dir.glob("*.onnx"))
if len(classifiers) != 1:
raise ValueError("expected exactly one ONNX classifier in output")
for relative in ("manifest.json", "evaluation.json"):
path = output_dir / relative
if not path.is_file():
raise ValueError(f"required evidence file is missing: output/{relative}")
self._write_checksums(run_dir)
classifier = classifiers[0]
candidate = {
"schema_version": 1,
"run_id": state.run_id,
"request_sha256": state.request_sha256,
"intended_release_id": canonical_request["intended_release_id"],
"run_kind": canonical_request["run_kind"],
# A smoke is permanently quarantined. A commercial attempt is
# still review-required; this runner never grants publication.
"commercial_promotion_eligible": False,
"publication_status": "quarantined"
if canonical_request["run_kind"] == RunKind.QUARANTINED_SMOKE.value
else "review_required",
"classifier": {
"path": classifier.relative_to(run_dir).as_posix(),
"sha256": _sha256_bytes(classifier.read_bytes()),
"size_bytes": classifier.stat().st_size,
},
"manifest_path": "output/manifest.json",
"evaluation_path": "output/evaluation.json",
"checksums_path": "checksums.sha256",
}
self._write_json_atomic(run_dir / "publish-candidate.json", candidate)
def _seal_local_result(
self, run_dir: Path, state: RunState, canonical_request: Mapping[str, object]
) -> None:
"""Seal a user-owned result without creating a catalog publication handoff."""
output_dir = run_dir / "output"
classifiers = sorted(output_dir.glob("*.onnx"))
checkpoints = sorted(output_dir.rglob("*.pt"))
if len(classifiers) != 1:
raise ValueError("expected exactly one ONNX classifier in output")
if len(checkpoints) != 1:
raise ValueError("expected exactly one private PT checkpoint in output")
for relative in ("manifest.json", "evaluation.json"):
if not (output_dir / relative).is_file():
raise ValueError(f"required evidence file is missing: output/{relative}")
def artifact(path: Path) -> dict[str, object]:
return {
"path": path.relative_to(run_dir).as_posix(),
"sha256": _sha256_bytes(path.read_bytes()),
"size_bytes": path.stat().st_size,
}
retraining_paths = sorted(
path
for path in output_dir.rglob("*")
if path.is_file() and path.suffix.lower() in {".pt", ".npy"}
)
public_bundle = self._write_public_bundle(
run_dir,
str(canonical_request["intended_release_id"]),
classifiers[0],
output_dir / "manifest.json",
output_dir / "evaluation.json",
)
private_bundle = self._write_private_bundle(
run_dir, str(canonical_request["intended_release_id"])
)
result = {
"schema_version": 1,
"run_id": state.run_id,
"request_sha256": state.request_sha256,
"intended_release_id": canonical_request["intended_release_id"],
"run_kind": RunKind.LOCAL_EXPERIMENT.value,
"publication_status": "local_only",
"classifier": artifact(classifiers[0]),
"checkpoint": artifact(checkpoints[0]),
"manifest_path": "output/manifest.json",
"evaluation_path": "output/evaluation.json",
"checksums_path": "checksums.sha256",
"retraining_artifacts": [artifact(path) for path in retraining_paths],
"public_bundle": artifact(public_bundle),
"private_retraining_bundle": artifact(private_bundle),
"resume_capability": "sealed_corpus_reuse_only",
}
self._write_json_atomic(run_dir / "training-result.json", result)
self._write_checksums(run_dir)
@staticmethod
def _write_private_bundle(run_dir: Path, release_id: str) -> Path:
"""Persist the full corpus/retraining state needed after host loss."""
bundle = run_dir / "exports" / f"{release_id}-private-retraining.tar.gz"
fixed = (
"state.json",
"request.json",
"resolved-config.yaml",
"sources.lock.json",
"runner-provenance.json",
"stage-evidence.json",
"image-stage-events.ndjson",
"image-provenance.json",
"resource-summary.json",
)
members = [run_dir / name for name in fixed if (run_dir / name).is_file()]
members.extend(
path
for path in sorted((run_dir / "output").rglob("*"))
if path.is_file() and not path.is_symlink()
)
corpus_integrity = {
path.relative_to(run_dir).as_posix(): _sha256_bytes(path.read_bytes())
for path in members
if path.relative_to(run_dir).as_posix().startswith("output/")
}
with tarfile.open(bundle, "w:gz") as archive:
for path in members:
archive.add(path, arcname=path.relative_to(run_dir).as_posix(), recursive=False)
integrity = (_canonical_json(corpus_integrity) + "\n").encode("utf-8")
info = tarfile.TarInfo("corpus-integrity.json")
info.size = len(integrity)
info.mtime = int(time.time())
with tempfile.TemporaryFile() as file:
file.write(integrity)
file.seek(0)
archive.addfile(info, file)
return bundle
@staticmethod
def _write_public_bundle(
run_dir: Path,
release_id: str,
classifier: Path,
manifest: Path,
evaluation: Path,
) -> Path:
exports = run_dir / "exports"
exports.mkdir(exist_ok=True)
bundle = exports / f"{release_id}-public.tar.gz"
members = {
"classifier.onnx": classifier,
"manifest.json": manifest,
"evaluation.json": evaluation,
"LICENSE": _PACKAGE_ROOT / "LICENSE",
"THIRD_PARTY_NOTICES.md": (
_PACKAGE_ROOT / "deploy" / "wakeword-training" / "THIRD_PARTY_NOTICES.md"
),
}
missing = [name for name, path in members.items() if not path.is_file()]
if missing:
raise ValueError(f"public bundle source is missing: {', '.join(missing)}")
checksums = "".join(
f"{_sha256_bytes(path.read_bytes())} {name}\n" for name, path in members.items()
).encode("utf-8")
with tarfile.open(bundle, "w:gz") as archive:
for name, path in members.items():
archive.add(path, arcname=f"{release_id}/{name}")
info = tarfile.TarInfo(f"{release_id}/SHA256SUMS")
info.size = len(checksums)
info.mtime = int(time.time())
with tempfile.TemporaryFile() as file:
file.write(checksums)
file.seek(0)
archive.addfile(info, file)
return bundle
def _write_checksums(self, run_dir: Path) -> None:
paths: list[Path] = []
for path in run_dir.rglob("*"):
excluded = {"state.json", "checksums.sha256", "publish-candidate.json"}
if not path.is_file() or path.name in excluded:
continue
if path.parent == run_dir and path.name.startswith("."):
continue
paths.append(path)
lines = [
f"{_sha256_bytes(path.read_bytes())} {path.relative_to(run_dir).as_posix()}"
for path in sorted(paths)
]
self._write_text_atomic(run_dir / "checksums.sha256", "\n".join(lines) + "\n")
@staticmethod
def _write_source_lock(path: Path, value: bytes) -> None:
WakewordTrainingRunner._write_bytes_atomic(path, value)
path.chmod(0o444)
@staticmethod
def _write_state(run_dir: Path, state: RunState) -> None:
WakewordTrainingRunner._write_json_atomic(run_dir / "state.json", asdict(state))
@staticmethod
def _write_json_atomic(path: Path, value: Mapping[str, Any]) -> None:
WakewordTrainingRunner._write_text_atomic(path, _canonical_json(value) + "\n")
@staticmethod
def _write_text_atomic(path: Path, value: str) -> None:
WakewordTrainingRunner._write_bytes_atomic(path, value.encode("utf-8"))
@staticmethod
def _write_bytes_atomic(path: Path, value: bytes) -> None:
path.parent.mkdir(parents=True, exist_ok=True)
descriptor, temporary_name = tempfile.mkstemp(prefix=f".{path.name}.", dir=path.parent)
temporary = Path(temporary_name)
try:
with os.fdopen(descriptor, "wb") as file:
file.write(value)
file.flush()
os.fsync(file.fileno())
os.replace(temporary, path)
finally:
temporary.unlink(missing_ok=True)
@staticmethod
def _append_logs(
run_dir: Path, result: ExecutionResult, redaction_patterns: tuple[str, ...]
) -> None:
# This wrapper never puts request values, credentials, or host paths into
# either log. Deployment-specific literal secret values can also be
# supplied as patterns and are redacted before persistence.
def redact(value: str) -> str:
for pattern in redaction_patterns:
if pattern:
value = value.replace(pattern, "[REDACTED]")
return value
(run_dir / "stdout.log").write_text(redact(result.stdout), encoding="utf-8")
(run_dir / "stderr.log").write_text(redact(result.stderr), encoding="utf-8")