diff --git a/pyproject.toml b/pyproject.toml index eccff52..f6a7ee5 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -10,7 +10,7 @@ requires-python = ">=3.10" dynamic = ["version"] dependencies = [ "ezmsg>=3.9.0", - "ezmsg-baseproc", + "ezmsg-baseproc>=1.9.0", "numpy>=1.26.0", "pynwb", "h5py", diff --git a/src/ezmsg/nwb/clockdriven.py b/src/ezmsg/nwb/clockdriven.py index cd8fd5a..3815165 100644 --- a/src/ezmsg/nwb/clockdriven.py +++ b/src/ezmsg/nwb/clockdriven.py @@ -2,6 +2,7 @@ from __future__ import annotations +import asyncio import os import typing @@ -59,6 +60,21 @@ class NWBClockDrivenProducer(BaseClockDrivenProducer[NWBClockDrivenSettings, NWB on :class:`NWBClockDrivenSettings`. """ + # Fields that can change without forcing a slicer teardown + reopen. The + # producer rebinds ``self.settings`` live and, for ``start_offset``, + # follows up with :meth:`seek`. Anything else — filepath, stream_key, + # reference_clock, reref_now, fs — routes through ``_request_reset`` and + # the next clock tick triggers a fresh :meth:`_reset_state`. + NONRESET_SETTINGS_FIELDS = frozenset({"start_offset", "playback_rate", "n_time"}) + + def update_settings(self, new_settings: NWBClockDrivenSettings) -> None: + old_offset = self.settings.start_offset + super().update_settings(new_settings) + # If super() queued a reset, _reset_state will re-seek from the new + # start_offset; only seek on the hot path. + if self._hash != -1 and new_settings.start_offset != old_offset: + self.seek(new_settings.start_offset) + @property def exhausted(self) -> bool: if self._state.slicer is None: @@ -67,6 +83,15 @@ def exhausted(self) -> bool: return False # Time-window streams don't track exhaustion via index return self._state.sample_idx >= self._state.n_total_samples + async def _areset_state(self, time_axis: LinearAxis) -> None: + """Offload the slow ``NWBSlicer`` open onto a worker thread so that + sibling subscribers (notably ``INPUT_SETTINGS`` for live scrubbing) + can keep running while the file is opened. First-message NWB I/O + can take 1.5–4 seconds per the writer-module docstring; without + this, the unit's event loop is held for the whole duration. + """ + await asyncio.to_thread(self._reset_state, time_axis) + def _reset_state(self, time_axis: LinearAxis) -> None: if self._state.slicer is not None: self._state.slicer.close() @@ -74,10 +99,9 @@ def _reset_state(self, time_axis: LinearAxis) -> None: # Idle mode: no file / stream configured yet. Leave the slicer as # None so ``_process`` short-circuits. A later settings push that - # sets ``filepath`` lands in the ``_RESET_FIELDS`` path in - # ``NWBClockDrivenUnit.on_settings``, which calls - # ``create_processor`` — a fresh producer then reruns this method - # and loads the file. + # sets ``filepath`` / ``stream_key`` falls outside + # ``NONRESET_SETTINGS_FIELDS``, so ``update_settings`` queues a reset + # and the next clock tick reruns this method and loads the file. if not self.settings.filepath or not self.settings.stream_key: self._state.template = None self._state.has_timestamps = False @@ -231,45 +255,7 @@ def __del__(self): self._state.slicer.close() -# Fields whose change semantically forces a full reopen of the NWB file -# (slicer teardown + rediscovery) and a fresh state. Anything else — start -# offset, playback rate, chunk size — is hot-updated so playback keeps its -# current position across scrubbing / speed changes / pause-resume. -_RESET_FIELDS: frozenset[str] = frozenset({"filepath", "stream_key", "reference_clock", "reref_now", "fs"}) - - class NWBClockDrivenUnit(BaseClockDrivenUnit[NWBClockDrivenSettings, NWBClockDrivenProducer]): """Clock-driven NWB unit that reads one stream synchronized to a shared clock.""" SETTINGS = NWBClockDrivenSettings - - @ez.subscriber(BaseClockDrivenUnit.INPUT_SETTINGS) - async def on_settings(self, msg: NWBClockDrivenSettings) -> None: - """Apply new settings, preserving playback position when possible. - - The base class unconditionally recreates the processor on every - settings push, which for us means closing and reopening the NWB - file and resetting the playback cursor to ``start_offset``. That's - the right behaviour when the file or stream changes, but not for - the common GUI interactions — scrubbing, speed changes, pause — - which just want to tweak position or emission rate without losing - context. We only recreate when a ``_RESET_FIELDS`` field actually - changed, and otherwise hot-swap ``self.processor.settings`` plus - issue a ``seek`` when ``start_offset`` moved. - """ - prev: NWBClockDrivenSettings = self.SETTINGS - processor: NWBClockDrivenProducer | None = getattr(self, "processor", None) - - needs_full_reset = processor is None or any(getattr(prev, f) != getattr(msg, f) for f in _RESET_FIELDS) - - self.apply_settings(msg) - - if needs_full_reset: - self.create_processor() - return - - # Hot-update path: swap settings on the live producer and seek if - # the user moved the playback cursor. - processor.settings = msg - if prev.start_offset != msg.start_offset: - processor.seek(msg.start_offset) diff --git a/src/ezmsg/nwb/iterator.py b/src/ezmsg/nwb/iterator.py index e74812e..77e4fb8 100644 --- a/src/ezmsg/nwb/iterator.py +++ b/src/ezmsg/nwb/iterator.py @@ -1,5 +1,6 @@ from __future__ import annotations +import asyncio import os import typing from collections import deque @@ -36,7 +37,6 @@ class NWBIteratorState: class NWBAxisArrayIterator(BaseStatefulProducer[NWBIteratorSettings, AxisArray, NWBIteratorState]): - def __init__(self, *args, **kwargs): super().__init__(*args, **kwargs) # Eagerly initialize state (load NWB file metadata) so that @@ -48,6 +48,14 @@ def __init__(self, *args, **kwargs): def exhausted(self) -> bool: return self._state.chunk_ix >= self._state.n_chunks and not self._state.deque + async def _areset_state(self) -> None: + """Offload the slow ``NWBSlicer`` open and chunk-table build onto + a worker thread so the unit's event loop can keep servicing other + async tasks during the multi-second first-open. See the matching + override on ``NWBClockDrivenProducer`` for context. + """ + await asyncio.to_thread(self._reset_state) + def _reset_state(self) -> None: self._state.n_chunks = 0 self._state.chunk_ix = 0 diff --git a/src/ezmsg/nwb/writer.py b/src/ezmsg/nwb/writer.py index 7aa126c..a67f865 100644 --- a/src/ezmsg/nwb/writer.py +++ b/src/ezmsg/nwb/writer.py @@ -66,12 +66,13 @@ """ import asyncio +import dataclasses import datetime import os import re import time import typing -from collections import defaultdict +import warnings from pathlib import Path from uuid import uuid4 @@ -79,7 +80,7 @@ import h5py import numpy as np import pynwb -from ezmsg.baseproc import BaseConsumer, BaseConsumerUnit +from ezmsg.baseproc import BaseConsumerUnit, BaseStatefulConsumer, processor_state from ezmsg.util.messages.axisarray import AxisArray from ezmsg.util.messages.util import replace from hdmf.backends.hdf5.h5_utils import H5DataIO @@ -106,46 +107,99 @@ class NWBSinkSettings(ez.Settings): expected_series: typing.Optional[typing.Union[str, os.PathLike]] = None -class NWBSinkConsumer(BaseConsumer[NWBSinkSettings, AxisArray]): +@dataclasses.dataclass +class SeriesState: + """Per-stream bookkeeping for one NWB series. + + ``shape`` is the trailing (non-time) shape used by + :meth:`NWBSinkConsumer._check_msg_consistency`. ``data`` and ``ts`` hold + the live ``h5py.Dataset`` (or pynwb container) references used by + :meth:`NWBSinkConsumer._process` to append rows. ``bytes_written`` + tracks cumulative append size for the file-split threshold. + """ + + shape: typing.Tuple[int, ...] + data: typing.Any = None + ts: typing.Any = None + bytes_written: int = 0 + + +@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 + nwbfile: typing.Optional[pynwb.NWBFile] = None + series: typing.Optional[typing.Dict[str, SeriesState]] = None + start_timestamp: float = 0.0 + split_count: int = 0 + + +class NWBSinkConsumer(BaseStatefulConsumer[NWBSinkSettings, AxisArray, NWBSinkState]): # Session start datetime. It should have a valid timezone and that should be UTC. shared_start_datetime: typing.Optional[datetime.datetime] = None # Session start time.time. It does not have a timezone but unqualified conversions assume local time. shared_t0: typing.Optional[float] = None shared_clock_type: typing.Optional[ReferenceClockType] = None + # Fields that can change without reconstructing the NWB file. Both are + # read live from ``self.settings`` inside ``_process``. Everything else + # (filepath, axis, inc_clock, meta_yaml, expected_series, overwrite_old) + # is consumed in ``_reset_state`` side effects; ``update_settings`` + # requests a reset so the next message closes the old file and reopens. + NONRESET_SETTINGS_FIELDS = frozenset({"recording", "split_bytes"}) + def __init__(self, *args, settings: typing.Optional[NWBSinkSettings] = None, **kwargs): super().__init__(*args, settings=settings, **kwargs) - - self._filepath = Path(self.settings.filepath) - self._overwrite_old = self.settings.overwrite_old - self._axis = self.settings.axis - self._recording = self.settings.recording - self._inc_clock = self.settings.inc_clock - self._meta_yaml = self.settings.meta_yaml - self._split_bytes = self.settings.split_bytes - - self._start_timestamp: float = 0.0 - self._split_count: int = 0 - self._stream_bytes = defaultdict(lambda: 0) + # ``_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. self._current_msg: typing.Optional[AxisArray] = None - self._io: typing.Optional[pynwb.NWBHDF5IO] = None - self._nwbfile: typing.Optional[pynwb.NWBFile] = None - self._datasets: typing.Dict[str, typing.Dict[str, typing.Any]] = {} + # Eagerly open the file so construction-time errors (file exists, + # permission denied, bad metadata yaml) surface immediately rather + # than on the first message. Match _hash_message() so the first + # inbound message does not re-trigger reset. + self._reset_state(None) + self._hash = 0 + + def _hash_message(self, message: typing.Optional[AxisArray]) -> int: + # Reset is driven exclusively by settings changes via + # ``_request_reset``; message identity does not force a rebuild. + return 0 + + def _reset_state(self, message: typing.Optional[AxisArray]) -> None: + """(Re)open the NWB file using the current settings. + + Called at construction (with ``None``) and again after + ``update_settings`` flags a reset, which happens whenever any non- + ``NONRESET_SETTINGS_FIELDS`` field changes. In the reset case the + prior file is flushed and closed before the new one opens. + """ + # Settings-triggered reset: flush the prior file before rebuilding. + if self._state.io is not None: + self.close(write=True) - # Normalize filepath and delete existing file if enabled. - self._check_filepath() + self._state.filepath = Path(self.settings.filepath) + self._state.series = {} + self._state.start_timestamp = 0.0 + self._state.split_count = 0 - # Create the self._nwbfile and self._io objects. Note: Nothing written to disk yet! + self._check_filepath() self._nwb_create_or_fail() if self.settings.expected_series is not None and Path(self.settings.expected_series).expanduser().exists(): expected_series = Path(self.settings.expected_series).expanduser() meta = load_dict_from_file(expected_series) _ = self.get_session_datetime(None) - self._start_timestamp = self.get_session_timestamp(None) + self._state.start_timestamp = self.get_session_timestamp(None) self._prep_from_meta(meta) def __del__(self): + if not hasattr(self, "_state"): + return self.close(write=False, log=False) async def _aprocess(self, message: AxisArray) -> None: @@ -159,51 +213,59 @@ def _process(self, message: AxisArray) -> None: if _HAS_SAMPLE_TRIGGER and isinstance(self._current_msg, SampleTriggerMessage): # SampleTriggerMessage. Rewrite as AxisArray. timestamp = self._current_msg.timestamp - if hasattr(self._current_msg, "period") and len(self._current_msg.period) > 0: - timestamp = timestamp + self._current_msg.period[0] + 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"], - axes={"time": AxisArray.Axis(gain=1.0, offset=timestamp)}, + 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: - targ_ax_ix = self._current_msg.get_axis_idx(self._axis) + 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=[self._axis] + self._current_msg.dims[:targ_ax_ix] + self._current_msg.dims[targ_ax_ix + 1 :], + dims=[axis] + self._current_msg.dims[:targ_ax_ix] + self._current_msg.dims[targ_ax_ix + 1 :], ) # Is this a new series? - b_new = self._io is None - b_new = b_new or self._current_msg.key not in self._datasets + 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 datasets but properties do not match previous dataset properties + # 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(): - b_final_write = hasattr(self._nwbfile, "epochs") and self._nwbfile.epochs is not None - b_final_write = b_final_write or (hasattr(self._nwbfile, "trials") and self._nwbfile.trials is not None) + 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 self._axis in ["time", "win"] or "time" in self._current_msg.axes: - targ_dim = self._axis if self._axis in ["time", "win"] else "time" + 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._start_timestamp = self.get_session_timestamp(t0) - if self._inc_clock == ReferenceClockType.MONOTONIC: - self._start_timestamp += time.monotonic() - time.time() + 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"]: @@ -216,36 +278,40 @@ def _process(self, message: AxisArray) -> None: self._flush_io(reopen=True) self._update_rate_for_current() - if self._recording and self._current_msg.data.size: + if self.settings.recording and self._current_msg.data.size: + axis = self.settings.axis timestamps = None - if self._axis in ["time", "win"] or "time" in self._current_msg.axes: - targ_dim = self._axis if self._axis in ["time", "win"] else "time" + 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._start_timestamp + timestamps = time_ax.data - self._state.start_timestamp else: timestamps = (np.arange(len(self._current_msg.data)) * time_ax.gain) + ( - time_ax.offset - self._start_timestamp + time_ax.offset - self._state.start_timestamp ) - if self._current_msg.key in ["epochs", "trials"]: - self._append_events(self._current_msg.key, timestamps, self._current_msg.data) + 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 = self._datasets[self._current_msg.key]["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 - self._stream_bytes[self._current_msg.key] += self._current_msg.data.nbytes + series_state.bytes_written += self._current_msg.data.nbytes # Write timestamps if timestamps is not None: - ts = self._datasets[self._current_msg.key]["ts"] + ts = series_state.ts ts.resize(len(ts) + len(timestamps), axis=0) ts[-len(timestamps) :] = timestamps - self._stream_bytes[self._current_msg.key] += timestamps.nbytes + series_state.bytes_written += timestamps.nbytes - if 0 < self._split_bytes <= sum(self._stream_bytes.values()) and "%d" not in str(self._filepath): - self._split_count += 1 + 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() @@ -254,12 +320,12 @@ def _process(self, message: AxisArray) -> None: @property def path_on_disk(self) -> Path: - fp = Path(self._filepath) - if self._split_bytes > 0: + 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._split_count:02}" + fp.suffix) + return fp.parent / (fp.stem + f"_{self._state.split_count:02}" + fp.suffix) else: return fp @@ -275,20 +341,24 @@ def get_session_datetime(self, try_t0: typing.Optional[float] = None) -> datetim Returns: Common session starttime among all instances of this class """ - if self.__class__.shared_clock_type is not None and self.__class__.shared_clock_type != self._inc_clock: + inc_clock = self.settings.inc_clock + if self.__class__.shared_clock_type is not None and self.__class__.shared_clock_type != inc_clock: raise ValueError( - f"All instances must share the same clock type. {self._inc_clock} != {self.__class__.shared_clock_type}" + f"All instances must share the same clock type. {inc_clock} != {self.__class__.shared_clock_type}" ) if self.__class__.shared_start_datetime is None: - if try_t0 is not None and self._inc_clock in [ + if try_t0 is not None and inc_clock in [ ReferenceClockType.SYSTEM, ReferenceClockType.MONOTONIC, ]: - if self._inc_clock == ReferenceClockType.MONOTONIC: + if inc_clock == ReferenceClockType.MONOTONIC: try_t0 = try_t0 - time.monotonic() + time.time() self.__class__.shared_start_datetime = datetime.datetime.fromtimestamp(try_t0, datetime.timezone.utc) else: self.__class__.shared_start_datetime = datetime.datetime.now(datetime.timezone.utc) + # Latch the clock type so subsequent instances are forced to + # match — otherwise the mismatch check above never fires. + self.__class__.shared_clock_type = inc_clock return self.__class__.shared_start_datetime def get_session_timestamp(self, try_t0: typing.Optional[float] = None) -> float: @@ -301,12 +371,13 @@ def get_session_timestamp(self, try_t0: typing.Optional[float] = None) -> float: Returns: Common session timestamp among all instances of this class. """ - if self.__class__.shared_clock_type is not None and self.__class__.shared_clock_type != self._inc_clock: + inc_clock = self.settings.inc_clock + if self.__class__.shared_clock_type is not None and self.__class__.shared_clock_type != inc_clock: raise ValueError( - f"All instances must share the same clock type. {self._inc_clock} != {self.__class__.shared_clock_type}" + f"All instances must share the same clock type. {inc_clock} != {self.__class__.shared_clock_type}" ) if self.__class__.shared_t0 is None: - if try_t0 is None or self._inc_clock in [ + if try_t0 is None or inc_clock in [ ReferenceClockType.SYSTEM, ReferenceClockType.MONOTONIC, ]: @@ -317,21 +388,27 @@ def get_session_timestamp(self, try_t0: typing.Optional[float] = None) -> float: "Clock type is UNKNOWN. Timestamps are relative to the first incoming timestamp " "but this value is NOT recoverable as it is not stored in the NWB file." ) + # Latch the clock type if get_session_datetime didn't already + # (the UNKNOWN branch above bypasses it). + if self.__class__.shared_clock_type is None: + self.__class__.shared_clock_type = inc_clock return self.__class__.shared_t0 def _check_filepath(self) -> None: """ - Check self._filepath. Update path if necessary. Check if the path already exists and potentially raise - an error if overwriting is disabled. + Normalize ``self._state.filepath`` (reading the raw path from + ``self.settings.filepath``). If the resolved path already exists, + delete it when ``overwrite_old`` is enabled or raise otherwise. """ _suffix = ".nwb" - if self._filepath.name.startswith("."): + filepath = Path(self._state.filepath) + if filepath.name.startswith("."): raise FileNotFoundError( - f"filepath {self._filepath} name begins with `.` -- cannot discriminate name from extension." + f"filepath {filepath} name begins with `.` -- cannot discriminate name from extension." ) - filepath = Path(self._filepath).expanduser() + filepath = filepath.expanduser() # If provided path is merely a directory then create a new filename. is_dir = (isinstance(filepath, Path) and filepath.is_dir()) or ( @@ -350,35 +427,36 @@ def _check_filepath(self) -> None: if not filepath.suffix: filepath = filepath.parent / (filepath.name + _suffix) - self._filepath = filepath + self._state.filepath = filepath if self.path_on_disk.exists(): age = (time.time() - os.path.getctime(self.path_on_disk)) / 60 ez.logger.info(f"File at {self.path_on_disk} is {age:.2f} minutes old.") - if self._overwrite_old: + if self.settings.overwrite_old: ez.logger.info("File will be overwritten.") self.path_on_disk.unlink(missing_ok=False) else: msg = "File exists but overwriting is disabled. Set overwrite_old=True to overwrite." ez.logger.error(msg) raise ValueError(msg) - self._filepath = filepath def _read_meta_dict(self) -> typing.Union[typing.Mapping, dict]: """ - Load the metadata from self._meta_yaml if that path exists, else load it from the default location. + Load the metadata from ``self.settings.meta_yaml`` if that path + exists, else load it from the default location. Returns: A dict containing the metadata for this NWB file. """ - if self._meta_yaml is None or not Path(self._meta_yaml).expanduser().exists(): + meta_yaml = self.settings.meta_yaml + if meta_yaml is None or not Path(meta_yaml).expanduser().exists(): default_path = Path(__file__).parent meta_dict = DeepDict() for yaml_name in ["nwb_metadata", "nwb_session"]: yaml_path = default_path / f"{yaml_name}.yaml" meta_dict = dict_deep_update(meta_dict, load_dict_from_file(yaml_path)) else: - yaml_path = Path(self._meta_yaml).expanduser() + yaml_path = Path(meta_yaml).expanduser() meta_dict = load_dict_from_file(yaml_path) return meta_dict @@ -433,7 +511,7 @@ 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._nwbfile.acquisition[key] + series = self._state.nwbfile.acquisition[key] series.timestamps.attrs["rate"] = ss["fs"] def close(self, write=False, log=True) -> None: @@ -445,49 +523,67 @@ def close(self, write=False, log=True) -> None: log: Set True to log the closing and deletion of the file. This must be kept False when calling from __del__. """ - if self._io is not None: - if write: - self._io.write(self._nwbfile) - src_str = f"{self._io.source}" - b_delete = sum(self._stream_bytes.values()) == 0 - for key in ["epochs", "trials"]: - if hasattr(self._nwbfile, key) and getattr(self._nwbfile, key) is not None: - b_delete = b_delete and len(getattr(self._nwbfile, key)) == 1 # EZNWB-START - self._io.close() - del self._nwbfile - del self._io - self._nwbfile = None - self._io = None + state = getattr(self, "_state", None) + if state is None or state.io is None: + return + 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 + io.close() + state.nwbfile = None + state.io = None + state.series = {} + 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"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}.") + ez.logger.info(f"Deleted empty file at {src_str}.") def toggle_recording(self, recording: typing.Optional[bool] = None): - self._recording = recording if recording is not None else not self._recording + """Deprecated. Send a ``NWBSinkSettings`` update with the desired + ``recording`` value instead — the update is routed through + :meth:`update_settings` and takes effect on the next message. + """ + warnings.warn( + "NWBSinkConsumer.toggle_recording is deprecated; publish a " + "NWBSinkSettings update with the desired `recording` value " + "instead.", + DeprecationWarning, + stacklevel=2, + ) + new_value = recording if recording is not None else not self.settings.recording + self.update_settings(dataclasses.replace(self.settings, recording=new_value)) def _check_msg_consistency(self) -> bool: + axis = self.settings.axis key = self._current_msg.key - in_ax = self._current_msg.axes[self._axis] + series_state = self._state.series[key] + in_ax = self._current_msg.axes[axis] b_rate_change = ( - self._axis in self._current_msg.axes - and "ts" in self._datasets[key] + axis in self._current_msg.axes + and series_state.ts is not None and not hasattr(in_ax, "data") - and self._datasets[key]["ts"].attrs["rate"] != 1 / in_ax.gain + and series_state.ts.attrs["rate"] != 1 / in_ax.gain ) - b_shape_change = self._datasets[key]["shape"] != self._current_msg.data.shape[1:] + b_shape_change = series_state.shape != self._current_msg.data.shape[1:] return not (b_rate_change or b_shape_change) def _update_rate_for_current(self): - if self._axis in ["time", "win"]: - time_ax = self._current_msg.axes[self._axis] + axis = self.settings.axis + if axis in ["time", "win"]: + time_ax = self._current_msg.axes[axis] if hasattr(time_ax, "data"): rate = 0.0 else: rate = 1 / time_ax.gain if time_ax.gain != 0 else 0 - self._datasets[self._current_msg.key]["ts"].attrs["rate"] = rate + self._state.series[self._current_msg.key].ts.attrs["rate"] = rate def _copy_nwb(self) -> typing.Tuple[pynwb.NWBFile, dict]: copy_keys = [ @@ -510,19 +606,20 @@ def _copy_nwb(self) -> typing.Tuple[pynwb.NWBFile, dict]: "stimulus_notes", "lab", ] - new_nwb_kwargs = {k: getattr(self._nwbfile, k) for k in copy_keys if hasattr(self._nwbfile, k)} + old_nwbfile = self._state.nwbfile + new_nwb_kwargs = {k: getattr(old_nwbfile, k) for k in copy_keys if hasattr(old_nwbfile, k)} new_nwb_kwargs["keywords"] = ( - self._nwbfile.keywords if isinstance(self._nwbfile.keywords, list) else self._nwbfile.keywords[:].tolist() + old_nwbfile.keywords if isinstance(old_nwbfile.keywords, list) else old_nwbfile.keywords[:].tolist() ) nwbfile = pynwb.NWBFile(identifier=str(uuid4()), **new_nwb_kwargs) - nwbfile.subject = pynwb.file.Subject(**self._nwbfile.subject.fields) + nwbfile.subject = pynwb.file.Subject(**old_nwbfile.subject.fields) meta = {} for key in ["epochs", "trials"]: - if hasattr(self._nwbfile, key) and getattr(self._nwbfile, key) is not None: + if hasattr(old_nwbfile, key) and getattr(old_nwbfile, key) is not None: meta[key] = {"fs": 0.0, "shape": (0, 1)} - for key, ds in self._datasets.items(): + for key, ss in self._state.series.items(): if key not in ["epochs", "trials"]: - meta[key] = {"fs": ds["ts"].attrs["rate"], "shape": (0,) + ds["shape"]} + meta[key] = {"fs": ss.ts.attrs["rate"], "shape": (0,) + ss.shape} return nwbfile, meta @@ -544,20 +641,21 @@ def _nwb_create_or_fail(self, nwbfile: typing.Optional[pynwb.NWBFile] = None) -> if "Subject" in meta_dict: nwbfile.subject = pynwb.file.Subject(**meta_dict["Subject"]) - if "%d" in str(self._filepath): + if "%d" in str(self._state.filepath): io_file = h5py.File( - name=self._filepath, + name=self._state.filepath, mode="w", driver="family", - memb_size=self._split_bytes, + memb_size=self.settings.split_bytes, ) io = pynwb.NWBHDF5IO(file=io_file, mode="w") else: io = pynwb.NWBHDF5IO(self.path_on_disk, "w") - self._io = io - self._nwbfile = nwbfile - self._stream_bytes = defaultdict(lambda: 0) + self._state.io = io + self._state.nwbfile = nwbfile + # Fresh series map; prior entries (if any) belong to the closed file. + self._state.series = {} def _flush_io(self, reopen: bool = True): """ @@ -567,33 +665,33 @@ def _flush_io(self, reopen: bool = True): * enable appending to our epochs/trials table (but only after it has an entry). * create the appendable datasets for our continuous data. """ - self._io.write(self._nwbfile) + self._state.io.write(self._state.nwbfile) if reopen: - if self._io: - self._io.close() - if "%d" in str(self._filepath): + if self._state.io: + self._state.io.close() + if "%d" in str(self._state.filepath): io_file = h5py.File( - name=self._filepath, + name=self._state.filepath, mode="a", driver="family", - memb_size=self._split_bytes, + memb_size=self.settings.split_bytes, ) io = pynwb.NWBHDF5IO(file=io_file, mode="a") else: io = pynwb.NWBHDF5IO(self.path_on_disk, "a") - self._io = io - self._nwbfile = self._io.read() + self._state.io = io + self._state.nwbfile = self._state.io.read() # Get references to our continuous datasets - for k, v in self._datasets.items(): - if k in self._nwbfile.acquisition: - series = self._nwbfile.acquisition[k] + for k, ss in self._state.series.items(): + if k in self._state.nwbfile.acquisition: + series = self._state.nwbfile.acquisition[k] if isinstance(series.data, H5DataIO): - v["data"] = series.data.dataset - v["ts"] = series.timestamps.dataset + ss.data = series.data.dataset + ss.ts = series.timestamps.dataset else: - v["data"] = series.data - v["ts"] = series.timestamps + ss.data = series.data + ss.ts = series.timestamps def _append_events( self, @@ -601,7 +699,8 @@ def _append_events( timestamps: typing.Iterable[float], data: typing.Iterable[typing.Iterable[str]], ): - fun = {"epochs": self._nwbfile.add_epoch, "trials": self._nwbfile.add_trial}[key] + nwbfile = self._state.nwbfile + fun = {"epochs": nwbfile.add_epoch, "trials": nwbfile.add_trial}[key] for ev_t, ev_str in zip(timestamps, data): fun(start_time=ev_t, stop_time=ev_t + 0, **{"label": ",".join(ev_str)}) @@ -611,15 +710,16 @@ def _prep_event_io(self): """ colname = "label" key = self._current_msg.key + nwbfile = self._state.nwbfile fun = { - "epochs": self._nwbfile.add_epoch_column, - "trials": self._nwbfile.add_trial_column, + "epochs": nwbfile.add_epoch_column, + "trials": nwbfile.add_trial_column, }[key] fun(name=colname, description=f"{colname} {key}") - self._datasets[key] = {"shape": self._current_msg.data.shape[1:]} + self._state.series[key] = SeriesState(shape=self._current_msg.data.shape[1:]) - table = {"epochs": self._nwbfile.epochs, "trials": self._nwbfile.trials}[key] + table = {"epochs": nwbfile.epochs, "trials": nwbfile.trials}[key] table.id.set_data_io(H5DataIO, {"maxshape": (None,)}) table.start_time.set_data_io(H5DataIO, {"maxshape": (None,)}) table.stop_time.set_data_io(H5DataIO, {"maxshape": (None,)}) @@ -637,9 +737,10 @@ def _prep_continuous_io(self): channel info is provided). """ key = self._current_msg.key - targ_ax_ix = self._current_msg.get_axis_idx(self._axis) + nwbfile = self._state.nwbfile + targ_ax_ix = self._current_msg.get_axis_idx(self.settings.axis) shape = self._current_msg.data.shape[:targ_ax_ix] + self._current_msg.data.shape[targ_ax_ix + 1 :] - self._datasets[key] = {"shape": shape} + self._state.series[key] = SeriesState(shape=shape) dataio = H5DataIO( shape=(0,) + shape, dtype=self._current_msg.data.dtype, @@ -661,21 +762,21 @@ def _prep_continuous_io(self): and hasattr(self._current_msg.axes["ch"], "data") and len(self._current_msg.axes["ch"].data) ): - b_first = self._nwbfile.electrodes is None or "label" not in self._nwbfile.electrodes.colnames + b_first = nwbfile.electrodes is None or "label" not in nwbfile.electrodes.colnames if b_first: - self._nwbfile.add_electrode_column(name="label", description="electrode label") + nwbfile.add_electrode_column(name="label", description="electrode label") dev_name = "unified device" - if dev_name in self._nwbfile.devices: - device = self._nwbfile.devices[dev_name] + if dev_name in nwbfile.devices: + device = nwbfile.devices[dev_name] else: - device = self._nwbfile.create_device(name=dev_name, description="created by ezmsg nwbsink") + device = nwbfile.create_device(name=dev_name, description="created by ezmsg nwbsink") el_grp_name = "unified electrode group" - if el_grp_name in self._nwbfile.electrode_groups: - el_grp = self._nwbfile.electrode_groups[el_grp_name] + if el_grp_name in nwbfile.electrode_groups: + el_grp = nwbfile.electrode_groups[el_grp_name] else: - el_grp = self._nwbfile.create_electrode_group( + el_grp = nwbfile.create_electrode_group( name=el_grp_name, description="electrode group created by ezmsg nwbsink", device=device, @@ -683,21 +784,23 @@ def _prep_continuous_io(self): ) if not b_first: self._flush_io(reopen=True) + # _flush_io swaps nwbfile out from under us; refresh. + nwbfile = self._state.nwbfile - el_df = self._nwbfile.electrodes.to_dataframe() + el_df = nwbfile.electrodes.to_dataframe() el_df = el_df[el_df["group"] == el_grp] for ll in self._current_msg.axes["ch"].data: if ll not in el_df["label"].values: - self._nwbfile.add_electrode(label=ll, location="unknown", group=el_grp) + nwbfile.add_electrode(label=ll, location="unknown", group=el_grp) - if type(self._nwbfile.electrodes.id.data) is list: + if type(nwbfile.electrodes.id.data) is list: for fn in ["id", "location", "group_name", "group", "label"]: - getattr(self._nwbfile.electrodes, fn).set_data_io(H5DataIO, {"maxshape": (None,), "chunks": True}) + getattr(nwbfile.electrodes, fn).set_data_io(H5DataIO, {"maxshape": (None,), "chunks": True}) - el_df = self._nwbfile.electrodes.to_dataframe() + el_df = nwbfile.electrodes.to_dataframe() el_df = el_df[el_df["group"] == el_grp] b_in = el_df["label"].isin(self._current_msg.axes["ch"].data) - el_tbl_region = self._nwbfile.create_electrode_table_region( + el_tbl_region = nwbfile.create_electrode_table_region( region=el_df[b_in].index.tolist(), description=f"electrodes for {key}", ) @@ -719,28 +822,12 @@ def _prep_continuous_io(self): conversion=1e-6, description=series_description, ) - self._nwbfile.add_acquisition(series) + nwbfile.add_acquisition(series) class NWBSink(BaseConsumerUnit[NWBSinkSettings, AxisArray, NWBSinkConsumer]): SETTINGS = NWBSinkSettings - INPUT_SETTINGS = ez.InputStream(NWBSinkSettings) - - @ez.subscriber(INPUT_SETTINGS) - async def on_settings(self, msg: NWBSinkSettings) -> None: - # Reset if settings _other than `recording`_ have changed. - b_reset = msg.filepath != self.SETTINGS.filepath - b_reset = b_reset or msg.overwrite_old != self.SETTINGS.overwrite_old - b_reset = b_reset or msg.axis != self.SETTINGS.axis - b_reset = b_reset or msg.inc_clock != self.SETTINGS.inc_clock - b_reset = b_reset or msg.meta_yaml != self.SETTINGS.meta_yaml - if b_reset: - self.apply_settings(msg) - self.create_processor() - elif msg.recording != self.SETTINGS.recording: - self.processor.toggle_recording(msg.recording) - async def shutdown(self) -> None: await super().shutdown() self.processor.close() diff --git a/tests/test_clockdriven.py b/tests/test_clockdriven.py index bd25124..9ff6bf6 100644 --- a/tests/test_clockdriven.py +++ b/tests/test_clockdriven.py @@ -1,10 +1,43 @@ """Tests for NWBClockDrivenProducer.""" -from ezmsg.util.messages.axisarray import AxisArray +import threading + +from ezmsg.util.messages.axisarray import AxisArray, LinearAxis from ezmsg.nwb.clockdriven import NWBClockDrivenProducer, NWBClockDrivenSettings from ezmsg.nwb.util import ReferenceClockType + +async def test_areset_state_runs_reset_in_worker_thread(): + """``_areset_state`` must offload sync ``_reset_state`` to a worker thread + so the unit's event loop isn't blocked by the multi-second NWB open.""" + main_tid = threading.get_ident() + seen_tids: list[int] = [] + + class Spy(NWBClockDrivenProducer): + def _reset_state(self, time_axis): + seen_tids.append(threading.get_ident()) + super()._reset_state(time_axis) + + # Idle-mode (empty filepath/stream_key) short-circuits the slow NWB + # I/O — we only want to assert the offload mechanism, not exercise + # pynwb here. + producer = Spy( + settings=NWBClockDrivenSettings( + fs=50.0, + filepath="", + stream_key="", + reference_clock=ReferenceClockType.UNKNOWN, + ) + ) + seen_tids.clear() # discard any eager-init invocation + + await producer._areset_state(LinearAxis(gain=0.05, offset=0.0)) + + assert seen_tids == [t for t in seen_tids if t != main_tid], "_reset_state ran on the main event-loop thread" + assert len(seen_tids) == 1, "_reset_state should run exactly once per _areset_state" + + # --- Rate-only continuous stream --- diff --git a/tests/test_integration.py b/tests/test_integration.py index 1948216..958f0ee 100644 --- a/tests/test_integration.py +++ b/tests/test_integration.py @@ -1,11 +1,15 @@ """Integration tests for ezmsg-nwb: ezmsg system tests and writer round-trip.""" +import asyncio import tempfile +import typing from dataclasses import field from pathlib import Path import ezmsg.core as ez import numpy as np +import pynwb +import pytest from ezmsg.baseproc.clock import Clock, ClockSettings from ezmsg.util.messagecodec import message_log from ezmsg.util.messagelogger import MessageLogger, MessageLoggerSettings @@ -15,6 +19,7 @@ from ezmsg.nwb import ( NWBAxisArrayIterator, NWBIteratorSettings, + NWBSink, NWBSinkConsumer, NWBSinkSettings, ReferenceClockType, @@ -226,10 +231,139 @@ def test_writer_recording_toggle(): inc_clock=ReferenceClockType.UNKNOWN, ) ) - assert sink._recording is False - sink.toggle_recording(True) - assert sink._recording is True - sink.toggle_recording() - assert sink._recording is False + assert sink.settings.recording is False + with pytest.warns(DeprecationWarning): + sink.toggle_recording(True) + assert sink.settings.recording is True + with pytest.warns(DeprecationWarning): + sink.toggle_recording() + assert sink.settings.recording is False sink.close(write=False) outpath.unlink(missing_ok=True) + + +# -- Settings-push through running graph -- + + +class EnableRecordingPubSettings(ez.Settings): + sink_filepath: Path + n_messages: int + chunk_samples: int + chunk_channels: int + + +class EnableRecordingPub(ez.Unit): + """Pushes a settings update that flips ``recording`` from ``False`` to + ``True``, waits for it to land at the sink, then publishes data. The + inverse direction (``True → False``) races: NWB's first-message I/O + takes seconds, so a fast settings flip beats every queued data message + to the consumer and we observe nothing on disk. Going ``False → True`` + pushes the slow side ahead of the fast side, so a settled sleep is + enough to guarantee ordering.""" + + SETTINGS = EnableRecordingPubSettings + + OUTPUT_SETTINGS = ez.OutputStream(NWBSinkSettings) + OUTPUT_SIGNAL = ez.OutputStream(AxisArray) + + @ez.publisher(OUTPUT_SIGNAL) + @ez.publisher(OUTPUT_SETTINGS) + async def go(self) -> typing.AsyncGenerator: + s = self.SETTINGS + + # Flip recording on. Same filepath → only ``recording`` differs → + # NONRESET path (no file rebuild), settings just rebinds. + yield ( + self.OUTPUT_SETTINGS, + NWBSinkSettings( + filepath=s.sink_filepath, + overwrite_old=True, + recording=True, + inc_clock=ReferenceClockType.UNKNOWN, + ), + ) + # Give the settings subscriber time to process before any data hits + # the data subscriber. Settings handling is fast (no I/O) so 1s is + # plenty. + await asyncio.sleep(1.0) + + for i in range(s.n_messages): + yield ( + self.OUTPUT_SIGNAL, + AxisArray( + data=np.ones((s.chunk_samples, s.chunk_channels), dtype=np.float32) * float(i), + dims=["time", "ch"], + axes={"time": AxisArray.TimeAxis(fs=1000.0, offset=i * 0.05)}, + key="K", + ), + ) + await asyncio.sleep(0.05) + + # Wait long enough for the slow first-message NWB write to finish + # before the runtime tears the graph down. + await asyncio.sleep(3.0) + raise ez.NormalTermination + + +class SettingsPushTestSettings(ez.Settings): + pub: EnableRecordingPubSettings + sink: NWBSinkSettings + + +class SettingsPushCollection(ez.Collection): + SETTINGS = SettingsPushTestSettings + + PUB = EnableRecordingPub() + SINK = NWBSink() + + def configure(self) -> None: + self.PUB.apply_settings(self.SETTINGS.pub) + self.SINK.apply_settings(self.SETTINGS.sink) + + def network(self) -> ez.NetworkDefinition: + return ( + (self.PUB.OUTPUT_SIGNAL, self.SINK.INPUT_SIGNAL), + (self.PUB.OUTPUT_SETTINGS, self.SINK.INPUT_SETTINGS), + ) + + +def test_settings_push_through_graph_enables_recording(): + """A NWBSinkSettings message published into ``NWBSink.INPUT_SETTINGS`` + while the graph is running should reach the consumer's + ``update_settings`` and propagate to subsequent writes.""" + outpath = Path(tempfile.gettempdir()) / "ezmsg_nwb_settings_push.nwb" + outpath.unlink(missing_ok=True) + + n_messages, chunk = 3, 10 + + system = SettingsPushCollection( + SettingsPushTestSettings( + pub=EnableRecordingPubSettings( + sink_filepath=outpath, + n_messages=n_messages, + chunk_samples=chunk, + chunk_channels=2, + ), + # Sink starts with recording=False — without the in-graph + # settings push, no data would land on disk. + sink=NWBSinkSettings( + filepath=outpath, + overwrite_old=True, + recording=False, + inc_clock=ReferenceClockType.UNKNOWN, + ), + ) + ) + ez.run(SYSTEM=system) + + assert outpath.exists(), ( + "If the in-graph settings update never reached the consumer, " + "recording stays False and the empty-file cleanup would have " + "deleted this path on shutdown." + ) + with pynwb.NWBHDF5IO(str(outpath), "r") as io: + nwbfile = io.read() + assert "K" in nwbfile.acquisition + assert len(nwbfile.acquisition["K"].data) == n_messages * chunk + + outpath.unlink(missing_ok=True) diff --git a/tests/test_iterator.py b/tests/test_iterator.py index 301f2a2..f9ad309 100644 --- a/tests/test_iterator.py +++ b/tests/test_iterator.py @@ -1,6 +1,7 @@ """Tests for NWBAxisArrayIterator.""" import math +import threading from collections import Counter import numpy as np @@ -8,6 +9,34 @@ from ezmsg.nwb import NWBAxisArrayIterator, NWBIteratorSettings, ReferenceClockType + +async def test_areset_state_runs_reset_in_worker_thread(test_nwb_path): + """``_areset_state`` must offload sync ``_reset_state`` to a worker + thread so the unit's event loop stays responsive during the NWB open.""" + main_tid = threading.get_ident() + seen_tids: list[int] = [] + + class Spy(NWBAxisArrayIterator): + def _reset_state(self): + seen_tids.append(threading.get_ident()) + super()._reset_state() + + producer = Spy( + NWBIteratorSettings( + filepath=test_nwb_path, + chunk_dur=1.0, + reference_clock=ReferenceClockType.UNKNOWN, + ) + ) + # Discard the eager sync invocation from __init__. + seen_tids.clear() + + await producer._areset_state() + + assert len(seen_tids) == 1 + assert seen_tids[0] != main_tid, "_reset_state ran on the main event-loop thread" + + # --- Stream discovery --- diff --git a/tests/test_writer.py b/tests/test_writer.py index 200c7aa..608e7ae 100644 --- a/tests/test_writer.py +++ b/tests/test_writer.py @@ -1,6 +1,65 @@ """Tests for NWB writer module.""" +import dataclasses +import tempfile +from pathlib import Path + +import numpy as np +import pynwb +import pytest +from ezmsg.util.messages.axisarray import AxisArray + from ezmsg.nwb import NWBSinkSettings, ReferenceClockType +from ezmsg.nwb.writer import NWBSinkConsumer + + +@pytest.fixture(autouse=True) +def _reset_nwbsink_shared(): + """The ``shared_*`` class attrs are deliberately sticky across instances + within a process; reset them between tests so session-start semantics + don't leak.""" + NWBSinkConsumer.shared_start_datetime = None + NWBSinkConsumer.shared_t0 = None + NWBSinkConsumer.shared_clock_type = None + yield + NWBSinkConsumer.shared_start_datetime = None + NWBSinkConsumer.shared_t0 = None + NWBSinkConsumer.shared_clock_type = None + + +def _make_continuous( + n: int = 50, + ch: int = 4, + fs: float = 1000.0, + offset: float = 0.0, + key: str = "TestStream", + ch_labels: list[str] | None = None, +) -> AxisArray: + axes = {"time": AxisArray.TimeAxis(fs=fs, offset=offset)} + if ch_labels is not None: + axes["ch"] = AxisArray.CoordinateAxis(np.asarray(ch_labels), dims=["ch"], unit="") + return AxisArray( + data=np.random.randn(n, ch).astype(np.float32), + dims=["time", "ch"], + axes=axes, + key=key, + ) + + +def _fresh_path(stem: str) -> Path: + p = Path(tempfile.gettempdir()) / f"ezmsg_nwb_{stem}.nwb" + p.unlink(missing_ok=True) + return p + + +def _sink(filepath: Path, **overrides) -> NWBSinkConsumer: + base = dict( + filepath=filepath, + overwrite_old=True, + inc_clock=ReferenceClockType.UNKNOWN, + ) + base.update(overrides) + return NWBSinkConsumer(settings=NWBSinkSettings(**base)) def test_sink_settings_defaults(): @@ -30,3 +89,435 @@ def test_sink_settings_custom(): assert settings.recording is False assert settings.inc_clock == ReferenceClockType.MONOTONIC assert settings.split_bytes == 1024 + + +# -- update_settings / reset machinery -- + + +def test_update_settings_recording_does_not_reset(): + path = _fresh_path("update_recording") + sink = _sink(path, recording=True) + io_before = sink._state.io + hash_before = sink._hash + + sink.update_settings(dataclasses.replace(sink.settings, recording=False)) + + assert sink._state.io is io_before, "recording-only change must not swap io" + assert sink._hash == hash_before, "recording-only change must not request reset" + assert sink.settings.recording is False + + sink.close(write=False) + path.unlink(missing_ok=True) + + +def test_update_settings_split_bytes_does_not_reset(): + path = _fresh_path("update_splitbytes") + sink = _sink(path, split_bytes=0) + io_before = sink._state.io + hash_before = sink._hash + + sink.update_settings(dataclasses.replace(sink.settings, split_bytes=1024)) + + assert sink._state.io is io_before + assert sink._hash == hash_before + assert sink.settings.split_bytes == 1024 + + sink.close(write=False) + path.unlink(missing_ok=True) + + +def test_update_settings_filepath_swaps_file(): + path_a = _fresh_path("update_filepath_a") + path_b = _fresh_path("update_filepath_b") + sink = _sink(path_a) + io_a = sink._state.io + + sink.update_settings(dataclasses.replace(sink.settings, filepath=path_b, overwrite_old=True)) + assert sink._hash == -1, "reset-field change must request reset" + + # Simulate what the stateful machinery does on the next message: + # hash mismatch triggers _reset_state, which closes A and opens B. + sink._reset_state(None) + sink._hash = sink._hash_message(None) + + assert sink._state.io is not io_a + assert sink._state.filepath.name == path_b.name + + sink.close(write=False) + for p in (path_a, path_b): + p.unlink(missing_ok=True) + + +def test_update_settings_axis_requires_reset(): + path = _fresh_path("update_axis") + sink = _sink(path) + hash_before = sink._hash + + sink.update_settings(dataclasses.replace(sink.settings, axis="win")) + + assert sink._hash == -1, "axis change must request reset" + assert hash_before == 0 + + sink.close(write=False) + path.unlink(missing_ok=True) + + +# -- recording gate -- + + +def test_recording_false_skips_data_write(): + path = _fresh_path("recording_gate") + sink = _sink(path, recording=False) + + msg = _make_continuous(n=100, ch=2, key="Gated") + sink._process(msg) + + # Series was created (prep ran), but no data should have been appended. + assert "Gated" in sink._state.series + ss = sink._state.series["Gated"] + assert ss.bytes_written == 0 + assert len(ss.data) == 0 + + sink.close(write=False) + path.unlink(missing_ok=True) + + +def test_recording_toggled_live_resumes_writes(): + path = _fresh_path("recording_live") + sink = _sink(path, recording=False) + + msg1 = _make_continuous(n=50, key="Live") + sink._process(msg1) + assert sink._state.series["Live"].bytes_written == 0 + + # Flip the recording flag with a NONRESET update. + sink.update_settings(dataclasses.replace(sink.settings, recording=True)) + + msg2 = _make_continuous(n=50, offset=0.05, key="Live") + sink._process(msg2) + assert sink._state.series["Live"].bytes_written > 0 + + sink.close(write=False) + path.unlink(missing_ok=True) + + +# -- error paths -- + + +def test_shape_mismatch_raises_and_closes(): + path = _fresh_path("shape_mismatch") + sink = _sink(path) + + sink._process(_make_continuous(n=50, ch=4, key="S")) + bad = _make_continuous(n=50, ch=8, key="S") # different trailing shape + + with pytest.raises(ValueError, match="changed shape"): + sink._process(bad) + assert sink._state.io is None, "close() should run on shape mismatch" + + path.unlink(missing_ok=True) + + +def test_str_data_without_event_key_raises(): + path = _fresh_path("str_no_event") + sink = _sink(path) + + msg = AxisArray( + data=np.array([["a"], ["b"]], dtype="0 and no "%d" in path, files land at _00.nwb, + # _01.nwb, ...; the nominal path itself is never created. + file_00 = nominal.parent / (nominal.stem + "_00" + nominal.suffix) + file_01 = nominal.parent / (nominal.stem + "_01" + nominal.suffix) + for p in (file_00, file_01): + p.unlink(missing_ok=True) + + threshold = 1_000_000 # ~1 MB + sink = _sink(nominal, split_bytes=threshold) + + # Each chunk: 1000 × 100 × float32 (= 400 KB data) + 1000 × float64 + # (= 8 KB timestamps) ≈ 408 KB. Three chunks → ~1.22 MB, crossing the + # 1 MB threshold and triggering a split right after the third chunk. + n_samples, n_ch = 1000, 100 + chunks = [_make_continuous(n=n_samples, ch=n_ch, fs=1000.0, offset=i * 1.0, key="Big") for i in range(3)] + for c in chunks: + sink._process(c) + + assert sink._state.split_count == 1, "expected exactly one split after ~1.2 MB" + assert file_00.exists(), "first split file should be on disk" + # _state.io now points at file_01 — confirm it's open and empty so far. + assert sink._state.io is not None + assert sink._state.series["Big"].bytes_written == 0 + + # Drop one more chunk into the new file, then close. + sink._process(_make_continuous(n=n_samples, ch=n_ch, fs=1000.0, offset=4.0, key="Big")) + sink.close(write=True) + + assert file_01.exists(), "second split file should be on disk" + + # Verify each file holds its expected slice. + with pynwb.NWBHDF5IO(str(file_00), "r") as io: + assert len(io.read().acquisition["Big"].data) == 3 * n_samples + with pynwb.NWBHDF5IO(str(file_01), "r") as io: + assert len(io.read().acquisition["Big"].data) == n_samples + + for p in (file_00, file_01): + p.unlink(missing_ok=True) + + +def test_multi_sink_share_session_start(): + """Two sinks with the same ``inc_clock`` should anchor session start once + and share it — that's the entire point of the ``shared_*`` class attrs.""" + path_a = _fresh_path("multi_session_a") + path_b = _fresh_path("multi_session_b") + + # Anchor the session via sink A's first message. + sink_a = _sink(path_a, inc_clock=ReferenceClockType.SYSTEM) + sink_a._process(_make_continuous(n=10, fs=100.0, offset=1_700_000_000.0, key="A")) + anchor_dt = NWBSinkConsumer.shared_start_datetime + anchor_t0 = NWBSinkConsumer.shared_t0 + assert anchor_dt is not None + assert anchor_t0 is not None + + # Sink B should reuse the same anchor without overwriting it. + sink_b = _sink(path_b, inc_clock=ReferenceClockType.SYSTEM) + sink_b._process(_make_continuous(n=10, fs=100.0, offset=1_700_000_500.0, key="B")) + + assert NWBSinkConsumer.shared_start_datetime == anchor_dt + assert NWBSinkConsumer.shared_t0 == anchor_t0 + + sink_a.close(write=True) + sink_b.close(write=True) + + # session_start_time on disk should match across files. + with pynwb.NWBHDF5IO(str(path_a), "r") as io_a, pynwb.NWBHDF5IO(str(path_b), "r") as io_b: + assert io_a.read().session_start_time == io_b.read().session_start_time + + for p in (path_a, path_b): + p.unlink(missing_ok=True) + + +def test_clock_type_mismatch_raises(): + """A second sink with a different ``inc_clock`` than the established + shared one must refuse — preventing inconsistent timestamps in the file.""" + path_a = _fresh_path("clock_a") + path_b = _fresh_path("clock_b") + + sink_a = _sink(path_a, inc_clock=ReferenceClockType.SYSTEM) + sink_a._process(_make_continuous(n=5, key="A")) + assert NWBSinkConsumer.shared_clock_type == ReferenceClockType.SYSTEM + + # Construction triggers _nwb_create_or_fail → get_session_datetime which + # checks the shared_clock_type. Mismatched clock must raise. + with pytest.raises(ValueError, match="share the same clock type"): + _sink(path_b, inc_clock=ReferenceClockType.MONOTONIC) + + sink_a.close(write=False) + for p in (path_a, path_b): + p.unlink(missing_ok=True) + + +def test_meta_yaml_roundtrip(tmp_path): + """Custom ``meta_yaml`` should populate NWBFile / Subject fields readable + after the file closes.""" + meta_path = tmp_path / "meta.yaml" + meta_path.write_text( + "NWBFile:\n" + " session_description: ezmsg-nwb meta_yaml test\n" + " experimenter:\n" + " - Tester, Test\n" + " institution: Test Institute\n" + "Subject:\n" + " subject_id: TestSubject001\n" + " species: Mus musculus\n" + " sex: U\n" + " age: P30D\n" + ) + nwb_path = tmp_path / "meta.nwb" + sink = _sink(nwb_path, meta_yaml=meta_path) + sink._process(_make_continuous(n=10, key="K")) + sink.close(write=True) + + with pynwb.NWBHDF5IO(str(nwb_path), "r") as io: + nwbfile = io.read() + assert nwbfile.session_description == "ezmsg-nwb meta_yaml test" + assert "Tester, Test" in list(nwbfile.experimenter) + assert nwbfile.institution == "Test Institute" + assert nwbfile.subject is not None + assert nwbfile.subject.subject_id == "TestSubject001" + assert nwbfile.subject.species == "Mus musculus" + + +def test_rate_change_raises_and_closes(): + """Same key, different sample rate → consistency check fires alongside + the shape check and tears the file down.""" + path = _fresh_path("rate_change") + sink = _sink(path) + + sink._process(_make_continuous(n=50, ch=4, fs=1000.0, key="K")) + bad = _make_continuous(n=50, ch=4, fs=2000.0, offset=0.05, key="K") + + with pytest.raises(ValueError, match="changed shape"): + sink._process(bad) + assert sink._state.io is None, "rate change must close the file" + + path.unlink(missing_ok=True) + + +@pytest.mark.xfail( + reason=( + "_prep_from_meta carries _current_msg.data.dtype across loop iterations " + "(or reads it as None at construction). Deferred per the writer cleanup TODO." + ), + raises=(AttributeError, TypeError, ValueError), + strict=False, +) +def test_expected_series_smoke(tmp_path): + """``expected_series`` should pre-allocate stream containers from a + metadata yaml so the file is fully shaped before any message arrives.""" + expected_yaml = tmp_path / "expected.yaml" + expected_yaml.write_text("PreAlloc:\n" " fs: 100.0\n" " shape: [-1, 4]\n") + nwb_path = tmp_path / "expected.nwb" + sink = NWBSinkConsumer( + settings=NWBSinkSettings( + filepath=nwb_path, + overwrite_old=True, + expected_series=expected_yaml, + inc_clock=ReferenceClockType.UNKNOWN, + ) + ) + assert "PreAlloc" in sink._state.series + sink.close(write=True) + + with pynwb.NWBHDF5IO(str(nwb_path), "r") as io: + nwbfile = io.read() + assert "PreAlloc" in nwbfile.acquisition + + +def test_sample_trigger_routes_to_epochs(): + """SampleTriggerMessage with a string label should land in nwbfile.epochs.""" + pytest.importorskip("ezmsg.baseproc") + from ezmsg.baseproc import SampleTriggerMessage + + path = _fresh_path("sample_trigger") + sink = _sink(path) + + # Prime the session start so the trigger's timestamp resolves to + # a positive offset rather than a near-zero "now" datetime. + NWBSinkConsumer.shared_t0 = 0.0 + NWBSinkConsumer.shared_clock_type = ReferenceClockType.UNKNOWN + + trigger = SampleTriggerMessage( + timestamp=1.5, + period=(0.0, 0.0), + value="trigger_a", + ) + sink._process(trigger) + sink.close(write=True) + + with pynwb.NWBHDF5IO(str(path), "r") as io: + nwbfile = io.read() + assert nwbfile.epochs is not None + labels = list(nwbfile.epochs.to_dataframe()["label"]) + # Writer always prepends the EZNWB-START sentinel; assert the + # trigger label is present alongside it. + assert "EZNWB-START" in labels + assert "trigger_a" in labels + + path.unlink(missing_ok=True)