Allow TelemetryStream AST nodes to subscribe to data on demand

Earlier we were subscribing to all streams, even if they were going to
be discarded in a coalesce.

This commit makes further progress towards implementing lazy coalesce.

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

Author: shsms

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

@@ -6,9 +6,10 @@
 from __future__ import annotations
 
 import logging
-from collections.abc import AsyncIterator, Callable
+from collections.abc import Callable, Coroutine
 from dataclasses import dataclass
 
+from frequenz.channels import Receiver
 from frequenz.quantities import Quantity
 from typing_extensions import TypeIs, override
 
@@ -25,10 +26,24 @@ class TelemetryStream(AstNode[QuantityT]):
     """A AST node that retrieves values from a component's telemetry stream."""
 
     source: str
-    stream: AsyncIterator[Sample[QuantityT] | Sample[Quantity]]
+    metric_fetcher: (
+        Callable[
+            [], Coroutine[None, None, Receiver[Sample[QuantityT] | Sample[Quantity]]]
+        ]
+        | None
+    ) = None
     create_method: Callable[[float], QuantityT]
+    _stream: Receiver[Sample[QuantityT] | Sample[Quantity]] | None = None
     _latest_sample: Sample[QuantityT] | None = None
 
+    def __post_init__(self) -> None:
+        """Validate at least one of stream or metric_fetcher is set."""
+        if self._stream is None and self.metric_fetcher is None:
+            raise ValueError(
+                "Either stream or metric_fetcher must be provided for "
+                + "TelemetryStream node."
+            )
+
     @property
     def latest_sample(self) -> Sample[QuantityT] | None:
         """Return the latest fetched sample for this component."""
@@ -48,7 +63,13 @@ def format(self, wrap: bool = False) -> str:
 
     async def fetch_next(self) -> None:
         """Fetch the next value for this component and store it internally."""
-        latest_sample = await anext(self.stream)
+        if self._stream is None:
+            await self._fetch_stream()
+        assert self._stream is not None
+
+        latest_sample = await anext(self._stream)
+        # pylint: disable-next=fixme
+        # TODO: convert to QuantityT if needed only at the end in the evaluator.
         if self._is_quantity_sample(latest_sample):
             assert latest_sample.value is not None
             self._latest_sample = Sample(
@@ -58,6 +79,14 @@ async def fetch_next(self) -> None:
         else:
             self._latest_sample = latest_sample
 
+    async def _fetch_stream(self) -> None:
+        """Subscribe to the telemetry stream for this component."""
+        if self._stream is not None:
+            return
+        if self.metric_fetcher is None:
+            raise RuntimeError("Metric fetcher is not set for TelemetryStream node.")
+        self._stream = await self.metric_fetcher()
+
     def _is_quantity_sample(
         self, sample: Sample[QuantityT] | Sample[Quantity]
     ) -> TypeIs[Sample[Quantity]]:

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

@@ -6,7 +6,7 @@
 from __future__ import annotations
 
 import logging
-from collections.abc import Callable
+from collections.abc import Callable, Coroutine
 from typing import Generic
 
 from frequenz.channels import Broadcast, Receiver
@@ -27,6 +27,18 @@
 _logger = logging.getLogger(__name__)
 
 
+def metric_fetcher(
+    formula: Formula[QuantityT],
+) -> Callable[[], Coroutine[None, None, Receiver[Sample[QuantityT]]]]:
+    """Fetch a receiver for the formula's output samples."""
+
+    async def fetcher(formula: Formula[QuantityT]) -> Receiver[Sample[QuantityT]]:
+        formula.start()
+        return formula.new_receiver()
+
+    return lambda: fetcher(formula)
+
+
 class Formula(BackgroundService, ReceiverFetcher[Sample[QuantityT]]):
     """A formula represented as an AST."""
 
@@ -178,7 +190,7 @@ def __init__(
         if isinstance(formula, Formula):
             self.root: AstNode[QuantityT] = _ast.TelemetryStream(
                 source=str(formula),
-                stream=formula.new_receiver(),
+                metric_fetcher=metric_fetcher(formula),
                 create_method=create_method,
             )
             self._streams.append(self.root)
@@ -197,7 +209,7 @@ def __add__(
         elif isinstance(other, Formula):
             right_node = _ast.TelemetryStream(
                 source=str(other),
-                stream=other.new_receiver(),
+                metric_fetcher=metric_fetcher(other),
                 create_method=self._create_method,
             )
             self._streams.append(right_node)
@@ -224,7 +236,7 @@ def __sub__(
         elif isinstance(other, Formula):
             right_node = _ast.TelemetryStream(
                 source=str(other),
-                stream=other.new_receiver(),
+                metric_fetcher=metric_fetcher(other),
                 create_method=self._create_method,
             )
             self._streams.append(right_node)
@@ -278,7 +290,7 @@ def coalesce(
             elif isinstance(item, Formula):
                 right_node = _ast.TelemetryStream(
                     source=str(item),
-                    stream=item.new_receiver(),
+                    metric_fetcher=metric_fetcher(item),
                     create_method=self._create_method,
                 )
                 right_nodes.append(right_node)
@@ -311,7 +323,7 @@ def min(
             elif isinstance(item, Formula):
                 right_node = _ast.TelemetryStream(
                     source=str(item),
-                    stream=item.new_receiver(),
+                    metric_fetcher=metric_fetcher(item),
                     create_method=self._create_method,
                 )
                 right_nodes.append(right_node)
@@ -344,7 +356,7 @@ def max(
             elif isinstance(item, Formula):
                 right_node = _ast.TelemetryStream(
                     source=str(item),
-                    stream=item.new_receiver(),
+                    metric_fetcher=metric_fetcher(item),
                     create_method=self._create_method,
                 )
                 right_nodes.append(right_node)

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

@@ -13,7 +13,7 @@
 from ...actor import Actor
 from .._base_types import QuantityT, Sample3Phase
 from . import _ast
-from ._formula import Formula
+from ._formula import Formula, metric_fetcher
 from ._formula_evaluator import synchronize_receivers
 
 _logger = logging.getLogger(__name__)
@@ -45,17 +45,17 @@ def __init__(
         self._components: list[_ast.TelemetryStream[QuantityT]] = [
             _ast.TelemetryStream(
                 source="phase_1",
-                stream=phase_1.new_receiver(),
+                metric_fetcher=metric_fetcher(phase_1),
                 create_method=phase_1._create_method,  # pylint: disable=protected-access
             ),
             _ast.TelemetryStream(
                 source="phase_2",
-                stream=phase_2.new_receiver(),
+                metric_fetcher=metric_fetcher(phase_2),
                 create_method=phase_2._create_method,  # pylint: disable=protected-access
             ),
             _ast.TelemetryStream(
                 source="phase_3",
-                stream=phase_3.new_receiver(),
+                metric_fetcher=metric_fetcher(phase_3),
                 create_method=phase_3._create_method,  # pylint: disable=protected-access
             ),
         ]

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

@@ -53,8 +53,6 @@ def __init__(  # pylint: disable=too-many-arguments
     @override
     async def _run(self) -> None:
         """Run the formula evaluator actor."""
-        if self._metric_fetcher is not None:
-            await self._metric_fetcher.subscribe()
         await synchronize_receivers(self._components)
 
         while True:

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

@@ -69,7 +69,7 @@ async def __call__(self) -> Sample[QuantityT] | QuantityT | None:
         for param in self.params:
             arg = await param.evaluate()
             match arg:
-                case Sample(value=value, timestamp=timestamp):
+                case Sample(timestamp, value):
                     if value is not None:
                         return arg
                     ts = timestamp

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

@@ -5,11 +5,15 @@
 
 from __future__ import annotations
 
-from collections.abc import Callable
+import logging
+from collections.abc import Callable, Coroutine
 from typing import Generic
 
+from frequenz.channels import Receiver
 from frequenz.client.common.microgrid.components import ComponentId
+from frequenz.quantities import Quantity
 
+from frequenz.sdk.timeseries import Sample
 from frequenz.sdk.timeseries._base_types import QuantityT
 
 from . import _ast, _token
@@ -20,6 +24,8 @@
 from ._peekable import Peekable
 from ._resampled_stream_fetcher import ResampledStreamFetcher
 
+_logger = logging.getLogger(__name__)
+
 
 def parse(
     *,
@@ -181,12 +187,19 @@ def _parse_primary(self) -> AstNode[QuantityT] | None:
         if token is None:
             return None
 
+        def make_component_stream_fetcher(
+            f: ResampledStreamFetcher, cid: ComponentId
+        ) -> Callable[[], Coroutine[None, None, Receiver[Sample[Quantity]]]]:
+            return lambda: f.fetch_stream(cid)
+
         if isinstance(token, _token.Component):
             _ = next(self._lexer)  # consume token
             comp = _ast.TelemetryStream(
                 span=token.span,
                 source=f"#{token.id}",
-                stream=self._telemetry_fetcher.fetch_stream(ComponentId(int(token.id))),
+                metric_fetcher=make_component_stream_fetcher(
+                    self._telemetry_fetcher, ComponentId(int(token.id))
+                ),
                 create_method=self._create_method,
             )
             self._components.append(comp)

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

@@ -7,10 +7,11 @@
 from frequenz.client.common.microgrid.components import ComponentId
 from frequenz.quantities import Quantity
 
+from frequenz.sdk.timeseries import Sample
+
 from ..._internal._channels import ChannelRegistry
 from ...microgrid._data_sourcing import ComponentMetricRequest, Metric
 from ...microgrid._old_component_data import TransitionalMetric
-from .. import Sample
 
 
 class ResampledStreamFetcher:
@@ -41,9 +42,7 @@ def __init__(
         )
         self._metric: Metric | TransitionalMetric = metric
 
-        self._pending_requests: list[ComponentMetricRequest] = []
-
-    def fetch_stream(
+    async def fetch_stream(
         self,
         component_id: ComponentId,
     ) -> Receiver[Sample[Quantity]]:
@@ -61,13 +60,8 @@ def fetch_stream(
             self._metric,
             None,
         )
-        self._pending_requests.append(request)
+        await self._resampler_subscription_sender.send(request)
+
         return self._channel_registry.get_or_create(
             Sample[Quantity], request.get_channel_name()
         ).new_receiver()
-
-    async def subscribe(self) -> None:
-        """Subscribe to all resampled component metric streams."""
-        for request in self._pending_requests:
-            await self._resampler_subscription_sender.send(request)
-        self._pending_requests.clear()

tests/microgrid/test_grid.py

@@ -65,7 +65,7 @@ async def test_grid_power_1(mocker: MockerFixture) -> None:
 
         grid_power_recv = grid.power.new_receiver()
 
-        grid_meter_recv = get_resampled_stream(
+        grid_meter_recv = await get_resampled_stream(
             grid._formula_pool._namespace,  # pylint: disable=protected-access
             mockgrid.meter_ids[0],
             Metric.AC_ACTIVE_POWER,
@@ -109,7 +109,7 @@ async def test_grid_power_2(mocker: MockerFixture) -> None:
         grid_power_recv = grid.power.new_receiver()
 
         component_receivers = [
-            get_resampled_stream(
+            await get_resampled_stream(
                 grid._formula_pool._namespace,  # pylint: disable=protected-access
                 component_id,
                 Metric.AC_ACTIVE_POWER,
@@ -160,7 +160,7 @@ async def test_grid_reactive_power_1(mocker: MockerFixture) -> None:
 
         grid_power_recv = grid.reactive_power.new_receiver()
 
-        grid_meter_recv = get_resampled_stream(
+        grid_meter_recv = await get_resampled_stream(
             grid._formula_pool._namespace,  # pylint: disable=protected-access
             mockgrid.meter_ids[0],
             Metric.AC_REACTIVE_POWER,
@@ -204,7 +204,7 @@ async def test_grid_reactive_power_2(mocker: MockerFixture) -> None:
         grid_power_recv = grid.reactive_power.new_receiver()
 
         component_receivers = [
-            get_resampled_stream(
+            await get_resampled_stream(
                 grid._formula_pool._namespace,  # pylint: disable=protected-access
                 component_id,
                 Metric.AC_REACTIVE_POWER,

tests/timeseries/_formulas/test_formula_composition.py

@@ -45,7 +45,7 @@ async def test_formula_composition(  # pylint: disable=too-many-locals
             grid = microgrid.grid()
             stack.push_async_callback(grid.stop)
 
-            grid_meter_recv = get_resampled_stream(
+            grid_meter_recv = await get_resampled_stream(
                 grid._formula_pool._namespace,  # pylint: disable=protected-access
                 mockgrid.meter_ids[0],
                 Metric.AC_ACTIVE_POWER,