Portainer add resource usage of containers (#155113)

Co-authored-by: Copilot <[email protected]>

Author: erwindouna

homeassistant/components/portainer/binary_sensor.py

@@ -4,9 +4,6 @@
 
 from collections.abc import Callable
 from dataclasses import dataclass
-from typing import Any
-
-from pyportainer.models.docker import DockerContainer
 
 from homeassistant.components.binary_sensor import (
     BinarySensorDeviceClass,
@@ -18,7 +15,7 @@
 from homeassistant.helpers.entity_platform import AddConfigEntryEntitiesCallback
 
 from . import PortainerConfigEntry
-from .coordinator import PortainerCoordinator
+from .coordinator import PortainerContainerData, PortainerCoordinator
 from .entity import (
     PortainerContainerEntity,
     PortainerCoordinatorData,
@@ -27,24 +24,31 @@
 
 
 @dataclass(frozen=True, kw_only=True)
-class PortainerBinarySensorEntityDescription(BinarySensorEntityDescription):
-    """Class to hold Portainer binary sensor description."""
+class PortainerContainerBinarySensorEntityDescription(BinarySensorEntityDescription):
+    """Class to hold Portainer container binary sensor description."""
+
+    state_fn: Callable[[PortainerContainerData], bool | None]
+
+
+@dataclass(frozen=True, kw_only=True)
+class PortainerEndpointBinarySensorEntityDescription(BinarySensorEntityDescription):
+    """Class to hold Portainer endpoint binary sensor description."""
 
-    state_fn: Callable[[Any], bool]
+    state_fn: Callable[[PortainerCoordinatorData], bool | None]
 
 
-CONTAINER_SENSORS: tuple[PortainerBinarySensorEntityDescription, ...] = (
-    PortainerBinarySensorEntityDescription(
+CONTAINER_SENSORS: tuple[PortainerContainerBinarySensorEntityDescription, ...] = (
+    PortainerContainerBinarySensorEntityDescription(
         key="status",
         translation_key="status",
-        state_fn=lambda data: data.state == "running",
+        state_fn=lambda data: data.container.state == "running",
         device_class=BinarySensorDeviceClass.RUNNING,
         entity_category=EntityCategory.DIAGNOSTIC,
     ),
 )
 
-ENDPOINT_SENSORS: tuple[PortainerBinarySensorEntityDescription, ...] = (
-    PortainerBinarySensorEntityDescription(
+ENDPOINT_SENSORS: tuple[PortainerEndpointBinarySensorEntityDescription, ...] = (
+    PortainerEndpointBinarySensorEntityDescription(
         key="status",
         translation_key="status",
         state_fn=lambda data: data.endpoint.status == 1,  # 1 = Running | 2 = Stopped
@@ -76,7 +80,7 @@ def _async_add_new_endpoints(endpoints: list[PortainerCoordinatorData]) -> None:
         )
 
     def _async_add_new_containers(
-        containers: list[tuple[PortainerCoordinatorData, DockerContainer]],
+        containers: list[tuple[PortainerCoordinatorData, PortainerContainerData]],
     ) -> None:
         """Add new container binary sensors."""
         async_add_entities(
@@ -113,12 +117,12 @@ def _async_add_new_containers(
 class PortainerEndpointSensor(PortainerEndpointEntity, BinarySensorEntity):
     """Representation of a Portainer endpoint binary sensor entity."""
 
-    entity_description: PortainerBinarySensorEntityDescription
+    entity_description: PortainerEndpointBinarySensorEntityDescription
 
     def __init__(
         self,
         coordinator: PortainerCoordinator,
-        entity_description: PortainerBinarySensorEntityDescription,
+        entity_description: PortainerEndpointBinarySensorEntityDescription,
         device_info: PortainerCoordinatorData,
     ) -> None:
         """Initialize Portainer endpoint binary sensor entity."""
@@ -141,13 +145,13 @@ def is_on(self) -> bool | None:
 class PortainerContainerSensor(PortainerContainerEntity, BinarySensorEntity):
     """Representation of a Portainer container sensor."""
 
-    entity_description: PortainerBinarySensorEntityDescription
+    entity_description: PortainerContainerBinarySensorEntityDescription
 
     def __init__(
         self,
         coordinator: PortainerCoordinator,
-        entity_description: PortainerBinarySensorEntityDescription,
-        device_info: DockerContainer,
+        entity_description: PortainerContainerBinarySensorEntityDescription,
+        device_info: PortainerContainerData,
         via_device: PortainerCoordinatorData,
     ) -> None:
         """Initialize the Portainer container sensor."""
@@ -164,6 +168,4 @@ def available(self) -> bool:
     @property
     def is_on(self) -> bool | None:
         """Return true if the binary sensor is on."""
-        return self.entity_description.state_fn(
-            self.coordinator.data[self.endpoint_id].containers[self.device_name]
-        )
+        return self.entity_description.state_fn(self.container_data)

homeassistant/components/portainer/button.py

@@ -12,7 +12,6 @@
     PortainerConnectionError,
     PortainerTimeoutError,
 )
-from pyportainer.models.docker import DockerContainer
 
 from homeassistant.components.button import (
     ButtonDeviceClass,
@@ -26,7 +25,11 @@
 
 from . import PortainerConfigEntry
 from .const import DOMAIN
-from .coordinator import PortainerCoordinator, PortainerCoordinatorData
+from .coordinator import (
+    PortainerContainerData,
+    PortainerCoordinator,
+    PortainerCoordinatorData,
+)
 from .entity import PortainerContainerEntity
 
 
@@ -64,7 +67,7 @@ async def async_setup_entry(
     coordinator = entry.runtime_data
 
     def _async_add_new_containers(
-        containers: list[tuple[PortainerCoordinatorData, DockerContainer]],
+        containers: list[tuple[PortainerCoordinatorData, PortainerContainerData]],
     ) -> None:
         """Add new container button sensors."""
         async_add_entities(
@@ -97,7 +100,7 @@ def __init__(
         self,
         coordinator: PortainerCoordinator,
         entity_description: PortainerButtonDescription,
-        device_info: DockerContainer,
+        device_info: PortainerContainerData,
         via_device: PortainerCoordinatorData,
     ) -> None:
         """Initialize the Portainer button entity."""

homeassistant/components/portainer/coordinator.py

@@ -2,6 +2,7 @@
 
 from __future__ import annotations
 
+import asyncio
 from collections.abc import Callable
 from dataclasses import dataclass
 from datetime import timedelta
@@ -13,7 +14,7 @@
     PortainerConnectionError,
     PortainerTimeoutError,
 )
-from pyportainer.models.docker import DockerContainer
+from pyportainer.models.docker import DockerContainer, DockerContainerStats
 from pyportainer.models.docker_inspect import DockerInfo, DockerVersion
 from pyportainer.models.portainer import Endpoint
 
@@ -39,11 +40,20 @@ class PortainerCoordinatorData:
     id: int
     name: str | None
     endpoint: Endpoint
-    containers: dict[str, DockerContainer]
+    containers: dict[str, PortainerContainerData]
     docker_version: DockerVersion
     docker_info: DockerInfo
 
 
+@dataclass(slots=True)
+class PortainerContainerData:
+    """Container data held by the Portainer coordinator."""
+
+    container: DockerContainer
+    stats: DockerContainerStats
+    stats_pre: DockerContainerStats | None
+
+
 class PortainerCoordinator(DataUpdateCoordinator[dict[int, PortainerCoordinatorData]]):
     """Data Update Coordinator for Portainer."""
 
@@ -72,7 +82,9 @@ def __init__(
             Callable[[list[PortainerCoordinatorData]], None]
         ] = []
         self.new_containers_callbacks: list[
-            Callable[[list[tuple[PortainerCoordinatorData, DockerContainer]]], None]
+            Callable[
+                [list[tuple[PortainerCoordinatorData, PortainerContainerData]]], None
+            ]
         ] = []
 
     async def _async_setup(self) -> None:
@@ -119,8 +131,6 @@ async def _async_update_data(self) -> dict[int, PortainerCoordinatorData]:
                 translation_key="cannot_connect",
                 translation_placeholders={"error": repr(err)},
             ) from err
-        else:
-            _LOGGER.debug("Fetched endpoints: %s", endpoints)
 
         mapped_endpoints: dict[int, PortainerCoordinatorData] = {}
         for endpoint in endpoints:
@@ -136,6 +146,47 @@ async def _async_update_data(self) -> dict[int, PortainerCoordinatorData]:
                 containers = await self.portainer.get_containers(endpoint.id)
                 docker_version = await self.portainer.docker_version(endpoint.id)
                 docker_info = await self.portainer.docker_info(endpoint.id)
+
+                container_map: dict[str, PortainerContainerData] = {}
+
+                container_stats_task = [
+                    (
+                        container,
+                        self.portainer.container_stats(
+                            endpoint_id=endpoint.id,
+                            container_id=container.id,
+                        ),
+                    )
+                    for container in containers
+                ]
+
+                container_stats_gather = await asyncio.gather(
+                    *[task for _, task in container_stats_task],
+                )
+                for (container, _), container_stats in zip(
+                    container_stats_task, container_stats_gather, strict=False
+                ):
+                    container_name = container.names[0].replace("/", " ").strip()
+
+                    # Store previous stats if available. This is used to calculate deltas for CPU and network usage
+                    # In the first call it will be None, since it has nothing to compare with
+                    # Added a walrus pattern to check if not None on prev_container, to keep mypy happy. :)
+                    container_map[container_name] = PortainerContainerData(
+                        container=container,
+                        stats=container_stats,
+                        stats_pre=(
+                            prev_container.stats
+                            if self.data
+                            and (prev_data := self.data.get(endpoint.id)) is not None
+                            and (
+                                prev_container := prev_data.containers.get(
+                                    container_name
+                                )
+                            )
+                            is not None
+                            else None
+                        ),
+                    )
             except PortainerConnectionError as err:
                 _LOGGER.exception("Connection error")
                 raise UpdateFailed(
@@ -155,10 +206,7 @@ async def _async_update_data(self) -> dict[int, PortainerCoordinatorData]:
                 id=endpoint.id,
                 name=endpoint.name,
                 endpoint=endpoint,
-                containers={
-                    container.names[0].replace("/", " ").strip(): container
-                    for container in containers
-                },
+                containers=container_map,
                 docker_version=docker_version,
                 docker_info=docker_info,
             )
@@ -179,7 +227,7 @@ def _async_add_remove_endpoints(
 
         # Surprise, we also handle containers here :)
         current_containers = {
-            (endpoint.id, container.id)
+            (endpoint.id, container.container.id)
             for endpoint in mapped_endpoints.values()
             for container in endpoint.containers.values()
         }

homeassistant/components/portainer/diagnostics.py

@@ -30,11 +30,11 @@ def _serialize_coordinator(coordinator: PortainerCoordinator) -> dict[str, Any]:
                 },
                 "containers": [
                     {
-                        "id": container.id,
-                        "names": list(container.names or []),
-                        "image": container.image,
-                        "state": container.state,
-                        "status": container.status,
+                        "id": container.container.id,
+                        "names": list(container.container.names or []),
+                        "image": container.container.image,
+                        "state": container.container.state,
+                        "status": container.container.status,
                     }
                     for container in endpoint_data.containers.values()
                 ],

homeassistant/components/portainer/entity.py

@@ -1,14 +1,17 @@
 """Base class for Portainer entities."""
 
-from pyportainer.models.docker import DockerContainer
 from yarl import URL
 
 from homeassistant.const import CONF_URL
 from homeassistant.helpers.device_registry import DeviceInfo
 from homeassistant.helpers.update_coordinator import CoordinatorEntity
 
 from .const import DEFAULT_NAME, DOMAIN
-from .coordinator import PortainerCoordinator, PortainerCoordinatorData
+from .coordinator import (
+    PortainerContainerData,
+    PortainerCoordinator,
+    PortainerCoordinatorData,
+)
 
 
 class PortainerCoordinatorEntity(CoordinatorEntity[PortainerCoordinator]):
@@ -47,21 +50,22 @@ class PortainerContainerEntity(PortainerCoordinatorEntity):
 
     def __init__(
         self,
-        device_info: DockerContainer,
+        device_info: PortainerContainerData,
         coordinator: PortainerCoordinator,
         via_device: PortainerCoordinatorData,
     ) -> None:
         """Initialize a Portainer container."""
         super().__init__(coordinator)
         self._device_info = device_info
-        self.device_id = self._device_info.id
+        self.device_id = self._device_info.container.id
         self.endpoint_id = via_device.endpoint.id
 
         # Container ID's are ephemeral, so use the container name for the unique ID
         # The first one, should always be unique, it's fine if users have aliases
         # According to Docker's API docs, the first name is unique
-        assert self._device_info.names, "Container names list unexpectedly empty"
-        self.device_name = self._device_info.names[0].replace("/", " ").strip()
+        names = self._device_info.container.names
+        assert names, "Container names list unexpectedly empty"
+        self.device_name = names[0].replace("/", " ").strip()
 
         self._attr_device_info = DeviceInfo(
             identifiers={
@@ -79,3 +83,8 @@ def __init__(
             ),
             translation_key=None if self.device_name else "unknown_container",
         )
+
+    @property
+    def container_data(self) -> PortainerContainerData:
+        """Return the coordinator data for this container."""
+        return self.coordinator.data[self.endpoint_id].containers[self.device_name]

homeassistant/components/portainer/icons.json

@@ -22,6 +22,9 @@
       "cpu_total": {
         "default": "mdi:cpu-64-bit"
       },
+      "cpu_usage_total": {
+        "default": "mdi:cpu-64-bit"
+      },
       "docker_version": {
         "default": "mdi:docker"
       },
@@ -34,9 +37,18 @@
       "kernel_version": {
         "default": "mdi:memory"
       },
+      "memory_limit": {
+        "default": "mdi:memory"
+      },
       "memory_total": {
         "default": "mdi:memory"
       },
+      "memory_usage": {
+        "default": "mdi:memory"
+      },
+      "memory_usage_percentage": {
+        "default": "mdi:memory"
+      },
       "operating_system": {
         "default": "mdi:chip"
       },