ezmsg-tools 0.3.2__tar.gz → 0.4.0__tar.gz

{ezmsg_tools-0.3.2 → ezmsg_tools-0.4.0}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: ezmsg-tools
-Version: 0.3.2
+Version: 0.4.0
 Summary: ezmsg namespace package containing tools for working with ezmsg graphs.
 Author-email: Chadwick Boulay <chadwick.boulay@gmail.com>
 License-Expression: MIT

{ezmsg_tools-0.3.2 → ezmsg_tools-0.4.0}/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"},
 ]

ezmsg_tools-0.4.0/src/ezmsg/tools/__version__.py

@@ -0,0 +1,24 @@
+# file generated by vcs-versioning
+# don't change, don't track in version control
+from __future__ import annotations
+
+__all__ = [
+    "__version__",
+    "__version_tuple__",
+    "version",
+    "version_tuple",
+    "__commit_id__",
+    "commit_id",
+]
+
+version: str
+__version__: str
+__version_tuple__: tuple[int | str, ...]
+version_tuple: tuple[int | str, ...]
+commit_id: str | None
+__commit_id__: str | None
+
+__version__ = version = '0.4.0'
+__version_tuple__ = version_tuple = (0, 4, 0)
+
+__commit_id__ = commit_id = None

ezmsg_tools-0.4.0/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

ezmsg_tools-0.4.0/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)

{ezmsg_tools-0.3.2 → ezmsg_tools-0.4.0}/src/ezmsg/tools/shmem/shmem.py

@@ -19,6 +19,12 @@ Upon receiving a data message, its metadata is checked, and if it does not match
 
 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 as np
 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 @@ class ShMemCircBuff(ez.Unit):
 
     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 @@ class ShMemCircBuff(ez.Unit):
         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 @@ class ShMemCircBuff(ez.Unit):
             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 @@ class ShMemCircBuff(ez.Unit):
         # 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 @@ class ShMemCircBuff(ez.Unit):
         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