diff --git a/.pre-commit-config.yaml b/.pre-commit-config.yaml index 5605562..9da46be 100644 --- a/.pre-commit-config.yaml +++ b/.pre-commit-config.yaml @@ -1,6 +1,6 @@ repos: - repo: https://github.com/astral-sh/ruff-pre-commit - rev: v0.8.3 + rev: v0.15.12 hooks: # Run the linter. - id: ruff diff --git a/README.md b/README.md index 637d8ea..c0e1b7e 100644 --- a/README.md +++ b/README.md @@ -11,6 +11,7 @@ Key features: * **NWB Reader** - Stream data from NWB files (local or remote) as AxisArray messages * **NWB Writer** - Write incoming AxisArray streams to NWB files with automatic container management * **Flexible clock handling** - Support for system, monotonic, and unknown reference clocks +* **Pipeline settings logging** - Automatically record every component's settings into a `pipeline_settings` intervals table inside the NWB file ## Installation @@ -48,6 +49,38 @@ from ezmsg.nwb import NWBIteratorUnit, NWBSink For general ezmsg tutorials and guides, visit [ezmsg.org](https://www.ezmsg.org). +### Pipeline settings table + +When `NWBSink` is used inside an `ez.run` pipeline (with `ezmsg>=3.9.0`), it +opens a `GraphContext` against the running graph server, snapshots the +settings of every component in its session, and subscribes to subsequent +settings change events. Each snapshot is flattened into dotted column names +(e.g. `MY.UNIT.MyUnitSettings.endpoint.host`) and appended as a row to a +`pipeline_settings` `TimeIntervals` table inside the NWB file, alongside an +`updated_component` column identifying which component triggered the +transition. Reading back is straightforward: + +```python +from pynwb import NWBHDF5IO + +with NWBHDF5IO(path, "r") as io: + nwbfile = io.read() + df = nwbfile.intervals["pipeline_settings"].to_dataframe() +``` + +Notes: + +* Settings logging is best-effort. If the writer cannot connect to the graph + server (e.g. when running the consumer outside of `ez.run`), it logs a + warning and continues writing data without the table. +* Settings values are sanitized for NWB storage: primitives, NumPy scalars, + enums, paths, and fixed-shape sequences/arrays are stored natively; + mappings and irregular structures are JSON-encoded; `None` becomes the + string `"None"`. +* If a settings update changes a column's shape (scalar↔array or rank + change), the writer rotates into a new file segment (`_01.nwb`, + `_02.nwb`, …) so each segment's table stays internally consistent. + ## Development We use [`uv`](https://docs.astral.sh/uv/getting-started/installation/) for development. diff --git a/pyproject.toml b/pyproject.toml index f6a7ee5..0ac85fb 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -10,7 +10,7 @@ requires-python = ">=3.10" dynamic = ["version"] dependencies = [ "ezmsg>=3.9.0", - "ezmsg-baseproc>=1.9.0", + "ezmsg-baseproc>=1.10.2", "numpy>=1.26.0", "pynwb", "h5py", @@ -34,6 +34,8 @@ lint = [ test = [ "pytest>=8.0.0", "pytest-asyncio>=0.24.0", + "pytest-xdist>=3.5.0", + "filelock>=3.0", ] docs = [ "sphinx>=8.0", @@ -59,6 +61,7 @@ packages = ["src/ezmsg"] [tool.pytest.ini_options] asyncio_mode = "auto" asyncio_default_fixture_loop_scope = "function" +addopts = "-n auto --dist=loadfile" [tool.ruff] line-length = 120 diff --git a/src/ezmsg/nwb/__init__.py b/src/ezmsg/nwb/__init__.py index 5aa2f5b..9456d1f 100644 --- a/src/ezmsg/nwb/__init__.py +++ b/src/ezmsg/nwb/__init__.py @@ -5,6 +5,21 @@ from .iterator import NWBAxisArrayIterator as NWBAxisArrayIterator from .iterator import NWBIteratorSettings as NWBIteratorSettings from .iterator import NWBIteratorState as NWBIteratorState +from .pipeline_settings import ( + NWBPipelineSettingsSink as NWBPipelineSettingsSink, +) +from .pipeline_settings import ( + NWBPipelineSettingsSinkConsumer as NWBPipelineSettingsSinkConsumer, +) +from .pipeline_settings import ( + NWBPipelineSettingsSinkSettings as NWBPipelineSettingsSinkSettings, +) +from .pipeline_settings import ( + PipelineSettingsTableCollection as PipelineSettingsTableCollection, +) +from .pipeline_settings import ( + PipelineSettingsTableCollectionSettings as PipelineSettingsTableCollectionSettings, +) from .reader import NWBIteratorUnit as NWBIteratorUnit from .slicer import NWBSlicer as NWBSlicer from .util import ReferenceClockType as ReferenceClockType diff --git a/src/ezmsg/nwb/pipeline_settings.py b/src/ezmsg/nwb/pipeline_settings.py new file mode 100644 index 0000000..c8c9199 --- /dev/null +++ b/src/ezmsg/nwb/pipeline_settings.py @@ -0,0 +1,648 @@ +"""Typed-column pipeline-settings sink for ezmsg-nwb. + +Phase 2 of the pipeline-settings story. Where the generic :class:`NWBSink` +lands :class:`PipelineSettingsEvent` messages in a JSON-encoded +``AnnotationSeries`` (one string column, fixed schema), this subclass +projects each event's ``structured_value`` into a +``pynwb.epoch.TimeIntervals`` table named ``pipeline_settings`` with one +column per dotted-key path. Columns get their values' native dtypes, so +analysis tools can read settings as a normal pandas DataFrame. + +Interval semantics +------------------ + +Each event opens a new interval. The PREVIOUS interval is closed (and +written to disk) at the moment the next event arrives — its ``stop_time`` +is the new event's ``timestamp``. The currently-active (open) interval is +held in memory and flushed at file close. + +This means: a non-graceful crash will lose the most recent open interval. +That's accepted as the cost of interval semantics — your collaborator +chose this trade-off. + +File rotation on schema change +------------------------------ + +If an incoming event introduces a new column, or changes the per-cell +shape of an existing column (scalar↔array, rank change, fixed-shape +mismatch), the sink rotates into a fresh file segment whose table opens +already populated with the new schema. Each segment's table is +internally consistent. Reading back means iterating ``_NN.nwb`` +files in order. + +Use the bundled :class:`PipelineSettingsTableCollection` to wire a +:class:`PipelineSettingsUnit` (from ezmsg-baseproc) to this sink in a +single component. +""" + +from __future__ import annotations + +import asyncio +import time +import typing + +import ezmsg.core as ez +import numpy as np +import pynwb +from ezmsg.baseproc import ( + INIT_FINAL_COMPONENT_ADDRESS, + PipelineSettingsEvent, + PipelineSettingsEventType, + PipelineSettingsProducerSettings, + PipelineSettingsUnit, + flatten_component_settings, + processor_state, +) +from ezmsg.util.messages.axisarray import AxisArray +from hdmf.backends.hdf5.h5_utils import H5DataIO + +from .writer import ( + NWBSink, + NWBSinkConsumer, + NWBSinkSettings, + NWBSinkState, +) + +# --------------------------------------------------------------------------- +# Settings + State +# --------------------------------------------------------------------------- + + +class NWBPipelineSettingsSinkSettings(NWBSinkSettings): + """Settings for :class:`NWBPipelineSettingsSink`. + + Inherits all file-level fields from :class:`NWBSinkSettings` + (``filepath``, ``overwrite_old``, ``inc_clock``, ``recording``, + ``split_bytes``, ``meta_yaml``, ``expected_series``, ``axis``). + """ + + pipeline_settings_table_name: str = "pipeline_settings" + """Name of the per-file ``TimeIntervals`` table to write into. Lives + under ``nwbfile.intervals[]``.""" + + +@processor_state +class NWBPipelineSettingsSinkState(NWBSinkState): + """State for :class:`NWBPipelineSettingsSinkConsumer`. + + Adds in-memory bookkeeping for the open interval. ``settings_columns`` + locks the table schema for the current file segment; + ``settings_state`` mirrors the most-recent flat settings dict so a + closed interval can be written without re-flattening. + """ + + settings_columns: typing.Optional[typing.List[str]] = None + settings_state: typing.Optional[typing.Dict[str, typing.Any]] = None + settings_active_since: typing.Optional[float] = None + settings_prev_component: str = "__init__" + pending_initial_state: typing.Optional[typing.Dict[str, typing.Any]] = None + """Per-component INITIAL events accumulate here while the startup + snapshot is in flight. ``None`` means we're not currently buffering; + the dict is populated by per-component INITIALs and flushed to one + merged anchor row on the ``INIT_FINAL_COMPONENT_ADDRESS`` sentinel + (or on the first non-INITIAL event, if the producer dropped the + sentinel).""" + pending_initial_first_seen: typing.Optional[float] = None + """Timestamp of the first INITIAL event in the current buffer — used + as the merged anchor row's ``start_time`` when we eventually flush.""" + + +# --------------------------------------------------------------------------- +# Consumer +# --------------------------------------------------------------------------- + + +class NWBPipelineSettingsSinkConsumer(NWBSinkConsumer): + """Adds typed-column ``TimeIntervals`` writing on top of :class:`NWBSinkConsumer`. + + Public method: :meth:`write_settings_event`. The unit subclass calls + it on each :class:`PipelineSettingsEvent`. Internally the consumer + detects schema-incompatible changes and rotates the file via + :meth:`_rotate_file`. + """ + + @classmethod + def get_state_type(cls) -> type: + return NWBPipelineSettingsSinkState + + def _reset_state(self, message: typing.Optional[AxisArray]) -> None: + super()._reset_state(message) + # Per-file pipeline-settings tracker; cleared on every reset so a + # fresh file starts empty. The next event re-seeds the table. + self._state.settings_columns = [] + self._state.settings_state = {} + self._state.settings_active_since = None + self._state.settings_prev_component = "__init__" + self._state.pending_initial_state = None + self._state.pending_initial_first_seen = None + + @property + def _settings_table_name(self) -> str: + return self.settings.pipeline_settings_table_name + + # ------------------------------------------------------------------ + # Public API: write one settings event + # ------------------------------------------------------------------ + + def write_settings_event(self, event: PipelineSettingsEvent) -> None: + """Project a :class:`PipelineSettingsEvent` into native columns and append/rotate. + + Aggregation rules: + + - **Per-component INITIAL events** at startup are buffered in + memory (``pending_initial_state``). The table is *not* + registered with the file yet — registering it per component + would force a schema-driven rotation per component. + - **The ``INIT_FINAL_COMPONENT_ADDRESS`` sentinel** flushes the + buffer as ONE merged anchor row, registers the table, and + opens an interval tracking the merged state. + - **Any non-INITIAL event arriving while the buffer is + populated** (producer dropped the sentinel) flushes the buffer + first as if the sentinel had arrived, then processes the + event normally. + - **Subsequent schema-compatible event** closes the prior open + interval (writes a row ``[prev.timestamp, this.timestamp]``) + and opens a new one. + - **Schema-incompatible event** rotates into a new file segment + whose table opens with the merged schema and a fresh anchor. + - **Close** flushes any pending initial buffer (merged anchor) + and the still-open interval as ``[active_since, close_time]``. + """ + with self._lock: + if self._state.io is None: + # Reset cleared state and a new file isn't open yet; drop. + return + + is_init_final = ( + event.event_type == PipelineSettingsEventType.INITIAL + and event.component_address == INIT_FINAL_COMPONENT_ADDRESS + ) + + if is_init_final: + self._flush_pending_initial(event.timestamp) + return + + value = event.structured_value if event.structured_value is not None else event.repr_value + flat = flatten_component_settings(event.component_address, value) + if not flat: + return + + is_initial = event.event_type == PipelineSettingsEventType.INITIAL + + # Buffer per-component INITIALs only while the table hasn't + # been registered yet. After the first anchor row exists, + # late-arriving INITIALs (e.g. a runtime new component) are + # processed as ordinary updates so they go through the + # rotation path if their schema diverges. + if is_initial and not self._state.settings_columns: + if self._state.pending_initial_state is None: + self._state.pending_initial_state = {} + self._state.pending_initial_first_seen = event.timestamp + self._state.pending_initial_state.update(flat) + self._state.settings_prev_component = event.component_address + return + + # Non-INITIAL event arrived but we never saw the sentinel — + # flush the buffer as a merged anchor before processing. + if self._state.pending_initial_state is not None: + self._flush_pending_initial(event.timestamp) + + if not self._state.settings_columns: + # No INITIAL events were ever buffered — first event is + # a one-component snapshot. Register the table with an + # anchor row so the state hits disk immediately. + self._state.settings_columns = list(flat.keys()) + self._state.settings_state = {col: "" for col in self._state.settings_columns} + self._state.settings_state.update(flat) + self._state.settings_active_since = event.timestamp + self._state.settings_prev_component = event.component_address + rel_t = self._settings_relative_time(event.timestamp) + self._register_settings_table_with_first_row( + start_time=rel_t, + stop_time=rel_t, + updated_component=event.component_address, + ) + return + + if self._settings_update_requires_rotation(flat): + self._rotate_file( + timestamp=event.timestamp, + next_settings_state=self._merged_settings_state(flat), + next_settings_prev_component=event.component_address, + ) + return + + self._validate_settings_columns(flat) + self._apply_settings_update(event.component_address, flat, event.timestamp) + + def _flush_pending_initial(self, sentinel_timestamp: float) -> None: + """Materialize the buffered per-component INITIAL state as one anchor row. + + Called either when the producer's ``INIT_FINAL_COMPONENT_ADDRESS`` + sentinel arrives, on a non-INITIAL event arriving with a + non-empty buffer (producer dropped the sentinel), or from + ``_close_locked`` if neither happened before close. + """ + pending = self._state.pending_initial_state + first_seen = self._state.pending_initial_first_seen + self._state.pending_initial_state = None + self._state.pending_initial_first_seen = None + if not pending: + return + # Anchor at the FIRST INITIAL's timestamp (when the snapshot + # started accumulating); the open interval tracks from there + # until the next non-INITIAL event closes it. + anchor_t = first_seen if first_seen is not None else sentinel_timestamp + self._state.settings_columns = list(pending.keys()) + self._state.settings_state = dict(pending) + self._state.settings_active_since = anchor_t + self._state.settings_prev_component = INIT_FINAL_COMPONENT_ADDRESS + rel_t = self._settings_relative_time(anchor_t) + self._register_settings_table_with_first_row( + start_time=rel_t, + stop_time=rel_t, + updated_component=INIT_FINAL_COMPONENT_ADDRESS, + ) + + # ------------------------------------------------------------------ + # Table prep + # ------------------------------------------------------------------ + + def _configure_appendable_table( + self, + table: typing.Any, + sample_values: typing.Optional[typing.Dict[str, typing.Any]] = None, + ) -> None: + """Configure a dynamic table so its columns remain appendable after flush/reopen.""" + table.id.set_data_io(H5DataIO, {"maxshape": (None,), "chunks": True}) + for col in table.colnames: + sample_value = None if sample_values is None else sample_values.get(col) + if sample_value is None: + col_shape = getattr(table[col].data, "shape", ()) + maxshape = (None,) + tuple(col_shape[1:]) if len(col_shape) > 1 else (None,) + elif np.isscalar(sample_value) or isinstance(sample_value, (str, bytes)): + maxshape = (None,) + else: + col_shape = np.asarray(sample_value).shape + maxshape = (None,) + tuple(col_shape) + table[col].set_data_io(H5DataIO, {"maxshape": maxshape, "chunks": True}) + + def _get_settings_table(self) -> typing.Any: + nwbfile = self._state.nwbfile + if nwbfile is None or nwbfile.intervals is None: + return None + try: + return nwbfile.intervals[self._settings_table_name] + except Exception: + return None + + def _register_settings_table_with_first_row( + self, + start_time: float, + stop_time: float, + updated_component: str, + ) -> None: + """Build the table with one row already populated, then register it. + + pynwb's ``VectorData`` columns can't be serialized while empty + (dtype inference fails). The cheapest workaround is to register + the table only when we have a closed interval to write — the + first row supplies dtypes for every column. Caller must have + ``self._state.settings_state`` filled with the row's values. + """ + intervals = pynwb.epoch.TimeIntervals( + name=self._settings_table_name, + description="Flattened ezmsg settings snapshots active over each logged interval", + ) + intervals.add_column( + name="updated_component", + description="component that triggered the snapshot transition", + ) + for column_name in self._state.settings_columns or []: + intervals.add_column(name=column_name, description="flattened ezmsg setting") + intervals.add_interval( + start_time=start_time, + stop_time=stop_time, + updated_component=updated_component, + **(self._state.settings_state or {}), + ) + self._state.nwbfile.add_time_intervals(intervals) + self._configure_appendable_table(intervals, self._state.settings_state) + # Flush so the file gains a valid NWB structure with our table + # populated and h5py datasets become reachable for subsequent + # in-place appends. + self._flush_io(reopen=True) + + # ------------------------------------------------------------------ + # Schema validation / shape tracking + # ------------------------------------------------------------------ + + def _settings_relative_time(self, timestamp: float) -> float: + """Convert a wall-clock timestamp into file-relative session time. + + Mirrors :meth:`NWBSinkConsumer.write_annotation`'s behavior: when + ``start_timestamp`` is unset (no data has anchored the file yet), + the wall-clock value is passed through unmodified. Once data has + set ``start_timestamp``, every subsequent row uses that baseline. + """ + if self._state.start_timestamp != 0.0: + return float(timestamp) - self._state.start_timestamp + return float(timestamp) + + def _settings_value_shape(self, value: typing.Any) -> typing.Tuple[int, ...]: + if value is None or np.isscalar(value) or isinstance(value, (str, bytes)): + return () + try: + return tuple(np.asarray(value).shape) + except Exception: + return () + + def _column_value_shape(self, column_name: str) -> typing.Tuple[int, ...]: + table = self._get_settings_table() + if table is not None and column_name in table.colnames: + data = table[column_name].data + if hasattr(data, "shape"): + shape = tuple(data.shape) + return shape[1:] if len(shape) > 1 else () + + state = self._state.settings_state or {} + if column_name in state: + return self._settings_value_shape(state[column_name]) + return () + + def _validate_settings_columns(self, flat_settings: typing.Dict[str, typing.Any]) -> None: + cols = self._state.settings_columns or [] + missing = [name for name in flat_settings if name not in cols] + if missing: + raise ValueError( + f"Received settings fields not present in settings table schema: {', '.join(sorted(missing))}" + ) + + def _settings_update_requires_rotation(self, flat_settings: typing.Dict[str, typing.Any]) -> bool: + cols = self._state.settings_columns or [] + for column_name, value in flat_settings.items(): + if column_name not in cols: + return True + current_shape = self._column_value_shape(column_name) + new_shape = self._settings_value_shape(value) + current_is_scalar = current_shape == () + new_is_scalar = new_shape == () + if current_is_scalar != new_is_scalar: + return True + if not current_is_scalar and current_shape != new_shape: + return True + return False + + def _merged_settings_state(self, flat_settings: typing.Dict[str, typing.Any]) -> typing.Dict[str, typing.Any]: + """Merge an incoming partial update over the current state for the next file.""" + cols = self._state.settings_columns or [] + state = self._state.settings_state or {} + merged: typing.Dict[str, typing.Any] = {column_name: state.get(column_name, "") for column_name in cols} + merged.update(flat_settings) + return merged + + # ------------------------------------------------------------------ + # Apply update / flush interval + # ------------------------------------------------------------------ + + def _apply_settings_update( + self, + component_address: str, + flat_settings: typing.Dict[str, typing.Any], + timestamp: float, + ) -> None: + """Close the current interval (writing a row) and open a new one.""" + self._flush_settings_interval(timestamp, self._state.settings_prev_component) + if not self._state.settings_state: + self._state.settings_state = {col: "" for col in (self._state.settings_columns or [])} + self._state.settings_state.update(flat_settings) + self._state.settings_active_since = timestamp + self._state.settings_prev_component = component_address + + def _flush_settings_interval(self, end_timestamp: float, updated_component: str) -> None: + """Append the currently-active settings snapshot as a closed interval row. + + On the first row in a file, registers the table with the NWB + structure (see :meth:`_register_settings_table_with_first_row`). + On subsequent rows, appends to the existing table. + """ + state = self._state + if not state.settings_state or state.settings_active_since is None: + return + + start_time = self._settings_relative_time(state.settings_active_since) + stop_time = self._settings_relative_time(end_timestamp) + if stop_time < start_time: + stop_time = start_time + + table = self._get_settings_table() + if table is None: + self._register_settings_table_with_first_row( + start_time=start_time, + stop_time=stop_time, + updated_component=updated_component, + ) + return + + table.add_interval( + start_time=start_time, + stop_time=stop_time, + updated_component=updated_component, + **state.settings_state, + ) + + # ------------------------------------------------------------------ + # File rotation on schema change + # ------------------------------------------------------------------ + + def _rotate_file( + self, + timestamp: float, + next_settings_state: typing.Optional[typing.Dict[str, typing.Any]], + next_settings_prev_component: str, + ) -> None: + """Close the current segment and open a new one tracking ``next_settings_state``. + + The current open interval is flushed to the OLD file as a closed + interval ``[active_since, timestamp]``. The NEW file opens with + no settings table registered yet — the next event (or close) + will lazily register it via :meth:`_register_settings_table_with_first_row`, + with ``next_settings_state`` as the open interval starting at + ``timestamp``. + """ + state = self._state + if state.settings_state and state.settings_active_since is not None: + self._flush_settings_interval(timestamp, state.settings_prev_component) + + next_settings_state = dict(next_settings_state or {}) + + # Prevent close() from re-appending the interval we just flushed. + state.settings_active_since = None + + # Advance to the next file segment before creating the replacement file. + state.split_count += 1 + self.path_on_disk.unlink(missing_ok=True) + + new_nwbfile, new_meta = self._copy_nwb() + self.close() + self._nwb_create_or_fail(nwbfile=new_nwbfile) + self._prep_from_meta(new_meta) + + if next_settings_state: + # Track the migrated state as a new open interval AND write + # an anchor row in the new segment immediately, mirroring the + # first-event behaviour on a fresh file. A crash between + # rotations should not lose the rotation snapshot. + self._state.settings_columns = list(next_settings_state.keys()) + self._state.settings_state = dict(next_settings_state) + self._state.settings_active_since = timestamp + self._state.settings_prev_component = next_settings_prev_component + rel_t = self._settings_relative_time(timestamp) + self._register_settings_table_with_first_row( + start_time=rel_t, + stop_time=rel_t, + updated_component=next_settings_prev_component, + ) + else: + self._state.settings_columns = [] + self._state.settings_state = {} + self._state.settings_active_since = None + self._state.settings_prev_component = next_settings_prev_component + + # ------------------------------------------------------------------ + # Close: flush the open interval; keep the file if the table has rows + # ------------------------------------------------------------------ + + def _close_locked(self, state: NWBPipelineSettingsSinkState, write: bool, log: bool) -> None: + if state.io is None: + return + # If we received per-component INITIAL events but never saw the + # sentinel before close, materialize the buffer as a merged anchor + # row now so the snapshot still hits disk. + if state.pending_initial_state is not None: + self._flush_pending_initial(time.time()) + # Flush any open pipeline-settings interval so it lands in the file. + if state.settings_state and state.settings_active_since is not None: + self._flush_settings_interval(time.time(), state.settings_prev_component) + state.settings_active_since = None + + # Compute b_delete here (parent's logic + settings-table check) so + # we can override the parent's potential delete-on-empty. + nwbfile = state.nwbfile + io = state.io + if write: + io.write(nwbfile) + src_str = f"{io.source}" + b_delete = sum(s.bytes_written for s in state.series.values()) == 0 + for key in ["epochs", "trials"]: + if hasattr(nwbfile, key) and getattr(nwbfile, key) is not None: + b_delete = b_delete and len(getattr(nwbfile, key)) == 1 # EZNWB-START + if state.annotation_ts and any(len(ts) > 0 for ts in state.annotation_ts.values()): + b_delete = False + # A populated pipeline-settings table is also "content". + settings_table = self._get_settings_table() + if settings_table is not None and len(settings_table.id) > 0: + b_delete = False + io.close() + state.nwbfile = None + state.io = None + state.series = {} + state.annotation_data = {} + state.annotation_ts = {} + state.settings_columns = [] + state.settings_state = {} + state.settings_prev_component = "__init__" + state.pending_initial_state = None + state.pending_initial_first_seen = None + if log: + ez.logger.info(f"Closed file at {src_str}") + if b_delete: + self.path_on_disk.unlink(missing_ok=True) + if log: + ez.logger.info(f"Deleted empty file at {src_str}.") + + +# --------------------------------------------------------------------------- +# Sink unit +# --------------------------------------------------------------------------- + + +class NWBPipelineSettingsSink(NWBSink): + """Specialized :class:`NWBSink` that writes :class:`PipelineSettingsEvent` + messages into a typed-column ``TimeIntervals`` table. + + Drop-in replacement for :class:`NWBSink` when used inside + :class:`PipelineSettingsTableCollection`. Inherits all of + ``NWBSink``'s data-recording behaviour; replaces ``on_annotation`` so + only :class:`PipelineSettingsEvent` traffic on ``INPUT_ANNOTATION`` + is consumed (other annotation-shaped messages are silently dropped — + use :class:`NWBSink` directly if you need the JSON-AnnotationSeries + fallback). + """ + + SETTINGS = NWBPipelineSettingsSinkSettings + + def create_processor(self) -> None: + """Construct the typed-column consumer rather than the generic one.""" + from ezmsg.baseproc.units import _close_previous + + _close_previous(getattr(self, "processor", None)) + self.processor = NWBPipelineSettingsSinkConsumer(settings=self.SETTINGS) + + @ez.subscriber(NWBSink.INPUT_ANNOTATION) + async def on_annotation(self, msg: typing.Any) -> None: + """Dispatch :class:`PipelineSettingsEvent` to the typed-column writer. + + Non-:class:`PipelineSettingsEvent` messages are silently ignored + — keep the typed semantics clean. Wrap with another sink (or run + a plain :class:`NWBSink` in parallel) if you need JSON + annotations alongside. + """ + if not isinstance(msg, PipelineSettingsEvent): + return + try: + await asyncio.to_thread(self.processor.write_settings_event, msg) + except Exception as exc: + ez.logger.warning( + f"{self.address} failed to write pipeline-settings event for {msg.component_address}: {exc}" + ) + + +# --------------------------------------------------------------------------- +# Bundled Collection +# --------------------------------------------------------------------------- + + +class PipelineSettingsTableCollectionSettings(ez.Settings): + producer: PipelineSettingsProducerSettings + sink: NWBPipelineSettingsSinkSettings + + +class PipelineSettingsTableCollection(ez.Collection): + """Bundle a :class:`PipelineSettingsUnit` with an :class:`NWBPipelineSettingsSink`. + + Accepts an external ``INPUT_SIGNAL`` (relayed to the sink) so the + Collection drops into a pipeline as a normal NWB sink, with settings + logging happening transparently alongside acquisition recording. + Also relays ``INPUT_SETTINGS`` so users can push :class:`NWBSinkSettings` + updates from outside the Collection. + """ + + SETTINGS = PipelineSettingsTableCollectionSettings + + PUB = PipelineSettingsUnit() + SINK = NWBPipelineSettingsSink() + + INPUT_SIGNAL = ez.InputTopic(AxisArray) + INPUT_SETTINGS = ez.InputTopic(NWBPipelineSettingsSinkSettings) + + def configure(self) -> None: + self.PUB.apply_settings(self.SETTINGS.producer) + self.SINK.apply_settings(self.SETTINGS.sink) + + def network(self) -> ez.NetworkDefinition: + return ( + (self.PUB.OUTPUT_SIGNAL, self.SINK.INPUT_ANNOTATION), + (self.INPUT_SIGNAL, self.SINK.INPUT_SIGNAL), + (self.INPUT_SETTINGS, self.SINK.INPUT_SETTINGS), + ) diff --git a/src/ezmsg/nwb/util.py b/src/ezmsg/nwb/util.py index 6e0fdf3..b5a90cb 100644 --- a/src/ezmsg/nwb/util.py +++ b/src/ezmsg/nwb/util.py @@ -10,7 +10,8 @@ class ReferenceClockType(Enum): def build_nwb_fname(metadata: DeepDict) -> str: + """Build a default NWB filename from the session metadata.""" fname_str = f"sub-{metadata['Subject']['subject_id']}" - ses = metadata["NWBFile"].get("session_id", metadata["NWBFile"]["session_start_time"].strftime("%Y%m%dT%H%M")) + ses = metadata["NWBFile"].get("session_id", metadata["NWBFile"]["session_start_time"].strftime("%Y%m%dT%H%M%S")) fname_str += f"_ses-{ses}" return f"{fname_str}_ephys.nwb" diff --git a/src/ezmsg/nwb/writer.py b/src/ezmsg/nwb/writer.py index e6e6cab..b15b227 100644 --- a/src/ezmsg/nwb/writer.py +++ b/src/ezmsg/nwb/writer.py @@ -70,6 +70,7 @@ import datetime import os import re +import threading import time import typing import warnings @@ -86,7 +87,10 @@ from hdmf.backends.hdf5.h5_utils import H5DataIO from neuroconv.utils import DeepDict, dict_deep_update, load_dict_from_file -from .util import ReferenceClockType, build_nwb_fname +from .util import ( + ReferenceClockType, + build_nwb_fname, +) try: from ezmsg.baseproc import SampleTriggerMessage @@ -126,9 +130,6 @@ class SeriesState: @processor_state class NWBSinkState: - # hash is required by the stateful machinery we'll migrate onto in - # step 3; carrying it now keeps the state class identical across the - # base-class switch. hash: int = -1 filepath: typing.Optional[Path] = None io: typing.Optional[pynwb.NWBHDF5IO] = None @@ -136,6 +137,12 @@ class NWBSinkState: series: typing.Optional[typing.Dict[str, SeriesState]] = None start_timestamp: float = 0.0 split_count: int = 0 + # Per-file map of annotation-series name → its h5py string dataset. + # Reset on every ``_reset_state`` so a fresh file starts empty; the + # entry is created on first sight of each ``table_name`` in an + # incoming annotation row. + annotation_data: typing.Optional[typing.Dict[str, typing.Any]] = None + annotation_ts: typing.Optional[typing.Dict[str, typing.Any]] = None class NWBSinkConsumer(BaseStatefulConsumer[NWBSinkSettings, AxisArray, NWBSinkState]): @@ -154,6 +161,10 @@ class NWBSinkConsumer(BaseStatefulConsumer[NWBSinkSettings, AxisArray, NWBSinkSt def __init__(self, *args, settings: typing.Optional[NWBSinkSettings] = None, **kwargs): super().__init__(*args, settings=settings, **kwargs) + # Serializes ``_process`` appends and ``write_annotation`` calls + # from the annotation subscriber; HDF5 writes from two threads to + # the same file are unsafe. + self._lock = threading.RLock() # ``_current_msg`` is scratch space shared between ``_process`` and # helpers like ``_prep_continuous_io`` / ``_prep_from_meta``. Deferred # migration into state — see _prep_from_meta cleanup follow-up. @@ -186,6 +197,8 @@ def _reset_state(self, message: typing.Optional[AxisArray]) -> None: self._state.series = {} self._state.start_timestamp = 0.0 self._state.split_count = 0 + self._state.annotation_data = {} + self._state.annotation_ts = {} self._check_filepath() self._nwb_create_or_fail() @@ -207,127 +220,132 @@ async def _aprocess(self, message: AxisArray) -> None: await asyncio.to_thread(self._process, message) def _process(self, message: AxisArray) -> None: - self._current_msg = message - - # Adjust incoming data - if _HAS_SAMPLE_TRIGGER and isinstance(self._current_msg, SampleTriggerMessage): - # SampleTriggerMessage. Rewrite as AxisArray. - timestamp = self._current_msg.timestamp - period = self._current_msg.period - if period is not None and len(period) > 0: - timestamp = timestamp + period[0] - # Wrap value in a 2D shape (1, 1) so ``_append_events`` iterates - # rows of label-iterables (matching the convention used by the - # plain ``key="epochs"`` AxisArray path). A 1D shape would make - # ``",".join(ev_str)`` join the characters of a single string. - self._current_msg = AxisArray( - data=np.array([[self._current_msg.value]]), - dims=["time", "ch"], - axes={"time": AxisArray.LinearAxis(gain=1.0, offset=timestamp)}, - key="epochs", - ) - elif not hasattr(self._current_msg, "data"): - return - else: - axis = self.settings.axis - targ_ax_ix = self._current_msg.get_axis_idx(axis) - if targ_ax_ix != 0: - self._current_msg = replace( - self._current_msg, - data=np.moveaxis(self._current_msg.data, targ_ax_ix, 0), - dims=[axis] + self._current_msg.dims[:targ_ax_ix] + self._current_msg.dims[targ_ax_ix + 1 :], + with self._lock: + self._current_msg = message + + # Adjust incoming data + if _HAS_SAMPLE_TRIGGER and isinstance(self._current_msg, SampleTriggerMessage): + # SampleTriggerMessage. Rewrite as AxisArray. + timestamp = self._current_msg.timestamp + period = self._current_msg.period + if period is not None and len(period) > 0: + timestamp = timestamp + period[0] + # Wrap value in a 2D shape (1, 1) so ``_append_events`` iterates + # rows of label-iterables (matching the convention used by the + # plain ``key="epochs"`` AxisArray path). A 1D shape would make + # ``",".join(ev_str)`` join the characters of a single string. + self._current_msg = AxisArray( + data=np.array([[self._current_msg.value]]), + dims=["time", "ch"], + axes={"time": AxisArray.LinearAxis(gain=1.0, offset=timestamp)}, + key="epochs", ) - - # Is this a new series? - b_new = self._state.io is None - b_new = b_new or self._current_msg.key not in self._state.series - - # If inc message key is in state.series but properties do not match previous dataset properties - # then close io and raise error - if not b_new and not self._check_msg_consistency(): - nwbfile = self._state.nwbfile - b_final_write = hasattr(nwbfile, "epochs") and nwbfile.epochs is not None - b_final_write = b_final_write or (hasattr(nwbfile, "trials") and nwbfile.trials is not None) - self.close(write=b_final_write) - raise ValueError("Data provided to NWBSink has changed shape. Closing NWB file.") - - if b_new: - # Use first incoming timestamp to set the session start time. - key = self._current_msg.key - axis = self.settings.axis - t0 = None - if axis in ["time", "win"] or "time" in self._current_msg.axes: - targ_dim = axis if axis in ["time", "win"] else "time" - if hasattr(self._current_msg.axes[targ_dim], "data"): - t0 = self._current_msg.axes[targ_dim].data[0] - else: - t0 = self._current_msg.axes[targ_dim].offset - _ = self.get_session_datetime(t0) - self._state.start_timestamp = self.get_session_timestamp(t0) - if self.settings.inc_clock == ReferenceClockType.MONOTONIC: - self._state.start_timestamp += time.monotonic() - time.time() - - # Create the container(s) for the new stream. - if key in ["epochs", "trials"]: - self._prep_event_io() - self._flush_io(reopen=True) - elif self._current_msg.data.dtype.type is np.str_: - raise ValueError(f"Cannot stream varlen str data to series {key}. Use 'epochs' or 'trials' instead.") + elif not hasattr(self._current_msg, "data"): + return else: - self._prep_continuous_io() - self._flush_io(reopen=True) - self._update_rate_for_current() - - if self.settings.recording and self._current_msg.data.size: - axis = self.settings.axis - timestamps = None - if axis in ["time", "win"] or "time" in self._current_msg.axes: - targ_dim = axis if axis in ["time", "win"] else "time" - time_ax = self._current_msg.axes[targ_dim] - if hasattr(time_ax, "data"): - timestamps = time_ax.data - self._state.start_timestamp - else: - timestamps = (np.arange(len(self._current_msg.data)) * time_ax.gain) + ( - time_ax.offset - self._state.start_timestamp + axis = self.settings.axis + targ_ax_ix = self._current_msg.get_axis_idx(axis) + if targ_ax_ix != 0: + self._current_msg = replace( + self._current_msg, + data=np.moveaxis(self._current_msg.data, targ_ax_ix, 0), + dims=[axis] + self._current_msg.dims[:targ_ax_ix] + self._current_msg.dims[targ_ax_ix + 1 :], ) - key = self._current_msg.key - if key in ["epochs", "trials"]: - self._append_events(key, timestamps, self._current_msg.data) - else: - series_state = self._state.series[key] - # Write data - dataset = series_state.data - dataset.resize(len(dataset) + len(self._current_msg.data), axis=0) - dataset[-len(self._current_msg.data) :] = self._current_msg.data - series_state.bytes_written += self._current_msg.data.nbytes - - # Write timestamps - if timestamps is not None: - ts = series_state.ts - ts.resize(len(ts) + len(timestamps), axis=0) - ts[-len(timestamps) :] = timestamps - series_state.bytes_written += timestamps.nbytes - - total_bytes = sum(s.bytes_written for s in self._state.series.values()) - if 0 < self.settings.split_bytes <= total_bytes and "%d" not in str(self._state.filepath): - self._state.split_count += 1 - self.path_on_disk.unlink(missing_ok=True) - new_nwbfile, new_meta = self._copy_nwb() - self.close() - self._nwb_create_or_fail(nwbfile=new_nwbfile) - self._prep_from_meta(new_meta) + # Is this a new series? + b_new = self._state.io is None + b_new = b_new or self._current_msg.key not in self._state.series + + # If inc message key is in state.series but properties do not match previous dataset properties + # then close io and raise error + if not b_new and not self._check_msg_consistency(): + nwbfile = self._state.nwbfile + b_final_write = hasattr(nwbfile, "epochs") and nwbfile.epochs is not None + b_final_write = b_final_write or (hasattr(nwbfile, "trials") and nwbfile.trials is not None) + self.close(write=b_final_write) + raise ValueError("Data provided to NWBSink has changed shape. Closing NWB file.") + + if b_new: + # Use first incoming timestamp to set the session start time. + key = self._current_msg.key + axis = self.settings.axis + t0 = None + if axis in ["time", "win"] or "time" in self._current_msg.axes: + targ_dim = axis if axis in ["time", "win"] else "time" + if hasattr(self._current_msg.axes[targ_dim], "data"): + t0 = self._current_msg.axes[targ_dim].data[0] + else: + t0 = self._current_msg.axes[targ_dim].offset + _ = self.get_session_datetime(t0) + self._state.start_timestamp = self.get_session_timestamp(t0) + if self.settings.inc_clock == ReferenceClockType.MONOTONIC: + self._state.start_timestamp += time.monotonic() - time.time() + + # Create the container(s) for the new stream. + if key in ["epochs", "trials"]: + self._prep_event_io() + self._flush_io(reopen=True) + elif self._current_msg.data.dtype.type is np.str_: + raise ValueError( + f"Cannot stream varlen str data to series {key}. Use 'epochs' or 'trials' instead." + ) + else: + self._prep_continuous_io() + self._flush_io(reopen=True) + self._update_rate_for_current() + + if self.settings.recording and self._current_msg.data.size: + axis = self.settings.axis + timestamps = None + if axis in ["time", "win"] or "time" in self._current_msg.axes: + targ_dim = axis if axis in ["time", "win"] else "time" + time_ax = self._current_msg.axes[targ_dim] + if hasattr(time_ax, "data"): + timestamps = time_ax.data - self._state.start_timestamp + else: + timestamps = (np.arange(len(self._current_msg.data)) * time_ax.gain) + ( + time_ax.offset - self._state.start_timestamp + ) + + key = self._current_msg.key + if key in ["epochs", "trials"]: + self._append_events(key, timestamps, self._current_msg.data) + else: + series_state = self._state.series[key] + # Write data + dataset = series_state.data + dataset.resize(len(dataset) + len(self._current_msg.data), axis=0) + dataset[-len(self._current_msg.data) :] = self._current_msg.data + series_state.bytes_written += self._current_msg.data.nbytes + + # Write timestamps + if timestamps is not None: + ts = series_state.ts + ts.resize(len(ts) + len(timestamps), axis=0) + ts[-len(timestamps) :] = timestamps + series_state.bytes_written += timestamps.nbytes + + total_bytes = sum(s.bytes_written for s in self._state.series.values()) + if 0 < self.settings.split_bytes <= total_bytes and "%d" not in str(self._state.filepath): + self._state.split_count += 1 + self.path_on_disk.unlink(missing_ok=True) + new_nwbfile, new_meta = self._copy_nwb() + self.close() + self._nwb_create_or_fail(nwbfile=new_nwbfile) + self._prep_from_meta(new_meta) @property def path_on_disk(self) -> Path: fp = Path(self._state.filepath) - if self.settings.split_bytes > 0: - if "%d" in str(fp): - return Path(re.sub("%d", "0", str(fp))) - else: - return fp.parent / (fp.stem + f"_{self._state.split_count:02}" + fp.suffix) - else: - return fp + if "%d" in str(fp): + return Path(re.sub("%d", "0", str(fp))) + # Suffix when ``split_bytes`` is configured (uniform "_NN" naming + # for the whole split-set, starting at _00) OR when a settings-driven + # rotation has incremented split_count past 0 (the original file kept + # its bare name, subsequent segments get "_01", "_02", ...). + if self.settings.split_bytes > 0 or self._state.split_count > 0: + return fp.parent / (fp.stem + f"_{self._state.split_count:02}" + fp.suffix) + return fp def get_session_datetime(self, try_t0: typing.Optional[float] = None) -> datetime.datetime: """ @@ -499,7 +517,7 @@ def _sanitize_shape( self._prep_event_io() for key, ss in meta.items(): - if key in ["epochs", "trials"]: + if key in ("epochs", "trials"): continue shape = _sanitize_shape(ss["shape"]) # Each stream gets its own dtype from the meta dict. Reading @@ -524,9 +542,10 @@ def _sanitize_shape( # Add the rate attribute to the timestamps series. Can only do this after flushing. for key, ss in meta.items(): - if key not in ["epochs", "trials"]: - series = self._state.nwbfile.acquisition[key] - series.timestamps.attrs["rate"] = ss["fs"] + if key in ("epochs", "trials"): + continue + series = self._state.nwbfile.acquisition[key] + series.timestamps.attrs["rate"] = ss["fs"] def close(self, write=False, log=True) -> None: """ @@ -540,6 +559,18 @@ def close(self, write=False, log=True) -> None: state = getattr(self, "_state", None) if state is None or state.io is None: return + # If __init__ raised before the lock was set up, fall through without + # locking; nothing else can be racing on this half-built instance. + lock = getattr(self, "_lock", None) + if lock is None: + self._close_locked(state, write=write, log=log) + return + with lock: + self._close_locked(state, write=write, log=log) + + def _close_locked(self, state: NWBSinkState, write: bool, log: bool) -> None: + if state.io is None: + return nwbfile = state.nwbfile io = state.io if write: @@ -549,10 +580,19 @@ def close(self, write=False, log=True) -> None: for key in ["epochs", "trials"]: if hasattr(nwbfile, key) and getattr(nwbfile, key) is not None: b_delete = b_delete and len(getattr(nwbfile, key)) == 1 # EZNWB-START + # Annotation rows are also "content" — keep the file even if no + # acquisition data was written. + if state.annotation_ts and any(len(ts) > 0 for ts in state.annotation_ts.values()): + b_delete = False + # empty file is the intended artifact even if no data was streamed. + if self.settings.expected_series is not None: + b_delete = False io.close() state.nwbfile = None state.io = None state.series = {} + state.annotation_data = {} + state.annotation_ts = {} if log: ez.logger.info(f"Closed file at {src_str}") if b_delete: @@ -715,6 +755,19 @@ def _flush_io(self, reopen: bool = True): ss.data = series.data ss.ts = series.timestamps + # Annotation series datasets share the same file handle; reopen + # invalidates them, so refresh from the live nwbfile. + if self._state.annotation_data is not None: + for name in list(self._state.annotation_data.keys()): + if name in self._state.nwbfile.acquisition: + series = self._state.nwbfile.acquisition[name] + self._state.annotation_data[name] = ( + series.data.dataset if isinstance(series.data, H5DataIO) else series.data + ) + self._state.annotation_ts[name] = ( + series.timestamps.dataset if isinstance(series.timestamps, H5DataIO) else series.timestamps + ) + def _append_events( self, key: str, @@ -726,6 +779,90 @@ def _append_events( for ev_t, ev_str in zip(timestamps, data): fun(start_time=ev_t, stop_time=ev_t + 0, **{"label": ",".join(ev_str)}) + # ------------------------------------------------------------------ + # Annotation series writer + # + # The unit's ``INPUT_ANNOTATION`` subscriber forwards anything with a + # ``flatten_for_table()`` method (e.g. ``PipelineSettingsEvent``) to + # ``write_annotation``. Each row becomes one entry in a per-file + # ``pynwb.misc.AnnotationSeries`` named after ``msg.table_name``, + # auto-created on first sight and appended to thereafter. + # ------------------------------------------------------------------ + + def _prep_annotation_series(self, name: str, description: str = "") -> None: + """Create an empty appendable ``AnnotationSeries`` and stash its datasets.""" + if self._state.annotation_data is None: + self._state.annotation_data = {} + self._state.annotation_ts = {} + if name in self._state.annotation_data: + return + + nwbfile = self._state.nwbfile + if nwbfile is None: + raise RuntimeError("Cannot prepare annotation series before NWB file is open.") + + dataio = H5DataIO( + data=np.array([], dtype=h5py.string_dtype()), + maxshape=(None,), + chunks=True, + ) + tsio = H5DataIO( + shape=(0,), + dtype=np.float64, + maxshape=(None,), + fillvalue=np.nan, + ) + series = pynwb.misc.AnnotationSeries( + name=name, + data=dataio, + timestamps=tsio, + description=description or f"Annotations written by NWBSink to '{name}'", + ) + nwbfile.add_acquisition(series) + + # Flush + reopen to materialize h5py datasets we can append to. + self._flush_io(reopen=True) + + live_series = self._state.nwbfile.acquisition[name] + data_ds = live_series.data.dataset if isinstance(live_series.data, H5DataIO) else live_series.data + ts_ds = ( + live_series.timestamps.dataset if isinstance(live_series.timestamps, H5DataIO) else live_series.timestamps + ) + self._state.annotation_data[name] = data_ds + self._state.annotation_ts[name] = ts_ds + + def write_annotation(self, table_name: str, timestamp: float, data: str) -> None: + """Append one row to the named ``AnnotationSeries``, creating it if needed. + + ``timestamp`` is wall-clock seconds. The value stored in the file is + ``timestamp - start_timestamp`` once a data message has set + ``start_timestamp`` (the file-relative convention NWB expects); + before then, the wall-clock value is passed through unmodified — the + sink has no baseline to subtract against, and inventing one (e.g. + from ``datetime.now()``) would silently desync from the eventual + data anchor. Annotations written before any data arrives therefore + carry whatever baseline their producer chose; downstream readers + should treat them accordingly. + """ + with self._lock: + if self._state.io is None: + # File closed (between reset and first message); silently drop. + return + if self._state.annotation_data is None or table_name not in self._state.annotation_data: + self._prep_annotation_series(table_name) + + if self._state.start_timestamp != 0.0: + relative_ts = float(timestamp) - self._state.start_timestamp + else: + relative_ts = float(timestamp) + + data_ds = self._state.annotation_data[table_name] + ts_ds = self._state.annotation_ts[table_name] + data_ds.resize(len(data_ds) + 1, axis=0) + data_ds[-1] = data + ts_ds.resize(len(ts_ds) + 1, axis=0) + ts_ds[-1] = relative_ts + def _prep_event_io(self): """ Prepare the NWB file to receive event data, either "epochs" or "trials". @@ -850,6 +987,48 @@ def _prep_continuous_io(self): class NWBSink(BaseConsumerUnit[NWBSinkSettings, AxisArray, NWBSinkConsumer]): SETTINGS = NWBSinkSettings + # Accept any object exposing the ``NWBPointRow`` shape (``table_name``, + # ``timestamp``, ``flatten_for_table`` returning ``{"data": str}``). + # Typed as ``object`` so the sink is open to multiple producer types + # (pipeline settings, user-defined annotation messages, …) without + # requiring a base class. Routing is by duck-typing in ``on_annotation``. + INPUT_ANNOTATION = ez.InputStream(object) + + @ez.subscriber(INPUT_ANNOTATION) + async def on_annotation(self, msg: typing.Any) -> None: + """Append one row to the per-table ``AnnotationSeries`` named by ``msg.table_name``. + + Ignores messages that don't quack like an ``NWBPointRow`` so other + traffic on this stream (if any) doesn't crash the sink. Any exception + from the writer is logged and swallowed — annotation logging is + best-effort and must not break primary acquisition recording. + """ + flatten = getattr(msg, "flatten_for_table", None) + table_name = getattr(msg, "table_name", None) + timestamp = getattr(msg, "timestamp", None) + if not callable(flatten) or not isinstance(table_name, str) or not isinstance(timestamp, (int, float)): + return + try: + row = flatten() + except Exception as exc: + ez.logger.warning(f"{self.address} annotation flatten_for_table raised: {exc}") + return + if row is None: + # Producer's signal that this is a control message (e.g. the + # ``INIT_FINAL_COMPONENT_ADDRESS`` sentinel from the pipeline- + # settings producer). No row to write; silently skip. + return + data = row.get("data") if isinstance(row, typing.Mapping) else None + if not isinstance(data, str): + ez.logger.warning( + f"{self.address} annotation row missing 'data: str' (got {type(data).__name__}); skipping." + ) + return + try: + await asyncio.to_thread(self.processor.write_annotation, table_name, float(timestamp), data) + except Exception as exc: + ez.logger.warning(f"{self.address} failed to write annotation to '{table_name}': {exc}") + async def shutdown(self) -> None: await super().shutdown() self.processor.close() diff --git a/tests/conftest.py b/tests/conftest.py index 7d23567..78fcfc2 100644 --- a/tests/conftest.py +++ b/tests/conftest.py @@ -4,13 +4,21 @@ import pytest from create_test_nwb import create_test_nwb +from filelock import FileLock @pytest.fixture(scope="session") def test_nwb_path(): - """Generate-and-cache synthetic NWB test file.""" + """Generate-and-cache synthetic NWB test file. + + Under pytest-xdist, ``scope="session"`` is per-worker, so each worker + would independently race to create the same file and collide on the + HDF5 write lock. The filelock serializes the check-and-create across + workers; the first builds it, the rest skip. + """ path = Path(__file__).parent / "data" / "test_synthetic.nwb" - if not path.exists(): - path.parent.mkdir(exist_ok=True) - create_test_nwb(path) + path.parent.mkdir(exist_ok=True) + with FileLock(str(path) + ".lock"): + if not path.exists(): + create_test_nwb(path) return path diff --git a/tests/test_integration.py b/tests/test_integration.py index 958f0ee..a0c5385 100644 --- a/tests/test_integration.py +++ b/tests/test_integration.py @@ -1,6 +1,7 @@ """Integration tests for ezmsg-nwb: ezmsg system tests and writer round-trip.""" import asyncio +import json import tempfile import typing from dataclasses import field @@ -14,7 +15,13 @@ from ezmsg.util.messagecodec import message_log from ezmsg.util.messagelogger import MessageLogger, MessageLoggerSettings from ezmsg.util.messages.axisarray import AxisArray -from ezmsg.util.terminate import TerminateOnTotal, TerminateOnTotalSettings +from ezmsg.util.terminate import ( + TerminateOnTimeout, + TerminateOnTimeoutSettings, + TerminateOnTotal, + TerminateOnTotalSettings, +) +from pynwb import NWBHDF5IO from ezmsg.nwb import ( NWBAxisArrayIterator, @@ -367,3 +374,297 @@ def test_settings_push_through_graph_enables_recording(): assert len(nwbfile.acquisition["K"].data) == n_messages * chunk outpath.unlink(missing_ok=True) + + +# -- Pipeline-settings table integration -- + + +def _make_writer_continuous_msg() -> AxisArray: + return AxisArray( + data=np.arange(6, dtype=float).reshape(3, 2), + dims=["time", "ch"], + axes={"time": AxisArray.TimeAxis(fs=100.0)}, + key="sig", + ) + + +def _make_writer_epochs_msg() -> AxisArray: + return AxisArray( + data=np.array([["a"], ["b"]], dtype="U"), + dims=["time", "ch"], + axes={"time": AxisArray.CoordinateAxis(np.array([0.0, 1.0]), dims=["time"], unit="s")}, + key="epochs", + ) + + +def test_writer_annotation_then_data_lands_in_acquisition(tmp_path): + """Annotations written before data should still materialize on close.""" + outpath = tmp_path / "ezmsg_nwb_annotation_then_data_test.nwb" + + sink = NWBSinkConsumer( + settings=NWBSinkSettings( + filepath=outpath, + overwrite_old=True, + inc_clock=ReferenceClockType.UNKNOWN, + ) + ) + sink.write_annotation("settings_annotations", timestamp=0.5, data='{"step": "init"}') + sink._process(_make_writer_continuous_msg()) + sink.write_annotation("settings_annotations", timestamp=1.5, data='{"step": "running"}') + sink.close(write=False) + + assert outpath.exists() + + with NWBHDF5IO(outpath, "r") as io: + nwbfile = io.read() + series = nwbfile.acquisition["settings_annotations"] + assert list(series.data[:]) == ['{"step": "init"}', '{"step": "running"}'] + + +def _find_free_port() -> int: + """Pick a free TCP port for an isolated GraphServer.""" + import socket + + with socket.socket(socket.AF_INET, socket.SOCK_STREAM) as s: + s.bind(("127.0.0.1", 0)) + return s.getsockname()[1] + + +def test_writer_pipeline_settings_event_via_graph(tmp_path): + """End-to-end: ``PipelineSettingsUnit`` publishes its session's settings + snapshot through the graph and ``NWBSink`` lands the events in a + ``settings_annotations`` AnnotationSeries. + + Uses the real producer (which opens a ``GraphContext`` against the + running ``GraphServer``) so this exercises Phase 1 end-to-end rather + than a mock publisher. We pre-start a ``GraphServer`` at a known free + port and pin both ``ez.run`` and the producer to that address so the + producer's ``GraphContext`` connects to the same server the runner is + using. (When ``ez.run`` is given an explicit ``graph_address``, + ``GraphService.ensure`` will *not* auto-start a server — it expects + one already listening at that address. That's why we start one here.) + """ + from ezmsg.baseproc import PipelineSettingsProducerSettings, PipelineSettingsUnit + from ezmsg.core.graphserver import GraphServer + + outpath = tmp_path / "ezmsg_nwb_pipeline_settings_via_graph.nwb" + outpath.unlink(missing_ok=True) + + graph_address = ("127.0.0.1", _find_free_port()) + + class _Settings(ez.Settings): + producer: PipelineSettingsProducerSettings + sink: NWBSinkSettings + term: TerminateOnTimeoutSettings + + class _Pipeline(ez.Collection): + SETTINGS = _Settings + + PUB = PipelineSettingsUnit() + SINK = NWBSink() + TERM = TerminateOnTimeout() + + def configure(self) -> None: + self.PUB.apply_settings(self.SETTINGS.producer) + self.SINK.apply_settings(self.SETTINGS.sink) + self.TERM.apply_settings(self.SETTINGS.term) + + def network(self) -> ez.NetworkDefinition: + # Fan PUB's events out to the sink AND the terminator: the + # terminator resets its idle clock on every event, then fires + # ``NormalTermination`` after the producer goes quiet (no more + # settings changes after the initial snapshot). + return ( + (self.PUB.OUTPUT_SIGNAL, self.SINK.INPUT_ANNOTATION), + (self.PUB.OUTPUT_SIGNAL, self.TERM.INPUT), + ) + + system = _Pipeline( + _Settings( + producer=PipelineSettingsProducerSettings(graph_address=graph_address), + sink=NWBSinkSettings( + filepath=outpath, + overwrite_old=True, + inc_clock=ReferenceClockType.UNKNOWN, + ), + term=TerminateOnTimeoutSettings(time=1.5), + ) + ) + + server = GraphServer() + server.start(graph_address) + try: + ez.run(SYSTEM=system, graph_address=graph_address) + finally: + server.stop() + + assert outpath.exists() + with NWBHDF5IO(str(outpath), "r") as io: + nwbfile = io.read() + series = nwbfile.acquisition["settings_annotations"] + rows = [json.loads(s) for s in series.data[:]] + + # PipelineSettingsUnit emits one INITIAL row per in-scope component; + # the session contains the Collection (SYSTEM) plus PUB + SINK + TERM. + assert all(r["event_type"] == "INITIAL" for r in rows) + components = {r["component"] for r in rows} + # Addresses use ``/`` as the separator (e.g. "SYSTEM/PUB"). Don't pin + # the root ("SYSTEM" only because we passed it that way to ez.run); + # confirm the snapshot covers each unit. + assert any(c.endswith("/PUB") for c in components) + assert any(c.endswith("/SINK") for c in components) + assert any(c.endswith("/TERM") for c in components) + + outpath.unlink(missing_ok=True) + + +def test_writer_event_append_after_reopen(tmp_path): + """Epoch rows should remain appendable after the initial write/reopen cycle.""" + outpath = tmp_path / "ezmsg_nwb_event_append_test.nwb" + + sink = NWBSinkConsumer( + settings=NWBSinkSettings( + filepath=outpath, + overwrite_old=True, + inc_clock=ReferenceClockType.UNKNOWN, + ) + ) + sink._process(_make_writer_epochs_msg()) + sink.close(write=False) + + assert outpath.exists() + + with NWBHDF5IO(outpath, "r") as io: + nwbfile = io.read() + df = nwbfile.epochs.to_dataframe() + + assert df["label"].tolist() == ["EZNWB-START", "a", "b"] + + +class _SinkSettingsPokerSettings(ez.Settings): + sink_filepath: Path + target_recording: bool = False + publish_after_s: float = 0.5 + terminate_after_s: float = 2.0 + + +class _SinkSettingsPoker(ez.Unit): + """Publish a ``NWBSinkSettings`` update to ``NWBSink.INPUT_SETTINGS`` + mid-run, then raise ``NormalTermination`` after a settle delay so the + producer has time to emit the resulting UPDATED event.""" + + SETTINGS = _SinkSettingsPokerSettings + OUTPUT_SETTINGS = ez.OutputStream(NWBSinkSettings) + + @ez.publisher(OUTPUT_SETTINGS) + async def poke(self) -> typing.AsyncGenerator: + s = self.SETTINGS + await asyncio.sleep(s.publish_after_s) + # Flip the ``recording`` flag — same filepath, NONRESET field, so + # the sink's file isn't disturbed; the only observable side-effect + # is the graph server recording a SettingsChangedEvent for SINK. + yield ( + self.OUTPUT_SETTINGS, + NWBSinkSettings( + filepath=s.sink_filepath, + overwrite_old=True, + recording=s.target_recording, + inc_clock=ReferenceClockType.UNKNOWN, + ), + ) + await asyncio.sleep(s.terminate_after_s) + raise ez.NormalTermination + + +def test_writer_pipeline_settings_updated_event_via_graph(tmp_path): + """End-to-end: ``PipelineSettingsUnit`` emits an UPDATED event when a + settings message is published into another unit's ``INPUT_SETTINGS`` + while the graph is running, and ``NWBSink`` lands it in the + ``settings_annotations`` series alongside the INITIAL snapshot. + + The mid-run change is driven by ``_SinkSettingsPoker`` — when its + ``NWBSinkSettings`` message lands at ``NWBSink.INPUT_SETTINGS``, the + backend reports the new value to the graph server, which broadcasts a + ``SettingsChangedEvent`` that the running ``PipelineSettingsProducer`` + receives via its subscription and forwards as an ``UPDATED`` event. + """ + from ezmsg.baseproc import PipelineSettingsProducerSettings, PipelineSettingsUnit + from ezmsg.core.graphserver import GraphServer + + outpath = tmp_path / "ezmsg_nwb_pipeline_settings_updated.nwb" + outpath.unlink(missing_ok=True) + + graph_address = ("127.0.0.1", _find_free_port()) + + class _Settings(ez.Settings): + producer: PipelineSettingsProducerSettings + sink: NWBSinkSettings + poker: _SinkSettingsPokerSettings + + class _Pipeline(ez.Collection): + SETTINGS = _Settings + + PUB = PipelineSettingsUnit() + SINK = NWBSink() + POKER = _SinkSettingsPoker() + + def configure(self) -> None: + self.PUB.apply_settings(self.SETTINGS.producer) + self.SINK.apply_settings(self.SETTINGS.sink) + self.POKER.apply_settings(self.SETTINGS.poker) + + def network(self) -> ez.NetworkDefinition: + return ( + (self.PUB.OUTPUT_SIGNAL, self.SINK.INPUT_ANNOTATION), + (self.POKER.OUTPUT_SETTINGS, self.SINK.INPUT_SETTINGS), + ) + + system = _Pipeline( + _Settings( + producer=PipelineSettingsProducerSettings(graph_address=graph_address), + sink=NWBSinkSettings( + filepath=outpath, + overwrite_old=True, + recording=True, + inc_clock=ReferenceClockType.UNKNOWN, + ), + poker=_SinkSettingsPokerSettings( + sink_filepath=outpath, + target_recording=False, + publish_after_s=0.5, + terminate_after_s=1.5, + ), + ) + ) + + server = GraphServer() + server.start(graph_address) + try: + ez.run(SYSTEM=system, graph_address=graph_address) + finally: + server.stop() + + assert outpath.exists() + with NWBHDF5IO(str(outpath), "r") as io: + nwbfile = io.read() + rows = [json.loads(s) for s in nwbfile.acquisition["settings_annotations"].data[:]] + + initial = [r for r in rows if r["event_type"] == "INITIAL"] + updated = [r for r in rows if r["event_type"] == "UPDATED"] + + # Initial snapshot covers the whole session. + assert len(initial) >= 3 + initial_components = {r["component"] for r in initial} + assert any(c.endswith("/PUB") for c in initial_components) + assert any(c.endswith("/SINK") for c in initial_components) + assert any(c.endswith("/POKER") for c in initial_components) + + # POKER's settings push lands as one (or more) UPDATED row whose + # component is SINK and whose ``recording`` field is the new value. + assert len(updated) >= 1 + sink_updates = [r for r in updated if r["component"].endswith("/SINK")] + assert sink_updates, f"expected an UPDATED row for SINK; got components {sorted({r['component'] for r in updated})}" + last = sink_updates[-1] + assert last["settings"]["recording"] is False + + outpath.unlink(missing_ok=True) diff --git a/tests/test_iterator.py b/tests/test_iterator.py index 2e40eb7..2ee46cf 100644 --- a/tests/test_iterator.py +++ b/tests/test_iterator.py @@ -515,7 +515,7 @@ def test_prefetch_worker_does_not_keep_iterator_alive(test_nwb_path): # Without forcing GC: refcount alone should be enough to drop the # iterator if the worker doesn't capture self. assert ref() is None, ( - "iterator survived `del` — something (most likely the prefetch worker) " "is holding a strong reference to self" + "iterator survived `del` — something (most likely the prefetch worker) is holding a strong reference to self" ) gc.collect() # belt-and-suspenders for any cyclic refs diff --git a/tests/test_pipeline_settings_table.py b/tests/test_pipeline_settings_table.py new file mode 100644 index 0000000..4b2d0a2 --- /dev/null +++ b/tests/test_pipeline_settings_table.py @@ -0,0 +1,616 @@ +"""Tests for the typed-column pipeline-settings sink (Phase 2).""" + +import asyncio +import socket +import typing +from pathlib import Path + +import ezmsg.core as ez +import numpy as np +import pytest +from ezmsg.baseproc import ( + INIT_FINAL_COMPONENT_ADDRESS, + PipelineSettingsEvent, + PipelineSettingsEventType, + PipelineSettingsProducerSettings, +) +from pynwb import NWBHDF5IO + +from ezmsg.nwb import ( + NWBPipelineSettingsSinkConsumer, + NWBPipelineSettingsSinkSettings, + PipelineSettingsTableCollection, + PipelineSettingsTableCollectionSettings, + ReferenceClockType, +) + + +@pytest.fixture(autouse=True) +def _reset_nwbsink_shared(): + NWBPipelineSettingsSinkConsumer.shared_start_datetime = None + NWBPipelineSettingsSinkConsumer.shared_t0 = None + NWBPipelineSettingsSinkConsumer.shared_clock_type = None + yield + NWBPipelineSettingsSinkConsumer.shared_start_datetime = None + NWBPipelineSettingsSinkConsumer.shared_t0 = None + NWBPipelineSettingsSinkConsumer.shared_clock_type = None + + +def _find_free_port() -> int: + with socket.socket(socket.AF_INET, socket.SOCK_STREAM) as s: + s.bind(("127.0.0.1", 0)) + return s.getsockname()[1] + + +def _make_event( + *, + component: str = "X.Y", + seq: int = 1, + timestamp: float = 1.0, + event_type: PipelineSettingsEventType = PipelineSettingsEventType.UPDATED, + structured: typing.Optional[dict] = None, +) -> PipelineSettingsEvent: + sv = structured if structured is not None else {"foo": 1, "bar": 2.5} + return PipelineSettingsEvent( + seq=seq, + timestamp=timestamp, + component_address=component, + event_type=event_type, + repr_value=sv, + structured_value=sv, + ) + + +def _sink_consumer(filepath: Path, **overrides) -> NWBPipelineSettingsSinkConsumer: + base = dict( + filepath=filepath, + overwrite_old=True, + inc_clock=ReferenceClockType.UNKNOWN, + ) + base.update(overrides) + return NWBPipelineSettingsSinkConsumer(settings=NWBPipelineSettingsSinkSettings(**base)) + + +# --------------------------------------------------------------------------- +# Single-event semantics +# --------------------------------------------------------------------------- + + +def test_single_event_creates_table_with_anchor_and_close_rows(tmp_path): + """First event eagerly writes an anchor row ``[t, t]`` so the initial + settings hit disk immediately. Close adds the open-interval row + ``[t, close_time]``.""" + outpath = tmp_path / "single_event.nwb" + sink = _sink_consumer(outpath) + + sink.write_settings_event(_make_event(structured={"alpha": 1, "beta": "hello"})) + sink.close(write=False) + + with NWBHDF5IO(outpath, "r") as io: + nwbfile = io.read() + table = nwbfile.intervals["pipeline_settings"] + df = table.to_dataframe() + + assert "X.Y.alpha" in df.columns + assert "X.Y.beta" in df.columns + # 2 rows: the anchor [1.0, 1.0] and the close-flush [1.0, close]. + assert len(df) == 2 + assert df["start_time"].iloc[0] == df["stop_time"].iloc[0] == 1.0 + assert df["start_time"].iloc[1] == 1.0 + assert df["stop_time"].iloc[1] >= 1.0 + assert (df["X.Y.alpha"] == 1).all() + assert (df["X.Y.beta"] == "hello").all() + assert (df["updated_component"] == "X.Y").all() + + +def test_two_events_close_first_interval(tmp_path): + """Two events plus close = anchor + 2 closed intervals = 3 rows.""" + outpath = tmp_path / "two_events.nwb" + sink = _sink_consumer(outpath) + + sink.write_settings_event(_make_event(component="A", seq=1, timestamp=10.0, structured={"x": 1})) + sink.write_settings_event(_make_event(component="A", seq=2, timestamp=20.0, structured={"x": 2})) + sink.close(write=False) + + with NWBHDF5IO(outpath, "r") as io: + nwbfile = io.read() + df = nwbfile.intervals["pipeline_settings"].to_dataframe() + + # Three rows in order: anchor [10, 10] @ values=1, closed [10, 20] @ + # values=1, open-on-close [20, close-time] @ values=2. + assert len(df) == 3 + assert df["A.x"].tolist() == [1, 1, 2] + assert df["start_time"].tolist()[:2] == [10.0, 10.0] + assert df["stop_time"].iloc[0] == 10.0 + assert df["stop_time"].iloc[1] == 20.0 + assert df["start_time"].iloc[2] == 20.0 + assert df["stop_time"].iloc[2] >= 20.0 + + +def test_native_dtypes_preserved(tmp_path): + """Floats stay floats, ints stay ints, strings stay strings, fixed-shape + lists become arrays.""" + outpath = tmp_path / "dtypes.nwb" + sink = _sink_consumer(outpath) + + sink.write_settings_event( + _make_event( + structured={ + "a_int": 7, + "a_float": 3.14, + "a_str": "x", + "a_list": [1.0, 2.0, 3.0], + } + ) + ) + sink.close(write=False) + + with NWBHDF5IO(outpath, "r") as io: + nwbfile = io.read() + df = nwbfile.intervals["pipeline_settings"].to_dataframe() + + assert df["X.Y.a_int"].iloc[-1] == 7 + assert df["X.Y.a_float"].iloc[-1] == pytest.approx(3.14) + assert df["X.Y.a_str"].iloc[-1] == "x" + np.testing.assert_array_equal(df["X.Y.a_list"].iloc[-1], np.array([1.0, 2.0, 3.0])) + + +def test_close_with_no_events_deletes_empty_file(tmp_path): + """A sink that received no settings events should still get the + parent's empty-file cleanup behavior.""" + outpath = tmp_path / "empty.nwb" + sink = _sink_consumer(outpath) + sink.close(write=False) + assert not outpath.exists() + + +def test_close_with_only_settings_keeps_file(tmp_path): + """A populated settings table is content — file should not be deleted + even if no acquisition data was written.""" + outpath = tmp_path / "settings_only.nwb" + sink = _sink_consumer(outpath) + sink.write_settings_event(_make_event()) + sink.close(write=False) + assert outpath.exists() + + +# --------------------------------------------------------------------------- +# Schema-compatible append +# --------------------------------------------------------------------------- + + +def test_compatible_update_appends_without_rotation(tmp_path): + """Same keys, same shapes → no rotation; rows accumulate in one table. + + Three events + close = 1 anchor + 3 closed-by-next-event/close = 4 rows. + """ + outpath = tmp_path / "compat.nwb" + rotated = tmp_path / "compat_01.nwb" + sink = _sink_consumer(outpath) + + sink.write_settings_event(_make_event(seq=1, timestamp=1.0, structured={"x": 1, "y": "a"})) + sink.write_settings_event(_make_event(seq=2, timestamp=2.0, structured={"x": 2, "y": "b"})) + sink.write_settings_event(_make_event(seq=3, timestamp=3.0, structured={"x": 3, "y": "c"})) + sink.close(write=False) + + assert outpath.exists() + assert not rotated.exists(), "no schema change → no rotation" + + with NWBHDF5IO(outpath, "r") as io: + nwbfile = io.read() + df = nwbfile.intervals["pipeline_settings"].to_dataframe() + + # Anchor row holds event 1's values; rows 1-3 are the proper closed + # intervals as each successive event arrives (and close). + assert df["X.Y.x"].tolist() == [1, 1, 2, 3] + assert df["X.Y.y"].tolist() == ["a", "a", "b", "c"] + + +# --------------------------------------------------------------------------- +# Schema rotation +# --------------------------------------------------------------------------- + + +def test_new_column_triggers_rotation(tmp_path): + """A new column key in a later event should rotate to a fresh file + segment whose table opens with both old + new columns.""" + outpath = tmp_path / "rotate_new_col.nwb" + rotated = tmp_path / "rotate_new_col_01.nwb" + sink = _sink_consumer(outpath) + + sink.write_settings_event(_make_event(seq=1, timestamp=1.0, structured={"x": 1})) + sink.write_settings_event(_make_event(seq=2, timestamp=2.0, structured={"x": 2, "y": 99})) + sink.close(write=False) + + assert outpath.exists() + assert rotated.exists() + + with NWBHDF5IO(outpath, "r") as io: + nwbfile = io.read() + df0 = nwbfile.intervals["pipeline_settings"].to_dataframe() + assert "X.Y.x" in df0.columns + assert "X.Y.y" not in df0.columns + + with NWBHDF5IO(rotated, "r") as io: + nwbfile = io.read() + df1 = nwbfile.intervals["pipeline_settings"].to_dataframe() + assert "X.Y.x" in df1.columns + assert "X.Y.y" in df1.columns + assert df1["X.Y.x"].iloc[-1] == 2 + assert df1["X.Y.y"].iloc[-1] == 99 + + +def test_scalar_to_array_triggers_rotation(tmp_path): + """Scalar→array transition for an existing column rotates files.""" + outpath = tmp_path / "rotate_scalar_array.nwb" + rotated = tmp_path / "rotate_scalar_array_01.nwb" + sink = _sink_consumer(outpath) + + sink.write_settings_event(_make_event(seq=1, timestamp=1.0, structured={"v": "foo"})) + sink.write_settings_event(_make_event(seq=2, timestamp=2.0, structured={"v": [1, 2, 3]})) + sink.close(write=False) + + assert outpath.exists() + assert rotated.exists() + + with NWBHDF5IO(outpath, "r") as io: + nwbfile = io.read() + df0 = nwbfile.intervals["pipeline_settings"].to_dataframe() + assert df0["X.Y.v"].iloc[-1] == "foo" + + with NWBHDF5IO(rotated, "r") as io: + nwbfile = io.read() + df1 = nwbfile.intervals["pipeline_settings"].to_dataframe() + np.testing.assert_array_equal(df1["X.Y.v"].iloc[-1], np.array([1, 2, 3])) + + +def test_rank_change_triggers_rotation(tmp_path): + """1-D → 2-D shape change rotates files.""" + outpath = tmp_path / "rotate_rank.nwb" + rotated = tmp_path / "rotate_rank_01.nwb" + sink = _sink_consumer(outpath) + + sink.write_settings_event(_make_event(seq=1, timestamp=1.0, structured={"v": [1, 2, 3]})) + sink.write_settings_event(_make_event(seq=2, timestamp=2.0, structured={"v": [[1, 2, 3]]})) + sink.close(write=False) + + assert outpath.exists() + assert rotated.exists() + + with NWBHDF5IO(rotated, "r") as io: + nwbfile = io.read() + df1 = nwbfile.intervals["pipeline_settings"].to_dataframe() + np.testing.assert_array_equal(df1["X.Y.v"].iloc[-1], np.array([[1, 2, 3]])) + + +def test_inner_dim_shape_change_triggers_rotation(tmp_path): + """Same rank, different inner-dim shape (e.g. (2,2) → (2,3)) rotates.""" + outpath = tmp_path / "rotate_inner.nwb" + rotated = tmp_path / "rotate_inner_01.nwb" + sink = _sink_consumer(outpath) + + sink.write_settings_event(_make_event(seq=1, timestamp=1.0, structured={"v": [[1, 2], [3, 4]]})) + sink.write_settings_event(_make_event(seq=2, timestamp=2.0, structured={"v": [[1, 2, 3], [4, 5, 6]]})) + sink.close(write=False) + + assert outpath.exists() + assert rotated.exists() + + with NWBHDF5IO(rotated, "r") as io: + nwbfile = io.read() + df = nwbfile.intervals["pipeline_settings"].to_dataframe() + np.testing.assert_array_equal(df["X.Y.v"].iloc[-1], np.array([[1, 2, 3], [4, 5, 6]])) + + +# --------------------------------------------------------------------------- +# Multi-component INITIAL snapshot — aggregation via INIT_FINAL sentinel +# --------------------------------------------------------------------------- + + +def _init_final_event(timestamp: float = 0.0) -> PipelineSettingsEvent: + """Build a sentinel PipelineSettingsEvent matching what the producer emits.""" + return PipelineSettingsEvent( + seq=999, + timestamp=timestamp, + component_address=INIT_FINAL_COMPONENT_ADDRESS, + event_type=PipelineSettingsEventType.INITIAL, + repr_value="", + structured_value=None, + ) + + +def test_initial_events_buffer_until_sentinel_then_merge(tmp_path): + """Per-component INITIAL events with disjoint columns should buffer + in memory; the sentinel flushes one merged anchor row that contains + every component's settings — no rotation required.""" + outpath = tmp_path / "initial_buffered.nwb" + rotated = tmp_path / "initial_buffered_01.nwb" + sink = _sink_consumer(outpath) + + sink.write_settings_event( + _make_event( + component="A", + seq=1, + timestamp=1.0, + event_type=PipelineSettingsEventType.INITIAL, + structured={"foo": 1}, + ) + ) + sink.write_settings_event( + _make_event( + component="B", + seq=2, + timestamp=2.0, + event_type=PipelineSettingsEventType.INITIAL, + structured={"bar": "hi"}, + ) + ) + sink.write_settings_event(_init_final_event(timestamp=3.0)) + sink.close(write=False) + + assert outpath.exists() + assert not rotated.exists(), "merged anchor → no rotation" + + with NWBHDF5IO(outpath, "r") as io: + nwbfile = io.read() + df = nwbfile.intervals["pipeline_settings"].to_dataframe() + + # Anchor row + close-flush row, both with the merged columns. + assert "A.foo" in df.columns + assert "B.bar" in df.columns + assert (df["A.foo"] == 1).all() + assert (df["B.bar"] == "hi").all() + # Anchor is at the FIRST INITIAL's timestamp (1.0), not the sentinel's. + assert df["start_time"].iloc[0] == 1.0 + assert df["stop_time"].iloc[0] == 1.0 + # close-flush extends through close. + assert df["start_time"].iloc[-1] == 1.0 + assert df["stop_time"].iloc[-1] >= 1.0 + assert df["updated_component"].iloc[0] == INIT_FINAL_COMPONENT_ADDRESS + + +def test_updated_event_after_buffered_initials_flushes_buffer(tmp_path): + """If an UPDATED event arrives without an INIT_FINAL sentinel + (producer dropped it), the buffer should flush as a merged anchor + before the UPDATED event closes the open interval.""" + outpath = tmp_path / "initial_no_sentinel.nwb" + sink = _sink_consumer(outpath) + + sink.write_settings_event( + _make_event( + component="A", + seq=1, + timestamp=1.0, + event_type=PipelineSettingsEventType.INITIAL, + structured={"x": 10}, + ) + ) + sink.write_settings_event( + _make_event( + component="B", + seq=2, + timestamp=2.0, + event_type=PipelineSettingsEventType.INITIAL, + structured={"y": 20}, + ) + ) + # No sentinel — go straight to an UPDATED. + sink.write_settings_event( + _make_event( + component="A", + seq=3, + timestamp=5.0, + event_type=PipelineSettingsEventType.UPDATED, + structured={"x": 11}, + ) + ) + sink.close(write=False) + + rotated = tmp_path / "initial_no_sentinel_01.nwb" + assert not rotated.exists(), "buffer-flush should not require a rotation" + + with NWBHDF5IO(outpath, "r") as io: + nwbfile = io.read() + df = nwbfile.intervals["pipeline_settings"].to_dataframe() + + # 3 rows: anchor, closed-by-UPDATED, close-flush. + assert len(df) == 3 + assert "A.x" in df.columns + assert "B.y" in df.columns + # Anchor + first interval carry initial values; close-flush has the + # UPDATED value. + assert df["A.x"].tolist() == [10, 10, 11] + assert df["B.y"].tolist() == [20, 20, 20] + + +def test_initial_buffered_then_close_without_sentinel(tmp_path): + """Buffered INITIALs but neither sentinel nor UPDATED ever arrived; + close should still flush them as a merged anchor.""" + outpath = tmp_path / "initial_only_close.nwb" + sink = _sink_consumer(outpath) + + sink.write_settings_event( + _make_event( + component="A", + seq=1, + timestamp=1.0, + event_type=PipelineSettingsEventType.INITIAL, + structured={"foo": 7}, + ) + ) + sink.close(write=False) + + assert outpath.exists() + with NWBHDF5IO(outpath, "r") as io: + nwbfile = io.read() + df = nwbfile.intervals["pipeline_settings"].to_dataframe() + assert "A.foo" in df.columns + assert (df["A.foo"] == 7).all() + + +def test_late_initial_after_anchor_takes_normal_path(tmp_path): + """An INITIAL event arriving AFTER the table is registered (e.g. a + runtime new component) should go through the normal update/rotation + path rather than the buffer.""" + outpath = tmp_path / "late_initial.nwb" + rotated = tmp_path / "late_initial_01.nwb" + sink = _sink_consumer(outpath) + + # Establish the table first via a buffered initial + sentinel. + sink.write_settings_event( + _make_event( + component="A", + seq=1, + timestamp=1.0, + event_type=PipelineSettingsEventType.INITIAL, + structured={"x": 1}, + ) + ) + sink.write_settings_event(_init_final_event(timestamp=2.0)) + + # A new component appears mid-run and announces itself as INITIAL — + # introduces a brand-new column, so this should rotate. + sink.write_settings_event( + _make_event( + component="LATE", + seq=2, + timestamp=5.0, + event_type=PipelineSettingsEventType.INITIAL, + structured={"y": "hi"}, + ) + ) + sink.close(write=False) + + assert outpath.exists() + assert rotated.exists() + + +def test_double_close_idempotent(tmp_path): + """Closing twice after writing settings should not error.""" + outpath = tmp_path / "double_close.nwb" + sink = _sink_consumer(outpath) + sink.write_settings_event(_make_event()) + sink.close(write=False) + sink.close(write=False) + assert outpath.exists() + + +# --------------------------------------------------------------------------- +# End-to-end through the graph: PipelineSettingsTableCollection +# --------------------------------------------------------------------------- + + +class _SinkSettingsPokerSettings(ez.Settings): + sink_filepath: Path + target_recording: bool = False + publish_after_s: float = 0.5 + terminate_after_s: float = 2.0 + + +class _SinkSettingsPoker(ez.Unit): + """Mid-run settings update on the sink, modeled on the Phase 1 + integration test. Triggers an UPDATED event from the producer.""" + + SETTINGS = _SinkSettingsPokerSettings + OUTPUT_SETTINGS = ez.OutputStream(NWBPipelineSettingsSinkSettings) + + @ez.publisher(OUTPUT_SETTINGS) + async def poke(self) -> typing.AsyncGenerator: + s = self.SETTINGS + await asyncio.sleep(s.publish_after_s) + yield ( + self.OUTPUT_SETTINGS, + NWBPipelineSettingsSinkSettings( + filepath=s.sink_filepath, + overwrite_old=True, + recording=s.target_recording, + inc_clock=ReferenceClockType.UNKNOWN, + ), + ) + await asyncio.sleep(s.terminate_after_s) + raise ez.NormalTermination + + +def test_collection_via_graph_records_initial_and_updated(tmp_path): + """End-to-end: PipelineSettingsTableCollection drops into a graph, + records INITIAL + an UPDATED row in the typed-column table.""" + from ezmsg.core.graphserver import GraphServer + + outpath = tmp_path / "collection_e2e.nwb" + outpath.unlink(missing_ok=True) + + graph_address = ("127.0.0.1", _find_free_port()) + + class _Settings(ez.Settings): + coll: PipelineSettingsTableCollectionSettings + poker: _SinkSettingsPokerSettings + + class _Pipeline(ez.Collection): + SETTINGS = _Settings + + COLL = PipelineSettingsTableCollection() + POKER = _SinkSettingsPoker() + + def configure(self) -> None: + self.COLL.apply_settings(self.SETTINGS.coll) + self.POKER.apply_settings(self.SETTINGS.poker) + + def network(self) -> ez.NetworkDefinition: + return ((self.POKER.OUTPUT_SETTINGS, self.COLL.INPUT_SETTINGS),) + + system = _Pipeline( + _Settings( + coll=PipelineSettingsTableCollectionSettings( + producer=PipelineSettingsProducerSettings(graph_address=graph_address), + sink=NWBPipelineSettingsSinkSettings( + filepath=outpath, + overwrite_old=True, + recording=True, + inc_clock=ReferenceClockType.UNKNOWN, + ), + ), + poker=_SinkSettingsPokerSettings( + sink_filepath=outpath, + target_recording=False, + publish_after_s=0.5, + terminate_after_s=1.5, + ), + ) + ) + + server = GraphServer() + server.start(graph_address) + try: + ez.run(SYSTEM=system, graph_address=graph_address) + finally: + server.stop() + + # With INIT_FINAL aggregation in place, the per-component INITIAL + # events buffer in memory and flush as ONE merged anchor row when + # the sentinel arrives — so we expect a single output file (no + # rotation from the snapshot). + base_stem = outpath.stem + files = sorted(tmp_path.glob(f"{base_stem}*.nwb")) + assert len(files) == 1, f"expected single output file; got {[f.name for f in files]}" + + with NWBHDF5IO(str(files[0]), "r") as io: + nwbfile = io.read() + df = nwbfile.intervals["pipeline_settings"].to_dataframe() + + # Anchor row should cover every component in the session. + # Column paths use dots (sanitize_settings_column_name converts the + # graph's "/" separators), so e.g. "SYSTEM.COLL.PUB.target_table". + cols = [c for c in df.columns if c not in ("start_time", "stop_time", "updated_component")] + assert any(".PUB." in c for c in cols), cols + assert any(".SINK." in c for c in cols), cols + assert any(".POKER." in c for c in cols), cols + + # The SINK's recording flag should land True (initial), then flip to + # False after the poker fires. The last row across all writes must + # show False on the SINK's recording column. + sink_recording_cols = [ + c for c in df.columns if c.endswith(".SINK.recording") and "POKER" not in c and "PUB" not in c + ] + assert sink_recording_cols, f"expected a SINK '.recording' column; got {list(df.columns)}" + for col in sink_recording_cols: + assert df[col].iloc[-1] in (False, "False", 0), f"expected {col}=False after poke; got {df[col].iloc[-1]!r}" diff --git a/tests/test_writer.py b/tests/test_writer.py index a99dde9..ad09262 100644 --- a/tests/test_writer.py +++ b/tests/test_writer.py @@ -1,6 +1,7 @@ """Tests for NWB writer module.""" import dataclasses +import json import tempfile from pathlib import Path @@ -8,9 +9,9 @@ import pynwb import pytest from ezmsg.util.messages.axisarray import AxisArray +from pynwb import NWBHDF5IO -from ezmsg.nwb import NWBSinkSettings, ReferenceClockType -from ezmsg.nwb.writer import NWBSinkConsumer +from ezmsg.nwb import NWBSinkConsumer, NWBSinkSettings, ReferenceClockType @pytest.fixture(autouse=True) @@ -480,7 +481,7 @@ def test_expected_series_smoke(tmp_path): With no ``dtype`` field in the yaml, the default is ``float64``. """ expected_yaml = tmp_path / "expected.yaml" - expected_yaml.write_text("PreAlloc:\n" " fs: 100.0\n" " shape: [-1, 4]\n") + expected_yaml.write_text("PreAlloc:\n fs: 100.0\n shape: [-1, 4]\n") nwb_path = tmp_path / "expected.nwb" sink = NWBSinkConsumer( settings=NWBSinkSettings( @@ -570,3 +571,137 @@ def test_sample_trigger_routes_to_epochs(): assert "trigger_a" in labels path.unlink(missing_ok=True) + + +# -- annotation series writer -- + + +def test_write_annotation_creates_series_and_appends(tmp_path): + """First annotation write should create the AnnotationSeries; subsequent + writes should append rows.""" + outpath = tmp_path / "ezmsg_nwb_annotation_test.nwb" + sink = NWBSinkConsumer( + settings=NWBSinkSettings( + filepath=outpath, + overwrite_old=True, + inc_clock=ReferenceClockType.UNKNOWN, + ) + ) + + sink.write_annotation("settings_annotations", timestamp=1.0, data='{"k": "a"}') + sink.write_annotation("settings_annotations", timestamp=2.0, data='{"k": "b"}') + sink.close(write=False) + + with NWBHDF5IO(outpath, "r") as io: + nwbfile = io.read() + series = nwbfile.acquisition["settings_annotations"] + assert list(series.data[:]) == ['{"k": "a"}', '{"k": "b"}'] + np.testing.assert_array_equal(series.timestamps[:], np.array([1.0, 2.0])) + + +def test_write_annotation_keeps_file_with_no_acquisition_data(tmp_path): + """A file with only annotations (no acquisition) should not be deleted on close.""" + outpath = tmp_path / "ezmsg_nwb_annotation_only_test.nwb" + sink = NWBSinkConsumer( + settings=NWBSinkSettings( + filepath=outpath, + overwrite_old=True, + inc_clock=ReferenceClockType.UNKNOWN, + ) + ) + sink.write_annotation("settings_annotations", timestamp=1.0, data="row") + sink.close(write=False) + assert outpath.exists() + + +def test_write_annotation_after_data_uses_data_baseline(tmp_path): + """Annotation timestamps stored in-file should be relative to the data anchor. + + With UNKNOWN clock, ``start_timestamp`` is latched off the first data + message's axis offset (here, 0). An annotation at wall-clock 1.0 should + therefore land at relative time 1.0 in the file. + """ + outpath = tmp_path / "ezmsg_nwb_annotation_after_data_test.nwb" + sink = NWBSinkConsumer( + settings=NWBSinkSettings( + filepath=outpath, + overwrite_old=True, + inc_clock=ReferenceClockType.UNKNOWN, + ) + ) + + sink._process( + AxisArray( + data=np.arange(6, dtype=float).reshape(3, 2), + dims=["time", "ch"], + axes={"time": AxisArray.TimeAxis(fs=100.0)}, + key="sig", + ) + ) + sink.write_annotation("settings_annotations", timestamp=1.0, data="row") + sink.close(write=False) + + with NWBHDF5IO(outpath, "r") as io: + nwbfile = io.read() + ts = nwbfile.acquisition["settings_annotations"].timestamps[:] + np.testing.assert_array_equal(ts, np.array([1.0])) + + +def test_write_annotation_multiple_tables_coexist(tmp_path): + """Two distinct table_names should produce two distinct AnnotationSeries.""" + outpath = tmp_path / "ezmsg_nwb_annotation_multi_test.nwb" + sink = NWBSinkConsumer( + settings=NWBSinkSettings( + filepath=outpath, + overwrite_old=True, + inc_clock=ReferenceClockType.UNKNOWN, + ) + ) + + sink.write_annotation("settings_annotations", timestamp=1.0, data="A") + sink.write_annotation("user_notes", timestamp=2.0, data="B") + sink.write_annotation("settings_annotations", timestamp=3.0, data="C") + sink.close(write=False) + + with NWBHDF5IO(outpath, "r") as io: + nwbfile = io.read() + assert list(nwbfile.acquisition["settings_annotations"].data[:]) == ["A", "C"] + assert list(nwbfile.acquisition["user_notes"].data[:]) == ["B"] + + +def test_write_annotation_serializes_json_payload(tmp_path): + """Sanity check: a self-describing JSON payload round-trips through the file.""" + outpath = tmp_path / "ezmsg_nwb_annotation_json_test.nwb" + sink = NWBSinkConsumer( + settings=NWBSinkSettings( + filepath=outpath, + overwrite_old=True, + inc_clock=ReferenceClockType.UNKNOWN, + ) + ) + + payload = json.dumps({"component": "X.Y", "event_type": "INITIAL", "seq": 0, "settings": {"foo": 1}}) + sink.write_annotation("settings_annotations", timestamp=0.5, data=payload) + sink.close(write=False) + + with NWBHDF5IO(outpath, "r") as io: + nwbfile = io.read() + round_tripped = json.loads(nwbfile.acquisition["settings_annotations"].data[0]) + assert round_tripped["component"] == "X.Y" + assert round_tripped["settings"] == {"foo": 1} + + +def test_close_is_idempotent_after_annotation(tmp_path): + """Closing twice after writing annotations should not error.""" + outpath = tmp_path / "ezmsg_nwb_annotation_double_close_test.nwb" + sink = NWBSinkConsumer( + settings=NWBSinkSettings( + filepath=outpath, + overwrite_old=True, + inc_clock=ReferenceClockType.UNKNOWN, + ) + ) + sink.write_annotation("settings_annotations", timestamp=1.0, data="row") + sink.close(write=False) + sink.close(write=False) + assert outpath.exists()