♻️ Refactor DiskUsage functionality to support efs-guardian (#6536)

Co-authored-by: Andrei Neagu <[email protected]>

Author: GitHK

packages/models-library/src/models_library/api_schemas_dynamic_sidecar/__init__.py

@@ -0,0 +1,9 @@
+from typing import Final
+
+from pydantic import parse_obj_as
+
+from ..rabbitmq_basic_types import RPCNamespace
+
+DYNAMIC_SIDECAR_RPC_NAMESPACE: Final[RPCNamespace] = parse_obj_as(
+    RPCNamespace, "dynamic-sidecar"
+)

packages/models-library/src/models_library/api_schemas_dynamic_sidecar/telemetry.py

@@ -1,9 +1,28 @@
 from abc import abstractmethod
-from pathlib import Path
-from typing import Protocol
+from enum import auto
+from typing import Any, Final, Protocol
 
-from models_library.projects_nodes_io import NodeID
-from pydantic import BaseModel, ByteSize, Field
+from pydantic import (
+    BaseModel,
+    ByteSize,
+    Field,
+    NonNegativeFloat,
+    NonNegativeInt,
+    root_validator,
+    validator,
+)
+
+from ..projects_nodes_io import NodeID
+from ..utils.enums import StrAutoEnum
+
+_EPSILON: Final[NonNegativeFloat] = 1e-16
+
+
+class MountPathCategory(StrAutoEnum):
+    HOST = auto()
+    STATES_VOLUMES = auto()
+    INPUTS_VOLUMES = auto()
+    OUTPUTS_VOLUMES = auto()
 
 
 class SDiskUsageProtocol(Protocol):
@@ -28,31 +47,71 @@ def percent(self) -> float:
         ...
 
 
+def _get_percent(used: float, total: float) -> float:
+    return round(used * 100 / (total + _EPSILON), 2)
+
+
 class DiskUsage(BaseModel):
     used: ByteSize = Field(description="used space")
     free: ByteSize = Field(description="remaining space")
 
     total: ByteSize = Field(description="total space = free + used")
-    used_percent: float = Field(
+    used_percent: NonNegativeFloat = Field(
         gte=0.00,
         lte=100.00,
         description="Percent of used space relative to the total space",
     )
 
+    @validator("free")
     @classmethod
-    def from_ps_util_disk_usage(
-        cls, ps_util_disk_usage: SDiskUsageProtocol
+    def _free_positive(cls, v: float) -> float:
+        if v < 0:
+            msg = f"free={v} cannot be a negative value"
+            raise ValueError(msg)
+        return v
+
+    @validator("used")
+    @classmethod
+    def _used_positive(cls, v: float) -> float:
+        if v < 0:
+            msg = f"used={v} cannot be a negative value"
+            raise ValueError(msg)
+        return v
+
+    @root_validator(pre=True)
+    @classmethod
+    def _check_total(cls, values: dict[str, Any]) -> dict[str, Any]:
+        total = values["total"]
+        free = values["free"]
+        used = values["used"]
+        if total != free + used:
+            msg = f"{total=} is different than the sum of {free=}+{used=} => sum={free+used}"
+            raise ValueError(msg)
+        return values
+
+    @classmethod
+    def from_efs_guardian(
+        cls, used: NonNegativeInt, total: NonNegativeInt
     ) -> "DiskUsage":
-        total = ps_util_disk_usage.free + ps_util_disk_usage.used
-        used_percent = round(ps_util_disk_usage.used * 100 / total, 2)
+        free = total - used
         return cls(
-            used=ByteSize(ps_util_disk_usage.used),
-            free=ByteSize(ps_util_disk_usage.free),
+            used=ByteSize(used),
+            free=ByteSize(free),
             total=ByteSize(total),
-            used_percent=used_percent,
+            used_percent=_get_percent(used, total),
         )
 
+    @classmethod
+    def from_ps_util_disk_usage(
+        cls, ps_util_disk_usage: SDiskUsageProtocol
+    ) -> "DiskUsage":
+        total = ps_util_disk_usage.free + ps_util_disk_usage.used
+        return cls.from_efs_guardian(ps_util_disk_usage.used, total)
+
+    def __hash__(self):
+        return hash((self.used, self.free, self.total, self.used_percent))
+
 
 class ServiceDiskUsage(BaseModel):
     node_id: NodeID
-    usage: dict[Path, DiskUsage]
+    usage: dict[MountPathCategory, DiskUsage]

packages/models-library/tests/test_api_schemas_dynamic_sidecar_telemetry.py

@@ -2,6 +2,7 @@
 import pytest
 from models_library.api_schemas_dynamic_sidecar.telemetry import DiskUsage
 from psutil._common import sdiskusage
+from pydantic import ByteSize, ValidationError
 
 
 def _assert_same_value(ps_util_disk_usage: sdiskusage) -> None:
@@ -27,3 +28,35 @@ def test_disk_usage_regression_cases(ps_util_disk_usage: sdiskusage):
 def test_disk_usage():
     ps_util_disk_usage = psutil.disk_usage("/")
     _assert_same_value(ps_util_disk_usage)
+
+
+def test_from_efs_guardian_constructor():
+    result = DiskUsage.from_efs_guardian(10, 100)
+    assert result.used == ByteSize(10)
+    assert result.free == ByteSize(90)
+    assert result.total == ByteSize(100)
+    assert result.used_percent == 10
+
+
+def test_failing_validation():
+    with pytest.raises(ValidationError) as exc:
+        assert DiskUsage.from_efs_guardian(100, 10)
+
+    assert "free=" in f"{exc.value}"
+    assert "negative value" in f"{exc.value}"
+
+    with pytest.raises(ValidationError) as exc:
+        assert DiskUsage(
+            used=-10,  # type: ignore
+            free=ByteSize(10),
+            total=ByteSize(0),
+            used_percent=-10,
+        )
+    assert "used=" in f"{exc.value}"
+    assert "negative value" in f"{exc.value}"
+
+    with pytest.raises(ValidationError) as exc:
+        DiskUsage(
+            used=ByteSize(10), free=ByteSize(10), total=ByteSize(21), used_percent=0
+        )
+    assert "is different than the sum of" in f"{exc.value}"

packages/service-library/src/servicelib/rabbitmq/rpc_interfaces/dynamic_sidecar/__init__.py

@@ -0,0 +1,22 @@
+import logging
+
+from models_library.api_schemas_dynamic_sidecar import DYNAMIC_SIDECAR_RPC_NAMESPACE
+from models_library.api_schemas_dynamic_sidecar.telemetry import DiskUsage
+from models_library.rabbitmq_basic_types import RPCMethodName
+from pydantic import parse_obj_as
+from servicelib.logging_utils import log_decorator
+from servicelib.rabbitmq import RabbitMQRPCClient
+
+_logger = logging.getLogger(__name__)
+
+
+@log_decorator(_logger, level=logging.DEBUG)
+async def update_disk_usage(
+    rabbitmq_rpc_client: RabbitMQRPCClient, *, usage: dict[str, DiskUsage]
+) -> None:
+    result = await rabbitmq_rpc_client.request(
+        DYNAMIC_SIDECAR_RPC_NAMESPACE,
+        parse_obj_as(RPCMethodName, "update_disk_usage"),
+        usage=usage,
+    )
+    assert result is None  # nosec

packages/service-library/src/servicelib/rabbitmq/rpc_interfaces/dynamic_sidecar/disk_usage.py

@@ -1,3 +0,0 @@
-from ._routing import get_main_router
-
-__all__: tuple[str, ...] = ("get_main_router",)

services/dynamic-sidecar/src/simcore_service_dynamic_sidecar/api/__init__.py

@@ -0,0 +1,3 @@
+from ._routing import get_main_router
+
+__all__: tuple[str, ...] = ("get_main_router",)

services/dynamic-sidecar/src/simcore_service_dynamic_sidecar/api/rest/__init__.py

@@ -7,15 +7,16 @@
 from fastapi import Depends, FastAPI, Request
 from fastapi.datastructures import State
 from servicelib.rabbitmq import RabbitMQClient
+from servicelib.rabbitmq._client_rpc import RabbitMQRPCClient
 
-from ..core import rabbitmq
-from ..core.settings import ApplicationSettings
-from ..models.schemas.application_health import ApplicationHealth
-from ..models.shared_store import SharedStore
-from ..modules.inputs import InputsState
-from ..modules.mounted_fs import MountedVolumes
-from ..modules.outputs import OutputsContext, OutputsManager
-from ..modules.prometheus_metrics import UserServicesMetrics
+from ...core import rabbitmq
+from ...core.settings import ApplicationSettings
+from ...models.schemas.application_health import ApplicationHealth
+from ...models.shared_store import SharedStore
+from ...modules.inputs import InputsState
+from ...modules.mounted_fs import MountedVolumes
+from ...modules.outputs import OutputsContext, OutputsManager
+from ...modules.prometheus_metrics import UserServicesMetrics
 
 
 def get_application(request: Request) -> FastAPI:
@@ -84,3 +85,9 @@ def get_rabbitmq_client(
     app: Annotated[FastAPI, Depends(get_application)]
 ) -> RabbitMQClient:
     return rabbitmq.get_rabbitmq_client(app)
+
+
+def get_rabbitmq_rpc_server(
+    app: Annotated[FastAPI, Depends(get_application)]
+) -> RabbitMQRPCClient:
+    return rabbitmq.get_rabbitmq_rpc_server(app)

services/dynamic-sidecar/src/simcore_service_dynamic_sidecar/api/_dependencies.py renamed to services/dynamic-sidecar/src/simcore_service_dynamic_sidecar/api/rest/_dependencies.py

@@ -5,8 +5,8 @@
 
 from fastapi import APIRouter, FastAPI
 
-from .._meta import API_VTAG
-from ..core.settings import ApplicationSettings
+from ..._meta import API_VTAG
+from ...core.settings import ApplicationSettings
 from . import (
     containers,
     containers_extension,

services/dynamic-sidecar/src/simcore_service_dynamic_sidecar/api/_routing.py renamed to services/dynamic-sidecar/src/simcore_service_dynamic_sidecar/api/rest/_routing.py

@@ -15,22 +15,22 @@
 from pydantic import parse_raw_as
 from servicelib.fastapi.requests_decorators import cancel_on_disconnect
 
-from ..core.docker_utils import docker_client
-from ..core.errors import (
+from ...core.docker_utils import docker_client
+from ...core.errors import (
     ContainerExecCommandFailedError,
     ContainerExecContainerNotFoundError,
     ContainerExecTimeoutError,
 )
-from ..core.settings import ApplicationSettings
-from ..core.validation import (
+from ...core.settings import ApplicationSettings
+from ...core.validation import (
     ComposeSpecValidation,
     parse_compose_spec,
     validate_compose_spec,
 )
-from ..models.schemas.containers import ContainersComposeSpec
-from ..models.shared_store import SharedStore
-from ..modules.container_utils import run_command_in_container
-from ..modules.mounted_fs import MountedVolumes
+from ...models.schemas.containers import ContainersComposeSpec
+from ...models.shared_store import SharedStore
+from ...modules.container_utils import run_command_in_container
+from ...modules.mounted_fs import MountedVolumes
 from ._dependencies import (
     get_container_restart_lock,
     get_mounted_volumes,