328 lines
15 KiB
Python
328 lines
15 KiB
Python
"""Snapshot selection and filesystem writers for runtime collection histories."""
|
|
|
|
from __future__ import annotations
|
|
|
|
import json
|
|
import logging
|
|
from contextlib import contextmanager
|
|
from pathlib import Path
|
|
import shutil
|
|
import tempfile
|
|
from typing import Any, Iterator
|
|
|
|
import numpy as np
|
|
|
|
from python_app.models.dataset_model import ResultCollection, ResultPayload, SweepCollection
|
|
from python_app.storage.npz.paths import collection_dir_name, sanitize_path_component
|
|
from python_app.storage.npz.serialize import serialize_result_collection, serialize_trace_collection, validate_trace_collection
|
|
|
|
logger = logging.getLogger(__name__)
|
|
|
|
|
|
@contextmanager
|
|
def atomic_snapshot_directory(destination: Path) -> Iterator[Path]:
|
|
"""Publish a complete export directory; discard staging files after a failed write."""
|
|
destination.parent.mkdir(parents=True, exist_ok=True)
|
|
if destination.exists():
|
|
raise FileExistsError(f"Snapshot directory already exists: {destination}")
|
|
staging = Path(tempfile.mkdtemp(prefix=f".{destination.name}.", dir=destination.parent))
|
|
try:
|
|
yield staging
|
|
if destination.exists():
|
|
raise FileExistsError(f"Snapshot directory already exists: {destination}")
|
|
staging.rename(destination)
|
|
finally:
|
|
if staging.exists():
|
|
shutil.rmtree(staging)
|
|
|
|
|
|
def _matching_collection(
|
|
history: list[SweepCollection], anchor: SweepCollection | ResultCollection,
|
|
used: set[int], *, same_timestamp: bool,
|
|
) -> SweepCollection | None:
|
|
"""Match preprocessing identities exactly and raw sweeps by their capture window.
|
|
|
|
Producers stamp raw/preprocessed collections at different instants. Their capture
|
|
window is copied unchanged; results keep the preprocessed timestamp. Old files
|
|
without timing fall back to the latest preceding collection with the same id.
|
|
"""
|
|
candidates = [
|
|
(index, item) for index, item in enumerate(history)
|
|
if index not in used and int(item.collection_id) == int(anchor.collection_id)
|
|
]
|
|
timestamp = int(anchor.monotonic_ns)
|
|
start_ns = int(getattr(anchor, "capture_start_ns", 0))
|
|
end_ns = int(getattr(anchor, "capture_end_ns", 0))
|
|
if same_timestamp and timestamp:
|
|
candidates = [(index, item) for index, item in candidates if int(item.monotonic_ns) == timestamp]
|
|
elif start_ns or end_ns:
|
|
candidates = [
|
|
(index, item) for index, item in candidates
|
|
if (int(item.capture_start_ns), int(item.capture_end_ns)) == (start_ns, end_ns)
|
|
]
|
|
elif timestamp:
|
|
candidates = [(index, item) for index, item in candidates if int(item.monotonic_ns) <= timestamp]
|
|
if not candidates:
|
|
return None
|
|
index, item = max(candidates, key=lambda pair: (int(pair[1].monotonic_ns), pair[0]))
|
|
used.add(index)
|
|
return item
|
|
|
|
|
|
def select_aligned_histories(
|
|
raw_history: list[SweepCollection],
|
|
preprocessed_history: list[SweepCollection],
|
|
result_history: list[ResultCollection],
|
|
last_n: int,
|
|
*,
|
|
include_unaligned: bool = False,
|
|
) -> tuple[list[SweepCollection], list[SweepCollection], list[ResultCollection], dict[str, Any]]:
|
|
"""Select aligned tails, or retain every stage independently for complete exports."""
|
|
if last_n <= 0:
|
|
raise ValueError("last_n must be > 0")
|
|
if include_unaligned:
|
|
return (
|
|
raw_history[-last_n:], preprocessed_history[-last_n:], result_history[-last_n:],
|
|
{"selection_mode": "independent_tail", "anchor_stage": "none",
|
|
"selected_collection_ids": [int(item.collection_id) for item in result_history[-last_n:]],
|
|
"aligned_key_count": 0},
|
|
)
|
|
|
|
selected_raw: list[SweepCollection] = []
|
|
selected_pre: list[SweepCollection] = []
|
|
selected_results = result_history[-last_n:]
|
|
used_raw: set[int] = set()
|
|
used_pre: set[int] = set()
|
|
if selected_results:
|
|
anchors = selected_results
|
|
anchor_stage = "results"
|
|
for result in reversed(selected_results):
|
|
pre = _matching_collection(preprocessed_history, result, used_pre, same_timestamp=True)
|
|
raw = _matching_collection(raw_history, pre or result, used_raw, same_timestamp=False)
|
|
if pre is not None:
|
|
selected_pre.append(pre)
|
|
if raw is not None:
|
|
selected_raw.append(raw)
|
|
selected_pre.reverse()
|
|
selected_raw.reverse()
|
|
elif preprocessed_history:
|
|
selected_pre = preprocessed_history[-last_n:]
|
|
anchors = selected_pre
|
|
anchor_stage = "preprocessed"
|
|
for pre in reversed(selected_pre):
|
|
raw = _matching_collection(raw_history, pre, used_raw, same_timestamp=False)
|
|
if raw is not None:
|
|
selected_raw.append(raw)
|
|
selected_raw.reverse()
|
|
else:
|
|
selected_raw = raw_history[-last_n:]
|
|
anchors = selected_raw
|
|
anchor_stage = "raw" if anchors else "none"
|
|
|
|
return selected_raw, selected_pre, selected_results, {
|
|
"selection_mode": f"{anchor_stage}_tail_with_optional_alignment",
|
|
"anchor_stage": anchor_stage,
|
|
"selected_collection_ids": [int(item.collection_id) for item in anchors],
|
|
"aligned_key_count": len(anchors),
|
|
"raw_missing_count": len(anchors) - len(selected_raw),
|
|
"preprocessed_missing_count": len(anchors) - len(selected_pre),
|
|
"result_missing_count": len(anchors) - len(selected_results),
|
|
}
|
|
|
|
|
|
def save_trace_history_binary(stage_dir: Path, history: list[SweepCollection], magic: int) -> None:
|
|
"""Write binary trace history with lightweight metadata sidecars."""
|
|
stage_dir.mkdir(parents=True, exist_ok=True)
|
|
logger.debug("Writing %d binary trace collection(s) to %s", len(history), stage_dir)
|
|
for index, collection in enumerate(history):
|
|
binary_path = stage_dir / f"{index:04d}.bin"
|
|
metadata_path = stage_dir / f"{index:04d}.json"
|
|
binary_path.write_bytes(serialize_trace_collection(collection, magic))
|
|
metadata_path.write_text(
|
|
json.dumps(
|
|
{
|
|
"collection_id": collection.collection_id,
|
|
"monotonic_ns": collection.monotonic_ns,
|
|
"capture_start_ns": int(collection.capture_start_ns),
|
|
"capture_end_ns": int(collection.capture_end_ns),
|
|
"trace_count": len(collection.traces),
|
|
# Also in the .bin trailer; repeated here so per-combo timing is
|
|
# readable without decoding the binary payload.
|
|
"traces": [
|
|
{
|
|
"input": int(trace.combo.input),
|
|
"output": int(trace.combo.output),
|
|
"capture_start_ns": int(trace.capture_start_ns),
|
|
"capture_end_ns": int(trace.capture_end_ns),
|
|
}
|
|
for trace in collection.traces
|
|
],
|
|
},
|
|
indent=2,
|
|
),
|
|
encoding="utf-8",
|
|
)
|
|
|
|
|
|
def save_result_history_binary(stage_dir: Path, history: list[ResultCollection]) -> None:
|
|
"""Write binary processed-result history with metadata sidecars."""
|
|
stage_dir.mkdir(parents=True, exist_ok=True)
|
|
logger.debug("Writing %d binary result collection(s) to %s", len(history), stage_dir)
|
|
for index, collection in enumerate(history):
|
|
binary_path = stage_dir / f"{index:04d}.bin"
|
|
metadata_path = stage_dir / f"{index:04d}.json"
|
|
binary_path.write_bytes(serialize_result_collection(collection))
|
|
metadata_path.write_text(
|
|
json.dumps(
|
|
{
|
|
"collection_id": collection.collection_id,
|
|
"monotonic_ns": collection.monotonic_ns,
|
|
"processing_duration_ns": int(collection.processing_duration_ns),
|
|
"collection_payload_count": len(collection.collection_payloads),
|
|
"block_count": len(collection.blocks),
|
|
},
|
|
indent=2,
|
|
),
|
|
encoding="utf-8",
|
|
)
|
|
|
|
|
|
def save_trace_history_numpy(
|
|
stage_dir: Path, history: list[SweepCollection], *, index_offset: int = 0
|
|
) -> None:
|
|
"""Write raw/preprocessed collections as a NumPy directory tree.
|
|
|
|
``index_offset`` continues the per-collection directory numbering across calls so a
|
|
streaming recorder can append successive chunks into the same stage directory
|
|
without colliding or restarting the index.
|
|
"""
|
|
for collection in history:
|
|
validate_trace_collection(collection)
|
|
stage_dir.mkdir(parents=True, exist_ok=True)
|
|
logger.debug("Writing %d NumPy trace collection(s) to %s", len(history), stage_dir)
|
|
for index, collection in enumerate(history, start=index_offset):
|
|
collection_dir = stage_dir / collection_dir_name(index, collection.collection_id, collection.monotonic_ns)
|
|
collection_dir.mkdir(parents=True, exist_ok=False)
|
|
|
|
traces_meta: list[dict[str, int | str]] = []
|
|
for trace in collection.traces:
|
|
tag = f"i{trace.combo.input}_o{trace.combo.output}"
|
|
freq = np.asarray(trace.frequency_hz, dtype=np.float32)
|
|
s11 = np.asarray(trace.s11, dtype=np.complex64)
|
|
s21 = np.asarray(trace.s21, dtype=np.complex64)
|
|
np.save(collection_dir / f"{tag}_freq.npy", freq)
|
|
np.save(collection_dir / f"{tag}_s11.npy", s11)
|
|
np.save(collection_dir / f"{tag}_s21.npy", s21)
|
|
traces_meta.append(
|
|
{
|
|
"input": int(trace.combo.input),
|
|
"output": int(trace.combo.output),
|
|
"points": int(freq.size),
|
|
# When each combo was measured, which in a switched matrix is
|
|
# spread across the collection window rather than aligned with it.
|
|
"capture_start_ns": int(trace.capture_start_ns),
|
|
"capture_end_ns": int(trace.capture_end_ns),
|
|
"freq_file": f"{tag}_freq.npy",
|
|
"s11_file": f"{tag}_s11.npy",
|
|
"s21_file": f"{tag}_s21.npy",
|
|
}
|
|
)
|
|
|
|
(collection_dir / "meta.json").write_text(
|
|
json.dumps(
|
|
{
|
|
"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),
|
|
"trace_count": len(collection.traces),
|
|
"traces": traces_meta,
|
|
},
|
|
indent=2,
|
|
),
|
|
encoding="utf-8",
|
|
)
|
|
|
|
|
|
def _save_result_payload(directory: Path, base_name: str, payload: ResultPayload) -> dict[str, Any]:
|
|
"""Persist every supported payload kind identically at collection and combo scope."""
|
|
meta: dict[str, Any] = {"kind": int(payload.kind), "name": payload.processing_name}
|
|
|
|
def save_array(field: str, array: np.ndarray) -> None:
|
|
filename = f"{base_name}_{field}.npy"
|
|
np.save(directory / filename, array, allow_pickle=False)
|
|
meta[f"{field}_file"] = filename
|
|
|
|
if payload.kind == 1:
|
|
freq = np.asarray(payload.frequency_hz, dtype=np.float32)
|
|
trace = np.asarray(payload.trace, dtype=np.complex64)
|
|
if freq.ndim != 1 or trace.ndim != 1 or freq.shape != trace.shape:
|
|
raise ValueError("Result trace frequency and values must be equally sized 1D arrays")
|
|
save_array("freq", freq)
|
|
save_array("trace", trace)
|
|
meta["points"] = int(freq.size)
|
|
elif payload.kind == 2:
|
|
save_array("scalar", np.asarray([payload.scalar_value], dtype=np.float32))
|
|
meta["scalar_value"] = float(payload.scalar_value)
|
|
elif payload.kind == 3:
|
|
x_axis = np.asarray(payload.image_x_axis, dtype=np.float32)
|
|
y_axis = np.asarray(payload.image_y_axis, dtype=np.float32)
|
|
image = np.asarray(payload.image, dtype=np.float32)
|
|
if x_axis.ndim != 1 or y_axis.ndim != 1 or image.shape != (y_axis.size, x_axis.size):
|
|
raise ValueError("Result image axis sizes must match its 2D matrix shape")
|
|
save_array("x_axis", x_axis)
|
|
save_array("y_axis", y_axis)
|
|
save_array("image", image)
|
|
meta.update(x_points=int(x_axis.size), y_points=int(y_axis.size))
|
|
elif payload.kind == 4:
|
|
table = np.asarray(payload.table, dtype=np.float32)
|
|
if table.ndim != 2:
|
|
raise ValueError("Result table payload must be a 2D matrix")
|
|
save_array("table", table)
|
|
meta.update(rows=int(table.shape[0]), columns=int(table.shape[1]))
|
|
else:
|
|
raise ValueError(f"Unsupported payload kind: {payload.kind}")
|
|
return meta
|
|
|
|
|
|
def save_result_history_numpy(
|
|
stage_dir: Path, history: list[ResultCollection], *, index_offset: int = 0
|
|
) -> None:
|
|
"""Write result payloads, preserving per-stage indices across streaming chunks."""
|
|
stage_dir.mkdir(parents=True, exist_ok=True)
|
|
logger.debug("Writing %d NumPy result collection(s) to %s", len(history), stage_dir)
|
|
for index, collection in enumerate(history, start=index_offset):
|
|
collection_dir = stage_dir / collection_dir_name(index, collection.collection_id, collection.monotonic_ns)
|
|
collection_dir.mkdir(parents=True, exist_ok=False)
|
|
collection_payload_meta = []
|
|
for payload_index, payload in enumerate(collection.collection_payloads):
|
|
name = sanitize_path_component(payload.processing_name or "processor")
|
|
base_name = f"collection_{payload_index:03d}_{name}_kind{payload.kind}"
|
|
collection_payload_meta.append(_save_result_payload(collection_dir, base_name, payload))
|
|
|
|
blocks_meta = []
|
|
for block_index, block in enumerate(collection.blocks):
|
|
block_dir = collection_dir / f"block_{block_index:03d}_i{block.combo.input}_o{block.combo.output}"
|
|
block_dir.mkdir(parents=True, exist_ok=False)
|
|
payload_meta = []
|
|
for payload_index, payload in enumerate(block.payloads):
|
|
name = sanitize_path_component(payload.processing_name or "processor")
|
|
base_name = f"{payload_index:03d}_{name}_kind{payload.kind}"
|
|
payload_meta.append(_save_result_payload(block_dir, base_name, payload))
|
|
blocks_meta.append({
|
|
"input": int(block.combo.input), "output": int(block.combo.output),
|
|
"payload_count": len(block.payloads), "dir": block_dir.name, "payloads": payload_meta,
|
|
})
|
|
|
|
(collection_dir / "meta.json").write_text(
|
|
json.dumps({
|
|
"collection_id": int(collection.collection_id),
|
|
"monotonic_ns": int(collection.monotonic_ns),
|
|
"processing_duration_ns": int(collection.processing_duration_ns),
|
|
"collection_payload_count": len(collection.collection_payloads),
|
|
"collection_payloads": collection_payload_meta,
|
|
"block_count": len(collection.blocks), "blocks": blocks_meta,
|
|
}, indent=2),
|
|
encoding="utf-8",
|
|
)
|