Add model manager as BackgroundService for handling machine learning models

Signed-off-by: Idlir Shkurti <[email protected]>

Author: idlir-shkurti-frequenz

src/frequenz/sdk/ml/__init__.py

@@ -0,0 +1,11 @@
+# License: MIT
+# Copyright © 2024 Frequenz Energy-as-a-Service GmbH
+
+"""Model interface."""
+
+from ._model_manager import ModelManager
+
+# Explicitly declare the public API.
+__all__ = [
+    "ModelManager",
+]

src/frequenz/sdk/ml/_model_manager.py

@@ -0,0 +1,136 @@
+# License: MIT
+# Copyright © 2024 Frequenz Energy-as-a-Service GmbH
+
+"""Load, update, monitor and retrieve machine learning models."""
+
+import asyncio
+import logging
+import pickle
+from asyncio import CancelledError
+from dataclasses import dataclass
+from pathlib import Path
+from typing import Generic, TypeVar, cast
+
+from frequenz.channels.file_watcher import EventType, FileWatcher
+
+from frequenz.sdk.actor import BackgroundService
+
+_logger = logging.getLogger(__name__)
+
+T = TypeVar("T")
+
+
+@dataclass
+class _Model(Generic[T]):
+    """Represent a machine learning model."""
+
+    data: T
+    path: Path
+
+
+class ModelManager(BackgroundService, Generic[T]):
+    """Load, update, monitor and retrieve machine learning models."""
+
+    def __init__(self, model_paths: dict[str, Path]):
+        """Initialize the model manager with the specified model paths.
+
+        Args:
+            model_paths: A dictionary of model keys and their corresponding file paths.
+        """
+        super().__init__()
+        self._models: dict[str, _Model[T]] = {}
+        self.model_paths = model_paths
+        self.load_models()
+
+    def load_models(self) -> None:
+        """Load the models from the specified paths."""
+        for key, path in self.model_paths.items():
+            self._models[key] = _Model(data=self._load(path), path=path)
+
+    def _load(self, path: Path) -> T:
+        """Load the model from the specified path.
+
+        Args:
+            path: The path to the model file.
+
+        Returns:
+            T: The loaded model data.
+
+        Raises:
+            FileNotFoundError: If the model file does not exist.
+        """
+        if not path.exists():
+            raise FileNotFoundError(f"The model path {path} does not exist.")
+        with path.open("rb") as file:
+            return cast(T, pickle.load(file))
+
+    def start(self) -> None:
+        """Start the model monitoring service by creating a background task."""
+        if not self.is_running:
+            task = asyncio.create_task(self.run())
+            self._tasks.add(task)
+            _logger.info("Started ModelManager service with task %s", task)
+
+    async def run(self) -> None:
+        """Monitor model file paths and reload models as necessary."""
+        model_paths = [model.path for model in self._models.values()]
+        file_watcher = FileWatcher(paths=list(model_paths))
+        _logger.info("Monitoring model paths for changes.")
+        async for event in file_watcher:
+            if event.type in (EventType.CREATE, EventType.MODIFY):
+                _logger.info("Model file %s modified, reloading...", event.path)
+                self.reload_model(Path(event.path))
+
+    def reload_model(self, path: Path) -> None:
+        """Reload the model from the specified path.
+
+        Args:
+            path: The path to the model file.
+        """
+        for key, model in self._models.items():
+            if model.path == path:
+                try:
+                    model.data = self._load(path)
+                    _logger.info("Successfully reloaded model from %s", path)
+                except Exception as e:  # pylint: disable=broad-except
+                    _logger.error("Failed to reload model from %s: %s", path, e)
+
+    async def stop(self, msg: str | None = None) -> None:
+        """Stop all model monitoring tasks with enhanced exception handling.
+
+        Args:
+            msg: An optional message to log when stopping the service.
+        """
+        _logger.info("Stopping ModelManager service: %s", msg)
+        # Attempt to cancel all running tasks
+        for task in list(self._tasks):
+            task.cancel()
+            try:
+                await task  # Wait for task to be cancelled
+            except CancelledError:
+                _logger.info("Task %s cancelled successfully", task)
+            except Exception as e:  # pylint: disable=broad-except
+                _logger.error("Error while cancelling task %s: %s", task, e)
+
+        # Call the parent stop method if it handles additional teardown
+        try:
+            await super().stop(msg)
+        except Exception as e:  # pylint: disable=broad-except
+            _logger.error("Error during stop in superclass: %s", e)
+
+    def get_model(self, key: str) -> T:
+        """Retrieve a loaded model by key.
+
+        Args:
+            key: The key of the model to retrieve.
+
+        Returns:
+            The loaded model data.
+
+        Raises:
+            KeyError: If the model with the specified key is not found.
+        """
+        try:
+            return self._models[key].data
+        except KeyError as exc:
+            raise KeyError(f"Model with key '{key}' is not found.") from exc

tests/model/__init__.py

@@ -0,0 +1,4 @@
+# License: MIT
+# Copyright © 2024 Frequenz Energy-as-a-Service GmbH
+
+"""Tests for the model package."""

tests/model/test_model_manager.py

@@ -0,0 +1,118 @@
+# License: MIT
+# Copyright © 2024 Frequenz Energy-as-a-Service GmbH
+
+"""Tests for machine learning model manager."""
+
+import pickle
+from dataclasses import dataclass
+from pathlib import Path
+from typing import Any
+from unittest.mock import AsyncMock, MagicMock, mock_open, patch
+
+import pytest
+
+from frequenz.sdk.ml import ModelManager
+
+
+@dataclass
+class MockModel:
+    """Mock model for unit testing purposes."""
+
+    data: int | str
+
+    def predict(self) -> int | str:
+        """Make a prediction based on the model data."""
+        return self.data
+
+
+async def test_model_manager_loading() -> None:
+    """Test loading models using ModelManager with direct configuration."""
+    model1 = MockModel("Model 1 Data")
+    model2 = MockModel("Model 2 Data")
+    pickled_model1 = pickle.dumps(model1)
+    pickled_model2 = pickle.dumps(model2)
+
+    model_paths = {
+        "model1": Path("path/to/model1.pkl"),
+        "model2": Path("path/to/model2.pkl"),
+    }
+
+    mock_files = {
+        "path/to/model1.pkl": mock_open(read_data=pickled_model1)(),
+        "path/to/model2.pkl": mock_open(read_data=pickled_model2)(),
+    }
+
+    def mock_open_func(file_path: Path, *__args: Any, **__kwargs: Any) -> Any:
+        """Mock open function to return the correct mock file object.
+
+        Args:
+            file_path: The path to the file to open.
+            *__args: Variable length argument list. This can be used to pass additional
+                    positional parameters typically used in file opening operations,
+                    such as `mode` or `buffering`.
+            **__kwargs: Arbitrary keyword arguments. This can include parameters like
+                    `encoding` and `errors`, common in file opening operations.
+
+        Returns:
+            Any: The mock file object.
+
+        Raises:
+            FileNotFoundError: If the file path is not in the mock files dictionary.
+        """
+        file_path_str = str(file_path)
+        if file_path_str in mock_files:
+            file_handle = MagicMock()
+            file_handle.__enter__.return_value = mock_files[file_path_str]
+            return file_handle
+        raise FileNotFoundError(f"No mock setup for {file_path_str}")
+
+    with patch("pathlib.Path.open", new=mock_open_func):
+        with patch.object(Path, "exists", return_value=True):
+            model_manager: ModelManager[MockModel] = ModelManager(
+                model_paths=model_paths
+            )
+
+        with patch(
+            "frequenz.channels.file_watcher.FileWatcher", new_callable=AsyncMock
+        ):
+            model_manager.start()  # Start the service
+
+            assert isinstance(model_manager.get_model("model1"), MockModel)
+            assert model_manager.get_model("model1").data == "Model 1 Data"
+            assert model_manager.get_model("model2").data == "Model 2 Data"
+
+            with pytest.raises(KeyError):
+                model_manager.get_model("key3")
+
+            await model_manager.stop()  # Stop the service to clean up
+
+
+async def test_model_manager_update() -> None:
+    """Test updating a model in ModelManager."""
+    original_model = MockModel("Original Data")
+    updated_model = MockModel("Updated Data")
+    pickled_original_model = pickle.dumps(original_model)
+    pickled_updated_model = pickle.dumps(updated_model)
+
+    model_paths = {"model1": Path("path/to/model1.pkl")}
+
+    mock_file = mock_open(read_data=pickled_original_model)
+    with (
+        patch("pathlib.Path.open", mock_file),
+        patch.object(Path, "exists", return_value=True),
+    ):
+        model_manager = ModelManager[MockModel](model_paths=model_paths)
+        with patch(
+            "frequenz.channels.file_watcher.FileWatcher", new_callable=AsyncMock
+        ):
+            model_manager.start()  # Start the service
+
+            assert model_manager.get_model("model1").data == "Original Data"
+
+            # Simulate updating the model file
+            mock_file.return_value.read.return_value = pickled_updated_model
+            with patch("pathlib.Path.open", mock_file):
+                model_manager.reload_model(Path("path/to/model1.pkl"))
+                assert model_manager.get_model("model1").data == "Updated Data"
+
+            await model_manager.stop()  # Stop the service to clean up