Switch to use the external component graph

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

Author: shsms

src/frequenz/sdk/microgrid/connection_manager.py

@@ -8,13 +8,19 @@
 component graph.
 """
 
+import asyncio
 import logging
 from abc import ABC, abstractmethod
 
 from frequenz.client.common.microgrid import MicrogridId
-from frequenz.client.microgrid import Location, MicrogridApiClient, MicrogridInfo
-
-from .component_graph import ComponentGraph, _MicrogridComponentGraph
+from frequenz.client.common.microgrid.components import ComponentId
+from frequenz.client.microgrid import (
+    Location,
+    MicrogridApiClient,
+    MicrogridInfo,
+)
+from frequenz.client.microgrid.component import Component, ComponentConnection
+from frequenz.microgrid_component_graph import ComponentGraph
 
 _logger = logging.getLogger(__name__)
 
@@ -51,7 +57,9 @@ def api_client(self) -> MicrogridApiClient:
 
     @property
     @abstractmethod
-    def component_graph(self) -> ComponentGraph:
+    def component_graph(
+        self,
+    ) -> ComponentGraph[Component, ComponentConnection, ComponentId]:
         """Get component graph.
 
         Returns:
@@ -101,7 +109,9 @@ def __init__(self, server_url: str) -> None:
         self._client = MicrogridApiClient(server_url)
         # To create graph from the API client we need await.
         # So create empty graph here, and update it in `run` method.
-        self._graph = _MicrogridComponentGraph()
+        self._graph: (
+            ComponentGraph[Component, ComponentConnection, ComponentId] | None
+        ) = None
 
         self._microgrid: MicrogridInfo
         """The microgrid information."""
@@ -130,12 +140,19 @@ def location(self) -> Location | None:
         return self._microgrid.location
 
     @property
-    def component_graph(self) -> ComponentGraph:
+    def component_graph(
+        self,
+    ) -> ComponentGraph[Component, ComponentConnection, ComponentId]:
         """Get component graph.
 
         Returns:
             component graph
+
+        Raises:
+            RuntimeError: If the microgrid is not initialized yet.
         """
+        if self._graph is None:
+            raise RuntimeError("Microgrid not initialized yet.")
         return self._graph
 
     async def _update_client(self, server_url: str) -> None:
@@ -156,7 +173,11 @@ async def _update_client(self, server_url: str) -> None:
 
     async def _initialize(self) -> None:
         self._microgrid = await self._client.get_microgrid_info()
-        await self._graph.refresh_from_client(self._client)
+        components, connections = await asyncio.gather(
+            self._client.list_components(),
+            self._client.list_connections(),
+        )
+        self._graph = ComponentGraph(set(components), set(connections))
 
 
 _CONNECTION_MANAGER: ConnectionManager | None = None

src/frequenz/sdk/timeseries/_grid_frequency.py

@@ -15,6 +15,7 @@
 from frequenz.quantities import Frequency, Quantity
 
 from .._internal._channels import ChannelRegistry
+from .._internal._graph_traversal import find_first_descendant_component
 from ..microgrid import connection_manager
 from ..microgrid._data_sourcing import ComponentMetricRequest
 from ..timeseries._base_types import Sample
@@ -54,7 +55,8 @@ def __init__(
         """
         if not source:
             component_graph = connection_manager.get().component_graph
-            source = component_graph.find_first_descendant_component(
+            source = find_first_descendant_component(
+                component_graph,
                 descendants=[Meter, Inverter, EvCharger],
             )
 

src/frequenz/sdk/timeseries/_voltage_streamer.py

@@ -19,6 +19,7 @@
 from frequenz.quantities import Quantity, Voltage
 
 from .._internal._channels import ChannelRegistry
+from .._internal._graph_traversal import find_first_descendant_component
 from ..timeseries._base_types import Sample, Sample3Phase
 
 if TYPE_CHECKING:
@@ -81,7 +82,8 @@ def __init__(
 
         if not source_component:
             component_graph = connection_manager.get().component_graph
-            source_component = component_graph.find_first_descendant_component(
+            source_component = find_first_descendant_component(
+                component_graph,
                 descendants=[Meter, Inverter, EvCharger],
             )
 

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

@@ -10,6 +10,7 @@
 from frequenz.client.microgrid.metrics import Metric
 from frequenz.quantities import Power
 
+from ...._internal._graph_traversal import is_battery_inverter
 from ....microgrid import connection_manager
 from ...formula_engine import FormulaEngine
 from ._fallback_formula_metric_fetcher import FallbackFormulaMetricFetcher
@@ -73,7 +74,7 @@ def generate(
         for bat_id in component_ids:
             inverters = set(
                 filter(
-                    component_graph.is_battery_inverter,
+                    is_battery_inverter,
                     component_graph.predecessors(bat_id),
                 )
             )
@@ -84,7 +85,7 @@ def generate(
                 )
 
             for inverter in inverters:
-                all_connected_batteries = component_graph.successors(inverter.id)
+                all_connected_batteries = set(component_graph.successors(inverter.id))
                 battery_ids = set(
                     map(lambda battery: battery.id, all_connected_batteries)
                 )

src/frequenz/sdk/timeseries/formula_engine/_formula_generators/_consumer_power_formula.py

@@ -9,6 +9,13 @@
 from frequenz.client.microgrid.metrics import Metric
 from frequenz.quantities import Power
 
+from ...._internal._graph_traversal import (
+    dfs,
+    is_battery_chain,
+    is_chp_chain,
+    is_ev_charger_chain,
+    is_pv_chain,
+)
 from ....microgrid import connection_manager
 from .._formula_engine import FormulaEngine
 from .._resampled_formula_builder import ResampledFormulaBuilder
@@ -43,10 +50,10 @@ def _are_grid_meters(self, grid_successors: set[Component]) -> bool:
         component_graph = connection_manager.get().component_graph
         return all(
             isinstance(successor, Meter)
-            and not component_graph.is_battery_chain(successor)
-            and not component_graph.is_chp_chain(successor)
-            and not component_graph.is_pv_chain(successor)
-            and not component_graph.is_ev_charger_chain(successor)
+            and not is_battery_chain(component_graph, successor)
+            and not is_chp_chain(component_graph, successor)
+            and not is_pv_chain(component_graph, successor)
+            and not is_ev_charger_chain(component_graph, successor)
             for successor in grid_successors
         )
 
@@ -107,17 +114,17 @@ def non_consumer_component(component: Component) -> bool:
             # If the component graph supports additional types of grid successors in the
             # future, additional checks need to be added here.
             return (
-                component_graph.is_battery_chain(component)
-                or component_graph.is_chp_chain(component)
-                or component_graph.is_pv_chain(component)
-                or component_graph.is_ev_charger_chain(component)
+                is_battery_chain(component_graph, component)
+                or is_chp_chain(component_graph, component)
+                or is_pv_chain(component_graph, component)
+                or is_ev_charger_chain(component_graph, component)
             )
 
         # join all non consumer components reachable from the different grid meters
         non_consumer_components: set[Component] = set()
         for grid_meter in grid_meters:
             non_consumer_components = non_consumer_components.union(
-                component_graph.dfs(grid_meter, set(), non_consumer_component)
+                dfs(component_graph, grid_meter, set(), non_consumer_component)
             )
 
         # push all grid meters
@@ -180,14 +187,14 @@ def consumer_component(component: Component) -> bool:
             # future, additional checks need to be added here.
             return (
                 isinstance(component, (Meter, Inverter))
-                and not component_graph.is_battery_chain(component)
-                and not component_graph.is_chp_chain(component)
-                and not component_graph.is_pv_chain(component)
-                and not component_graph.is_ev_charger_chain(component)
+                and not is_battery_chain(component_graph, component)
+                and not is_chp_chain(component_graph, component)
+                and not is_pv_chain(component_graph, component)
+                and not is_ev_charger_chain(component_graph, component)
             )
 
         component_graph = connection_manager.get().component_graph
-        consumer_components = component_graph.dfs(grid, set(), consumer_component)
+        consumer_components = dfs(component_graph, grid, set(), consumer_component)
 
         if not consumer_components:
             _logger.warning(

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

@@ -18,6 +18,12 @@
 from frequenz.client.microgrid.metrics import Metric
 
 from ...._internal._channels import ChannelRegistry
+from ...._internal._graph_traversal import (
+    is_battery_inverter,
+    is_chp,
+    is_ev_charger,
+    is_pv_inverter,
+)
 from ....microgrid import connection_manager
 from ....microgrid._data_sourcing import ComponentMetricRequest
 from ..._base_types import QuantityT
@@ -138,7 +144,7 @@ def _get_grid_component_successors(self) -> set[Component]:
         if not grid_successors:
             raise ComponentNotFound("No components found in the component graph.")
 
-        return grid_successors
+        return set(grid_successors)
 
     @abstractmethod
     def generate(
@@ -183,7 +189,7 @@ def _get_metric_fallback_components(
             if isinstance(component, Meter):
                 fallbacks[component] = self._get_meter_fallback_components(component)
             else:
-                predecessors = graph.predecessors(component.id)
+                predecessors = set(graph.predecessors(component.id))
                 if len(predecessors) == 1:
                     predecessor = predecessors.pop()
                     if self._is_primary_fallback_pair(predecessor, component):
@@ -213,11 +219,11 @@ def _get_meter_fallback_components(self, meter: Component) -> set[Component]:
 
         # All fallbacks has to be of the same type and category.
         if (
-            all(graph.is_pv_inverter(c) for c in successors)
-            or all(graph.is_battery_inverter(c) for c in successors)
-            or all(graph.is_ev_charger(c) for c in successors)
+            all(is_pv_inverter(c) for c in successors)
+            or all(is_battery_inverter(c) for c in successors)
+            or all(is_ev_charger(c) for c in successors)
         ):
-            return successors
+            return set(successors)
         return set()
 
     def _is_primary_fallback_pair(
@@ -246,9 +252,9 @@ def _is_primary_fallback_pair(
 
         # fmt: off
         return (
-            graph.is_pv_inverter(fallback) and graph.is_pv_meter(primary)
-            or graph.is_chp(fallback) and graph.is_chp_meter(primary)
-            or graph.is_ev_charger(fallback) and graph.is_ev_charger_meter(primary)
-            or graph.is_battery_inverter(fallback) and graph.is_battery_meter(primary)
+            is_pv_inverter(fallback) and graph.is_pv_meter(primary.id)
+            or is_chp(fallback) and graph.is_chp_meter(primary.id)
+            or is_ev_charger(fallback) and graph.is_ev_charger_meter(primary.id)
+            or is_battery_inverter(fallback) and graph.is_battery_meter(primary.id)
         )
         # fmt: on

src/frequenz/sdk/timeseries/formula_engine/_formula_generators/_producer_power_formula.py

@@ -10,6 +10,7 @@
 from frequenz.client.microgrid.metrics import Metric
 from frequenz.quantities import Power
 
+from ...._internal._graph_traversal import dfs, is_chp_chain, is_pv_chain
 from ....microgrid import connection_manager
 from .._formula_engine import FormulaEngine
 from ._fallback_formula_metric_fetcher import FallbackFormulaMetricFetcher
@@ -52,11 +53,12 @@ def generate(  # noqa: DOC502
 
         component_graph = connection_manager.get().component_graph
         # if in the future we support additional producers, we need to add them to the lambda
-        producer_components = component_graph.dfs(
+        producer_components = dfs(
+            component_graph,
             self._get_grid_component(),
             set(),
-            lambda component: component_graph.is_pv_chain(component)
-            or component_graph.is_chp_chain(component),
+            lambda component: is_pv_chain(component_graph, component)
+            or is_chp_chain(component_graph, component),
         )
 
         if not producer_components:

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

@@ -9,6 +9,7 @@
 from frequenz.client.microgrid.metrics import Metric
 from frequenz.quantities import Power
 
+from ...._internal._graph_traversal import dfs, is_pv_chain
 from ....microgrid import connection_manager
 from .._formula_engine import FormulaEngine
 from ._fallback_formula_metric_fetcher import FallbackFormulaMetricFetcher
@@ -48,10 +49,11 @@ def generate(  # noqa: DOC502
         if component_ids:
             pv_components = component_graph.components(set(component_ids))
         else:
-            pv_components = component_graph.dfs(
+            pv_components = dfs(
+                component_graph,
                 self._get_grid_component(),
                 set(),
-                component_graph.is_pv_chain,
+                lambda c: is_pv_chain(component_graph, c),
             )
 
         if not pv_components:

tests/microgrid/fixtures.py

@@ -12,12 +12,17 @@
 from typing import AsyncIterator
 
 from frequenz.channels import Sender
-from frequenz.client.microgrid.component import ComponentCategory
+from frequenz.client.common.microgrid.components import ComponentId
+from frequenz.client.microgrid.component import (
+    Component,
+    ComponentCategory,
+    ComponentConnection,
+)
+from frequenz.microgrid_component_graph import ComponentGraph
 from pytest_mock import MockerFixture
 
 from frequenz.sdk import microgrid
 from frequenz.sdk.microgrid._power_distributing import ComponentPoolStatus
-from frequenz.sdk.microgrid.component_graph import _MicrogridComponentGraph
 from frequenz.sdk.timeseries import ResamplerConfig2
 
 from ..timeseries.mock_microgrid import MockMicrogrid
@@ -42,7 +47,9 @@ async def new(
         cls,
         component_category: ComponentCategory,
         mocker: MockerFixture,
-        graph: _MicrogridComponentGraph | None = None,
+        graph: (
+            ComponentGraph[Component, ComponentConnection, ComponentId] | None
+        ) = None,
         grid_meter: bool | None = None,
     ) -> _Mocks:
         """Initialize the mocks."""
@@ -95,7 +102,7 @@ async def _mocks(
     mocker: MockerFixture,
     component_category: ComponentCategory,
     *,
-    graph: _MicrogridComponentGraph | None = None,
+    graph: ComponentGraph[Component, ComponentConnection, ComponentId] | None = None,
     grid_meter: bool | None = None,
 ) -> AsyncIterator[_Mocks]:
     """Initialize the mocks."""

tests/microgrid/test_grid.py

@@ -5,20 +5,20 @@
 
 from contextlib import AsyncExitStack
 
+import frequenz.microgrid_component_graph as gr
 from frequenz.client.common.microgrid import MicrogridId
 from frequenz.client.common.microgrid.components import ComponentId
 from frequenz.client.microgrid.component import (
+    Component,
     ComponentCategory,
     ComponentConnection,
     GridConnectionPoint,
     Meter,
-    UnspecifiedComponent,
 )
 from frequenz.client.microgrid.metrics import Metric
 from frequenz.quantities import Current, Power, Quantity, ReactivePower
 from pytest_mock import MockerFixture
 
-import frequenz.sdk.microgrid.component_graph as gr
 from frequenz.sdk import microgrid
 from frequenz.sdk.timeseries import Fuse
 from tests.utils.graph_generator import GraphGenerator
@@ -38,7 +38,9 @@ async def test_grid_2(mocker: MockerFixture) -> None:
     components = {grid_1, meter_2}
     connections = {ComponentConnection(source=grid_1.id, destination=meter_2.id)}
 
-    graph = gr._MicrogridComponentGraph(  # pylint: disable=protected-access
+    graph = gr.ComponentGraph[
+        Component, ComponentConnection, ComponentId
+    ](  # pylint: disable=protected-access
         components=components, connections=connections
     )