Replace busy-wait publishers in ResampleConcat and ResampleUnit with event-driven publishing

src/ezmsg/sigproc/resample.py

@@ -416,19 +416,45 @@ class ResampleUnit(BaseConsumerUnit[ResampleSettings, AxisArray, ResampleProcess
     :class:`~ezmsg.sigproc.resampleconcat.ResampleConcat`, which fuses both steps).
     """
 
+    async def initialize(self) -> None:
+        await super().initialize()
+        self._wake = asyncio.Event()
+
+    @ez.subscriber(BaseConsumerUnit.INPUT_SIGNAL)
+    async def on_signal(self, message: AxisArray) -> None:
+        await self.processor.__acall__(message)
+        self._wake.set()
+
     @ez.subscriber(INPUT_REFERENCE)
     async def on_reference(self, message: AxisArray):
         self.processor.push_reference(message)
+        self._wake.set()
 
     @ez.publisher(OUTPUT_SIGNAL)
     @ez.publisher(OUTPUT_REFERENCE)
     async def gen_resampled(self):
         while True:
-            result: AxisArray = next(self.processor)
-            if np.prod(result.data.shape) > 0:
+            # Wake on a push from either input. In prescribed-rate mode with a
+            # finite max_chunk_delay, also wake when the delay budget elapses so
+            # the wall-clock extrapolation in `__next__` can fire without input.
+            # Reference-driven mode never becomes ready by wall-clock, so there
+            # a pure event wait suffices.
+            timeout = self.SETTINGS.max_chunk_delay if self.SETTINGS.resample_rate is not None else np.inf
+            if np.isfinite(timeout):
+                try:
+                    await asyncio.wait_for(self._wake.wait(), timeout=timeout)
+                except asyncio.TimeoutError:
+                    pass
+            else:
+                await self._wake.wait()
+            # Clear before draining: a set() landing mid-drain re-arms the event
+            # so the next wait() returns immediately instead of losing the push.
+            self._wake.clear()
+            while True:
+                result: AxisArray = next(self.processor)
+                if np.prod(result.data.shape) == 0:
+                    break
                 yield self.OUTPUT_SIGNAL, result
                 ref_out = self.processor.state.reference_output
                 if ref_out is not None:
                     yield self.OUTPUT_REFERENCE, ref_out
-            else:
-                await asyncio.sleep(0)

src/ezmsg/sigproc/resampleconcat.py

@@ -26,7 +26,6 @@
 
 from __future__ import annotations
 
-import asyncio
 import typing
 
 import ezmsg.core as ez
@@ -152,19 +151,32 @@ class ResampleConcat(ez.Unit):
     async def initialize(self) -> None:
         self.processor = ResampleConcatProcessor(self.SETTINGS)
 
+    def _drain(self) -> typing.Iterator[AxisArray]:
+        """Yield every chunk the processor can currently produce.
+
+        Event-driven draining after each push is lossless here: the composed
+        resampler is always reference-driven (``resample_rate=None``), and in
+        that mode output readiness only ever changes on new input -- the
+        wall-clock ``max_chunk_delay`` extrapolation applies to prescribed-rate
+        mode only. A single ``next()`` consumes all currently-eligible
+        reference values, so this loop runs at most twice.
+        """
+        while True:
+            result = next(self.processor)
+            if result is None or np.prod(result.data.shape) == 0:
+                return
+            yield result
+
     @ez.subscriber(INPUT_REFERENCE)
-    async def on_reference(self, message: AxisArray) -> None:
+    @ez.publisher(OUTPUT_SIGNAL)
+    async def on_reference(self, message: AxisArray) -> typing.AsyncGenerator:
         self.processor.push_reference(message)
+        for out in self._drain():
+            yield self.OUTPUT_SIGNAL, out
 
     @ez.subscriber(INPUT_SIGNAL)
-    async def on_signal(self, message: AxisArray) -> None:
-        self.processor.push_signal(message)
-
     @ez.publisher(OUTPUT_SIGNAL)
-    async def output(self) -> typing.AsyncGenerator:
-        while True:
-            result = next(self.processor)
-            if result is not None and np.prod(result.data.shape) > 0:
-                yield self.OUTPUT_SIGNAL, result
-            else:
-                await asyncio.sleep(0)
+    async def on_signal(self, message: AxisArray) -> typing.AsyncGenerator:
+        self.processor.push_signal(message)
+        for out in self._drain():
+            yield self.OUTPUT_SIGNAL, out

tests/integration/ezmsg/test_resample_system.py

@@ -0,0 +1,128 @@
+"""Integration tests for the ResampleUnit's event-driven publisher.
+
+The publisher waits on an asyncio.Event set by both input handlers instead of
+polling. Two wake paths need proving through a live graph:
+
+* reference-driven mode: output flows on pushes alone (pure event wait), and
+  the graph terminates once sources go quiet;
+* prescribed-rate mode with a finite ``max_chunk_delay``: the wall-clock
+  extrapolation in ``ResampleProcessor.__next__`` must still fire when input
+  stops, which requires the timed wake.
+"""
+
+import os
+import typing
+
+import ezmsg.core as ez
+import numpy as np
+from ezmsg.util.messagecodec import message_log
+from ezmsg.util.messagelogger import MessageLogger, MessageLoggerSettings
+from ezmsg.util.messages.axisarray import AxisArray
+from ezmsg.util.terminate import TerminateOnTimeout, TerminateOnTimeoutSettings
+
+from ezmsg.sigproc.resample import ResampleSettings, ResampleUnit
+from tests.helpers.util import get_test_fn
+
+
+def _mk(fs: float, offset: float, n: int, n_ch: int, key: str) -> AxisArray:
+    """Build a [time, ch] message where ``data[:, j] == time + 1000*j``."""
+    t = offset + np.arange(n) / fs
+    data = t[:, None] + np.arange(n_ch)[None, :] * 1000.0
+    return AxisArray(
+        data=data,
+        dims=["time", "ch"],
+        axes={"time": AxisArray.LinearAxis(gain=1 / fs, offset=offset, unit="s")},
+        key=key,
+    )
+
+
+class DualSourceSettings(ez.Settings):
+    n_msgs: int = 20
+    chunk: int = 30
+    fs_ref: float = 100.0
+    fs_sig: float = 99.7
+    emit_reference: bool = True
+
+
+class DualSource(ez.Unit):
+    """Emit interleaved reference/signal chunks, then go quiet."""
+
+    SETTINGS = DualSourceSettings
+
+    OUTPUT_REFERENCE = ez.OutputStream(AxisArray)
+    OUTPUT_SIGNAL = ez.OutputStream(AxisArray)
+
+    @ez.publisher(OUTPUT_REFERENCE)
+    @ez.publisher(OUTPUT_SIGNAL)
+    async def produce(self) -> typing.AsyncGenerator:
+        s = self.SETTINGS
+        for i in range(s.n_msgs):
+            if s.emit_reference:
+                yield self.OUTPUT_REFERENCE, _mk(s.fs_ref, i * s.chunk / s.fs_ref, s.chunk, 1, "ref")
+            yield self.OUTPUT_SIGNAL, _mk(s.fs_sig, i * s.chunk / s.fs_sig, s.chunk, 2, "sig")
+
+
+def _run_graph(resample_settings: ResampleSettings, source_settings: DualSourceSettings, term_time: float) -> list:
+    test_filename = get_test_fn(None)
+    comps = {
+        "SRC": DualSource(source_settings),
+        "RESAMPLE": ResampleUnit(resample_settings),
+        "LOG": MessageLogger(MessageLoggerSettings(output=test_filename)),
+        "TERM": TerminateOnTimeout(TerminateOnTimeoutSettings(time=term_time)),
+    }
+    conns = (
+        (comps["SRC"].OUTPUT_REFERENCE, comps["RESAMPLE"].INPUT_REFERENCE),
+        (comps["SRC"].OUTPUT_SIGNAL, comps["RESAMPLE"].INPUT_SIGNAL),
+        (comps["RESAMPLE"].OUTPUT_SIGNAL, comps["LOG"].INPUT_MESSAGE),
+        (comps["LOG"].OUTPUT_MESSAGE, comps["TERM"].INPUT),
+    )
+    ez.run(components=comps, connections=conns)
+    messages = [_ for _ in message_log(test_filename)]
+    os.remove(test_filename)
+    return messages
+
+
+def test_resample_system_reference_driven():
+    """Pushes alone must wake the publisher; no polling loop exists to find output."""
+    messages = _run_graph(
+        ResampleSettings(axis="time", resample_rate=None, buffer_duration=4.0),
+        DualSourceSettings(),
+        term_time=1.0,
+    )
+    assert messages, "ResampleUnit published no output in reference-driven mode."
+    cat = AxisArray.concatenate(*messages, dim="time")
+    t = cat.axes["time"].value(np.arange(cat.data.shape[0]))
+    assert np.all(np.diff(t) > 0), "Output time axis is not monotonic."
+    # Signal was built with data[:, 0] == time, so resampling onto the
+    # reference grid must reproduce the grid itself.
+    assert np.max(np.abs(cat.data[:, 0] - t)) < 1e-6
+
+
+def test_resample_system_prescribed_rate():
+    """Prescribed-rate mode publishes on signal pushes alone (no reference input).
+
+    Note: this does NOT assert the ``max_chunk_delay`` wall-clock extrapolation
+    past the end of input. That path is currently unreachable at the processor
+    level regardless of how the unit polls or wakes: after a drain the source
+    buffer retains only ~2 samples, and ``ResampleProcessor.__next__`` returns
+    early on ``src.available() < 3`` before evaluating ``b_project``. The
+    unit's timed wake preserves the intended trigger; if the processor guard is
+    reworked, extend this test to assert output beyond the input end.
+    """
+    n_msgs, chunk, fs = 5, 30, 100.0
+    max_chunk_delay = 0.2
+    messages = _run_graph(
+        ResampleSettings(
+            axis="time",
+            resample_rate=90.0,
+            buffer_duration=4.0,
+            max_chunk_delay=max_chunk_delay,
+        ),
+        DualSourceSettings(n_msgs=n_msgs, chunk=chunk, fs_sig=fs, emit_reference=False),
+        term_time=3 * max_chunk_delay,
+    )
+    assert messages, "ResampleUnit published no output in prescribed-rate mode."
+    cat = AxisArray.concatenate(*messages, dim="time")
+    t = cat.axes["time"].value(np.arange(cat.data.shape[0]))
+    assert np.all(np.diff(t) > 0), "Output time axis is not monotonic."
+    assert np.max(np.abs(cat.data[:, 0] - t)) < 1e-6

tests/integration/ezmsg/test_resampleconcat_system.py

@@ -0,0 +1,101 @@
+"""Integration test for the ResampleConcat unit's event-driven publishing.
+
+The unit publishes from its two subscriber handlers (no polling loop), so this
+test verifies that outputs actually flow through a live pub/sub graph when
+reference and signal chunks arrive interleaved, and that the graph terminates
+once the sources go quiet (i.e., nothing depends on a background publisher).
+"""
+
+import os
+import typing
+
+import ezmsg.core as ez
+import numpy as np
+from ezmsg.util.messagecodec import message_log
+from ezmsg.util.messagelogger import MessageLogger, MessageLoggerSettings
+from ezmsg.util.messages.axisarray import AxisArray
+from ezmsg.util.terminate import TerminateOnTimeout, TerminateOnTimeoutSettings
+
+from ezmsg.sigproc.resampleconcat import ResampleConcat, ResampleConcatSettings
+from tests.helpers.util import get_test_fn
+
+
+def _mk(fs: float, offset: float, n: int, n_ch: int, key: str) -> AxisArray:
+    """Build a [time, ch] message where ``data[:, j] == time + 1000*j``."""
+    t = offset + np.arange(n) / fs
+    data = t[:, None] + np.arange(n_ch)[None, :] * 1000.0
+    return AxisArray(
+        data=data,
+        dims=["time", "ch"],
+        axes={
+            "time": AxisArray.LinearAxis(gain=1 / fs, offset=offset, unit="s"),
+            "ch": AxisArray.CoordinateAxis(
+                data=np.array([f"{key}{i}" for i in range(n_ch)]), dims=["ch"], unit="label"
+            ),
+        },
+        key=key,
+    )
+
+
+class DualSourceSettings(ez.Settings):
+    n_msgs: int = 20
+    chunk: int = 30
+    fs_ref: float = 100.0
+    fs_sig: float = 99.7
+
+
+class DualSource(ez.Unit):
+    """Emit interleaved reference/signal chunks, then go quiet."""
+
+    SETTINGS = DualSourceSettings
+
+    OUTPUT_REFERENCE = ez.OutputStream(AxisArray)
+    OUTPUT_SIGNAL = ez.OutputStream(AxisArray)
+
+    @ez.publisher(OUTPUT_REFERENCE)
+    @ez.publisher(OUTPUT_SIGNAL)
+    async def produce(self) -> typing.AsyncGenerator:
+        s = self.SETTINGS
+        for i in range(s.n_msgs):
+            yield self.OUTPUT_REFERENCE, _mk(s.fs_ref, i * s.chunk / s.fs_ref, s.chunk, 4, "h1")
+            yield self.OUTPUT_SIGNAL, _mk(s.fs_sig, i * s.chunk / s.fs_sig, s.chunk, 2, "h2")
+
+
+def test_resampleconcat_system(test_name: str | None = None):
+    test_filename = get_test_fn(test_name)
+
+    comps = {
+        "SRC": DualSource(DualSourceSettings()),
+        "RSC": ResampleConcat(
+            ResampleConcatSettings(
+                axis="time",
+                concat_axis="ch",
+                buffer_duration=4.0,
+                label_a="_h1",
+                label_b="_h2",
+            )
+        ),
+        "LOG": MessageLogger(MessageLoggerSettings(output=test_filename)),
+        "TERM": TerminateOnTimeout(TerminateOnTimeoutSettings(time=1.0)),
+    }
+    conns = (
+        (comps["SRC"].OUTPUT_REFERENCE, comps["RSC"].INPUT_REFERENCE),
+        (comps["SRC"].OUTPUT_SIGNAL, comps["RSC"].INPUT_SIGNAL),
+        (comps["RSC"].OUTPUT_SIGNAL, comps["LOG"].INPUT_MESSAGE),
+        (comps["LOG"].OUTPUT_MESSAGE, comps["TERM"].INPUT),
+    )
+    ez.run(components=comps, connections=conns)
+
+    messages: list[AxisArray] = [_ for _ in message_log(test_filename)]
+    os.remove(test_filename)
+
+    assert messages, "ResampleConcat unit published no output."
+    cat = AxisArray.concatenate(*messages, dim="time")
+    assert cat.data.shape[1] == 6, "Expected 4 (reference) + 2 (signal) = 6 channels."
+
+    t = cat.axes["time"].value(np.arange(cat.data.shape[0]))
+    # Output time axis must be strictly monotonic even though both handlers publish.
+    assert np.all(np.diff(t) > 0), "Output time axis is not monotonic."
+    # Side A is the reference gathered on the output grid; side B is resampled onto it.
+    assert np.max(np.abs(cat.data[:, 0] - t)) < 1e-9
+    assert np.max(np.abs(cat.data[:, 4] - t)) < 1e-6