Files
radar_system/python_app/storage/npz/store.py
T
2026-10-05 16:46:28 +03:00

622 lines
28 KiB
Python

"""Concrete :class:`StoreApi` implementation 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
import tempfile
from typing import Any
from zipfile import BadZipFile
import numpy as np
from python_app.models.dataset_model import ComboKey, ResultCollection, SweepCollection, TraceData
from python_app.storage.npz.paths import radar_key_from_config, sanitize_path_component
from python_app.storage.npz.serialize import PREPROC_MAGIC, RAW_MAGIC, serialize_trace_collection, validate_trace_collection
from python_app.storage.npz.snapshot_numpy import (
save_result_history_binary,
save_result_history_numpy,
save_trace_history_binary,
save_trace_history_numpy,
select_aligned_histories,
atomic_snapshot_directory,
)
from python_app.storage.npz.vna_history_json import build_vna_history_payload
from python_app.storage.store_api import StoreApi
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:
"""Append received raw/preprocessed/result collections to a snapshot dir in chunks.
A streaming alternative to :meth:`NpzStore.save_runtime_snapshot_numpy` for long
recordings: the caller flushes small chunks as measurements arrive and frees them,
so memory stays flat instead of buffering every measurement. The on-disk layout is
identical (``raw``/``preprocessed``/``results`` subtrees), and per-collection
directory indices continue across chunks via a running offset per stage.
"""
__slots__ = ("_dir", "_raw_written", "_preprocessed_written", "_result_written")
def __init__(self, snapshot_dir: Path) -> None:
self._dir = snapshot_dir
self._raw_written = 0
self._preprocessed_written = 0
self._result_written = 0
@property
def directory(self) -> Path:
return self._dir
@property
def result_count(self) -> int:
"""Number of result collections written to disk so far."""
return self._result_written
def append(
self,
raw: list[SweepCollection],
preprocessed: list[SweepCollection],
results: list[ResultCollection],
) -> None:
"""Write a chunk without requiring equal stage lengths, continuing each index."""
# Stage a whole chunk before publishing complete collection directories. A
# failed write leaves earlier chunks intact and the next retry uses the same indices.
chunk_dir = Path(tempfile.mkdtemp(prefix=".chunk-", dir=self._dir))
published: list[Path] = []
try:
save_trace_history_numpy(chunk_dir / "raw", raw, index_offset=self._raw_written)
save_trace_history_numpy(
chunk_dir / "preprocessed", preprocessed, index_offset=self._preprocessed_written
)
save_result_history_numpy(chunk_dir / "results", results, index_offset=self._result_written)
for stage in ("raw", "preprocessed", "results"):
stage_dir = self._dir / stage
stage_dir.mkdir(exist_ok=True)
for source in sorted((chunk_dir / stage).iterdir()):
destination = stage_dir / source.name
if destination.exists():
raise FileExistsError(f"Collection directory already exists: {destination}")
source.rename(destination)
published.append(destination)
except BaseException:
for destination in published:
shutil.rmtree(destination)
raise
finally:
shutil.rmtree(chunk_dir)
self._raw_written += len(raw)
self._preprocessed_written += len(preprocessed)
self._result_written += len(results)
class NpzStore(StoreApi):
"""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
self._root_dir.mkdir(parents=True, exist_ok=True)
logger.debug("NpzStore rooted at %s", self._root_dir)
@staticmethod
def _vna_json_output_dir(output_root_dir: Path, output_stem: str) -> Path:
"""Return per-export directory for VNA history JSON files."""
return output_root_dir / output_stem
def save_set(self, kind: str, radar_key: str, set_name: str, collection: SweepCollection) -> None:
"""Persist named preprocess set as NPZ and metadata JSON."""
self._validate_set_name(set_name)
validate_trace_collection(collection, require_finite=True)
set_dir = self._set_dir(kind, radar_key)
set_dir.mkdir(parents=True, exist_ok=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}"
s11_key = f"s11_{suffix}"
s21_key = f"s21_{suffix}"
payload[freq_key] = np.asarray(trace.frequency_hz, dtype=np.float32)
payload[s11_key] = np.asarray(trace.s11, dtype=np.complex64)
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,
"s11_key": s11_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 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)
payload[_SET_METADATA_KEY] = np.asarray(metadata_json)
npz_tmp = meta_tmp = backup_path = None
archive_replaced = False
try:
npz_tmp = self._temporary_file(npz_path)
meta_tmp = self._temporary_file(meta_path)
backup_path = self._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)
meta_tmp.write_text(metadata_json, encoding="utf-8")
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
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)
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 = Path(tempfile.mkdtemp(prefix=".sets-", dir=self._root_dir))
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)
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."""
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")
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),
s11=np.asarray(arrays[combo["s11_key"]], dtype=np.complex64),
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
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)
output_path.parent.mkdir(parents=True, exist_ok=True)
temporary = self._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 set %s/%s/%s bundle to %s", kind, radar_key, set_name, output_path)
return output_path
def save_runtime_snapshot(
self,
output_dir: Path,
raw_history: list[SweepCollection],
preprocessed_history: list[SweepCollection],
result_history: list[ResultCollection],
last_n: int,
) -> Path:
"""Save historical runtime collections using binary on-disk format."""
if last_n <= 0:
raise ValueError("last_n must be > 0")
output_dir.mkdir(parents=True, exist_ok=True)
timestamp = datetime.now(timezone.utc).strftime("%Y%m%d_%H%M%S")
snapshot_dir = output_dir / f"snapshot_{timestamp}"
with atomic_snapshot_directory(snapshot_dir) as staging:
save_trace_history_binary(staging / "raw", raw_history[-last_n:], RAW_MAGIC)
save_trace_history_binary(staging / "preprocessed", preprocessed_history[-last_n:], PREPROC_MAGIC)
save_result_history_binary(staging / "results", result_history[-last_n:])
logger.info("Saved binary runtime snapshot (last_n=%d) to %s", last_n, snapshot_dir)
return snapshot_dir
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],
last_n: int,
*,
name_prefix: str = "",
include_unaligned: bool = False,
) -> tuple[Path, dict[str, Any]]:
"""Save historical runtime collections in NumPy tree format."""
if last_n <= 0:
raise ValueError("last_n must be > 0")
snapshot_stem = _compose_snapshot_stem(snapshot_name, name_prefix)
output_root_dir.mkdir(parents=True, exist_ok=True)
snapshot_dir = output_root_dir / snapshot_stem
if snapshot_dir.exists():
raise FileExistsError(f"Snapshot directory already exists: {snapshot_dir}")
selected_raw, selected_preprocessed, selected_results, selection_summary = select_aligned_histories(
raw_history,
preprocessed_history,
result_history,
last_n,
include_unaligned=include_unaligned,
)
if selection_summary.get("anchor_stage") == "results":
expected = min(int(last_n), len(result_history))
if len(selected_results) != expected:
raise RuntimeError(
"Snapshot results selection invariant failed: "
f"expected={expected}, selected={len(selected_results)}"
)
with atomic_snapshot_directory(snapshot_dir) as staging:
save_trace_history_numpy(staging / "raw", selected_raw)
save_trace_history_numpy(staging / "preprocessed", selected_preprocessed)
save_result_history_numpy(staging / "results", selected_results)
selection_summary["raw_count"] = len(selected_raw)
selection_summary["preprocessed_count"] = len(selected_preprocessed)
selection_summary["result_count"] = len(selected_results)
selection_summary["snapshot_stem"] = snapshot_stem
selection_summary["snapshot_dir"] = str(snapshot_dir)
logger.info(
"Saved NumPy runtime snapshot to %s (raw=%d preprocessed=%d results=%d, anchor=%s)",
snapshot_dir,
len(selected_raw),
len(selected_preprocessed),
len(selected_results),
selection_summary.get("anchor_stage"),
)
return snapshot_dir, selection_summary
def save_runtime_vna_history_json(
self,
output_root_dir: Path,
output_name: str,
raw_history: list[SweepCollection],
preprocessed_history: list[SweepCollection],
result_history: list[ResultCollection],
last_n: int,
*,
input_index: int = 0,
output_index: int = 0,
channel: str = "s21",
primary_stage: str = "preprocessed",
include_unaligned: bool = False,
) -> tuple[Path, dict[str, Any]]:
"""Save runtime history as vna_system-compatible JSON file."""
if last_n <= 0:
raise ValueError("last_n must be > 0")
output_stem = output_name.strip() or datetime.now(timezone.utc).strftime("snapshot_%Y%m%d_%H%M%S")
output_stem = sanitize_path_component(output_stem)
output_root_dir.mkdir(parents=True, exist_ok=True)
output_dir = self._vna_json_output_dir(output_root_dir, output_stem)
output_path = output_dir / f"{output_stem}_vna_bscan_history.json"
selected_raw, selected_preprocessed, selected_results, selection_summary = select_aligned_histories(
raw_history,
preprocessed_history,
result_history,
last_n,
include_unaligned=include_unaligned,
)
payload = build_vna_history_payload(
selected_raw,
selected_preprocessed,
selected_results,
input_index=input_index,
output_index=output_index,
channel=channel,
primary_stage=primary_stage,
)
with atomic_snapshot_directory(output_dir) as staging:
(staging / output_path.name).write_text(json.dumps(payload, ensure_ascii=False, indent=2), encoding="utf-8")
summary: dict[str, Any] = dict(selection_summary)
summary["raw_count"] = len(selected_raw)
summary["preprocessed_count"] = len(selected_preprocessed)
summary["result_count"] = len(selected_results)
summary["raw_record_count"] = int(payload.get("raw_record_count", 0))
summary["preprocessed_record_count"] = int(payload.get("preprocessed_record_count", 0))
summary["sweep_count"] = len(payload.get("sweep_history", []))
summary["input_index"] = int(input_index)
summary["output_index"] = int(output_index)
summary["channel"] = str(channel)
summary["primary_stage"] = str(primary_stage)
summary["output_stem"] = output_stem
summary["output_dir"] = str(output_dir)
summary["output_path"] = str(output_path)
logger.info(
"Saved VNA history JSON to %s (input=%d output=%d channel=%s sweeps=%d)",
output_path,
int(input_index),
int(output_index),
channel,
summary["sweep_count"],
)
return output_path, summary
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],
last_n: int,
*,
channel: str = "s21",
primary_stage: str = "preprocessed",
name_prefix: str = "",
include_unaligned: bool = False,
) -> tuple[list[Path], dict[str, Any]]:
"""Save one runtime VNA-history JSON per available combo."""
if last_n <= 0:
raise ValueError("last_n must be > 0")
output_stem = _compose_snapshot_stem(output_name, name_prefix)
output_root_dir.mkdir(parents=True, exist_ok=True)
output_dir = self._vna_json_output_dir(output_root_dir, output_stem)
selected_raw, selected_preprocessed, selected_results, selection_summary = select_aligned_histories(
raw_history,
preprocessed_history,
result_history,
last_n,
include_unaligned=include_unaligned,
)
combos = sorted(
{
(int(trace.combo.input), int(trace.combo.output))
for collection in [*selected_raw, *selected_preprocessed]
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}_{channel}_vna_bscan_history.json"
)
if output_path.exists():
raise FileExistsError(f"Output JSON file already exists: {output_path}")
payloads.append(
build_vna_history_payload(
selected_raw,
selected_preprocessed,
selected_results,
input_index=input_index,
output_index=output_index,
channel=channel,
primary_stage=primary_stage,
)
)
output_paths.append(output_path)
with atomic_snapshot_directory(output_dir) as staging:
for output_path, payload in zip(output_paths, payloads, strict=True):
(staging / output_path.name).write_text(
json.dumps(payload, ensure_ascii=False, indent=2), encoding="utf-8"
)
summary: dict[str, Any] = dict(selection_summary)
summary["raw_count"] = len(selected_raw)
summary["preprocessed_count"] = len(selected_preprocessed)
summary["result_count"] = len(selected_results)
summary["preprocessed_record_count"] = int(sum(int(payload.get("preprocessed_record_count", 0)) for payload in payloads))
summary["combo_count"] = len(combos)
summary["combos"] = [list(combo) for combo in combos]
summary["channel"] = str(channel)
summary["primary_stage"] = str(primary_stage)
summary["output_stem"] = output_stem
summary["output_dir"] = str(output_dir)
summary["output_paths"] = [str(path) for path in output_paths]
logger.info(
"Saved %d VNA history JSON file(s) to %s (channel=%s)", len(output_paths), output_dir, channel
)
return output_paths, summary
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}")
@staticmethod
def _temporary_file(destination: Path) -> Path:
"""Create a unique temporary file beside its atomic-replacement destination."""
with tempfile.NamedTemporaryFile(
dir=destination.parent, prefix=f".{destination.name}.", suffix=".tmp", delete=False
) as handle:
path = Path(handle.name)
return path
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
__all__ = ["NpzStore", "radar_key_from_config"]