Implement lazy subscription for Coalesce function

Signed-off-by: Sahas Subramanian <[email protected]>

Author: shsms

src/frequenz/sdk/timeseries/formulas/_functions.py

@@ -7,6 +7,7 @@
 
 import abc
 import asyncio
+import logging
 from dataclasses import dataclass, field
 from datetime import datetime
 from typing import Generic
@@ -15,8 +16,11 @@
 from typing_extensions import override
 
 from .._base_types import QuantityT, Sample
+from ._ast import Constant
 from ._base_ast_node import AstNode, NodeSynchronizer
 
+_logger = logging.getLogger(__name__)
+
 
 @dataclass(kw_only=True)
 class FunCall(AstNode[QuantityT]):
@@ -99,6 +103,8 @@ def from_string(
 class Coalesce(Function[QuantityT]):
     """A function that returns the first non-None argument."""
 
+    num_subscribed: int = 0
+
     @property
     @override
     def name(self) -> str:
@@ -110,22 +116,70 @@ async def __call__(self) -> Sample[QuantityT] | QuantityT | None:
         """Return the first non-None argument."""
         ts: datetime | None = None
 
-        args = await self._synchronizer.evaluate(self.params)
-        for arg in args:
+        if self.num_subscribed == 0:
+            await self._subscribe_next()
+
+        args = await self._synchronizer.evaluate(
+            self.params[: self.num_subscribed], sync_to_first_node=True
+        )
+        for ctr, arg in enumerate(args, start=1):
             match arg:
                 case Sample(timestamp, value):
                     if value is not None:
+                        # Found a non-None value, unsubscribe from subsequent params
+                        if ctr < self.num_subscribed:
+                            await self._unsubscribe_after(ctr)
                         return arg
                     ts = timestamp
                 case Quantity():
+                    # Found a non-None value, unsubscribe from subsequent params
+                    if ctr < self.num_subscribed:
+                        await self._unsubscribe_after(ctr)
                     if ts is not None:
                         return Sample(timestamp=ts, value=arg)
                     return arg
                 case None:
                     continue
+        # Don't have a non-None value yet, subscribe to the next parameter for
+        # next time and return None for now, unless the next value is a constant.
+        next_value: Sample[QuantityT] | QuantityT | None = None
+        await self._subscribe_next()
+
+        if isinstance(self.params[self.num_subscribed - 1], Constant):
+            next_value = await self.params[self.num_subscribed - 1].evaluate()
+        if isinstance(next_value, Sample):
+            return next_value
+
         if ts is not None:
-            return Sample(timestamp=ts, value=None)
-        return None
+            return Sample(timestamp=ts, value=next_value)
+        return next_value
+
+    @override
+    async def subscribe(self) -> None:
+        """Subscribe to the first parameter if not already subscribed."""
+        if self.num_subscribed == 0:
+            await self._subscribe_next()
+
+    async def _subscribe_next(self) -> None:
+        """Subscribe to the next parameter."""
+        if self.num_subscribed < len(self.params):
+            _logger.debug(
+                "Coalesce subscribing to param %d: %s",
+                self.num_subscribed + 1,
+                self.params[self.num_subscribed],
+            )
+            await self.params[self.num_subscribed].subscribe()
+            self.num_subscribed += 1
+
+    async def _unsubscribe_after(self, index: int) -> None:
+        """Unsubscribe from parameters after the given index."""
+        for param in self.params[index:]:
+            _logger.debug(
+                "Coalesce unsubscribing from param: %s",
+                param,
+            )
+            await param.unsubscribe()
+        self.num_subscribed = index
 
 
 class Max(Function[QuantityT]):

tests/timeseries/_formulas/test_formula_composition.py

@@ -230,11 +230,29 @@ async def test_formula_composition_min_max(self, mocker: MockerFixture) -> None:
             stack.push_async_callback(formula_max.stop)
             formula_max_rx = formula_max.new_receiver()
 
+            assert (
+                str(formula_min)
+                == "[grid_power_min]("
+                + "MIN([grid_power](COALESCE(#4, #7)), [chp_power](COALESCE(#5, #7, 0.0)))"
+                + ")"
+            )
+            assert (
+                str(formula_max)
+                == "[grid_power_max]("
+                + "MAX([grid_power](COALESCE(#4, #7)), [chp_power](COALESCE(#5, #7, 0.0)))"
+                + ")"
+            )
+
+            await mockgrid.mock_resampler.send_meter_power([100.0, 200.0])
+            await mockgrid.mock_resampler.send_chp_power([None])
+            # Because it got None for CHP (#5), it will then subscribe to the meter (#7).
+            # So we have to send again so it uses the meter value.
             await mockgrid.mock_resampler.send_meter_power([100.0, 200.0])
             await mockgrid.mock_resampler.send_chp_power([None])
 
             # Test min
             min_pow = await formula_min_rx.receive()
+            min_pow = await formula_min_rx.receive()
             assert (
                 min_pow
                 and min_pow.value
@@ -243,17 +261,21 @@ async def test_formula_composition_min_max(self, mocker: MockerFixture) -> None:
 
             # Test max
             max_pow = await formula_max_rx.receive()
+            max_pow = await formula_max_rx.receive()
             assert (
                 max_pow
                 and max_pow.value
                 and max_pow.value.isclose(Power.from_watts(200.0))
             )
 
+            await mockgrid.mock_resampler.send_meter_power([-100.0, -200.0])
+            await mockgrid.mock_resampler.send_chp_power([None])
             await mockgrid.mock_resampler.send_meter_power([-100.0, -200.0])
             await mockgrid.mock_resampler.send_chp_power([None])
 
             # Test min
             min_pow = await formula_min_rx.receive()
+            min_pow = await formula_min_rx.receive()
             assert (
                 min_pow
                 and min_pow.value
@@ -262,6 +284,7 @@ async def test_formula_composition_min_max(self, mocker: MockerFixture) -> None:
 
             # Test max
             max_pow = await formula_max_rx.receive()
+            max_pow = await formula_max_rx.receive()
             assert (
                 max_pow
                 and max_pow.value

tests/timeseries/_formulas/test_formulas.py

@@ -4,6 +4,7 @@
 """Tests for the Formula implementation."""
 
 import asyncio
+import logging
 from collections import OrderedDict
 from collections.abc import Callable
 from datetime import datetime, timedelta
@@ -22,6 +23,8 @@
     ResampledStreamFetcher,
 )
 
+_logger = logging.getLogger(__name__)
+
 
 @pytest.fixture
 def event_loop_policy() -> async_solipsism.EventLoopPolicy:
@@ -40,6 +43,7 @@ async def run_test(  # pylint: disable=too-many-locals
         io_pairs: list[tuple[list[float | None], float | None]],
     ) -> None:
         """Run a formula test."""
+        _logger.debug("TESTING FORMULA: %s", formula_str)
         channels: OrderedDict[ComponentId, Broadcast[Sample[Quantity]]] = OrderedDict()
         for comp_id in component_ids:
             channels[ComponentId(comp_id)] = Broadcast(
@@ -67,6 +71,7 @@ def stream_recv(comp_id: ComponentId) -> Receiver[Sample[Quantity]]:
 
             for io_pair in io_pairs:
                 io_input, io_output = io_pair
+                now += timedelta(seconds=1)
                 _ = await asyncio.gather(
                     *[
                         chan.new_sender().send(
@@ -90,6 +95,7 @@ def stream_recv(comp_id: ComponentId) -> Receiver[Sample[Quantity]]:
                         and next_val.value.base_value == io_output
                     )
                 tests_passed += 1
+                _logger.debug("%s: Passed inputs: %s", tests_passed, io_input)
             assert tests_passed == len(io_pairs)
 
     async def test_simple(self) -> None:
@@ -287,6 +293,17 @@ async def test_max_min_coalesce(self) -> None:
                 ([10.0, 12.0, 15.0], 25.0),
             ],
         )
+        await self.run_test(
+            "#2 + COALESCE(#4, #5, 0.0)",
+            "[f](#2 + COALESCE(#4, #5, 0.0))",
+            [2, 4, 5],
+            [
+                ([10.0, 12.0, 15.0], 22.0),
+                ([10.0, None, 15.0], None),
+                ([10.0, None, 15.0], 25.0),
+                ([10.0, None, None], 10.0),
+            ],
+        )
         await self.run_test(
             "MIN(#2, #4) + COALESCE(#5, 0.0)",
             "[f](MIN(#2, #4) + COALESCE(#5, 0.0))",
@@ -370,6 +387,8 @@ def stream_recv(comp_id: int) -> Receiver[Sample[Quantity]]:
 
         assert str(formula) == expected
 
+        _logger.debug("TESTING FORMULA: %s", expected)
+
         result_chan = formula.new_receiver()
         await asyncio.sleep(0.1)
         now = datetime.now()
@@ -394,6 +413,7 @@ def stream_recv(comp_id: int) -> Receiver[Sample[Quantity]]:
                     and next_val.value.base_value == io_output
                 )
             tests_passed += 1
+            _logger.debug("%s: Passed inputs: %s", tests_passed, io_input)
         await formula.stop()
         assert tests_passed == len(io_pairs)
 
@@ -562,7 +582,9 @@ async def test_coalesce(self) -> None:
             lambda c2, c4, c5: c2.coalesce([c4, c5]),
             "[l2](COALESCE([0](#0), [1](#1), [2](#2)))",
             [
+                ([None, 12.0, 15.0], None),
                 ([None, 12.0, 15.0], 12.0),
+                ([None, None, 15.0], None),
                 ([None, None, 15.0], 15.0),
                 ([10.0, None, 15.0], 10.0),
                 ([None, None, None], None),
@@ -574,9 +596,14 @@ async def test_coalesce(self) -> None:
             lambda c2, c4, c5: (c2 * 5.0).coalesce([c4 / 2.0, c5]),
             "[l2](COALESCE([0](#0) * 5.0, [1](#1) / 2.0, [2](#2)))",
             [
+                ([None, 12.0, 15.0], None),
                 ([None, 12.0, 15.0], 6.0),
+                ([None, None, 15.0], None),
                 ([None, None, 15.0], 15.0),
                 ([10.0, None, 15.0], 50.0),
+                ([None, None, 15.0], None),
+                ([None, None, 15.0], None),
+                ([None, None, 15.0], 15.0),
                 ([None, None, None], None),
             ],
         )

tests/timeseries/_formulas/test_formulas_3_phase.py

@@ -82,6 +82,7 @@ def stream_recv(comp_id: int) -> Receiver[Sample[Quantity]]:
         )
         builder = make_builder(p3_formula_1, p3_formula_2)
         formula = builder.build("l2")
+
         receiver = formula.new_receiver()
 
         await asyncio.sleep(0.1)
@@ -154,7 +155,9 @@ async def test_composition(self) -> None:
         await self.run_test(
             [
                 ([(4.0, 9.0, 16.0), (2.0, 13.0, 4.0)], (4.0, 9.0, 16.0)),
+                ([(-5.0, 10.0, None), (10.0, 5.0, None)], (-5.0, 10.0, None)),
                 ([(-5.0, 10.0, None), (10.0, 5.0, None)], (-5.0, 10.0, 0.0)),
+                ([(None, 2.0, 3.0), (2.0, None, 4.0)], (None, 2.0, 3.0)),
                 ([(None, 2.0, 3.0), (2.0, None, 4.0)], (2.0, 2.0, 3.0)),
             ],
             lambda f1, f2: f1.coalesce(