Make MovingWindow a BackgroundService

This gives the moving window a common interface and removes some
bookkeeping from the implementation, as well as making tests more
correct (when using it as an async context manager).

Signed-off-by: Leandro Lucarella <[email protected]>

Author: llucax

benchmarks/timeseries/periodic_feature_extractor.py

@@ -12,6 +12,8 @@
 from __future__ import annotations
 
 import asyncio
+import collections.abc
+import contextlib
 import logging
 from datetime import datetime, timedelta, timezone
 from functools import partial
@@ -27,19 +29,23 @@
 from frequenz.sdk.timeseries._quantities import Quantity
 
 
-async def init_feature_extractor(period: int) -> PeriodicFeatureExtractor:
+@contextlib.asynccontextmanager
+async def init_feature_extractor(
+    period: int,
+) -> collections.abc.AsyncIterator[PeriodicFeatureExtractor]:
     """Initialize the PeriodicFeatureExtractor class."""
     # We only need the moving window to initialize the PeriodicFeatureExtractor class.
     lm_chan = Broadcast[Sample[Quantity]]("lm_net_power")
-    moving_window = MovingWindow(
+    async with MovingWindow(
         timedelta(seconds=1), lm_chan.new_receiver(), timedelta(seconds=1)
-    )
-
-    await lm_chan.new_sender().send(Sample(datetime.now(tz=timezone.utc), Quantity(0)))
+    ) as moving_window:
+        await lm_chan.new_sender().send(
+            Sample(datetime.now(tz=timezone.utc), Quantity(0))
+        )
 
-    # Initialize the PeriodicFeatureExtractor class with a period of period seconds.
-    # This works since the sampling period is set to 1 second.
-    return PeriodicFeatureExtractor(moving_window, timedelta(seconds=period))
+        # Initialize the PeriodicFeatureExtractor class with a period of period seconds.
+        # This works since the sampling period is set to 1 second.
+        yield PeriodicFeatureExtractor(moving_window, timedelta(seconds=period))
 
 
 def _calculate_avg_window(
@@ -211,22 +217,22 @@ async def main() -> None:
 
     # create a random ndarray with 29 days -5 seconds of data
     days_29_s = 29 * DAY_S
-    feature_extractor = await init_feature_extractor(10)
-    data = rng.standard_normal(days_29_s)
-    run_benchmark(data, 4, feature_extractor)
-
-    days_29_s = 29 * DAY_S + 3
-    data = rng.standard_normal(days_29_s)
-    run_benchmark(data, 4, feature_extractor)
-
-    # create a random ndarray with 29 days +5 seconds of data
-    data = rng.standard_normal(29 * DAY_S + 5)
-
-    feature_extractor = await init_feature_extractor(7 * DAY_S)
-    # TEST one day window and 6 days distance. COPY (Case 3)
-    run_benchmark(data, DAY_S, feature_extractor)
-    # benchmark one day window and 6 days distance. NO COPY (Case 1)
-    run_benchmark(data[: 28 * DAY_S], DAY_S, feature_extractor)
+    async with init_feature_extractor(10) as feature_extractor:
+        data = rng.standard_normal(days_29_s)
+        run_benchmark(data, 4, feature_extractor)
+
+        days_29_s = 29 * DAY_S + 3
+        data = rng.standard_normal(days_29_s)
+        run_benchmark(data, 4, feature_extractor)
+
+        # create a random ndarray with 29 days +5 seconds of data
+        data = rng.standard_normal(29 * DAY_S + 5)
+
+    async with init_feature_extractor(7 * DAY_S) as feature_extractor:
+        # TEST one day window and 6 days distance. COPY (Case 3)
+        run_benchmark(data, DAY_S, feature_extractor)
+        # benchmark one day window and 6 days distance. NO COPY (Case 1)
+        run_benchmark(data[: 28 * DAY_S], DAY_S, feature_extractor)
 
 
 logging.basicConfig(level=logging.DEBUG)

src/frequenz/sdk/timeseries/_moving_window.py

@@ -16,7 +16,7 @@
 from frequenz.channels import Broadcast, Receiver, Sender
 from numpy.typing import ArrayLike
 
-from .._internal._asyncio import cancel_and_await
+from ..actor._background_service import BackgroundService
 from ._base_types import UNIX_EPOCH, Sample
 from ._quantities import Quantity
 from ._resampling import Resampler, ResamplerConfig
@@ -25,7 +25,7 @@
 _logger = logging.getLogger(__name__)
 
 
-class MovingWindow:
+class MovingWindow(BackgroundService):
     """
     A data window that moves with the latest datapoints of a data stream.
 
@@ -72,22 +72,21 @@ async def run() -> None:
 
             send_task = asyncio.create_task(send_mock_data(resampled_data_sender))
 
-            window = MovingWindow(
+            async with MovingWindow(
                 size=timedelta(seconds=5),
                 resampled_data_recv=resampled_data_receiver,
                 input_sampling_period=timedelta(seconds=1),
-            )
+            ) as window:
+                time_start = datetime.now(tz=timezone.utc)
+                time_end = time_start + timedelta(seconds=5)
 
-            time_start = datetime.now(tz=timezone.utc)
-            time_end = time_start + timedelta(seconds=5)
+                # ... wait for 5 seconds until the buffer is filled
+                await asyncio.sleep(5)
 
-            # ... wait for 5 seconds until the buffer is filled
-            await asyncio.sleep(5)
-
-            # return an numpy array from the window
-            array = window[time_start:time_end]
-            # and use it to for example calculate the mean
-            mean = array.mean()
+                # return an numpy array from the window
+                array = window[time_start:time_end]
+                # and use it to for example calculate the mean
+                mean = array.mean()
 
         asyncio.run(run())
         ```
@@ -112,19 +111,18 @@ async def run() -> None:
 
             # create a window that stores two days of data
             # starting at 1.1.23 with samplerate=1
-            window = MovingWindow(
+            async with MovingWindow(
                 size=timedelta(days=2),
                 resampled_data_recv=resampled_data_receiver,
                 input_sampling_period=timedelta(seconds=1),
-            )
-
-            # wait for one full day until the buffer is filled
-            await asyncio.sleep(60*60*24)
+            ) as window:
+                # wait for one full day until the buffer is filled
+                await asyncio.sleep(60*60*24)
 
-            # create a polars series with one full day of data
-            time_start = datetime(2023, 1, 1, tzinfo=timezone.utc)
-            time_end = datetime(2023, 1, 2, tzinfo=timezone.utc)
-            series = pl.Series("Jan_1", window[time_start:time_end])
+                # create a polars series with one full day of data
+                time_start = datetime(2023, 1, 1, tzinfo=timezone.utc)
+                time_end = datetime(2023, 1, 2, tzinfo=timezone.utc)
+                series = pl.Series("Jan_1", window[time_start:time_end])
 
         asyncio.run(run())
         ```
@@ -137,6 +135,8 @@ def __init__(  # pylint: disable=too-many-arguments
         input_sampling_period: timedelta,
         resampler_config: ResamplerConfig | None = None,
         align_to: datetime = UNIX_EPOCH,
+        *,
+        name: str | None = None,
     ) -> None:
         """
         Initialize the MovingWindow.
@@ -154,22 +154,21 @@ def __init__(  # pylint: disable=too-many-arguments
             align_to: A datetime object that defines a point in time to which
                 the window is aligned to modulo window size. For further
                 information, consult the class level documentation.
-
-        Raises:
-            asyncio.CancelledError: when the task gets cancelled.
+            name: The name of this moving window. If `None`, `str(id(self))` will be
+                used. This is used mostly for debugging purposes.
         """
         assert (
             input_sampling_period.total_seconds() > 0
         ), "The input sampling period should be greater than zero."
         assert (
             input_sampling_period <= size
         ), "The input sampling period should be equal to or lower than the window size."
+        super().__init__(name=name)
 
         self._sampling_period = input_sampling_period
 
         self._resampler: Resampler | None = None
         self._resampler_sender: Sender[Sample[Quantity]] | None = None
-        self._resampler_task: asyncio.Task[None] | None = None
 
         if resampler_config:
             assert (
@@ -191,12 +190,14 @@ def __init__(  # pylint: disable=too-many-arguments
             align_to=align_to,
         )
 
+    async def start(self) -> None:
+        """Start the MovingWindow.
+
+        This method starts the MovingWindow tasks.
+        """
         if self._resampler:
             self._configure_resampler()
-
-        self._update_window_task: asyncio.Task[None] = asyncio.create_task(
-            self._run_impl()
-        )
+        self._tasks.add(asyncio.create_task(self._run_impl(), name="update-window"))
 
     @property
     def sampling_period(self) -> timedelta:
@@ -228,12 +229,6 @@ async def _run_impl(self) -> None:
 
         _logger.error("Channel has been closed")
 
-    async def stop(self) -> None:
-        """Cancel the running tasks and stop the MovingWindow."""
-        await cancel_and_await(self._update_window_task)
-        if self._resampler_task:
-            await cancel_and_await(self._resampler_task)
-
     def _configure_resampler(self) -> None:
         """Configure the components needed to run the resampler."""
         assert self._resampler is not None
@@ -247,7 +242,9 @@ async def sink_buffer(sample: Sample[Quantity]) -> None:
         self._resampler.add_timeseries(
             "avg", resampler_channel.new_receiver(), sink_buffer
         )
-        self._resampler_task = asyncio.create_task(self._resampler.resample())
+        self._tasks.add(
+            asyncio.create_task(self._resampler.resample(), name="resample")
+        )
 
     def __len__(self) -> int:
         """

src/frequenz/sdk/timeseries/_periodic_feature_extractor.py

@@ -82,28 +82,27 @@ class PeriodicFeatureExtractor:
         from frequenz.sdk import microgrid
         from datetime import datetime, timedelta, timezone
 
-        moving_window = MovingWindow(
+        async with MovingWindow(
             size=timedelta(days=35),
             resampled_data_recv=microgrid.logical_meter().grid_power.new_receiver(),
             input_sampling_period=timedelta(seconds=1),
-        )
-
-        feature_extractor = PeriodicFeatureExtractor(
-            moving_window = moving_window,
-            period=timedelta(days=7),
-        )
+        ) as moving_window:
+            feature_extractor = PeriodicFeatureExtractor(
+                moving_window=moving_window,
+                period=timedelta(days=7),
+            )
 
-        now = datetime.now(timezone.utc)
+            now = datetime.now(timezone.utc)
 
-        # create a daily weighted average for the next 24h
-        avg_24h = feature_extractor.avg(
-            now,
-            now + timedelta(hours=24),
-            weights=[0.1, 0.2, 0.3, 0.4]
-        )
+            # create a daily weighted average for the next 24h
+            avg_24h = feature_extractor.avg(
+                now,
+                now + timedelta(hours=24),
+                weights=[0.1, 0.2, 0.3, 0.4]
+            )
 
-        # create a daily average for Thursday March 23 2023
-        th_avg_24h = feature_extractor.avg(datetime(2023, 3, 23), datetime(2023, 3, 24))
+            # create a daily average for Thursday March 23 2023
+            th_avg_24h = feature_extractor.avg(datetime(2023, 3, 23), datetime(2023, 3, 24))
         ```
     """
 

tests/timeseries/test_moving_window.py

@@ -75,15 +75,17 @@ def init_moving_window(
 async def test_access_window_by_index() -> None:
     """Test indexing a window by integer index"""
     window, sender = init_moving_window(timedelta(seconds=1))
-    await push_logical_meter_data(sender, [1])
-    assert np.array_equal(window[0], 1.0)
+    async with window:
+        await push_logical_meter_data(sender, [1])
+        assert np.array_equal(window[0], 1.0)
 
 
 async def test_access_window_by_timestamp() -> None:
     """Test indexing a window by timestamp"""
     window, sender = init_moving_window(timedelta(seconds=1))
-    await push_logical_meter_data(sender, [1])
-    assert np.array_equal(window[UNIX_EPOCH], 1.0)
+    async with window:
+        await push_logical_meter_data(sender, [1])
+        assert np.array_equal(window[UNIX_EPOCH], 1.0)
 
 
 async def test_access_window_by_int_slice() -> None:
@@ -94,35 +96,39 @@ async def test_access_window_by_int_slice() -> None:
     since the push_lm_data function is starting with the same initial timestamp.
     """
     window, sender = init_moving_window(timedelta(seconds=14))
-    await push_logical_meter_data(sender, range(0, 5))
-    assert np.array_equal(window[3:5], np.array([3.0, 4.0]))
+    async with window:
+        await push_logical_meter_data(sender, range(0, 5))
+        assert np.array_equal(window[3:5], np.array([3.0, 4.0]))
 
-    data = [1, 2, 2.5, 1, 1, 1, 2, 2, 1, 1, 1, 1, 1, 1]
-    await push_logical_meter_data(sender, data)
-    assert np.array_equal(window[5:14], np.array(data[5:14]))
+        data = [1, 2, 2.5, 1, 1, 1, 2, 2, 1, 1, 1, 1, 1, 1]
+        await push_logical_meter_data(sender, data)
+        assert np.array_equal(window[5:14], np.array(data[5:14]))
 
 
 async def test_access_window_by_ts_slice() -> None:
     """Test accessing a subwindow with a timestamp slice"""
     window, sender = init_moving_window(timedelta(seconds=5))
-    await push_logical_meter_data(sender, range(0, 5))
-    time_start = UNIX_EPOCH + timedelta(seconds=3)
-    time_end = time_start + timedelta(seconds=2)
-    assert np.array_equal(window[time_start:time_end], np.array([3.0, 4.0]))  # type: ignore
+    async with window:
+        await push_logical_meter_data(sender, range(0, 5))
+        time_start = UNIX_EPOCH + timedelta(seconds=3)
+        time_end = time_start + timedelta(seconds=2)
+        assert np.array_equal(window[time_start:time_end], np.array([3.0, 4.0]))  # type: ignore
 
 
 async def test_access_empty_window() -> None:
     """Test accessing an empty window, should throw IndexError"""
     window, _ = init_moving_window(timedelta(seconds=5))
-    with pytest.raises(IndexError, match=r"^The buffer is empty\.$"):
-        _ = window[42]
+    async with window:
+        with pytest.raises(IndexError, match=r"^The buffer is empty\.$"):
+            _ = window[42]
 
 
 async def test_window_size() -> None:
     """Test the size of the window."""
     window, sender = init_moving_window(timedelta(seconds=5))
-    await push_logical_meter_data(sender, range(0, 20))
-    assert len(window) == 5
+    async with window:
+        await push_logical_meter_data(sender, range(0, 20))
+        assert len(window) == 5
 
 
 # pylint: disable=redefined-outer-name
@@ -136,21 +142,20 @@ async def test_resampling_window(fake_time: time_machine.Coordinates) -> None:
     output_sampling = timedelta(seconds=2)
     resampler_config = ResamplerConfig(resampling_period=output_sampling)
 
-    window = MovingWindow(
+    async with MovingWindow(
         size=window_size,
         resampled_data_recv=channel.new_receiver(),
         input_sampling_period=input_sampling,
         resampler_config=resampler_config,
-    )
-
-    stream_values = [4.0, 8.0, 2.0, 6.0, 5.0] * 100
-    for value in stream_values:
-        timestamp = datetime.now(tz=timezone.utc)
-        sample = Sample(timestamp, Quantity(float(value)))
-        await sender.send(sample)
-        await asyncio.sleep(0.1)
-        fake_time.shift(0.1)
-
-    assert len(window) == window_size / output_sampling
-    for value in window:  # type: ignore
-        assert 4.9 < value < 5.1
+    ) as window:
+        stream_values = [4.0, 8.0, 2.0, 6.0, 5.0] * 100
+        for value in stream_values:
+            timestamp = datetime.now(tz=timezone.utc)
+            sample = Sample(timestamp, Quantity(float(value)))
+            await sender.send(sample)
+            await asyncio.sleep(0.1)
+            fake_time.shift(0.1)
+
+        assert len(window) == window_size / output_sampling
+        for value in window:  # type: ignore
+            assert 4.9 < value < 5.1