"""
It is possible to move data from ezmsg to non-ezmsg processes using shared memory. This module contains the ezmsg half
of that communication. The non-ezmsg half is in the .shmem_mirror module.
The same `shmem_name` must be passed to both the ShMemCircBuff and the EZShmMirror objects!
The ShMemCircBuff class is a sink node that receives AxisArray messages and writes them to a shared memory buffer.
Upon initialization, or upon receiving updated settings with a different shmem_name value, the node creates a shared
memory object located at {shorten_shmem_name(shmem_name)} to hold the metadata initialized with placeholder values
(e.g., srate = -1).
Additionally, the node has a convenience handle to the metadata via
`self.STATE.meta_struct = ShmemArrMeta.from_buffer(shmem.buf)`.
Upon receiving a data message, its metadata is checked, and if it does not match the shmem metadata
(which will always be true for the first message) then the node first updates the metadata, then it (re-)creates
a shared memory buffer to hold the data, located at shorten_shmem_name(f"{shmem_name}/buffer{buffer_generation}"),
where `buffer_generation` is an integer that tracks how many times the buffer has been reset. This corresponds to the
same integer stored in the metadata.
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
import base64
import ctypes
import hashlib
import multiprocessing.connection
import time
import typing
from multiprocessing.shared_memory import SharedMemory
import ezmsg.core as ez
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"
[docs]
def to_bytes(data: typing.Any) -> bytes:
if isinstance(data, bool):
return data.to_bytes(2, byteorder=BYTEORDER, signed=False)
elif isinstance(data, int):
return np.int64(data).to_bytes(UINT64_SIZE, BYTEORDER, signed=False)
[docs]
def shorten_shmem_name(long_name: str) -> str:
"""
Convert a potentially long shared memory name to a shorter, fixed-length name.
Args:
long_name: The original, potentially long shared memory name
Returns:
A shortened, deterministic name suitable for shared memory
"""
if long_name is None:
return None
# Create a hash of the original name
hash_obj = hashlib.sha256(long_name.encode("utf-8"))
# Convert to URL-safe base64 and limit to 20 characters (plus 'sm_' prefix)
# The 'sm_' prefix helps identify this as a shared memory name
short_name = "sm_" + base64.urlsafe_b64encode(hash_obj.digest()).decode("ascii")[:20]
return short_name
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
[docs]
class ShmemVersionError(RuntimeError):
"""A shmem segment was written by an incompatible build.
Not recoverable and not transient: upgrade both ends together.
"""
[docs]
class ShMemCircBuffSettings(ez.Settings):
shmem_name: typing.Optional[str]
buf_dur: float
conn: typing.Optional[multiprocessing.connection.Connection] = None
axis: typing.Optional[str] = None
"""Dimension to buffer along. ``None`` follows the message's ``stream_dim``.
The ring is a history of the stream, so this has to be the dimension
messages accumulate along; buffering a static one would store the same
elements over and over. Only the producer reliably knows which that is --
it is ``time`` on a raw signal and ``win`` downstream of a windowing stage.
The old default of ``"time"`` was silently wrong for the latter. It is
*present* in a ``(win, time, ch)`` message, so nothing rejected it: the
window count ended up inside ``frame_shape`` -- reallocating the buffer
whenever the window count jittered -- and the reported sample rate was the
within-window rate, a 10x error in the viewer's time base for a 10-sample
window. Set explicitly only for a producer that declares no ``stream_dim``.
"""
[docs]
class ShMemCircBuffState(ez.State):
meta_shmem: typing.Optional[SharedMemory] = None
meta_struct: typing.Optional[ShmemArrMeta] = None
buffer_shmem: typing.Optional[SharedMemory] = None
buffer_arr: typing.Optional[npt.NDArray] = None
meta_hash: int = -1
# The dimension currently buffered along; see ShMemCircBuffSettings.axis.
buff_axis: typing.Optional[str] = None
# 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, purpose: str = "") -> SharedMemory:
"""
Create a shared memory object, retrying if necessary.
Args:
name: The name of the shared memory object.
size: The size of the shared memory object.
purpose: What this segment is for, for the log line. Names are hashed
to fit the platform's length limit, so without this a reader cannot
tell the data ring from the metadata blob -- and a shape change
recreates both, back to back.
Returns: The SharedMemory object.
"""
t0 = time.time()
n_attempts = 0
while True:
try:
result = SharedMemory(
name=name,
create=True,
size=size,
)
break
except FileExistsError:
n_attempts += 1
tmp_shmem = SharedMemory(
name=name,
create=False,
)
tmp_shmem.close()
tmp_shmem.unlink()
retried = f" after {n_attempts} stale-name retries," if n_attempts else ""
ez.logger.info(f"Created {purpose or 'shmem'} ({size} bytes) at {name}{retried} in {time.time() - t0:.3f} s.")
return result
[docs]
class ShMemCircBuff(ez.Unit):
SETTINGS = ShMemCircBuffSettings
STATE = ShMemCircBuffState
INPUT_SIGNAL = ez.InputStream(AxisArray)
INPUT_SETTINGS = ez.InputStream(ShMemCircBuffSettings)
[docs]
async def initialize(self) -> None:
# Prophylactic cleanup. These should mostly be a no-ops
# because the shared memory objects should not exist yet.
self._cleanup_buffer()
self._cleanup_meta()
# Create the metadata shared memory object.
# Our SETTINGS need valid values for shmem_name and buf_dur.
# Even then, the meta_struct will be invalid until we receive
# a data packet.
self._reset_meta()
[docs]
@ez.subscriber(INPUT_SETTINGS)
def on_settings(self, msg: ShMemCircBuffSettings) -> None:
b_reset_meta = msg.shmem_name != self.SETTINGS.shmem_name
b_reset_buff = msg.buf_dur != self.SETTINGS.buf_dur
b_reset_buff = b_reset_buff or msg.axis != self.SETTINGS.axis
self.apply_settings(msg)
if b_reset_buff or b_reset_meta:
# First we destroy the data buffer, because it is no
# longer valid with the new settings.
self._cleanup_buffer()
# It will be recreated with the next data packet.
if b_reset_meta:
# Destroy the metadata and its shared memory object because the name has changed.
self._cleanup_meta()
# Then we reset the metadata to the new name.
self._reset_meta(reset_generation=False)
# Do not reset the buffer. We will wait for a new data packet.
[docs]
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")
self.SETTINGS.conn.close()
def _cleanup_meta(self):
"""
Destroy the metadata and its shared memory object.
This is called during initialization and shutdown,
or if the SETTINGS name has changed.
"""
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:
self.STATE.meta_shmem.close()
try:
self.STATE.meta_shmem.unlink()
except FileNotFoundError:
pass
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.
This is called during initialization and shutdown,
as well as whenever the incoming data changes its statistics
requiring a new buffer.
"""
if self.SETTINGS.conn is not None:
self.SETTINGS.conn.send("buffer cleanup")
if self.STATE.meta_struct is not None:
# Mark the metadata as invalid
self.STATE.meta_struct.bvalid = False
# Destroy the buffer
self.STATE.buffer_arr = None
# Destroy the shared memory object
if self.STATE.buffer_shmem is not None:
self.STATE.buffer_shmem.close()
try:
self.STATE.buffer_shmem.unlink()
except FileNotFoundError:
pass
if self.STATE.buffer_shmem is not None:
del self.STATE.buffer_shmem
self.STATE.buffer_shmem = None
def _reset_meta(self, reset_generation: bool = True) -> None:
"""
Crete the metadata shared memory object.
This is called during initialization and whenever the SETTINGS.shmem_name changes.
"""
if self.SETTINGS.conn is not None:
self.SETTINGS.conn.send("meta reset")
# Create the metadata shared memory object.
meta_size = int(ctypes.sizeof(ShmemArrMeta))
short_name = shorten_shmem_name(self.SETTINGS.shmem_name)
self.STATE.meta_shmem = _persist_create_shmem(short_name, meta_size, purpose="shmem header")
if self.SETTINGS.shmem_name is None:
# If the name is None, then we need to get the name from the shared memory object.
self.SETTINGS.shmem_name = self.STATE.meta_shmem.name
# 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
# The ring rolls the buffered axis to the front (see on_message), and
# meta.shape already describes that order, so dims must too -- a reader
# given the sender's original order would have to know to re-roll it,
# which is knowledge it has no way to arrive at.
buff_axis = self.STATE.buff_axis
rolled_dims = [buff_axis] + [d for d in msg.dims if d != buff_axis]
blob, dropped = encode_aux(rolled_dims, msg.axes, msg.attrs, msg.key, buff_axis, stream_dim=msg.stream_dim)
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), purpose=f"stream metadata gen {generation}")
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 _resolve_axis(self, msg: AxisArray) -> typing.Optional[str]:
"""The dimension to buffer along, or None if this message has none.
An explicit setting wins so an operator can still drive a producer that
declares nothing; otherwise the message decides.
"""
axis = self.SETTINGS.axis if self.SETTINGS.axis is not None else msg.stream_dim
if axis is None:
axis = "time" if "time" in msg.dims else None
return axis
def _n_frames_for_axis(self, axis: AxisBase) -> int:
"""
Utility function to calculate the number of frames to allocate for the buffer.
Args:
axis: The axis object containing the metadata for the axis along
which we are buffering.
Returns: number of frames we should buffer based on the axis and settings.
"""
if hasattr(axis, "data"):
fs = 1 / np.median(np.diff(axis.data)) if len(axis.data) > 1 else 100.0
else:
fs = 1 / axis.gain
return int(np.ceil(self.SETTINGS.buf_dur * fs))
def _get_msg_meta(self, msg: AxisArray) -> typing.Tuple[bytes, float, int, typing.Tuple[int, ...]]:
"""
Utility function to extract relevant metadata from the incoming message.
Args:
msg: The incoming AxisArray message.
Returns:
A tuple of metadata extracted from the message.
msg_dtype, msg_srate, n_frames, frame_shape
"""
buff_axis = self.STATE.buff_axis
ax_idx = msg.get_axis_idx(buff_axis)
axis = msg.axes[buff_axis]
n_frames = self._n_frames_for_axis(axis)
frame_shape = msg.data.shape[:ax_idx] + msg.data.shape[ax_idx + 1 :]
data = np.moveaxis(msg.data, ax_idx, 0)
msg_dtype = data.dtype.char.encode("utf8")
msg_srate = 1 / axis.gain if hasattr(axis, "gain") else 0.0
return msg_dtype, msg_srate, n_frames, frame_shape
def _update_meta_if_needed(self, msg: AxisArray) -> bool:
"""
Update the metadata structure if the incoming message has different metadata.
Args:
msg: The incoming AxisArray message.
Returns: True if the metadata was updated, False otherwise.
"""
# Extract the metadata from the incoming message
msg_dtype, msg_srate, n_frames, frame_shape = self._get_msg_meta(msg)
# Get its hash for quick comparison, and we will reuse the hash.
new_hash = hash(
(
msg_dtype,
msg_srate,
n_frames,
)
+ frame_shape
+ (msg.key,)
)
b_update = self.STATE.meta_hash != new_hash
if b_update:
if self.SETTINGS.conn is not None:
self.SETTINGS.conn.send("begin update")
self.STATE.meta_struct.bvalid = False
self.STATE.meta_struct.dtype = msg_dtype
self.STATE.meta_struct.srate = msg_srate
self.STATE.meta_struct.ndim = 1 + len(frame_shape)
self.STATE.meta_struct.shape[: self.STATE.meta_struct.ndim] = (n_frames,) + frame_shape
self.STATE.meta_struct.key = msg.key
self.STATE.meta_struct.write_index = 0
self.STATE.meta_struct.wrap_counter = 0
self.STATE.meta_hash = new_hash
if self.SETTINGS.conn is not None:
self.SETTINGS.conn.send("meta updated")
return b_update
def _reset_buffer(self, msg: AxisArray) -> None:
"""
Reset the buffer to accommodate the new metadata and new message.
Args:
msg: The incoming AxisArray message.
"""
self.STATE.meta_struct.buffer_generation += 1
msg_dtype, msg_srate, n_frames, frame_shape = self._get_msg_meta(msg)
buff_size = int(n_frames * np.prod(frame_shape) * msg.data.itemsize)
buff_shm_name = self.SETTINGS.shmem_name + "/buffer" + str(self.STATE.meta_struct.buffer_generation)
short_name = shorten_shmem_name(buff_shm_name)
self.STATE.buffer_shmem = _persist_create_shmem(
short_name,
buff_size,
purpose=f"data ring gen {self.STATE.meta_struct.buffer_generation} "
f"({'x'.join(str(d) for d in (n_frames,) + frame_shape)})",
)
self.STATE.buffer_arr = np.ndarray(
self.STATE.meta_struct.shape[: self.STATE.meta_struct.ndim],
dtype=np.dtype(self.STATE.meta_struct.dtype.decode("utf8")),
buffer=self.STATE.buffer_shmem.buf[:],
)
self.STATE.meta_struct.write_index = 0
self.STATE.meta_struct.wrap_counter = 0
self.STATE.meta_struct.bvalid = True
if self.SETTINGS.conn is not None:
self.SETTINGS.conn.send("buffer reset")
[docs]
@ez.task
async def check_continue(self):
while True:
if self.SETTINGS.conn is not None and self.SETTINGS.conn.poll():
obj = self.SETTINGS.conn.recv()
if obj == "quit":
self.shutdown()
break
else:
print(f"Unhandled object received on connection: {obj}")
else:
await asyncio.sleep(0.05)
raise ez.NormalTermination
[docs]
@ez.subscriber(INPUT_SIGNAL, zero_copy=True)
async def on_message(self, msg: AxisArray):
# Sanity check the input
if not isinstance(msg, AxisArray):
return
buff_axis = self._resolve_axis(msg)
if buff_axis is None or buff_axis not in msg.dims:
return
if buff_axis != self.STATE.buff_axis:
# The dimension the stream accumulates along changed under us -- a
# windowing stage inserted upstream, say. The buffer describes the
# old one, so it cannot be appended to.
self.STATE.buff_axis = buff_axis
self._cleanup_buffer()
ax_idx = msg.get_axis_idx(buff_axis)
data = np.moveaxis(msg.data, ax_idx, 0)
# Check if we need to update the metadata, and if so, reset the buffer.
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
if write_stop > self.STATE.buffer_arr.shape[0]:
overflow = write_stop - self.STATE.buffer_arr.shape[0]
self.STATE.buffer_arr[self.STATE.meta_struct.write_index :] = data[: n_samples - overflow]
self.STATE.buffer_arr[:overflow] = data[n_samples - overflow :]
self.STATE.meta_struct.write_index = overflow
self.STATE.meta_struct.wrap_counter += 1
else:
self.STATE.buffer_arr[self.STATE.meta_struct.write_index : write_stop] = data[:]
self.STATE.meta_struct.write_index = write_stop