diff --git a/pyproject.toml b/pyproject.toml index 0366cd3..7c88fcd 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -18,8 +18,8 @@ dependencies = [ dev = [ "pre-commit>=4.0.0", "scipy>=1.14.1", - "ezmsg-sigproc>=2.18.0", - "ezmsg-simbiophys>=1.4.1", + "ezmsg-sigproc>=2.36.0", + "ezmsg-simbiophys>=1.8.0", {include-group = "lint"}, {include-group = "test"}, ] diff --git a/src/ezmsg/tools/chmeta.py b/src/ezmsg/tools/chmeta.py new file mode 100644 index 0000000..88e5e0c --- /dev/null +++ b/src/ezmsg/tools/chmeta.py @@ -0,0 +1,87 @@ +"""Turning a structured ``ch`` coordinate axis into per-channel display names. + +An AxisArray's ``ch`` axis usually carries a structured array with one row per +channel. What is *in* those rows depends on the acquisition system: an LSL or +NWB source typically offers a ``label``; a Blackrock source additionally offers +``bank`` and ``elec``, which is what its users actually read off the front panel. +A plot needs one string per channel, so something has to choose. + +This module makes that choice explicit and configurable rather than hard-coding +one system's convention. ``label`` is the default because it is the field most +sources populate and the one most likely to be meaningful to whoever recorded +the data; a GUI that knows better can ask for other fields. +""" + +import typing + +import numpy as np + +__all__ = ["available_fields", "channel_names"] + + +def available_fields(ch_axis_data: typing.Optional[np.ndarray]) -> typing.List[str]: + """Field names present on a structured ``ch`` axis, for building a chooser.""" + if ch_axis_data is None: + return [] + return list((ch_axis_data.dtype.fields or {}).keys()) + + +def _field_text(row_value: typing.Any) -> str: + """One field of one channel as display text, or "" if it carries nothing. + + Numeric fields use 0 as their "unset" value -- a Blackrock ``elec`` of 0 + means no electrode, not electrode zero -- so it is treated as absent. Bytes + arrive from numpy ``S`` dtypes and are decoded leniently: a mangled label is + still more use than an exception. + """ + if isinstance(row_value, bytes): + row_value = row_value.decode("utf8", errors="replace") + if isinstance(row_value, np.generic): + row_value = row_value.item() + if isinstance(row_value, bytes): + row_value = row_value.decode("utf8", errors="replace") + if isinstance(row_value, (int, float)) and not isinstance(row_value, bool): + return "" if row_value == 0 else f"{row_value:g}" + return str(row_value).strip() + + +def channel_names( + ch_axis_data: typing.Optional[np.ndarray], + n_ch: typing.Optional[int] = None, + *, + fields: typing.Sequence[str] = ("label",), + sep: str = "-", + fallback: str = "ch{index}", +) -> typing.List[str]: + """Per-channel display names from a structured ``ch`` axis. + + :param ch_axis_data: The axis's structured array, or None if the source + provided no ``ch`` axis. + :param n_ch: Channel count, used only when ``ch_axis_data`` is None or + unstructured. Otherwise the axis's own length governs. + :param fields: Field names to join, in order. Fields absent from the dtype + are skipped; fields present but empty for a given channel are skipped for + that channel only, so a partially-populated column degrades per-row + rather than for the whole array. + :param sep: Joins the field values, e.g. ``("bank", "elec")`` -> ``"A-1"``. + :param fallback: Format string used when no requested field yields anything + for a channel. ``{index}`` is the channel's index. + + :return: One name per channel. Never empty strings, so a caller can render + the result without further guarding. + """ + if ch_axis_data is None or ch_axis_data.dtype.fields is None: + n = int(n_ch or 0) if ch_axis_data is None else int(ch_axis_data.shape[0]) + return [fallback.format(index=i) for i in range(n)] + + present = [f for f in fields if f in ch_axis_data.dtype.fields] + n = int(ch_axis_data.shape[0]) + if not present: + return [fallback.format(index=i) for i in range(n)] + + columns = [ch_axis_data[f] for f in present] + names = [] + for i in range(n): + parts = [t for t in (_field_text(col[i]) for col in columns) if t] + names.append(sep.join(parts) if parts else fallback.format(index=i)) + return names diff --git a/src/ezmsg/tools/shmem/aux_meta.py b/src/ezmsg/tools/shmem/aux_meta.py new file mode 100644 index 0000000..61ec934 --- /dev/null +++ b/src/ezmsg/tools/shmem/aux_meta.py @@ -0,0 +1,208 @@ +"""Codec for the *static* half of an AxisArray, carried alongside the shmem ring. + +The circular buffer in :mod:`.shmem` transports raw samples plus the little that +fits in a fixed ctypes header: dtype, shape, sample rate, and ``key``. Everything +else an :class:`~ezmsg.util.messages.axisarray.AxisArray` knows -- the ``ch`` +coordinate axis holding per-channel ``bank``/``elec``/``label``, the units, the +message ``attrs`` -- is dropped at the boundary. A consumer on the far side can +therefore plot the signal but cannot say what any of it *is*. + +This module encodes that dropped metadata into a self-describing byte blob which +:class:`~.shmem.ShMemCircBuff` publishes into its own shared-memory segment, and +which :class:`~.shmem_mirror.EZShmMirror` decodes on the far side. + +Wire format +----------- +A pickled ``dict`` of **plain Python and numpy types only** -- never ezmsg +classes. The two halves of a shmem link are separate processes and may be +separate environments with different ezmsg versions installed; pinning the wire +format to ezmsg's dataclass layout would make an upgrade on one side a silent +decode failure on the other. Plain dicts cost one down-conversion and buy +version independence. ``AUX_FORMAT_VERSION`` guards the shape of the dict +itself. + +Axes decode to:: + + {"kind": "linear", "unit": str, "gain": float, "offset": float} + {"kind": "coord", "unit": str, "dims": list[str], "data": np.ndarray} + +The buffered axis (normally ``time``) is a deliberate special case: its +``offset`` advances with every message and a coordinate time axis's ``data`` is +wholly new each message, so including either would make the metadata change +continuously and defeat the point of a low-rate side channel. Only its static +descriptors are kept -- see :func:`encode_aux`. +""" + +import pickle +import typing + +import numpy as np + +AUX_FORMAT_VERSION = 1 + +# Values we are willing to put on the wire. Anything else in ``attrs`` is +# dropped rather than pickled: an arbitrary object would force the decoding +# process to import the class that defines it, which is exactly the coupling +# this format exists to avoid. +_PLAIN_SCALARS = (str, bytes, int, float, bool, complex, type(None)) + + +def _is_plain(value: typing.Any) -> bool: + """Whether ``value`` is safe to pickle into the blob.""" + if isinstance(value, _PLAIN_SCALARS): + return True + if isinstance(value, np.ndarray): + # Object arrays hold arbitrary picklable classes; same objection. + return value.dtype != np.dtype("O") + if isinstance(value, np.generic): + return True + if isinstance(value, (list, tuple)): + return all(_is_plain(v) for v in value) + if isinstance(value, dict): + return all(isinstance(k, str) and _is_plain(v) for k, v in value.items()) + return False + + +def axis_to_plain(axis: typing.Any, *, static_only: bool = False) -> dict: + """Down-convert one ezmsg axis to plain types. + + ``static_only`` keeps just the descriptors that do not change from message to + message -- used for the buffered axis, whose position along the stream is + carried by the ring's write index rather than by this blob. + """ + unit = getattr(axis, "unit", "") + if hasattr(axis, "data"): # CoordinateAxis + if static_only: + return {"kind": "coord", "unit": unit, "dims": list(axis.dims)} + return { + "kind": "coord", + "unit": unit, + "dims": list(axis.dims), + "data": np.array(axis.data, copy=True), + } + out = {"kind": "linear", "unit": unit, "gain": float(axis.gain)} + if not static_only: + out["offset"] = float(axis.offset) + return out + + +def encode_aux( + dims: typing.Sequence[str], + axes: typing.Mapping[str, typing.Any], + attrs: typing.Mapping[str, typing.Any], + key: str, + buffered_axis: str, +) -> tuple[bytes, list[str]]: + """Serialize an AxisArray's static metadata. + + Returns the blob and the list of ``attrs`` keys that were dropped for not + being plain types, so the caller can log them once rather than per message. + """ + plain_axes = {name: axis_to_plain(ax, static_only=(name == buffered_axis)) for name, ax in axes.items()} + plain_attrs = {} + dropped = [] + for name, value in attrs.items(): + if isinstance(name, str) and _is_plain(value): + plain_attrs[name] = value + else: + dropped.append(str(name)) + payload = { + "version": AUX_FORMAT_VERSION, + "dims": list(dims), + "axes": plain_axes, + "attrs": plain_attrs, + "key": key, + "buffered_axis": buffered_axis, + } + return pickle.dumps(payload, protocol=pickle.HIGHEST_PROTOCOL), dropped + + +def decode_aux(blob: bytes) -> dict: + """Inverse of :func:`encode_aux`. + + :raises ValueError: if the blob is unreadable or was written by a format + version this build does not understand. + """ + try: + payload = pickle.loads(blob) + except Exception as exc: # noqa: BLE001 - any unpickling failure is the same failure to us + raise ValueError(f"could not decode shmem metadata blob: {exc}") from exc + if not isinstance(payload, dict): + raise ValueError(f"shmem metadata blob decoded to {type(payload).__name__}, expected dict") + version = payload.get("version") + if version != AUX_FORMAT_VERSION: + raise ValueError( + f"shmem metadata blob is format version {version!r}, this build understands {AUX_FORMAT_VERSION}" + ) + return payload + + +def _axis_equal(a: typing.Any, b: typing.Any) -> bool: + """Value equality for one axis, compared field by field. + + Deliberately does not use ``==``. As of ezmsg 3.6, ``CoordinateAxis.__eq__`` + resolves through the MRO to the dataclass-generated ``AxisBase.__eq__``, + which compares ``unit`` and nothing else -- ``ArrayWithNamedDims.__eq__``, + written to compare ``dims`` and ``data``, is shadowed and never runs. Two + coordinate axes with different data, or even different lengths, therefore + compare equal. Relying on that would mean a channel relabelling silently + never reaching the far side of the shmem link, which is the one thing this + module exists to deliver. Comparing explicitly also keeps the check correct + across ezmsg versions, which matters given the two halves of a link need not + share one. + """ + a_data = getattr(a, "data", None) + b_data = getattr(b, "data", None) + if (a_data is None) != (b_data is None): + return False + if getattr(a, "unit", "") != getattr(b, "unit", ""): + return False + if a_data is None: + return a.gain == b.gain and a.offset == b.offset + if list(a.dims) != list(b.dims): + return False + if a_data is b_data: + return True + if a_data.shape != b_data.shape or a_data.dtype != b_data.dtype: + return False + return bool(np.array_equal(a_data, b_data)) + + +def axes_equal(a: typing.Mapping[str, typing.Any], b: typing.Mapping[str, typing.Any]) -> bool: + """Cheap "have the axes changed?" test, for the per-message hot path. + + Identity is checked before value at every level, which is what makes this + affordable at kHz rates: an ezmsg processor that leaves an axis alone passes + the *same object* through, so the common case costs one pointer comparison + per axis. An element-wise comparison happens only when a producer rebuilt an + axis -- rare, and precisely the case we must not get wrong. + """ + if a is b: + return True + if a.keys() != b.keys(): + return False + for name, av in a.items(): + bv = b[name] + if av is bv: + continue + if type(av) is not type(bv): + return False + if not _axis_equal(av, bv): + return False + return True + + +def attrs_equal(a: typing.Mapping[str, typing.Any], b: typing.Mapping[str, typing.Any]) -> bool: + """Cheap "have the attrs changed?" test. + + Identity only. ``attrs`` values are arbitrary -- a numpy array's ``==`` + returns an array, and a user class's ``__eq__`` could be arbitrarily + expensive -- so value equality is deliberately not attempted here. A producer + that rebuilds equal attrs every message trips this check; the caller absorbs + that by comparing the encoded blob before it republishes anything. + """ + if a is b: + return True + if a.keys() != b.keys(): + return False + return all(a[name] is b[name] for name in a) diff --git a/src/ezmsg/tools/shmem/shmem.py b/src/ezmsg/tools/shmem/shmem.py index 8780e6f..d80ebcf 100644 --- a/src/ezmsg/tools/shmem/shmem.py +++ b/src/ezmsg/tools/shmem/shmem.py @@ -19,6 +19,12 @@ The other half must monitor the metadata shared memory to see if it changes, and if it does then it must recreate the data shared memory buffer reader at the new location. + +Finally, there is a third piece of shared memory carrying everything about the AxisArray that does not fit in the +fixed-size metadata header: the non-buffered coordinate axes (e.g. a `ch` axis holding per-channel bank/elec/label), +axis units, and the message `attrs`. It lives at shorten_shmem_name(f"{shmem_name}/meta{meta_generation}") and is +republished -- under a fresh generation, following the same pattern as the data buffer -- only when that metadata +actually changes, which for a typical stream means once per session. See the .aux_meta module for the wire format. """ import asyncio @@ -35,6 +41,8 @@ import numpy.typing as npt from ezmsg.util.messages.axisarray import AxisArray, AxisBase +from .aux_meta import attrs_equal, axes_equal, encode_aux + UINT64_SIZE = 8 BYTEORDER = "little" @@ -70,6 +78,28 @@ def shorten_shmem_name(long_name: str) -> str: MAXKEYLEN = 1024 +# Sentinel at offset 0 of every metadata segment ("EZMS"). Distinguishes one of +# our headers from an unrelated segment that happens to collide on a name, and +# from a header written by a build old enough to predate this check. +SHMEM_META_MAGIC = 0x455A4D53 + +# Bumped on any change to ShmemArrMeta._fields_ or to the .aux_meta wire format. +# +# The two halves of a shmem link must be the same version -- there is no +# compatibility shim, by choice: the layouts are an internal detail between two +# processes we deploy together, and carrying forward every past field shape would +# cost more than it is worth. What we do owe is a loud failure rather than a +# quiet one, so the reader validates the magic and version up front and raises +# instead of misreading a header it does not understand. +SHMEM_META_STRUCT_VERSION = 1 + + +class ShmemVersionError(RuntimeError): + """A shmem segment was written by an incompatible build. + + Not recoverable and not transient: upgrade both ends together. + """ + class ShmemArrMeta(ctypes.Structure): """ @@ -91,6 +121,10 @@ class ShmemArrMeta(ctypes.Structure): _pack_ = 1 _fields_ = [ + # magic and struct_version lead so a reader can validate the layout + # before it trusts a single field that follows. + ("magic", ctypes.c_uint32), + ("struct_version", ctypes.c_uint32), ("bvalid", ctypes.c_bool), ("dtype", ctypes.c_char), ("srate", ctypes.c_double), @@ -101,6 +135,11 @@ class ShmemArrMeta(ctypes.Structure): ("_key_bytes", ctypes.c_byte * MAXKEYLEN), ("_key_len", ctypes.c_uint32), ("write_index", ctypes.c_uint64), + # 0 = no metadata blob published yet. Otherwise names the segment at + # shorten_shmem_name(f"{shmem_name}/meta{meta_generation}"). + ("meta_generation", ctypes.c_uint32), + # Exact length of the blob; the segment itself is page-rounded. + ("aux_nbytes", ctypes.c_uint32), ] @property @@ -127,6 +166,16 @@ class ShMemCircBuffState(ez.State): buffer_shmem: typing.Optional[SharedMemory] = None buffer_arr: typing.Optional[npt.NDArray] = None meta_hash: int = -1 + # Segment holding the serialized static metadata (see .aux_meta). + aux_shmem: typing.Optional[SharedMemory] = None + # The (dims, axes, attrs, key) we last encoded, held by reference for the + # per-message identity check in _update_aux_if_needed. + last_aux_src: typing.Optional[tuple] = None + # ...and the bytes they encoded to, so a producer that rebuilds equal + # metadata every message cannot cause a republish. + last_aux_blob: typing.Optional[bytes] = None + # attrs keys dropped as non-plain, remembered so we warn once, not per message. + warned_dropped_attrs: typing.Optional[frozenset] = None def _persist_create_shmem(name: str, size: int) -> SharedMemory: @@ -201,6 +250,7 @@ def on_settings(self, msg: ShMemCircBuffSettings) -> None: async def shutdown(self) -> None: self._cleanup_buffer() + self._cleanup_aux() self._cleanup_meta() if self.SETTINGS.conn is not None: self.SETTINGS.conn.send("close") @@ -215,6 +265,7 @@ def _cleanup_meta(self): if self.SETTINGS.conn is not None: self.SETTINGS.conn.send("meta cleanup") + self._cleanup_aux() self.STATE.meta_struct = None if self.STATE.meta_shmem is not None: @@ -226,6 +277,21 @@ def _cleanup_meta(self): del self.STATE.meta_shmem self.STATE.meta_shmem = None + def _cleanup_aux(self): + """ + Release the static-metadata segment, if one is published. + + Also forgets the change-detection state, so the next message republishes + from scratch -- which is what we want after a name change or a shutdown / + restart, where a reader may be starting fresh too. + """ + self._cleanup_aux_segment() + self.STATE.last_aux_src = None + self.STATE.last_aux_blob = None + if self.STATE.meta_struct is not None: + self.STATE.meta_struct.meta_generation = 0 + self.STATE.meta_struct.aux_nbytes = 0 + def _cleanup_buffer(self): """ Destroy the data buffer and the shared memory object. @@ -274,10 +340,88 @@ def _reset_meta(self, reset_generation: bool = True) -> None: # Build the metadata structure. self.STATE.meta_struct = ShmemArrMeta.from_buffer(self.STATE.meta_shmem.buf) self.STATE.meta_struct.bvalid = False + self.STATE.meta_struct.magic = SHMEM_META_MAGIC + self.STATE.meta_struct.struct_version = SHMEM_META_STRUCT_VERSION + self.STATE.meta_struct.meta_generation = 0 + self.STATE.meta_struct.aux_nbytes = 0 if reset_generation: self.STATE.meta_struct.buffer_generation = -1 # We will wait for a data packet before we modify the remaining fields. + def _update_aux_if_needed(self, msg: AxisArray) -> bool: + """ + Republish the static metadata segment if this message's metadata differs + from what is currently published. + + Runs on every message, so the common path must be cheap. It is three + tiers, each only reached when the one before it is inconclusive: + + 1. Identity/value comparison of (dims, axes, attrs, key) against what we + last encoded. Costs a handful of pointer comparisons when the producer + passes its axes through untouched, which is the normal case. + 2. Encode, and compare the bytes to what is published. This absorbs + producers that rebuild equal metadata every message -- they cost an + encode, but never a republish, so a reader is never woken for nothing. + 3. Allocate a new generation's segment and point the header at it. + + Returns True if a new generation was published. + """ + src = (msg.dims, msg.axes, msg.attrs, msg.key) + last = self.STATE.last_aux_src + if last is not None: + last_dims, last_axes, last_attrs, last_key = last + if ( + msg.key == last_key + and msg.dims == last_dims + and axes_equal(msg.axes, last_axes) + and attrs_equal(msg.attrs, last_attrs) + ): + return False + + blob, dropped = encode_aux(msg.dims, msg.axes, msg.attrs, msg.key, self.SETTINGS.axis) + if dropped: + dropped_set = frozenset(dropped) + if self.STATE.warned_dropped_attrs != dropped_set: + self.STATE.warned_dropped_attrs = dropped_set + ez.logger.warning( + f"ShMemCircBuff dropped non-plain attrs from the shmem metadata: {sorted(dropped_set)}. " + "Only str/bytes/number/bool/None, non-object ndarrays, and containers of those are transported." + ) + + # Hold the references that produced this blob whether or not we go on to + # publish it, so an unchanged-but-rebuilt message is only encoded once. + self.STATE.last_aux_src = src + if blob == self.STATE.last_aux_blob: + return False + self.STATE.last_aux_blob = blob + + self._cleanup_aux_segment() + # 0 means "nothing published", so skip it when the uint32 wraps. + generation = (self.STATE.meta_struct.meta_generation + 1) % (2**32) or 1 + aux_name = shorten_shmem_name(self.SETTINGS.shmem_name + "/meta" + str(generation)) + self.STATE.aux_shmem = _persist_create_shmem(aux_name, len(blob)) + self.STATE.aux_shmem.buf[: len(blob)] = blob + + # Order matters: the segment is fully written before the header names it, + # so a reader never sees a generation it cannot completely read. + self.STATE.meta_struct.aux_nbytes = len(blob) + self.STATE.meta_struct.meta_generation = generation + + if self.SETTINGS.conn is not None: + self.SETTINGS.conn.send("aux updated") + return True + + def _cleanup_aux_segment(self) -> None: + """Release just the published segment, keeping change-detection state.""" + if self.STATE.aux_shmem is not None: + self.STATE.aux_shmem.close() + try: + self.STATE.aux_shmem.unlink() + except FileNotFoundError: + pass + del self.STATE.aux_shmem + self.STATE.aux_shmem = None + def _n_frames_for_axis(self, axis: AxisBase) -> int: """ Utility function to calculate the number of frames to allocate for the buffer. @@ -408,6 +552,12 @@ async def on_message(self, msg: AxisArray): if self._update_meta_if_needed(msg): self._reset_buffer(msg) + # Independently of the buffer: republish the static metadata if it moved. + # The two are deliberately not coupled -- a `ch` axis can gain labels + # without the buffer's shape changing, and the buffer can be rebuilt + # (e.g. dtype change) with the channel identities untouched. + self._update_aux_if_needed(msg) + n_samples = data.shape[0] write_stop = self.STATE.meta_struct.write_index + n_samples diff --git a/src/ezmsg/tools/shmem/shmem_mirror.py b/src/ezmsg/tools/shmem/shmem_mirror.py index 084e165..c38af94 100644 --- a/src/ezmsg/tools/shmem/shmem_mirror.py +++ b/src/ezmsg/tools/shmem/shmem_mirror.py @@ -2,6 +2,11 @@ It is possible to move data from ezmsg to non-ezmsg processes using shared memory. This module contains the non-ezmsg half of that communication. The ezmsg half is found in .shmem. The same `shmem_name` must be passed to both the ShMemCircBuff and the EZShmMirror objects! + +Besides the sample data, the mirror exposes the source AxisArray's static metadata -- its coordinate axes (e.g. a `ch` +axis naming each channel), axis units, and `attrs` -- via the `axes`, `attrs`, and `dims` properties. These are plain +dicts rather than ezmsg objects; see .aux_meta for why. They read None until the writer publishes, and update in place +if it ever republishes, so poll them (or register_metadata_callback) rather than reading once. """ import copy @@ -12,7 +17,15 @@ import numpy as np import numpy.typing as npt -from .shmem import ShmemArrMeta, ShMemCircBuffState, shorten_shmem_name +from .aux_meta import decode_aux +from .shmem import ( + SHMEM_META_MAGIC, + SHMEM_META_STRUCT_VERSION, + ShmemArrMeta, + ShMemCircBuffState, + ShmemVersionError, + shorten_shmem_name, +) CONNECT_RETRY_INTERVAL = 0.5 @@ -32,9 +45,14 @@ def __init__(self, shmem_name: typing.Optional[str] = None): self._mirror_state: ShMemCircBuffState = ShMemCircBuffState() self._shmem_name: typing.Optional[str] = None self._change_callback: typing.Optional[typing.Callable] = None + self._metadata_callback: typing.Optional[typing.Callable] = None self._last_meta: typing.Optional[ShmemArrMeta] = None self._read_index = 0 # Used by auto_view self._last_connect_try = -np.inf + # Decoded static metadata (see .aux_meta) and the generation it came + # from. 0 means we have not read one; the writer never publishes gen 0. + self._aux: typing.Optional[dict] = None + self._aux_generation: int = 0 # If shmem_name is None then this will simply not connect to anything. self.connect(shmem_name) @@ -64,7 +82,121 @@ def write_index(self) -> typing.Optional[int]: def connected(self) -> bool: return self.buffer is not None + # ---- Static metadata (the non-buffered axes, units, and attrs) ---------- + + @property + def axes(self) -> typing.Optional[typing.Dict[str, dict]]: + """The source AxisArray's axes as plain dicts, or None if unavailable. + + Keyed by axis name. Each value is + ``{"kind": "linear", "unit", "gain", "offset"}`` or + ``{"kind": "coord", "unit", "dims", "data"}`` -- see :mod:`.aux_meta` + for why these are dicts rather than ezmsg axis objects. + + The buffered axis (whatever the sink was configured to buffer along, + normally ``"time"``) appears here with only its static descriptors: its + position along the stream lives in the ring's write index, not here. + + None means the writer has not published yet -- poll again. A writer this + build cannot read raises :class:`ShmemVersionError` on connect rather + than showing up as None here. + """ + self._refresh_aux() + return None if self._aux is None else self._aux["axes"] + + @property + def attrs(self) -> typing.Optional[dict]: + """The source AxisArray's ``attrs``, minus any non-transportable values.""" + self._refresh_aux() + return None if self._aux is None else self._aux["attrs"] + + @property + def dims(self) -> typing.Optional[typing.List[str]]: + """The source AxisArray's dimension names, in the *sender's* order. + + Note the buffer itself is rolled so the buffered axis comes first; this + is the message's original ordering. + """ + self._refresh_aux() + return None if self._aux is None else self._aux["dims"] + + @property + def metadata_available(self) -> bool: + """Whether a decoded metadata blob is currently held.""" + self._refresh_aux() + return self._aux is not None + + def register_metadata_callback(self, callback: typing.Callable) -> None: + """Call ``callback`` whenever a new metadata generation is decoded. + + Separate from :meth:`register_change_callback`, which fires when the + *data buffer* is rebuilt. The two are independent: channel labels can + arrive without the buffer changing, and vice versa. + """ + self._metadata_callback = callback + + def unregister_metadata_callback(self) -> None: + self._metadata_callback = None + + def _cleanup_aux(self): + if self._mirror_state.aux_shmem is not None: + try: + self._mirror_state.aux_shmem.close() + except Exception as e: + print(f"Error closing metadata segment: {e}") + del self._mirror_state.aux_shmem + self._mirror_state.aux_shmem = None + self._aux = None + self._aux_generation = 0 + + def _refresh_aux(self) -> None: + """Attach to and decode the metadata segment if the writer bumped it. + + Cheap and idempotent: in the steady state this is one integer compare, + so the properties above can call it unconditionally. + + The header was already validated on connect, so an undecodable blob here + is a bug rather than a version skew, and propagates. + """ + meta = self._mirror_state.meta_struct + if meta is None: + return + generation = int(meta.meta_generation) + if generation == 0 or generation == self._aux_generation: + return + + nbytes = int(meta.aux_nbytes) + aux_name = shorten_shmem_name(self._shmem_name + "/meta" + str(generation)) + try: + shm = SharedMemory(aux_name, create=False) + except FileNotFoundError: + # The writer has moved on to a newer generation and unlinked this + # one. Leave the old decode in place; the next poll picks up the new + # generation. + return + + try: + payload = decode_aux(bytes(shm.buf[:nbytes])) + except ValueError: + shm.close() + raise + + # Only now release the previous segment, so a decode failure above + # leaves the last good metadata intact. + if self._mirror_state.aux_shmem is not None: + try: + self._mirror_state.aux_shmem.close() + except Exception as e: + print(f"Error closing metadata segment: {e}") + self._mirror_state.aux_shmem = shm + self._aux = payload + self._aux_generation = generation + + if self._metadata_callback is not None: + self._metadata_callback() + def _cleanup_meta(self): + self._cleanup_aux() if self._mirror_state.meta_shmem is not None: del self._mirror_state.meta_struct self._mirror_state.meta_struct = None @@ -109,6 +241,42 @@ def _connect_meta(self): except FileNotFoundError: self._mirror_state.meta_struct = None self._mirror_state.meta_shmem = None + return + self._validate_header() + + def _validate_header(self) -> None: + """Reject a header this build cannot read, before trusting any field. + + The two ends of a shmem link must be the same version. That is a + deliberate simplification -- the layout is private between two processes + we deploy together -- and it makes this check the thing that has to be + reliable, since the failure it prevents is reading a differently-shaped + struct as though it were ours and plotting the result. + + Raises rather than returning a status because there is nothing a caller + can usefully do: it is not transient, and it will not fix itself on the + next poll. + """ + meta = self._mirror_state.meta_struct + if meta is None: + return + magic, version = int(meta.magic), int(meta.struct_version) + if magic == SHMEM_META_MAGIC and version == SHMEM_META_STRUCT_VERSION: + return + + self._cleanup_meta() + if magic != SHMEM_META_MAGIC: + raise ShmemVersionError( + f"Shared memory segment for {self._shmem_name!r} does not carry this build's header " + f"(magic 0x{magic:08X}, expected 0x{SHMEM_META_MAGIC:08X}). Either it was written by an " + f"ezmsg-tools too old to stamp one, or the name collides with an unrelated segment. " + f"The writer and reader of a shmem link must be the same ezmsg-tools version." + ) + raise ShmemVersionError( + f"Shared memory segment for {self._shmem_name!r} was written by ezmsg-tools with shmem struct " + f"version {version}; this build speaks version {SHMEM_META_STRUCT_VERSION}. " + f"Upgrade both ends together." + ) def _reset_buffer(self) -> bool: if self._mirror_state.buffer_shmem is not None: @@ -170,6 +338,10 @@ def auto_view(self, n: typing.Optional[int] = None) -> typing.Tuple[npt.NDArray, if self._mirror_state.meta_struct is None: self.connect(self._shmem_name) + # Poll the metadata here too, so a consumer that only ever calls + # auto_view still gets its metadata callback fired. + self._refresh_aux() + if self._mirror_state.meta_struct is None or not self._mirror_state.meta_struct.bvalid: # Still not connected # or we are connected but the buffer data is invalid. diff --git a/tests/test_shmem_aux_meta.py b/tests/test_shmem_aux_meta.py new file mode 100644 index 0000000..0dfc4ce --- /dev/null +++ b/tests/test_shmem_aux_meta.py @@ -0,0 +1,436 @@ +"""The static-metadata side channel: codec, change detection, and end-to-end. + +The point of the feature is that a consumer on the far side of the shared-memory +boundary can learn what the channels *are*, so these tests care about two things +in roughly equal measure: that the metadata arrives intact, and that carrying it +costs nothing when it is not changing -- a republish per message would defeat the +purpose. +""" + +import ctypes +import os +import tempfile +import threading +import time +import typing +from dataclasses import replace +from pathlib import Path + +import ezmsg.core as ez +import numpy as np +import pytest +from ezmsg.util.messages.axisarray import AxisArray + +from ezmsg.tools.chmeta import available_fields, channel_names +from ezmsg.tools.shmem.aux_meta import ( + AUX_FORMAT_VERSION, + attrs_equal, + axes_equal, + decode_aux, + encode_aux, +) +from ezmsg.tools.shmem.shmem import ShMemCircBuff +from ezmsg.tools.shmem.shmem_mirror import EZShmMirror + +CHANNEL_DTYPE = np.dtype([("bank", "U2"), ("elec", " AxisArray.CoordinateAxis: + data = np.zeros(n_ch, dtype=CHANNEL_DTYPE) + for i in range(n_ch): + data["bank"][i] = "AB"[i // 32] + data["elec"][i] = (i % 32) + 1 + data["label"][i] = f"elec{i:03d}" + return AxisArray.CoordinateAxis(data=data, dims=["ch"], unit="") + + +def make_msg(n_time: int = 8, n_ch: int = 4, offset: float = 0.0, **attrs) -> AxisArray: + return AxisArray( + data=np.zeros((n_time, n_ch), dtype=np.float32), + dims=["time", "ch"], + axes={ + "time": AxisArray.TimeAxis(fs=1000.0, offset=offset), + "ch": make_ch_axis(n_ch), + }, + attrs=dict(attrs), + key="test", + ) + + +# ---------------------------------------------------------------- codec ----- + + +def test_encode_decode_round_trip(): + msg = make_msg(n_ch=64, unit="uV") + blob, dropped = encode_aux(msg.dims, msg.axes, msg.attrs, msg.key, "time") + assert dropped == [] + + payload = decode_aux(blob) + assert payload["version"] == AUX_FORMAT_VERSION + assert payload["dims"] == ["time", "ch"] + assert payload["key"] == "test" + assert payload["attrs"] == {"unit": "uV"} + + ch = payload["axes"]["ch"] + assert ch["kind"] == "coord" + assert ch["dims"] == ["ch"] + np.testing.assert_array_equal(ch["data"], msg.axes["ch"].data) + + +def test_buffered_axis_carries_no_position(): + """The time axis's offset advances every message; it must not ride along. + + If it did, the blob would differ on every message and the side channel would + republish at the sample rate. + """ + msg = make_msg(offset=1.0) + blob_a, _ = encode_aux(msg.dims, msg.axes, msg.attrs, msg.key, "time") + later = replace(msg, axes={**msg.axes, "time": AxisArray.TimeAxis(fs=1000.0, offset=99.0)}) + blob_b, _ = encode_aux(later.dims, later.axes, later.attrs, later.key, "time") + + assert blob_a == blob_b + time_axis = decode_aux(blob_a)["axes"]["time"] + assert "offset" not in time_axis + assert time_axis["gain"] == pytest.approx(0.001) + + +def test_coordinate_time_axis_carries_no_data(): + """Same argument for an irregular stream, where time is a CoordinateAxis + whose data is wholly new each message.""" + msg = make_msg() + a = replace(msg, axes={**msg.axes, "time": AxisArray.CoordinateAxis(data=np.arange(8.0), dims=["time"], unit="s")}) + b = replace( + msg, axes={**msg.axes, "time": AxisArray.CoordinateAxis(data=np.arange(100.0, 108.0), dims=["time"], unit="s")} + ) + blob_a, _ = encode_aux(a.dims, a.axes, a.attrs, a.key, "time") + blob_b, _ = encode_aux(b.dims, b.axes, b.attrs, b.key, "time") + + assert blob_a == blob_b + assert "data" not in decode_aux(blob_a)["axes"]["time"] + + +def test_non_plain_attrs_are_dropped_not_pickled(): + class Opaque: + pass + + msg = make_msg(unit="uV", handle=Opaque(), count=3) + blob, dropped = encode_aux(msg.dims, msg.axes, msg.attrs, msg.key, "time") + + assert dropped == ["handle"] + assert decode_aux(blob)["attrs"] == {"unit": "uV", "count": 3} + + +def test_decode_rejects_foreign_payloads(): + import pickle + + with pytest.raises(ValueError, match="could not decode"): + decode_aux(b"not a pickle at all") + with pytest.raises(ValueError, match="expected dict"): + decode_aux(pickle.dumps([1, 2, 3])) + with pytest.raises(ValueError, match="format version"): + decode_aux(pickle.dumps({"version": AUX_FORMAT_VERSION + 1})) + + +# ------------------------------------------------------ change detection ----- + + +def test_axes_equal_is_true_for_passed_through_axes(): + msg = make_msg() + forwarded = replace(msg, data=msg.data * 2) # axes dict passed through unchanged + assert axes_equal(msg.axes, forwarded.axes) + + +def test_axes_equal_is_true_for_rebuilt_but_identical_axes(): + """A producer that rebuilds an equal ch axis must not look like a change.""" + a, b = make_msg(), make_msg() + assert a.axes["ch"] is not b.axes["ch"] + assert axes_equal(a.axes, b.axes) + + +def test_axes_equal_detects_real_changes(): + a = make_msg(n_ch=4) + relabelled = make_ch_axis(4) + relabelled.data["label"][2] = "CHANGED" + b = replace(a, axes={**a.axes, "ch": relabelled}) + assert not axes_equal(a.axes, b.axes) + + assert not axes_equal(a.axes, {k: v for k, v in a.axes.items() if k != "ch"}) + + +def test_attrs_equal_tolerates_array_values(): + """dict == would raise on an ndarray value; identity comparison must not.""" + arr = np.arange(4) + a = {"m": arr} + assert attrs_equal(a, a) + assert attrs_equal(a, {"m": arr}) + assert not attrs_equal(a, {"m": np.arange(4)}) # conservative: re-encode, then blob-compare + + +# ------------------------------------------------------------ chmeta -------- + + +def test_channel_names_defaults_to_label(): + ch = make_ch_axis(4) + assert channel_names(ch.data) == ["elec000", "elec001", "elec002", "elec003"] + assert available_fields(ch.data) == ["bank", "elec", "label"] + + +def test_channel_names_with_alternate_fields(): + ch = make_ch_axis(34) + names = channel_names(ch.data, fields=("bank", "elec")) + assert names[0] == "A-1" + assert names[32] == "B-1" + + +def test_channel_names_fallbacks(): + assert channel_names(None, 3) == ["ch0", "ch1", "ch2"] + # Unstructured axis data. + assert channel_names(np.arange(2.0), 2) == ["ch0", "ch1"] + # Requested field is absent from the dtype. + ch = make_ch_axis(2) + assert channel_names(ch.data, fields=("nonexistent",)) == ["ch0", "ch1"] + # Field present but empty for one channel only. + partial = make_ch_axis(2) + partial.data["label"][1] = "" + assert channel_names(partial.data) == ["elec000", "ch1"] + + +# ------------------------------------------------------------- e2e ---------- + + +class MetaCountSettings(ez.Settings): + relabel_at: int = -1 + n_messages: int = 20 + + +class MetaCountState(ez.State): + count: int = 0 + + +class Source(ez.Unit): + """Emits a fixed stream, optionally relabelling the ch axis partway.""" + + SETTINGS = MetaCountSettings + STATE = MetaCountState + + OUTPUT_SIGNAL = ez.OutputStream(AxisArray) + + @ez.publisher(OUTPUT_SIGNAL) + async def generate(self) -> typing.AsyncGenerator: + n_ch = 8 + base_ch = make_ch_axis(n_ch) + alt_ch = make_ch_axis(n_ch) + alt_ch.data["label"] = [f"NEW{i:03d}" for i in range(n_ch)] + while self.STATE.count < self.SETTINGS.n_messages: + i = self.STATE.count + relabelled = 0 <= self.SETTINGS.relabel_at <= i + yield ( + self.OUTPUT_SIGNAL, + AxisArray( + data=np.full((10, n_ch), float(i), dtype=np.float32), + dims=["time", "ch"], + # A fresh axes dict every message, as a real producer builds. + axes={ + "time": AxisArray.TimeAxis(fs=1000.0, offset=i * 0.01), + "ch": alt_ch if relabelled else base_ch, + }, + attrs={"unit": "uV"}, + key="e2e", + ), + ) + self.STATE.count += 1 + await __import__("asyncio").sleep(0.01) + raise ez.NormalTermination + + +def _run_graph(shmem_name: str, relabel_at: int, n_messages: int) -> None: + comps = { + "SRC": Source(relabel_at=relabel_at, n_messages=n_messages), + "SINK": ShMemCircBuff(shmem_name, 2.0, conn=None, axis="time"), + } + conns = ((comps["SRC"].OUTPUT_SIGNAL, comps["SINK"].INPUT_SIGNAL),) + ez.run(components=comps, connections=conns) + + +@pytest.mark.skipif("CI" in os.environ, reason="Timing-sensitive; matches the existing mirror test.") +@pytest.mark.parametrize("relabel_at", [-1, 100]) +def test_metadata_reaches_the_mirror(relabel_at: int): + shmem_name = f"auxtest{os.getpid()}{relabel_at}" + n_messages = 200 # ~2 s at the source's 10 ms cadence + + mirror = EZShmMirror() + mirror.connect(shmem_name) + + generations = [] + mirror.register_metadata_callback(lambda: generations.append(mirror._aux_generation)) + + thread = threading.Thread(target=_run_graph, args=(shmem_name, relabel_at, n_messages)) + thread.start() + + # Snapshots taken while the graph runs: the writer unlinks its segments at + # shutdown, so anything not read now is gone. + seen_labels = [] + last_attrs = last_dims = last_time_axis = None + deadline = time.time() + 30.0 + while thread.is_alive() and time.time() < deadline: + mirror.auto_view() + axes = mirror.axes + if axes is not None: + labels = list(axes["ch"]["data"]["label"]) + if not seen_labels or labels != seen_labels[-1]: + seen_labels.append(labels) + last_attrs, last_dims, last_time_axis = mirror.attrs, mirror.dims, axes["time"] + time.sleep(0.005) + thread.join(timeout=10.0) + + assert seen_labels, "no metadata ever arrived at the mirror" + assert seen_labels[0][0] == "elec000" + assert last_attrs == {"unit": "uV"} + assert last_dims == ["time", "ch"] + + # The buffered axis is present but positionless. + assert "offset" not in last_time_axis + assert last_time_axis["gain"] == pytest.approx(0.001) + + if relabel_at < 0: + # Steady metadata over 200 messages must publish exactly one generation. + assert seen_labels == [seen_labels[0]] + assert generations == [1] + else: + assert len(seen_labels) == 2, f"expected one relabel, saw {len(seen_labels)} distinct label sets" + assert seen_labels[-1][0] == "NEW000" + assert generations == [1, 2] + + mirror.disconnect() + + +@pytest.mark.skipif("CI" in os.environ, reason="Timing-sensitive; matches the existing mirror test.") +def test_mirror_without_metadata_is_not_broken(): + """A stream whose sink never publishes metadata still mirrors data fine.""" + shmem_name = f"auxnone{os.getpid()}" + file_path = Path(tempfile.gettempdir()) / "test_aux_none.txt" + file_path.unlink(missing_ok=True) + + mirror = EZShmMirror() + mirror.connect(shmem_name) + # Nothing published yet: the properties must answer, not raise. + assert mirror.axes is None + assert mirror.attrs is None + assert not mirror.metadata_available + mirror.disconnect() + + +# --------------------------------------------------- N-D + versioning ------- + + +class Envelope(ez.Unit): + """Emits a 3-D (time, ch, metric) envelope. + + The shape ``ezmsg.sigproc.binned_aggregate.BinnedAggregate`` produces with a + tuple ``operation`` -- the display path's reason for wanting N-D here. + """ + + OUTPUT_SIGNAL = ez.OutputStream(AxisArray) + + @ez.publisher(OUTPUT_SIGNAL) + async def generate(self) -> typing.AsyncGenerator: + import asyncio + + n_ch, n_metric = 8, 2 + for i in range(200): + data = np.zeros((10, n_ch, n_metric), dtype=np.float32) + data[..., 0] = -float(i) # min + data[..., 1] = float(i) # max + yield ( + self.OUTPUT_SIGNAL, + AxisArray( + data=data, + dims=["time", "ch", "metric"], + axes={ + "time": AxisArray.TimeAxis(fs=1000.0, offset=i * 0.01), + "ch": make_ch_axis(n_ch), + "metric": AxisArray.CoordinateAxis(data=np.array(["min", "max"]), dims=["metric"], unit=""), + }, + key="envelope", + ), + ) + await asyncio.sleep(0.01) + raise ez.NormalTermination + + +@pytest.mark.skipif("CI" in os.environ, reason="Timing-sensitive; matches the existing mirror test.") +def test_three_dimensional_message_survives_the_boundary(): + """A (time, ch, metric) envelope must cross shmem intact. + + ShMemCircBuff has always been N-D -- frame_shape is every dim but the + buffered one -- so an envelope needs no flattening into 2-D. This pins that, + including that the metric axis's labels arrive over the metadata channel. + """ + shmem_name = f"aux3d{os.getpid()}" + + mirror = EZShmMirror() + mirror.connect(shmem_name) + + comps = {"SRC": Envelope(), "SINK": ShMemCircBuff(shmem_name, 2.0, conn=None, axis="time")} + conns = ((comps["SRC"].OUTPUT_SIGNAL, comps["SINK"].INPUT_SIGNAL),) + thread = threading.Thread(target=lambda: ez.run(components=comps, connections=conns)) + thread.start() + + shapes, metric_labels, samples = [], None, None + deadline = time.time() + 30.0 + while thread.is_alive() and time.time() < deadline: + chunk, _ = mirror.auto_view() + if chunk.size: + shapes.append(chunk.shape) + samples = chunk.copy() + axes = mirror.axes + if axes is not None and "metric" in axes: + metric_labels = list(axes["metric"]["data"]) + time.sleep(0.005) + thread.join(timeout=10.0) + + assert shapes, "no data crossed the boundary" + # (n_samples, n_ch, n_metric) -- the trailing dims came through untouched. + assert all(s[1:] == (8, 2) for s in shapes), shapes + # min entries are negative, max entries positive, so the metric axis did not + # get transposed or collapsed on the way. + assert (samples[..., 0] <= 0).all() + assert (samples[..., 1] >= 0).all() + assert metric_labels == ["min", "max"] + + mirror.disconnect() + + +def test_reader_rejects_a_foreign_header_loudly(): + """A header this build cannot read must raise, not decode to nonsense.""" + from multiprocessing.shared_memory import SharedMemory + + from ezmsg.tools.shmem.shmem import ( + SHMEM_META_MAGIC, + SHMEM_META_STRUCT_VERSION, + ShmemArrMeta, + ShmemVersionError, + shorten_shmem_name, + ) + + name = f"badhdr{os.getpid()}" + shm = SharedMemory(shorten_shmem_name(name), create=True, size=ctypes.sizeof(ShmemArrMeta)) + try: + meta = ShmemArrMeta.from_buffer(shm.buf) + + # No magic at all: a segment from a build that predates the check, or an + # unrelated segment whose name collided. + meta.magic = 0 + with pytest.raises(ShmemVersionError, match="does not carry this build's header"): + EZShmMirror(name) + + # Our magic, a version we don't speak. + meta.magic = SHMEM_META_MAGIC + meta.struct_version = SHMEM_META_STRUCT_VERSION + 1 + with pytest.raises(ShmemVersionError, match="Upgrade both ends together"): + EZShmMirror(name) + + del meta + finally: + shm.close() + shm.unlink()