Make FormulaEngine the owner of the run task and output channel

... instead of the `LogicalMeter`.  Also add a `name` parameter to the
`FormulaEngine` and the formula builders.

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

Author: shsms

src/frequenz/sdk/timeseries/logical_meter/_formula_engine.py

@@ -6,10 +6,11 @@
 from __future__ import annotations
 
 import asyncio
+import logging
 from datetime import datetime
 from typing import Dict, List, Optional, Set, Tuple
 
-from frequenz.channels import Receiver
+from frequenz.channels import Broadcast, Receiver
 
 from .. import Sample
 from ._formula_steps import (
@@ -23,6 +24,8 @@
     Subtractor,
 )
 
+logger = logging.Logger(__name__)
+
 _operator_precedence = {
     "(": 0,
     "/": 1,
@@ -40,17 +43,24 @@ class FormulaEngine:
     """
 
     def __init__(
-        self, steps: List[FormulaStep], metric_fetchers: Dict[str, MetricFetcher]
+        self,
+        name: str,
+        steps: List[FormulaStep],
+        metric_fetchers: Dict[str, MetricFetcher],
     ) -> None:
         """Create a `FormulaEngine` instance.
 
         Args:
+            name: A name for the formula.
             steps: Steps for the engine to execute, in post-fix order.
             metric_fetchers: Fetchers for each metric stream the formula depends on.
         """
+        self._name = name
         self._steps = steps
         self._metric_fetchers = metric_fetchers
         self._first_run = True
+        self._channel = Broadcast[Sample](self._name)
+        self._task = None
 
     async def _synchronize_metric_timestamps(
         self, metrics: Set[asyncio.Task[Optional[Sample]]]
@@ -93,12 +103,13 @@ async def _synchronize_metric_timestamps(
                 metric_ts = next_val.timestamp
             if metric_ts > latest_ts:
                 raise RuntimeError(
-                    "Unable to synchronize timestamps of resampled metrics"
+                    "Unable to synchronize resampled metric timestamps, "
+                    f"for formula: {self._name}"
                 )
         self._first_run = False
         return latest_ts
 
-    async def apply(self) -> Sample:
+    async def _apply(self) -> Sample:
         """Fetch the latest metrics, apply the formula once and return the result.
 
         Returns:
@@ -118,7 +129,9 @@ async def apply(self) -> Sample:
         )
 
         if pending or any(res.result() is None for res in iter(ready_metrics)):
-            raise RuntimeError("Some resampled metrics didn't arrive")
+            raise RuntimeError(
+                f"Some resampled metrics didn't arrive, for formula: {self._name}"
+            )
 
         if self._first_run:
             metric_ts = await self._synchronize_metric_timestamps(ready_metrics)
@@ -133,10 +146,40 @@ async def apply(self) -> Sample:
         # if all steps were applied and the formula was correct, there should only be a
         # single value in the evaluation stack, and that would be the formula result.
         if len(eval_stack) != 1:
-            raise RuntimeError("Formula application failed.")
+            raise RuntimeError(f"Formula application failed: {self._name}")
 
         return Sample(metric_ts, eval_stack[0])
 
+    async def _run(self) -> None:
+        sender = self._channel.new_sender()
+        while True:
+            try:
+                msg = await self._apply()
+            except asyncio.CancelledError:
+                logger.exception("FormulaEngine task cancelled: %s", self._name)
+                break
+            except Exception as err:  # pylint: disable=broad-except
+                logger.warning(
+                    "Formula application failed: %s. Error: %s", self._name, err
+                )
+            else:
+                await sender.send(msg)
+
+    def new_receiver(self) -> Receiver[Sample]:
+        """Create a new receiver that streams the output of the formula engine.
+
+        Args:
+            name: An optional name for the receiver.
+            max_size: The size of the receiver's buffer.
+
+        Returns:
+            A receiver that streams output `Sample`s from the formula engine.
+        """
+        if self._task is None:
+            self._task = asyncio.create_task(self._run())
+
+        return self._channel.new_receiver()
+
 
 class FormulaBuilder:
     """Builds a post-fix formula engine that operates on `Sample` receivers.
@@ -161,10 +204,13 @@ class FormulaBuilder:
         add the values and return the result.
     """
 
-    def __init__(
-        self,
-    ) -> None:
-        """Create a `FormulaBuilder` instance."""
+    def __init__(self, name: str) -> None:
+        """Create a `FormulaBuilder` instance.
+
+        Args:
+            name: A name for the formula being built.
+        """
+        self._name = name
         self._build_stack: List[FormulaStep] = []
         self._steps: List[FormulaStep] = []
         self._metric_fetchers: Dict[str, MetricFetcher] = {}
@@ -242,4 +288,4 @@ def build(self) -> FormulaEngine:
         while self._build_stack:
             self._steps.append(self._build_stack.pop())
 
-        return FormulaEngine(self._steps, self._metric_fetchers)
+        return FormulaEngine(self._name, self._steps, self._metric_fetchers)

src/frequenz/sdk/timeseries/logical_meter/_formula_generators/_battery_power_formula.py

@@ -32,7 +32,7 @@ async def generate(
             FormulaGenerationError: If a battery has a non-inverter predecessor
                 in the component graph.
         """
-        builder = self._get_builder(ComponentMetricId.ACTIVE_POWER)
+        builder = self._get_builder("battery-power", ComponentMetricId.ACTIVE_POWER)
         component_graph = microgrid.get().component_graph
         battery_inverters = list(
             comp

src/frequenz/sdk/timeseries/logical_meter/_formula_generators/_battery_soc_formula.py

@@ -62,7 +62,7 @@ async def generate(
             FormulaGenerationError: If a battery has a non-inverter predecessor
                 in the component graph.
         """
-        builder = self._get_builder(ComponentMetricId.ACTIVE_POWER)
+        builder = self._get_builder("soc", ComponentMetricId.ACTIVE_POWER)
         component_graph = microgrid.get().component_graph
         inv_bat_pairs = {
             comp: component_graph.successors(comp.component_id)

src/frequenz/sdk/timeseries/logical_meter/_formula_generators/_formula_generator.py

@@ -44,10 +44,11 @@ def __init__(
         self._namespace = namespace
 
     def _get_builder(
-        self, component_metric_id: ComponentMetricId
+        self, name: str, component_metric_id: ComponentMetricId
     ) -> ResampledFormulaBuilder:
         builder = ResampledFormulaBuilder(
             self._namespace,
+            name,
             self._channel_registry,
             self._resampler_subscription_sender,
             component_metric_id,

src/frequenz/sdk/timeseries/logical_meter/_formula_generators/_grid_power_formula.py

@@ -23,7 +23,7 @@ async def generate(
         Raises:
             ComponentNotFound: when the component graph doesn't have a `GRID` component.
         """
-        builder = self._get_builder(ComponentMetricId.ACTIVE_POWER)
+        builder = self._get_builder("grid-power", ComponentMetricId.ACTIVE_POWER)
         component_graph = microgrid.get().component_graph
         grid_component = next(
             (

src/frequenz/sdk/timeseries/logical_meter/_formula_generators/_pv_power_formula.py

@@ -21,7 +21,7 @@ async def generate(self) -> FormulaEngine:
         Raises:
             ComponentNotFound: if there are no PV inverters in the component graph.
         """
-        builder = self._get_builder(ComponentMetricId.ACTIVE_POWER)
+        builder = self._get_builder("pv-power", ComponentMetricId.ACTIVE_POWER)
 
         component_graph = microgrid.get().component_graph
         pv_inverters = list(

src/frequenz/sdk/timeseries/logical_meter/_logical_meter.py

@@ -10,7 +10,7 @@
 import uuid
 from typing import Dict, List, Type
 
-from frequenz.channels import Broadcast, Receiver, Sender
+from frequenz.channels import Receiver, Sender
 
 from ...actor import ChannelRegistry, ComponentMetricRequest
 from ...microgrid import ComponentGraph
@@ -63,40 +63,21 @@ def __init__(
         # meter to use when communicating with the resampling actor.
         self._namespace = f"logical-meter-{uuid.uuid4()}"
         self._component_graph = component_graph
-        self._output_channels: Dict[str, Broadcast[Sample]] = {}
+        self._engines: Dict[str, FormulaEngine] = {}
         self._tasks: List[asyncio.Task[None]] = []
 
     async def _engine_from_formula_string(
         self, formula: str, metric_id: ComponentMetricId, nones_are_zeros: bool
     ) -> FormulaEngine:
         builder = ResampledFormulaBuilder(
             self._namespace,
+            formula,
             self._channel_registry,
             self._resampler_subscription_sender,
             metric_id,
         )
         return await builder.from_string(formula, nones_are_zeros)
 
-    async def _run_formula(
-        self, formula: FormulaEngine, sender: Sender[Sample]
-    ) -> None:
-        """Run the formula repeatedly and send the results to a channel.
-
-        Args:
-            formula: The formula to run.
-            sender: A sender for sending the formula results to.
-        """
-        while True:
-            try:
-                msg = await formula.apply()
-            except asyncio.CancelledError:
-                logger.exception("LogicalMeter task cancelled")
-                break
-            except Exception as err:  # pylint: disable=broad-except
-                logger.warning("Formula application failed: %s", err)
-            else:
-                await sender.send(msg)
-
     async def start_formula(
         self,
         formula: str,
@@ -116,40 +97,29 @@ async def start_formula(
             A Receiver that streams values with the formulas applied.
         """
         channel_key = formula + component_metric_id.value
-        if channel_key in self._output_channels:
-            return self._output_channels[channel_key].new_receiver()
+        if channel_key in self._engines:
+            return self._engines[channel_key].new_receiver()
 
         formula_engine = await self._engine_from_formula_string(
             formula, component_metric_id, nones_are_zeros
         )
-        out_chan = Broadcast[Sample](channel_key)
-        self._output_channels[channel_key] = out_chan
-        self._tasks.append(
-            asyncio.create_task(
-                self._run_formula(formula_engine, out_chan.new_sender())
-            )
-        )
-        return out_chan.new_receiver()
+        self._engines[channel_key] = formula_engine
+
+        return formula_engine.new_receiver()
 
     async def _get_formula_stream(
         self,
         channel_key: str,
         generator: Type[FormulaGenerator],
     ) -> Receiver[Sample]:
-        if channel_key in self._output_channels:
-            return self._output_channels[channel_key].new_receiver()
+        if channel_key in self._engines:
+            return self._engines[channel_key].new_receiver()
 
-        formula_engine = await generator(
+        engine = await generator(
             self._namespace, self._channel_registry, self._resampler_subscription_sender
         ).generate()
-        out_chan = Broadcast[Sample](channel_key)
-        self._output_channels[channel_key] = out_chan
-        self._tasks.append(
-            asyncio.create_task(
-                self._run_formula(formula_engine, out_chan.new_sender())
-            )
-        )
-        return out_chan.new_receiver()
+        self._engines[channel_key] = engine
+        return engine.new_receiver()
 
     async def grid_power(self) -> Receiver[Sample]:
         """Fetch the grid power for the microgrid.

src/frequenz/sdk/timeseries/logical_meter/_resampled_formula_builder.py

@@ -16,9 +16,10 @@
 class ResampledFormulaBuilder(FormulaBuilder):
     """Provides a way to build a FormulaEngine from resampled data streams."""
 
-    def __init__(
+    def __init__(  # pylint: disable=too-many-arguments
         self,
         namespace: str,
+        formula_name: str,
         channel_registry: ChannelRegistry,
         resampler_subscription_sender: Sender[ComponentMetricRequest],
         metric_id: ComponentMetricId,
@@ -28,6 +29,7 @@ def __init__(
         Args:
             namespace: The unique namespace to allow reuse of streams in the data
                 pipeline.
+            formula_name: A name for the formula.
             channel_registry: The channel registry instance shared with the resampling
                 and the data sourcing actors.
             resampler_subscription_sender: A sender to send metric requests to the
@@ -38,7 +40,7 @@ def __init__(
         self._resampler_subscription_sender = resampler_subscription_sender
         self._namespace = namespace
         self._metric_id = metric_id
-        super().__init__()
+        super().__init__(formula_name)
 
     async def _get_resampled_receiver(
         self, component_id: int, metric_id: ComponentMetricId

tests/timeseries/test_formula_engine.py

@@ -48,7 +48,7 @@ async def run_test(
     ) -> None:
         """Run a formula test."""
         channels: Dict[str, Broadcast[Sample]] = {}
-        builder = FormulaBuilder()
+        builder = FormulaBuilder("test_formula")
         for token in Tokenizer(formula):
             if token.type == TokenType.COMPONENT_METRIC:
                 if token.value not in channels:
@@ -76,7 +76,8 @@ async def run_test(
                     ]
                 )
             )
-            assert (await engine.apply()).value == io_output
+            next_val = await engine._apply()  # pylint: disable=protected-access
+            assert (next_val).value == io_output
             tests_passed += 1
         assert tests_passed == len(io_pairs)
 

tests/timeseries/test_logical_meter.py

@@ -38,6 +38,7 @@ async def _get_resampled_stream(  # pylint: disable=too-many-arguments
         # pylint: disable=protected-access
         builder = ResampledFormulaBuilder(
             logical_meter._namespace,
+            "",
             channel_registry,
             request_sender,
             metric_id,