ezmsg-baseproc 1.0__py3-none-any.whl

ezmsg/baseproc/units.py

@@ -0,0 +1,282 @@
+"""Base Unit classes for ezmsg integration."""
+
+import math
+import traceback
+import typing
+from abc import ABC, abstractmethod
+
+import ezmsg.core as ez
+from ezmsg.util.generator import GenState
+from ezmsg.util.messages.axisarray import AxisArray
+
+from .composite import CompositeProcessor
+from .processor import BaseConsumer, BaseProducer, BaseTransformer
+from .protocols import MessageInType, MessageOutType, SettingsType
+from .stateful import BaseAdaptiveTransformer, BaseStatefulConsumer, BaseStatefulTransformer
+from .util.message import SampleMessage
+from .util.profile import profile_subpub
+from .util.typeresolution import resolve_typevar
+
+# --- Type variables for Unit classes ---
+ProducerType = typing.TypeVar("ProducerType", bound=BaseProducer)
+ConsumerType = typing.TypeVar("ConsumerType", bound=BaseConsumer | BaseStatefulConsumer)
+TransformerType = typing.TypeVar(
+    "TransformerType",
+    bound=BaseTransformer | BaseStatefulTransformer | CompositeProcessor,
+)
+AdaptiveTransformerType = typing.TypeVar("AdaptiveTransformerType", bound=BaseAdaptiveTransformer)
+
+
+def get_base_producer_type(cls: type) -> type:
+    return resolve_typevar(cls, ProducerType)
+
+
+def get_base_consumer_type(cls: type) -> type:
+    return resolve_typevar(cls, ConsumerType)
+
+
+def get_base_transformer_type(cls: type) -> type:
+    return resolve_typevar(cls, TransformerType)
+
+
+def get_base_adaptive_transformer_type(cls: type) -> type:
+    return resolve_typevar(cls, AdaptiveTransformerType)
+
+
+# --- Base classes for ezmsg Unit with specific processing capabilities ---
+class BaseProducerUnit(ez.Unit, ABC, typing.Generic[SettingsType, MessageOutType, ProducerType]):
+    """
+    Base class for producer units -- i.e. units that generate messages without consuming inputs.
+    Implement a new Unit as follows:
+
+    class CustomUnit(BaseProducerUnit[
+        CustomProducerSettings,    # SettingsType
+        AxisArray,                 # MessageOutType
+        CustomProducer,            # ProducerType
+    ]):
+        SETTINGS = CustomProducerSettings
+
+    ... that's all!
+
+    Where CustomProducerSettings, and CustomProducer are custom implementations of ez.Settings,
+    and BaseProducer or BaseStatefulProducer, respectively.
+    """
+
+    INPUT_SETTINGS = ez.InputStream(SettingsType)
+    OUTPUT_SIGNAL = ez.OutputStream(MessageOutType)
+
+    async def initialize(self) -> None:
+        self.create_producer()
+
+    def create_producer(self) -> None:
+        # self.producer: ProducerType
+        """Create the producer instance from settings."""
+        producer_type = get_base_producer_type(self.__class__)
+        self.producer = producer_type(settings=self.SETTINGS)
+
+    @ez.subscriber(INPUT_SETTINGS)
+    async def on_settings(self, msg: SettingsType) -> None:
+        """
+        Receive a settings message, override self.SETTINGS, and re-create the producer.
+        Child classes that wish to have fine-grained control over whether the
+        core producer resets on settings changes should override this method.
+
+        Args:
+            msg: a settings message.
+        """
+        self.apply_settings(msg)  # type: ignore
+        self.create_producer()
+
+    @ez.publisher(OUTPUT_SIGNAL)
+    async def produce(self) -> typing.AsyncGenerator:
+        while True:
+            out = await self.producer.__acall__()
+            if out is not None:  # and math.prod(out.data.shape) > 0:
+                yield self.OUTPUT_SIGNAL, out
+
+
+class BaseProcessorUnit(ez.Unit, ABC, typing.Generic[SettingsType]):
+    """
+    Base class for processor units -- i.e. units that process messages.
+    This is an abstract base class that provides common functionality for consumer and transformer
+    units. You probably do not want to inherit from this class directly as you would need to define
+    a custom implementation of `create_processor`.
+    Refer instead to BaseConsumerUnit or BaseTransformerUnit.
+    """
+
+    INPUT_SETTINGS = ez.InputStream(SettingsType)
+
+    async def initialize(self) -> None:
+        self.create_processor()
+
+    @abstractmethod
+    def create_processor(self) -> None: ...
+
+    @ez.subscriber(INPUT_SETTINGS)
+    async def on_settings(self, msg: SettingsType) -> None:
+        """
+        Receive a settings message, override self.SETTINGS, and re-create the processor.
+        Child classes that wish to have fine-grained control over whether the
+        core processor resets on settings changes should override this method.
+
+        Args:
+            msg: a settings message.
+        """
+        self.apply_settings(msg)  # type: ignore
+        self.create_processor()
+
+
+class BaseConsumerUnit(
+    BaseProcessorUnit[SettingsType],
+    ABC,
+    typing.Generic[SettingsType, MessageInType, ConsumerType],
+):
+    """
+    Base class for consumer units -- i.e. units that receive messages but do not return results.
+    Implement a new Unit as follows:
+
+    class CustomUnit(BaseConsumerUnit[
+        CustomConsumerSettings,    # SettingsType
+        AxisArray,                 # MessageInType
+        CustomConsumer,            # ConsumerType
+    ]):
+        SETTINGS = CustomConsumerSettings
+
+    ... that's all!
+
+    Where CustomConsumerSettings and CustomConsumer are custom implementations of:
+    - ez.Settings for settings
+    - BaseConsumer or BaseStatefulConsumer for the consumer implementation
+    """
+
+    INPUT_SIGNAL = ez.InputStream(MessageInType)
+
+    def create_processor(self):
+        # self.processor: ConsumerType[SettingsType, MessageInType, StateType]
+        """Create the consumer instance from settings."""
+        consumer_type = get_base_consumer_type(self.__class__)
+        self.processor = consumer_type(settings=self.SETTINGS)
+
+    @ez.subscriber(INPUT_SIGNAL, zero_copy=True)
+    async def on_signal(self, message: MessageInType):
+        """
+        Consume the message.
+        Args:
+            message: Input message to be consumed
+        """
+        await self.processor.__acall__(message)
+
+
+class BaseTransformerUnit(
+    BaseProcessorUnit[SettingsType],
+    ABC,
+    typing.Generic[SettingsType, MessageInType, MessageOutType, TransformerType],
+):
+    """
+    Base class for transformer units -- i.e. units that transform input messages into output messages.
+    Implement a new Unit as follows:
+
+    class CustomUnit(BaseTransformerUnit[
+        CustomTransformerSettings,    # SettingsType
+        AxisArray,                    # MessageInType
+        AxisArray,                    # MessageOutType
+        CustomTransformer,            # TransformerType
+    ]):
+        SETTINGS = CustomTransformerSettings
+
+    ... that's all!
+
+    Where CustomTransformerSettings and CustomTransformer are custom implementations of:
+    - ez.Settings for settings
+    - One of these transformer types:
+      * BaseTransformer
+      * BaseStatefulTransformer
+      * CompositeProcessor
+    """
+
+    INPUT_SIGNAL = ez.InputStream(MessageInType)
+    OUTPUT_SIGNAL = ez.OutputStream(MessageOutType)
+
+    def create_processor(self):
+        # self.processor: TransformerType[SettingsType, MessageInType, MessageOutType, StateType]
+        """Create the transformer instance from settings."""
+        transformer_type = get_base_transformer_type(self.__class__)
+        self.processor = transformer_type(settings=self.SETTINGS)
+
+    @ez.subscriber(INPUT_SIGNAL, zero_copy=True)
+    @ez.publisher(OUTPUT_SIGNAL)
+    @profile_subpub(trace_oldest=False)
+    async def on_signal(self, message: MessageInType) -> typing.AsyncGenerator:
+        result = await self.processor.__acall__(message)
+        if result is not None:  # and math.prod(result.data.shape) > 0:
+            yield self.OUTPUT_SIGNAL, result
+
+
+class BaseAdaptiveTransformerUnit(
+    BaseProcessorUnit[SettingsType],
+    ABC,
+    typing.Generic[SettingsType, MessageInType, MessageOutType, AdaptiveTransformerType],
+):
+    INPUT_SAMPLE = ez.InputStream(SampleMessage)
+    INPUT_SIGNAL = ez.InputStream(MessageInType)
+    OUTPUT_SIGNAL = ez.OutputStream(MessageOutType)
+
+    def create_processor(self) -> None:
+        # self.processor: AdaptiveTransformerType[SettingsType, MessageInType, MessageOutType, StateType]
+        """Create the adaptive transformer instance from settings."""
+        adaptive_transformer_type = get_base_adaptive_transformer_type(self.__class__)
+        self.processor = adaptive_transformer_type(settings=self.SETTINGS)
+
+    @ez.subscriber(INPUT_SIGNAL, zero_copy=True)
+    @ez.publisher(OUTPUT_SIGNAL)
+    @profile_subpub(trace_oldest=False)
+    async def on_signal(self, message: MessageInType) -> typing.AsyncGenerator:
+        result = await self.processor.__acall__(message)
+        if result is not None:  # and math.prod(result.data.shape) > 0:
+            yield self.OUTPUT_SIGNAL, result
+
+    @ez.subscriber(INPUT_SAMPLE)
+    async def on_sample(self, msg: SampleMessage) -> None:
+        await self.processor.apartial_fit(msg)
+
+
+# Legacy class
+class GenAxisArray(ez.Unit):
+    STATE = GenState
+
+    INPUT_SIGNAL = ez.InputStream(AxisArray)
+    OUTPUT_SIGNAL = ez.OutputStream(AxisArray)
+    INPUT_SETTINGS = ez.InputStream(ez.Settings)
+
+    async def initialize(self) -> None:
+        self.construct_generator()
+
+    # Method to be implemented by subclasses to construct the specific generator
+    def construct_generator(self):
+        raise NotImplementedError
+
+    @ez.subscriber(INPUT_SETTINGS)
+    async def on_settings(self, msg: ez.Settings) -> None:
+        """
+        Update unit settings and reset generator.
+        Note: Not all units will require a full reset with new settings.
+        Override this method to implement a selective reset.
+
+        Args:
+            msg: Instance of SETTINGS object.
+        """
+        self.apply_settings(msg)
+        self.construct_generator()
+
+    @ez.subscriber(INPUT_SIGNAL, zero_copy=True)
+    @ez.publisher(OUTPUT_SIGNAL)
+    @profile_subpub(trace_oldest=False)
+    async def on_signal(self, message: AxisArray) -> typing.AsyncGenerator:
+        try:
+            ret = self.STATE.gen.send(message)
+            if math.prod(ret.data.shape) > 0:
+                yield self.OUTPUT_SIGNAL, ret
+        except (StopIteration, GeneratorExit):
+            ez.logger.debug(f"Generator closed in {self.address}")
+        except Exception:
+            ez.logger.info(traceback.format_exc())

ezmsg/baseproc/util/__init__.py

@@ -0,0 +1 @@
+# Utility modules for ezmsg-baseproc

ezmsg/baseproc/util/asio.py

@@ -0,0 +1,138 @@
+import asyncio
+import contextlib
+import inspect
+import threading
+from concurrent.futures import ThreadPoolExecutor
+from typing import Any, Coroutine, TypeVar
+
+T = TypeVar("T")
+
+
+class CoroutineExecutionError(Exception):
+    """Custom exception for coroutine execution failures"""
+
+    pass
+
+
+def run_coroutine_sync(coroutine: Coroutine[Any, Any, T], timeout: float = 30) -> T:
+    """
+    Executes an asyncio coroutine synchronously, with enhanced error handling.
+
+    Args:
+        coroutine: The asyncio coroutine to execute
+        timeout: Maximum time in seconds to wait for coroutine completion (default: 30)
+
+    Returns:
+        The result of the coroutine execution
+
+    Raises:
+        CoroutineExecutionError: If execution fails due to threading or event loop issues
+        TimeoutError: If execution exceeds the timeout period
+        Exception: Any exception raised by the coroutine
+    """
+
+    def run_in_new_loop() -> T:
+        """
+        Creates and runs a new event loop in the current thread.
+        Ensures proper cleanup of the loop.
+        """
+        new_loop = asyncio.new_event_loop()
+        asyncio.set_event_loop(new_loop)
+        try:
+            return new_loop.run_until_complete(asyncio.wait_for(coroutine, timeout=timeout))
+        finally:
+            with contextlib.suppress(Exception):
+                # Clean up any pending tasks
+                pending = asyncio.all_tasks(new_loop)
+                for task in pending:
+                    task.cancel()
+                new_loop.run_until_complete(asyncio.gather(*pending, return_exceptions=True))
+            new_loop.close()
+
+    try:
+        loop = asyncio.get_running_loop()
+    except RuntimeError:
+        try:
+            return asyncio.run(asyncio.wait_for(coroutine, timeout=timeout))
+        except Exception as e:
+            raise CoroutineExecutionError(f"Failed to execute coroutine: {str(e)}") from e
+
+    if threading.current_thread() is threading.main_thread():
+        if not loop.is_running():
+            try:
+                return loop.run_until_complete(asyncio.wait_for(coroutine, timeout=timeout))
+            except Exception as e:
+                raise CoroutineExecutionError(f"Failed to execute coroutine in main loop: {str(e)}") from e
+        else:
+            with ThreadPoolExecutor() as pool:
+                try:
+                    future = pool.submit(run_in_new_loop)
+                    return future.result(timeout=timeout)
+                except Exception as e:
+                    raise CoroutineExecutionError(f"Failed to execute coroutine in thread: {str(e)}") from e
+    else:
+        try:
+            future = asyncio.run_coroutine_threadsafe(coroutine, loop)
+            return future.result(timeout=timeout)
+        except Exception as e:
+            raise CoroutineExecutionError(f"Failed to execute coroutine threadsafe: {str(e)}") from e
+
+
+class SyncToAsyncGeneratorWrapper:
+    """
+    A wrapper for synchronous generators to be used in an async context.
+    """
+
+    def __init__(self, gen):
+        self._gen = gen
+        self._closed = False
+        # Prime the generator to ready for first send/next call
+        try:
+            is_not_primed = inspect.getgeneratorstate(self._gen) is inspect.GEN_CREATED
+        except AttributeError as e:
+            raise TypeError("The provided generator is not a valid generator object") from e
+        if is_not_primed:
+            try:
+                next(self._gen)
+            except StopIteration:
+                self._closed = True
+            except Exception as e:
+                raise RuntimeError(f"Failed to prime generator: {e}") from e
+
+    async def asend(self, value):
+        if self._closed:
+            raise StopAsyncIteration("Generator is closed")
+        try:
+            return await asyncio.to_thread(self._gen.send, value)
+        except StopIteration as e:
+            self._closed = True
+            raise StopAsyncIteration("Generator is closed") from e
+        except Exception as e:
+            raise RuntimeError(f"Error while sending value to generator: {e}") from e
+
+    async def __anext__(self):
+        if self._closed:
+            raise StopAsyncIteration("Generator is closed")
+        try:
+            return await asyncio.to_thread(self._gen.__next__)
+        except StopIteration as e:
+            self._closed = True
+            raise StopAsyncIteration("Generator is closed") from e
+        except Exception as e:
+            raise RuntimeError(f"Error while getting next value from generator: {e}") from e
+
+    async def aclose(self):
+        if self._closed:
+            return
+        try:
+            await asyncio.to_thread(self._gen.close)
+        except Exception as e:
+            raise RuntimeError(f"Error while closing generator: {e}") from e
+        finally:
+            self._closed = True
+
+    def __aiter__(self):
+        return self
+
+    def __getattr__(self, name):
+        return getattr(self._gen, name)

ezmsg/baseproc/util/message.py

@@ -0,0 +1,31 @@
+import time
+import typing
+from dataclasses import dataclass, field
+
+from ezmsg.util.messages.axisarray import AxisArray
+
+
+@dataclass(unsafe_hash=True)
+class SampleTriggerMessage:
+    timestamp: float = field(default_factory=time.time)
+    """Time of the trigger, in seconds. The Clock depends on the input but defaults to time.time"""
+
+    period: tuple[float, float] | None = None
+    """The period around the timestamp, in seconds"""
+
+    value: typing.Any = None
+    """A value or 'label' associated with the trigger."""
+
+
+@dataclass
+class SampleMessage:
+    trigger: SampleTriggerMessage
+    """The time, window, and value (if any) associated with the trigger."""
+
+    sample: AxisArray
+    """The data sampled around the trigger."""
+
+
+def is_sample_message(message: typing.Any) -> typing.TypeGuard[SampleMessage]:
+    """Check if the message is a SampleMessage."""
+    return hasattr(message, "trigger")

ezmsg/baseproc/util/profile.py

@@ -0,0 +1,171 @@
+import functools
+import logging
+import os
+import time
+import typing
+from pathlib import Path
+
+import ezmsg.core as ez
+
+HEADER = "Time,Source,Topic,SampleTime,PerfCounter,Elapsed"
+
+
+def get_logger_path() -> Path:
+    # Retrieve the logfile name from the environment variable
+    logfile = os.environ.get("EZMSG_PROFILE", None)
+
+    # Determine the log file path, defaulting to "ezprofiler.log" if not set
+    logpath = Path(logfile or "ezprofiler.log")
+
+    # If the log path is not absolute, prepend it with the user's home directory and ".ezmsg/profile"
+    if not logpath.is_absolute():
+        logpath = Path.home() / ".ezmsg" / "profile" / logpath
+
+    return logpath
+
+
+def _setup_logger(append: bool = False) -> logging.Logger:
+    logpath = get_logger_path()
+    logpath.parent.mkdir(parents=True, exist_ok=True)
+
+    write_header = True
+    if logpath.exists() and logpath.is_file():
+        if append:
+            with open(logpath) as f:
+                first_line = f.readline().rstrip()
+            if first_line == HEADER:
+                write_header = False
+            else:
+                # Remove the file if appending, but headers do not match
+                ezmsg_logger = logging.getLogger("ezmsg")
+                ezmsg_logger.warning(
+                    "Profiling header mismatch: please make sure to use the same version of ezmsg for all processes."
+                )
+                logpath.unlink()
+        else:
+            # Remove the file if not appending
+            logpath.unlink()
+
+    # Create a logger with the name "ezprofile"
+    _logger = logging.getLogger("ezprofile")
+
+    # Set the logger's level to EZMSG_LOGLEVEL env var value if it exists, otherwise INFO
+    _logger.setLevel(os.environ.get("EZMSG_LOGLEVEL", "INFO").upper())
+
+    # Create a file handler to write log messages to the log file
+    fh = logging.FileHandler(logpath)
+    fh.setLevel(logging.DEBUG)  # Set the file handler log level to DEBUG
+
+    # Add the file handler to the logger
+    _logger.addHandler(fh)
+
+    # Add the header if writing to new file or if header matched header in file.
+    if write_header:
+        _logger.debug(HEADER)
+
+    # Set the log message format
+    formatter = logging.Formatter("%(asctime)s,%(message)s", datefmt="%Y-%m-%dT%H:%M:%S%z")
+    fh.setFormatter(formatter)
+
+    return _logger
+
+
+logger = _setup_logger(append=True)
+
+
+def _process_obj(obj, trace_oldest: bool = True):
+    samp_time = None
+    if hasattr(obj, "axes") and ("time" in obj.axes or "win" in obj.axes):
+        axis = "win" if "win" in obj.axes else "time"
+        ax = obj.get_axis(axis)
+        len = obj.data.shape[obj.get_axis_idx(axis)]
+        if len > 0:
+            idx = 0 if trace_oldest else (len - 1)
+            if hasattr(ax, "data"):
+                samp_time = ax.data[idx]
+            else:
+                samp_time = ax.value(idx)
+            if ax == "win" and "time" in obj.axes:
+                if hasattr(obj.axes["time"], "data"):
+                    samp_time += obj.axes["time"].data[idx]
+                else:
+                    samp_time += obj.axes["time"].value(idx)
+    return samp_time
+
+
+def profile_method(trace_oldest: bool = True):
+    """
+    Decorator to profile a method by logging its execution time and other details.
+
+    Args:
+        trace_oldest (bool): If True, trace the oldest sample time; otherwise, trace the newest.
+
+    Returns:
+        Callable: The decorated function with profiling.
+    """
+
+    def profiling_decorator(func: typing.Callable):
+        @functools.wraps(func)
+        def wrapped_func(caller, *args, **kwargs):
+            start = time.perf_counter()
+            res = func(caller, *args, **kwargs)
+            stop = time.perf_counter()
+            source = ".".join((caller.__class__.__module__, caller.__class__.__name__))
+            topic = f"{caller.address}"
+            samp_time = _process_obj(res, trace_oldest=trace_oldest)
+            logger.debug(
+                ",".join(
+                    [
+                        source,
+                        topic,
+                        f"{samp_time}",
+                        f"{stop}",
+                        f"{(stop - start) * 1e3:0.4f}",
+                    ]
+                )
+            )
+            return res
+
+        return wrapped_func if logger.level == logging.DEBUG else func
+
+    return profiling_decorator
+
+
+def profile_subpub(trace_oldest: bool = True):
+    """
+    Decorator to profile a subscriber-publisher method in an ezmsg Unit
+    by logging its execution time and other details.
+
+    Args:
+        trace_oldest (bool): If True, trace the oldest sample time; otherwise, trace the newest.
+
+    Returns:
+        Callable: The decorated async task with profiling.
+    """
+
+    def profiling_decorator(func: typing.Callable):
+        @functools.wraps(func)
+        async def wrapped_task(unit: ez.Unit, msg: typing.Any = None):
+            source = ".".join((unit.__class__.__module__, unit.__class__.__name__))
+            topic = f"{unit.address}"
+            start = time.perf_counter()
+            async for stream, obj in func(unit, msg):
+                stop = time.perf_counter()
+                samp_time = _process_obj(obj, trace_oldest=trace_oldest)
+                logger.debug(
+                    ",".join(
+                        [
+                            source,
+                            topic,
+                            f"{samp_time}",
+                            f"{stop}",
+                            f"{(stop - start) * 1e3:0.4f}",
+                        ]
+                    )
+                )
+                start = stop
+                yield stream, obj
+
+        return wrapped_task if logger.level == logging.DEBUG else func
+
+    return profiling_decorator