Add StreamingActiveLearner: learner loop driven by streamed data

examples/active_learn/streaming/run_me.py

@@ -0,0 +1,86 @@
+# This is a *draft* example of the StreamingActiveLearner
+# It cuts some corners in order to deliver a bare-bones test.
+#     Namely: StreamingActiveLearner currently requires a stop criterion.
+#         so, I have a dummy one that always returns true.
+#         (and technically the stop criterion is usually executable, but
+#         did this for brevity)
+#
+#     All tasks are function tasks. Currently, the only way to fetch state in
+#         the model is if they are function tasks returning dictionaries. If these were
+#         executable tasks, the developer would be responsible for their own data transfer
+#         management.
+#
+# A more "real" example of the StreamingActiveLearner would eventually replace
+# this quick test script.
+
+import asyncio
+from concurrent.futures import ProcessPoolExecutor
+
+from radical.asyncflow import WorkflowEngine
+from rhapsody.backends import ConcurrentExecutionBackend
+
+from rose.al.streaming_learner import StreamingActiveLearner
+from rose.learner import IterationState
+from rose.metrics import MEAN_SQUARED_ERROR_MSE
+
+
+async def rose_al():
+    engine = await ConcurrentExecutionBackend(ProcessPoolExecutor())
+    asyncflow = await WorkflowEngine.create(engine)
+
+    acl = StreamingActiveLearner(asyncflow, batch_size=2)
+
+    # Define and register the simulation task
+    @acl.simulation_task(as_executable=False)
+    async def simulation(window, *args):
+        print(f"Start simulation: {window}")
+        return {"train_window": window}
+
+    # Define and register the training task
+    @acl.training_task(as_executable=False)
+    async def training(train_window, *args):
+        print(f"Start training: {train_window}")
+        out = []
+        for i in train_window:
+            out.append(i * 2)
+        return out
+
+    # Define and register the active learning task
+    @acl.active_learn_task(as_executable=False)
+    async def active_learn(from_train, *args):
+        print(f"Start learning: {from_train}")
+        out = []
+        for i in from_train:
+            out.append(i * 3)
+        return {"sum": sum(out)}
+
+    # Defining the stop criterion with a metric (MSE in this case)
+    @acl.as_stop_criterion(as_executable=False, metric_name=MEAN_SQUARED_ERROR_MSE, threshold=0.1)
+    async def check_mse(*args):
+        return 0.01
+
+    # dummy data source
+    async def data_source():
+        for i in range(10):
+            yield i
+            await asyncio.sleep(1)
+
+    acl.attach_source(data_source())
+
+    # model callback
+    async def on_model_callback(state: IterationState):
+        print(f"Publish: {state.sum}")
+
+    acl.on_model_ready(on_model_callback)
+
+    # start learner
+    async for _ in acl.start():
+        await asyncio.sleep(0)
+
+    acl.stop()
+
+    await acl.shutdown()
+
+
+if __name__ == "__main__":
+    asyncio.run(rose_al())

rose/al/__init__.py

@@ -1,8 +1,10 @@
 from rose.al.active_learner import ParallelActiveLearner, SequentialActiveLearner
 from rose.al.selector import AlgorithmSelector
+from rose.al.streaming_learner import StreamingActiveLearner
 
 __all__ = [
     "ParallelActiveLearner",
     "SequentialActiveLearner",
+    "StreamingActiveLearner",
     "AlgorithmSelector",
 ]

rose/al/streaming_learner.py

@@ -0,0 +1,268 @@
+import asyncio
+import inspect
+from collections.abc import AsyncIterator, Callable
+from typing import Any
+
+from radical.asyncflow import WorkflowEngine
+
+from ..learner import IterationState, Learner, LearnerConfig
+
+
+class StreamingActiveLearner(Learner):
+    """Active learner driven by streamed data instead of a simulation task.
+
+    Data items arrive via :meth:`feed` or :meth:`attach_source` and are
+    collected into windows of ``batch_size`` items (a partial window is
+    flushed after ``max_wait`` seconds). Each window triggers one learning
+    iteration: training -> active learning -> criterion. The window is
+    passed as the first positional argument to the training task, and the
+    previous iteration's active-learn result is appended as a dependency
+    (for warm starts).
+
+    Unlike SequentialActiveLearner, a met stop criterion does not end the
+    loop: it marks the model as publishable (``state.should_stop`` is True
+    and ``on_model_ready`` callbacks fire) while consumption continues.
+    The loop ends when :meth:`stop` is called or all attached sources are
+    exhausted.
+
+    Example::
+
+        learner = StreamingActiveLearner(asyncflow, batch_size=10)
+        learner.attach_source(sensor_stream())
+
+        async for state in learner.start():
+            if state.should_stop:
+                publish(state)
+    """
+
+    _END = object()  # sentinel: an attached source finished
+    _WAKE = object()  # sentinel: stop() unblocking the collector
+
+    def __init__(
+        self,
+        asyncflow: WorkflowEngine,
+        batch_size: int = 1,
+        max_wait: float | None = None,
+        conflate: bool = False,
+        sources: AsyncIterator[Any] | list[AsyncIterator[Any]] | None = None,
+    ) -> None:
+        """Initialize the Streaming Active Learner.
+
+        Args:
+            asyncflow: The workflow engine instance used to manage async tasks.
+            batch_size: Number of streamed items per learning window.
+            max_wait: Flush a partial window after this many seconds of
+                waiting for more items. None waits for a full window.
+            conflate: If True, drop the oldest backlog at ingestion time and
+                keep only the newest ~``batch_size`` items when iterations
+                are slower than the stream ("latest wins"); the internal
+                queue stays bounded.
+            sources: One or more async iterators to consume as data
+                sources; equivalent to calling :meth:`attach_source` for
+                each.
+        """
+        super().__init__(asyncflow, register_and_submit=True)
+        if batch_size < 1:
+            raise ValueError(f"batch_size must be >= 1, got {batch_size}")
+        if max_wait is not None and max_wait <= 0:
+            raise ValueError(f"max_wait must be positive or None, got {max_wait}")
+        self.batch_size = batch_size
+        self.max_wait = max_wait
+        self.conflate = conflate
+
+        self._queue: asyncio.Queue[Any] = asyncio.Queue()
+        self._open_sources = 0
+        self._exhausted = False
+        self._started = False
+        self._sources: list[asyncio.Task] = []
+        self._pending_sources: list[AsyncIterator[Any]] = []
+        self._model_callbacks: list[Callable[[IterationState], Any]] = []
+        self._pending_config: LearnerConfig | None = None
+
+        if sources is not None:
+            for source in sources if isinstance(sources, list) else [sources]:
+                self.attach_source(source)
+
+    async def feed(self, item: Any) -> None:
+        """Feed a single data item into the learner's stream."""
+        await self._put(item)
+
+    async def _put(self, item: Any) -> None:
+        """Enqueue a data item; when conflating, drop the oldest backlog first so the queue stays
+        bounded to roughly ``batch_size`` items even while an iteration is running."""
+        if self.conflate:
+            while self._queue.qsize() >= self.batch_size:
+                self._drain_sentinel(self._queue.get_nowait())
+        await self._queue.put(item)
+
+    def attach_source(self, source: AsyncIterator[Any]) -> None:
+        """Attach an async iterator as a data source.
+
+        Sources attached before :meth:`start` are only consumed once the learner loop runs. The loop
+        ends once all attached sources are exhausted and the queue is drained; learners fed only via
+        :meth:`feed` run until :meth:`stop` is called.
+        """
+        self._open_sources += 1
+        if self._started:
+            self._start_pump(source)
+        else:
+            self._pending_sources.append(source)
+
+    def _start_pump(self, source: AsyncIterator[Any]) -> None:
+        async def pump() -> None:
+            try:
+                async for item in source:
+                    await self._put(item)
+            finally:
+                await self._queue.put(self._END)
+
+        self._sources.append(asyncio.ensure_future(pump()))
+
+    def on_model_ready(self, callback: Callable[[IterationState], Any]) -> None:
+        """Register a callback fired whenever the stop criterion is met.
+
+        In streaming mode the criterion acts as a publish gate, not a terminal condition. Callbacks
+        receive the IterationState and may be sync or async.
+        """
+        self._model_callbacks.append(callback)
+
+    def set_next_config(self, config: LearnerConfig) -> None:
+        """Set configuration to apply from the next window on."""
+        self._pending_config = config
+
+    def stop(self) -> None:
+        """Signal the learner to stop and unblock the window collector."""
+        super().stop()
+        self._queue.put_nowait(self._WAKE)
+
+    def _drain_sentinel(self, item: Any) -> bool:
+        """Process a sentinel item; return True if it was one."""
+        if item is self._WAKE:
+            return True
+        if item is self._END:
+            self._open_sources -= 1
+            if self._open_sources <= 0:
+                self._exhausted = True
+            return True
+        return False
+
+    async def _collect(self) -> list[Any]:
+        """Collect the next window; empty means stopped or exhausted."""
+        window: list[Any] = []
+        while len(window) < self.batch_size:
+            if self.is_stopped or (self._exhausted and self._queue.empty()):
+                break
+            try:
+                timeout = self.max_wait if window else None
+                item = await asyncio.wait_for(self._queue.get(), timeout)
+            except asyncio.TimeoutError:
+                break  # flush partial window
+            if not self._drain_sentinel(item):
+                window.append(item)
+
+        if self.conflate:
+            while not self._queue.empty():
+                item = self._queue.get_nowait()
+                if not self._drain_sentinel(item):
+                    window.append(item)
+            window = window[-self.batch_size :]
+
+        return window
+
+    async def start(
+        self, initial_config: LearnerConfig | None = None
+    ) -> AsyncIterator[IterationState]:
+        """Consume the stream and yield an IterationState per window.
+
+        Args:
+            initial_config: Optional LearnerConfig; can be replaced between
+                windows via set_next_config().
+
+        Yields:
+            IterationState per processed window, with ``window_size`` in
+            its state dict. ``should_stop`` marks criterion-met (model
+            ready) states; the loop itself keeps running.
+        """
+        if not self.training_function or not self.active_learn_function:
+            raise ValueError("Training and Active Learning functions must be set!")
+
+        self._started = True
+        for source in self._pending_sources:
+            self._start_pump(source)
+        self._pending_sources.clear()
+
+        config = initial_config
+        acl_task: Any = None
+        _stop_reason = "stream_exhausted"
+
+        try:
+            i = 0
+            while True:
+                window = await self._collect()
+                if self.is_stopped:
+                    _stop_reason = "stopped"
+                    break
+                if not window:
+                    break  # sources exhausted
+
+                if self._pending_config is not None:
+                    config = self._pending_config
+                    self._pending_config = None
+
+                self.clear_state()
+
+                train_cfg = self._get_iteration_task_config(
+                    self.training_function, config, "training", i
+                )
+                train_cfg["args"] = (window, *train_cfg["args"])
+                train_task = self._register_task(train_cfg, deps=acl_task)
+                train_result = await train_task
+
+                acl_cfg = self._get_iteration_task_config(
+                    self.active_learn_function, config, "active_learn", i
+                )
+                acl_task = self._register_task(acl_cfg, deps=train_task)
+                acl_result = await acl_task
+
+                if self.is_stopped:
+                    _stop_reason = "stopped"
+                    break
+                self._extract_state_from_result(train_result)
+                self._extract_state_from_result(acl_result)
+
+                metric_value: float | None = None
+                should_stop = False
+                if self.criterion_function:
+                    crit_cfg = self._get_iteration_task_config(
+                        self.criterion_function, config, "criterion", i
+                    )
+                    stop_result = await self._register_task(crit_cfg)
+                    if self.is_stopped:
+                        _stop_reason = "stopped"
+                        break
+                    should_stop, metric_value = self._check_stop_criterion(stop_result)
+
+                self.register_state("window_size", len(window))
+                self._iteration_state = self.build_iteration_state(
+                    iteration=i,
+                    metric_value=metric_value,
+                    should_stop=should_stop,
+                    current_config=config,
+                )
+
+                self._notify_trackers_iteration(self._iteration_state)
+                if should_stop:
+                    for cb in self._model_callbacks:
+                        result = cb(self._iteration_state)
+                        if inspect.isawaitable(result):
+                            await result
+
+                yield self._iteration_state
+                i += 1
+        except Exception:
+            _stop_reason = "error"
+            raise
+        finally:
+            self._notify_trackers_stop(self._iteration_state, _stop_reason)
+            for task in self._sources:
+                task.cancel()