typing and placement refactor

Author: Andrei Neagu

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

@@ -1,7 +1,7 @@
 from abc import abstractmethod
+from enum import auto
 from typing import Any, Final, Protocol
 
-from models_library.projects_nodes_io import NodeID
 from pydantic import (
     BaseModel,
     ByteSize,
@@ -12,9 +12,19 @@
     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):
     @property
     @abstractmethod
@@ -104,4 +114,4 @@ def __hash__(self):
 
 class ServiceDiskUsage(BaseModel):
     node_id: NodeID
-    usage: dict[str, DiskUsage]
+    usage: dict[MountPathCategory, DiskUsage]

services/dynamic-sidecar/src/simcore_service_dynamic_sidecar/modules/notifications/_notifications_system_monitor.py

@@ -1,13 +1,20 @@
 from fastapi import FastAPI
-from models_library.api_schemas_dynamic_sidecar.telemetry import DiskUsage
+from models_library.api_schemas_dynamic_sidecar.telemetry import (
+    DiskUsage,
+    MountPathCategory,
+)
 from models_library.projects_nodes_io import NodeID
 from models_library.users import UserID
 
 from ._notifier import Notifier
 
 
 async def publish_disk_usage(
-    app: FastAPI, *, user_id: UserID, node_id: NodeID, usage: dict[str, DiskUsage]
+    app: FastAPI,
+    *,
+    user_id: UserID,
+    node_id: NodeID,
+    usage: dict[MountPathCategory, DiskUsage]
 ) -> None:
     notifier: Notifier = Notifier.get_from_app_state(app)
     await notifier.notify_service_disk_usage(

services/dynamic-sidecar/src/simcore_service_dynamic_sidecar/modules/notifications/_notifier.py

@@ -16,6 +16,7 @@
 )
 from models_library.api_schemas_dynamic_sidecar.telemetry import (
     DiskUsage,
+    MountPathCategory,
     ServiceDiskUsage,
 )
 from models_library.api_schemas_webserver.socketio import SocketIORoomStr
@@ -33,7 +34,10 @@ def __init__(self, sio_manager: socketio.AsyncAioPikaManager):
         self._sio_manager = sio_manager
 
     async def notify_service_disk_usage(
-        self, user_id: UserID, node_id: NodeID, usage: dict[str, DiskUsage]
+        self,
+        user_id: UserID,
+        node_id: NodeID,
+        usage: dict[MountPathCategory, DiskUsage],
     ) -> None:
         await self._sio_manager.emit(
             SOCKET_IO_SERVICE_DISK_USAGE_EVENT,

services/dynamic-sidecar/src/simcore_service_dynamic_sidecar/modules/system_monitor/_disk_usage.py

@@ -2,17 +2,18 @@
 import logging
 from dataclasses import dataclass, field
 from datetime import timedelta
-from enum import auto
 from functools import cached_property
 from pathlib import Path
 from typing import Final
 
 import psutil
 from fastapi import FastAPI
-from models_library.api_schemas_dynamic_sidecar.telemetry import DiskUsage
+from models_library.api_schemas_dynamic_sidecar.telemetry import (
+    DiskUsage,
+    MountPathCategory,
+)
 from models_library.projects_nodes_io import NodeID
 from models_library.users import UserID
-from models_library.utils.enums import StrAutoEnum
 from servicelib.background_task import start_periodic_task, stop_periodic_task
 from servicelib.logging_utils import log_context
 from servicelib.utils import logged_gather
@@ -27,13 +28,6 @@
 _logger = logging.getLogger(__name__)
 
 
-class MountPathCategory(StrAutoEnum):
-    HOST = auto()
-    STATES_VOLUMES = auto()
-    INPUTS_VOLUMES = auto()
-    OUTPUTS_VOLUMES = auto()
-
-
 _SUPPORTED_ITEMS: Final[set[str]] = {
     MountPathCategory.HOST,
     MountPathCategory.STATES_VOLUMES,
@@ -74,7 +68,7 @@ class DiskUsageMonitor:
     _monitor_task: asyncio.Task | None = None
 
     # tracked disk usage
-    _last_usage: dict[str, DiskUsage] = field(default_factory=dict)
+    _last_usage: dict[MountPathCategory, DiskUsage] = field(default_factory=dict)
     _usage_overwrite: dict[str, DiskUsage] = field(
         default_factory=dict,
         metadata={
@@ -141,7 +135,7 @@ def _get_grouped_usage_to_folder_names(
             usage_to_folder_names[disk_usage].add(folder_name)
         return usage_to_folder_names
 
-    async def _publish_disk_usage(self, usage: dict[str, DiskUsage]):
+    async def _publish_disk_usage(self, usage: dict[MountPathCategory, DiskUsage]):
         await publish_disk_usage(
             self.app, user_id=self.user_id, node_id=self.node_id, usage=usage
         )
@@ -158,7 +152,7 @@ async def _monitor(self) -> None:
         )
 
         # compute new version of DiskUsage for FE, only 1 label for each unique disk usage entry
-        usage: dict[str, DiskUsage] = {}
+        usage: dict[MountPathCategory, DiskUsage] = {}
 
         normalized_paths = self._normalized_monitored_paths
 

services/dynamic-sidecar/tests/unit/test_modules_notifier.py

@@ -25,6 +25,7 @@
 )
 from models_library.api_schemas_dynamic_sidecar.telemetry import (
     DiskUsage,
+    MountPathCategory,
     ServiceDiskUsage,
 )
 from models_library.api_schemas_webserver.socketio import SocketIORoomStr
@@ -45,7 +46,6 @@
 )
 from simcore_service_dynamic_sidecar.modules.system_monitor._disk_usage import (
     DiskUsageMonitor,
-    MountPathCategory,
 )
 from socketio import AsyncServer
 from tenacity import AsyncRetrying

services/dynamic-sidecar/tests/unit/test_modules_system_monitor__disk_usage.py

@@ -10,7 +10,10 @@
 import pytest
 from faker import Faker
 from fastapi import FastAPI
-from models_library.api_schemas_dynamic_sidecar.telemetry import DiskUsage
+from models_library.api_schemas_dynamic_sidecar.telemetry import (
+    DiskUsage,
+    MountPathCategory,
+)
 from models_library.projects_nodes_io import NodeID
 from models_library.services_types import RunID
 from models_library.users import UserID
@@ -21,7 +24,6 @@
 from simcore_service_dynamic_sidecar.modules.mounted_fs import MountedVolumes
 from simcore_service_dynamic_sidecar.modules.system_monitor._disk_usage import (
     DiskUsageMonitor,
-    MountPathCategory,
     _get_monitored_paths,
     get_usage,
 )