diff --git a/examples/active_learn/streaming/run_me.py b/examples/active_learn/streaming/run_me.py new file mode 100644 index 0000000..51a7498 --- /dev/null +++ b/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()) diff --git a/rose/al/__init__.py b/rose/al/__init__.py index fd2e905..70576f8 100644 --- a/rose/al/__init__.py +++ b/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", ] diff --git a/rose/al/streaming_learner.py b/rose/al/streaming_learner.py new file mode 100644 index 0000000..153f2fe --- /dev/null +++ b/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() diff --git a/tests/unit/test_streaming_learner.py b/tests/unit/test_streaming_learner.py new file mode 100644 index 0000000..da8eaa5 --- /dev/null +++ b/tests/unit/test_streaming_learner.py @@ -0,0 +1,180 @@ +"""Unit tests for StreamingActiveLearner: windowing, publish-gate criterion, stop() and source- +exhaustion termination.""" + +import asyncio +from unittest.mock import AsyncMock, MagicMock + +import pytest +from radical.asyncflow import WorkflowEngine + +from rose.al.streaming_learner import StreamingActiveLearner + + +def make_learner(batch_size=2, max_wait=None, conflate=False, criterion_results=None, sources=None): + """Create a learner with mocked task submission. + + Training results echo the received window; criterion results are taken from the given list + (metric, compared as '< 1.0'). + """ + learner = StreamingActiveLearner( + MagicMock(spec=WorkflowEngine), + batch_size=batch_size, + max_wait=max_wait, + conflate=conflate, + sources=sources, + ) + learner.training_function = {"func": AsyncMock(), "args": (), "kwargs": {}, "decor_kwargs": {}} + learner.active_learn_function = { + "func": AsyncMock(), + "args": (), + "kwargs": {}, + "decor_kwargs": {}, + } + metrics = iter(criterion_results or ()) + if criterion_results is not None: + learner.criterion_function = { + "func": AsyncMock(), + "args": (), + "kwargs": {}, + "decor_kwargs": {}, + "operator": "<", + "threshold": 1.0, + "metric_name": "test_metric", + } + + windows = [] + + async def mock_reg(task_obj, deps=None): + if task_obj.get("metric_name"): + return next(metrics) + if task_obj["args"]: # training task: window is first arg + windows.append(task_obj["args"][0]) + return "result" + + learner._register_task = AsyncMock(side_effect=mock_reg) + return learner, windows + + +@pytest.mark.asyncio +async def test_windows_batched_and_source_exhaustion(): + learner, windows = make_learner(batch_size=2) + + async def source(): + for i in range(4): + yield i + + learner.attach_source(source()) + + states = [state async for state in learner.start()] + + assert windows == [[0, 1], [2, 3]] + assert [s.iteration for s in states] == [0, 1] + assert all(s.window_size == 2 for s in states) + + +def test_sources_at_construction_need_no_event_loop(): + # pumps are deferred to start(), so construction works outside a loop + async def source(): + for i in range(4): + yield i + + learner, windows = make_learner(batch_size=2, sources=source()) + + async def run(): + states = [state async for state in learner.start()] + assert windows == [[0, 1], [2, 3]] + assert len(states) == 2 + + asyncio.run(run()) + + +@pytest.mark.asyncio +async def test_criterion_is_publish_gate_not_termination(): + learner, _ = make_learner(batch_size=1, criterion_results=[0.5, 2.0]) + published = [] + learner.on_model_ready(published.append) + + await learner.feed("a") + await learner.feed("b") + + states = [] + + async def run(): + async for state in learner.start(): + states.append(state) + if len(states) == 2: + learner.stop() + + await asyncio.wait_for(run(), timeout=5.0) + + # criterion met on first window (0.5 < 1.0) but loop continued + assert [s.should_stop for s in states] == [True, False] + assert len(published) == 1 and published[0] is states[0] + + +@pytest.mark.asyncio +async def test_stop_unblocks_empty_queue(): + learner, _ = make_learner() + + async def run(): + async for _ in learner.start(): + pytest.fail("no data was fed, no state expected") + + task = asyncio.ensure_future(run()) + await asyncio.sleep(0.1) + learner.stop() + await asyncio.wait_for(task, timeout=5.0) + assert learner.is_stopped + + +@pytest.mark.asyncio +async def test_max_wait_flushes_partial_window(): + learner, windows = make_learner(batch_size=10, max_wait=0.05) + + await learner.feed(1) + await learner.feed(2) + + async def run(): + async for _ in learner.start(): + learner.stop() + + await asyncio.wait_for(run(), timeout=5.0) + assert windows == [[1, 2]] + + +@pytest.mark.asyncio +async def test_conflate_keeps_newest_items(): + learner, windows = make_learner(batch_size=2, conflate=True) + for i in range(6): + await learner.feed(i) + + async def run(): + async for _ in learner.start(): + learner.stop() + + await asyncio.wait_for(run(), timeout=5.0) + assert windows == [[4, 5]] + + +@pytest.mark.asyncio +async def test_missing_tasks_raise(): + learner = StreamingActiveLearner(MagicMock(spec=WorkflowEngine)) + with pytest.raises(ValueError): + async for _ in learner.start(): + pass + + +def test_invalid_params_raise(): + with pytest.raises(ValueError): + StreamingActiveLearner(MagicMock(spec=WorkflowEngine), batch_size=0) + with pytest.raises(ValueError): + StreamingActiveLearner(MagicMock(spec=WorkflowEngine), max_wait=-1.0) + + +@pytest.mark.asyncio +async def test_conflate_bounds_queue_at_ingestion(): + learner, _ = make_learner(batch_size=2, conflate=True) + for i in range(100): + await learner.feed(i) + # backlog is dropped at ingestion, not at window collection + assert learner._queue.qsize() <= 2