From 4a1207b34854f737d5b2b81761464051719d49fd Mon Sep 17 00:00:00 2001 From: Chadwick Boulay Date: Thu, 3 Sep 2026 18:53:31 -0400 Subject: [PATCH] Build the iterators on ezmsg-baseproc's producer classes This was the only ezmsg source package still on `ez.Unit` + `GenState`; blackrock, lsl, neo and nwb all use `BaseStatefulProducer` and `BaseProducerUnit`. Four things follow from joining them. **The file open no longer blocks the event loop.** `pyxdf.load_xdf` reads and decodes the entire file, and it ran inside a synchronous `initialize()`, so graph startup stalled every other unit in the process for as long as that took. It is now a state reset, and `_areset_state` puts it on a worker thread. This deliberately diverges from ezmsg-neo and ezmsg-nwb, which call `_reset_state()` eagerly from `__init__` and therefore still pay the first open on the loop -- nwb's own test says as much ("Discard the eager sync invocation from __init__"), so its `_areset_state` only covers reopens after a settings change. Nothing in this package's public surface reads stream metadata before the first chunk, so there is nothing to lose by waiting, and a test asserts both halves: construction touches no file, and the load lands off the loop. **Settings can change at runtime.** `BaseProducerUnit` brings `INPUT_SETTINGS` and delegates to `update_settings`; previously settings were read once in `initialize` and a running unit could not be retargeted. `playback_rate` and `self_terminating` are listed in `NONRESET_SETTINGS_FIELDS` because they belong to the unit's pacing rather than to the reader, so changing them must not throw away a loaded file. **The two near-identical unit bodies collapse** into `_XDFUnitBase`, which owns the playback clock and end-of-file handling. Both publishers are named `produce`, which matters: ezmsg collects publishers per attribute, so keeping the old `pub_chunk`/`pub_multi` names left `BaseProducerUnit.produce` running alongside them -- two publishers draining one producer, neither stopping. That hung the new graph tests until the names matched. **Templates are built in one place.** `_build_template` and `_with_time` replace the copy in each iterator, so the chunk_dim declaration and the fingerprint priming exist once rather than twice. `XDFIterator` itself is unchanged except for a public `exhausted` property and `print` becoming `ez.logger.info` -- a library writing to stdout is a defect regardless. The keyword API is preserved: `BaseProducer.__init__` forwards `**kwargs` to the settings type, so `XDFAxisArrayIterator(filepath=..., select=...)` still works, alongside `settings=`. `XDFMultiIteratorUnitSettings` is kept as an alias of the settings type that moved to `iter.py`. 27 tests pass, up from 20. The seven new ones cover the producer contract (lazy open, off-loop load, both construction forms, which settings reopen the file) and both units end to end in a real `ez.run` graph -- which is what caught the publisher-name collision. The async test is driven with `asyncio.run` rather than written `async def`: the pytest config names `asyncio_mode` but pytest-asyncio is not installed, so an async test would be silently never awaited and would pass without running. --- pyproject.toml | 3 +- src/ezmsg/xdf/iter.py | 374 +++++++++++++++++++++++++--------------- src/ezmsg/xdf/source.py | 178 +++++++++---------- tests/test_iter.py | 121 ++++++++++++- 4 files changed, 435 insertions(+), 241 deletions(-) diff --git a/pyproject.toml b/pyproject.toml index 02d0879..45243a7 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -10,6 +10,7 @@ dynamic = ["version"] dependencies = [ # 3.10.0b2 for AxisArray.chunk_dim and CoordinateAxis.fingerprint. "ezmsg>=3.10.0b2", + "ezmsg-baseproc>=1.11.0", "numpy>=2.0.2", "pyxdf>=1.16.8", ] @@ -63,7 +64,7 @@ select = ["E", "F", "I", "W"] [tool.ruff.lint.isort] known-first-party = ["ezmsg.xdf"] -known-third-party = ["ezmsg"] +known-third-party = ["ezmsg", "ezmsg.baseproc"] [tool.uv.sources] # Uncomment to use development version of ezmsg from git diff --git a/src/ezmsg/xdf/iter.py b/src/ezmsg/xdf/iter.py index 3d84766..e204351 100644 --- a/src/ezmsg/xdf/iter.py +++ b/src/ezmsg/xdf/iter.py @@ -1,9 +1,16 @@ +import asyncio +import os import queue +import typing +from dataclasses import field from pathlib import Path +import ezmsg.core as ez import numpy as np import numpy.typing as npt import pyxdf +from ezmsg.baseproc.protocols import processor_state +from ezmsg.baseproc.stateful import BaseStatefulProducer from ezmsg.util.messages.axisarray import AxisArray from ezmsg.util.messages.util import replace @@ -143,7 +150,7 @@ def _scan_file(self): self._streams = [self._streams[stream_names.index(_)] for _ in self._select] self._metadata = {k: self._metadata[k] for k in self._select} - print( + ez.logger.info( f"Imported {len(self._streams)} streams from {self._filepath} " f"spanning {xdf_dur:.2f} s beginning at t={xdf_t0:.2f}." ) @@ -156,12 +163,17 @@ def stream_meta(self) -> list[dict] | dict: def n_chunks(self) -> int: return self._n_chunks + @property + def exhausted(self) -> bool: + """True once every chunk boundary has been handed out.""" + return self._chunk_ix >= self._n_chunks + def __iter__(self): self._chunk_ix = 0 return self def __next__(self) -> dict[str, tuple[npt.NDArray, npt.NDArray]]: - if self._chunk_ix >= self.n_chunks: + if self.exhausted: raise StopIteration else: out_dict = {} @@ -190,156 +202,236 @@ def labels_from_strm(strm: dict) -> list[str]: return labels -class XDFAxisArrayIterator(XDFIterator): - def __init__(self, *args, select: str, **kwargs): - """ - This Iterator loads only a single stream and yields a single :obj:`AxisArray` object per chunk. +def _build_template(stream: dict, name: str, n_ch: int, fs: float) -> AxisArray: + """The message every chunk of *stream* is a `replace` of. - Args: - *args: - select: Unlike :obj:`XDFIterator`, this must be a single string, the name of the stream to select. - **kwargs: + Built once per stream so the `ch` axis object -- and the fingerprint cached + on it -- is shared by every message, which is what makes priming cheap. + """ + labels = labels_from_strm(stream) + time_ax = ( + AxisArray.TimeAxis(fs=fs, offset=0.0) + if fs + else AxisArray.CoordinateAxis(data=np.array([]), dims=["time"], unit="s") + ) + ch_ax = AxisArray.CoordinateAxis(data=np.array(labels), dims=["ch"]) + # Compute the channel fingerprint once, now. It is cached on the axis and + # pickled with it, and every message from this stream reuses this same axis + # object, so one checksum covers the whole file. Left cold it would be + # computed by the first stateful consumer in this process -- and, since + # unpickling builds a new axis object per message, by the first consumer in + # every other process, on every message. + ch_ax.fingerprint + return AxisArray( + data=np.zeros((0, n_ch), dtype=stream["time_series"].dtype), + dims=["time", "ch"], + axes={"time": time_ax, "ch": ch_ax}, + key=name, + # Messages accumulate along `time`, whether the stream is regular or + # carries per-sample timestamps; `ch` describes the stream itself. + chunk_dim="time", + ) + + +def _with_time(template: AxisArray, data: npt.NDArray, tvec: npt.NDArray, fallback_t: float) -> AxisArray: + """A chunk message: the template's data replaced, and its time axis advanced. + + An irregular stream carries every timestamp; a regular one carries only where + the chunk starts, since its gain says the rest. + """ + time_ax = template.axes["time"] + if isinstance(time_ax, AxisArray.CoordinateAxis): + t_kwargs = {"data": tvec if len(tvec) else np.array([])} + else: + t_kwargs = {"offset": tvec[0] if len(tvec) else fallback_t} + return replace( + template, + data=data, + axes={**template.axes, "time": replace(time_ax, **t_kwargs)}, + ) + + +class XDFIteratorSettings(ez.Settings): + """Settings shared by both AxisArray iterators. + + ``playback_rate`` and ``self_terminating`` belong to the unit rather than to + the reader, and are listed in :attr:`NONRESET_SETTINGS_FIELDS` so changing + either does not reopen the file. + """ + + filepath: typing.Union[os.PathLike, str] + select: str = "" + chunk_dur: float = 1.0 + start_time: float | None = None + stop_time: float | None = None + rezero: bool = True + playback_rate: float | None = None + self_terminating: bool = False + """ + If True, the unit will raise a :obj:`ez.NormalTermination` exception when the file is exhausted. + Note, however, that this will terminate the pipeline even if the data published by this unit are still in transit, + which will lead to the pipeline output being truncated before it has finished processing the stream. + `self_terminating` should only be used when it is not important that the pipeline finish processing data, such + as during prototyping and testing. + """ + + +class XDFMultiIteratorSettings(XDFIteratorSettings): + select: set[str] | None = None + force_single_sample: set = field(default_factory=set) + + +@processor_state +class XDFIteratorState: + reader: XDFIterator | None = None + template: AxisArray | None = None + + +@processor_state +class XDFMultiIteratorState: + reader: XDFIterator | None = None + templates: dict | None = None + pubqueue: queue.SimpleQueue | None = None + + +class _XDFProducerBase: + """Shared plumbing for the two AxisArray producers. + + The file load is deliberately *not* run from ``__init__``. ezmsg-neo and + ezmsg-nwb both reset eagerly there and so pay the whole open on the event + loop during ``initialize``; here the first ``__acall__`` triggers + ``_areset_state``, which puts it on a worker thread. Nothing in this + package's public surface reads stream metadata before the first chunk, so + there is nothing to lose by waiting. + """ + + NONRESET_SETTINGS_FIELDS = frozenset({"playback_rate", "self_terminating"}) + + async def _areset_state(self) -> None: + """Offload the sync open onto a worker thread. + + ``pyxdf.load_xdf`` reads and decodes the entire file, which for a + several-hundred-megabyte recording is seconds of pure CPU and I/O. On the + event loop that stalls every other unit in the process. """ - kwargs["select"] = set((select,)) - super().__init__(*args, **kwargs) - _sel = [_ for _ in self._select][0] - labels = labels_from_strm(self._streams[0]) - if self._metadata[_sel].get("nominal_srate", None): - time_ax = AxisArray.TimeAxis(fs=self._metadata[_sel]["nominal_srate"], offset=0) - else: - time_ax = AxisArray.CoordinateAxis(data=np.array([]), dims=["time"], unit="s") - ch_ax = AxisArray.CoordinateAxis(data=np.array(labels), dims=["ch"]) - # Compute the channel fingerprint once, now. It is cached on the axis and - # pickled with it, and every message from this stream reuses this same axis - # object, so one checksum covers the whole file. Left cold it would be - # computed by the first stateful consumer in this process -- and, since - # unpickling builds a new axis object per message, by the first consumer in - # every other process, on every message. - ch_ax.fingerprint - self._template = AxisArray( - data=np.zeros((0, len(labels)), dtype=self._streams[0]["time_series"].dtype), - dims=["time", "ch"], - axes={ - "time": time_ax, - "ch": ch_ax, - }, - key=self._streams[0]["info"]["name"][0], - # Messages accumulate along `time`, whether the stream is regular or - # carries per-sample timestamps; `ch` describes the stream itself. - chunk_dim="time", + await asyncio.to_thread(self._reset_state) + + def _build_reader(self, select: set[str] | None) -> XDFIterator: + return XDFIterator( + filepath=self.settings.filepath, + select=select, + chunk_dur=self.settings.chunk_dur, + start_time=self.settings.start_time, + stop_time=self.settings.stop_time, + rezero=self.settings.rezero, ) + +class XDFAxisArrayIterator( + _XDFProducerBase, + BaseStatefulProducer[XDFIteratorSettings, AxisArray, XDFIteratorState], +): + """Loads a single stream and produces one :obj:`AxisArray` per chunk. + + ``select`` must be a single stream name, unlike :obj:`XDFIterator`. + """ + + @property + def exhausted(self) -> bool: + reader = self._state.reader + return reader is not None and reader.exhausted + + def _reset_state(self) -> None: + reader = self._build_reader({self.settings.select}) + meta = reader.stream_meta[self.settings.select] + self._state.reader = reader + self._state.template = _build_template( + reader._streams[0], + name=reader._streams[0]["info"]["name"][0], + n_ch=meta["channel_count"], + fs=meta["nominal_srate"], + ) + + async def _produce(self) -> AxisArray | None: + reader = self._state.reader + try: + chunk_dict = next(reader) + except StopIteration: + return None + data, tvec = chunk_dict.get(self.settings.select, (None, None)) + if data is None: + return None + return _with_time(self._state.template, data, tvec, reader._last_time) + def __next__(self) -> AxisArray: - result: AxisArray | None = None - chunk_dict = super().__next__() - # Should only be 1 in self._select. If there are more then we overwrite with the last. - for strm_name in self._select: - if strm_name in chunk_dict: - data, tvec = chunk_dict[strm_name] - if isinstance(self._template.axes["time"], AxisArray.CoordinateAxis): - t_kwargs = {"data": tvec} - else: - t_kwargs = {"offset": tvec[0] if len(tvec) else self._last_time} - result = replace( - self._template, - data=data, - axes={ - **self._template.axes, - "time": replace( - self._template.axes["time"], - **t_kwargs, - ), - }, - ) + result = self() + if result is None: + raise StopIteration return result -class XDFMultiAxArrIterator(XDFIterator): - def __init__(self, *args, force_single_sample: set = set(), **kwargs): - """ - This Iterator loads multiple streams and yields a :obj:`AxisArray` object per iteration, - but the stream source might different between chunks. +class XDFMultiAxArrIterator( + _XDFProducerBase, + BaseStatefulProducer[XDFMultiIteratorSettings, AxisArray, XDFMultiIteratorState], +): + """Loads multiple streams and produces one :obj:`AxisArray` per iteration. - Args: - *args: - force_single_sample: Use this to identify irregular-rate streams that might conceivably have more than one - event within the defined chunk_dur, for which :obj:`AxisArray` cannot represent timestamps properly. - **kwargs: - """ - super().__init__(*args, **kwargs) - self._force_single_sample = force_single_sample - stream_names = [_["info"]["name"][0] for _ in self._streams] - - # Create template messages for each stream - self._templates = {} - for stream_name, stream_meta in self._metadata.items(): - stream = self._streams[stream_names.index(stream_name)] - labels = labels_from_strm(stream) - fs = stream_meta["nominal_srate"] - time_ax = ( - AxisArray.TimeAxis(fs=fs, offset=0.0) - if fs - else AxisArray.CoordinateAxis(data=np.array([]), dims=["time"], unit="s") - ) - ch_ax = AxisArray.CoordinateAxis(data=np.array(labels), dims=["ch"]) - ch_ax.fingerprint # primed once per stream -- see the single-stream iterator above - self._templates[stream_name] = AxisArray( - data=np.zeros((0, stream_meta["channel_count"]), dtype=stream["time_series"].dtype), - dims=["time", "ch"], - axes={ - "time": time_ax, - "ch": ch_ax, - }, - key=stream_name, - chunk_dim="time", + Which stream a given message came from varies; read ``.key``. Returns + ``None`` when a chunk held nothing for any stream, and raises + ``StopIteration`` only once the file is done. + + ``force_single_sample`` names irregular-rate streams that may carry more than + one event within ``chunk_dur``, which :obj:`AxisArray` cannot represent as a + single message with correct timestamps; those are split one event per + message. + """ + + @property + def exhausted(self) -> bool: + reader = self._state.reader + if reader is None: + return False + return reader.exhausted and self._state.pubqueue.empty() + + def _reset_state(self) -> None: + reader = self._build_reader(self.settings.select) + stream_names = [_["info"]["name"][0] for _ in reader._streams] + self._state.reader = reader + self._state.pubqueue = queue.SimpleQueue() + self._state.templates = { + name: _build_template( + reader._streams[stream_names.index(name)], + name=name, + n_ch=meta["channel_count"], + fs=meta["nominal_srate"], ) - self._pubqueue: queue.SimpleQueue[AxisArray] = queue.SimpleQueue() + for name, meta in reader.stream_meta.items() + } - def __next__(self) -> AxisArray | None: - if self._pubqueue.empty(): - chunk_dict = super().__next__() - for k, template in self._templates.items(): - if k in chunk_dict and len(chunk_dict[k][1]) > 0: - data, tvec = chunk_dict[k] - if k in self._force_single_sample: - if isinstance(template.axes["time"], AxisArray.CoordinateAxis): - t_kwargs = {"data": np.array([])} - else: - t_kwargs = {"offset": 0.0} - for ix, _t in enumerate(tvec): - if "data" in t_kwargs: - t_kwargs["data"] = np.array([_t]) - else: - t_kwargs["offset"] = _t - self._pubqueue.put_nowait( - replace( - template, - data=data[ix : ix + 1], - axes={ - **template.axes, - "time": replace(template.axes["time"], **t_kwargs), - }, - ) - ) - else: - if isinstance(template.axes["time"], AxisArray.CoordinateAxis): - t_kwargs = {"data": tvec if len(tvec) else np.array([])} - else: - t_kwargs = {"offset": tvec[0] if len(tvec) else self._last_time} - self._pubqueue.put_nowait( - replace( - template, - data=data, - axes={ - **template.axes, - "time": replace( - template.axes["time"], - **t_kwargs, - ), - }, - ) - ) + def _enqueue_chunk(self, chunk_dict: dict) -> None: + reader = self._state.reader + for name, template in self._state.templates.items(): + if name not in chunk_dict or len(chunk_dict[name][1]) == 0: + continue + data, tvec = chunk_dict[name] + if name in self.settings.force_single_sample: + for ix, stamp in enumerate(tvec): + self._state.pubqueue.put_nowait(_with_time(template, data[ix : ix + 1], np.array([stamp]), stamp)) + else: + self._state.pubqueue.put_nowait(_with_time(template, data, tvec, reader._last_time)) + + async def _produce(self) -> AxisArray | None: + if self._state.pubqueue.empty(): + try: + self._enqueue_chunk(next(self._state.reader)) + except StopIteration: + return None try: - return self._pubqueue.get_nowait() + return self._state.pubqueue.get_nowait() except queue.Empty: return None + + def __next__(self) -> AxisArray | None: + if self.exhausted: + raise StopIteration + return self() diff --git a/src/ezmsg/xdf/source.py b/src/ezmsg/xdf/source.py index db2bd7e..bc4a12e 100644 --- a/src/ezmsg/xdf/source.py +++ b/src/ezmsg/xdf/source.py @@ -1,14 +1,30 @@ import asyncio -import os import time import typing -from dataclasses import field import ezmsg.core as ez -from ezmsg.util.generator import GenState +from ezmsg.baseproc.units import BaseProducerUnit from ezmsg.util.messages.axisarray import AxisArray -from .iter import XDFAxisArrayIterator, XDFMultiAxArrIterator +from .iter import ( + XDFAxisArrayIterator, + XDFIteratorSettings, + XDFMultiAxArrIterator, + XDFMultiIteratorSettings, +) + +# The settings types moved to `iter.py` so the producers can own them, but they +# were importable from here first. +XDFMultiIteratorUnitSettings = XDFMultiIteratorSettings + +__all__ = [ + "PlaybackClock", + "XDFIteratorSettings", + "XDFIteratorUnit", + "XDFMultiIteratorSettings", + "XDFMultiIteratorUnit", + "XDFMultiIteratorUnitSettings", +] class PlaybackClock: @@ -48,109 +64,75 @@ def step(self) -> None: time.sleep(self._get_duration()) -class XDFIteratorSettings(ez.Settings): - filepath: typing.Union[os.PathLike, str] - select: str - chunk_dur: float = 1.0 - start_time: float | None = None - stop_time: float | None = None - rezero: bool = True - playback_rate: float | None = None - self_terminating: bool = False - """ - If True, the unit will raise a :obj:`ez.NormalTermination` exception when the file is exhausted. - Note, however, that this will terminate the pipeline even if the data published by this unit are still in transit, - which will lead to the pipeline output being truncated before it has finished processing the stream. - `self_terminating` should only be used when it is not important that the pipeline finish processing data, such - as during prototyping and testing. - """ +class _XDFUnitBase: + """Playback pacing and end-of-file handling, shared by both units. + Note both subclasses name their publisher ``produce``: that is the name + ``BaseProducerUnit`` uses, and ezmsg collects publishers per attribute, so a + differently named one would run *alongside* the base class\'s rather than + replacing it -- two publishers draining one producer, neither stopping. -class XDFIteratorUnit(ez.Unit): - STATE = GenState - SETTINGS = XDFIteratorSettings + The producer supplies chunks as fast as they can be sliced out of memory; + ``playback_rate`` is what turns that into a paced stream, and it is a + property of the unit rather than of the reader. + """ - OUTPUT_SIGNAL = ez.OutputStream(AxisArray) OUTPUT_TERM = ez.OutputStream(typing.Any) - def initialize(self) -> None: - self.construct_generator() - - def construct_generator(self): - self.STATE.gen = XDFAxisArrayIterator( - filepath=self.SETTINGS.filepath, - select=self.SETTINGS.select, - chunk_dur=self.SETTINGS.chunk_dur, - start_time=self.SETTINGS.start_time, - stop_time=self.SETTINGS.stop_time, - rezero=self.SETTINGS.rezero, + async def initialize(self) -> None: + await super().initialize() + self._clock = ( + PlaybackClock(rate=self.SETTINGS.playback_rate, step_dur=self.SETTINGS.chunk_dur) + if self.SETTINGS.playback_rate is not None + else None ) - if self.SETTINGS.playback_rate is not None: - self._clock = PlaybackClock(rate=self.SETTINGS.playback_rate, step_dur=self.SETTINGS.chunk_dur) - else: - self._clock = None - @ez.publisher(OUTPUT_SIGNAL) - async def pub_chunk(self) -> typing.AsyncGenerator: - try: - while True: - if self._clock is not None: - await self._clock.astep() - msg = next(self.STATE.gen) - if msg.data.size > 0: - yield self.OUTPUT_SIGNAL, msg - else: - await asyncio.sleep(0) - except StopIteration: - ez.logger.debug(f"File ({self.SETTINGS.filepath} :: {self.SETTINGS.select}) exhausted.") - if self.SETTINGS.self_terminating: - raise ez.NormalTermination - yield self.OUTPUT_TERM, True - - -class XDFMultiIteratorUnitSettings(XDFIteratorSettings): - select: set[str] | None = None # Override with a default - force_single_sample: set = field(default_factory=set) - - -class XDFMultiIteratorUnit(ez.Unit): - STATE = GenState - SETTINGS = XDFMultiIteratorUnitSettings + async def _finish(self) -> typing.AsyncGenerator: + ez.logger.debug(f"File ({self.SETTINGS.filepath} :: {self.SETTINGS.select}) exhausted.") + if self.SETTINGS.self_terminating: + raise ez.NormalTermination + yield self.OUTPUT_TERM, True + + +class XDFIteratorUnit( + _XDFUnitBase, + BaseProducerUnit[XDFIteratorSettings, AxisArray, XDFAxisArrayIterator], +): + SETTINGS = XDFIteratorSettings OUTPUT_SIGNAL = ez.OutputStream(AxisArray) - OUTPUT_TERM = ez.OutputStream(typing.Any) - def initialize(self) -> None: - self.construct_generator() - - def construct_generator(self): - self.STATE.gen = XDFMultiAxArrIterator( - filepath=self.SETTINGS.filepath, - select=self.SETTINGS.select, - chunk_dur=self.SETTINGS.chunk_dur, - start_time=self.SETTINGS.start_time, - stop_time=self.SETTINGS.stop_time, - rezero=self.SETTINGS.rezero, - force_single_sample=self.SETTINGS.force_single_sample, - ) - if self.SETTINGS.playback_rate is not None: - self._clock = PlaybackClock(rate=self.SETTINGS.playback_rate, step_dur=self.SETTINGS.chunk_dur) - else: - self._clock = None + @ez.publisher(OUTPUT_SIGNAL) + async def produce(self) -> typing.AsyncGenerator: + while not self.producer.exhausted: + if self._clock is not None: + await self._clock.astep() + msg = await self.producer.__acall__() + if msg is not None and msg.data.size > 0: + yield self.OUTPUT_SIGNAL, msg + else: + await asyncio.sleep(0) + async for out in self._finish(): + yield out + + +class XDFMultiIteratorUnit( + _XDFUnitBase, + BaseProducerUnit[XDFMultiIteratorSettings, AxisArray, XDFMultiAxArrIterator], +): + SETTINGS = XDFMultiIteratorSettings + + OUTPUT_SIGNAL = ez.OutputStream(AxisArray) @ez.publisher(OUTPUT_SIGNAL) - async def pub_multi(self) -> typing.AsyncGenerator: - try: - while True: - if self._clock is not None: - await self._clock.astep() - msg = next(self.STATE.gen) - if msg is not None: - yield self.OUTPUT_SIGNAL, msg - else: - await asyncio.sleep(0) - except StopIteration: - ez.logger.debug(f"File ({self.SETTINGS.filepath} :: {self.SETTINGS.select}) exhausted.") - if self.SETTINGS.self_terminating: - raise ez.NormalTermination - yield self.OUTPUT_TERM, True + async def produce(self) -> typing.AsyncGenerator: + while not self.producer.exhausted: + if self._clock is not None: + await self._clock.astep() + msg = await self.producer.__acall__() + if msg is not None: + yield self.OUTPUT_SIGNAL, msg + else: + await asyncio.sleep(0) + async for out in self._finish(): + yield out diff --git a/tests/test_iter.py b/tests/test_iter.py index 708942c..deb67c3 100644 --- a/tests/test_iter.py +++ b/tests/test_iter.py @@ -10,12 +10,21 @@ import math import pickle +import ezmsg.core as ez import numpy as np import pytest from conftest import EEG_STREAM, MARKER_STREAM from ezmsg.util.messages.axisarray import AxisArray +from ezmsg.util.messages.util import replace as replace_settings -from ezmsg.xdf.iter import XDFAxisArrayIterator, XDFIterator, XDFMultiAxArrIterator +from ezmsg.xdf.iter import ( + XDFAxisArrayIterator, + XDFIterator, + XDFIteratorSettings, + XDFMultiAxArrIterator, + XDFMultiIteratorSettings, +) +from ezmsg.xdf.source import XDFIteratorUnit, XDFMultiIteratorUnit def eeg_messages(path, **kwargs) -> list[AxisArray]: @@ -185,3 +194,113 @@ def test_the_samples_span_more_than_one_chunk(self, test_xdf_path): """Otherwise the reader's chunk stitching is never exercised.""" assert EEG_STREAM.n_samples > 32 assert math.ceil(EEG_STREAM.n_samples / 32) > 1 + + +class TestTheProducerContract: + """The producers are `BaseStatefulProducer`s, which means the file open is a + state reset rather than construction work.""" + + def test_construction_does_not_read_the_file(self, test_xdf_path): + """Deliberately unlike ezmsg-neo and ezmsg-nwb, which reset eagerly in + ``__init__`` and so pay the whole open on the event loop during + ``initialize``. Nothing here reads stream metadata before the first + chunk, so there is nothing to lose by waiting.""" + it = XDFAxisArrayIterator(filepath=test_xdf_path, select=EEG_STREAM.name) + assert it._state.reader is None + assert next(it) is not None + assert it._state.reader is not None + + def test_the_file_load_runs_off_the_event_loop(self, test_xdf_path): + """``pyxdf.load_xdf`` reads and decodes the whole file. On the event loop + that stalls every other unit in the process. + + Driven with ``asyncio.run`` rather than an async test, because this repo + has no pytest-asyncio -- the pytest config names ``asyncio_mode`` but the + plugin is not installed, so an ``async def`` test is silently never + awaited and passes without running. + """ + import asyncio + import threading + + seen: list[int] = [] + + class Spy(XDFAxisArrayIterator): + def _reset_state(self): + seen.append(threading.get_ident()) + super()._reset_state() + + producer = Spy(filepath=test_xdf_path, select=EEG_STREAM.name) + assert not seen, "construction should not have opened anything" + + loop_tid: list[int] = [] + + async def drive(): + loop_tid.append(threading.get_ident()) + await producer.__acall__() + + asyncio.run(drive()) + + assert len(seen) == 1 + assert seen[0] != loop_tid[0], "_reset_state ran on the event-loop thread" + + def test_settings_arrive_as_keywords_or_as_a_settings_object(self, test_xdf_path): + by_kwargs = XDFAxisArrayIterator(filepath=test_xdf_path, select=EEG_STREAM.name, chunk_dur=0.5) + by_settings = XDFAxisArrayIterator( + settings=XDFIteratorSettings(filepath=test_xdf_path, select=EEG_STREAM.name, chunk_dur=0.5) + ) + assert by_kwargs.settings == by_settings.settings + + def test_pacing_settings_do_not_reopen_the_file(self, test_xdf_path): + """``playback_rate`` and ``self_terminating`` belong to the unit, so + changing either must not throw away a loaded file.""" + it = XDFAxisArrayIterator(filepath=test_xdf_path, select=EEG_STREAM.name) + next(it) + reader = it._state.reader + it.update_settings(replace_settings(it.settings, playback_rate=2.0, self_terminating=True)) + next(it) + assert it._state.reader is reader + + def test_changing_the_file_does_reopen(self, test_xdf_path): + it = XDFAxisArrayIterator(filepath=test_xdf_path, select=EEG_STREAM.name) + next(it) + reader = it._state.reader + it.update_settings(replace_settings(it.settings, chunk_dur=0.25)) + next(it) + assert it._state.reader is not reader + + +class TestTheUnitsInAGraph: + @staticmethod + def _run(unit_cls, settings) -> list[AxisArray]: + collected: list[AxisArray] = [] + + class Collector(ez.Unit): + INPUT_SIGNAL = ez.InputStream(AxisArray) + + @ez.subscriber(INPUT_SIGNAL) + async def on_msg(self, msg: AxisArray) -> None: + collected.append(msg) + + src, sink = unit_cls(settings), Collector() + ez.run(SRC=src, SINK=sink, connections=((src.OUTPUT_SIGNAL, sink.INPUT_SIGNAL),)) + return collected + + def test_the_single_stream_unit_publishes_the_whole_file(self, test_xdf_path): + msgs = self._run( + XDFIteratorUnit, + XDFIteratorSettings(filepath=test_xdf_path, select=EEG_STREAM.name, self_terminating=True), + ) + assert msgs, "no messages published" + assert sum(m.data.shape[0] for m in msgs) == EEG_STREAM.n_samples + assert all(m.chunk_dim == "time" for m in msgs) + assert all("_fingerprint" in m.axes["ch"].__dict__ for m in msgs) + + def test_the_multi_stream_unit_publishes_both_streams(self, test_xdf_path): + msgs = self._run( + XDFMultiIteratorUnit, + XDFMultiIteratorSettings(filepath=test_xdf_path, self_terminating=True), + ) + assert {m.key for m in msgs} == {EEG_STREAM.name, MARKER_STREAM.name} + eeg = [m for m in msgs if m.key == EEG_STREAM.name] + assert sum(m.data.shape[0] for m in eeg) == EEG_STREAM.n_samples + assert all(m.chunk_dim == "time" for m in msgs)