471 lines
21 KiB
Python
471 lines
21 KiB
Python
"""Preprocess sets and datasets backed by NPZ files."""
|
|
|
|
from __future__ import annotations
|
|
|
|
from contextlib import suppress
|
|
from datetime import datetime, timezone
|
|
import json
|
|
import logging
|
|
import os
|
|
from pathlib import Path
|
|
import shutil
|
|
from typing import Any
|
|
from uuid import uuid4
|
|
from zipfile import BadZipFile
|
|
|
|
import numpy as np
|
|
|
|
from python_app.file_io import make_directory, rename_no_replace, sync_directory, temporary_file
|
|
from python_app.models.dataset_model import ComboKey, ResultCollection, SweepCollection, TraceData
|
|
from python_app.file_io import sanitize_path_component
|
|
from python_app.orchestration.shm.decoder import RAW_MAGIC
|
|
from python_app.storage.npz.serialize import serialize_trace_collection, validate_trace_collection
|
|
from python_app.storage.npz.snapshot_numpy import (
|
|
save_result_history_numpy,
|
|
save_trace_history_numpy,
|
|
save_processing_settings,
|
|
atomic_snapshot_directory,
|
|
)
|
|
from python_app.storage.npz.vna_history_json import build_vna_history_payload
|
|
|
|
logger = logging.getLogger(__name__)
|
|
_SET_METADATA_KEY = "metadata_json"
|
|
|
|
|
|
def _compose_snapshot_stem(name: str, name_prefix: str) -> str:
|
|
"""Compose a sanitized snapshot stem from an optional prefix and the operator's name.
|
|
|
|
The name falls back to a UTC timestamp when blank, so a save is always uniquely named.
|
|
"""
|
|
base = name.strip() or datetime.now(timezone.utc).strftime("snapshot_%Y%m%d_%H%M%S")
|
|
return sanitize_path_component(f"{name_prefix}_{base}" if name_prefix else base)
|
|
|
|
|
|
class SnapshotStreamWriter:
|
|
"""Publish complete collections incrementally, retaining them if a later write fails."""
|
|
|
|
def __init__(self, snapshot_dir: Path) -> None:
|
|
self._dir = snapshot_dir
|
|
self._counts = {"raw": 0, "preprocessed": 0, "results": 0}
|
|
|
|
@property
|
|
def directory(self) -> Path:
|
|
return self._dir
|
|
|
|
@property
|
|
def result_count(self) -> int:
|
|
return self._counts["results"]
|
|
|
|
def append(
|
|
self,
|
|
raw: list[SweepCollection],
|
|
preprocessed: list[SweepCollection],
|
|
results: list[ResultCollection],
|
|
*, processing_settings: dict[int, dict] | None = None,
|
|
) -> None:
|
|
"""Append a chunk; after failure, only the unwritten collections may be retried."""
|
|
save_processing_settings(self._dir, processing_settings)
|
|
for stage, history in (("raw", raw), ("preprocessed", preprocessed)):
|
|
stage_dir = self._dir / stage
|
|
stage_dir.mkdir(exist_ok=True)
|
|
for collection in history:
|
|
save_trace_history_numpy(stage_dir, [collection], index_offset=self._counts[stage])
|
|
self._counts[stage] += 1
|
|
stage_dir = self._dir / "results"
|
|
stage_dir.mkdir(exist_ok=True)
|
|
for result in results:
|
|
save_result_history_numpy(stage_dir, [result], index_offset=self.result_count)
|
|
self._counts["results"] += 1
|
|
|
|
|
|
class NpzStore:
|
|
"""Persist preprocess sets and runtime snapshots using NumPy files."""
|
|
|
|
def __init__(self, root_dir: Path) -> None:
|
|
"""Create store rooted at `root_dir`."""
|
|
self._root_dir = root_dir
|
|
make_directory(self._root_dir, durable=True)
|
|
logger.debug("NpzStore rooted at %s", self._root_dir)
|
|
|
|
def save_set(
|
|
self, kind: str, radar_key: str, set_name: str, collection: SweepCollection,
|
|
*, metadata: dict[str, Any] | None = None,
|
|
) -> None:
|
|
"""Persist a named set with optional JSON metadata embedded in the NPZ."""
|
|
self._validate_set_name(set_name)
|
|
validate_trace_collection(collection, require_finite=True)
|
|
if metadata is not None and not isinstance(metadata, dict):
|
|
raise ValueError("Set user metadata must be a JSON object")
|
|
set_dir = self._set_dir(kind, radar_key)
|
|
make_directory(set_dir, durable=True)
|
|
|
|
npz_path = set_dir / f"{set_name}.npz"
|
|
meta_path = set_dir / f"{set_name}.json"
|
|
|
|
payload: dict[str, Any] = {}
|
|
combo_records: list[dict[str, str | int]] = []
|
|
|
|
for trace in collection.traces:
|
|
suffix = f"i{trace.combo.input}_o{trace.combo.output}"
|
|
freq_key = f"freq_{suffix}"
|
|
s21_key = f"s21_{suffix}"
|
|
payload[freq_key] = np.asarray(trace.frequency_hz, dtype=np.float32)
|
|
payload[s21_key] = np.asarray(trace.s21, dtype=np.complex64)
|
|
combo_records.append(
|
|
{
|
|
"input": trace.combo.input,
|
|
"output": trace.combo.output,
|
|
"capture_start_ns": int(trace.capture_start_ns),
|
|
"capture_end_ns": int(trace.capture_end_ns),
|
|
"freq_key": freq_key,
|
|
"s21_key": s21_key,
|
|
}
|
|
)
|
|
|
|
meta = {
|
|
"collection_id": int(collection.collection_id),
|
|
"monotonic_ns": int(collection.monotonic_ns),
|
|
"capture_start_ns": int(collection.capture_start_ns),
|
|
"capture_end_ns": int(collection.capture_end_ns),
|
|
"combos": combo_records,
|
|
"metadata": {} if metadata is None else metadata,
|
|
}
|
|
|
|
# Metadata in the NPZ makes a single atomic replacement authoritative. The
|
|
# JSON remains a readable sidecar and legacy NPZ/JSON pairs remain loadable.
|
|
metadata_json = json.dumps(meta, indent=2, allow_nan=False)
|
|
payload[_SET_METADATA_KEY] = np.asarray(metadata_json)
|
|
npz_tmp = meta_tmp = backup_path = None
|
|
archive_replaced = False
|
|
try:
|
|
npz_tmp = temporary_file(npz_path)
|
|
meta_tmp = temporary_file(meta_path)
|
|
backup_path = temporary_file(npz_path) if npz_path.exists() else None
|
|
if backup_path is not None:
|
|
backup_path.unlink()
|
|
os.link(npz_path, backup_path)
|
|
with npz_tmp.open("wb") as npz_file:
|
|
np.savez(npz_file, **payload)
|
|
npz_file.flush()
|
|
os.fsync(npz_file.fileno())
|
|
with meta_tmp.open("w", encoding="utf-8") as meta_file:
|
|
meta_file.write(metadata_json)
|
|
meta_file.flush()
|
|
os.fsync(meta_file.fileno())
|
|
npz_tmp.replace(npz_path)
|
|
archive_replaced = True
|
|
meta_tmp.replace(meta_path)
|
|
except BaseException:
|
|
logger.exception("Failed to save set %s/%s/%s; rolling back temp files", kind, radar_key, set_name)
|
|
if archive_replaced:
|
|
if backup_path is None:
|
|
npz_path.unlink(missing_ok=True)
|
|
else:
|
|
try:
|
|
backup_path.replace(npz_path)
|
|
except OSError:
|
|
logger.exception("Set rollback failed; previous archive retained at %s", backup_path)
|
|
backup_path = None # retain the recovery copy rather than deleting it below
|
|
raise
|
|
sync_directory(set_dir)
|
|
raise
|
|
finally:
|
|
for tmp_path in (npz_tmp, meta_tmp, backup_path):
|
|
if tmp_path is None:
|
|
continue
|
|
with suppress(OSError):
|
|
tmp_path.unlink(missing_ok=True)
|
|
sync_directory(set_dir)
|
|
logger.info(
|
|
"Saved set %s/%s/%s (%d traces) to %s", kind, radar_key, set_name, len(combo_records), npz_path
|
|
)
|
|
|
|
def save_sets(self, radar_key: str, set_name: str, collections: dict[str, SweepCollection]) -> None:
|
|
"""Save related sets together, restoring their previous files if any save fails."""
|
|
self._validate_set_name(set_name)
|
|
for collection in collections.values():
|
|
validate_trace_collection(collection, require_finite=True)
|
|
backup_dir = self._root_dir / f".sets-{uuid4().hex}"
|
|
backup_dir.mkdir()
|
|
backups: dict[Path, Path | None] = {}
|
|
retain_backups = False
|
|
try:
|
|
for kind in collections:
|
|
set_dir = self._set_dir(kind, radar_key)
|
|
for extension in ("npz", "json"):
|
|
destination = set_dir / f"{set_name}.{extension}"
|
|
backup = backup_dir / f"{kind}.{extension}" if destination.exists() else None
|
|
if backup is not None:
|
|
os.link(destination, backup)
|
|
backups[destination] = backup
|
|
try:
|
|
for kind, collection in collections.items():
|
|
self.save_set(kind, radar_key, set_name, collection)
|
|
except BaseException as save_error:
|
|
retain_backups = True # preserve recovery files if rollback itself is interrupted
|
|
rollback_failed = False
|
|
for destination, backup in backups.items():
|
|
try:
|
|
if backup is None:
|
|
destination.unlink(missing_ok=True)
|
|
else:
|
|
backup.replace(destination)
|
|
except OSError:
|
|
rollback_failed = True
|
|
logger.exception("Failed to restore previous set file %s", destination)
|
|
for directory in {path.parent for path in backups if path.parent.exists()}:
|
|
try:
|
|
sync_directory(directory)
|
|
except OSError:
|
|
rollback_failed = True
|
|
logger.exception("Failed to sync restored set directory %s", directory)
|
|
retain_backups = rollback_failed
|
|
if retain_backups:
|
|
raise RuntimeError(
|
|
f"Saving set group '{set_name}' failed and rollback was incomplete; "
|
|
f"previous files are retained at {backup_dir}"
|
|
) from save_error
|
|
raise
|
|
finally:
|
|
if not retain_backups:
|
|
with suppress(OSError):
|
|
shutil.rmtree(backup_dir)
|
|
|
|
def load_set(self, kind: str, radar_key: str, set_name: str) -> SweepCollection:
|
|
"""Load named preprocess set from NPZ representation."""
|
|
return self.load_set_with_metadata(kind, radar_key, set_name)[0]
|
|
|
|
def load_set_with_metadata(
|
|
self, kind: str, radar_key: str, set_name: str,
|
|
) -> tuple[SweepCollection, dict[str, Any]]:
|
|
"""Load a named set and its user metadata; older sets have empty metadata."""
|
|
self._validate_set_name(set_name)
|
|
set_dir = self._set_dir(kind, radar_key)
|
|
npz_path = set_dir / f"{set_name}.npz"
|
|
meta_path = set_dir / f"{set_name}.json"
|
|
|
|
if not npz_path.exists():
|
|
logger.error("Missing set files for %s/%s/%s", kind, radar_key, set_name)
|
|
raise FileNotFoundError(f"Missing set files for {kind}/{radar_key}/{set_name}")
|
|
|
|
set_label = f"{kind}/{radar_key}/{set_name}"
|
|
try:
|
|
with np.load(npz_path, allow_pickle=False) as arrays:
|
|
meta = json.loads(
|
|
str(arrays[_SET_METADATA_KEY].item())
|
|
if _SET_METADATA_KEY in arrays
|
|
else meta_path.read_text(encoding="utf-8")
|
|
)
|
|
if not isinstance(meta, dict):
|
|
raise ValueError("Set metadata must be a JSON object")
|
|
metadata = meta.get("metadata", {})
|
|
if not isinstance(metadata, dict):
|
|
raise ValueError("Set user metadata must be a JSON object")
|
|
json.dumps(metadata, allow_nan=False)
|
|
traces: list[TraceData] = []
|
|
for combo in meta["combos"]:
|
|
traces.append(
|
|
TraceData(
|
|
combo=ComboKey(input=combo["input"], output=combo["output"]),
|
|
frequency_hz=np.asarray(arrays[combo["freq_key"]], dtype=np.float32),
|
|
s21=np.asarray(arrays[combo["s21_key"]], dtype=np.complex64),
|
|
capture_start_ns=combo.get("capture_start_ns", 0),
|
|
capture_end_ns=combo.get("capture_end_ns", 0),
|
|
)
|
|
)
|
|
|
|
collection = SweepCollection(
|
|
collection_id=meta["collection_id"],
|
|
monotonic_ns=meta["monotonic_ns"],
|
|
traces=traces,
|
|
capture_start_ns=meta.get("capture_start_ns", 0),
|
|
capture_end_ns=meta.get("capture_end_ns", 0),
|
|
)
|
|
validate_trace_collection(collection, require_finite=True)
|
|
logger.debug("Loaded set %s (%d traces)", set_label, len(traces))
|
|
return collection, metadata
|
|
except (OSError, BadZipFile, EOFError, KeyError, TypeError, ValueError) as exc:
|
|
logger.exception("Corrupted preprocess set %s", set_label)
|
|
raise RuntimeError(f"Corrupted preprocess set {set_label}: {exc}") from exc
|
|
|
|
def list_sets(self, kind: str, radar_key: str) -> list[str]:
|
|
"""List available set names for `(kind, radar_key)`."""
|
|
set_dir = self._set_dir(kind, radar_key)
|
|
if not set_dir.exists():
|
|
return []
|
|
return sorted(path.stem for path in set_dir.glob("*.npz"))
|
|
|
|
def has_combo_coverage(self, kind: str, radar_key: str, set_name: str, combos: list[ComboKey]) -> bool:
|
|
"""Validate that named set covers all required switch combinations."""
|
|
collection = self.load_set(kind, radar_key, set_name)
|
|
existing = {(trace.combo.input, trace.combo.output) for trace in collection.traces}
|
|
required = {(combo.input, combo.output) for combo in combos}
|
|
return required.issubset(existing)
|
|
|
|
def export_set_bundle(self, kind: str, radar_key: str, set_name: str, output_path: Path) -> Path:
|
|
"""Export named set as binary collection bundle for C++ preprocessing stage."""
|
|
collection = self.load_set(kind, radar_key, set_name)
|
|
return self._export_collection_bundle(collection, output_path)
|
|
|
|
def _export_collection_bundle(self, collection: SweepCollection, output_path: Path) -> Path:
|
|
"""Publish a complete binary trace bundle using one atomic replacement."""
|
|
validate_trace_collection(collection, require_finite=True)
|
|
output_path.parent.mkdir(parents=True, exist_ok=True)
|
|
temporary = temporary_file(output_path)
|
|
try:
|
|
temporary.write_bytes(serialize_trace_collection(collection, RAW_MAGIC))
|
|
temporary.replace(output_path)
|
|
finally:
|
|
temporary.unlink(missing_ok=True)
|
|
logger.info("Exported %d traces to bundle %s", len(collection.traces), output_path)
|
|
return output_path
|
|
|
|
def snapshot_directory(
|
|
self, output_root_dir: Path, snapshot_name: str, *, name_prefix: str = ""
|
|
) -> Path:
|
|
"""Return the directory a snapshot save would create for these inputs.
|
|
|
|
Lets callers check ``.exists()`` *before* acquiring data (e.g. the disk
|
|
recorder arms a run only once the destination is free), so a name clash is
|
|
reported up front instead of after the measurements are collected. Note a
|
|
blank ``snapshot_name`` resolves to a fresh timestamp each call, so this is
|
|
meaningful only for explicit names — exactly the case that can collide.
|
|
"""
|
|
return output_root_dir / _compose_snapshot_stem(snapshot_name, name_prefix)
|
|
|
|
def create_snapshot_stream(
|
|
self, output_root_dir: Path, snapshot_name: str, *, name_prefix: str = ""
|
|
) -> SnapshotStreamWriter:
|
|
"""Create an empty snapshot directory and return a chunked stream writer for it.
|
|
|
|
Raises ``FileExistsError`` if the directory already exists (same guard as
|
|
:meth:`save_runtime_snapshot_numpy`), so a name clash is reported at arm time.
|
|
"""
|
|
snapshot_dir = self.snapshot_directory(output_root_dir, snapshot_name, name_prefix=name_prefix)
|
|
output_root_dir.mkdir(parents=True, exist_ok=True)
|
|
if snapshot_dir.exists():
|
|
raise FileExistsError(f"Snapshot directory already exists: {snapshot_dir}")
|
|
snapshot_dir.mkdir(parents=True, exist_ok=False)
|
|
logger.info("Opened streaming snapshot directory %s", snapshot_dir)
|
|
return SnapshotStreamWriter(snapshot_dir)
|
|
|
|
def save_runtime_snapshot_numpy(
|
|
self,
|
|
output_root_dir: Path,
|
|
snapshot_name: str,
|
|
raw_history: list[SweepCollection],
|
|
preprocessed_history: list[SweepCollection],
|
|
result_history: list[ResultCollection],
|
|
*,
|
|
name_prefix: str = "",
|
|
processing_settings: dict[int, dict] | None = None,
|
|
) -> Path:
|
|
"""Save exactly the supplied collections and their processing revisions."""
|
|
snapshot_dir = self.snapshot_directory(output_root_dir, snapshot_name, name_prefix=name_prefix)
|
|
with atomic_snapshot_directory(snapshot_dir) as staging:
|
|
save_processing_settings(staging, processing_settings)
|
|
save_trace_history_numpy(staging / "raw", raw_history)
|
|
save_trace_history_numpy(staging / "preprocessed", preprocessed_history)
|
|
save_result_history_numpy(staging / "results", result_history)
|
|
logger.info("Saved NumPy runtime snapshot to %s", snapshot_dir)
|
|
return snapshot_dir
|
|
|
|
def save_runtime_vna_history_json_batch(
|
|
self,
|
|
output_root_dir: Path,
|
|
output_name: str,
|
|
raw_history: list[SweepCollection],
|
|
preprocessed_history: list[SweepCollection],
|
|
result_history: list[ResultCollection],
|
|
*,
|
|
primary_stage: str = "preprocessed",
|
|
name_prefix: str = "",
|
|
) -> list[Path]:
|
|
"""Save one runtime VNA-history JSON per available combo."""
|
|
output_stem = _compose_snapshot_stem(output_name, name_prefix)
|
|
output_root_dir.mkdir(parents=True, exist_ok=True)
|
|
output_dir = output_root_dir / output_stem
|
|
|
|
combos = sorted(
|
|
{
|
|
(int(trace.combo.input), int(trace.combo.output))
|
|
for collection in [*raw_history, *preprocessed_history]
|
|
for trace in collection.traces
|
|
}
|
|
)
|
|
if not combos:
|
|
logger.warning("No raw/preprocessed combos found in runtime history for VNA JSON batch export")
|
|
raise ValueError("No matching raw/preprocessed traces were found in runtime history for any combo.")
|
|
|
|
output_paths: list[Path] = []
|
|
payloads: list[dict[str, Any]] = []
|
|
for input_index, output_index in combos:
|
|
output_path = output_dir / (
|
|
f"{output_stem}_i{input_index}_o{output_index}_s21_vna_bscan_history.json"
|
|
)
|
|
if output_path.exists():
|
|
raise FileExistsError(f"Output JSON file already exists: {output_path}")
|
|
payloads.append(
|
|
build_vna_history_payload(
|
|
raw_history,
|
|
preprocessed_history,
|
|
result_history,
|
|
input_index=input_index,
|
|
output_index=output_index,
|
|
primary_stage=primary_stage,
|
|
)
|
|
)
|
|
output_paths.append(output_path)
|
|
|
|
output_existed = output_dir.exists()
|
|
output_dir.mkdir(exist_ok=True)
|
|
published: list[Path] = []
|
|
try:
|
|
# Publish complete files without overwriting a concurrent export.
|
|
for output_path, payload in zip(output_paths, payloads, strict=True):
|
|
temporary = temporary_file(output_path)
|
|
try:
|
|
with temporary.open("w", encoding="utf-8") as handle:
|
|
json.dump(payload, handle, ensure_ascii=False, indent=2, allow_nan=False)
|
|
handle.flush()
|
|
os.fsync(handle.fileno())
|
|
rename_no_replace(temporary, output_path)
|
|
published.append(output_path)
|
|
finally:
|
|
temporary.unlink(missing_ok=True)
|
|
sync_directory(output_dir)
|
|
except BaseException:
|
|
for path in published:
|
|
path.unlink()
|
|
if not output_existed:
|
|
with suppress(OSError):
|
|
output_dir.rmdir()
|
|
raise
|
|
|
|
logger.info(
|
|
"Saved %d S21 VNA history JSON file(s) to %s", len(output_paths), output_dir
|
|
)
|
|
return output_paths
|
|
|
|
def _set_dir(self, kind: str, radar_key: str) -> Path:
|
|
"""Return directory for set kind and radar key."""
|
|
self._validate_set_name(kind)
|
|
self._validate_set_name(radar_key)
|
|
return self._root_dir / kind / radar_key
|
|
|
|
@staticmethod
|
|
def _validate_set_name(value: str) -> None:
|
|
"""Keep caller-provided set components inside their storage directory."""
|
|
if not value or value in {".", ".."} or Path(value).name != value or "\\" in value:
|
|
raise ValueError(f"Invalid set path component: {value!r}")
|
|
|
|
def preview_png_dir(self, kind: str, set_name: str, radar_key: str) -> Path:
|
|
"""Return directory for preview PNGs of one set/radar variant.
|
|
|
|
Lives under ``preview_png/`` inside the store so saved graphs sit next to
|
|
the data they describe, grouped by set name then radar variant. The set
|
|
name is operator-supplied, so it is sanitized; ``radar_key`` is already
|
|
filesystem-safe by construction.
|
|
"""
|
|
return self._root_dir / "preview_png" / kind / sanitize_path_component(set_name) / radar_key
|