diff --git a/src/detectmatelibrary/common/core.py b/src/detectmatelibrary/common/core.py index 62c0524..4377d4d 100644 --- a/src/detectmatelibrary/common/core.py +++ b/src/detectmatelibrary/common/core.py @@ -11,11 +11,16 @@ from detectmatelibrary.tools.logging import logger, setup_logging -from typing import Any, Dict, List, Protocol +from detectmatelibrary.utils.persistency import EventPersistency +from detectmatelibrary.utils import persistency + +from typing import Any, Dict, List, Protocol, Callable setup_logging() +# Persistency ################################################################## + class _Stoppable(Protocol): """Structural type for objects ``Component`` will stop on context-manager @@ -30,6 +35,71 @@ class _Stoppable(Protocol): def stop(self) -> None: ... +class PersistencyOp: + @staticmethod + def __get_persistency(instance: object) -> EventPersistency | None: + ep = getattr(instance, "persistency", None) + if ep is None: + logger.debug("No persistency configured, nothing to export") + return ep + + @staticmethod + def __apply( + instance: object, + op: Callable[[EventPersistency, Any, dict[str, Any] | None], bytes | None], + path: Any = None, + storage_options: dict[str, Any] | None = None, + ) -> bytes | None: + + if (ep := PersistencyOp.__get_persistency(instance)) is None: + return None + + saver = getattr(instance, "saver", None) + if saver is not None: + with saver.locked(): + return op(ep, path, storage_options) + return op(ep, path, storage_options) + + @staticmethod + def save( + instance: object, + path: str | None = None, + storage_options: dict[str, Any] | None = None, + ) -> bytes | None: + """Save this component's EventPersistency state. + + When path is None, returns the state as bytes (zip archive). + When path is given, writes to that fsspec URI and returns None. + Returns None if no persistency is configured. Thread-safe when a + PersistencySaver is running: acquires the saver lock before saving + (guards against the background save timer and concurrent ingest). + """ + + return PersistencyOp.__apply( + instance=instance, path=path, storage_options=storage_options, op=persistency.save + ) + + @staticmethod + def load( + instance: object, + path: str | bytes, + storage_options: dict[str, Any] | None = None, + ) -> None: + """Restore this component's EventPersistency state. + + path may be an fsspec URI string or bytes returned by + export_state(). No-op if no persistency is configured. Thread-safe + when a PersistencySaver is running: acquires the saver lock before + loading (guards against the background save timer). + """ + + PersistencyOp.__apply( + instance=instance, path=path, storage_options=storage_options, op=persistency.load + ) + + +# Train operations ################################################################## + class TrainBuffer: def __init__(self) -> None: self.buffer: list[BaseSchema | list[BaseSchema]] = [] @@ -50,6 +120,8 @@ def __iter__(self) -> "TrainBuffer": return self +# Core component skeleton structure ################################################ + class CoreConfig(BasicConfig): start_id: int = 10 data_use_training: int | None = None @@ -106,6 +178,8 @@ def __exit__(self, *_: Any) -> None: self.saver.stop() +# Core component ################################################ + class CoreComponent(Component): """Base class for all components in the system.""" def __init__( @@ -129,56 +203,18 @@ def __init__( self.buffer_train = TrainBuffer() def export_state( - self, - path: str | None = None, - storage_options: dict[str, Any] | None = None, + self, path: str | None = None, storage_options: dict[str, Any] | None = None, ) -> bytes | None: - """Save this component's EventPersistency state. - - When path is None, returns the state as bytes (zip archive). - When path is given, writes to that fsspec URI and returns None. - Returns None if no persistency is configured. Thread-safe when a - PersistencySaver is running: acquires the saver lock before saving - (guards against the background save timer and concurrent ingest). - """ - # ponytail: local import keeps common.core free of a persistency - # package import at module load (see _Stoppable). - from detectmatelibrary.utils import persistency - - ep = getattr(self, "persistency", None) - if ep is None: - logger.debug(f"{self.name}: no persistency configured, nothing to export") - return None - saver = getattr(self, "saver", None) - if saver is not None: - with saver.locked(): - return persistency.save(ep, path, storage_options) - return persistency.save(ep, path, storage_options) + return PersistencyOp.save( + instance=self, path=path, storage_options=storage_options + ) def import_state( - self, - path: str | bytes, - storage_options: dict[str, Any] | None = None, + self, path: str | bytes, storage_options: dict[str, Any] | None = None ) -> None: - """Restore this component's EventPersistency state. - - path may be an fsspec URI string or bytes returned by - export_state(). No-op if no persistency is configured. Thread-safe - when a PersistencySaver is running: acquires the saver lock before - loading (guards against the background save timer). - """ - from detectmatelibrary.utils import persistency - - ep = getattr(self, "persistency", None) - if ep is None: - logger.debug(f"{self.name}: no persistency configured, nothing to import") - return - saver = getattr(self, "saver", None) - if saver is not None: - with saver.locked(): - persistency.load(ep, path, storage_options) - else: - persistency.load(ep, path, storage_options) + return PersistencyOp.load( + instance=self, path=path, storage_options=storage_options + ) def update_state(self, state: StatesL) -> None: self.fitlogic.update_state(state) @@ -186,6 +222,9 @@ def update_state(self, state: StatesL) -> None: def get_state(self) -> str: return self.fitlogic.get_last_state() + def get_window_size(self) -> int: + return self.data_buffer.get_window_size() + def process(self, data: BaseSchema | bytes) -> BaseSchema | bytes | None: is_byte, data = SchemaPipeline.preprocess(self.input_schema(), data) logger.debug(f"<<{self.name}>> received:\n{data}") diff --git a/src/detectmatelibrary/utils/data_buffer.py b/src/detectmatelibrary/utils/data_buffer.py index aaa5d6c..c42f2fd 100644 --- a/src/detectmatelibrary/utils/data_buffer.py +++ b/src/detectmatelibrary/utils/data_buffer.py @@ -106,3 +106,7 @@ def flush(self) -> Any: if self.buffer: clear = self.mode != BufferMode.WINDOW return self._process_and_clear(self.buffer, clear=clear) + + def get_window_size(self) -> int: + """Return the output size of the buffer.""" + return 1 if self.size is None else self.size diff --git a/tests/test_common/test_core.py b/tests/test_common/test_core.py index 114a74b..92bc88d 100644 --- a/tests/test_common/test_core.py +++ b/tests/test_common/test_core.py @@ -114,12 +114,14 @@ def run(self, input_, output_) -> bool: class DummyComponentWithBuffer(CoreComponent): - def __init__(self, name: str, config: MockConfig = MockConfig()) -> None: + def __init__( + self, name: str, size: int, config: MockConfig = MockConfig() + ) -> None: super().__init__( name=name, config=config, type_="DummyWithBuffer", - args_buffer=ArgsBuffer(mode="batch", size=3, process_function=sum), + args_buffer=ArgsBuffer(mode=BufferMode.BATCH, size=size), ) @@ -132,6 +134,10 @@ def test_initialize(self) -> None: assert isinstance(config, BasicConfig) assert config.get_config() == expected + def test_get_window(self) -> None: + assert DummyComponentWithBuffer("test_1", 2).get_window_size() == 2 + assert DummyComponentWithBuffer("test_2", 3).get_window_size() == 3 + def test_initialize_dict(self) -> None: config = MockConfig.from_dict( { diff --git a/tests/test_utils/test_data_buffer.py b/tests/test_utils/test_data_buffer.py index 784c849..f6af67d 100644 --- a/tests/test_utils/test_data_buffer.py +++ b/tests/test_utils/test_data_buffer.py @@ -10,6 +10,7 @@ def test_no_buf_mode(self): buf.add(1) buf.add(2) assert results == [1, 2] + assert buf.get_window_size() == 1 def test_batch_mode(self): buf = DataBuffer(ArgsBuffer(mode=BufferMode.BATCH, process_function=sum, size=3)) @@ -20,6 +21,7 @@ def test_batch_mode(self): assert result == 6 # Buffer should be empty after processing assert len(buf.buffer) == 0 + assert buf.get_window_size() == 3 def test_batch_mode2(self): results = [] @@ -32,6 +34,7 @@ def test_batch_mode2(self): assert results == [[1, 1, 1]] # Buffer should be empty after processing assert len(buf.buffer) == 0 + assert buf.get_window_size() == 3 def test_window_mode(self): buf = DataBuffer(ArgsBuffer(mode=BufferMode.WINDOW, process_function=sum, size=2)) @@ -42,6 +45,7 @@ def test_window_mode(self): assert buf.add(5) == 7 # Buffer should contain last two elements assert list(buf.buffer) == [2, 5] + assert buf.get_window_size() == 2 def test_flush_batch(self): results = []