"""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"]